perf(moe): wire qwen2_moe to the fused decode-MoE kernel #286

Problem / Background

PR #285 (issue #282) flipped MLXCEL_FUSED_MOE to on by default and, during M5 validation, found the fused single-token decode-MoE kernel (#268) is only wired into four model paths: qwen3_moe, qwen3_next, dots1, and gemma4. Several other MoE families reuse the shared SwitchGLU for the expert matmul but never call forward_fused_kernel, so they always fall back to gather_qmm regardless of the flag.

qwen2_moe (which serves qwen1.5-moe / Qwen2-MoE) is one of them. It was listed in #282's validation set but could not be validated, because it never dispatches the kernel (0 dispatches confirmed via a dispatch trace during #282). The doc's "Models covered" list overstated coverage and was corrected in #285.

Scope

Wire qwen2_moe's SparseMoeBlock to dispatch the fused kernel for single-token decode, mirroring qwen3_moe / dots1. This is low-risk and small:

• qwen2_moe::SparseMoeBlock already holds experts: SwitchGLU (src/models/qwen2_moe.rs:398), and SwitchGLU already exposes forward_fused_kernel.
• The block's forward (src/models/qwen2_moe.rs:436-443) currently always calls the non-fused self.experts.forward(...) + moe_weighted_sum. Replace only the routed-expert computation with the standard fused-gated branch, leaving the shared-expert addition (lines 445-456) intact.

Concrete change (mirrors dots1.rs / qwen3_moe.rs):

let expert_out_sum = {
    let fused = if mlxcel_core::array_shape(&x_flat)[0] == 1
        && crate::models::switch_layers::fused_moe_enabled()
    {
        self.experts
            .forward_fused_kernel(&x_flat, &topk_indices, &scores)
            .map(|out| mlxcel_core::reshape(&out, &[1, hidden_dim]))
    } else {
        None
    };
    match fused {
        Some(out) => out,
        None => {
            let expert_out = self.experts.forward(&x_flat, &topk_indices);
            crate::models::switch_layers::moe_weighted_sum(
                &expert_out, &scores, mlxcel_core::array_dtype(&x_flat),
            )
        }
    }
};

The kernel auto-falls-back for any unsupported config, so this is safe. Qwen1.5-MoE-A2.7B is 4-bit affine, group_size 64, which should satisfy the kernel's gating (affine 4/8-bit gate/up, 4/6/8-bit down, matching group_size, single-token decode).

Related unwired families (same gap, optional sibling work)

The same one-branch pattern applies to: mixtral, olmoe, minimax, phimoe, lfm2, and qwen3_vl_moe (the last imports SwitchGLU from qwen3_moe but keeps its own non-fused MoE forward). These can be addressed in this issue or in sibling issues; each must be validated independently. This issue's primary deliverable is qwen2_moe.

Acceptance Criteria

• qwen2_moe::SparseMoeBlock dispatches the fused kernel for single-token decode when MLXCEL_FUSED_MOE is enabled (default) and the config qualifies; the shared-expert path is unchanged.
• Validate on a real checkpoint (mlx-community/Qwen1.5-MoE-A2.7B-Chat-4bit): greedy temp-0 output is byte-identical or within the documented f16 jitter class vs MLXCEL_FUSED_MOE=0 (the gather_qmm baseline), with no NaN / no garbage. Confirm the kernel actually dispatches (not a silent fallback) rather than inferring it from byte-identical output.
• Measure decode tok/s fused-on vs off; confirm no regression and record the gain.
• Update docs/benchmark_results/fused-moe-decode-kernel-design.md "Models covered" to move qwen2_moe from the not-wired list to the wired list.
• Add or extend tests where practical.

Technical Considerations / References

• Established pattern and shared helpers: src/models/qwen3_moe.rs, src/models/dots1.rs, src/models/qwen3_next.rs (the fused_moe_enabled() gate + experts.forward_fused_kernel(...) with the moe_weighted_sum fallback), and src/models/switch_layers.rs (fused_moe_enabled, SwitchGLU::forward_fused_kernel, gating/fallback conditions).
• 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 accepted (not a bug). See the design doc.
• Context: perf(moe): validate fused decode-MoE on M5, then enable MLXCEL_FUSED_MOE by default #282 (M5 validation methodology), perf(moe): default-on MLXCEL_FUSED_MOE, validated on M5 (#282) #285 (the default flip), perf(models): general MoE decode gap vs mlx-lm (dots.llm1 0.51x, qwen3_moe 0.69x, nemotron MoE) #268 (the kernel series).