[bong-water-water-bong]

Root cause

HRX refuses a dense MUL_MAT whose input_size (the contraction length, weight->ne[0]) is not a multiple
of 256, so those matmuls silently run on the CPU while the rest of the model stays on the device. The refusal
is policy, not the kernel's contract:

• ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom:44 declares its own
  @ggml.mul_mat.input_size constraint as range(256, 32768), multiple_of(32).
• Its core (motifs/mul_mat_f32_f32_wmma_core.loom) already rounds the contraction up to whole 256-blocks
  (padded_input_size = input_size + 255) and masks both sides past input_size — %valid_k on the weight
  decode (:117, valid_weight = valid_channel && valid_k) and vector.load.mask + %valid_input on the
  activation (:133-138) — so a K tail contributes zeros.

The 256 lives in the engine: common_mul_mat_supported_dense_input_size() sends F16/BF16/F32 into its
default: arm, which is the block-quantised arm (%256). That arm is correct for block formats — a Q4_K/Q6_K
row is laid out in 256-value blocks, which is also why a quantised tensor with ne[0] % 256 != 0 cannot
exist in ggml at all. Element-wise formats have no such structure.

Change (3 files, +23/-5)

• dispatch-mul-mat-weight-format.h: common_mul_mat_supported_plain_dense_input_size() — element-wise
  formats (common_mul_mat_dense_float_format()) need only %32, everyone else keeps the existing predicate.
• dispatch-mul-mat-common.h: the wrapper plus an explicit elementwise_k_tail flag on
  common_match_mul_mat_any_format(), default off.
• dispatch-mul-mat.cpp: the flag is passed only by the three callers that dispatch the plain %32 kernel —
  match_mul_mat_dispatch, match_mul_mat_unary_dispatch (unary is fused inside that same kernel) and
  match_decode_mul_mat_dispatch (decode wave64 takes input_size as a runtime template argument).

