[wyanzhao]

Fixes #4498.

QMV currently falls back to the generic kernel when the output size is not divisible by 8, even when the input meets the qmv_fast alignment requirement. This is particularly expensive for narrow, dependent projections such as the N=1 and N=4 calls in autoregressive decode.

This PR adds affine_qmv_fast_rows and fp_qmv_fast_rows. In the final SIMD group, reads beyond the output tile reuse the last valid weight row, and stores are limited to valid rows, leaving the reduction loop unchanged. The new paths are used when qmv is selected for an aligned input whose output size is not divisible by 8. They cover affine, MXFP4, MXFP8 and NVFP4, including NVFP4 calls with a global scale. Gathered QMV keeps its existing kernels. The extra instantiations add about 1.3 MB to mlx.metallib, including 0.31 MB for the FP variants.

While testing the new path, I found that the affine bias term accumulated sum(x) in the input type before widening it. That loses small contributions for float16 and bfloat16 inputs. The first commit widens each value before accumulation for 2- through 6-bit quantization; the 8-bit path already did this.

Performance

Measured on an Apple M5 Max with bfloat16 inputs. Baseline uses the generic kernel and Candidate uses affine_qmv_fast_rows; both include the precision fix. The ratio is the geometric mean of paired Baseline/Candidate time ratios from A/A-calibrated runs, so values above 1 favor Candidate.

Workload	Baseline	Candidate	Ratio	95% CI
Chained QMV: N=1, K=2560, 8-bit, group size 64	11.37 µs/call	8.22 µs/call	1.3793	[1.3559, 1.4030]
Qwen3.8-Flash-Next layers 4–7, context 128	1.428 ms/step	1.185 ms/step	1.2040	[1.1972, 1.2108]
Qwen3.8-Flash-Next layers 4–7, context 4096	1.526 ms/step	1.287 ms/step	1.1872	[1.1817, 1.1928]

For the model-shaped benchmark, I used layers 4–7, one complete four-layer cycle with three GDN layers followed by one QSA layer. The benchmark configuration cannot run the full model, so I did not measure whole-model generation.

The following MXFP4 results compare main with this PR applied to the same main revision on an Apple M5 Max. They measure synchronized dependent operator chains with BF16 inputs. Times are medians per call; ratios are paired geometric means of Baseline/Candidate, with 95% confidence intervals.

Workload	Baseline	Candidate	Ratio	95% CI
MXFP4 / BF16, N=5, K=1024, batch=2	9.298 µs/call	8.657 µs/call	1.0734	[1.0671, 1.0798]
MXFP4 / BF16, N=12, K=4096	10.482 µs/call	8.890 µs/call	1.1792	[1.1774, 1.1811]

Correctness

test_qmv_affine_bias_sum_precision fails on main for the affected float16 and bfloat16 cells below 8 bits and passes on this branch. test_qmv_fast_rows covers output sizes 1, 3, 5, 12 and 17 across 2- through 8-bit affine quantization, including batched weights. Its M=3 cases use qmv_wide on M3 and newer GPUs, so I could not exercise their affine_qmv_fast_rows route on this machine.

test_fp_qmv_fast_rows and test_fp_qmv_fast_rows_global_scale pass in native and JIT builds, covering N = 1, 3, 5, 7, 8, 9, 12, 17, all three FP quantization modes, FP32/FP16/BF16, batched inputs and weights, and NVFP4 global scales; aligned cases also match padded full-tile outputs bitwise.

For the layer benchmark, Baseline and Candidate produced bitwise-identical outputs over eight decode steps at both context lengths. Every changed call was also within half a bfloat16 step of a float64 reference.

Validation and quick reproduction

All four new tests also pass with MLX_METAL_JIT=ON. The full test_quantized run has the same 1273 failing subtests on main and this branch; 1188 are in test_qmv_wide.

The reproducer in #4498 shows the original output-size cliff. To reproduce the N=1 row above, set IN_DIM = 2560, GROUP_SIZE = 64, BITS = 8, and include out_dim = 1, then run it once against main and once against this branch. It prints differential median time rather than the paired ratio in the table, and results depend on the machine and build.

The layer benchmark uses Qwen4ExpDecoderLayer layers 4–7 from mlx-vlm 1ecf1ecdd28af102eded679be0daa5c76ab2a068 and tensors from mlx-community/Qwen3.8-Flash-Next-4bit revision 07b5dc6c54600a359b87f1e53e7adf6351c72a2c. It restores a cache prefetched with 128 or 4096 synthetic tokens before each sample and times 64 decode steps through Qwen4ExpBatchInvariantForward._layer.

---

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: AI assistance was used in preparing this contribution, including the implementation, tests and benchmark runs. I am responsible for the contribution.

[ursk]

Reporter of #4498 here. Confirming on a different machine, Mac Studio M3 Ultra (96 GB), builds from source at the merge-base (d9add9d) and this PR's head (7ee2bd0).

The issue's reproducer (chained quantized_matmul, K=10240, 4-bit, group 64, µs/call, two interleaved rounds):

out_dim	main	this PR
1	45	27
2	61	26
3	78	26
4	95	27
5–7	~80	~26
8	26	26
12	35	27
16	25	27

The cliff is gone and the curve is flat from 1 to 16. The N=4 shape is Qwen3.8-Flash-Next's per-block gate (96 calls/token), which is where we found it.

Numerics vs a float64 reference on dequantized weights: 8-bit is bit-identical to main; every 4-bit shape I tried is 2–3× closer to the reference on this branch, including a 320-row shape that doesn't take the new path, so that's the widened bias accumulation.

[zcbenz]

Should fp_qmv_fast also get equivalent updates?

[wyanzhao]

Yes, I added the same row-tail handling as the affine version in 862681ab3. test_fp_qmv_fast_rows and test_fp_qmv_fast_rows_global_scale pass in native and JIT builds, covering N = 1, 3, 5, 7, 8, 9, 12, 17, MXFP4/MXFP8/NVFP4, FP32/FP16/BF16, batching and NVFP4 global scales, with bitwise comparisons against padded full-tile outputs for aligned inputs.