Repository navigation
Fix non transposed affine qmm dispatch logic - #4392
Conversation
|
Looks good to me! We probably have the same issue in Regarding mlx-lm, ml-explore/mlx-lm#1781 proposed to add a prefill size as an argument, so it is more controllable now! |
Yep! I wonder if I'm missing something too. Now that you pointed out, this was a fix for a bug introduced in #4202 which the author for #4051 mentioned as a comment on that PR. Also now that i look into this more, the qmm_n_nax kernels for non-affine quantisations aren't instantiated, which is why the |
|
@RohanGautam There is no active works on NAX kernels, the main reason I was hesitated about enabling it was because I was not sure if I would get time to look into issues if anything goes wrong, if you are confident to look into possible followup issues, then let's merge. |
…sses v0.32.2 (2026-08-25) is the newest tag, and it predates several correctness fixes in the quantized matmul path the Gemma 4 packed-weight work builds on: the non-transposed affine qmm dispatch (ml-explore/mlx#4392), fp qmm corruption when the quantized dimension is not a multiple of 32 (#3912), and two sorted `gather_qmm` fixes (#4009, #3922). It also predates D512 vector attention (#4459) and the fused NAX head-dim-256 prefill path (#4416), which cover Gemma 4's full-attention and sliding layers respectively. `gather_qmm` gained a `global_scale` parameter between the mode and `sorted_indices`, so the positional call bound a bool to `std::optional<array>` and stopped compiling. That argument and `sorted_indices` are named now, since the next insertion would otherwise bind silently. No speedup, and none was available: decode on a 12B is memory-bandwidth bound. 51.8 ms/token measured against a 44.0 ms roofline for 24 GB of bf16 weights on an M4 Max is 85% of theoretical, so no kernel change can move it. Reading 6.3 GB of packed 4-bit instead is the only lever, which is #27. Verified: full preflight, plus Llama-3.2-1B, Qwen2.5-0.5B, gemma-2-2b and Phi-3-mini at 4-bit and gemma-4-12B at bf16 and 4-bit all still decode. Refs #27
3 participants
The path for using the non-transposed NAX
qmmkernel was unreachable in the current implementation (because of the conditionif (has_nax_kernel && transpose && (K % 64 == 0) ...- lets only the transposed case through).We needed to add some more conditions:
benchmarks
The benchmarks are run with
mlx_lmlinked to the different mlx versions. The non-transposed path is exercised for large prefills (when chunked), for the quantized kv cache computations. We usemlx_lm.generate(notbenchmark, to use the kv cache quantisation options) . The command run ismlx_lm.generate --model <model> -p - --ignore-chat-template --max-tokens 8 --temp 0 --kv-bits 4 --quantized-kv-start 0 < <file withprompt of desired size>.txtThe 4096 one has no changes because of
mlx-lmusing N-1 tokens instead of N tokens. There is only one quantised kv cache to attend over, and this is the wrong size for the non-transposed qmm to kick in (it does not have a ragged-K variant). The last chunk for all others also go through the normal path, just the middle ones don't and are sped up with this PR. This boundary effect can be seen below, and is perhaps something to address in mlx-lm down the line.