refactor(moe): migrate qwen3_moe onto the shared SwitchGLU #1884

Problem / Background

qwen3_moe carries its own SwitchLinear, SwitchGLU and forward_fused_kernel instead of the shared ones in src/models/switch_layers.rs, and the copy has fallen behind: the shared function gained a dff upper bound that the copy never received. Where two implementations do the same thing, the shared one should be used and optimization concentrated there. Patching the missing guard into the copy would drift again the next time the shared function changes, which is how this copy fell behind: both functions landed together in 7fd4dcac (#275), and the bound was then added to the shared function only, in d846fd4f (#311), and made backend-aware in 37931d7f (#643).

This is a type migration, not a call-site change. The two quantized_parts accessors return different types: qwen3_moe.rs:205 returns a 5-tuple (&MlxArray, &MlxArray, &MlxArray, i32, i32), while switch_layers.rs:422 returns QuantizedSwitchLinearRef<'_> and declines when biases are absent or a per-expert global_scale is present. qwen3_vl_moe also imports the copy (qwen3_vl_moe.rs:40) and builds it as struct literals, so it has to move in the same change.

Current Behavior

The copy: SwitchLinear at qwen3_moe.rs:140-224, SwitchGLU at :228-386 (with private gather_sort at :329 and scatter_unsort at :368), its loaders at :1497-1540, called from SparseMoeBlock::from_weights at :1486. The shared type: SwitchLinear at switch_layers.rs:334-445, SwitchGLU at :913-1216, free gather_sort at :1220 and scatter_unsort at :1245.

Guard-by-guard comparison of forward_fused_kernel (switch_layers.rs:1032-1114 against qwen3_moe.rs:286-326):

Both call mlxcel_core::fused_moe_expert_kernel with the same argument order (x, idx, gate w/s/b, up w/s/b, down w/s/b, scores, din, dff, k, gate bits, down bits, gate group_size): shared :1094-1113, copy :322-325. The shared comment for the bound, switch_layers.rs:1068-1073: "Large experts: gather_qmm already saturates the GPU, so the two-kernel fused path's all-cores advantage disappears and its extra dispatch + global-memory activation staging becomes a net loss." Defaults are FUSED_MOE_MAX_DFF_METAL = 4096 (:118) and FUSED_MOE_MAX_DFF_CUDA = 8192 (:139), overridable by MLXCEL_FUSED_MOE_MAX_DFF. The docs currently claim coverage that does not exist: switch_layers.rs:148-150 says the cap governs Qwen3-MoE, and the module doc at :17-19 lists Qwen3Moe as a user of the shared type.

The defect is latent, not realized. Every qwen3_moe checkpoint on the development host is affine and below both bounds, so neither missing guard fires, and the migration should be behavior-preserving on all of them:

What each gap would take to matter:

• dff bound (the realizable gap). A qwen3_moe or qwen3_vl_moe variant with dff above 4096 on Metal (8192 on CUDA) would take the fused path in the regime measured as a net loss (phi-3.5-moe at Dff 6400 regresses on M1 Ultra, switch_layers.rs:116). MLXCEL_FUSED_MOE_MAX_DFF is silently ignored for both families today.
• Mode check (not reachable through the copy). The copy cannot represent a non-affine plane at all: its Quantization (qwen3_moe.rs:80-84) has no mode, its loader requires .biases (:1525), which MLX block-float modes do not ship, and even its gather_qmm path hardcodes "affine" (:185). The generic text loader repacks NVFP4 only for gemma4 (sanitize.rs:1643-1648). So a native mxfp4 or nvfp4 qwen3_moe checkpoint fails to load today rather than reaching the fused kernel. The benefit of migrating is that any future non-affine support is covered by the shared guard from the start.

Related: #1859. Its thread notes that qwen3_moe.rs:286 defines its own forward_fused_kernel rather than only calling the shared one, so a wired-family count keyed on definitions misleads. (The 17 vs 18 count difference it reconciled comes from laguna_layers.rs, added in #1739.)

Proposed Solution

1. qwen3_moe.rs: delete the local SwitchLinear, SwitchGLU (including quantized_parts, forward_fused_kernel, gather_sort, scatter_unsort) and their loaders. SparseMoeBlock.experts (:418) becomes crate::models::switch_layers::SwitchGLU, built with SwitchGLU::from_weights(weights, &format!("{prefix}.switch_mlp"), args.group_size(), args.bits()). The call sites in forward (:476-477, :485) and forward_profiled (:547-548, :564) keep their method names and need no logic change. Drop the now-unused validate_expert_quantization_params import (:27), since SwitchLinear::from_stacked_parts carries the bound (switch_layers.rs:546-547).
2. qwen3_vl_moe.rs: import the shared SwitchGLU (:40) and replace load_switch_glu / load_switch_linear (:641-688) with SwitchGLU::from_weights(weights, prefix, config.group_size(), config.bits()). Its forward_fused_kernel call at :613-614 is unchanged.
3. Use from_weights (affine), not from_weights_with_mode. Non-affine qwen3_moe support is out of scope.
4. Fix the docs: remove Qwen3MoE and Qwen3VLMoE from validate_expert_quantization_params "Used by" (switch_layers.rs:203-206), add Qwen3VLMoe to the module list (:17-19), and update the family counts at :172 and :243.

Rejected: adding the dff bound to the copy. It fixes today's gap and leaves the drift that caused it.

Scope

In scope: src/models/qwen3_moe.rs, src/models/qwen3_vl_moe.rs, src/models/qwen3_moe_tests.rs, src/models/qwen3_vl_moe_tests.rs, doc comments and tests in src/models/switch_layers.rs.

Out of scope, candidates for the same consolidation (no issue filed yet): qwen3_next.rs:1055 SwitchGLU plus :1095 forward_fused_kernel, which is the same SwiGLU kernel with the same two guards missing, reached by Qwen3Next, Qwen3.5 and the qwen3_omni_moe talker. gemma4.rs:618 drives the GeGLU kernel (fused_moe_geglu_kernel) and has the mode check but not the dff bound, so consolidating it first needs a GeGLU SwitchGluActivation. Ten more families keep a local SwitchGLU with no fused path, so no guard can drift there: deepseek.rs:245, deepseek_v2.rs:695, deepseek_v3.rs:639, deepseek_v32.rs:703, ernie4_5_moe.rs:311, exaone_moe.rs:208, glm4_moe.rs:542, glm4_moe_lite.rs:528, hunyuan_moe.rs:309, llama4.rs:382. Also out of scope: non-affine qwen3_moe loading, and any change to the shared forward_fused_kernel logic.

Implementation Notes

• Reuse: switch_layers::SwitchGLU::from_weights, forward, forward_fused_kernel and the free gather_sort / scatter_unsort. No qwen3-specific expert code should remain.
• Forward equivalence: shared SwitchGLU::forward (:945-969) runs the same ops as the copy (qwen3_moe.rs:236-279): the same do_sort threshold of 64, the same sort and unsort, and compiled_swiglu_activation. Shared SwitchLinear::forward passes the stored mode, which is "affine" via from_weights (:413), and skips the global-scale multiply because global_scale is None.
• Loader differences: the shared loader stores bits inferred from the tensor shapes (switch_layers.rs:586-597), while the copy stores the declared bits (qwen3_moe.rs:1520). On all four host checkpoints the shapes solve to 4, the declared value (for example 30B gate [128,768,256] with scales [128,768,32], and 480B down [160,6144,320] with scales [160,6144,40]), so the stored triple is identical. The shared loader adds refusals only for loads that were going to abort anyway (:599-629), plus an unstacked experts.{idx} fallback (:495-508).
• Error text: a quantized plane without .biases now fails with the validate_quantization_biases message (mlxcel-core/src/layers.rs:1552-1572) instead of Weight not found: ...biases. Update any assertion that matches the old text.
• Visibility: the shared fields are private (switch_layers.rs:914-919). Neither family reads them; both only call forward and forward_fused_kernel. pub mod qwen3_moe (models/mod.rs:225) currently exports the removed types, but the only in-tree users are qwen3_vl_moe.rs:40 and qwen3_moe_tests.rs:24.
• Tests to retarget: qwen3_moe_switch_linear_rejects_quantization_params_that_would_abort_gather_qmm (qwen3_moe_tests.rs:214-254) and the load_switch_linear guard test (qwen3_vl_moe_tests.rs:56-87) should drive switch_layers::SwitchLinear::from_weights(&weights, PREFIX, group_size, bits). Keep the HOSTILE_QUANT_PARAMS loop, the honest positive control and the bf16 control.
• New decline test (qwen3_moe_tests.rs): build honest affine 4-bit gate, up and down planes with mlxcel_core::quantize_weights (hidden 64, group size 64, 2 experts stacked with stack_owned, as in laguna_tests.rs:1150-1170) at dff 8256 (129 groups, above both defaults), plus a {prefix}.gate.weight router. Load them through SparseMoeBlock::from_weights and assert experts.forward_fused_kernel returns None for one token with 64 elements. Return early with a message when MLXCEL_FUSED_MOE_MAX_DFF is set. Add a positive control at dff 64 on the same geometry that asserts Some, but run it only when mlxcel_core::custom_kernels_available() is true: on a backend without a port the kernel throw ends the process (refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803).
• Non-affine decline test (switch_layers.rs tests): an mxfp4 plane (group size 32, bits 4, from quantize_weights_with_mode) loaded with SwitchGLU::from_weights_with_mode(.., "mxfp4") returns None from forward_fused_kernel. It declines at quantized_parts because there are no biases.

Acceptance Criteria

• grep -n "struct SwitchGLU\|struct SwitchLinear\|fn quantized_parts\|fn forward_fused_kernel\|fn gather_sort" src/models/qwen3_moe.rs src/models/qwen3_vl_moe.rs prints nothing, and both SparseMoeBlocks hold switch_layers::SwitchGLU. Done in PR refactor(moe): migrate qwen3_moe onto the shared SwitchGLU #1888: the grep prints nothing; both blocks hold switch_layers::SwitchGLU.
• Greedy token ids on qwen3-30b-a3b-4bit (64 tokens, -t 0, compared by [token ids ...] rather than text) are byte-identical between the base commit and the branch, both on the default fused path and with MLXCEL_FUSED_MOE=0. Done in PR refactor(moe): migrate qwen3_moe onto the shared SwitchGLU #1888: identical base b8d10fb1 vs branch, fused and MLXCEL_FUSED_MOE=0 (GB10).
• The fused kernel is still dispatched after migration: the decode trace shows [QWEN3_MOE_SPARSE] path=fused tokens=1 and no path=gather_qmm tokens=1. This is required because greedy ids on this checkpoint agree with the kernel on and off (switch_layers.rs:41-44, :1019-1021), so an id match alone cannot detect a silent fallback. Done in PR refactor(moe): migrate qwen3_moe onto the shared SwitchGLU #1888: 144 path=fused tokens=1, 0 path=gather_qmm tokens=1 on both binaries; with MLXCEL_FUSED_MOE_MAX_DFF=512 the branch trace is 144 gather_qmm while the base stays fused.
• Greedy token ids on qwen3-vl-30b-a3b-4bit (text-only prompt) are byte-identical before and after on the default path. Done in PR refactor(moe): migrate qwen3_moe onto the shared SwitchGLU #1888: identical on the default path and with MLXCEL_FUSED_MOE=0, CLI and mlxcel-server.
• Decode tok/s on qwen3-30b-a3b-4bit is unchanged: n >= 3 per arm after warm-up, min to max reported for each arm, and the ranges overlap. Done in PR refactor(moe): migrate qwen3_moe onto the shared SwitchGLU #1888: n = 15 interleaved per arm; base 66.21 to 87.76, branch 61.27 to 86.28; ranges overlap (Mann-Whitney p = 0.21).
• Served path: mlxcel-server on qwen3-30b-a3b-4bit at temperature 0 returns the same completion text before and after. Both CLI and server reach SparseMoeBlock through Qwen3MoeModel::forward or forward_batched_impl (qwen3_moe.rs:1582-1592). Done in PR refactor(moe): migrate qwen3_moe onto the shared SwitchGLU #1888: identical content and reasoning on qwen3-30b-a3b-4bit and qwen3-vl-30b-a3b-4bit.
• The new decline tests and the retargeted fix(models): bound quantization params in the family-local MoE expert loaders #958 guard tests pass, and the doc comments listed in Proposed Solution step 4 are updated. Done in PR refactor(moe): migrate qwen3_moe onto the shared SwitchGLU #1888: 38 of 38 tests under the three modules pass; both new decline tests fail against a mutant with the bound disabled.
• make verify (macOS) or make verify-fmt verify-clippy verify-test-cuda (Linux) passes. Partial in PR refactor(moe): migrate qwen3_moe onto the shared SwitchGLU #1888: cargo fmt --check and cargo clippy --profile test-fast --features cuda -p mlxcel --lib --tests --no-deps -- -D warnings are clean on Linux; verify-clippy (metal,accelerate) cannot run on Linux and the workspace-wide verify-test-cuda was not run.

Verification

make test-fast-cuda FILTER=models::qwen3_moe      # macOS: make test-fast FILTER=...
make test-fast-cuda FILTER=models::qwen3_vl_moe
make test-fast-cuda FILTER=models::switch_layers
make verify-fmt verify-clippy verify-test-cuda    # macOS: make verify

make release-cuda                                 # run everything below on the base commit and on the branch
M=models/mlx/qwen3-30b-a3b-4bit
P="Explain how a hash map handles collisions."

# Token ids, fused on and off
MLXCEL_PRINT_TOKEN_IDS=1 ./target/release/mlxcel generate -m $M -p "$P" -n 64 -t 0 2>&1 | grep '^\[token ids'
MLXCEL_FUSED_MOE=0 MLXCEL_PRINT_TOKEN_IDS=1 ./target/release/mlxcel generate -m $M -p "$P" -n 64 -t 0 2>&1 | grep '^\[token ids'

# Fused dispatch still taken (first count > 0, second count 0)
MLXCEL_PROFILE_BLOCKS=1 MLXCEL_PROFILE_QWEN3_MOE_DETAIL=1 ./target/release/mlxcel generate -m $M -p "$P" -n 4 -t 0 2>&1 | grep -c 'path=fused tokens=1'
MLXCEL_PROFILE_BLOCKS=1 MLXCEL_PROFILE_QWEN3_MOE_DETAIL=1 ./target/release/mlxcel generate -m $M -p "$P" -n 4 -t 0 2>&1 | grep -c 'path=gather_qmm tokens=1'

# Throughput, n >= 3 per arm
for i in 1 2 3; do ./target/release/mlxcel-bench-decode -m $M -p "$P" -n 128 --ignore-eos --warmup-tokens 20; done

On GB10, single-run decode is bimodal by up to about 25 percent, so compare per-arm ranges rather than single numbers.