From b215976a7353ee1edda2d21a97edc42bca74b053 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 17 Apr 2026 22:50:23 +0200 Subject: [PATCH] =?UTF-8?q?fix:=20disable=20temporal=20ops=20until=20cuBLA?= =?UTF-8?q?S=20aux=20rewrite=20=E2=80=94=20corrupts=20graph=5Faux?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Temporal ops (mamba2, regime_dropout, isv_temporal_route, risk_budget) between graph_forward and graph_aux replay corrupt graph_aux on Hopper. ISV signal update + hold enforcement + aux_frequency still active. Temporal pipeline will be properly addressed in the cuBLAS aux rewrite. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/trainers/dqn/fused_training.rs | 24 +++++--------------- 1 file changed, 6 insertions(+), 18 deletions(-) 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 {