From 12af7e6c34b3c75b120594f19555f90d43999184 Mon Sep 17 00:00:00 2001 From: Peng Xingchen <112612131+stellaxcpeng@users.noreply.github.com> Date: Fri, 17 Jul 2026 12:05:57 +0800 Subject: [PATCH] [NPU] Fix DSA top-k seed buffer shape for MTP IndexShare (#31005) --- .../srt/hardware_backend/npu/attention/ascend_backend.py | 8 ++++++++ python/sglang/srt/layers/attention/dsa/dsa_indexer.py | 3 ++- 2 files changed, 10 insertions(+), 1 deletion(-) 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 62d3b974a..e95072db6 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -44,6 +44,13 @@ logger = logging.getLogger(__name__) FULL_ATTENTION_WINDOW = 2147483647 +def _expand_dsa_sparse_indices(topk_indices: torch.Tensor) -> torch.Tensor: + """Expand [T, K] to [T, 1, K] for NPU sparse attention.""" + if topk_indices.dim() == 2: + return topk_indices.unsqueeze(-2) + return topk_indices + + def _reshape_kv_for_fia_nz( tensor: torch.Tensor, num_heads: int, head_dim: int, page_size: int ) -> torch.Tensor: @@ -1053,6 +1060,7 @@ class AscendAttnBackend(AttentionBackend): actual_seq_lengths_kv, ) else: + topk_indices = _expand_dsa_sparse_indices(topk_indices) attn_out, _, _ = torch_npu.npu_sparse_flash_attention( query=q_nope, key=k_nope, diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index e6a1b8331..196639f18 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -2342,7 +2342,8 @@ class Indexer(MultiPlatformOp): sparse_count=self.index_topk, sparse_mode=3, ) - return topk_indices[0] + # Keep DSA top-k as [T, K]; NPU attention expands it when needed. + return topk_indices[0].squeeze(1) def do_npu_cp_balance_indexer( self,