Support MTP for Qwen3.5

python/sglang/srt/hardware_backend/npu/memory_pool_npu.py

@@ -15,6 +15,30 @@
     from sglang.srt.layers.radix_attention import RadixAttention
 
 
+def _init_npu_conv_state(
+    conv_state_in, conv_state_shape, speculative_num_draft_tokens: Optional[int] = None
+):
+    extra_conv_len = 0
+    if speculative_num_draft_tokens is not None:
+        extra_conv_len = speculative_num_draft_tokens - 1
+
+    # conv_state shape (layers, pool_size, conv_wind + draft_step, dim) for conv1d ascendc ops require dim as last dim
+    conv_state = [
+        torch.zeros(
+            size=(
+                conv_state_in.shape[0],
+                conv_state_in.shape[1],
+                conv_shape[1] + extra_conv_len,
+                conv_shape[0],
+            ),
+            dtype=conv_state_in.dtype,
+            device=conv_state_in.device,
+        )
+        for conv_shape in conv_state_shape
+    ]
+    return conv_state
+
+
 class NPUMHATokenToKVPool(MHATokenToKVPool):
 
     def __init__(

python/sglang/srt/layers/attention/attention_registry.py

@@ -192,7 +192,6 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
             HybridLinearAttnBackend,
             Mamba2AttnBackend,
         )
-        from sglang.srt.layers.attention.linear.gdn_backend import GDNAttnBackend
         from sglang.srt.layers.attention.linear.kda_backend import KDAAttnBackend
         from sglang.srt.layers.attention.linear.lightning_backend import (
             LightningAttentionBackend,
@@ -202,6 +201,13 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
         )
         from sglang.srt.utils import is_blackwell, is_npu
 
+        if is_npu():
+            from sglang.srt.hardware_backend.npu.attention.ascend_gdn_backend import (
+                AscendGDNAttnBackend as GDNAttnBackend,
+            )
+        else:
+            from sglang.srt.layers.attention.linear.gdn_backend import GDNAttnBackend
+
         check_environments()
         initialize_linear_attn_config(runner.server_args)
         if runner.hybrid_gdn_config is not None:

python/sglang/srt/layers/attention/fla/fused_gdn_gating.py

@@ -67,3 +67,69 @@ def fused_gdn_gating(
         num_warps=1,
     )
     return g, beta_output
+
+
+@triton.jit
+def fused_gdn_gating_kernel_without_sigmoid_kernel(
+    g,
+    A_log,
+    a,
+    dt_bias,
+    batch,
+    seq_len,
+    NUM_HEADS: tl.constexpr,
+    beta: tl.constexpr,
+    threshold: tl.constexpr,
+    BLK_BATCHES: tl.constexpr,
+    BLK_HEADS: tl.constexpr,
+):
+    i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    batch_off = i_b * BLK_BATCHES + tl.arange(0, BLK_BATCHES)
+    head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS)
+    head_mask = head_off < NUM_HEADS
+    a_off = (
+        batch_off[:, None] * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off[None, :]
+    )
+    a_mask = (batch_off[:, None] < batch) & head_mask[None, :]
+    blk_A_log = tl.load(A_log + head_off, mask=head_mask)
+    blk_bias = tl.load(dt_bias + head_off, mask=head_mask)
+    blk_a = tl.load(a + a_off, mask=a_mask)
+    x = blk_a.to(tl.float32) + blk_bias.to(tl.float32)
+    softplus_x = tl.where(
+        beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x
+    )
+    blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x
+    tl.store(g + a_off, blk_g.to(g.dtype.element_ty), mask=a_mask)
+
+
+def fused_gdn_gating_kernel_without_sigmoid(
+    A_log: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    dt_bias: torch.Tensor,
+    beta: float = 1.0,
+    threshold: float = 20.0,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    batch, num_heads = a.shape
+    seq_len = 1
+    g = torch.empty_like(a, dtype=torch.float32)
+    num_cores = 48  # num_vectorcore of NPU
+    NUM_BLK_BATCHES = triton.cdiv(num_cores, triton.cdiv(num_heads, 8))
+    BLK_BATCHES = triton.cdiv(batch, NUM_BLK_BATCHES)
+    grid = (NUM_BLK_BATCHES, seq_len, triton.cdiv(num_heads, 8))
+    fused_gdn_gating_kernel_without_sigmoid_kernel[grid](
+        g,
+        A_log,
+        a,
+        dt_bias,
+        batch,
+        seq_len,
+        num_heads,
+        beta,
+        threshold,
+        BLK_BATCHES=BLK_BATCHES,
+        BLK_HEADS=8,
+        num_warps=1,
+    )
+    g = g.unsqueeze(0)
+    return g, b

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

@@ -22,13 +22,19 @@
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
 from sglang.srt.speculative.spec_info import SpecInput
-from sglang.srt.utils import is_cpu
+from sglang.srt.utils import is_cpu, is_npu
 
 if not is_cpu():
     from sglang.srt.layers.attention.fla.chunk_delta_h import (
         CHUNK_SIZE as FLA_CHUNK_SIZE,
     )
 
+if is_npu():
+    from sgl_kernel_npu.mamba.mamba_state_update_triton import (
+        conv_state_rollback,
+        move_intermediate_cache_dynamic_h_block,
+    )
+
 logger = logging.getLogger(__name__)
 
 
@@ -142,6 +148,7 @@ def __init__(self, model_runner: ModelRunner):
         self.req_to_token_pool: HybridReqToTokenPool = model_runner.req_to_token_pool
         self.forward_metadata: ForwardMetadata = None
         self.state_indices_list = []
+        self.state_indices_list_gdn = []
         self.query_start_loc_list = []
         self.retrieve_next_token_list = []
         self.retrieve_next_sibling_list = []
@@ -409,6 +416,14 @@ def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
                     (i + 1,), self.pad_slot_id, dtype=torch.int32, device=self.device
                 )
             )
