Skip to content

perf(moe): wire qwen3_next fused decode-MoE (qwen3.5/3.6) (#268 step 3a) - #279

Merged
inureyes merged 2 commits into
mainfrom
perf/issue-268-fused-moe-qwen3next
Jun 14, 2026
Merged

inureyes merged 2 commits into
mainfrom
perf/issue-268-fused-moe-qwen3next

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

Step 3a of #268: route qwen3_next's MoE decode through the fused single-token kernel, covering qwen3.5 / qwen3.6 (qwen3_5_moe) and qwen3_next itself. They went through a separate SparseMoeBlock/SwitchGLU copy that had no fused dispatch.

What changed

qwen3_next carries its own SwitchLinear/SwitchGLU (distinct from switch_layers and qwen3_moe). This adds:

  • quantized_parts accessor on qwen3_next::SwitchLinear,
  • forward_fused_kernel on qwen3_next::SwitchGLU (mirrors qwen3_moe),
  • the fused dispatch in SparseMoeBlock::forward on single-token decode behind MLXCEL_FUSED_MOE.

Experts are 4-bit affine — fully supported by the existing kernel, no kernel change.

Result (qwen3.5-35b-a3b, M1 Ultra)

Decode 68.7 → 74.7 tok/s (+8.7%).

Correctness — within f16 jitter class (not byte-identical)

Unlike qwen3_moe/dots (byte-identical on the prompts tested), qwen3.5/3.6 greedy output is within the f16 jitter class, not bit-exact. The fused kernel accumulates the GEMV in f32 with a different reduction order than gather_qmm's tiling, perturbing the MoE block output by abs RMS ~1e-4 (max ~4e-4, both < 5e-3) — measured directly with temporary instrumentation. On a deep hybrid that is enough to flip a near-tie greedy token (~token 16 in tests); qwen3-30b happens not to hit one through 128 tokens. This is the design harness's documented "byte-identical or within f16 jitter class" allowance.

The rel_rms looks like ~1% only because the MoE block output is small-magnitude (~8e-3); the perturbation to the O(1) residual stream is ~1e-4, negligible.

cargo fmt/clippy clean. Off by default (MLXCEL_FUSED_MOE), single-token decode only.

Scope

gemma4-26b (its own Experts/gather_sort path) and nemotron-h (C++ fused_moe_forward) are separate MoE paths — follow-ups.

Refs #268.

Route qwen3_next's MoE decode through the fused single-token kernel (#268),
covering qwen3.5 and qwen3.6 (model_type qwen3_5_moe) and qwen3_next itself.

qwen3_next carries its own SwitchLinear/SwitchGLU copies (separate from
switch_layers and qwen3_moe), so this adds a quantized_parts accessor to
qwen3_next::SwitchLinear and a forward_fused_kernel to qwen3_next::SwitchGLU
mirroring qwen3_moe, then dispatches to it from SparseMoeBlock::forward on
single-token decode behind MLXCEL_FUSED_MOE. Experts are 4-bit affine, fully
supported by the existing kernel.

Decode 68.7 -> 74.7 tok/s (+8.7%) on qwen3.5-35b-a3b (M1 Ultra).

Unlike qwen3_moe/dots (byte-identical on the prompts tested), qwen3.5/3.6
greedy output is within the f16 jitter class, not bit-exact: the fused kernel
reorders the GEMV accumulation (f32) vs gather_qmm's tiling, perturbing the MoE
block output by abs RMS ~1e-4 (max ~4e-4, both < 5e-3). On a deep hybrid that
is enough to flip a near-tie greedy token (~token 16 in tests); qwen3-30b
happens not to hit one through 128 tokens. This is the design harness's
documented "byte-identical OR within f16 jitter class" allowance. Off by
default; opt-in via MLXCEL_FUSED_MOE.
@inureyes inureyes added type:performance Performance improvements area:models Model architectures, weights, loading, metadata area:inference Generation, sampling, decoding (incl. speculative, DRY) platform:macos macOS (Apple Silicon) specific priority:medium Medium priority labels Jun 14, 2026
@inureyes
inureyes merged commit c3fc093 into main Jun 14, 2026
5 checks passed
@inureyes
inureyes deleted the perf/issue-268-fused-moe-qwen3next branch June 14, 2026 08:14
inureyes added a commit that referenced this pull request Jun 14, 2026
…#283)

* docs(moe): document MLXCEL_FUSED_MOE flags, per-model gains, and M5 follow-up (#268)

Add a usage/flags section to the fused decode-MoE design doc: the
MLXCEL_FUSED_MOE / MLXCEL_FUSED_MOE_SGY / MLXCEL_FUSED_MOE_RELU2 env vars, the
measured per-model decode gains on M1 Ultra (gemma4 +13%, qwen3.5 +8.7%, dots
+4.7%, qwen3-30b +3.5%, nemotron ~0%), the f16-jitter-class parity caveat, and
the list of covered models. Mark roadmap steps 3-4 done (#278/#279/#280/#281)
and add step 5: validate on M5, then decide on flipping the flag default-on.

* docs(moe): reference the M5 validation issue (#282)
@inureyes inureyes added this to the 0.3 milestone Jun 21, 2026
@inureyes inureyes self-assigned this Aug 31, 2026
@inureyes inureyes added the status:done Completed label Sep 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:inference Generation, sampling, decoding (incl. speculative, DRY) area:models Model architectures, weights, loading, metadata platform:macos macOS (Apple Silicon) specific priority:medium Medium priority status:done Completed type:performance Performance improvements

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant