iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Attention] Avoid redundant mask compute in GDN metadata build (#52078)

Browse Source 
Signed-off-by: Xin Yang <xyangx@amazon.com>
This commit is contained in:
Xin Yang
2026-08-20 11:52:42 +08:00
committed by GitHub
parent 4f66bc3e0d
commit 11baa0ebd8
1 changed files with 11 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/v1/attention/backends/gdn_attn.py
+11 -13
View File
@@ -228,20 +228,17 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
)
spec_sequence_masks_cpu: torch.Tensor | None = None
if (
not self.use_spec_decode
or num_decode_draft_tokens_cpu is None
or num_decode_draft_tokens_cpu[num_decode_draft_tokens_cpu >= 0]
.sum()
.item()
== 0
):
if not self.use_spec_decode or num_decode_draft_tokens_cpu is None:
spec_sequence_masks = None
num_spec_decodes = 0
else:
spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0
num_spec_decodes = spec_sequence_masks_cpu.sum().item()
if num_spec_decodes == 0:
if (
num_spec_decodes == 0
or num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item()
== 0
):
spec_sequence_masks = None
spec_sequence_masks_cpu = None
else:
@@ -265,10 +262,11 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
else:
query_lens = query_start_loc[1:] - query_start_loc[:-1]
assert spec_sequence_masks_cpu is not None
non_spec_sequence_masks_cpu = ~spec_sequence_masks_cpu
query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
# Use CPU tensors to avoid CPU-GPU sync
non_spec_query_lens_cpu = query_lens_cpu[~spec_sequence_masks_cpu]
non_spec_query_lens_cpu = query_lens_cpu[non_spec_sequence_masks_cpu]
num_decodes = (non_spec_query_lens_cpu == 1).sum().item()
# Exclude zero-length padded sequences from prefill count.
num_zero_len = (non_spec_query_lens_cpu == 0).sum().item()
@@ -330,7 +328,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
spec_sequence_masks_cpu, : self.num_spec + 1
]
non_spec_state_indices_tensor = block_table_tensor[
~spec_sequence_masks_cpu, 0
non_spec_sequence_masks_cpu, 0
]
spec_query_start_loc = torch.zeros(
@@ -349,7 +347,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
device=query_start_loc.device,
)
torch.cumsum(
query_lens[~spec_sequence_masks_cpu],
query_lens[non_spec_sequence_masks_cpu],
dim=0,
out=non_spec_query_start_loc[1:],
)
@@ -358,7 +356,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
dtype=torch.int32,
)
torch.cumsum(
query_lens_cpu[~spec_sequence_masks_cpu],
query_lens_cpu[non_spec_sequence_masks_cpu],
dim=0,
out=non_spec_query_start_loc_cpu[1:],
)