[bong-water-water-bong]

hrx: admit element-wise dense MUL_MAT contractions at input_size % 4 (engine ggml-org#991)

The gap

The plain dense MUL_MAT path declared its contraction as range(input_size, 256, 32768), multiple_of(input_size, 32) —
in mul_mat_f32_f32_wmma's ggml.mul_mat.input_size config, in the assume of its core, in both entry points of
mul_mat_f32_f32_decode, and (this is the one that actually bites) in the element-wise packet reader motifs
f16_f16 / bf16_f16 / f32_f16, each of which builds view<[bounded_input_size]x{bf16,f16,f32}> under a 256 contract.
common_mul_mat_supported_plain_dense_input_size repeated the same >= 256 && % 32 == 0 for element-wise weights. So any
contraction that is not 32-aligned, or is narrower than 256, went to the CPU.

That is the last class in the Zyphra ports: BlackMamba's ssm_dt.weight is [92, 2944], Zamba v1's is [232, 3712].

Why % 4 is the real requirement

Reading the kernels and the motifs, not the contract:

• the element-wise reader takes four values at a time and the motif's vector.load is unmasked — but its caller
  gates each packet with k < input_size, so the four values are inside the row exactly when input_size % 4 == 0;
• the activation side is vector.mask.range [input_k to input_size), i.e. zero-exact at element granularity;
• the block loop is already clamped: padded_input_size = input_size + 255, block_count, valid_block,
  safe_block = select(valid_block, block, 0) and select(valid_block, sum, 0) on the contribution.

32 and 256 were contracts nobody's arithmetic needed. K % 4 != 0 stays refused on purpose — such a row's last packet
would read up to three values past the row, which is correct only by accident of where the next row happens to be. A test
case asserts that boundary.

What changes

• common_mul_mat_supported_plain_dense_input_size: element-wise formats need input_size >= 4 && % 4 == 0. Block
  formats are untouched and keep their %256/%128/%64/%32 from common_mul_mat_supported_dense_input_size().
• range 256 -> 4, multiple_of 32/256 -> 4 in the five places above. Weakening a declared contract is safe for every
  existing caller: a block-format contraction is always a multiple of 256, hence also of 4.

Who enforces what, after this. The three motifs are shared with the MoE/gated paths, so for those paths the 256
guarantee moves from the JIT to the dispatcher (they stay gated at %32/%256). That is a real coupling — the same one I
flagged on #204: a predicate that grows looser in one path changes what can reach another path's kernels. Here the only new
exposure is a shape no dispatcher admits today, and the motif's arithmetic is caller-gated per packet, so the change is
safe; it is stated rather than left implicit.

Evidence

Three builds, one ticket, and every arm asserted rather than printed: the five padded cases must read OK, the k=90
control must read not supported, no arm may log Loom compilation failed, and every model arm must emit its bench row and
its PPL estimate. failed=0.

test-backend-ops -b HRX0 -o MUL_MAT (this branch, before #207):

bf16 m=2944 k=92  n=8  OK        bf16 m=3712 k=232 n=8  OK        bf16 m=512 k=260 n=8  OK
bf16 m=2944 k=92  n=1  OK        bf16 m=3712 k=232 n=1  OK        f32  m=512 k=260 n=8  OK
bf16 m=512  k=90  n=8  not supported   <- the boundary: 90 % 4 == 2

BlackMamba-2.8B (this change alone; the base arm is the same binary tree with the batched dispatch disabled, i.e.
origin behaviour for MUL_MAT):

arm	nodes off HRX	class split	graph splits	pp512 (-r 5)	tg64 (-r 5)	PPL
base	90 / 1179	MUL_MAT 18, MUL_MAT_ID 54, GLU 18	73	788.69 ± 12.47	38.38 ± 2.18	12.5702
this PR	72 / 1179	MUL_MAT_ID 54, GLU 18	37	1235.54 ± 38.65	46.89 ± 1.24	12.5702

dt_proj's 18 nodes are gone, and the remaining 72 are the MoE classes that fork #204 covers. PPL is unchanged at the
precision printed; the per-shape correctness evidence is the test-backend-ops table.

Zamba-7B-v1 needs this and #207, because its dt_proj nodes are also batched (ne[2] = 2):

arm	nodes off HRX	graph splits	pp512 (-r 5)	tg64 (-r 5)	PPL
this change alone	241 / 5048 (MUL_MAT 152, CONT 76, FLASH_ATTN_EXT 13)	179	68.98 ± 0.58	5.88 ± 0.23	20.0688
local merge with #207	13 / 5048 (FLASH_ATTN_EXT 13)	27	195.10 ± 4.29	6.64 ± 0.27	20.0704

The middle row is the honest one for this PR: it moves v1 not at all, because a batched node is still refused before any
size rule. The last row is measured on a local merge of this branch with #207 (git merge 1f2768431, never pushed) so
the reviewer can see what lands once both are in; #207 must merge first for that number to apply.

What this does not do

• K % 4 != 0 stays on the CPU, with a test asserting it.
• mul_mat_id_swiglu_f32_f32_wmma, mul_mat_id_decode_f32 and mul_mat_dual_q4_f32_decode keep their %32 contracts —
  they belong to the MoE/gated classes, and widening them without a shape that needs it would be an unmeasured change.
• No engine build; the fork built here is the evidence.

Notes for the reviewer

• My first attempt at this was broken in a way that looked like success, and the PR would have shipped a JIT-time
  failure as a placement win. I widened the predicate and the op/core/decode assumes, and the arm happily printed a
  placement line (72 off, dt_proj gone) and a PPL — while the log was full of
  Loom compilation failed … predicate 'range' on 'bounded_input_size' requires range [256, 32768]. The contract that
  actually rejected the spec lives in the three motifs, not in the files I had read; the verdict fields in my greps were
  blank, and I read blank as printed. Take 2 asserts on the evidence instead of printing it, which is why the table above
  has failed=0 next to it rather than a summary line.
• Nothing in either kernel's instruction stream changed. This is a contract-and-predicate change, and the proof is the
  device-vs-CPU rows plus the model placement.

[bong-water-water-bong]

#208 ready for review at 21dfef4

Engine ggml-org#991, the padded-tile half — and it turned out to be a contract, not a kernel: no instruction changed.

The gate was not where I first looked. The plain dense path declared range(256, 32768), multiple_of(32) in five places, and the one that actually rejects an unaligned contraction is the element-wise packet reader motif (f16_f16, bf16_f16, f32_f16), which builds view<[input_size]x{bf16,f16,f32}> under a 256 contract and does an unmasked 4-wide load. The caller gates each packet with k < input_size, so the real requirement is % 4 — the four loaded values stay inside the row exactly then.

Measured, three builds, one ticket, and every arm asserted (five padded cases must read OK, the k=90 control must read not supported, no arm may log Loom compilation failed, every model arm must emit bench + PPL): failed=0.

• test-backend-ops: k=92 (real BlackMamba [92,2944]), k=232 (real v1 [232,3712]), k=260 all OK vs the CPU reference; k=90 (%4 == 2) not supported, which is the boundary this PR deliberately keeps.
• BlackMamba-2.8B, this change alone: 90 -> 72 nodes off HRX (MUL_MAT 18 -> 0), graph splits 73 -> 37, pp512 788.69 ± 12.47 -> 1235.54 ± 38.65, tg64 38.38 ± 2.18 -> 46.89 ± 1.24, PPL unchanged at 12.5702. The remaining 72 are the MoE classes hrx: let element-wise expert weights use the 32-value packet in MUL_MAT_ID (engine #991) #204 covers.
• Zamba-7B-v1 needs this and hrx: run a batched dense MUL_MAT as one dispatch per plane (engine #991) #207 (its dt_proj is also batched): this change alone leaves it at 241/5048, and a local merge with hrx: run a batched dense MUL_MAT as one dispatch per plane (engine #991) #207 gives 13/5048 (FLASH_ATTN_EXT only), splits 179 -> 27, pp512 68.98 -> 195.10, tg64 5.88 -> 6.64, PPL 20.0688 -> 20.0704. The isolated row is in the body too, because it is the one that applies to this PR on its own.

Two things a reviewer should weigh, both in the body: the three motifs are shared with the MoE/gated paths, so their 256-guarantee moves from the JIT to the dispatcher (the coupling I flagged on #204), and my first attempt at this shipped a JIT-time failure that looked like a placement win — the arm printed placement and PPL while the log was full of Loom compilation failed; the grep verdicts were blank and I read blank as printed. Take 2 asserts instead of printing.

[bong-water-water-bong]

#208 re-based onto 4a375bf2d and re-verified — now at 9972afa1b

Three things changed since the first push, all in the body:

1. Rebased onto the real head, which already carries hrx: let element-wise expert weights use the 32-value packet in MUL_MAT_ID (engine #991) #204/graph: keep the fused attention op on its layer's device (engine #1078) #205/convert: Zamba2-VL (Zamba2 language model + Qwen2.5-VL tower), so the published GGUFs have an authoring site (engine#1098) #206/hrx: run a batched dense MUL_MAT as one dispatch per plane (engine #991) #207, so the numbers are no longer from a local merge.
2. The hrx: let element-wise expert weights use the 32-value packet in MUL_MAT_ID (engine #991) #204 coupling is closed in code, not left as a comment. common_mul_mat_id_supported_dense_input_size keeps >= 256 && % 32 == 0 for element-wise formats: hrx: let element-wise expert weights use the 32-value packet in MUL_MAT_ID (engine #991) #204 routes element-wise experts through the predicate this PR widens, but the five MUL_MAT_ID kernels declare range(256,32768), multiple_of(32) and none of the widened dense motifs is bound there — inheriting the looser rule would be a JIT failure, not a fallback. Pinned from both sides by test: MUL_MAT_ID k=1472 OK, k=92 and k=1480 not supported.
3. Merged-base results: BlackMamba places completely — 0 of 1179 nodes off HRX, 1 graph split, PPL 12.5406 (its speed rows that run were ±312 t/s, i.e. taken under other lanes' load, so not quoted; the isolated same-tree A/B stands at pp512 788.69 -> 1235.54, tg64 38.38 -> 46.89). Zamba v1: 13 of 5048 off HRX, FLASH_ATTN_EXT only, 27 splits (from 179), PPL 20.0704.

Two harness notes, because they are the same lesson twice: the final arm failed its own assertions on a regex bug in my script (a trailing comma the MUL_MAT_ID case format does not have), not on the code — the logs show the verdicts exactly as intended, and I re-asserted the saved logs offline to PASS before writing the body. That is the harness doing its job: a loud failure instead of a quiet pass. The explanation of how the first attempt shipped a JIT-time failure that looked like a placement win is in the body, since it is the failure mode a reviewer should know about for this area.

[bong-water-water-bong]

VERDICT: READY at 9972afa1 on the code — the relaxation is scoped to the path whose motifs justify it, and the path it is not applied to is the interesting half. Device row owed.

The decision I most wanted to see is the negative one. common_mul_mat_id_supported_dense_input_size() keeps %32 for element-wise formats, with the reason written down:

// The five MUL_MAT_ID kernels declare `range(256, 32768), multiple_of(32)` … Sharing the looser predicate would let
// an element-wise expert reach a kernel whose own config rejects it, which is a JIT failure, not a fallback.

That is exactly the failure mode a contract relaxation invites — the dispatcher and the kernel disagreeing about the same
contraction — and refusing to touch the ID path while relaxing the dense one is the conservative split. Compare it with what
engine #204 did on the ID side last night (which loosened %256 → %32 there, because those kernels declare 32), and the
two together are consistent rather than contradictory.

The relaxation itself is justified in the same place it is made: the element-wise reader loads a 4-value packet and the
caller gates each packet with k < input_size, so a packet starts inside the row exactly when input_size % 4 == 0; the 256
floor and the 256 multiple were a contract, not a requirement of the code. And the change is coherent across every place the
contract is declared — mul_mat_f32_f32_wmma's config, its core's assume, both decode entry points, and all three packet
reader motifs (f16_f16, bf16_f16, f32_f16) — which matters because a config/assume mismatch is a JIT failure, not a slow
path.

What I cannot verify and would not sign off without: test-backend-ops on HRX0 at an input_size that is %4 but not
%32 (the new cases are the guard). A relaxation here either works or produces wrong values, and only the device distinguishes
those; the GPU has been held all night. If the hw-gate label goes on the next pin bump, this is one of the two things worth
pointing it at.

[bong-water-water-bong]

Correction to a claim in my standing line — the strict run caught it.

I wrote "BlackMamba places completely: 0 of 1179 nodes off HRX, 1 graph split". That is the decode graph only. A
GGML_HRX_STRICT=1 run on the same binary aborts on:

placement_observer: GGML_HRX_STRICT=1: "node_31" is MUL_MAT f32 ne=[2944,512,1,1] and is not placed on HRX;
18 of 1179 nodes off HRX (MUL_MAT x 18), 37 splits

So the 512-token prefill ubatch still has the 18 dt_proj nodes off (37 splits). The placement line I quoted was the
last of several printed per graph — the first line in that log is the 18-off prefill graph and I took tail -1. The PR
body now says this instead, and the strict abort is the check that cannot be mis-read, which is why I ran it.

What it means for this PR's claim, narrowly: the % 4 relaxation places K=92 and K=232 at the token counts my test
cases cover (8, and token 1) and not at 512 tokens, where the prefill ubatch selects a different dispatch/kernel path
that still carries a 256 contract. My test cases used n=8, which is the wrong token count to pin the model's real prefill
shape — the same family as the earlier shared-vs-paired weight mistake: a test that proves a shape the model does not run is
not a test.

Next step, in progress: reproduce it as test_mul_mat(BF16, F32, 2944, 512, 92) and either close it here or narrow this
PR's claim to what it actually does.

[bong-water-water-bong]

Follow-up on the correction above — the reproduction is done and it names the gate, which is not the K contract:

PROBE prefill-small-k: w(t=30 bf16 ne=[92,2944,1,1] c=1)  in(t=0 f32 ne=[92,512,1,1] c=0)  out(t=0 f32 ne=[2944,512,1,1] c=1)

At the 512-token prefill ubatch, dt_proj's activation is a non-contiguous view (c=0). The same contraction with a
dense activation is fine: test_mul_mat(bf16, f32, m=2944, n=512, k=92) and m=3712, n=512, k=232 are both OK on the
device against the CPU reference, with zero JIT rejections — so this PR's predicate is correct at the model's real token
count too, and the 18 nodes that GGML_HRX_STRICT=1 reports are refused by the matcher's !input->contiguous guard,
before any size rule.

That means the % 4 work is complete for what it claims and the remaining prefill fallback is a separate unit:
either the matcher accepts a view whose rows are dense (binding it with a row stride) or the dispatch emits a densifying
copy transient — the same mechanism common_prepare_f16_input already uses for the f16 activation path. I have not
attempted it here, and I am not claiming it.

I pushed 3f0942f91 on top to keep the coverage that produced this: n=512 cases for both shapes. They pass, and they
are the reason the difference between "predicate is wrong" and "the node is a strided view" is now a measurement rather
than a guess.