perf(spec-decode): fuse the unfiltered per-row-noise draw into two launches (ATOM_SPEC_BONUS_FUSED)

atom/model_ops/bonus_sample.py

@@ -0,0 +1,170 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+"""`ATOM_SPEC_BONUS_FUSED`: the unfiltered temperature sample in two vocab-tiled
+launches instead of five.
+
+`Sampler._temperature_sample` + `_apply_greedy_tokens` run, per call,
+`mixed_sample_outer_exponential` (one program per row walking the whole vocab),
+then `topk_select` (two launches), `<=` and `where`. Here:
+
+1. `_bonus_tile_kernel` [(row, vocab tile)]: the tile's max of x = logit / T and
+   its best race key exp(x - tile_max) / (E + eps) with the lowest id among equal
+   keys, plus the tile's raw-logit argmax (lowest id) for greedy rows.
+2. `_bonus_pick_kernel` [(row,)]: greedy rows (T <= the sampler's greedy
+   epsilon) take the global raw-logit argmax, lowest id on ties -- exactly
+   `topk_select(tie="low")`; sampling rows rescale every tile key to the global
+   max and take the best.
+
+The noise is the same `exponential_` draw `_temperature_sample` makes (same
+generator, same tensor shape), and the race is the same: argmax_i p_i / (E_i + eps)
+with p = softmax(logit / max(T, 1e-5)) (aiter's clamp), so the drawn token has
+exactly the same distribution. The fp32 rounding of the keys differs from the
+aiter kernel's per-thread running-max rescale, so on a near-exact tie of the two
+best keys (relative gap ~1e-7) the pick can differ; greedy rows are bit-identical.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+_BONUS_BLOCK_V = 4096
+
+
+@triton.jit
+def _bonus_tile_kernel(
+    logits_ptr,
+    stride_l,
+    exp_ptr,
+    stride_e,
+    temps_ptr,
+    eps,
+    part_m_ptr,  # [n_tiles, rows] fp32
+    part_key_ptr,  # [n_tiles, rows] fp32
+    part_idx_ptr,  # [n_tiles, rows] int32
+    part_gmax_ptr,  # [n_tiles, rows] fp32
+    part_gidx_ptr,  # [n_tiles, rows] int32
+    rows,
+    vocab,
+    BLOCK_V: tl.constexpr,
+):
+    r = tl.program_id(0)
+    t = tl.program_id(1)
+    offs = t * BLOCK_V + tl.arange(0, BLOCK_V)
+    mask = offs < vocab
+    x = tl.load(
+        logits_ptr + r.to(tl.int64) * stride_l + offs, mask=mask, other=-float("inf")
+    ).to(tl.float32)
+    temp = tl.load(temps_ptr + r).to(tl.float32)
+    inv_t = 1.0 / tl.maximum(temp, 1e-5)
+    xs = x * inv_t
+    m = tl.max(xs, axis=0)
+    e = tl.load(exp_ptr + r.to(tl.int64) * stride_e + offs, mask=mask, other=1.0) + eps
+    # A tile the required channel masked to -inf has m = -inf, and
+    # exp(-inf - -inf) is NaN. Shifting by 0 instead keys every id 0: the tile
+    # loses to any tile with a finite logit and its idx stays in range.
+    m_safe = tl.where(m == -float("inf"), 0.0, m)
+    key = tl.where(mask, tl.exp(xs - m_safe) / e, -1.0)
+    best = tl.max(key, axis=0)
+    idx = tl.min(tl.where((key == best) & mask, offs, vocab), axis=0)
+    gmax = tl.max(x, axis=0)
+    gidx = tl.min(tl.where((x == gmax) & mask, offs, vocab), axis=0)
+    o = t * rows + r
+    tl.store(part_m_ptr + o, m)
+    tl.store(part_key_ptr + o, best)
+    tl.store(part_idx_ptr + o, idx)
+    tl.store(part_gmax_ptr + o, gmax)
+    tl.store(part_gidx_ptr + o, gidx)
+
+
+@triton.jit
+def _bonus_pick_kernel(
+    out_ptr,  # [rows] int32
+    temps_ptr,
+    greedy_eps,
+    part_m_ptr,
+    part_key_ptr,
+    part_idx_ptr,
+    part_gmax_ptr,
+    part_gidx_ptr,
+    rows,
+    n_tiles,
+    vocab,
+    BLOCK_T: tl.constexpr,
+):
+    r = tl.program_id(0)
+    ti = tl.arange(0, BLOCK_T)
+    tm = ti < n_tiles
+    o = ti * rows + r
+    temp = tl.load(temps_ptr + r).to(tl.float32)
+    if temp <= greedy_eps:
+        g = tl.load(part_gmax_ptr + o, mask=tm, other=-float("inf"))
+        gi = tl.load(part_gidx_ptr + o, mask=tm, other=vocab)
+        gm = tl.max(g, axis=0)
+        tok = tl.min(tl.where((g == gm) & tm, gi, vocab), axis=0)
+    else:
+        m = tl.load(part_m_ptr + o, mask=tm, other=-float("inf"))
+        k = tl.load(part_key_ptr + o, mask=tm, other=-1.0)
+        ki = tl.load(part_idx_ptr + o, mask=tm, other=vocab)
+        mm = tl.max(m, axis=0)
+        # An all -inf row would otherwise compute exp(-inf - -inf) here too.
+        mm = tl.where(mm == -float("inf"), 0.0, mm)
+        ks = tl.where(tm, k * tl.exp(m - mm), -1.0)
+        best = tl.max(ks, axis=0)
+        tok = tl.min(tl.where((ks == best) & tm, ki, vocab), axis=0)
+    tl.store(out_ptr + r, tok.to(tl.int32))
+
+
+def bonus_temperature_sample(
+    logits: torch.Tensor,  # [rows, V]
+    temperatures: torch.Tensor,  # [rows] fp32, pre-clamped to SAMPLER_EPS
+    exponential: torch.Tensor,  # [rows, V] fp32 Exp(1) noise (unit inner stride)
+    eps: float,
+    greedy_eps: float,
+) -> torch.Tensor:
+    rows, vocab = logits.shape
+    out = torch.empty(rows, dtype=torch.int, device=logits.device)
+    if rows == 0:
+        return out
+    assert logits.stride(1) == 1 and exponential.stride(1) == 1
+    n_tiles = triton.cdiv(vocab, _BONUS_BLOCK_V)
+    f32 = {"dtype": torch.float32, "device": logits.device}
+    i32 = {"dtype": torch.int32, "device": logits.device}
+    part_m = torch.empty((n_tiles, rows), **f32)
+    part_key = torch.empty((n_tiles, rows), **f32)
+    part_idx = torch.empty((n_tiles, rows), **i32)
+    part_gmax = torch.empty((n_tiles, rows), **f32)
+    part_gidx = torch.empty((n_tiles, rows), **i32)
+    temps = temperatures.contiguous()
+    _bonus_tile_kernel[(rows, n_tiles)](
+        logits,
+        logits.stride(0),
+        exponential,
+        exponential.stride(0),
+        temps,
+        eps,
+        part_m,
+        part_key,
+        part_idx,
+        part_gmax,
+        part_gidx,
+        rows,
+        vocab,
+        BLOCK_V=_BONUS_BLOCK_V,
+        num_warps=8,
+    )
+    _bonus_pick_kernel[(rows,)](
+        out,
+        temps,
+        greedy_eps,
+        part_m,
+        part_key,
+        part_idx,
+        part_gmax,
+        part_gidx,
+        rows,
+        n_tiles,
+        vocab,
+        BLOCK_T=max(16, triton.next_power_of_2(n_tiles)),
+        num_warps=1,
+    )
+    return out

atom/model_ops/sampler.py

@@ -9,6 +9,8 @@
 from aiter.ops.triton.topk import topk
 from torch import nn
 
+from atom.utils import envs
+
 # Try to import aiter top-k/top-p sampling ops
 try:
     import aiter.ops.sampling  # noqa: F401
@@ -65,6 +67,13 @@ class Sampler(nn.Module):
     def __init__(self):
         super().__init__()
         self.eps = SAMPLER_EPS
+        # `ATOM_SPEC_BONUS_FUSED`: read once, so it must be set before the
+        # server starts. None keeps the aiter spelling in `forward`.
+        self._fused_temperature_sample = None
+        if envs.ATOM_SPEC_BONUS_FUSED:
+            from atom.model_ops.bonus_sample import bonus_temperature_sample
+
+            self._fused_temperature_sample = bonus_temperature_sample
 
     def sample_verification_tokens(
         self,
@@ -144,6 +153,16 @@ def forward(
 
         # No Top-K Top-P parameters, perform temperature-based sampling
         if not self._needs_filtering(top_ks, top_ps):
+            if self._fused_temperature_sample is not None and needs_independent_noise:
+                # The same noise draw as `_temperature_sample`, then the race
+                # and the greedy override in two vocab-tiled launches.
+                num_tokens, vocab_size = logits.shape
+                exponential = torch.empty(
+                    (num_tokens, vocab_size), dtype=torch.float, device=logits.device
+                ).exponential_(1)
+                return self._fused_temperature_sample(
+                    logits, temperatures, exponential, self.eps, SAMPLER_EPS
+                )
             sampled = self._temperature_sample(
                 logits, temperatures, needs_independent_noise=needs_independent_noise
             )

atom/utils/envs.py

@@ -1015,6 +1015,20 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float:
     "ATOM_ENABLE_RELAXED_MTP": lambda: (
         os.getenv("ATOM_ENABLE_RELAXED_MTP", "0").lower() == "1"
     ),
+    # "1": an unfiltered, per-row-noise draw (the speculative bonus token, and
+    # every draft-conditioned verification row of a non-greedy batch) runs as
+    # two vocab-tiled Triton launches (`model_ops/bonus_sample.py`) instead of
+    # `mixed_sample_outer_exponential` + `topk_select` + `<=` + `where` (five
+    # launches, the first with one program per row). It draws the same
+    # `exponential_` noise and runs the same race, argmax_i p_i / (E_i + eps),
+    # so the sample has the same distribution; greedy rows (temperature at or
+    # below the sampler's epsilon) are bit-identical, lowest id on ties.
+    # Sampling rows can differ from the aiter kernel only when the two best
+    # race keys sit within ~1e-7 relative of each other (fp32 rounding order).
+    # Rows with top-k/top-p, and batches without per-row noise, keep the aiter
+    # path. Read when the `Sampler` is built, so it must be set before the
+    # server starts.
+    "ATOM_SPEC_BONUS_FUSED": lambda: os.getenv("ATOM_SPEC_BONUS_FUSED", "0") == "1",
     # --- Atomesh ---
     # Build atomesh when installing ATOM from source.
     "ATOM_MESH_BUILD": lambda: os.getenv("ATOM_MESH_BUILD", "0") == "1",

docs/environment_variables.md

@@ -307,6 +307,7 @@ that ever changes, so a fusion left inert by an unrecognised layout says so.
 | Variable | Type | Default | Description |
 |----------|------|---------|-------------|
 | **ATOM_ENABLE_RELAXED_MTP** | bool | 0 (false) | Accept a draft token when it lands in the target's top 10 within 0.6 of the top logit, instead of requiring the argmax. Intended for quantized MTP heads, whose drafts are right about the region and wrong about the exact winner often enough that strict acceptance throws away usable tokens. Read once at `rejection_sampler` import, so it must be set before the server starts. |
+| **ATOM_SPEC_BONUS_FUSED** | bool | 0 (false) | Draw the unfiltered, per-row-noise speculative sample (the bonus token, and the verification rows of a non-greedy batch) in two vocab-tiled Triton launches (`model_ops/bonus_sample.py`) instead of `mixed_sample_outer_exponential` + `topk_select` + `<=` + `where` (five launches). Same `exponential_` noise and same race, argmax_i p_i / (E_i + eps), so the distribution is unchanged; greedy rows are bit-identical (lowest id on ties); a sampling row can differ from the aiter kernel only if its two best race keys are within ~1e-7 relative (fp32 rounding order). Rows with top-k/top-p, and batches without per-row noise, keep the aiter path. Read when the `Sampler` is built, so set it before the server starts. |
 
 ## Engram (DeepSeek-V4.1)
 

docs/model_ops_guide.md

@@ -496,6 +496,8 @@ def forward(self, logits, temperatures) -> sampled_tokens:
 - `greedy_sample()`: `aiter.ops.triton.topk.topk(logits, 1)`
 - `random_sample()`: `aiter.ops.triton.softmax.softmax(logits)` followed by exponential sampling and `topk`.
 
+**Fused unfiltered draw (`ATOM_SPEC_BONUS_FUSED=1`, default off).** For an unfiltered batch with per-row noise (the speculative bonus token, and the verification rows of a non-greedy batch), `bonus_temperature_sample` (`bonus_sample.py`) replaces `mixed_sample_outer_exponential` + `topk_select` + `<=` + `where` with two vocab-tiled launches: a `(row, tile)` kernel keeps each tile's best race key and raw-logit argmax, and a `(row,)` kernel combines the tiles. The noise is the same `exponential_` draw and the race is the same, so the distribution is unchanged; greedy rows are bit-identical (lowest id on ties); a sampling row can differ only when its two best race keys are within ~1e-7 relative. Rows with top-k/top-p keep the aiter path.
+
 ### Required channel (`required_channel.py`)
 
 Holds a request with a `RequiredChannel` (`atom/sampling_params.py`; today only
@@ -572,6 +574,7 @@ ATOM uses fused kernels to reduce memory traffic by combining multiple operation
 | `topK.py` | `rocm_aiter_topk_softmax`, `rocm_aiter_grouped_topk`, `init_aiter_topK_meta_data` |
 | `triton_topk_softmax.py` | `triton_topk_softmax` -- renormalized softmax top-k, bit-identical to AITER's `topk_softmax` |
 | `sampler.py` | `Sampler` — unified greedy/random sampling |
+| `bonus_sample.py` | `bonus_temperature_sample` — the unfiltered per-row-noise draw in two vocab-tiled Triton launches (`ATOM_SPEC_BONUS_FUSED`) |
 | `rejection_sampler.py` | `RejectionSampler` — speculative decoding rejection sampling |
 | `required_channel.py` | `constrain_logits`, `advance_states` — `tool_choice="required"` channel mask (Kimi-K3) |
 | `base_config.py` | `QuantizeMethodBase` abstract class |