[XXXXRT666]

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: Benchmarking, Testing, Profiling, Heuristic Searching

---

Based on #4796 for benchmarking

mx.qqmm now supports mode="affine", with a new biases argument for quantized affine weights. x is quantized to int8 to use NAX.

On chips with NAX, a new kernel runs int8 x int8 MMAs. It is used for 4, 6 and 8 bits weights when M >= 32, the group size is 64 or more, and there are at least 1.5 output tiles per GPU core. Otherwise qqmm dequantizes the 8-bit x and runs the quantized_matmul. CPU and CUDA raise an error for affine.

Accuracy

M = 1024, K = N = 4096, x and w from normal distribution, bf16. RMS error relative to the RMS of x @ w in float32:

w	quantized_matmul	qqmm
affine 4-bit, gs 64	9.10%	9.12%
affine 4-bit, gs 128	10.0%	10.0%
affine 6-bit, gs 64	2.23%	2.30%
affine 8-bit, gs 64	0.76%	0.95%
affine 8-bit, gs 128	0.81%	1.03%
nvfp4	-	13.4%
mxfp8	-	3.77%

Performance

Measured on M5 Pro (16 cores), M5 Max (40 cores) and M6 (12 cores), bf16, group size 64.

Shapes (K x N): 4096x4096, 4096x14336, 14336x4096, 8192x2048.

Speedup of qqmm over quantized_matmul, geomean over the 4 shapes (>1 is faster):

M5 Pro

M	4-bit	6-bit	8-bit
16	0.98x	0.98x	0.96x
64	1.28x	1.12x	1.00x
256	1.20x	1.12x	1.17x
1024	1.21x	1.11x	1.19x
4096	1.21x	1.12x	1.19x

M5 Max

M	4-bit	6-bit	8-bit
16	0.97x	0.95x	0.97x
64	1.08x	1.00x	0.94x
256	1.21x	1.11x	1.18x
1024	1.23x	1.06x	1.20x
4096	1.18x	1.09x	1.18x

M6

M	4-bit	6-bit	8-bit
16	0.98x	0.98x	0.99x
64	1.19x	1.04x	1.04x
256	1.31x	1.27x	1.26x
1024	1.37x	1.31x	1.34x
4096	1.60x	1.38x	1.52x

At M = 16 qqmm uses the quantized_matmul kernels.

Per shape at M = 1024, microseconds (quantized_matmul / qqmm affine / qqmm nvfp4 or mxfp8):

4-bit, with qqmm nvfp4:

K x N	M5 Pro	M5 Max	M6
4096x4096	1324 / 1066 / 1562	557 / 448 / 658	1669 / 1274 / 1919
4096x14336	4638 / 3533 / 5180	1865 / 1444 / 2124	5884 / 4147 / 6360
14336x4096	4654 / 3823 / 5668	1936 / 1614 / 2420	7914 / 5114 / 7287
8192x2048	1299 / 1183 / 1709	597 / 498 / 744	1731 / 1431 / 2076

8-bit, with qqmm mxfp8:

K x N	M5 Pro	M5 Max	M6
4096x4096	1313 / 1097 / 1731	548 / 455 / 721	1670 / 1317 / 2117
4096x14336	4618 / 3616 / 5629	1855 / 1477 / 2306	5875 / 4295 / 6828
14336x4096	4674 / 3896 / 6460	1936 / 1633 / 2673	8107 / 5743 / 8422
8192x2048	1287 / 1194 / 1936	591 / 505 / 812	1915 / 1468 / 2325

[zcbenz]

It will take a while before I can look into this, but is this the same approach taken in #4722/#4627?

[XXXXRT666]

It will take a while before I can look into this, but is this the same approach taken in #4722/#4627?

Almost, I think it's better to be implemented in qqmm

4722 converts the affine weights to symmetric int8 with one scale per row, so they are rounded a second time. Here the MMA uses the original codes, and the per-group scales and biases are applied in the kernel.