qmm_splitk dispatch causes 1.5-1.8x regression vs fp16 at M=96, M=128 on NAX hardware

[marian-gg]

TL;DR

mx.quantized_matmul is 1.5–1.8× slower than mx.matmul (fp16) at M=96 and
M=128 on common transformer FFN-down shapes (N=2048, K=8192) on M5 Max NAX
hardware. At M ≥ 160 it's at parity (~1.0× of fp16) as expected. The cause is
qmm_splitk's aggressive split-K threshold: it forces a split-K reduction
even when the underlying NAX qmm_t_nax would be faster as a single kernel
launch.

Fix candidate: bump the 512 threshold in mlx/backend/metal/quantized.cpp:793
(or remove split-K from the NAX path entirely — NAX qmm is already fast at
small M).

Reproducer

See bench/bench_splitk_boundary.py in this repo. Sweeps M on FFN-down shape,
prints split_k value and timing ratios.

Observed on M5 Max (g17s, NAX), macOS 26.5, MLX 0.31.2:

M=  64  fp16: 0.300 ms   qmm: 0.307 ms   qmm/fp16: 1.02x    (split_k=4)
M=  96  fp16: 0.265 ms   qmm: 0.418 ms   qmm/fp16: 1.58x    (split_k=2) ← regression
M= 128  fp16: 0.272 ms   qmm: 0.486 ms   qmm/fp16: 1.79x    (split_k=2) ← regression
M= 160  fp16: 0.354 ms   qmm: 0.354 ms   qmm/fp16: 1.00x    (split_k=1)
M= 192  fp16: 0.341 ms   qmm: 0.349 ms   qmm/fp16: 1.02x    (split_k=1)
M= 256  fp16: 0.349 ms   qmm: 0.381 ms   qmm/fp16: 1.09x    (split_k=1)

Root cause

In mlx/backend/metal/quantized.cpp:787-805:

void qmm_splitk(...) {
  int bm = 32, bn = 32;
  int n_tiles = (N + bn - 1) / bn;
  int m_tiles = (M + bm - 1) / bm;
  int current_tgs = n_tiles * m_tiles;
  int split_k = std::max(1, 512 / current_tgs);
  ...
  if (split_k <= 1) {
    return qmm(...);  // fall back to non-split kernel
  }
  // ... else allocate intermediate buffer, launch split-K kernel, sum-reduce
}

For FFN-down (N=2048):

• n_tiles = 64
• M=64 → m_tiles=2, current_tgs=128, split_k=4
• M=96 → m_tiles=3, current_tgs=192, split_k=2
• M=128 → m_tiles=4, current_tgs=256, split_k=2
• M=160 → m_tiles=5, current_tgs=320, split_k=1 (integer div: 512/320=1)
• M=192+ → split_k=1

The split-K path allocates an intermediate [split_k, M, N] buffer and then
runs a sum-along-axis-0 reduce after the matmul. On the NAX-capable path,
qmm_t_nax is so fast that this extra allocation + reduce step is more
overhead than the parallelism gain. The split_k = 2 cases at M=96, M=128
clearly demonstrate this: the regular qmm() would run in ~half the time.

Suggested fix

Option 1 (minimal): bump the threshold so split-K only triggers in extreme
small-N cases:

int split_k = std::max(1, 128 / current_tgs);  // was 512

Option 2 (NAX-aware): bypass split-K entirely on NAX:

if (metal::is_nax_available() && transpose) {
    return qmm(x, w, scales, biases, out, true, group_size, bits, M, N, K, d, s, mode);
}

Option 3 (full re-tuning): empirically sweep the threshold per (N, K) shape on
representative hardware. The 512 was likely tuned for the pre-NAX kernel.

Impact

This regression hits exactly the prefill lengths most chatbot prompts produce
(roughly 96-160 tokens after system prompt + short user message). Fixing it
recovers ~1.5× prefill throughput at that operating point for free, on any
NAX-capable Apple Silicon (M3/M4/M5 Pro+).

Hardware tested

Apple M5 Max (applegpu_g17s), 128 GB unified memory, macOS 26.5, Xcode 26.5
Metal Toolchain 32023.883.

