Skip to content

perf(moe): GeGLU fused decode-MoE; wire gemma4 (+13%) (#268 step 3b) - #281

Merged
inureyes merged 1 commit into
mainfrom
perf/issue-268-fused-moe-gemma4
Jun 14, 2026
Merged

inureyes merged 1 commit into
mainfrom
perf/issue-268-fused-moe-gemma4

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

Step 3b of #268: extend the two-kernel fused decode-MoE to GeGLU experts and route gemma4's MoE decode through it. gemma4's experts use gelu-tanh-approx(gate) * up, not SwiGLU.

What changed

  • The gate/up kernel gains an act template arg: 0 = SwiGLU (silu), 1 = GeGLU (gelu tanh approx, matching compiled_geglu_approx_activation).
  • The shared two-kernel body is factored into run_fused_moe_two_kernel, exposed as two FFIs: fused_moe_expert_kernel (silu — unchanged for qwen3_moe/dots/qwen3_next) and the new fused_moe_geglu_kernel.
  • SwitchGeGLU gains forward_fused_kernel; Experts::forward dispatches to it on single-token decode behind MLXCEL_FUSED_MOE.

Result (gemma-4-26b-a4b-it, M1 Ultra)

Decode 73.8 → 83.2 tok/s (+13%) — the largest fused-MoE win so far. gemma4 previously used the compiled SwitchGeGLU path (3 gather_qmm + geglu); the fused kernel reads each weight once across all cores and folds the activation + score weighting into the GEMV epilogues. Its small experts (Dff=704) make the MoE a meaningful decode fraction.

Validation

  • Greedy temp-0 output byte-identical to the compiled-switch path (with the chat template) on gemma-4-26b-a4b-it across multiple prompts. (Raw no-chat-template prompts make this instruct model degenerate, so parity must be checked with the template.)
  • 4-bit affine, single-token decode, off by default (MLXCEL_FUSED_MOE).
  • cargo fmt + cargo clippy clean.

Refs #268.

Extend the two-kernel fused decode-MoE to GeGLU experts and route gemma4's MoE
decode through it. gemma4's experts use gelu-tanh-approx(gate) * up, not SwiGLU.

The gate/up kernel gains an `act` template arg (0 = SwiGLU/silu, 1 = GeGLU/gelu
tanh approx, matching compiled_geglu_approx_activation); the shared two-kernel
body is factored into run_fused_moe_two_kernel and exposed as two FFIs,
fused_moe_expert_kernel (silu, unchanged for qwen3_moe/dots/qwen3_next) and the
new fused_moe_geglu_kernel. SwitchGeGLU gains a forward_fused_kernel and
Experts::forward dispatches to it on single-token decode behind MLXCEL_FUSED_MOE.

gemma4's decode previously used the compiled SwitchGeGLU path (3 gather_qmm +
geglu); the fused kernel reads each weight once across all cores and folds the
activation + score weighting into the GEMV epilogues. Decode 73.8 -> 83.2 tok/s
(+13%) on gemma-4-26b-a4b-it (M1 Ultra), the largest fused-MoE win so far
(small experts, Dff=704, make the MoE a meaningful decode fraction).

Validated greedy temp-0 byte-identical to the compiled-switch path with the chat
template on gemma-4-26b-a4b-it; 4-bit affine, off by default. fmt + clippy clean.
@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 103b512 into main Jun 14, 2026
5 checks passed
@inureyes
inureyes deleted the perf/issue-268-fused-moe-gemma4 branch June 14, 2026 09:58
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