perf(moe): attribute the MoE decode gap at the M = 1 expert gather #1859

Problem / Background

Two MoE families reach gather_qmm with M = 1 on every decode token and never reach the fused decode-MoE kernel. That much is a code fact, verified below against 60341873. Whether it is the dominant term in a measured decode gap is not established, and this issue asks for a profiled attribution before any change is written.

The shape. SwitchGLU::forward (src/models/switch_layers.rs:939-963) expands x with two expand_dims at switch_layers.rs:946-947, producing [n, 1, 1, K], so gather_qmm sees M = 1 and B = n * top_k. The GB10 attribution record for #1799 states this and names its consequence on CUDA. docs/benchmark_results/laguna-dflash-verify-cost-gb10-2026-09-11.md:96: "SwitchGLU::forward expands x to [n, 1, 1, K], so gather_qmm sees M = 1, B = 8n and the routed experts take qmm_sm80 at every width including classic". qmm_sm80 is the prefill GEMM; docs/benchmark_results/qmm-sm121-tile-tuning-gb10-2026-07-10.md:24 records that single-sequence decode is supposed to take qmv. So on CUDA the routed experts are already known to run a tiled GEMM at one row per (token, expert) pair during classic decode, on a different backend and a different checkpoint from the observation below.

What is specific to these two families is that neither has the fused decode-MoE kernel as an escape hatch. Eighteen families wire SwitchGLU::forward_fused_kernel (grep -rln forward_fused_kernel src/models/, excluding the definition and a test file); deepseek_v4 and inkling are not among them.

• deepseek_v4 does not use SwitchGLU at all. It carries its own LimitedSwitchGlu::forward (src/models/deepseek_v4_moe.rs:242-252) with the same two expand_dims at deepseek_v4_moe.rs:243-244, and it passes sorted: false unconditionally. The module doc is explicit at deepseek_v4_moe.rs:49-50: "The fused MoE decode kernel (MLXCEL_FUSED_MOE) is deliberately not wired: it has no clamp stage."
• inkling calls SwitchGLU::forward_with_expert_scales (src/models/inkling/mlp.rs:217-222). That method has no fused-kernel branch at all: it delegates to SwitchGLU::forward only when both sidecars are None (switch_layers.rs:982-984) and otherwise runs the unsorted gather (switch_layers.rs:986-989), documented as deliberate at switch_layers.rs:970-972.

Sorting does not rescue decode either. SwitchGLU::forward sorts only when n_tokens * top_k >= 64 (switch_layers.rs:940-944), which one decode token never reaches, and both forward_with_expert_scales and LimitedSwitchGlu::forward are unsorted by construction.

Corroborating observation (Metal, one session, unprofiled)

An Apple Silicon session on a Mac Studio M3 Ultra benchmarked deepseek-v4-flash-4bit, inkling-small-mlx-4bit and inkling-small-nvfp4 against mlx-vlm on 2026-09-08 and reported it today. Matched shape on both sides: 512-token prompt, 128 generated, warmup pass, EOS suppressed, no chat template. mlxcel decode landed at 16 to 23 percent of the mlx-vlm baseline while prefill was 55 to 82 percent, consistent across both architectures and across text and image modalities. Decode being the affected phase while prefill is much less so points away from the GEMMs and toward the per-token expert gather, which is why this is recorded rather than dropped. It is a number plus an inference, not a mechanism, so it supports the section above rather than standing on its own.

That session reports both checkpoints as 256-expert top-6. Neither checkpoint is present on the GB10 host (models/mlx holds 190 checkpoints, none matching), so the expert counts and routing widths are unverified here and should be read from each config.json at re-measurement time.

Limits, which are severe

