[chelnnexy]

Summary

Speculative decoding draws the bonus token, plus every draft-conditioned verify row of a non-greedy batch (needs_independent_noise=True), with no top-k/top-p filter. On main that draw is five launches on the critical path of every verify step:

mixed_sample_outer_exponential (one program per row walking the whole 163840 vocab) → topk_select (2 launches) → <= → where

With ATOM_SPEC_BONUS_FUSED=1 the sampler calls aiter's fused_mixed_sample (ROCm/aiter#7416) instead: the same race and greedy override in two vocab-tiled Triton launches, on the same exponential_ noise. Default off. With the flag off, the kernels and their order are exactly as on main.

Depends on ROCm/aiter#7416. Until ATOM's aiter has fused_mixed_sample, setting the flag only warns once and keeps the aiter path, so this PR is safe to merge before or after it; the flag does nothing useful before.

What changes

• atom/model_ops/sampler.py
  • __init__ reads the flag once and imports aiter.ops.triton.fusions.fused_mixed_sample. An aiter without it: one warning, aiter path.
  • forward calls fused_mixed_sample(logits, noise, temperatures, eps=self.eps, greedy_eps=SAMPLER_EPS) only for an unfiltered batch with per-row noise. Top-k/top-p, shared-noise and all-greedy batches keep the aiter path.
  • aiter ships tuned configs for gfx942 and gfx950 only, and raises FileNotFoundError before launching elsewhere. The sampler then runs the aiter race on the noise it already drew (so the RNG stream is the flag-off one) and keeps the aiter path for later steps, with one warning.
  • Both per-row-noise paths draw their noise through one helper, get_per_row_exponential; the aiter race takes it through _race.
• atom/utils/envs.py, docs/environment_variables.md, docs/model_ops_guide.md: the flag, its dependency, and why it is off by default.
• tests/test_bonus_sample_wiring.py (CPU, no AITER or Triton needed).

The kernel, its tests and its tuned configs are in ROCm/aiter#7416, not here.

Measurements

Image rocm/atom-dev:nightly_202611011508-rocm10.0 (ROCm 10.0.0, torch 2.11.0+rocm10.0.0, triton 3.8.0), aiter 0dd271600 with #7416's 6 files at 6074827, this head 0a4e7ec. MI355X (gfx950) and MI300X (gfx942).

Tests. aiter's test_fused_mixed_sample.py: 39 passed on each arch. ATOM's wiring and sampler tests: 71 passed.

ATOM's Sampler, flag off vs on, on the same noise (sampler_check.py: 20 seeds × 256 rows, vocab 163840, bf16 logits). Greedy rows that differ: 0 on each arch. Sampling rows that differ, per temperature:

temperature	rows (each arch)	differ, MI355X	differ, MI300X
1e-10	860	0	0
0.01	860	0	0
0.1	860	0	0
0.3	860	0	0
0.7	840	0	0
1	840	0	0

Sampler.forward, noise draw included, median µs at T = 0.7:

rows	MI355X off	MI355X on	change	MI300X off	MI300X on	change
1	68.5	50.0	-27.0%	161.8	120.3	-25.6%
8	70.5	50.7	-28.1%	161.5	121.2	-25.0%
16	70.8	50.8	-28.2%	161.4	119.9	-25.7%
32	75.2	50.2	-33.2%	162.9	119.8	-26.5%
64	78.3	50.2	-35.9%	160.3	119.8	-25.3%
128	92.3	59.9	-35.1%	160.6	120.3	-25.1%
256	120.6	97.0	-19.6%	202.8	157.7	-22.2%

Kernel-level numbers (stage by stage, tuned configs) are in ROCm/aiter#7416.

End to end with the fused path armed: Kimi-K3 + DSpark on MI355X, recipes/Agentic-Kimi-K3.md §1 at CONC=4 (DCP1) without --spec-decode-acceptance-length, ATOM_SPEC_BONUS_FUSED toggled alone. Each run: 200 sampling requests (T 0.7 / 1.0, half with n=2) and 64 tool_choice="required" requests (T 0.7).

metric	flag off	flag on
failures (errors, empty, wrong n, broken tool calls)	0	0
arithmetic answers correct	300/300	300/300
MTP acceptance	65.7%	64.6% (64.8% in a repeat)

With the flag on, aiter's fused_mixed_sample ran 17,080 times during the requests (1 to 49 rows, BLOCK_V 1024 at 1 row and 2048 above, from the gfx950 config), with no fallback. The acceptance difference is within run-to-run noise (earlier pairs: 63.1 / 64.7%, 64.8 / 65.2%).

Why default off

Sampling rows are not bit-exact with the aiter kernel: they can differ when a row's two best race keys are within a few 1e-7 relative (fp32 rounding order), unlike ATOM_DSPARK_FUSED_MARKOV_SAMPLE. It also needs an aiter that has fused_mixed_sample.

Testing

