Skip to content

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

Merged
inureyes merged 7 commits into
mainfrom
refactor/issue-1884-qwen3-moe-switchglu
Sep 14, 2026
Merged

inureyes merged 7 commits into
mainfrom
refactor/issue-1884-qwen3-moe-switchglu

Conversation

@inureyes

@inureyes inureyes commented Sep 14, 2026 •

Copy link
Copy Markdown
Member

Summary

qwen3_moe kept its own SwitchLinear, SwitchGLU and forward_fused_kernel, and qwen3_vl_moe (which also builds the Qwen3-Omni-MoE thinker) built that copy by hand. The copy never got the fused-kernel dff bound the shared type gained in #311 and #643, so MLXCEL_FUSED_MOE_MAX_DFF was silently ignored. Both families now hold switch_layers::SwitchGLU (SwitchGLU::from_weights, affine) and the copy is deleted rather than patched.

  • Loader changes: bits are inferred from tensor shapes (the copy stored the configured value), a missing .biases plane fails through validate_quantization_biases instead of Weight not found, and unstacked experts.{idx} checkpoints now reach the shared stacking fallback. Every expert plane of the four local Qwen3-MoE checkpoints (safetensors headers) solves to 4-bit with biases shaped like the scales.
  • Tests: the fix(models): bound quantization params in the family-local MoE expert loaders #958 guard tests drive each family's SparseMoeBlock::from_weights; new tests pin the Dff 8256 decline through both loaders (positive control at Dff 64 only when custom_kernels_available(), refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803) and the mxfp4 decline.
  • Docs: switch_layers.rs, docs/environment-variables.md and the fused-MoE design note now say which families read the bound, fix the family counts, and qualify the qwen3-30b-a3b greedy parity claim.

Verification (GB10 sm_121, release --features cuda, base b8d10fb1)

  • Greedy ids (64 tokens, -t 0, as ids): identical base vs branch on qwen3-30b-a3b-4bit, qwen3-vl-30b-a3b-4bit and the qwen3-omni-30b-a3b-instruct-4bit thinker, default fused and MLXCEL_FUSED_MOE=0. Fused and gather diverge from each other at token 39 (qwen3-30b) and 26 (omni) on both binaries, so the fused match also rules out a silent fallback.
  • Trace (MLXCEL_PROFILE_BLOCKS=1 MLXCEL_PROFILE_QWEN3_MOE_DETAIL=1, -n 4): 144 path=fused tokens=1, 0 path=gather_qmm tokens=1 on both binaries.
  • The bound now applies: with MLXCEL_FUSED_MOE_MAX_DFF=512 the base trace stays 144 fused and the branch is 144 path=gather_qmm tokens=1; omni base ids equal its fused ids and branch ids equal its MLXCEL_FUSED_MOE=0 ids; VL 64-token runs stay at 2.31 to 2.35 s on base (kernel launched) and drop to 1.28 to 1.33 s on the branch.
  • mlxcel-server, temperature 0: identical content and reasoning base vs branch on qwen3-30b and VL; the VL completion is coherent and opens with the CLI's tokens.
  • Decode tok/s, qwen3-30b (mlxcel-bench-decode -n 128 --ignore-eos --warmup-tokens 20), 15 interleaved runs per arm: base 66.21 to 87.76 (median 78.00), branch 61.27 to 86.28 (median 73.63). Decode here is bimodal; high-mode means 81.06 vs 79.65, Mann-Whitney p = 0.21.
  • 38 of 38 tests under models::qwen3_moe, models::qwen3_vl_moe, models::switch_layers pass, one process each; both new decline tests fail against a mutant with the bound disabled. Scoped clippy (-p mlxcel --lib --tests --features cuda -D warnings) and fmt are clean. Ids, trace, cap and server checks were repeated with binaries rebuilt at the final commit.
  • Not run: metal,accelerate gate (Linux host), workspace-wide verify-test-cuda. The OpenXLA feature compile CI failure also fails on main at b8d10fb1 (unused imports in models/mod.rs and server modules) and is unrelated.

Follow-up candidates

