[XXXXRT666]

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: Test, Benchmark, Heuristic Search, Fill in the Data in PR Descriptions

Few-row and mid-row quantized_matmul on GPUs with NAX is slower than it has to be. The FP formats will not take the NAX split-K route. Affine has no qmv_nax for 6–16 rows.

Over 36 projection shapes from 8 models and 5 formats, measured against upstream on three GPUs:

• FP formats at M = 17–256: 1.3–2.6× faster, up to 5.2× for single shapes.
• affine 4/8-bit at M = 6–16: 1.00–1.64× faster.
• M = 1: unchanged.

Fix #4720

Changes

1. NAX split-K for FP formats. New fp_qmm_t_nax_splitk kernel.
2. Staged schedule for fp qmm_t_nax. Adapted from [Metal] Add staged affine_qmm_t_nax kernels with split-k #4682
3. qmv_nax for narrower FP shapes. The NAX limit for N and K goes from 2048 to 1024. nvfp4 keeps 2048 because qmv_wide is faster there.
4. Split-K count scaled with GPU core count. qmm_nax_k_partitions caps the grid at min(256, 6 × cores × max(K, 2048) / 2048) threadgroups, so long K can still split on small GPUs.
5. affine_qmv_nax for 6–16 rows. Adapted from fp_qmv_nax.

Dispatch changes

Case	Before	After
FP, M >= 8 with M*N >=16384, or M >= qmv batch limit, split-K conditions met	qmv_wide or non-NAX qmm_t_splitk	NAX split-K
FP formats, fp16/bf16 NAX qmm	2×2 simdgroup schedule	staged schedule
FP formats, M = 6–16, smaller of N and K in [1024, 2048]	qmv_wide or split-K	qmv_nax
affine 3/4/5/6/8-bit, M = 6–16	qmv_wide / split-K / qmv	affine_qmv_nax

Performance

Measured on M5 Pro (16 core), M5 Max (40 core) and M6 (12 core), macOS 27.0.1

Shapes: q/o, k/v, gate/up, down for Qwen3-0.6B, Llama-3.2-1B, Llama-3.2-3B, Qwen3-4B, Qwen3-8B / Llama-3-8B, Qwen3-14B, Qwen3-32B and Llama-3-70B (35 shapes), plus a 4096 × 151936 LM head. That is 5,400 pairs in total.

Speedup by M

M5 Pro

Mode	M=1	M=6	M=8	M=12	M=16	M=17	M=32	M=64	M=128	M=256
affine 4-bit	1.00	1.51	1.46	1.64	1.24	1.02	1.04	1.05	1.07	1.08
affine 8-bit	1.00	1.14	1.28	1.50	1.24	1.02	1.02	1.04	1.06	1.07
mxfp4	1.00	1.11	1.17	1.26	1.34	2.55	2.54	2.17	1.59	1.42
nvfp4	1.00	1.00	1.01	1.16	1.25	2.11	2.21	1.98	1.49	1.32
mxfp8	1.00	1.10	1.14	1.24	1.30	2.26	2.23	2.24	1.70	1.54

M5 Max

Mode	M=1	M=6	M=8	M=12	M=16	M=17	M=32	M=64	M=128	M=256
affine 4-bit	1.00	1.27	1.31	1.39	1.20	1.00	1.00	0.99	1.00	1.01
affine 8-bit	1.00	1.04	1.22	1.33	1.25	1.00	1.01	1.00	1.00	0.99
mxfp4	1.00	1.00	1.06	1.12	1.18	1.96	1.98	1.99	1.62	1.44
nvfp4	1.00	1.01	1.00	1.03	1.10	1.67	1.77	1.84	1.49	1.32
mxfp8	1.00	1.04	1.08	1.13	1.19	1.72	1.71	1.93	1.75	1.53

M6

