quantized_matmul takes 14- 41% longer than dequantize + bf16 matmul at prefill sizes (> 512 tokens) #4621

☑️ I understand it is strictly prohibited to use AI to write issues.

Summary

at prefill sizes mx.quantized_matmul with affine 4-bit / 8-bit weights is slower than dequantising the same weights to bf16 inside the call and running a plain matmul.
The gap holds at both bit widths, so it looks like the quantised kernel's path at large M rather than the cost of dequantisation. at M = 1 to 32 (decode), the quantised kernel is faster, as expected.

Env used

• MLX 0.32.3, mlx-lm 0.32.0
• Apple M2 Max, 64 GB, macOS 26.5.2

Repro

Weight shape 12288 × 4096 (the MLP up-projection of Qwen3.5 9B), group size 64, bf16 activations, mean of 20 runs after 3 warm-up runs

import time

import mlx.core as mx

D, F, GROUP = 4096, 12288, 64


def bench(fn, repeats=20):
    for _ in range(3):
        mx.eval(fn())
    start = time.perf_counter()
    for _ in range(repeats):
        mx.eval(fn())
    return (time.perf_counter() - start) / repeats * 1000


mx.random.seed(0)
w = (mx.random.normal((F, D)) * 0.02).astype(mx.bfloat16)
for bits in (4, 8):
    wq, scales, biases = mx.quantize(w, group_size=GROUP, bits=bits)
    dense = mx.dequantize(wq, scales, biases, group_size=GROUP, bits=bits)
    for tokens in (1, 32, 512, 2048, 8192):
        x = mx.random.normal((1, tokens, D)).astype(mx.bfloat16)
        q = bench(lambda: mx.quantized_matmul(x, wq, scales, biases, transpose=True, group_size=GROUP, bits=bits))
        jit = bench(lambda: x @ mx.dequantize(wq, scales, biases, group_size=GROUP, bits=bits).T)
        ref = bench(lambda: x @ dense.T)
        print(tokens, bits, f"{q:.2f}", f"{jit:.2f}", f"{ref:.2f}")

Results

The gap grows with the number of tokens

Effect on a model

Swapping each 4-bit projection of a model to dequantise-then-matmul during prefill only (and decode stays on quantized_matmul) raised prefill throughput at 8k tokens from 479 to 612 tokens/s on Qwen3.5 9B 4-bit and from 137 to 179 tokens/s on Qwen3.8 27B 4-bit, with peak memory unchanged.
on short chunks (under about 512 tokens) the fixed cost of dequantising every weight made it slower, so we only use it from 512 tokens up

Potential duplicate - or not

#4498 covers the decode side: quantized_matmul slower than dense bf16 at batch 1 when the output dimension is not a multiple of 8. The shapes here are multiples of 8, so this looks like a separate path

Question

Is the quantised kernel expected to fall behind dequantise-plus-matmul at large M - or is there a dispatch or tiling choice that could close it? If a heuristic in quantized_matmul picked the dequantise path above some M, callers would get the prefill speed-up without a workaround