qwen3_next.rs (same kernel copy and missing guards), gemma4.rs (GeGLU, no bound), NemotronH's opt-in MLXCEL_FUSED_MOE_RELU2 path, ten families with a local SwitchGLU and no fused path, a cached bound instead of a per-call env read, and pre-existing shared-loader gaps found in review (no cross-plane or router-vs-expert-count check; the unstacked fallback does not compare shapes before stack).

Closes #1884

qwen3_moe kept its own SwitchLinear, SwitchGLU and forward_fused_kernel, and the copy had fallen behind the shared type: the fused-kernel dff upper bound added to switch_layers in #311 and made backend-aware in #643 never reached it, so MLXCEL_FUSED_MOE_MAX_DFF was silently ignored for Qwen3-MoE and Qwen3-VL-MoE.

Both families now hold switch_layers::SwitchGLU, loaded with SwitchGLU::from_weights (affine). The local types, their loaders and qwen3_vl_moe's struct-literal builders are gone, so the fused path carries every shared guard (dff bound, mode, activation). The shared loader stores bit widths inferred from the tensor shapes and reports a missing .biases plane through validate_quantization_biases; both are 4-bit and fully biased on every local checkpoint.

The #958 guard tests now drive each family's SparseMoeBlock loader. New tests pin the dff 8256 decline through the real qwen3_moe loader (positive control at dff 64 only when a kernel port exists) and the non-affine decline in switch_layers. The switch_layers docs that claimed Qwen3-MoE already honoured the bound are corrected, with qwen3_next and gemma4 named as the families that still do not.

Refs #1884
Move the honest affine SwiGLU expert-plane builder into switch_layers as a test helper so both families' decline tests quantize the same geometry, and add the Qwen3-VL-MoE counterpart: Dff 8256 declines the fused kernel through the family's own block loader, Dff 64 still dispatches it when a kernel port exists.

Refs #1884
The switch_layers note on the fused dff bound now names the shared fused-kernel callers instead of pointing at the module list, which mixes in families with no fused dispatch, and names the paths that still bypass the bound (qwen3_next, gemma4, and NemotronH's opt-in MLXCEL_FUSED_MOE_RELU2 kernel). The storage-point narrative names ernie4_5_moe_vl and nemotron_h as the remaining hand-built expert planes, and the shared plane-builder test helper lists the Qwen3 block-loader tests.

docs/environment-variables.md and the fused decode design note record that qwen3_moe and qwen3_vl_moe now read the bound.

Refs #1884
The fused-kernel docs said greedy output on qwen3-30b-a3b agrees exactly with the kernel on and off. Measured for #1884 on GB10 at b8d10fb, the chat-templated prompt "Explain how a hash map handles collisions." agrees for 39 generated tokens and then diverges, deterministically on both paths and identically before and after the migration. The claim now records it as prompt-dependent.

Refs #1884
@inureyes inureyes added status:review Under review type:refactor Code restructuring without changing functionality priority:medium Medium priority area:models Model architectures, weights, loading, metadata area:core mlxcel-core: MLX FFI, primitives, KV cache, layers labels Sep 14, 2026
Review follow-ups for #1884. The Qwen3-Omni-MoE thinker is built as a Qwen3VLMoeModel, so it moved onto the shared SwitchGLU too and now reads MLXCEL_FUSED_MOE_MAX_DFF; the switch_layers note and the environment-variable table say so, and the MLXCEL_FUSED_MOE row lists every family that actually dispatches the fused kernel. The qwen3_moe loader comment cites #929 and #973 for the shared loader's refusals instead of #975, which covers a DeepSeek-only helper.

Refs #1884
@inureyes inureyes added status:done Completed and removed status:review Under review labels Sep 14, 2026
@inureyes
inureyes merged commit c1234f3 into main Sep 14, 2026
15 of 16 checks passed
@inureyes
inureyes deleted the refactor/issue-1884-qwen3-moe-switchglu branch September 14, 2026 11:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:models Model architectures, weights, loading, metadata priority:medium Medium priority status:done Completed type:refactor Code restructuring without changing functionality

Projects

None yet

Development

Successfully merging this pull request may close these issues.

refactor(moe): migrate qwen3_moe onto the shared SwitchGLU

1 participant