Add a BM=16 qmm tile for the small-M decode dead-zone#3863
Conversation
This workflow is designed to build a CMake project on multiple platforms, including Windows and Ubuntu, using different compilers.
This workflow installs Python dependencies, runs linting with flake8, and executes tests with pytest on push and pull request events.
Past the qmv batch-limit, affine qmm rounds M up to its 32-row tile, so M in [limit,16] all pay the M=32 price (flat plateau). Add a 16-row tile for M<=16: ~1.65x on M in [10,16] across 4-/2-bit and both projection and lm_head shapes; M>=17 and every other path byte-identical. Correctness exact vs dequantize-then- matmul; test_quantized.py green (32 passed, 2896 subtests). Refs ml-explore#3852
|
Can you please provide the benchmark you ran? I think part of this should be handled by |
|
Sure. Script is at the bottom; here's the full M sweep across the dispatch boundary, plus a measurement of the Where the boundary actually is. On this M3 Max ( 4-bit, cold harness (hot agrees within ~1%; min of 15 trials × 30 reps):
M < 10 is the unchanged qmv/qmv_wide path and M ≥ 17 is the unchanged Would
It's linear in M by construction, so it diverges exactly where this PR helps — at M=16 it's 2.7× the The flip side is probably the more useful finding: with a 16-row tile available the crossover moves down. Scope, to be explicit: the win band is benchmark scriptimport time
import mlx.core as mx
mx.random.seed(0)
K, N, GS = 5120, 17408, 128
REPS, TRIALS = 30, 15
MS = list(range(1, 21)) + [24, 32, 33]
def make_bank(bits):
w = mx.random.normal((N, K)).astype(mx.bfloat16)
b = mx.quantize(w, group_size=GS, bits=bits)
mx.eval(b)
return b
def timeit(run):
run()
run()
ts = []
for _ in range(TRIALS):
t0 = time.perf_counter()
run()
ts.append((time.perf_counter() - t0) / REPS)
return min(ts)
def bench(M, bits, banks):
nb = len(banks)
xs = [mx.random.normal((1, M, K)).astype(mx.bfloat16) for _ in range(REPS)]
mx.eval(xs)
def run():
outs = [
mx.quantized_matmul(
xs[i], *banks[i % nb], transpose=True, group_size=GS, bits=bits
)
for i in range(REPS)
]
mx.eval(outs)
t = timeit(run)
mx.clear_cache()
return t
def bench_split(M, bits, banks, chunk=5):
"""What qmv_wide would cost above its batch limit: it tiles M into
ceil(M/5) blocks and each block re-reads the weights."""
nb = len(banks)
parts = [min(chunk, M - i) for i in range(0, M, chunk)]
xs = [
[mx.random.normal((1, p, K)).astype(mx.bfloat16) for p in parts]
for _ in range(REPS)
]
mx.eval(xs)
def run():
outs = []
for i in range(REPS):
for x in xs[i]:
outs.append(
mx.quantized_matmul(
x, *banks[i % nb], transpose=True, group_size=GS, bits=bits
)
)
mx.eval(outs)
t = timeit(run)
mx.clear_cache()
return t
info = mx.device_info()
print(f"mlx {mx.__version__} | {info['device_name']} | {info['architecture']}")
for bits in (4, 2):
hot = [make_bank(bits)] # one weight reused: cache-warm
cold = [make_bank(bits) for _ in range(REPS)] # ~1.4 GB working set
print(f"=== {bits}-bit ===")
for M in MS:
print(f"{M:>4} hot {bench(M, bits, hot) * 1e3:7.3f} cold {bench(M, bits, cold) * 1e3:7.3f}")
for M in (10, 12, 16):
th, tc = bench_split(M, bits, hot), bench_split(M, bits, cold)
print(f" M={M} as {-(-M // 5)}x<=5 rows: hot {th * 1e3:7.3f} cold {tc * 1e3:7.3f}")
del hot, cold
mx.clear_cache()"cold" = 30 distinct weight banks round-robin (~1.4 GB) so the weights can't sit in cache; that distinction matters at M=1 (it's what hid the 2-bit advantage in #3852) but not in this band, where the two harnesses agree to ~1%. |
|
Correction to my numbers above, plus the sweep behind the crossover claim — which turned out to contradict part of what I said. The correction. The
Same conclusion, but the emulated row shouldn't stand. (Absolute times in this run are ~10% above my first table — different build and session — so only compare within a table.) Sweep. Both routings measured at the same M in one process, 7 shapes × {4,2}-bit, M ∈ [2,20] ∪ {24,32}, cold (31–64 distinct weight banks, ~1.5 GB, so nothing is cache-resident — that matters because
Where I was wrong. I said the limit looked "2–3 too high". That holds for the MLP-sized layers — at M=9, the last M currently routed to qmv_wide, qmm is 1.34× / 1.62× / 1.42× / 1.56× faster on 4096→14336, 3584→18944, 5120→17408, lm_head. It does not hold for the two smaller shapes: at 4096→4096 and 2048→8192, What is clean is that the crossover tracks K·N and nothing else I varied. The two shapes with identical K·N — 4096→4096 and 2048→8192 — land on the same crossover of 13, while the current tiers give them different limits (12 and 10). And 4-bit and 2-bit agree on every shape, i.e. both kernels are element-bound rather than byte-bound in this band, which is consistent with a limit keyed on K·N rather than on And a scope limit on this PR that the same run exposes. Comparing
So the tile buys ~1.5× on MLP- and lm_head-class layers and nothing at all below ~58M elements — the small shapes are flat in M across the whole band, so they're not tile-quantization-bound in the first place. My original numbers were both large shapes, which overstated how general the win is. Worth adding to the PR body if you'd like; it doesn't change the change, but it does narrow the claim. I'm not proposing a retune off this. It's one part (M3 Max, gen-15 non- |
2 participants
Summary
Affine
quantized_matmulat small M (past the qmv batch-limit, up to 32) always pays for a full 32-row tile: theqmmpath uses a singleBM=32tile and launchesceil(M/32)M-tiles, so M=10 and M=32 take the same absolute time. This shows up as a flat plateau — see #3852 (@ARahim3): qmm is flat from M=10 to M=32, then steps at 33 / 65 / 97 (one step per 32-row tile).This adds a
BM=16tile and selects it forM <= 16in affine mode, recovering the[qmv-limit, 16]band that was rounding up to 32.Mechanism
qmm()launchesgrid = (ceil(N/32), ceil(M/32), B), so for M∈[1,32] that's one M-tile and the kernel computes 32 rows regardless of the real M. A fine M-sweep (M3 Max, 5120→17408, 4-bit) confirms the tile quantization is the whole story: flat ~0.51 ms for M=10→32, then +95% at M=33, +49% at M=65 — i.e. time =ceil(M/32)× per-tile cost. At M=32 the op runs at ~92 GB/s, far under the ~400 GB/s roofline, so it's tile/compute-bound with real headroom (not bandwidth-bound) — halving the tile actually helps rather than just re-reading weights.Change
qmm()(host):bm = (mode == "affine" && M <= 16) ? 16 : 32, and threadBMinto the kernel name + template instantiation. Thebm == 32path is byte-identical to before, so all existing kernels resolve unchanged.quantized.metal: two_bm-suffixed instantiation macros, plusaffine_qmm_t/affine_qmm_ninstantiated atBM=16(BK=BN=32).The
BM=16tile reuses the existingqmm_t_impl/qmm_n_impl(already templated onBM), so there is no new kernel logic.BM=8is intentionally not added:BM/wm = 8/2 = 4rows per simdgroup breaks the 8-row MMA fragment tiling (TM = BM/wm/8), soBM=16(TM=1) is the smallest valid tile.Results (M3 Max, min of 15×30, baseline = BM=32 plateau at M=24)
M≥17 is unchanged; the
qmv/qmv_widepaths (M below the batch-limit) are untouched.The win is size-dependent. Comparing
BM=16(M=16) againstBM=32(M=17) inside one run, across shapes:So it pays ~1.5× on MLP- and lm_head-class layers and is neutral below ~58M weight elements: the smaller shapes are already flat in M across the whole band, i.e. not tile-quantization-bound to begin with. The two headline shapes above are both large, so they overstate how general the win is.
Correctness
max|Δ| = 0) for 4-bit and 2-bit at M = 10 / 12 / 16 / 17 / 32.python/tests/test_quantized.py: 32 passed, 2896 subtests passed.Scope
n > 1sampling (M in the[limit, 16]band), on layers large enough to be tile-bound — see the size table above; neutral on smaller ones. It deliberately does not touch theqmv_wideslope where single-stream speculative-decode verify (M = 2–6) lives — that path is already well-tuned (also noted in 2-bit quantized_matmul loses its size advantage over 4-bit at M >= 3 #3852).[limit, 16]band itself is set byget_qmv_batch_limit. A sweep of where theqmv_wide↔qmmcrossover actually sits (7 shapes, both bit-widths) is in the thread below; it does not change this PR, but it does show the current limit sits above the crossover on large layers and below it on small ones.Refs #3852.