[MRV2] Make FP32 Gumbel sampling more accurate (#45996)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Jun 18, 2026 at 12:42 UTC
16908e132e10f75af93049e865130f8987573f5d
2 files changed
+240
-8
tests/v1/worker/test_gpu_gumbel_sample.py
new
+227
@@ -0,0 +1,227 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the Model Runner V2 Gumbel-max sampling kernel.
4
+
5
+Accuracy: define a target categorical distribution as a non-negative int64
6
+count tensor summing to N, turn it into logits (= log(count)), sample many
7
+times with `gumbel_sample`, and check the empirical distribution matches.
8
+
9
+The count tensor is deliberately heavy-tailed (one dominant token, the rest
10
+~18 logits below). That tail is the sensitive part: the fp32 Gumbel noise must
11
+reach ~18 to ever sample it. A flat distribution would keep every token within
12
+a few logits of the top and would not exercise the noise tail at all.
13
+"""
14
+
15
+import math
16
+
17
+import pytest
18
+import torch
19
+
20
+pytest.importorskip("triton")
21
+if not torch.cuda.is_available():
22
+ pytest.skip("CUDA required for Gumbel sampler tests", allow_module_level=True)
23
+
24
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
25
+
26
+DEVICE = "cuda"
27
+VOCAB_SIZE = 200_000
28
+NUM_SAMPLES = 500_000
29
+# Dominant token is exp(HEAD_LOG_GAP)x larger than the unit-count tail, so the
30
+# tail sits ~HEAD_LOG_GAP logits below the top.
31
+HEAD_LOG_GAP = 18.0
32
+# 10-sigma band: a correct sampler effectively never trips it.
33
+Z_TOLERANCE = 10.0
34
+
35
+
36
+def _make_heavy_tailed_counts(seed: int = 1234) -> torch.Tensor:
37
+ """Non-negative int64 counts of shape [VOCAB_SIZE]; target prob = counts/N."""
38
+ gen = torch.Generator(device=DEVICE).manual_seed(seed)
39
+ counts = torch.randint(
40
+ 1, 4, (VOCAB_SIZE,), generator=gen, dtype=torch.int64, device=DEVICE
41
+ )
42
+ counts[0] = round(math.exp(HEAD_LOG_GAP)) # dominant token
43
+ return counts
44
+
45
+
46
+def _counts_to_logits(counts: torch.Tensor) -> torch.Tensor:
47
+ # softmax(log(count)) == count / sum(count); count 0 -> logit -inf -> prob 0.
48
+ return counts.double().log().to(torch.float32)
49
+
50
+
51
+def _sample(
52
+ logits_1d: torch.Tensor,
53
+ num_samples: int,
54
+ *,
55
+ use_fp64: bool = False,
56
+ temperature: float = 1.0,
57
+) -> torch.Tensor:
58
+ """Sample `num_samples` tokens from one logit vector.
59
+
60
+ Fixed seed with a distinct `pos` per sample gives independent draws; the
61
+ logits are broadcast with a 0-stride view to avoid materializing
62
+ [num_samples, vocab_size].
63
+ """
64
+ vocab_size = logits_1d.shape[0]
65
+ logits = logits_1d.unsqueeze(0).expand(num_samples, vocab_size)
66
+ idx_mapping = torch.zeros(num_samples, dtype=torch.int32, device=DEVICE)
67
+ temp = torch.tensor([temperature], dtype=torch.float32, device=DEVICE)
68
+ seed = torch.tensor([0xABCD], dtype=torch.int64, device=DEVICE)
69
+ pos = torch.arange(num_samples, dtype=torch.int64, device=DEVICE)
70
+ return gumbel_sample(
71
+ logits,
72
+ idx_mapping,
73
+ temp,
74
+ seed,
75
+ pos,
76
+ apply_temperature=True,
77
+ use_fp64=use_fp64,
78
+ )
79
+
80
+
81
+def _z_score(observed: int, expected: float, num_trials: int) -> float:
82
+ p = expected / num_trials
83
+ return (observed - expected) / math.sqrt(num_trials * p * (1 - p))
84
+
85
+
86
+def _sample_histogram(
87
+ logits_1d: torch.Tensor, num_samples: int, *, chunk: int = 1_000_000
88
+) -> torch.Tensor:
89
+ """Histogram of `num_samples` draws, accumulated in chunks.
90
+
91
+ Chunking keeps the kernel's per-sample scratch ([chunk, num_blocks]) bounded
92
+ so a large sample count does not blow up memory.
93
+ """
94
+ vocab_size = logits_1d.shape[0]
95
+ hist = torch.zeros(vocab_size, dtype=torch.float64, device=DEVICE)
96
+ for start in range(0, num_samples, chunk):
97
+ size = min(chunk, num_samples - start)
98
+ logits = logits_1d.unsqueeze(0).expand(size, vocab_size)
99
+ idx_mapping = torch.zeros(size, dtype=torch.int32, device=DEVICE)
100
+ temp = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
101
+ seed = torch.tensor([0xABCD], dtype=torch.int64, device=DEVICE)
102
+ pos = torch.arange(start, start + size, dtype=torch.int64, device=DEVICE)
103
+ out = gumbel_sample(
104
+ logits, idx_mapping, temp, seed, pos, apply_temperature=True
105
+ )
106
+ hist += torch.bincount(out, minlength=vocab_size).double()
107
+ return hist
108
+
109
+
110
+# ----------------------------- Accuracy ------------------------------------
111
+
112
+
113
+@pytest.mark.parametrize("use_fp64", [False, True])
114
+def test_sampling_matches_target_distribution(use_fp64: bool):
115
+ counts = _make_heavy_tailed_counts()
116
+ total = counts.sum().item()
117
+ logits = _counts_to_logits(counts)
118
+
119
+ sampled = _sample(logits, NUM_SAMPLES, use_fp64=use_fp64)
120
+ assert sampled.min() >= 0 and sampled.max() < VOCAB_SIZE
121
+
122
+ # The dominant token (index 0) and the aggregate tail are the two
123
+ # statistically resolvable bins (individual tail tokens are far below the
124
+ # ~5/N detectability floor). The tail mass is small but well above noise,
125
+ # and it lives beyond the fp32 Gumbel cap -- the regime sensitive to noise
126
+ # precision -- so matching it is the meaningful check.
127
+ tail_prob = (total - counts[0].item()) / total
128
+ tail_count = (sampled != 0).sum().item()
129
+ z = _z_score(tail_count, NUM_SAMPLES * tail_prob, NUM_SAMPLES)
130
+ assert abs(z) < Z_TOLERANCE, (
131
+ f"sampled tail mass {tail_count / NUM_SAMPLES:.3e} != target "
132
+ f"{tail_prob:.3e} (z={z:.2f})"
133
+ )
134
+
135
+
136
+def test_full_vocab_distribution_fidelity():
137
+ """The sampled distribution matches the target across the WHOLE vocab.
138
+
139
+ A near-flat count tensor makes every one of the 200K bins individually
140
+ measurable. With ~20 samples/bin, a goodness-of-fit over all bins checks
141
+ that no part of the vocab is over- or under-represented (the heavy-tailed
142
+ test above only resolves head vs aggregate tail). Empirically the fp32
143
+ sampler is as faithful here as torch.multinomial; the residual error is the
144
+ multinomial sampling-noise floor, not the kernel.
145
+ """
146
+ gen = torch.Generator(device=DEVICE).manual_seed(2024)
147
+ counts = torch.randint(
148
+ 500, 1500, (VOCAB_SIZE,), generator=gen, dtype=torch.int64, device=DEVICE
149
+ )
150
+ total = counts.sum().item()
151
+ logits = _counts_to_logits(counts)
152
+
153
+ num_samples = 4_000_000
154
+ hist = _sample_histogram(logits, num_samples)
155
+
156
+ # Diversity: essentially every token must be reachable (no starved region).
157
+ coverage = (hist > 0).sum().item() / VOCAB_SIZE
158
+ assert coverage > 0.99, f"only {coverage:.4f} of the vocab was ever sampled"
159
+
160
+ # Goodness-of-fit across all bins (each has expected count >= ~10).
161
+ expected = (counts.double() / total) * num_samples
162
+ chi2 = (((hist - expected) ** 2) / expected).sum().item()
163
+ df = VOCAB_SIZE - 1
164
+ assert chi2 < df + 10 * math.sqrt(2 * df), f"chi2={chi2:.0f}, df={df}"
165
+
166
+
167
+# ----------------------------- Edge cases ----------------------------------
168
+
169
+
170
+def test_greedy_temperature_zero_returns_argmax():
171
+ """temperature == 0 skips Gumbel noise and returns the exact argmax."""
172
+ torch.manual_seed(0)
173
+ num_reqs = 128
174
+ logits = torch.randn(num_reqs, VOCAB_SIZE, device=DEVICE, dtype=torch.float32)
175
+ idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE)
176
+ temp = torch.zeros(num_reqs, dtype=torch.float32, device=DEVICE)
177
+ seed = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)
178
+ pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)
179
+
180
+ sampled = gumbel_sample(
181
+ logits, idx_mapping, temp, seed, pos, apply_temperature=True
182
+ )
183
+ assert torch.equal(sampled, logits.argmax(dim=-1))
184
+
185
+
186
+def test_zero_count_tokens_are_never_sampled():
187
+ """Count 0 -> -inf logit -> probability 0; must never be selected."""
188
+ counts = _make_heavy_tailed_counts(seed=7)
189
+ zeroed = torch.arange(1, VOCAB_SIZE, 2, device=DEVICE) # odd indices (not head)
190
+ counts[zeroed] = 0
191
+ logits = _counts_to_logits(counts)
192
+
193
+ sampled = _sample(logits, NUM_SAMPLES)
194
+ assert sampled.min() >= 0 and sampled.max() < VOCAB_SIZE
195
+ assert not torch.isin(sampled, zeroed).any(), "sampled a zero-probability token"
196
+
197
+
198
+def test_single_nonzero_token_is_always_sampled():
199
+ """A lone finite logit must win every draw, regardless of its index."""
200
+ counts = torch.zeros(VOCAB_SIZE, dtype=torch.int64, device=DEVICE)
201
+ counts[123_456] = 1000
202
+ logits = _counts_to_logits(counts)
203
+
204
+ sampled = _sample(logits, 10_000)
205
+ assert (sampled == 123_456).all()
206
+
207
+
208
+@pytest.mark.parametrize("vocab_size", [1, 999, 1024, 4097])
209
+def test_vocab_size_not_multiple_of_block(vocab_size: int):
210
+ """Per-block tail masking for non-block-aligned vocab; all bins measurable."""
211
+ gen = torch.Generator(device=DEVICE).manual_seed(vocab_size)
212
+ counts = torch.randint(
213
+ 20, 200, (vocab_size,), generator=gen, dtype=torch.int64, device=DEVICE
214
+ )
215
+ total = counts.sum().item()
216
+ logits = _counts_to_logits(counts)
217
+ num_samples = max(40 * vocab_size, 50_000)
218
+
219
+ sampled = _sample(logits, num_samples)
220
+ assert sampled.min() >= 0 and sampled.max() < vocab_size
221
+
222
+ observed = torch.bincount(sampled, minlength=vocab_size).double()
223
+ expected = (counts.double() / total) * num_samples
224
+ chi2 = (((observed - expected) ** 2) / expected).sum().item()
225
+ df = vocab_size - 1
226
+ if df >= 1:
227
+ assert chi2 < df + 10 * math.sqrt(2 * df), f"chi2={chi2:.1f}, df={df}"
vllm/v1/worker/gpu/sample/gumbel.py
+13
-8
@@ -2,18 +2,16 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import torch
4
5
-from vllm.triton_utils import HAS_TRITON, tl, triton
5
+from vllm.triton_utils import HAS_TRITON, tl, tldevice, triton
6
7
-# Smallest positive normal fp32 value. Used to clamp the uniform draw so that
8
-# `log(u)` cannot produce -inf (and thus `-log(-log(u))` stays finite).
7
+# Smallest positive value produced by Triton's fp32 `tl.rand`. Used to clamp
8
+# zero draws before the flipped Gumbel transform below.
9
#
10
# Triton requires globals accessed from `@triton.jit` functions to be wrapped
11
# in `tl.constexpr(...)`. We can only do that when Triton is actually
12
# available — on the CPU worker path `tl` is a placeholder whose `constexpr`
13
# attribute is `None`, and `tl.constexpr(...)` would crash at import time.
14
-_FP32_TINY = (
15
- tl.constexpr(float.fromhex("0x1p-126")) if HAS_TRITON else float.fromhex("0x1p-126")
16
-)
14
+_TL_RAND_MIN = tl.constexpr(4.6566127342e-10) if HAS_TRITON else 4.6566127342e-10
15
16
17
@triton.jit
@@ -131,10 +129,17 @@ def gumbel_block_argmax(
129
130
if USE_FP64:
131
u = tl_rand64(gumbel_seed, block, includes_zero=False)
132
+ gumbel_noise = -tl.log(-tl.log(u))
133
else:
134
u = tl.rand(gumbel_seed, block)
136
- u = tl.maximum(u, _FP32_TINY)
137
- gumbel_noise = -tl.log(-tl.log(u))
135
+ u = tl.maximum(u, _TL_RAND_MIN)
136
+ # Draw the large-noise tail (which decides the argmax winner) from u -> 0,
137
+ # where fp32 has fine resolution, instead of u -> 1, where fp32 spacing is
138
+ # ~2**-24. The naive `-log(-log(u))` puts the winning tail at u -> 1,
139
+ # hard-capping the noise at ~16.6 and coarsely quantizing it; using
140
+ # `log1p(-u)` == `log(1 - u)` keeps the tail in the well-resolved region.
141
+ # Note `1 - u` would lose precision for small u, so `log1p` is required.
142
+ gumbel_noise = -tl.log(-tldevice.log1p(-u))
143
144
# Apply gumbel noise.
145
logits = tl.where(mask, logits + gumbel_noise, float("-inf"))