foxhunt/docs/superpowers/specs/2026-04-08-h100-graph-optimizations.md

H100 Graph-Level Optimizations: 4 Deferred Fixes

Goal: Implement the 4 remaining H100 performance optimizations (Fixes 2,3,4,7) saving ~13s/epoch. Combined with the 5 already-implemented fixes, total epoch time drops from 47s to ~16s.

Architecture: CUDA graph management improvements — eliminate re-capture, async overlap, kernel fusion, mega-graph. All GPU-side, no CPU path.

Tech Stack: CUDA driver API (cuGraphExecUpdate, cuStreamWaitEvent, cuEventRecord), cudarc, cuBLAS

---

Fix 2: Eliminate Validation Graph Re-capture (saves ~4s/epoch)

Problem

compute_validation_loss() calls evaluator.invalidate_dqn_graph() every epoch (metrics.rs:510), forcing full CUDA graph re-capture of the backtest forward pass.

Solution

The evaluator has its OWN weight buffers (eval_weights, eval_branching) that the graph was captured against. Instead of invalidating, DtoD copy new weights INTO those existing buffers. The graph's captured pointers remain valid — no re-capture needed.

Implementation

In metrics.rs, replace evaluator.invalidate_dqn_graph() with:

evaluator.update_weights(online_weights, branching_weights)?;

In gpu_backtest_evaluator.rs, add:

pub fn update_weights(
    &mut self,
    online: &DuelingWeightSet,
    branching: &BranchingWeightSet,
) -> Result<(), MLError> {
    // DtoD copy into evaluator's own weight buffers (same addresses graph was captured with)
    self.copy_weights_from(online, branching)?;
    // No invalidation — topology unchanged, data updated in-place
    Ok(())
}

The evaluator already has copy_weights_from or similar (check for sync_weights or upload_weights methods). If not, add DtoD copies for each weight tensor using the same pattern as sync_one but into the evaluator's buffers.

Files

• Modify: crates/ml/src/trainers/dqn/trainer/metrics.rs — replace invalidate call
• Modify: crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs — add update_weights method

---

Fix 3: Async Validation on Separate Stream (saves ~8s overlap)

Problem

Validation runs synchronously on the main stream, blocking the next epoch's experience collection.

Solution

Create a dedicated validation stream + event. Validation runs on stream 2 while the next epoch's experience collection runs on stream 1.

Implementation

Add fields to DQNTrainer (mod.rs):

pub(crate) validation_stream: Option<Arc<CudaStream>>,
pub(crate) validation_event: Option<cudarc::driver::sys::CUevent>,
pub(crate) pending_val_loss: Option<f64>,

Initialize in constructor:

validation_stream: if let Some(ref s) = cuda_stream {
    Some(Arc::new(s.context().new_stream()
        .map_err(|e| anyhow::anyhow!("validation stream: {e}"))?))
} else { None },
validation_event: if cuda_stream.is_some() {
    let mut event: cudarc::driver::sys::CUevent = std::ptr::null_mut();
    unsafe { cudarc::driver::sys::cuEventCreate(&mut event, 0); }
    Some(event)
} else { None },
pending_val_loss: None,

Modify epoch loop (training_loop.rs):

// At epoch START: collect pending validation from previous epoch
if let Some(val) = self.pending_val_loss.take() {
    val_loss = val;
}

// At epoch END (after training, before next epoch):
// Record event on main stream
if let (Some(ref val_stream), Some(event)) = (&self.validation_stream, self.validation_event) {
    let main = self.cuda_stream.as_ref().unwrap();
    unsafe {
        cudarc::driver::sys::cuEventRecord(event, main.cu_stream());
        cudarc::driver::sys::cuStreamWaitEvent(val_stream.cu_stream(), event, 0);
    }
    // Launch validation on stream 2 (non-blocking on main)
    let val_result = self.compute_validation_loss().await;
    self.pending_val_loss = val_result.ok();
    // Main stream continues to next epoch immediately
}

The backtest evaluator needs to use the validation stream for its cuBLAS calls. Add a set_stream() method or pass the stream to evaluate_dqn_graphed(). The evaluator's internal cuBLAS handle must be associated with the validation stream.

If the evaluator creates its own cuBLAS handle at construction, it can be bound to the validation stream at init time. The graph captured on this stream replays on this stream — no cross-stream issues.

Files

• Modify: crates/ml/src/trainers/dqn/trainer/mod.rs — add fields
• Modify: crates/ml/src/trainers/dqn/trainer/training_loop.rs — async launch
• Modify: crates/ml/src/trainers/dqn/trainer/metrics.rs — pass validation stream
• Modify: crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs — accept stream param

---

Fix 4: Aux GEMM in graph_aux (saves ~0.5s/epoch)

Problem

Exposure aux GEMM runs as 3 separate kernel launches outside graph_aux. The aux_weight scalar is passed by value — graph captures the VALUE, not a pointer.

Solution

1. Store aux_weight in a GPU-resident scalar buffer CudaSlice<f32> of size 1
2. Modify exposure_aux_grad_kernel signature: float aux_weight → const float* aux_weight_ptr
3. Before graph_aux replay, async HtoD the current weight value into the scalar buffer
4. Move the aux GEMM call from run_full_step() into submit_aux_ops()

Implementation

CUDA kernel change (dqn_utility_kernels.cu):

