[yitingw1]

Motivation

Every sampler in csrc/kernels/sample_kernels.cu can return a token id outside [0, vocab) — negative, or past the end of the row — and exponential_kernel in the same file can write one element past its row. All of it comes from one shape: the row's last vector is partial when N % VecSize != 0, and the buffer bound meant to cover for that is enforced per dword, so it cannot separate element N-1 from element N at a 2-byte dtype.

Surfaced as a side finding while reviewing ROCm/ATOM#2914, reported there at vocab 1000 and 4097 on gfx950; reproduced here on gfx942, so it is not arch-specific.

Current production vocabularies happen to sit in the safe region (129280, 151936, 163840 are all multiples of the vector width and far above one block's span), which is why this has gone unnoticed. Any model with a small vocabulary, any draft head, or any vocabulary that is not a multiple of the vector width gets a silently wrong token.

Technical Details

Root cause

All three *_impl functions share one software-pipelined prologue/tail structure, and both defects live in it:

int k          = threadIdx.x * vec_size_i;
int vec_stride = BlockSize * vec_size_i;
vec_i vec_pre;                                   // <-- left undefined below
if(k < N) { vec_pre = load(...); k += vec_stride; }
for(; k < N; k += vec_stride) { ... loop(); }
if((k - vec_stride) < N) { loop(); }             // <-- tail

// inside loop(), with no bound on i:
thread_kvp.key = k - vec_stride + i;

A — undefined registers and a negative index. A thread whose first element is already past N fails the prologue check, so vec_pre stays undefined and k is never advanced. At the tail, k - vec_stride is negative, < N passes, and the body runs on undefined registers and writes a negative key.

Fires when N <= (BlockSize - 1) * VecSize.

B — out-of-range lanes in the last vector. When N % VecSize != 0 the row's last vector is partial and its trailing lanes fall outside the row. What they pick up differs by which descriptor they come from, and that decides how easily the defect fires — not whether it is present:

• buffer_e (the caller's noise) is built with extent exactly N, so a lane past N yields 0. That 0 lands in the divisor, exp(x - max) / (0 + eps) is enormous, and it wins the argmax for any logits. This is the outer-exponential path.
• buffer_i (the logits) has its extent rounded up to a dword (oob_i), so a lane past N yields 0 at a 4-byte dtype and, for the first lane past an odd N at a 2-byte dtype, real memory. Either way it only wins when it becomes the row maximum, i.e. when every real logit in the row is negative — a -inf channel mask produces exactly that. This is the path greedy_sample and the inner-exponential kernels take.

Fires when N % VecSize != 0.

Write side: exponential_kernel

The same partial-last-vector shape, on the store rather than the load. exponential_kernel's last store covers k .. k + VecSize - 1, and its descriptor extent is rounded up to a dword by the same oob_o idiom. At a 2-byte dtype with N % VecSize != 0 that puts element N inside the extent, so the store writes one element past the row.

The descriptor cannot fix this, and the rounding is not an oversight. The buffer bound is enforced per dword, so at an odd N elements N-1 and N share one and the hardware can only keep or drop both. Measured both ways on a 2-row bf16 buffer:

oob_o	row fully written	nothing written past the row
rounded up to a dword (today)	yes	no — row 1's first element is clobbered
exactly N	no — one real element is lost	yes

So the rounding is what protects the last valid element, and the price is the spill. The fix masks the tail in software instead: a whole vector takes the original vectorised store, and the one thread per row whose vector straddles N writes its in-range elements scalar. The philox draw is unchanged, so the noise sequence is identical.

Which kernel gets which

On the read side, all five entry points are affected by both mechanisms. They differ only in how easily B fires:

impl	reached from	BlockSize x VecSize	A fires when	B fires when
random_sample_outer_exponential_impl	mixed_sample_outer_exponential, random_sample_outer_exponential (T != 0)	1024 x 4	N <= 4092	N % 4 != 0, any logits — the 0 is in the divisor
random_sample_impl	mixed_sample, random_sample (T != 0)	1024 x 4	N <= 4092	N % 4 != 0 and the row is all-negative
argmax_impl	greedy_sample	1024 x 16	N <= 16368	N % 16 != 0 and the row is all-negative
argmax_impl	mixed_sample_outer_exponential, mixed_sample at T == 0	1024 x 4	N <= 4092	N % 4 != 0 and the row is all-negative

Measured on the unpatched tree, M=4, T=1.0, every logit -10.0: mixed_sample and random_sample return [4098, 4097, 4099, 4099] at N=4097 and go out of range at N=4095 and N=50001 as well, in both fp32 and bf16.

Fix

Both are fixed by neutralising the lanes past N, not by shortening the loops:

• has_pre records whether the prologue actually loaded, and the tail is gated on it.
• Inside the bodies, lanes at or past N are neutralised instead of using their loaded value, and the argmax update is additionally gated on (base + i) < N. Both halves of the neutralised pair are set, so no key at or past N reaches the reduction.
• argmax_impl neutralises with -INFINITY, not -FLT_MAX (thanks @junhaha666). It compares raw logits, so -FLT_MAX is a value the input can take: a row whose only finite logit is -FLT_MAX tied with the {0, -FLT_MAX} sentinel and lost it on ArgMax's smaller-key rule, returning token 0. Reproduced on fp32 at N = 17 / 4095 / 16385 for both greedy_sample and the temperature == 0 mixed path, and now covered by the check. The two sampling impls keep -FLT_MAX: their candidates are exp(..)/u >= 0 so no real value can reach the sentinel, and -inf there would turn thread_kvp.value * ratio into NaN whenever the running-max rescale underflows ratio to 0.

Keeping the trip count a compile-time constant matters: a first version used for(int i = 0; i < vec_size_i && (base + i) < N; i++) and greedy_sample (VecSize=16) measured 5x slower — 10.0 -> 58.3 us at M=64, N=151936 — because the data-dependent bound stops the body being unrolled. A neutralised lane can never displace a real one: aiter::ArgMax prefers the smaller key on equal values, and the neutral value is now unreachable by any real candidate on every path.

The out-of-bounds lanes still take part in the running max and (for NeedSum) the sum. That is argmax-neutral — the online-softmax rescale is self-consistent for any running max, and sum_softmax is a per-block scalar that divides every candidate equally — but -FLT_MAX keeps them out of both anyway, so the reduction no longer depends on what the out-of-bounds read returned.

Test Plan

op_tests/test_sample.py gains check_vocab_edges(), run after the perf sweep — it asserts, and it draws from the same global RNG the sweep does, so running it first would both hide the perf tables on a regressed tree and change the inputs every table is measured on.

It sweeps 15 vocabularies x {bf16, fp32} x {normal, all-negative} rows across all five entry points plus the T == 0 dispatch into argmax_impl, asserting against a torch reference wherever the noise is caller-supplied or the path is an argmax, and a range check for the two that draw their own noise. fp32 is swept because oob_i's dword rounding makes the 2-byte case data-dependent while fp32 reads a hard 0.

The existing l_n = [129280, 151936][-1:] is left alone: both of its values are multiples of the vector width and far above one block's span, so neither mechanism is reachable from the perf sweep by construction — which is why the dedicated check exists rather than more sweep shapes.

A second check, check_exponential_row_bounds(), covers the store side: it fills a two-row buffer with a sentinel, runs aiter.exponential on row 0 only, and asserts both that no sentinel is left in row 0 (every element was written) and that row 1 is untouched (nothing was written past). Both halves are needed — each of the two descriptor settings satisfies exactly one of them.

python op_tests/test_sample.py -m 1 64 512

Test Result

One MI300X (gfx942), rocm/atom-dev:latest, torch 2.10.0+rocm7.2.4.

Correctness. On the unpatched tree the new check fails immediately:

AssertionError: greedy_sample N=256 all_negative=False
                out=tensor([-13547, -15593, -433, -926], dtype=torch.int32)

With the fix, all 15 vocabularies x both row signs pass, and the existing sweep reports err = 0 on every row:

[aiter] vocab-edge check passed over [256, 1000, 2048, 4092, 4095, 4096, 4097,
        4099, 8193, 16368, 16383, 16385, 50001, 163840, 163841]

Store side, two-row buffer, row 0 filled by aiter.exponential:

dtype	N	row 0 fully written	row 1 intact (before)	row 1 intact (after)
fp32	4095 / 4097 / 16385	yes	yes	yes
bf16	4094 / 4096	yes	yes	yes
bf16	4095	yes	no	yes
bf16	4097	yes	no	yes
bf16	16385	yes	no	yes

Read side, representative before/after, M=4, bf16, T=1.0, eps=1e-10:

op	vocab	before	after
mixed_sample_outer_exponential	1000	[-1262, -1262, -3093, -1262]	in range, matches fp64 reference
mixed_sample_outer_exponential	4097	[4097, 4097, 4098, 4097]	in range, matches fp64 reference
greedy_sample	4097	[-1864, -11483, -11783, -8300]	matches argmax
greedy_sample	163841, all-negative row	[163842, 163842, 163842, 163841]	matches argmax

A range check alone is not a sufficient detector — at N=256 the unpatched kernel returned [-1262, -1262, -1262, 0], where the in-range 0 is also the wrong token. The new check compares against a reference for that reason.

No behaviour change in the previously-correct region. Same seed, same inputs, mixed_sample_outer_exponential output is bit-identical before and after:

vocab	before	after
4096	[1577, 4028, 379, 2909]	[1577, 4028, 379, 2909]
8192	[8124, 379, 4642, 497]	[8124, 379, 4642, 497]
50000	[34579, 37439, 27804, 48357]	[34579, 37439, 27804, 48357]
163840	[160052, 34517, 134986, 1179]	[160052, 34517, 134986, 1179]

Performance. N=151936, same JIT build system on both arms (not against the shipped .so), median of 7 timed runs per arm per shape with the observed min/max:

op	M	before (us)	after (us)	delta	per-arm spread
greedy_sample	1	9.930	10.363	+4.4%	<0.5%
greedy_sample	64	11.503	11.869	+3.2%	<0.5%
greedy_sample	512	36.068	36.316	+0.7%	~2.8%
mixed_sample_outer_exponential	1	46.895	47.195	+0.6%	~0.1%
mixed_sample_outer_exponential	64	47.889	48.238	+0.7%	~0.1%
mixed_sample_outer_exponential	512	169.514	168.804	-0.4%	~3%
mixed_sample	1	78.411	79.450	+1.3%	~0.1%
mixed_sample	64	79.603	80.581	+1.2%	~0.05%
mixed_sample	512	184.297	183.665	-0.3%	~2%

greedy_sample at small M costs a real 3-4%, not noise — the per-arm spread there is under 0.5%, so the delta sits well outside it. An earlier revision of this PR reported "within run-to-run spread" off single samples per arm; that claim was wrong and is withdrawn. Everything else is at or under 1.3%.

That 3-4% is the per-lane bounds select in a 16-wide unrolled body, and the alternatives are worse: shortening the loop instead costs 5x (above), and branching the whole-vector case around the select — only the row's last vector can straddle N — measured 10.97 us at M=1 against the shipped 10.363, so it is not taken.

black --check and ruff check are clean on op_tests/test_sample.py.

Known Limitations

1. Measured on gfx942 only. The original report was on gfx950 and the defect is in arch-independent C++, but a gfx950 run would be worth having.
2. check_vocab_edges() runs on the default dtype (bf16) only. VecSize is dtype-independent here, so the boundaries do not move, but fp16/fp32 are not swept.
3. The out-of-range read itself is unchanged — this PR stops its value from reaching the result, it does not remove the read. Note oob_i rounds the logits descriptor up to a dword, so at a 2-byte dtype the descriptor genuinely extends past the tensor's last element by 2 bytes; that is pre-existing and out of scope here.
4. In argmax_impl the tail's (k - vec_stride) < N clause is now always true given has_pre, and is kept only as a defensive restatement.

PR Self-Review

Reviewed revision

5243300

Behavior and intent

Every sampler in sample_kernels.cu could return a token id outside [0, vocab) — negative, or past the end of the row — and the existing behavior is insufficient because nothing rejects or corrects it: the kernel writes the bad index and the caller consumes it as a token. Two independent defects in the shared prologue/tail pipeline produce it. (A) A thread whose first element is already past N fails the prologue check, so its vector registers stay undefined and k is never advanced; at the tail, k - vec_stride is negative, the < N guard passes, and the body runs on undefined registers and writes a negative key. (B) When N % VecSize != 0 the row's last vector is partial and the lanes past N fall outside the row; for the outer-exponential kernels the noise descriptor ends exactly at N, so the 0 it yields lands in the divisor and exp(x - max) / (0 + eps) wins the argmax for any logits, while every other entry point takes its out-of-range value from the logits descriptor and is therefore hit only when the whole row is negative. A third defect was raised in review and is also fixed: argmax_impl neutralised with -FLT_MAX, a value raw fp32 logits can equal, so such a logit tied with the sentinel and lost it on the smaller-key rule; it now uses -INFINITY. After this PR the lanes past N are neutralised in both key and value and the tail is gated on whether the prologue actually loaded, so the result no longer depends on what the out-of-range read returned. A fourth defect of the same shape sits on the write side: exponential_kernel's partial last store reaches one element past the row at a 2-byte dtype and clobbers the next row, and the descriptor cannot mask it (shrinking the extent to N drops a real element instead, since at an odd N elements N-1 and N share a dword), so the tail is masked in software. Behavior in the previously-correct region is bit-identical (table under Test Result).

Validation and unverified scope

Measured on one MI300X (gfx942, one GPU), image rocm/atom-dev:latest, torch 2.10.0+rocm7.2.4, HIP_VISIBLE_DEVICES=0, with this branch's tree on PYTHONPATH and module_sample JIT-rebuilt from source on both arms (so the comparison is not against the image's prebuilt .so, which differs by up to 17% on its own).

• python op_tests/test_sample.py -m 1 64 512 with the fix: vocab-edge check passed over [256, 1000, 2048, 4092, 4095, 4096, 4097, 4099, 8193, 16368, 16383, 16385, 50001, 163840, 163841], and every err column in all three sweep tables is 0.
• The same command against a pristine origin/main csrc/kernels/sample_kernels.cu, same JIT build system: exits 1 at the first case — AssertionError: greedy_sample torch.bfloat16 N=256 all_negative=False out=tensor([-15771, -15467, -3562, -12315]). The new check is therefore known to fail without the fix, not only to pass with it.
• Per-category detection power against that same pristine build, recording every failure instead of asserting on the first, to confirm each slice of the new coverage pulls its weight rather than riding on the first failure: T == 0 3-12 of 15 vocabularies per category, range-only inner-exponential 4-6, everything else 10-14; no category scored 0.
• Bit-identical output before/after at vocab 4096, 8192, 50000 and 163840 with the same seed (table under Test Result), which is the no-regression evidence for current production vocabularies.
• Perf table under Test Result: median of 7 timed runs per arm per shape, min/max recorded. The earlier revision of this PR used one sample per arm and concluded "no systematic regression"; that was wrong — greedy_sample at M=1/64 costs a real 3-4%, which the repeated runs put well outside the <0.5% per-arm spread. The claim is corrected rather than softened. An earlier version of this fix that shortened the loop instead of neutralising the lanes measured 5x slower on greedy_sample; that version is not what is committed here.
• Store side: check_exponential_row_bounds passes at fp32 and bf16 for N = 4094 / 4095 / 4096 / 4097 / 16385; against a pristine origin/main kernel the three odd-N bf16 cases fail (row 1 clobbered), and an experimental build with the descriptor shrunk to N instead leaves one element of row 0 unwritten — which is why the fix is in software rather than in the extent.
• black --check and ruff check clean on op_tests/test_sample.py.

Unverified: gfx950 (the arch of the original report in ROCm/ATOM#2914 — the defect is in arch-independent C++ and reproduces on gfx942, but no gfx950 run was done); dtypes other than bf16 and fp32; multi-GPU; the rest of op_tests/ beyond test_sample.py; CI results at this head; generated ISA was not diffed.

PR-specific question

Question 3 — which architecture, dtype and shape boundaries select the changed kernel path, and what invariant prevents unsupported inputs from reaching it.

There is no arch gate: these are plain HIP kernels with no gfx branch and no arch-specific dispatch, so the changed code is live on every supported card. The dtype axis is VLLM_DISPATCH_FLOATING_TYPES_rmTorch on the input tensor; VecSize is fixed per call site (16 for greedy_sample, sizeof(float4)/sizeof(float) = 4 for the rest) and does not vary with dtype, so the shape boundaries below are the same for fp16/bf16/fp32. The shape axis is N (the vocabulary) against BlockSize * VecSize, which is 16384 for greedy_sample and 4096 for the other four entry points.

The honest answer to the second half is that no invariant prevented unsupported inputs from reaching it — that is the bug. There is no size check, no assert and no documented minimum N: a vocabulary at or below (BlockSize - 1) * VecSize, or any vocabulary not a multiple of VecSize, was accepted and silently produced a corrupt index. Current production vocabularies (129280, 151936, 163840) sit in the safe region only by coincidence of being multiples of 4 and far above one block's span, which is why this went unnoticed. After this PR every N is handled and no invariant is required.

Changed op tests

op_tests/test_sample.py (the only op test this PR changes) gains check_vocab_edges() and check_exponential_row_bounds(), both called from the sweep driver after the perf tables.

What it covers: 15 vocabularies x {bf16, fp32} x {normal, all-negative}. It asserts greedy_sample equals argmax; that mixed_sample_outer_exponential and random_sample_outer_exponential equal the softmax/exponential race computed in torch on the same noise; that both mixed kernels at temperature = 0 equal argmax (that is the argmax_impl<VecSize=4> dispatch, which greedy_sample never reaches — it only ever enters argmax_impl 16 elements wide); and that mixed_sample and random_sample return an in-range index, those two drawing their own noise so only the range is checkable. The vocabularies straddle both block spans and both mechanisms: 256 / 1000 / 2048 / 4092 / 16368 for the undefined-register path, 4095 / 4097 / 4099 / 8193 / 16383 / 16385 / 50001 / 163841 for the partial-vector path, 4096 / 163840 as in-region controls. The all-negative rows are what make mechanism B visible on every path except the outer-exponential one.

Each category was confirmed to have independent detection power by running the check against a pristine origin/main module_sample and recording every failure rather than asserting on the first: the T == 0 categories catch 3-12 of 15 vocabularies each (they had no coverage at all before), the range-only inner-exponential categories catch 4-6, and the rest 10-14.

Why it is not redundant: the file's existing coverage is the perf sweep over l_n = [129280, 151936][-1:], and both of those values are multiples of the vector width and far above one block's span, so neither mechanism is reachable from it by construction — adding shapes to l_n would also only report a non-zero err column rather than fail, since checkAllclose returns a ratio and does not raise. The sweep's test_mixed_sample does exercise T == 0, but only at N=151936, which is in the safe region. check_exponential_row_bounds is likewise the only coverage of aiter.exponential at a vocab that is not a multiple of the vector width: the sweep calls it only at N=151936 and only on a single row, so neither the spill nor a dropped tail element would be visible there. A separate asserting check is the only way to gate this in CI. A range assertion alone would also have been insufficient: at N=256 the unpatched kernel returns [-1262, -1262, -1262, 0], where the in-range 0 is also the wrong token, which is why the check compares against a reference wherever the noise is caller-supplied.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
ci:windows	Windows JIT op smoke test on the gfx1201 ROCm runner
ci:multi-gpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 7373 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[aiter-bot]

READY FOR HUMAN REVIEW

PR Self-Review

Revision 52433004b68f · mode shadow

Author self-review accepted. Readiness report:

Author answers

• PASS — current-SHA answers are present and structurally specific

Show the author's self-review answers

Behavior and intent

Every sampler in `sample_kernels.cu` could return a token id outside `[0, vocab)` — negative, or past the end of the row — and the existing behavior is insufficient because nothing rejects or corrects it: the kernel writes the bad index and the caller consumes it as a token. Two independent defects in the shared prologue/tail pipeline produce it. (A) A thread whose first element is already past `N` fails the prologue check, so its vector registers stay undefined *and* `k` is never advanced; at the tail, `k - vec_stride` is negative, the `< N` guard passes, and the body runs on undefined registers and writes a negative key. (B) When `N % VecSize != 0` the row's last vector is partial and the lanes past `N` fall outside the row; for the outer-exponential kernels the noise descriptor ends exactly at `N`, so the 0 it yields lands in the divisor and `exp(x - max) / (0 + eps)` wins the argmax for any logits, while every other entry point takes its out-of-range value from the logits descriptor and is therefore hit only when the whole row is negative. A third defect was raised in review and is also fixed: `argmax_impl` neutralised with `-FLT_MAX`, a value raw fp32 logits can equal, so such a logit tied with the sentinel and lost it on the smaller-key rule; it now uses `-INFINITY`. After this PR the lanes past `N` are neutralised in both key and value and the tail is gated on whether the prologue actually loaded, so the result no longer depends on what the out-of-range read returned. A fourth defect of the same shape sits on the *write* side: `exponential_kernel`'s partial last store reaches one element past the row at a 2-byte dtype and clobbers the next row, and the descriptor cannot mask it (shrinking the extent to `N` drops a real element instead, since at an odd `N` elements `N-1` and `N` share a dword), so the tail is masked in software. Behavior in the previously-correct region is bit-identical (table under Test Result).

Validation and unverified scope

Measured on one MI300X (gfx942, one GPU), image `rocm/atom-dev:latest`, torch 2.10.0+rocm7.2.4, `HIP_VISIBLE_DEVICES=0`, with this branch's tree on `PYTHONPATH` and `module_sample` JIT-rebuilt from source on both arms (so the comparison is not against the image's prebuilt `.so`, which differs by up to 17% on its own).
- `python op_tests/test_sample.py -m 1 64 512` with the fix: `vocab-edge check passed over [256, 1000, 2048, 4092, 4095, 4096, 4097, 4099, 8193, 16368, 16383, 16385, 50001, 163840, 163841]`, and every `err` column in all three sweep tables is 0.
- The same command against a pristine `origin/main` `csrc/kernels/sample_kernels.cu`, same JIT build system: exits 1 at the first case — `AssertionError: greedy_sample torch.bfloat16 N=256 all_negative=False out=tensor([-15771, -15467, -3562, -12315])`. The new check is therefore known to fail without the fix, not only to pass with it.
- Per-category detection power against that same pristine build, recording every failure instead of asserting on the first, to confirm each slice of the new coverage pulls its weight rather than riding on the first failure: `T == 0` 3-12 of 15 vocabularies per category, range-only inner-exponential 4-6, everything else 10-14; no category scored 0.
- Bit-identical output before/after at vocab 4096, 8192, 50000 and 163840 with the same seed (table under Test Result), which is the no-regression evidence for current production vocabularies.
- Perf table under Test Result: median of 7 timed runs per arm per shape, min/max recorded. The earlier revision of this PR used one sample per arm and concluded "no systematic regression"; that was wrong — `greedy_sample` at M=1/64 costs a real 3-4%, which the repeated runs put well outside the <0.5% per-arm spread. The claim is corrected rather than softened. An earlier version of this fix that shortened the loop instead of neutralising the lanes measured 5x slower on `greedy_sample`; that version is not what is committed here.
- Store side: `check_exponential_row_bounds` passes at fp32 and bf16 for N = 4094 / 4095 / 4096 / 4097 / 16385; against a pristine `origin/main` kernel the three odd-N bf16 cases fail (row 1 clobbered), and an experimental build with the descriptor shrunk to `N` instead leaves one element of row 0 unwritten — which is why the fix is in software rather than in the extent.
- `black --check` and `ruff check` clean on `op_tests/test_sample.py`.

Unverified: gfx950 (the arch of the original report in ROCm/ATOM#2914 — the defect is in arch-independent C++ and reproduces on gfx942, but no gfx950 run was done); dtypes other than bf16 and fp32; multi-GPU; the rest of `op_tests/` beyond `test_sample.py`; CI results at this head; generated ISA was not diffed.

PR-specific question

Question 3 — which architecture, dtype and shape boundaries select the changed kernel path, and what invariant prevents unsupported inputs from reaching it.

There is no arch gate: these are plain HIP kernels with no `gfx` branch and no arch-specific dispatch, so the changed code is live on every supported card. The dtype axis is `VLLM_DISPATCH_FLOATING_TYPES_rmTorch` on the input tensor; `VecSize` is fixed per call site (16 for `greedy_sample`, `sizeof(float4)/sizeof(float) = 4` for the rest) and does not vary with dtype, so the shape boundaries below are the same for fp16/bf16/fp32. The shape axis is `N` (the vocabulary) against `BlockSize * VecSize`, which is 16384 for `greedy_sample` and 4096 for the other four entry points.

The honest answer to the second half is that **no invariant prevented unsupported inputs from reaching it** — that is the bug. There is no size check, no assert and no documented minimum `N`: a vocabulary at or below `(BlockSize - 1) * VecSize`, or any vocabulary not a multiple of `VecSize`, was accepted and silently produced a corrupt index. Current production vocabularies (129280, 151936, 163840) sit in the safe region only by coincidence of being multiples of 4 and far above one block's span, which is why this went unnoticed. After this PR every `N` is handled and no invariant is required.

Changed op tests

`op_tests/test_sample.py` (the only op test this PR changes) gains `check_vocab_edges()` and `check_exponential_row_bounds()`, both called from the sweep driver after the perf tables.

What it covers: 15 vocabularies x {bf16, fp32} x {normal, all-negative}. It asserts `greedy_sample` equals `argmax`; that `mixed_sample_outer_exponential` and `random_sample_outer_exponential` equal the softmax/exponential race computed in torch on the same noise; that both mixed kernels at `temperature = 0` equal `argmax` (that is the `argmax_impl<VecSize=4>` dispatch, which `greedy_sample` never reaches — it only ever enters `argmax_impl` 16 elements wide); and that `mixed_sample` and `random_sample` return an in-range index, those two drawing their own noise so only the range is checkable. The vocabularies straddle both block spans and both mechanisms: `256 / 1000 / 2048 / 4092 / 16368` for the undefined-register path, `4095 / 4097 / 4099 / 8193 / 16383 / 16385 / 50001 / 163841` for the partial-vector path, `4096 / 163840` as in-region controls. The all-negative rows are what make mechanism B visible on every path except the outer-exponential one.

Each category was confirmed to have independent detection power by running the check against a pristine `origin/main` `module_sample` and recording every failure rather than asserting on the first: the `T == 0` categories catch 3-12 of 15 vocabularies each (they had no coverage at all before), the range-only inner-exponential categories catch 4-6, and the rest 10-14.

Why it is not redundant: the file's existing coverage is the perf sweep over `l_n = [129280, 151936][-1:]`, and both of those values are multiples of the vector width and far above one block's span, so **neither mechanism is reachable from it by construction** — adding shapes to `l_n` would also only report a non-zero `err` column rather than fail, since `checkAllclose` returns a ratio and does not raise. The sweep's `test_mixed_sample` does exercise `T == 0`, but only at N=151936, which is in the safe region. `check_exponential_row_bounds` is likewise the only coverage of `aiter.exponential` at a vocab that is not a multiple of the vector width: the sweep calls it only at N=151936 and only on a single row, so neither the spill nor a dropped tail element would be visible there. A separate asserting check is the only way to gate this in CI. A range assertion alone would also have been insufficient: at `N=256` the unpatched kernel returns `[-1262, -1262, -1262, 0]`, where the in-range `0` is also the wrong token, which is why the check compares against a reference wherever the noise is caller-supplied.

PR state

• PASS — pull request is marked ready for review

Copilot findings

• PASS — no active Copilot finding lacks a substantive author reply
• ADVISORY — Copilot did not review the current head SHA

Code and op-test hygiene

• PASS — no committed artifact or conflict marker detected

Required CI

• PASS — Black: passed
• PASS — Ruff: passed
• PASS — Repository dependency: passed
• PASS — Aiter Test Gate: passed

_{Shadow mode reports what would block but keeps the status successful. Set AITER_PR_SELF_REVIEW_MODE=enforce only after reviewing real false positives.}

[junhaha666]

@aiter-bot review

[junhaha666]

[P2] Use negative infinity for the argmax sentinel

argmax_impl initializes each thread's candidate with a synthetic key/value pair:

kvp thread_kvp{0, -FLT_MAX};

The masked out-of-range lanes also use -FLT_MAX. This is not a true neutral element for FP32 inputs. If the largest real logit is exactly -FLT_MAX and token 0 is -inf, the real candidate cannot beat the synthetic {0, -FLT_MAX} because ArgMax resolves equal values in favor of the smaller key.

For example:

logits = torch.full((1, N), float("-inf"), device="cuda")
logits[0, expected_token] = -torch.finfo(torch.float32).max

The reference argmax returns expected_token, but both greedy_sample and the temperature == 0 mixed path return token 0. This was reproduced on the current PR revision:

N=17     expected=5      greedy=0  mixed(T=0)=0
N=4095   expected=4094   greedy=0  mixed(T=0)=0
N=16385  expected=16384  greedy=0  mixed(T=0)=0

Please use negative infinity as the neutral value in argmax_impl, including for masked lanes:

kvp thread_kvp{0, -INFINITY};
kvp tmp_kvp{0, -INFINITY};

tmp_kvp.value =
    in_range ? opus::cast<float>(vec_pre[i]) : -INFINITY;

This allows a real -FLT_MAX value to beat the sentinel while preserving the expected smallest-key tie-breaking behavior when all valid logits are -inf.

Please also add an FP32 regression test covering this case.

[aiter-bot]

ROCm/aiter PR #7373 — [HIP] [Bugfix] sample kernels: keep token ids inside [0, vocab)

Fixes two defects in the shared prologue/tail pipeline of the three sampler impls in csrc/kernels/sample_kernels.cu — a thread whose first element is past N running the tail on undefined registers with a negative index, and out-of-bounds lanes of a partial last vector winning the argmax — by gating the tail on whether the prologue loaded and neutralising lanes at or past N with -FLT_MAX, and adds a vocab-edge check to op_tests/test_sample.py that fails on the unpatched tree.

Review (advisory): ⚠️ NEEDS WORK
Validation (deterministic): NOT RUN — REVIEW_AUTO_VALIDATE=0 (auto-validation disabled); the required target op_tests/test_sample.py existed for the changed runtime surface csrc/kernels/sample_kernels.cu — an environment gap, not a PR defect
Perf (advisory): NOT RUN — no idle GPU on this box (all 8 at ~98% with active jobs) and no validation report; the PR's own base-vs-head table (one MI300X, module_sample JIT-rebuilt on both arms, one timed run per arm per shape) is the only timing evidence and is unverified here

⚠️ [inferred] The partial-vector edge this PR fixes on the load side is still live on the store side of exponential_kernel in the same file (csrc/kernels/sample_kernels.cu:749, unchanged): at N % 4 != 0 the last store spans past the row, and at odd N with a 2-byte dtype the dword-rounded descriptor puts the lane at N inside the buffer extent, so it plausibly writes the next row's first noise value (per-lane buffer bounds, the granularity the PR's own N=4097 load evidence exhibits) — or drops the whole partial store if bounds are per-instruction — and check_vocab_edges never calls aiter.exponential, so nothing in this PR covers it. Author must confirm the bounded store's behaviour past N at a 2-byte dtype (or mask the store tail the same way) and add aiter.exponential at an odd N to the vocab-edge check.
📝 [verified] The "within run-to-run spread; no systematic regression" conclusion for the changed greedy_sample path in csrc/kernels/sample_kernels.cu rests on one timed run per arm per shape (the PR states this; only greedy_sample was repeated, on head alone), and the +3.9% greedy_sample M=1 delta (9.94 → 10.32 us) exceeds the ~0.7% head-only spread the author measured — so a real ~4% regression on the hottest decode shape could hide inside a single-sample comparison. Author must re-run the greedy_sample rows with several samples per arm, or soften the no-regression claim.

[yitingw1]

Thanks — reproduced and fixed in 0a1dca6.

[P2] -FLT_MAX sentinel. Confirmed exactly as you describe, on fp32 at all three N:

 dtype       N  expected   greedy  mixed T=0
float32     17         5        0          0   <- before
float32   4095      4094        0          0
float32  16385     16384        0          0

argmax_impl now uses -INFINITY for both the thread sentinel and the masked lanes, and the three cases return the expected token.

One note on scope: I kept -FLT_MAX in random_sample_impl and random_sample_outer_exponential_impl. Their candidates are exp(x - max) / u >= 0, so no real value can reach the sentinel there, and -inf would be actively harmful — thread_kvp.value * ratio becomes NaN whenever the running-max rescale underflows ratio to 0. So the change is scoped to the one impl that compares raw logits.

The fp32 regression test is in check_vocab_edges, using -torch.finfo(dtype).max per dtype rather than a fixed -FLT_MAX — at bf16 that constant rounds to -inf, which is the one case where token 0 is the right answer, so a fixed value would have made the bf16 row vacuous. Verified it fails on a pristine origin/main kernel (fp32 rows) and passes with the fix.

---

While re-checking the perf numbers I also had to withdraw a claim of my own. @aiter-bot was right that one sample per arm could hide a real regression: with 7 runs per arm and a per-arm spread under 0.5%, greedy_sample costs +4.4% at M=1 and +3.2% at M=64 — real, not noise. Everything else is at or under 1.3%. The table and the conclusion are updated; "within run-to-run spread" is gone.

That cost is the per-lane bounds select in the 16-wide body. I tried hoisting it — only the row's last vector can straddle N, so the whole-vector case could branch around the select — and it came out worse (10.97 us at M=1 vs 10.363 shipped), so the straight-line select is what ships.

---

Also confirming @aiter-bot's inferred finding on exponential_kernel, which I did not fix here. oob_o rounds the output descriptor up to a dword, so at a 2-byte dtype with odd N the lane at index N sits inside the extent and the last store writes one element past the row:

  dtype      N   oob_o   next-row sentinel survived?
float32   4095    4095   yes
float32   4097    4097   yes
bfloat16  4095    4096   NO  - row 1 clobbered: [0.71484375, -7.0, -7.0, -7.0]
bfloat16  4097    4098   NO  - row 1 clobbered: [1.2265625,  -7.0, -7.0, -7.0]
bfloat16 16385   16386   NO  - row 1 clobbered: [0.2314453125, -7.0, -7.0, -7.0]

I left it out of this PR deliberately: it is a different kernel and a different defect class (out-of-bounds write, not a wrong index), it is pre-existing, and the obvious fix — bounding the descriptor at N — is not obviously safe, since store_vector's bounds apply per chunk and a naive change risks silently dropping valid tail elements. That is a worse failure than the current spill, so it wants its own PR with its own before/after. Happy to open it, or leave it to whoever owns that path.

[yitingw1]

Changing my mind on scope — the exponential_kernel store overrun is now fixed in this PR (5243300), not deferred as I said above. The PR is retitled accordingly.

Two reasons:

1. The stated justification turned out to be the weak half. I said it was "a different kernel and a different defect class", but the real reason was the clause after it — that bounding the descriptor at N was not obviously safe. I went and measured that, and it resolved into exactly the same mechanism this PR is already about:

oob_o	row fully written	nothing written past the row
rounded up to a dword (today)	yes	no — row 1's first element clobbered
exactly N	no — one real element lost	yes

The buffer bound is enforced per dword, so at an odd N elements N-1 and N share one and the hardware can only keep or drop both. That is the same fact this PR's read side already turns on — oob_i's rounding is what lets a bf16 row's last element be read, and the price there is one extra lane, which the PR neutralises. On the store side the price is a write past the row, which cannot be neutralised, so the tail is masked in software. Splitting them would have meant explaining the same dword-granularity fact twice.

2. It was going to cost a merge conflict for no benefit. Both changes land in op_tests/test_sample.py in the same place, so the second PR would have conflicted with the first on a file neither of them is really about.

What changed in the kernel is three lines: a whole vector takes the original vectorised store, and the one thread per row whose vector straddles N writes its in-range elements scalar. The philox draw is untouched, so the noise sequence is identical.

Verification is in the description; the short version is that check_exponential_row_bounds asserts both halves — row 0 fully written and row 1 untouched — because each of the two descriptor settings satisfies exactly one of them, and a test checking only one would have passed on the broken build either way.