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 c05988994..6a7fa7127 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -737,6 +737,7 @@ class AscendAttnBackend(AttentionBackend): causal=causal, ) elif sum(forward_batch.extend_prefix_lens_cpu) > 0: + num_token_padding = q.shape[0] q, k, v = [ data[: forward_batch.num_token_non_padded_cpu] for data in [q, k, v] ] @@ -824,6 +825,17 @@ class AscendAttnBackend(AttentionBackend): attn_output = attn_output.reshape( [-1, layer.tp_q_head_num, layer.v_head_dim] ) + if num_token_padding != forward_batch.num_token_non_padded_cpu: + attn_output = torch.cat( + [ + attn_output, + attn_output.new_zeros( + num_token_padding - attn_output.shape[0], + *attn_output.shape[1:], + ), + ], + dim=0, + ) else: assert ( layer.qk_head_dim != layer.v_head_dim