Mode	M=1	M=6	M=8	M=12	M=16	M=17	M=32	M=64	M=128	M=256
affine 4-bit	1.00	1.11	1.23	1.44	1.21	1.02	1.04	1.07	1.09	1.08
affine 8-bit	1.00	1.00	1.16	1.33	1.28	1.00	1.02	1.05	1.10	1.08
mxfp4	1.00	1.04	1.10	1.18	1.26	2.18	2.27	2.10	1.64	1.47
nvfp4	1.00	1.00	1.00	1.10	1.19	1.85	2.00	1.94	1.55	1.44
mxfp8	1.00	1.01	1.04	1.13	1.20	1.57	1.62	1.77	1.67	1.53

Speedup per model (M = 6–32, five quant formats)

Model	Shapes	M5 Pro	M5 Max	M6
Qwen3-0.6B	5	1.65	1.28	1.48
Llama-3.2-1B	4	1.30	1.20	1.23
Llama-3.2-3B	4	1.39	1.23	1.27
Qwen3-4B	5	1.34	1.19	1.23
Qwen3-8B / Llama-3-8B	6	1.33	1.21	1.20
Qwen3-14B	4	1.39	1.26	1.23
Qwen3-32B	4	1.29	1.24	1.18
Llama-3-70B	4	1.41	1.27	1.24
LM head (151936 vocab)	1	1.18	1.15	1.09

Count of cases (M = 6–256, 1,620 per GPU)

GPU	≥ 1.2× faster	< 0.9×	Largest speedup
M5 Pro	672	1	5.23× (mxfp8, 2048×2048, M = 32)
M5 Max	651	11	3.80× (mxfp4, 9728×2560, M = 64)
M6	591	0	4.63× (mxfp8, 1024×1024, M = 256)

Per-shape times, Qwen3-8B / Llama-3-8B and LM head (µs, upstream → branch)

M5 Pro

Shape (K×N)	Mode	M=8	M=16	M=17	M=32	M=64	M=128
q/o 4096x4096	affine 4-bit	56.0 → 47.6	62.3 → 50.0	67.9 → 66.6	73.1 → 70.9	109 → 98	194 → 174
q/o 4096x4096	affine 8-bit	82.9 → 71.1	89.3 → 70.6	88.4 → 83.9	94.5 → 87.2	116 → 104	194 → 174
q/o 4096x4096	mxfp4	47.1 → 47.2	51.5 → 50.7	243 → 69	204 → 73	394 → 101	208 → 186
q/o 4096x4096	nvfp4	50.4 → 50.4	55.7 → 56.0	185 → 76	198 → 81	374 → 111	214 → 198
k/v 4096x1024	affine 4-bit	25.4 → 13.9	17.5 → 14.5	20.8 → 20.9	20.9 → 21.1	29.0 → 28.9	54.6 → 54.7
k/v 4096x1024	affine 8-bit	28.7 → 21.1	23.4 → 21.6	24.2 → 24.5	25.4 → 25.6	30.4 → 30.5	53.8 → 53.9
k/v 4096x1024	mxfp4	28.3 → 13.9	51.6 → 14.9	64.8 → 21.9	56.9 → 22.0	102 → 29	196 → 55
k/v 4096x1024	nvfp4	23.5 → 23.5	41.3 → 16.6	49.5 → 23.9	56.2 → 24.3	98.2 → 31.6	189 → 58
gate/up 4096x12288	affine 4-bit	244 → 124	167 → 135	199 → 201	210 → 208	305 → 300	523 → 517
gate/up 4096x12288	affine 8-bit	274 → 197	269 → 200	270 → 274	275 → 281	312 → 311	519 → 514
gate/up 4096x12288	mxfp4	128 → 126	142 → 140	248 → 195	253 → 208	355 → 303	603 → 541
gate/up 4096x12288	nvfp4	140 → 140	152 → 150	255 → 215	257 → 228	359 → 340	629 → 588
down 12288x4096	affine 4-bit	156 → 123	163 → 133	185 → 186	199 → 196	286 → 283	526 → 525
down 12288x4096	affine 8-bit	261 → 206	269 → 206	275 → 275	292 → 289	322 → 324	537 → 532
down 12288x4096	mxfp4	127 → 124	135 → 132	614 → 188	595 → 196	1158 → 289	664 → 554
down 12288x4096	nvfp4	137 → 136	148 → 148	583 → 207	570 → 214	1110 → 303	689 → 643
LM head 4096x151936	affine 4-bit	2930 → 1361	1672 → 1491	1947 → 1952	2044 → 2041	3101 → 3115	6170 → 6169
LM head 4096x151936	affine 8-bit	3425 → 2234	2970 → 2256	2988 → 2979	3025 → 3009	3238 → 3294	6266 → 6604
LM head 4096x151936	mxfp4	1369 → 1369	1505 → 1504	2421 → 2079	2374 → 2166	3695 → 3512	7248 → 7292
LM head 4096x151936	nvfp4	1542 → 1544	1678 → 1690	2556 → 2312	2581 → 2443	4050 → 4006	7933 → 7958

