A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 07:30 UTC
[GDN] Use CPU tensors to build GDN metadata (#34498)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 13, 2026 at 01:24 UTC
0916e7960bddf565c33153d2cf753e799de105b7
2 files changed
+12
-9
vllm/v1/attention/backends/gdn_attn.py
+10
-7
@@ -206,21 +206,24 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
206
assert spec_sequence_masks_cpu is not None
207
query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
208
209
- non_spec_query_lens = query_lens[~spec_sequence_masks]
210
- num_decodes = (non_spec_query_lens == 1).sum().item()
209
+ # Use CPU tensors to avoid CPU-GPU sync
210
+ non_spec_query_lens_cpu = query_lens_cpu[~spec_sequence_masks_cpu]
211
+ num_decodes = (non_spec_query_lens_cpu == 1).sum().item()
212
# Exclude zero-length padded sequences from prefill count.
212
- num_zero_len = (non_spec_query_lens == 0).sum().item()
213
- num_prefills = non_spec_query_lens.size(0) - num_decodes - num_zero_len
213
+ num_zero_len = (non_spec_query_lens_cpu == 0).sum().item()
214
+ num_prefills = non_spec_query_lens_cpu.size(0) - num_decodes - num_zero_len
215
num_decode_tokens = num_decodes
215
- num_prefill_tokens = non_spec_query_lens.sum().item() - num_decode_tokens
216
+ num_prefill_tokens = (
217
+ non_spec_query_lens_cpu.sum().item() - num_decode_tokens
218
+ )
219
num_spec_decode_tokens = (
217
- query_lens.sum().item() - num_prefill_tokens - num_decode_tokens
220
+ query_lens_cpu.sum().item() - num_prefill_tokens - num_decode_tokens
221
)
222
223
if num_prefills == 0 and num_decodes == 0:
224
spec_token_size = min(
225
num_spec_decodes * (self.num_spec + 1),
223
- query_start_loc[-1].item(),
226
+ query_start_loc_cpu[-1].item(),
227
)
228
spec_token_indx = torch.arange(
229
spec_token_size,
vllm/v1/attention/backends/utils.py
+2
-2
@@ -775,10 +775,10 @@ def compute_causal_conv1d_metadata(
775
MAX_NUM_PROGRAMS
776
).fill_(PAD_SLOT_ID)
777
778
- batch_ptr[0:mlist_len].copy_(mlist)
778
+ batch_ptr[0:mlist_len].copy_(mlist, non_blocking=True)
779
token_chunk_offset_ptr[ # type: ignore
780
0:mlist_len
781
- ].copy_(offsetlist)
781
+ ].copy_(offsetlist, non_blocking=True)
782
nums_dict[BLOCK_M]["batch_ptr"] = batch_ptr
783
nums_dict[BLOCK_M]["token_chunk_offset_ptr"] = token_chunk_offset_ptr # type: ignore
784