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 11, 01:15 UTC
[SpecDecode] Simplified alternative padded-speculation acceptance rate fix (#29845)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Dec 22, 2025 at 16:06 UTC
de71747655668d268998264365dc9bbc8e58d410
8 files changed
+62
-25
tests/v1/spec_decode/test_eagle.py
+8
-2
@@ -306,10 +306,16 @@ def test_prepare_inputs_padded():
306
307
proposer = _create_proposer("eagle", num_speculative_tokens)
308
309
- output_metadata, token_indices_to_sample = proposer.prepare_inputs_padded(
310
- common_attn_metadata, spec_decode_metadata, valid_sampled_tokens_count
309
+ output_metadata, token_indices_to_sample, num_rejected_tokens_gpu = (
310
+ proposer.prepare_inputs_padded(
311
+ common_attn_metadata, spec_decode_metadata, valid_sampled_tokens_count
312
+ )
313
)
314
315
+ # Verify num_rejected_tokens_gpu is calculated correctly
316
+ expected_num_rejected = torch.tensor([1, 0, 2], dtype=torch.int32, device=device)
317
+ assert torch.equal(num_rejected_tokens_gpu, expected_num_rejected)
318
+
319
assert output_metadata.max_query_len == 3
320
assert torch.equal(output_metadata.query_start_loc, expected_query_start_loc)
321
assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample)
vllm/v1/attention/backends/mla/common.py
+15
-9
@@ -564,6 +564,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
564
self.dcp_rank = 0
565
self.dcp_local_block_size = parallel_config.cp_kv_cache_interleave_size
566
self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size
567
+ self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
568
569
# Don't try to access the runner on AMD
570
if self.aot_schedule:
@@ -727,8 +728,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
728
def _build_decode(
729
self,
730
block_table_tensor: torch.Tensor,
730
- seq_lens_cpu: torch.Tensor,
731
seq_lens_device: torch.Tensor,
732
+ max_seq_len: int,
733
query_start_loc_cpu: torch.Tensor,
734
query_start_loc_device: torch.Tensor,
735
num_decode_tokens: int,
@@ -778,13 +779,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
779
query_start_loc = common_attn_metadata.query_start_loc
780
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
781
seq_lens = common_attn_metadata.seq_lens
781
- seq_lens_cpu = common_attn_metadata.seq_lens_cpu
782
dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens
783
- dcp_local_seq_lens_cpu = common_attn_metadata.dcp_local_seq_lens_cpu
784
-
785
- query_seq_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
786
-
787
- num_computed_tokens_cpu = common_attn_metadata.seq_lens_cpu - query_seq_lens_cpu
783
784
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
785
split_decodes_and_prefills(
@@ -799,6 +794,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
794
795
prefill_metadata = None
796
if num_prefills > 0:
797
+ num_computed_tokens_cpu = common_attn_metadata.num_computed_tokens_cpu
798
+
799
reqs_start = num_decodes # prefill_start
800
801
context_lens_cpu = num_computed_tokens_cpu[reqs_start:num_reqs]
@@ -995,13 +992,22 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
992
dcp_tot_seq_lens_device = None
993
if self.dcp_world_size > 1:
994
dcp_tot_seq_lens_device = seq_lens[:num_decodes]
998
- seq_lens_cpu = dcp_local_seq_lens_cpu
995
seq_lens = dcp_local_seq_lens
996
997
+ # After DCP distribution, the maximum number of tokens for any rank is
998
+ # ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size,
999
+ # and I is cp_kv_cache_interleave_size.
1000
+ # This eliminates GPU->CPU sync while minimizing workspace
1001
+ # over-allocation.
1002
+ num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size
1003
+ max_seq_len = (
1004
+ (max_seq_len + num_partitions - 1) // num_partitions
1005
+ ) * self.cp_kv_cache_interleave_size
1006
+
1007
decode_metadata = self._build_decode(
1008
block_table_tensor=block_table_tensor[:num_decodes, ...],
1003
- seq_lens_cpu=seq_lens_cpu[:num_decodes],
1009
seq_lens_device=seq_lens[:num_decodes],
1010
+ max_seq_len=max_seq_len,
1011
query_start_loc_cpu=query_start_loc_cpu[: num_decodes + 1],
1012
query_start_loc_device=query_start_loc[: num_decodes + 1],
1013
num_decode_tokens=num_decode_tokens,
vllm/v1/attention/backends/mla/flashattn_mla.py
+2
-3
@@ -169,8 +169,8 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
169
def _build_decode(
170
self,
171
block_table_tensor: torch.Tensor,
172
- seq_lens_cpu: torch.Tensor,
172
seq_lens_device: torch.Tensor,
173
+ max_seq_len: int,
174
query_start_loc_cpu: torch.Tensor,
175
query_start_loc_device: torch.Tensor,
176
num_decode_tokens: int,
@@ -178,7 +178,6 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
178
) -> FlashAttnMLADecodeMetadata:
179
query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
180
max_query_len = query_lens_cpu.max().item()
181
- max_seq_len = seq_lens_cpu.max().item()
181
182
# For Flash Attention MLA + full cudagraph
183
max_num_splits = 0
@@ -193,7 +192,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
192
max_num_splits = 1
193
194
scheduler_metadata = self._schedule_decode(
196
- num_reqs=seq_lens_cpu.numel(),
195
+ num_reqs=seq_lens_device.shape[0],
196
cu_query_lens=query_start_loc_device,
197
max_query_len=max_query_len,
198
seqlens=seq_lens_device,
vllm/v1/attention/backends/mla/flashmla.py
+1
-1
@@ -143,8 +143,8 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
143
def _build_decode(
144
self,
145
block_table_tensor: torch.Tensor,
146
- seq_lens_cpu: torch.Tensor,
146
seq_lens_device: torch.Tensor,
147
+ max_seq_len: int,
148
query_start_loc_cpu: torch.Tensor,
149
query_start_loc_device: torch.Tensor,
150
num_decode_tokens: int,
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+1
-1
@@ -106,8 +106,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
106
def _build_decode(
107
self,
108
block_table_tensor: torch.Tensor,
109
- seq_lens_cpu: torch.Tensor,
109
seq_lens_device: torch.Tensor,
110
+ max_seq_len: int,
111
query_start_loc_cpu: torch.Tensor,
112
query_start_loc_device: torch.Tensor,
113
num_decode_tokens: int,
vllm/v1/spec_decode/eagle.py
+23
-3
@@ -236,6 +236,7 @@ class EagleProposer:
236
common_attn_metadata: CommonAttentionMetadata,
237
sampling_metadata: SamplingMetadata,
238
mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
239
+ num_rejected_tokens_gpu: torch.Tensor | None = None,
240
) -> torch.Tensor:
241
num_tokens = target_token_ids.shape[0]
242
batch_size = next_token_ids.shape[0]
@@ -414,6 +415,17 @@ class EagleProposer:
415
common_attn_metadata.query_start_loc_cpu = torch.from_numpy(
416
self.token_arange_np[: batch_size + 1]
417
).clone()
418
+
419
+ # In padded drafter batch, we need to adjust the sequence lengths
420
+ # to remove the "padding" (i.e. rejected tokens).
421
+ # Only apply this adjustment when we have rejected tokens
422
+ # (i.e., not the first proposal).
423
+ if self.num_speculative_tokens > 1 and num_rejected_tokens_gpu is not None:
424
+ common_attn_metadata.seq_lens -= num_rejected_tokens_gpu
425
+ # Invalidate the CPU-side shadows to avoid H<>D sync.
426
+ common_attn_metadata._seq_lens_cpu = None
427
+ common_attn_metadata._num_computed_tokens_cpu = None
428
+
429
for token_index in range(self.num_speculative_tokens - 1):
430
# Update the inputs.
431
# cast to int32 is crucial when eagle model is compiled.
@@ -628,13 +640,14 @@ class EagleProposer:
640
common_attn_metadata: CommonAttentionMetadata,
641
spec_decode_metadata: SpecDecodeMetadata,
642
valid_sampled_tokens_count: torch.Tensor,
631
- ) -> tuple[CommonAttentionMetadata, torch.Tensor]:
643
+ ) -> tuple[CommonAttentionMetadata, torch.Tensor, torch.Tensor]:
644
"""
645
This function is used to prepare the inputs for speculative decoding
646
It updates the common_attn_metadata for speculative decoding,
647
but does not consider the rejected tokens. Instead, all tokens
648
are included as inputs to the speculator, with the rejected tokens
649
used as padding and filtered out later by `token_indices_to_sample`.
650
+ No blocking CPU operations should be introduced in this function.
651
"""
652
num_reqs = common_attn_metadata.num_reqs
653
device = valid_sampled_tokens_count.device
@@ -642,14 +655,17 @@ class EagleProposer:
655
token_indices_to_sample = torch.empty(
656
(num_reqs,), dtype=torch.int32, device=device
657
)
658
+ num_rejected_tokens_gpu = torch.empty(
659
+ (num_reqs,), dtype=torch.int32, device=device
660
+ )
661
646
- # Kernel grid: one program per request (row)
662
grid = (num_reqs,)
663
eagle_prepare_inputs_padded_kernel[grid](
664
spec_decode_metadata.cu_num_draft_tokens,
665
valid_sampled_tokens_count,
666
common_attn_metadata.query_start_loc,
667
token_indices_to_sample,
668
+ num_rejected_tokens_gpu,
669
num_reqs,
670
)
671
@@ -674,7 +690,11 @@ class EagleProposer:
690
dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,
691
)
692
677
- return spec_common_attn_metadata, token_indices_to_sample
693
+ return (
694
+ spec_common_attn_metadata,
695
+ token_indices_to_sample,
696
+ num_rejected_tokens_gpu,
697
+ )
698
699
def propose_tree(
700
self,
vllm/v1/spec_decode/utils.py
+2
@@ -23,6 +23,7 @@ def eagle_prepare_inputs_padded_kernel(
23
valid_sampled_tokens_count_ptr, # [num_reqs]
24
query_start_loc_gpu_ptr, # [num_reqs + 1]
25
token_indices_to_sample_ptr, # [num_reqs] (output)
26
+ num_rejected_tokens_gpu_ptr, # [num_reqs] (output)
27
num_reqs, # tl.int32
28
):
29
"""
@@ -56,6 +57,7 @@ def eagle_prepare_inputs_padded_kernel(
57
58
index_to_sample = q_last_tok_idx - num_rejected_tokens
59
tl.store(token_indices_to_sample_ptr + req_idx, index_to_sample)
60
+ tl.store(num_rejected_tokens_gpu_ptr + req_idx, num_rejected_tokens)
61
62
63
@triton.jit
vllm/v1/worker/gpu_model_runner.py
+10
-6
@@ -3534,6 +3534,7 @@ class GPUModelRunner(
3534
next_token_ids, valid_sampled_tokens_count
3535
)
3536
3537
+ num_rejected_tokens_gpu = None
3538
if spec_decode_metadata is None:
3539
token_indices_to_sample = None
3540
# input_ids can be None for multimodal models.
@@ -3564,12 +3565,14 @@ class GPUModelRunner(
3565
else:
3566
target_hidden_states = hidden_states[token_indices]
3567
else:
3567
- common_attn_metadata, token_indices_to_sample = (
3568
- self.drafter.prepare_inputs_padded(
3569
- common_attn_metadata,
3570
- spec_decode_metadata,
3571
- valid_sampled_tokens_count,
3572
- )
3568
+ (
3569
+ common_attn_metadata,
3570
+ token_indices_to_sample,
3571
+ num_rejected_tokens_gpu,
3572
+ ) = self.drafter.prepare_inputs_padded(
3573
+ common_attn_metadata,
3574
+ spec_decode_metadata,
3575
+ valid_sampled_tokens_count,
3576
)
3577
total_num_tokens = common_attn_metadata.num_actual_tokens
3578
# When padding the batch, token_indices is just a range
@@ -3600,6 +3603,7 @@ class GPUModelRunner(
3603
sampling_metadata=sampling_metadata,
3604
common_attn_metadata=common_attn_metadata,
3605
mm_embed_inputs=mm_embed_inputs,
3606
+ num_rejected_tokens_gpu=num_rejected_tokens_gpu,
3607
)
3608
3609
return draft_token_ids