Address PR #239 follow-up: GQA limit, GGUF k_eq_v, DRY expert rename, tests - #240
Merged
Conversation
1. Update GQA rewrite rule MAX_HEAD_DIM from 256 to 512 — latest ORT supports head_dim=512 for Gemma4 full-attention layers 2. Set attention_k_eq_v in GGUF postprocessor when num_global_key_value_heads differs from num_key_value_heads, fixing build_from_gguf for 26b/31b 3. Extract _remap_moe_expert_weights() shared helper to avoid duplicating expert rename + router scale folding in CausalLM and multimodal paths 4. Add targeted preprocess_weights tests for both Gemma4CausalLMModel and Gemma4Model (multimodal) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Performance Comparison
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR addresses follow-ups from PR #239 across optimization rewrite rules, GGUF config extraction, Gemma4 weight preprocessing, and adds targeted unit tests to prevent regressions.
Changes:
- Increase the Attention→GroupQueryAttention rewrite rule’s head-dim limit from 256 to 512.
- Set
attention_k_eq_vin the Gemma4 GGUF config postprocessor whennum_global_key_value_headsis detected. - DRY Gemma4 MoE weight preprocessing by extracting a shared
_remap_moe_expert_weights()helper and add new preprocess_weights tests.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
src/mobius/rewrite_rules/_group_query_attention.py |
Raises _MAX_GQA_HEAD_DIM to allow GQA fusion for larger head dims (now 512). |
src/mobius/integrations/gguf/_config_mapping.py |
Infers attention_k_eq_v from GGUF-derived per-layer KV head metadata for Gemma4. |
src/mobius/models/gemma4.py |
Extracts shared MoE expert rename + router scale folding helper used by both Gemma4 preprocess paths. |
src/mobius/models/gemma4_test.py |
Adds unit tests validating expert-weight renames and router scale folding behavior for Gemma4 models. |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
37 tasks
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
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.
Addresses 4 review comments from PR #239:
1. GQA rewrite rule head_dim limit (#1)
Updated
_MAX_GQA_HEAD_DIMfrom 256 to 512 in the Attention→GQA rewrite rule to match latest ORT support for head_dim=512.2. GGUF path missing
attention_k_eq_v(#2)Added
attention_k_eq_v=Trueto Gemma4 GGUF postprocessor whennum_global_key_value_headsis detected. Fixesbuild_from_gguf()for 26b-a4b/31b checkpoints.3. DRY expert weight rename (#3)
Extracted
_remap_moe_expert_weights()shared helper used by bothGemma4CausalLMModel.preprocess_weights()andGemma4Model.preprocess_weights(). Single source of truth for expert rename + router scale folding.4. Missing multimodal preprocess_weights test (#4)
Added
gemma4_test.pywith 5 targeted tests covering bothGemma4CausalLMModelandGemma4Modelweight preprocessing: expert rename, router scale folding, and per_expert_scale passthrough.Testing