M5 Max

Shape (K×N)	Mode	M=8	M=16	M=17	M=32	M=64	M=128
q/o 4096x4096	affine 4-bit	32.7 → 27.7	34.8 → 28.0	37.3 → 36.4 †	39.3 → 39.0 †	52.4 → 52.9 †	89.8 → 89.8 †
q/o 4096x4096	affine 8-bit	46.4 → 40.4	47.8 → 41.6	48.3 → 49.3 †	51.4 → 54.2 †	59.8 → 59.6 †	90.1 → 91.9 †
q/o 4096x4096	mxfp4	26.0 → 26.2	29.6 → 27.5	102 → 37	87.5 → 39.5	162 → 53	109 → 93
q/o 4096x4096	nvfp4	28.4 → 28.7	29.8 → 30.0	77.5 → 43.1	90.2 → 43.3	165 → 58	105 → 102
k/v 4096x1024	affine 4-bit	16.5 → 11.5	15.4 → 12.4	17.2 → 17.2 †	17.2 → 17.4 †	20.3 → 20.4 †	28.1 → 28.0 †
k/v 4096x1024	affine 8-bit	18.4 → 13.5	19.5 → 14.3	19.6 → 18.2 †	20.3 → 18.7 †	22.2 → 20.7 †	27.7 → 27.7 †
k/v 4096x1024	mxfp4	14.7 → 11.6	24.1 → 12.7	29.4 → 18.1	28.3 → 18.3	46.0 → 21.5	83.6 → 28.6
k/v 4096x1024	nvfp4	11.9 → 12.2	19.6 → 13.9	23.4 → 19.9	28.5 → 19.7	45.3 → 23.1	80.9 → 31.4
gate/up 4096x12288	affine 4-bit	108 → 69	86.9 → 69.5	86.2 → 89.1 †	95.9 → 98.4 †	133 → 138 †	238 → 239 †
gate/up 4096x12288	affine 8-bit	122 → 108	150 → 106	146 → 152 †	149 → 154 †	160 → 171 †	240 → 253 †
gate/up 4096x12288	mxfp4	67.3 → 67.2	69.8 → 71.1	108 → 78	111 → 85	163 → 135	271 → 248
gate/up 4096x12288	nvfp4	74.6 → 71.1	74.4 → 77.7	112 → 87	114 → 95	168 → 143	300 → 258
down 12288x4096	affine 4-bit	89.8 → 69.2	92.4 → 73.4	96.1 → 98.3 †	99.1 → 100.5 †	125 → 129 †	244 → 248 †
down 12288x4096	affine 8-bit	143 → 107	146 → 108	146 → 148 †	152 → 151 †	164 → 162 †	236 → 239 †
down 12288x4096	mxfp4	67.7 → 69.1	72.3 → 73.1	257 → 100	252 → 101	471 → 132	347 → 253
down 12288x4096	nvfp4	73.7 → 73.7	79.4 → 77.0	265 → 109	260 → 113	481 → 141	345 → 284
LM head 4096x151936	affine 4-bit	1173 → 619	775 → 650	850 → 851 †	878 → 880 †	1268 → 1377 †	2498 → 2818 †
LM head 4096x151936	affine 8-bit	1381 → 1156	1492 → 1161	1504 → 1519 †	1525 → 1548 †	1652 → 1656 †	3200 → 3414 †
LM head 4096x151936	mxfp4	606 → 611	647 → 648	1046 → 861	1073 → 921	1598 → 1420	3173 → 2934
LM head 4096x151936	nvfp4	672 → 676	720 → 708	1136 → 965	1155 → 1020	1742 → 1651	3459 → 3375

