Repository navigation
perf(models): general MoE decode gap vs mlx-lm (dots.llm1 0.51x, qwen3_moe 0.69x, nemotron MoE) #268
Description
Problem / Background
On M1 Ultra, dots.llm1.inst-mixed-4-6bit decodes at ~13.2 tok/s in mlxcel versus ~25.7 tok/s in mlx-lm (0.51x, a ~2x gap). This is NOT a Mamba2 hybrid. dots.llm1 is a DeepSeek-V3-style MoE (sigmoid router with e_score_correction_bias selection-only bias, unbiased gather, norm_topk_prob, routed_scaling_factor) with mixed 4-bit/6-bit expert quantization (v_proj/down_proj at 6-bit). For comparison on the same host and harness, the other MoE family (LFM2-8B-A1B, lfm2-moe) is at parity (~1.00x) and dense granite-3.3-2b is at ~0.99x, so the gap is specific to this model's routing / attention / mixed-precision-gather path, not a general MoE problem.
Measured on an Apple M1 Ultra (128 GB, macOS 26.5.1, MLX commit a6ec712) against mlx-lm 0.31.3. Decode throughput was measured with mlxcel-bench-decode (warmup pass + measured pass in one process, --no-chat-template) versus scripts/bench_mlxlm.py (mlx-lm's own prompt_tps/generation_tps), identical checkpoint, raw continuation prompt, 100 decode tokens, best-of-3 with cooldown, temp-0 greedy.
Proposed Solution
Investigate and root-cause the ~2x decode gap. Suspects to profile:
- SwitchGLU /
gather_qmmwith mixed 4/6-bit experts (the per-tensor bit inference added in PR feat(models): add dots.llm1 (dots1) MoE #263 /switch_layers.rs). - The attention path (whether dots.llm1 uses MLA and how the absorbed/decompressed projections are implemented vs mlx-lm).
- Router overhead (group_mask_scores, sigmoid + correction-bias selection).
- Any extra
eval()/ synchronization points or avoidable astype/graph nodes in the decode step.
Use scripts/bench_decode.sh / mlxcel-bench-decode and scripts/bench_mlxlm.py; the mlxcel-gpu-profiling skill may help localize the hotspot.
Acceptance Criteria
- Dominant cause of the ~2x decode gap identified and documented.
- Gap closed meaningfully (target parity-ish with mlx-lm).
- Greedy temp-0 output unchanged.
Activity
perf(models): gate nemotron_h Mamba2 per-mixer eval to M5 Max (#271)
inureyes commentedon Jun 13, 2026
Investigation: this is a general MoE-decode gap, not dots.llm1-specific
Benchmarked several MoE families on M1 Ultra (decode tok/s,
mlxcel-bench-decodevsmlx-lm0.31.3, raw prompt, 100 tok after warmup, temp 0):The gap scales with MoE compute and appears across three different MoE implementations (qwen3_moe, the dots1 DeepSeek-V3 router, nemotron's local
SwitchMLP).lfm2_moe(smallest active set) is at parity. So this is a general MoE-decode issue; dots.llm1 is just the worst case.Ruled out
eval()calls inqwen3_moe.rs/switch_layers.rsare inforward_profiled(profiling-only) and test code. Productionforwardis clean.dots1.rshas zeroeval(). So this is not a perf(models): gate Mamba2 mixer eval boundary to M5 Max #266-style fusion-breaking sync.src/models/switch_layers.rsSwitchGLU::forwardmatchesmlx-lm'sswitch_layers.pyexactly: samedo_sort = total >= 64threshold (so decode takes the unsorted path on both), samegather_qmm(transpose=True, sorted_indices=do_sort).e_score_correction_biasselection-only + unbiased gather +norm_topk_prob+routed_scaling_factor) andmoe_weighted_summatch the mlx-lm reference op-for-op.Localization (qwen3-30b-a3b,
MLXCEL_PROFILE_BLOCKS=1 MLXCEL_PROFILE_QWEN3_MOE_DETAIL=1)Per MoE layer (profile mode forces per-step eval, so absolute ms are inflated; the split is the signal):
gather_qmm~39%, router ~25%, combine ~20%, residual/overhead ~16%.Router + combine together (~45% of MoE) are many tiny ops on
[1, n_experts]/[1, top_k]tensors. On a single decode token these are latency/dispatch-bound, not throughput-bound.Conclusion / next steps
The production MoE decode path is structurally identical to mlx-lm and uses the same MLX
gather_qmmkernel (commit a6ec712), yet mlxcel is slower in proportion to MoE compute. The gap is execution efficiency in the decode MoE (per-op dispatch on small tensors + the gather_qmm path), not the algorithm. This needs a Metal-trace comparison against mlx-lm to attribute the difference (gather_qmm kernel time vs op-dispatch overhead vs lazy-graph fusion), then likely a fused router/combine path or a fused decode-MoE kernel. It is a deliberate kernel/fusion effort, not a one-line fix, so it is parked here with this diagnosis rather than rushed.Related: the nemotron-h decode gap (#267 follow-up) is the same MoE bottleneck; qwen3_moe is affected too. Whatever fix lands here should be validated across all three.
[-]perf(models): dots.llm1 decodes at ~0.51x of mlx-lm (MoE/MLA path)[/-][+]perf(models): general MoE decode gap vs mlx-lm (dots.llm1 0.51x, qwen3_moe 0.69x, nemotron MoE)[/+]inureyes commentedon Jun 13, 2026
GPU profiling: the MoE gap is GPU-bound dtype-conversion plumbing (773 AsType/token)
Profiled
qwen3-30b-a3b-4bitdecode on M1 Ultra (the fast-iterating MoE proxy; 0.69× of mlx-lm) using themlxcel-gpu-profilinghooks.Step 0 — it's GPU-bound, not CPU/dispatch
MLXCEL_PROFILE_PIPELINE_DETAIL:async_eval(GPU encode+execute) 20.3 ms/tok vsforward(Rust graph build) 0.8 ms/tok vsitem_wait0.14 ms. So ~95% of decode is GPU; the Rust/FFI graph-build is negligible. Command-buffer batching (MLX_MAX_OPS_PER_BUFFERswept 100–1000) is flat at ~47 tok/s, so it is not command-buffer-gap bound either.Step 1 — not at the memory wall
~3B active params at 4-bit ≈ 1.7 GB streamed/token; 47 tok/s × 1.7 GB ≈ 80 GB/s effective, ~16–20% of M1 Ultra GEMV peak (~450 GB/s). The GPU is mostly idle — it spends its 20.3 ms on many small kernels, not on streaming weights. (mlx-lm at 68.7 tok/s ≈ 117 GB/s, also well under the wall — MoE per-expert GEMVs underutilize the GPU on both, but mlxcel worse.)
Step 2 — the concentrated excess: AsType
One-decode-token op histogram (
MLXCEL_EXPORT_DECODE_DOTvs mlx-lmmx.export_to_dot), 48 layers:mlx-lm emits zero dtype conversions in its decode graph; mlxcel emits 773/token (16/layer) — ~80% of the total node excess. They are not from mlxcel Rust code (≤6 explicit
astypein the whole hot path, all in the unrelated KV-cache turbo module); they are inserted by the MLX C++ quantized-op layer. With 289 quantized ops/token (145 QuantizedMatmul + 144 GatherQMM), that is ≈2.7 AsType per quantized op — i.e. the activations entering/leavingquantized_matmul/gather_qmmare a different dtype than the quantized params, so MLX inserts conversions. mlx-lm keeps the whole decode in one dtype and pays none.This is the same root cause the
mlxcel-gpu-profilingskill documented for Gemma3n (#762: "AsType from f16/bf16 mixing"), here at larger magnitude because MoE op density is higher.Fix direction
Trace the decode dtype flow around the quantized ops and make the activation dtype match the quantized params so MLX inserts no
AsType(mlx-lm's invariant). Likely the samebf16/f16mixing class as #762 / [[bf16-quant-scales]]; start at the most frequent site (the per-layer q/k/v/o + router quantized matmuls and the gate/up/down gather_qmm).Caveat (from #762): MLX fuses elementwise
AsTypechains at eval, so the static node count overstates GPU kernels — Gemma3n recovered only ~2.6% from AsType removal. So validate the runtime payoff per change, do not assume 773→0 maps to the full 0.69×→1.0×. But it is the clear lead, and MoE density is higher than Gemma3n, so it warrants a measured attempt.Validation harness: re-bench qwen3-30b-a3b (cheap, 16 GB) for each change, then confirm on dots.llm1 + nemotron-h; greedy temp-0 output byte-identical; RMS < 5e-3 for any fused-kernel substitution.
17 remaining items
perf(moe): wire qwen2_moe to the fused decode-MoE kernel (#308)
perf(moe): wire mixtral to the fused decode-MoE kernel (#311)
perf(moe): implement the fused decode-MoE kernel on CUDA
perf(core): hardware-gated MLX_MAX_OPS_PER_BUFFER decode default (#360)
perf(moe): wire minimax to the fused decode-MoE kernel (#390)
inureyes commentedon Jun 22, 2026
Resolution
The MoE decode-gap work for this issue is complete and shipped default-on. Summarizing for closure.
What landed
A fused single-token decode-MoE Metal kernel, built and tuned over a series of merged PRs:
docs/benchmark_results/moe-decode-gap-investigation.md): the gap is GPU-bound diffuse per-op overhead in the single-token MoE step, not a forced-eval bug.docs/benchmark_results/fused-moe-decode-kernel-design.md).fused_moe_expert_kernelfoundation, correctness-validated byte-identical.gather_qmm(the key finding: a single fully-fused launch cannot win because the gate/up→down dependency pins each expert to one threadgroup; staging the swiglu activation in global memory across two dispatches lets every GEMV row run as an independent simdgroup across all cores).MLXCEL_FUSED_MOE_RELU2flag (neutral on nemotron, whose gap is the SSM/attention path).MLXCEL_FUSED_MOEflag documentation.MLXCEL_FUSED_MOEby default (closing perf(moe): validate fused decode-MoE on M5, then enable MLXCEL_FUSED_MOE by default #282).Subsequently the same kernel was wired to more families: qwen2_moe (#308), mixtral (#311), CUDA (#319), and minimax (#390), plus olmoe.
Per-model decode gains (M1 Ultra, single-token, default-on)
Acceptance criteria
Residual
The headline
dots.llm1figure improves from 0.51x to 0.53x of mlx-lm. The deeper finding (documented in the investigation report and the nemotron-h #284 profile) is that the remaining gap is fundamental single-token GPU underutilization across a diffuse set of small per-op overheads, with no remaining single high-value lever; an AsType-reduction campaign caps at roughly +2.6% on this class. Further single-token MoE decode speedup would need a different approach (e.g. batching/speculative decode) tracked separately, not more fusion.Closing as completed.
perf(models): fused single-launch xIELU Metal kernel for Apertus (#414)
feat(models): add Kuaishou Klear (Klear) text model support (#1044)
docs(moe): record that fused decode-MoE byte-identity is checkpoint-d…
perf(moe): route granite-4.0-h single-token MoE through the fused ker…