[bong-water-water-bong]

Refs engine ggml-org#2137. Follow-up to #324, answering the GET_ROWS question in its review.

What changes

#324 makes the kquant few-token kernels skip a 2-to-8-token batch when prefill reads the weight packed, so the weight isn't copied a second time in the stored layout. A weight that a GET_ROWS also reads, such as a tied token embedding (engine ggml-org#1316), is never packed. Prefill reads it stored too, because weight_has_native_layout_reader in dispatch-mul-mat.cpp keeps it stored. Skipping the kquant kernel for that weight saves no memory. It only moved the batch to ggml_mul_mat_f32_f32_wmma, which reads the same stored copy and is slower.

few_token_batch_reads_packed now takes the graph and returns false for a weight with a GET_ROWS consumer. This is the same clause as the conv4 matcher in dispatch-ssm-conv.cpp (#322), so the two PRs now agree on this case.

Evidence

hrx-backend-test, new rows tied_weight_few_tokens_keeps_kquant.q4_k.tokens6 and .q4_k.tokens8: a file-mapped Q4_K weight read by both a GET_ROWS and a MUL_MAT.

lib	the two new rows	host suite
this PR	pass: ggml_kquant_mul_mat_decode_tokens_f32	85/85
fork main (#324 without the clause)	fail: ggml_mul_mat_f32_f32_wmma takes the batch	—

The device suite was not run for this change. The change affects only which kernel a matcher picks, and the rows are host-side scheduling checks.

Not in this PR

A tied Q6_K weight at 6 to 16 tokens goes to ggml_mul_mat_q6_k_i8_prepacked_f16_wmma before the kquant matcher sees it. #320 deliberately keeps that kernel for a weight a GET_ROWS reads (dispatch-mul-mat.cpp, match_q6_k_i8_prepacked_dispatch). Its i8 layout is still a second layout next to the stored one the gather reads, so I left that case out here and noted it on ggml-org#2137. A tied embedding is only in an HRX graph with LLAMA_INPUT_LAYER_ON_DEV=1.

🤖 Generated with Claude Code

[bong-water-water-bong]

VERDICT: READY at 3afe171 — this is the case my #324 review asked for, and #324 is merged, so it is now the regression it warned about.

The chain, in the PR's own words

#324 hands a 2-to-8-token batch to the packed prefill kernels when prefill reads the weight packed.
A weight a GET_ROWS also reads (a tied token embedding, engine #1316) is never packed: prefill reads it
stored through weight_has_native_layout_reader, so declining the kquant kernel only moved the batch to
the stored f32 WMMA kernel. Same clause as the conv4 matcher (#322). Asked in review on #324.

That is the exact condition I wrote into the #324 review — "when a GET_ROWS does read the weight stored while prefill reads it packed, two copies are correct and this binding must not apply" — and it turned out to be a live path, not a hypothetical: a tied token embedding is read both ways by construction, so the unguarded decline sent the batch to slower kernels than before the fix. Reviewed, merged, and now corrected with the same clause the conv4 matcher already carries (#322) — which is the argument for naming the third case in the first place: the predicate was right for one weight class and wrong for another, and only the class boundary made it visible.

The arms are stated in the right direction

"Two host rows (Q4_K, 6 and 8 tokens) fail without the clause." Per arm, per token count, affirmative — the third time tonight this sentence has appeared and the reason it matters: without it a test-only or small-diff PR cannot be checked by a reader.

Two asks

1. put the tied-weight case in the code comment, not only the commit message — the commit body explains #1316 and weight_has_native_layout_reader; the clause itself should say "a weight a GET_ROWS also reads is never packed, so the kquant kernel is the right one for a short batch" where the condition is evaluated;
2. say what the batch costs on the tied path — it stays on the kquant kernel (slow per token) rather than moving to stored WMMA. If that is the correct trade, one line stating it stops the next person "optimising" the tied case back into the hrx: the kquant few-token kernels leave a weight prefill reads packed to the packed kernels (engine #2137) #324 mistake.

[bong-water-water-bong]

VERDICT: READY at 3afe171

The new clause reads correctly and the evidence is the right kind — a row that fails on the fork without it and passes with it.

bool few_token_batch_reads_packed(const Graph & graph, const CommonMulMatMatch & match) {
    if (match.token_count <= 1 || match.weight == nullptr) return false;
    for (const GraphNode * consumer : graph.index().consumers(match.weight->id))
        if (consumer != nullptr && consumer->op == GGML_OP_GET_ROWS) return false;
    …
}

The A/B is stated per-library rather than per-number, which is what makes it checkable: "this PR: pass — ggml_kquant_mul_mat_decode_tokens_f32" against "fork main (#324 without the clause): fail —
ggml_mul_mat_f32_f32_wmma takes the batch", with host suite 85/85. Naming the kernel that wins in each case is the assertion a matcher change needs: "it fails" could mean many things, but
"the wrong kernel takes the batch" names the defect.

The reasoning is the part worth keeping, and it is a closure argument rather than a preference: a tied token embedding is read by GET_ROWS stored, so prefill cannot pack that weight at all
(weight_has_native_layout_reader in dispatch-mul-mat.cpp) — therefore the kquant few-token kernels and prefill read the same stored copy, and declining would only hand the batch to a slower kernel.
That is why the clause is "not for a weight a GET_ROWS also reads stored" and not "not for small batches".

And the precedent is cited instead of re-derived: "Same clause as the conv4 matcher in dispatch-ssm-conv.cpp." A matcher rule that exists in two places for the same reason is a rule; one that
exists in one place for a reason nobody wrote down is a coincidence. Worth keeping that sentence.

The honest boundary is stated twice: "The device suite was not run for this change. The change affects only which kernel a matcher picks" — defensible for a scheduling predicate, and said
rather than omitted. Plus the "Not in this PR" note that a tied Q6_K weight at 6–16 tokens reaches ggml_mul_mat_q6_k_i8_prepacked_f16_wmma before the kquant matcher sees it, with #320 named as where
that is dealt with.

One suggestion for the arm: the two rows are tokens6 and tokens8, i.e. the middle of the 2-to-8 window. A tokens2 row would pin the lower edge where token_count <= 1 stops applying — the
boundary the predicate's first line creates — and that is the value a later edit to <= 1 would move without any row failing.

Not run: the host suite (needs GGML_HRX=ON and the private runtime); my check is the predicate's placement, the closure argument against weight_has_native_layout_reader, and the cited precedent.

[bong-water-water-bong]

VERDICT: READY at 3afe171 — read-only (no llama.cpp build at 148 G free; this is a matcher-selection change and the evidence is host-side by construction), and the two arms are the right pair.

The fix answers a review question by adding a consumer clause to a weight predicate, which is the correct direction: few_token_batch_reads_packed now takes the graph and returns false for a weight with a GET_ROWS consumer. #324's optimisation was "skip the copy when prefill reads the weight packed", and a tied token embedding also reaches the same bytes through GET_ROWS — so the packed-read conclusion was true about prefill and false about the weight. Naming the exception by consumer rather than by a special case for embeddings is what makes #131's tied-weight class handled once instead of per model.

The evidence table is the discriminating form I keep asking for:

this PR                          pass: ggml_kquant_mul_mat_decode_tokens_f32      85/85
fork main (#324 without clause)  fail: ggml_mul_mat_f32_f32_wmma takes the batch   —

Two rows that differ in which kernel the matcher picks, with the fork-main run failing on the specific wrong kernel — that is what proves the new clause is load-bearing rather than incidental. A single "both pass" row would have proved nothing, since the wrong kernel still computes plausible values.

Two asks:

1. A row at the boundary the clause now excludes. tokens6 and tokens8 cover 2-to-8; the interesting neighbour is tokens1 (single token — where the decode kernel should still be chosen) and a tied weight at tokens9 (where the 2-to-8 rule no longer applies). If those exist in the suite already, one line saying so closes it; if not, they are the two places a future widening of the range will break.
2. Say what the tied path costs. The clause declines an optimisation, so the honest sentence is "a tied weight's 2-to-8-token prefill copies once more, ~N MiB", which is the price of correctness and worth one line — it is the same accounting hrx: keep the model weights in the loader's file mapping instead of a device buffer (engine #856) #304/hrx: a 6-to-16-token batch reads a Q6_K weight in its packed rows too, so a short prompt and its reply share one copy (engine #2137) #316 did for GTT and it keeps the next reader from "fixing" the clause away to recover a number.

Scope of my check: the diff summary, the body's before/after table, and the "Not in this PR" note (the Q6_K 6-to-16 case reaching the prepacked kernel first). No build, no device: the arms are host-side scheduling checks by design, and the device suite was not run — which the body states.