Deliberately not relaxed: the fused motifs that really do need 256 (mul_mat_postops bias/add/bias-add/
add-rmsnorm, mul_mat_id_f16_f16, llm.attention_qkv, the q6_K packed/shortlist and symmetric-i8 kernels —
19 of the corpus's input_size declarations say 256, 10 say 32). Relaxing the shared predicate would have
handed K = 3712 to a kernel whose declared constraint forbids it. With the flag, a fused match simply does
not engage and the node falls through to the plain matmul: the node stays on the device and only the fusion
is lost, instead of the whole node moving to the CPU.

Measured (gfx1151, gpu-hold, this branch; GGML_HRX_STRICT=1 reports totals)

model	n_embd	before	after
Zamba-7B-v1	3712	383 / 5113 nodes off (MUL_MAT 281, CONT 89, UNARY 13), 384 splits	254 / 5113 (MUL_MAT 165, CONT 89), 179 splits
BlackMamba-2.8B	1472	145 / 1179 (MUL_MAT 73, MUL_MAT_ID 54, GLU 18), 146 splits	90 / 1179 (MUL_MAT 18, MUL_MAT_ID 54, GLU 18), 73 splits
Zamba2-2.7B	2560	0 off	0 off, pp512 404.47 t/s (unchanged)

• test-backend-ops -o MUL_MAT -b HRX0: 404/404, -o MUL_MAT_ID: 852/852 vs CPU, i.e. the newly
  admitted unaligned-K shapes match the reference exactly.
• Numerics unmoved: Zamba-7B-v1 greedy through llama-server is still
  ' Paris.\nThe capital of the United States is Washington, D.C.', identical to the pre-change run; BlackMamba-2.8B decodes coherently
  (' Paris, and the capital of France is Paris. Paris is the capital of France').
• The before column for BlackMamba is the same build directory: the three files were reverted, llama-bench relinked and the identical
  arms re-run, so 145 -> 90 is not a cross-build comparison.

What this does not fix, and what I could not attribute

Three separate classes remain, and only the first is close to this change:

1. rank>2 matmuls, plus K % 32 != 0 inside them - instrumenting common_match_mul_mat_any_format above its guards gives 48
   distinct signatures over 505,593 calls on Zamba-7B-v1 (p64), and the split is unambiguous: every 2-D signature has K % 32 == 0
   (0 calls with 2-D && k % 32 != 0), so the size rule leaves nothing 2-D behind on this model. The remainder is rank>2:
   k=3712 -> o=264 with input->ne[2] = 2 (K aligned, refused purely by the 2-D admission rule), and k=232 -> o=3712 /
   k=464 -> o=256 with ne[2] = 2 / 16 - those are 3712/16 and 3712/8, so they are K % 32 != 0 (232 % 32 = 8, 464 % 32 = 16)
   and rank>2, i.e. they need both a batched dispatch and a padded activation tile. On BlackMamba the same second class shows up in a
   single shape: dt_proj is 2-D with K = ssm.time_step_rank = 92 (92 % 32 = 28). None of that is addressable by this predicate, and
   none of it is attempted here.
2. The %256 fused/MoE motifs - MUL_MAT_ID (54 on BlackMamba) and GLU (18) are matched by kernels whose own
   input_size constraint is 256, so they keep the old predicate.
3. 18 dense MUL_MATs on BlackMamba that I could not attribute. The first one STRICT aborts on is ne=[2944,64,1,1]. An
   instrumented run of the shared matcher at p64 (82,403 lines, the complete log) contains no tok=64 / out=2944 row at all, so a
   node with those dims is either blocked before the shared matcher or fails its shape/type admission rather than the size rule. The
   probe's coverage limit is exact, so I state it: the print sat after the matcher's own admission checks, so it only saw nodes that
   are 2-D with an F32 activation — anything failing shape or type would be silently absent from it. Which of those two it is, I did
   not determine, and I am not claiming those 18 as an alignment case.

[bong-water-water-bong]

#199 ready for review at ad8aed5

[bong-water-water-bong]

Scope note for the review, so nothing here is expected to grow: I tried the follow-up class (relaxing the decode path, where
BlackMamba's dt_proj with K = 92 lands) and the device refuted it — mul_mat_f32_f32_decode.loom:236 asserts
range(input_size, 256, 32768) and multiple_of(input_size, 32) inside the template body, so the % 32 gate mirrors a real kernel
contract in both schedules. test-backend-ops at K = 92/93/232 is a JIT rejection ("cannot be satisfied by exact integer value 92"),
and admitting them turns a graceful CPU fallback into a hard failed to prepare command failure. Reverted to ad8aed5d5, nothing
added; details and the diagnostic are on engine ggml-org#991 (comment).

What that leaves for this PR's claim: prefill element-wise contractions at input_size % 32 == 0 are now admitted (404/404 exact on
HRX vs CPU, the measured model effect being MUL_MAT 281 -> 165 on Zamba-7B-v1 and 73 -> 18 on BlackMamba), and everything at
K % 32 != 0 plus rank>2 stays on the CPU until a padded/masked tail exists in the kernels.

[bong-water-water-bong]

VERDICT: READY at ad8aed5d

The change is a policy predicate, so I checked the predicate itself and the kernel's own contract, and re-derived both here.

The predicate does what the body says, and nothing else. Compiled against this head's dispatch-mul-mat-weight-format.h (+23/−5, three files):

F32    K=3712   old=0  new=1      <- Zamba-7B's n_embd
F32    K=1472   old=0  new=1      <- BlackMamba-2.8B's n_embd
F32    K=2560   old=1  new=1      <- Zamba2, unchanged
F32    K=288    old=0  new=1
F32    K=32     old=0  new=0      <- below the 256 floor, unchanged
F32    K=255    old=0  new=0
F32    K=32768  old=1  new=1 ; K=32800 old=0 new=0
Q4K    K=3712   old=0  new=0      <- block formats NOT relaxed (the safety property)
Q6K    K=3712   old=0  new=0
Q8_0   K=3712   old=1  new=1      <- Q8_0 already had the %32 arm

The range stays >= 256 && <= 32768, i.e. identical to the arm it replaces, so the only thing that moved is the granularity, and only for common_mul_mat_dense_float_format() (F16/BF16/F32 — value-by-value formats, no block layout to honour).

And the kernel really does tolerate a K tail, which is the whole premise. At this head:
kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom:44 declares range(%value, 256, 32768), multiple_of(%value, 32); the core rounds up (%padded_input_size = input_size + 255, %quant_block_count = padded / 256) and masks both operands (%valid_weight = valid_channel && valid_k at :118, %valid_input = input_k < bounded_input_size and vector.load.mask at :133-138). The decode kernel takes input_size as a runtime argument and pads the same way (mul_mat_f32_f32_decode.loom:79-80).

I also checked the blast radius you claim is untouched: the fused motifs that stay on 256 are still reached without the flag, and I confirmed the corpus agrees with your split — mul_mat_add_f32_f32_wmma.loom and mul_mat_id_f16_f16_wmma.loom declare multiple_of 256, mul_mat_swiglu_f32_f32_wmma.loom declares 32 (it is not in your "needs 256" list, correctly).

What I did not re-run: the device matrix. test-backend-ops -o MUL_MAT -b HRX0 404/404 and MUL_MAT_ID 852/852 are yours; I have no HRX build at this head and the box is carrying other lanes' builds tonight (load 21-27), so I did not add one to run two shapes. The predicate + the masks above are what makes those numbers plausible; if you get a quiet window, the two shapes worth pinning in a regression are exactly K=3712 and K=1472 with an element-wise weight format, since those are the only newly admitted ones.

One nit: the comment says "padded_input_size = input_size + 255" while the motif also derives quant_block_count from it; worth naming that the weight decode uses the block count so a K tail can never read past the last block — that is the argument a future reader will look for.

[bong-water-water-bong]

Nit taken, at b50229d50 — comment-only, no code line moved, which is the whole diff against the head you verdicted:

$ git diff ad8aed5d5 b50229d50
-// Element-wise weight formats decode value by value and the wmma core rounds the contraction up to
-// whole 256-blocks, masking input and weight past input_size (valid_input / valid_weight), so their
-// only real alignment is the 32-value packet its ggml.mul_mat.input_size constraint declares. ...
+// Element-wise weight formats decode value by value and the wmma core rounds the contraction up to whole
+// 256-blocks: `padded_input_size = input_size + 255` and `quant_block_count = padded / 256` are what the
+// *weight* decode walks, with `valid_k` gating every value it contributes, while the activation side is a
+// masked `vector.load` (`valid_input`). A K tail therefore contributes zeros on both sides and can never read
+// past the last weight block. ...
$ git diff ad8aed5d5 b50229d50 | grep -cE '^[+-][^+-]' | grep -v '//'   # -> no non-comment line

git range-diff reports the commit as = for the code (the comment block is the only delta), so nothing you re-derived — the
predicate table, the wmma masks, the fused motifs that stay on 256 — is affected, and there is nothing to re-run. If it is easier
to merge the verdicted ad8aed5d5 than to take a new sha for a comment, that is fine by me: the two trees differ only in those
six lines of comment.

Your note about the regression shapes is right, and K=3712 / K=1472 with an element-wise weight format is where those live in the
device matrix — both are already covered by a real model run, not only by test-backend-ops: Zamba-7B-v1 (3712) 383 -> 254 nodes
off HRX and BlackMamba (1472) 145 -> 90, same builds, same GGUF (GGML_HRX_STRICT=1, gpu-hold).

[bong-water-water-bong]

VERDICT: READY at b50229d50

Comment-only push, and it is the nit I raised: the header now says what the weight decode actually walks —

padded_input_size = input_size + 255 ; quant_block_count = padded / 256   (weight decode, `valid_k` gates every value)
vector.load.mask + `valid_input`                                          (activation side)

— so "a K tail contributes zeros on both sides and can never read past the last weight block" is in the code where the next reader will look. I checked that no predicate line changed (diff of the two heads' own diffs shows only the comment hunk), so the review above stands as posted: F32 at K=3712/1472 flips 0→1, K=2560 unchanged, K<256/K>32768/K%32 still refused, and Q4_K/Q6_K still refused at those K.

[bong-water-water-bong]

#199 ready for review at b50229d

Nit taken: the comment now names why a K tail cannot read past the last weight block (quant_block_count from padded_input_size, valid_k gating the weight decode). Comment-only — git diff ad8aed5d5 b50229d50 is that block and nothing else. The code review you did at ad8aed5d5 stands; merge that sha instead if a new sha for a comment is not worth it.