[BUG] qmm_splitk is chosen over qmm_nax on M5 for N=4096 projections M=16-64 #4720

☑️ I understand it is strictly prohibited to use AI to write issues.

Describe the bug

• For M at or above the qmv cutoff, with transposed weights and a non-batched call, eval_gpu always calls qmm_splitk.
• qmm_splitk uses 32x32 blocks and a plain Metal kernel, not NAX. It picks split_k = max(1, 512 / (N-tiles * M-tiles)).
• It falls to qmm, and so qmm_nax, only when split_k is 1:
  • gate/up (N = 14336): split_k = 1, so it uses qmm_nax
  • down and q/o (N = 4096): split_k is 2-4, so they use split-K
  • k/v (N = 1024): split_k is up to 16
• A split-K PR ([Metal][Performance]: Add split-K for quantized matmul (small M) #3120) was benchmarked on an M3 Max, which has no NAX.

To Reproduce

import time
import mlx.core as mx

K, N = 14336, 4096  
ws = [mx.quantize(mx.random.normal((N, K)).astype(mx.float16), group_size=64, bits=4) for _ in range(8)]
mx.eval([t for w in ws for t in w])
wbytes = N * K * 0.5 + N * K / 64 * 4  

for M in (8, 16, 32, 36, 48, 64):
    x = mx.random.normal((M, K)).astype(mx.float16)
    run = lambda: mx.eval([mx.quantized_matmul(x, *w, transpose=True, group_size=64, bits=4) for w in ws])
    run(); run()
    ts = []
    for _ in range(9):
        t = time.perf_counter(); run(); ts.append(time.perf_counter() - t)
    print(f"M={M:3d}  {wbytes / (min(ts) / len(ws)) / 1e9:6.0f} GB/s")

Expected behavior

• On a NAX applegpu, down (14336->4096) and q/o (4096->4096) for M = 16-64 should be as fast as the same call going through qmm_nax:
  • forced qmm_nax on this M5 Pro: down 131 / 120 / 90 / 91 GB/s at M = 16 / 32 / 48 / 64, q/o 104 / 100 / 78 / 73
  • the current dispatch gives down 55 / 56 / 28 / 29 GB/s and q/o 60 / 45 / 25 / 25 at the same M
  • for comparison, the gate/up shape (4096->14336), which does reach qmm_nax today, gets 143-155 GB/s for M = 16-32
• Whole-model step time should not roughly double between M = 32 and M = 36. Qwen3-8B 4-bit, real tokens, 1,024-token cache:
  • current: 57.8 ms at M = 32, 96.7 ms at M = 36, so tokens/s falls from 554 to 372 while batch grows
  • with split-K skipped and the qmv cutoff at 8: 39.2 ms at M = 32, 53.0 ms at M = 36
• M = 1-4 should stay as it is unaffected by the test, and qmv is the better path there.
• The qmv/qmm cutoff should be near where qmv_wide and qmm_nax actually cross. On 4096-class shapes, that is about M = 6-8.

Desktop (please complete the following information):

• OS Version: macOS 27.0.1
• Hardware: Apple M5 Pro 20 GPU cores 48 GB
• MLX version (which numbers came from which build):
  • git HEAD b6642e60 (0.32.5.dev20261017), built from source
  • mlx-lm: 0.31.3 (used for the Qwen3-8B whole-model runs)

Additional context

Related history:

• qmm_splitk's 512-threadgroup target routes M<=32 transposed qmm to the non-NAX split-K kernel on M5 (M=33 is 21% faster than M=32) #4198 split-K displacing the NAX qmm path on M5 Max, single op N = 6656 and a 30B-stack. Closed with "...currently we are only focusing on real cases happened during inference".
• [Metal][Performance]: Add split-K for quantized matmul (small M) #3120 added split-K, benchmarked on an M3 Max (25-30% gain there, no NAX).
• Raise qmv batch limit for large matrices on M5-class GPUs #3791 generation-17 qmv cutoffs.
• Add Neural Accelerator Support #2772 NAX support.

I used Claude to help me root cause the diff here. Close this issue if you don't like that. I can include more reproducers/comparisons if you want.