perf(moe): wire olmoe to the fused decode-MoE kernel #302

Parent epic: #307. Sibling: #286 (qwen2_moe, same pattern).

Background

The fused single-token decode-MoE kernel (#268 series) 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. Several 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 issue wires olmoe.

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).

Scope

File: src/models/olmoe.rs. SparseMoeBlock::forward is at line 299 (expert call at 337, moe_weighted_sum at 339). Routing: softmax over top-k, no shared expert. olmoe experts are SwiGLU (SiLU), so the kernel runs with the default swiglu activation (act=0); no GeGLU/relu2 template is needed.

Implementation class: Group B (needs SwitchGLU consolidation first). olmoe defines its OWN local SwitchGLU (line 173) that does NOT expose forward_fused_kernel. Before the one-branch dispatch can work, migrate the experts to the shared crate::models::switch_layers::SwitchGLU (preferred: removes duplication) OR add a forward_fused_kernel method to the local struct. Prefer migration. Verify there is no olmoe-specific behavior in the local struct before removing it.

Validation checkpoint: an OLMoE-1B-7B 4-bit mlx-community checkpoint (small, about 4 GB). Confirm the exact mlx-community ID before use; do not fabricate it.

Acceptance Criteria

• olmoe's MoE block dispatches the fused kernel for single-token decode when MLXCEL_FUSED_MOE is enabled (default) and the config qualifies; the non-fused fallback path is unchanged.
• The local SwitchGLU/SwitchLinear is migrated to the shared switch_layers::SwitchGLU (or a forward_fused_kernel method is added), the duplicate is removed where migrated, and the existing model forward stays correct.
• INTEGRATION: the change is wired into the actual olmoe forward decode path (not an orphaned helper). Confirm the kernel actually dispatches at runtime via a dispatch trace, not inferred from byte-identical output.
• Validate on the OLMoE-1B-7B 4-bit checkpoint: greedy temp-0 output is byte-identical or within the documented f16 jitter class vs the MLXCEL_FUSED_MOE=0 (gather_qmm) baseline, with no NaN and no garbage.
• Measure decode tok/s fused-on vs off; record the gain; confirm no regression.
• Update the "Models covered" list in docs/benchmark_results/fused-moe-decode-kernel-design.md to move olmoe to the wired list.
• Add or extend tests where practical.

Technical Considerations / References

• Pattern files: src/models/qwen3_moe.rs, src/models/dots1.rs, src/models/qwen3_next.rs. Shared helpers: src/models/switch_layers.rs (fused_moe_enabled at 42, forward_fused_kernel at 276, moe_weighted_sum at 419).
• Kernel gating: affine 4/8-bit gate/up, 4/6/8-bit down, matching group_size, single-token decode only; the kernel auto-falls-back for any unsupported config.
• f16 jitter caveat: the fused kernel accumulates the GEMV in f32 with a different reduction order than gather_qmm, so within-jitter greedy token flips are expected and accepted, not a bug.
• Related: 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), perf(moe): wire qwen2_moe to the fused decode-MoE kernel #286 (qwen2_moe, same pattern).