Introducing OptimizedMoeTransform - #1190
Open
ochougul wants to merge 21 commits into
Open
Conversation
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
ochougul
commented
Jul 21, 2026
|
|
||
| hidden_states = hidden_states + self.shared_experts(residuals) | ||
| return hidden_states | ||
| QEffPrefillChunkedGlm4MoeMoE = QEffGlm4MoeMoE |
ochougul
commented
Jul 21, 2026
Comment on lines
+590
to
+600
| if hasattr(self.experts, "gate_up_proj"): | ||
| self.moe_weights = build_canonical_expert_weights( | ||
| gate_up=self.experts.gate_up_proj, | ||
| down=self.experts.down_proj, | ||
| fused=True, | ||
| fused_split_dim=1, | ||
| transpose_gate_up=True, | ||
| transpose_down=True, | ||
| ) | ||
| delete_module_attrs(self.experts, "gate_up_proj", "down_proj") | ||
| else: |
Contributor
Author
There was a problem hiding this comment.
this if-else is not needed, as we always know what weights model is going to have, let's fix it instead of handling un-necessary cases.
ochougul
commented
Jul 21, 2026
Comment on lines
+445
to
+446
| order = torch.argsort(out, dim=1) | ||
| last_positions = order[:, -state_len:] |
ochougul
commented
Jul 21, 2026
Comment on lines
+567
to
+570
| mask_causal = torch.ones(chunk_size, chunk_size, dtype=torch.bool) | ||
| for i in range(chunk_size): | ||
| for j in range(i, chunk_size): | ||
| mask_causal[i, j] = True | ||
| for j in range(i + 1): | ||
| mask_causal[i, j] = False |
ochougul
commented
Jul 21, 2026
| g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) | ||
| # mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0) | ||
| mask = mask_causal | ||
| mask = mask_causal.to(device=query.device) |
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
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
1 participant
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.
This PR standardizes QEff MoE model wrappers around a single explicit
QEffMoEBlockMixin contract.
Key changes:
Enforces every QEffMoEBlockMixin subclass to explicitly declare
supported_moe_flavours.
Collapses DeepSeek V3 MoE into a single QEffDeepseekV3MoE block and
removes the separate prefill-only class.
Adds explicit MoE flavour declarations for DeepSeek V3 and Grok1.
Introduces a single Gemma4 text MoE block backed by QEffMoEBlockMixin.
Updates external MoE mapper wiring to source support metadata from the
QEff block classes.
Removes Gemma4’s prefill-only expert mapping in favor of the shared MoE
flavour path.
Impact For New MoE Models
When adding a new MoE model, the model should expose one QEff MoE block that
inherits QEffMoEBlockMixin and explicitly declares:
supported_moe_flavours = (...)
The block should implement the common variation points:
This keeps flavour selection, export config, and future optimizations
centralized instead of spreading model-specific prefill/decode paths across
multiple classes.
Future Scope
The longer-term direction is to support all MoE flavours across all MoE
models where the model architecture allows it, especially:
This PR makes that easier by making supported flavour declarations explicit
and discoverable.
Tests
Added/updated coverage for:
Validation run:
pytest tests/transformers/models/test_moe_prefill_blocked.py -k 'moe_block
or deepseek or grok or gemma4'
pytest tests/unit_test/transforms/test_transform_accuracy.py -k
'OptimizedMoE or deepseek or external_mapper'
pytest tests/unit_test/models/test_model_quickcheck.py -k 'moe_prefill'