From d73ef139e6a155778eee25ebd3f128fa4b62091c Mon Sep 17 00:00:00 2001 From: "Iluno, Chelsea" Date: Mon, 2 Nov 2026 16:38:33 +0000 Subject: [PATCH 1/9] perf(spec-decode): fuse the unfiltered per-row-noise draw into two launches (ATOM_SPEC_BONUS_FUSED) The speculative bonus token, and every draft-conditioned verification row of a non-greedy batch, is an unfiltered draw with per-row noise. It ran as mixed_sample_outer_exponential + topk_select + <= + where. With ATOM_SPEC_BONUS_FUSED=1 (default off) it runs as two vocab-tiled Triton launches: a (row, tile) kernel keeping each tile's best race key and raw-logit argmax, and a (row,) kernel combining the tiles. Same exponential_ noise, same race argmax_i p_i / (E_i + eps), so the distribution is unchanged; greedy rows are bit-identical (lowest id on ties). Rows with top-k/top-p and batches without per-row noise keep the aiter path. Tested on MI350X (gfx950) and on CPU under the Triton interpreter; Sampler routing checked against a fake aiter. --- atom/model_ops/bonus_sample.py | 165 ++++++++++++++++++++++++++++++ atom/model_ops/sampler.py | 19 ++++ atom/utils/envs.py | 14 +++ docs/environment_variables.md | 1 + docs/model_ops_guide.md | 3 + tests/test_bonus_sample.py | 109 ++++++++++++++++++++ tests/test_bonus_sample_wiring.py | 118 +++++++++++++++++++++ tests/triton_interpreter.py | 50 +++++++++ 8 files changed, 479 insertions(+) create mode 100644 atom/model_ops/bonus_sample.py create mode 100644 tests/test_bonus_sample.py create mode 100644 tests/test_bonus_sample_wiring.py create mode 100644 tests/triton_interpreter.py diff --git a/atom/model_ops/bonus_sample.py b/atom/model_ops/bonus_sample.py new file mode 100644 index 0000000000..f6b8c400c0 --- /dev/null +++ b/atom/model_ops/bonus_sample.py @@ -0,0 +1,165 @@ +# 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 + key = tl.where(mask, tl.exp(xs - m) / 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) + 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) + tok = tl.minimum(tok, vocab - 1) + 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 diff --git a/atom/model_ops/sampler.py b/atom/model_ops/sampler.py index c69cb618ff..04f643a961 100644 --- a/atom/model_ops/sampler.py +++ b/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 ) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 10f7d4349a..4624f75216 100644 --- a/atom/utils/envs.py +++ b/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", diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 57cdbaed64..3a6df9428a 100644 --- a/docs/environment_variables.md +++ b/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) diff --git a/docs/model_ops_guide.md b/docs/model_ops_guide.md index e5cb9dfa40..74d7de039c 100644 --- a/docs/model_ops_guide.md +++ b/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 | diff --git a/tests/test_bonus_sample.py b/tests/test_bonus_sample.py new file mode 100644 index 0000000000..a6b99f7b19 --- /dev/null +++ b/tests/test_bonus_sample.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""`ATOM_SPEC_BONUS_FUSED`: the two-launch draw against the race it replaces. + +Runs on CPU tensors under Triton's interpreter (`triton_interpreter`), so it +needs Triton but no GPU and no AITER. + +The reference is float64 `argmax_i softmax(x / max(T, 1e-5))_i / (E_i + eps)` on +the noise the kernel is given -- the definition of +`mixed_sample_outer_exponential` -- and, for a greedy row (T at or below the +epsilon), the lowest-id argmax of the raw logits (`topk_select(tie="low")`). +Greedy rows must match exactly, ties planted. A sampling row must match unless +the reference's two best keys are a near tie (relative gap < 1e-5), which the +fp32 rounding of either kernel may order either way. +""" + +import pytest +import torch +from triton_interpreter import in_the_interpreter, interpreter_test + +EPS = 1e-10 + +test_under_the_interpreter = interpreter_test(__file__) + + +def _bonus_sample(): + # Imported inside the child only: with the interpreter off, Triton would + # build compiled kernels that nothing here can launch on CPU tensors. + from atom.model_ops import bonus_sample + + return bonus_sample + + +def _reference_pick(logits_row, temp, noise_row): + """(token, relative gap between the two best keys) for one sampling row.""" + t = max(float(temp), 1e-5) + p = torch.softmax(logits_row.double() / t, dim=0) + key = p / (noise_row.double() + EPS) + top2 = torch.topk(key, 2) + gap = float((top2.values[0] - top2.values[1]) / top2.values[0]) + return int(top2.indices[0]), int(top2.indices[1]), gap + + +@in_the_interpreter +@pytest.mark.parametrize("seed,vocab", [(1, 163840), (2, 50000), (3, 4097), (4, 1000)]) +def test_matches_the_race_it_replaces(seed, vocab): + gen = torch.Generator().manual_seed(seed) + rows = 6 + logits = (torch.randn(rows, vocab, generator=gen) * 4).to(torch.bfloat16) + temps = torch.tensor([1.0, 0.7, EPS, 1.3, EPS, 0.05], dtype=torch.float32) + # Exact ties on the greedy rows: the lowest id must win. + logits[2, 777 % vocab] = logits[2].float().max() + logits[2, 12345 % vocab] = logits[2].float().max() + logits[4, 5] = 50.0 + logits[4, vocab - 1] = 50.0 + noise = torch.empty(rows, vocab).exponential_(1, generator=gen) + + out = _bonus_sample().bonus_temperature_sample(logits, temps, noise, EPS, EPS) + + assert out.dtype == torch.int32 and out.shape == (rows,) + for r in range(rows): + if temps[r] <= EPS: + expected = int(torch.nonzero(logits[r] == logits[r].max())[0]) + assert int(out[r]) == expected + continue + best, runner_up, gap = _reference_pick(logits[r], temps[r], noise[r]) + if int(out[r]) != best: + assert gap < 1e-5 and int(out[r]) == runner_up, (r, gap) + + +@in_the_interpreter +def test_every_row_uses_its_own_temperature_and_noise(): + vocab = 4097 + gen = torch.Generator().manual_seed(7) + logits = torch.randn(3, vocab, generator=gen).to(torch.bfloat16) + logits[:] = logits[0] # identical logits: only T and the noise differ + temps = torch.tensor([1.0, 1.0, EPS]) + noise = torch.empty(3, vocab).exponential_(1, generator=gen) + out = _bonus_sample().bonus_temperature_sample(logits, temps, noise, EPS, EPS) + assert int(out[2]) == int(logits[0].float().argmax()) + for r in (0, 1): + best, runner_up, gap = _reference_pick(logits[r], temps[r], noise[r]) + assert int(out[r]) == best or (gap < 1e-5 and int(out[r]) == runner_up) + + +@in_the_interpreter +def test_a_row_below_the_clamp_samples_like_the_clamped_temperature(): + # aiter clamps T to 1e-5 inside the race; a row just above the greedy + # epsilon is a sampling row whose draw is (almost surely) the argmax. + vocab = 4097 + gen = torch.Generator().manual_seed(11) + logits = torch.randn(1, vocab, generator=gen).to(torch.bfloat16) + noise = torch.empty(1, vocab).exponential_(1, generator=gen) + temps = torch.tensor([1e-7]) + out = _bonus_sample().bonus_temperature_sample(logits, temps, noise, EPS, EPS) + assert int(out[0]) == int(logits[0].float().argmax()) + + +@in_the_interpreter +def test_no_rows_returns_an_empty_result(): + out = _bonus_sample().bonus_temperature_sample( + torch.empty(0, 4097, dtype=torch.bfloat16), + torch.empty(0), + torch.empty(0, 4097), + EPS, + EPS, + ) + assert out.shape == (0,) and out.dtype == torch.int32 diff --git a/tests/test_bonus_sample_wiring.py b/tests/test_bonus_sample_wiring.py new file mode 100644 index 0000000000..d53cea6ae3 --- /dev/null +++ b/tests/test_bonus_sample_wiring.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""`ATOM_SPEC_BONUS_FUSED`: which `Sampler.forward` paths take the fused draw. + +The sampler needs AITER at import, which a CPU runner does not have, so it is +loaded here against a fake `aiter` under a private module name and the fakes are +removed again. Nothing runs a kernel: the aiter ops and the fused draw are +recording stand-ins. +""" + +import importlib.util +import pathlib +import sys +from unittest.mock import MagicMock + +import pytest +import torch + +pytest.importorskip("triton") + +SAMPLER_PATH = ( + pathlib.Path(__file__).resolve().parents[1] / "atom" / "model_ops" / "sampler.py" +) + + +@pytest.fixture +def load_sampler(monkeypatch): + for name in ( + "aiter", + "aiter.ops", + "aiter.ops.sampling", + "aiter.ops.triton", + "aiter.ops.triton.softmax", + "aiter.ops.triton.topk", + ): + monkeypatch.setitem(sys.modules, name, MagicMock()) + stub = sys.modules["aiter"] + stub.mixed_sample_outer_exponential = MagicMock(name="aiter_race") + stub.topk_select = MagicMock( + side_effect=lambda x, k, tie: ( + None, + torch.zeros(x.shape[0], 1, dtype=torch.int32), + ) + ) + + def load(fused: bool): + monkeypatch.setenv("ATOM_SPEC_BONUS_FUSED", "1" if fused else "0") + spec = importlib.util.spec_from_file_location( + "_sampler_under_test", SAMPLER_PATH + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module, stub + + return load + + +def _batch(rows=4, vocab=64): + return torch.randn(rows, vocab), torch.ones(rows) + + +def test_flag_off_keeps_the_aiter_path(load_sampler): + module, aiter = load_sampler(fused=False) + sampler = module.Sampler() + assert sampler._fused_temperature_sample is None + logits, temps = _batch() + sampler(logits, temps, needs_independent_noise=True) + aiter.mixed_sample_outer_exponential.assert_called_once() + + +def test_flag_on_routes_the_per_row_noise_draw(load_sampler, monkeypatch): + module, aiter = load_sampler(fused=True) + sampler = module.Sampler() + seen = [] + + def fake(logits, temps, noise, eps, greedy_eps): + seen.append((logits, temps, noise, eps, greedy_eps)) + return torch.full((logits.shape[0],), 7, dtype=torch.int32) + + sampler._fused_temperature_sample = fake + logits, temps = _batch() + out = sampler(logits, temps, needs_independent_noise=True) + + assert out.tolist() == [7] * 4 + aiter.mixed_sample_outer_exponential.assert_not_called() + ((_, _, noise, eps, greedy_eps),) = seen + assert noise.shape == logits.shape and noise.dtype == torch.float32 + assert eps == sampler.eps and greedy_eps == module.SAMPLER_EPS + + +@pytest.mark.parametrize( + "kwargs", + [ + {"needs_independent_noise": False}, + {"needs_independent_noise": True, "top_ks": 5}, + {"needs_independent_noise": True, "top_ps": 0.9}, + ], +) +def test_flag_on_leaves_other_batches_to_aiter(load_sampler, kwargs): + module, _ = load_sampler(fused=True) + sampler = module.Sampler() + sampler._fused_temperature_sample = MagicMock(name="fused") + sampler._topk_topp_sample = MagicMock( + return_value=torch.zeros(4, dtype=torch.int32) + ) + logits, temps = _batch() + sampler(logits, temps, **kwargs) + sampler._fused_temperature_sample.assert_not_called() + + +def test_all_greedy_batches_never_reach_the_fused_draw(load_sampler): + module, _ = load_sampler(fused=True) + sampler = module.Sampler() + sampler._fused_temperature_sample = MagicMock(name="fused") + logits, temps = _batch() + sampler(logits, temps, all_greedy=True, needs_independent_noise=True) + sampler._fused_temperature_sample.assert_not_called() diff --git a/tests/triton_interpreter.py b/tests/triton_interpreter.py new file mode 100644 index 0000000000..dd8283c63b --- /dev/null +++ b/tests/triton_interpreter.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: MIT +# Test helper: drive a Triton kernel on CPU tensors, without a GPU. +# +# Triton's interpreter has to be on before `triton.language` is imported (its +# `tl.max`, `tl.sum`, ... are decorated at import), so it cannot be switched on +# inside a process that has already imported Triton. A kernel test module +# therefore has two halves: with TRITON_INTERPRET=1 in the environment it holds +# the real tests; without it, it holds one test that runs the same file again in +# a child process that has the variable set: +# +# test_under_the_interpreter = interpreter_test(__file__) +# +# @in_the_interpreter +# def test_a_kernel(): ... + +import os +import subprocess +import sys + +import pytest + + +def interpreting() -> bool: + return os.environ.get("TRITON_INTERPRET") == "1" + + +in_the_interpreter = pytest.mark.skipif( + not interpreting(), + reason="runs in the child process of test_under_the_interpreter", +) + + +def interpreter_test(test_file: str, timeout: float = 1800.0): + """A test that re-runs `test_file` under Triton's CPU interpreter.""" + + @pytest.mark.skipif(interpreting(), reason="this is the child process") + def test_under_the_interpreter(): + pytest.importorskip("triton") + env = dict(os.environ, TRITON_INTERPRET="1") + done = subprocess.run( + [sys.executable, "-m", "pytest", test_file, "-q", "-p", "no:cacheprovider"], + env=env, + check=False, + capture_output=True, + text=True, + timeout=timeout, + ) + assert done.returncode == 0, done.stdout[-6000:] + done.stderr[-2000:] + + return test_under_the_interpreter From 43066047d62d9b6114668021aeece74af46b9bcf Mon Sep 17 00:00:00 2001 From: chelnnexy Date: Mon, 2 Nov 2026 10:57:05 -0600 Subject: [PATCH 2/9] test(bonus-sample): determinism + late-tile reduction coverage Add two interpreter tests to test_bonus_sample.py: - same (logits, temps, noise) returns the same token every call (the kernels hold no state of their own); - a dominating race key in the final vocab tile beats tile 0, exercising the per-tile running max and the (row,) kernel's k * exp(m - mm) rescale. --- tests/test_bonus_sample.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/test_bonus_sample.py b/tests/test_bonus_sample.py index a6b99f7b19..8e062b3080 100644 --- a/tests/test_bonus_sample.py +++ b/tests/test_bonus_sample.py @@ -107,3 +107,34 @@ def test_no_rows_returns_an_empty_result(): EPS, ) assert out.shape == (0,) and out.dtype == torch.int32 + + +@in_the_interpreter +def test_same_inputs_give_the_same_output(): + # The draw is in `Sampler.forward`; the kernels carry no state of their own, + # so the same (logits, temps, noise) must give the same token every call. + vocab = 50000 + gen = torch.Generator().manual_seed(3) + logits = (torch.randn(5, vocab, generator=gen) * 4).to(torch.bfloat16) + temps = torch.tensor([1.0, 0.7, 0.9, 1.2, 0.5]) + noise = torch.empty(5, vocab).exponential_(1, generator=gen) + f = _bonus_sample().bonus_temperature_sample + first = f(logits, temps, noise, EPS, EPS) + assert torch.equal(first, f(logits, temps, noise, EPS, EPS)) + + +@in_the_interpreter +def test_winner_in_a_late_tile_survives_the_reduction(): + # One dominating race key in the final vocab tile must beat tile 0: the + # per-tile running max and the `(row,)` kernel's k * exp(m - mm) rescale + # have to reconstruct the global key across tiles. + vocab = 163840 # 40 tiles of 4096 + gen = torch.Generator().manual_seed(5) + logits = (torch.randn(1, vocab, generator=gen) * 2).to(torch.bfloat16) + noise = torch.full((1, vocab), 10.0) + win = vocab - 3 + logits[0, win] = 60.0 + noise[0, win] = 1e-3 # huge key exp(0) / ~1e-3 against a flat field + temps = torch.tensor([1.0]) + out = _bonus_sample().bonus_temperature_sample(logits, temps, noise, EPS, EPS) + assert int(out[0]) == win From e34da58ddd3bad575a4f10aecaa498e5c142cf8c Mon Sep 17 00:00:00 2001 From: chelnnexy Date: Tue, 3 Nov 2026 16:09:13 +0000 Subject: [PATCH 3/9] fix(bonus-sample): guard all -inf tiles and drop the result clamp With tool_choice="required" the bonus and verification rows are masked to -inf everywhere but the forced id, so 39 of 40 vocab tiles have a tile max of -inf and exp(-inf - -inf) put NaN into the race keys. The forced id only won because tl.max drops NaN. Shift by 0 when the tile max (and the pick kernel's global max) is -inf, so no NaN reaches a reduction. Remove tl.minimum(tok, vocab - 1): with the guard it can no longer change a result, and it turned an escaped sentinel into a valid-looking token id. Add tests for forced-channel rows and for partly and fully masked rows. --- atom/model_ops/bonus_sample.py | 9 +++++-- tests/test_bonus_sample.py | 45 ++++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/atom/model_ops/bonus_sample.py b/atom/model_ops/bonus_sample.py index f6b8c400c0..f67fb8fffc 100644 --- a/atom/model_ops/bonus_sample.py +++ b/atom/model_ops/bonus_sample.py @@ -59,7 +59,11 @@ def _bonus_tile_kernel( 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 - key = tl.where(mask, tl.exp(xs - m) / e, -1.0) + # 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) @@ -102,10 +106,11 @@ def _bonus_pick_kernel( 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) - tok = tl.minimum(tok, vocab - 1) tl.store(out_ptr + r, tok.to(tl.int32)) diff --git a/tests/test_bonus_sample.py b/tests/test_bonus_sample.py index 8e062b3080..0adbed4f17 100644 --- a/tests/test_bonus_sample.py +++ b/tests/test_bonus_sample.py @@ -138,3 +138,48 @@ def test_winner_in_a_late_tile_survives_the_reduction(): temps = torch.tensor([1.0]) out = _bonus_sample().bonus_temperature_sample(logits, temps, noise, EPS, EPS) assert int(out[0]) == win + + +@in_the_interpreter +@pytest.mark.filterwarnings("error::RuntimeWarning") +@pytest.mark.parametrize("forced", [5, 100000, 163837]) +def test_a_forced_channel_row_samples_the_forced_id(forced): + # `constrain_logits` writes -inf over every id but the forced one, so all + # tiles but one are entirely -inf: their tile max is -inf and the race + # must not compute exp(-inf - -inf) = NaN on the way to the forced id. + vocab = 163840 + gen = torch.Generator().manual_seed(forced) + logits = torch.full((2, vocab), float("-inf")) + logits[:, forced] = 3.0 + temps = torch.tensor([1.0, EPS]) + noise = torch.empty(2, vocab).exponential_(1, generator=gen) + out = _bonus_sample().bonus_temperature_sample( + logits.to(torch.bfloat16), temps, noise, EPS, EPS + ) + assert out.tolist() == [forced, forced] + + +@in_the_interpreter +@pytest.mark.filterwarnings("error::RuntimeWarning") +def test_masked_ids_never_win_and_every_id_is_in_range(): + # Whole -inf tiles next to tiles with scattered -inf ids, and one row that + # is -inf throughout: no NaN, no -inf id drawn, and with no clamp on the + # result every id must already be in range. + vocab = 50000 + gen = torch.Generator().manual_seed(9) + logits = (torch.randn(4, vocab, generator=gen) * 4).to(torch.bfloat16) + logits[:, : 3 * 4096] = float("-inf") + banned = torch.randint(0, vocab, (256,), generator=gen) + logits[:, banned] = float("-inf") + logits[3] = float("-inf") + temps = torch.tensor([1.0, 0.7, EPS, 1.0]) + noise = torch.empty(4, vocab).exponential_(1, generator=gen) + out = _bonus_sample().bonus_temperature_sample(logits, temps, noise, EPS, EPS) + assert ((out >= 0) & (out < vocab)).all() + for r in range(3): + assert logits[r, int(out[r])] != float("-inf") + if temps[r] <= EPS: + assert int(out[r]) == int(torch.nonzero(logits[r] == logits[r].max())[0]) + continue + best, runner_up, gap = _reference_pick(logits[r], temps[r], noise[r]) + assert int(out[r]) == best or (gap < 1e-5 and int(out[r]) == runner_up) From 6a9bc8f80a97b134c874e3ce07774e28fbf5a366 Mon Sep 17 00:00:00 2001 From: chelnnexy Date: Tue, 3 Nov 2026 16:13:45 +0000 Subject: [PATCH 4/9] refactor(sampler): one helper for the per-row exponential draw _temperature_sample and the ATOM_SPEC_BONUS_FUSED branch of forward each spelled the (num_tokens, vocab_size) exponential_(1) draw out by hand. Both now call get_per_row_exponential, so the two paths cannot drift apart and keep consuming the RNG identically. --- atom/model_ops/sampler.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/atom/model_ops/sampler.py b/atom/model_ops/sampler.py index 04f643a961..42c5a231d6 100644 --- a/atom/model_ops/sampler.py +++ b/atom/model_ops/sampler.py @@ -62,6 +62,17 @@ def get_per_token_exponential(vocab_size: int, device) -> torch.Tensor: ) +def get_per_row_exponential(num_tokens: int, vocab_size: int, device) -> torch.Tensor: + """Fresh Exp(1) noise of shape (num_tokens, vocab_size), one row per token. + + The one draw both per-row-noise paths take (`_temperature_sample` and the + `ATOM_SPEC_BONUS_FUSED` draw), so they consume the RNG identically. + """ + return torch.empty( + (num_tokens, vocab_size), dtype=torch.float, device=device + ).exponential_(1) + + class Sampler(nn.Module): def __init__(self): @@ -157,9 +168,9 @@ def forward( # 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) + exponential = get_per_row_exponential( + num_tokens, vocab_size, logits.device + ) return self._fused_temperature_sample( logits, temperatures, exponential, self.eps, SAMPLER_EPS ) @@ -207,9 +218,7 @@ def _temperature_sample( num_tokens, vocab_size = logits.shape sampled_tokens = torch.empty(num_tokens, dtype=torch.int, device=logits.device) if needs_independent_noise: - exponential = torch.empty( - (num_tokens, vocab_size), dtype=torch.float, device=logits.device - ).exponential_(1) + exponential = get_per_row_exponential(num_tokens, vocab_size, logits.device) else: exponential = get_per_token_exponential(vocab_size, logits.device).expand( num_tokens, vocab_size From a8347e84b277e7bd3228ab3add5c1af81578135b Mon Sep 17 00:00:00 2001 From: chelnnexy Date: Tue, 3 Nov 2026 16:15:43 +0000 Subject: [PATCH 5/9] test(bonus-sample): run the wiring tests without Triton pytest.importorskip("triton") skipped the whole file on a CPU CI runner without Triton, although no test here launches a kernel. The fixture now fakes atom.model_ops.bonus_sample the way it already fakes aiter, so all of it runs on any CPU runner, and a new test checks that the flag wires Sampler to bonus_temperature_sample. --- tests/test_bonus_sample_wiring.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/tests/test_bonus_sample_wiring.py b/tests/test_bonus_sample_wiring.py index d53cea6ae3..31a870ce74 100644 --- a/tests/test_bonus_sample_wiring.py +++ b/tests/test_bonus_sample_wiring.py @@ -5,20 +5,20 @@ The sampler needs AITER at import, which a CPU runner does not have, so it is loaded here against a fake `aiter` under a private module name and the fakes are -removed again. Nothing runs a kernel: the aiter ops and the fused draw are -recording stand-ins. +removed again. `atom.model_ops.bonus_sample` (Triton) is faked the same way, so +these tests need neither AITER nor Triton and run on any CPU runner. Nothing +runs a kernel: the aiter ops and the fused draw are recording stand-ins. """ import importlib.util import pathlib import sys +import types from unittest.mock import MagicMock import pytest import torch -pytest.importorskip("triton") - SAMPLER_PATH = ( pathlib.Path(__file__).resolve().parents[1] / "atom" / "model_ops" / "sampler.py" ) @@ -43,6 +43,9 @@ def load_sampler(monkeypatch): torch.zeros(x.shape[0], 1, dtype=torch.int32), ) ) + fused = types.ModuleType("atom.model_ops.bonus_sample") + fused.bonus_temperature_sample = MagicMock(name="fused_draw") + monkeypatch.setitem(sys.modules, "atom.model_ops.bonus_sample", fused) def load(fused: bool): monkeypatch.setenv("ATOM_SPEC_BONUS_FUSED", "1" if fused else "0") @@ -53,6 +56,7 @@ def load(fused: bool): spec.loader.exec_module(module) return module, stub + load.fused_draw = fused.bonus_temperature_sample return load @@ -60,6 +64,11 @@ def _batch(rows=4, vocab=64): return torch.randn(rows, vocab), torch.ones(rows) +def test_flag_on_wires_the_fused_draw(load_sampler): + module, _ = load_sampler(fused=True) + assert module.Sampler()._fused_temperature_sample is load_sampler.fused_draw + + def test_flag_off_keeps_the_aiter_path(load_sampler): module, aiter = load_sampler(fused=False) sampler = module.Sampler() From c70d5bf7caa1a76307e869013f14dcb6c42560c7 Mon Sep 17 00:00:00 2001 From: chelnnexy Date: Tue, 3 Nov 2026 16:16:03 +0000 Subject: [PATCH 6/9] docs(bonus-sample): say why ATOM_SPEC_BONUS_FUSED defaults off Sampling rows are not bit-exact with the aiter kernel (fp32 near-ties), and on MI355X the gain shrinks with the verify-row count: about -25% up to 128 rows but none at 256 (about Kimi-K3 DSpark at concurrency 64). State that in the flag's comment and in environment_variables.md. --- atom/utils/envs.py | 5 ++++- docs/environment_variables.md | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 4624f75216..1a5f08ebbc 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -1027,7 +1027,10 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: # 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. + # server starts. Off by default: sampling rows are not bit-exact with the + # aiter kernel (above), and the gain shrinks with the verify-row count -- + # on MI355X about -25% at <= 128 rows but none at 256 (~Kimi-K3 DSpark at + # concurrency 64) -- so it stays opt-in until that shape gains too. "ATOM_SPEC_BONUS_FUSED": lambda: os.getenv("ATOM_SPEC_BONUS_FUSED", "0") == "1", # --- Atomesh --- # Build atomesh when installing ATOM from source. diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 3a6df9428a..819c25bc24 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -307,7 +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. | +| **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. Off by default because sampling rows are not bit-exact with the aiter kernel, and the gain shrinks as verify rows grow: on MI355X about −25% up to 128 rows, but none at 256 rows (about Kimi-K3 DSpark at concurrency 64). | ## Engram (DeepSeek-V4.1) From 5f35fbb98b78ae2a14fc1ff5d0e47c30b07d19d3 Mon Sep 17 00:00:00 2001 From: chelnnexy Date: Wed, 4 Nov 2026 02:50:13 +0000 Subject: [PATCH 7/9] docs(bonus-sample): name what would turn ATOM_SPEC_BONUS_FUSED on by default Per review: a row-count gate (fused at or below the crossover, aiter above) would remove the speed reason for keeping the flag off, and aiter's fused_mixed_sample (ROCm/aiter#7416), already faster at 256 rows, would remove it without one. Say so in the env var doc and the envs.py comment. The test module docstring now says why a plain pytest run reports most tests as skipped: they run in the interpreter child, which fails if any of them fails or if it collects nothing. --- atom/utils/envs.py | 5 ++++- docs/environment_variables.md | 2 +- tests/test_bonus_sample.py | 5 +++++ 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 1a5f08ebbc..fce51a5e7e 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -1030,7 +1030,10 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: # server starts. Off by default: sampling rows are not bit-exact with the # aiter kernel (above), and the gain shrinks with the verify-row count -- # on MI355X about -25% at <= 128 rows but none at 256 (~Kimi-K3 DSpark at - # concurrency 64) -- so it stays opt-in until that shape gains too. + # concurrency 64) -- so it stays opt-in until that shape gains too. A + # row-count gate (fused at or below the crossover, aiter above) would remove + # the speed reason; aiter's fused_mixed_sample (ROCm/aiter#7416), faster + # than the aiter path at 256 rows too, would remove it without one. "ATOM_SPEC_BONUS_FUSED": lambda: os.getenv("ATOM_SPEC_BONUS_FUSED", "0") == "1", # --- Atomesh --- # Build atomesh when installing ATOM from source. diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 819c25bc24..8e42c0ff55 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -307,7 +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. Off by default because sampling rows are not bit-exact with the aiter kernel, and the gain shrinks as verify rows grow: on MI355X about −25% up to 128 rows, but none at 256 rows (about Kimi-K3 DSpark at concurrency 64). | +| **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. Off by default because sampling rows are not bit-exact with the aiter kernel, and the gain shrinks as verify rows grow: on MI355X about −25% up to 128 rows, but none at 256 rows (about Kimi-K3 DSpark at concurrency 64). What would change the default: a row-count gate (fused at or below the crossover, aiter above) removes the speed reason, and aiter's `fused_mixed_sample` (ROCm/aiter#7416), which is faster than the aiter path at 256 rows on MI355X and MI300X, removes it without a gate; the near-tie difference would then be the only reason left. | ## Engram (DeepSeek-V4.1) diff --git a/tests/test_bonus_sample.py b/tests/test_bonus_sample.py index 0adbed4f17..ecc0d410d2 100644 --- a/tests/test_bonus_sample.py +++ b/tests/test_bonus_sample.py @@ -6,6 +6,11 @@ Runs on CPU tensors under Triton's interpreter (`triton_interpreter`), so it needs Triton but no GPU and no AITER. +A plain `pytest` run reports most tests as skipped: those are marked +`in_the_interpreter` and run in the `TRITON_INTERPRET=1` child that +`test_under_the_interpreter` starts, which fails if any of them fails or +if the child collects nothing. + The reference is float64 `argmax_i softmax(x / max(T, 1e-5))_i / (E_i + eps)` on the noise the kernel is given -- the definition of `mixed_sample_outer_exponential` -- and, for a greedy row (T at or below the From 0a4e7ec2397060f20ef8b378340dea0cf80c13f2 Mon Sep 17 00:00:00 2001 From: chelnnexy Date: Thu, 5 Nov 2026 14:11:58 +0000 Subject: [PATCH 8/9] refactor(spec-decode): take the fused draw from aiter's fused_mixed_sample Per review: the kernel now lives in aiter as fused_mixed_sample (ROCm/aiter#7416), so ATOM keeps only its side. With ATOM_SPEC_BONUS_FUSED=1 the sampler calls fused_mixed_sample(logits, noise, temperatures, eps=..., greedy_eps=...) on the same per-row exponential draw. Remove the in-tree bonus_sample.py, its kernel tests and tests/triton_interpreter.py. aiter's op subtracts the max before scaling by 1/T, so a sampling row differs from mixed_sample_outer_exponential only on a near tie of a few 1e-7 relative. The in-tree kernel scaled first, and that window grew like 1/T. An aiter without fused_mixed_sample keeps the aiter race and warns once. On a GPU aiter has no tuned config for (it ships gfx942 and gfx950), the op raises FileNotFoundError before launching; the sampler then runs the aiter race on the noise it already drew, so the RNG stream matches the flag-off path, and keeps that path for later steps. The wiring tests cover both fallbacks and the argument order. --- atom/model_ops/bonus_sample.py | 170 -------------------------- atom/model_ops/sampler.py | 59 ++++++++-- atom/utils/envs.py | 31 +++-- docs/environment_variables.md | 2 +- docs/model_ops_guide.md | 3 +- tests/test_bonus_sample.py | 190 ------------------------------ tests/test_bonus_sample_wiring.py | 71 ++++++++--- tests/triton_interpreter.py | 50 -------- 8 files changed, 119 insertions(+), 457 deletions(-) delete mode 100644 atom/model_ops/bonus_sample.py delete mode 100644 tests/test_bonus_sample.py delete mode 100644 tests/triton_interpreter.py diff --git a/atom/model_ops/bonus_sample.py b/atom/model_ops/bonus_sample.py deleted file mode 100644 index f67fb8fffc..0000000000 --- a/atom/model_ops/bonus_sample.py +++ /dev/null @@ -1,170 +0,0 @@ -# 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 diff --git a/atom/model_ops/sampler.py b/atom/model_ops/sampler.py index 42c5a231d6..b69a640eb2 100644 --- a/atom/model_ops/sampler.py +++ b/atom/model_ops/sampler.py @@ -73,6 +73,22 @@ def get_per_row_exponential(num_tokens: int, vocab_size: int, device) -> torch.T ).exponential_(1) +def _load_fused_mixed_sample(): + """aiter's `fused_mixed_sample` (ROCm/aiter#7416), or None if this aiter has none.""" + try: + from aiter.ops.triton.fusions.fused_mixed_sample import fused_mixed_sample + except ImportError: + warnings.warn( + "ATOM_SPEC_BONUS_FUSED=1 needs aiter's fused_mixed_sample " + "(ROCm/aiter#7416), which this aiter does not have; the draw keeps " + "mixed_sample_outer_exponential.", + UserWarning, + stacklevel=2, + ) + return None + return fused_mixed_sample + + class Sampler(nn.Module): def __init__(self): @@ -82,9 +98,7 @@ def __init__(self): # 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 + self._fused_temperature_sample = _load_fused_mixed_sample() def sample_verification_tokens( self, @@ -166,14 +180,32 @@ def forward( 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. + # and the greedy override in aiter's two vocab-tiled launches. num_tokens, vocab_size = logits.shape exponential = get_per_row_exponential( num_tokens, vocab_size, logits.device ) - return self._fused_temperature_sample( - logits, temperatures, exponential, self.eps, SAMPLER_EPS - ) + try: + return self._fused_temperature_sample( + logits, + exponential, + temperatures, + eps=self.eps, + greedy_eps=SAMPLER_EPS, + ) + except FileNotFoundError: + # aiter tunes fused_mixed_sample per arch and raises before + # launching where it has no config. Keep its HIP race, on + # this step's noise, for this and every later step. + warnings.warn( + "ATOM_SPEC_BONUS_FUSED: aiter's fused_mixed_sample has no " + "config for this GPU; using mixed_sample_outer_exponential.", + UserWarning, + stacklevel=2, + ) + self._fused_temperature_sample = None + sampled = self._race(logits, temperatures, exponential) + return _apply_greedy_tokens(logits, temperatures, sampled) sampled = self._temperature_sample( logits, temperatures, needs_independent_noise=needs_independent_noise ) @@ -216,13 +248,24 @@ def _temperature_sample( which preserves the existing run-to-run determinism optimization. """ num_tokens, vocab_size = logits.shape - sampled_tokens = torch.empty(num_tokens, dtype=torch.int, device=logits.device) if needs_independent_noise: exponential = get_per_row_exponential(num_tokens, vocab_size, logits.device) else: exponential = get_per_token_exponential(vocab_size, logits.device).expand( num_tokens, vocab_size ) + return self._race(logits, temperatures, exponential) + + def _race( + self, + logits: torch.Tensor, + temperatures: torch.Tensor, + exponential: torch.Tensor, + ) -> torch.Tensor: + """aiter's exponential race on the given noise, one token per row.""" + sampled_tokens = torch.empty( + logits.shape[0], dtype=torch.int, device=logits.device + ) mixed_sample_outer_exponential( sampled_tokens, logits, exponential, temperatures, eps=self.eps ) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index fce51a5e7e..4b5ee36f45 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -1017,23 +1017,20 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: ), # "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. Off by default: sampling rows are not bit-exact with the - # aiter kernel (above), and the gain shrinks with the verify-row count -- - # on MI355X about -25% at <= 128 rows but none at 256 (~Kimi-K3 DSpark at - # concurrency 64) -- so it stays opt-in until that shape gains too. A - # row-count gate (fused at or below the crossover, aiter above) would remove - # the speed reason; aiter's fused_mixed_sample (ROCm/aiter#7416), faster - # than the aiter path at 256 rows too, would remove it without one. + # aiter's `fused_mixed_sample` (ROCm/aiter#7416, two vocab-tiled Triton + # launches) 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 a few 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. Needs an aiter that has + # `fused_mixed_sample`; without it, or on a GPU it has no tuned config for, + # the draw keeps the aiter path and warns once. Read when the `Sampler` is + # built, so it must be set before the server starts. Off by default: + # sampling rows are not bit-exact with the aiter kernel (above). "ATOM_SPEC_BONUS_FUSED": lambda: os.getenv("ATOM_SPEC_BONUS_FUSED", "0") == "1", # --- Atomesh --- # Build atomesh when installing ATOM from source. diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 8e42c0ff55..c91258f19a 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -307,7 +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. Off by default because sampling rows are not bit-exact with the aiter kernel, and the gain shrinks as verify rows grow: on MI355X about −25% up to 128 rows, but none at 256 rows (about Kimi-K3 DSpark at concurrency 64). What would change the default: a row-count gate (fused at or below the crossover, aiter above) removes the speed reason, and aiter's `fused_mixed_sample` (ROCm/aiter#7416), which is faster than the aiter path at 256 rows on MI355X and MI300X, removes it without a gate; the near-tie difference would then be the only reason left. | +| **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) with aiter's `fused_mixed_sample` (ROCm/aiter#7416, two vocab-tiled Triton launches) 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 a few 1e-7 relative (fp32 rounding order). Rows with top-k/top-p, and batches without per-row noise, keep the aiter path. Needs an aiter with `fused_mixed_sample`; without it, or on a GPU without a tuned config for it (aiter ships gfx942 and gfx950), the draw keeps the aiter path and warns once. Read when the `Sampler` is built, so set it before the server starts. Off by default because sampling rows are not bit-exact with the aiter kernel. | ## Engram (DeepSeek-V4.1) diff --git a/docs/model_ops_guide.md b/docs/model_ops_guide.md index 74d7de039c..646ec9870b 100644 --- a/docs/model_ops_guide.md +++ b/docs/model_ops_guide.md @@ -496,7 +496,7 @@ 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. +**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), the sampler calls aiter's `fused_mixed_sample` (ROCm/aiter#7416) instead of `mixed_sample_outer_exponential` + `topk_select` + `<=` + `where`: two vocab-tiled launches, a `(row, tile)` kernel that keeps each tile's best race key and raw-logit argmax, and a `(row,)` kernel that 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 a few 1e-7 relative. Rows with top-k/top-p keep the aiter path, and so does every row when aiter has no `fused_mixed_sample` or no tuned config for the GPU. ### Required channel (`required_channel.py`) @@ -574,7 +574,6 @@ 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 | diff --git a/tests/test_bonus_sample.py b/tests/test_bonus_sample.py deleted file mode 100644 index ecc0d410d2..0000000000 --- a/tests/test_bonus_sample.py +++ /dev/null @@ -1,190 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""`ATOM_SPEC_BONUS_FUSED`: the two-launch draw against the race it replaces. - -Runs on CPU tensors under Triton's interpreter (`triton_interpreter`), so it -needs Triton but no GPU and no AITER. - -A plain `pytest` run reports most tests as skipped: those are marked -`in_the_interpreter` and run in the `TRITON_INTERPRET=1` child that -`test_under_the_interpreter` starts, which fails if any of them fails or -if the child collects nothing. - -The reference is float64 `argmax_i softmax(x / max(T, 1e-5))_i / (E_i + eps)` on -the noise the kernel is given -- the definition of -`mixed_sample_outer_exponential` -- and, for a greedy row (T at or below the -epsilon), the lowest-id argmax of the raw logits (`topk_select(tie="low")`). -Greedy rows must match exactly, ties planted. A sampling row must match unless -the reference's two best keys are a near tie (relative gap < 1e-5), which the -fp32 rounding of either kernel may order either way. -""" - -import pytest -import torch -from triton_interpreter import in_the_interpreter, interpreter_test - -EPS = 1e-10 - -test_under_the_interpreter = interpreter_test(__file__) - - -def _bonus_sample(): - # Imported inside the child only: with the interpreter off, Triton would - # build compiled kernels that nothing here can launch on CPU tensors. - from atom.model_ops import bonus_sample - - return bonus_sample - - -def _reference_pick(logits_row, temp, noise_row): - """(token, relative gap between the two best keys) for one sampling row.""" - t = max(float(temp), 1e-5) - p = torch.softmax(logits_row.double() / t, dim=0) - key = p / (noise_row.double() + EPS) - top2 = torch.topk(key, 2) - gap = float((top2.values[0] - top2.values[1]) / top2.values[0]) - return int(top2.indices[0]), int(top2.indices[1]), gap - - -@in_the_interpreter -@pytest.mark.parametrize("seed,vocab", [(1, 163840), (2, 50000), (3, 4097), (4, 1000)]) -def test_matches_the_race_it_replaces(seed, vocab): - gen = torch.Generator().manual_seed(seed) - rows = 6 - logits = (torch.randn(rows, vocab, generator=gen) * 4).to(torch.bfloat16) - temps = torch.tensor([1.0, 0.7, EPS, 1.3, EPS, 0.05], dtype=torch.float32) - # Exact ties on the greedy rows: the lowest id must win. - logits[2, 777 % vocab] = logits[2].float().max() - logits[2, 12345 % vocab] = logits[2].float().max() - logits[4, 5] = 50.0 - logits[4, vocab - 1] = 50.0 - noise = torch.empty(rows, vocab).exponential_(1, generator=gen) - - out = _bonus_sample().bonus_temperature_sample(logits, temps, noise, EPS, EPS) - - assert out.dtype == torch.int32 and out.shape == (rows,) - for r in range(rows): - if temps[r] <= EPS: - expected = int(torch.nonzero(logits[r] == logits[r].max())[0]) - assert int(out[r]) == expected - continue - best, runner_up, gap = _reference_pick(logits[r], temps[r], noise[r]) - if int(out[r]) != best: - assert gap < 1e-5 and int(out[r]) == runner_up, (r, gap) - - -@in_the_interpreter -def test_every_row_uses_its_own_temperature_and_noise(): - vocab = 4097 - gen = torch.Generator().manual_seed(7) - logits = torch.randn(3, vocab, generator=gen).to(torch.bfloat16) - logits[:] = logits[0] # identical logits: only T and the noise differ - temps = torch.tensor([1.0, 1.0, EPS]) - noise = torch.empty(3, vocab).exponential_(1, generator=gen) - out = _bonus_sample().bonus_temperature_sample(logits, temps, noise, EPS, EPS) - assert int(out[2]) == int(logits[0].float().argmax()) - for r in (0, 1): - best, runner_up, gap = _reference_pick(logits[r], temps[r], noise[r]) - assert int(out[r]) == best or (gap < 1e-5 and int(out[r]) == runner_up) - - -@in_the_interpreter -def test_a_row_below_the_clamp_samples_like_the_clamped_temperature(): - # aiter clamps T to 1e-5 inside the race; a row just above the greedy - # epsilon is a sampling row whose draw is (almost surely) the argmax. - vocab = 4097 - gen = torch.Generator().manual_seed(11) - logits = torch.randn(1, vocab, generator=gen).to(torch.bfloat16) - noise = torch.empty(1, vocab).exponential_(1, generator=gen) - temps = torch.tensor([1e-7]) - out = _bonus_sample().bonus_temperature_sample(logits, temps, noise, EPS, EPS) - assert int(out[0]) == int(logits[0].float().argmax()) - - -@in_the_interpreter -def test_no_rows_returns_an_empty_result(): - out = _bonus_sample().bonus_temperature_sample( - torch.empty(0, 4097, dtype=torch.bfloat16), - torch.empty(0), - torch.empty(0, 4097), - EPS, - EPS, - ) - assert out.shape == (0,) and out.dtype == torch.int32 - - -@in_the_interpreter -def test_same_inputs_give_the_same_output(): - # The draw is in `Sampler.forward`; the kernels carry no state of their own, - # so the same (logits, temps, noise) must give the same token every call. - vocab = 50000 - gen = torch.Generator().manual_seed(3) - logits = (torch.randn(5, vocab, generator=gen) * 4).to(torch.bfloat16) - temps = torch.tensor([1.0, 0.7, 0.9, 1.2, 0.5]) - noise = torch.empty(5, vocab).exponential_(1, generator=gen) - f = _bonus_sample().bonus_temperature_sample - first = f(logits, temps, noise, EPS, EPS) - assert torch.equal(first, f(logits, temps, noise, EPS, EPS)) - - -@in_the_interpreter -def test_winner_in_a_late_tile_survives_the_reduction(): - # One dominating race key in the final vocab tile must beat tile 0: the - # per-tile running max and the `(row,)` kernel's k * exp(m - mm) rescale - # have to reconstruct the global key across tiles. - vocab = 163840 # 40 tiles of 4096 - gen = torch.Generator().manual_seed(5) - logits = (torch.randn(1, vocab, generator=gen) * 2).to(torch.bfloat16) - noise = torch.full((1, vocab), 10.0) - win = vocab - 3 - logits[0, win] = 60.0 - noise[0, win] = 1e-3 # huge key exp(0) / ~1e-3 against a flat field - temps = torch.tensor([1.0]) - out = _bonus_sample().bonus_temperature_sample(logits, temps, noise, EPS, EPS) - assert int(out[0]) == win - - -@in_the_interpreter -@pytest.mark.filterwarnings("error::RuntimeWarning") -@pytest.mark.parametrize("forced", [5, 100000, 163837]) -def test_a_forced_channel_row_samples_the_forced_id(forced): - # `constrain_logits` writes -inf over every id but the forced one, so all - # tiles but one are entirely -inf: their tile max is -inf and the race - # must not compute exp(-inf - -inf) = NaN on the way to the forced id. - vocab = 163840 - gen = torch.Generator().manual_seed(forced) - logits = torch.full((2, vocab), float("-inf")) - logits[:, forced] = 3.0 - temps = torch.tensor([1.0, EPS]) - noise = torch.empty(2, vocab).exponential_(1, generator=gen) - out = _bonus_sample().bonus_temperature_sample( - logits.to(torch.bfloat16), temps, noise, EPS, EPS - ) - assert out.tolist() == [forced, forced] - - -@in_the_interpreter -@pytest.mark.filterwarnings("error::RuntimeWarning") -def test_masked_ids_never_win_and_every_id_is_in_range(): - # Whole -inf tiles next to tiles with scattered -inf ids, and one row that - # is -inf throughout: no NaN, no -inf id drawn, and with no clamp on the - # result every id must already be in range. - vocab = 50000 - gen = torch.Generator().manual_seed(9) - logits = (torch.randn(4, vocab, generator=gen) * 4).to(torch.bfloat16) - logits[:, : 3 * 4096] = float("-inf") - banned = torch.randint(0, vocab, (256,), generator=gen) - logits[:, banned] = float("-inf") - logits[3] = float("-inf") - temps = torch.tensor([1.0, 0.7, EPS, 1.0]) - noise = torch.empty(4, vocab).exponential_(1, generator=gen) - out = _bonus_sample().bonus_temperature_sample(logits, temps, noise, EPS, EPS) - assert ((out >= 0) & (out < vocab)).all() - for r in range(3): - assert logits[r, int(out[r])] != float("-inf") - if temps[r] <= EPS: - assert int(out[r]) == int(torch.nonzero(logits[r] == logits[r].max())[0]) - continue - best, runner_up, gap = _reference_pick(logits[r], temps[r], noise[r]) - assert int(out[r]) == best or (gap < 1e-5 and int(out[r]) == runner_up) diff --git a/tests/test_bonus_sample_wiring.py b/tests/test_bonus_sample_wiring.py index 31a870ce74..d3eeb5207d 100644 --- a/tests/test_bonus_sample_wiring.py +++ b/tests/test_bonus_sample_wiring.py @@ -5,9 +5,9 @@ The sampler needs AITER at import, which a CPU runner does not have, so it is loaded here against a fake `aiter` under a private module name and the fakes are -removed again. `atom.model_ops.bonus_sample` (Triton) is faked the same way, so -these tests need neither AITER nor Triton and run on any CPU runner. Nothing -runs a kernel: the aiter ops and the fused draw are recording stand-ins. +removed again. aiter's `fused_mixed_sample` is faked the same way, so these +tests need neither AITER nor Triton and run on any CPU runner. Nothing runs a +kernel: the aiter ops and the fused draw are recording stand-ins. """ import importlib.util @@ -22,6 +22,7 @@ SAMPLER_PATH = ( pathlib.Path(__file__).resolve().parents[1] / "atom" / "model_ops" / "sampler.py" ) +FUSED_MODULE = "aiter.ops.triton.fusions.fused_mixed_sample" @pytest.fixture @@ -33,6 +34,7 @@ def load_sampler(monkeypatch): "aiter.ops.triton", "aiter.ops.triton.softmax", "aiter.ops.triton.topk", + "aiter.ops.triton.fusions", ): monkeypatch.setitem(sys.modules, name, MagicMock()) stub = sys.modules["aiter"] @@ -43,12 +45,15 @@ def load_sampler(monkeypatch): torch.zeros(x.shape[0], 1, dtype=torch.int32), ) ) - fused = types.ModuleType("atom.model_ops.bonus_sample") - fused.bonus_temperature_sample = MagicMock(name="fused_draw") - monkeypatch.setitem(sys.modules, "atom.model_ops.bonus_sample", fused) + fused = types.ModuleType(FUSED_MODULE) + fused.fused_mixed_sample = MagicMock(name="fused_draw") + monkeypatch.setitem(sys.modules, FUSED_MODULE, fused) - def load(fused: bool): + def load(fused: bool, aiter_has_op: bool = True): monkeypatch.setenv("ATOM_SPEC_BONUS_FUSED", "1" if fused else "0") + if not aiter_has_op: + # an aiter without the op: the import raises ImportError + monkeypatch.setitem(sys.modules, FUSED_MODULE, None) spec = importlib.util.spec_from_file_location( "_sampler_under_test", SAMPLER_PATH ) @@ -56,7 +61,7 @@ def load(fused: bool): spec.loader.exec_module(module) return module, stub - load.fused_draw = fused.bonus_temperature_sample + load.fused_draw = fused.fused_mixed_sample return load @@ -64,11 +69,21 @@ def _batch(rows=4, vocab=64): return torch.randn(rows, vocab), torch.ones(rows) -def test_flag_on_wires_the_fused_draw(load_sampler): +def test_flag_on_wires_aiters_fused_mixed_sample(load_sampler): module, _ = load_sampler(fused=True) assert module.Sampler()._fused_temperature_sample is load_sampler.fused_draw +def test_flag_on_without_the_op_keeps_the_aiter_path(load_sampler): + module, aiter = load_sampler(fused=True, aiter_has_op=False) + with pytest.warns(UserWarning, match="fused_mixed_sample"): + sampler = module.Sampler() + assert sampler._fused_temperature_sample is None + logits, temps = _batch() + sampler(logits, temps, needs_independent_noise=True) + aiter.mixed_sample_outer_exponential.assert_called_once() + + def test_flag_off_keeps_the_aiter_path(load_sampler): module, aiter = load_sampler(fused=False) sampler = module.Sampler() @@ -78,24 +93,42 @@ def test_flag_off_keeps_the_aiter_path(load_sampler): aiter.mixed_sample_outer_exponential.assert_called_once() -def test_flag_on_routes_the_per_row_noise_draw(load_sampler, monkeypatch): +def test_flag_on_routes_the_per_row_noise_draw(load_sampler): module, aiter = load_sampler(fused=True) sampler = module.Sampler() - seen = [] - - def fake(logits, temps, noise, eps, greedy_eps): - seen.append((logits, temps, noise, eps, greedy_eps)) - return torch.full((logits.shape[0],), 7, dtype=torch.int32) - - sampler._fused_temperature_sample = fake + fused = load_sampler.fused_draw + fused.side_effect = lambda logits, *a, **k: torch.full( + (logits.shape[0],), 7, dtype=torch.int32 + ) logits, temps = _batch() out = sampler(logits, temps, needs_independent_noise=True) assert out.tolist() == [7] * 4 aiter.mixed_sample_outer_exponential.assert_not_called() - ((_, _, noise, eps, greedy_eps),) = seen + ((args, kwargs),) = fused.call_args_list + got_logits, noise, got_temps = args + assert got_logits is logits and got_temps is temps assert noise.shape == logits.shape and noise.dtype == torch.float32 - assert eps == sampler.eps and greedy_eps == module.SAMPLER_EPS + assert kwargs == {"eps": sampler.eps, "greedy_eps": module.SAMPLER_EPS} + + +def test_no_config_for_the_gpu_falls_back_on_the_same_noise(load_sampler): + module, aiter = load_sampler(fused=True) + sampler = module.Sampler() + load_sampler.fused_draw.side_effect = FileNotFoundError("no config") + logits, temps = _batch() + with pytest.warns(UserWarning, match="no config"): + sampler(logits, temps, needs_independent_noise=True) + + ((fused_args, _),) = load_sampler.fused_draw.call_args_list + ((race_args, race_kwargs),) = aiter.mixed_sample_outer_exponential.call_args_list + assert race_args[2] is fused_args[1], "the fallback must reuse the drawn noise" + assert race_kwargs == {"eps": sampler.eps} + assert sampler._fused_temperature_sample is None + + sampler(logits, temps, needs_independent_noise=True) + assert load_sampler.fused_draw.call_count == 1 + assert aiter.mixed_sample_outer_exponential.call_count == 2 @pytest.mark.parametrize( diff --git a/tests/triton_interpreter.py b/tests/triton_interpreter.py deleted file mode 100644 index dd8283c63b..0000000000 --- a/tests/triton_interpreter.py +++ /dev/null @@ -1,50 +0,0 @@ -# SPDX-License-Identifier: MIT -# Test helper: drive a Triton kernel on CPU tensors, without a GPU. -# -# Triton's interpreter has to be on before `triton.language` is imported (its -# `tl.max`, `tl.sum`, ... are decorated at import), so it cannot be switched on -# inside a process that has already imported Triton. A kernel test module -# therefore has two halves: with TRITON_INTERPRET=1 in the environment it holds -# the real tests; without it, it holds one test that runs the same file again in -# a child process that has the variable set: -# -# test_under_the_interpreter = interpreter_test(__file__) -# -# @in_the_interpreter -# def test_a_kernel(): ... - -import os -import subprocess -import sys - -import pytest - - -def interpreting() -> bool: - return os.environ.get("TRITON_INTERPRET") == "1" - - -in_the_interpreter = pytest.mark.skipif( - not interpreting(), - reason="runs in the child process of test_under_the_interpreter", -) - - -def interpreter_test(test_file: str, timeout: float = 1800.0): - """A test that re-runs `test_file` under Triton's CPU interpreter.""" - - @pytest.mark.skipif(interpreting(), reason="this is the child process") - def test_under_the_interpreter(): - pytest.importorskip("triton") - env = dict(os.environ, TRITON_INTERPRET="1") - done = subprocess.run( - [sys.executable, "-m", "pytest", test_file, "-q", "-p", "no:cacheprovider"], - env=env, - check=False, - capture_output=True, - text=True, - timeout=timeout, - ) - assert done.returncode == 0, done.stdout[-6000:] + done.stderr[-2000:] - - return test_under_the_interpreter From f47858e5f1b220882939bc962ff32fc256b5e0fa Mon Sep 17 00:00:00 2001 From: chelnnexy Date: Fri, 6 Nov 2026 13:21:22 +0000 Subject: [PATCH 9/9] fix(spec-decode): degrade, not fail, when aiter's fused_mixed_sample cannot take the draw Two failure modes of ATOM_SPEC_BONUS_FUSED=1 surfaced at the first speculative sample instead of degrading to the aiter path: - an aiter whose fused_mixed_sample signature no longer takes the call (renamed or dropped eps / greedy_eps, a new required argument) raised TypeError. _load_fused_mixed_sample now binds the call against the op's signature and treats a mismatch like a missing op: warn once, keep mixed_sample_outer_exponential. - an input the op refuses (its ValueError preconditions) crashed a run that works with the flag off. It now takes the same warn-and-fall-back path as a missing tuned config, on the noise already drawn. The wiring tests cover three signature mismatches and the ValueError fallback, and bind the call to the real op wherever aiter has it. The env doc states that batches sharing one noise row are out of scope by choice. --- atom/model_ops/sampler.py | 46 +++++++++++++------- atom/utils/envs.py | 11 +++-- docs/environment_variables.md | 2 +- docs/model_ops_guide.md | 2 +- tests/test_bonus_sample_wiring.py | 71 +++++++++++++++++++++++++++---- 5 files changed, 101 insertions(+), 31 deletions(-) diff --git a/atom/model_ops/sampler.py b/atom/model_ops/sampler.py index b69a640eb2..eecdc230db 100644 --- a/atom/model_ops/sampler.py +++ b/atom/model_ops/sampler.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. +import inspect import warnings import torch @@ -74,19 +75,29 @@ def get_per_row_exponential(num_tokens: int, vocab_size: int, device) -> torch.T def _load_fused_mixed_sample(): - """aiter's `fused_mixed_sample` (ROCm/aiter#7416), or None if this aiter has none.""" + """aiter's `fused_mixed_sample` (ROCm/aiter#7416), or None if this aiter has + none, or has one that `Sampler.forward`'s call does not bind to.""" try: from aiter.ops.triton.fusions.fused_mixed_sample import fused_mixed_sample except ImportError: - warnings.warn( - "ATOM_SPEC_BONUS_FUSED=1 needs aiter's fused_mixed_sample " - "(ROCm/aiter#7416), which this aiter does not have; the draw keeps " - "mixed_sample_outer_exponential.", - UserWarning, - stacklevel=2, - ) - return None - return fused_mixed_sample + reason = "which this aiter does not have" + else: + try: + # The call in `Sampler.forward`: three positionals, two keywords. + inspect.signature(fused_mixed_sample).bind( + None, None, None, eps=0.0, greedy_eps=0.0 + ) + return fused_mixed_sample + except (TypeError, ValueError) as e: + reason = f"whose signature this ATOM does not match ({e})" + warnings.warn( + f"ATOM_SPEC_BONUS_FUSED=1 needs aiter's fused_mixed_sample " + f"(ROCm/aiter#7416), {reason}; the draw keeps " + "mixed_sample_outer_exponential.", + UserWarning, + stacklevel=2, + ) + return None class Sampler(nn.Module): @@ -193,13 +204,16 @@ def forward( eps=self.eps, greedy_eps=SAMPLER_EPS, ) - except FileNotFoundError: - # aiter tunes fused_mixed_sample per arch and raises before - # launching where it has no config. Keep its HIP race, on - # this step's noise, for this and every later step. + except (FileNotFoundError, ValueError) as e: + # Both raise before either kernel launches: no tuned config + # for this arch, or an input precondition (2-D, f32 noise + # and temperatures, contiguous last dim, shapes) that this + # call no longer meets. Keep its HIP race, on this step's + # noise, for this and every later step. warnings.warn( - "ATOM_SPEC_BONUS_FUSED: aiter's fused_mixed_sample has no " - "config for this GPU; using mixed_sample_outer_exponential.", + f"ATOM_SPEC_BONUS_FUSED: aiter's fused_mixed_sample " + f"refused this draw ({type(e).__name__}: {e}); using " + "mixed_sample_outer_exponential.", UserWarning, stacklevel=2, ) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 4b5ee36f45..e684b48b2e 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -1025,10 +1025,13 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: # 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 a few 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. Needs an aiter that has - # `fused_mixed_sample`; without it, or on a GPU it has no tuned config for, - # the draw keeps the aiter path and warns once. Read when the `Sampler` is + # each other (fp32 rounding order). Rows with top-k/top-p keep the aiter + # path. So, by choice of scope, do batches that share one noise row + # (`needs_independent_noise=False`, ordinary non-speculative sampling): + # aiter accepts that stride-0 broadcast, but it is not measured here. Needs + # an aiter whose `fused_mixed_sample` takes this call; without it, on a GPU + # it has no tuned config for, or for an input it refuses (ValueError), the + # draw keeps the aiter path and warns once. Read when the `Sampler` is # built, so it must be set before the server starts. Off by default: # sampling rows are not bit-exact with the aiter kernel (above). "ATOM_SPEC_BONUS_FUSED": lambda: os.getenv("ATOM_SPEC_BONUS_FUSED", "0") == "1", diff --git a/docs/environment_variables.md b/docs/environment_variables.md index c91258f19a..fc632d68e9 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -307,7 +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) with aiter's `fused_mixed_sample` (ROCm/aiter#7416, two vocab-tiled Triton launches) 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 a few 1e-7 relative (fp32 rounding order). Rows with top-k/top-p, and batches without per-row noise, keep the aiter path. Needs an aiter with `fused_mixed_sample`; without it, or on a GPU without a tuned config for it (aiter ships gfx942 and gfx950), the draw keeps the aiter path and warns once. Read when the `Sampler` is built, so set it before the server starts. Off by default because sampling rows are not bit-exact with the aiter kernel. | +| **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) with aiter's `fused_mixed_sample` (ROCm/aiter#7416, two vocab-tiled Triton launches) 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 a few 1e-7 relative (fp32 rounding order). Rows with top-k/top-p keep the aiter path; so, by choice of scope, do batches that share one noise row (ordinary non-speculative sampling), which aiter accepts but this flag does not measure. Needs an aiter whose `fused_mixed_sample` takes this call; without it, on a GPU without a tuned config for it (aiter ships gfx942 and gfx950), or for an input it refuses (`ValueError`), the draw keeps the aiter path and warns once. Read when the `Sampler` is built, so set it before the server starts. Off by default because sampling rows are not bit-exact with the aiter kernel. | ## Engram (DeepSeek-V4.1) diff --git a/docs/model_ops_guide.md b/docs/model_ops_guide.md index 646ec9870b..e1388f7609 100644 --- a/docs/model_ops_guide.md +++ b/docs/model_ops_guide.md @@ -496,7 +496,7 @@ 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), the sampler calls aiter's `fused_mixed_sample` (ROCm/aiter#7416) instead of `mixed_sample_outer_exponential` + `topk_select` + `<=` + `where`: two vocab-tiled launches, a `(row, tile)` kernel that keeps each tile's best race key and raw-logit argmax, and a `(row,)` kernel that 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 a few 1e-7 relative. Rows with top-k/top-p keep the aiter path, and so does every row when aiter has no `fused_mixed_sample` or no tuned config for the GPU. +**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), the sampler calls aiter's `fused_mixed_sample` (ROCm/aiter#7416) instead of `mixed_sample_outer_exponential` + `topk_select` + `<=` + `where`: two vocab-tiled launches, a `(row, tile)` kernel that keeps each tile's best race key and raw-logit argmax, and a `(row,)` kernel that 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 a few 1e-7 relative. Rows with top-k/top-p keep the aiter path, as do batches that share one noise row (a scoping choice: aiter accepts the broadcast, but it is not measured here), and so does every row when aiter has no `fused_mixed_sample`, one whose signature does not take this call, no tuned config for the GPU, or refuses an input (`ValueError`). ### Required channel (`required_channel.py`) diff --git a/tests/test_bonus_sample_wiring.py b/tests/test_bonus_sample_wiring.py index d3eeb5207d..65a02fb54b 100644 --- a/tests/test_bonus_sample_wiring.py +++ b/tests/test_bonus_sample_wiring.py @@ -7,13 +7,16 @@ loaded here against a fake `aiter` under a private module name and the fakes are removed again. aiter's `fused_mixed_sample` is faked the same way, so these tests need neither AITER nor Triton and run on any CPU runner. Nothing runs a -kernel: the aiter ops and the fused draw are recording stand-ins. +kernel: the aiter ops and the fused draw are recording stand-ins. The one +exception binds the call to the real op's signature where aiter has the op, so +a change in #7416's signature fails here and not only in production. """ import importlib.util import pathlib import sys import types +import warnings from unittest.mock import MagicMock import pytest @@ -25,6 +28,16 @@ FUSED_MODULE = "aiter.ops.triton.fusions.fused_mixed_sample" +@pytest.fixture(scope="module") +def real_fused_mixed_sample(): + """aiter's real op, imported before `load_sampler` fakes aiter; else skip.""" + try: + from aiter.ops.triton.fusions.fused_mixed_sample import fused_mixed_sample + except ImportError as e: # no aiter, or one without the op (CPU runners) + pytest.skip(f"aiter's fused_mixed_sample is not importable: {e!r}") + return fused_mixed_sample + + @pytest.fixture def load_sampler(monkeypatch): for name in ( @@ -45,12 +58,14 @@ def load_sampler(monkeypatch): torch.zeros(x.shape[0], 1, dtype=torch.int32), ) ) - fused = types.ModuleType(FUSED_MODULE) - fused.fused_mixed_sample = MagicMock(name="fused_draw") - monkeypatch.setitem(sys.modules, FUSED_MODULE, fused) + fused_module = types.ModuleType(FUSED_MODULE) + fused_module.fused_mixed_sample = MagicMock(name="fused_draw") + monkeypatch.setitem(sys.modules, FUSED_MODULE, fused_module) - def load(fused: bool, aiter_has_op: bool = True): + def load(fused: bool, aiter_has_op: bool = True, op=None): monkeypatch.setenv("ATOM_SPEC_BONUS_FUSED", "1" if fused else "0") + if op is not None: + fused_module.fused_mixed_sample = op if not aiter_has_op: # an aiter without the op: the import raises ImportError monkeypatch.setitem(sys.modules, FUSED_MODULE, None) @@ -61,7 +76,7 @@ def load(fused: bool, aiter_has_op: bool = True): spec.loader.exec_module(module) return module, stub - load.fused_draw = fused.fused_mixed_sample + load.fused_draw = fused_module.fused_mixed_sample return load @@ -84,6 +99,37 @@ def test_flag_on_without_the_op_keeps_the_aiter_path(load_sampler): aiter.mixed_sample_outer_exponential.assert_called_once() +def _renamed_eps(input, exponentials, temperatures, sample_eps=1e-10, greedy_eps=0.0): + pass + + +def _no_greedy_eps(input, exponentials, temperatures, eps=1e-10): + pass + + +def _new_required(input, exponentials, temperatures, seed, eps=1e-10, greedy_eps=0.0): + pass + + +@pytest.mark.parametrize("op", [_renamed_eps, _no_greedy_eps, _new_required]) +def test_flag_on_with_a_changed_signature_keeps_the_aiter_path(load_sampler, op): + module, aiter = load_sampler(fused=True, op=op) + with pytest.warns(UserWarning, match="signature"): + sampler = module.Sampler() + assert sampler._fused_temperature_sample is None + logits, temps = _batch() + sampler(logits, temps, needs_independent_noise=True) + aiter.mixed_sample_outer_exponential.assert_called_once() + + +def test_real_fused_mixed_sample_takes_the_call(real_fused_mixed_sample, load_sampler): + module, _ = load_sampler(fused=True, op=real_fused_mixed_sample) + with warnings.catch_warnings(): + warnings.simplefilter("error") + sampler = module.Sampler() + assert sampler._fused_temperature_sample is real_fused_mixed_sample + + def test_flag_off_keeps_the_aiter_path(load_sampler): module, aiter = load_sampler(fused=False) sampler = module.Sampler() @@ -112,12 +158,19 @@ def test_flag_on_routes_the_per_row_noise_draw(load_sampler): assert kwargs == {"eps": sampler.eps, "greedy_eps": module.SAMPLER_EPS} -def test_no_config_for_the_gpu_falls_back_on_the_same_noise(load_sampler): +@pytest.mark.parametrize( + "error", + [ + FileNotFoundError("no config"), # no tuned config for this arch + ValueError("the last dim of input and exponentials must be contiguous"), + ], +) +def test_no_config_for_the_gpu_falls_back_on_the_same_noise(load_sampler, error): module, aiter = load_sampler(fused=True) sampler = module.Sampler() - load_sampler.fused_draw.side_effect = FileNotFoundError("no config") + load_sampler.fused_draw.side_effect = error logits, temps = _batch() - with pytest.warns(UserWarning, match="no config"): + with pytest.warns(UserWarning, match=type(error).__name__): sampler(logits, temps, needs_independent_noise=True) ((fused_args, _),) = load_sampler.fused_draw.call_args_list