Small-L fused expert gather for fine-grained MoE decode — measured expert-overlap data on GLM-5.2 (34% adjacent-token, depth-graded to 52%)

[anicaise-ai]

Context

On fine-grained MoE models (GLM-5.2: top-8 of 256 routed experts, 78 layers), decode-time forwards at small sequence lengths (L=2–4) cost nearly L× the L=1 forward on Apple Silicon, because each position routes to a mostly-disjoint expert set and routed-expert weights dominate memory traffic. This makes speculative decoding (including GLM-5.2's own MTP head, which Zhipu designed for it) structurally unprofitable at batch 1 — the verify forward eats the speedup.

Measured on M3 Ultra 512 GB, GLM-5.2 mxfp4 (via mlx-lm 0.31.3 + the community glm_moe_dsa port):

• L=1 decode forward: ~51 ms/token (≈19.7 tok/s end-to-end)
• L=2 verify forward: ~112 ms ≈ 2.2× L=1 (mean over ~700 speculative cycles)
• MTP head accept rate: 80–92 % greedy — draft quality is not the limiter
• Long prefill amortizes experts fully (~164 tok/s cold at L≈13k), so this is specifically a small-L gap

The data: adjacent-token expert overlap is substantial and depth-graded

We instrumented the router on real greedy decode traffic (256 tokens × 75 MoE layers). Adjacent tokens share far more experts than independent streams would (~3 % expected at random):

MoE layers  3-17 : mean 21.1%  (min  3.8%, max 35.3%)
MoE layers 18-32 : mean 32.2%  (min 23.6%, max 43.1%)
MoE layers 33-47 : mean 32.5%  (min 26.4%, max 38.3%)
MoE layers 48-62 : mean 39.8%  (min 31.8%, max 42.8%)
MoE layers 63-77 : mean 44.7%  (min 37.1%, max 52.1%)
overall: mean 34.1%, median 37.5% → 13.3 unique experts per L=2 pair (floor 1.66× vs 2.2× measured)

Two implications:

1. There is ~25 % of the L=2 verify cost recoverable by a fused small-L expert gather: compute the union of routed experts across the L positions, stream each selected expert's weights once, apply to whichever positions routed to it (instead of the current per-position gather paths in gather_qmm-based switch layers).
2. The win is depth-weighted: the last ~30 layers already sit at 40–52 % overlap, so even a partial implementation fused only on deep layers captures most of the recoverable bytes.

With verify at its 1.6× floor, MTP@90 % accept turns into a ~+10–15 % net decode win on these models — and CUDA stacks (sglang/vllm grouped-GEMM) already realize MTP gains on the same architecture, so this is an ecosystem gap rather than a conceptual wall.

Ask

Is a small-L (2–8) union-of-experts fused path in gather_qmm / the quantized switch-layer path something the team would consider? Happy to:

• run benchmarks on M3 Ultra 512 GB with quick turnaround (reproducible harness: fixed prompts, temp-0 diff protocol, per-cycle timing),
• share the raw per-layer overlap distributions and the 3-line env-gated router instrumentation,
• test against a working GLM-5.2 MTP implementation (feat: GLM-5.2 native MTP speculative decoding (external head) jundot/omlx#2076) that is currently parity-bound by exactly this cost (analysis in Fused small-L expert gather: fine-grained MoE (GLM-5.2 8/256) makes L=2 verify forwards cost ~2.2x L=1, capping speculative decoding jundot/omlx#2077).

[katlun-lgtm]

This is really well characterized — the depth-graded overlap curve (34% mean rising to ~52% in the deep layers) is the part that makes it actionable, since it says a partial implementation fused on just the last ~30 layers captures most of the recoverable bytes.

One framing that might help scope the kernel: the small-L union path is essentially a tiny grouped/segmented GEMM — sort the L·top_k (position, expert) pairs by expert, stream each unique expert's weights once, matmul over the ragged set of positions that routed to it, scatter back. The interesting Metal question is the raggedness: at L=2–8 the per-expert groups are tiny (often 1–2 positions), so a grouped-GEMM that pads to max-group-size would hand a lot of the win back. Does the win survive once you account for gather/scatter + padding at these tiny group sizes, or does it need a bespoke kernel that keeps each expert's position list variable-length?

I've done some quantized-matmul kernel work in mlx core (the mxfp4 matvec correctness fix, #3804), so I'd be happy to help prototype the small-L union gather on the gather_qmm side. I'm on an M3 Max 64GB, not a 512GB Ultra, so I can't reproduce GLM-5.2 — but I could build and iterate the kernel on a smaller fine-grained MoE that fits 64GB and share the small-L overlap + timing, and you'd validate the real GLM-5.2 win on the Ultra. That splits it fairly cleanly along the hardware we each have.

Two questions to make that useful:

1. You put ~25% of L=2 verify cost as recoverable — is that derived purely from the expert-byte overlap (13.3 unique / 1.66× vs the 2.2× measured), or does it already net out the non-expert costs (attention, router, shared/dense MLP) that a fused gather can't touch? Trying to pin the realistic ceiling before writing a kernel.
2. What would you consider a faithful smaller proxy for GLM-5.2's routing granularity (top-8/256) — something like OLMoE (8/64) or Qwen3-MoE (8/128)? I'd rather prototype on a model whose adjacent-token overlap profile actually resembles yours than pick one that flatters the result.

Happy to share the reproducible small-rig harness the other direction too.

[anicaise-ai]

Thanks — these are exactly the right two questions, and answering them properly made us re-measure. Full disclosure up front: the re-measurement moves the target, and I'd rather hand you that before you write Metal than after.

Q1: the ~25% was pure expert-byte arithmetic — and it no longer holds on current mlx

The 25% figure was derived purely from the byte overlap: 13.3 unique / 16 slots → 1.66× "floor" applied to the whole forward, vs the 2.2× measured. It did not net out attention, router, shared/dense MLP. Your instinct to pin that down was right.

So today we isolated the actual production path: one real GLM-5.2 MoE layer's routed bank (mxfp4, 256×3×6144×2048, ~20 MB/expert), loaded into mlx-lm's SwitchGLU/gather_qmm exactly as the model runs it, mlx 0.32.0, M3 Ultra 512 GB. Indices follow the measured 34% adjacent-overlap chain. "Union floor" = the same bank hit as a single position routed to the union set (streams each unique expert once — what an ideal fused kernel would approach at the byte level):

config	uniques	time	vs L=1	union floor	recoverable
L=1 top-8	8	0.55 ms	1.00×	—	—
L=2 chained overlap	13	0.72 ms	1.30×	1.16×	~10%
L=3	18	0.85 ms	1.55×	1.42×	~8%
L=4	23	1.04 ms	1.89×	1.71×	~10%
L=6	32	1.40 ms	2.54×	2.23×	~12%

Two conclusions we did not expect when we opened this thread:

1. The routed-bytes share of the thesis holds: L=1 switch time × 75 MoE layers ≈ 39 ms of the ~51 ms full decode forward, i.e. ~77% of decode is the routed path (matches the config-derived byte budget, ~72%).
2. But current gather_qmm is already within ~10% of the union-bytes floor at every small L. The 2.2× L=2 forward we published was measured 2026-07-02 on mlx-lm 0.31.3 + the community glm_moe_dsa port. Most of that gap has since been closed elsewhere in the stack — on today's mlx the switch path scales 1.30× at L=2, not 2.2×. A union-dedup kernel's realistic end-to-end ceiling is now single-digit % at L=2, maybe 8–10% at the L=4–8 depths (where GLM's MTP head was actually trained to operate), not 25%.

This also answers your ragged-vs-padded question at the byte level: with the existing path within ~10% of floor, padding strategy can't move much. The residual worth chasing is absolute small-row bandwidth: L=1 streams 160 MB in 0.55 ms ≈ 290 GB/s effective (~36% of the machine's peak), improving to ~430 GB/s at L=2–6. A kernel that lifts small-row gather-matvec bandwidth helps every decode, spec or not — that's closer to #3757 (gather_qqmm) territory than to union dedup.

Where the multi-row cost actually lives now, per production traces (oMLX Lightning MTP on this box): a verify cycle at M≈2–3 still costs ~100 ms vs ~40–50 ms single-row — the residual ~2× is in the DSA attention/indexer path and per-cycle sync, not expert bytes (cf. jundot/omlx#2160, where verify rows landing on the prefill indexer kernel cost 5×). The ecosystem gap we described is real, but it has moved out of the expert gather.

Q2: measured answer — use Qwen3.6-35B-A3B, and here is its overlap curve

Neither OLMoE (top-8/64) nor Qwen3-30B (8/128) matches GLM's granularity, but Qwen3.6-35B-A3B does exactly: top-8 of 256 routed experts (tiny experts, 2048×512), and it fits your 64 GB (mxfp8 ≈ 37 GB, 4-bit ≈ 19 GB comfortable). We ran the same router instrumentation on it today — 256-token greedy decode, one prose + one code prompt, top-k recorded at all 40 MoE layers:

prose:  layers 0-7: 24.2%   8-15: 36.6%   16-23: 38.1%   24-31: 37.7%   32-39: 35.0%
        overall mean 34.3%, median 36.5% → 13.3 uniques/pair
code:   layers 0-7: 24.4%   8-15: 37.9%   16-23: 39.7%   24-31: 39.7%   32-39: 35.1%
        overall mean 35.4%, median 37.3% → 13.2 uniques/pair

Mean adjacent overlap 34.9%, 13.2 unique experts per L=2 pair — statistically identical to GLM-5.2 (34.1%, 13.3). The routing granularity (8/256) appears to determine the overlap statistics almost entirely. Shape differs mildly: GLM ramps monotonically with depth (21%→45–52%), Qwen plateaus at 35–40% from mid-depth — but for kernel iteration the group-size distribution is what matters, and it matches.

One caveat so the proxy doesn't mislead: on 35B-A3B the routed experts are only ~30% of active bytes per token (vs ~75% on GLM-5.2), so end-to-end speedups won't transfer — but per-switch-layer timings and overlap/raggedness stats will, which is what you need to iterate a kernel.

Happy to share both scripts (the 3-line router recorder + the bank microbench, they're short) and to run any kernel candidate against real GLM-5.2 banks on the M3 Ultra with quick turnaround. And if after the numbers above you'd rather aim at absolute small-row gather bandwidth than union dedup — we'd gladly bench that too; it's the bigger prize on this hardware.

[katlun-lgtm]

Thanks for re-measuring instead of defending the headline — that's the useful kind of answer, and it lines up with something I found from the dispatch side, so let me close the loop from both ends.

Your "L=2 is already within ~10% of the union-bytes floor" is exactly what the dispatch path predicts. The fused dedup kernel already exists — affine_gather_qmm_rhs (backend/metal/kernels/quantized.h) streams each unique expert once over expert-sorted rows — but small-L decode never reaches it. The fast path fires only on M==1 && B>=16 && right_sorted && B/E>=4 (backend/metal/quantized.cpp), and with E=256 that needs B≥1024, i.e. L≥128, to satisfy B/E>=4. So L=2–8 falls through to per-row gather_qmv, which re-streams the shared experts with no dedup — at L=2 that's ~16 slot-reads for your 13 uniques, which is most of the ~10–12% above floor you measured. (SwitchGLU also only sorts at indices.size>=64 — switch_layers.py:181 — so small-L never even builds the right_sorted layout; and there's a // TODO: Tune 16 and 4 here right above the guard.) So the ~10% is real and it's a guard/threshold relaxation, not a new kernel — but I agree ~10% doesn't justify a dedicated one.

Independent check on the traffic ceiling: I ran the same router instrumentation on Qwen3-30B-A3B (8/128, what I could run at the time) over 256 greedy-decode tokens — recoverable L=2 18.8%, L=4 35.4%, L=8 51.6% (bytes, so it's the union-floor ceiling your wall-time sits ~10% under). One divergence worth flagging: Qwen3-30B's overlap is flat across depth (~18–20% in every layer group), not depth-graded like GLM's climb to ~52% — so "fuse only the deep layers" is model-specific. Your Qwen3.6-35B-A3B (8/256) is the cleaner proxy and I'll switch to it.

So — sold on the pivot. Absolute small-row gather-matvec bandwidth (your L=1 ≈ 290 GB/s, ~36% of peak) is the bigger prize and squarely the lane I want to be in — closer to #3757 gather_qqmm than to union dedup, same family as #3804. I'll take you up on both scripts, iterate a kernel candidate on the 35B-A3B banks locally (64 GB), and hand you candidates to bench against real GLM-5.2 on the Ultra.

Two things before I start:

1. At L=1, is the 36%-of-peak dominated by the mxfp4 dequant throughput, the non-coalesced expert-row gather, or 8-row threadgroup under-occupancy? A trace that separates them would tell me where the kernel should push.
2. Does Add gather_qqmm #3757 gather_qqmm already target the M=1 decode shape, or only M>1 prefill? If it's prefill-only, the decode small-row path is the open slot; if not, I'd rather not duplicate it.

[anicaise-ai]

Great find on affine_gather_qmm_rhs and the B/E>=4 guard — that closes the "why are we near floor yet at 36% of peak" question from the dispatch side, and your flat-profile Qwen3-30B curve is a useful contrast (confirms depth-grading is model-specific, so partial-depth fusion isn't a portable strategy).

Both questions answered with measurements:

Q1: it's occupancy, by a wide margin

Decomposition on the real GLM-5.2 layer-40 gate bank (256×2048×6144, mxfp4 gs32), M3 Ultra, one experiment per hypothesis:

(a) k-scaling — occupancy (single position, k unique experts in one gather_qmm):

k	time	effective BW	% of ~819 GB/s peak
1	0.34 ms	20 GB/s	2.4%
8 (the decode case)	0.35 ms	155 GB/s	19%
32	0.55 ms	389 GB/s	48%
128	1.53 ms	558 GB/s	68%
256	2.67 ms	640 GB/s	78%

Time is nearly flat from k=1 to k=8 — at decode size the GPU is row-starved, not bandwidth-bound. Two corroborating points: a contiguous plain-qmv on 2048 rows (6.7 MB ≈ 8 µs of streaming at peak) still takes 0.24 ms, i.e. a ~30× latency/occupancy floor over the byte time; and gathering 8× the same expert (working set ÷8) costs the same as 8 distinct ones — time doesn't follow bytes at this scale.

(b) gather vs contiguous — indirection: same bytes, same 16384 rows, only the indirection differs: 0.274 ms gathered vs 0.238 ms contiguous → ~14% penalty. Real but secondary.

(c) dequant mode: identical shapes/indices, random weights: affine gs32 208 GB/s, mxfp4 195 GB/s (~6% slower), and mxfp8 (2× bytes) reaches 278 GB/s effective in only 1.36× the time — more bytes per row actually helps the utilization, which again points at occupancy rather than dequant throughput.

So the ranking is: occupancy ≫ gather indirection (~14%) > mxfp4 dequant (~6%). For the kernel that means the win is in getting more useful rows in flight per dispatch at M=1/k=8 — fusing the three GLU matrices into one dispatch (3× rows), processing several experts per threadgroup, and shaving the fixed dispatch floor — rather than in the memory pattern per se. One measurement caveat: these are mx.eval-per-call timings, so the absolute floor includes command-buffer submission; inside a full-graph decode step the CPU-side share shrinks, but the k-scaling curve is a GPU-side property and matches what we see end-to-end (75 switch layers ≈ 39 ms of a 51 ms forward).

Q2: #3757 does not cover the M=1 decode shape yet

Per the PR description, current gather_qqmm is a bare-minimum implementation: on Metal it falls back to qmv (CUDA to qmm_naive), with "specialized fast kernels in future" explicitly deferred. So the decode small-row slot is open — you wouldn't be duplicating, and given (a) above, a small-row kernel that attacks occupancy would slot exactly into the gap #3757 leaves.

Scripts

All three are here, standalone, with a README: https://github.com/anicaise-ai/moe-small-l-bench
(moe_router_overlap.py — the router recorder; bench_gather_glm_smallL.py — L=1..6 vs union floor on a real bank; gather_qmv_decomposition.py — the Q1 decomposition above.) Paths at the top point at local model dirs; the overlap recorder defaults to Qwen3.6-35B-A3B so it runs as-is on your 64 GB.

Standing offer unchanged: send kernel candidates (branch or patch) and we'll run them against the real GLM-5.2 banks on the M3 Ultra, usually same-day.

[katlun-lgtm]

That decomposition settles it — flat time from k=1 to 8 with a 30× floor on even the contiguous case is occupancy, full stop, and it changes what's worth building. Thanks for re-measuring twice now; this thread keeps getting more useful.

Reading the dispatch code against your numbers makes the mechanism concrete: gather_qmv launches (M, ceil(N/8), B) threadgroups of 64 threads, so your contiguous 2048-row probe got 256 threadgroups — nowhere near enough to fill an 80-core part — and the k=8 case still only 2048. All the parallelism comes from N/8 × B; the K=6144 reduction stays serial inside each threadgroup. That's why time doesn't follow bytes at this scale.

So the candidate I'll prototype is split-K for gather_qmv: grid (M, ceil(N/bn), B·S) with each threadgroup reducing a K/S slice and a small combine pass. mlx already has exactly this machinery on the qmm side (qmm_splitk), so it's an extension of an existing pattern rather than a new invention. At S=4–8 the k=8 decode case lands at 8–16k threadgroups, and your own k-scaling curve says the GPU reaches ~50–68% of peak once it gets that many rows in flight. Gate/up stacking composes on top (same x, same indices — 2× N per dispatch, half the dispatches) and needs no kernel at all.

Plan: metal_kernel prototype against a synthetic 256×2048×6144 mxfp4 bank on my M3 Max first, using your decomposition harness as the baseline so the numbers are directly comparable. If it clears ~1.5× on the M=1/k=8 shape I'll port it into quantized.cpp behind an occupancy guard and take you up on the GLM-5.2 validation offer. Will report back either way — including if the prototype embarrasses the theory.

[katlun-lgtm]

Promised to report back even if the prototype embarrassed the theory — it did, and the failure is more informative than a win would have been.

I built the split-K gather-qmv as a metal_kernel prototype (affine-4bit stand-in, correctness at 5e-6 against the reference) and swept it properly: scalar and uint4-vectorized variants, 1–8 rows per threadgroup, split factors 1–16, up to 65k threadgroups on the k=8 decode shape. Everything lands within ±20% of the stock gather_qmm call (0.455 ms, ~166 GB/s on my M3 Max). More threadgroups actually made it slightly worse. So threadgroup-count occupancy is not the limiter, at least on this part.

What made your numbers click for me afterward: take the marginal rate of your own k-scaling instead of the averages — (k=256 − k=128) works out to ~8.9 µs per expert ≈ 750 GB/s, and (k=128 − k=32) ≈ 655 GB/s. Streaming beyond the first ~0.3 ms runs near peak. The "19% of peak at k=8" is a fixed per-dispatch latency/ramp floor, and no grid shaping inside the dispatch can remove a fixed floor — which is exactly what my flat sweep measured from the other side.

The good news is your suggestion list already contained the thing that does work: amortization. My M3 Max kscan has k=8 at 0.583 ms and k=16 at 0.831 ms — 2× the bytes in one dispatch for 1.43× the time. That makes gate/up stacking (one [2N, K] bank per expert, same x, same indices) worth ~1.4× on that pair with zero Metal — it's a SwitchGLU-level change. Down-proj stays separate for the data dependency, so a MoE layer goes from 3 dispatches to 2, with the big one better amortized.

Next step on my side: patch stacked gate_up into mlx-lm's SwitchGLU and measure end-to-end decode on Qwen3-30B-A3B (fits my 64 GB). If the ~1.4×-on-the-pair holds end-to-end it should be worth roughly 15–20% of MoE-projection time on your GLM-5.2 numbers too — and that one's easy for you to validate since it's a Python-level patch, no kernel build.

[katlun-lgtm]

Follow-through on the stacking idea, and it closes the loop on this whole line of inquiry — mostly by deflating it, including my own last theory.

Three measurements:

1. Gate/up stacking end-to-end: patched mlx-lm's SwitchGLU to use one [E, 2H, K] bank (one gather dispatch instead of two, greedy output identical) and ran Qwen3-30B-A3B decode — 102.2 vs 102.5 tok/s. Flat.

2. The interesting one: an in-graph chain of 16 GLM-shaped MoE layers (2048×6144 banks, k=8), split vs stacked, all inside single eval graphs. 450 vs 439 µs/layer — 1.024×, bitwise identical outputs. So inside one command buffer the dispatches pipeline and the ~0.3 ms "per-dispatch floor" we were both reasoning from is an artifact of eval-per-call timing. My marginal-rate reading of your k-scan was right about there being a fixed component, wrong about it surviving into a real decode graph.

3. Which reframes the envelope: in-graph, the existing gather path streams 107 MB/layer in 450 µs ≈ 238 GB/s on my 400 GB/s part, and your GLM end-to-end (~0.52 ms/layer for ~160 MB) lands at the same ~60-75% of peak on the Ultra. Split-K, threadgroup count up to 65k, vectorized loads, dual accumulator chains, and dispatch fusion are all now measured no-ops on this shape.

So the honest summary for anyone landing on this thread: at batch-1 decode, current mlx gather_qmm/qmv is already within ~1.3× of its practical in-graph envelope on these shapes, the L=2 verify cost is mostly genuine unique bytes, and the recoverable prize for any small-L restructuring is far smaller than the numbers at the top of this discussion suggested. What's left — pushing the 4-bit gather inner loop from ~60% toward 80%+ of peak — is a qmv streaming-efficiency question that doesn't depend on L or expert overlap at all, and probably belongs to the maintainers' qmv tuning rather than a new kernel.

Your instinct to re-measure twice was the right call both times; happy to have the negative results on the record next to the data. If you want the SwitchGLU stacking patch anyway (it's harmless, bitwise-identical, and one dispatch fewer), it's ~20 lines — but I wouldn't claim perf for it.

[anicaise-ai]

Confirming your results on the 80-core M3 Ultra, on the real GLM-5.2 mxfp4 banks — and closing this out from our side.

First, a correction we owe the thread: the ~0.3 ms/dispatch "floor" entered this discussion through our eval-per-call decomposition, and your in-graph probe shows it was a timing artifact, not a hardware property. Worth stating explicitly since we propagated it.

Replication of your probe #2 on production weights (ingraph_chain_glm.py): 16-step in-graph chain of real GLM-5.2 gate+up banks (W(256, 2048, 768) mxfp4, k=8, 107 MB touched/step), single mx.eval per iteration, stacked output bitwise-identical to split (max|diff| = 0). Two runs, M3 Ultra 80-core (819 GB/s peak):

timing mode	time	effective BW
eval-per-call, 1 gather dispatch (k=8, 53 MB)	0.327 / 0.323 ms	~165 GB/s (20% peak)
in-graph chain, split (2 dispatches/step)	210 µs/step (105 µs/dispatch)	509 GB/s (62% peak)
in-graph chain, stacked (1 dispatch/step)	211 / 214 µs/step	501–507 GB/s

split/stacked ratio = 0.984× / 0.997× — same conclusion as your 1.024× on the M3 Max: within noise of 1.0, stacking buys nothing once you're inside one command buffer. And the eval-per-call floor reproduces here too (~0.33 ms on an 80-core, vs your 0.455 ms on the Max) while the same dispatch costs 105 µs in-graph — a ~3.1× measurement artifact, scale-invariant across both parts.

So the full picture holds on production weights: batch-1 decode gather_qmm is already at ~62% of peak in realistic graphs, small-L restructuring (fused gather, split-K, stacking, union-dedup) has no recoverable headroom, and whatever remains between 62% and peak is generic qmv tuning that doesn't depend on L or expert overlap. No need for the SwitchGLU stacking patch on our side — thanks for offering it, and for running the negatives to ground and reporting them. We'll point omlx#2077 at this thread so it can be closed with a clean rationale.

[katlun-lgtm]

Agreed on all of it — and thanks for replicating the in-graph probe on production weights; that's what makes this a trustworthy negative rather than just my M3 Max noise. Good outcome even as a "don't build it": ~62%-of-peak-in-real-graphs is the number to walk away with, and it kills a kernel that'd have been a lot of Metal for nothing. If anyone picks up the generic qmv gap, I'll ping this thread.