M6

Shape (K×N)	Mode	M=8	M=16	M=17	M=32	M=64	M=128
q/o 4096x4096	affine 4-bit	88.4 → 77.5	96.8 → 80.9	98.5 → 100.0	112 → 106	153 → 144	258 → 253
q/o 4096x4096	affine 8-bit	162 → 135	168 → 132	168 → 175	185 → 188	209 → 198	259 → 244
q/o 4096x4096	mxfp4	74.8 → 75.4	78.5 → 79.9	287 → 100	255 → 107	477 → 143	278 → 222
q/o 4096x4096	nvfp4	81.1 → 80.9	85.2 → 84.9	224 → 108	247 → 116	465 → 154	278 → 230
k/v 4096x1024	affine 4-bit	29.8 → 22.6	29.8 → 23.6	31.7 → 31.6	33.7 → 33.7	43.7 → 43.6	82.9 → 83.2
k/v 4096x1024	affine 8-bit	39.0 → 35.4	45.7 → 36.8	46.4 → 46.8	48.1 → 48.2	55.3 → 55.4	79.5 → 79.2
k/v 4096x1024	mxfp4	32.8 → 22.2	59.7 → 23.2	75.4 → 30.7	76.6 → 33.6	132 → 44	243 → 77
k/v 4096x1024	nvfp4	26.5 → 26.6	48.6 → 25.7	58.4 → 33.6	75.2 → 35.6	129 → 47	239 → 83
gate/up 4096x12288	affine 4-bit	306 → 217	269 → 222	268 → 272	286 → 288	360 → 360	663 → 664
gate/up 4096x12288	affine 8-bit	380 → 380	498 → 384	488 → 502	510 → 504	514 → 524	670 → 675
gate/up 4096x12288	mxfp4	203 → 207	208 → 210	316 → 259	331 → 266	440 → 364	733 → 659
gate/up 4096x12288	nvfp4	220 → 221	223 → 225	328 → 277	344 → 288	456 → 386	741 → 668
down 12288x4096	affine 4-bit	269 → 248	282 → 255	285 → 282	297 → 296	392 → 388	691 → 693
down 12288x4096	affine 8-bit	498 → 401	509 → 400	518 → 513	534 → 544	572 → 587	742 → 750
down 12288x4096	mxfp4	234 → 237	244 → 249	724 → 270	735 → 288	1414 → 379	869 → 694
down 12288x4096	nvfp4	248 → 249	264 → 265	703 → 288	700 → 316	1342 → 397	859 → 754
LM head 4096x151936	affine 4-bit	3677 → 2558	3029 → 2618	3070 → 3094	3166 → 3252	3885 → 3885	7723 → 7670
LM head 4096x151936	affine 8-bit	4558 → 4554	5914 → 4573	5884 → 5910	5978 → 5993	6097 → 6079	7729 → 7740
LM head 4096x151936	mxfp4	2370 → 2369	2421 → 2417	3164 → 2906	3238 → 2952	4356 → 3992	8334 → 7927
LM head 4096x151936	nvfp4	2558 → 2545	2616 → 2618	3355 → 3120	3452 → 3183	4455 → 4326	8372 → 8056