# H100 GPU Pipeline Optimization Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Reduce DQN epoch time from 2 min to <10s on H100 by eliminating GPU idle time from blocking synchronize calls, maximizing batch size, and enabling multi-stream parallelism.

**Architecture:** Replace `stream.synchronize()` with CUDA events for async scalar readback. Let `AutoBatchSizer` drive batch size (removes static cap). Capture experience collection as a CUDA Graph. Fork child streams for independent branch/target forward passes.

**Tech Stack:** Rust, cudarc 0.17 (CudaEvent, CudaStream::fork), CUDA 12.4, H100 PCIe 80GB.

---

## File Structure

| File | Action | Change |
|------|--------|--------|
| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Modify | CUDA events in `replay_adam_and_readback`, async q-stats, multi-stream forward |
| `crates/ml/src/cuda_pipeline/batched_forward.rs` | Modify | Multi-stream branch dispatch |
| `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` | Modify | CUDA Graph capture for timestep loop |
| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | Modify | Async training loop, event-based phase transitions |
| `config/gpu/h100.toml` | Modify | Remove static batch_size cap |
| `crates/ml/src/trainers/dqn/trainer/constructor.rs` | Modify | Let AutoBatchSizer drive batch_size |
| `infra/k8s/argo/precompute-features-template.yaml` | Done | Already committed (data-dir/SYMBOL alignment) |

---

## Phase 1: Eliminate GPU Idle Time (biggest impact)

### Task 1: CUDA Events replace per-step stream.synchronize()

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:2865-2893`

The `replay_adam_and_readback()` method calls `stream.synchronize()` on every training step (341×/epoch). This is the primary bottleneck — H100 finishes both CUDA graphs in microseconds then sits idle.

- [ ] **Step 1: Add event and async readback state to GpuDqnTrainer struct**

In `gpu_dqn_trainer.rs`, add fields to the `GpuDqnTrainer` struct (near line 479):

```rust
/// CUDA event for async scalar readback (replaces stream.synchronize per step)
readback_event: Option<Arc<CudaEvent>>,
/// Host-side buffers for deferred readback (loss, grad_norm from previous step)
pending_loss: f32,
pending_grad_norm: f32,
/// Whether there's an in-flight readback to collect
readback_pending: bool,
```

Initialize in `new()`:
```rust
readback_event: None,
pending_loss: 0.0,
pending_grad_norm: 0.0,
readback_pending: false,
```

- [ ] **Step 2: Rewrite replay_adam_and_readback to use events**

Replace the synchronous `replay_adam_and_readback` method:

```rust
/// Collect scalars from the PREVIOUS step (non-blocking poll),
/// then launch grad_norm + adam for the CURRENT step and record an event.
pub fn replay_adam_and_readback(&mut self) -> Result<FusedTrainScalars, MLError> {
    // 1. Collect previous step's scalars (if any)
    let prev_scalars = if self.readback_pending {
        // Poll: is the previous step's event done?
        if let Some(ref event) = self.readback_event {
            if !event.is_complete() {
                // Still in-flight — wait on just this event (not full stream)
                event.synchronize().map_err(|e|
                    MLError::ModelError(format!("readback event sync: {e}"))
                )?;
            }
        }
        // Now safe to read the scalars from the previous step
        self.readback_pending = false;
        FusedTrainScalars {
            total_loss: self.pending_loss,
            grad_norm: self.pending_grad_norm,
        }
    } else {
        FusedTrainScalars { total_loss: 0.0, grad_norm: 0.0 }
    };

    // 2. Launch current step: grad_norm + adam
    self.compute_grad_norm_outside_graph()?;
    self.replay_adam()?;

    // 3. Issue async DtoH for current step's scalars
    unsafe {
        cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
            (&mut self.pending_loss as *mut f32).cast(),
            self.total_loss_buf.raw_ptr(),
            std::mem::size_of::<f32>(),
            *self.stream.cu_stream(),
        );
        cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
            (&mut self.pending_grad_norm as *mut f32).cast(),
            self.grad_norm_f32_buf.raw_ptr(),
            std::mem::size_of::<f32>(),
            *self.stream.cu_stream(),
        );
    }

    // 4. Record event marking when the async DtoH completes
    self.readback_event = Some(self.stream.record_event(None).map_err(|e|
        MLError::ModelError(format!("readback event record: {e}"))
    )?);
    self.readback_pending = true;

    // Return PREVIOUS step's scalars (current step's will be collected next call)
    self.scalars_readback_host = [prev_scalars.total_loss, prev_scalars.grad_norm];
    Ok(prev_scalars)
}

