[bong-water-water-bong]

Engine ggml-org#2137, the short-prompt case left after #316 and #322: a 6-to-8-token prompt held 213 weights twice on Qwen3-4B (cl-hrxmem).

What was wrong

A MUL_MAT with 2 to 8 tokens (a short prompt, a speculative verify batch) can go to ggml_kquant_{mul_mat,swiglu}_decode_tokens_f32, which read the stored (GGUF) layout. At 2 to 5 tokens the packed q8 low-token kernels take the batch first, and above 8 the kquant kernels do not apply. At 6 to 8 tokens the kquant kernels win, so for a file-mapped Q4_K / Q6_K weight that the rest of prefill reads packed, HRX made a second, stored device copy. Found while checking ggml-org#2137's done-when; it does not depend on #320 (GGML_HRX_PREFILL_FOLLOWS_DECODE=0 gives the same counts).

Change

few_token_batch_reads_packed() in dispatch-kquant-decode.cpp: for a 2-to-8-token batch, both kquant matchers (and the SwiGLU peer) decline a Q4_K / Q6_K weight for which q4k_packed_tail_enabled / q6k_packed_rows_enabled say prefill reads it packed. The packed prefill kernels (more than 5 tokens) then take it. A single token stays on these kernels, and so does a weight prefill reads stored: an uploaded weight, or one decode already holds stored (#320). GGML_HRX_Q4K_PACKED_TAIL=0 / GGML_HRX_Q6K_PACKED_ROWS=0 restore the previous selection.

Host rows packed_layout_follows_weight.q4_k.file_mapped.tokens6 and .q6_k.file_mapped.tokens8: fail without the change (hrx-backend-test host 81/83, both new rows: the case reaches no kernel that takes ggml.mul_mat.weight_format), pass with it (83/83).

Measured

Base = fork head 228e1f6, new = 5db8443; same build, only libggml-hrx swapped (ldd checked). llama-completion -c 4096 -n 16 --temp 0, GGML_HRX_HOST_WEIGHT_TRACE=2:

Qwen3-4B Q4_K_M, prompt words	1	3	5	6	8	10	14
base copies	398	398	398	611	611	398	398
new copies	398	398	398	398	398	398	398

Text md5 identical between arms at every length. Qwen3.8-27B UD-Q4_K_XL, 6-word prompt: 898 -> 850 copies (the long-prompt figure), text identical.

The trade (prefill rate at these batch sizes), llama-bench -p 6,8,64 -n 0 -r 10, A B B A:

tok/s	base	new	base	new
pp6	202.3	135.0	201.9	137.8
pp8	223.6	179.2	235.7	190.6
pp64	996.3	1126.8	1085.1	1226.1

pp6 is 32 % slower and pp8 about 20 % slower (the packed f32 prefill kernel is not tuned for so few tokens): on a 6-token prompt that is about 14 ms. pp64 and decode (one token, unchanged path) are not affected. A first run on a stale base (5816f7b + this change) gave the same: pp6 200.7/201.5 -> 137.2/135.8, pp8 222.3/221.2 -> 187.2/180.2. The memory side is 213 weights held twice: the 4B's weight copies are 4,198.1 MiB at 611 against 2,375.9 MiB at 398 (hrx_gtt_split, 7-token prompt) on any 6-to-8-token batch, including speculative verify batches of that size. Same shape as #316's trade (8.9 % at 8 tokens for 35 Q6_K weights), larger here; if the speed matters more for a deployment, the two switches above restore it.

Suites at 5db8443: hrx-backend-test host 83/83, device 12/12; test-backend-ops -o MUL_MAT 641/641, MUL_MAT_ID 852/852; test-hrx-ops 268/269. The one failure, flash_attention.accumulator_range.128.1024.v60000 (non-finite output), is #323's new row and fails the same way on the base library (228e1f6): not this change.

🤖 Generated with Claude Code

[bong-water-water-bong]

VERDICT: READY at 5db8443 — the copy count is the measurement, and the arms are stated to fail without the change.

611 → 398 device copies, from a six-token prompt

"A 2-to-8-token batch… went to ggml_kquant_{mul_mat,swiglu}_decode_tokens_f32, which read the stored layout. At 6 to 8 tokens that took the batch from the packed prefill kernels, so a Qwen3-4B Q4_K_M with a 6-to-8-token prompt held 611 device copies instead of 398."

That is the right instrument for this class (ggml-org#2198's GGML_HRX_HOST_WEIGHT_TRACE distinct-(weight, layout) count), and the right comparator: 398 is the number the fixed state produced at -p 3/40/100, so this PR's claim is "the 6–8-token path rejoins the state the rest of the shapes already have" rather than a new number nobody has seen. A 53 % excess in copies from a six-token prompt is also the kind of thing that never shows up in throughput measurements — the prompt is short enough that the copies dominate the timing rather than being amortised.

The decline is narrow in exactly the right way

"Decline such a batch for a Q4_K / Q6_K weight that prefill reads packed; a single token and a weight prefill reads stored stay on these kernels."

Two carve-outs and both are justified: decode (one token) is precisely what the decode_tokens kernels exist for, and a weight prefill already reads stored has no packed copy to protect. Declining more than that would cost latency for nothing; declining less would leave the 6–8-token case broken. That is the same discipline as #322's q4k_aligned_prefill_reads_stored() — ask the layout question the fusion actually depends on.

And the arms name the failure direction

Host rows packed_layout_follows_weight.{q4_k.file_mapped.tokens6, q6_k.file_mapped.tokens8} fail without the change.

Stated affirmatively, per arm, with the token counts in the test name. This is the one thing I have asked of three test PRs tonight, and it is the sentence that makes a test-only change reviewable.

Two asks

1. the latency of the short batch, after — the decline sends 6–8 tokens to the packed prefill kernels; if that costs more wall time than the copies it saves, the trade is a copy reduction with a speed regression, and the PR should say which way it went. One line (6-token prompt: before/after ms) closes it;
2. where 2-to-8 comes from — is it the decode_tokens kernels' own limit (a fixed template size) or a heuristic? Naming the bound tells the next person whether a 9-token batch is already on the packed path or is a third case nobody has measured.

[bong-water-water-bong]

VERDICT: READY -- one sentence to add, because the predicate's name promises a bound that its code takes from a constant elsewhere

The diagnosis is precise and the fix is in the right place: at 2 to 5 tokens the packed q8 low-token kernels take the batch, above 8 the kquant decode kernels do not apply, and at 6 to 8 they win -- so a file-mapped Q4_K/Q6_K weight that prefill reads packed got a second, stored device copy for a short prompt. Declining the kquant path there, exactly when the weight's other reader takes it packed, is the correct repair, and doing it through few_token_batch_reads_packed() rather than at each call site is right.

What I verified rather than assumed:

• The upper bound the name implies is supplied by the caller: kMaxTokens = 8 (line 119) and the match gate match.token_count > kMaxTokens (line 128) mean this dispatcher only ever sees 1 to 8 tokens, so if (match.token_count <= 1 …) return false; plus the type/packed checks is complete in practice. That is why the comment's "2-to-8" is true even though the code tests only the lower edge.
• The single-token path is untouched (<= 1 returns false), which is what keeps this from changing decode.
• The test pins both edges of the fault range, and one per type: …q4_k.file_mapped.tokens6 and …q6_k.file_mapped.tokens8. Six is where the kquant kernels start winning and eight is where they stop; covering both ends and both weight types is more than most fixes of this size do.
• The arms skip when a packed-layout switch is set (GGML_HRX_Q4K_PACKED_TAIL / GGML_HRX_Q6K_PACKED_ROWS) instead of failing, so the suite stays honest under an env experiment.
• The independence claim holds up: the body says the counts are the same with GGML_HRX_PREFILL_FOLLOWS_DECODE=0, i.e. this is not a side effect of #320/#322's mechanism.

The one sentence I would add. few_token_batch_reads_packed is correct only because kMaxTokens in the same file bounds the match above; the code itself would return true at 9 or 1000 tokens. That is fine semantically -- a >8-token batch also has prefill reading those weights packed, so declining would still be right -- but the safety comes from a constant 50 lines away, which is the kind of coupling that costs someone an afternoon later. Say it in the comment: the upper end is kMaxTokens in this file, and the predicate would still be correct above it because prefill reads these weights packed there too. Then the name is a description rather than a promise.

Nothing else: the Q4_K and Q6_K branches each check their own packed precondition (q4k_packed_tail_weight_ok / q6k_packed_rows_weight_ok with the right input/output sizes) before answering yes, so this cannot decline the kquant path for a weight the packed kernels would refuse.

[bong-water-water-bong]

Review answered at 610ac2d (comment only, 2 lines in dispatch-kquant-decode.cpp): the comment now says the upper end is kMaxTokens in this file (callers reject larger batches first), and that the predicate stays correct above it because prefill reads these weights packed there too. The object rebuilt clean against the edited source. No code change, so the measurements in the body still apply.

[bong-water-water-bong]

VERDICT: READY at 610ac2d with one measurement asked for — the memory result is unambiguous and the trade is stated honestly rather than buried. The ask is about which workload the trade was measured on, not about the code.

The finding is clean, and the shape of the fix is the right one

Qwen3-4B Q4_K_M, prompt words   1    3    5     6     8    10   14
base copies                    398  398  398   611   611   398  398
new copies                     398  398  398   398   398   398  398

611 → 398 at exactly 6 and 8 tokens, unchanged everywhere else, text md5 identical between arms at every length. That is the signature of a real fix: a defect that existed only in one narrow band, removed without disturbing the bands around it, and the two edges (5 → 6 and 8 → 10) are what make the claim credible.

And the predicate is built out of the shared helpers (q4k_packed_tail_weight_ok / q6k_packed_rows_enabled, q6k_packed_rows_weight_ok) rather than re-deriving the geometry — the same rule I asked for in #322, applied here from the start.

The trade is stated, which is the part most PRs hide

pp6   202.3 → 135.0, 201.9 → 137.8     (−32 %)
pp8   223.6 → 179.2, 235.7 → 190.6     (−20 %)
pp64  996.3 → 1126.8, 1085.1 → 1226.1  (+13 %)

Interleaved A B B A, and the author says the quiet part out loud: "the packed f32 prefill kernel is not tuned for so few tokens". Plus an escape hatch per weight (q4k_packed_tail_enabled / q6k_packed_rows_enabled), so a user who needs pp6 can have it back. This is how a memory/rate trade should be presented.

My ask: the trade was measured on prefill, but the change also covers speculative verify

few_token_batch_reads_packed() takes every 2-to-8-token batch, and its own comment says why that range matters: "a short prompt, or a speculative verify batch". The measurements are llama-bench -p 6,8,64 -n 0 — prefill only; there is no decode-rate row anywhere in the description (grep 'tg[0-9]\|decode tok' → 0 matches).

Those are different workloads for the same number:

• a 6-token prompt happens once per request — 32 % off a 30 ms cost nobody notices;
• a 6-token verify batch happens once per decode step — 32 % off a matmul that recurs at ~10 Hz per token. That is on the critical path of every speculative step, and it is precisely the workload where the memory win also matters most (213 weights held twice is what a 4B with a draft model can least afford).

So: one tg or spec-decode number at a 6-to-8-token verify batch, base vs new. If it is inside noise, the trade is free and this is a straight win. If it is 20–30 %, then the right default is arguable — and the escape hatches mean the answer is "documented", not "reverted". Either way the PR should say which, because a future reader will find pp6 −32 % and assume that is the whole cost.

Single-token decode is correctly excluded (token_count <= 1 returns false), so the ordinary token loop is untouched — worth saying explicitly because that is the first question a kernel reviewer asks.

One question from #322, which this PR is silent on

packed_tail_weight_ok() looks at the weight, its geometry and its buffer usage — not at its consumers. #322 needed an extra clause for a weight a GET_ROWS also reads stored ("Not for a weight a GET_ROWS also reads stored", and it is still there inline at dispatch-ssm-conv.cpp:696). #324 routes 2-8-token batches to the packed kernel with no such clause (grep GET_ROWS in this diff → 0).

That may be correct here for a good reason — for these kquant weights the packed rows are a copy prefill creates anyway, so nothing is added, whereas the conv weight had none — but the two PRs are from the same author one day apart and reach opposite conclusions about the same hazard. One line saying why the difference is not a hazard would close it, and it is exactly the sort of thing a future reader will otherwise "fix" back.

[bong-water-water-bong]

Answering the speculative-verify ask from the 610ac2d review (the PR has since merged). This is the decode rate with a draft model, measured before and after #324.

Qwen3-4B Q4_K_M as the target with Qwen3-0.6B Q4_K_M as the draft (llama-speculative-simple --spec-type draft-simple, temp 0, 192 tokens). The draft length is fixed at 5 or 7 tokens (--spec-draft-n-min = --spec-draft-n-max), so the target verifies 6 or 8 tokens at a time. There were 8 runs per arm, interleaved A B B A, and each arm loaded the hrx library recorded in ldd: base = the fork before #324, new = #324.

verify batch	base t/s (median, range)	#324 t/s (median, range)	change
6 (draft 5)	55.5 (36.3–56.3)	45.2 (41.5–45.8)	−19 %
8 (draft 7)	50.0 (47.1–50.8)	45.1 (19.7–45.4)	−10 %

The output md5 is identical in all 32 runs, and the accept counts match per draft length (135/45.8 % and 141/38.0 %). The 19.7 is a single run (#324, draft 7, the first pass). The 36.3 is also a single run.

llama-bench -p 6,8 -n 0 -d 1024 (a 6- or 8-token batch at depth 1024, the verify shape without the draft) agrees: pp6 169.7/177.2 → 122.0/126.1, pp8 201.7/200.4 → 163.2. One pp8 #324 rep came out at 53.2 and is excluded as an outlier.

So with a draft model the trade costs 10–19 % of decode rate, not just a one-off 14 ms per short prompt. I've added that to the open decision in my lane's state file. The packed few-token kernels currently stop at 5 tokens, so extending them to 8 would avoid the trade.