• CPU, ATOM's Pre Checkin environment (.github/scripts/run_unit_tests.sh): 11124 passed, 0 failed. pytest tests/test_bonus_sample_wiring.py → 9 passed: flag off keeps the aiter ops; the flag wires aiter's fused_mixed_sample; the argument order; filtered, shared-noise and all-greedy batches never reach it; an aiter without the op, and a GPU without a config, fall back (the latter on the same noise, and for later steps). black and ruff are clean on the changed files.
• GPU, MI355X and MI300X: the tests, the Sampler check and (MI355X) the end-to-end run above.

Notes

• This may conflict textually with feat(spec-decode): probability-ratio rejection sampling behind draft_probs #2582 (sampler.py, envs.py). There is no semantic overlap: feat(spec-decode): probability-ratio rejection sampling behind draft_probs #2582 changes the verify rule, this PR changes the unfiltered draw.
• The out-of-range ids from aiter mixed_sample_outer_exponential at vocab 1000 and 4097 (vocab not a multiple of 4, or below 4096) are fixed upstream in [HIP] [Bugfix] sample_kernels: partial-vector and prologue-tail defects aiter#7373. Vocab 163840 (production) was never affected.

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2914 --add-label <label>

[chelnnexy]

Validation on MI355X (gfx950)

Ran on an 8×MI355X host (smci355-ccs-aus), image rocm/atom-dev:nightly_202611011508-rocm10.0 (torch 2.11.0+rocm10.0, triton 3.8.0, aiter).

Unit tests

