[NPU] fix extend_seq_lens_cpu shape in eager mode (#36843)

This commit is contained in:
heziiop
2026-09-04 09:03:30 +08:00
committed by GitHub
parent 6147a54ddf
commit 9ed2721c6d
2 changed files with 15 additions and 9 deletions
@@ -2021,9 +2021,12 @@ class AscendAttnBackend(AttentionBackend):
)
if forward_batch.forward_mode.is_draft_extend_v2():
actual_seq_lengths = (
np.array(forward_batch.extend_seq_lens_cpu).cumsum().tolist()
)
extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
if not self.graph_mode:
extend_seq_lens_cpu = extend_seq_lens_cpu[
: forward_batch._original_batch_size
]
actual_seq_lengths = np.array(extend_seq_lens_cpu).cumsum().tolist()
else:
actual_seq_lengths = np.arange(
self.speculative_num_draft_tokens,
@@ -2031,6 +2034,10 @@ class AscendAttnBackend(AttentionBackend):
self.speculative_num_draft_tokens,
)
if not self.graph_mode:
actual_bs = len(actual_seq_lengths)
actual_seq_lengths_kv = actual_seq_lengths_kv[:actual_bs]
is_swa_layer = layer.sliding_window_size != -1
if (
is_swa_layer