diff --git a/atom/model_ops/sampler.py b/atom/model_ops/sampler.py index c69cb618ff..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 @@ -9,6 +10,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 @@ -60,11 +63,53 @@ 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) + + +def _load_fused_mixed_sample(): + """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: + 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): 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: + self._fused_temperature_sample = _load_fused_mixed_sample() def sample_verification_tokens( self, @@ -144,6 +189,37 @@ 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 aiter's two vocab-tiled launches. + num_tokens, vocab_size = logits.shape + exponential = get_per_row_exponential( + num_tokens, vocab_size, logits.device + ) + try: + return self._fused_temperature_sample( + logits, + exponential, + temperatures, + eps=self.eps, + greedy_eps=SAMPLER_EPS, + ) + 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( + 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, + ) + 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 ) @@ -186,15 +262,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 = 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 ) + 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 10f7d4349a..e684b48b2e 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -1015,6 +1015,26 @@ 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 + # 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 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", # --- 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..fc632d68e9 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) 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 e5cb9dfa40..e1388f7609 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), 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`) Holds a request with a `RequiredChannel` (`atom/sampling_params.py`; today only diff --git a/tests/test_bonus_sample_wiring.py b/tests/test_bonus_sample_wiring.py new file mode 100644 index 0000000000..65a02fb54b --- /dev/null +++ b/tests/test_bonus_sample_wiring.py @@ -0,0 +1,213 @@ +# 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. 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. 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 +import torch + +SAMPLER_PATH = ( + pathlib.Path(__file__).resolve().parents[1] / "atom" / "model_ops" / "sampler.py" +) +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 ( + "aiter", + "aiter.ops", + "aiter.ops.sampling", + "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"] + 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), + ) + ) + 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, 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) + 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 + + load.fused_draw = fused_module.fused_mixed_sample + return load + + +def _batch(rows=4, vocab=64): + return torch.randn(rows, vocab), torch.ones(rows) + + +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 _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() + 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): + module, aiter = load_sampler(fused=True) + sampler = module.Sampler() + 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() + ((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 kwargs == {"eps": sampler.eps, "greedy_eps": module.SAMPLER_EPS} + + +@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 = error + logits, temps = _batch() + with pytest.warns(UserWarning, match=type(error).__name__): + 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( + "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()