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;