// Change parameter from scalar to pointer:
extern "C" __global__ void exposure_aux_grad_kernel(
    ...
    const float* __restrict__ aux_weight_ptr,  // was: float aux_weight
    ...
) {
    float aux_weight = *aux_weight_ptr;  // device read
    ...
}

Rust changes (gpu_dqn_trainer.rs):

// Add field:
exposure_aux_weight_gpu: CudaSlice<f32>,  // [1] stable address

// Allocate:
let exposure_aux_weight_gpu = stream.alloc_zeros::<f32>(1)?;

// In launch method: pass pointer instead of value
.arg(&self.exposure_aux_weight_gpu.raw_ptr())  // was: .arg(&aux_weight)

Update weight before replay (fused_training.rs):

// In run_full_step(), BEFORE graph_aux replay:
let weight_val = [self.exposure_aux_weight as f32];
unsafe {
    cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
        self.trainer.exposure_aux_weight_gpu.raw_ptr(),
        weight_val.as_ptr().cast(),
        4, // sizeof(f32)
        self.stream.cu_stream(),
    );
}

Move aux GEMM into submit_aux_ops (fused_training.rs): Cut the aux GEMM block from run_full_step() (line ~826-835) and paste at the end of submit_aux_ops().

Files

• Modify: crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu — kernel sig change
• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs — scalar buffer + launch update
• Modify: crates/ml/src/trainers/dqn/fused_training.rs — HtoD update + move aux into submit_aux_ops

---

Fix 7: Mega-Graph Fusion (saves ~0.3s/epoch)

Problem

4 graph launches per training step: spectral → forward → aux → adam = 6ms overhead/step.

Solution

Merge spectral + forward + aux into one graph_mega. Conditional ops (vaccine, causal) stay outside. Result: 2 launches per step (mega + adam) = 3ms overhead.

Implementation

Add field (fused_training.rs):

graph_mega: Option<RawCudaGraph>,

New capture method:

fn capture_graph_mega(&mut self, agent: &mut DQNAgentType, gpu_batch: &GpuBatch) -> Result<()> {
    self.trainer.sync_all_streams()?;
    let _ = self.stream.context().check_err();

    let cu_stream = self.stream.cu_stream();
    let mut graph: cuda_sys::CUgraph = std::ptr::null_mut();
    let mut exec: cuda_sys::CUgraphExec = std::ptr::null_mut();

    unsafe { cuda_sys::cuStreamBeginCapture_v2(cu_stream, THREAD_LOCAL); }

    // Phase 1: Spectral norm
    self.trainer.apply_spectral_norm(&self.online_dueling, &self.online_branching)?;
    
    // Phase 2: Forward + backward (upload + forward + loss + backward)
    self.trainer.submit_forward_ops_main()?;
    
    // Phase 3: Aux ops (clip + EMA + attention + IQL + IQN + CQL + aux GEMM)
    self.submit_aux_ops(agent, gpu_batch)?;

    unsafe { cuda_sys::cuStreamEndCapture(cu_stream, &mut graph); }
    unsafe { cuda_sys::cuGraphInstantiateWithFlags(&mut exec, graph, 0); }

    self.graph_mega = Some(RawCudaGraph { exec, graph });
    self.graph_spectral = None;  // superseded
    self.graph_forward = None;   // superseded
    self.graph_aux = None;       // superseded
    Ok(())
}

Update run_full_step:

if let Some(ref graph) = self.graph_mega {
    // Update dynamic values via HtoD (read by graph kernels from device pointers)
    self.pre_replay_state_update(agent);
    self.update_aux_weight_gpu()?;  // Fix 4: aux weight scalar
    graph.launch(self.stream.cu_stream())?;
} else {
    // Steps 0-1: ungraphed (same as current)
    self.trainer.apply_spectral_norm(...)?;
    self.trainer.train_step_gpu(...)?;
    self.submit_aux_ops(agent, gpu_batch)?;
    
    // Capture mega-graph at step 2
    if self.steps_since_varmap_sync == 2 {
        self.capture_graph_mega(agent, gpu_batch)?;
    }
}

// Steps that stay OUTSIDE the mega-graph:
// - Conditional ops (vaccine, causal) — step-dependent
// - Exposure aux targets DtoD copy — data-dependent  
// - compute_grad_norm_outside_graph — already outside
// - replay_adam — separate graph (different frequency possible)

Files

• Modify: crates/ml/src/trainers/dqn/fused_training.rs — new capture + updated run_full_step

---

Combined Impact

Fix	Savings	Mechanism
2: No re-capture	-4s/epoch	DtoD into existing buffers
3: Async validation	-8s/epoch	Separate stream overlap
4: Aux in graph	-0.5s/epoch	GPU scalar + kernel sig change
7: Mega-graph	-0.3s/epoch	4 launches → 2 launches
Total	-12.8s/epoch

With the 5 already-implemented fixes (-18s), total savings: -30.8s/epoch (47s → ~16s).

50 trials × 40 epochs × 16s = 8.9 hours (was 26 hours — 66% reduction, $51 saved per run)

Testing

• cargo test -p ml --lib — 900+ pass, 0 regressions
• GPU smoke test on RTX 3050 — no SIGSEGV from graph changes
• Verify validation still produces correct Sharpe (async doesn't corrupt)
• Verify mega-graph replay matches ungraphed execution (step 0 vs step 3+)