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

fix: ISV+plan kernels run pre-mega-replay, not inside graph_mega

Browse Source 
NVRTC-compiled ISV kernels conflict with cuBLAS GEMM internal stream
scheduling when captured in the same CUDA Graph as cuBLAS ops.
Moving ISV+plan to pre-replay phase: they run BEFORE graph_mega
replay on every training step. Their outputs (gamma_mod_buf,
branch_gate_buf, gated_h_s2_buf, predicted_error_buf, plan_params_buf)
are stable device buffers consumed by the captured graph.
Same result, no deadlock. ISV updates every step (not every 50).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-17 10:36:23 +02:00
parent 99f1687972
commit 341346af09
2 changed files with 26 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -7877,17 +7877,13 @@ impl GpuDqnTrainer {
// ── 1. Forward (cuBLAS SGEMM — Pass 1 + 2, no Pass 3) ──────────
self.launch_cublas_forward()?;
// ── 1b. ISV + plan heads inside graph capture.
// Pinned device-mapped pointers are stable (address captured, values change).
// Same pattern as total_loss_dev_ptr, grad_norm_dev_ptr, etc.
{
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)?;
}
// ── 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).
// ── 2+3. Loss + gradient (blended MSE + C51 via c51_alpha ramp) ─
self.launch_curiosity_inference()?;

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

@@ -875,6 +875,25 @@ 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}"))?;
self.trainer.fill_gamma_buf()
.map_err(|e| anyhow::anyhow!("Fill gamma buf (pre-mega): {e}"))?;
}
// HER donor computation (outside graph — indices vary per step)
if let Some(ref mut her) = self.gpu_her {
use crate::cuda_pipeline::gpu_her::HerGpuStrategy;