ggml-hrx: MUL_MAT_ID at input sizes that are a multiple of 32 (gpt-oss experts on HRX) + decode input row stride

[bong-water-water-bong]

gpt-oss-20b's MXFP4 experts now run on HRX. Their input size is 2880, a multiple of 32 but not of 256, as is BlackMamba's 1152. On gpt-oss: pp512 25.8 → 527 tok/s, tg128 12.6 → 24.8 tok/s, with correct text.

Commits

1. MUL_MAT_ID WMMA kernels take input sizes that are a multiple of 32. This follows the pattern mul_mat_id_swiglu_f32_f32_wmma already uses.
   • In AMD's shared motifs/mul_mat_id_f32_f32_wmma_core.loom: loop over whole 256-value tiles with the tail masked. Weights past input_size decode as zero; activations use vector.load.mask. The row stride comes from ggml_dequant_weight_row_bytes.
   • That's about 20 lines, plus mul(32) in the plain, postops and postops_next_rmsnorm configs. Their file headers are unchanged.
   • This relaxes ggml-hrx: MXFP4 weights (shared dequantizer, exact E8M0 scale, known-answer test) #65's guard: the MUL_MAT_ID matcher goes back to the per-format dense rule. ggml-hrx: MXFP4 weights (shared dequantizer, exact E8M0 scale, known-answer test) #65 required input_size % 256 there, which also turned away the decode and gated kernels; both already took multiples of 32.
2. Decode kernels read input row t at t * input_size.
   • @ggml_mul_mat_decode_load_f32_block in AMD's ops/mul_mat_f32_f32_decode.loom viewed the input as [token_count] x [ceil(K/256)] x 256. For K % 256 != 0, every token or input row after the first was read from the wrong offset: rows 3072 apart instead of 2880.
   • The fix indexes the input as [token_count] x [input_size] (4 loads); the existing masks cover the partial tile.
   • Not reachable on the current pin (cebcd70): the dense decode kernels take one token (common_is_supported_decode_token_count is == 1), and MUL_MAT_ID at K % 256 != 0 isn't admitted there because of ggml-hrx: MXFP4 weights (shared dequantizer, exact E8M0 scale, known-answer test) #65's guard. It becomes reachable with commit 1, where mul_mat_id_decode_f32_wave64 gets input_rows 4 (the down projection) or 2-4 tokens at K = 2880. That's why it ships here rather than on its own.
3. tests/test-hrx-decode-stride.cpp (new):
   • MUL_MAT at 2/3/4/8 tokens and MUL_MAT_ID at input_rows 1/4 with 1-2 tokens;
   • K = 2880, 1152 and 2048 (control); MXFP4, Q8_0, Q4_0;
   • 32 experts, 4 routes, against the CPU backend; every case must be planned on HRX.
4. tests/test-hrx-mul-mat-id-k32.cpp (new): MUL_MAT_ID at K = 2880, 1152 and 2816 (no tail); MXFP4, Q8_0, Q4_0; 1 token (decode kernel) and 7/40 tokens (WMMA); against the CPU.

Before (commit 1 without commit 2, HRX0):

• A gpt-oss-shaped probe on mul_mat_id_decode_f32_wave64 at K = 2880 gave NMSE 20-55 with input_rows 4 or 2 tokens, for MXFP4 and Q8_0 alike. K = 2048 was fine for every format; at 2880, 1-row single-token cases were fine.
• gpt-oss greedy degenerated: "The sky is blue in the sense that the sky is a blue color. The sky is a blue. The sky is a blue. ..."

After, on this branch's own build (HRX0, balanced power mode):

• test-hrx-decode-stride 72 cases, 0 failures; test-hrx-mul-mat-id-k32 all OK; test-hrx-mxfp4 bit-exact.
• test-backend-ops MUL_MAT 322/322 and MUL_MAT_ID 108/108, three runs each; full suite 1051/1051.
• MUL_MAT_ID check cases (iree-test-loom, hrx-system 244cd38, --sanitizer=access, MXFP4): K = 1152 exact, plus a random differential against a scalar reference at 2880 and 1152. Expert 1 is the last expert, and the cases cover its last row and the last token, so an unmasked tail would read past the buffers; the sanitizer found nothing.
• gpt-oss-20b greedy is correct:
  • sky: "The sky appears blue because the Earth's atmosphere scatters sunlight..."
  • colors: red, green and blue;
  • haiku: "Snowflakes drift, hush, / Frosted breath on silent trees— / Winter's quiet breath."
  • Greedy repeats are identical.
• llama-bench pp512 527 ± 26, tg128 24.8 ± 0.5 (the earlier hrx-tq run). On this tree: 404 ± 137 and 23.9 ± 2.2, with other jobs on the box. Before: 25.8 and 12.6 with the experts on the CPU.

Not changed here

• KLD vs CPU is 0.028 with or without this change. Perplexity doesn't touch the decode path.
• It looks like the f16-accumulation drift in the WMMA cores, tracked in engine bugfix: centos 7, gcc (GCC) 11.2.1 20220127 (Red Hat 11.2.1-9) ggml-org/llama.cpp#284 (bcloud-1b).
  • With the iree-test-loom constant case above, the unchanged MUL_MAT_ID kernel already gives 77.3 instead of 66 at K = 2816. Don't read that as a tail problem.
• gpt-oss still has 193 graph splits: ADD_ID, SWIGLU_OAI and FLASH_ATTN with sinks run on the CPU. That's next.

Pinned models with sizes that are a multiple of 32 but not of 256 (scan of 72 local GGUFs): gpt-oss-20b (2880), BlackMamba-1.5B (1152), Zamba-7B (3712), qwen2.5-0.5B (896), and Qwen3.8-Flash-Next expert FFN 640. On the pin, none of them reach the decode stride bug. ZAYA1-8B is all multiples of 256.

🤖 Generated with Claude Code

[bong-water-water-bong]

Review (PR-Agent duty): approve.

• g32 (AMD's MUL_MAT_ID WMMA core). The K loop pads to whole 256 tiles. Weight loads are guarded by k < input_size, and activation loads are masked with vector.mask.range, so the tail contributes zero. The row stride comes from ggml_dequant_weight_row_bytes, so it's right for 32-value block formats. Only the three configs change, mul(256) to mul(32). The guard from ggml-hrx: MXFP4 weights (shared dequantizer, exact E8M0 scale, known-answer test) #65 is relaxed back to the per-format dense rule.
• Stride fix (AMD's decode loader). The input is viewed as [tokens]x[input_size] and indexed [token, k] with the existing masks, so row t starts at t * input_size, not t * ceil(K/256)*256. This bug is unreachable on cebcd70: dense decode is token_count == 1, and MMID at K % 256 != 0 wasn't admitted there. g32 makes it reachable, which is why both land together.
• Evidence.
  • Before: round-3 probe NMSE 20-55 at rows 4 / tokens 2, and degenerate gpt-oss text.
  • After: test-hrx-decode-stride 72/72, test-hrx-mul-mat-id-k32 OK, MUL_MAT 322/322 and MUL_MAT_ID 108/108 x3, full suite 1051/1051, test-hrx-mxfp4 bit-exact. gpt-oss-20b answers correctly with identical greedy repeats; pp512 25.8 -> 527, tg128 12.6 -> 24.8 (round 4, balanced mode).
• Edits stay in AMD's files with no notice change; our tests are in our own files.

KLD 0.028 vs CPU is pre-existing and tracked in engine ggml-org#284 (f16 accumulation). Merging when the hosted jobs pass, then pinning.