[Kernel] Use Murmur3 RNG for Gumbel sampling (#51367)
Signed-off-by: freyfwt <110148731+freyfwt@users.noreply.github.com> Signed-off-by: freyfwt <freytian1996@gmail.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
freyfwt committed
Sep 16, 2026 at 14:52 UTC
18553c5c9ff02acb84bd0f8698e90609e1aec2e9
5 files changed
+216
-26
tests/entrypoints/llm/test_struct_output_generate.py
+1
-1
@@ -948,7 +948,7 @@ def test_structured_output_batched_with_non_structured_outputs_requests(
948
prompts = [structured_outputs_prompt, non_structured_outputs_prompt]
949
sampling_params = [
950
SamplingParams(
951
- temperature=1.0,
951
+ temperature=0,
952
max_tokens=400,
953
structured_outputs=StructuredOutputsParams(json=sample_json_schema),
954
),
tests/entrypoints/openai/chat_completion/test_include_reasoning.py
+8
-4
@@ -45,7 +45,8 @@ async def test_include_reasoning_true_non_streaming(client: openai.AsyncOpenAI):
45
response = await client.chat.completions.create(
46
model=MODEL_NAME,
47
messages=MESSAGES,
48
- max_tokens=200,
48
+ max_tokens=512,
49
+ temperature=0,
50
extra_body={"include_reasoning": True},
51
)
52
@@ -63,7 +64,8 @@ async def test_include_reasoning_false_non_streaming(client: openai.AsyncOpenAI)
64
response = await client.chat.completions.create(
65
model=MODEL_NAME,
66
messages=MESSAGES,
66
- max_tokens=200,
67
+ max_tokens=512,
68
+ temperature=0,
69
extra_body={"include_reasoning": False},
70
)
71
@@ -99,7 +101,8 @@ async def test_include_reasoning_true_streaming(client: openai.AsyncOpenAI):
101
stream = await client.chat.completions.create(
102
model=MODEL_NAME,
103
messages=MESSAGES,
102
- max_tokens=200,
104
+ max_tokens=512,
105
+ temperature=0,
106
stream=True,
107
extra_body={"include_reasoning": True},
108
)
@@ -130,7 +133,8 @@ async def test_include_reasoning_false_streaming(client: openai.AsyncOpenAI):
133
stream = await client.chat.completions.create(
134
model=MODEL_NAME,
135
messages=MESSAGES,
133
- max_tokens=200,
136
+ max_tokens=512,
137
+ temperature=0,
138
stream=True,
139
extra_body={"include_reasoning": False},
140
)
tests/entrypoints/openai/completion/test_completion.py
+2
-10
@@ -353,7 +353,7 @@ async def test_parallel_no_streaming(client: openai.AsyncOpenAI, model_name: str
353
354
prompt = "What is an LLM?"
355
n = 3
356
- max_tokens = 50 # we want some to finish earlier than others
356
+ max_tokens = 50
357
358
# High temperature to maximize chance of unique completions.
359
completion = await client.completions.create(
@@ -371,16 +371,12 @@ async def test_parallel_no_streaming(client: openai.AsyncOpenAI, model_name: str
371
num_completions = len(completion.choices)
372
assert num_completions == n, f"Num completions {num_completions} but expected {n}."
373
completion_repeats: dict[str, int] = {}
374
- output_token_lengths = set()
374
for idx, choice in enumerate(completion.choices):
375
# Assert correct completion index & some finish reason.
376
assert choice.index == idx, f"Index {choice.index} but expected {idx}."
377
assert choice.finish_reason is not None, "None finish_reason is invalid."
378
text = choice.text
379
completion_repeats[text] = completion_repeats.get(text, 0) + 1
381
- output_token_lengths.add(len(choice.logprobs.tokens))
382
- # Assert subrequests finished at different times
383
- assert len(output_token_lengths) > 1
380
# Assert `n` unique completions
381
num_unique = len(completion_repeats)
382
if num_unique != n:
@@ -403,7 +399,7 @@ async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str):
399
400
prompt = "What is an LLM?"
401
n = 3
406
- max_tokens = 50 # we want some to finish earlier than others
402
+ max_tokens = 50
403
404
stream = await client.completions.create(
405
model=model_name,
@@ -427,19 +423,15 @@ async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str):
423
f"Expected {n} completions with valid indices and finish_reason."
424
)
425
completion_repeats: dict[str, int] = {}
430
- chunk_lengths = set()
426
for chunk in chunks:
427
chunk_len = len(chunk)
428
# Assert correct number of completion tokens
434
- chunk_lengths.add(chunk_len)
429
assert chunk_len <= max_tokens, (
430
f"max_tokens={max_tokens} but chunk len is {chunk_len}."
431
)
432
text = "".join(chunk)
433
completion_repeats[text] = completion_repeats.get(text, 0) + 1
434
print(text)
441
- # Assert subrequests finished at different times
442
- assert len(chunk_lengths) > 1
435
# Assert `n` unique completions
436
num_unique = len(completion_repeats)
437
if num_unique != n:
tests/v1/worker/test_gpu_gumbel_sample.py
+105
-3
@@ -21,7 +21,12 @@ 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
24
+from vllm.triton_utils import tl, triton
25
+from vllm.v1.worker.gpu.sample.gumbel import (
26
+ _uniform64_from_random53,
27
+ gumbel_sample,
28
+ murmur3_hash32,
29
+)
30
31
DEVICE = "cuda"
32
VOCAB_SIZE = 200_000
@@ -33,6 +38,103 @@ HEAD_LOG_GAP = 18.0
38
Z_TOLERANCE = 10.0
39
40
41
+@triton.jit
42
+def _murmur3_known_answer_kernel(
43
+ seeds_ptr,
44
+ positions_ptr,
45
+ offsets_ptr,
46
+ hashes_ptr,
47
+ domain_hashes_ptr,
48
+ numel: tl.constexpr,
49
+ BLOCK_SIZE: tl.constexpr,
50
+):
51
+ idx = tl.arange(0, BLOCK_SIZE)
52
+ mask = idx < numel
53
+ seed = tl.load(seeds_ptr + idx, mask=mask)
54
+ pos = tl.load(positions_ptr + idx, mask=mask)
55
+ offset = tl.load(offsets_ptr + idx, mask=mask)
56
+ hash32 = murmur3_hash32(seed, pos, offset).to(tl.int64)
57
+ domain_hash32 = murmur3_hash32(seed, pos, offset, domain=0x9E3779B9).to(tl.int64)
58
+ tl.store(hashes_ptr + idx, hash32, mask=mask)
59
+ tl.store(domain_hashes_ptr + idx, domain_hash32, mask=mask)
60
+
61
+
62
+@triton.jit
63
+def _uniform64_endpoint_kernel(
64
+ random53_ptr,
65
+ uniform_ptr,
66
+ numel: tl.constexpr,
67
+ BLOCK_SIZE: tl.constexpr,
68
+):
69
+ idx = tl.arange(0, BLOCK_SIZE)
70
+ mask = idx < numel
71
+ random53 = tl.load(random53_ptr + idx, mask=mask).to(tl.uint64)
72
+ tl.store(uniform_ptr + idx, _uniform64_from_random53(random53), mask=mask)
73
+
74
+
75
+def test_murmur3_known_answers():
76
+ seeds = torch.tensor(
77
+ [0, 1, -1, 0x123456789ABCDEF, -(1 << 63), (1 << 63) - 1],
78
+ dtype=torch.int64,
79
+ device=DEVICE,
80
+ )
81
+ positions = torch.tensor(
82
+ [0, 2, 17, 123456789, 1 << 30, (1 << 30) + 7],
83
+ dtype=torch.int64,
84
+ device=DEVICE,
85
+ )
86
+ offsets = torch.tensor(
87
+ [0, 3, 1023, 199999, 4096, 0xFFFFFFFF],
88
+ dtype=torch.int64,
89
+ device=DEVICE,
90
+ )
91
+ hashes = torch.empty_like(seeds)
92
+ domain_hashes = torch.empty_like(seeds)
93
+
94
+ _murmur3_known_answer_kernel[(1,)](
95
+ seeds,
96
+ positions,
97
+ offsets,
98
+ hashes,
99
+ domain_hashes,
100
+ numel=seeds.numel(),
101
+ BLOCK_SIZE=8,
102
+ )
103
+
104
+ assert hashes.tolist() == [
105
+ 0x8134CDF8,
106
+ 0x23B53F91,
107
+ 0xC8C13C14,
108
+ 0x5DDE974F,
109
+ 0xF1A9355B,
110
+ 0xD952F9AA,
111
+ ]
112
+ assert domain_hashes.tolist() == [
113
+ 0xD28AD962,
114
+ 0xE3AAF818,
115
+ 0xD223C51C,
116
+ 0x0547D2FE,
117
+ 0x3AED160F,
118
+ 0x49E8EA7F,
119
+ ]
120
+
121
+
122
+def test_uniform64_excludes_endpoints():
123
+ random53 = torch.tensor([0, (1 << 53) - 1], dtype=torch.int64, device=DEVICE)
124
+ uniform = torch.empty(2, dtype=torch.float64, device=DEVICE)
125
+
126
+ _uniform64_endpoint_kernel[(1,)](
127
+ random53,
128
+ uniform,
129
+ numel=random53.numel(),
130
+ BLOCK_SIZE=2,
131
+ )
132
+
133
+ assert torch.isfinite(uniform).all()
134
+ assert ((uniform > 0.0) & (uniform < 1.0)).all()
135
+ assert uniform.tolist() == [2.0**-54, 1.0 - 2.0**-53]
136
+
137
+
138
def _make_heavy_tailed_counts(seed: int = 1234) -> torch.Tensor:
139
"""Non-negative int64 counts of shape [VOCAB_SIZE]; target prob = counts/N."""
140
gen = torch.Generator(device=DEVICE).manual_seed(seed)
@@ -180,7 +282,7 @@ def test_full_vocab_distribution_fidelity():
282
283
284
def test_drafting_uses_a_separate_noise_stream():
183
- """is_drafting salts the Philox offset: same inputs, different draws.
285
+ """is_drafting salts the RNG position: same inputs, different draws.
286
287
The draft proposal and the residual resample after a rejection must be
288
independent. They key noise by the same (seed, pos), so only the salt keeps
@@ -189,7 +291,7 @@ def test_drafting_uses_a_separate_noise_stream():
291
tests/v1/spec_decode/test_rejection_sampler_utils.py for the distributional
292
consequence.
293
192
- Relocating the offset must not distort the draw either, so both streams are
294
+ Salting the position must not distort the draw either, so both streams are
295
checked against the target's far-tail mass, which sits HEAD_LOG_GAP logits
296
below the head -- the regime where fp32 Gumbel precision matters.
297
"""
vllm/v1/worker/gpu/sample/gumbel.py
+100
-8
@@ -2,10 +2,10 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import torch
4
5
-from vllm.triton_utils import HAS_TRITON, tl, tldevice, triton
5
+from vllm.triton_utils import HAS_TRITON, tl, triton
6
7
-# Smallest positive value produced by Triton's fp32 `tl.rand`. Used to clamp
8
-# zero draws before the flipped Gumbel transform below.
7
+# Smallest positive value produced by Triton's fp32 `tl.rand`. Used by the
8
+# retained Philox helper for rejection sampling.
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
@@ -87,6 +87,97 @@ def tl_rand32(seed, offset, includes_zero: tl.constexpr):
87
return u
88
89
90
+@triton.jit
91
+def _murmur3_rotl32(value, shift: tl.constexpr):
92
+ return (value << shift) | (value >> (32 - shift))
93
+
94
+
95
+@triton.jit
96
+def _murmur3_mix(h, key):
97
+ key *= 0xCC9E2D51
98
+ key = _murmur3_rotl32(key, 15)
99
+ key *= 0x1B873593
100
+ h ^= key
101
+ h = _murmur3_rotl32(h, 13)
102
+ return h * 5 + 0xE6546B64
103
+
104
+
105
+@triton.jit
106
+def _murmur3_fmix32(h):
107
+ h ^= h >> 16
108
+ h *= 0x85EBCA6B
109
+ h ^= h >> 13
110
+ h *= 0xC2B2AE35
111
+ return h ^ (h >> 16)
112
+
113
+
114
+@triton.jit
115
+def murmur3_hash32(seed, pos, offset, domain: tl.constexpr = 0):
116
+ seed = seed.to(tl.int64)
117
+ pos = pos.to(tl.int64)
118
+ offset = offset.to(tl.uint32)
119
+ # Keep the request-wide prefix scalar until the token offset is mixed in.
120
+ h = (seed ^ seed).to(tl.uint32)
121
+ h ^= domain
122
+ h = _murmur3_mix(h, (seed & 0xFFFFFFFF).to(tl.uint32))
123
+ h = _murmur3_mix(h, ((seed >> 32) & 0xFFFFFFFF).to(tl.uint32))
124
+ h = _murmur3_mix(h, (pos & 0xFFFFFFFF).to(tl.uint32))
125
+ h = _murmur3_mix(h, offset)
126
+ return _murmur3_fmix32(h ^ 16)
127
+
128
+
129
+@triton.jit
130
+def murmur3_uniform32(seed, pos, offset):
131
+ random32 = murmur3_hash32(seed, pos, offset)
132
+ # Split the uint32 before converting to fp32 so backends without a native
133
+ # uint32-to-float conversion can still use all 32 source bits. Both 16-bit
134
+ # halves convert exactly; their sum is the correctly rounded fp32 value of
135
+ # (random32 + 0.5) * 2**-32. In particular, the u -> 0 winning tail keeps
136
+ # the full 32-bit source resolution instead of being truncated to 24 bits.
137
+ hi16 = (random32 >> 16).to(tl.int32)
138
+ lo16 = (random32 & 0xFFFF).to(tl.int32)
139
+ return (
140
+ hi16.to(tl.float32) * 1.52587890625e-05
141
+ + (lo16.to(tl.float32) + 0.5) * 2.3283064365386963e-10
142
+ )
143
+
144
+
145
+@triton.jit
146
+def _uniform64_from_random53(random53):
147
+ uniform = (random53.to(tl.float64) + 0.5) * 1.1102230246251565e-16
148
+ # The largest midpoint rounds to 1.0 in fp64; keep the uniform open without
149
+ # relying on a near-one literal that Triton may materialize in fp32.
150
+ return tl.where(uniform == 1.0, uniform - 1.1102230246251565e-16, uniform)
151
+
152
+
153
+@triton.jit
154
+def murmur3_uniform64(seed, pos, offset):
155
+ lo = murmur3_hash32(seed, pos, offset).to(tl.uint64)
156
+ hi = murmur3_hash32(seed, pos, offset, domain=0x9E3779B9).to(tl.uint64)
157
+ random53 = ((hi << 32) | lo) >> 11
158
+ return _uniform64_from_random53(random53)
159
+
160
+
161
+@triton.jit
162
+def _log1p_neg_stable(value):
163
+ # Preserve precision for the positive Gumbel tail without relying on a
164
+ # backend-specific libdevice log1p. The degree-8 series has absolute error
165
+ # below 6e-7 on [0, 0.25]; subtraction is well-conditioned elsewhere for
166
+ # the part of the distribution that can win the argmax.
167
+ polynomial = 1.0 / 8.0
168
+ polynomial = 1.0 / 7.0 + value * polynomial
169
+ polynomial = 1.0 / 6.0 + value * polynomial
170
+ polynomial = 1.0 / 5.0 + value * polynomial
171
+ polynomial = 1.0 / 4.0 + value * polynomial
172
+ polynomial = 1.0 / 3.0 + value * polynomial
173
+ polynomial = 1.0 / 2.0 + value * polynomial
174
+ polynomial = 1.0 + value * polynomial
175
+ series = -value * polynomial
176
+
177
+ direct = tl.log(tl.maximum(1.0 - value, 5.960464477539063e-08))
178
+ return tl.where(value < 0.25, series, direct)
179
+
180
+
181
@triton.jit
182
def gumbel_noised_argmax(
183
logits,
@@ -117,14 +208,15 @@ def gumbel_noised_argmax(
208
if temp != 0.0:
209
if IS_DRAFTING:
210
pos = pos + _DRAFT_NOISE_SALT
120
- gumbel_seed = tl.randint(seed, pos)
211
if USE_FP64:
122
- u = tl_rand64(gumbel_seed, keys, includes_zero=False)
212
+ u = murmur3_uniform64(seed, pos, keys)
213
gumbel_noise = -tl.log(-tl.log(u))
214
else:
125
- u = tl_rand32(gumbel_seed, keys, includes_zero=False)
126
- # log1p keeps the winning tail at u -> 0, where fp32 resolves it.
127
- gumbel_noise = -tl.log(-tldevice.log1p(-u))
215
+ u = murmur3_uniform32(seed, pos, keys)
216
+ # Draw the large-noise tail (which decides the argmax winner) from
217
+ # u -> 0, where fp32 has fine resolution. Avoid backend-specific
218
+ # log1p while preserving precision in the winning tail.
219
+ gumbel_noise = -tl.log(-_log1p_neg_stable(u))
220
logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
221
222
return tl.max(logits, axis=0, return_indices=True)