perf(moe): wire mixtral to the fused decode-MoE kernel #301

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

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/mixtral.rs. SparseMoeBlock::forward is at line 281 (expert call at 313, moe_weighted_sum at 315). Routing: softmax over top-k, no shared expert. mixtral 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). mixtral defines its OWN local SwitchGLU (line 156) and a local SwitchLinear that do 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 mixtral-specific behavior in the local struct before removing it.

Validation checkpoint: mlx-community/Mixtral-8x7B-Instruct-v0.1-4bit (in the project regression set; about 24 GB at 4-bit, fits 128 GB).

Acceptance Criteria

• mixtral'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 mixtral 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 mlx-community/Mixtral-8x7B-Instruct-v0.1-4bit: 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 mixtral 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).