[Metal] Add staged affine_qmm_t_nax kernels with split-k

[wyanzhao]

This PR adds a split-K NAX route and a 16-row simdgroup schedule for transposed affine quantized_matmul on M5. For LiquidAI d1-3B with locally converted 4-bit weights, an 85-token forward pass takes 14.5 ms instead of 30.1 ms on upstream main: 2.07× (95% CI [2.05, 2.09]) on M5 Max (40-core GPU). d1 returns decisions in one forward pass without generating tokens, making short-input latency a useful target. Gains are smaller when longer inputs or larger batches make the matmuls compute-bound.

At M = 8 to a few hundred rows, narrow projections (o_proj, down_proj, GQA k/v_proj) often produce too few output tiles to fill the GPU. They use qmv_wide, the non-NAX qmm_t_splitk, or an underfilled qmm_t_nax launch. qmv_wide also re-reads weights for each tile of at most 5 rows.

The new affine_qmm_t_nax_splitk kernel writes fp32 partials and sums them with the existing steel_gemm_splitk_accum kernel. The route requires a two-way K split that fits within 256 threadgroups. K is split at most 4 ways, with at least 256 columns per partition; K below 512 uses unsplit NAX.

For half/bfloat16 inputs with aligned N and 8-byte-aligned activation offsets, batched and unbatched qmm_t_nax and the new split-K kernel use four 16 x 64 simdgroups per 64 x 64 tile instead of a 2 x 2 grid of 32 x 32. Packed weights, scale/bias and activation fragments for the next K step are prefetched into registers. Weights are then dequantized into the other threadgroup buffer, leaving one barrier between K steps. The schedule uses flat (128, 1, 1) threadgroups.

The loader reuses QuantizedBlockLoader's thread mapping and dequantize for bits 1–6 and 8 and group sizes 32/64/128. The MMA order is unchanged, so outputs are bit-identical at the same split count.

For a single weight matrix, transposed affine mode on NAX, half/bfloat16, N % 64 == 0 and K % 64 == 0, dispatch changes as follows:

Shape	Before	After
8 <= M < qmv batch limit, M * N >= 16384, route taken	qmv / qmv_wide	NAX qmm: split-K for K >= 512, unsplit below
M >= qmv batch limit, route taken	qmm_t_splitk, or an unsplit qmm_t_nax	NAX qmm: split-K for K >= 512, unsplit below
M >= qmv batch limit, route not taken (wide outputs, large M)	qmm_t_nax	qmm_t_nax with the new schedule

Dispatch is unchanged for float32, unaligned N or K, fp modes and M < 8. Below the qmv batch limit, shapes with M * N < 16384 or K that does not split in two also keep their existing route. Above the limit, shapes still selected for qmm_t_splitk stay there.

Float32 retains the old schedule because double buffering would exceed 32 KB of threadgroup memory. Unaligned N and activations at offsets that are not 8-byte aligned also keep the old schedule; the split-K route copies those activations first. Gather kernels and fp modes are unchanged.

Performance

Measured on M5 Max (40-core GPU). Ratios are upstream main time divided by this PR's time: paired geometric means with 95% CIs from A/A calibration followed by 6 alternating rounds in the same process. Latencies are the median times for each configuration.

I timed one text-only forward pass on each prompt, including the output head on the last row:

Model	Prompt tokens	Baseline (main)	Candidate (this PR)	Ratio	95% CI
LiquidAI d1-3B (4-bit)	85	30.1 ms	14.5 ms	2.07	[2.05, 2.09]
LiquidAI d1-3B (4-bit)	643	93.5 ms	88.3 ms	1.060	[1.055, 1.065]

Operator measurements use bfloat16 activations and 4-bit affine weights, group size 64, in a dependent chain of quantized_matmul calls rotating through 8 weight sets.

Shape (K x N, M)	Baseline kernel	Ratio	95% CI
2048 x 2048, M = 22	qmv_wide	2.17	[2.11, 2.24]
4096 x 4096, M = 16	qmv_wide	2.33	[2.25, 2.41]
768 x 4096, M = 16	qmv_wide	2.09	[1.99, 2.18]
256 x 4096, M = 16 (unsplit)	qmv_wide	2.03	[2.01, 2.04]
10752 x 2048, M = 85	qmm_t_splitk	3.36	[3.34, 3.39]
14336 x 4096, M = 33	qmm_t_splitk	4.35	[4.30, 4.40]
256 x 4096, M = 40 (unsplit)	qmm_t_splitk	2.32	[2.27, 2.37]
17408 x 5120, M = 56	qmm_t_nax	2.09	[1.97, 2.21]
5120 x 17408, M = 28	qmm_t_nax	1.38	[1.37, 1.39]

The compute-bound case at 4096 x 4096, M = 512 shows no regression: 1.048 [1.033, 1.062]. The 4096 x 14336, M = 8 control keeps qmv_wide: 1.003 [0.966, 1.042].

