[Bugfix][MRV2][Spec Decode] Reject draft slots that were never proposed (#58784)
Signed-off-by: zixi-qi <zixi@inferact.ai> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
qizixi committed
Sep 28, 2026 at 13:07 UTC
fedbc3b564102efb38f3697acb69920c3ac912be
2 files changed
+158
-2
tests/v1/spec_decode/test_rejection_sampler_utils.py
+80
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import math
5
+from types import SimpleNamespace
6
7
import pytest
8
import torch
@@ -10,6 +11,10 @@ from vllm.v1.watermarking import GumbelWatermarker
11
from vllm.v1.watermarking.spec_decode import watermarked_rejection_sample
12
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
13
from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
14
+from vllm.v1.worker.gpu.spec_decode.rejection_sampler import (
15
+ RejectionSampler,
16
+ gather_draft_sampled,
17
+)
18
from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
19
rejection_sample,
20
)
@@ -669,6 +674,81 @@ def test_placeholder_blocks_later_draft_tokens(use_block_verification: bool):
674
)
675
676
677
+def test_verify_rejects_unproposed_drafts():
678
+ """Zero drafts padded onto a step starting in the prefill must be rejected."""
679
+ torch.manual_seed(0)
680
+ device = "cuda"
681
+ num_trials = 256
682
+ K = 8
683
+
684
+ # Row 0: token 0 is very unlikely. Rows >= 1 follow the zero inputs and
685
+ # predict token 0. The never-written draft logits row is uniform.
686
+ target_logits_1d = torch.randn(VOCAB_SIZE, device=device)
687
+ target_logits_1d[0] = -20.0
688
+ inputs = _build_rejection_sample_inputs(
689
+ target_logits_1d,
690
+ torch.zeros(VOCAB_SIZE, device=device),
691
+ K,
692
+ temperature=1.0,
693
+ num_trials=num_trials,
694
+ )
695
+ target_rest = torch.randn(VOCAB_SIZE, device=device)
696
+ target_rest[0] = 20.0
697
+ inputs["target_logits"].view(num_trials, K + 1, VOCAB_SIZE)[:, 1:] = target_rest
698
+ inputs["draft_sampled"].view(num_trials, K + 1)[:, 1:] = 0
699
+
700
+ # Each request's first logit is the last prefill token, as in the first
701
+ # decode step after a P/D remote KV load.
702
+ first_pos = inputs["pos"].view(num_trials, K + 1)[:, 0]
703
+ prefill_len = first_pos + 1
704
+ rejection_sampler = object.__new__(RejectionSampler)
705
+ rejection_sampler.sampler = SimpleNamespace(
706
+ apply_sampling_params=lambda logits, *args: logits,
707
+ sampling_states=SimpleNamespace(
708
+ temperature=SimpleNamespace(gpu=inputs["temperature"]),
709
+ seeds=SimpleNamespace(gpu=inputs["seed"]),
710
+ ),
711
+ use_fp64_gumbel=False,
712
+ )
713
+ rejection_sampler.num_speculative_steps = K
714
+ rejection_sampler.synthetic_conditional_rates = None
715
+ rejection_sampler.use_block_verification = True
716
+ rejection_sampler.watermark_key = None
717
+
718
+ def verify() -> tuple[torch.Tensor, torch.Tensor]:
719
+ draft_sampled, pos = gather_draft_sampled(
720
+ inputs["draft_sampled"],
721
+ inputs["pos"],
722
+ torch.arange(num_trials * (K + 1), device=device),
723
+ inputs["expanded_idx_mapping"],
724
+ inputs["expanded_local_pos"],
725
+ prefill_len,
726
+ )
727
+ _, sampled, num_sampled = rejection_sampler._verify(
728
+ inputs["target_logits"],
729
+ inputs["draft_logits"],
730
+ draft_sampled,
731
+ pos,
732
+ inputs["cu_num_logits"],
733
+ inputs["idx_mapping"],
734
+ None,
735
+ inputs["expanded_idx_mapping"],
736
+ inputs["expanded_local_pos"],
737
+ None,
738
+ )
739
+ return sampled, num_sampled
740
+
741
+ sampled, num_sampled = verify()
742
+ assert (num_sampled == 1).all()
743
+ # Resampled from the target at row 0, not from a residual.
744
+ assert (sampled[:, 0] != 0).all()
745
+
746
+ # Once past the prefill, block verification accepts the same zero drafts.
747
+ prefill_len.copy_(first_pos)
748
+ _, num_sampled = verify()
749
+ assert (num_sampled == K + 1).all()
750
+
751
+
752
@pytest.mark.parametrize(
753
"num_speculative_steps,temperature",
754
[
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+78
-2
@@ -74,6 +74,76 @@ def _flatten_sampled_kernel(
74
tl.store(flat_sampled_ptr + start_idx + i, token_id)
75
76
77
+@triton.jit
78
+def _gather_draft_sampled_kernel(
79
+ # [num_logits]
80
+ draft_sampled_ptr,
81
+ # [num_logits]
82
+ pos_ptr,
83
+ # [num_tokens]
84
+ input_ids_ptr,
85
+ # [num_tokens]
86
+ positions_ptr,
87
+ # [num_logits]
88
+ logits_indices_ptr,
89
+ # [num_logits]
90
+ expanded_idx_mapping_ptr,
91
+ # [num_logits]
92
+ expanded_local_pos_ptr,
93
+ # [max_num_reqs]
94
+ prefill_len_ptr,
95
+ num_logits,
96
+ BLOCK_SIZE: tl.constexpr,
97
+):
98
+ block = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
99
+ mask = block < num_logits
100
+ token_idx = tl.load(logits_indices_ptr + block, mask=mask)
101
+ token_id = tl.load(input_ids_ptr + token_idx, mask=mask)
102
+ pos = tl.load(positions_ptr + token_idx, mask=mask)
103
+ req_state_idx = tl.load(expanded_idx_mapping_ptr + block, mask=mask)
104
+ local_pos = tl.load(expanded_local_pos_ptr + block, mask=mask)
105
+ prefill_len = tl.load(prefill_len_ptr + req_state_idx, mask=mask)
106
+ # Drafts are only proposed after sampling past the prefill, so any draft
107
+ # slots in a step starting within the prefill are placeholders (e.g. the
108
+ # padded first step after a P/D remote KV load).
109
+ is_placeholder = (local_pos > 0) & (pos - local_pos < prefill_len)
110
+ token_id = tl.where(is_placeholder, -1, token_id)
111
+ tl.store(draft_sampled_ptr + block, token_id, mask=mask)
112
+ tl.store(pos_ptr + block, pos, mask=mask)
113
+
114
+
115
+def gather_draft_sampled(
116
+ input_ids: torch.Tensor,
117
+ positions: torch.Tensor,
118
+ logits_indices: torch.Tensor,
119
+ expanded_idx_mapping: torch.Tensor,
120
+ expanded_local_pos: torch.Tensor,
121
+ prefill_len: torch.Tensor,
122
+) -> tuple[torch.Tensor, torch.Tensor]:
123
+ """Gather the input token and position of each logits row.
124
+
125
+ Draft rows of requests that have not yet sampled past their prefill are
126
+ set to -1 so that the rejection kernels reject them.
127
+ """
128
+ num_logits = logits_indices.shape[0]
129
+ draft_sampled = input_ids.new_empty(num_logits)
130
+ pos = positions.new_empty(num_logits)
131
+ BLOCK_SIZE = 1024
132
+ _gather_draft_sampled_kernel[(triton.cdiv(num_logits, BLOCK_SIZE),)](
133
+ draft_sampled,
134
+ pos,
135
+ input_ids,
136
+ positions,
137
+ logits_indices,
138
+ expanded_idx_mapping,
139
+ expanded_local_pos,
140
+ prefill_len,
141
+ num_logits,
142
+ BLOCK_SIZE=BLOCK_SIZE,
143
+ )
144
+ return draft_sampled, pos
145
+
146
+
147
class RejectionSampler:
148
def __init__(
149
self,
@@ -296,8 +366,14 @@ class RejectionSampler:
366
# that num_nans is computed before applying penalties and temperature.
367
num_nans = get_num_nans(logits) if self.sampler.compute_nans else None
368
299
- draft_sampled = input_batch.input_ids[input_batch.logits_indices]
300
- pos = input_batch.positions[input_batch.logits_indices]
369
+ draft_sampled, pos = gather_draft_sampled(
370
+ input_batch.input_ids,
371
+ input_batch.positions,
372
+ input_batch.logits_indices,
373
+ input_batch.expanded_idx_mapping,
374
+ input_batch.expanded_local_pos,
375
+ self.sampler.req_states.prefill_len.gpu,
376
+ )
377
378
max_num_logprobs = self.sampler.sampling_states.max_num_logprobs(
379
input_batch.idx_mapping_np