[bong-water-water-bong]

hrx: let element-wise expert weights use the 32-value packet in MUL_MAT_ID, not the dense table's 256 default (engine ggml-org#991)

BlackMamba-2.8B loses 90 of 1179 nodes to the CPU, and 72 of them are one MoE block pattern:
MUL_MAT_ID x 54, GLU x 18. Its experts are ffn_gate_exps/ffn_up_exps = [1472, 3872, 8] and
ffn_down_exps = [3872, 1472, 8], so the contractions are 1472 and 3872: both multiples of 32, neither a
multiple of 256. The 18 GLU nodes are collateral — its matcher already accepts rank>2, but its operands are the
expert matmuls' outputs, which sat on the CPU.

Root cause

common_mul_mat_id_supported_dense_input_size() forwarded straight to the dense table
(common_mul_mat_supported_dense_input_size()), whose default: branch is % 256. F16, BF16 and F32 are not
among its listed cases, so element-wise expert weights inherited the block-quant default. (IQ4_NL, MXFP4, Q8_0,
Q4_0 and friends are listed and already get % 32 — which is why some MoE models place fine and these do not.)

That refusal is policy, not hardware. The three matchers that consult this predicate
(dispatch-mul-mat-id.cpp, dispatch-mul-mat-id-decode.cpp, dispatch-gated-mul-mat-id.cpp) bind exactly five
kernels, and all five declare or assume multiple_of(32):

kernel	contract
ggml_mul_mat_id_f32_f32_wmma	config.decl … input_size … multiple_of(32)
ggml_mul_mat_id_postops_f32_f32_wmma	same
ggml_mul_mat_id_postops_next_rmsnorm_f32_f32_wmma	same
ggml_mul_mat_id_swiglu_f32_f32_wmma	same
ggml_mul_mat_id_decode_f32_wave64	index.assume %input_size … multiple_of(32) (mul_mat_id_decode_f32.loom:66)

and the wmma core already handles a K tail the way the dense prefill core does — padded_input_size = input_size + 255,
quant_block_count = padded / 256, a valid_k gate on every weight value, and masked activation loads
(motifs/mul_mat_id_f32_f32_wmma_core.loom:94-95, 180, 210-211).

The change

One line: the MUL_MAT_ID predicate now goes through common_mul_mat_supported_plain_dense_input_size(), i.e.
element-wise formats get % 32 and block formats keep the granularity their layout requires — the same split the plain
dense path took in #199. Because only F16/BF16/F32 verdicts change, no quantized format's behaviour moves.

Plus three test-backend-ops cases at the real shapes, next to the existing gpt-oss k = 2880 ones: F16 [1472→3872],
F16 [3872→1472], BF16 [1472→3872].

Verification

test-backend-ops -o MUL_MAT_ID -b HRX0 — the same tree and the same test list in both arms, the "before" build differing
only in this one predicate (reverted), so the counts are comparable rather than two different case sets:

arm	result	the three new shapes
before (predicate reverted)	852 / 852 passed	not supported [HRX0]
after (this change)	855 / 855 passed	OK — known-answer vs the CPU backend

BlackMamba-2.8B GGML_HRX_STRICT=1 (one 512-token prefill chunk):

arm	off HRX	breakdown	splits
before	90 / 1179	MUL_MAT 18, MUL_MAT_ID 54, GLU 18	73
after	18 / 1179	MUL_MAT 18	37

72 nodes moved onto the device and the graph lost 36 of its 73 splits. Exactly the predicted set: the 54 expert matmuls,
plus the 18 GLU nodes that were only on the CPU because their operands were. The remaining 18 are BlackMamba's dt_proj
(K = 92, below the 256 floor as well as unaligned) — a separate padded-tile issue, not this one.

Numerics: the model's PPL moves, and by how much

Moving real work onto the device changes accumulation order and precision, so I measured it rather than asserting it was
harmless. BlackMamba-2.8B, same file, expert matmuls on HRX vs on the CPU:

	PPL, 2 chunks, run 1	run 2
expert matmuls on CPU	11.3926	11.3926
expert matmuls on HRX	11.4083	11.4083

Both configurations are deterministic (each repeated bit-identically), so this is a real +0.14 % shift, not run noise.
In the currency the fleet uses for "did numerics move", against the CPU-expert run's own logits over 2048 tokens:

Mean KLD 0.000103 +/- 0.000083, Same top p 99.902 +/- 0.098 %

For scale: this model's own Q8_0 quant arm scores KLD 0.0012 and its Q6_K 0.0073 — this change is an order of magnitude
smaller than the tightest arm in its own table. The size is also a correctness argument: had the K tail been dropped (the
failure mode a padded/masked contraction can have), the error would be ~13 % of the contraction, not 0.14 % of the PPL.

