[Model Runner V2][Spec Decode] Use fp32 uniform threshold for acceptance (#46878)
Giancarlo Delfin committed
Jun 27, 2026 at 16:09 UTC
b6caeb5a0966103c6df22f019270d66233e1b687
2 files changed
+12
-5
vllm/v1/worker/gpu/sample/gumbel.py
+9
-2
@@ -73,6 +73,14 @@ def tl_rand64(seed, offset, includes_zero: tl.constexpr):
73
return u
74
75
76
+@triton.jit
77
+def tl_rand32(seed, offset, includes_zero: tl.constexpr):
78
+ u = tl.rand(seed, offset)
79
+ if not includes_zero:
80
+ u = tl.maximum(u, _TL_RAND_MIN)
81
+ return u
82
+
83
+
84
@triton.jit
85
def gumbel_block_argmax(
86
logits,
@@ -131,8 +139,7 @@ def gumbel_block_argmax(
139
u = tl_rand64(gumbel_seed, block, includes_zero=False)
140
gumbel_noise = -tl.log(-tl.log(u))
141
else:
134
- u = tl.rand(gumbel_seed, block)
135
- u = tl.maximum(u, _TL_RAND_MIN)
142
+ u = tl_rand32(gumbel_seed, block, includes_zero=False)
143
# Draw the large-noise tail (which decides the argmax winner) from u -> 0,
144
# where fp32 has fine resolution, instead of u -> 1, where fp32 spacing is
145
# ~2**-24. The naive `-log(-log(u))` puts the winning tail at u -> 1,
vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+3
-3
@@ -3,7 +3,7 @@
3
import torch
4
5
from vllm.triton_utils import tl, tldevice, triton
6
-from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand64
6
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand32
7
8
9
@triton.jit
@@ -243,7 +243,7 @@ def _rejection_kernel(
243
244
if SYNTHETIC_MODE:
245
pos = tl.load(pos_ptr + logit_idx)
246
- u = tl_rand64(seed, pos, includes_zero=False)
246
+ u = tl_rand32(seed, pos, includes_zero=False)
247
rate = tl.load(synthetic_conditional_rates_ptr + i)
248
# -1 is used for padded draft token ids that should be rejected.
249
accepted &= (u < rate) & (draft_sampled >= 0)
@@ -272,7 +272,7 @@ def _rejection_kernel(
272
)
273
target_log_prob = target_logit - target_lse
274
pos = tl.load(pos_ptr + logit_idx)
275
- u = tl_rand64(seed, pos, includes_zero=False)
275
+ u = tl_rand32(seed, pos, includes_zero=False)
276
if HAS_DRAFT_LOGITS:
277
draft_logit = tl.load(
278
draft_logits_ptr