[NPU][eagle3] support qwen eagle3 on NPU (#14820)

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)

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,
         )
 

python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py

@@ -889,6 +889,66 @@ class AscendAttnBackend(AttentionBackend):
                     layer, forward_batch.out_cache_loc, k, v
                 )
 
+        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,
+                )
+
+            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:
             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(
@@ -979,7 +1039,8 @@ class AscendAttnBackend(AttentionBackend):
                     [
                         attn_output,
                         attn_output.new_zeros(
-                        num_token_padding - attn_output.shape[0], *attn_output.shape[1:]
+                            num_token_padding - attn_output.shape[0],
+                            *attn_output.shape[1:],
                         ),
                     ],
                     dim=0,

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()

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

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):
+    def _get_update_attr_name(self, model_runner, forward_batch):
-        if self.bs < get_attention_tp_size():
+        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):
+    def _get_update_attr_type(self, model_runner, forward_batch):
-        if self.bs < get_attention_tp_size():
+        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_name = self._get_update_attr_name(
-        self.update_attr_type = self._get_update_attr_type(self.model_runner)
+            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():

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,

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,6 +750,13 @@ 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
         if get_bool_env_var("SGLANG_USE_MODELSCOPE"):
@@ -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(

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 = (

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()

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