A wrong turn worth recording

I first read ops/mul_mat_id_f16_f16_wmma.loom — which declares multiple_of(256), and whose core computes
quant_block_count = index.div %input_size, 256 with no valid_k, no padded_input_size and no masked load — and
concluded the 256 was load-bearing and that relaxing a predicate would silently truncate the contraction by 13 %. That
file is not bound by any of these matchers; it is a different path. The question that settles it is which kernel
does this matcher bind, which is answered at dispatch-mul-mat-id.cpp:172, and asking it should come before believing
the first multiple_of that grep finds.

[bong-water-water-bong]

#204 ready for review at d9837f2

Engine ggml-org#991, MoE half. One predicate: element-wise expert weights were inheriting the dense table's % 256 default, so
BlackMamba's 1472/3872-wide expert matmuls were refused by the dispatcher while every kernel the matchers bind declares
or assumes multiple_of(32) and the wmma core already pads and masks the K tail.

Measured, same tree and same test list in both arms (the "before" build differs only in this predicate):

• test-backend-ops -o MUL_MAT_ID -b HRX0: 852/852 -> 855/855, the three real-shape cases not supported [HRX0] -> OK.
• BlackMamba-2.8B: 90/1179 -> 18/1179 nodes off HRX, graph splits 73 -> 37.
• Numerics, which is the part a reviewer should weigh: PPL 11.3926 -> 11.4083, deterministic across repeats
  (11.3926/11.3926 and 11.4083/11.4083), i.e. a real +0.14 % shift from f16-vs-f32 accumulation. Against the CPU-expert
  run's own logits over 2048 tokens: Mean KLD 0.000103 +/- 0.000083, Same top p 99.902 +/- 0.098 % — an order of
  magnitude tighter than this model's own Q8_0 arm (0.0012). Had a K tail been dropped the error would be ~13 % of the
  contraction, so 0.14 % of PPL is the precision signature, not a masking bug.

Two things I got wrong on the way, both recovered by asking "which kernel does the matcher bind?" rather than believing the
first multiple_of that grep finds: I first read ops/mul_mat_id_f16_f16_wmma.loom (declares % 256, exact div 256, no
masking) and nearly reported the constraint as load-bearing; and my first numerics A/B used
GGML_HRX_DISABLE_DISPATCH=common.mul_mat_id.f32_f32_wmma, which has no contrast at all because a second, higher-priority
Fused registration claims the same nodes — that one is worth knowing for anyone debugging this area.

[bong-water-water-bong]

VERDICT: READY at c9adbe12 — the fix flips exactly the verdicts it should and nothing else, which I proved with a device-free probe against both versions of the predicate rather than by reading them.

The bug, reproduced: common_mul_mat_id_supported_dense_input_size() for element-wise experts, compiled from prs-204^ and from this head, called directly:

BEFORE (prs-204^)                        AFTER (this PR)
F16   K=1472   id_predicate=0            F16   K=1472   id_predicate=1
F16   K=3872   id_predicate=0            F16   K=3872   id_predicate=1
BF16  K=1472   id_predicate=0            BF16  K=1472   id_predicate=1
BF16  K=3872   id_predicate=0            BF16  K=3872   id_predicate=1
Q8_0  K=1472   id_predicate=1            Q8_0  K=1472   id_predicate=1     <- unchanged
Q4_0  K=1472   id_predicate=1            Q4_0  K=1472   id_predicate=1     <- unchanged

That is the body's claim exactly: BlackMamba's 1472 / 3872 were refused for F16/BF16 while the quantized formats were already fine, which is why some MoE models place and these do not.

The negative control holds too — a contraction that is not a multiple of 32, or below the 256 floor, is still refused:

F16 K=1500 (%32 != 0)  -> 0
F16 K=100  (< 256)     -> 0
F16 K=3840 (%256)      -> 1

And the fix is not a new policy, it is the dense path's own rule applied to the ID path, which is the part I most wanted to check. The dense MUL_MAT matcher already chooses between the two predicates on exactly this distinction:

!(elementwise_k_tail ? common_is_supported_plain_dense_input_size(format, input_size)
                     : common_is_supported_dense_input_size(format, input_size)) || …

and common_mul_mat_supported_plain_dense_input_size() gives % 32 for common_mul_mat_dense_float_format and defers to the block table otherwise — with the reasoning (wmma pads to 256-blocks, valid_k gates the tail, activation is a masked vector.load) documented at the helper itself. So MUL_MAT_ID inheriting % 256 was the anomaly, and pointing it at the same predicate is the smallest correct change.

