[Model Runner V2][Spec Decode] Use log1p to compute residual during rejection sampling (#46665)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Jun 25, 2026 at 18:46 UTC
c53994e1348bac3496aafb88e9e731124a00a8a7
1 file changed
+7
-5
vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+7
-5
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import torch
4
5
-from vllm.triton_utils import tl, triton
5
+from vllm.triton_utils import tl, tldevice, triton
6
from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand64
7
8
@@ -387,14 +387,16 @@ def _resample_kernel(
387
draft_lse = tl.load(draft_rejected_logsumexp_ptr + req_idx)
388
target_log_probs = target_logits - target_lse
389
draft_log_probs = draft_logits - draft_lse
390
- # Compute the residual: max(p(x) - q(x), 0)
391
- # Equivalent log form: log(max(exp(log_p(x)) - exp(log_q(x)), 0))
390
+ # Compute the residual:
391
+ # r(x) = max(p(x) - q(x), 0)
392
+ # Gumbel sampling needs logits, so we compute it in log space:
393
+ # log(r(x)) = log(max(exp(log_p(x)) - exp(log_q(x)), 0))
394
# The more numerically stable form is:
393
- # log(max(exp(a) - exp(b), 0)) = a + log(max(1 - exp(b - a), 0))
395
+ # log(max(exp(a) - exp(b), 0)) = a + log(max(1 - exp(b - a), 0))
396
ratio = tl.exp(draft_log_probs - target_log_probs)
397
residual_logits = tl.where(
398
ratio < 1.0,
397
- target_log_probs + tl.log(1 - ratio),
399
+ target_log_probs + tldevice.log1p(-ratio),
400
float("-inf"),
401
).to(tl.float32)
402
else: