[Perf][Spec Decode] Skip bonus-logit reduction for short drafts in rejection sampling (#61451)

Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com>

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,

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_{<i}) - M_s(x|x_{<i}), 0). Tokens at the
         # first and last (bonus) positions aren't needed for this computation.
@@ -1138,7 +1148,7 @@ def rejection_sample(
         expanded_local_pos,
         temperature,
         vocab_size,
-        num_speculative_steps,
+        num_logits,
         BLOCK_SIZE=VOCAB_BLOCK_SIZE,
         HAS_DRAFT_LOGITS=has_draft_logits,
     )
@@ -1209,7 +1219,7 @@ def rejection_sample(
                 expanded_local_pos,
                 temperature,
                 vocab_size,
-                num_speculative_steps,
+                num_logits,
                 vocab_num_blocks,
                 BLOCK_SIZE=VOCAB_BLOCK_SIZE,
                 PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,