K3: add identity-restore diagnostic (inject verifier's own K/V) to isolate restore machinery from f_theta accuracy

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

Author: cursoragent

inference_engine/v04/cross_model_dlm_verifier.py

@@ -456,6 +456,64 @@ def _patched_forward(
 # ---------------------------------------------------------------------------
 
 
+@torch.no_grad()
+def capture_verifier_own_kv(
+    verifier_model: Any, input_ids: torch.Tensor,
+) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
+    """Capture the verifier's OWN pre-norm per-layer K/V via k_proj /
+    v_proj forward hooks (identity-restoration diagnostic).
+
+    Returns ``(k_layers, v_layers)``: per-layer lists where element
+    ``i`` is ``[B, T, kv_heads_i, head_dim_i]`` (heterogeneous per
+    layer, matching the f_θ output layout). Layers whose ``v_proj`` is
+    ``None`` (Gemma 4 full-attention K==V) take V from the k_proj
+    output, mirroring the model's own behaviour.
+
+    Injecting these at evicted positions reproduces exactly what the
+    verifier would have computed under full attention — so cross-model
+    recall under identity restoration should match the oracle. This
+    isolates "is the K/V Restoration machinery correct?" (this helper)
+    from "is f_θ accurate enough?" (the trained projection).
+    """
+    layers = get_verifier_decoder(verifier_model).layers
+    n = len(layers)
+    k_cap: List[Optional[torch.Tensor]] = [None] * n
+    v_cap: List[Optional[torch.Tensor]] = [None] * n
+    v_shared: List[int] = []
+    handles = []
+    for i, layer in enumerate(layers):
+        attn = layer.self_attn
+
+        def _kh(_m, _inp, out, idx=i):
+            k_cap[idx] = out.detach()
+
+        def _vh(_m, _inp, out, idx=i):
+            v_cap[idx] = out.detach()
+
+        handles.append(attn.k_proj.register_forward_hook(_kh))
+        if getattr(attn, "v_proj", None) is not None:
+            handles.append(attn.v_proj.register_forward_hook(_vh))
+        else:
+            v_shared.append(i)
+    try:
+        verifier_model(input_ids=input_ids, use_cache=False)
+    finally:
+        for h in handles:
+            h.remove()
+    for i in v_shared:
+        v_cap[i] = k_cap[i]
+    if any(k is None for k in k_cap) or any(v is None for v in v_cap):
+        raise RuntimeError("verifier own-K/V capture missing some layers")
+    k_layers: List[torch.Tensor] = []
+    v_layers: List[torch.Tensor] = []
+    for i, layer in enumerate(layers):
+        hd = layer.self_attn.head_dim
+        b, t, kvdim = k_cap[i].shape
+        k_layers.append(k_cap[i].view(b, t, kvdim // hd, hd))
+        v_layers.append(v_cap[i].view(b, t, kvdim // hd, hd))
+    return k_layers, v_layers
+
+
 def _capture_drafter_kv(
     *, verifier_model: Any, drafter: Any, input_ids: torch.Tensor,
 ) -> KVCapture:

scripts/research/k3_integrated_niah_eval.py

@@ -107,6 +107,14 @@ def parse_args() -> argparse.Namespace:
         help="Skip the full-attention oracle baseline (saves time but "
              "loses the |delta vs oracle| gate signal).",
     )
+    ap.add_argument(
+        "--identity-restore", action="store_true",
+        help="Diagnostic: restore evicted positions with the verifier's "
+             "OWN true pre-norm K/V instead of the f_θ projection. Under "
+             "this mode cross-model recall should match the oracle — it "
+             "isolates 'is the restoration machinery correct?' from 'is "
+             "f_θ accurate enough?'.",
+    )
     return ap.parse_args()
 
 
@@ -167,6 +175,20 @@ def main() -> int:
           f"(sink={args.sink_size}, window={args.window_size})",
           file=sys.stderr)
 
+    if args.identity_restore:
+        # Diagnostic: restore evicted positions with the verifier's own
+        # true pre-norm K/V (not f_θ). Validates the restoration
+        # machinery independent of f_θ accuracy.
+        from inference_engine.v04.cross_model_dlm_verifier import (
+            capture_verifier_own_kv,
+        )
+        cross_verifier.project_drafter_kv = (
+            lambda ids: capture_verifier_own_kv(verifier, ids)
+        )
+        print("[k3-integrated] IDENTITY-RESTORE diagnostic enabled "
+              "(evicted K/V come from verifier's own k_proj/v_proj)",
+              file=sys.stderr)
+
     # ---------- NIAH dataset ----------
     samples: List[NIAHSample] = make_niah_dataset(
         n_samples=args.n_samples,
@@ -307,6 +329,7 @@ def _oracle_step(cur):
             "max_new_tokens": args.max_new_tokens,
             "seed": args.seed,
             "skip_oracle": bool(args.skip_oracle),
+            "identity_restore": bool(args.identity_restore),
             "prompt_token_lens": seq_lens,
         },
         "results": {