Repository navigation
refactor(moe): migrate qwen3_moe onto the shared SwitchGLU - #1888
Merged
Merged
Conversation
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
7 of 8 tasks
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
qwen3_moekept its ownSwitchLinear,SwitchGLUandforward_fused_kernel, andqwen3_vl_moe(which also builds the Qwen3-Omni-MoE thinker) built that copy by hand. The copy never got the fused-kerneldffbound the shared type gained in #311 and #643, soMLXCEL_FUSED_MOE_MAX_DFFwas silently ignored. Both families now holdswitch_layers::SwitchGLU(SwitchGLU::from_weights, affine) and the copy is deleted rather than patched..biasesplane fails throughvalidate_quantization_biasesinstead ofWeight not found, and unstackedexperts.{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.SparseMoeBlock::from_weights; new tests pin the Dff 8256 decline through both loaders (positive control at Dff 64 only whencustom_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.switch_layers.rs,docs/environment-variables.mdand 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, baseb8d10fb1)-t 0, as ids): identical base vs branch onqwen3-30b-a3b-4bit,qwen3-vl-30b-a3b-4bitand theqwen3-omni-30b-a3b-instruct-4bitthinker, default fused andMLXCEL_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.MLXCEL_PROFILE_BLOCKS=1 MLXCEL_PROFILE_QWEN3_MOE_DETAIL=1,-n 4): 144path=fused tokens=1, 0path=gather_qmm tokens=1on both binaries.MLXCEL_FUSED_MOE_MAX_DFF=512the base trace stays 144 fused and the branch is 144path=gather_qmm tokens=1; omni base ids equal its fused ids and branch ids equal itsMLXCEL_FUSED_MOE=0ids; 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.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.models::qwen3_moe,models::qwen3_vl_moe,models::switch_layerspass, 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.metal,accelerategate (Linux host), workspace-wideverify-test-cuda. TheOpenXLA feature compileCI failure also fails onmainatb8d10fb1(unused imports inmodels/mod.rsand 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-inMLXCEL_FUSED_MOE_RELU2path, ten families with a localSwitchGLUand 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 beforestack).Closes #1884