K3: Gemma4-faithful cross-model restore forward (per-layer KV, v_norm, RoPE unsqueeze_dim=2, v_proj-None, evicted slicing) + gemma4 helpers import + tests

Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

Author: cursoragent

inference_engine/v04/cross_model_dlm_verifier.py

@@ -247,14 +247,15 @@ def _validate_dimensions(self) -> None:
     @torch.no_grad()
     def project_drafter_kv(
         self, input_ids: torch.Tensor,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
         """Run the drafter forward over input_ids, project K/V through f_θ.
 
         Returns
         -------
-        (verifier_k, verifier_v) tensors of shape
-        ``[B, T, verifier_num_layers, verifier_num_kv_heads, verifier_head_dim]``
-        on the f_θ device.
+        (verifier_k, verifier_v): per-layer lists of length
+        ``verifier_num_layers``, element ``i`` shaped
+        ``[B, T, layer_kv_heads[i], verifier_head_dim]`` on the f_θ
+        device. Per-layer KV-head counts can differ (Gemma 4).
 
         These are the per-position-per-verifier-layer K/V that the
         cross-model verifier injects at evicted positions during its
@@ -309,9 +310,9 @@ def forward(
         if not evicted_positions:
             return self.verifier_model(input_ids=input_ids, use_cache=False)
 
-        # f_θ projection
-        verifier_k_full, verifier_v_full = self.project_drafter_kv(input_ids)
-        # verifier_k_full shape: [B, T, L_v, num_kv_heads_v, head_dim_v]
+        # f_θ projection → per-layer lists (layers can have different
+        # KV-head counts on Gemma 4).
+        verifier_k_layers, verifier_v_layers = self.project_drafter_kv(input_ids)
 
         # Patch verifier attention forwards to inject K/V at evicted
         # positions. Restore originals after the forward.
@@ -325,8 +326,8 @@ def forward(
                     attn,
                     layer_idx=layer_idx,
                     evicted_positions=evicted_positions,
-                    verifier_k_at_layer=verifier_k_full[:, :, layer_idx],
-                    verifier_v_at_layer=verifier_v_full[:, :, layer_idx],
+                    verifier_k_at_layer=verifier_k_layers[layer_idx],
+                    verifier_v_at_layer=verifier_v_layers[layer_idx],
                     apply_rotary_pos_emb=apply_rotary_pos_emb,
                     eager_attention_forward=eager_attention_forward,
                     all_attention_functions=all_attention_functions,
@@ -351,47 +352,72 @@ def _make_patched_forward(
         instead of using the verifier's own k_proj / v_proj at those
         positions.
 
-        The patched forward replicates the standard verifier attention
-        layer (Q, K, V projections + RoPE + GQA + softmax) with one
-        change: after K, V are computed at every position, K and V at
-        evicted positions are OVERWRITTEN with the f_θ-projected values
-        (after k_norm + RoPE applied to match the standard pipeline).
+        Mirrors Gemma 4's ``Gemma4TextAttention.forward`` exactly (RoPE
+        applied per-tensor on the ``[B, T, H, D]`` layout with
+        ``unsqueeze_dim=2``; ``q_norm`` / ``k_norm`` / ``v_norm``;
+        ``v_proj`` is ``None`` on full-attention layers where V = raw K)
+        with one change: at evicted positions K/V come from the
+        f_θ-projected values (after k_norm + RoPE for K, after v_norm
+        for V), so the verifier attends over full context despite
+        holding only sink+window in its local cache.
         """
         def _patched_forward(
             hidden_states: torch.Tensor,
             position_embeddings: Tuple[torch.Tensor, torch.Tensor],
             attention_mask: Optional[torch.Tensor] = None,
+            shared_kv_states: Any = None,
             past_key_values=None,
             cache_position=None,
             **kwargs,
         ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
-            B, T, _ = hidden_states.shape
-
-            input_shape = (B, T)
-            hidden_shape = (*input_shape, -1, attn_module.head_dim)
-
-            query_states = attn_module.q_proj(hidden_states).view(*hidden_shape).transpose(1, 2)
-            key_states = attn_module.k_proj(hidden_states).view(*hidden_shape).transpose(1, 2)
-            value_states = attn_module.v_proj(hidden_states).view(*hidden_shape).transpose(1, 2)
+            input_shape = hidden_states.shape[:-1]   # [B, T]
+            head_dim = attn_module.head_dim
+            hidden_shape = (*input_shape, -1, head_dim)
+            cos, sin = position_embeddings
 
+            # Query (Gemma 4: norm → RoPE on [B,T,Hq,D] → transpose)
+            query_states = attn_module.q_proj(hidden_states).view(hidden_shape)
             query_states = attn_module.q_norm(query_states)
-            key_states = attn_module.k_norm(key_states)
+            query_states = apply_rotary_pos_emb(
+                query_states, cos, sin, unsqueeze_dim=2,
+            )
+            query_states = query_states.transpose(1, 2)   # [B, Hq, T, D]
 
-            cos, sin = position_embeddings
-            query_states, key_states = apply_rotary_pos_emb(
-                query_states, key_states, cos, sin,
+            # Key / Value. Full-attention layers have v_proj=None ⇒ the
+            # value is the raw k_proj output (pre k_norm), per Gemma 4.
+            k_lin = attn_module.k_proj(hidden_states).view(hidden_shape)
+            if getattr(attn_module, "v_proj", None) is not None:
+                v_lin = attn_module.v_proj(hidden_states).view(hidden_shape)
+            else:
+                v_lin = k_lin
+            key_states = attn_module.k_norm(k_lin)
+            key_states = apply_rotary_pos_emb(
+                key_states, cos, sin, unsqueeze_dim=2,
             )
+            key_states = key_states.transpose(1, 2)        # [B, Hkv, T, D]
+            value_states = attn_module.v_norm(v_lin).transpose(1, 2)  # [B, Hkv, T, D]
 
             # Inject f_θ K/V at evicted positions.
             # verifier_k_at_layer shape: [B, T, num_kv_heads_v, head_dim_v]
-            # K/V from k_proj also at all T positions; we overwrite the
-            # evicted slice with f_θ output (after k_norm + RoPE).
+            # (pre-norm pre-RoPE). prepare_restored_attention_kv applies
+            # k_norm + RoPE to K; V must be v_norm'd here to match the
+            # local branch (Gemma 4 runs V through v_norm).
             if evicted_positions:
+                idx = torch.tensor(
+                    evicted_positions, device=key_states.device, dtype=torch.long,
+                )
+                cap_k_pre = verifier_k_at_layer.index_select(1, idx).to(
+                    device=key_states.device, dtype=key_states.dtype,
+                )
+                cap_v_pre = verifier_v_at_layer.index_select(1, idx).to(
+                    device=value_states.device, dtype=value_states.dtype,
+                )
+                cap_v_norm = attn_module.v_norm(cap_v_pre)
                 key_states, value_states = prepare_restored_attention_kv(
                     K_local=key_states,
                     V_local=value_states,
-                    captured_K_pre_norm=verifier_k_at_layer,
-                    captured_V=verifier_v_at_layer,
+                    captured_K_pre_norm=cap_k_pre,
+                    captured_V=cap_v_norm,
                     evicted_positions=evicted_positions,
                     k_norm=attn_module.k_norm,
                     position_embeddings=(cos, sin),

scripts/research/k3_integrated_niah_eval.py

@@ -125,7 +125,7 @@ def main() -> int:
 
     # ---------- Verifier (CUDA bf16) ----------
     from transformers import AutoModelForCausalLM, AutoTokenizer
-    from transformers.models.gemma3.modeling_gemma3 import (  # type: ignore
+    from transformers.models.gemma4.modeling_gemma4 import (  # type: ignore
         apply_rotary_pos_emb, eager_attention_forward, ALL_ATTENTION_FUNCTIONS,
     )
 

tests/inference_engine/v04/test_cross_model_dlm_verifier.py

@@ -77,7 +77,9 @@ def __init__(self) -> None:
         self.o_proj = nn.Linear(32, 32, bias=False)
         self.q_norm = nn.Identity()
         self.k_norm = nn.Identity()
+        self.v_norm = nn.Identity()   # Gemma 4 runs V through v_norm
         self.head_dim = 8
+        self.num_key_value_groups = 1
         self.scaling = 8 ** -0.5
         self.attention_dropout = 0.0
         self.sliding_window = None
@@ -217,11 +219,15 @@ def test_returns_correct_shape(self):
         B, T = 1, 6
         ids = torch.randint(0, 64, (B, T), dtype=torch.long)
         v_k, v_v = v.project_drafter_kv(ids)
-        assert tuple(v_k.shape) == (
-            B, T, f_cfg.verifier_num_layers,
-            f_cfg.verifier_num_kv_heads, f_cfg.verifier_head_dim,
+        # Per-layer list contract (layers may have heterogeneous KV heads).
+        assert len(v_k) == f_cfg.verifier_num_layers
+        assert len(v_v) == f_cfg.verifier_num_layers
+        per_layer = (
+            B, T, f_cfg.verifier_num_kv_heads, f_cfg.verifier_head_dim,
         )
-        assert tuple(v_v.shape) == tuple(v_k.shape)
+        for ko, vo in zip(v_k, v_v):
+            assert tuple(ko.shape) == per_layer
+            assert tuple(vo.shape) == per_layer
 
 
 class TestNoEvictPath:
@@ -268,6 +274,52 @@ def _counted(ids):
         assert calls["drafter"] == 0
 
 
+def _gemma4_style_rope(x, cos, sin, unsqueeze_dim=2):
+    """Mirror Gemma 4's apply_rotary_pos_emb(x, cos, sin, unsqueeze_dim)."""
+    c = cos.unsqueeze(unsqueeze_dim)
+    s = sin.unsqueeze(unsqueeze_dim)
+    half = x.shape[-1] // 2
+    rot = torch.cat([-x[..., half:], x[..., :half]], dim=-1)
+    return x * c + rot * s
+
+
+def _synthetic_eager(module, q, k, v, mask, dropout=0.0, scaling=1.0,
+                     sliding_window=None, **kw):
+    """Minimal eager attention returning [B, T, H, D] like HF's."""
+    attn = torch.matmul(q, k.transpose(-1, -2)) * scaling
+    if mask is not None:
+        attn = attn + mask[..., : k.shape[-2]]
+    attn = attn.softmax(dim=-1)
+    out = torch.matmul(attn, v)             # [B, H, T, D]
+    return out.transpose(1, 2).contiguous(), None   # [B, T, H, D]
+
+
+class TestPatchedForwardRestore:
+    """Exercise the Gemma 4-style patched attention forward end-to-end on
+    the synthetic verifier with real evictions, so signature/shape/RoPE
+    regressions are caught without the 26B model."""
+
+    def test_forward_with_eviction_runs_and_injects(self):
+        f_cfg = _tiny_f_theta_config()
+        f_theta = FThetaProjection(f_cfg)
+        drafter = DFlashDrafter(_tiny_drafter_config())
+        verifier = _SyntheticVerifier()
+        verifier.embed_tokens = torch.nn.Embedding(64, 16)
+        verifier.get_input_embeddings = lambda: verifier.embed_tokens
+        v = CrossModelDLMRestoredVerifier(
+            verifier_model=verifier, drafter=drafter, f_theta=f_theta,
+            sink_size=1, window_size=2,    # sink+window = 3
+        )
+        B, T = 1, 6                        # evicted positions [1, 2, 3]
+        ids = torch.randint(0, 64, (B, T), dtype=torch.long)
+        out = v.forward(
+            ids,
+            apply_rotary_pos_emb=_gemma4_style_rope,
+            eager_attention_forward=_synthetic_eager,
+        )
+        assert out.logits.shape == (B, T, 64)
+
+
 class TestExports:
 
     def test_module_exposes_classes(self):

tests/inference_engine/v04/test_f_theta.py

@@ -77,19 +77,42 @@ def test_forward_k_shape(self):
         B, T = 2, 7
         x = torch.randn(B, T, c.encoder_in_features)
         y = m.forward_k(x)
-        assert tuple(y.shape) == (
-            B, T, c.verifier_num_layers, c.verifier_num_kv_heads, c.verifier_head_dim,
-        )
+        assert isinstance(y, list)
+        assert len(y) == c.verifier_num_layers
+        for layer_out in y:
+            assert tuple(layer_out.shape) == (
+                B, T, c.verifier_num_kv_heads, c.verifier_head_dim,
+            )
 
     def test_forward_v_shape(self):
         c = _tiny_config()
         m = FThetaProjection(c)
         B, T = 1, 3
         x = torch.randn(B, T, c.encoder_in_features)
         y = m.forward_v(x)
-        assert tuple(y.shape) == (
-            B, T, c.verifier_num_layers, c.verifier_num_kv_heads, c.verifier_head_dim,
+        assert isinstance(y, list)
+        assert len(y) == c.verifier_num_layers
+        for layer_out in y:
+            assert tuple(layer_out.shape) == (
+                B, T, c.verifier_num_kv_heads, c.verifier_head_dim,
+            )
+
+    def test_heterogeneous_layer_kv_heads(self):
+        """Per-layer KV-head counts (Gemma 4: 8 sliding, 4 full)."""
+        c = FThetaConfig(
+            drafter_num_layers=2, drafter_num_kv_heads=2, drafter_head_dim=4,
+            verifier_num_layers=4, verifier_num_kv_heads=8, verifier_head_dim=8,
+            rank=16, verifier_layer_kv_heads=(8, 4, 8, 4),
         )
+        assert c.layer_kv_dims == (64, 32, 64, 32)
+        m = FThetaProjection(c)
+        B, T = 1, 3
+        x = torch.randn(B, T, c.encoder_in_features)
+        y = m.forward_k(x)
+        assert [t.shape[2] for t in y] == [8, 4, 8, 4]
+        # JSON round-trip preserves the per-layer field
+        c2 = FThetaConfig.from_json_dict(c.to_json_dict())
+        assert c2 == c
 
     def test_forward_k_rejects_wrong_rank(self):
         c = _tiny_config()
@@ -122,9 +145,12 @@ def test_returns_paired_k_v(self):
             for _ in range(c.drafter_num_layers)
         ]
         k_out, v_out = m.forward_kv_pack(k_per_layer, v_per_layer)
-        expected = (B, T, c.verifier_num_layers, c.verifier_num_kv_heads, c.verifier_head_dim)
-        assert tuple(k_out.shape) == expected
-        assert tuple(v_out.shape) == expected
+        assert len(k_out) == c.verifier_num_layers
+        assert len(v_out) == c.verifier_num_layers
+        per_layer = (B, T, c.verifier_num_kv_heads, c.verifier_head_dim)
+        for ko, vo in zip(k_out, v_out):
+            assert tuple(ko.shape) == per_layer
+            assert tuple(vo.shape) == per_layer
 
     def test_rejects_wrong_layer_count(self):
         c = _tiny_config()
@@ -162,8 +188,10 @@ def test_consistency_with_explicit_concat(self):
         with torch.no_grad():
             k_out_direct = m.forward_k(k_concat)
             v_out_direct = m.forward_v(v_concat)
-        assert torch.allclose(k_out_pack, k_out_direct, atol=1e-6)
-        assert torch.allclose(v_out_pack, v_out_direct, atol=1e-6)
+        for kp, kd in zip(k_out_pack, k_out_direct):
+            assert torch.allclose(kp, kd, atol=1e-6)
+        for vp, vd in zip(v_out_pack, v_out_direct):
+            assert torch.allclose(vp, vd, atol=1e-6)
 
 
 class TestParameterCount:
@@ -218,8 +246,10 @@ def test_save_and_load_preserves_outputs(self, tmp_path):
             y_k_2 = m2.forward_k(x_k)
             y_v_2 = m2.forward_v(x_v)
 
-        assert torch.allclose(y_k_1, y_k_2)
-        assert torch.allclose(y_v_1, y_v_2)
+        for a, b in zip(y_k_1, y_k_2):
+            assert torch.allclose(a, b)
+        for a, b in zip(y_v_1, y_v_2):
+            assert torch.allclose(a, b)
 
     def test_load_rejects_missing_config(self, tmp_path):
         # Write only weights, no config
@@ -267,12 +297,8 @@ def test_gradients_flow_for_k_path(self):
         m = FThetaProjection(c)
         B, T = 1, 3
         x = torch.randn(B, T, c.encoder_in_features, requires_grad=False)
-        target = torch.randn(
-            B, T, c.verifier_num_layers,
-            c.verifier_num_kv_heads, c.verifier_head_dim,
-        )
-        out = m.forward_k(x)
-        loss = ((out - target) ** 2).mean()
+        out = m.forward_k(x)   # list of [B, T, H, D]
+        loss = sum(((o) ** 2).mean() for o in out)
         loss.backward()
         # encoder_k should have a grad
         assert m.encoder_k.weight.grad is not None