OnnxDebugSession: windowed-cache (cache_window=C) support via identity placement

The session debugs windowed artifacts under identity slot placement —
slot j = position j, nothing ever evicted, hard cap at C committed
rows. Identity placement satisfies the windowed exporter's host
contract (committed slots filled in slot order, never wrapping), so
the exporter's own equivalence guarantee makes debug output
token-identical to the unbounded export's; the new tests pin that
parity rather than assume it.

- _feeds: windowed bindings are exactly C + n_new wide (the protocol
  requirement — the staging scatter indices are baked at [C, C+n_new)),
  committed rows at identity slots [0, base), zeros elsewhere (the
  never-written band is masked to weight exactly 0.0; staging is
  overwritten in-graph). Overrun raises naming the identity-placement
  cap and the way out (slices, or an unbounded export).
- capture_attention: the one real remap. The wire key axis is
  C + n_new; the session splices committed [0, base) + staging
  [C, C+n_new) back into the unbounded base + n_new layout, so
  key index == absolute position on every backend and protocol and
  probe_attention/AttentionProbe need zero changes. The dropped
  never-written band is asserted to carry weight exactly 0.0 first —
  a free runtime check of the windowed writtenness mask (an emission
  bug there is exactly the class this backend exists to catch).
- step / residual snapshots / self-consistency / asserts / debug_value
  / probe_compiled / probe_residual / probe_layer_diff: unchanged —
  snapshots are (n_new, d) row tensors; slot layout never touches them.
- Host eviction policies stay out of scope by design: probing a
  post-eviction state needs the host's own slot contents (a possible
  later affordance), not a policy reimplementation in torchwright.

Tests: test_windowed_cache_models_are_rejected (the NotImplementedError
pin) is replaced by three windowed tests — probe-clean + one-shot and
prefill/decode-seam parity against the unbounded export, attention
probes translated at base=0 and base>0, and the loud cap error.
Docs: module docstring + CLAUDE.md caveat updated.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: physicsrob

CLAUDE.md

@@ -425,7 +425,14 @@ Requirements and caveats:
   put it on a hot path.
 - Fetching all snapshots costs `n_pos × d × 2·n_layers` floats per
   run — probe very long prefills in slices.
-- Windowed-cache models (`cache_window=C`) are not supported yet.
+- Windowed-cache models (`cache_window=C`) are debugged under
+  identity slot placement (slot j = position j, nothing evicted),
+  capped at `C` committed rows — that satisfies the exporter's host
+  contract, so debug output is token-identical to the unbounded
+  export's.  Probe attention key indices are absolute positions on
+  both protocols (the session splices out the windowed wire layout).
+  Debug longer runs in slices, or export an unbounded variant; the
+  session never reproduces a host eviction policy.
 
 ## probe_compiled — full oracle comparison
 

tests/debug/test_onnx_debug_session.py

@@ -317,22 +317,158 @@ def test_headless_kind_probe_and_residual(tmp_path):
     assert rp.layers, "output node never materialised in any snapshot"
 
 
-def test_windowed_cache_models_are_rejected(tmp_path):
-    onnx_path = str(tmp_path / "windowed.onnx")
-    out1, pos1 = _build_headless_graph()
+# ---------------------------------------------------------------------------
+# Windowed-cache (cache_window=C) support — identity slot placement
+# ---------------------------------------------------------------------------
+
+_C = 8  # cache_window for every windowed test below
+
+
+@pytest.fixture(scope="module")
+def windowed_pair(tmp_path_factory):
+    """The same graph exported both ways: (unbounded_path, windowed_path)."""
+    tmpdir = tmp_path_factory.mktemp("onnx_debug_windowed")
+    out, pos = _build_headless_graph()
+    unbounded = str(tmpdir / "unbounded.onnx")
+    windowed = str(tmpdir / "windowed.onnx")
     compile_headless_to_onnx(
-        out1,
-        pos1,
-        onnx_path,
+        out, pos, unbounded, d=256, d_head=D_HEAD, max_seq_len=32, verbose=False
+    )
+    compile_headless_to_onnx(
+        out,
+        pos,
+        windowed,
         d=256,
         d_head=D_HEAD,
         max_seq_len=32,
-        cache_window=8,
+        cache_window=_C,
         verbose=False,
     )
