K3 Block B + C: f_theta projection + cross-model DLMRestoredVerifier (P0)

Per user 'go P0' directive 2026-06-09 after architectural observation
that PR #102's Mac MLX spec decode eval doesn't exercise the Kakeya
inference engine's core architecture (sink+window verifier + dLM
proposer K/V Restoration).

This PR ships the foundational engine code for the integrated
Kakeya inference architecture per ADR 0008 §11.3:

  verifier (Gemma 4 26B-A4B):
    └─ holds only sink+window local KV cache (sink=4 + window=64)
    └─ at evicted positions, takes K/V supplied by proposer (via f_θ)

  drafter (DFlash 0.4B, alignment-trained baseline):
    └─ runs full forward over committed prefix per step
    └─ K/V at every layer at every position captured
    └─ K/V projected through f_θ into verifier K/V space, injected at
       evicted positions

Three new files
---------------

inference_engine/v04/f_theta.py (~290 LOC)

  FThetaConfig dataclass + FThetaProjection nn.Module.

  Architecture: shared encoder + per-verifier-layer decoders, low-rank
  factorisation:

    drafter_kv_input [B, T, drafter_layers * drafter_kv_dim]
              ↓ encoder Linear(in, rank)
    rep [B, T, rank]
              ↓ per-verifier-layer decoders (30 × Linear(rank, verifier_kv_dim))
    output [B, T, num_verifier_layers, num_kv_heads_v, head_dim_v]

  Default rank=256. Production K3 config (Gemma 4 26B-A4B + DFlash 0.4B):
    encoder:   2 × 5×256 × 256 = 655k params
    decoders:  2 × 30 × 256 × 2048 = 31.5M params
    Total:     ~32M params (vs drafter 430M, verifier 26B)

  Separate K and V projections (different downstream roles).

  Save/load: save_pretrained(dir) writes f_theta_config.json +
  f_theta_weights.pt; from_pretrained(dir, dtype, device) loads back.

inference_engine/v04/cross_model_dlm_verifier.py (~270 LOC)

  CrossModelDLMRestoredVerifier wrapper. Construction validates
  drafter + verifier dimensions match the f_θ config (rejects
  drafter-vs-verifier-vs-f_θ mismatch loudly at __init__).

  forward(input_ids, apply_rotary_pos_emb, eager_attention_forward):
    1. compute_evicted_positions(T, sink, window)
    2. If no evicted (T <= sink+window): plain verifier forward
    3. Drafter forward via _capture_drafter_kv (forward hooks on
       k_proj/v_proj at each drafter layer)
    4. f_θ.forward_kv_pack(drafter_K_per_layer, drafter_V_per_layer)
       → verifier K, V at every (layer, position)
    5. Patch each verifier layer's self_attn.forward to:
       a. Run standard q/k/v_proj + q_norm/k_norm + RoPE
       b. At evicted positions, REPLACE k, v with f_θ output (after
          k_norm + RoPE applied via prepare_restored_attention_kv)
       c. Standard attention compute path through eager_attention_forward
    6. Run verifier forward → logits
    7. Restore original attention forwards (try/finally)

  Two scope-outs (recorded inline):
    * MLX verifier path: this module patches HF transformers
      attention. Mac MLX integration is a follow-up PR (instrument
      mlx_lm Gemma 4 model directly, not via attention monkey-patch).
    * Speculative decoding accept/reject loop: separate inference
      engine concern. PR #93's DFlashProposer + mlx_verify_block
      handles the spec-decode side; combining with this module's
      K/V Restoration is a separate integration step.

  Drafter K/V capture (_capture_drafter_kv): instruments DFlashDrafter's
  internal layer.self_attn.k_proj / v_proj via forward hooks. NOTE
  inline that the first-iteration synthetic-context capture (zero
  hidden as drafter input) is plumbing-validation; product-meaningful
  K/V values require conditioning on verifier aux hiddens, which is
  the next integration step (after f_θ training validates the
  projection alone).

