From 387ce136d71c1d4e84d712829350d6fc99b58fe9 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Fri, 17 Apr 2026 22:31:22 +0200
Subject: [PATCH] =?UTF-8?q?fix:=20temporal=20ops=20outside=20graph=5Fforwa?=
 =?UTF-8?q?rd=20=E2=80=94=20ungraphed=20per-step?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adding mamba2/predictive_coding/regime_dropout/isv_temporal_route/
risk_budget to submit_forward_ops_main broke graph_forward replay
on Hopper (same graph structure issue). Moved to run_full_step
between graph_forward replay and aux ops — ungraphed but still
runs every step, enriching h_s2 before IQL/IQN/attention.

graph_forward keeps the working kernel set (cuBLAS + ISV + plan).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs   | 25 +------------------
 crates/ml/src/trainers/dqn/fused_training.rs  | 18 +++++++++++++
 2 files changed, 19 insertions(+), 24 deletions(-)

diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
index a684de44b..aad9eca7c 100644
--- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
+++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
@@ -7874,36 +7874,13 @@ impl GpuDqnTrainer {
         // ── 1. Forward (cuBLAS SGEMM — Pass 1 + 2, no Pass 3) ──────────
         self.launch_cublas_forward()?;
 
-        // ── 1b. Temporal + ISV + plan pipeline (full training forward) ──
+        // ── 1b. ISV + plan heads (inside graph capture — matches working SHA) ──
         {
             let batch_size = self.config.batch_size;
-
-            // Mamba2: temporal scan enriches h_s2 with history
-            self.mamba2_step(batch_size)?;
-
-            // Predictive coding: temporal smoothness loss on enriched trunk
-            self.compute_predictive_coding_loss(batch_size)?;
-
-            // Regime-conditioned dropout on h_s2
-            self.apply_regime_dropout(batch_size, true)?;
-
-            // ISV forward: encoder MLP → branch gate + gamma mod
             self.launch_isv_forward()?;
-
-            // ISV temporal routing: per-feature temporal_weight for next mamba2
-            self.launch_isv_temporal_route()?;
-
-            // ISV feature gate: modulate h_s2 features by regime embedding
             self.launch_isv_feature_gate(batch_size)?;
-
-            // Recursive confidence: predict own TD-error from h_s2
             self.launch_recursive_confidence_forward(batch_size)?;
-
-            // Trade plan: h_s2 → 6 plan parameters
             self.launch_trade_plan_forward(batch_size)?;
-
-            // Risk budget: h_s2 → R ∈ (0,1) before Q-value computation
-            self.risk_budget_forward(batch_size)?;
         }
 
         // ── 2+3. Loss + gradient (blended MSE + C51 via c51_alpha ramp) ─
diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs
index 9f3b2b81d..caacf660b 100644
--- a/crates/ml/src/trainers/dqn/fused_training.rs
+++ b/crates/ml/src/trainers/dqn/fused_training.rs
@@ -948,6 +948,24 @@ impl FusedTrainingCtx {
                 }
             }
 
+            // ── Step 2b: Temporal pipeline (outside graph — runs ungraphed) ────
+            // These enrich h_s2 with temporal context before aux ops.
+            // Cannot be inside graph_forward — adding kernels to the captured
+            // graph breaks replay on Hopper.
+            {
+                let bs = self.batch_size;
+                self.trainer.mamba2_step(bs)
+                    .map_err(|e| anyhow::anyhow!("mamba2_step: {e}"))?;
+                self.trainer.compute_predictive_coding_loss(bs)
+                    .map_err(|e| anyhow::anyhow!("predictive_coding: {e}"))?;
+                self.trainer.apply_regime_dropout(bs, true)
+                    .map_err(|e| anyhow::anyhow!("regime_dropout: {e}"))?;
+                self.trainer.launch_isv_temporal_route()
+                    .map_err(|e| anyhow::anyhow!("isv_temporal_route: {e}"))?;
+                self.trainer.risk_budget_forward(bs)
+                    .map_err(|e| anyhow::anyhow!("risk_budget: {e}"))?;
+            }
+
             // ── Step 2c: HER donor computation (outside graph_aux) ───────────
             if let Some(ref mut her) = self.gpu_her {
                 use crate::cuda_pipeline::gpu_her::HerGpuStrategy;