• Metal only, one run per cell, not profiled, no kernel-level attribution.
• The two Inkling rows are not a quantization comparison. inkling-small-mlx-4bit carries recipe: experts_only while the NVFP4 build quantizes a different scope, so they are one checkpoint under two non-comparable recipes and must not be presented as 4-bit against NVFP4.
• Measured on mlxcel b8fdddfa, which is 124 commits behind 60341873, against MLX pin 9a795735. Main is now on 81ba1c6a, and the Metal quantized kernels changed in between. 8ad62d4b (fix(mlx): bump MLX pin to 81ba1c6a so mxfp8 block scales round up #1772) rewrote 150 lines of src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp, including the tile choice, which went from int bm = 64; to int bm = (transpose && M <= 32) ? 32 : 64; under the upstream comment "Use smaller bm when one block covers all of M". Decode is M = 1, so that switch fires on exactly the shape this issue is about and did not exist when the measurement was taken. Eight diff lines in the same file touch gather_qmm: gather_qmm_t_nax and gather_qmm_n_nax gained a global_scale parameter and a _hgs kernel-name suffix. Do not write this off as "Metal was untouched". The mlxcel-side commits in the window are largely CUDA and ROCm, but the pin bump reaches every backend, and src/models/switch_layers.rs itself moved 231 lines (Laguna NVFP4 and Kimi K3 SiTU, from feat(models): add the Laguna family with compressed-tensors NVFP4 experts #1739 and feat(models): add the Kimi K3 text backbone #1741).
• The one thing that did not move: the x_exp expansion and the do_sort threshold are byte-identical across b8fdddfa..60341873, and the deepseek_v4 and inkling call sites are unchanged. The hypothesis is still live on main; the kernels underneath it are not the ones that were measured. That is a reason to re-measure first, not a reason to discount the lead.

Prior art, so it is not re-derived

• perf(models): general MoE decode gap vs mlx-lm (dots.llm1 0.51x, qwen3_moe 0.69x, nemotron MoE) #268 (closed, docs/benchmark_results/moe-decode-gap-investigation.md) is the general MoE decode gap against mlx-lm on Metal: worst case dots.llm1 at 0.51x, qwen3_moe 0.69x, nemotron-h 0.60x. It concluded the residual is diffuse per-op execution overhead with no single fusable lever, and it named a fused decode-MoE kernel as the highest-potential path. Every family in that table has that kernel today. 0.16 to 0.23x sits far outside its range, which is consistent with these two families lacking the escape hatch the others received, and is why this is not simply a duplicate of perf(models): general MoE decode gap vs mlx-lm (dots.llm1 0.51x, qwen3_moe 0.69x, nemotron MoE) #268.
• perf(moe): batched decode scales 1.55x against 3.25x on dense #1616 (closed) found MoE batched decode scaling to 1.55x against 3.25x on dense, traced it to the fused kernel being a B=1-only path, and already analysed the do_sort threshold and the one-row-per-(token, k) expert traffic. Its rejected alternative (forcing supports_batching() false for MoE) should not be relitigated.
• perf(speculative): attribute the Laguna DFlash verify round's fixed and per-row cost on GB10 #1799 is the source of the M = 1, B = 8n finding quoted above. perf(cuda): evaluate raising the MLX graph ops budget on GB10 (sm_121) #1798 is the adjacent GB10 investigation from the same sweep, on the MLX graph ops budget; it is useful for the host, method and noise-floor conventions used there, but it is not where the gather-shape finding came from.

Proposed Solution

Measure before changing anything, in this order.

1. Re-measure on current main, on both backends. Repeat the reporting session's matched shape on 60341873 or later on Metal, and add a CUDA arm on GB10, which requires fetching both checkpoints. Report n of at least 3 with per-arm min-to-max ranges and treat overlapping ranges as no effect, following the convention in docs/benchmark_results/cuda-graph-budget-gb10-2026-09-12.md. Include one family that does have the fused kernel as a control, so a general regression is distinguishable from a family-specific one.
2. Attribute the gap by kernel, not by inference. Confirm which kernel the routed experts actually take at decode on each backend (MLXCEL_TRACE_ARCH prints the quantized-matmul path on CUDA, per docs/environment-variables.md:504), and split the per-token cost into router, expert gather and combine with the mlxcel-gpu-profiling hooks already used in docs/benchmark_results/moe-decode-gap-investigation.md. The deliverable is a committed record under docs/benchmark_results/ naming the kernels with their per-token launch counts and milliseconds.
3. Only then decide what to build. Candidates, in the order perf(models): general MoE decode gap vs mlx-lm (dots.llm1 0.51x, qwen3_moe 0.69x, nemotron MoE) #268 and perf(moe): batched decode scales 1.55x against 3.25x on dense #1616 already ranked them: give these two families a fused path (which for deepseek_v4 means a clamp stage the current kernel lacks, per deepseek_v4_moe.rs:49-50), deduplicate expert planes across the gathered rows, or re-tune the do_sort threshold. None should start before step 2 is committed.

Named risk. The obvious wrong turn is accepting the M = 1 hypothesis because it is tidy and fits the numbers. The #1798 and #1799 investigations both opened with a plausible mechanism that the source then contradicted: #1799's own record corrects two of its issue's premises at laguna-dflash-verify-cost-gb10-2026-09-11.md:96, including a supposed fp_qmv to qmm_sm80 family switch on the dense projections that turned out not to exist on that pairing because those projections are plain bf16. This issue requires a profiled attribution; a mechanism that merely fits is not an answer.

Scope

In scope: a committed benchmark and profile record under docs/benchmark_results/; the kernel-level attribution on Metal and on CUDA; the decision about what, if anything, to build, written down with its evidence.

Out of scope: any edit to src/models/switch_layers.rs, src/models/deepseek_v4_moe.rs or src/models/inkling/mlp.rs before step 2 is committed; batched decode, which #1616 owns; the bench_mlxlm.py memory budget, filed separately.

Acceptance Criteria

• A record under docs/benchmark_results/ reports decode and prefill against the mlx-vlm baseline for both families on current main, on Metal and on CUDA, with host, MLX pin, tree commit, n and per-arm ranges.
• The record names the kernel the routed experts take at decode on each backend, with per-token launch counts and milliseconds, and gives the router / gather / combine split.
• The record includes a fused-kernel family as a control, so a family-specific gap is distinguishable from a general one.
• The record states explicitly whether the M = 1 gather shape is or is not the dominant term, naming the measurement that settles it.
• Any resulting code change is reached by the real decode path used by mlxcel-server and mlxcel generate, not by a bench-only path, and either preserves greedy temp-0 token identity or documents exactly where it does not.
• A measured negative closes this issue, provided the measurement meets the criteria above.

Verification

# Neither checkpoint is on the GB10 host; fetch before the CUDA arm.
ls models/mlx | grep -iE "deepseek-v4|inkling"

make release-cuda
# Metal: cargo build --release --features metal,accelerate --bin mlxcel-bench-decode

# Paired arms, same shape on both sides, n >= 3 after a discarded warm-up.
./target/release/mlxcel-bench-decode --model models/mlx/<ckpt> --max-tokens 128 --warmup-tokens 20
MODELS_DIR=models/mlx ./scripts/bench_mlxlm.py models/mlx/<ckpt>

# Which quantized-matmul path the experts actually take (CUDA).
MLXCEL_TRACE_ARCH=1 ./target/release/mlxcel-bench-decode --model models/mlx/<ckpt> --max-tokens 4

# Control: a family that does have the fused kernel, with it forced off.
MLXCEL_FUSED_MOE=0 ./target/release/mlxcel-bench-decode --model models/mlx/qwen3-30b-a3b-4bit --max-tokens 128