From cd6efcb947d01bc761b31f7b579dcc7f1372abf3 Mon Sep 17 00:00:00 2001
From: AndyLi429 <68410213+AndyLi429@users.noreply.github.com>
Date: Tue, 9 Jun 2026 15:47:23 +0800
Subject: [PATCH] [NPU][Bugfix] fix MTP accuracy regression on Qwen3 hybrid
 models (#27202)

---
 .../ascend_hybrid_linear_attn_backend.py      | 40 +++++++++++++++----
 .../hardware_backend/npu/memory_pool_npu.py   |  5 +++
 2 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py
index f11aecc11..01ec313ca 100644
--- a/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py
+++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py
@@ -157,16 +157,11 @@ class AscendMambaAttnBackendBase(MambaAttnBackendBase):
                 )
         elif forward_mode.is_target_verify():
             ssm_state_indices = torch.arange(
-                len(mamba_indices[: bs - num_padding]) * spec_info.draft_token_num,
+                bs * spec_info.draft_token_num,
                 dtype=torch.int32,
                 device=mamba_indices.device,
             )
-            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
+            self.state_indices_list_gdn[bs - 1].copy_(ssm_state_indices)
             if num_padding == 0:
                 self.query_start_loc_list[bs - 1].copy_(
                     self.cached_cuda_graph_verify_query_start_loc[: bs + 1]
@@ -269,6 +264,28 @@ class AscendHybridLinearAttnBackend(HybridLinearAttnBackend):
         )
 
         draft_token_num = intermediate_state_cache.shape[2]
+        if mamba_track_indices is not None:
+            assert mamba_steps_to_track is not None
+            mamba_track_indices = mamba_track_indices.to(torch.int64)
+            mamba_steps_to_track = mamba_steps_to_track.to(torch.int64)
+
+            move_intermediate_cache(
+                ssm_states,
+                intermediate_state_cache,
+                mamba_track_indices,
+                src_indices_tensor,
+                mamba_steps_to_track,
+            )
+
+            track_mask = mamba_steps_to_track >= 0
+            # Track conv state from the verify-time window before rolling back
+            # the working slot; NPU does not keep per-step conv intermediates.
+            track_indices = mamba_track_indices[track_mask]
+            if track_indices.numel() > 0:
+                conv_states[:, track_indices] = conv_states[
+                    :, dst_indices_tensor[track_mask]
+                ]
+
         if dst_indices_tensor.numel() > 0:
             conv_state_rollback(
                 conv_states,
@@ -276,6 +293,15 @@ class AscendHybridLinearAttnBackend(HybridLinearAttnBackend):
                 last_steps,
                 draft_token_num,
             )
+
+        if mamba_track_indices is not None and mamba_track_indices.numel() > 0:
+            conv_state_rollback(
+                conv_states,
+                mamba_track_indices,
+                mamba_steps_to_track,
+                draft_token_num,
+            )
+
         return
 
     def update_verify_buffers_to_fill_after_draft(
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 145993f7e..c9f55ba12 100644
--- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py
+++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py
@@ -109,6 +109,11 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
                     self.k_buffer.append(k_buffer_layer)
                     self.v_buffer.append(v_buffer_layer)
 
+    def _init_kv_copy_and_warmup(self):
+        # implementation relies on self.data_strides / self.data_ptrs, which the
+        # NPU paged buffer layout never builds.
+        self._kv_copy_config = None
+
     # for disagg
     def get_contiguous_buf_infos(self):
         # layer_num x [seq_len, head_num, head_dim]