Repository navigation
hrx: score each key once above 2048 keys, instead of on every value lane (engine #991) - #214
Conversation
…ane (engine ggml-org#991) Zamba2-2.7B's prefill collapsed to 0.065 of its pp512 throughput at -p 4096, and the cause was a kernel selection cliff rather than the model or the graph: `common.attention_strided_f32_f16` picks its kernel by context length, and past 2048 keys the rows kernel can no longer keep its score row in workgroup memory, so every FLASH_ATTN_EXT node fell to the strided kernel, whose workgroup has one lane per value dimension and each of those lanes recomputes the whole q.k dot product for every key - nkv*dv*d instead of nkv*(d + dv), which is 80x the arithmetic at d = dv = 160 and worse at 224. Measured on one binary, Zamba2-2.7B BF16, `-p 512,2048,2560,4096`: default matcher pp512 395.38 pp2048 270.59 (0.684) pp2560 72.19 (0.183) pp4096 25.69 (0.065) ONEBIT_HRX_ATTN_PER_LANE=1 (strided) pp512 87.48 The second row is the point: at 512 tokens, with the same shapes and the same arithmetic, the strided kernel costs 78% of the throughput, so the kernel and not the context length is the cost. The cliff is exactly at the 2048-key switch (0.684 -> 0.183 between neighbouring lengths). This adds `ggml_attention_rows_blocked_f32_f16`: the rows kernel walking the keys in blocks of 2048 with the flash-attention online softmax - running max and sum carried across blocks, a dv-wide output row in workgroup memory rescaled by exp(old_max - new_max) per block, 1 KB query row + 8 KB score row + 1 KB output row. Keys are padded to whole blocks by clamping the key index and forcing the padded score to the lowest f32, so their probability is exactly zero in both the sum and the value pass and no load ever leaves the buffer. `ggml_attention_rows_f32_f16` is untouched: 2048 keys or fewer keeps the old kernel, so nothing below the switch changes and the A/B above is attributable. Why clamping instead of `min(nkv - block_start, 2048)` as the loop bound: loom cannot invert a subtraction, so `block_start + slot < nkv` was not provable and the JIT refused the kernel outright with SUBRANGE/024 on the key, mask and value loads - a hard `failed to prepare command`, not a CPU fallback. `index.min %j0, %last_key` states the bound as a fact about a min result, which is what the compiler can use. Found with ggml-hrx-compile-kernel, which compiles a single specialization on the CPU from the text loom source, so the diagnosis took seconds instead of a GPU ticket; verified there at kv 2560/3000/4096/5000 before rebuilding. Cases: six new FLASH_ATTN_EXT cases at kv 2560 and 4096 (hsv 64, 160 and 256; 160 is Zamba2's own head size). The suite's longest case was exactly 2048 - the last length that still took the rows kernel - which is why no gate ever saw this. On gfx1151 all six pass against the CPU reference, the chosen specialization is deterministic (1024 workgroups x 64 lanes, 10256 B of workgroup memory), and the model then measures pp2560 72.19 -> 190.93 and pp4096 25.69 -> 131.28 with the <= 2048 path unchanged. The remaining decline with context is the honest O(nkv) attention cost, not a cliff.
|
#214 ready for review at Three things a reviewer should get without digging:
For anyone doing loom work: I hit |
bong-water-water-bong
left a comment
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at ffef39f18 on the parts that can be proved without a device — the six new cases exist and pass on the CPU reference, the dispatch three-way is right, the new kernel's online-softmax/padding scheme is sound and follows the file's existing conventions — and the HRX numbers stay the author's, for the reason at the bottom.
The new cases are real and green here. Built at this head (GGML_HRX=OFF, CPU), then:
./bin/test-backend-ops test -o FLASH_ATTN_EXT -b CPU -> 5149/5149 tests passed, Backend CPU: OK
FLASH_ATTN_EXT(hsk=64,hsv=64,nh=8,nr23=[4,1],kv=2560,nb=32,…,type_K=f16,type_V=f16): OK
FLASH_ATTN_EXT(hsk=160,hsv=160,…kv=2560…): OK (160 is Zamba2's own head size)
FLASH_ATTN_EXT(hsk=256,hsv=256,…kv=2560…): OK
FLASH_ATTN_EXT(hsk=160,hsv=160,…kv=4096…): OK (+ the kv=4096 hs=64/256 pair from the same loop)
so the shapes are accepted by the op, the reference agrees, and the case set is not decorative. The claim that motivated them also checks out at the base: in tests/test-backend-ops.cpp at base214 the FLASH_ATTN_EXT kv list is {1024, 2048} — 2048 is the largest, i.e. exactly the last length that still took the rows kernel, which is why no gate ever saw the cliff.
The dispatch cannot change anything at or below 2048 keys. Three-way, and the old branch is preserved verbatim for the range that had it:
per_lane || d > 256 || dv > 256 -> strided (unchanged, and still the escape hatch)
else nkv <= 2048 -> rows (untouched kernel, so the A/B is attributable)
else -> rows_blocked
The blocked kernel's own index.assume bounds (d, dv in [1,256]) are the same ones the shipped rows kernel already carries and the same convention every kernel in this file uses for its maximum (range(%token_count, 1, 2048) with 2048-token launches), so nothing new is being asserted to the compiler at the top of the range.
The kernel's correctness argument holds when read against the code: per 2048-key block it computes scores, takes the block max, folds it into the running max, rescales the running sum and the running output row by exp(old_max − new_max), accumulates p·v from workgroup memory, and divides by the final sum once at the end — the standard online-softmax recurrence, with o_row explicitly zero-initialised first (the comment says why: the first rescale multiplies a defined value). Padding is done by clamping the key index (index.min %j0, %last_key) and forcing the padded score to the lowest f32, so the padded slots' probability is exactly zero in both the sum and the value pass; the mask load uses the clamped index too, and its value cannot matter because the score is overwritten. Workgroup memory is 1 KB query + 8 KB scores + 1 KB output = 10 KB of the 64 KB budget, and the barriers are where they have to be: after the query/output init, and between the score pass and the value pass (the value pass reads all 2048 scores, not just its own lane's strided subset).
What I did not run: the HRX side — the blocked kernel on the device, the A/B table (395.38 → 87.48 t/s at 512 with ONEBIT_HRX_ATTN_PER_LANE=1, 270.59 → 72.19 at 2560, 25.69 at 4096) and the claim that the result is bit-comparable to the rows kernel at lengths that fit. Those need GGML_HRX=ON, which on this box means an ExternalProject build of hrx + loom + IREE, and then a GPU window: tonight's c31 gate hold (TISSUES-ALL rule 1) excludes both. I would rather hand the queue an honest partial than hold a kernel PR for an hour of build. The re-derivation for whoever has the window is one command pair: test-backend-ops -b HRX0 -o FLASH_ATTN_EXT on this head (the six cases above are the ones that exercise it), then the same with ONEBIT_HRX_ATTN_PER_LANE=1 to reproduce the cliff.
One nit, style only: the dispatch comment carries the measured A/B numbers inline ("pp2048/pp512 = 0.68, pp2560 = 0.18, pp4096 = 0.065"), which is good, but the numbers are from Zamba2-2.7B and the comment does not say so where they could be mistaken for this box's general law — one clause naming the model would keep the next reader from generalising them.
hrx: score each key once above 2048 keys, instead of on every value lane (engine ggml-org#991)
Part of engine ggml-org#991, and the follow-up to #213: that one put Zamba-7B-v1's decode attention fully on HRX; this one removes a
15× prefill cliff that every long-context model on this box was paying for.
The bug
common.attention_strided_f32_f16(dispatch-small-rows.cpp) picks its kernel by context length:The rows kernel scores each key once and keeps the whole 2048-entry score row in workgroup memory, which is exactly what limits it
to 2048 keys. Past that, every node fell to the strided kernel: workgroup size =
dv, one lane per value dimension, and every lanerecomputes the whole q·k dot product for every key, keeping only its own share of the result. That is
nkv·dv·dinstead ofnkv·(d + dv)— 80× the arithmetic at d = dv = 160, 112× at 224.Measured A/B on one binary, Zamba2-2.7B BF16,
-p 512,2048,2560,4096:ONEBIT_HRX_ATTN_PER_LANE=1(strided at every length)The second column is the point: at 512 tokens, with the same shapes and the same arithmetic, the strided kernel costs 78 % of the
throughput, so the kernel and not the context length is the cost. The cliff is exactly at the 2048-key switch — 0.684 → 0.183 between
two neighbouring lengths. This is the mechanism behind the collapse another lane measured across four Zyphra quants (0.049-0.073 of
pp512 at -p 4096);
Zamba-7B-v1was unaffected only because its 464-wide value head is refused by this registration entirely.The change (4 files, +354/−3)
ops/small_rows_f32.loom: new exportggml_attention_rows_blocked_f32_f16, the rows kernel walking the keys in blocks of 2048with the flash-attention online softmax — running max and sum carried across blocks, a
dv-wide output row in workgroup memoryrescaled by
exp(old_max − new_max)per block, then accumulated. Workgroup memory: 1 KB query row + 8 KB score row + 1 KB output row= 10 KB of the 64 KB budget (10256 B in the compiled hsaco).
index.min %j0, %last_key) and forcing the padded score to the lowestf32, so their probability is exactly zero in both the sum and the value pass, and no load can leave the buffer.
ggml_attention_rows_f32_f16is untouched: 2048 keys or fewer keeps the old kernel, so nothing below the switch changes and theA/B above is attributable to this change alone.
stridedstays fordordvabove 256, the head sizes only it admits.manifest.json: one export entry (mechanical copy of the rows entry; the catalogstatic_assertinGGML_HRX_KERNEL_REFresolving at build is the well-formedness check).
dispatch-small-rows.cpp: three-way selection, with the measured numbers in the comment so nobody has to re-derive whystridedexists at all.
tests/test-backend-ops.cpp: six newFLASH_ATTN_EXTcases atkv = 2560, 4096×hsv = 64, 160, 256(160 is Zamba2's ownhead size). The suite's longest existing case was exactly 2048 — the last length that still took the rows kernel — which is why no
gate ever saw this. These are the first cases above 2048 keys.
Why the padding/index clamp rather than
min(nkv − block_start, 2048)as the loop boundBecause loom cannot invert a subtraction, so
block_start + slot < nkvwas not provable and the JIT refused the kernel outright:— on the key, mask and value loads.
index.minbounds its result, not the algebra around it, so a clamp states the bound as a fact thecompiler can use. (This is also why the first version failed at
kv=2560but passed at 4096: at 4096 every block is full, so thebound was trivially constant.)
Worth knowing for reviewers of any loom change: a refused kernel is a hard
failed to prepare command, not a CPU fallback, and onesuch kernel id takes the whole model down (
failed to run prompt warmup). The diagnosis above came fromggml-hrx-compile-kernel,which compiles a single specialization on the CPU from the text loom source — no GPU ticket, seconds per iteration instead of ~40
minutes queued. I verified the fixed kernel there at kv 2560/3000/4096/5000 with zero diagnostics before rebuilding.
Device evidence (gfx1151,
pi-zarch-rb-blocked2, gpu-hold, 166 s on the token)OKagainst the CPU reference (tolerance 5e-4 nMSE), includinghsv=160and bothkvvalues;Loom compilation failedanywhere in the run (now a hard assertion in the harness, because that failure is silent otherwise);group_segment_fixed_size10256 B;-p 512,2048,2560,4096 -r 2:pp4096/pp512 goes 0.065 → 0.355. The ≤2048 path is byte-identical code, so the pp512/pp2048 differences are measurement noise on a
box other lanes were using, not a claim of improvement — the prefill rows above 2048 are the claim. What remains of the decline with
context (0.746 → 0.517 → 0.355) is the honest linear attention cost, not a cliff.
Not in this PR
A tail-block optimisation: a block that is 25 % full still scans all 2048 slots (that padding is what makes the bounds provable, and its
cost is at most one padded block per (query, head)). The prefill V-tail for the other attention kernel (
flash_attention_f32_f16_wmma,which publishes whole 64-wide tiles) is separate kernel work; and that file's value predicate caps at
kPrefillMaxHeadSize = 576while both ops declarerange(value_head_size, 64, 512), so V = 576 is claimed and then JIT-rejected —unreachable on this box (no model has V = 576) and left for the prefill work rather than as review churn here.