• pytest tests/test_bonus_sample_wiring.py → 6 passed.
• pytest tests/test_bonus_sample.py (Triton CPU interpreter) → 9 passed, 1 skipped, including two added here: determinism (same (logits, temps, noise) → same token) and a dominating race key in the final vocab tile beating tile 0 (exercises the (row,) kernel's k * exp(m - mm) cross-tile rescale).

GPU bit-exact (real Triton kernels on CUDA, not the interpreter)

Compared bonus_temperature_sample against the float64 definition of the race it replaces — argmax_i softmax(x / max(T, 1e-5))_i / (E_i + eps), and for a greedy row the lowest-id raw-logit argmax — across vocab 1000 → 163840, 24 rows with planted greedy ties:

PASS (24/24 rows) — greedy rows bit-exact, sampling rows matched with no near-tie exceptions needed at these seeds.

End-to-end (both decode bands)

Full GSM8K 1319, 5-shot, 64 concurrency, ATOM_SPEC_BONUS_FUSED toggled alone, on DCP1 (interactive band) and DCP8 (conc-8 band):

Band	metric	OFF	ON
DCP1	acceptance / toks-fwd	61.70% / 4.09	62.03% / 4.10
DCP1	GSM8K flexible / strict	0.9416 / 0.9424	0.9530 / 0.9522
DCP8	acceptance / toks-fwd	61.85% / 4.09	61.89% / 4.09
DCP8	GSM8K flexible / strict	0.9515 / 0.9507	0.9507 / 0.9515

Important nuance: the fused path is gated on needs_independent_noise (non-greedy / n>1 fan-out). GSM8K is greedy (do_sample=False, n=1), so the fused draw is never taken on this workload — ON and OFF run identical code, and the ON-vs-OFF spread above is run-to-run noise between separate server launches (the OFF baseline alone moved 61.70→61.85 / 0.9416→0.9515). So GSM8K only shows "no regression when enabled"; it is not a neutrality probe for the fused draw. Unlike the sibling ATOM_DSPARK_FUSED_MARKOV_SAMPLE (which samples the draft block every step), this flag isn't exercised by greedy decoding.

The decisive correctness evidence is the GPU bit-exact 24/24 and the 3200/3200 agreement above, which run the actual fused path.

[yitingw1]

@valar-cop review this pr

[valar-cop]

some changes needed..

e2e with the fused path armed (temperature>0 or n>1)? and a gfx942 run for the new kernel?

[valar-cop]

all -inf tile here gives -inf - -inf = NaN.. reachable with required tool_choice, please guard m == -inf and add a test

[chelnnexy]

Confirmed on MI355X (gfx950) and MI300X (gfx942): the current kernel gives the same NaN result on both. The fix follows @yitingw1's version: guard the tile max and remove the clamp. Full results are in the top-level comment.

Before the fix (PR head 680d7359), gpu_check.py:

[1] NaN partial keys on a forced row: 78 / 80   (unpatched: >0 expected = bug present; patched: must be 0)
[2] forced rows wrong: 0 / 7 
[3] masked rows: 0 unexplained mismatches, 0 near-ties, 0 out-of-range ids, over 1600 rows
[4] all -inf rows out of range: 0 / 4 [163839, 163839, 0, 163839]
FAIL (1 check(s))

There are NaN keys in 78 of 80 tiles for a forced row. The forced id still wins only because tl.max drops NaN. Separately, on a row that is -inf throughout, the sampling result is 163839: that is the vocab sentinel after tl.minimum(tok, vocab - 1), which shows the clamp is what's hiding it.

Fix:

@@ _bonus_tile_kernel
-    key = tl.where(mask, tl.exp(xs - m) / e, -1.0)
+    m_safe = tl.where(m == -float("inf"), 0.0, m)
+    key = tl.where(mask, tl.exp(xs - m_safe) / e, -1.0)
@@ _bonus_pick_kernel
         mm = tl.max(m, axis=0)
+        mm = tl.where(mm == -float("inf"), 0.0, mm)
         ks = tl.where(tm, k * tl.exp(m - mm), -1.0)
@@
-    tok = tl.minimum(tok, vocab - 1)

• Cost: one scalar where per program in each kernel. The per-element work is unchanged.
• Rows with no -inf logits: results are unchanged. Checked against the current kernel on 80 rows at vocab 1000 to 163840.
• All--inf tile: keys every id at 0, so it loses to any tile with a finite logit, and its idx stays in range.
• Pick kernel: the mm guard covers a row that is -inf throughout. That row now returns id 0 instead of the clamped sentinel.
• The clamp: with the guard in place it can no longer change a result, so it's removed. Any future escape will now show up as an out-of-range id instead of a plausible but wrong token.

Tests in tests/test_bonus_sample.py, with RuntimeWarning turned into an error:

• test_a_forced_channel_row_samples_the_forced_id: forced id at 5 / 100000 / 163837, for both T=1.0 and greedy rows.
• test_masked_ids_never_win_and_every_id_is_in_range: whole -inf tiles plus scattered -inf ids, compared against the float64 race, plus one row that is -inf throughout.

Results under the interpreter: the current kernel fails 4 of these, and with the fix the file gives 13 passed, 1 skipped. black and ruff report nothing.

After the fix, gpu_check.py:

MI355X (gfx950):

[1] NaN partial keys on a forced row: 0 / 80   (unpatched: >0 expected = bug present; patched: must be 0)
[2] forced rows wrong: 0 / 7 
[3] masked rows: 0 unexplained mismatches, 0 near-ties, 0 out-of-range ids, over 1600 rows
[4] all -inf rows out of range: 0 / 4 [0, 0, 0, 0]
PASS

MI300X (gfx942):

[1] NaN partial keys on a forced row: 0 / 80   (unpatched: >0 expected = bug present; patched: must be 0)
[2] forced rows wrong: 0 / 7 
[3] masked rows: 0 unexplained mismatches, 0 near-ties, 0 out-of-range ids, over 1600 rows
[4] all -inf rows out of range: 0 / 4 [0, 0, 0, 0]
PASS

Pushed as e34da58. The branch is also rebased onto main, which brings in the loader fix from #2936.

[valar-cop]

why default off? the sibling ATOM_DSPARK_FUSED_MARKOV_SAMPLE is on by default..

[chelnnexy]

It stays off on purpose. The reason is now in the flag's comment and in environment_variables.md (c70d5bf):

• Not bit-exact: sampling rows can differ from the aiter kernel on fp32 near-ties (~1e-7 relative).
• The gain shrinks with verify rows: on MI355X it's about −25% up to 128 rows, but nothing at 256 (about Kimi-K3 DSpark at concurrency 64).

ATOM_DSPARK_FUSED_MARKOV_SAMPLE doesn't have that trade-off. The tuned aiter version (see the aiter thread) is faster at 256 rows too; switching to it is the point to revisit the default.

[yitingw1]

+1 on valar-cop's NaN finding — it is reachable exactly as described. Two refinements that change what the fix should be, plus a few things not yet raised.

1. The NaN is real and reachable, but it does not produce a wrong token today. Worth knowing before you fix it, so you don't change a result that is currently correct. required_channel.py:154 stores neg_inf across the whole row except the forced id, so at vocab 163840 with _BONUS_BLOCK_V=4096, 39 of 40 tiles are entirely -inf: m = -inf, exp(xs - m) = exp(NaN) = NaN, and key == best is false everywhere so idx = vocab. But tl.max defaults to PropagateNan.NONE (standard.py:189 → core.maximum → semantic.py:380 → create_maxnumf), so in _bonus_pick_kernel the NaN tiles lose to the one finite tile and the right token still wins. The correctness of this path is resting on a Triton default that nothing in the file mentions.

The greedy half is genuinely safe and needs no guard: -inf == -inf is true, so gmax/gidx never go NaN, even for an all--inf row (required_channel.py:161 notes the non-forced branch has no all--inf guard). Your bit-exact 24/24 greedy evidence is unaffected.

2. So the thing to remove is the clamp, not just the NaN. tok = tl.minimum(tok, vocab - 1) is what turns a sentinel escape into a plausible token id instead of an obvious bad index. If it is unreachable it is dead code; if it is reachable it silently returns the last token. Either way, guarding m == -inf and keeping the clamp leaves the landmine armed for whenever that Triton default shifts or someone passes PropagateNan.ALL. Suggest:

m = tl.max(xs, axis=0)
m_safe = tl.where(m == -float("inf"), 0.0, m)
key = tl.where(mask, tl.exp(xs - m_safe) / e, -1.0)

An all--inf tile then contributes key = 0, loses cleanly, writes an in-range idx, and no NaN or out-of-range value ever reaches part_key / part_idx — at which point the clamp is provably dead and can go.

3. No test uses a -inf logit. All five cases in test_bonus_sample.py build logits with torch.randn. A row that is -inf everywhere except forced, with T=1.0, asserting the fused draw returns forced, pins both the current behaviour and the fix above.

4. The microbench percentages have a cost excluded from the denominator that is paid on the same critical path every step. Excluding the exponential_ draw is correct for the A/B — it is identical on both arms, so the 15–22 µs delta is honest. But at rows=64 that tensor is 64 × 163840 × 4 B = 41.9 MB, the same order as the 44 µs the fused path takes, so −24.8% to −34.5% overstates the step-level impact. Adding one row for the noise draw's own time (you already measured it) lets the table speak for itself.

5. rows in this path is verify rows, not batch size — sampler.py:96 sets it to the batch's total draft tokens, so your own Kimi-K3 DSpark MTP at concurrency 64 is around 256, 4× past the table's top row. The trend is favourable (−24.8% → −34.5% as rows grow), so nothing is hidden; one rows=256 point would just pin the production shape.

6. The aiter side finding will get lost where it is. "mixed_sample_outer_exponential returned out-of-range token ids at vocab 1000 and 4097" is an out-of-range return from a production op, filed as the second bullet of a Notes section in an unrelated ATOM PR. It deserves its own aiter issue. It also carries weight for this PR: it is why the vocab 1000 / 4097 cases can only compare against the float64 reference rather than against aiter, which is the right call but is not currently connected to the small-vocab coverage claim.

7. On "why default off" — the answer is already in the PR body, it just isn't wired to the question. Sampling rows are not bit-exact (the ~1e-7 near-tie divergence), unlike ATOM_DSPARK_FUSED_MARKOV_SAMPLE, and the GSM8K table never exercises the fused path (which you say outright, and that disclosure is good). Default-off is the right call; worth stating that reason in the env var doc.

---

Things I checked and found correct, so nobody re-derives them: the two-stage rescale identity (key_t * exp(m_t - mm) = exp(xs - mm)/e); underflow is safe (exp(m - mm) <= 1, a flushed tile could not have won); tie-break equals tie="low" (tiles are id-ordered, min id within tile, min tile index across); the greedy half correctly uses raw logits, matching _apply_greedy_tokens whose probs parameter is actually called with logits at forward:151; return dtype matches (topk_select's indices_type defaults to int32); RNG consumption is identical; r.to(tl.int64) * stride_l is properly widened; and cudagraph is not a concern since _temperature_sample already allocates on this same path.

[chelnnexy]

MI355X (gfx950) and MI300X (gfx942) results for the review asks

Summary

• NaN on forced rows: both GPUs show NaN keys in 78 of 80 tiles on the current kernel, and none after the fix (see the reply on the bonus_sample.py:58 thread).
• Agreement with main: 0 mismatches on both GPUs, including rows masked to -inf. Each case is 3200 rows, with the same noise on both paths.
• Cost of the fix: none measurable. Fused timings before and after are within 1% on gfx950 and within 2% on gfx942, where main moved by a similar amount between the two runs.
• End to end on MI355X: Kimi-K3 + DSpark at T > 0 (including n=2) and tool_choice="required", flag OFF vs ON. 0 failures and 300/300 correct both times. MTP acceptance was 63.1% (OFF) vs 64.7% (ON). A first pair of runs gave 64.8% vs 65.2%, so the gap is within run-to-run noise.
• Speed at production batch size (@yitingw1 fix the example code cmd #5): the speedup shrinks as rows grow and disappears at 256 verify rows, which is roughly Kimi-K3 DSpark at concurrency 64. On gfx950 the fused draw takes 78.8 µs against main's 78.3 µs (+0.7%); on gfx942 it is 144.7 vs 149.7 µs (−3.3%). The fix doesn't cause this: the unpatched kernel shows the same (+0.9% on gfx950). The gain is real up to 128 rows: −24 to −36% on gfx950 and −27 to −31% on gfx942. It may be worth tuning the tile size or warp count for large row counts, or gating the fused path on row count.

All runs have the NaN fix applied, now pushed as e34da58 (see the reply on the bonus_sample.py:58 thread). bonus_sample.py and its tests are identical to what was tested here, and the rebase onto main didn't touch them. Each section names the GPU, image and versions it ran on.

MI355X (gfx950)

Image: rocm/atom-dev:nightly_202611011508-rocm10.0, torch/triton/aiter: 2.11.0+rocm10.0.0 / 3.8.0 / amd-aiter 0.1.28.dev158+g0dd271600

End-to-end run that actually reaches the fused path (@valar-cop)

Kimi-K3 + DSpark, recipes/Agentic-Kimi-K3.md §1 at CONC=4 (DCP1), without --spec-decode-acceptance-length, with ATOM_SPEC_BONUS_FUSED toggled on its own. Each run sends 200 sampling requests (T = 0.7 / 1.0, half of them n=2, top_p=1, top_k=-1) and 64 tool_choice="required" requests (T = 0.7):

metric                         off          on
wall_s                        64.3        63.1
sampling_requests              200         200
sampling_errors                  0           0
sampling_wrong_n                 0           0
sampling_choices               300         300
sampling_empty                   0           0
sampling_correct               300         300
required_requests               64          64
required_errors                  0           0
required_broken                  0           0
mtp_acceptance              0.6311      0.6474
hard_failures                    0           0

Note: before the rebase, the branch didn't start in this image. The weight loader failed with undefined symbol: omp_set_num_threads, which main fixed in #2936. The table above is from a run with #2936 applied and no LD_PRELOAD. An earlier pair of runs using LD_PRELOAD=libgomp.so.1 instead gave the same result: 0 failures, 300/300, acceptance 64.8% / 65.2%. The branch is now rebased onto main, which includes #2936.

Kernel: -inf checks before and after the fix

Before the fix (PR head, expected to fail: this is the bug):

[1] NaN partial keys on a forced row: 78 / 80   (unpatched: >0 expected = bug present; patched: must be 0)
[2] forced rows wrong: 0 / 7 
[3] masked rows: 0 unexplained mismatches, 0 near-ties, 0 out-of-range ids, over 1600 rows
[4] all -inf rows out of range: 0 / 4 [163839, 163839, 0, 163839]
FAIL (1 check(s))

After the fix:

[1] NaN partial keys on a forced row: 0 / 80   (unpatched: >0 expected = bug present; patched: must be 0)
[2] forced rows wrong: 0 / 7 
[3] masked rows: 0 unexplained mismatches, 0 near-ties, 0 out-of-range ids, over 1600 rows
[4] all -inf rows out of range: 0 / 4 [0, 0, 0, 0]
PASS

Agreement with main, same noise (3200 rows per case)

case                            same  near-tie  UNEXPLAINED
mask=none T=1.0                 3200         0            0
mask=none T=0.7                 3200         0            0
mask=none T=greedy              3200         0            0
mask=partial T=1.0              3200         0            0
mask=partial T=0.7              3200         0            0
mask=partial T=greedy           3200         0            0
mask=forced T=1.0               3200         0            0
mask=forced T=0.7               3200         0            0
mask=forced T=greedy            3200         0            0

Microbench including the noise draw and rows up to 256 (@yitingw1 #4, #5)

rows is verify rows; Kimi-K3 DSpark at concurrency 64 is about 256. step = exponential_ noise draw + the draw itself, which is what each verify step actually pays.

 rows     main    fused   delta    noise  step main  step fused  step delta  fused(forced)
    1    60.08    44.76  -25.5%     7.20      67.28       51.96      -22.8%          44.56
    4    61.42    44.60  -27.4%     7.72      69.14       52.32      -24.3%          44.92
    8    62.36    44.68  -28.4%     8.28      70.64       52.96      -25.0%          44.84
   16    62.60    44.64  -28.7%    10.96      73.56       55.60      -24.4%          44.60
   32    64.08    45.04  -29.7%    14.16      78.24       59.20      -24.3%          44.84
   64    66.88    43.16  -35.5%    19.16      86.04       62.32      -27.6%          43.42
  128    70.84    54.06  -23.7%    29.56     100.40       83.62      -16.7%          54.00
  256    78.28    78.80    0.7%    50.44     128.72      129.24        0.4%          79.06

Fused-kernel timings before and after the NaN fix (median µs, every row within ±1%):

 rows   fused before   fused after
    1          44.84         44.76
    4          44.60         44.60
    8          44.92         44.68
   16          44.82         44.64
   32          45.12         45.04
   64          43.56         43.16
  128          54.36         54.06
  256          79.10         78.80

MI300X (gfx942): the kernel run (@valar-cop)

Image: rocm/atom-dev:nightly_202611011508-rocm10.0, torch/triton/aiter: 2.11.0+rocm10.0.0 / 3.8.0 / amd-aiter 0.1.28.dev158+g0dd271600

Kernel: -inf checks before and after the fix

Before the fix (PR head, expected to fail: this is the bug):

[1] NaN partial keys on a forced row: 78 / 80   (unpatched: >0 expected = bug present; patched: must be 0)
[2] forced rows wrong: 0 / 7 
[3] masked rows: 0 unexplained mismatches, 0 near-ties, 0 out-of-range ids, over 1600 rows
[4] all -inf rows out of range: 0 / 4 [163839, 163839, 0, 163839]
FAIL (1 check(s))

After the fix:

[1] NaN partial keys on a forced row: 0 / 80   (unpatched: >0 expected = bug present; patched: must be 0)
[2] forced rows wrong: 0 / 7 
[3] masked rows: 0 unexplained mismatches, 0 near-ties, 0 out-of-range ids, over 1600 rows
[4] all -inf rows out of range: 0 / 4 [0, 0, 0, 0]
PASS

Agreement with main and microbench

== Agreement: fused vs main (same noise), 16 rows/batch ==
case                            same  near-tie  UNEXPLAINED
mask=none T=1.0                 3200         0            0
mask=none T=0.7                 3200         0            0
mask=none T=greedy              3200         0            0
mask=partial T=1.0              3200         0            0
mask=partial T=0.7              3200         0            0
mask=partial T=greedy           3200         0            0
mask=forced T=1.0               3200         0            0
mask=forced T=0.7               3200         0            0
mask=forced T=greedy            3200         0            0

== Microbench (median us) ==
noise = the exponential_ draw both paths pay first; step = noise + draw.
rows is verify rows (total draft tokens): Kimi-K3 DSpark at conc 64 ~ 256.
 rows     main    fused   delta    noise  step main  step fused  step delta  fused(forced)
    1   136.64    94.40  -30.9%    14.14     150.79      108.54      -28.0%          93.94
    4   137.31    94.56  -31.1%    13.78     151.09      108.34      -28.3%          94.90
    8   137.31    94.46  -31.2%    14.10     151.41      108.56      -28.3%          95.50
   16   136.72    94.62  -30.8%    16.52     153.25      111.15      -27.5%          95.32
   32   137.33    95.14  -30.7%    22.35     159.68      117.50      -26.4%          94.94
   64   136.72    94.58  -30.8%    32.01     168.73      126.59      -25.0%          95.34
  128   136.74    99.63  -27.1%    48.59     185.34      148.22      -20.0%          94.74
  256   149.70   144.70   -3.3%    80.04     229.74      224.74       -2.2%         132.60

PASS

End-to-end (optional): GPT-OSS-120B, 1 GPU, all requests n=2 at T = 0.7 / 1.0

Kimi-K3 is gfx950-only, so here the fused draw is reached through n>1 fan-out rather than speculative decoding.

(not run: optional on MI300X, and no local gpt-oss-120b model was available)

[valar-cop]

a few small things..

[valar-cop]

stub bonus_sample in the fixture like aiter, drop importorskip so this runs on cpu?

[chelnnexy]

Done in a8347e8. The fixture now fakes atom.model_ops.bonus_sample the same way it fakes aiter, and the importorskip is gone. A new test checks that the flag wires Sampler to bonus_temperature_sample.

With Triton blocked (as on a CPU runner) all 7 tests pass. Before, the whole file was skipped.

[valar-cop]

still allocate 5 buffers per call on the hot path.. cache them keyed by (n_tiles, rows)?

[chelnnexy]

I tried it and measured it, and kept torch.empty.

A cached buffer, keyed by (device, stream) and grown on demand, made the draw slower: about +2.2 µs per call at every row count on MI355X. The main path stayed within ±0.4 µs between the two runs, so this is not run-to-run noise. The cost is in Python: looking up the current stream builds a Stream object on every call, and slicing views out of the buffer creates as many tensor objects as torch.empty does. The memory itself is not allocated per call. PyTorch's caching allocator already hands back the same blocks, so there is nothing to save there.

MI355X, median µs of 300 calls:

rows	main path	fused, torch.empty	fused, cached buffer
1	60.08 → 59.64	44.76	46.88 (+2.12)
4	61.42 → 61.08	44.60	46.80 (+2.20)
8	62.36 → 62.12	44.68	46.96 (+2.28)
16	62.60 → 62.24	44.64	46.66 (+2.02)
32	64.08 → 64.12	45.04	47.36 (+2.32)
64	66.88 → 66.52	43.16	45.36 (+2.20)
128	70.84 → 70.66	54.06	57.08 (+3.02)
256	78.28 → 78.28	78.80	81.44 (+2.64)

So the five torch.empty calls stay.

[valar-cop]

this is a second impl of aiter's mixed_sample_outer_exponential.. why not put it in aiter and call here?

[chelnnexy]

Agreed, and it's moved: ROCm/aiter#7416 adds it to aiter as fused_mixed_sample (Triton, in aiter/ops/triton/fusions/, with configs tuned per arch, a unit test and a benchmark).

It uses the same race and the same noise as mixed_sample_outer_exponential (that's why the agreement check is exact), but it isn't the same kernel: aiter's kernel runs one workgroup per row over the whole vocab, while this one tiles each row across the grid and folds in the greedy override (topk_select + <= + where), so the draw takes 2 launches instead of 5. In the aiter version, greedy rows also skip the race and the noise read.

With the tuned configs, the aiter version is faster than the HIP path at every row count we measured, including 256 rows (−33% on MI355X, −27% on MI300X), where the kernel here is break-even.

#2914 doesn't need to wait for it: the kernel here stays behind the default-off flag. Once a release with fused_mixed_sample is in ATOM's image, a small follow-up switches sampler.py to it and deletes atom/model_ops/bonus_sample.py.

Related: the existing aiter kernel's out-of-range ids at small vocab sizes, from the PR description, were fixed today in ROCm/aiter#7373.

[valar-cop]

same exponential_ draw spelled twice, fused path relies on same rng.. one helper for both?

[chelnnexy]

Done in 6a9bc8f. get_per_row_exponential(num_tokens, vocab_size, device) in sampler.py is now the only place the per-row draw happens. _temperature_sample and the fused branch both call it, so they can't drift apart and they consume the RNG identically.

[chelnnexy]

Follow-up to the review (3 commits on top of e34da58)

Review comment	Answer
Stub bonus_sample in the wiring fixture, drop importorskip	a8347e8
exponential_ draw written twice	6a9bc8f: get_per_row_exponential
Why default off	c70d5bf: reason documented
Five buffers allocated per call	Measured a cached buffer: +2.2 µs per call on MI355X, so torch.empty stays (numbers in the thread)
Move into aiter	ROCm/aiter#7416 (fused_mixed_sample, tuned for gfx950 and gfx942). #2914 doesn't depend on it

@yitingw1's points

• move from private repo to ROCm #1–update the link to the repo #3 (the NaN is real; remove the clamp; add a -inf test): done in e34da58, see the reply on the NaN thread.
• Update README.md and pyproject.toml #4–fix the example code cmd #5 (include the noise-draw cost; add a rows=256 point): in the results comment above. The draw here is break-even at 256 rows on MI355X, which is also why the flag stays off. The tuned aiter version is 33% faster than the HIP path at 256 rows.
• update readme for deepseek #6 (the aiter out-of-range ids at vocab 1000 / 4097): fixed upstream in [HIP] [Bugfix] sample_kernels: partial-vector and prologue-tail defects aiter#7373, thanks.
• support gpt oss #7 (state the default-off reason in the env doc): c70d5bf.

Validation

bonus_sample.py and its kernel tests are unchanged from e34da58 (GPU results in the comment above). The three commits change sampler.py, the wiring tests and docs.

• CPU: the full CI unit suite (.github/scripts/run_unit_tests.sh) passes on the new head, and black and ruff are clean.
• MI355X, no LD_PRELOAD: the server starts on the rebased branch (main's fix(loader): resolve OpenMP symbol through torch #2936 loader fix is in). The end-to-end run with the follow-up passes with the flag OFF and ON:

metric                         off          on
wall_s                        64.5        68.2
sampling_requests              200         200
sampling_errors                  0           0
sampling_wrong_n                 0           0
sampling_choices               300         300
sampling_empty                   0           0
sampling_correct               300         300
required_requests               64          64
required_errors                  0           0
required_broken                  0           0
mtp_acceptance               0.648      0.6493
hard_failures                    0           0

Kernel checks:

[1] NaN partial keys on a forced row: 0 / 80   (unpatched: >0 expected = bug present; patched: must be 0)
[2] forced rows wrong: 0 / 7 
[3] masked rows: 0 unexplained mismatches, 0 near-ties, 0 out-of-range ids, over 1600 rows
[4] all -inf rows out of range: 0 / 4 [0, 0, 0, 0]
PASS

That GPU run also included the cached-buffer commit, which was then dropped. The draw it measured is the same draw.

[yitingw1]

The previous round is fully addressed, as far as I can tell — the m_safe guard and the dropped clamp, the -inf tests, the noise column in the microbench, the rows=256 point, the "Why default off" section, and ROCm/aiter#7373 for the out-of-range ids. Nothing from that list is still open.

Two notes from re-reading it with fresh eyes, neither blocking.

Coverage reads thinner than it is. pytest tests/test_bonus_sample.py reports 8 passed, 13 skipped, which looks like most of the file is off. It isn't — the 13 run in the TRITON_INTERPRET=1 child, and I confirmed that directly (13 passed there). Since interpreter_test asserts the child's returncode == 0, and pytest returns 5 when it collects nothing, the empty-child hole is closed too. Worth a line in the test module's docstring, since the headline number is what a reviewer sees. I also mutated the two subtle guards to check they're pinned: dropping m_safe reds 4 tests, and flipping the greedy tie-break to the highest id reds 2.

A row-count gate would answer the second of the two default-off reasons. The crossover is now pinned on both arches (~256 rows), so rows <= N → fused, else aiter, would make the flag safe to turn on by default as far as performance is concerned — it would stop being a flag that is slower on the production shape, which is what currently argues against anyone enabling it. That leaves the first reason, the ~1e-7 near-tie divergence, as the only thing keeping it opt-in, which is a much narrower question than the two reasons together. If you'd rather keep both, it's worth saying in the env var doc that the gate is the missing piece, so the next person knows what would change the default.

Nit: self.eps and SAMPLER_EPS are passed as two arguments but are the same constant — one is the race denominator, the other the greedy threshold. Fine today; the signature just suggests they can diverge.

[chelnnexy]

Thanks for re-reading it, and for running the interpreter child and the two mutations.

• Coverage: 5f35fbb adds a paragraph to the test module docstring. It says that a plain run reports most tests as skipped because they run in the TRITON_INTERPRET=1 child, and that the child fails the run if any of them fails or if it collects nothing.
• Row-count gate: I'd keep both reasons for now. 5f35fbb says in the env var doc (and in the envs.py comment) what would change the default. A row-count gate removes the speed reason. So does ATOM moving to aiter's fused_mixed_sample ([Triton/Gluon] [Kernel] Add fused_mixed_sample: tiled exponential-race sampling with greedy override aiter#7416), which is faster than the aiter path at 256 rows too, on MI355X (51.6 µs against 77.7 µs) and MI300X (99.1 µs against 133.0 µs), with no gate. The near-tie difference would then be the only reason left.
• eps and SAMPLER_EPS: they are kept as two arguments because they are two roles, as in main's path. self.eps is the race denominator, which main passes to mixed_sample_outer_exponential as eps. SAMPLER_EPS is the greedy threshold, which main's _apply_greedy_tokens compares against. Both are SAMPLER_EPS today.

[yitingw1]

Thanks — the three points are answered, and I'll withdraw the eps one: two roles that happen to share a constant is the right reading, and it matches main's path.

You were also right that fused_mixed_sample removes the speed reason without a gate. I should have led with that last round instead of proposing the weaker fix — "Move into aiter → ROCm/aiter#7416" was already in your 21:42 table.

Having read both, I think the scope question comes before anything else here.

By line count, 410 of this PR's 592 additions are the part ROCm/aiter#7416 now covers: bonus_sample.py (170), its kernel tests (190), and triton_interpreter.py (50), which exists only because the kernel lives in ATOM. The remaining ~180 — sampler.py, the env var, the wiring tests — is ATOM-side work that has to exist whichever kernel it calls. And the swap is one line: the two wrappers take the same arguments, only reordered.

bonus_temperature_sample(logits, temperatures, exponential, eps, greedy_eps)
fused_mixed_sample(input, exponentials, temperatures, eps=…, greedy_eps=…, out=…)

Two things I found while comparing them argue the same way — the in-tree copy isn't just a duplicate, it's the weaker of the two.

1. The ~1e-7 window holds for #7416's scaling order, not for this one. #2914 computes xs = x*inv_t; m = max(xs); exp(xs - m_safe); #7416 computes m = max(x); exp((x - m_safe)*inv_t). Algebraically identical; not in fp32, because #2914 subtracts between numbers of magnitude |x|/T. Replicating both orders on CPU (V=163840, 400 draws per point, float64 reference, worst relative error among the keys that decide the argmax):

T	#2914	#7416
1.0	9.5e-08	9.5e-08
0.7	3.3e-06	3.9e-07
0.3	1.0e-05	1.8e-07
0.1	2.3e-05	8.0e-08
0.01	1.4e-04	5.6e-08

They agree exactly at T = 2.0 / 1.0 / 0.5, where x*inv_t is exact — which is what identifies that multiply as the source. Elsewhere the window grows like 1/T. Not a correctness issue: both are valid draws from the same race, and top-2 gaps below 2.3e-5 occur in about 1 row in 20000. But ~1e-7 is the number a reader uses to size the risk of enabling the flag, and at T=0.1 it is off by ~300×. Your own PR body on #7416 gives this reordering as its reason.

2. Stage 1 here computes both paths for every row and stores 5 partials; #7416 branches on greedy_eps in stage 1 and stores 3. Since you measured the buffer allocation at only +2.2 µs, the extra path looks like the better candidate for the missing gain at 256 rows — the shape where #7416 is −33%, and the shape that currently keeps this flag off. I haven't measured this; it's an inference from your numbers.

→ Change: land the kernel in #7416 and reduce this PR to the ATOM side — sampler.py calling fused_mixed_sample, the env var, the wiring tests. If you'd rather land this first and swap later, please say so in the description, so the in-tree copy doesn't outlive its purpose.

One thing I couldn't settle: #7416's _get_config raises FileNotFoundError on an arch with no tuned config, and only gfx950/gfx942 ship one; this kernel's fixed BLOCK_V=4096, num_warps=8 runs anywhere. If ATOM needs an arch outside those two, that's a real argument for the in-tree copy and I'd drop the scope point.

[chelnnexy]

Thanks — agreed on all of it.

• Scope: 0a4e7ec calls aiter's fused_mixed_sample ([Triton/Gluon] [Kernel] Add fused_mixed_sample: tiled exponential-race sampling with greedy override aiter#7416) and drops bonus_sample.py, its kernel tests and triton_interpreter.py. Left: sampler.py, the flag, docs, wiring tests (+119 / −457).
• ~1e-7 window: right, it held only for #7416's order. Docs now say "a few 1e-7". Through ATOM's Sampler on the same noise (5,120 rows, T from 1e-10 to 1), 0 rows differ from the aiter race on MI355X and MI300X.
• Five partials: gone with the in-tree kernel; #7416 branches on greedy_eps and keeps three.
• Arch without a config: the op raises FileNotFoundError before launching; the sampler falls back to the aiter race on the noise already drawn (same RNG stream as flag off) and warns once. So other archs just keep main's path — no in-tree copy needed. Covered by the wiring tests.
• 256 rows: now faster too — Sampler.forward −19.6% at 256 and −27 to −36% below on MI355X, −22 to −26% on MI300X. MI355X end to end: 17,080 fused draws, 300/300 correct, 0 failures.

Depends on #7416 landing in ATOM's aiter; until then the flag just warns. Description updated.