/// Flush: block until the last in-flight readback completes.
/// Called once at epoch end.
pub fn flush_readback(&mut self) -> Result<FusedTrainScalars, MLError> {
    if !self.readback_pending {
        return Ok(FusedTrainScalars { total_loss: 0.0, grad_norm: 0.0 });
    }
    if let Some(ref event) = self.readback_event {
        if !event.is_complete() {
            event.synchronize().map_err(|e|
                MLError::ModelError(format!("flush readback sync: {e}"))
            )?;
        }
    }
    self.readback_pending = false;
    Ok(FusedTrainScalars {
        total_loss: self.pending_loss,
        grad_norm: self.pending_grad_norm,
    })
}
```

- [ ] **Step 3: Update training loop to flush at epoch end**

In `training_loop.rs`, after the training step loop (after the `for chunk_start in ...` loop ends), add:

```rust
// Flush the last in-flight readback
if let Some(ref mut fused) = self.fused_ctx {
    let final_scalars = fused.flush_readback()?;
    // Use final_scalars for epoch-end logging if needed
}
```

- [ ] **Step 4: Verify compile**

Run: `SQLX_OFFLINE=true cargo check -p ml --lib`
Expected: compiles clean

- [ ] **Step 5: Run local smoke test to validate loss trajectory**

Run: `FOXHUNT_TEST_DATA=test_data/futures-baseline FOXHUNT_FEATURE_CACHE_DIR=test_data/feature-cache SQLX_OFFLINE=true cargo test -p ml --lib -- smoke_tests::training_stability::test_50_epoch_convergence --ignored --nocapture`

Expected: test passes, loss values are finite, Sharpe is reasonable. Loss trajectory may differ slightly from synchronous path (1-step lag on PER priorities) but should converge similarly.

- [ ] **Step 6: Commit**

```bash
git commit -m "perf: CUDA events replace per-step stream.synchronize in training loop"
```

---

### Task 2: Async Q-stats readback

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:3640-3660`

The `reduce_current_q_stats()` calls `memcpy_dtoh` every 50 steps — implicit sync. Use the same event pattern.

- [ ] **Step 1: Add q-stats event and pending buffer**

Add to `GpuDqnTrainer` struct:

```rust
/// Pending q-stats from previous reduction (deferred readback)
q_stats_event: Option<Arc<CudaEvent>>,
q_stats_pending: [f32; 5],
q_stats_ready: bool,
```

Initialize:
```rust
q_stats_event: None,
q_stats_pending: [0.0; 5],
q_stats_ready: false,
```

- [ ] **Step 2: Rewrite reduce_current_q_stats for async readback**

Replace the synchronous `memcpy_dtoh` in `reduce_current_q_stats()`:

```rust
pub fn reduce_current_q_stats(&mut self) -> Result<QValueStatsResult, MLError> {
    // Collect previous reduction if ready
    let result = if self.q_stats_ready {
        if let Some(ref event) = self.q_stats_event {
            if !event.is_complete() {
                event.synchronize().map_err(|e|
                    MLError::ModelError(format!("q_stats event sync: {e}"))
                )?;
            }
        }
        self.q_stats_ready = false;
        let h = self.q_stats_pending;
        QValueStatsResult {
            avg_max_q: h[0] as f64, q_min: h[1], q_max: h[2], q_mean: h[3], q_var: h[4],
        }
    } else {
        QValueStatsResult { avg_max_q: 0.0, q_min: 0.0, q_max: 0.0, q_mean: 0.0, q_var: 0.0 }
    };

    // Launch new reduction kernel
    // ... (existing kernel launch code stays the same) ...

    // Async readback
    unsafe {
        cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
            self.q_stats_pending.as_mut_ptr().cast(),
            self.q_stats_buf.raw_ptr(),
            5 * std::mem::size_of::<f32>(),
            *self.stream.cu_stream(),
        );
    }
    self.q_stats_event = Some(self.stream.record_event(None).map_err(|e|
        MLError::ModelError(format!("q_stats event record: {e}"))
    )?);
    self.q_stats_ready = true;

    Ok(result)
}
```

- [ ] **Step 3: Verify compile + smoke test**

Run: `SQLX_OFFLINE=true cargo check -p ml --lib`
Run: smoke test (same as Task 1 Step 5)

