metal : skip the empty half of the mul_mm_id token tile, load iq2/iq3 codebooks as uint32

[masterFoad]

Overview

This makes kernel_mul_mm_id skip the upper half of the token tile when an expert does not fill it, and reads iq2/iq3 codebook entries as whole words instead of individual bytes.

In MoE prefill, each mul_mm_id threadgroup multiplies a 32-row weight tile against a 32-row token tile, but an expert often receives far fewer than 32 tokens per ubatch (ne21*ne20/ne02: 8 for a 256-expert top-8 model at -ub 256, and 2 to 4 for a 128-expert top-8 model at n=32..64). The padded rows still incur work.

The patch splits the token tile into two NR1H = 16 halves and skips the upper half when nr1 <= NR1H. On the tensor-ops path, this uses one matmul2d call per half. On the simdgroup path, the two simdgroups responsible for rows 16..31 idle through sg_active. Padding rows in the B tile are also zeroed once instead of being restaged from a clamped row on every k step.

I noticed this while profiling prefill on Tiel-Coder-35B-A3B, where the expert matmul accounted for around half of prefill time while most expert tiles were only partially populated.

The dequantize_iq2_xxs/iq2_xs/iq3_xxs/iq3_s/iq2_s functions now read each codebook entry using one or two uint32 loads instead of 4 to 8 uint8 loads. The ksigns_iq2xs lookup is also replaced by the identity encoded by that table, k | (parity(k) << 7). The result is bit-exact. This relies on little-endian layout, which covers Metal targets. It is the same general idea as #27370 for q8_0.

This part is independent of the tile change and I can split it into a separate PR if that is easier to review.

Also included: the tB extents were (NR1, NK) for a [NR1][NK] row-major tile, i.e. reversed. This was harmless while NR1 == NK == 32, but needed fixing once the tile was split into (NK, NR1H) halves.

Tested on my 14inch MacBook Pro M5 24 GB using the tensor-ops path, and on my 16inch MacBook Pro M1 Pro 32 GB with a 16-core GPU using the simdgroup path. On the M1 Pro, Metal reports tensor API disabled for pre-M5 and pre-A19 devices.

On the M5, using the same fixed coding-agent replay as #28086 (Tiel-Coder-35B-A3B IQ3_XXS, b10488 tree), cold 5628-token prefill improved from 455 to 508 tok/s (−10.4% time), warm suffix prefill improved by about 10%, decode by 3 to 4%, and perplexity remained unchanged at 3.3969. Time to first token on that prompt dropped from 12.4 s to 11.1 s. The full 13-turn session improved by about 3% wall-clock because that workload spends roughly 70% of its time in decode.

On the M1 Pro, llama-bench on Qwen3-30B-A3B UD-IQ3_XXS improved pp4096 from 298.5 to 324.1 tok/s at -ub 256 (+8.6%), and from 346.8 to 362.2 tok/s at -ub 512 (+4.4%). tg64 was unchanged at 45.9 tok/s, and perplexity remained unchanged at 8.4328. The kernel-level gain below is roughly twice the end-to-end gain because the expert matmul accounts for less than half of prefill time on that model.

Additional information

Who this helps. The half-tile skip fires for any (expert, tile) with 16 or fewer routed rows. Average rows per expert are ubatch * n_expert_used / n_expert, so 256-expert top-8 models such as Qwen3.6-35B-A3B and the DeepSeek-V3 family are in this regime on every tile at -ub 256. 128-expert top-8 models such as Qwen3-30B-A3B and GLM-4.5-Air hit it on partial last tiles at -ub 512 and on most tiles at -ub 256. Routing skew increases the number of underfilled tiles relative to the uniform average. The dequantization change applies to iq2/iq3 models generally, both dense and MoE.

Where the gain lands. On the simdgroup path used by pre-M5 Apple GPUs, all four simdgroups run simdgroup_multiply_accumulate over the full tile regardless of nr1. Idling the upper pair therefore removes roughly half of the matrix work on a half-empty tile.

On the M5 tensor-ops path, the padded rows were already relatively cheap and the kernel is more sensitive to occupancy and A-tile dequantization, which is what the dequantize.h change targets. In practice, the simdgroup path benefits across quant types, while the tensor-ops path sees the larger gain on iq2/iq3 and is approximately neutral on q4_0.

Upstream test-backend-ops perf cases, q4_0, M1 Pro 16-core GPU, upstream/master = 7798007, min of 3:

case	master (µs)	PR (µs)	delta
128 experts top-8, m=768, n=8	301.4	298.4	−1.0%
n=32	2703.9	1696.0	−37.3%
n=64	3192.1	1961.6	−38.6%
n=128	3258.6	2023.3	−37.9%
n=256	3311.9	2622.1	−20.8%
n=512	4829.7	4292.9	−11.1%
32 experts top-4, m=1792, n=8	339.9	339.8	0.0%
n=32	1832.9	1135.0	−38.1%
n=64	1942.4	1204.1	−38.0%
n=128	1964.1	1553.7	−20.9%
n=256	2853.8	2506.2	−12.2%
n=512	4816.2	4471.6	−7.2%

