Repository navigation
Larger M tile for the NAX affine quantized matmul during prefill #4692
Description
On M5 chips, prefill of quantized models runs quantized_matmul with hundreds to thousands of rows (M). For M > 64 the NAX kernel uses a 64x64x64 tile. A 128x64x128 tile is about 5% faster on most of these shapes, and 4 to 10% faster end to end.
Measurements (M5 Pro)
Per call, bf16, 4 bit, group size 32. x is M x K, the weight maps K to N:
| K to N | M | old (ms) | new (ms) |
|---|---|---|---|
| 5120 to 17408 | 1024 | 7.29 | 6.72 |
| 5120 to 17408 | 4096 | 30.84 | 28.36 |
| 17408 to 5120 | 4096 | 29.42 | 27.46 |
Over a sweep of about 2270 shapes and types, the rule below picks the new tile in 1283. Their median ratio (new / old) is 0.948 and the worst point 1.050, well within the noise of these small kernels.
End to end, mlx_lm.benchmark -p 4096, prompt tokens per second, two repetitions:
| Model (4-bit) | prefill step | gain |
|---|---|---|
| Llama-3.1-8B-Instruct | 2048 | +6.5%, +5.7% |
| Qwen3-8B | 2048 | +4.1%, +3.6% |
| Qwen3-32B (hidden size 5120, as in the shapes above) | 2048 | +9.9%, +7.0% |
The first 32B repetition started with two slow stock trials, so +7% is the safer figure. Peak memory is unchanged.
Proposal
Use the 128x64x128 tile in qmm_nax for affine, transposed, fp16/bf16 weights when K % 128 == 0, K >= 512, M > 64 and M % 128 == 0 || M % 128 > 64 || M >= 2048. The last condition avoids M just above a multiple of 128, where the larger tile is mostly padding and was up to 1.29x slower. Everything else keeps the current path.
- Bit-identical to the current tile: the accumulation order over K does not change.
- BK=128 needs its own block loader, used only when BK > 64; existing tiles keep the current loader.
- Cost: 168 new kernel instantiations,
mlx.metallib+6 MB (+2.8%).
Open question on testing
Correctness is tested against a float32 reference and against the default tile. The new tile is bit-identical to the old one, so a test cannot see from the output which kernel ran. Is that enough for you, or would you like a routing check (for example a small read-only has_kernel(name) on metal::Device for a C++ test)?
Related
#4204 tried a related change for M >= 4096 on non-NAX devices.
Branch: jvmenen/mlx@main...qmm-large-m-tile-v2
Would a PR along these lines be welcome?
Activity
jvmenen commentedon Oct 21, 2026
Follow-up now that #4682 is on main. In short runs its staged
qmm_t_naxkernel beats the tile proposed here. In long prefills it does not.Setup: M5 Pro (20-core GPU, 64 GB), Qwen3.8-27B (4-bit body with group size 32, some 8-bit layers), bf16 activations.
gandbetain float32 ingated_delta_updatewhen the activations are bf16 #4725). Each variant runs in its own process, in the order A B B A.Isolated matmul, M = 4096, chained calls, about 1.2 GB of weight copies (ms per call):
Each cell covers four runs. The B column leaves out one process that ran about 12% slow on every shape (for example 26.3 ms on gate/up 4-bit). Output is identical across A, B and v0.32.4. Here #4682 is 4 to 15% faster than this tile.
Sustained load, about 60 s of an MLP-shaped loop (ms per iteration):
End to end, cold prefill on an mlx-lm based server (time to first token, s):
In-process without the server, the 32K prefill took 72.6 / 73.2 s on A and 75.8 / 75.5 s on B. Per module over that prefill, B is 4-6% slower in the MLP, the linear attention and the full attention alike. B's 50K runs reached thermal pressure "moderate"; A's did not.
My reading, not yet measured: the staged kernel draws more power and reaches the sustained power or thermal limit sooner. Short benchmarks then favour it, while prefills longer than about 10 s end up slower. I can add
powermetricsnumbers if that helps. I closed this yesterday on the short-run numbers alone and have reopened it: for long prefills the larger tile is still the faster option on this machine.zcbenz commentedon Oct 23, 2026
We don't consider GPU throttling when benchmarking kernels.