- [ ] **Step 4: Commit**

```bash
git commit -m "perf: async q-stats readback via CUDA events"
```

---

### Task 3: Dynamic batch sizing — remove static cap

**Files:**
- Modify: `config/gpu/h100.toml:3`
- Modify: `crates/ml/src/trainers/dqn/trainer/constructor.rs` (AutoBatchSizer logic)

- [ ] **Step 1: Remove static batch_size from h100.toml**

Change `config/gpu/h100.toml` line 3:

```toml
# batch_size is auto-computed by AutoBatchSizer from available VRAM.
# Do NOT set a static value here — let the sizer maximize GPU utilization.
# batch_size = 1024  # REMOVED: was capping H100 far below optimal
```

Actually: the TOML profile must still have the field for the parser. Set it to 0 as a sentinel meaning "auto":

```toml
batch_size = 0  # 0 = auto-compute from VRAM (AutoBatchSizer)
```

- [ ] **Step 2: Update constructor to honor batch_size=0 as auto**

In `constructor.rs`, find where `hyperparams.batch_size` is set from the profile. After the profile applies, if `batch_size == 0`, let AutoBatchSizer compute it:

```rust
// After profile.apply_to(&mut hyperparams):
if hyperparams.batch_size == 0 {
    // Auto-compute: AutoBatchSizer already ran and set the VRAM ceiling.
    // Use the ceiling as the batch size, capped at a sane max (16384).
    hyperparams.batch_size = auto_batch_ceiling.min(16384);
    info!("AutoBatchSizer: batch_size auto-computed to {}", hyperparams.batch_size);
}
```

- [ ] **Step 3: Verify the auto-computed value is reasonable**

Run locally (RTX 3050): the sizer should still pick a small batch (64-128) for 4GB VRAM.
Check H100 would get: look at the `AutoBatchSizer: GPU VRAM ceiling = 2085808` log — that's the theoretical max. The actual batch should be much smaller (constrained by kernel shared memory and practical gradient noise). Cap at 8192 or 16384.

- [ ] **Step 4: Commit**

```bash
git commit -m "perf: dynamic batch sizing — AutoBatchSizer drives batch_size on H100"
```

---

### Task 4: Experience collection → training phase overlap

**Files:**
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs:1196`

- [ ] **Step 1: Replace stream.synchronize() with event after experience collection**

At line 1196, change:

```rust
// BEFORE:
stream.synchronize()
    .map_err(|e| MLError::DeviceError(format!("experience sync: {e}")))?;

// AFTER:
let experience_done_event = stream.record_event(None)
    .map_err(|e| MLError::DeviceError(format!("experience event record: {e}")))?;
```

Then before the first training step that samples from the replay buffer, wait on the event:

```rust
// Before first replay buffer sample:
if !experience_done_event.is_complete() {
    experience_done_event.synchronize()
        .map_err(|e| MLError::DeviceError(format!("experience event wait: {e}")))?;
}
```

This allows CPU-side training setup (lock acquisition, pre-sampling initialization) to overlap with the tail end of experience collection.

- [ ] **Step 2: Verify + commit**

```bash
git commit -m "perf: event-based experience→training phase transition"
```

---

### Task 5: fxcache PVC regeneration

- [ ] **Step 1: Regenerate PVC cache with aligned data-dir**

The precompute template already uses `data-dir/SYMBOL`. Just re-run:

```bash
./scripts/argo-precompute.sh ES.FUT --pool ci-compile-cpu
```

- [ ] **Step 2: Verify training hits fxcache**

Submit a quick 2-epoch baseline and check logs for "fxcache hit":

```bash
argo submit -n foxhunt --from=wftmpl/train-baseline-rl -p train-epochs=2
```

Check: `argo logs <name> | grep fxcache`

- [ ] **Step 3: Commit (if any template changes needed)**

---

## Phase 2: Multi-Stream Parallelism

### Task 6: Multi-stream branch dispatch in forward pass

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs:265-284`
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (add child streams)

- [ ] **Step 1: Add 3 child streams to CublasForward**

In `batched_forward.rs`, add to the struct:

```rust
/// Child streams for parallel branch dispatch (forked from main training stream)
branch_streams: [Arc<CudaStream>; 3],
```

Create them in the constructor by forking:

```rust
let branch_streams = [
    stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 0: {e}")))?,
    stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 1: {e}")))?,
    stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 2: {e}")))?,
];
```

