Repository navigation
epic(moe): wire the remaining unwired MoE families to the fused decode-MoE kernel #307
Description
Background
The fused single-token decode-MoE kernel (#268 series, steps 2b-3b) and the default-on flip (#282, PR #285) are wired into only four model paths: qwen3_moe, qwen3_next, dots1, gemma4. #286 wires qwen2_moe. Seven more MoE families reuse a SwitchGLU for the expert matmul but never call forward_fused_kernel, so they always fall back to gather_qmm regardless of MLXCEL_FUSED_MOE. This epic tracks wiring those remaining families.
The canonical one-branch pattern lives in src/models/qwen3_moe.rs, src/models/dots1.rs, and src/models/qwen3_next.rs: gate on single-token decode plus fused_moe_enabled() (src/models/switch_layers.rs:42), call experts.forward_fused_kernel(&x_flat, &topk_indices, &scores) (src/models/switch_layers.rs:276), and fall back to experts.forward(...) plus moe_weighted_sum(...) (src/models/switch_layers.rs:419). All target families use SwiGLU (SiLU) experts, so the kernel runs with the default swiglu activation (act=0); no GeGLU/relu2 template is needed. Kernel gating: affine 4/8-bit gate/up, 4/6/8-bit down, matching group_size, single-token decode only; it auto-falls-back for any unsupported config.
Two implementation classes
This is the key technical finding, verified in-tree.
Group A (direct one-branch wiring): the model's experts already use the shared crate::models::switch_layers::SwitchGLU, which exposes forward_fused_kernel. Wiring is the small one-branch change from #286/dots1.
minimax(importscrate::models::switch_layers::SwitchGLU).lfm2(importssuper::switch_layers::SwitchGLU; only the LFM2-MoE variant has the MoE block).qwen3_vl_moe(importssuper::qwen3_moe::SwitchGLU, which hasforward_fused_kernel).
Group B (needs SwitchGLU consolidation first): the model defines its OWN local SwitchGLU struct that does NOT expose forward_fused_kernel. These must migrate to the shared switch_layers::SwitchGLU (preferred: removes duplication) OR have a forward_fused_kernel method added to the local struct, BEFORE the one-branch dispatch works.
mixtral(localSwitchGLUatsrc/models/mixtral.rs:156).olmoe(localSwitchGLUatsrc/models/olmoe.rs:173).phimoe(localSwitchGLUatsrc/models/phimoe.rs:202).
Sub-issues
- perf(moe): wire qwen2_moe to the fused decode-MoE kernel #286 qwen2_moe (existing, primary deliverable)
- perf(moe): wire mixtral to the fused decode-MoE kernel #301 mixtral
- perf(moe): wire olmoe to the fused decode-MoE kernel #302 olmoe
- perf(moe): wire phimoe to the fused decode-MoE kernel #303 phimoe
- perf(moe): wire minimax to the fused decode-MoE kernel #304 minimax
- perf(moe): wire lfm2 (LFM2-MoE) to the fused decode-MoE kernel #305 lfm2 (LFM2-MoE)
- perf(moe): wire qwen3_vl_moe to the fused decode-MoE kernel #306 qwen3_vl_moe
Out of scope / related
- perf(nemotron-h): decode gap is MoE-block op-density (routed + shared expert), not SSM/attention #284: nemotron-h decode gap is shared-expert plus op-density, a different lever, tracked separately.
- perf(models): general MoE decode gap vs mlx-lm (dots.llm1 0.51x, qwen3_moe 0.69x, nemotron MoE) #268: umbrella MoE-gap issue.
- perf(moe): validate fused decode-MoE on M5, then enable MLXCEL_FUSED_MOE by default #282 and perf(moe): default-on MLXCEL_FUSED_MOE, validated on M5 (#282) #285: default-on with auto-fallback.
Definition of done
- All sub-issues closed.
- Every wired family validated at runtime, or its hardware blocker documented (minimax: MiniMax-Text-01 is about 456B params and does not fit 128 GB even at 4-bit; same class as hunyuan-moe in fix: two models regressed to FAIL in the 0.2.1 sweep (hunyuan-moe-a13b-bf16, qwen2.5-vl-3b-4bit) #294).
- The design doc "Models covered" list in
docs/benchmark_results/fused-moe-decode-kernel-design.mdreflects reality. - Default-on with auto-fallback still holds.
Activity
inureyes commentedon Jun 16, 2026
Epic implementation complete
All sub-issues implemented, validated on real checkpoints (M1 Ultra), and merged. Executed as a single serial chain (one worktree) because every unit edits the same "Models covered" doc list, which would otherwise conflict under parallel merges.
Sub-issues (6/6 merged)
Key outcome: the fused kernel is not a universal win, so it is now self-gating
Real-checkpoint measurement showed the two-kernel fused decode path wins only while
gather_qmmunderutilizes the GPU (small experts) and loses for large experts (phi-3.5-moe Dff 6400: -5.9%; mixtral Dff 14336: -21%) wheregather_qmmalready saturates. PR #311 addedMLXCEL_FUSED_MOE_MAX_DFF(default 4096, env-tunable):forward_fused_kerneldeclines above it and the caller falls back togather_qmm. The 4096 default was calibrated against measured break-even data and verified to keep every winner (including the pre-existing qwen3_next at Dff 2560, +8.7%) while excluding both regressions. Confirmed: with the guard mixtral/phimoe match their gather_qmm baselines, and raising the bound re-arms the slow path.Each Group B migration (qwen2_moe, mixtral, phimoe, olmoe) removed a duplicated local
SwitchGLU/SwitchLinearin favor of the sharedswitch_layersone, which gained two additive loaders along the way: a per-expert stacking fallback for theexperts.{idx}layout (#286) and overridable projection leaf names for Mixtral'sw1/w2/w3convention (#301).Pre-existing bugs found while validating (separate from the MoE work)
q_norm/k_normwere applied after the head reshape against hidden-sized weights. Fixed by normalizing before the reshape (matches ml-explore/mlx-lm). OLMoE now runs.Verification
every_variant_is_registered_for_archand theswitch_layerssuite (including the new stacking and projection-name tests) pass.Totals: 6 sub-issues across 1 wave; 1 follow-up bugfix merged (#317); 1 follow-up bug filed (#318).
docs(benchmarks): record fused decode-MoE gains for the newly wired M…
docs(benchmarks): record fused decode-MoE gains for the newly wired M…