gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: disable temporal ops until cuBLAS aux rewrite — corrupts graph_aux

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-17 22:50:23 +02:00
parent 6252e93bac
commit b215976a73
1 changed files with 6 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -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 {