diff --git a/docs/platforms/ascend_npu_qwen3_examples.md b/docs/platforms/ascend_npu_qwen3_examples.md index 3c35b1ba9..787dfb6a2 100644 --- a/docs/platforms/ascend_npu_qwen3_examples.md +++ b/docs/platforms/ascend_npu_qwen3_examples.md @@ -23,6 +23,35 @@ ASCEND_RT_VISIBLE_DEVICES=0,1,2,3 python -m sglang.launch_server \ --mem-fraction-static 0.8 ``` +#### Running Qwen3-32B on 1 x Atlas 800I A3 with Qwen3-32B-Eagle3. + +Model weights could be found [here](https://huggingface.co/Qwen/Qwen3-32B) + +Speculative model weights could be found [here](https://huggingface.co/Zhihu-ai/Zhi-Create-Qwen3-32B-Eagle3) + +```shell +export SGLANG_SET_CPU_AFFINITY=1 +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export STREAMS_PER_DEVICE=32 +export HCCL_OP_EXPANSION_MODE=AIV +export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 +export SGLANG_ENABLE_SPEC_V2=1 + +python -m sglang.launch_server \ + --device npu \ + --attention-backend ascend \ + --trust-remote-code \ + --tp-size 4 \ + --model-path Qwen/Qwen3-32B \ + --port 30111 \ + --mem-fraction-static 0.8 \ + --speculative-algorithm EAGLE3 \ + --speculative-draft-model-path Qwen/Qwen3-32B-Eagle3 \ + --speculative-num-steps 1 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 2 +``` + #### Running Qwen3-30B-A3B MOE on 1 x Atlas 800I A3. Model weights could be found [here](https://huggingface.co/Qwen/Qwen3-30B-A3B) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index c8963fbc5..c85062666 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -236,8 +236,14 @@ class ModelConfig: server_args: ServerArgs, model_path: str = None, model_revision: str = None, + is_draft_model: bool = False, **kwargs, ): + quantization = ( + server_args.speculative_draft_model_quantization + if is_draft_model + else server_args.quantization + ) return ModelConfig( model_path=model_path or server_args.model_path, trust_remote_code=server_args.trust_remote_code, @@ -247,7 +253,7 @@ class ModelConfig: is_embedding=server_args.is_embedding, enable_multimodal=server_args.enable_multimodal, dtype=server_args.dtype, - quantization=server_args.quantization, + quantization=quantization, hybrid_kvcache_ratio=server_args.hybrid_kvcache_ratio, model_impl=server_args.model_impl, sampling_defaults=server_args.sampling_defaults, @@ -255,6 +261,7 @@ class ModelConfig: override_config_file=server_args.decrypted_config_file, language_only=server_args.language_only, encoder_only=server_args.encoder_only, + is_draft_model=is_draft_model, **kwargs, ) diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index aee30fac7..ee3b41402 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -889,102 +889,163 @@ class AscendAttnBackend(AttentionBackend): layer, forward_batch.out_cache_loc, k, v ) - c_kv, k_rope = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) - if is_fia_nz(): - k_rope_cache = _reshape_kv_for_fia_nz( - k_rope, layer.tp_k_head_num, self.qk_rope_head_dim, self.page_size - ) - c_kv_cache = _reshape_kv_for_fia_nz( - c_kv, layer.tp_v_head_num, self.kv_lora_rank, self.page_size - ) - else: - k_rope_cache = k_rope.view( - -1, layer.tp_k_head_num, self.page_size, self.qk_rope_head_dim - ) - c_kv_cache = c_kv.view( - -1, layer.tp_v_head_num, self.page_size, self.kv_lora_rank - ) + if not self.use_mla: + k_cache = forward_batch.token_to_kv_pool.get_key_buffer( + layer.layer_id + ).view(-1, self.page_size, layer.tp_k_head_num * layer.qk_head_dim) + v_cache = forward_batch.token_to_kv_pool.get_value_buffer( + layer.layer_id + ).view(-1, self.page_size, layer.tp_v_head_num * layer.v_head_dim) + query = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim).contiguous() + if not self.graph_mode: + num_token_padding = query.shape[0] + query = query[: forward_batch.num_token_non_padded_cpu] + if self.forward_metadata.seq_lens_cpu_int is None: + actual_seq_lengths_kv = self.forward_metadata.seq_lens_cpu_list + else: + actual_seq_lengths_kv = ( + self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist() + ) + if forward_batch.forward_mode.is_draft_extend(): + actual_seq_lengths = ( + np.array(forward_batch.extend_seq_lens_cpu).cumsum().tolist() + ) + else: + actual_seq_lengths = np.arange( + self.speculative_num_draft_tokens, + self.speculative_num_draft_tokens + query.shape[0], + self.speculative_num_draft_tokens, + ) - q_nope = q.view(-1, layer.tp_q_head_num, self.kv_lora_rank).contiguous() - q_rope = q_rope.view(-1, layer.tp_q_head_num, self.qk_rope_head_dim) - if not self.graph_mode: - num_token_padding = q.shape[0] - q_nope = q_nope[: forward_batch.num_token_non_padded_cpu] - q_rope = q_rope[: forward_batch.num_token_non_padded_cpu] - if self.forward_metadata.seq_lens_cpu_int is None: - actual_seq_lengths_kv = self.forward_metadata.seq_lens_cpu_list + attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score( + query, + k_cache, + v_cache, + block_table=self.forward_metadata.block_tables, + block_size=self.page_size, + num_heads=layer.tp_q_head_num, + num_key_value_heads=layer.tp_k_head_num, + input_layout="TND", + atten_mask=self.mtp_mask, + scale=layer.scaling, + actual_seq_lengths=actual_seq_lengths, + actual_seq_lengths_kv=actual_seq_lengths_kv, + sparse_mode=3, + ) + attn_output = attn_output.view(-1, layer.tp_q_head_num * layer.v_head_dim) + if ( + not self.graph_mode + and forward_batch.num_token_non_padded_cpu != num_token_padding + ): + attn_output = torch.cat( + [ + attn_output, + attn_output.new_zeros( + num_token_padding - forward_batch.num_token_non_padded_cpu, + *attn_output.shape[1:], + ), + ], + dim=0, + ) + return attn_output else: - actual_seq_lengths_kv = ( - self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist() - ) - if forward_batch.forward_mode.is_draft_extend(): - actual_seq_lengths = ( - np.array(forward_batch.extend_seq_lens_cpu).cumsum().tolist() - ) - else: - actual_seq_lengths = np.arange( - self.speculative_num_draft_tokens, - self.speculative_num_draft_tokens + q_nope.shape[0], - self.speculative_num_draft_tokens, - ) + c_kv, k_rope = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) + if is_fia_nz(): + k_rope_cache = _reshape_kv_for_fia_nz( + k_rope, layer.tp_k_head_num, self.qk_rope_head_dim, self.page_size + ) + c_kv_cache = _reshape_kv_for_fia_nz( + c_kv, layer.tp_v_head_num, self.kv_lora_rank, self.page_size + ) + else: + k_rope_cache = k_rope.view( + -1, layer.tp_k_head_num, self.page_size, self.qk_rope_head_dim + ) + c_kv_cache = c_kv.view( + -1, layer.tp_v_head_num, self.page_size, self.kv_lora_rank + ) - workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace( - q_nope, - c_kv_cache, - c_kv_cache, - query_rope=q_rope, - key_rope=k_rope_cache, - num_heads=layer.tp_q_head_num, - num_key_value_heads=layer.tp_k_head_num, - input_layout="TND", - scale=layer.scaling, - antiquant_mode=0, - antiquant_scale=None, - block_table=self.forward_metadata.block_tables, - block_size=self.page_size, - sparse_mode=3, - atten_mask=self.mtp_mask, - actual_seq_lengths=actual_seq_lengths, - actual_seq_lengths_kv=actual_seq_lengths_kv, - ) - attn_output = torch.empty_like(q_nope, dtype=q.dtype, device=q.device) - softmax_lse = torch.empty(1, dtype=q.dtype, device=q.device) - torch_npu.npu_fused_infer_attention_score.out( - q_nope, - c_kv_cache, - c_kv_cache, - query_rope=q_rope, - key_rope=k_rope_cache, - num_heads=layer.tp_q_head_num, - num_key_value_heads=layer.tp_k_head_num, - input_layout="TND", - scale=layer.scaling, - antiquant_mode=0, - antiquant_scale=None, - block_table=self.forward_metadata.block_tables, - block_size=self.page_size, - sparse_mode=3, - atten_mask=self.mtp_mask, - actual_seq_lengths=actual_seq_lengths, - actual_seq_lengths_kv=actual_seq_lengths_kv, - workspace=workspace, - out=[attn_output, softmax_lse], - ) - attn_output = attn_output.view(-1, layer.tp_q_head_num * layer.v_head_dim) - if ( - not self.graph_mode - and forward_batch.num_token_non_padded_cpu != num_token_padding - ): - attn_output = torch.cat( - [ - attn_output, - attn_output.new_zeros( - num_token_padding - attn_output.shape[0], *attn_output.shape[1:] - ), - ], - dim=0, + q_nope = q.view(-1, layer.tp_q_head_num, self.kv_lora_rank).contiguous() + q_rope = q_rope.view(-1, layer.tp_q_head_num, self.qk_rope_head_dim) + if not self.graph_mode: + num_token_padding = q.shape[0] + q_nope = q_nope[: forward_batch.num_token_non_padded_cpu] + q_rope = q_rope[: forward_batch.num_token_non_padded_cpu] + if self.forward_metadata.seq_lens_cpu_int is None: + actual_seq_lengths_kv = self.forward_metadata.seq_lens_cpu_list + else: + actual_seq_lengths_kv = ( + self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist() + ) + if forward_batch.forward_mode.is_draft_extend(): + actual_seq_lengths = ( + np.array(forward_batch.extend_seq_lens_cpu).cumsum().tolist() + ) + else: + actual_seq_lengths = np.arange( + self.speculative_num_draft_tokens, + self.speculative_num_draft_tokens + q_nope.shape[0], + self.speculative_num_draft_tokens, + ) + + workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace( + q_nope, + c_kv_cache, + c_kv_cache, + query_rope=q_rope, + key_rope=k_rope_cache, + num_heads=layer.tp_q_head_num, + num_key_value_heads=layer.tp_k_head_num, + input_layout="TND", + scale=layer.scaling, + antiquant_mode=0, + antiquant_scale=None, + block_table=self.forward_metadata.block_tables, + block_size=self.page_size, + sparse_mode=3, + atten_mask=self.mtp_mask, + actual_seq_lengths=actual_seq_lengths, + actual_seq_lengths_kv=actual_seq_lengths_kv, ) - return attn_output + attn_output = torch.empty_like(q_nope, dtype=q.dtype, device=q.device) + softmax_lse = torch.empty(1, dtype=q.dtype, device=q.device) + torch_npu.npu_fused_infer_attention_score.out( + q_nope, + c_kv_cache, + c_kv_cache, + query_rope=q_rope, + key_rope=k_rope_cache, + num_heads=layer.tp_q_head_num, + num_key_value_heads=layer.tp_k_head_num, + input_layout="TND", + scale=layer.scaling, + antiquant_mode=0, + antiquant_scale=None, + block_table=self.forward_metadata.block_tables, + block_size=self.page_size, + sparse_mode=3, + atten_mask=self.mtp_mask, + actual_seq_lengths=actual_seq_lengths, + actual_seq_lengths_kv=actual_seq_lengths_kv, + workspace=workspace, + out=[attn_output, softmax_lse], + ) + attn_output = attn_output.view(-1, layer.tp_q_head_num * layer.v_head_dim) + if ( + not self.graph_mode + and forward_batch.num_token_non_padded_cpu != num_token_padding + ): + attn_output = torch.cat( + [ + attn_output, + attn_output.new_zeros( + num_token_padding - attn_output.shape[0], + *attn_output.shape[1:], + ), + ], + dim=0, + ) + return attn_output def forward_decode_graph( self, diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py index 895de235c..92308ca46 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py @@ -37,6 +37,9 @@ class EAGLEDraftExtendNpuGraphRunner(EAGLEDraftExtendCudaGraphRunner): def _create_graph(self): return torch.npu.NPUGraph() + def _cache_loc_dtype(self): + return torch.int32 + def _capture_init(self, run_once_fn): for _ in range(2): torch.npu.synchronize() diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py index 11ce91489..adca2c6c2 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py @@ -17,11 +17,13 @@ from __future__ import annotations import logging import threading -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Dict, Union +import numpy as np import torch -from sglang.srt.configs.model_config import is_deepseek_nsa +from sglang.srt.configs.model_config import AttentionArch, is_deepseek_nsa +from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( EAGLEDraftCudaGraphRunner, @@ -46,6 +48,19 @@ if is_npu(): class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner): def __init__(self, eagle_worker: EAGLEWorker): super().__init__(eagle_worker) + self.update_attr_name = None + self.update_attr_type = None + self._init_arch_map() + + def _init_arch_map(self): + self.attr_name: Dict[str, str] = { + AttentionArch.MLA: "actual_seq_lengths_kv", + AttentionArch.MHA: "context_lens", + } + self.attr_type: Dict[str, Union[list, torch.Tensor]] = { + AttentionArch.MLA: [], + AttentionArch.MHA: torch.Tensor(), + } def _create_graph(self): return torch.npu.NPUGraph() @@ -63,12 +78,27 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner): out = run_once_fn() return out + def _get_update_attr_name(self, model_runner): + if self.bs < get_attention_tp_size(): + return self.attr_name[AttentionArch.MLA] + return self.attr_name[model_runner.model_config.attention_arch] + + def _get_update_attr_type(self, model_runner): + if self.bs < get_attention_tp_size(): + return self.attr_type[AttentionArch.MLA] + return self.attr_type[model_runner.model_config.attention_arch] + def _replay_update(self, seq_lens): + if isinstance(self.update_attr_type, torch.Tensor): + seq_lens = torch.from_numpy(np.array(seq_lens).astype(np.int32)) + self.graphs[self.bs].update( - cpu_update_input=[{"actual_seq_lengths_kv": seq_lens}] + cpu_update_input=[{self.update_attr_name: seq_lens}] ) def _replay(self, forward_batch: ForwardBatch): + self.update_attr_name = self._get_update_attr_name(self.model_runner) + self.update_attr_type = self._get_update_attr_type(self.model_runner) if not is_deepseek_nsa(self.model_runner.model_config.hf_config): seq_lens = forward_batch.seq_lens_cpu.tolist() + [0] * ( self.bs - self.raw_bs @@ -79,3 +109,6 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner): thread.join() else: self.graphs[self.bs].replay() + + def _cache_loc_dtype(self): + return torch.int32 diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py index aa9241a77..838c5727d 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py @@ -77,13 +77,19 @@ class NPUGraphRunner(CudaGraphRunner): out = run_once_fn() return out - def _get_update_attr_name(self, model_runner): - if self.bs < get_attention_tp_size(): + def _get_update_attr_name(self, model_runner, forward_batch): + if ( + self.bs < get_attention_tp_size() + or forward_batch.forward_mode.is_target_verify() + ): return self.attr_name[AttentionArch.MLA] return self.attr_name[model_runner.model_config.attention_arch] - def _get_update_attr_type(self, model_runner): - if self.bs < get_attention_tp_size(): + def _get_update_attr_type(self, model_runner, forward_batch): + if ( + self.bs < get_attention_tp_size() + or forward_batch.forward_mode.is_target_verify() + ): return self.attr_type[AttentionArch.MLA] return self.attr_type[model_runner.model_config.attention_arch] @@ -139,8 +145,12 @@ class NPUGraphRunner(CudaGraphRunner): self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids) self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions) - self.update_attr_name = self._get_update_attr_name(self.model_runner) - self.update_attr_type = self._get_update_attr_type(self.model_runner) + self.update_attr_name = self._get_update_attr_name( + self.model_runner, forward_batch + ) + self.update_attr_type = self._get_update_attr_type( + self.model_runner, forward_batch + ) # Replay if not is_deepseek_nsa(self.model_runner.model_config.hf_config): if forward_batch.forward_mode.is_target_verify(): diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index 943476786..47ec111ba 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -89,7 +89,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): if self.store_dtype != self.dtype: cache_k = cache_k.view(self.store_dtype) cache_v = cache_v.view(self.store_dtype) - + loc = loc.to(torch.int32) torch_npu._npu_reshape_and_cache( key=cache_k, value=cache_v, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 4498f4214..f731ad432 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -111,6 +111,8 @@ QUANTIZATION_CHOICES = [ "modelslim", # for NPU ] +SPECULATIVE_DRAFT_MODEL_QUANTIZATION_CHOICES = [*QUANTIZATION_CHOICES, "unquant"] + ATTENTION_BACKEND_CHOICES = [ # Common "triton", @@ -431,6 +433,7 @@ class ServerArgs: speculative_attention_mode: str = "prefill" speculative_moe_runner_backend: Optional[str] = None speculative_moe_a2a_backend: Optional[str] = None + speculative_draft_model_quantization: Optional[str] = None # Speculative decoding (ngram) speculative_ngram_min_match_window_size: int = 1 @@ -747,8 +750,15 @@ class ServerArgs: # TODO: when extra_buffer is more verified, we can set the default path based on # [overlap, non-overlap] self.mamba_scheduler_strategy = "no_buffer" + # In speculative scenario: + # - If `speculative_draft_model_quantization` is specified, the draft model uses this quantization method. + # - Otherwise, the draft model defaults to the same quantization as the target model. + if self.speculative_draft_model_quantization is None: + self.speculative_draft_model_quantization = self.quantization + elif self.speculative_draft_model_quantization == "unquant": + self.speculative_draft_model_quantization = None - # Handle ModelScope model downloads + # Handle ModelScope model downloads if get_bool_env_var("SGLANG_USE_MODELSCOPE"): if not os.path.exists(self.model_path): from modelscope import snapshot_download @@ -3399,6 +3409,13 @@ class ServerArgs: default=ServerArgs.speculative_moe_a2a_backend, help="Choose the backend for MoE A2A in speculative decoding", ) + parser.add_argument( + "--speculative-draft-model-quantization", + type=str, + choices=SPECULATIVE_DRAFT_MODEL_QUANTIZATION_CHOICES, + default=ServerArgs.speculative_draft_model_quantization, + help="The quantization method for speculative model.", + ) # Speculative decoding (ngram) parser.add_argument( diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index cd9d171fe..3f5070aec 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -88,7 +88,8 @@ class EAGLEDraftCudaGraphRunner: self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64) self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32) self.out_cache_loc = torch.zeros( - (self.max_num_token * self.speculative_num_steps,), dtype=torch.int64 + (self.max_num_token * self.speculative_num_steps,), + dtype=self._cache_loc_dtype(), ) self.positions = torch.zeros((self.max_num_token,), dtype=torch.int64) self.mrope_positions = torch.zeros( @@ -132,6 +133,9 @@ class EAGLEDraftCudaGraphRunner: f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}" ) + def _cache_loc_dtype(self): + return torch.int64 + def can_run(self, forward_batch: ForwardBatch): if self.require_mlp_tp_gather: cuda_graph_bs = ( diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 0cfe2ef8f..c884bef8e 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -92,7 +92,9 @@ class EAGLEDraftExtendCudaGraphRunner: with torch.device(model_runner.device): self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64) self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32) - self.out_cache_loc = torch.ones((self.max_num_token,), dtype=torch.int64) + self.out_cache_loc = torch.ones( + (self.max_num_token,), dtype=self._cache_loc_dtype() + ) self.positions = torch.zeros((self.max_num_token,), dtype=torch.int64) self.mrope_positions = torch.zeros( (3, self.max_num_token), dtype=torch.int64 @@ -204,6 +206,9 @@ class EAGLEDraftExtendCudaGraphRunner: def _create_graph(self): return torch.cuda.CUDAGraph() + def _cache_loc_dtype(self): + return torch.int64 + def _capture_init(self, run_once_fn): for _ in range(2): torch.cuda.synchronize() diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 7c13ddc08..3894c2176 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -468,8 +468,6 @@ def assign_extend_cache_locs_func( return out_cache_loc elif _is_npu: - import sgl_kernel_npu # noqa: F401 - out_cache_loc = torch.empty( (batch_size * draft_token_num,), dtype=torch.int32, @@ -482,6 +480,5 @@ def assign_extend_cache_locs_func( end_offset, out_cache_loc, ) - out_cache_loc = out_cache_loc.to(dtype=torch.int64) return out_cache_loc