Repository navigation
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
github-actions
Bot
changed the title
MHYangAMD
changed the title
github-actions
Bot
changed the title
A second fp32 prefill selector alongside top_k_per_row_prefill, emitting int32 indices and optionally values, over M=1..4096, N=2048..1M, k up to 2048. `sampled` names the mechanism: estimate a threshold by sampling, filter a candidate set against it, select exactly, and send the rows whose candidate set came out unusable -- too few, or overflowed -- to a separate exact fallback. It is not a port of DeepSelect. It borrows two things, both cited where they are used: the "distort" fp32-to-monotone-uint32 trick, and a lesson about a __shfl_down tree leaving its result in lane 0 only. The architecture differs -- DeepSelect's v3_fp32 is one kernel, this is nine. The .cu and the three headers under csrc/include/topk_sampled/ are generated from a separate repository by its scripts/export_aiter_op.py. Regenerate rather than edit them; the banner in each says so, and the generator is what keeps them consistent with the source of truth. aiter/jit/optCompilerConfig.json names the .cu by path, so it moves with it. op_tests/test_topk_per_row.py gains a --prefill_backend axis so that one harness, one data generator and one timer drive both selectors. A number measured through a different harness cannot be compared with one that is not: the two disagreed by 18% at M=256 when the harness, the timer and the data generator all differed at once, and the cause was not attributable until they were pinned.
Three things in one file, and the floor is the part worth reading. **The floor.** `sampled` takes stride0 >= SAMPLED_MIN_STRIDE0 = 131072, ahead of the FlyDSL one-block path, which serves everything below it. Measured per shape, k=2048, fp32, both ops under one @perftest on the same data (flydsl_us / sampled_us, above 1.00 meaning `sampled` is faster): M \ N 49152 65536 131072 262144 1 0.92 1.06 1.68 2.71 8 0.78 0.90 1.38 2.15 64 0.72 0.83 1.17 1.57 256 0.67 0.79 0.99 1.14 512 0.72 0.77 1.08 1.27 1024 0.89 0.96 1.07 1.30 2048 0.85 0.89 1.02 1.37 4096 0.83 0.92 1.04 1.32 The crossover is a function of stride0 and not of the row count, so there is no wide/narrow split. M=256 at 131072 is 0.99, a wash, left on the simple rule rather than carved out. An earlier version of this floor was 32768/49152 with a row-count split, fitted against the mb/ob path. That measurement was correct when taken and is now obsolete: FlyDSL sits in front of mb/ob, so a floor fitted against mb/ob is a comparison with an op this function no longer reaches at those widths. Ordering matters as much as the number -- placed after the FlyDSL check, this routing was unreachable for every shape tested, M=4096 N=65536 included. **Validation.** Bad arguments raised AITER_CHECK, which calls std::abort() unless g_aiter_can_throw is set, and only the ctypes bridge sets it -- so k above the Phase C cap, stride1 != 1, a short workspace, stride0 = 0 and k <= 0 all killed the caller's process. They now raise ValueError in Python. **Memoisation.** The per-call shape lookups are pure functions of their arguments and were not cheap through the binding: topk_use_mulblocks 6.34 us, topk_ob_workspace_size 6.74 us, topk_mb_workspace_size 5.06 us, plus the two AVO queries at 4.86 and 4.80. On a small shape that was most of the call -- 13.1 us of host time in front of a 2.8 us kernel. topk_use_mulblocks keeps TOPK_FORCE_PATH and TOPK_DISPATCH_FACTOR in its cache key rather than assuming nobody sets them; two dict lookups against a 6.34 us call keeps them live. op_tests/test_topk_prefill_dispatch.py asserts only whether `sampled` ran. Three paths can serve a prefill now, so "not sampled" no longer implies mb/ob, and asserting which of the others took it made the test fail for a dispatch that was behaving correctly. It reads SAMPLED_MIN_STRIDE0 rather than repeating its value.
The routing rule is a single threshold on total work, `rows * width >= 2**28`,
and it is sharp. Across M=256..4096 the crossover landed on exactly 2**28 every
time -- 256x1M, 512x512K, 1024x256K, 2048x128K, 4096x64K -- with `sampled`
fastest at and above it and another backend fastest below.
Safe to ship because of what the rule routes on, not because of the constant:
`sampled` is only given a cell where it measured fastest of EVERY backend that
can serve it. A cell that changes hands was already going to be at least as slow
under the previous rule, so no cell can regress. That is the
A/B-against-the-rule-in-place test `topk_select_backend` asks for rather than the
per-cell oracle.
Deliberately conservative in four places:
- `_BACKENDS_BY_TIE` lists it under no promise only. Its tie direction is not
measured, and this file already records what happens otherwise: `tie='high'`
once fell through to `decode`, which ties the opposite way, silently.
- `_NONDETERMINISTIC` includes it because its repeat-stability is NOT measured,
not because it was found unstable. A `deterministic` caller never reaches it.
`decode` is the cautionary tale there -- a radix select being a pure function
of the row is the obvious guess and it is wrong.
- `_PREFERENCE` lists it last, so the existing fallback order is unchanged
except where nothing else can serve.
- Below 2**28 it stays out even where it won. M=128 took N=65536 and N=131072
then lost N=262144 through N=1048576, which no monotone rule in the work size
expresses; those two cells stay on the old rule. The threshold is inert for
M <= 128 anyway, since at the spec's widest N of 1M that reaches only 2**27.
`_available` admits it on the dtype half of the predicate only. The row-count
half cannot be asked there -- that set is memoized without the row count -- so
`topk_select_backend` asks `topk_sampled_supports` with the rows it has. Same
shape as the streaming selector's two-block-width case immediately above it.
`AITER_DISABLE_TOPK_SAMPLED=1` withdraws it from the available set, which is what
`bench/select_ab_sweep.py --mode ab` toggles. Changing it mid-process needs both
`_available.cache_clear()` and `_choose.cache_clear()`.
**Fitted on torch.randn and nothing else.** Top-k time on this machine depends on
the value distribution as well as the shape, and every cell behind this constant
was measured on one distribution. Re-fit before trusting it on real data. The
docstring on `_sampled_takes` says so where someone changing the constant will
read it.
Verified: routing fires at and only at the threshold, checked either side at ten
shapes. Correctness through `topk_select` against torch.topk on the five routed
crossover cells plus two unrouted controls, index uniqueness and value multiset,
0 failed. op_tests unchanged from the pre-change control -- test_topk_select's
2 failed + 2 errors are its own pre-existing pytest signature problem, confirmed
by running it on a stashed tree before any of this. test_topk_per_row_stable ALL
PASS, spec_low_n 32 cells 0 failed.
3414344 to
7f6f204
Compare
…hind. black reformatted three files -- the compile_ops decorator line for top_k_per_row_prefill_sampled, and the argument list of the sampled branch in _dispatch. ruff fixed the import block ordering in aiter/ops/topk.py (I001). Two ruff findings needed a decision rather than --fix: EXE001 on op_tests/test_topk_prefill_dispatch.py, which had a shebang without the executable bit. Dropped the shebang rather than chmod +x, because that is what the siblings do: none of the other test_topk*.py files carry one, and only 12 of the 148 files in op_tests do. UP031 on the two assertion messages in that same file, which I had written with percent formatting. Now f-strings. That changed the line lengths enough for black to want the file again, so black and ruff were run alternately until both converged -- they disagree, and one pass of each is not enough to prove it. Also removes a stale comment above the floor constants. The re-fit replaced the constants but this comment sat above them and survived, so the file carried both the old 32K/48K story and the new 131072 one and they contradicted each other. ruff check . passes; black --check . leaves all 1242 files unchanged. Nothing here is a logic change: routing still resolves M=4096 N=65536 to sampled and N=32768 to plain, spec_low_n and stress --smoke still pass, and the pytest result is unchanged from before (test_topk_select s 2 failed + 2 errors are its own pre-existing signature problem).
The one-block radix select re-scanned the whole row on every pass. Pass 0 has
to (the pivot is unknown), but passes 1..n-1 only need the candidates that
survived, and the machinery to stage them was already in the file and simply
never wired up:
* filter_and_histogram_for_one_block()'s compact branch, commented
"not used by one-block"
* the single-block overload of set_buf_pointers(), which had no callers
* the per-row scratch, sized by topk_plain_workspace_size(), allocated on
the Python side, passed in, offset inside the kernel -- and never written
The call site hardcoded out_buf/out_idx_buf to nullptr and passed row_len as
previous_len on every pass, so all of that was dead. This wires it up,
following aiter::ob::radix_kernel's buffer switching:
* set_buf_pointers() picks in_buf/out_buf per pass; pass 0 and pass 1 read
the original row (pass 0 cannot compact), later passes read what the
previous pass staged
* previous_len comes from counter.previous_len once a pass has staged
* the last_filter emits from the compacted buffer instead of the full row
* a candidate set larger than buf_len is left unstaged and the next pass
falls back to the full row, exactly as radix_kernel does
Measured on MI355X (gfx950, fp32) at M=4096 N=131072 k=2048, indices only:
HBM reads drop from 3.07x the logits array to 1.97x (rocprofv3 FETCH_SIZE
6.58 GB -> 4.24 GB), i.e. the intended 3 passes -> 2, and the kernel goes
922 us -> 819 us (1.13x). The win grows with the row length: 1.26x at
M=512/N=131072, 1.31x at M=4096/N=262144.
Two instantiations opt out.
STABLE is a correctness constraint: last_filter_stable_scan/_fast recover an
element's original index from its position in the buffer they scan, and take
no index buffer, so they can only run over the original row.
WRITE_TOPK_VALUES is a performance one. Compacting moves the emit of the
definite winners out of the final pass into the pass-1 scan loop, where each
costs an LDS atomic plus a dependent scattered store; with values that is two
stores per winner and the ~2k extra dependent stores per row cost more than
the full-row pass they save (922 us -> 1282 us). Traffic is identical either
way (FETCH_SIZE 4.240 GB both, WRITE_SIZE +33 MB) and occupancy is unchanged
at 8 waves/SIMD with no spills, so this is store/atomic serialisation in the
scan loop rather than bandwidth. Aggregating the emit counters per wave would
likely lift the restriction.
COMPACT is a template parameter decided on the host rather than a runtime test
on row_len: the pass loop is fully unrolled and every buffer pointer folds per
pass, so making the decision inside the kernel costs more than compaction
saves (it regressed every N measured, e.g. 819 us -> 1110 us at N=131072).
The threshold is the measured crossover -- compaction replaces one full-row
scan (cost ~ len) with staging plus ~k dependent emits (cost ~ k) on top of a
fixed per-pass cost that a short row cannot amortise, so it needs the row to
be long both absolutely and relative to k.
Also drops a latent race: counter.previous_len was assigned from counter.len
immediately after choose_bucket with no barrier in between, and choose_bucket
is written by whichever thread owns the crossing bucket, not thread 0. It read
the newly chosen bucket's size rather than the length of the buffer the pass
produced, which is what the next pass needs. Harmless while the field was
unused; taking it from the register copy read before choose_bucket fixes both.
Follow-up to the previous commit, which had to restrict compaction to the
indices-only instantiation and to rows of at least 32768 because the pass-1
scan was much slower than its traffic justified (1.97x the logits array read,
but only 5.11 TB/s against the 7.14 TB/s the un-compacted kernel sustains).
That was blamed on the single-address LDS atomics behind counter.filter_cnt /
counter.out_cnt. Disassembly says otherwise: LLVM's AMDGPUAtomicOptimizer
already converts every one of them into the standard wave aggregation --
s_mov_b64 s[68:69], exec ; mask of the hitting lanes
v_mbcnt_lo_u32_b32 v12, s68, 0 ; this lane's prefix
v_mbcnt_hi_u32_b32 v12, s69, v12
v_cmp_eq_u32_e32 vcc, 0, v12 ; elect the first active lane
s_and_saveexec_b64 s[66:67], vcc
s_bcnt1_i32_b64 s55, s[68:69] ; how many lanes hit
ds_add_rtn_u32 v13, v29, v13 ; ONE atomic for the wave
v_readfirstlane_b32 s55, v13
v_add_u32_e32 v36, s55, v12 ; slot = wave base + prefix
-- 212 ds_add_rtn_u32 against 212 s_bcnt1_i32_b64, i.e. every returning LDS
atomic is already aggregated. Hand-writing ballot+popcount would have changed
nothing.
The actual stall is one line away from those atomics:
global_load_dword v13, v[36:37], off ; in_idx_buf[idx]
.LBB10_66:
s_waitcnt vmcnt(0) ; <-- drains EVERYTHING
global_store_dword v[36:37], v13, off
`out_idx_buf[pos] = in_idx_buf ? in_idx_buf[idx] : idx` compiles to a gather,
and the wait that feeds the dependent store sits on the *merge* path of the
ternary, so it executes even when in_idx_buf is null. vmcnt has no per-access
form here, so it drains all outstanding vector memory -- including the four
global_load_dwordx4 that vectorized_process keeps in flight. Every wave
holding a candidate therefore stalls the whole streaming scan, and pass 1 is
the pass that still reads the full row.
Hoisting the index source out of the loop, exactly as last_filter already does
it, removes the gather from the pass-1 specialisation entirely. Store-blocking
drains in the compacting instantiation drop from 94 to 50 (indices only) and
to 3 (with values), and the wave aggregation that was already there starts
paying off:
M=4096 N=131072 k=2048 before -> after (baseline for reference)
prefill, indices only 819us -> 689us (926us)
prefill, with values 923us -> 741us (924us)
plain 922us -> 740us (925us)
achieved bandwidth, indices only 5.11 -> 5.88 TB/s (7.06 baseline)
Traffic is unchanged at 1.97x the logits array (rocprofv3 FETCH_SIZE 4.240 GB),
confirming this is purely the overlap coming back rather than less work.
With the scan unblocked the WRITE_TOPK_VALUES restriction is gone: writing
values now costs 52us over indices-only instead of 359us, and both beat the
un-compacted kernel by a wide margin, so `plain` -- which always passes a
values tensor -- benefits too. Only STABLE still opts out, and that remains a
correctness constraint rather than a performance one.
The length gate is re-derived on the new cost model and loosens from
len >= 32768 && len >= 24*k to len >= 16384 && len >= 16*k, which is the
measured crossover across k in {128, 512, 1024, 2048}.
…g them
The compact path wrote each surviving candidate straight out with
`out_buf[atomicAdd(...)] = value`, and at the widths its gate admits the
candidates are far too sparse for that to be a sensible write. At the customer
shape -- 4096 rows of 131072 at k=2048 -- pass 1 keeps 1556 of 131072 elements,
1.19%, which is 0.19 candidate lanes per 64-lane wave per 16-element step. So
almost every staged store was a single lane writing four bytes, and the write
path saw about ten sub-line transactions for every line's worth of useful data.
Measured rather than inferred, by ablation: deleting the two staged stores (an
incorrect build, kept only to bound the prize) takes the kernel from 705.1us to
627.6us, and `rocprofv3` shows FETCH_SIZE unchanged at 4.240 GB across both, so
the 77.5us is the write path and not a read-modify-write. The staged data is
0.051 GB, so it was moving at 658 GB/s against this part's measured 6642 GB/s
write peak.
Collecting the candidates in LDS and flushing them once per block turns that
into one contiguous run written by all 1024 threads with a full exec mask.
`kStageCap` is 2*BlockSize; at every shape the gate admits, `buf_len` bounds
the candidate count below that, and the overflow route is kept for the rest
because correctness only needs the positions to be a permutation of
[0, current_len), not any particular order. 4*BlockSize was measured and scores
identically, so the smaller LDS footprint stands.
Two things this is NOT, both checked:
* It is not a wider store. No `global_store_dwordx4` is generated and none
was wanted: with 1.5 candidates per thread there is nothing to widen, which
is why the batching has to be per block. What was missing was density
across the exec mask, not bytes per instruction.
* It is not the 50 remaining `s_waitcnt vmcnt(0)`-before-store drains that
the previous round left behind. Those are all in the `in_idx_buf != nullptr`
specialisation, i.e. the pass that reads the compacted buffer rather than
the row -- at this shape 32 KB against 1 MB, under 3% of the traffic, so
making that pass free would be worth at most 21us of 690. The drain count
is unchanged at 50 and that is fine.
Also tried and rejected: giving the winners emitted in the same scan
(`pb < kth_value_bits` -> `out_idx[pos]`) the same treatment. It is worth a
further 20us on the indices-only build and costs the WRITE_TOPK_VALUES one
1.16x (748.8us -> 873.8us), so it is not a net win.
Customer shape, batched, host dispatch amortised out, both JIT modules rebuilt:
topk_select / topk_plain(vals=None) / prefill(values=None) 703.2 -> 677.0us
topk_plain(vals) / prefill(values) 751.2 -> 723.3us
1.039x on both instantiations, which is what the change had to do -- the
`topk_select` router sends the customer shape to `plain`, and `topk_plain`
compiles this file into a second JIT module of its own.
4.317 GB in 677.0us is 6376 GB/s, 85.1% of the 7490 GB/s deliverable read peak
`roofline-calibrate` measures on this part, up from 81.6%. Traffic is unchanged
(FETCH 4.240 -> 4.234 GB, WRITE 0.077 -> 0.077 GB). VGPRs 50 -> 60, LDS 33992
-> 50376 bytes/block, occupancy still 8 waves/SIMD, no spills.
`topk_select`'s plain branch allocated a values tensor per call and threw it
away -- the comment said so ("Write-only scratch: the caller never sees these
values"). Whatever `topk_select` returns is gathered from the input at the
selected indices, because that is the only form that survives the reorderings
`sorted` and `sorted_index` apply, so the kernel's own values were never read
by anyone.
They were not free. The one-block radix kernel has had a WRITE_TOPK_VALUES
template parameter all along and `radix_topk_dispatch` already picks the
false build when handed a null `out`; only `topk_plain`'s own signature
required a tensor. Writing k values per row costs 52us of the 741us the kernel
takes at 4096 rows of 131072 at k=2048.
So `topk_out` becomes optional, and the waste goes away for every caller of
`topk_plain` rather than just this one. Measured through `topk_select`, same
build, toggling only the predicate:
M=4096 N=131072 k=2048 747.7us -> 706.2us (0.944x)
M=1024 N=131072 k=2048 204.0us -> 191.4us (0.939x)
M=4096 N= 32768 k=2048 233.3us -> 219.2us (0.940x)
Small cells are launch-bound and flat, as expected -- the saving is the write,
so it scales with rows*k.
Only the radix path can skip the values; the block-sort fallback under
`AdaptiveTopK` writes them unconditionally and has no such build. Which path
runs is `should_use_topk_radix`, so that is what decides whether nullopt is
allowed, asked two ways that cannot disagree in silence: `topk_plain.py`
exposes `topk_plain_values_optional` for the caller to check beforehand, and
`topk_plain` itself AITER_CHECKs the same condition and refuses rather than
writing through a null. `topk_plain_batches_ragged_rows` was already a mirror
of that predicate, so both now sit on one `_uses_radix` helper instead of two
copies of the formula.
The routing constants were re-fitted against a fresh 572-cell sweep, since
this is again a change to what `_dispatch` hands a backend. Nothing moves: a
search over 70560 combinations of plain's five terms finds no candidate that
beats the shipped rule without regressing a cell, and the boundary cells are
rejected by margins far wider than 44us -- extending the band to 262144 still
loses 4.957x at 256 rows of 262144 at k=1024 (stream 106.6us against plain
528.4us). A uniform per-call saving on the cells plain already serves does not
move a crossover it loses by 5x.
github-actions
Bot
changed the title
`aiter/ops/topk_select.py` passes `black --check` at 7d9c2d1^ and fails at 7d9c2d1, so the terminal-window stream commit introduced this on its way in and the style CI (.github/workflows/pre-checks.yaml, psf/black@stable) would reject the branch as it stands. Reformatting only, no behaviour: the call now fits on one line.
_dispatch's sampled branch built both ends of the per-row window from scratch each time: torch.zeros(rows) for the starts and torch.full((rows,)) for the ends, two allocations and two fill kernels for a pair of constants. Measured at 8.1us of device time at 4096 rows, and a fifth of the whole call at 64 -- calling the kernel directly rather than through this entry measured 1.20x at m=64 n=524288 and 1.13x at m=64 n=1048576, and this pair was that gap. Neither end needs building. row_lens already IS the per-row end -- _full_rows when the caller gave no `end`, the caller's own tensor otherwise -- and topk_select has already refused anything that is not int32 [rows], so the .to(torch.int32) beside it was a no-op and the torch.full rebuilt, per call, the tensor _full_rows is cached to avoid. _zero_rows is that same cache at the other end of the range. Measured on gfx950 over the 390-cell (M, N) grid at k=2048, seed 0, profiler device time: sampled through this entry is a median 1.130x faster at N >= 131072 and 1.268x at its best (m=32 n=1048576). The arms this does not touch -- plain, decode, stream -- read 1.001x over 360 arm-cells, which is the control.
Regenerated from topk-prefill-avo by scripts/export_aiter_op.py; the source
change is on that repo's csrc/topk_shape.hip.hpp.
derive_sample_s_for_n's rule 1 accepted the first sample count whose candidate
window fit under Phase C's cap:
if (candidate_hi(K, S, N, m) <= PHASE_C_CAP_MAX) return S;
candidate_hi is the THREE-SIGMA upper edge, so that accepts a plan which
overflows on any row past 3 sigma -- and taking the FIRST S that fits takes the
one with the least headroom there is. Measured at k=2048 on gfx950: N=524288
chose S=5440, putting the edge at 99.1% of the cap and leaving 3.09 sigma, and
N=1048576 chose S=10816 for 99.7% and 3.04 sigma. Every other width, and both of
rule 0's, sat at 5.7 to 15.2 sigma.
What that cost: a row whose candidate set comes out unusable goes to
phase_d_fallback, which gives it a single workgroup for a four-pass exact
select, so it costs a flat ~350us whatever M is -- m=32 n=524288 goes 39us to
388us for ONE row of 32, m=4096 goes 1881us to 2243us for one of 4096. Over 5
seeds and ~30000 rows, N=524288 was the only width that tripped it at all, at
1.2% of rows under rule 1 against 0.005% under rule 0.
The fix is to require the window to fit inside CAP_SAFE_FILL of the cap instead.
That fraction is not fitted: derive_shape_params already used 0.85 of the same
cap when it clamps the margin, and this names it so both uses share it. N=524288
moves to S=6528 (84.9% of the cap, 4.94 sigma) and N=1048576 to S=13056;
N <= 262144, already at 51.9% and 72.6%, is untouched, as is rule 0.
Measured after, 8 seeds x 70656 rows: 1 fallback row in total, and that one is
under rule 0 at m=4096 where the rate was already ~1/30000. M <= 32 at
N >= 524288 went from 1.2% of rows to 0 of 384. The larger S costs nothing --
m=16 n=524288 measured 36.26us before and 33.84us after, m=32 39.5us and
38.62us -- because phase_a is latency-bound at small M rather than
throughput-bound, while the tighter threshold leaves phase_b fewer candidates to
write and phase_c fewer to select. Rule 0's cells moved 0.1% to 0.8%, inside
this box's run-to-run spread.
op_tests/test_topk_select.py test_lds_sizing, test_invariants and
test_half_invariants all hold; the source repo's benchmark_topk --mode verify
is VERDICT PASS at N=524288 and N=262144.
_sampled_takes had one door, rows * width >= 2**28, and that door cannot reach
the wide few-row corner at all: 128 rows of the spec's widest N is 2**27. Its
own docstring said as much ("the rule is inert for M <= 128 anyway"). That
corner is where sampled now wins by the most.
Re-measured on gfx950 over the 390-cell (M, N) grid at k=2048, seed 0, forcing
each backend that serves the cell: sampled is the fastest arm on every cell of
N >= 131072 bar one, by up to 2.96x (m=128 n=1048576, decode 387.5us against
131.1us). Below N=32768 it is 1.05x to 1.80x off the best, so the threshold that
fits the data is on the WIDTH, not on the product, and 131072 is where the two
regimes separate.
Scored by simulating candidate rules over the measured table -- the A/B against
the rule in place that topk_select_backend asks for, not only the per-cell
oracle. The simulated current rule reproduces the sweep's recorded dispatch on
390 of 390 cells, so the comparison is exact:
rule green mean-vs-oracle total time >2% slower
current (work >= 2**28) 65 1.256 1.000x 0
+ width >= 131072 73 1.095 0.933x 1
The one regression was m=32 n=524288, and it was the fallback that the preceding
CAP_SAFE_FILL commit removed; with that in place there is no row floor and no
regression. Measured end to end, not simulated: the width gate moved 42 cells at
a median 0.585x (best 0.317x, worst 0.786x), dropping the row floor moved 72
more at a median 0.711x (best 0.481x, worst 0.925x), and no cell in either step
was slower. Cells whose dispatch did not change read a median 1.001x, which is
the control.
Efficiency against the published PR 5686 pipe101 floor, N >= 131072: 33 of 156
cells at or above 60% before, 40 after. m=128 n=1048576 goes 20% to 62%.
op_tests/test_topk_select.py test_lds_sizing, test_invariants and
test_half_invariants all hold, including on cells this newly routes to sampled
(m=128 n=262144 k=2048, m=1024 n=1048576 k=2048).
valarLip
left a comment
valarLip
left a comment
There was a problem hiding this comment.
I re-reviewed at f443413b and reproduced two issues on MI355X (gfx950). Both come from a cache outside this diff, and both reach topk_select itself. The rest is performance, test coverage and docs. Inline comments point at the lines.
1. topk_select can still return wrong indices: the stride fix covers only the single-row case
The new rewrite (input.as_strided((1, width), (width, 1)) for rows == 1) works around the symptom in one caller. The cause is _run_compiled in aiter/ops/flydsl/kernels/tensor_shim.py. It keeps the first CompiledFunction on exe._cf and reuses it for every later call, which bypasses FlyDSL's own cache key: that key distinguishes which axis has stride 1. (tensor_shim.py isn't in this diff, so there's no inline comment for it.)
Repro through the public topk_per_row_argmax, with a fresh process per case:
| first call in the process | then 3× dense (64, 4096) |
|---|---|
dense (64, 4096) |
0/64 rows wrong |
randn(64, 1) (strides (1, 1)) |
63/64 rows wrong, every time |
randn(4096, 1).t() (strides (1, 1)) |
63/64 rows wrong, every time |
The same happens through topk_select alone, in one process:
| call | rows wrong |
|---|---|
topk_select(randn(64, 4096), 1) in a clean process |
0/64 |
topk_select(randn(64, 1), 1) first, then topk_select(randn(64, 4096), 1) |
63/64 |
… then topk_select(randn(64, 4096), 8) / (64, 32768), 512 |
0/64 (other launchers) |
A (rows > 1, 1) input with k = 1 is contiguous with strides (1, 1), routes to argmax, and is not covered by the rows == 1 rewrite. After it, every k = 1 call in the process is wrong. flydsl_top_k_per_row_decode and topk_per_row_argmax called directly are exposed the same way. Keying _cf on the stride pattern, or not caching past FlyDSL's own dispatch for these launchers, would fix it at the source. The rows == 1 rewrite would then be unnecessary.
2. The cross-device guard doesn't make a second GPU work
The same _cf cache has no device in its key, so a launcher stays bound to the GPU of its first launch. The new guard (with torch.cuda.device(dev): return topk_select(...)) switches the current device, but the cached function is still the one loaded on the first GPU. test_cross_device runs in a subprocess and says why: "a process that already ran these selectors on the current GPU cannot run them on another one."
Repro (two GPUs visible, one process):
cuda:0 (current): 64x32768 k=512 -> 0/64 rows wrong
cuda:1 (other) : hipModuleLaunchKernel ... failed with 'hipErrorInvalidDevice'
Any multi-GPU process that calls topk_select from one thread on several devices (e.g. a test harness, or a single-process multi-GPU runner) hits this. The fix belongs with §1: key the cached function on the device too.
3. Performance
- Row-strided uniform input is routed to plain, and plain copies it.
_plain_takesand_plain_k2048_lds_tail_takesdon't look atstride0, whiletopk_plainrunsx = x.contiguous()whenever the call isn't ranged. Example: gfx950, k=2048, a4096 x 131072fp32 view sliced from a wider buffer lands in the LDS-tail band (rows >= 3000). Each call first copies ~2 GiB (read + write), then runs a ~400 µs kernel, and peak memory doubles.sampled, which this band yields to, reads rows at their pitch in place. The plain band should requirestride0 == width, or plain should take the pitch. starts = torch.zeros_like(row_lens)on every ragged plain call. That is an allocation and a fill launch per call, plus an extra node under graph capture._zero_rows, added in this PR for exactly this, isn't used here.build_topk_per_row_radix_stream_module_for_widthisn't cached. Every stream call re-runs_window_cap/_resolve_lds, with their nested budget loop whenlds_planis None, before reaching the@cached builder. All its arguments are hashable, so anlru_cachewould remove that per-call Python cost on a backend that exists for small, launch-bound shapes.- The terminal build reserves LDS it never touches. With
terminal, winners go straight tooutput_idx, andkeep_key/keep_colare never read or written.SharedStoragestill holds2 * k * 4B for them (16 KiB at k = 2048), and_resolve_ldscounts that against the budget. The occupancy and unroll that_resolve_ldspicks for terminal builds are therefore sized for memory that isn't used.
4. Two routing tables for one decision
top_k_per_row_prefill (topk.py) decides "sampled or not" from _SAMPLED_MIN_STRIDE0 = {"gfx950": 131072}. topk_select decides from _sampled_takes: the work door, _SAMPLED_MIN_WIDTH = 131072 again, and yielding to plain's k=2048 LDS tail. They already disagree. At gfx950, 4096 x 131072, k=2048, prefill takes sampled, while topk_select takes plain's LDS tail, which this PR's own comment measured as faster (405 vs 420 µs). Tuning one side will silently leave the other stale. Sharing one predicate, or at least the constant, would keep them in step.
5. Tests, docs, dead code
-
_NotServedmakes some tests pass silently.test_strided_layoutsrecords it as "not served" and logs it._single_row_stride_failures(break) andtest_terminal_window(continue) drop it with no log at all.
If a backend stops serving those shapes (say
_availablenarrows plain, or_lds_budgets(arch)drops stream), those tests still pass, and CI can't tell coverage was lost. At minimum, log it the waytest_strided_layoutsdoes; better, fail when the current arch is one that is expected to serve the shape. -
The
deterministicdocstring is stale. It says it "only excludesplain" and "costs up to 1.9x whereplainwould have won", but_NONDETERMINISTICis now{"plain", "sampled"}. On wide rows,deterministic=Truealso gives upsampled, which is often the fastest backend there, so the documented cost understates it. -
"This entry always instantiates RAGGED=true" (in
topk_sampled_supportsand intopk_common.hip.hpp) is no longer true.top_k_per_row_prefill_sampledtakesragged(defaulttrue), and_sampled_runpassesFalsefor dense rows. Correctness now rests on three things:_sampled_runforcingraggedwhensize(1) != stride0;- the
ragged_smallpromotion (N % FP32_EPT != 0); - the phase-B/C tail handling.
The comment should name those, since someone trusting the current text could remove the flip and read past the row.
-
topk_plain_batches_ragged_rowsis now unused. Its only caller moved totopk_plain_values_optional, and both bodies arereturn _uses_radix(width, topk). Two public names answering one question will drift; keep one. -
A third fp32 sortable-key encoding.
fp32_to_sortable_bits/sortable_key<NAN_HIGH>/key_to_fp32intopk_common.hip.hppre-implement whattwiddle_in<NanHigh>(topk_per_row_kernels.cu) andaiter::radix_traitsalready encode, with the same NaN-above-+inf convention. A later change to NaN or −0.0 ordering that misses one copy would makesampledand the radix paths disagree.
6. A question: the cached bound buffers across streams
_cached_bounds (an lru_cache keyed by rows, width, device) and _cached_zero_rows are created and filled on whichever stream happens to be current at the first call. Later calls on other streams read them with no record_stream and no event. So a second stream could read the ends before the first stream's fill completes. Also, an lru_cache eviction returns the memory to the creating stream's pool while another stream's kernel may still be reading it. Single-stream deployments can't hit this. Is multi-stream use in scope for topk_select? If so, filling them under a sync, or calling record_stream on use, would close it.
…axis _run_compiled keeps the first CompiledFunction on exe._cf and reuses it for every later call, past FlyDSL's own cache key, which holds a tensor's first stride-1 axis static. A first call on a (64, 1) column -- strides (1, 1) -- built a kernel that reads rows at a pitch of 1, and every later dense call through that launcher returned 63 of 64 rows wrong. The cache held no device either, so after one call on cuda:0 an input on cuda:1 failed with hipErrorInvalidDevice. The TopK launchers now go through run_topk_launcher, which hands _run_compiled a specialization key of each tensor's dtype, rank and first stride-1 axis -- what FlyDSL's key reads from a tensor -- and gives each further GPU its own copy of the launcher, whose in-process caches start empty and reload from disk. Extents stay out of the key, as they are out of FlyDSL's, so dense calls of any row count share one build and one disk cache file; a second GPU reads the first one's files. The shared _run_compiled is unchanged. topk_select's rows == 1 rewrite, which worked around one case of this, is gone. test_unit_stride_first gives each FlyDSL route a first call with a unit stride on the leading axis, then dense rows; it fails on all five routes without the stride half of the key. test_cross_device now runs in the test's own process, on the current GPU before and after the other one; it fails without the per-GPU copy. Building the key costs 1.6-2.0 us per call on the argmax, small_k and stream launchers' shapes (MI355X).
topk_plain copied any row-strided uniform call dense before selecting, and topk_select's router sent such calls to plain: a 4096 x 131072 fp32 view at k=2048 moved 2 GiB per call ahead of a ~400 us selection, and doubled the peak memory. The register copies and the LDS tail -- gfx950's k=2048 one-block kernels, and the register window elsewhere -- now take the row pitch. It is an int after k, in padding the kernel arguments already had, so no other argument moves; an int64 after len moved them and cost the few-row register copies up to 4%. topk_plain reads indices-only, largest-first fp32 rows in place where topk_plain_reads_pitch says one of those kernels serves the call, and copies them as before elsewhere; the launchers refuse a pitch that would reach a kernel reading rows width apart. On gfx950 the router sends the row-strided calls plain would still copy to the backends that read them in place. Measured on MI355X, device time, the 590 row-strided cells the router sent to plain (rows 1..16384, widths 2048..1M, k 2..4096; three rounds): 0.601x geomean, at most 1.006x; 4096 x 131072 k=2048 went from 1140 to 408 us. The 165 of them plain now declines run 0.78x, at most 0.99x. Dense rows over the kernels' reach (360 cells, NaN-high off and on, three rounds, the two-level cells over eight processes) are unchanged within run-to-run spread. Dense routing is identical over 105,192 configurations on gfx950 and simulated gfx942; on gfx942 row-strided routing is unchanged too.
Each ragged plain call allocated and filled torch.zeros_like(row_lens) for its row starts, one allocation and one fill launch per call and one more node under graph capture. It takes them from _zero_rows instead.
build_topk_per_row_radix_stream_module_for_width reran _window_cap and _resolve_lds, with the nested budget loop when lds_plan is None, on every stream call before it reached the cached builder. Its arguments are all hashable, so it is cached too.
A terminal build writes its winners straight to output_idx and never touches keep_key / keep_col, but SharedStorage reserved 2 * k * 4 B for them (16 KiB at k=2048) and _resolve_lds counted it, so the occupancy and unroll chosen for terminal builds were sized for memory nothing used. Terminal builds keep one slot -- the traced stores that never run there still need a target -- the budget counts the size the buffer has, and the window cap a terminal build is checked against is asked as terminal. At k=2048 and 1024 threads a terminal build's LDS goes from 81920 B to 65544 B. Non-terminal builds resolve as before. On gfx950 the cells topk_select runs as a terminal build (ragged rows, k=2048; 113 of the sweep) keep their tiling and their terminal decision, and their kernel time is 0.957x of the old sizing over 38 of them (MI355X, profiler, best of three processes per side), 1.005x at worst.
top_k_per_row_prefill routed sampled from _SAMPLED_MIN_STRIDE0, and topk_select repeated the same 131072 as its own _SAMPLED_MIN_WIDTH. Both now ask _sampled_min_width(device), so the door cannot be retuned on one side only. Routing is unchanged. The rest of topk_select's rule stays its own, and the prefill docstring now says where the two still differ and why. Past the door, topk_select gives gfx950 k=2048 uniform rows at 3000+ (4096 x 131072 among them) to plain's LDS-tail kernel, which serves no row bounds, and prefill always passes them. Below it, topk_select also opens sampled on total work and at some gfx950 k=2048 shapes; those doors were measured against topk_select's backends, not against the per-row kernels prefill falls back to.
test_terminal_window dropped a route its backend no longer served with no log, and test_nan_order did the same off gfx950. Both run their routes on gfx950, where an unserved one now fails; test_nan_order logs the routes it skips elsewhere, as test_strided_layouts does.
The deterministic docstring said it only excludes plain and costs up to 1.9x. _NONDETERMINISTIC is now plain and sampled, so it gives up sampled on wide rows too, and both costs were re-measured on gfx950 (MI355X, topk_select with and without deterministic alternated in-process, answers checked): up to 3.4x where plain would have won (49 cells; k=2048, m=4096 n=16384, 75 against 258us on stream) and up to 3.7x where sampled would have (72 wide-row cells; k=2048, m=64 n=1M, 57 against 210us on decode).
The comments in topk_sampled_supports and topk_common.hip.hpp said the aiter entry always instantiates RAGGED=true; _sampled_run passes ragged=False for dense rows. They now name the three things the odd-pitch argument rests on: _sampled_run running a row-strided call as ragged, the ragged_small build for a width that is not a multiple of FP32_EPT, and the fused path's own tail handling -- phase_b_filter_coop picking up a row's partial last vector and phase_c_select_contig's fallback re-reading such a row through the bounds-checked loader.
topk_plain_batches_ragged_rows, which main has, and topk_plain_values_optional, which this PR added, had the same body and answered what is one question: whether the radix path serves the geometry, which decides both whether a rowStarts/rowEnds call stays on the batched launcher and whether topk_out may be None. main's name stays, its docstring now says both, and the PR's name is gone; its callers ask main's.
twiddle_in<NanHigh> in topk_per_row_kernels.cu (twice), the two places that read a radix key back as a float, and sortable_key / key_to_fp32 in topk_common.hip.hpp each encoded the same order. They now come from csrc/include/topk_fp32_key.h, whose static_asserts hold the radix and sampled keys to one order -- the radix key the sampled one reversed, NaN-high moving nothing but NaNs and ranking every NaN above +inf, a NaN ranked by its sign otherwise -- and every key to an exact inverse, NaN payloads included. A change to either encoding that breaks that fails to compile. The radix NaN-high still canonicalises a NaN to its magnitude, where sampled rotates the key space: the radix tail reads a bound back as a float, so it cannot move other keys. Each call site keeps its own expression form, and the device ISA is identical on gfx942 and gfx950 in all four translation units (topk_per_row with and without the plain kernels, topk_plain, sampled).
_cached_bounds was keyed by rows, width and device: filled on the stream of the first call, read from any other with no event, and an eviction returned the memory to the first stream's pool while another stream could still read it. It is keyed by the current stream too (a raw stream id, under 0.1 us a call, where record_stream on the pair at every use costs about 3 us), so a pair is read only after its fill and freed only to the pool of the one stream that reads it. _cached_zero_rows, shared by every stream, now waits for each growth's fill before handing the buffer out; it grows by doubling, so that wait is rare and no call pays a fence. With one stream stalled while it builds them, a call on a second stream read the first one's unfilled bounds and returned wrong rows in 4 of 5 trials keyed by device only, and 0 of 5 keyed by stream; it read the stale zero buffer without the growth wait and zeros with it.
…serves On gfx942 stream serves k up to 256 at 32768 columns (gfx950 up to 2048), so the k=512 stream route was not served there, and as gfx942 is one of the arches expected to serve every route, test_unit_stride_first failed. k=256 is served on both. Without the launch key the route still returns wrong rows on both arches; on gfx942 the decode route does not reproduce the stale build, so there the test covers stream, small_k and argmax.
Off gfx950, topk_select routes uniform rows at a pitch past their width to plain wherever it routes their dense twin, and topk_plain copied every such call its uniform kernels could not read at that pitch, at twice the input's memory. plain's ranged form -- every row from column 0 to the width -- reads them in place at the pitch. topk_select now takes it wherever both forms run plain's one-block kernel (_plain_reads_pitch_ranged), where it is the kernel the copy would feed, minus the copy. The ranged form hands the launcher the pitch as the row length, and wherever the multi-block kernel takes either form the pitch decides its speed, so those keep the copy. C++ picks radix or block-sort from the width, not the pitch, so the ranged form stays on the batched radix launcher. gfx950 is unchanged: it declines these calls to stream, which measured faster than the copy on all 165 cells; plain's ranged form beat stream on 136 of them but ran up to 1.19x slower on the other 29. test_plain_pitch_in_place watches Tensor.contiguous through topk_select and fails on a copy wherever either predicate says plain reads the pitch; with the ranged branch removed it fails on gfx942.
The comment over the gfx942 below-128-CU k=2048 bands named one part; the CU count it already gives is what the bands key on.
…_stream_module terminal had been added between labelled and block_threads, so a caller passing block_threads or a later argument by position, as main allowed, would have set terminal instead. It is the last parameter now; every caller in aiter names it.
|
@aiter-bot review |
ROCm/aiter PR #5686 — [HIP] [FlyDSL] [JIT] Feat/topk per row gfx950Adds a fifth Review (advisory): 🔴 HIGH RISK 🔴 [verified] The PR no longer applies to main —
📝 [verified] The small-reject Triton kernel's hand-written fp32 key transform (NaN remapped to 0x7FFFFFFF, claimed to match the stream key exactly) has no NaN coverage — test_nan_order's routes in op_tests/test_topk_select.py include no stream k=2048 entry for the 2049..2056-column band it serves, and the strided-layouts route over that band checks values only, so a NaN-order regression in this kernel would pass CI. Reviewer should ask for a NaN row through the small-reject route, or drop the untested remap claim. |
_launcher_on shared the first GPU's launcher with a further one whenever the installed FlyDSL lacked any of the cache attributes it resets, which brings back the cross-device fault it exists to prevent, and its copy kept the GPU target FlyDSL 0.4.0 resolves once with the signature (_sig, _backend_target), so a second GPU of another architecture would have built for the first. The copy now starts every piece of per-device state afresh, the target included, and a launcher that cannot be copied raises on the further GPU instead. test_launcher_per_device checks both on one GPU, by asking for a device id nothing is launched on: CI's runners have one GPU, where test_cross_device returns early. Sharing the launcher, or keeping _sig and _backend_target, fails it.
The gfx950 k=2048 small-reject kernel maps a NaN of either sign above +inf, as the stream key does, but test_nan_order had no route through it. It now has one per row band of _stream_small_reject (8 x 2050, 1024 x 2055, 4096 x 2056).
…does They log what failed rather than assert: on gfx942 stream does not serve k=512 at 1024 threads, so test_full_rows_graphs cannot dispatch that shape there, on main or here.
…before Re-measured every routing change against the choice the rules make without it, in device time of replayed graphs (MI355X, randn, answers checked; dense rows and rows ending at random lengths), and kept each door only where it held: - ragged k=2048 rows: decode below plain's row floor up to 65535 columns and `sampled` from 65536; `sampled` had taken them from 16385 and lost to decode by up to 2.6x when rows end well short of the width. - whole-row write without values: 640..8192 rows only; at 5..512 rows and at 16384 it ran 1.02x-1.08x of stream. - `sampled`: k in 512..4096, the work door only from 768 to below 4096, and ragged rows past 2048 below 262144 columns only inside that range. Outside it lost by up to 1.49x (k=96..256), 1.12x (k=512), 1.51x (k=4096) and 3.36x (k=6144). - gfx942 k=2048 plain bands: below 128 CUs from 128 rows through 20000 columns, with 128 CUs or more the uniform register-window triplets only; elsewhere plain lost to decode by up to 2.55x. Other targets take the general rules. The k=2048 bands now sit in one place per architecture: _route_gfx950_k2048 tries gfx950's in order and ends on the general tail, and _plain_gfx942_k2048_takes holds gfx942's, so _plain_takes and _sampled_takes keep only the general rules. The two Triton kernels that served stream's gfx950 k=2048 rows from inside _dispatch and the entry are backends of their own, `stream_reject` and `whole_row`, so the router names the kernel that runs and the tests drive them by name.
and from 32768 for rows with their own bounds. At 16384 columns the build ran 1.01x-1.06x of the re-scan one at 64..4096 rows (k=512 and 1024), and bounded rows, which mostly end short of the padded length, 1.05x-1.06x at 16384 and up to 1.03x at 24576; from those lengths up it is 0.74x-0.98x.
…12..2048 top_k_per_row_prefill and topk_select each opened their own door into `sampled`, in topk.py and topk_select.py. topk_select's now sits beside prefill's in topk.py, and the two share the width door and the k range. Re-timed against main with rows that end far short of their pitch -- a preallocated one, which the earlier measurements never had (MI355X, device time of replayed graphs). `sampled` takes as long whatever the rows' live length, so there it ran up to 5.3x slower than the kernels that follow it: - topk_select gives it uniform rows only; - top_k_per_row_prefill, which always takes row bounds, at most 64 rows, where it ran 0.24x-1.00x of the per-row kernels with rows reaching the pitch or ending in [k, 4k]; past 64 rows those rows ran up to 3.5x slower; - both take k from 512 to 2048. Below 512 one such row ran 1.41x slower through prefill (k=128); above 2048 a row whose sample undershoots k falls back often enough to cost 2.11x on uniform rows (k=4096, 64 rows of 1M, 1 of 8 seeds). The dispatch test covers the prefill bounds. op_tests/op_benchmarks/hip/bench_topk_select_routing.py times topk_select and top_k_per_row_prefill per call over grids of shapes and row-end patterns (replayed CUDA graphs, every answer checked against torch.topk, the backend each call ran) and compares two trees cell by cell; the routing comments cite it rather than individual runs.
Thanks for the review. Each point below, as of 1. Rebase / the two end=None designsMerged current main (with #6833) and took its end=None design: argmax, small_k, decode and stream take the empty 2.
|
| rows \ width | 2048 | 2050 | 2056 | 4096 | 8192 | 16384 | 32768 | 65536 | 131072 |
|---|---|---|---|---|---|---|---|---|---|
| 1 | 1.00 stream | 0.48 stream_reject | 0.64 plain | 0.58 plain | 0.52 plain | 0.51 plain | 0.58 plain | 0.43 plain | 0.61 sampled |
| 8 | 1.00 stream | 0.48 stream_reject | 0.70 plain | 0.65 plain | 0.55 plain | 0.51 plain | 0.59 plain | 0.43 plain | 0.62 sampled |
| 64 | 1.00 stream | 0.48 stream_reject | 0.74 plain | 0.68 plain | 0.57 plain | 0.51 plain | 0.57 plain | 0.41 plain | 0.69 sampled |
| 256 | 1.00 stream | 0.49 stream_reject | 0.62 plain | 0.64 plain | 0.44 plain | 0.51 plain | 0.42 plain | 0.37 plain | 0.63 sampled |
| 1024 | 0.95 whole_row | 0.37 stream_reject | 0.55 plain | 0.46 plain | 0.53 plain | 0.57 plain | 0.43 plain | 0.37 plain | 0.49 sampled |
| 4096 | 0.86 whole_row | 0.27 stream_reject | 0.41 plain | 0.38 plain | 0.49 plain | 0.54 plain | 0.43 plain | 0.45 plain | 0.42 plain |
Rows ending uniformly in [k, width]
| rows \ width | 2048 | 2050 | 2056 | 4096 | 8192 | 16384 | 32768 | 65536 | 131072 |
|---|---|---|---|---|---|---|---|---|---|
| 1 | 0.98 stream | 0.89 stream | 1.01 stream | 0.89 stream | 1.00 decode | 1.00 decode | 1.00 decode | 1.00 decode | 0.99 decode |
| 8 | 1.01 stream | 0.88 stream | 0.89 stream | 0.89 stream | 1.00 decode | 1.00 decode | 1.00 decode | 1.01 decode | 1.01 decode |
| 64 | 1.02 stream | 0.89 stream | 0.89 stream | 0.91 stream | 1.00 decode | 1.00 decode | 1.00 decode | 1.00 decode | 1.00 decode |
| 256 | 1.00 stream | 0.89 stream | 0.90 stream | 0.90 stream | 0.83 plain | 0.93 plain | 0.86 plain | 0.87 plain | 1.00 stream |
| 1024 | 1.00 stream | 0.84 stream | 0.87 stream | 0.75 stream | 0.91 plain | 0.98 plain | 0.96 plain | 0.86 plain | 1.00 stream |
| 4096 | 1.00 stream | 0.69 stream | 0.69 stream | 0.70 stream | 0.99 plain | 1.04 plain | 0.97 plain | 0.89 plain | 1.00 stream |
Rows ending uniformly in [k, 4k] (far short of the pitch)
| rows \ width | 2048 | 2050 | 2056 | 4096 | 8192 | 16384 | 32768 | 65536 | 131072 |
|---|---|---|---|---|---|---|---|---|---|
| 1 | 1.01 stream | 0.89 stream | 1.01 stream | 0.90 stream | 0.99 decode | 1.00 decode | 0.99 decode | 1.01 decode | 1.00 decode |
| 8 | 1.01 stream | 0.88 stream | 0.89 stream | 0.89 stream | 1.00 decode | 1.00 decode | 1.00 decode | 1.01 decode | 1.00 decode |
| 64 | 1.00 stream | 0.88 stream | 0.89 stream | 0.91 stream | 1.00 decode | 1.00 decode | 1.00 decode | 1.01 decode | 1.01 decode |
| 256 | 1.00 stream | 0.90 stream | 0.90 stream | 0.90 stream | 0.84 plain | 0.85 plain | 0.88 plain | 0.89 plain | 1.01 stream |
| 1024 | 1.00 stream | 0.85 stream | 0.88 stream | 0.76 stream | 0.90 plain | 0.91 plain | 0.93 plain | 0.89 plain | 1.00 stream |
| 4096 | 1.00 stream | 0.68 stream | 0.69 stream | 0.70 stream | 1.00 plain | 1.00 plain | 1.03 plain | 1.02 plain | 1.00 stream |
5. NaN coverage for the small-reject kernel
Added. The small-reject kernel is now its own backend, stream_reject, and test_nan_order drives it by name through three routes, one per row band of _stream_small_reject (8x2050, 1024x2055, 4096x2056, k=2048), plus whole_row (1024x2048). Each route uses NaNs of both signs and several payloads, infinities, signed zeros and NaNs past a row's end, and checks indices in range, no repeats, and the selected set against the documented order. The remap claim holds: the Triton key is the stream key (_ord_unsigned) without its final sign flip, -0 collapses to +0 in both, and every NaN goes to the top of each encoding (0x7FFFFFFF signed, 0xFFFFFFFF unsigned), so the two order identically.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist