2-bit quantized_matmul loses its size advantage over 4-bit at M >= 3

[ARahim3]

While getting speculative decoding to run on a 2-bit 27B checkpoint I noticed that multi-token verify steps were far more expensive than the weight sizes suggested, so I benchmarked quantized_matmul directly. At M=1, 2-bit is ~1.6x faster than 4-bit, as you'd expect from the bytes. But each extra row costs about 45% of the whole M=1 call, so by M=3 the two bit widths cost the same absolute time, and it stays that way:

shape (K→N), g128	bits	M1	M2	M3	M4	M5	M8	M10	M32
5120→17408	2	0.121 ms	1.30×	1.82×	2.34×	2.86×	4.53×	7.3×	7.3×
5120→17408	4	0.198 ms	0.99×	1.13×	1.43×	1.75×	2.84×	4.5×	4.5×
5120→248320 (lm_head)	2	1.539 ms	1.39×	1.98×	2.53×	3.14×	5.11×
5120→248320 (lm_head)	4	2.725 ms	1.01×	1.14×	1.47×	1.85×	3.03×

(M4 Pro applegpu_g16s, mlx 0.32.0 wheel, macOS 15.6. Medians of 5 trials, 20 warm calls per eval. Multipliers are relative to that row's own M1. Absolute numbers at the first shape, M3: 2-bit 0.221 ms vs 4-bit 0.224 ms.)

Summing this over every linear in a real model (prism-ml/Ternary-Bonsai-27B-mlx-2bit, all shapes times their counts) gives a per-forward qmm bill of 34.1 / 45.4 / 63.0 / 80.2 / 98.3 ms at widths 1–5, and measured end-to-end forwards match those numbers within a few percent, so this is the whole story for multi-token decode on that model. Concretely: baseline decode is ~25 tok/s (41 ms/forward, ~89% of the bandwidth roofline, so M=1 is in good shape), but a width-3 speculative verify costs 71.5 ms where a 4-bit-like slope would give ~45 ms. That gap is roughly the difference between the ~1.2x speculative speedup I measure on this checkpoint and the 1.6–2.1x the same loop gets on 8-bit models on the same machine.

I saw #3553 and #3839, but both measure the small-M slope at 4-bit. The thing I wanted to flag here is the bit-width side of it: the slope fully erases 2-bit's storage advantage exactly in the window speculative verify lives in (M = draft+1 = 2–6). Everything that widens decode a little (spec decode, n>1 sampling, small continuous batches) sits in this window.

A second, smaller thing from the same sweep: past the qmv batch limit the qmm path is flat at 0.887 ms for every M from 10 to 32 (both bit widths, 5120→17408), so M=10 pays the M=32 price. Different issue really, but it shows up in the same table.

Things I checked before filing:

• Dispatch: gen 16 takes qmv_wide for affine at M >= 2 (use_qmv_wide), up to get_qmv_batch_limit (10–12 at these dims), then qmm. So the table above is qmv_wide's slope plus qmm's plateau.
• qmv_wide already dequantizes each weight sub-chunk once and reuses it across the M rows, so the marginal row looks FMA/issue-bound rather than bandwidth-bound: the ~0.05 ms/row at 5120→17408 works out to ~89M FMA/row at ~1.8 Tfma/s, about 40% of this GPU's scalar fp32 peak.
• I tried writing my own skinny-M kernel with mx.fast.metal_kernel (dequant once, M register accumulators) and could not beat qmv_wide at any M — it's well tuned. Half-precision arithmetic ran at identical speed (no 2x half rate on M-series), and math_mode: "fast" was a no-op. So I don't have a downstream workaround; whatever headroom exists needs in-tree tuning, which is why I'm reporting it instead.

Is there room in qmv_wide's per-row inner loop (vectorized activation loads, more columns per simdgroup, dual-issue scheduling) to push the marginal row closer to peak? Even ~60% of peak would take the 27B width-3 verify from 71.5 to ~55 ms. And would a small-M tile variant of the qmm path (BM = 8/16) make sense for the M in [batch_limit, 32] dead zone?

One more long-shot idea while I'm here: a lot of published "2-bit" checkpoints are actually ternary repacked as affine (values exactly {-s, 0, +s}, biases == -scales). A denser ternary storage mode (~1.6 b/weight vs 2.25 effective) would cut M=1 decode bytes another ~25% independent of all the above.

Repro (self-contained):

import time
import mlx.core as mx

mx.random.seed(0)
REPS, TRIALS = 20, 5

def time_qmm(K, N, M, bits, group_size=128):
    w = mx.random.normal((N, K)).astype(mx.bfloat16)
    wq, sc, bi = mx.quantize(w, group_size=group_size, bits=bits)
    xs = [mx.random.normal((1, M, K)).astype(mx.bfloat16) for _ in range(REPS)]
    mx.eval(wq, sc, bi, xs)
    def run():
        outs = [mx.quantized_matmul(x, wq, sc, bi, transpose=True,
                                    group_size=group_size, bits=bits) for x in xs]
        mx.eval(outs)
    run(); run()  # warm
    ts = []
    for _ in range(TRIALS):
        t0 = time.perf_counter(); run(); ts.append((time.perf_counter() - t0) / REPS)
    mx.clear_cache()
    return sorted(ts)[TRIALS // 2]

print(mx.__version__, mx.device_info()["device_name"])
for bits in (2, 4):
    t1 = None
    for m in (1, 2, 3, 4, 5, 8, 10, 16, 32):
        t = time_qmm(5120, 17408, m, bits)
        t1 = t1 or t
        print(f"bits={bits} M={m}: {t*1e3:6.3f} ms ({t/t1:4.2f}x M1)")

Environment: mlx 0.32.0 (Metal), macOS 15.x / Darwin 24.6.0, Apple M4 Pro (20-core GPU, applegpu_g16s), 48 GB, Python 3.12.