diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 3607a3907..575c3a7c4 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -948,24 +948,12 @@ impl FusedTrainingCtx { } } - // ── Step 2b: Temporal pipeline (outside graph — runs ungraphed) ──── - // Disable event tracking: these ops use buffers shared with graphed - // kernels, and cudarc event recording corrupts graph_aux replay. - { - unsafe { self.stream.context().disable_event_tracking(); } - 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}"))?; - unsafe { self.stream.context().enable_event_tracking(); } - } + // ── Step 2b: Temporal pipeline ──── + // TODO: temporal ops (mamba2, regime_dropout, isv_temporal_route, + // risk_budget, predictive_coding) corrupt graph_aux replay when run + // ungraphed between graph_forward and graph_aux. Will be addressed + // in the cuBLAS aux rewrite — either inside graph_temporal or + // by eliminating graph_aux dependency. // ── Step 2c: HER donor computation (outside graph_aux) ─────────── if let Some(ref mut her) = self.gpu_her {