The test cases pin the real geometry ([3872, 512, 1472], [1472, 512, 3872] for F16 and one BF16), not a round number, so a future regression in this predicate fails test-backend-ops on the shape that matters.

Not verified by me: the placement itself (90 of 1179 nodes, MUL_MAT_ID x 54 moving to HRX) and the end-to-end on BlackMamba-2.8B — that needs the device and the model. What I can say is that the predicate now returns true for those contractions and the kernels' own multiple_of(32) contract is what makes it safe.

[bong-water-water-bong]

Correction to the sha in my verdict above: the head is d9837f2ea (single commit, 2 files, +16/−1). The c9adbe12 I quoted is the post-image blob of tests/test-backend-ops.cpp (index 5aea24db9b8..c9adbe12021 in the diff), not a commit — I read the index line instead of the head. The branch has not moved since the d9837f2e this PR announced, so the verdict stands unchanged at d9837f2ea; only the label was wrong.

Recorded here rather than edited so the gate names a real commit. I am re-deriving the 852/852 -> 855/855 matrix on the device at this head now (HRX build, -o MUL_MAT_ID -b HRX0, the exact device name — -b HRX still matches nothing and exits 0, engine ggml-org#657) and will post the counts as a follow-up.

[bong-water-water-bong]

Follow-up on the counts I said I would post: I did not get the device run. The HRX build succeeded (clean, test-backend-ops linked, the two fork-side warnings are pre-existing), but the gpu-hold -l "pi prs" claim sat queued from 05:29 to my 40-minute timeout without ever taking the token — the queue has five other claimants behind it (hw-gate run_42004722505, demo-img-hrx2, mosaic-tile-small2, nll-probe, pi-hrx-mem), so this is contention, not a stall I could act on. Retrying now would put the same claim at the back of the same queue.

So my READY at d9837f2ea stands on what I actually verified, which is the predicate and not the device:

BEFORE (prs-204^)                        AFTER (this PR)
F16   K=1472   id_predicate=0            F16   K=1472   id_predicate=1
F16   K=3872   id_predicate=0            F16   K=3872   id_predicate=1
BF16  K=1472   id_predicate=0            BF16  K=1472   id_predicate=1
BF16  K=3872   id_predicate=0            BF16  K=3872   id_predicate=1
Q8_0  K=1472   id_predicate=1            Q8_0  K=1472   id_predicate=1     <- unchanged
Q4_0  K=1472   id_predicate=1            Q4_0  K=1472   id_predicate=1     <- unchanged
F16 K=1500 (%32 != 0) -> 0    F16 K=100 (< 256) -> 0    F16 K=3840 (%256) -> 1

and the PR's own 852/852 -> 855/855 is arithmetic I can check from the case list (the three added cases are exactly the difference: 852 passing + 3 not supported → 855 passing), but the device counts remain the author's. For anyone who picks this up: the run is test-backend-ops test -o MUL_MAT_ID -b HRX0 with IREE_HAL_AMDGPU_LIBHSA_PATH set, and -b HRX (family name) matches no device and exits 0 — engine ggml-org#657.

[bong-water-water-bong]

Review note for the interaction with the padded-tile change (fork #207's sibling, 1bit/hrx-padded-k-tail @ ad732f951), so this is decided before both land rather than after:

This PR routes element-wise expert weights to common_mul_mat_supported_plain_dense_input_size(), which is exactly the
right rule today (it returns % 32 for element-wise formats). The padded-tile change then widens that same function to
% 4, because the plain dense kernels' mask granularity is the 4-value packet and every residual Zyphra contraction is
% 4 == 0.

The five MUL_MAT_ID kernels this PR binds are not widened by that change — they still declare/assume
multiple_of(32). So after both merge, an element-wise expert with K % 4 == 0 && K % 32 != 0 would pass the predicate and
then bind a kernel whose config declares % 32, which is a JIT-specialisation mismatch rather than a fallback. No model in
my set has such an expert (BlackMamba's are 1472 and 3872, both % 32 == 0), so this is a latent hazard, not a present
break — but it means one of two things should be true before or soon after both land:

1. this route keeps % 32 for element-wise experts (a separate helper beside the padded one), or
2. the five MUL_MAT_ID kernels' % 32 contracts are widened to % 4 in the same step, with the same kind of proof this
   PR already carries (test-backend-ops at the unaligned shape).

I am not proposing to change this PR: its rule is correct for the kernels it binds. Recording it here because "the dense path
grew a looser contract" silently changing a different path's admission is exactly the kind of coupling that costs someone a
night later.