From fe4b29d3911acc3ec36c410e286791b7915b8d38 Mon Sep 17 00:00:00 2001 From: AndyLi429 <68410213+AndyLi429@users.noreply.github.com> Date: Sat, 30 May 2026 15:07:39 +0800 Subject: [PATCH] [Bugfix] Fix Ascend NPU CP attention for batch size > 1 (#26705) --- .../npu/attention/ascend_backend.py | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) 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 690a8535a..68b37fa7b 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -820,12 +820,15 @@ class AscendAttnBackend(AttentionBackend): """ cp_meta = forward_batch.attn_cp_metadata - # Split Q into prev/next halves per zigzag pattern. - # torch.chunk(q, 2) gives ceil(n/2) and floor(n/2), matching - # actual_seq_q_prev and actual_seq_q_next. - q_prev, q_next = torch.chunk(q, 2, dim=0) - q_prev = q_prev.contiguous().reshape(-1, layer.tp_q_head_num, layer.qk_head_dim) - q_next = q_next.contiguous().reshape(-1, layer.tp_q_head_num, layer.qk_head_dim) + # Local tokens are laid out [all_seqs_prev, all_seqs_next]; split at + # total_q_prev_tokens rather than the midpoint to support bs > 1. + split = cp_meta.total_q_prev_tokens + q_prev = ( + q[:split].contiguous().reshape(-1, layer.tp_q_head_num, layer.qk_head_dim) + ) + q_next = ( + q[split:].contiguous().reshape(-1, layer.tp_q_head_num, layer.qk_head_dim) + ) k_cache_paged = k_cache.view( -1, self.page_size, layer.tp_k_head_num * layer.qk_head_dim @@ -847,8 +850,8 @@ class AscendAttnBackend(AttentionBackend): sparse_mode=3, next_tokens=0, scale=layer.scaling, - actual_seq_lengths=[cp_meta.actual_seq_q_prev], - actual_seq_lengths_kv=[cp_meta.kv_len_prev], + actual_seq_lengths=np.cumsum(cp_meta.actual_seq_q_prev_list).tolist(), + actual_seq_lengths_kv=cp_meta.kv_len_prev_list, ) attn_out_next, _ = torch.ops.npu.npu_fused_infer_attention_score( @@ -864,8 +867,8 @@ class AscendAttnBackend(AttentionBackend): sparse_mode=3, next_tokens=0, scale=layer.scaling, - actual_seq_lengths=[cp_meta.actual_seq_q_next], - actual_seq_lengths_kv=[cp_meta.kv_len_next], + actual_seq_lengths=np.cumsum(cp_meta.actual_seq_q_next_list).tolist(), + actual_seq_lengths_kv=cp_meta.kv_len_next_list, ) attn_out = torch.cat([attn_out_prev, attn_out_next], dim=0)