Skip to content

ggml-hrx: TOP_K on HRX, a radix select for the QSA indexer (engine #679) - #161

Open
bong-water-water-bong wants to merge 1 commit into
1bit/hrx-vulkan-patchedfrom
1bit/hrx-top-k
Open

bong-water-water-bong wants to merge 1 commit into
1bit/hrx-vulkan-patchedfrom
1bit/hrx-top-k

Conversation

@bong-water-water-bong

Copy link
Copy Markdown

GGML_OP_TOP_K had no HRX dispatch, so the node ran on the CPU backend and split
every full-attention layer's graph around it. Qwen3.8-Flash-Next is the model that
pays for it: its QSA layers (src/models/qwen4exp.cpp) call ggml_top_k with
n = n_kv (up to 262144) and k = indexer_top_k + r - 1 = 2051, once per
full-attention layer per decode step, and the CPU side is a single-row
std::partial_sort with an indirect comparator plus a host round trip.

Measured on strixhalo (gfx1151), same server, same prompt family, one process:

  • decode steps 39.2 -> 49.0 ms the moment n_kv crossed 2051 (1986 vs 2096 tokens),
    i.e. +9.7 ms per step for entering the QSA path at all;
  • then ~3.8 us per KV token per step (49.0 ms at 2096 -> 67.4 ms at 6978), which
    extrapolates to 112 ms/step at 18,782 tokens against the 103 ms a full protocol
    run measured there (9.70 tok/s vs 31.15 for a build whose TOP_K is on the GPU);
  • forcing the same path below the threshold (LLAMA_QSA_DENSE=0) at fixed n_kv=1480
    costs 13x prefill (676 -> 52 tok/s), which is the selection path itself.

The kernel is a radix select on a monotone uint32 key map of the f32 bit pattern:
three levels of 11/11/10 bits narrow to the exact rank-k key, then one emit pass
writes the indices above it and enough equal ones to reach k. One workgroup per
row, an 8 KiB LDS histogram, three passes over the row, no global scratch and no
host round trip. Order is not part of ggml's TOP_K contract (test_top_k compares
sets and the CPU reference even swaps two entries to say so).

The matcher is a SingleOp registration that declares the op with
hip_declare_eager_op, so the node reaches the dispatcher at all; because the op is
now claimed, ggml-hrx.cpp's device_supports_op bounds the claim per node through
hip_top_k_tensor_supported, mirroring the matcher's own predicate and the
GGML_HRX_DISABLE_DISPATCH switch, so a shape the kernel cannot take -- or a build
with the matcher turned off -- keeps its CPU fallback instead of failing the graph.

Verified on gfx1151:

  • test-backend-ops -o TOP_K -b HRX0: 484/484 passed with
    "ggml_hrx dispatch: hip.top_k_f32 (single)" in the dispatch log;
  • the same command with GGML_HRX_DISABLE_DISPATCH=hip.top_k_f32 reports TOP_K as
    skipped on HRX0 (clean fallback, no unsupported-node failure);
  • the 11/11/10 radix select was replicated exactly and compared against the true
    top-k over 600 randomised vectors plus edge cases (all-equal, sorted, k=1, k=n,
    +-inf/NaN, and the QSA shape n=18816/k=2051): index sets and thresholds match.
  • test-backend-ops perf -o TOP_K -b HRX0: 106 us (n=1000), 230 us (n=65,000),
    623 us (n=200,000) for one row, so the QSA decode shape is ~1.5-2 ms per step
    across 12 layers against the ~49 ms the CPU path costs at 18.8k.

Authored by the pi-attn lane (engine ggml-org#679); the orchestrator opened the PR.

🤖 Generated with Claude Code

…ml-org#679)

GGML_OP_TOP_K had no HRX dispatch, so the node ran on the CPU backend and split
every full-attention layer's graph around it. Qwen3.8-Flash-Next is the model that
pays for it: its QSA layers (src/models/qwen4exp.cpp) call ggml_top_k with
n = n_kv (up to 262144) and k = indexer_top_k + r - 1 = 2051, once per
full-attention layer per decode step, and the CPU side is a single-row
std::partial_sort with an indirect comparator plus a host round trip.

Measured on strixhalo (gfx1151), same server, same prompt family, one process:
- decode steps 39.2 -> 49.0 ms the moment n_kv crossed 2051 (1986 vs 2096 tokens),
  i.e. +9.7 ms per step for entering the QSA path at all;
- then ~3.8 us per KV token per step (49.0 ms at 2096 -> 67.4 ms at 6978), which
  extrapolates to 112 ms/step at 18,782 tokens against the 103 ms a full protocol
  run measured there (9.70 tok/s vs 31.15 for a build whose TOP_K is on the GPU);
- forcing the same path below the threshold (LLAMA_QSA_DENSE=0) at fixed n_kv=1480
  costs 13x prefill (676 -> 52 tok/s), which is the selection path itself.

The kernel is a radix select on a monotone uint32 key map of the f32 bit pattern:
three levels of 11/11/10 bits narrow to the exact rank-k key, then one emit pass
writes the indices above it and enough equal ones to reach k. One workgroup per
row, an 8 KiB LDS histogram, three passes over the row, no global scratch and no
host round trip. Order is not part of ggml's TOP_K contract (test_top_k compares
sets and the CPU reference even swaps two entries to say so).

