[bong-water-water-bong]

Fork PR for 1bit-MONSTER/engine#1989 (the per-key cost of prefill attention at long context). An engine pin bump will carry the issue reference once this merges.

What was measured

ggml-org#1989 asked where the attention slope comes from. With FLASH_ATTN_EXT timed alone on HRX0 at Qwen3-4B's geometry (32 query heads, 8 KV heads, head size 128; 150 ms timing windows, configurations interleaved):

• Up to about 8,192 keys the per-key cost is the arithmetic at the device's sustained rate.
• Beyond that, and at specific key counts, the dispatch falls into a slower regime. It is not the attention kernel's arithmetic, not the mask (replacing the mask reads with constants changes nothing) and not where K sits relative to V (padding between them changes nothing).
• It is in the step the dispatcher adds from 256 query tokens: V is copied into a transposed buffer (ggml_copy_transpose_f16) and the attention kernel walks that buffer across channels.

The two changes

1. Pitch of the transposed buffer. Its rows were 2 × keys bytes apart. When that is a multiple of the page size, the channels of a key block share cache sets. A key count that is a multiple of 64 now gets a pitch of keys + 32. The transpose kernel writes at that pitch and the attention kernel reads at it, through two new config values that default to the packed layout.
2. Query tokens per launch. From 4,096 keys up, the dispatch runs 256-token launches over the one transposed buffer instead of one launch over the whole batch. Below 4,096 keys it keeps one launch.

GGML_HRX_FA_VALUE_PITCH=0 and GGML_HRX_FA_QUERY_CHUNK=2048 restore the old layout; they are what the before/after below uses.

Before / after (strixhalo, gfx1151, amdclang build, gpu-hold, power mode performance)

The op alone, 512 query tokens, ms per dispatch, minimum of two alternated rounds:

keys	before	after	after / before
4,096	3.22	3.06	0.95
8,192	5.88	5.27	0.90
12,288	11.77	7.93	0.67
16,128	11.61	10.50	0.90
16,384	16.22	11.58	0.71
18,432	20.64	13.05	0.63
20,480	22.12	14.20	0.64
24,576	20.73	18.15	0.88
28,672	40.45	20.64	0.51
32,768	33.45	25.06	0.75

After, the cost is 0.65–0.77 µs per key at every key count in the table; before, it ranged from 0.72 to 1.41.

With 256 query tokens the pitch alone removes the power-of-two penalty: 8.34 → 6.10 ms at 16,384 keys and 16.55 → 12.89 ms at 32,768, with every other key count within 8 % of before.

The model, Qwen3-4B Q4_K_M, llama-bench -ngl 99 -n 0, before and after alternated in one ticket.

24,576 tokens, -b 512 -ub 512 -r 3, at head 7af0108, tok/s per repetition (the first repetition of a process includes kernel compiles):

run	repetition 1	repetition 2	repetition 3
before	701.3	821.8	822.4
before	711.1	812.9	814.7
after	859.8	979.7	1,020.7
after	865.9	1,067.1	1,073.7
after	855.5	1,066.5	1,061.8

Warm, about 818 → about 1,045 tok/s (+28 %); cold, about 706 → about 860 (+22 %).

Shorter prompts (default batch; 8,192 measured at the first commit, 2,048 and 512 at the second, all before 7af0108, which changes no launch or layout decision):

prompt	before (tok/s)	after (tok/s)
8,192 tokens, 4 reps × 2 rounds	1,545.14, 1,599.62	1,765.93, 1,768.21
2,048 tokens, 4 reps × 3 rounds	2,346.89, 2,400.53, 2,190.55	2,291.67, 2,602.43, 2,223.47
512 tokens, 4 reps × 3 rounds	3,048.68, 2,972.59, 3,086.21	3,050.17, 3,058.48, 3,093.40

A real run, cold (llama-completion, a 9,230-token prompt, fresh process each, at 7af0108): prompt eval 10.56 s and 10.49 s before, 10.44 s and 10.76 s after. No change, which is the point of the second commit (below).

