Repository navigation
hrx: cooperative bucket scan, single emit pass and vec4 row reads for hip_top_k_f32 (engine #679) - #195
Merged
bong-water-water-bong merged 3 commits intoNov 4, 2026
Conversation
Author
bong-water-water-bong
left a comment
bong-water-water-bong
left a comment
Author
There was a problem hiding this comment.
VERDICT: READY on the code — the kernel compiles for gfx1151 and the three rewrites each preserve the index contract as
read; the 484/484 device gate is the attn lane's run, not mine.
Reviewed at b2b545d06 (1 file, +101/-28 in ggml/src/ggml-hrx/hip/kernels/hip_top_k_f32.hip).
What I actually ran:
$ amdclang++ -x hip --offload-arch=gfx1151 -O3 -DNDEBUG \
--rocm-device-lib-path=$(amdclang++ --print-resource-dir)/../../../amdgcn/bitcode \
-c ggml/src/ggml-hrx/hip/kernels/hip_top_k_f32.hip
rc 0, 0 warnings
(and that flag is required on this box — without it the compile fails on the device library, the same finding as #185.)
What I could not reproduce, stated plainly: the device gate. test-backend-ops -o TOP_K -b HRX0 484/484 twice plus
the GGML_HRX_DISABLE_DISPATCH=hip.top_k_f32 control is the attn lane's gpu-hold run at this head; the build tree it
was taken in did not survive the reboot, and a fresh GGML_HRX=ON configure still fails on the hrx package here (the
two routes that work on this box are a CPU-only configure and replaying recorded compile commands, which is exactly what
the rc 0 above is). So the correctness half of this review is reading the code against the old kernel, not a run.
The reading, in the order the risk sits:
- The suffix scan.
lane_sum[lane] = chunkthenv = lane_sum[lane + off]read before the write to
lane_sum[lane], with a sync on both sides — a correct in-place reverse Hillis-Steele, and it does not depend on
wave width (which is the defect the body says the shuffle version had).suffix_excl = suffix_incl - chunk, so the
crossing predicatesuffix_incl >= remaining && suffix_excl < remainingselects exactly the highest bucket whose
inclusive suffix count reachesremaining— the same bucket the serial descending loop picks, ties included, because
both stop at the first bucket where the cumulative count first reachesremaining. above_slots = top_k - remainingcannot wrap:remainingstarts attop_kand each level sets it to
remaining - cumulative <= remaining, so it is monotonically non-increasing and never exceedstop_k. If the
crossing lane were somehow not found,take_foundstays at its fallbackremainingand the emit degrades to "all
equal-group", not to an out-of-bounds slot. That is the same bound the two extra stores buy.- One emit pass, duplicate-free. Strictly-above goes to
[0, above_slots)and the equal group to
[above_slots, top_k); the groups are disjoint by construction of the threshold key, each element is visited once, so
thetiesvariant's "no repeated index" requirement still holds without the second scan. - The vec4 visit. Groups cover
[0, n4*4)and the scalar tail starts atn4*4 + lane, so no element is visited
twice and none is skipped; within a group the four indices are ascending, and lane-stride ordering was never part of
the contract (the histogram is order-insensitive and the emit uses atomics on disjoint slot ranges). The alignment
guard covers the row base, which already includes the pitch, so an unaligned pitch simply takes the scalar path.
Nits, none blocking:
topk_visittakesf(uint32_t, float)but the histogram lambda wastes the index on(void) i;. Either give the
helper a value-only overload or drop the parameter; as written a reader wonders which call site needed the index.- The body's three round labels (
r3_d0,r1_d0,r1_d3) refer to a design document that is not in this repo or in
the fork; the measured rows cannot be traced from the PR alone.
…2 (engine ggml-org#679) Two changes to the radix select from #161, both aimed at a measured defect rather than guessed at. 1. The descending bucket scan ran SERIAL in lane 0: 3 levels x 2,048 buckets = 6,144 dependent LDS reads plus integer work, independent of n. That is the ~100 us fixed cost test-backend-ops perf shows at n=1000, paid by all 12 full-attention layers of Qwen3.8-Flash-Next on every decode step. Now each lane owns 8 consecutive buckets, the 256 per-lane chunk sums get a reverse (suffix) inclusive scan as an LDS Hillis-Steele tree (8 steps, 2 syncs each), the crossing lane is the one with suffix_incl >= remaining > suffix_excl, and only that lane walks its own 8 buckets. Same bucket, same residue, same tie behaviour as the serial scan. 2. The emit phase ran TWO full passes over the row: one for keys strictly above the rank-k key, one for keys equal to it. The split is already known when the level loop ends -- exactly top_k - remaining keys are above and the other "remaining" are equal -- so both groups can be written in one pass into disjoint slot ranges, [0, above_slots) and [above_slots, top_k). One less pass over the row; the per-group slot guard keeps the result duplicate-free, which is what the harness's ties variant requires. Deliberately not a warp-shuffle scan. The first attempt used __shfl_down_sync and returned wrong index sets (ERR 4.0, then 0.667 with a range guard) because shfl returns the caller's own value for an out-of-range source lane and the wave width is not something this kernel should depend on. The LDS tree is wave-agnostic. Measured (test-backend-ops perf -o TOP_K -b HRX0, arms interleaved A/B/A/B, best of 3 rounds per shape, same binary, only this file differs): n=1000 rows=1, k=400 106 -> 11 us 9.36x n=1000 rows=16,k=400 106 -> 12 us 9.05x n=65000 rows=1, k=400 225 -> 133 us 1.70x n=65000 rows=16,k=400 328 -> 211 us 1.55x n=200000 rows=1, k=400 617 -> 423 us 1.46x n=200000 rows=16,k=400 865 -> 620 us 1.40x geomean 2.77x, worst 1.40x; baseline's own per-round spread 1.01x Correctness: test-backend-ops -o TOP_K -b HRX0 484/484 with hip.top_k_f32 still matched, the GGML_HRX_DISABLE_DISPATCH control still declining every case as unsupported (529 cases), CPU reference 529/529. Registers: 17 VGPRs (was 8), 26 SGPRs, 0 spills, occupancy 16 waves/SIMD, LDS 9,228 B (was 8,204). Only the .hip changed; the matcher, its shape bounds and the device_supports_op claim are byte-identical, and both arms report "HIP add-on ABI check: 336 kernels match their code objects". Not measured, stated: the model-level row (18.8k decode, kernel ON vs GGML_HRX_DISABLE_DISPATCH=hip.top_k_f32) needs the 76 GiB Flash-Next file and ~92 GiB of headroom; the owner's serve holds ~35 GiB, so it waits. Worth knowing before anyone reads this as a decode win: top_k at 18.8k is ~50 us x 12 layers, about 0.8% of a 79 ms step, so these are op-level gains -- the decode residue is the O(n_kv) mask rebuild, not the sort.
The cooperative scan and the single emit pass (previous commit) removed the fixed cost and one pass; what is left is memory latency, not arithmetic. One pass over a 200,000-element row is 200000/256 = 781 iterations per lane and costs ~143 ns each, i.e. ~300 cycles: exactly one un-hidden DRAM latency per iteration, because a lane has a single outstanding load and the workgroup is only 4 waves at rows=1. Adding passes or removing them moves that number very little -- that is what the rejected compaction experiment measured. So give each lane four floats in flight per iteration. Every pass over the row (the three radix histograms and the emit) now goes through one topk_visit helper that loads float4 when the row base is 16-byte aligned -- the row base is src + row * column_count, so the guard covers both the pointer and the pitch -- with a scalar tail for column_count % 4 and a fully scalar fallback. Index order is unchanged (i*4 + 0..3 is ascending), so slot numbering and the ties rule are untouched; the ABI, the matcher bounds and the device_supports_op claim are not touched at all. Same-hold A/B against the previous commit, arms interleaved A/B/A/B, best of 3 rounds per shape, the fork's own test-backend-ops perf: shape before us/run after us/run ratio n=1000, rows=1, k=400 12/12/12 12/11/11 1.03x n=1000, rows=16, k=400 12/12/12 12/11/11 1.07x n=65000, rows=1, k=400 132/132/127 57/57/55 2.30x n=65000, rows=16, k=400 221/224/211 83/81/81 2.61x n=200000, rows=1, k=400 422/426/414 169/167/165 2.51x n=200000, rows=16, k=400 666/636/622 217/216/212 2.93x geomean 1.909x, worst 1.03x. Cumulative against the merged pin (1991650) on the same gate, same protocol: geomean 5.436x, worst 4.23x (the pin's own per-round spread reaches 1.65x, so the minimum over interleaved rounds is the estimator and the raw values are above). Verification: test-backend-ops -o TOP_K -b HRX0: 484/484, "ggml_hrx dispatch: hip.top_k_f32 (single)" taken. The bounded claim still declines the 45 shapes it is not meant to serve: with GGML_HRX_DISABLE_DISPATCH=hip.top_k_f32 all 529 cases report "not supported [HRX0]" and no case fails, and the CPU reference passes 529/529. Whole-backend guard: the full HRX0 suite's 4289 per-case outcomes are identical between the two arms (both carry this pin's pre-existing 320 [VIEW]/TOPK_MOE failures, which are flaky in their error magnitude and are stripped before the comparison). Registers 18 VGPRs / 34 SGPRs, 0 spills, 0 scratch, occupancy 16 waves/SIMD, LDS 8,228 B (unchanged). Not measured here: the end-to-end model row. It needs the 76 GiB qwen4exp file at -c 19456 and ~92 GiB of available memory, which this box cannot currently give while another lane's serve is resident; the op-level shapes this kernel serves for QSA (rows=1, k=2051, n=18816) are not in the harness sweep at all, so that row is the outstanding acceptance item and is not being claimed. Co-authored-by: pi-attn <pi-attn@1bit-MONSTER>
…ils (engine ggml-org#679) The rank-k resolve publishes `bucket_found` and `take_found` from the one lane that crosses the descending suffix scan, and every lane reads them after the barrier. The invariant that exactly one lane crosses holds today -- the histogram of elements matching the resolved prefix sums to >= remaining, so the scan always finds that lane -- but nothing in the kernel enforces it, and the failure mode is not a wrong answer, it is an out-of-bounds write: a stale `take_found` larger than `top_k` makes `remaining` exceed `top_k`, `above_slots = top_k - remaining` wraps to a huge value, and the emit's `slot < above_slots` guard then passes for every element in the row. Reset both values at the top of each level (two stores in lane 0, next to the histogram clear that already ends in the same barrier) so `remaining <= top_k` holds unconditionally and `above_slots` cannot wrap. Also fixes a comment on `topk_vec4_ok` that claimed the 16-byte check needs "both the base pointer and the row pitch": the row base already includes the pitch, so testing the pointer is sufficient (an unaligned pitch just sends the following rows down the scalar path). Verified on gfx1151 after the change: `test-backend-ops -o TOP_K -b HRX0` 484/484 twice with `ggml_hrx dispatch: hip.top_k_f32 (single)`.
bong-water-water-bong
force-pushed
the
1bit/hrx-top-k-coop-scan
branch
from
November 4, 2026 13:13
b2b545d to
bbeea63
Compare
Author
|
Rebased onto |
bong-water-water-bong
left a comment
bong-water-water-bong
left a comment
Author
There was a problem hiding this comment.
VERDICT: READY
Head bbeea63aa: the PR's own file is unchanged since my READY at b2b545d06 —
git diff b2b545d06 bbeea63aa -- ggml/src/ggml-hrx/hip/kernels/hip_top_k_f32.hip is empty, and the three-dot diff
against 1bit/hrx-vulkan-patched is still that one file (+113/-28). The commits listed after my review are the base
branch moving (convert_hf_to_gguf.py's --dspark flag, which is not this PR's).
Holds from that review: the kernel compiles for gfx1151 with amdclang, the three rewrites (cooperative bucket scan,
single emit pass, vec4 row reads) each preserve the index contract as read against the CPU reference, and the lane's
484/484 device gate is theirs, not mine.
This was referenced Nov 4, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR: ggml-hrx — cooperative bucket scan, single emit pass, and vec4 row reads for
hip_top_k_f32(engine ggml-org#679)Branch
1bit/hrx-top-k-coop-scan@b2b545d06, base1bit/hrx-vulkan-patched@d8bd0ff60(current tip),2 commits, 1 file, +101/-28:
ggml/src/ggml-hrx/hip/kernels/hip_top_k_f32.hip. Nothing else in the tree moves(the matcher, its shape bounds and the
device_supports_opclaim are byte-identical to the pin).Rebased onto the current base tip on 2026-11-03; the rebase was verified conflict-free and the resulting kernel file is
byte-identical to the pre-rebase version, and the op gate was then re-run on the rebound head itself (see
"Verified on this head"). Commit 3 (
b2b545d06) bounds the emit against a broken crossing-lane invariant -- see"Hardening" below; it is two stores per level and changes no result (484/484 before and after).
Follow-up to #161, which put
GGML_OP_TOP_Kon HRX and closed the QSA cliff and the prefill gap. This removes the threecosts that dispatch left behind: a serial scan, a second emit pass, and one outstanding load per lane.
Re-verified 2026-11-04 after the c16 hardware reboot (this head, unchanged)
The head below is the head that was measured and gated:
b2b545d06, and the branch has not moved since. Correctness wasre-run on it again today on a freshly built
test-backend-ops(build tree atb2b545d06, gfx1151, onegpu-hold -l attnclaim, die 51 C, 0 gfxhub faults):
The perf rows in this body were taken on this same head on 2026-11-03 with the interleaved A/B protocol described in
"How the numbers were taken"; they are not re-quoted from a different build, and I did not re-run the A/B today because
this box drifts 13-15 % per hour (measured, recorded in the lane state) and a single sweep is the thing that protocol
exists to defeat.
What it changes
The descending bucket scan was serial in lane 0 — 3 levels x 2,048 buckets = 6,144 dependent LDS reads plus
integer work, independent of
n. That is the ~100 µs fixed costtest-backend-ops perfshows atn=1000, and all12 full-attention layers of Qwen3.8-Flash-Next pay it on every decode step.
Now each lane owns
kPerLane = 8consecutive buckets; the 256 per-lane chunk sums get a reverse (suffix) inclusivescan as an LDS Hillis-Steele tree (8 steps, 2 syncs each); the crossing lane is the one with
suffix_incl >= remaining > suffix_excl, and only that lane walks its own 8 buckets. Same bucket, same residue, sametie behaviour as the serial scan — both select the highest bucket whose inclusive suffix count reaches the rank.
The emit phase ran two full passes over the row (keys strictly above the rank-k key, then keys equal to it). The
split is already known when the level loop ends — exactly
top_k - remainingkeys are above, the otherremainingare equal — so both groups are written in one pass into disjoint slot ranges,
[0, above_slots)and[above_slots, top_k). The per-group slot guard keeps the result duplicate-free, which is what the harness'stiesvariant requires. Five passes over the row become four.
Each lane had exactly one load in flight. One pass over a 200,000-element row is 781 iterations per lane and
costs ~143 ns each — ~300 cycles, i.e. one un-hidden DRAM latency per iteration, because the workgroup is only 4
waves at
rows=1and nothing else can be outstanding. Every pass over the row (the three histograms and the emit)now goes through one
topk_visithelper that loadsfloat4when the row base is 16-byte aligned — the base issrc + row * column_count, so the guard covers both the pointer and the pitch — with a scalar tail forcolumn_count % 4and a fully scalar fallback. Index order is unchanged (i*4 + 0..3is ascending), so slotnumbering and the ties rule are untouched.
Measured
test-backend-ops perf -o TOP_K -b HRX0, same binary, same model, arms interleaved A/B/A/B, best of 3 rounds pershape (see "How the numbers were taken" — a single sweep is not usable here). Commit 1 against #161:
geomean 2.77x, worst 1.40x; the baseline's own per-round spread in the accepted run was 1.01x, so the deltas are
not drift. The small-
nrows are the fixed cost disappearing; those two changes remove work per pass, which is why thelarge-
nrows gain only 1.4-1.7x there.Commit 2 (vec4 row reads) against commit 1, same hold, same protocol:
geomean 1.909x, worst 1.03x. Cumulative against the pin, same gate: geomean 5.436x, worst 4.23x (the pin's own
per-round spread reaches 1.65x in that run, hence the best-of-rounds estimator).
One direction tried and rejected (so nobody re-spends it)
After commit 1 the remaining cost looked like "too many passes", so I implemented candidate compaction: after level 0,
write the higher-bucket indices straight into dst and collect the selected bucket's members into an LDS candidate list,
then run levels 1/2 and the rest of the emit over ≤ 2048 candidates (4 passes → 2). It was rejected on measurement:
0.973x geomean, worst 0.95x against commit 1, on a quiet box with per-round spreads ~1.0x. The passes it removes are
cheap — the row is cache-warm by then — and its restructured base was ~5-10% slower to begin with. That failure is what
identified memory-level parallelism as the real limit and produced commit 2. The algorithm was proved in
scripts/ref_topk.py(both paths against a sort reference, 616 cases, 0 failures) before any GPU time; the candidate iskept in
~/geak-tasks/tattn-top-k/rounds/for the record.Correctness
test-backend-ops -o TOP_K -b HRX0: 484/484,GGML_HRX_LOG_DISPATCH=1still showinghip.top_k_f32 (single).Note the harness has 529 TOP_K cases in total: the other 45 are shapes the bounded claim declines by design
(
k > n, non-f32,nabove the template range) and they correctly fall to the CPU backend. 484 of 529 is theexpected result, not a regression.
GGML_HRX_DISABLE_DISPATCH=hip.top_k_f32: all 529 cases reportnot supported [HRX0]and the harness exits 1 with"no test case ran" — that is the correct behaviour (unclaimed node, real graphs put it on CPU), and it is what the
control asserts rather than
rc == 0.test-backend-ops -b HRX0on both arms — identical 4289 per-case outcomes(3968/4288 passed on both; the 320 failing cases are pre-existing
[VIEW]/TOPK_MOEfailures at this pin, unrelated,and their
ERRmagnitude even varies run to run). Only one.hipchanged, the matcher is byte-identical, and botharms report
HIP add-on ABI check: 336 kernels match their code objects.8,228 B/block (commit 1 was 17 VGPRs, LDS 9,228 B). Every row above was re-verified for the head that is pushed:
484/484 with the matcher taken, control declining all 529, CPU 529/529, guards identical.
How the numbers were taken (and why the obvious protocol does not work)
A byte-identical candidate measured 0.59x-2.05x per shape on a single 81-shape sweep, and 0.62x-1.95x with 3 sweeps
per arm in two phases: the disturbance on a shared box is one-sided (another lane's load or a thermal event makes one
sweep ~1.8x slower for every shape at once), so a phase-ordered or single-sample protocol measures the box. Interleaving
the arms round by round and taking the best of N per shape makes an identical candidate read 0.985x, and the raw
per-round values are printed beside every row so the contamination stays visible.
Hardening (commit 3,
b2b545d06)bucket_found/take_foundare published by the one lane that crosses the descending suffix scan and read by every laneafter the barrier. The invariant that exactly one lane crosses always holds (the histogram of elements matching the
resolved prefix sums to
>= remaining), but if it ever failed the failure mode is not a wrong answer, it is anout-of-bounds write: a stale
take_found > top_kmakesremaining > top_k,above_slots = top_k - remainingwraps,and the emit's
slot < above_slotsguard passes for every element of the row. Both values are now reset at the top ofeach level (two stores in lane 0, beside the histogram clear that already ends in the same barrier), so
remaining <= top_kholds unconditionally. Also corrected a comment that claimed the vec4 alignment check needs "both the base pointer and the
row pitch" - the row base already includes the pitch.
Verified on this head (
b2b545d06, 2026-11-03)Rebuilt on the rebound base (
ninja test-backend-ops llama-server), then on gfx1151:Perf on the same build (
test-backend-ops perf -o TOP_K -b HRX0, single row unless noted), against the #161 pin measuredon 2026-10-31 with the same
k=400shapes:Geomean 5.8x over those six shapes; the fixed cost that all 12 QSA layers pay every decode step drops from ~100 us to
~10 us (the launch floor), which is the cooperative bucket scan doing what it was for.
What is not claimed
top_kat 18.8k context is ~50 us x 12 layers ≈ 0.8% of a ~79 ms decode step, so themodel row should be flat or slightly better. That row (18,782 tokens, 512 greedy, median of 3,
-c 19456 --no-repack)needs the 76.33 GiB Qwen3.8-Flash-Next file and a GTT pool of >= 88 GiB (
GTT total - in_use - 16 GiB reserve, theguardrail from engine Update README.md ggml-org/llama.cpp#680/hrx: refuse a launch that would plan past the GTT ceiling (engine #680) #166); with the owner's Bonsai serve holding ~32 GiB of GTT, every other tenant's budget has
to stay under ~18 GiB, and it has not. It is armed and not yet measured, so this PR rests on op-level numbers plus
the backend-equality guard.
rows=1it now runs at roughlythe per-workgroup latency limit (see commit 2), while
rows=16still does 16x the work in ~1.3x the time. Splittinga long row across workgroups remains the biggest lever for the QSA decode shape, but it needs a global histogram and a
device-scope counter — an extra kernel binding for scratch or a cooperative launch — so it is a matcher-level
change outside this PR's contract (the ABI, the shape bounds and the
device_supports_opclaim are untouched here).rows=1,k=2051,n=2051/18816/262144) are not in the harness sweep(
k=2051never appears in it);k=400is the proxy above and the model row is their real evidence.Repro
Those commands, the frozen
baseline_src/,meta.json, the round log (round 0's protocol validation, r1_d0, r1_d3,r2_d1 rejected, r3_d0) and the reference model
scripts/ref_topk.pylive in~/geak-tasks/tattn-top-k/onstrixhalo.
rocprofv3records nothing for HRX (fork PR #154), so the loop is benchmark plus static ISA.Head on the current base, re-gated there
Base moved while this PR was open (
d8bd0ff60->a5c32fbf8, oneconvert_hf_to_gguf.pycommit that does not touch thistree; the three commits here are byte-identical after the rebase). Heads are now
723bb93e1,008bc0749,bbeea63aaon base
a5c32fbf8, and the op gate was re-run on that head rather than carried over:GGML_HRX_LOG_DISPATCH=1 test-backend-ops test -o TOP_K -b HRX0->ggml_hrx dispatch: hip.top_k_f32 (single),484/484 tests passed,
Backend HRX0: OK, twice,rc=0(~/lb/aval/attn-topk-regate2/20261104T131313Z).The handful of
ne=[524299,1,2,1]shapes the matcher declines are the same ones as on the pre-rebase head.pin, worst 4.23x, interleaved A/B best-of-3), the rebase does not change a byte of
hip_top_k_f32.hip, and the boxdrifts 13-15 %/hour, so a fresh single sweep would be worse evidence than the interleaved one. Said plainly instead of
quoting a new number.
HIP_DEVICE_LIB_PATH=.../_rocm_sdk_core/lib/llvm/amdgcn/bitcodefor the.hsacostep(otherwise andclang fails with "cannot find ROCm device library" - that is hrx: give amdclang++ the device-library bitcode path so a changed HIP kernel compiles #185's subject).