From f264ba00cb5b7ae931f8435f1536fdaa91abc0cb Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 17 Apr 2026 11:53:29 +0200 Subject: [PATCH] fix: ISV+plan kernels BACK inside graph_mega (matching working 11b1a1ca9) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reverted to the structure that WORKED on H100: ISV kernels inside submit_forward_ops_main (captured in graph_mega). Removed duplicate pre-replay calls. The working version had ISV in graph_mega — removing them was wrong. Added plan_noise_inject to the same block. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 16 ++++++++------- crates/ml/src/trainers/dqn/fused_training.rs | 20 ++----------------- 2 files changed, 11 insertions(+), 25 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index d7f5b1593..1a164e286 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -7877,13 +7877,15 @@ impl GpuDqnTrainer { // ── 1. Forward (cuBLAS SGEMM — Pass 1 + 2, no Pass 3) ────────── self.launch_cublas_forward()?; - // ── 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). + // ── 1b. ISV + plan heads inside graph capture (same as working 11b1a1ca9) ── + { + 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)?; + } // ── 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 9d471fde3..d2edf6cc9 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -875,24 +875,8 @@ 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}"))?; - // fill_gamma_buf is inside graph_mega (submit_forward_ops_main) — NOT here. - // adaptive_gamma is a pinned scalar read by the captured kernel. - } + // ISV + plan heads are INSIDE graph_mega (submit_forward_ops_main). + // No pre-replay needed — they're captured in the graph. // HER donor computation (outside graph — indices vary per step) if let Some(ref mut her) = self.gpu_her {