[Perf] DiffusionGemma: one-pass sampler statistics kernel (#58226)
Signed-off-by: Matt Mastracci <matthew@mastracci.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Matt Mastracci committed
Sep 25, 2026 at 12:14 UTC
353e491efd3a23057004a3ab68fcfbda4abbee13
4 files changed
+386
-57
tests/v1/sample/test_diffusion_gemma_reads.py
+13
-4
@@ -12,7 +12,9 @@ from vllm.model_executor.models.diffusion_gemma import (
12
DiffusionGemmaRequestStates,
13
_compiled_sample_step,
14
_concat_logprob_stashes,
15
+ _denoise_temperature,
16
_mask_rows_to_allowed,
17
+ sample_row_stats_reference,
18
)
19
from vllm.platforms import current_platform
20
from vllm.v1.outputs import LogprobsTensors
@@ -132,8 +134,18 @@ def _denoise_once(
134
decode_idx = torch.arange(n, dtype=torch.int64, device=device)
135
sampled = torch.zeros(n, CL, dtype=torch.int32, device=device)[:, :width]
136
num_sampled = torch.zeros(n, dtype=torch.int32, device=device)
135
- _compiled_sample_step(
137
+ temp = _denoise_temperature(states.step, decode_slots, float(MAX_STEPS), 0.5, 1.0)
138
+ argmax, sample, entropy, probs = sample_row_stats_reference(
139
torch.zeros(n * width, VOCAB, device=device),
140
+ temp,
141
+ width,
142
+ embed_dtype if compute_sc else None,
143
+ )
144
+ _compiled_sample_step(
145
+ sample.view(n, width),
146
+ argmax.view(n, width),
147
+ entropy.view(n, width),
148
+ probs.view(n, width, -1) if probs is not None else None,
149
decode_slots,
150
decode_idx,
151
decode_slots,
@@ -159,9 +171,6 @@ def _denoise_once(
171
sampled,
172
num_sampled,
173
torch.zeros(MAX_REQS, CL, dtype=torch.int64, device=device),
162
- max_denoising_steps=float(MAX_STEPS),
163
- t_min=0.5,
164
- t_max=1.0,
174
confidence_threshold=0.1,
175
vocab_size=VOCAB,
176
CL=width,
tests/v1/sample/test_diffusion_sampler_stats.py
new
+129
@@ -0,0 +1,129 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""The fused row statistics of the DiffusionGemma sampler against the PyTorch
4
+reference: argmax, entropy and softmax must match. A zero temperature must
5
+sample the argmax and a positive one must sample from the distribution."""
6
+
7
+import math
8
+
9
+import pytest
10
+import torch
11
+
12
+from vllm.model_executor.models.diffusion_gemma_sampler import (
13
+ sample_row_stats,
14
+ sample_row_stats_reference,
15
+)
16
+
17
+pytestmark = pytest.mark.skipif(
18
+ not torch.cuda.is_available(), reason="the fused kernel needs a GPU"
19
+)
20
+
21
+
22
+@pytest.mark.parametrize("vocab", [5000, 9001])
23
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
24
+def test_stats_match_reference(vocab: int, dtype: torch.dtype):
25
+ torch.manual_seed(0)
26
+ num_decode, cl = 9, 4
27
+ logits = (torch.randn(num_decode * cl, vocab, device="cuda") * 3).to(dtype)
28
+ temps = torch.tensor([0.0, 0.7, 2.0, 1.0, 0.0, 0.3, 1.5, 1.0, 0.0], device="cuda")
29
+ argmax, sample, entropy, probs = sample_row_stats(
30
+ logits, temps, cl, seed=1234, probs_dtype=torch.bfloat16
31
+ )
32
+ ref_argmax, _, ref_entropy, ref_probs = sample_row_stats_reference(
33
+ logits, temps, cl, probs_dtype=torch.bfloat16
34
+ )
35
+ assert torch.equal(argmax, ref_argmax)
36
+ torch.testing.assert_close(entropy, ref_entropy, atol=2e-3, rtol=1e-3)
37
+ torch.testing.assert_close(probs.float(), ref_probs.float(), atol=5e-3, rtol=2e-2)
38
+ assert probs.dtype == torch.bfloat16
39
+ assert probs.shape == logits.shape
40
+ # A padded (all-zero) row is uniform: max entropy, argmax 0.
41
+ zero = torch.zeros(1, vocab, device="cuda")
42
+ a0, _, e0, _ = sample_row_stats(zero, temps[:1] + 1.0, 1, 1, None)
43
+ assert a0.item() == 0
44
+ assert abs(e0.item() - torch.log(torch.tensor(float(vocab))).item()) < 1e-3
45
+
46
+
47
+def test_zero_temperature_is_greedy_and_probs_optional():
48
+ torch.manual_seed(0)
49
+ logits = torch.randn(12, 4097, device="cuda")
50
+ temps = torch.zeros(12, device="cuda")
51
+ argmax, sample, entropy, probs = sample_row_stats(logits, temps, 1, 7, None)
52
+ assert torch.equal(sample, argmax)
53
+ assert probs is None
54
+ assert torch.equal(argmax, logits.argmax(dim=-1))
55
+ # Greedy entropy is that of the reference's clamped temperature.
56
+ _, _, ref_entropy, _ = sample_row_stats_reference(logits, temps, 1, None)
57
+ torch.testing.assert_close(entropy, ref_entropy, atol=2e-3, rtol=1e-3)
58
+
59
+
60
+def test_positive_temperature_samples_from_the_distribution():
61
+ torch.manual_seed(0)
62
+ rows, vocab = 2000, 4099
63
+ # A peaked row samples its argmax almost always and a flat row rarely.
64
+ peaked = torch.zeros(rows, vocab, device="cuda")
65
+ peaked[:, 17] = 16.0 # p(17) = e^16 / (e^16 + 4098) > 0.999
66
+ flat = torch.zeros(rows, vocab, device="cuda")
67
+ temps = torch.ones(rows, device="cuda")
68
+ _, s_peaked, _, _ = sample_row_stats(peaked, temps, 1, 99, None)
69
+ _, s_flat, _, _ = sample_row_stats(flat, temps, 1, 99, None)
70
+ assert (s_peaked == 17).float().mean().item() > 0.99
71
+ assert s_flat.unique().numel() > rows // 2
72
+ # The same seed repeats the draws and another seed changes them.
73
+ _, s_again, _, _ = sample_row_stats(flat, temps, 1, 99, None)
74
+ _, s_other, _, _ = sample_row_stats(flat, temps, 1, 100, None)
75
+ assert torch.equal(s_flat, s_again)
76
+ assert not torch.equal(s_flat, s_other)
77
+
78
+
79
+def test_empty_batch():
80
+ logits = torch.empty(0, 128, device="cuda")
81
+ temps = torch.empty(0, device="cuda")
82
+ argmax, sample, entropy, probs = sample_row_stats(
83
+ logits, temps, 1, 1, torch.bfloat16
84
+ )
85
+ assert argmax.numel() == sample.numel() == entropy.numel() == 0
86
+ assert probs is not None and probs.shape == (0, 128)
87
+
88
+
89
+def test_masked_logits_keep_a_finite_entropy():
90
+ """top_k/top_p leave -inf logits. Masked columns gave NaN entropy, and a
91
+ row with one live column has zero entropy."""
92
+ logits = torch.full((4, 5000), float("-inf"), device="cuda")
93
+ logits[:, 3] = 0.0
94
+ logits[2, 9] = 0.0 # two live columns: entropy log 2
95
+ temps = torch.tensor([1.0, 0.0, 1.0, 0.7], device="cuda")
96
+ argmax, sample, entropy, probs = sample_row_stats(
97
+ logits, temps, 1, 5, torch.float32
98
+ )
99
+ assert torch.isfinite(entropy).all()
100
+ assert entropy[0].item() < 1e-6 and entropy[1].item() < 1e-6
101
+ assert abs(entropy[2].item() - math.log(2.0)) < 1e-4
102
+ assert torch.equal(argmax[[0, 1, 3]], torch.tensor([3, 3, 3], device="cuda"))
103
+ assert probs[0, 3].item() == 1.0 and probs[0].sum().item() == 1.0
104
+ _, _, ref_entropy, _ = sample_row_stats_reference(logits, temps, 1, None)
105
+ assert torch.isfinite(ref_entropy).all()
106
+
107
+
108
+def test_leading_masked_block_keeps_a_finite_entropy():
109
+ """top_k leaves most of the vocabulary -inf, so whole leading blocks can be
110
+ masked before the first live logit. The running max was -inf through
111
+ them, and -inf - -inf poisoned the row with NaN."""
112
+ logits = torch.full((3, 9001), float("-inf"), device="cuda")
113
+ logits[:, 8500] = 0.0 # the only live column, past the first block
114
+ logits[1, 8600] = 0.0 # two live columns: entropy log 2
115
+ logits[2, 4] = -1.0 # a live column in the first block as well
116
+ temps = torch.tensor([1.0, 1.0, 0.7], device="cuda")
117
+ argmax, sample, entropy, probs = sample_row_stats(
118
+ logits, temps, 1, 5, torch.float32
119
+ )
120
+ _, _, ref_entropy, ref_probs = sample_row_stats_reference(
121
+ logits, temps, 1, torch.float32
122
+ )
123
+ assert torch.isfinite(entropy).all() and torch.isfinite(probs).all()
124
+ assert entropy[0].item() < 1e-6
125
+ assert abs(entropy[1].item() - math.log(2.0)) < 1e-4
126
+ torch.testing.assert_close(entropy, ref_entropy, atol=1e-4, rtol=1e-4)
127
+ torch.testing.assert_close(probs, ref_probs, atol=1e-5, rtol=1e-4)
128
+ assert argmax.tolist() == [8500, 8500, 8500]
129
+ assert sample[0].item() == 8500
vllm/model_executor/models/diffusion_gemma.py
+70
-53
@@ -37,6 +37,10 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
37
from vllm.model_executor.layers.vocab_parallel_embedding import (
38
ParallelLMHead,
39
)
40
+from vllm.model_executor.models.diffusion_gemma_sampler import (
41
+ sample_row_stats,
42
+ sample_row_stats_reference,
43
+)
44
from vllm.model_executor.models.gemma4 import Gemma4Model
45
from vllm.model_executor.models.gemma4_mm import (
46
Gemma4DummyInputsBuilder,
@@ -447,9 +451,27 @@ def _mask_rows_to_allowed(
451
# eager: about twice as slow for a self-conditioned step.
452
@torch._dynamo.config.patch(recompile_limit=64)
453
@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
454
+def _denoise_temperature(
455
+ step_tensor: torch.Tensor,
456
+ slots: torch.Tensor,
457
+ max_denoising_steps: float,
458
+ t_min: float,
459
+ t_max: float,
460
+) -> torch.Tensor:
461
+ """The schedule's temperature for each slot at its current step."""
462
+ steps_f = step_tensor[slots].float()
463
+ remaining = (max_denoising_steps - steps_f).clamp(min=1.0)
464
+ return t_min + (t_max - t_min) * (remaining / max_denoising_steps)
465
+
466
+
467
def _compiled_sample_step(
451
- # Logits from the model [num_decode * CL, vocab]
452
- logits: torch.Tensor,
468
+ # Per-position statistics of the temperature-scaled logits, from
469
+ # sample_row_stats: [num_decode, CL] each, and the softmax
470
+ # [num_decode, CL, vocab] in the embedding dtype when compute_sc.
471
+ new_tokens: torch.Tensor,
472
+ argmax_tokens: torch.Tensor,
473
+ token_entropy: torch.Tensor,
474
+ probs: torch.Tensor | None,
475
# Request mapping
476
decode_slots: torch.Tensor, # [num_decode] int64 → slot indices
477
decode_idx: torch.Tensor, # [num_decode] int64 → position in num_reqs
@@ -475,9 +497,6 @@ def _compiled_sample_step(
497
num_sampled: torch.Tensor, # [num_reqs]
498
draft_tokens: torch.Tensor, # [max_num_reqs, >=CL]
499
# Scalar config
478
- max_denoising_steps: float,
479
- t_min: float,
480
- t_max: float,
500
confidence_threshold: float,
501
vocab_size: int,
502
CL: int,
@@ -492,46 +511,15 @@ def _compiled_sample_step(
511
tp_size: int,
512
tp_group_name: str,
513
compute_sc: bool = True,
495
- # Constrained reads: ``logits`` has one column per allowed id and
496
- # ``embed_weight`` holds those ids' rows, so every argmax is an index
497
- # into ``allowed``. ``vocab_size`` stays the real vocabulary for renoise.
498
- allowed: torch.Tensor | None = None,
499
-) -> torch.Tensor:
500
- """Compiled decode step: temperature → Gumbel sample → probs/confidence →
501
- accept/renoise → convergence, all as vectorized PyTorch ops.
502
-
503
- Returns the temperature-scaled logits ``[num_decode, CL, vocab]`` so the
504
- caller can compute logprobs outside the compiled region."""
514
+) -> None:
515
+ """Compiled decode step: confidence → accept/renoise → convergence, as
516
+ vectorized PyTorch ops over [num_decode, CL] tensors. The per-position
517
+ statistics (argmax, Gumbel-max sample, entropy, softmax) come from one
518
+ pass over the logits in ``sample_row_stats``."""
519
num_decode = decode_slots.shape[0]
520
device = decode_slots.device
521
508
- # ---- Phase 1: Temperature schedule ----
509
- steps_f = step_tensor[decode_slots].float()
510
- remaining = (max_denoising_steps - steps_f).clamp(min=1.0)
511
- temp = t_min + (t_max - t_min) * (remaining / max_denoising_steps)
512
-
513
- # ---- Phase 2: Temperature scaling + Gumbel-max sampling ----
514
- logits_3d = logits.reshape(num_decode, CL, -1).float()
515
- scaled = logits_3d / temp[:, None, None].clamp(min=1e-10)
516
-
517
- # Gumbel-max trick: argmax(logits/T + Gumbel) ~ sample from softmax(logits/T)
518
- u = torch.rand_like(scaled).clamp(min=1e-20)
519
- gumbel = -torch.log(-torch.log(u))
520
- # Zero noise when temp==0 (greedy)
521
- noisy = scaled + gumbel * (temp[:, None, None] > 0).float()
522
- new_tokens = noisy.view(-1, noisy.shape[-1]).argmax(dim=-1).view(num_decode, CL)
523
- argmax_tokens = (
524
- scaled.view(-1, scaled.shape[-1]).argmax(dim=-1).view(num_decode, CL)
525
- )
526
- if allowed is not None:
527
- new_tokens = allowed[new_tokens]
528
- argmax_tokens = allowed[argmax_tokens]
529
-
530
- # ---- Phase 3: Probs, self-conditioning, confidence ----
531
- log_probs = scaled.log_softmax(dim=-1)
532
- probs = log_probs.exp()
533
-
534
- token_entropy = -(probs * log_probs).sum(dim=-1) # [num_decode, CL]
522
+ # ---- Phase 3: Confidence ----
523
# A canvas truncated near max_model_len is zero-padded up to CL by the
524
# caller; those padded rows are uniform (max entropy, argmax 0), so they
525
# never trigger early convergence and are stable, and only the real
@@ -626,6 +614,7 @@ def _compiled_sample_step(
614
# [.., vocab] probs avoids a giant persistent buffer.
615
sc_keep = (is_denoise & ~is_encoder_phase[decode_slots])[:, None, None]
616
if compute_sc:
617
+ assert probs is not None
618
# Self-conditioning soft embed = probs @ embed_tokens.weight. Under
619
# tensor parallelism the embedding is vocab-sharded ([vocab/tp,
620
# hidden]) while probs spans the full vocab, so each rank multiplies
@@ -664,8 +653,6 @@ def _compiled_sample_step(
653
# ---- Phase 7: Copy canvas → draft_tokens for all slots ----
654
draft_tokens[all_slots, :CL] = canvas[all_slots]
655
667
- return scaled
668
-
656
657
class DiffusionGemmaRequestStates:
658
"""Pre-allocated GPU tensors for DiffusionGemma per-request state.
@@ -1513,9 +1500,10 @@ class DiffusionSampler:
1500
run_end = run_start
1501
while run_end < num_decode and widths_np[order[run_end]] == W:
1502
run_end += 1
1516
- # ~10 transient fp32 copies of [tile * W, vocab] inside the step
1517
- # (eager peaks at ~8; pad for allocator overhead and small tensors).
1518
- budget = max(1, int(free * 0.5) // max(W * self.vocab_size * 4 * 10, 1))
1503
+ # Transient [tile * W, vocab] tensors: the softmax in the embedding
1504
+ # dtype for self-conditioning, plus the fp32 scaled logits when
1505
+ # logprobs are wanted (pad for allocator overhead).
1506
+ budget = max(1, int(free * 0.5) // max(W * self.vocab_size * 4 * 3, 1))
1507
for t0 in range(run_start, run_end, budget):
1508
sel_np = order[t0 : min(t0 + budget, run_end)]
1509
n = len(sel_np)
@@ -1547,8 +1535,40 @@ class DiffusionSampler:
1535
)
1536
)
1537
1550
- scaled = _compiled_sample_step(
1551
- tile_logits,
1538
+ temp = _denoise_temperature(
1539
+ states.step,
1540
+ tile_slots,
1541
+ float(states.max_denoising_steps),
1542
+ self.t_min,
1543
+ self.t_max,
1544
+ )
1545
+ probs_dtype = self.embed_weight.dtype if compute_sc else None
1546
+ if tile_logits.is_cuda:
1547
+ seed = int(torch.randint(0, 2**31 - 1, (1,)).item())
1548
+ stats = sample_row_stats(tile_logits, temp, W, seed, probs_dtype)
1549
+ else:
1550
+ stats = sample_row_stats_reference(
1551
+ tile_logits, temp, W, probs_dtype
1552
+ )
1553
+ argmax_rows, sample_rows, entropy_rows, probs_rows = stats
1554
+ if allowed is not None:
1555
+ # K-space picks back to token ids. The softmax stays K wide
1556
+ # for the K-row self-conditioning matmul.
1557
+ argmax_rows = allowed[argmax_rows]
1558
+ sample_rows = allowed[sample_rows]
1559
+ probs = probs_rows.view(n, W, -1) if probs_rows is not None else None
1560
+ # Only the logprob stash reads the scaled logits.
1561
+ scaled = None
1562
+ if want_logprobs:
1563
+ scaled = tile_logits.float().view(n, W, -1) / temp[
1564
+ :, None, None
1565
+ ].clamp(min=1e-10)
1566
+
1567
+ _compiled_sample_step(
1568
+ sample_rows.view(n, W),
1569
+ argmax_rows.view(n, W),
1570
+ entropy_rows.view(n, W),
1571
+ probs,
1572
tile_slots,
1573
decode_idx[sel],
1574
all_slots,
@@ -1573,9 +1593,6 @@ class DiffusionSampler:
1593
num_sampled,
1594
self.req_states.draft_tokens,
1595
# Config
1576
- max_denoising_steps=float(states.max_denoising_steps),
1577
- t_min=self.t_min,
1578
- t_max=self.t_max,
1596
confidence_threshold=self.confidence_threshold,
1597
vocab_size=vocab_size,
1598
CL=W,
@@ -1586,7 +1603,6 @@ class DiffusionSampler:
1603
tp_size=self.tp_size,
1604
tp_group_name=self.tp_group_name,
1605
compute_sc=compute_sc,
1589
- allowed=allowed,
1606
)
1607
1608
# Stash newly converged logprobs, including reads that emit now.
@@ -1596,6 +1612,7 @@ class DiffusionSampler:
1612
)
1613
just_converged = converged_mask & ~is_committing[sel]
1614
if just_converged.any():
1615
+ assert scaled is not None
1616
flat_logits = scaled.reshape(-1, scaled.shape[-1])
1617
argmax_tokens = scaled.argmax(dim=-1)
1618
raw_flat: torch.Tensor | None = None
vllm/model_executor/models/diffusion_gemma_sampler.py
new
+174
@@ -0,0 +1,174 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""One-pass row statistics for the DiffusionGemma denoise sampler.
4
+
5
+The step needs, for every canvas position, the argmax of the
6
+temperature-scaled logits, a Gumbel-max sample from them, the entropy of
7
+their softmax, and the softmax itself in the model dtype for the
8
+self-conditioning matmul. As PyTorch ops these are several passes over a
9
+``[rows, vocab]`` fp32 tensor plus a same-sized noise tensor. The Triton
10
+kernel reads each row once with an online max and sum, draws the noise
11
+inline, and writes the probabilities once.
12
+"""
13
+
14
+import torch
15
+
16
+from vllm.triton_utils import tl, triton
17
+
18
+
19
+@triton.jit
20
+def _row_stats_kernel(
21
+ logits_ptr,
22
+ logits_stride,
23
+ temps_ptr,
24
+ argmax_ptr,
25
+ sample_ptr,
26
+ entropy_ptr,
27
+ probs_ptr,
28
+ probs_stride,
29
+ seed,
30
+ V,
31
+ CL,
32
+ WRITE_PROBS: tl.constexpr,
33
+ BLOCK: tl.constexpr,
34
+):
35
+ row = tl.program_id(0)
36
+ temp = tl.load(temps_ptr + row // CL).to(tl.float32)
37
+ inv = 1.0 / tl.maximum(temp, 1e-10)
38
+ noise = (temp > 0).to(tl.float32)
39
+ base = logits_ptr + row.to(tl.int64) * logits_stride
40
+ rng_base = row.to(tl.int64) * V
41
+
42
+ # The running max stays in raw logit units and the temperature is applied
43
+ # to differences from it. Scaling the row first breaks at a zero
44
+ # temperature: the row is multiplied by 1e10, fp32 spacing there is about
45
+ # a thousand, and exp(x - max) becomes 0 or inf for the max element itself
46
+ # whenever the two sides round differently.
47
+ m = float("-inf") # running max of the raw row
48
+ z = 0.0 # sum exp(d), d = (x - m) * inv
49
+ s = 0.0 # sum d exp(d)
50
+ best = float("-inf")
51
+ best_idx = 0
52
+ best_noisy = float("-inf")
53
+ best_noisy_idx = 0
54
+ for start in range(0, V, BLOCK):
55
+ offs = start + tl.arange(0, BLOCK)
56
+ mask = offs < V
57
+ x = tl.load(base + offs, mask=mask, other=float("-inf")).to(tl.float32)
58
+
59
+ block_max = tl.max(x, 0)
60
+ new_m = tl.maximum(m, block_max)
61
+ # Until a block holds a finite logit the running max stays -inf, and
62
+ # top_k/top_p leave whole blocks masked. x - new_m would then be
63
+ # -inf - -inf = NaN and poison z for the rest of the row, so subtract
64
+ # 0 instead: every d is still -inf and every e 0.
65
+ safe_m = tl.where(new_m > float("-inf"), new_m, 0.0)
66
+ # Rescale the running sums to the new max. While m is -inf, z and s
67
+ # are 0, so alpha and the shift are zeroed rather than computed from
68
+ # -inf - -inf.
69
+ alpha = tl.where(m > float("-inf"), tl.exp((m - safe_m) * inv), 0.0)
70
+ shift = tl.where(z > 0, (m - safe_m) * inv * z, 0.0)
71
+ d = (x - safe_m) * inv
72
+ e = tl.exp(d)
73
+ z = z * alpha + tl.sum(e, 0)
74
+ # A -inf logit (top_k/top_p, or the row's padding) has d = -inf and
75
+ # e = 0, and -inf * 0 is NaN. The where drops it from the sum.
76
+ s = (s + shift) * alpha + tl.sum(tl.where(e > 0, d * e, 0.0), 0)
77
+ m = new_m
78
+
79
+ block_idx = tl.argmax(x, 0)
80
+ take = block_max > best
81
+ best_idx = tl.where(take, start + block_idx, best_idx)
82
+ best = tl.where(take, block_max, best)
83
+
84
+ u = tl.rand(seed, rng_base + offs)
85
+ u = tl.maximum(u, 1e-20)
86
+ gumbel = -tl.log(-tl.log(u))
87
+ noisy = tl.where(mask, x * inv + gumbel * noise, float("-inf"))
88
+ noisy_max = tl.max(noisy, 0)
89
+ noisy_idx = tl.argmax(noisy, 0)
90
+ take_n = noisy_max > best_noisy
91
+ best_noisy_idx = tl.where(take_n, start + noisy_idx, best_noisy_idx)
92
+ best_noisy = tl.where(take_n, noisy_max, best_noisy)
93
+
94
+ # With d = (x - m) * inv: p = e^d / z, log Z = log z in d units, and
95
+ # H = -sum p (d - log z) = log z - s / z.
96
+ # tl.store casts to the pointee dtype. mypy type-checks the kernel body
97
+ # as Python and rejects .to() on the int accumulators.
98
+ tl.store(argmax_ptr + row, best_idx)
99
+ tl.store(sample_ptr + row, best_noisy_idx)
100
+ tl.store(entropy_ptr + row, tl.log(z) - s / z)
101
+
102
+ if WRITE_PROBS:
103
+ out = probs_ptr + row.to(tl.int64) * probs_stride
104
+ for start in range(0, V, BLOCK):
105
+ offs = start + tl.arange(0, BLOCK)
106
+ mask = offs < V
107
+ x = tl.load(base + offs, mask=mask, other=float("-inf")).to(tl.float32)
108
+ p = tl.exp((x - m) * inv) / z
109
+ tl.store(out + offs, p.to(probs_ptr.dtype.element_ty), mask=mask)
110
+
111
+
112
+def sample_row_stats(
113
+ logits: torch.Tensor,
114
+ temps: torch.Tensor,
115
+ canvas_len: int,
116
+ seed: int,
117
+ probs_dtype: torch.dtype | None,
118
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
119
+ """Argmax, Gumbel-max sample, entropy and (optionally) softmax per row.
120
+
121
+ ``logits`` is ``[rows, vocab]`` and row ``i`` uses temperature
122
+ ``temps[i // canvas_len]``. A zero temperature is greedy: the sample is
123
+ the argmax, and the entropy is the reference's, which clamps the
124
+ temperature at 1e-10.
125
+ """
126
+ rows, vocab = logits.shape
127
+ device = logits.device
128
+ argmax = torch.empty(rows, dtype=torch.int64, device=device)
129
+ sample = torch.empty(rows, dtype=torch.int64, device=device)
130
+ entropy = torch.empty(rows, dtype=torch.float32, device=device)
131
+ if probs_dtype is not None:
132
+ probs = torch.empty(rows, vocab, dtype=probs_dtype, device=device)
133
+ probs_arg, probs_stride = probs, probs.stride(0)
134
+ else:
135
+ probs = None
136
+ probs_arg, probs_stride = entropy, 0
137
+ if rows == 0:
138
+ return argmax, sample, entropy, probs
139
+ _row_stats_kernel[(rows,)](
140
+ logits,
141
+ logits.stride(0),
142
+ temps,
143
+ argmax,
144
+ sample,
145
+ entropy,
146
+ probs_arg,
147
+ probs_stride,
148
+ seed,
149
+ vocab,
150
+ canvas_len,
151
+ WRITE_PROBS=probs is not None,
152
+ BLOCK=4096,
153
+ num_warps=8,
154
+ )
155
+ return argmax, sample, entropy, probs
156
+
157
+
158
+def sample_row_stats_reference(
159
+ logits: torch.Tensor,
160
+ temps: torch.Tensor,
161
+ canvas_len: int,
162
+ probs_dtype: torch.dtype | None,
163
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
164
+ """The PyTorch form of the same statistics, for tests and non-CUDA runs."""
165
+ temp = temps.repeat_interleave(canvas_len).float()[:, None]
166
+ scaled = logits.float() / temp.clamp(min=1e-10)
167
+ u = torch.rand_like(scaled).clamp(min=1e-20)
168
+ noisy = scaled + (-torch.log(-torch.log(u))) * (temp > 0).float()
169
+ log_probs = scaled.log_softmax(dim=-1)
170
+ probs = log_probs.exp()
171
+ # Masked (-inf) columns: 0 * -inf is NaN, so they are dropped from the sum.
172
+ entropy = -torch.where(probs > 0, probs * log_probs, 0.0).sum(dim=-1)
173
+ out_probs = probs.to(probs_dtype) if probs_dtype is not None else None
174
+ return scaled.argmax(dim=-1), noisy.argmax(dim=-1), entropy, out_probs