docs: H100 graph-level optimization spec — 4 deferred fixes (~13s/epoch)

Fix 2: Eliminate validation graph re-capture (DtoD into same buffers) Fix 3: Async validation on separate CUDA stream (overlap with experience) Fix 4: Aux GEMM in graph_aux (GPU scalar + kernel sig change) Fix 7: Mega-graph fusion (spectral+forward+aux → 2 launches) Combined with 5 implemented fixes: 47s → ~16s/epoch Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-08 08:48:36 +02:00
parent bb1b9a703d
commit fdc8f30c8d
1 changed files with 270 additions and 0 deletions
--- a/docs/superpowers/specs/2026-04-08-h100-graph-optimizations.md
+++ b/docs/superpowers/specs/2026-04-08-h100-graph-optimizations.md
@@ -0,0 +1,270 @@
+# 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:
+```rust
+evaluator.update_weights(online_weights, branching_weights)?;
+```
+
+In `gpu_backtest_evaluator.rs`, add:
+```rust
+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`):
+```rust
+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:
+```rust
+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`):
+```rust
+// 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`):
+```cuda
+// 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`):
+```rust
+// 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`):
+```rust
+// 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`):
+```rust
+graph_mega: Option<RawCudaGraph>,
+```
+
+**New capture method:**
+```rust
+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:**
+```rust
+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+)