diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py
index 09b3a83223..769481b06b 100644
--- a/tests/v1/sample/test_diffusion_gemma_reads.py
+++ b/tests/v1/sample/test_diffusion_gemma_reads.py
@@ -12,7 +12,9 @@ from vllm.model_executor.models.diffusion_gemma import (
     DiffusionGemmaRequestStates,
     _compiled_sample_step,
     _concat_logprob_stashes,
+    _denoise_temperature,
     _mask_rows_to_allowed,
+    sample_row_stats_reference,
 )
 from vllm.platforms import current_platform
 from vllm.v1.outputs import LogprobsTensors
@@ -132,8 +134,18 @@ def _denoise_once(
     decode_idx = torch.arange(n, dtype=torch.int64, device=device)
     sampled = torch.zeros(n, CL, dtype=torch.int32, device=device)[:, :width]
     num_sampled = torch.zeros(n, dtype=torch.int32, device=device)
-    _compiled_sample_step(
+    temp = _denoise_temperature(states.step, decode_slots, float(MAX_STEPS), 0.5, 1.0)
+    argmax, sample, entropy, probs = sample_row_stats_reference(
         torch.zeros(n * width, VOCAB, device=device),
+        temp,
+        width,
+        embed_dtype if compute_sc else None,
+    )
+    _compiled_sample_step(
+        sample.view(n, width),
+        argmax.view(n, width),
+        entropy.view(n, width),
+        probs.view(n, width, -1) if probs is not None else None,
         decode_slots,
         decode_idx,
         decode_slots,
@@ -159,9 +171,6 @@ def _denoise_once(
         sampled,
         num_sampled,
         torch.zeros(MAX_REQS, CL, dtype=torch.int64, device=device),
-        max_denoising_steps=float(MAX_STEPS),
-        t_min=0.5,
-        t_max=1.0,
         confidence_threshold=0.1,
         vocab_size=VOCAB,
         CL=width,
diff --git a/tests/v1/sample/test_diffusion_sampler_stats.py b/tests/v1/sample/test_diffusion_sampler_stats.py
new file mode 100644
index 0000000000..589c28c0a6
--- /dev/null
+++ b/tests/v1/sample/test_diffusion_sampler_stats.py
@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The fused row statistics of the DiffusionGemma sampler against the PyTorch
+reference: argmax, entropy and softmax must match. A zero temperature must
+sample the argmax and a positive one must sample from the distribution."""
+
+import math
+
+import pytest
+import torch
+
+from vllm.model_executor.models.diffusion_gemma_sampler import (
+    sample_row_stats,
+    sample_row_stats_reference,
+)
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="the fused kernel needs a GPU"
+)
+
+
+@pytest.mark.parametrize("vocab", [5000, 9001])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_stats_match_reference(vocab: int, dtype: torch.dtype):
+    torch.manual_seed(0)
+    num_decode, cl = 9, 4
+    logits = (torch.randn(num_decode * cl, vocab, device="cuda") * 3).to(dtype)
+    temps = torch.tensor([0.0, 0.7, 2.0, 1.0, 0.0, 0.3, 1.5, 1.0, 0.0], device="cuda")
+    argmax, sample, entropy, probs = sample_row_stats(
+        logits, temps, cl, seed=1234, probs_dtype=torch.bfloat16
+    )
+    ref_argmax, _, ref_entropy, ref_probs = sample_row_stats_reference(
+        logits, temps, cl, probs_dtype=torch.bfloat16
+    )
+    assert torch.equal(argmax, ref_argmax)
+    torch.testing.assert_close(entropy, ref_entropy, atol=2e-3, rtol=1e-3)
+    torch.testing.assert_close(probs.float(), ref_probs.float(), atol=5e-3, rtol=2e-2)
+    assert probs.dtype == torch.bfloat16
+    assert probs.shape == logits.shape
+    # A padded (all-zero) row is uniform: max entropy, argmax 0.
+    zero = torch.zeros(1, vocab, device="cuda")
+    a0, _, e0, _ = sample_row_stats(zero, temps[:1] + 1.0, 1, 1, None)
+    assert a0.item() == 0
+    assert abs(e0.item() - torch.log(torch.tensor(float(vocab))).item()) < 1e-3
+
+
+def test_zero_temperature_is_greedy_and_probs_optional():
+    torch.manual_seed(0)
+    logits = torch.randn(12, 4097, device="cuda")
+    temps = torch.zeros(12, device="cuda")
+    argmax, sample, entropy, probs = sample_row_stats(logits, temps, 1, 7, None)
+    assert torch.equal(sample, argmax)
+    assert probs is None
+    assert torch.equal(argmax, logits.argmax(dim=-1))
+    # Greedy entropy is that of the reference's clamped temperature.
+    _, _, ref_entropy, _ = sample_row_stats_reference(logits, temps, 1, None)
+    torch.testing.assert_close(entropy, ref_entropy, atol=2e-3, rtol=1e-3)
+
+
+def test_positive_temperature_samples_from_the_distribution():
+    torch.manual_seed(0)
+    rows, vocab = 2000, 4099
+    # A peaked row samples its argmax almost always and a flat row rarely.
+    peaked = torch.zeros(rows, vocab, device="cuda")
+    peaked[:, 17] = 16.0  # p(17) = e^16 / (e^16 + 4098) > 0.999
+    flat = torch.zeros(rows, vocab, device="cuda")
+    temps = torch.ones(rows, device="cuda")
+    _, s_peaked, _, _ = sample_row_stats(peaked, temps, 1, 99, None)
+    _, s_flat, _, _ = sample_row_stats(flat, temps, 1, 99, None)
+    assert (s_peaked == 17).float().mean().item() > 0.99
+    assert s_flat.unique().numel() > rows // 2
+    # The same seed repeats the draws and another seed changes them.
+    _, s_again, _, _ = sample_row_stats(flat, temps, 1, 99, None)
+    _, s_other, _, _ = sample_row_stats(flat, temps, 1, 100, None)
+    assert torch.equal(s_flat, s_again)
+    assert not torch.equal(s_flat, s_other)
+
+
+def test_empty_batch():
+    logits = torch.empty(0, 128, device="cuda")
+    temps = torch.empty(0, device="cuda")
+    argmax, sample, entropy, probs = sample_row_stats(
+        logits, temps, 1, 1, torch.bfloat16
+    )
+    assert argmax.numel() == sample.numel() == entropy.numel() == 0
+    assert probs is not None and probs.shape == (0, 128)
+
+
+def test_masked_logits_keep_a_finite_entropy():
+    """top_k/top_p leave -inf logits. Masked columns gave NaN entropy, and a
+    row with one live column has zero entropy."""
+    logits = torch.full((4, 5000), float("-inf"), device="cuda")
+    logits[:, 3] = 0.0
+    logits[2, 9] = 0.0  # two live columns: entropy log 2
+    temps = torch.tensor([1.0, 0.0, 1.0, 0.7], device="cuda")
+    argmax, sample, entropy, probs = sample_row_stats(
+        logits, temps, 1, 5, torch.float32
+    )
+    assert torch.isfinite(entropy).all()
+    assert entropy[0].item() < 1e-6 and entropy[1].item() < 1e-6
+    assert abs(entropy[2].item() - math.log(2.0)) < 1e-4
+    assert torch.equal(argmax[[0, 1, 3]], torch.tensor([3, 3, 3], device="cuda"))
+    assert probs[0, 3].item() == 1.0 and probs[0].sum().item() == 1.0
+    _, _, ref_entropy, _ = sample_row_stats_reference(logits, temps, 1, None)
+    assert torch.isfinite(ref_entropy).all()
+
+
+def test_leading_masked_block_keeps_a_finite_entropy():
+    """top_k leaves most of the vocabulary -inf, so whole leading blocks can be
+    masked before the first live logit. The running max was -inf through
+    them, and -inf - -inf poisoned the row with NaN."""
+    logits = torch.full((3, 9001), float("-inf"), device="cuda")
+    logits[:, 8500] = 0.0  # the only live column, past the first block
+    logits[1, 8600] = 0.0  # two live columns: entropy log 2
+    logits[2, 4] = -1.0  # a live column in the first block as well
+    temps = torch.tensor([1.0, 1.0, 0.7], device="cuda")
+    argmax, sample, entropy, probs = sample_row_stats(
+        logits, temps, 1, 5, torch.float32
+    )
+    _, _, ref_entropy, ref_probs = sample_row_stats_reference(
+        logits, temps, 1, torch.float32
+    )
+    assert torch.isfinite(entropy).all() and torch.isfinite(probs).all()
+    assert entropy[0].item() < 1e-6
+    assert abs(entropy[1].item() - math.log(2.0)) < 1e-4
+    torch.testing.assert_close(entropy, ref_entropy, atol=1e-4, rtol=1e-4)
+    torch.testing.assert_close(probs, ref_probs, atol=1e-5, rtol=1e-4)
+    assert argmax.tolist() == [8500, 8500, 8500]
+    assert sample[0].item() == 8500
diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py
index a700107ddb..1ae1cc7116 100644
--- a/vllm/model_executor/models/diffusion_gemma.py
+++ b/vllm/model_executor/models/diffusion_gemma.py
@@ -37,6 +37,10 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
 )
+from vllm.model_executor.models.diffusion_gemma_sampler import (
+    sample_row_stats,
+    sample_row_stats_reference,
+)
 from vllm.model_executor.models.gemma4 import Gemma4Model
 from vllm.model_executor.models.gemma4_mm import (
     Gemma4DummyInputsBuilder,
@@ -447,9 +451,27 @@ def _mask_rows_to_allowed(
 # eager: about twice as slow for a self-conditioned step.
 @torch._dynamo.config.patch(recompile_limit=64)
 @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
+def _denoise_temperature(
+    step_tensor: torch.Tensor,
+    slots: torch.Tensor,
+    max_denoising_steps: float,
+    t_min: float,
+    t_max: float,
+) -> torch.Tensor:
+    """The schedule's temperature for each slot at its current step."""
+    steps_f = step_tensor[slots].float()
+    remaining = (max_denoising_steps - steps_f).clamp(min=1.0)
+    return t_min + (t_max - t_min) * (remaining / max_denoising_steps)
+
+
 def _compiled_sample_step(
-    # Logits from the model [num_decode * CL, vocab]
-    logits: torch.Tensor,
+    # Per-position statistics of the temperature-scaled logits, from
+    # sample_row_stats: [num_decode, CL] each, and the softmax
+    # [num_decode, CL, vocab] in the embedding dtype when compute_sc.
+    new_tokens: torch.Tensor,
+    argmax_tokens: torch.Tensor,
+    token_entropy: torch.Tensor,
+    probs: torch.Tensor | None,
     # Request mapping
     decode_slots: torch.Tensor,  # [num_decode] int64 → slot indices
     decode_idx: torch.Tensor,  # [num_decode] int64 → position in num_reqs
@@ -475,9 +497,6 @@ def _compiled_sample_step(
     num_sampled: torch.Tensor,  # [num_reqs]
     draft_tokens: torch.Tensor,  # [max_num_reqs, >=CL]
     # Scalar config
-    max_denoising_steps: float,
-    t_min: float,
-    t_max: float,
     confidence_threshold: float,
     vocab_size: int,
     CL: int,
@@ -492,46 +511,15 @@ def _compiled_sample_step(
     tp_size: int,
     tp_group_name: str,
     compute_sc: bool = True,
-    # Constrained reads: ``logits`` has one column per allowed id and
-    # ``embed_weight`` holds those ids' rows, so every argmax is an index
-    # into ``allowed``. ``vocab_size`` stays the real vocabulary for renoise.
-    allowed: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """Compiled decode step: temperature → Gumbel sample → probs/confidence →
-    accept/renoise → convergence, all as vectorized PyTorch ops.
-
-    Returns the temperature-scaled logits ``[num_decode, CL, vocab]`` so the
-    caller can compute logprobs outside the compiled region."""
+) -> None:
+    """Compiled decode step: confidence → accept/renoise → convergence, as
+    vectorized PyTorch ops over [num_decode, CL] tensors. The per-position
+    statistics (argmax, Gumbel-max sample, entropy, softmax) come from one
+    pass over the logits in ``sample_row_stats``."""
     num_decode = decode_slots.shape[0]
     device = decode_slots.device
 
-    # ---- Phase 1: Temperature schedule ----
-    steps_f = step_tensor[decode_slots].float()
-    remaining = (max_denoising_steps - steps_f).clamp(min=1.0)
-    temp = t_min + (t_max - t_min) * (remaining / max_denoising_steps)
-
-    # ---- Phase 2: Temperature scaling + Gumbel-max sampling ----
-    logits_3d = logits.reshape(num_decode, CL, -1).float()
-    scaled = logits_3d / temp[:, None, None].clamp(min=1e-10)
-
-    # Gumbel-max trick: argmax(logits/T + Gumbel) ~ sample from softmax(logits/T)
-    u = torch.rand_like(scaled).clamp(min=1e-20)
-    gumbel = -torch.log(-torch.log(u))
-    # Zero noise when temp==0 (greedy)
-    noisy = scaled + gumbel * (temp[:, None, None] > 0).float()
-    new_tokens = noisy.view(-1, noisy.shape[-1]).argmax(dim=-1).view(num_decode, CL)
-    argmax_tokens = (
-        scaled.view(-1, scaled.shape[-1]).argmax(dim=-1).view(num_decode, CL)
-    )
-    if allowed is not None:
-        new_tokens = allowed[new_tokens]
-        argmax_tokens = allowed[argmax_tokens]
-
-    # ---- Phase 3: Probs, self-conditioning, confidence ----
-    log_probs = scaled.log_softmax(dim=-1)
-    probs = log_probs.exp()
-
-    token_entropy = -(probs * log_probs).sum(dim=-1)  # [num_decode, CL]
+    # ---- Phase 3: Confidence ----
     # A canvas truncated near max_model_len is zero-padded up to CL by the
     # caller; those padded rows are uniform (max entropy, argmax 0), so they
     # never trigger early convergence and are stable, and only the real
@@ -626,6 +614,7 @@ def _compiled_sample_step(
     # [.., vocab] probs avoids a giant persistent buffer.
     sc_keep = (is_denoise & ~is_encoder_phase[decode_slots])[:, None, None]
     if compute_sc:
+        assert probs is not None
         # Self-conditioning soft embed = probs @ embed_tokens.weight. Under
         # tensor parallelism the embedding is vocab-sharded ([vocab/tp,
         # hidden]) while probs spans the full vocab, so each rank multiplies
@@ -664,8 +653,6 @@ def _compiled_sample_step(
     # ---- Phase 7: Copy canvas → draft_tokens for all slots ----
     draft_tokens[all_slots, :CL] = canvas[all_slots]
 
-    return scaled
-
 
 class DiffusionGemmaRequestStates:
     """Pre-allocated GPU tensors for DiffusionGemma per-request state.
@@ -1513,9 +1500,10 @@ class DiffusionSampler:
             run_end = run_start
             while run_end < num_decode and widths_np[order[run_end]] == W:
                 run_end += 1
-            # ~10 transient fp32 copies of [tile * W, vocab] inside the step
-            # (eager peaks at ~8; pad for allocator overhead and small tensors).
-            budget = max(1, int(free * 0.5) // max(W * self.vocab_size * 4 * 10, 1))
+            # Transient [tile * W, vocab] tensors: the softmax in the embedding
+            # dtype for self-conditioning, plus the fp32 scaled logits when
+            # logprobs are wanted (pad for allocator overhead).
+            budget = max(1, int(free * 0.5) // max(W * self.vocab_size * 4 * 3, 1))
             for t0 in range(run_start, run_end, budget):
                 sel_np = order[t0 : min(t0 + budget, run_end)]
                 n = len(sel_np)
@@ -1547,8 +1535,40 @@ class DiffusionSampler:
                     )
                 )
 
-                scaled = _compiled_sample_step(
-                    tile_logits,
+                temp = _denoise_temperature(
+                    states.step,
+                    tile_slots,
+                    float(states.max_denoising_steps),
+                    self.t_min,
+                    self.t_max,
+                )
+                probs_dtype = self.embed_weight.dtype if compute_sc else None
+                if tile_logits.is_cuda:
+                    seed = int(torch.randint(0, 2**31 - 1, (1,)).item())
+                    stats = sample_row_stats(tile_logits, temp, W, seed, probs_dtype)
+                else:
+                    stats = sample_row_stats_reference(
+                        tile_logits, temp, W, probs_dtype
+                    )
+                argmax_rows, sample_rows, entropy_rows, probs_rows = stats
+                if allowed is not None:
+                    # K-space picks back to token ids. The softmax stays K wide
+                    # for the K-row self-conditioning matmul.
+                    argmax_rows = allowed[argmax_rows]
+                    sample_rows = allowed[sample_rows]
+                probs = probs_rows.view(n, W, -1) if probs_rows is not None else None
+                # Only the logprob stash reads the scaled logits.
+                scaled = None
+                if want_logprobs:
+                    scaled = tile_logits.float().view(n, W, -1) / temp[
+                        :, None, None
+                    ].clamp(min=1e-10)
+
+                _compiled_sample_step(
+                    sample_rows.view(n, W),
+                    argmax_rows.view(n, W),
+                    entropy_rows.view(n, W),
+                    probs,
                     tile_slots,
                     decode_idx[sel],
                     all_slots,
@@ -1573,9 +1593,6 @@ class DiffusionSampler:
                     num_sampled,
                     self.req_states.draft_tokens,
                     # Config
-                    max_denoising_steps=float(states.max_denoising_steps),
-                    t_min=self.t_min,
-                    t_max=self.t_max,
                     confidence_threshold=self.confidence_threshold,
                     vocab_size=vocab_size,
                     CL=W,
@@ -1586,7 +1603,6 @@ class DiffusionSampler:
                     tp_size=self.tp_size,
                     tp_group_name=self.tp_group_name,
                     compute_sc=compute_sc,
-                    allowed=allowed,
                 )
 
                 # Stash newly converged logprobs, including reads that emit now.
@@ -1596,6 +1612,7 @@ class DiffusionSampler:
                     )
                     just_converged = converged_mask & ~is_committing[sel]
                     if just_converged.any():
+                        assert scaled is not None
                         flat_logits = scaled.reshape(-1, scaled.shape[-1])
                         argmax_tokens = scaled.argmax(dim=-1)
                         raw_flat: torch.Tensor | None = None
diff --git a/vllm/model_executor/models/diffusion_gemma_sampler.py b/vllm/model_executor/models/diffusion_gemma_sampler.py
new file mode 100644
index 0000000000..9b15233791
--- /dev/null
+++ b/vllm/model_executor/models/diffusion_gemma_sampler.py
@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""One-pass row statistics for the DiffusionGemma denoise sampler.
+
+The step needs, for every canvas position, the argmax of the
+temperature-scaled logits, a Gumbel-max sample from them, the entropy of
+their softmax, and the softmax itself in the model dtype for the
+self-conditioning matmul. As PyTorch ops these are several passes over a
+``[rows, vocab]`` fp32 tensor plus a same-sized noise tensor. The Triton
+kernel reads each row once with an online max and sum, draws the noise
+inline, and writes the probabilities once.
+"""
+
+import torch
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _row_stats_kernel(
+    logits_ptr,
+    logits_stride,
+    temps_ptr,
+    argmax_ptr,
+    sample_ptr,
+    entropy_ptr,
+    probs_ptr,
+    probs_stride,
+    seed,
+    V,
+    CL,
+    WRITE_PROBS: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0)
+    temp = tl.load(temps_ptr + row // CL).to(tl.float32)
+    inv = 1.0 / tl.maximum(temp, 1e-10)
+    noise = (temp > 0).to(tl.float32)
+    base = logits_ptr + row.to(tl.int64) * logits_stride
+    rng_base = row.to(tl.int64) * V
+
+    # The running max stays in raw logit units and the temperature is applied
+    # to differences from it. Scaling the row first breaks at a zero
+    # temperature: the row is multiplied by 1e10, fp32 spacing there is about
+    # a thousand, and exp(x - max) becomes 0 or inf for the max element itself
+    # whenever the two sides round differently.
+    m = float("-inf")  # running max of the raw row
+    z = 0.0  # sum exp(d), d = (x - m) * inv
+    s = 0.0  # sum d exp(d)
+    best = float("-inf")
+    best_idx = 0
+    best_noisy = float("-inf")
+    best_noisy_idx = 0
+    for start in range(0, V, BLOCK):
+        offs = start + tl.arange(0, BLOCK)
+        mask = offs < V
+        x = tl.load(base + offs, mask=mask, other=float("-inf")).to(tl.float32)
+
+        block_max = tl.max(x, 0)
+        new_m = tl.maximum(m, block_max)
+        # Until a block holds a finite logit the running max stays -inf, and
+        # top_k/top_p leave whole blocks masked. x - new_m would then be
+        # -inf - -inf = NaN and poison z for the rest of the row, so subtract
+        # 0 instead: every d is still -inf and every e 0.
+        safe_m = tl.where(new_m > float("-inf"), new_m, 0.0)
+        # Rescale the running sums to the new max. While m is -inf, z and s
+        # are 0, so alpha and the shift are zeroed rather than computed from
+        # -inf - -inf.
+        alpha = tl.where(m > float("-inf"), tl.exp((m - safe_m) * inv), 0.0)
+        shift = tl.where(z > 0, (m - safe_m) * inv * z, 0.0)
+        d = (x - safe_m) * inv
+        e = tl.exp(d)
+        z = z * alpha + tl.sum(e, 0)
+        # A -inf logit (top_k/top_p, or the row's padding) has d = -inf and
+        # e = 0, and -inf * 0 is NaN. The where drops it from the sum.
+        s = (s + shift) * alpha + tl.sum(tl.where(e > 0, d * e, 0.0), 0)
+        m = new_m
+
+        block_idx = tl.argmax(x, 0)
+        take = block_max > best
+        best_idx = tl.where(take, start + block_idx, best_idx)
+        best = tl.where(take, block_max, best)
+
+        u = tl.rand(seed, rng_base + offs)
+        u = tl.maximum(u, 1e-20)
+        gumbel = -tl.log(-tl.log(u))
+        noisy = tl.where(mask, x * inv + gumbel * noise, float("-inf"))
+        noisy_max = tl.max(noisy, 0)
+        noisy_idx = tl.argmax(noisy, 0)
+        take_n = noisy_max > best_noisy
+        best_noisy_idx = tl.where(take_n, start + noisy_idx, best_noisy_idx)
+        best_noisy = tl.where(take_n, noisy_max, best_noisy)
+
+    # With d = (x - m) * inv: p = e^d / z, log Z = log z in d units, and
+    # H = -sum p (d - log z) = log z - s / z.
+    # tl.store casts to the pointee dtype. mypy type-checks the kernel body
+    # as Python and rejects .to() on the int accumulators.
+    tl.store(argmax_ptr + row, best_idx)
+    tl.store(sample_ptr + row, best_noisy_idx)
+    tl.store(entropy_ptr + row, tl.log(z) - s / z)
+
+    if WRITE_PROBS:
+        out = probs_ptr + row.to(tl.int64) * probs_stride
+        for start in range(0, V, BLOCK):
+            offs = start + tl.arange(0, BLOCK)
+            mask = offs < V
+            x = tl.load(base + offs, mask=mask, other=float("-inf")).to(tl.float32)
+            p = tl.exp((x - m) * inv) / z
+            tl.store(out + offs, p.to(probs_ptr.dtype.element_ty), mask=mask)
+
+
+def sample_row_stats(
+    logits: torch.Tensor,
+    temps: torch.Tensor,
+    canvas_len: int,
+    seed: int,
+    probs_dtype: torch.dtype | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """Argmax, Gumbel-max sample, entropy and (optionally) softmax per row.
+
+    ``logits`` is ``[rows, vocab]`` and row ``i`` uses temperature
+    ``temps[i // canvas_len]``. A zero temperature is greedy: the sample is
+    the argmax, and the entropy is the reference's, which clamps the
+    temperature at 1e-10.
+    """
+    rows, vocab = logits.shape
+    device = logits.device
+    argmax = torch.empty(rows, dtype=torch.int64, device=device)
+    sample = torch.empty(rows, dtype=torch.int64, device=device)
+    entropy = torch.empty(rows, dtype=torch.float32, device=device)
+    if probs_dtype is not None:
+        probs = torch.empty(rows, vocab, dtype=probs_dtype, device=device)
+        probs_arg, probs_stride = probs, probs.stride(0)
+    else:
+        probs = None
+        probs_arg, probs_stride = entropy, 0
+    if rows == 0:
+        return argmax, sample, entropy, probs
+    _row_stats_kernel[(rows,)](
+        logits,
+        logits.stride(0),
+        temps,
+        argmax,
+        sample,
+        entropy,
+        probs_arg,
+        probs_stride,
+        seed,
+        vocab,
+        canvas_len,
+        WRITE_PROBS=probs is not None,
+        BLOCK=4096,
+        num_warps=8,
+    )
+    return argmax, sample, entropy, probs
+
+
+def sample_row_stats_reference(
+    logits: torch.Tensor,
+    temps: torch.Tensor,
+    canvas_len: int,
+    probs_dtype: torch.dtype | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """The PyTorch form of the same statistics, for tests and non-CUDA runs."""
+    temp = temps.repeat_interleave(canvas_len).float()[:, None]
+    scaled = logits.float() / temp.clamp(min=1e-10)
+    u = torch.rand_like(scaled).clamp(min=1e-20)
+    noisy = scaled + (-torch.log(-torch.log(u))) * (temp > 0).float()
+    log_probs = scaled.log_softmax(dim=-1)
+    probs = log_probs.exp()
+    # Masked (-inf) columns: 0 * -inf is NaN, so they are dropped from the sum.
+    entropy = -torch.where(probs > 0, probs * log_probs, 0.0).sum(dim=-1)
+    out_probs = probs.to(probs_dtype) if probs_dtype is not None else None
+    return scaled.argmax(dim=-1), noisy.argmax(dim=-1), entropy, out_probs