[Performance] gather_qmm at M=1 streams 4-bit expert weights at ~half the bandwidth of M=4 (MoE decode)
Summary
On an M3 Ultra, mx.gather_qmm with a single token (M=1) streams 4-bit expert weights at roughly half the bandwidth the same call reaches at M=4, and well below what a plain read of the same bytes achieves. For MoE decode this matters: on GLM-5.3-Flash (288 experts, top-8, hidden 4096, intermediate 2048, 42 MoE layers) the routed-expert matmuls are ~40% of each decode step.
Numbers
One MoE layer's worth of random 4-bit experts (mx.quantize, group 64, affine), bf16 activations, top-8 routing, timed with mx.synchronize() around mx.eval, median of 60. "Effective GB/s" = expert weight bytes touched (14.2 MB per expert incl. scales/biases) ÷ time.
| call | time | effective bandwidth |
|---|---|---|
M=1, gather_qmm ×3 (up, gate, down), isolated |
0.439 ms | 258 GB/s |
| M=1, same, 42 layers chained without intermediate sync (launch cost amortised) | 0.292 ms / layer | 387 GB/s |
| M=4, isolated | 0.919 ms | 493 GB/s |
| M=4, chained ×42 | 0.836 ms / layer | 542 GB/s |
M=1, 8 separate quantized_matmul calls instead of gather |
1.891 ms | 60 GB/s |
reference: sum() over 512 MB bf16 |
454–514 GB/s |
sorted_indices=True makes no difference at M=1 (0.442 ms). The per-expert loop is launch-bound, so the gathered kernel is structurally the right thing; it's the M=1 occupancy of the gathered kernel that looks weak. At the model level, a cached forward over S tokens costs 1.23× / 1.67× / 2.46× / 3.69× a one-token step for S = 2 / 3 / 5 / 8, which is the same slack seen from the other side.
Environment
M3 Ultra, 512 GB; macOS 26.4; mlx 0.32.2; Python 3.13.
Script
moe_bench.pyimport statistics, time, random
import mlx.core as mx
HID, INTER, E, TOPK, GS, BITS = 4096, 2048, 288, 8, 64, 4
mx.random.seed(0)
def qexperts(out_dim, in_dim):
ws, ss, bs = [], [], []
for _ in range(E):
w, s, b = mx.quantize(mx.random.normal((out_dim, in_dim)).astype(mx.bfloat16), group_size=GS, bits=BITS)
ws.append(w); ss.append(s); bs.append(b)
out = tuple(mx.stack(t) for t in (ws, ss, bs)); mx.eval(out); return out
gate, up, down = qexperts(INTER, HID), qexperts(INTER, HID), qexperts(HID, INTER)
bytes_per_expert = sum(int(t.nbytes) for t in (*gate, *up, *down)) / E
def gq(x, wsb, idx, sort=False):
w, s, b = wsb
return mx.gather_qmm(x, w, s, b, rhs_indices=idx, transpose=True, group_size=GS, bits=BITS, sorted_indices=sort)
def moe(x, idx): # SwitchGLU order: up, gate, silu(gate)*up, down
x = mx.expand_dims(x, (-2, -3))
u, g = gq(x, up, idx), gq(x, gate, idx)
return gq((g * mx.sigmoid(g)) * u, down, idx).squeeze(-2)
def timeit(fn, *a, n=60):
for _ in range(8): mx.eval(fn(*a))
ts = []
for _ in range(n):
mx.synchronize(); t0 = time.perf_counter(); mx.eval(fn(*a)); mx.synchronize(); ts.append((time.perf_counter() - t0) * 1e3)
return statistics.median(ts)
for M in (1, 4):
x = mx.random.normal((M, HID)).astype(mx.bfloat16)
idx = mx.array([sorted(random.sample(range(E), TOPK)) for _ in range(M)]); mx.eval(x, idx)
ms = timeit(moe, x, idx); tb = M * TOPK * bytes_per_expert
print(f"M={M} isolated: {ms:.3f} ms -> {tb/ms/1e6:.0f} GB/s")
idxs = [mx.array([sorted(random.sample(range(E), TOPK)) for _ in range(M)]) for _ in range(42)]; mx.eval(idxs)
def chain(x):
h = x
for i in idxs:
h = h + moe(h, i).sum(1) * 0.0
return h
ms = timeit(chain, x, n=10) / 42
print(f"M={M} chained x42: {ms:.3f} ms/layer -> {tb/ms/1e6:.0f} GB/s")Question
Is there a better dispatch for the M=1 gathered case (e.g. a gather_qmv fast path, or something like the qmv_fast alignment conditions that the gathered kernel doesn't take), or is a batch-1 fused expert kernel (gate+up+activation+down in one launch) the way to get closer to the bus? Happy to test patches on this machine.
Source: ml-explore/mlx