From 341346af098bacbfab6d8fb7e2694f2ce92ea533 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Fri, 17 Apr 2026 10:36:23 +0200
Subject: [PATCH] fix: ISV+plan kernels run pre-mega-replay, not inside
 graph_mega

NVRTC-compiled ISV kernels conflict with cuBLAS GEMM internal stream
scheduling when captured in the same CUDA Graph as cuBLAS ops.
Moving ISV+plan to pre-replay phase: they run BEFORE graph_mega
replay on every training step. Their outputs (gamma_mod_buf,
branch_gate_buf, gated_h_s2_buf, predicted_error_buf, plan_params_buf)
are stable device buffers consumed by the captured graph.
Same result, no deadlock. ISV updates every step (not every 50).

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

diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
index e6f6636f2..d7f5b1593 100644
--- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
+++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
@@ -7877,17 +7877,13 @@ impl GpuDqnTrainer {
         // ── 1. Forward (cuBLAS SGEMM — Pass 1 + 2, no Pass 3) ──────────
         self.launch_cublas_forward()?;
 
-        // ── 1b. ISV + plan heads inside graph capture.
-        //        Pinned device-mapped pointers are stable (address captured, values change).
-        //        Same pattern as total_loss_dev_ptr, grad_norm_dev_ptr, etc.
-        {
-            let batch_size = self.config.batch_size;
-            self.launch_isv_forward()?;
-            self.launch_isv_feature_gate(batch_size)?;
-            self.launch_recursive_confidence_forward(batch_size)?;
-            self.launch_trade_plan_forward(batch_size)?;
-            self.launch_plan_noise_inject(batch_size)?;
-        }
+        // ── 1b. ISV + plan heads run BEFORE graph_mega replay (pre-replay phase
+        //        in fused_training.rs). Their outputs (gamma_mod_buf, branch_gate_buf,
+        //        gated_h_s2_buf, predicted_error_buf, plan_params_buf) are stable device
+        //        buffers consumed by kernels inside the graph.
+        //        NOT captured inside graph_mega — ISV kernel + cuBLAS backward interaction
+        //        causes graph replay deadlock (investigated, root cause: graph node scheduling
+        //        conflict between ISV NVRTC kernels and cuBLAS GEMM internal streams).
 
         // ── 2+3. Loss + gradient (blended MSE + C51 via c51_alpha ramp) ─
         self.launch_curiosity_inference()?;
diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs
index d503d1519..d4d844b66 100644
--- a/crates/ml/src/trainers/dqn/fused_training.rs
+++ b/crates/ml/src/trainers/dqn/fused_training.rs
@@ -875,6 +875,25 @@ impl FusedTrainingCtx {
             // Adam step counter update (outside graph — value captured by address)
             self.trainer.adam_step_async();
 
+            // ISV + plan heads: run BEFORE graph_mega replay, AFTER batch upload.
+            // Their outputs are stable device buffers consumed by the graph.
+            // Running these outside graph_mega avoids NVRTC-cuBLAS graph scheduling conflict.
+            {
+                let bs = self.batch_size;
+                self.trainer.launch_isv_forward()
+                    .map_err(|e| anyhow::anyhow!("ISV forward (pre-mega): {e}"))?;
+                self.trainer.launch_isv_feature_gate(bs)
+                    .map_err(|e| anyhow::anyhow!("ISV feature gate (pre-mega): {e}"))?;
+                self.trainer.launch_recursive_confidence_forward(bs)
+                    .map_err(|e| anyhow::anyhow!("Recursive confidence (pre-mega): {e}"))?;
+                self.trainer.launch_trade_plan_forward(bs)
+                    .map_err(|e| anyhow::anyhow!("Trade plan forward (pre-mega): {e}"))?;
+                self.trainer.launch_plan_noise_inject(bs)
+                    .map_err(|e| anyhow::anyhow!("Plan noise (pre-mega): {e}"))?;
+                self.trainer.fill_gamma_buf()
+                    .map_err(|e| anyhow::anyhow!("Fill gamma buf (pre-mega): {e}"))?;
+            }
+
             // HER donor computation (outside graph — indices vary per step)
             if let Some(ref mut her) = self.gpu_her {
                 use crate::cuda_pipeline::gpu_her::HerGpuStrategy;