[Performance] gather_qmm at M=1 streams 4-bit expert weights at ~half the bandwidth of M=4 (MoE decode)

import 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")