-    out2, pos2 = _build_headless_graph()
-    with pytest.raises(NotImplementedError, match="windowed"):
-        OnnxDebugSession(onnx_path, out2, pos2)
+    return unbounded, windowed
+
+
+def test_windowed_session_matches_unbounded(windowed_pair):
+    """Identity placement satisfies the exporter's host contract, so the
+    windowed equivalence guarantee makes parity with the unbounded export
+    a theorem — this pins it: probe clean, one-shot outputs equal, and a
+    debug=True step across the prefill/decode seam (base > 0) equal too."""
+    unbounded_path, windowed_path = windowed_pair
+    iv = {
+        "a": torch.tensor([[1.0], [2.0], [3.0], [-4.0]]),
+        "b": torch.tensor([[5.0], [6.0], [-7.0], [8.0]]),
+    }
+
+    out_w, pos_w = _build_headless_graph()
+    sess_w = OnnxDebugSession(windowed_path, out_w, pos_w)
+    report = probe_compiled(sess_w, out_w, iv, n_pos=4, atol=1e-2)
+    assert report.first_divergent is None, report.format_short()
+
+    out_u, pos_u = _build_headless_graph()
+    sess_u = OnnxDebugSession(unbounded_path, out_u, pos_u)
+
+    prefill = sess_w.build_prefill(iv, 4)
+    one_shot_w = sess_w(prefill, debug=True)
+    one_shot_u = sess_u(prefill, debug=True)
+    assert torch.allclose(one_shot_w, one_shot_u, atol=1e-4), (
+        f"windowed one-shot diverged from unbounded: max diff "
+        f"{float((one_shot_w - one_shot_u).abs().max()):.3e}"
+    )
+
+    # Seam: 2 committed rows, then a 2-row debug step at base=2 — the
+    # binding now has a never-written committed band [2, C).
+    past = sess_w.empty_past()
+    head, past = sess_w.step(prefill[:2], past)
+    tail, _ = sess_w.step(prefill[2:], past, debug=True)
+    seam = torch.cat([head, tail])
+    assert torch.allclose(seam, one_shot_w, atol=1e-4), (
+        f"windowed seam diverged from one-shot: max diff "
+        f"{float((seam - one_shot_w).abs().max()):.3e}"
+    )
+
+
+def test_windowed_attention_probe_translated(tmp_path):
+    """The one real windowed remap: the wire key axis is C + n_new, the
+    probe's key axis must stay base + n_new with key index == absolute
+    position.  Pins prefill (base=0) and decode (base>0, never-written
+    band inside the binding) against the unbounded export's probe."""
+
+    def build():
+        torch.manual_seed(0)
+        x = create_input("x", 4)
+        attn = Attn(
+            query_in=x,
+            key_in=x,
+            value_in=x,
+            query_matrix=torch.randn(4, 4),
+            key_matrix=torch.randn(4, 4),
+            value_matrix=torch.randn(4, 4),
+            output_matrix=torch.randn(4, 4),
+        )
+        return attn, create_pos_encoding()
+
+    unbounded = str(tmp_path / "unbounded.onnx")
+    windowed = str(tmp_path / "windowed.onnx")
+    out, pos = build()
+    compile_headless_to_onnx(
+        out, pos, unbounded, d=256, d_head=D_HEAD, max_seq_len=32, verbose=False
+    )
+    compile_headless_to_onnx(
+        out,
+        pos,
+        windowed,
+        d=256,
+        d_head=D_HEAD,
+        max_seq_len=32,
+        cache_window=_C,
+        verbose=False,
+    )
+    attn_u, pos_u = build()
+    sess_u = OnnxDebugSession(unbounded, attn_u, pos_u)
+    attn_w, pos_w = build()
+    sess_w = OnnxDebugSession(windowed, attn_w, pos_w)
+
+    torch.manual_seed(1)
+    inp = torch.randn(3, 4)
+
+    # Prefill probe (base=0: the dropped band is the whole committed window).
+    ap_u = probe_attention(sess_u, inp, attn_u, query_pos=2)
+    ap_w = probe_attention(sess_w, inp, attn_w, query_pos=2)
+    assert ap_w.weights.shape == ap_u.weights.shape  # n_keys = 3, both
+    assert torch.allclose(ap_w.weights, ap_u.weights, atol=1e-5)
+    assert torch.allclose(ap_w.logits, ap_u.logits, atol=1e-4)
+
+    # Decode probe (base=2: never-written band [2, C) inside the binding).
+    past_u = sess_u.empty_past()
+    _, past_u = sess_u.step(inp[:2], past_u)
+    past_w = sess_w.empty_past()
+    _, past_w = sess_w.step(inp[:2], past_w)
+    ap_u2 = probe_attention(
+        sess_u, inp[2:], attn_u, query_pos=0, past_len=2, past_kvs=past_u
+    )
+    ap_w2 = probe_attention(
+        sess_w, inp[2:], attn_w, query_pos=0, past_len=2, past_kvs=past_w
+    )
+    assert ap_w2.weights.shape == ap_u2.weights.shape  # n_keys = 2 + 1, both
+    assert torch.allclose(ap_w2.weights, ap_u2.weights, atol=1e-5)
+
+
+def test_windowed_cap_error(windowed_pair):
+    """Identity placement never evicts: committed rows are capped at C,
+    and the overrun error names the cap and the way out."""
+    _unbounded_path, windowed_path = windowed_pair
+    out, pos = _build_headless_graph()
+    sess = OnnxDebugSession(windowed_path, out, pos)
+
+    iv = {"a": torch.ones(_C, 1) * 0.5, "b": torch.ones(_C, 1) * 0.25}
+    prefill = sess.build_prefill(iv, _C)  # exactly C rows: allowed
+    past = sess.empty_past()
+    _, past = sess.step(prefill, past)
+
+    one_more = sess.build_prefill(
+        {"a": torch.tensor([[0.5]]), "b": torch.tensor([[0.25]])}, 1
+    )
+    with pytest.raises(RuntimeError, match="cache_window"):
+        sess.step(one_more, past)
 
 
 # ---------------------------------------------------------------------------

