MLXLLM Gemma4Text: add MoE block (router + experts) for the text path#364
Open
neuromechanist wants to merge 1 commit into
Open
MLXLLM Gemma4Text: add MoE block (router + experts) for the text path#364neuromechanist wants to merge 1 commit into
neuromechanist wants to merge 1 commit into
Conversation
Brings the MLXVLM Gemma4 MoE implementation to the LLM text path so the 26B-A4B (MoE) checkpoint loads through MLXLLM, not just MLXVLM. Adds Gemma4TextRouter, Gemma4TextExperts (SwitchGLU), the enable_moe_block config + dense/sparse decoder branch, and the fused-expert weight remap in sanitize. The E-series KV-sharing it pairs with is already upstream (ml-explore#342). Ported from MLXVLM/Models/Gemma4.swift (ml-explore PRs ml-explore#180, ml-explore#228).
aleroot
reviewed
Jun 26, 2026
| } | ||
| sanitized[k] = v | ||
|
|
||
| // Remap .experts.down_proj -> .experts.switch_glu.down_proj.weight |
Contributor
There was a problem hiding this comment.
The original sanitized[k] = value pass-through runs before the remap checks. For MoE keys, the original key (e.g. .experts.down_proj) is inserted into the dictionary alongside the remapped key (.experts.switch_glu.down_proj.weight). The continue only prevents the second pass-through it doesn't remove the first assignment. MoE keys should be remapped instead of passed through, not in addition.
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
2 participants
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.
The MLXLLM
Gemma4Texttext path can load dense and (since #342) KV-shared Gemma4 checkpoints, but not MoE ones (26B-A4B): loading fails withUnhandled keys ["experts", "router", "post_feedforward_layernorm_1", "post_feedforward_layernorm_2", "pre_feedforward_layernorm_2"]. The complete MoE implementation already exists inMLXVLM/Models/Gemma4.swift(#180, #228); this brings it to the LLM text path.Adds:
Gemma4TextRouter,Gemma4TextExperts(SwitchGLU), theenable_moe_blockconfig + dense/sparse decoder branch, and the fused-expert weight remap insanitize(experts.gate_up_proj/down_proj->switch_glu.*). Pairs with the KV-sharing already merged in #342.Verified by loading + generating
mlx-community/gemma-4-26b-a4b-it-4bitthrough the LLM text path (correct output at native context). Ported from ml-explore's own MLXVLM implementation.