From 11baa0ebd88480b0a3e30138e4ee189420cd33be Mon Sep 17 00:00:00 2001 From: Xin Yang <105740670+xyang16@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:52:42 -0700 Subject: [PATCH] [Attention] Avoid redundant mask compute in GDN metadata build (#52078) Signed-off-by: Xin Yang --- vllm/v1/attention/backends/gdn_attn.py | 24 +++++++++++------------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 1a07a6957e..0e843e31ba 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -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:], )