torchwright/debug/onnx_debug.py

@@ -40,7 +40,15 @@
 * Fetching every residual snapshot costs
   ``n_pos × d × 2·n_layers`` floats on host per run — fine for
   debug-sized runs; probe very long prefills in slices.
-* Windowed-cache models (``cache_window=C``) are not supported yet.
+* Windowed-cache models (``cache_window=C``) are debugged under
+  **identity slot placement**: slot ``j`` = position ``j``, nothing is
+  ever evicted, and runs are capped at ``C`` committed rows (debug
+  longer runs in slices, or export an unbounded variant).  Identity
+  placement satisfies the exporter's host contract, so the windowed
+  equivalence guarantee applies and debug output is token-identical to
+  the unbounded export's.  The session never reproduces a host eviction
+  policy — probing a post-eviction state needs the host's own slot
+  contents and is out of scope here.
 """
 
 from __future__ import annotations
@@ -138,12 +146,18 @@ def __init__(
         import onnxruntime as ort
 
         sidecar = _sidecar_or_raise(onnx_path)
-        if sidecar.get("cache_window") is not None:
-            raise NotImplementedError(
-                "OnnxDebugSession does not support windowed-cache models yet "
-                "(cache_window is set in the sidecar); export an unbounded-"
-                "protocol model for debugging"
-            )
+        # Windowed-cache artifacts are debugged under IDENTITY slot
+        # placement: slot j = position j, nothing ever evicted.  That
+        # satisfies the exporter's host contract (committed slots filled
+        # in slot order, never wrapping), so by the exporter's own
+        # equivalence guarantee the debug output is token-identical to
+        # the unbounded export's — at the cost of a hard cap at C
+        # committed rows (enforced in _feeds).
+        self._cache_window: Optional[int] = (
+            int(sidecar["cache_window"])
+            if sidecar.get("cache_window") is not None
+            else None
+        )
 
         self._onnx_path = onnx_path
         self._kind: str = sidecar["kind"]  # "token" | "headless"
@@ -332,9 +346,25 @@ def _feeds(self, prefill: torch.Tensor, base: int, past: Optional[tuple]) -> dic
         (valid under stride bucketing) so debug runs never materialize
         the full stride buffer; old static-dim exports get full-width
         zero-padded feeds instead.
+
+        Windowed-cache artifacts bind exactly ``C + n_new`` wide (the
+        protocol requirement): committed rows at identity slots
+        ``[0, base)``, never-written slots ``[base, C)`` zeroed (the
+        writtenness mask gives them weight exactly 0.0), staging tail
+        ``[C, C+n_new)`` zeroed (the in-graph scatter overwrites it).
         """
         n_new = int(prefill.shape[0])
