diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index 0768380582..4dd43503da 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -367,6 +367,128 @@ def test_placeholder_draft_token_rejected(): assert (recovered >= 0).all() and (recovered < VOCAB_SIZE).all() +@pytest.mark.parametrize("has_draft_logits", [True, False]) +@pytest.mark.parametrize("num_placeholders", [1, 2]) +def test_block_verification_placeholder_truncates_block( + has_draft_logits: bool, num_placeholders: int +): + """Trailing placeholder (-1) draft ids end the block early. The last real + draft token must then be verified against the closed-form threshold rather + than a residual mass derived from the placeholder position, whose logits + describe a token that was never drafted. + """ + torch.manual_seed(42) + device = "cuda" + num_trials = 20 * VOCAB_SIZE + K = 3 + temperature = 1.0 + + target_logits_1d = torch.randn(VOCAB_SIZE, device=device) / temperature + draft_logits_1d = torch.randn(VOCAB_SIZE, device=device) + + inputs = _build_rejection_sample_inputs( + target_logits_1d, + draft_logits_1d, + K, + temperature=temperature, + num_trials=num_trials, + ) + if not has_draft_logits: + inputs["draft_logits"] = None + inputs["draft_sampled"].view(num_trials, K + 1)[:, K + 1 - num_placeholders :] = -1 + + sampled, num_sampled = rejection_sample( + **inputs, num_speculative_steps=K, use_block_verification=True + ) + + num_real = K - num_placeholders + assert (num_sampled <= num_real + 1).all(), ( + "Accepted a draft token at or after a placeholder." + ) + target_probs = torch.softmax(target_logits_1d, dim=0) + for pos in range(num_real + 1): + accepted_mask = num_sampled >= pos + 1 + _assert_distribution_match( + sampled[accepted_mask, pos], target_probs, device, label=f"position {pos}" + ) + + +def test_greedy_placeholder_emits_target_argmax(): + """Greedy sampling skips resampling and relies on the rejection kernel + storing the target argmax at the rejected position. A placeholder must not + bypass that store, or the output slot is returned uninitialized. + """ + torch.manual_seed(0) + device = "cuda" + num_trials = 512 + K = 3 + + target_logits_1d = torch.randn(VOCAB_SIZE, device=device) + draft_logits_1d = torch.randn(VOCAB_SIZE, device=device) + + inputs = _build_rejection_sample_inputs( + target_logits_1d, + draft_logits_1d, + K, + temperature=0.0, + num_trials=num_trials, + ) + target_argmax = int(target_logits_1d.argmax()) + draft_sampled = inputs["draft_sampled"].view(num_trials, K + 1) + # Accept the first draft so that the rejection lands on the placeholder. + draft_sampled[:, 1] = target_argmax + draft_sampled[:, 2:] = -1 + + sampled, num_sampled = rejection_sample(**inputs, num_speculative_steps=K) + + assert torch.equal(num_sampled, torch.full_like(num_sampled, 2)) + steps = torch.arange(K + 1, device=device).unsqueeze(0) + emitted = sampled[steps < num_sampled.unsqueeze(1)] + assert (emitted == target_argmax).all(), ( + "Greedy sampling emitted a token that is not the target argmax" + ) + + +@pytest.mark.parametrize("use_block_verification", [False, True]) +def test_placeholder_blocks_later_draft_tokens(use_block_verification: bool): + """A placeholder is not necessarily the final draft. Nothing at or after + one may be accepted, even when a valid draft follows it. + """ + torch.manual_seed(0) + device = "cuda" + num_trials = 4 * VOCAB_SIZE + K = 3 + temperature = 1.0 + + # An identical draft is accepted with probability ~1, so any token after + # the placeholder would be accepted unless it is explicitly blocked. The + # draft logits are stored pre-temperature, so pass the unscaled base. + draft_logits_1d = torch.randn(VOCAB_SIZE, device=device) + target_logits_1d = draft_logits_1d / 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)[:, 2] = -1 + + _, num_sampled = rejection_sample( + **inputs, + num_speculative_steps=K, + use_block_verification=use_block_verification, + ) + + # Only the first draft is verifiable, plus one resampled token. + assert (num_sampled <= 2).all(), "Accepted a draft token past a placeholder." + assert (num_sampled == 2).float().mean().item() > 0.9, ( + "The first draft was rarely accepted; the test is not exercising " + "acceptance past the placeholder." + ) + + @pytest.mark.parametrize( "num_speculative_steps,temperature", [ 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 caab4c69a7..3ea9f35ee8 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -55,7 +55,7 @@ def _compute_global_residual_mass( target_local_max_stride, target_local_sumexp_ptr, target_local_sumexp_stride, - draft_sampled_ptr, + draft_token, logit_idx, vocab_num_blocks, PADDED_VOCAB_NUM_BLOCKS: tl.constexpr, @@ -74,7 +74,6 @@ def _compute_global_residual_mass( # One-hot draft. M_s is a point mass at this draft token # so the residual mass reduces to the closed form: # p * (1 - M_b(draft_token)). - draft_token = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64) target_lse = _compute_global_logsumexp( target_local_max_ptr, target_local_max_stride, @@ -347,35 +346,40 @@ def _compute_cumulative_log_p_kernel( if temp == 0.0: return - log_p = 0.0 + log_p = tl.zeros((), tl.float32) for step in range(num_draft_tokens): logit_idx = start_idx + step draft_token = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64) - target_logprob, draft_logprob, _, _ = _compute_global_logprobs_and_logsumexp( - draft_token, - True, # mask - logit_idx, - req_state_idx, - step, - temp, - target_logits_ptr, - target_logits_stride, - target_local_max_ptr, - target_local_max_stride, - target_local_sumexp_ptr, - target_local_sumexp_stride, - draft_logits_ptr, - draft_logits_stride_0, - draft_logits_stride_1, - draft_local_max_ptr, - draft_local_max_stride, - draft_local_sumexp_ptr, - draft_local_sumexp_stride, - vocab_num_blocks, - PADDED_VOCAB_NUM_BLOCKS, - HAS_DRAFT_LOGITS, - ) - log_p = tl.minimum(log_p + (target_logprob - draft_logprob), 0.0) + # -1 placeholder tokens can never be accepted. Skip their reductions + # and carry the last valid cumulative value. + if draft_token >= 0: + target_logprob, draft_logprob, _, _ = ( + _compute_global_logprobs_and_logsumexp( + draft_token, + True, # mask + logit_idx, + req_state_idx, + step, + temp, + target_logits_ptr, + target_logits_stride, + target_local_max_ptr, + target_local_max_stride, + target_local_sumexp_ptr, + target_local_sumexp_stride, + draft_logits_ptr, + draft_logits_stride_0, + draft_logits_stride_1, + draft_local_max_ptr, + draft_local_max_stride, + draft_local_sumexp_ptr, + draft_local_sumexp_stride, + vocab_num_blocks, + PADDED_VOCAB_NUM_BLOCKS, + HAS_DRAFT_LOGITS, + ) + ) + log_p = tl.minimum(log_p + (target_logprob - draft_logprob), 0.0) tl.store(cumulative_log_p_ptr + logit_idx, log_p) @@ -406,6 +410,8 @@ def _compute_local_residual_mass_kernel( draft_local_sumexp_ptr, draft_local_sumexp_stride, # [num_logits] + draft_sampled_ptr, + # [num_logits] expanded_idx_mapping_ptr, # [num_logits] expanded_local_pos_ptr, @@ -425,6 +431,11 @@ def _compute_local_residual_mass_kernel( # first and last (bonus) positions aren't needed for this computation. return + if tl.load(draft_sampled_ptr + logit_idx + 1) < 0: + # -1 placeholder token. The rejection kernel treats the preceding token + # as the end of the block, so this position's residual mass is unused. + return + req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx).to(tl.int64) temp = tl.load(temp_ptr + req_state_idx).to(tl.float32) if temp == 0.0: @@ -539,41 +550,23 @@ def _rejection_kernel( accepted_length = tl.zeros((), tl.int64) target_lse = 0.0 draft_lse = 0.0 - accepted = True + verifying = True for i in range(num_draft_tokens): logit_idx = start_idx + i draft_sampled = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64) - pos = tl.load(pos_ptr + logit_idx) - u = tl_rand32(seed, pos, includes_zero=False) - if USE_BLOCK_VERIFICATION and not is_greedy: - # Block verification (Sun et al., 2024): https://arxiv.org/abs/2403.10444 - prefix_joint_ratio = tl.exp( - tl.load(cumulative_log_p_ptr + logit_idx).to(tl.float32) - ) - if i < num_draft_tokens - 1: - residual_mass = _compute_global_residual_mass( - local_residual_mass_ptr, - local_residual_mass_stride, - prefix_joint_ratio, - target_logits_ptr, - target_logits_stride, - target_local_max_ptr, - target_local_max_stride, - target_local_sumexp_ptr, - target_local_sumexp_stride, - draft_sampled_ptr, - logit_idx + 1, - vocab_num_blocks, - PADDED_VOCAB_NUM_BLOCKS, - HAS_DRAFT_LOGITS, - ) - denom = residual_mass + 1.0 - prefix_joint_ratio - h = tl.where(denom > 0.0, residual_mass / denom, 1.0) - else: - h = prefix_joint_ratio - accepted_length = tl.where(u <= h, i + 1, accepted_length) - tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled) - elif accepted: + # -1 is used for placeholder 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) + if not is_greedy: + # A -1 placeholder ends verification. Greedy is excluded because it + # stores the target argmax upon first rejection, so it rejects the + # placeholder via `accepted` instead. + verifying &= is_valid_draft + + if verifying: + pos = tl.load(pos_ptr + logit_idx) + u = tl_rand32(seed, pos, includes_zero=False) if is_greedy: # Greedy sampling. Accept IFF draft matches target argmax. # NOTE: Target argmax is stored directly so that resampling @@ -589,20 +582,51 @@ def _rejection_kernel( ) if SYNTHETIC_MODE: rate = tl.load(synthetic_conditional_rates_ptr + i) - # -1 is used for padded draft token ids that should be rejected. - accepted &= (u < rate) & (draft_sampled >= 0) + accepted = u < rate else: - accepted &= target_argmax == draft_sampled + accepted = target_argmax == draft_sampled + accepted &= is_valid_draft + verifying = accepted + accepted_length += accepted tl.store( sampled_ptr + req_idx * sampled_stride + i, draft_sampled if accepted else target_argmax, ) + elif USE_BLOCK_VERIFICATION: + # Block verification (Sun et al., 2024): https://arxiv.org/abs/2403.10444 + prefix_joint_ratio = tl.exp( + tl.load(cumulative_log_p_ptr + logit_idx).to(tl.float32) + ) + next_draft_token = tl.load( + draft_sampled_ptr + logit_idx + 2, + mask=i < num_draft_tokens - 1, + other=-1, + ).to(tl.int64) + if next_draft_token >= 0: + residual_mass = _compute_global_residual_mass( + local_residual_mass_ptr, + local_residual_mass_stride, + prefix_joint_ratio, + target_logits_ptr, + target_logits_stride, + target_local_max_ptr, + target_local_max_stride, + target_local_sumexp_ptr, + target_local_sumexp_stride, + next_draft_token, + logit_idx + 1, + vocab_num_blocks, + PADDED_VOCAB_NUM_BLOCKS, + HAS_DRAFT_LOGITS, + ) + denom = residual_mass + 1.0 - prefix_joint_ratio + h = tl.where(denom > 0.0, residual_mass / denom, 1.0) + else: + h = prefix_joint_ratio + accepted_length = tl.where(u <= h, i + 1, accepted_length) + tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled) else: # Speculative decoding (Leviathan et al., 2023): https://arxiv.org/abs/2211.17192 - # -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_logprob, draft_logprob, target_lse, draft_lse = ( _compute_global_logprobs_and_logsumexp( draft_sampled, @@ -631,14 +655,15 @@ def _rejection_kernel( ) if SYNTHETIC_MODE: rate = tl.load(synthetic_conditional_rates_ptr + i) - accepted &= u < rate + accepted = u < rate else: # Probability ratio test: p(x) > u * q(x) # Equivalent log form: log_p(x) > log(u) + log_q(x) - accepted &= target_logprob > tl.log(u) + draft_logprob - accepted &= is_valid_draft + accepted = target_logprob > tl.log(u) + draft_logprob + verifying = accepted + accepted_length += accepted tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled) - accepted_length += accepted + tl.store(rejected_steps_ptr + req_idx, accepted_length) if USE_BLOCK_VERIFICATION and not is_greedy and accepted_length < num_draft_tokens: # Compute the target and draft log exponential sums for the @@ -722,6 +747,13 @@ def _resample_kernel( # the target argmax is already in the sampled tensor. return + rejected_draft_token = tl.load( + draft_sampled_ptr + resample_token_idx + 1, + mask=not is_bonus, + other=0, + ) + is_valid_rejected_draft = rejected_draft_token >= 0 + block_idx = tl.program_id(1) block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = block < vocab_size @@ -732,8 +764,9 @@ def _resample_kernel( ).to(tl.float32) # Compute the residual logits to resample the rejected token from. - if is_bonus: - # Bonus token (no rejections). Directly use the target logits. + if is_bonus or not is_valid_rejected_draft: + # Bonus token (no rejections) or -1 placeholder token. In either case, + # directly use the target logits. residual_logits = target_logits elif HAS_DRAFT_LOGITS: # draft_logits is stored pre-temperature, so apply scale first. @@ -783,7 +816,6 @@ def _resample_kernel( # p_tau * M_b(x) / Z otherwise # Therefore p_tau is a constant that cancels under normalization, # and does not need to be applied. - rejected_draft_token = tl.load(draft_sampled_ptr + resample_token_idx + 1) residual_logits = tl.where( block != rejected_draft_token, target_logits, @@ -1036,6 +1068,7 @@ def rejection_sample( draft_local_max.stride(0), draft_local_sumexp, draft_local_sumexp.stride(0), + draft_sampled, expanded_idx_mapping, expanded_local_pos, temperature,