For the mlx-lm models below, I timed one forward pass over N random tokens with a fresh KV cache, including the output head on all N rows. These use the same paired protocol.

Model	Prompt tokens	main	This PR	Ratio	95% CI
LFM2.5-1.2B-Instruct (4-bit)	32	8.6 ms	4.3 ms	2.00	[1.94, 2.07]
Llama-3.2-1B-Instruct (4-bit)	128	16.9 ms	8.6 ms	1.94	[1.89, 1.99]
Llama-3.2-1B-Instruct (4-bit)	512	35.9 ms	31.5 ms	1.14	[1.14, 1.15]
Qwen3.5-4B (4-bit)	32	22.3 ms	13.1 ms	1.71	[1.70, 1.71]
Qwen3.5-4B (4-bit)	128	39.4 ms	31.2 ms	1.28	[1.18, 1.38]
Qwen3.6-27B (4-bit)	32	139.5 ms	81.4 ms	1.70	[1.65, 1.76]
Qwen3.6-27B (4-bit)	128	204.3 ms	189.3 ms	1.10	[1.05, 1.15]
Qwen3.6-35B-A3B (4-bit, MoE)	128	84.1 ms	73.1 ms	1.13	[1.09, 1.16]

Correctness

Both kernels accumulate in fp32; split-K rounds once after summing the partials. On rerouted shapes, mean error against an fp32 reference stays within 1.1x of the output rounding.

test_qmm_mid_m checks every split count and both route boundaries, including K that does not split and K below 512. It covers bits 2/3/4/6/8, group sizes 32/64/128, float16/bfloat16 and M up to 257. test_qmm_t_batched_and_strided checks 3-D quantized weights, strided activations and odd-offset activations.

I also compared the new schedule with the first commit's kernels on 336 float16/bfloat16 cases: bits 1–6 and 8, group sizes 32/64/128 and M from 8 to 300. All were within fp32-reference tolerance, and outputs were bit-identical when the split count was unchanged.

Tests

• python -m unittest test_quantized (python/tests), M5 Max, MLX_ENABLE_TF32=0: 55 tests OK (1 skipped). With default TF32, 1270 subtests fail on main and 1268 fail with this PR. The 1268 shared failures are fp32 cases. The two additional failures on main are half/bfloat16 cases of test_qmm_splitk_precision; they fail on main with TF32 both on and off.
• Built and tested with MLX_METAL_JIT=ON as well.

Quick reproduction

import time, mlx.core as mx

glue = mx.compile(lambda a, b, z: a + b.reshape(-1)[:1].reshape(1, 1) * z)

def per_call(K, N, M, reps=64):
    sets = [mx.quantize((0.02 * mx.random.normal((N, K))).astype(mx.bfloat16), group_size=64, bits=4)
            for _ in range(8)]
    x = mx.random.normal((M, K)).astype(mx.bfloat16)
    z = mx.zeros((1, 1), dtype=mx.bfloat16)
    mx.eval(sets, x, z)

    def chain(r, qmm=True):  # each call depends on the previous one through `glue`
        y = x
        for i in range(r):
            wq, s, b = sets[i % 8]
            y = glue(x, y, z)
            if qmm:
                y = mx.quantized_matmul(y, wq, s, b, transpose=True, group_size=64, bits=4)
        mx.synchronize()
        t = time.perf_counter(); mx.eval(y); mx.synchronize()
        return time.perf_counter() - t

    def slope(qmm):
        chain(8, qmm)
        return min(chain(reps + 8, qmm) - chain(8, qmm) for _ in range(3)) / reps

    return (slope(True) - slope(False)) * 1e6

for K, N, M in [(2048, 2048, 22), (4096, 4096, 16), (10752, 2048, 85), (14336, 4096, 33),
                (17408, 5120, 56), (5120, 17408, 28)]:
    print(K, N, M, f"{per_call(K, N, M):.1f} us")

Run it on a build of the base commit (06eb7483f03bca49a80e596f4089d02a507a2cf9) and on a build of this PR (6e7ab39b4b1042d53860507e2e00e4854a152658). It prints microseconds per call. The tables above come from same-process paired runs; this script is a quick operator benchmark and does not reproduce their confidence intervals.

---

• ☑️ 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 implementation, tests and benchmark runs. I am responsible for the contribution.

[zcbenz]

Thanks for trying this approach, Metal's doc does not recommend double buffer (i.e. explicit software pipeline), it is amazing we actually can get 2x speed by doing it.

[zcbenz]

Will you also add staged implementation for fp_qmm in another PR as followup? If you do it is probably going to have an overlap with #4712

[wyanzhao]

I'll wait for #4712 to land, then benchmark the fp4/fp8 fp_qmm paths at larger M on M5. Based on the remaining bottlenecks, I'll decide whether to add a staged implementation in a follow-up PR.