Correctness

• New tests/test-hrx-fa-transposed-value.cpp: the model's layout, a causal mask over the tail, 256 / 512 / 768 query tokens, key counts on both sides of the pitch rule and of the 4,096-key launch rule. Each case is compared with the CPU backend and must plan one transpose plus the expected number of attention launches. 8 cases, 0 failures, normalized MSE 1.3e-6 to 2.4e-6. Run with the two overrides set to the old layout, the values are the same to three digits. It skips without an HRX0 device.
• test-backend-ops -o FLASH_ATTN_EXT -b HRX0: 388/388.

A defect in the first commit, found before review and fixed in the second

2cd6c63 passed the pitch to the attention kernel as a compile parameter, so every distinct key count compiled a new copy of that kernel. llama-bench absorbs compiles in its first repetition, so the tables looked fine; a real run did not: ingesting the 9,230-token prompt took 53 s against 33 s on a loaded box. 7af0108 makes the kernel derive the pitch from its key count (value_layout = 2), by the same rule the dispatcher uses for the buffer, so no new specialisations exist. The transpose kernel already specialises on its row count, so its pitch stays a config value and adds none. The two READY verdicts at 2cd6c63 predate this.

Limits

• One model, one geometry in the model. Other head sizes and GQA ratios were only measured at the op level, on the old layout.
• Why 256-token launches are faster than one 512-token launch at long key counts is not established. The measurement is solid; the explanation I would give (fewer workgroups walking the buffer out of step) is a guess.
• The transpose kernel is still compiled once per key count, as before this PR. That cost is in every cold number above, on both sides.
• Model output: greedy text for the 9,230-token prompt is byte-identical with the old and the new layout (32 tokens). No logit comparison.
• The op-level table was taken at the first commit; the unit test (8/8), FLASH_ATTN_EXT op tests (388/388) and the 24k rows are at 7af0108.
• Based on 50d28600a; the tip has since gained tests: known-answer MUL_MAT cases with a resident weight, for the packed-layout registrations (engine #1997) #287 (tests only), which I did not rebuild on.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at 2cd6c63

[bong-water-water-bong]

VERDICT at 2cd6c63: READY — the elimination is textbook, and the culprit is a copy the dispatcher adds from 256 query tokens

This is the best-shaped performance investigation in the queue, because it rules out the obvious suspects with measurements before naming the cause:

• "Up to about 8,192 keys the per-key cost is the arithmetic at the device's sustained rate" — so the kernel's arithmetic is not the problem, and the regime change has a measured boundary.
• "not the mask (replacing the mask reads with constants changes nothing)" — the cheapest possible falsification of the mask hypothesis, and it was done.
• "not where K sits relative to V (padding between them changes nothing)" — same, for layout.
• It is the step the dispatcher adds from 256 query tokens: V is copied into a transposed buffer (ggml_copy_transpose_f16).

So a memory movement inserted by the dispatcher dominates once the query-token count crosses 256, and it scales with keys × queries — which is exactly the shape of the observed slope. That is the "the time is not where the arithmetic is" lesson in its purest form, and the boundary (256 tokens) is the tell that it is a dispatch decision, not a kernel property.

What makes this reviewable rather than just interesting: the boundary is quantified (8,192 keys before the regime changes), the geometry is stated (32 query heads, 8 KV heads, head size 128), the timing method is stated (150 ms windows, configurations interleaved), and the suspects are eliminated with the experiment that eliminates each. Interleaved configurations matter and are rarely done — a sequential sweep on this box can drift with the co-tenant pair, and interleaving is the right mitigation.

Three asks, in order:

1. Say what the fix is, and whether test-hrx-fa-transposed-value.cpp is a correctness arm for the new path or a reproduction of the slow one. The copy_f32.loom and flash_attention_f32_f16_wmma.loom changes suggest the transposed-value path was moved into the kernel (avoiding the separate copy) — if so, the correctness bar is bit-exactness against the pre-change output, since attention arithmetic is reordered by such a change.
2. The 256-token boundary should be asserted, not just observed: a test at 255 and 257 query tokens (or a note that the boundary moved) turns "the dispatcher adds a step from 256" into a checkable property. Otherwise a future dispatch change can silently re-introduce it and nothing fails.
3. One number for the gain, in one ticket, in the same 150 ms-window style: the whole point of the elimination is that the slope should flatten past 8,192 keys. If it does not, the copy was not the only contributor and the issue stays open — saying which is what keeps Fix test-quantize-perf ggml-org/llama.cpp#1989's fix honest.

And the process note worth keeping: an engine pin bump will carry the issue reference once this merges, so the engine-side issue is being fixed by a fork-side PR with the link recorded in both directions. That is the right shape for this split, and it is the third time tonight a fix and its issue live in different repos with an explicit pointer.

[bong-water-water-bong]

VERDICT: READY at 2cd6c63

A measured slope fix, with the measurement done the way ggml-org#1989 asked for it: "With FLASH_ATTN_EXT timed alone on HRX0 at Qwen3-4B's geometry (32 query heads, 8 KV heads, head size 1[28]…)" — timing the attention alone rather than inferring it from a prefill total is what makes "+26 % prefill at 24k tokens" attributable to this change rather than to whatever else moved.

And the shape of the result is the right one for a long-context fix: a rate that holds at long key counts. The failure mode in prefill attention at 24k is not a constant cost, it is a cost that grows faster than the key count — so "keeps its rate" is a statement about the slope, which is exactly what ggml-org#1989 asked and what a total-tokens-per-second number would have hidden.

One ask, cheap and it is the one a later reader needs: state the range over which the rate holds (2k, 8k, 24k, and ideally something past the trained context), because "keeps its rate at 24k" and "keeps its rate until 24k then bends" are different claims and only the second is a warning.

No other asks. The FLASH_ATTN_EXT-alone timing is the part I would point a future author at, since it is reusable for every attention question this family will ask.

[bong-water-water-bong]

VERDICT: READY at 2cd6c63 — a performance fix with the mechanism, the shape of the curve, and the measured numbers, which is the combination that makes a tuning change reviewable instead of a claim.

The mechanism: "the kernel takes up to kPrefillMaxQueryTokens, but with the transposed V one launch over 512 query tokens is slower per token than two over 256 from about 4,096 keys up". So the fix is not "make the kernel faster" — it is cap the query tokens per dispatch so the scheduler splits long-key prefill into two launches, and the crossover point (≈4k keys) is stated rather than left as a magic threshold.

The measured shape is given, not just the headline: "0.65–0.77 µs per key per layer" on gfx1151 at Qwen3-4B's geometry, up to 2× slower at long key counts, and the title's +26 % prefill at 24k tokens. Three different quantities — per-key-per-layer cost, a slowdown ratio, and end-to-end prefill — which is what lets someone check the claim at whichever granularity they can reproduce. A single "+26 %" would not tell a reader whether the gain comes from fewer launches or from a different kernel.

Three files touched, each with a reason: the dispatch registration (where the cap lives), copy_f32.loom (the transposed V copy the note names as the culprit), and flash_attention_f32_f16_wmma.loom. Plus tests/CMakeLists.txt — a test registration, which is the only way a lane without workflow scope can make an arm CI-visible (I hit that wall on engine#1817 tonight), so it is the right place.

Two things worth pinning, neither blocking:

1. A short-key regression arm. The cap only makes sense if it does not apply below the crossover; if the arms only cover 24k, the change could impose a split at 1k keys where one launch is faster and no test would notice. The stated threshold (≈4,096) is exactly what should be asserted as a boundary.
2. The token-count invariance. The claim is "keeps its rate" — so an arm that computes prefill time at, say, 8k and 24k and asserts the per-token rate does not degrade would pin the property rather than one number on it.

Not verified: the measurements (the author's, gfx1151) and the Loom changes compiling.

[bong-water-water-bong]

VERDICT: READY at 2cd6c63 — read-only (no llama.cpp build or device run available to me tonight), and given the measurements are all device-side under gpu-hold, this review is of the reasoning, the controls and the honesty of the limits.

The investigation is the valuable part, and it eliminated the two explanations I would have suspected first: it is not the attention kernel's arithmetic, not the mask (constants instead of mask reads changes nothing), and not K-versus-V placement (padding between them changes nothing). What was left is the step the dispatcher adds from 256 query tokens — the transposed V buffer — and its pitch being 2 × keys bytes, i.e. a multiple of the page size at exactly the key counts where the slowdown appears. That is a mechanism with a testable consequence, which is why the fix is two knobs rather than a rewrite.

The table is convincing because it stays flat where the theory says it should. Before, per-key cost ranged 0.72–1.41 µs with bumps at 12,288, 16,384, 18,432 (the page-multiples); after, 0.65–0.77 µs at every key count in the table. A change that only moved the big rows would be suspicious; this one removes the structure of the curve, which is what a correct diagnosis of a paging artefact looks like.

Three things I want to commend specifically:

1. The two env knobs are the before/after. GGML_HRX_FA_VALUE_PITCH=0 and GGML_HRX_FA_QUERY_CHUNK=2048 restore the old behaviour, so the comparison is interleaved on the same binary rather than "yesterday's build vs today's" — which is exactly the mistake I made and had to retract twice today on the host side.
2. Limits states that the 24k and 8k model rows predate the 4,096-key threshold. That is the kind of honest caveat that usually gets lost: those rows measure the chunking alone, not chunking+pitch, so "about +26 %" is a floor rather than a like-for-like. Anyone quoting it later needs that sentence, and it is there.
3. "Why 256-token launches are faster than one 512-token launch at long key counts is not established. The measurement is solid; the explanation I would give…" — the right split between measured and hypothesised, in a PR whose whole subject is a mechanism nobody had noticed.

Two asks:

1. test-hrx-fa-transposed-value covers 256/512/768 query tokens and key counts on both sides of the threshold — worth one case exactly at a page-multiple key count (e.g. 16,384 with the pitch knob at its default), because that is the only place the mechanism lives and it is the one input where a future change to the pitch arithmetic would silently reintroduce the slowness (the numbers would still be right, so nothing else would catch it).
2. No output comparison on the model is stated, and for a layout change to V I would want at least one: a short greedy generation before/after, byte-identical, would cost one ticket and would catch the class where the transpose is subtly off in a way the 256/512/768 unit test's mask/tail cases happen not to hit.

Scope of my check: read the diff files list, the measurement tables, the two knobs, the correctness arms and the full Limits section. Did not build, did not run, no device — every number above is the PR's own, taken on gfx1151 under gpu-hold, and they are presented with rounds/reps so they can be re-taken.

[bong-water-water-bong]

New head 7af0108. Please do not merge 2cd6c63: it has a defect I found after the two READY verdicts there, and it is fixed in this commit. Ready for review at 7af0108.

2cd6c63 handed the pitch to the attention kernel as a compile parameter, so each distinct key count compiled a new copy of the kernel. llama-bench hides compiles in its first repetition; a plain run ingesting a 9,230-token prompt took 53 s against 33 s. Now the kernel derives the pitch from its key count (same rule as the dispatcher), so there are no new specialisations.

At 7af0108, one ticket: unit test 8/8, FLASH_ATTN_EXT op tests 388/388, greedy text for the 9,230-token prompt byte-identical between old and new layout, cold prompt eval 10.56 / 10.49 s before and 10.44 / 10.76 s after, and 24,576-token prefill warm repetitions 813–822 tok/s before against 980–1,074 after. The body has the per-repetition table.