PR-K1.E: NIAH validation harness for v0.4 GA gate (a)

Final PR of the K1 series. Implements the empirical NIAH validation
harness that makes ADR 0008 §11.8 gate (a) — '>= 95% mid-context
recall at 100k context' — falsifiable on Mac M4 hardware.

Files:

  inference_engine/v04/niah_eval.py  (369 lines)
    Pure-PyTorch harness for needle-in-haystack evaluation. Three
    verifier configurations evaluated on the same NIAH samples:

      (a) Full-attention oracle — model.forward, upper bound,
          target ~1.0
      (b) v0.3 sink+window — Gemma3 forward with 4D sink+window
          mask, confirms the regression v0.4 fixes; target ~0.17
          per the 2026-06-06 A/B benchmark
      (c) v0.4 DLMRestoredVerifier — the new architecture; gate
          target >= 0.95

    Components:
      * NIAHSample / make_niah_dataset — synthetic haystack with
        random needle codes (15 prefixes x 4-digit numbers, same
        vocab as cross_attn_toy_prototype). Needle inserted outside
        the first/last 4 padding lines so neither sink (4 lines)
        nor a small trailing window can plausibly catch it from
        positional luck.
      * recall_predicate — substring match (the answer code appears
        anywhere in the decoded continuation), case-sensitive.
      * NIAHEvalResult / aggregate_recall — recall %, mean/median
        latency, per-sample decoded text + correctness flag.
      * make_sink_window_4d_mask — 4D additive mask helper for the
        v0.3 baseline. Uses dtype's finfo.min (not -inf) for masked
        positions because bf16 attention kernels NaN-propagate from
        -inf in some implementations (matches the convention in
        cross_attn_toy_prototype.py).
      * greedy_decode_oracle / greedy_decode_sink_window /
        greedy_decode_v04 — three decoders, one per configuration.
        v0.3 wraps the 4D mask in the Gemma3 {full_attention,
        sliding_attention} dict for proper layer-type dispatch.
      * evaluate — the high-level eval-loop scaffold, takes a
        decode_fn(sample) -> (text, latency_s) callable.

  inference_engine/v04/__init__.py
    Re-exports the K1.E public API alongside K1.A-D primitives.

  tests/inference_engine/v04/test_niah_eval.py  (393 lines, 39 cases,
  all <0.05 s on Linux CI)
    Pure-PyTorch harness logic tests — no HF dependency. Covers:

      * TestMakeNIAHDataset (12 cases) — sample format, needle
        position invariants (>= 4 from start), reproducibility
        under fixed seed, validation raises (n_samples<=0,
        min>max, lines<10, empty prefixes, code_min>max), custom
        prefix sets.

      * TestRecallPredicate (6 cases) — exact match, no match,
        partial fail, case sensitivity, substring anywhere, empty.

      * TestAggregateRecall (8 cases) — all-correct, none-correct,
        partial, mean/median latency (odd and even count),
        decoded/latency length mismatch raises, empty samples
        raises, per-sample decoded text preserved.

      * TestSinkWindow4DMask (7 cases) — shape; allowed positions
        get 0.0 for every (q, k); masked positions get finfo.min;
        bf16 uses finite (not -inf); softmax(scores + mask) zeros
        forbidden keys and allowed weights sum to 1; negative
        seq_len/sink raises.

      * TestEvaluateOrchestration (3 cases) — eval loop calls
        decode_fn per sample; partial correctness propagates;
        latency stats computed.

  scripts/research/k1e_niah_validation.py  (260 lines)
    Runnable script that loads google/gemma-3-1b-it, generates a
    NIAH dataset, and evaluates the three configurations. CLI
    flags for n-samples, haystack size, sink/window, max-new-
    tokens, seed, plus skip flags for individual configs.

    Outputs schema_version=1 JSON with per-config results
    (recall, latency, per-sample decoded text + correctness) plus
    a 'gate' block that records:
      * v04_vs_oracle_delta (target ~0)
      * v04_recall_ge_0_95 (the gate predicate)
      * v04_within_5pct_of_oracle
      * v04_vs_v03_improvement
      * v04_dominates_v03

  scripts/review_pr_k1e_on_mac.sh  (96 lines)
    Mac M4 reviewer aid. Default 20 samples × 60-80 padding lines
    (~1-2k tokens per prompt). Time budget: ~30-60 min total.
    Env knobs for context length scaling (HAYSTACK_MIN/MAX),
    sink/window override (SINK/WINDOW), and config skipping
    (SKIP_V03, SKIP_V04).

    Documents three acceptance signals to look for in the report:
      * v0.4 recall >> v0.3 recall (target >= +50pp)
      * v0.4 recall close to oracle (target within 5pp)
      * v0.3 recall ≈ 0.17 (regression confirmed; matches the
        2026-06-06 A/B benchmark for sanity)

