From 663e5356404332d6c764180aef41cc07015cd34a Mon Sep 17 00:00:00 2001 From: cjackal <44624812+cjackal@users.noreply.github.com> Date: Sat, 17 Oct 2026 04:42:48 +0900 Subject: [PATCH] [Perf][Spec Decode] Skip bonus-logit reduction for short drafts in rejection sampling (#61451) Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com> --- .../test_rejection_sampler_utils.py | 85 +++++++++++++++++++ .../spec_decode/rejection_sampler_utils.py | 30 ++++--- 2 files changed, 105 insertions(+), 10 deletions(-) diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index ab94ec54d2..b93cd2b306 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -1157,6 +1157,91 @@ def test_chunked_requests_match_full_batch(has_draft_logits: bool): assert torch.equal(chunked_sampled[valid], sampled[valid]) +@pytest.mark.parametrize("use_block_verification", [False, True]) +@pytest.mark.parametrize("has_draft_logits", [False, True]) +def test_adaptive_draft_lengths_match_and_bonus_oracle( + use_block_verification: bool, has_draft_logits: bool +): + torch.manual_seed(7) + device = "cuda" + K = 3 + num_drafts = [3, 1, 0, 2, 3, 2] + temperatures = [1.0, 0.0, 0.0, 0.6, 0.0, 1.0] + greedy_num_matches = {1: 1, 2: 0, 4: 2} + num_reqs = len(num_drafts) + + num_logits_per_req = torch.tensor(num_drafts, device=device) + 1 + cu_num_logits = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device) + cu_num_logits[1:] = num_logits_per_req.cumsum(0) + num_logits = int(cu_num_logits[-1]) + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) + expanded_idx_mapping = idx_mapping.repeat_interleave(num_logits_per_req) + expanded_local_pos = ( + torch.arange(num_logits, dtype=torch.int32, device=device) + - cu_num_logits[expanded_idx_mapping] + ) + + draft_logits = torch.randn(num_reqs, K, VOCAB_SIZE, device=device) + target_logits = torch.randn(num_logits, VOCAB_SIZE, device=device) + # Allocate at exactly num_logits so a lookahead past the batch is OOB. + draft_sampled = torch.randint(VOCAB_SIZE, (num_logits,), device=device) + target_argmax = target_logits.argmax(dim=-1) + for req_idx, num_matches in greedy_num_matches.items(): + lo = int(cu_num_logits[req_idx]) + for step in range(num_drafts[req_idx]): + token = target_argmax[lo + step] + if step >= num_matches: + token = (token + 1) % VOCAB_SIZE + draft_sampled[lo + step + 1] = token + + inputs = dict( + target_logits=target_logits, + draft_logits=draft_logits if has_draft_logits else None, + draft_sampled=draft_sampled, + cu_num_logits=cu_num_logits, + pos=torch.arange(num_logits, dtype=torch.int32, device=device), + idx_mapping=idx_mapping, + expanded_idx_mapping=expanded_idx_mapping, + expanded_local_pos=expanded_local_pos, + temperature=torch.tensor(temperatures, device=device), + seed=torch.arange(num_reqs, dtype=torch.int64, device=device), + num_speculative_steps=K, + use_block_verification=use_block_verification, + ) + sampled, num_sampled = rejection_sample(**inputs) + + for req_idx, num_matches in greedy_num_matches.items(): + lo = int(cu_num_logits[req_idx]) + n = num_matches + 1 + assert int(num_sampled[req_idx]) == n, f"greedy request {req_idx}" + assert torch.equal(sampled[req_idx, :n], target_argmax[lo : lo + n]), ( + f"greedy request {req_idx}" + ) + + for req_idx in range(num_reqs): + lo, hi = int(cu_num_logits[req_idx]), int(cu_num_logits[req_idx + 1]) + req_inputs = dict(inputs) + for name in ( + "target_logits", + "draft_sampled", + "pos", + "expanded_idx_mapping", + "expanded_local_pos", + ): + req_inputs[name] = inputs[name][lo:hi] + req_inputs["cu_num_logits"] = cu_num_logits[req_idx : req_idx + 2] - lo + req_inputs["idx_mapping"] = idx_mapping[req_idx : req_idx + 1] + + req_sampled, req_num_sampled = rejection_sample(**req_inputs) + + n = int(req_num_sampled[0]) + assert n <= num_drafts[req_idx] + 1 + assert int(num_sampled[req_idx]) == n, f"request {req_idx} mismatch" + assert torch.equal(sampled[req_idx, :n], req_sampled[0, :n]), ( + f"request {req_idx} mismatch" + ) + + def _residual_recovery_logits( inputs: dict, num_sampled: torch.Tensor, 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 f35468d62b..060c64caa4 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -194,6 +194,16 @@ def _compute_global_logprobs_and_logsumexp( return target_log_prob, draft_log_prob, target_lse, draft_lse +@triton.jit +def _is_bonus_logit(expanded_local_pos_ptr, logit_idx, num_logits): + next_local_pos = tl.load( + expanded_local_pos_ptr + logit_idx + 1, + mask=logit_idx + 1 < num_logits, + other=0, + ) + return next_local_pos == 0 + + @triton.jit def _compute_local_logits_stats_kernel( # [num_logits, num_blocks] @@ -225,17 +235,17 @@ def _compute_local_logits_stats_kernel( # [max_num_reqs] temp_ptr, vocab_size, - num_speculative_steps, + num_logits, BLOCK_SIZE: tl.constexpr, HAS_DRAFT_LOGITS: tl.constexpr, ): logit_idx = tl.program_id(0).to(tl.int64) - draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx) - if draft_step_idx >= num_speculative_steps: - # Bonus token. Max/argmax and summed exponentials are not needed. + if _is_bonus_logit(expanded_local_pos_ptr, logit_idx, num_logits): + # Max/argmax and summed exponentials are not needed. return + draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx) 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) @@ -423,17 +433,17 @@ def _compute_local_residual_mass_kernel( # [max_num_reqs] temp_ptr, vocab_size, - num_speculative_steps, + num_logits, vocab_num_blocks, BLOCK_SIZE: tl.constexpr, PADDED_VOCAB_NUM_BLOCKS: tl.constexpr, ): logit_idx = tl.program_id(0).to(tl.int64) # Adaptive verification can end before num_speculative_steps. - if logit_idx + 1 >= tl.num_programs(0): - return draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx) - if draft_step_idx == 0 or draft_step_idx >= num_speculative_steps: + if draft_step_idx == 0 or _is_bonus_logit( + expanded_local_pos_ptr, logit_idx, num_logits + ): # The acceptance threshold, h, looks one position ahead and sums # over: max(p_i * M_b(x|x_{