From 0916e7960bddf565c33153d2cf753e799de105b7 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Fri, 13 Feb 2026 01:24:45 -0800 Subject: [PATCH] [GDN] Use CPU tensors to build GDN metadata (#34498) Signed-off-by: Woosuk Kwon --- vllm/v1/attention/backends/gdn_attn.py | 17 ++++++++++------- vllm/v1/attention/backends/utils.py | 4 ++-- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index c7a41abe5fd..3f76f3e248e 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -206,21 +206,24 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] assert spec_sequence_masks_cpu is not None query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] - non_spec_query_lens = query_lens[~spec_sequence_masks] - num_decodes = (non_spec_query_lens == 1).sum().item() + # Use CPU tensors to avoid CPU-GPU sync + non_spec_query_lens_cpu = query_lens_cpu[~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 == 0).sum().item() - num_prefills = non_spec_query_lens.size(0) - num_decodes - num_zero_len + num_zero_len = (non_spec_query_lens_cpu == 0).sum().item() + num_prefills = non_spec_query_lens_cpu.size(0) - num_decodes - num_zero_len num_decode_tokens = num_decodes - num_prefill_tokens = non_spec_query_lens.sum().item() - num_decode_tokens + num_prefill_tokens = ( + non_spec_query_lens_cpu.sum().item() - num_decode_tokens + ) num_spec_decode_tokens = ( - query_lens.sum().item() - num_prefill_tokens - num_decode_tokens + query_lens_cpu.sum().item() - num_prefill_tokens - num_decode_tokens ) if num_prefills == 0 and num_decodes == 0: spec_token_size = min( num_spec_decodes * (self.num_spec + 1), - query_start_loc[-1].item(), + query_start_loc_cpu[-1].item(), ) spec_token_indx = torch.arange( spec_token_size, diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index e0aa2c988a2..eda50155dda 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -775,10 +775,10 @@ def compute_causal_conv1d_metadata( MAX_NUM_PROGRAMS ).fill_(PAD_SLOT_ID) - batch_ptr[0:mlist_len].copy_(mlist) + batch_ptr[0:mlist_len].copy_(mlist, non_blocking=True) token_chunk_offset_ptr[ # type: ignore 0:mlist_len - ].copy_(offsetlist) + ].copy_(offsetlist, non_blocking=True) nums_dict[BLOCK_M]["batch_ptr"] = batch_ptr nums_dict[BLOCK_M]["token_chunk_offset_ptr"] = token_chunk_offset_ptr # type: ignore