Address PR #239 follow-up: GQA limit, GGUF k_eq_v, DRY expert rename, 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>

Author: 3 people

src/mobius/integrations/gguf/_config_mapping.py

@@ -465,6 +465,9 @@ def _gemma4_postprocess(
         else 1_000_000.0,
         global_partial_rotary_factor=global_partial_rotary_factor,
         num_global_key_value_heads=num_global_key_value_heads,
+        # attention_k_eq_v: derive from per-layer KV head counts. When
+        # full-attention layers use fewer KV heads, V = K (no v_proj).
+        attention_k_eq_v=num_global_key_value_heads is not None,
         final_logit_softcapping=float(final_logit_softcapping or 0.0),
         attn_logit_softcapping=float(attn_logit_softcapping or 0.0),
         num_kv_shared_layers=int(num_kv_shared_layers)

src/mobius/models/gemma4.py

@@ -55,6 +55,39 @@
 if TYPE_CHECKING:
     from mobius.components._attention import GQAContext
 
+
+# ---------------------------------------------------------------------------
+# Shared weight preprocessing helpers
+# ---------------------------------------------------------------------------
+
+
+def _remap_moe_expert_weights(
+    state_dict: dict[str, torch.Tensor],
+    config: Gemma4Config,
+) -> None:
+    """Rename HF MoE expert weights and fold router scale in-place.
+
+    Shared by ``Gemma4CausalLMModel`` and ``Gemma4Model`` to avoid
+    duplicating the rename/fold logic.
+    """
+    # experts.gate_up_proj → fc1_experts_weights
+    # experts.down_proj    → fc2_experts_weights
+    for key in list(state_dict.keys()):
+        if ".experts.gate_up_proj" in key:
+            new_key = key.replace(".experts.gate_up_proj", ".fc1_experts_weights")
+            state_dict[new_key] = state_dict.pop(key)
+        elif ".experts.down_proj" in key:
+            new_key = key.replace(".experts.down_proj", ".fc2_experts_weights")
+            state_dict[new_key] = state_dict.pop(key)
+
+    # Fold hidden_size^-0.5 into router.scale
+    if config.enable_moe_block:
+        scale_factor = float(config.hidden_size**-0.5)
+        for key in list(state_dict.keys()):
+            if ".router.scale" in key and ".per_expert_scale" not in key:
+                state_dict[key] = state_dict[key] * scale_factor
+
+
 # ---------------------------------------------------------------------------
 # Scale-free RMSNorm (Gemma4RMSNorm with with_scale=False)
 # ---------------------------------------------------------------------------
@@ -1653,27 +1686,8 @@ def preprocess_weights(
                 for i in range(num_layers):
                     shard = value[:, i * per_layer_dim : (i + 1) * per_layer_dim]
                     state_dict[f"model.embed_tokens_per_layer.{i}.weight"] = shard
-        # Map HF expert weight names to our 3D stacked parameter names.
-        # HF stores: layers.N.experts.gate_up_proj [E, 2*inter, H]
-        #             layers.N.experts.down_proj     [E, H, inter]
-        # We store:  layers.N.fc1_experts_weights   [E, 2*inter, H]
-        #             layers.N.fc2_experts_weights   [E, H, inter]
-        for key in list(state_dict.keys()):
-            if ".experts.gate_up_proj" in key:
-                new_key = key.replace(".experts.gate_up_proj", ".fc1_experts_weights")
-                state_dict[new_key] = state_dict.pop(key)
-            elif ".experts.down_proj" in key:
-                new_key = key.replace(".experts.down_proj", ".fc2_experts_weights")
-                state_dict[new_key] = state_dict.pop(key)
-        # Fold hidden_size^-0.5 into router.scale.
-        # The router computes: x_normed * scale * hidden_size^-0.5.
-        # We pre-multiply scale by hidden_size^-0.5 here so the forward only needs
-        # x_normed * self.scale, avoiding float-constant name collisions across layers.
-        if self.config.enable_moe_block:
-            scale_factor = float(self.config.hidden_size**-0.5)
-            for key in list(state_dict.keys()):
-                if ".router.scale" in key:
-                    state_dict[key] = state_dict[key] * scale_factor
+        # Map HF expert weight names and fold router scale
+        _remap_moe_expert_weights(state_dict, self.config)
         return super().preprocess_weights(state_dict)
 
 
@@ -2150,22 +2164,7 @@ def preprocess_weights(
             else:
                 renamed[key] = value
 
-        # 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
+        # Map HF expert weight names and fold router scale
+        _remap_moe_expert_weights(renamed, self.config)
 
         return renamed

src/mobius/models/gemma4_test.py

@@ -0,0 +1,119 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Tests for Gemma4 preprocess_weights — MoE expert rename and router scale."""
+
+from __future__ import annotations
+
+import torch
+
+from mobius._configs import Gemma4Config
+from mobius.models.gemma4 import Gemma4CausalLMModel, Gemma4Model
+
+
+def _tiny_gemma4_config(**overrides) -> Gemma4Config:
+    """Create a minimal Gemma4Config for preprocess_weights tests."""
+    from mobius._configs import VisionConfig
+
+    defaults = dict(
+        model_type="gemma4",
+        vocab_size=256,
+        hidden_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        intermediate_size=128,
+        head_dim=16,
+        global_head_dim=32,
+        hidden_act="gelu",
+        enable_moe_block=True,
+        num_local_experts=4,
+        num_experts_per_tok=2,
+        moe_intermediate_size=32,
+        layer_types=["sliding_attention", "full_attention"],
+        attention_k_eq_v=True,
+        num_global_key_value_heads=1,
+        vision=VisionConfig(
+            hidden_size=32,
+            intermediate_size=64,
+            num_hidden_layers=1,
+            num_attention_heads=2,
+            image_size=16,
+            patch_size=4,
+        ),
+    )
+    defaults.update(overrides)
+    return Gemma4Config(**defaults)
+
+
+class TestGemma4CausalLMPreprocessWeights:
+    """Test Gemma4CausalLMModel.preprocess_weights."""
+
+    def test_expert_weight_rename(self):
+        config = _tiny_gemma4_config()
+        model = Gemma4CausalLMModel(config)
+
+        fake_sd = {
+            "model.layers.0.experts.gate_up_proj": torch.zeros(4, 64, 64),
+            "model.layers.0.experts.down_proj": torch.zeros(4, 64, 32),
+        }
+        result = model.preprocess_weights(fake_sd)
+
+        assert "model.layers.0.fc1_experts_weights" in result
+        assert "model.layers.0.fc2_experts_weights" in result
+        assert "model.layers.0.experts.gate_up_proj" not in result
+
+    def test_router_scale_folding(self):
+        config = _tiny_gemma4_config()
+        model = Gemma4CausalLMModel(config)
+
+        scale_val = torch.tensor([2.0])
+        fake_sd = {"model.layers.0.router.scale": scale_val.clone()}
+        result = model.preprocess_weights(fake_sd)
+
+        expected = 2.0 * (64**-0.5)  # hidden_size=64
+        assert abs(result["model.layers.0.router.scale"].item() - expected) < 1e-6
+
+
+class TestGemma4ModelPreprocessWeights:
+    """Test Gemma4Model.preprocess_weights (multimodal path)."""
+
+    def test_expert_weight_rename(self):
+        config = _tiny_gemma4_config()
+        model = Gemma4Model(config)
+
+        fake_sd = {
+            "model.language_model.layers.0.experts.gate_up_proj": torch.zeros(4, 64, 64),
+            "model.language_model.layers.0.experts.down_proj": torch.zeros(4, 64, 32),
+        }
+        result = model.preprocess_weights(fake_sd)
+
+        assert "decoder.model.layers.0.fc1_experts_weights" in result
+        assert "decoder.model.layers.0.fc2_experts_weights" in result
+
+    def test_router_scale_folding(self):
+        config = _tiny_gemma4_config()
+        model = Gemma4Model(config)
+
+        scale_val = torch.tensor([2.0])
+        fake_sd = {"model.language_model.layers.0.router.scale": scale_val.clone()}
+        result = model.preprocess_weights(fake_sd)
+
+        expected = 2.0 * (64**-0.5)
+        key = "decoder.model.layers.0.router.scale"
+        assert key in result
+        assert abs(result[key].item() - expected) < 1e-6
+
+    def test_per_expert_scale_not_folded(self):
+        """router.per_expert_scale should NOT be multiplied by scale_factor."""
+        config = _tiny_gemma4_config()
+        model = Gemma4Model(config)
+
+        fake_sd = {
+            "model.language_model.layers.0.router.per_expert_scale": torch.ones(4),
+        }
+        result = model.preprocess_weights(fake_sd)
+
+        key = "decoder.model.layers.0.router.per_expert_scale"
+        assert key in result
+        assert torch.allclose(result[key], torch.ones(4))

src/mobius/rewrite_rules/_group_query_attention.py

@@ -49,11 +49,10 @@
     RewriteRuleSet,
 )
 
-# CUDA EP's GroupQueryAttention kernel enforces MAX_HEAD_SIZE = 256.
-# Rewriting Attention → GQA for layers with head_dim > 256 would produce
-# a node that fails at runtime.  We skip those nodes so they stay as
-# standard Attention ops (which can use the MHA code-path when
-# q_num_heads == kv_num_heads after KV expansion).
+# CUDA EP's GroupQueryAttention kernel historically enforced MAX_HEAD_SIZE = 256.
+# Keep the rewrite-rule limit conservative until GQA fusion is gated on a
+# runtime/EP capability check. This avoids emitting GroupQueryAttention nodes
+# with head_dim=512 that can still fail on released ORT builds.
 _MAX_GQA_HEAD_DIM = 256