[bong-water-water-bong]

hrx: run a batched dense MUL_MAT as one dispatch per plane (engine ggml-org#991)

The gap

A MUL_MAT whose operands carry batch planes — ne[2] * ne[3] > 1 — is the same 2-D contraction repeated per plane, and
it is how the SSM/Mamba ports write their small projections (Zamba v1's x_proj is w bf16 ne=[3712,264,2,1] against
in f32 ne=[3712,512,2,1]). Every format-agnostic matcher refused it in common_match_mul_mat_any_format on
common_is_2d(*input), before any size rule, so those nodes went to the CPU and took their CONT consumers with them.

This is not a kernel gap: ggml_mul_mat_f32_f32_wmma (prefill) and ggml_mul_mat_f32_f32_decode_wave64 (token 1) already
compute a 2-D tile from a buffer base, so the planes can simply be bound one dispatch at a time — the mechanism
dispatch-gated-delta-net.cpp already uses per token ({ id, t * nb[2], byte_count - t * nb[2] }).

What changes

• common_match_mul_mat_any_format takes allow_batched: the activation and output may carry planes when the output's
  planes match the activation's. ggml pairs the operands' planes when both carry them and broadcasts when only the
  activation does, so the weight may be one matrix shared by every plane or one matrix per plane — and in the
  paired case the factorisation must match too (ne[2]/ne[3] equal), because pairing follows that index order, not the
  plane product.
• common_plane_offset maps a plane index to its byte offset in a contiguous value.
• build_batched_mul_mat_dispatch emits one dispatch per plane: activation, weight and output bindings offset to that
  plane, the token count is the per-plane one. The weight offset is zero unless the weight is per-plane.
• Two registrations — common.mul_mat.f32_f32_wmma_batched (82) and common.mul_mat.f32_f32_decode_batched (62) — one
  above each plain dense registration they mirror.

The K contract is deliberately unchanged: still the plain element-wise predicate, so K % 32 != 0 stays on the CPU (see
"what this does not do" below).

Evidence

One binary (1f2768431), two arms. The "off" arm sets GGML_HRX_DISABLE_DISPATCH=batched, which matches exactly the two
registrations added here, so before/after differ in the dispatch decision only.

test-backend-ops -b HRX0 -o MUL_MAT — OK means HRX matched the CPU reference:

supported / refused        off: 373 / 1204        on: 526 / 1051      (+153 cases admitted)

The five new real-shape cases (the sixth is the boundary below):

case	off	on
f32/f32 m=264 n=8 k=3712 nr=[2,1] (shared weight)	not supported	OK
f32/f32 m=464 n=8 k=256 nr=[16,1] (shared weight, d_state planes)	not supported	OK
f32/f32 m=264 n=1 k=3712 nr=[2,1] (token 1 → decode registration)	not supported	OK
bf16/f32 m=264 n=8 k=3712 bs=[2,1] (one matrix per plane, the model's shape)	not supported	OK
bf16/f32 m=264 n=1 k=3712 bs=[2,1] (per-plane weight, decode)	not supported	OK
bf16/f32 m=3712 n=2 k=232 bs=[2,1] (unaligned contraction)	not supported	not supported

The rest of the +153 are existing upstream batched cases (nr>1, and bs>1 with an F32 activation), so this admits the
class generally rather than only the shapes I added.

Zamba-7B-v1 (BF16), placement and speed, gpu-hold, llama-bench -v -p 512 -n 64 -r 5:

arm	nodes off HRX	class split	graph splits	pp512 t/s	tg64 t/s	PPL (4 chunks)
on	89 / 5048	MUL_MAT 76, FLASH_ATTN_EXT 13	179	79.97 ± 1.82	6.67 ± 0.05	20.0718
off	241 / 5048	MUL_MAT 152, CONT 76, FLASH_ATTN_EXT 13	179	63.60 ± 6.00	5.89 ± 0.66	20.0688

• 241 → 89 nodes off HRX (−152). The 76 MUL_MAT that moved are exactly the K=3712 → o=264 family, and all 76
  CONT nodes were their consumers (CONT 76 → 0). GGML_HRX_LOG_DISPATCH=1 lists both new registrations as matched.
• Graph splits do not change (179), and I am calling that out rather than burying it: the splits are now driven by the
  remaining 76 MUL_MAT + 13 FLASH_ATTN_EXT, so "fewer splits" is not a benefit of this change.
• Numerics: PPL 20.0718 vs 20.0688 is +0.015 %, and each arm reproduced to 1e-4 across separate runs
  (on: 20.0715, 20.0718; off: 20.0688, 20.0688), i.e. a real but tiny f32-vs-device accumulation shift. The per-shape
  evidence is the test-backend-ops table, not the PPL.

What this does not do

The remaining 76 nodes are dt_proj, w ne=[232,3712,2,1], i.e. K % 32 == 8: they are refused by the plain K predicate
that this PR intentionally leaves alone. That is the padded-tile item (K % 32 != 0 / K < 256), and until it lands
GGML_HRX_STRICT=1 still aborts on this model — so no strict-clean claim is made here.

Notes for the reviewer

• My first version of this change covered only the shared-weight case, because my tests used nr>1 with bs=1. The model
  pairs the planes (bs>1), so the first version passed test-backend-ops and moved nothing on the model (241 → 241).
  A test that proves a shape nothing emits is not a test; both shapes are covered now, which is why the table above has a
  shared-weight and a per-plane block.
• Not verified here: no engine build, no other model. This only touches the HRX dispatch path, and the test-backend-ops
  rows above are the correctness evidence.

[bong-water-water-bong]

#207 ready for review at 1f27684

Engine ggml-org#991, the batched half. rank>2 activations were refused by common_is_2d before any size rule; they are the same 2-D contraction per plane, so the fix binds one dispatch per plane (activation, weight and output offset to that plane) instead of inventing a kernel. ggml pairs the planes when both operands carry them, so the weight may be one shared matrix or one matrix per plane — the model does the second, which is why the first version of this change moved nothing (241 -> 241) while passing the tests.

Measured on one binary, two arms (GGML_HRX_DISABLE_DISPATCH=batched as the off arm):

• test-backend-ops -o MUL_MAT: supported 373 -> 526, the six new real-shape rows flip not supported -> OK except the unaligned k=232 one, which stays refused by design.
• Zamba-7B-v1: 241 -> 89 nodes off HRX (MUL_MAT 152 -> 76, CONT 76 -> 0). Graph splits unchanged at 179 — stated in the body rather than buried, since splits are now driven by the remaining nodes.
• pp512 63.60 ± 6.00 -> 79.97 ± 1.82, tg64 5.89 ± 0.66 -> 6.67 ± 0.05 (-r 5), PPL 20.0688 -> 20.0718 (+0.015 %, each arm reproduced to 1e-4 across runs).
• The 76 nodes left are dt_proj at K % 32 == 8: the padded-tile item, not this PR. GGML_HRX_STRICT=1 still aborts on this model, so no strict-clean claim.

[bong-water-water-bong]

VERDICT: READY at 1f276843 on the code — the broadcast case is handled where it would have gone wrong, and the acceptance gate refuses the combinations it cannot express. No device run from me.

Read the two places that can be wrong. The acceptance is precise: output planes must equal input planes, and the weight must be either a single matrix (broadcast) or one matrix per plane (paired) —

const int64_t plane_count        = input->ne[2] * input->ne[3];
const int64_t weight_plane_count = weight->ne[2] * weight->ne[3];
… weight_plane_count == 1 || weight_plane_count == plane_count …
match.plane_count    = allow_batched ? plane_count : 1;
match.batched_weight = allow_batched && paired_weight;

and the emission gets the one case that silently corrupts results right:

const size_t input_offset  = common_plane_offset(*match.input,  plane);
const size_t output_offset = common_plane_offset(*match.output, plane);
const size_t weight_offset = match.batched_weight ? common_plane_offset(*match.weight, plane) : 0;

A broadcast weight keeping offset 0 on every plane is exactly the thing a per-plane loop gets wrong first, and plane_count is ne[2]*ne[3] with ggml's absent dims as 1, so a plain 2-D matmul takes the same path with plane_count == 1 rather than needing a separate branch. common_plane_offset decomposes the plane as ne[2]-major with nb[2]/nb[3] strides, which matches ggml's row-major batch enumeration.

The new test-backend-ops cases are the right instrument — the batched shapes are where this is either right or obviously wrong.

Not verified by me: the device. Running the matrix needs an engine build against this pin on the HRX route, and the GPU queue has been 8–9 deep all night; PR ggml-org#1104 (the bump that would carry this into an engine build) is still open. If you want the row, the moment is right after ggml-org#1104 merges with a hw-gate label.