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 5, 12:15 UTC
[Models][GDN] Remove GPU/CPU syncs in `GDNAttentionMetadata.build` during speculative decoding (#38047)
Signed-off-by: Lukas Geiger <lukas.geiger94@gmail.com>
Lukas Geiger committed
Apr 6, 2026 at 16:39 UTC
f40d9879f2dfe4d878b77768ad30935ea4e42b1f
1 file changed
+13
-9
vllm/v1/attention/backends/gdn_attn.py
+13
-9
@@ -253,7 +253,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
253
)
254
# Filter by spec_sequence_masks to exclude padded sequences
255
spec_state_indices_tensor = block_table_tensor[
256
- spec_sequence_masks, : self.num_spec + 1
256
+ spec_sequence_masks_cpu, : self.num_spec + 1
257
]
258
non_spec_state_indices_tensor = None
259
# Padded sequences are always at the back, so the first
@@ -264,7 +264,9 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
264
non_spec_query_start_loc_cpu = None
265
else:
266
spec_token_masks = torch.repeat_interleave(
267
- spec_sequence_masks, query_lens
267
+ spec_sequence_masks,
268
+ query_lens,
269
+ output_size=query_start_loc_cpu[-1].item(),
270
)
271
index = torch.argsort(spec_token_masks, stable=True)
272
num_non_spec_tokens = num_prefill_tokens + num_decode_tokens
@@ -272,10 +274,10 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
274
spec_token_indx = index[num_non_spec_tokens:]
275
276
spec_state_indices_tensor = block_table_tensor[
275
- spec_sequence_masks, : self.num_spec + 1
277
+ spec_sequence_masks_cpu, : self.num_spec + 1
278
]
279
non_spec_state_indices_tensor = block_table_tensor[
278
- ~spec_sequence_masks, 0
280
+ ~spec_sequence_masks_cpu, 0
281
]
282
283
spec_query_start_loc = torch.zeros(
@@ -284,7 +286,9 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
286
device=query_start_loc.device,
287
)
288
torch.cumsum(
287
- query_lens[spec_sequence_masks], dim=0, out=spec_query_start_loc[1:]
289
+ query_lens[spec_sequence_masks_cpu],
290
+ dim=0,
291
+ out=spec_query_start_loc[1:],
292
)
293
non_spec_query_start_loc = torch.zeros(
294
query_lens.size(0) - num_spec_decodes + 1,
@@ -292,7 +296,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
296
device=query_start_loc.device,
297
)
298
torch.cumsum(
295
- query_lens[~spec_sequence_masks],
299
+ query_lens[~spec_sequence_masks_cpu],
300
dim=0,
301
out=non_spec_query_start_loc[1:],
302
)
@@ -307,7 +311,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
311
)
312
313
assert num_accepted_tokens is not None
310
- num_accepted_tokens = num_accepted_tokens[spec_sequence_masks]
314
+ num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu]
315
316
chunk_indices: torch.Tensor | None = None
317
chunk_offsets: torch.Tensor | None = None
@@ -331,8 +335,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
335
336
if num_prefills > 0:
337
has_initial_state = context_lens_tensor > 0
334
- if spec_sequence_masks is not None:
335
- has_initial_state = has_initial_state[~spec_sequence_masks]
338
+ if spec_sequence_masks_cpu is not None:
339
+ has_initial_state = has_initial_state[~spec_sequence_masks_cpu]
340
assert non_spec_query_start_loc_cpu is not None
341
nums_dict, batch_ptr, token_chunk_offset_ptr = (
342
compute_causal_conv1d_metadata(