From f22f3f948883721b7e70ca2e2c41b3ddba665299 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Thu, 21 May 2026 11:39:29 +0200
Subject: [PATCH] fix(intervention-b): recapture train graph on boundary-state
 change
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Task 1's attn_pool gate (commit 7dd953aa) was inside the captured-graph
region, so the host branch's decision was frozen at step-2 capture time.
In Sequential mode this would lock the captured graph to whichever
boundary state happened to hold at step 2, breaking subsequent file-
boundary detection — violating pearl_no_host_branches_in_captured_graph.

Fix: track the boundary state at capture time as a new
`train_graph_boundary_state: Option<bool>` field. On each step_batched,
if the current `last_seen_file_boundary` diverges from the captured
state, drop the cached graph and re-capture on the next step.

- Random mode: gate is unconditionally true; state never diverges;
  zero recaptures.
- Sequential mode: ~8 recaptures per epoch (one per file boundary).
  Each recapture costs the same as the original step-2 capture (~50ms).
  Negligible.

Verified clean cargo check -p ml-alpha --all-targets.
---
 crates/ml-alpha/src/trainer/perception.rs | 28 +++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs
index 57b35c2c1..349be838d 100644
--- a/crates/ml-alpha/src/trainer/perception.rs
+++ b/crates/ml-alpha/src/trainer/perception.rs
@@ -548,6 +548,14 @@ pub struct PerceptionTrainer {
     /// call runs uncaptured (warmup); second call captures; third+
     /// replays. Eliminates ~155 individual kernel launches per step.
     train_graph: Option<CudaGraph>,
+    /// Snapshot of `last_seen_file_boundary` at the time `train_graph` was
+    /// captured. The captured graph records the attn_pool decision at
+    /// capture time (per `pearl_no_host_branches_in_captured_graph`); if
+    /// the current boundary state diverges from the captured one, the
+    /// graph is dropped and re-captured on the next step. In Random mode
+    /// this never changes (always `true`), so no re-capture. In Sequential
+    /// mode this triggers a re-capture per file boundary (~8 per epoch).
+    train_graph_boundary_state: Option<bool>,
     /// Captured CUDA Graph for the inference (`forward_only`) path. Same
     /// three-state machine as `train_graph`: first call eager (warmup),
     /// second call captures, third+ replays. Mirrors the forward chain
@@ -1339,6 +1347,7 @@ impl PerceptionTrainer {
             logit_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS)?,
             stg_labels: unsafe { MappedF32Buffer::new(k * cfg.n_batch * N_HORIZONS) }.map_err(|e| anyhow::anyhow!("stg_labels: {e}"))?,
             train_graph: None,
+            train_graph_boundary_state: None,
             forward_graph: None,
             forward_warmed: false,
             cublas_warmed: false,
@@ -1649,8 +1658,22 @@ impl PerceptionTrainer {
             }
         }
 
-        // ── 2. Three-state machine: warmup (first call), capture
-        //       (second call), replay (third+).
+        // ── 2. Four-state machine: warmup (first), capture (second),
+        //       replay (third+), recapture-on-boundary-change.
+        //       The captured graph records the attn_pool decision at
+        //       capture time (per `pearl_no_host_branches_in_captured_graph`).
+        //       In Sequential mode the decision toggles between sequences
+        //       within vs at file boundaries, so the cached graph is
+        //       invalidated whenever `last_seen_file_boundary` diverges
+        //       from the captured value. Random mode never diverges
+        //       (always `true`), so no recapture.
+        if let Some(captured_state) = self.train_graph_boundary_state {
+            if captured_state != self.last_seen_file_boundary {
+                // Boundary state changed — drop cached graph, fall through to recapture.
+                self.train_graph = None;
+                self.train_graph_boundary_state = None;
+            }
+        }
         if self.train_graph.is_some() {
             self.train_graph
                 .as_ref()
@@ -1687,6 +1710,7 @@ impl PerceptionTrainer {
                     "train end_capture returned None — no work captured"
                 ))?;
             self.train_graph = Some(graph);
+            self.train_graph_boundary_state = Some(self.last_seen_file_boundary);
         }
 
         // ── 3. Sync + read mapped-pinned loss.