- [ ] **Step 2: Dispatch branches on separate streams**

In `forward_online_raw()`, replace the sequential `for d in 0..3` loop (lines 270-284):

```rust
// Fork: dispatch each branch on its own stream
for d in 0..3 {
    let branch_stream = &self.branch_streams[d];
    // Wait for trunk (h_s2) to be ready on this branch stream
    let trunk_event = self.stream.record_event(None)
        .map_err(|e| MLError::ModelError(format!("trunk event: {e}")))?;
    branch_stream.wait(&trunk_event)
        .map_err(|e| MLError::ModelError(format!("branch wait trunk: {e}")))?;

    let n_d = branch_sizes[d];
    let w_fc_idx = branch_w_base[d];

    // Hidden layer
    self.gemmex_bf16_on_stream(branch_stream, w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], self.adv_h, b, self.shared_h2, "h_bd")?;
    self.launch_add_bias_relu_bf16_on_stream(branch_stream, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?;

    // Output layer
    let adv_out_ptr = b_logits_ptr + logit_byte_offset;
    self.gemmex_bf16_to_f32_on_stream(branch_stream, w_ptrs[w_fc_idx + 2], branch_h_ptrs[d], adv_out_ptr, n_d * na, b, self.adv_h, "adv_logits")?;
    self.launch_add_bias_f32_on_stream(branch_stream, adv_out_ptr, w_ptrs[w_fc_idx + 3], n_d * na, b)?;

    logit_byte_offset += (b * n_d * na * std::mem::size_of::<f32>()) as u64;
}

// Join: main stream waits for all branches to finish
for d in 0..3 {
    let branch_done = self.branch_streams[d].record_event(None)
        .map_err(|e| MLError::ModelError(format!("branch done event: {e}")))?;
    self.stream.wait(&branch_done)
        .map_err(|e| MLError::ModelError(format!("join branch {d}: {e}")))?;
}
```

Note: `gemmex_bf16_on_stream` and `launch_add_bias_relu_bf16_on_stream` are new variants that take an explicit stream parameter. Implement them by extracting the cuBLAS handle → set stream → call GemmEx → restore stream. cuBLAS handles are NOT stream-specific — `cublasSetStream()` switches the active stream per-call.

- [ ] **Step 3: Apply same pattern to f32 forward (line 325) and backward (line 461)**

Repeat the fork/join pattern for both the f32 experience collection forward and the backward pass branch loops.

- [ ] **Step 4: Verify compile + smoke test**

Run smoke tests — output logits must be identical (branches write to non-overlapping memory regions).

- [ ] **Step 5: Commit**

```bash
git commit -m "perf: multi-stream branch dispatch — 3 advantage heads in parallel"
```

---

### Task 7: Multi-stream target/online forward parallelism

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:4307-4323`

- [ ] **Step 1: Add a second stream for Double DQN pass**

Add to `GpuDqnTrainer`:

```rust
/// Second stream for parallel target/online next-state forward passes
double_dqn_stream: Arc<CudaStream>,
```

Fork in constructor.

- [ ] **Step 2: Dispatch Pass 2 and Pass 3 on separate streams**

In `submit_forward_ops()`, after Pass 1 (online forward on current states):

```rust
// Pass 2 + Pass 3 are independent: both read next_states, write to different outputs.
// Fork: record event after stochastic depth (Pass 1 done)
let pass1_done = self.stream.record_event(None)?;
self.double_dqn_stream.wait(&pass1_done)?;

// Pass 2: target forward on next_states (main stream)
cublas.forward_target_raw(
    &self.stream, self.ptrs.next_states_buf, &tg_w_ptrs, ...
)?;

// Pass 3: online forward on next_states (double_dqn_stream)
cublas.forward_online_raw(
    &self.double_dqn_stream, self.ptrs.next_states_buf, &on_w_ptrs, ...
)?;

// Join before loss kernel
let target_done = self.stream.record_event(None)?;
let online_done = self.double_dqn_stream.record_event(None)?;
self.stream.wait(&online_done)?;
// main stream now has both pass 2 + pass 3 complete
```

- [ ] **Step 3: Handle CUDA Graph capture with multi-stream**

CUDA Graph capture with fork/join requires `CU_STREAM_CAPTURE_MODE_THREAD_LOCAL` (already used at line 3734). The fork/join events within the graph are captured as graph dependencies. Verify this works with cudarc — may need raw `cuStreamWaitEvent` calls during capture.

- [ ] **Step 4: Verify + smoke test + commit**

```bash
git commit -m "perf: parallel target/online forward on separate CUDA streams"
```

---

### Task 8: CUDA Graph for experience collection

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs:1506`

