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,