[Perf][GDN] Slice pure spec-decode rows instead of a host-mask gather (#58763)
Signed-off-by: Luca Motz <luca.motz@icloud.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>
Luca Motz committed
Oct 1, 2026 at 08:53 UTC
87a4bf664f4b2344eced5a10add977b750d4b7d5
2 files changed
+27
-7
tests/v1/attention/test_gdn_metadata_builder.py
+16
-2
@@ -65,6 +65,17 @@ GDN_BUILD_TEST_CASES = {
65
expected_num_prefill_tokens=0,
66
expected_num_spec_decodes=2,
67
),
68
+ # Padded (CUDA graph) sequences trail the spec decodes
69
+ "pure_spec_decode_with_padding": GDNBuildTestCase(
70
+ seq_lens=[50, 30, 16],
71
+ query_lens=[3, 3, 0],
72
+ num_decode_draft_tokens=[2, 2, -1],
73
+ num_speculative_tokens=2,
74
+ expected_num_decodes=0,
75
+ expected_num_prefills=0,
76
+ expected_num_prefill_tokens=0,
77
+ expected_num_spec_decodes=2,
78
+ ),
79
# No speculative config at all — standard decode path
80
"pure_regular_decode": GDNBuildTestCase(
81
seq_lens=[40, 30, 20],
@@ -143,8 +154,9 @@ def _create_gdn_builder(
154
model_name="Qwen/Qwen3.5-0.8B",
155
block_size=BLOCK_SIZE,
156
)
146
- if full_cuda_graph:
147
- vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE
157
+ vllm_config.compilation_config.cudagraph_mode = (
158
+ CUDAGraphMode.FULL_AND_PIECEWISE if full_cuda_graph else CUDAGraphMode.NONE
159
+ )
160
if num_speculative_tokens > 0:
161
vllm_config.speculative_config = SpeculativeConfig(
162
method="ngram",
@@ -197,6 +209,8 @@ def test_gdn_build_classification(test_case: GDNBuildTestCase):
209
assert meta.num_prefills == test_case.expected_num_prefills
210
assert meta.num_prefill_tokens == test_case.expected_num_prefill_tokens
211
assert meta.num_spec_decodes == test_case.expected_num_spec_decodes
212
+ if meta.spec_state_indices_tensor is not None:
213
+ assert len(meta.spec_state_indices_tensor) == meta.num_spec_decodes
214
215
216
@pytest.mark.parametrize("mamba_cache_mode", ["none", "align"])
vllm/v1/attention/backends/gdn_attn.py
+11
-5
@@ -361,9 +361,10 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
361
non_spec_token_indx = torch.empty(
362
0, dtype=torch.int32, device=query_start_loc.device
363
)
364
- # Filter by spec_sequence_masks to exclude padded sequences
364
+ # Padded sequences trail the spec decodes, so slice them off
365
+ # rather than gather with the host mask (an H2D copy + a kernel).
366
spec_state_indices_tensor = block_table_tensor[
366
- spec_sequence_masks_cpu, : self.num_spec + 1
367
+ :num_spec_decodes, : self.num_spec + 1
368
]
369
non_spec_state_indices_tensor = None
370
# Padded sequences are always at the back, so the first
@@ -421,7 +422,10 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
422
)
423
424
assert num_accepted_tokens is not None
424
- num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu]
425
+ if num_prefills == 0 and num_decodes == 0:
426
+ num_accepted_tokens = num_accepted_tokens[:num_spec_decodes]
427
+ else:
428
+ num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu]
429
430
chunk_indices: torch.Tensor | None = None
431
chunk_offsets: torch.Tensor | None = None
@@ -626,10 +630,12 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
630
non_spec_indices = blk_table[:, 0]
631
if m.num_prefills > 0:
632
prefill_indices = non_spec_indices[m.num_decodes :]
633
+ elif m.num_prefills == 0:
634
+ # Same as build(): padded sequences trail the spec decodes.
635
+ spec_indices = blk_table[: m.num_spec_decodes, : self.num_spec + 1]
636
else:
637
spec_indices = blk_table[masks, : self.num_spec + 1]
631
- if m.num_prefills > 0:
632
- non_spec_indices = prefill_indices = blk_table[~masks, 0]
638
+ non_spec_indices = prefill_indices = blk_table[~masks, 0]
639
640
if self._stage_spec_decode(
641
m.num_prefills, m.num_decodes, m.num_spec_decodes, m.num_spec_decode_tokens