n ≤ 8 is effectively unchanged within ±1%. The gain tapers as tiles fill. For the 32-expert top-4 case, n=128 averages 16 rows per expert, exactly the half-tile boundary, which is why it behaves more like the n=256 cases than the n=32..64 cases.

Same cases on the M5, tensor-ops path, single pass:

case	master (µs)	PR (µs)	delta
128 experts top-8, n=32	1118.0	1104.6	−1.2%
n=64	1293.4	1266.5	−2.1%
n=128	1387.2	1374.7	−0.9%
n=256	1514.5	1528.5	+0.9%
n=512	2033.5	2022.2	−0.6%
32 experts top-4, n=32	768.2	800.3	+4.2%
n=64	815.2	804.8	−1.3%
n=128	862.6	875.3	+1.5%
n=256	1189.0	1170.3	−1.6%
n=512	1978.3	2034.3	+2.8%

These q4_0 tensor-ops results are within run-to-run noise on this setup (about ±5% for a single pass). Forcing the simdgroup path on the M5 with GGML_METAL_TENSOR_DISABLE=1 produces the same shape as the M1 Pro results, with slightly larger gains: about −42% at n=32..128.

iq3_xxs / iq2_s at the production shape, M5 tensor-ops path, master = 0eadefe, min of 3 (the 256-expert top-8 ffn_*_exps shape used by Tiel, with 8 rows per expert at n=256): iq3_xxs improves by 15.6% (2223.6 → 1877.7 µs), and iq2_s by 10.2% (2115.8 → 1900.8 µs). Smaller batches improve more: iq3_xxs is −21.4% at n=64 and −17.6% at n=128. Upstream has no perf cases for these types, so these measurements use the same case shapes with type_a changed.

Trade-off to be aware of: at large n on the tensor-ops path, the half-tile skip no longer triggers and splitting the operation into two calls adds a small cost for iq2_s: n=1024 is +3.6% and n=2048 is +7.0%. iq3_xxs remains neutral there. Default ubatch sizes are in the regime where this patch wins.

I also tried gating the split so full tiles use a single full-size matmul2d call. That avoided the large-n two-call case in isolation, but was about 3% slower in real prefill for the reason below, so I left the current implementation as-is. I can revisit this if the large-n iq2_s regression is considered more important.

Rejected alternative: a second full-NR1 matmul2d descriptor for full tiles, allowing one call instead of two, was neutral in the microbenchmark but about 3% slower in real prefill. Keeping cooperative tensors of both shapes live raises the accumulator count per thread from 16 to 24 floats and appears to reduce occupancy.

Benchmark harness: stock test-backend-ops perf draws the expert IDs once and then re-times the same expert set, so MoE measurements quickly become warm-cache. This PR adds a reinit_perf_iter hook to test_case, called between timed iterations and never inside them, which test_mul_mat_id uses to redraw the IDs. This is the rotating-ID change mentioned in #28086. Every kernel number above uses the same rotating-ID behavior on both the master and PR sides. I can split this into a separate PR if that is easier to review.

Test methodology

• test-backend-ops test -o MUL_MAT_ID and -o MUL_MAT against the CPU backend: 810/810 and 1221/1221 on master = 7798007, tested on the M5 using both the tensor-ops path and GGML_METAL_TENSOR_DISABLE=1, and on the M1 Pro.
• Kernel timings: test-backend-ops perf -b MTL0 -o MUL_MAT_ID -p type_a=q4_0, using the same binary modulo the kernel files, with the rotating-ID harness enabled on both sides. All trees were confirmed with CMAKE_BUILD_TYPE=Release.
• End-to-end M5: fixed replay of a multi-turn coding-agent transcript, 13 generations, ~5.6K-token system prompt, mean of 3 warm repetitions, plus a cold 5628-token prefill probe with cache_prompt: false. Same server configuration as metal : fix idle threads in mul_mv_iq3_xxs for ne00 < 1024 #28086.
• End-to-end M1 Pro: llama-bench -m Qwen3-30B-A3B-UD-IQ3_XXS.gguf -p 4096 -n 64 -ub 256,512 -fa 1, with master and PR built using the same flags.
• Perplexity: llama-perplexity ... -c 2048 -b 2048 --chunks 4 -ngl 999 -fa on -ctk q8_0 -ctv q8_0, unchanged at 3.3969 on Tiel on the M5 and 8.4328 on Qwen3-30B-A3B UD-IQ3_XXS on the M1 Pro.

Related: #25377 (64x8 dense mul_mm tile for small batches, same diagnosis on the dense kernel), #27370 (packed q8_0 dequant), and #26223 (touches the same kernel's B-tile staging; whichever lands second will need a small merge).

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES, AI was used during the implementation and profiling, guided and reviewed by me.

[ggml-gh-bot]

Hi @masterFoad, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• Multiple open PRs from a new contributor: We limit new contributors (those without a previously merged PR) to 1 open PR at a time. You currently have 3 open PRs.

• AI-generated content: While code is allowed to be generated by AI, please write the PR description and commit messages on your own without the help of AI.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.