Repository navigation
[BUG] test_qmm_splitk_precision fails on Apple silicon generation 17 or later (M=16) #4676
Description
☑️ I understand it is strictly prohibited to use AI to write issues.
Describe the bug
On Apple Silicon generation 17 or later, the test case test_qmm_splitk_precision fails for affine weights, M=16, with bf16 and fp16, in an unmodified build of MLX.
To Reproduce
On an M5/M6-class machine, run from python/tests:
python -m unittest test_quantized.TestQuantized.test_qmm_splitk_precision
The bug will not reproduce on M1/M2-class CI (limit 10, so M=16 takes split-K path and test passes)
Expected behavior
Test passes in an unmodified build, independent of the GPU generation.
Desktop (please complete the following information):
- OS Version: macOS 27.0
- Version: built from source at 77bf1fa
- Xcode: 27
- Deployment target: 27
- Devices: M5 Air, M6 mini
Additional context
This is a test-scope issue, not a kernel regression. The batched-qmv route dequantizes weights in fp32 while the test's reference is fp16/bf16-dequantized, so qmv compares as ~1.5 against a bound calibrated for the split-K path. With an fp32-dequantized reference the numbers invert: qmv 1.000, split-K ~1.5. The test was added in PR #4641.
The qmv path is selected when M is below get_qmv_batch_limit at K=2560, N=1024; split-K runs at or above it (QuantizedMatmul::eval_gpu).
| GPU | Limit at K=2560, N=1024 |
|---|---|
generation 17 or later, not d (M5, M6) |
25 |
generation 15 and 16, not d |
15 |
generation 13 and 14, not d |
10 |
generation 13 or later, d |
18 |
generation below 13, d / not d |
18 / 12 |
With M=16, the split-K path is not selected on generation 17 or later. M=25 is at or above every limit in the table, so the split-K path is always selected.
This is a one line fix of the M list: [16, 33, 65] -> [25, 33, 65], or an alternative fix is K,N > 4096 so the limit drops to 13/12.
Patch for adjusted M is prepared and will follow as PR.
Measurements on M5 Air (M5, applegpu_g17g, macOS 27.0, unmodified 77bf1fa01 build)
| Result | |
|---|---|
| Test as in upstream, before | fails at affine M = 16, bf16 and fp16 (0.0017171 against 0.0012368; 0.00021316 against 0.00015514): FAILED (failures=2) |
| Test with the change, after | OK, default environment and MLX_ENABLE_TF32=0 |
| Error over one rounding, affine, M = 1 / 8 / 16 / 24 | 1.50 / 1.52 / 1.53 / 1.51 (bf16); 1.55 / 1.52 / 1.51 / 1.52 (fp16) |
| The same, M = 25 / 32 / 33 / 65 | 1.000 in every cell, both dtypes |
| mxfp4, M = 1 ... 65 | 1.000 in every cell |
Measurements on M6 (M6, applegpu_g18g, macOS 27.0, same build and patch hash) before FAILED (failures=2), after OK in both environments, the ratio table the same. All 32 error and rounding figures are equal to the last digit on the M5 and the M6 (the quantized matmul is deterministic and runs the same arithmetic on both here).
Disclosure
I wrote this report and AI tooling assisted with the measurements and with preparing the patch.
Activity