[NPU] fix extend_seq_lens_cpu shape in eager mode (#36843)
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user