-        if base + n_new > self._cache_stride:
+        if self._cache_window is not None:
+            if base + n_new > self._cache_window:
+                raise RuntimeError(
+                    f"windowed-cache debug overrun: base {base} + n_new "
+                    f"{n_new} exceeds cache_window {self._cache_window}.  "
+                    f"This debug session uses identity slot placement and "
+                    f"never evicts, so a windowed artifact is debuggable up "
+                    f"to cache_window committed positions — debug longer "
+                    f"runs in slices, or export an unbounded variant"
+                )
+        elif base + n_new > self._cache_stride:
             raise RuntimeError(
                 f"static cache overrun: base {base} + n_new {n_new} exceeds "
                 f"cache_stride {self._cache_stride}"
@@ -349,11 +379,12 @@ def _feeds(self, prefill: torch.Tensor, base: int, past: Optional[tuple]) -> dic
             feeds["inputs"] = (
                 prefill.detach().cpu().numpy().astype(np.float32, copy=False)
             )
-        bind_width = (
-            self._static_slot_dim
-            if self._static_slot_dim is not None
-            else base + n_new
-        )
+        if self._cache_window is not None:
+            bind_width = self._cache_window + n_new
+        elif self._static_slot_dim is not None:
+            bind_width = self._static_slot_dim
+        else:
+            bind_width = base + n_new
         for i, nh in enumerate(self._per_layer_n_heads):
             buf = np.zeros((bind_width, nh, self._d_head), dtype=np.float32)
             if past is not None and base > 0:
@@ -573,6 +604,16 @@ def capture_attention(
         ``(n_heads, n_queries, n_keys)`` with ``n_keys = past + new`` —
         the same shape the in-process ``attention_capture`` hook
         produces.
+
+        On windowed-cache artifacts the wire key axis is ``C + n_new``
+        (committed window + staging tail); under this session's identity
+        placement the committed keys ``[0, base)`` already sit at their
+        positions and the new rows live at ``[C, C+n_new)``, so the two
+        bands are spliced back into the unbounded ``base + n_new`` layout
+        — key index == absolute position on every backend and protocol.
+        The dropped never-written band ``[base, C)`` is asserted to carry
+        weight exactly 0.0 first; anything else is a windowed-mask
+        emission bug.
         """
         if prefill.ndim == 1:
             prefill = prefill.reshape(-1, 1)
@@ -581,4 +622,22 @@ def capture_attention(
         weights_np, logits_np = self._session.run(
             [f"l{layer_index}_weights", f"l{layer_index}_logits_masked"], feeds
         )
+        if self._cache_window is not None:
+            C = self._cache_window
+            dead = weights_np[:, :, base:C]
+            if dead.size and float(np.abs(dead).max()) != 0.0:
+                raise RuntimeError(
+                    f"windowed mask leak at layer {layer_index}: never-written "
+                    f"committed slots [{base}, {C}) received attention weight "
+                    f"(max {float(np.abs(dead).max()):.3e}, expected exactly "
+                    f"0.0) — the writtenness mask in the windowed preamble is "
+                    f"broken (ONNX emission bug in "
+                    f"torchwright/compiler/export.py — D1: stop and report)"
+                )
+            weights_np = np.concatenate(
+                [weights_np[:, :, :base], weights_np[:, :, C:]], axis=2
+            )
+            logits_np = np.concatenate(
+                [logits_np[:, :, :base], logits_np[:, :, C:]], axis=2
+            )
         return torch.from_numpy(weights_np), torch.from_numpy(logits_np)