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

fix: ISV+plan kernels BACK inside graph_mega (matching working 11b1a1ca9)

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-17 11:53:29 +02:00
parent 2f4a669392
commit f264ba00cb
2 changed files with 11 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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()?;

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

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