Skip to content

perf(moe): wire phimoe to the fused decode-MoE kernel - #312

Merged
inureyes merged 1 commit into
mainfrom
perf/issue-303-phimoe-fused-kernel
Jun 16, 2026
Merged

inureyes merged 1 commit into
mainfrom
perf/issue-303-phimoe-fused-kernel

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

  • Removes PhiMoE's local SwitchLinear / SwitchGLU (plus gather_sort / scatter_unsort helpers) and retypes SparseMoeBlock.experts to the shared crate::models::switch_layers::SwitchGLU.
  • Wires the single-token fused decode dispatch in SparseMoeBlock::forward, matching the pattern used by qwen3_moe, gemma4, and lfm2.
  • Keeps sanitize_weights intact so unstacked experts.{i}.w1/w2/w3 community checkpoints still stack correctly.

What changed

  • src/models/phimoe.rs: removed ~190 lines of local SwitchLinear / SwitchGLU; retypes experts field; updates SparseMoeBlock::from_weights to call shared SwitchGLU::from_weights(weights, prefix+".switch_mlp", args.group_size(), args.bits()); adds fused dispatch block in SparseMoeBlock::forward.
  • src/models/phimoe_tests.rs: new file with four config-parsing tests (4-bit checkpoint, head_dim, num_kv_heads, unquantized defaults).
  • src/models/mod.rs: registers phimoe_tests.
  • docs/benchmark_results/fused-moe-decode-kernel-design.md: moves phimoe to the wired models list; count seven -> eight.

Weight-loading parity note

Pre-stacked checkpoints (mlx-community/Phi-3.5-MoE-instruct-4bit) have weights under block_sparse_moe.switch_mlp.{gate,up,down}_proj.{weight,scales,biases}. The shared SwitchGLU::from_weights reads exactly those keys when called with prefix {layer}.block_sparse_moe.switch_mlp, so loading is identical to the old local implementation. The sanitize_weights stacking path for unstacked community checkpoints is untouched.

Test plan

Closes #303

@inureyes inureyes added type:refactor Code restructuring without changing functionality type:performance Performance improvements priority:medium Medium priority area:models Model architectures, weights, loading, metadata area:inference Generation, sampling, decoding (incl. speculative, DRY) platform:macos macOS (Apple Silicon) specific status:review Under review labels Jun 16, 2026
Migrate PhiMoE's local SwitchLinear and SwitchGLU to the shared implementations in switch_layers.rs and add the single-token fused dispatch.

The local SwitchLinear enum (Quantized/Regular variants with gather_qmm/gather_mm dispatch) and the local SwitchGLU struct (with gather_sort/scatter_unsort helpers) are removed from phimoe.rs. SparseMoeBlock.experts now carries crate::models::switch_layers::SwitchGLU, loaded via SwitchGLU::from_weights(weights, prefix+".switch_mlp", group_size, bits) -- the same prefix the old local call used.

sanitize_weights is left entirely intact: it still handles the unstacked experts.{i}.w1/w2/w3 layout from community checkpoints and stacks them into block_sparse_moe.switch_mlp.{gate,up,down}_proj.*. Pre-stacked checkpoints (already in the target shape, as verified on mlx-community/Phi-3.5-MoE-instruct-4bit) pass through unchanged.

SparseMoeBlock::forward gains the same fused dispatch pattern used by qwen3_moe, gemma4, and lfm2: when the token count is 1 and MLXCEL_FUSED_MOE is not disabled, forward_fused_kernel is called and the [hidden_dim] result is reshaped to [1, hidden_dim] before the block reassembly; the gather_qmm+moe_weighted_sum path is the fallback.

Docs: phimoe moved from the not-yet-wired list to Models covered in fused-moe-decode-kernel-design.md; count updated from seven to eight.

Tests: phimoe_tests.rs added with four config-parsing tests covering the 4-bit checkpoint layout, head_dim derivation, num_kv_heads accessor, and unquantized defaults. Registered in models/mod.rs.
@inureyes

Copy link
Copy Markdown
Member Author

Rebased onto main to pick up the Dff guard (#311). phi-3.5-moe's expert intermediate is 6400, above the default MLXCEL_FUSED_MOE_MAX_DFF (4096), so the kernel declines and decode stays on gather_qmm. M1 Ultra validation: guard-on 71.8-72.3 tok/s == MLXCEL_FUSED_MOE=0 baseline 72.6 tok/s (falls back, no regression); MLXCEL_FUSED_MOE_MAX_DFF=99999 re-arms the fused path at 68.8 tok/s (the -5.9% the guard avoids). The migration removes the local SwitchGLU/SwitchLinear duplication; the dispatch arms only if the bound is raised.

@inureyes
inureyes force-pushed the perf/issue-303-phimoe-fused-kernel branch from e498e4a to cf068df Compare June 16, 2026 13:51
@inureyes inureyes added status:done Completed and removed status:review Under review labels Jun 16, 2026
@inureyes
inureyes merged commit 7689336 into main Jun 16, 2026
5 checks passed
@inureyes
inureyes deleted the perf/issue-303-phimoe-fused-kernel branch June 16, 2026 13:53
@inureyes inureyes added this to the 0.3 milestone Jun 21, 2026
@inureyes inureyes self-assigned this Aug 31, 2026
@inureyes inureyes removed the type:refactor Code restructuring without changing functionality label Sep 7, 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.

perf(moe): wire phimoe to the fused decode-MoE kernel

1 participant