[Spec Decode] Reject placeholder (-1) draft tokens in rejection sampler (#46533)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jun 23, 2026 at 20:32 UTC
7ee4d220097db4b397e55fd4ad58caf6a7977c5b
4 files changed
+75
-3
tests/v1/sample/test_rejection_sampler.py
+33
@@ -1150,3 +1150,36 @@ def test_synthetic_all_rejected(all_greedy: bool):
1150
for row in result:
1151
assert row[0] != PLACEHOLDER_TOKEN_ID
1152
assert (row[1:] == PLACEHOLDER_TOKEN_ID).all()
1153
+
1154
+
1155
+def test_placeholder_draft_token_rejected_random(rejection_sampler):
1156
+ """A placeholder draft id (-1) must be rejected in non-greedy sampling
1157
+ without indexing the probability tensors by the invalid id.
1158
+ """
1159
+ vocab_size = 100
1160
+ spec_tokens = [[1, vocab_size - 1, PLACEHOLDER_TOKEN_ID]]
1161
+ output_tokens = [[1, vocab_size - 1, 7, 9]]
1162
+
1163
+ temperature = torch.ones(1, dtype=torch.float32, device=DEVICE_TYPE)
1164
+ metadata = create_sampling_metadata(
1165
+ all_greedy=False,
1166
+ temperature=temperature,
1167
+ generators={0: torch.Generator(device=DEVICE_TYPE).manual_seed(0)},
1168
+ )
1169
+ logits = create_logits_tensor(output_tokens, vocab_size=vocab_size)
1170
+ bonus_token_tensor = torch.tensor([output_tokens[0][-1]], device=logits.device)
1171
+ spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)
1172
+
1173
+ mock_sampler_output(rejection_sampler, bonus_token_tensor)
1174
+ output = rejection_sampler(
1175
+ spec_decode_metadata,
1176
+ draft_probs=None,
1177
+ logits=logits,
1178
+ sampling_metadata=metadata,
1179
+ )
1180
+ sampled = output.sampled_token_ids
1181
+
1182
+ assert sampled[0, 0].item() == 1
1183
+ assert sampled[0, 1].item() == vocab_size - 1
1184
+ recovered = sampled[0, 2].item()
1185
+ assert 0 <= recovered < vocab_size
tests/v1/spec_decode/test_rejection_sampler_utils.py
+29
@@ -280,3 +280,32 @@ def test_synthetic_rejection_sample(
280
f"Step {i}: observed rate {observed_rate:.4f} deviates from "
281
f"expected rate {expected_rate:.4f} by more than {deviation_tol}."
282
)
283
+
284
+
285
+def test_placeholder_draft_token_rejected():
286
+ """A placeholder draft id (-1) must be rejected without reading the logit
287
+ tensors out of bounds, for any sampling method.
288
+ """
289
+ torch.manual_seed(0)
290
+ device = "cuda"
291
+ num_trials = 64
292
+ K = 1
293
+ temperature = 0.6
294
+
295
+ target_logits_1d = torch.randn(VOCAB_SIZE, device=device) / temperature
296
+ draft_logits_1d = torch.randn(VOCAB_SIZE, device=device) / temperature
297
+
298
+ inputs = _build_rejection_sample_inputs(
299
+ target_logits_1d,
300
+ draft_logits_1d,
301
+ K,
302
+ temperature=temperature,
303
+ num_trials=num_trials,
304
+ )
305
+ inputs["draft_sampled"].view(num_trials, K + 1)[:, 1:] = -1
306
+
307
+ sampled, num_sampled = rejection_sample(**inputs, num_speculative_steps=K)
308
+
309
+ assert torch.equal(num_sampled, torch.ones_like(num_sampled))
310
+ recovered = sampled[:, 0]
311
+ assert (recovered >= 0).all() and (recovered < VOCAB_SIZE).all()
vllm/v1/sample/rejection_sampler.py
+6
-2
@@ -748,7 +748,8 @@ def rejection_greedy_sample_kernel(
748
if SYNTHETIC_MODE:
749
uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
750
rate = tl.load(synthetic_conditional_rates_ptr + pos)
751
- accepted = uniform_prob < rate
751
+ # -1 is used for padded draft token ids that should be rejected.
752
+ accepted = (uniform_prob < rate) and draft_token_id >= 0
753
token_id = draft_token_id if accepted else target_argmax_id
754
rejected = not accepted
755
else:
@@ -805,7 +806,10 @@ def rejection_random_sample_kernel(
806
if not rejected:
807
draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
808
uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
808
- if SYNTHETIC_MODE:
809
+ if draft_token_id < 0:
810
+ # -1 is used for padded draft token ids that should be rejected.
811
+ accepted = False
812
+ elif SYNTHETIC_MODE:
813
rate = tl.load(synthetic_conditional_rates_ptr + pos)
814
accepted = uniform_prob < rate
815
else:
vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+7
-1
@@ -245,7 +245,8 @@ def _rejection_kernel(
245
pos = tl.load(pos_ptr + logit_idx)
246
u = tl_rand64(seed, pos, includes_zero=False)
247
rate = tl.load(synthetic_conditional_rates_ptr + i)
248
- accepted &= u < rate
248
+ # -1 is used for padded draft token ids that should be rejected.
249
+ accepted &= (u < rate) & (draft_sampled >= 0)
250
else:
251
accepted &= target_argmax == draft_sampled
252
tl.store(
@@ -253,6 +254,10 @@ def _rejection_kernel(
254
draft_sampled if accepted else target_argmax,
255
)
256
else:
257
+ # -1 is used for padded draft token ids that should be rejected.
258
+ is_valid_draft = draft_sampled >= 0
259
+ # Avoid possible OOB ptr access.
260
+ draft_sampled = tl.maximum(0, draft_sampled)
261
target_logit = tl.load(
262
target_logits_ptr + logit_idx * target_logits_stride + draft_sampled
263
).to(tl.float32)
@@ -296,6 +301,7 @@ def _rejection_kernel(
301
# Probability ratio test: p(x) > u * q(x)
302
# Equivalent log form: log_p(x) > log(u) + log_q(x)
303
accepted &= target_log_prob > tl.log(u) + draft_log_prob
304
+ accepted &= is_valid_draft
305
tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
306
rejected_step += accepted
307
tl.store(rejected_steps_ptr + req_idx, rejected_step)