[SpecDecode] Restore residual-logits comments in _resample_kernel (#58166)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Lucas Wilkinson committed
Sep 22, 2026 at 11:07 UTC
5c121e962a31ed0b5fe70d743b6e194686d2b473
1 file changed
+17
-2
vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+17
-2
@@ -806,7 +806,10 @@ def _resample_kernel(
806
other=float("-inf"),
807
).to(tl.float32)
808
809
+ # Compute the residual logits to resample the rejected token from.
810
if is_bonus or not is_valid_rejected_draft:
811
+ # Bonus token (no rejections) or -1 placeholder token. In either case,
812
+ # directly use the target logits.
813
residual_logits = target_logits
814
elif HAS_DRAFT_LOGITS:
815
# draft_logits is stored pre-temperature, so apply scale first.
@@ -836,7 +839,12 @@ def _resample_kernel(
839
)
840
target_log_probs += log_p_tau
841
draft_log_probs = draft_logits - draft_lse
839
- # Compute log(max(p - q, 0)) without subtracting probabilities.
842
+ # Compute the residual:
843
+ # r(x) = max(p(x) - q(x), 0)
844
+ # Gumbel sampling needs logits, so we compute it in log space:
845
+ # log(r(x)) = log(max(exp(log_p(x)) - exp(log_q(x)), 0))
846
+ # The more numerically stable form is:
847
+ # log(max(exp(a) - exp(b), 0)) = a + log(max(1 - exp(b - a), 0))
848
ratio = tl.exp(draft_log_probs - target_log_probs)
849
residual_logits = tl.where(
850
ratio < 1.0,
@@ -844,13 +852,20 @@ def _resample_kernel(
852
float("-inf"),
853
).to(tl.float32)
854
else:
847
- # The block-verification factor is constant and cancels on normalization.
855
+ # One-hot draft. The residual is just the target distribution with
856
+ # the rejected draft token probability zeroed out.
857
+ # NOTE: During block verification, the residual becomes:
858
+ # 0 if x == rejected_draft_token
859
+ # p_tau * M_b(x) / Z otherwise
860
+ # Therefore p_tau is a constant that cancels under normalization,
861
+ # and does not need to be applied.
862
residual_logits = tl.where(
863
block != rejected_draft_token,
864
target_logits,
865
float("-inf"),
866
).to(tl.float32)
867
868
+ # Resample the rejected/bonus token.
869
if WATERMARK:
870
# Padded and greedy rows retain the stock draw.
871
is_watermarked = (