perf(restored): ADR 0015 item #1 - memory-efficient restoration prefill

Unblock long-context multi-tenant for the restored-S5 path:
- forward() gains logits_to_keep (default 0 = unchanged; spec-decode still
  gets all-position logits). Bench passes 1 (only last logits needed).
- capture_verifier_own_kv skips the [B,T,vocab] logits entirely (logits_to_keep=1;
  it only needs K/V via hooks).
- restored K/V cast to verifier compute dtype right after f_theta (was held
  full-length in fp32; cast-at-use was already bf16, so numerically identical).
- bench --attn-impl (default eager for exact repro; sdpa for long context routes
  the patched forward to SDPA, killing the O(T^2) eager score materialisation).

41 v04 unit tests still pass (logits_to_keep default preserves behavior).

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

Author: cursoragent

inference_engine/v04/cross_model_dlm_verifier.py

@@ -310,9 +310,16 @@ def forward(
         eager_attention_forward: Callable,
         all_attention_functions: Optional[Any] = None,
         capture_kv: Optional[list] = None,
+        logits_to_keep: int = 0,
     ):
         """Run a verifier forward with f_θ-mediated K/V Restoration.
 
+        ``logits_to_keep`` is forwarded to the verifier model: 0 (default)
+        computes logits for all positions (needed by spec-decode block
+        verification); set 1 to compute only the last position's logits
+        (decode/prefill that only needs the next token) — this avoids the
+        ``[B, T, vocab]`` materialisation that dominates long-context memory.
+
         Steps:
           1. Compute evicted positions from sink+window per ADR §11.7.
           2. Drafter forward + f_θ projection → verifier K/V at every
@@ -335,11 +342,23 @@ def forward(
         # needed — run the verifier directly. This is the trivial case
         # for short prompts, e.g. T=8 with sink=4 + window=64.
         if not evicted_positions:
-            return self.verifier_model(input_ids=input_ids, use_cache=False)
+            return self.verifier_model(
+                input_ids=input_ids, use_cache=False,
+                logits_to_keep=logits_to_keep,
+            )
 
         # 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)
+        # Hold the restored K/V at the verifier's compute dtype (they are
+        # cast to it at injection anyway): keeps the full-length [B,T,kv,D]
+        # tensors from doubling long-context memory when f_θ runs in fp32.
+        try:
+            vdtype = next(self.verifier_model.parameters()).dtype
+            verifier_k_layers = [k.to(vdtype) for k in verifier_k_layers]
+            verifier_v_layers = [v.to(vdtype) for v in verifier_v_layers]
+        except StopIteration:
+            pass
 
         # S5: for exact layers, replace the f_θ-restored K/V with the
         # verifier's OWN true K/V (the recall-critical full-attention
@@ -375,7 +394,10 @@ def forward(
                     all_attention_functions=all_attention_functions,
                     capture_kv=capture_kv,
                 )
-            return self.verifier_model(input_ids=input_ids, use_cache=False)
+            return self.verifier_model(
+                input_ids=input_ids, use_cache=False,
+                logits_to_keep=logits_to_keep,
+            )
         finally:
             for layer_idx, layer in enumerate(layers):
                 layer.self_attn.forward = originals[layer_idx]
@@ -548,7 +570,9 @@ def _vh(_m, _inp, out, idx=i):
         else:
             v_shared.append(i)
     try:
-        verifier_model(input_ids=input_ids, use_cache=False)
+        # Only K/V are needed (captured via k_proj/v_proj hooks); skip the
+        # [B, T, vocab] logits materialisation entirely.
+        verifier_model(input_ids=input_ids, use_cache=False, logits_to_keep=1)
     finally:
         for h in handles:
             h.remove()

scripts/research/k3_cuda_multitenant_parallel_bench.py

@@ -36,7 +36,7 @@
 def _ar_batched(model, ids_bt, gen_tokens, device, eos_ids):
     """Batched AR decode. ids_bt: [N, T]. Returns (per_row_tokens, decode_s)."""
     N = ids_bt.size(0)
-    out = model(input_ids=ids_bt, use_cache=True)
+    out = model(input_ids=ids_bt, use_cache=True, logits_to_keep=1)
     cache = out.past_key_values
     nxt = out.logits[:, -1, :].argmax(-1)              # [N]
     gen = [[int(nxt[i].item())] for i in range(N)]
@@ -62,7 +62,7 @@ def _restored_prefill_batched(restored, ids_bt, helpers):
     from transformers.cache_utils import DynamicCache
     n_layers = len(_decoder_layers(restored.verifier_model))
     capture: list = [None] * n_layers
-    out = restored.forward(ids_bt, capture_kv=capture, **helpers)
+    out = restored.forward(ids_bt, capture_kv=capture, logits_to_keep=1, **helpers)
     logits = out.logits if hasattr(out, "logits") else out
     if any(c is None for c in capture):
         raise RuntimeError("restored prefill did not capture all layers "
@@ -111,6 +111,10 @@ def main() -> int:
     ap.add_argument("--sink", type=int, default=4)
     ap.add_argument("--window", type=int, default=64)
     ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--attn-impl", default="eager",
+                    help="verifier attn_implementation: 'eager' (default, exact "
+                         "repro) or 'sdpa'/'flash_attention_2' (memory-efficient "
+                         "prefill — required for long context, ADR 0015 item #1).")
     ap.add_argument("--output", default=None)
     args = ap.parse_args()
 
@@ -134,7 +138,7 @@ def main() -> int:
     print(f"[mt] loading verifier {args.verifier_id}", file=sys.stderr, flush=True)
     tok = AutoTokenizer.from_pretrained(args.verifier_id)
     verifier = AutoModelForCausalLM.from_pretrained(
-        args.verifier_id, dtype=dtype, attn_implementation="eager",
+        args.verifier_id, dtype=dtype, attn_implementation=args.attn_impl,
     ).to(device).eval()
     for p in verifier.parameters():
         p.requires_grad_(False)
@@ -233,6 +237,7 @@ def recall(tokens, ans):
             "verifier_id": args.verifier_id, "drafter_id": args.drafter_id,
             "haystack_lines": args.haystack_lines, "modal_prompt_len": modal_len,
             "gen_tokens": args.gen_tokens, "sink": args.sink, "window": args.window,
+            "attn_impl": args.attn_impl,
             "batch_sizes": batch_sizes, "exact_layers": exact_layers,
             "note": ("per-session binding via batched decode (each row = a "
                      "session with its own KV-cache row); recall-preserving S5 "