diff --git a/tests/v1/sample/test_rejection_sampler.py b/tests/v1/sample/test_rejection_sampler.py index 10c4d448f7..b02e53af8d 100644 --- a/tests/v1/sample/test_rejection_sampler.py +++ b/tests/v1/sample/test_rejection_sampler.py @@ -1150,3 +1150,36 @@ def test_synthetic_all_rejected(all_greedy: bool): for row in result: assert row[0] != PLACEHOLDER_TOKEN_ID assert (row[1:] == PLACEHOLDER_TOKEN_ID).all() + + +def test_placeholder_draft_token_rejected_random(rejection_sampler): + """A placeholder draft id (-1) must be rejected in non-greedy sampling + without indexing the probability tensors by the invalid id. + """ + vocab_size = 100 + spec_tokens = [[1, vocab_size - 1, PLACEHOLDER_TOKEN_ID]] + output_tokens = [[1, vocab_size - 1, 7, 9]] + + temperature = torch.ones(1, dtype=torch.float32, device=DEVICE_TYPE) + metadata = create_sampling_metadata( + all_greedy=False, + temperature=temperature, + generators={0: torch.Generator(device=DEVICE_TYPE).manual_seed(0)}, + ) + logits = create_logits_tensor(output_tokens, vocab_size=vocab_size) + bonus_token_tensor = torch.tensor([output_tokens[0][-1]], device=logits.device) + spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits) + + mock_sampler_output(rejection_sampler, bonus_token_tensor) + output = rejection_sampler( + spec_decode_metadata, + draft_probs=None, + logits=logits, + sampling_metadata=metadata, + ) + sampled = output.sampled_token_ids + + assert sampled[0, 0].item() == 1 + assert sampled[0, 1].item() == vocab_size - 1 + recovered = sampled[0, 2].item() + assert 0 <= recovered < vocab_size diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index 28f6044de8..613bd846e6 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -280,3 +280,32 @@ def test_synthetic_rejection_sample( f"Step {i}: observed rate {observed_rate:.4f} deviates from " f"expected rate {expected_rate:.4f} by more than {deviation_tol}." ) + + +def test_placeholder_draft_token_rejected(): + """A placeholder draft id (-1) must be rejected without reading the logit + tensors out of bounds, for any sampling method. + """ + torch.manual_seed(0) + device = "cuda" + num_trials = 64 + K = 1 + temperature = 0.6 + + target_logits_1d = torch.randn(VOCAB_SIZE, device=device) / temperature + draft_logits_1d = torch.randn(VOCAB_SIZE, device=device) / temperature + + inputs = _build_rejection_sample_inputs( + target_logits_1d, + draft_logits_1d, + K, + temperature=temperature, + num_trials=num_trials, + ) + inputs["draft_sampled"].view(num_trials, K + 1)[:, 1:] = -1 + + sampled, num_sampled = rejection_sample(**inputs, num_speculative_steps=K) + + assert torch.equal(num_sampled, torch.ones_like(num_sampled)) + recovered = sampled[:, 0] + assert (recovered >= 0).all() and (recovered < VOCAB_SIZE).all() diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py index 8b4d8c9dce..1324191be7 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/vllm/v1/sample/rejection_sampler.py @@ -748,7 +748,8 @@ def rejection_greedy_sample_kernel( if SYNTHETIC_MODE: uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos) rate = tl.load(synthetic_conditional_rates_ptr + pos) - accepted = uniform_prob < rate + # -1 is used for padded draft token ids that should be rejected. + accepted = (uniform_prob < rate) and draft_token_id >= 0 token_id = draft_token_id if accepted else target_argmax_id rejected = not accepted else: @@ -805,7 +806,10 @@ def rejection_random_sample_kernel( if not rejected: draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos) uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos) - if SYNTHETIC_MODE: + if draft_token_id < 0: + # -1 is used for padded draft token ids that should be rejected. + accepted = False + elif SYNTHETIC_MODE: rate = tl.load(synthetic_conditional_rates_ptr + pos) accepted = uniform_prob < rate else: diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 0cfbdf4182..92294e6c7e 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -245,7 +245,8 @@ def _rejection_kernel( pos = tl.load(pos_ptr + logit_idx) u = tl_rand64(seed, pos, includes_zero=False) rate = tl.load(synthetic_conditional_rates_ptr + i) - accepted &= u < rate + # -1 is used for padded draft token ids that should be rejected. + accepted &= (u < rate) & (draft_sampled >= 0) else: accepted &= target_argmax == draft_sampled tl.store( @@ -253,6 +254,10 @@ def _rejection_kernel( draft_sampled if accepted else target_argmax, ) else: + # -1 is used for padded draft token ids that should be rejected. + is_valid_draft = draft_sampled >= 0 + # Avoid possible OOB ptr access. + draft_sampled = tl.maximum(0, draft_sampled) target_logit = tl.load( target_logits_ptr + logit_idx * target_logits_stride + draft_sampled ).to(tl.float32) @@ -296,6 +301,7 @@ def _rejection_kernel( # Probability ratio test: p(x) > u * q(x) # Equivalent log form: log_p(x) > log(u) + log_q(x) accepted &= target_log_prob > tl.log(u) + draft_log_prob + accepted &= is_valid_draft tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled) rejected_step += accepted tl.store(rejected_steps_ptr + req_idx, rejected_step)