A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 05:45 UTC
[MRV2] Use FP64 for Gumbel noise (#37798)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Mar 22, 2026 at 12:28 UTC
a5e9d511defe2d2dc2dd270674fc197542fc0169
2 files changed
+24
-8
vllm/v1/worker/gpu/sample/gumbel.py
+20
-4
@@ -49,6 +49,21 @@ def apply_temperature(
49
)
50
51
52
+@triton.jit
53
+def tl_rand64(seed, offset, includes_zero: tl.constexpr):
54
+ lo, hi, _, _ = tl.randint4x(seed, offset)
55
+ lo = lo.to(tl.uint32, bitcast=True).to(tl.uint64)
56
+ hi = hi.to(tl.uint32, bitcast=True).to(tl.uint64)
57
+ r = (hi << 32) | lo
58
+
59
+ # 1 / 2**64
60
+ scale = 5.421010862427522170037e-20
61
+ u = r.to(tl.float64) * scale
62
+ if not includes_zero:
63
+ u = tl.maximum(u, 2.2250738585072014e-308) # float64 tiny
64
+ return u
65
+
66
+
67
@triton.jit
68
def _gumbel_sample_kernel(
69
local_argmax_ptr,
@@ -95,15 +110,16 @@ def _gumbel_sample_kernel(
110
mask=mask,
111
)
112
113
+ logits = logits.to(tl.float64)
114
if temp != 0.0:
115
# Calculate the seed for gumbel noise.
116
seed = tl.load(seeds_ptr + req_state_idx)
117
pos = tl.load(pos_ptr + token_idx)
118
gumbel_seed = tl.randint(seed, pos)
119
104
- # Generate gumbel noise in FP32.
105
- u = tl.rand(gumbel_seed, block)
106
- u = tl.maximum(u, 1e-7)
120
+ # tl.rand returns fp32, so build a true fp64 uniform from 64 random
121
+ # bits before applying the double-log transform.
122
+ u = tl_rand64(gumbel_seed, block, includes_zero=False)
123
gumbel_noise = -tl.log(-tl.log(u))
124
125
# Apply gumbel noise.
@@ -128,7 +144,7 @@ def gumbel_sample(
144
BLOCK_SIZE = 1024
145
num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
146
local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64)
131
- local_max = logits.new_empty(num_tokens, num_blocks, dtype=torch.float32)
147
+ local_max = logits.new_empty(num_tokens, num_blocks, dtype=torch.float64)
148
_gumbel_sample_kernel[(num_tokens, num_blocks)](
149
local_argmax,
150
local_argmax.stride(0),
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+4
-4
@@ -6,7 +6,7 @@ from vllm.triton_utils import tl, triton
6
from vllm.v1.outputs import LogprobsTensors
7
from vllm.v1.worker.gpu.input_batch import InputBatch
8
from vllm.v1.worker.gpu.metrics.logits import get_num_nans
9
-from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
9
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample, tl_rand64
10
from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
11
from vllm.v1.worker.gpu.sample.output import SamplerOutput
12
from vllm.v1.worker.gpu.sample.sampler import Sampler
@@ -211,12 +211,12 @@ def _probabilistic_rejection_kernel(
211
else:
212
target_prob = tl.load(
213
target_probs_ptr + logit_idx * target_probs_stride + draft_sampled
214
- )
214
+ ).to(tl.float64)
215
draft_prob = tl.load(
216
draft_probs_ptr + logit_idx * draft_probs_stride + draft_sampled
217
- )
217
+ ).to(tl.float64)
218
pos = tl.load(pos_ptr + logit_idx)
219
- u = tl.sum(tl.rand(seed, pos + tl.arange(0, 1)))
219
+ u = tl_rand64(seed, pos, includes_zero=False)
220
accepted &= target_prob > u * draft_prob
221
tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
222
rejected_step += accepted