Combined v04 test status after this PR:
  tests/inference_engine/v04/ has 163 cases (32 K1.A + 39 K1.B +
  32 K1.C + 21 K1.D + 39 K1.E), all <0.20 s on Linux CI, no HF
  model download.

Empirical run plan:
  After this PR + #71/#72/#73 are merged, run on Mac M4:
    bash scripts/review_pr_k1e_on_mac.sh
  ~30-60 min for the default 1-2k-token run. Report goes to
  results/research/k1e_niah_<stamp>.json. The K1 phase is then
  empirically validated and ADR 0008 §11.8 gate (a) becomes a
  measured fact at the run's context length.

  Larger context lengths (16k, 64k, 100k) require either explicit
  --haystack-* flags + sufficient memory or a beefier device. A
  100k-token oracle forward on Gemma 3-1B-it bf16 needs ~10 GB
  for the KV cache; v0.4 DLMRestoredVerifier sustained memory is
  constant in context so it scales further than the oracle.

What's after K1:
  K2 — cross-model proposer + verifier (e.g., Gemma 3-1B proposer +
       Gemma 3-4B verifier) with a learned per-layer linear
       projection f_θ trained to align proposer K/V representation
       to verifier K/V space (ADR 0008 §11.7 phase K2).
  K3 — production scale (Gemma 4-2B-MDLM proposer + Gemma 4-9B
       verifier) on long-context corpus.
  K4 — KakeyaLattice composition for byte-level KV compression on
       top of the Restoration architecture.
  K5 — default flip + docs.

Stacking notes:
  Logical base is #73 K1.D. After #71 -> #72 -> #73 are merged in
  order, the cumulative diff on main is the complete K1 phase
  implementation. base_branch is set to main for tooling reasons.
  Merge order: #71 -> #72 -> #73 -> this PR.

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

Author: cursoragent

inference_engine/v04/__init__.py

@@ -43,6 +43,19 @@
     slice_position_embeddings,
 )
 from inference_engine.v04.dlm_restored_verifier import DLMRestoredVerifier
+from inference_engine.v04.niah_eval import (
+    DEFAULT_NEEDLE_PREFIXES,
+    NIAHEvalResult,
+    NIAHSample,
+    aggregate_recall,
+    evaluate,
+    greedy_decode_oracle,
+    greedy_decode_sink_window,
+    greedy_decode_v04,
+    make_niah_dataset,
+    make_sink_window_4d_mask,
+    recall_predicate,
+)
 
 __all__ = [
     # K1.A — capture
@@ -58,4 +71,16 @@
     "slice_position_embeddings",
     # K1.D — end-to-end wrapper
     "DLMRestoredVerifier",
+    # K1.E — NIAH validation harness
+    "DEFAULT_NEEDLE_PREFIXES",
+    "NIAHEvalResult",
+    "NIAHSample",
+    "aggregate_recall",
+    "evaluate",
+    "greedy_decode_oracle",
+    "greedy_decode_sink_window",
+    "greedy_decode_v04",
+    "make_niah_dataset",
+    "make_sink_window_4d_mask",
+    "recall_predicate",
 ]