This is the most complex task. The 100-iteration host loop dispatches 400+ kernels. Capturing as a graph eliminates all host→GPU latency.

- [ ] **Step 1: Make per-timestep parameters GPU-resident**

The loop currently passes `epsilon`, step counter, and state pointers from the host. These must become GPU buffers:

```rust
/// GPU-resident parameters for graph-captured timestep loop
epsilon_buf: CudaSlice<f32>,       // [1] — updated via DtoH before graph launch
step_counter_buf: CudaSlice<u32>,  // [1] — incremented by a tiny CUDA kernel per step
```

- [ ] **Step 2: Create a step-increment kernel**

Tiny kernel that increments `step_counter_buf[0]` and computes `epsilon = eps_end + (eps_start - eps_end) * exp(-step / decay)` on GPU. No host involvement.

- [ ] **Step 3: Capture the 100-timestep loop as a single graph**

```rust
stream.begin_capture(CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?;
for t in 0..timesteps {
    self.launch_experience_state_gather()?;
    self.launch_cublas_forward_online()?;
    self.launch_compute_expected_q()?;
    self.launch_experience_action_select()?;  // reads epsilon_buf
    self.launch_experience_env_step()?;
    self.launch_step_increment()?;  // increments step_counter, updates epsilon
}
let graph = stream.end_capture(CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH)?;
```

Before each epoch: `memcpy_htod` the initial epsilon and step counter, then `graph.launch()`.

- [ ] **Step 4: Handle dynamic epsilon**

Epsilon decays across episodes within an epoch. The step-increment kernel computes the new epsilon from the step counter using the exponential decay formula, writing to `epsilon_buf`. The action selection kernel reads `epsilon_buf` directly — no host involvement.

- [ ] **Step 5: Verify collected experiences match sequential path**

Run 1 epoch with both paths (graph and sequential), compare states/actions/rewards for the first 100 experiences. Must be numerically identical given the same Philox RNG seed.

- [ ] **Step 6: Commit**

```bash
git commit -m "perf: CUDA Graph capture for 100-timestep experience collection loop"
```

---

### Task 9: Lock contention removal

**Files:**
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs:1336-1360`

- [ ] **Step 1: Pre-sample all batches for the epoch in one lock acquisition**

Instead of acquiring/releasing the read lock every 32 steps, sample ALL batches at epoch start:

```rust
let all_batches = {
    let agent = self.agent.read().await;
    let buffer = agent.memory();
    (0..num_training_steps)
        .map(|_| buffer.sample(self.current_batch_size))
        .collect::<Vec<_>>()
};
// Lock released. Training loop iterates over pre-sampled batches.
```

This is safe because the GPU PER buffer is immutable during training (new experiences are only added in Phase 2, which is finished before training starts).

- [ ] **Step 2: Remove PREFETCH_K chunking**

The `for chunk_start in (0..num_training_steps).step_by(PREFETCH_K)` loop becomes a simple `for batch in all_batches.iter()`.

- [ ] **Step 3: Verify + commit**

```bash
git commit -m "perf: pre-sample all training batches in single lock acquisition"
```

---

## Validation

### Task 10: H100 benchmark

- [ ] **Step 1: Rebuild and upload binary to PVC**

Compile in-cluster via the baseline template.

- [ ] **Step 2: Run 10-epoch baseline and measure**

```bash
argo submit -n foxhunt --from=wftmpl/train-baseline-rl -p train-epochs=10
```

Check:
- Epoch time: target <10s (from 2 min)
- `nvidia-smi`: SM utilization >85%, memory bandwidth >50%, TDP >250W
- Loss trajectory: finite, converging
- "fxcache hit" in logs

- [ ] **Step 3: Compare loss with pre-optimization baseline**

The 7-epoch metrics from the terminated run serve as reference:
- Epoch 4: Sharpe 0.06, PF 0.94
- Epoch 6: Sharpe -1.12, PF 1.17

Post-optimization should show similar or better convergence (larger batches may improve gradient quality).

- [ ] **Step 4: Run full 50-epoch baseline if metrics look good**

- [ ] **Step 5: Save metrics to memory**