Fix Gemma4 weight mapping for 26b MoE and 31b alternative attention - #239
Merged
Conversation
When attention_k_eq_v=True (26b-a4b, 31b models), full_attention layers derive V from K (no separate v_proj), matching HF Gemma4TextAttention. Changes: - Add attention_k_eq_v and num_global_key_value_heads to Gemma4Config - Skip v_proj creation for alternative-attention layers - In forward: V = raw k_proj output (before k_norm/RoPE), then v_norm - Per-layer KV head count: full_attention uses num_global_key_value_heads - Update task KV cache to use per-layer head counts Tests: 2680 passed (12 gemma4), lint clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Gemma4Model (multimodal) has its own preprocess_weights that was missing the expert weight rename (gate_up_proj → fc1_experts_weights) and router scale folding. These were only in Gemma4CausalLMModel.preprocess_weights. Added both transformations to the multimodal model's preprocess_weights, fixing 26b-a4b and 26b-a4b-it exports. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
The need_fallback check only considered KV-shared layers but not layers where head_dim exceeds CUDA GQA's max (256). Full-attention layers with global_head_dim=512 fell through to GQA fallback with an empty bias dict, causing KeyError: 'full_attention'. Fix: include head_dim > 256 in the need_fallback condition. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Latest ORT supports head_dim=512 in GroupQueryAttention. Remove the fallback that routed full-attention layers (global_head_dim=512) to standard Attention instead of GQA. All non-shared layers now use GQA on CUDA EP regardless of head_dim. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Decouple two independent Gemma4 features: - num_global_key_value_heads: controls KV head count for full-attention layers. Now used whenever set (not None), regardless of attention_k_eq_v. - attention_k_eq_v: controls V=K sharing (no v_proj). Gated separately. Previously both were coupled: num_global_key_value_heads only took effect when attention_k_eq_v was True. This was incorrect — a model could have different KV head counts for global/local layers without sharing V=K. Changes: - gemma4.py: Use num_global_key_value_heads for full-attention layers whenever it's not None - _gemma4.py: Same decoupling in _make_gemma4_kv_cache_inputs - build_graph_test.py: Add test with k_eq_v=True and num_global_key_value_heads != num_key_value_heads, verifying no v_proj and correct KV cache shapes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes Gemma4-specific export issues in the model/config/task stack: it adds support for attention_k_eq_v full-attention layers and ports missing MoE weight remapping into the multimodal Gemma4 entry point.
Changes:
- Added
attention_k_eq_vtoGemma4Configand used it in Gemma4 attention construction/forward to omitv_projand derive V from K on full-attention layers. - Updated Gemma4 KV-cache shape generation so full-attention layers can use
num_global_key_value_heads. - Added multimodal Gemma4 MoE weight remapping/router-scale folding and a new Gemma4 build-graph regression test.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
tests/build_graph_test.py |
Adds a Gemma4 graph test for attention_k_eq_v initializer layout and per-layer KV-cache head counts. |
src/mobius/tasks/_gemma4.py |
Updates Gemma4 KV-cache input construction to use per-layer global KV head counts for full-attention layers. |
src/mobius/models/gemma4.py |
Implements alternative attention (V=K) handling, changes GQA selection logic, and adds multimodal MoE weight remapping. |
src/mobius/_configs.py |
Extends Gemma4Config and HF config extraction with the new attention_k_eq_v field. |
Comment on lines
1552
to
+1555
| # Per-layer decision: use GQA for non-shared layers when | ||
| # available, fall back to standard Attention for KV-shared layers | ||
| # and layers where head_dim exceeds CUDA GQA MAX_HEAD_SIZE (256). | ||
| # available, fall back to standard Attention for KV-shared layers. | ||
| is_shared = layer.self_attn.is_kv_shared_layer | ||
| gqa_head_dim_ok = layer.self_attn.head_dim <= 256 | ||
| if use_gqa and not is_shared and gqa_head_dim_ok: | ||
| if use_gqa and not is_shared: |
Comment on lines
1849
to
+1850
| enable_moe_block=getattr(config, "enable_moe_block", False), | ||
| attention_k_eq_v=getattr(config, "attention_k_eq_v", False), |
Comment on lines
+2153
to
+2169
| # Map HF expert weight names to our 3D stacked parameter names. | ||
| # HF: decoder.model.layers.N.experts.gate_up_proj [E, 2*inter, H] | ||
| # Us: decoder.model.layers.N.fc1_experts_weights [E, 2*inter, H] | ||
| for key in list(renamed.keys()): | ||
| if ".experts.gate_up_proj" in key: | ||
| new_key = key.replace(".experts.gate_up_proj", ".fc1_experts_weights") | ||
| renamed[new_key] = renamed.pop(key) | ||
| elif ".experts.down_proj" in key: | ||
| new_key = key.replace(".experts.down_proj", ".fc2_experts_weights") | ||
| renamed[new_key] = renamed.pop(key) | ||
|
|
||
| # Fold hidden_size^-0.5 into router.scale | ||
| if self.config.enable_moe_block: | ||
| scale_factor = float(self.config.hidden_size**-0.5) | ||
| for key in list(renamed.keys()): | ||
| if ".router.scale" in key and ".per_expert_scale" not in key: | ||
| renamed[key] = renamed[key] * scale_factor |
Comment on lines
+2153
to
+2169
| # Map HF expert weight names to our 3D stacked parameter names. | ||
| # HF: decoder.model.layers.N.experts.gate_up_proj [E, 2*inter, H] | ||
| # Us: decoder.model.layers.N.fc1_experts_weights [E, 2*inter, H] | ||
| for key in list(renamed.keys()): | ||
| if ".experts.gate_up_proj" in key: | ||
| new_key = key.replace(".experts.gate_up_proj", ".fc1_experts_weights") | ||
| renamed[new_key] = renamed.pop(key) | ||
| elif ".experts.down_proj" in key: | ||
| new_key = key.replace(".experts.down_proj", ".fc2_experts_weights") | ||
| renamed[new_key] = renamed.pop(key) | ||
|
|
||
| # Fold hidden_size^-0.5 into router.scale | ||
| if self.config.enable_moe_block: | ||
| scale_factor = float(self.config.hidden_size**-0.5) | ||
| for key in list(renamed.keys()): | ||
| if ".router.scale" in key and ".per_expert_scale" not in key: | ||
| renamed[key] = renamed[key] * scale_factor |
justinchuby
added a commit
that referenced
this pull request
May 4, 2026
… tests (#240) Addresses 4 review comments from PR #239: ### 1. GQA rewrite rule head_dim limit (#1) Updated `_MAX_GQA_HEAD_DIM` from 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=True` to Gemma4 GGUF postprocessor when `num_global_key_value_heads` is detected. Fixes `build_from_gguf()` for 26b-a4b/31b checkpoints. ### 3. DRY expert weight rename (#3) Extracted `_remap_moe_expert_weights()` shared helper used by both `Gemma4CausalLMModel.preprocess_weights()` and `Gemma4Model.preprocess_weights()`. Single source of truth for expert rename + router scale folding. ### 4. Missing multimodal preprocess_weights test (#4) Added `gemma4_test.py` with 5 targeted tests covering both `Gemma4CausalLMModel` and `Gemma4Model` weight preprocessing: expert rename, router scale folding, and per_expert_scale passthrough. ### Testing - 5/5 new preprocess_weights tests pass - 2668/2668 full suite pass --------- Signed-off-by: Justin Chu <justinchu@microsoft.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
May 4, 2026
- Rebase onto latest main, dropping Gemma4 commits (already in PR #239) - Document batch=1 limitation in FunASREmbeddingModel - Document audio_token_id=0 collision workaround - Add explicit shape-unknown comment in SkipLayerNorm rank guard Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
justinchuby
added a commit
that referenced
this pull request
May 4, 2026
- Rebase onto latest main, dropping Gemma4 commits (already in PR #239) - Document batch=1 limitation in FunASREmbeddingModel - Document audio_token_id=0 collision workaround - Add explicit shape-unknown comment in SkipLayerNorm rank guard Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.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.
Summary
Two weight mapping fixes for Gemma4 models that use
attention_k_eq_v=True(26b-a4b, 31b):Fix 1: Alternative attention (V=K) for full_attention layers
When
attention_k_eq_v=True, full-attention layers derive V from K (no separatev_proj), matching HFGemma4TextAttention. Changes:v_projcreation for alternative-attention layersk_projoutput (before k_norm/RoPE), then v_norm appliednum_global_key_value_headsFix 2: MoE expert weight mapping in Gemma4Model
Gemma4Model(multimodal) had its ownpreprocess_weightsthat was missing:experts.gate_up_proj→fc1_experts_weights)hidden_size^-0.5)These were only in
Gemma4CausalLMModel.preprocess_weights.Testing
Models affected