[性能] gather_qmm 在 M=1 串流中使用 4 位专家权重,大约为 M=4 (MoE 解码) 的一半带宽

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((INTER, INTER)).astype(mx.bfloat16) for _ in range(8): mx.eval(mx.expand_dims(x, (-2, -3))) t0 = time.perf_counter() for _ in range(42): mx.eval(mx.expand_dims(moe(x, random.randint(0, E - 1)), (-2, -3))) t1 = time.perf_counter() print(f"M={M}, {timeit(moe, x, random.randint(0, E - 1))} ms, {bytes_per_expert * E / (t1 - t0):.2f} GB/s")