Reference

This was found while building a custom NAX-accelerated fused dequant+matmul
kernel for testing. The custom kernel doesn't use split-K and lands at
1.39× faster than mx.quantized_matmul at M=128 on the same shape,
specifically because it avoids this dispatch overhead. See
bench/bench_splitk_boundary.py and bench/bench_smollm2_real.py in this
repo for full reproducers.

[ima-helikoptaaa]

@zcbenz @marian-gg
Tried to reproduce the issue on the following hardware

Hardware/software:

• Apple M3 Max (Mac15,11), 36 GB
• macOS 26.2, Xcode 26.2, Metal Toolchain 32023.864
• MLX built from main @ 2165dc0

No regression found, Option 2 (NAX-aware): bypass split-K entirely on NAX, risks regressing on M3/M4. Will be safer to do sweep on each category of machine as M3, M4 and M5 all have different implementations of NAX

shape: N=2048, K=8192, group_size=64, bits=4
    M   fp16(ms)    qmm(ms)   qmm/fp16   split_k
   64      0.501      0.479      0.96x         4
   96      0.633      0.600      0.95x         2
  128      0.813      0.769      0.94x         2
  160      0.914      0.920      1.01x         1
  192      1.072      1.036      0.97x         1
  256      1.318      1.309      0.99x         1

Benchmarking script

import time
import mlx.core as mx

N, K = 2048, 8192
GROUP_SIZE, BITS = 64, 4

WARMUP_ITERS = 20
TIMED_ITERS = 100
REPEATS = 5

MS = [64, 96, 128, 160, 192, 256]


def compute_split_k(M, N, K, group_size):
    bm, bn = 32, 32
    n_tiles = (N + bn - 1) // bn
    m_tiles = (M + bm - 1) // bm
    current_tgs = n_tiles * m_tiles
    split_k = max(1, 512 // current_tgs)
    split_k = min(split_k, K // group_size)
    while split_k > 1 and (K % (split_k * group_size) != 0):
        split_k -= 1
    return split_k


def time_min(fn):
    for _ in range(WARMUP_ITERS):
        mx.eval(fn())
    mx.synchronize()

    best = float("inf")
    for _ in range(REPEATS):
        mx.synchronize()
        t0 = time.perf_counter()
        for _ in range(TIMED_ITERS):
            mx.eval(fn())
        mx.synchronize()
        elapsed = (time.perf_counter() - t0) * 1000 / TIMED_ITERS
        best = min(best, elapsed)
    return best


def main():
    w_fp16 = mx.random.normal((N, K)).astype(mx.float16)
    w_q, scales, biases = mx.quantize(w_fp16, group_size=GROUP_SIZE, bits=BITS)
    mx.eval(w_fp16, w_q, scales, biases)

    print("global warmup...")
    for M in MS:
        x = mx.random.normal((M, K)).astype(mx.float16)
        for _ in range(5):
            mx.eval(x @ w_fp16.T)
            mx.eval(
                mx.quantized_matmul(
                    x, w_q, scales, biases,
                    transpose=True, group_size=GROUP_SIZE, bits=BITS,
                )
            )
    mx.synchronize()

    print(f"\nshape: N={N}, K={K}, group_size={GROUP_SIZE}, bits={BITS}")
    print(f"{'M':>5} {'fp16(ms)':>10} {'qmm(ms)':>10} {'qmm/fp16':>10}  {'split_k':>8}")
    for M in MS:
        x = mx.random.normal((M, K)).astype(mx.float16)
        mx.eval(x)

        t_fp16 = time_min(lambda: x @ w_fp16.T)
        t_qmm = time_min(
            lambda: mx.quantized_matmul(
                x, w_q, scales, biases,
                transpose=True, group_size=GROUP_SIZE, bits=BITS,
            )
        )
        ratio = t_qmm / t_fp16
        sk = compute_split_k(M, N, K, GROUP_SIZE)
        marker = "  <-- regression" if ratio > 1.2 else ""
        print(f"{M:>5} {t_fp16:>10.3f} {t_qmm:>10.3f} {ratio:>9.2f}x  {sk:>8}{marker}")


if __name__ == "__main__":
    main()