[Attention] Avoid redundant mask compute in GDN metadata build (#52078)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Aug 19, 2026 at 20:52 UTC
11baa0ebd88480b0a3e30138e4ee189420cd33be
1 file changed
+11
-13
vllm/v1/attention/backends/gdn_attn.py
+11
-13
@@ -228,20 +228,17 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
228
)
229
230
spec_sequence_masks_cpu: torch.Tensor | None = None
231
- if (
232
- not self.use_spec_decode
233
- or num_decode_draft_tokens_cpu is None
234
- or num_decode_draft_tokens_cpu[num_decode_draft_tokens_cpu >= 0]
235
- .sum()
236
- .item()
237
- == 0
238
- ):
231
+ if not self.use_spec_decode or num_decode_draft_tokens_cpu is None:
232
spec_sequence_masks = None
233
num_spec_decodes = 0
234
else:
235
spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0
236
num_spec_decodes = spec_sequence_masks_cpu.sum().item()
244
- if num_spec_decodes == 0:
237
+ if (
238
+ num_spec_decodes == 0
239
+ or num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item()
240
+ == 0
241
+ ):
242
spec_sequence_masks = None
243
spec_sequence_masks_cpu = None
244
else:
@@ -265,10 +262,11 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
262
else:
263
query_lens = query_start_loc[1:] - query_start_loc[:-1]
264
assert spec_sequence_masks_cpu is not None
265
+ non_spec_sequence_masks_cpu = ~spec_sequence_masks_cpu
266
query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
267
268
# Use CPU tensors to avoid CPU-GPU sync
271
- non_spec_query_lens_cpu = query_lens_cpu[~spec_sequence_masks_cpu]
269
+ non_spec_query_lens_cpu = query_lens_cpu[non_spec_sequence_masks_cpu]
270
num_decodes = (non_spec_query_lens_cpu == 1).sum().item()
271
# Exclude zero-length padded sequences from prefill count.
272
num_zero_len = (non_spec_query_lens_cpu == 0).sum().item()
@@ -330,7 +328,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
328
spec_sequence_masks_cpu, : self.num_spec + 1
329
]
330
non_spec_state_indices_tensor = block_table_tensor[
333
- ~spec_sequence_masks_cpu, 0
331
+ non_spec_sequence_masks_cpu, 0
332
]
333
334
spec_query_start_loc = torch.zeros(
@@ -349,7 +347,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
347
device=query_start_loc.device,
348
)
349
torch.cumsum(
352
- query_lens[~spec_sequence_masks_cpu],
350
+ query_lens[non_spec_sequence_masks_cpu],
351
dim=0,
352
out=non_spec_query_start_loc[1:],
353
)
@@ -358,7 +356,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
356
dtype=torch.int32,
357
)
358
torch.cumsum(
361
- query_lens_cpu[~spec_sequence_masks_cpu],
359
+ query_lens_cpu[non_spec_sequence_masks_cpu],
360
dim=0,
361
out=non_spec_query_start_loc_cpu[1:],
362
)