scripts/research/k3_f_theta_train.py (~310 LOC)

  Training pipeline for f_θ on CUDA:

    1. Load Gemma 4 26B-A4B verifier (transformers bf16, sdpa)
    2. Load DFlash drafter (PR #93's DFlashDrafter from
       models/dflash-kakeya-baseline)
    3. Data collection: for each prompt in PROMPTS (same 64-prompt
       corpus as PR #93's alignment_train), run greedy AR generation
       to gen_len tokens, capture per-layer per-position K/V via
       hooks on k_proj/v_proj of both models
    4. Train f_θ with MSE loss across (layer, position) pairs,
       AdamW lr=1e-3, weight_decay=0.01, gradient clip 1.0
    5. Save checkpoint at --save (default results/research/f_theta_v1)

  Memory budget: at T=512, ~128 MB per sequence cached on GPU. 64
  sequences ≈ 8 GB. Fits H200 80 GB easily.

  Validation: report initial vs final loss; reduction factor.

inference_engine/v04/__init__.py: re-exports the new public surface
(FThetaConfig, FThetaProjection, CrossModelDLMRestoredVerifier,
CrossModelLayerMapping).

Tests (Linux CI: 27 new tests)
-----------------------------

tests/inference_engine/v04/test_f_theta.py (21 tests):
  TestFThetaConfig (4): dim properties + JSON round-trip
  TestForwardShapes (4): forward_k/v shape contract + input validation
  TestForwardKVPack (3): KVCapture-style input + consistency vs explicit concat
  TestParameterCount (2): tiny + production param count locked in
  TestSaveLoadRoundTrip (4): save+load preserves outputs; missing-file errors
  TestDeviceDtypeDispatch (2): to(dtype), from_pretrained dtype override
  TestGradientFlow (1): gradients flow through encoder + decoders separately
                       (K path doesn't update V weights and vice versa)

tests/inference_engine/v04/test_cross_model_dlm_verifier.py (6 tests):
  TestConstruction (3): dimension validation rejects mismatch; valid
                       construction succeeds; negative sink/window raises
  TestProjectDrafterKV (1): output shape contract
  TestNoEvictPath (1): short prompt (T <= sink+window) doesn't invoke drafter
  TestExports (1): module + namespace re-exports

Tests: 354 passing (336 pre-existing + 21 f_theta + 6 cross-model;
       12 research/ unchanged from PR #102).

What this PR does NOT yet do (deferred to follow-up PRs)
--------------------------------------------------------

1. Train f_θ on real data — requires vast.ai GPU time.
   scripts/research/k3_f_theta_train.py is the runnable trainer.
   Once trained, the checkpoint goes to a follow-up PR with the
   evidence (training report + integrated NIAH ladder evidence).

2. End-to-end integrated NIAH ladder evidence — needs:
   * trained f_θ checkpoint (step 1)
   * cross-model DLMRestoredVerifier reviewer aid (off-the-shelf K1.E
     NIAH harness needs a small adapter to use this verifier wrapper)
   * vast.ai run producing the evidence JSON

3. Mac MLX integration — instruments mlx_lm Gemma 4 model directly
   (different surgical approach than HF transformers attention
   monkey-patch). Follow-up PR.

4. _capture_drafter_kv proper aux-conditioning — current synthetic
   zero-hidden capture is plumbing only. The proper path passes
   verifier aux hiddens into the drafter (DFlash architecture),
   captures K/V from THAT forward. Adds a method to DFlashDrafter
   in a follow-up.

These are the remaining items on the K3 critical path; this PR
establishes the engine API surface they all depend on.

Stack
-----

Off main (post #93 + #99 + #94 + #100 + #101 + #102 merged).
Independent of any other open PR.

Outstanding work after this PR:
  Step 5 — K2.A backport PR (P2)
  Step 6 — alignment training corpus expansion (P2)
  P0 cont. — f_θ training run + integrated NIAH evidence
  P0 cont. — Mac MLX integration of cross-model DLMRestoredVerifier

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

Author: cursoragent

inference_engine/v04/__init__.py

@@ -49,6 +49,11 @@
     DFlashDrafter,
     DFlashProposer,
 )
+from inference_engine.v04.f_theta import FThetaConfig, FThetaProjection
+from inference_engine.v04.cross_model_dlm_verifier import (
+    CrossModelDLMRestoredVerifier,
+    CrossModelLayerMapping,
+)
 from inference_engine.v04.kv_compressor import (
     IdentityCompressor,
     KakeyaLatticeCompressor,
@@ -122,4 +127,11 @@
     "DFlashConfig",
     "DFlashDrafter",
     "DFlashProposer",
+    # K3 Block C — f_θ K/V projection
+    "FThetaConfig",
+    "FThetaProjection",
+    # K3 Block B — cross-model DLMRestoredVerifier with f_θ-mediated
+    # K/V Restoration (the integrated Kakeya inference architecture)
+    "CrossModelDLMRestoredVerifier",
+    "CrossModelLayerMapping",
 ]