# 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+)
