From 241ae17b25430e8599267bbe17881eb652fd2773 Mon Sep 17 00:00:00 2001
From: Hexq0210 <893781835@qq.com>
Date: Fri, 19 Dec 2025 22:41:17 +0800
Subject: [PATCH]  [NPU] bugfix for chunkedprefill (#15166)

---
 .../hardware_backend/npu/attention/ascend_backend.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

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