[Attention] Move context_lens_tensor compute into GDN prefill path (#51913)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Aug 12, 2026 at 00:35 UTC
10b7766a901c06ca35ca9fd995d5d43b3a585622
1 file changed
+1
-1
vllm/v1/attention/backends/gdn_attn.py
+1
-1
@@ -177,7 +177,6 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
177
178
query_start_loc = m.query_start_loc
179
query_start_loc_cpu = m.query_start_loc_cpu
180
- context_lens_tensor = m.compute_num_computed_tokens()
180
nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
181
block_table_tensor = mamba_get_block_table_tensor(
182
m.block_table_tensor,
@@ -389,6 +388,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
388
)
389
390
if num_prefills > 0:
391
+ context_lens_tensor = m.compute_num_computed_tokens()
392
has_initial_state = context_lens_tensor > 0
393
if spec_sequence_masks_cpu is not None:
394
has_initial_state = has_initial_state[~spec_sequence_masks_cpu]