+            self.state_indices_list_gdn.append(
+                torch.full(
+                    ((i + 1) * draft_token_num,),
+                    self.pad_slot_id,
+                    dtype=torch.int32,
+                    device=self.device,
+                )
+            )
             self.query_start_loc_list.append(
                 torch.zeros((i + 2,), dtype=torch.int32, device=self.device)
             )
@@ -462,6 +477,8 @@ def _capture_metadata(
         forward_mode: ForwardMode,
         spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
     ):
+        mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices)
+        self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
         if forward_mode.is_decode_or_idle():
             self.query_start_loc_list[bs - 1].copy_(
                 self.cached_cuda_graph_decode_query_start_loc[: bs + 1]
@@ -470,10 +487,17 @@ def _capture_metadata(
             self.query_start_loc_list[bs - 1].copy_(
                 self.cached_cuda_graph_verify_query_start_loc[: bs + 1]
             )
+            start_indices = mamba_indices * spec_info.draft_token_num
+            offset = torch.arange(
+                spec_info.draft_token_num, device=start_indices.device
+            )
+            ranges = start_indices.unsqueeze(1) + offset
+            ssm_state_indices = ranges.flatten().to(torch.int32)
+            self.state_indices_list_gdn[bs - 1][
+                : len(mamba_indices) * spec_info.draft_token_num
+            ].copy_(ssm_state_indices)
         else:
             raise ValueError(f"Invalid forward mode: {forward_mode=}")
-        mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices)
-        self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
 
         # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask
         if forward_mode.is_target_verify() and spec_info.topk > 1:
@@ -491,6 +515,7 @@ def _capture_metadata(
             return ForwardMetadata(
                 query_start_loc=self.query_start_loc_list[bs - 1],
                 mamba_cache_indices=self.state_indices_list[bs - 1],
+                mamba_cache_indices_gdn=self.state_indices_list_gdn[bs - 1],
             )
 
     def _replay_metadata(
@@ -507,7 +532,7 @@ def _replay_metadata(
         # Make sure forward metadata is correctly handled for padding reqs
         req_pool_indices[bs - num_padding :] = 0
         mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices)
-        mamba_indices[bs - num_padding :] = -1
+        mamba_indices[bs - num_padding :] = 0
         self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
         if forward_mode.is_decode_or_idle():
             if num_padding == 0:
@@ -522,6 +547,20 @@ def _replay_metadata(
                     bs - num_padding
                 )
         elif forward_mode.is_target_verify():
+            start_indices = (
+                mamba_indices[: bs - num_padding] * spec_info.draft_token_num
+            )
+            offset = torch.arange(
+                spec_info.draft_token_num, device=start_indices.device
+            )
+            ranges = start_indices.unsqueeze(1) + offset
+            ssm_state_indices = ranges.flatten().to(torch.int32)
+            self.state_indices_list_gdn[bs - 1][
+                : len(mamba_indices[: bs - num_padding]) * spec_info.draft_token_num
+            ].copy_(ssm_state_indices)
+            self.state_indices_list_gdn[bs - 1][
+                len(mamba_indices[: bs - num_padding]) * spec_info.draft_token_num :
+            ] = 0
             if num_padding == 0:
                 self.query_start_loc_list[bs - 1].copy_(
                     self.cached_cuda_graph_verify_query_start_loc[: bs + 1]
@@ -556,10 +595,11 @@ def _replay_metadata(
             return ForwardMetadata(
                 query_start_loc=self.query_start_loc_list[bs - 1],
                 mamba_cache_indices=self.state_indices_list[bs - 1],
+                mamba_cache_indices_gdn=self.state_indices_list_gdn[bs - 1],
             )
 
     def get_cuda_graph_seq_len_fill_value(self):
-        return 1  # Mamba attn does not use seq lens to index kv cache
+        return 0  # Mamba attn does not use seq lens to index kv cache
 
     def get_cpu_graph_seq_len_fill_value(self):
         return 1
@@ -960,6 +1000,23 @@ def update_mamba_state_after_mtp_verify(
         ssm_states = mamba_caches.temporal
         intermediate_state_cache = mamba_caches.intermediate_ssm
         intermediate_conv_window_cache = mamba_caches.intermediate_conv_window[0]
+        if is_npu():
+            valid_state_indices = state_indices_tensor.to(torch.int64)  # [N]
+            last_steps = accepted_steps.to(torch.int64)  # [N]
+
+            move_intermediate_cache_dynamic_h_block(
+                ssm_states, intermediate_state_cache, valid_state_indices, last_steps
+            )
+
+            draft_token_num = intermediate_state_cache.shape[2]
+            if valid_state_indices.numel() > 0:
+                conv_state_rollback(
+                    conv_states,
+                    valid_state_indices,
+                    last_steps,
+                    draft_token_num,
+                )
+            return
 
         # Use fully fused kernel that handles masking internally
         # This avoids separate nonzero() and index_select() calls
@@ -992,3 +1049,8 @@ def update_mamba_state_after_mtp_verify(
                 mamba_track_indices,
                 mamba_steps_to_track,
             )
+
+    def update_verify_buffers_to_fill_after_draft(
+        self, spec_info: SpecInput, cuda_graph_bs: Optional[int]
+    ):
+        pass

python/sglang/srt/layers/attention/mamba/mamba2_metadata.py

@@ -27,6 +27,7 @@
 class ForwardMetadata:
     query_start_loc: torch.Tensor
     mamba_cache_indices: torch.Tensor
+    mamba_cache_indices_gdn: Optional[torch.Tensor] = None
     # For topk > 1 eagle
     retrieve_next_token: Optional[torch.Tensor] = None
     retrieve_next_sibling: Optional[torch.Tensor] = None

python/sglang/srt/layers/layernorm.py

@@ -84,6 +84,7 @@
 
 if _is_npu:
     import torch_npu
+    from sgl_kernel_npu.norm.add_rmsnorm_bias import add_gemma_rms_norm
 
 
 class RMSNorm(MultiPlatformOp):
@@ -512,11 +513,13 @@ def forward_npu(
         if residual is not None:
             if post_residual_addition is not None:
                 residual = residual + post_residual_addition
-            x = x + residual
-            residual = x
+            norm_out, residual = add_gemma_rms_norm(
+                x, self.weight, residual, self.variance_epsilon
+            )
+            return norm_out, residual
 
         x, _ = torch_npu.npu_gemma_rms_norm(x, self.weight, self.variance_epsilon)
-        return x if residual is None else (x, residual)
+        return x
 
     def forward_xpu(
         self,

python/sglang/srt/mem_cache/memory_pool.py

@@ -258,6 +258,15 @@ def __init__(
                 for conv_shape in conv_state_shape
             ]
 
+            if _is_npu:
+                from sglang.srt.hardware_backend.npu.memory_pool_npu import (
+                    _init_npu_conv_state,
+                )
+
+                conv_state = _init_npu_conv_state(
+                    conv_state[0], conv_state_shape, speculative_num_draft_tokens
+                )
+
             if _is_cpu and _cpu_has_amx_support:
                 from sglang.srt.layers.amx_utils import _init_amx_conv_state
 

python/sglang/srt/models/qwen3_5_mtp.py

@@ -15,6 +15,7 @@
 """Inference-only Qwen3_5 MTP model."""
 
 import logging
+import os
 from typing import Iterable, Optional, Tuple
 
 import torch
@@ -31,7 +32,8 @@
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.qwen3_5 import Qwen3_5ForCausalLM
-from sglang.srt.utils import add_prefix
+from sglang.srt.server_args import get_global_server_args
+from sglang.srt.utils import add_prefix, is_npu
 
 logger = logging.getLogger(__name__)
 
@@ -53,6 +55,9 @@ def __init__(
         # The MTP model is unquantized in the nvfp4 checkpoint.
         if quant_config and quant_config.get_name() == "modelopt_fp4":
             quant_config = None
+        if get_global_server_args().speculative_draft_model_quantization is None:
+            quant_config = None
+        self.quant_config = quant_config
 
         self.config = config
         self.tp_size = get_tensor_model_parallel_world_size()
@@ -118,6 +123,10 @@ def forward(
         input_embeds: Optional[torch.Tensor] = None,
         **kwargs,
     ):
+        if is_npu() and self.quant_config is None:
+            # ascend mtp unquant
+            os.environ["SGLANG_DEEPEP_BF16_DISPATCH"] = "1"
+            os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "0"
         assert input_embeds is None
         input_embeds = forward_batch.mm_input_embeds
         if (
@@ -149,6 +158,10 @@ def forward(
                 forward_batch,
                 hidden_states,
             )
+        if is_npu() and self.quant_config is None:
+            # ascend mtp unquant
+            os.environ["SGLANG_DEEPEP_BF16_DISPATCH"] = "0"
+            os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "1"
 
         return self.logits_processor(
             input_ids, hidden_states, self.lm_head, forward_batch