metal : add 64x8 mul_mm tile variant for small batches (q4_0)

[discobot]

Overview

mat-mv kernels are compute-bound and scale ~linearly with the batch size, while the 64x32 mul_mm tile zero-pads 75% of the batch dimension at bs = 8.

Route q4_0 x f32 to a 64x8 tile for bs 5..16:

• new kernel_mul_mm_nr8 (64x8 tiles, all 4 simdgroups along the src0 rows)
• q4_0 hands over from mul_mv_ext to mul_mm at bs 5 instead of 9

bs = 1 decode unaffected

Performance (M4 Pro 64GB)

test-backend-ops perf, MUL_MAT q4_0 x f32, 4096x14336:

bs	master us/tok	PR us/tok	speedup
4	131.6	131.6	(keeps mv_ext)
5	155.4	97.7	1.59x
8	127.9	63.3	2.02x

llama-bench, gemma-4-26B-A4B q4_0: pp5 -8%, pp8 -19%, pp13 -8%, tg unchanged.

Correctness

test-backend-ops -o MUL_MAT -b MTL0: 1134/1134 vs CPU (includes q4_0 at bs 2..16 routed through the new kernel).

Additional information

Closes #25250
this problem encountered in #24768
Only tested on M4 Pro.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES - I use cursor with claude

[ggml-gh-bot]

Hi @discobot, thanks for your contribution!

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

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

---

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

[discobot]

updated to match the template, rebased on top of fresh main

[discobot]

this enables more efficient speculative decoding on metal

[discobot]

@ggerganov friendly ping

[forforever73]

Thanks a lot for tackling this :) the small-batch waste in the 5–16 range is a real pain point and a thin tile is a clean way at it. My main hesitation is maintainability, each new tile ends up being a near-copy of mul_mm gated by another hardcoded predicate like ne1<=16, and those predicates don't really generalize — e.g. I tried the same thin tile on f16 + long K and it was slower than 64×32 at small n, so the q4_0-tuned cutoff wouldn't carry over.

So I've been poking at a PoC that makes the tile a template parameter picked from a small shape-indexed table, instead of a new kernel per shape — the optimal tile turns out to be continuously tunable. It's basically the same approach as #25750 , just applied to mul_mm.

[discobot]

Love your approach, and agree that maintaining this zoo would be hard.
Let me benchmark your solution vs this one.

[discobot]

@forforever73
benchmarked both branches on M4 Pro, same session, r=3.

q4_0 4096x14336, us/token:
bs	this PR	tile-poc
4	127.8	129.4
5	97.7	 159.5
8	66.6	 131.1

f16 4096x14336, us/token:
bs	this PR (= master here)	tile-poc
4	148.6	                   127.9
5	171.8	                  154.4
8	154.4	                  144.0

End to end, gemma-4-26B-A4B q4_0: pp8 44.7 ms this PR vs 56.7 tile-poc.

so you win on f16, I win on q4_0 small bs, the poc tunes the tiles but keeps the old mv_ext -> mm.
for q4_0 at bs 5..8 it still goes down the mv_ext path, so your tuned tiles never run there.
moving the switch is where most of the win comes from

so best possible solution seems to be your table + per-dtype switch point + thin q4_0 rows

[forforever73]

In the parameterized version, those three naturally fold into one. The thin q4_0 tile is simply one row in the table, and the switch point (ne11_mm_min) is another per-dtype parameter (like use_nr8), so both become part of the same shape-indexed lookup rather than separate knobs. PoC only tuned the tile while leaving the switch unchanged, which is why q4_0 at bs 5–8 never reaches the tuned tiles.

I've already started working on this and hope to have a draft PR up in a few days. Happy to iterate on it together if you're interested.

[discobot]

@forforever73

thank you, let me know if I can be of any help
eagerly waiting for this, because it unblocks speculative decoding for gemma on macs

[ggerganov]

I'm just noticing this PR. Will take a look.