[inureyes]

Summary

Wire qwen2_moe (Qwen1.5-MoE / Qwen2-MoE) into the fused single-token decode-MoE Metal kernel (#268), mirroring qwen3_moe / dots1. For single-token decode with MLXCEL_FUSED_MOE on (the default), the routed-expert weighted sum runs through experts.forward_fused_kernel; the proven gather_qmm SwitchGLU + moe_weighted_sum path stays as the fallback for prefill, when the flag is off, and for any config the kernel does not support. The shared-expert (sigmoid-gated) addition is unchanged.

Group-B correction to the issue

The issue assumed qwen2_moe's experts already used the shared SwitchGLU exposing forward_fused_kernel. That was wrong: qwen2_moe defined its own local SwitchGLU and SwitchLinear that exposed neither forward_fused_kernel nor quantized_parts, and the local SwitchLinear::Quantized had no mode field. So this is a Group-B migration (the same class as #301/#302/#303): migrate qwen2_moe to crate::models::switch_layers::SwitchGLU, then add the one-branch fused dispatch. The local SwitchGLU / SwitchLinear and their loaders are removed; a repo-wide grep confirmed nothing else referenced them.

Weight-loading parity

The shared SwitchLinear loader was stacked-only ({prefix}.weight holding all experts), but the two checkpoint conventions differ: qwen3-moe ships experts pre-stacked under switch_mlp, while Qwen1.5-MoE / Qwen2-MoE ship them unstacked under experts.{idx}.{proj} (verified against the local models/qwen1.5-moe-a2.7b-4bit checkpoint, which is fully per-expert). The local qwen2_moe loader handled the per-expert layout via a stacking fallback; the shared loader did not, so loading qwen2_moe through it unchanged would have errored.

To preserve exactly which weights load, SwitchLinear::from_weights_with_mode gains a per-expert stacking fallback (stack_individual_experts) that converts a {root}.switch_mlp.{proj} prefix to {root}.experts.{idx}.{proj} and stacks contiguously from index 0. It runs only when the pre-stacked tensor is absent, so it never changes behavior for an already-loadable checkpoint, and it now also serves the other unstacked MoE families slated for migration. The per-tensor bit-width inference is shared across both paths; for qwen1.5-moe's uniform 4-bit it yields the same bits the local loader hardcoded (args.bits()), so there is no silent change to loaded weights, only a strictly more robust path for mixed-precision exports.

What changed

• src/models/qwen2_moe.rs: SparseMoeBlock.experts now typed crate::models::switch_layers::SwitchGLU; construction calls the shared SwitchGLU::from_weights(weights, "{moe}.switch_mlp", group_size, bits); the routed-expert sum in SparseMoeBlock::forward uses the fused-gated branch with a moe_weighted_sum fallback; local SwitchGLU / SwitchLinear definitions and loaders removed.
• src/models/switch_layers.rs: from_weights_with_mode refactored into a shared from_stacked_parts builder plus a new stack_individual_experts per-expert fallback; module Used by: doc updated to include Qwen2Moe; two unit tests added (per-expert quantized stacking, missing-weight error).
• docs/benchmark_results/fused-moe-decode-kernel-design.md: qwen2_moe moved from the not-wired list to the wired list.

Test plan

• cargo check --lib --tests --features metal,accelerate
• cargo clippy --lib --tests --features metal,accelerate -- -D warnings (clean)
• Real-checkpoint validation on mlx-community/Qwen1.5-MoE-A2.7B-Chat-4bit: confirm the kernel actually dispatches (not a silent fallback), temp-0 greedy parity vs MLXCEL_FUSED_MOE=0 within the documented f16 jitter class (no NaN / no garbage), and decode tok/s fused-on vs off. Pending the orchestrator's verification step.

Closes #286