Decouple attention_k_eq_v from num_global_key_value_heads

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>

Author: justinchuby

src/mobius/models/gemma4.py

@@ -607,11 +607,12 @@ def __init__(
         self._use_alternative_attention = (
             getattr(config, "attention_k_eq_v", False) and not is_sliding
         )
-        self.num_key_value_heads = (
-            (config.num_global_key_value_heads or config.num_key_value_heads)
-            if self._use_alternative_attention
-            else config.num_key_value_heads
-        )
+        # Full-attention layers use num_global_key_value_heads when set,
+        # independent of the k_eq_v flag.
+        if not is_sliding and config.num_global_key_value_heads is not None:
+            self.num_key_value_heads = config.num_global_key_value_heads
+        else:
+            self.num_key_value_heads = config.num_key_value_heads
 
         # KV sharing: layers >= first_kv_shared_layer_idx borrow K,V from source
         self.is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx > 0

src/mobius/tasks/_gemma4.py

@@ -54,8 +54,8 @@ def _make_gemma4_kv_cache_inputs(
     ``num_hidden_layers - num_kv_shared_layers`` entries are created.
 
     All layers use the original ``num_key_value_heads`` except
-    alternative-attention (``attention_k_eq_v``) full layers which use
-    ``num_global_key_value_heads`` (fewer KV heads).
+    full-attention layers when ``num_global_key_value_heads`` is set
+    (fewer KV heads, independent of the ``attention_k_eq_v`` flag).
     The Attention op supports GQA head counts natively.  CUDA EP limitations
     are tracked in microsoft/onnxruntime#28195 and #28196.
     """
@@ -75,13 +75,13 @@ def _make_gemma4_kv_cache_inputs(
     for i in range(num_kv_layers):
         layer_type = layer_types[i] if i < len(layer_types) else "sliding_attention"
         hd = global_head_dim if layer_type == "full_attention" else local_head_dim
-        # Alternative attention (k_eq_v) uses fewer KV heads for full_attention
-        use_alt = getattr(config, "attention_k_eq_v", False) and layer_type == "full_attention"
-        kv_heads = (
-            (config.num_global_key_value_heads or config.num_key_value_heads)
-            if use_alt
-            else config.num_key_value_heads
-        )
+        # Full-attention layers use num_global_key_value_heads when set,
+        # independent of the k_eq_v flag.
+        is_full = layer_type == "full_attention"
+        if is_full and config.num_global_key_value_heads is not None:
+            kv_heads = config.num_global_key_value_heads
+        else:
+            kv_heads = config.num_key_value_heads
         past_key = builder.input(
             f"past_key_values.{i}.key",
             dtype=config.dtype,

tests/build_graph_test.py

@@ -1415,6 +1415,74 @@ def test_gemma4_kv_shared_layer_tracing(self):
         assert "present.1.key" not in output_names
         assert "present.1.value" not in output_names
 
+    def test_gemma4_k_eq_v_with_global_kv_heads(self):
+        """Verify attention_k_eq_v removes v_proj and num_global_key_value_heads sets KV cache shapes.
+
+        Config: attention_k_eq_v=True, num_key_value_heads=4 (sliding),
+        num_global_key_value_heads=2 (full). Full-attention layers should:
+        - Have no v_proj initializer (V=K)
+        - Use num_global_key_value_heads=2 for KV cache shapes
+        Sliding layers should use num_key_value_heads=4.
+        """
+        from mobius._configs import Gemma4Config
+        from mobius.models.gemma4 import Gemma4CausalLMModel
+        from mobius.tasks._gemma4 import Gemma4TextCausalLMTask
+
+        config = Gemma4Config(
+            num_hidden_layers=2,
+            hidden_size=64,
+            intermediate_size=128,
+            num_attention_heads=4,
+            num_key_value_heads=4,
+            head_dim=16,
+            vocab_size=256,
+            rms_norm_eps=1e-6,
+            hidden_act="silu",
+            attn_qk_norm=True,
+            # Layer 0: sliding, Layer 1: full (k_eq_v + global heads)
+            layer_types=["sliding_attention", "full_attention"],
+            sliding_window=8,
+            global_head_dim=16,
+            global_rope_theta=10_000.0,
+            global_partial_rotary_factor=0.25,
+            final_logit_softcapping=0.0,
+            hidden_size_per_layer_input=0,
+            image_token_id=255999,
+            pad_token_id=0,
+            tie_word_embeddings=True,
+            attention_k_eq_v=True,
+            num_global_key_value_heads=2,
+        )
+        module = Gemma4CausalLMModel(config)
+        task = Gemma4TextCausalLMTask()
+        pkg = task.build(module, config)
+        decoder = pkg["model"]
+
+        # Check initializer names: full-attention layer (1) should have no v_proj
+        init_names = set(decoder.graph.initializers)
+        # Sliding layer 0 has k_proj, v_proj
+        assert "model.layers.0.self_attn.k_proj.weight" in init_names
+        assert "model.layers.0.self_attn.v_proj.weight" in init_names
+        # Full layer 1 has k_proj but NO v_proj (k_eq_v: V=K)
+        assert "model.layers.1.self_attn.k_proj.weight" in init_names
+        assert "model.layers.1.self_attn.v_proj.weight" not in init_names
+
+        # KV cache shapes:
+        # Layer 0 (sliding): num_key_value_heads=4
+        # Layer 1 (full): num_global_key_value_heads=2
+        input_shapes = {i.name: list(i.shape) for i in decoder.graph.inputs}
+        # Layer 0: kv_heads=4
+        layer0_key_shape = input_shapes["past_key_values.0.key"]
+        assert layer0_key_shape[1] == 4, (
+            f"Sliding layer 0 should have 4 KV heads, got {layer0_key_shape[1]}"
+        )
+        # Layer 1: kv_heads=2 (num_global_key_value_heads)
+        layer1_key_shape = input_shapes["past_key_values.1.key"]
+        assert layer1_key_shape[1] == 2, (
+            f"Full layer 1 should have 2 KV heads "
+            f"(num_global_key_value_heads), got {layer1_key_shape[1]}"
+        )
+
     def test_blip2_vision_language_graph(self):
         """Build BLIP-2 with ViT + Q-Former + LLM 3-model split."""
         config = _base_config(