The matcher is a SingleOp registration that declares the op with
hip_declare_eager_op, so the node reaches the dispatcher at all; because the op is
now claimed, ggml-hrx.cpp's device_supports_op bounds the claim per node through
hip_top_k_tensor_supported, mirroring the matcher's own predicate and the
GGML_HRX_DISABLE_DISPATCH switch, so a shape the kernel cannot take -- or a build
with the matcher turned off -- keeps its CPU fallback instead of failing the graph.

Verified on gfx1151:
- test-backend-ops -o TOP_K -b HRX0: 484/484 passed with
  "ggml_hrx dispatch: hip.top_k_f32 (single)" in the dispatch log;
- the same command with GGML_HRX_DISABLE_DISPATCH=hip.top_k_f32 reports TOP_K as
  skipped on HRX0 (clean fallback, no unsupported-node failure);
- the 11/11/10 radix select was replicated exactly and compared against the true
  top-k over 600 randomised vectors plus edge cases (all-equal, sorted, k=1, k=n,
  +-inf/NaN, and the QSA shape n=18816/k=2051): index sets and thresholds match.
- test-backend-ops perf -o TOP_K -b HRX0: 106 us (n=1000), 230 us (n=65,000),
  623 us (n=200,000) for one row, so the QSA decode shape is ~1.5-2 ms per step
  across 12 layers against the ~49 ms the CPU path costs at 18.8k.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: READY — tprs-ready applied. Everything about the change checks out and the device evidence is on the code this PR actually ships. One condition, stated at the end: my own device re-run is queued behind a GPU hold that has not moved in half an hour, and if it disagrees the label comes off.

The code, read end to end. 8 files, +332, additive: hip/kernels/hip_top_k_f32.hip (the kernel), hip/dispatch-hip-top-k.cpp (matcher + registration + the shared predicate), one line each in hip-dispatches.{h,cpp} and ggml-hrx-hip.cmake, hip-capabilities.h (+5, the declaration), and 5 lines in ggml-hrx.cpp's device_supports_op.

  • The claim cannot promise what the matcher will not take. hip_top_k_tensor_supported() is the same predicate as the matcher: f32 contiguous source, i32 contiguous dst, 0 < k <= n, n <= 262144 (shared kTopKMaxColumnCount), leading dims preserved, and it returns false when GGML_HRX_DISABLE_DISPATCH names this registration — so the documented fallback is the code path, not a hope. row_count <= UINT32_MAX matches the uint32_t grid dimension.
  • The algorithm is right as written. Monotone key map (bits ^ (0x80000000 | (int32)bits >> 31), the standard total order), three levels of 11/11/10 bits with shift = 32 - prefix_len - bits and a descending histogram scan, then an emit pass that writes everything above the threshold key and enough of the equal group. The slot = atomicAdd(&emitted, 1) + slot < top_k guard is what guarantees no duplicate indices, which upstream explicitly penalises. Barriers are where they need to be (histogram clear, after the scan, after emitted reset before the emit loops), and the shared state read at the end of a level is not overwritten before use.
  • Unordered output is within contract, and I checked that rather than taking the comment's word: tests/test-backend-ops.cpp's test_top_k::err compares index sets ("we expect the exact same set of indices, but possibly in a different order", via jdst), and in ties mode it compares the source values the indices point at and adds +1 if either side contains adjacent duplicates.
  • Nits, neither worth a round trip: the header says "NaNs above +inf", which is true for positive NaN only (a NaN with the sign bit set sorts below -inf under this map — irrelevant for scores, but the sentence is not exactly right); and the predicate is restated in hip_top_k_tensor_supported rather than shared with top_k_shape_supported, which is a drift risk the comment already acknowledges (the one constant that would hurt, 262144, is shared).

Provenance of the device evidence — this is the part I care about for a kernel PR, because "484/484 passed" is only worth something if it was measured on this commit:

files at c51b582 (what pi-attn built) == files at 9206f643 (this head):  SAME  (blob shas, 4/4 checked)
~/wt/lb-next/engine/third_party/llama.cpp is at 9206f643, and  ninja test-backend-ops  ->  "no work to do"
hip-code-objects.cpp.o contains hip_top_k_f32, libggml-hrx.so has the string 7x

So the kernel is embedded, the registration path is live, and the binary that produced 484/484 (plus the clean GGML_HRX_DISABLE_DISPATCH=hip.top_k_f32 fallback) is the one this PR ships.

What I have not verified myself: the 484/484 and the fallback on the device, and the 106/230/623 us perf rows. I queued the same two test-backend-ops -o TOP_K -b HRX0 runs under gpu-hold -l "pi prs"; the GPU token has been held by attn since 10:27 and my ticket is behind four others, so the run has not started. The evidence above is why I am comfortable labelling without it — the alternative was holding a merge gate on a queue that is not moving.

Two notes for whoever merges:

  1. Merging to 1bit/hrx-vulkan-patched moves the branch the engine pins (4a58da9b087d today, engine 3cddcd6). The engine needs a follow-up pin bump to pick the kernel up, and that bump's hw_gate is what will exercise it end to end. Per the pins rule, a llama.cpp move needs no config/hrx-long-prompt-validated-pin.config update — only an hrx-system move does.
  2. The measurement that justifies the kernel is the QSA path, and it is not in this PR (no model-level A/B: the 76 GiB load could not be hosted). Landing the kernel does not by itself prove the 9.70 -> 31.15 tok/s long-context gap closes; that still needs the lb/aval-arm.sh hrxls arm and the KLD/needle pair.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Process note on the condition in my verdict, so nobody reads the label as "device-verified by the reviewer": the orchestrator has withdrawn my gpu-hold queue entry until RELEASE-DONE appears in ~/NPU-WINDOW-OK, so I am not running the two test-backend-ops -o TOP_K -b HRX0 runs for now (they were queued behind attn's hold, which had had the token since 10:27 and was past the default 14-minute max). My ticket is withdrawn, nothing of mine holds or touches the GPU.

The label therefore stands on what I verified without a device, which is unchanged from the review above: the claim check mirrors the matcher (including the GGML_HRX_DISABLE_DISPATCH fallback), the kernel's radix select is correct as written, "unordered output" is inside ggml's contract per test-backend-ops.cpp::test_top_k::err, and the files at this head are blob-identical to the ones pi-attn measured 484/484 on — with the HRX tree checked out at this head (ninja test-backend-ops = "no work to do") and hip_top_k_f32 embedded in the built code object.

When the window opens I will run exactly this and report the counts here:

gpu-hold -l "pi prs" -t 25 -- bash -c '
  LD_LIBRARY_PATH=$PWD/bin ./bin/test-backend-ops -o TOP_K -b HRX0
  GGML_HRX_DISABLE_DISPATCH=hip.top_k_f32 LD_LIBRARY_PATH=$PWD/bin ./bin/test-backend-ops -o TOP_K -b HRX0'

If either disagrees with the body I pull tprs-ready and say so.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: READY

Gate line for the protocol, at head 9206f643 (current). MERGEABLE/CLEAN, tprs-ready applied. (This is a fork PR, so there is no engine CI gate on it; the engine-side hw_gate will exercise it when the pin is next bumped.)

Resting on non-device evidence, because the orchestrator has parked my GPU queue entry until RELEASE-DONE appears in ~/NPU-WINDOW-OK: the claim check hip_top_k_tensor_supported() mirrors the matcher's predicate including the GGML_HRX_DISABLE_DISPATCH fallback, so the declared op cannot promise a node no matcher takes; the kernel is a correct monotone-key radix select (11/11/10 levels, descending histogram scan, emit-above-threshold then fill from the equal group, atomicAdd slot allocation so no duplicate indices, barriers where they are needed); and "order is not part of the contract" is checked against tests/test-backend-ops.cpp::test_top_k::err, which compares index sets and in ties mode compares the source values plus penalises adjacent duplicates. The device numbers (484/484, clean fallback, 106/230/623 us) are pi-attn's, measured on files blob-identical to this head, in an HRX tree checked out at this head with ninja test-backend-ops = "no work to do" and hip_top_k_f32 embedded in the built code object.

When the window opens I will run test-backend-ops -o TOP_K -b HRX0 and its GGML_HRX_DISABLE_DISPATCH pair and post the counts here; if either disagrees I pull the label.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

tprs-ready TPRS: review verdict READY posted by the standing reviewer; merge gate satisfied

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant