From 190488e9a84ba34661618aab5668239e2ce22f0e Mon Sep 17 00:00:00 2001
From: Liwansi <62291011+Liwansi@users.noreply.github.com>
Date: Thu, 21 May 2026 14:44:25 +0800
Subject: [PATCH] [NPU] Support chunk prefill for Qwen3.5/Qwen3.6 models
 (#25839)

---
 .../hardware_backend/npu/attention/ascend_gdn_backend.py | 9 +--------
 python/sglang/srt/mem_cache/memory_pool.py               | 7 +++++++
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py
index cdb1dadb7..e15d3d239 100644
--- a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py
+++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py
@@ -350,14 +350,7 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase):
                     ssm_states.dtype, copy=False
                 )
                 ssm_states[cache_indices] = last_recurrent_state
-            if not forward_batch.spec_algorithm.is_none():
-                last_recurrent_state = last_recurrent_state.transpose(-1, -2).to(
-                    ssm_states.dtype, copy=False
-                )
-            else:
-                last_recurrent_state = last_recurrent_state.to(
-                    ssm_states.dtype, copy=False
-                )
+            last_recurrent_state = last_recurrent_state.to(ssm_states.dtype, copy=False)
             ssm_states[cache_indices] = last_recurrent_state
             if h is not None:
                 self._track_mamba_state_extend(
diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py
index 9f06d99c1..b1353902b 100644
--- a/python/sglang/srt/mem_cache/memory_pool.py
+++ b/python/sglang/srt/mem_cache/memory_pool.py
@@ -289,6 +289,13 @@ class MambaPool:
                 device=device,
             )
             if speculative_num_draft_tokens is not None:
+                if _is_npu:
+                    temporal_state = temporal_state.transpose(-1, -2)
+                    temporal_state_shape = (
+                        *temporal_state_shape[:-2],
+                        temporal_state_shape[-1],
+                        temporal_state_shape[-2],
+                    )
                 # Cache intermediate SSM states per draft token during target verify
                 # Shape: [num_layers, size + 1, speculative_num_draft_tokens, HV, K, V]
                 intermediate_ssm_state_cache = torch.zeros(