# H100 Maximum Performance Implementation 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:** Maximize H100 GPU utilization by eliminating all synchronization stalls, memory waste, kernel launch overhead, and compute inefficiency in the fused CUDA training pipeline.

**Architecture:** Ten fixes ordered by dependency: first fix the seg_tree race (P0 blocker), then eliminate sync points, then attack launch overhead and memory bandwidth, finally optimize compute and stream parallelism.

**Tech Stack:** Rust 1.85, CUDA 12.4 (SM 9.0, cudarc 0.19), cuBLAS, half (bf16)

---

### Task 1: Segment Tree AtomicAdd Delta Propagation

**Context:** `seg_tree_update` and `seg_tree_insert` kernels in `crates/ml-dqn/src/seg_tree_kernel.cu` use non-atomic `tree[node] = tree[2*node] + tree[2*node+1]` — 8192 threads race on shared internal nodes. Hangs on H100.

**Files:**
- Modify: `crates/ml-dqn/src/seg_tree_kernel.cu:50-61,83-93`

- [ ] **Step 1: Fix `seg_tree_update` propagation**

Replace the propagation loop (lines 53-61):
```cuda
    // Write leaf and propagate sums to root
    int leaf = capacity + (int)idx;
    tree[leaf] = pa;

    int node = leaf >> 1;
    while (node >= 1) {
        tree[node] = tree[2 * node] + tree[2 * node + 1];
        node >>= 1;
    }
```
with atomicAdd delta propagation:
```cuda
    // Write leaf and propagate delta to root via atomicAdd.
    // Each thread computes delta = new_leaf - old_leaf and adds it to every
    // ancestor. atomicAdd is commutative+associative, so concurrent threads
    // produce correct sums without synchronization barriers.
    int leaf = capacity + (int)idx;
    float old_pa = tree[leaf];
    tree[leaf] = pa;
    float delta = pa - old_pa;

    int node = leaf >> 1;
    while (node >= 1) {
        atomicAdd(&tree[node], delta);
        node >>= 1;
    }
```

- [ ] **Step 2: Fix `seg_tree_insert` propagation**

Same fix for `seg_tree_insert` (lines 85-93):
```cuda
    // Write leaf and propagate delta to root via atomicAdd (race-free).
    int leaf = capacity + (int)idx;
    float old_pa = tree[leaf];
    tree[leaf] = pa;
    float delta = pa - old_pa;

    int node = leaf >> 1;
    while (node >= 1) {
        atomicAdd(&tree[node], delta);
        node >>= 1;
    }
```

- [ ] **Step 3: Build and test**

Run: `SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | tail -5`
Expected: Clean compilation.

Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep 'test result'`
Expected: `test result: ok. 19 passed; 0 failed;`

- [ ] **Step 4: Commit**

```bash
git add crates/ml-dqn/src/seg_tree_kernel.cu
git commit -m "fix: race-free segment tree propagation via atomicAdd deltas

seg_tree_update and seg_tree_insert used non-atomic tree[node] =
tree[2*node] + tree[2*node+1] with 8192 threads racing on shared
internal nodes. On H100 (132 SMs), all threads execute simultaneously
causing data races and hangs.

Replace with atomicAdd delta propagation: each thread computes
delta = new_leaf - old_leaf, then atomicAdd to every ancestor.
Commutative+associative = race-free. O(log n) per thread."
```

---

### Task 2: Non-Blocking Readback + Async Diversity Loss

**Context:** `replay_adam_and_readback()` has `event.synchronize()` every step. `readback_diversity_loss()` has `cuStreamSynchronize` per epoch. Both are unnecessary — the CPU readback values are for logging, not control flow.

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`

- [ ] **Step 1: Make readback non-blocking**

In `gpu_dqn_trainer.rs`, find the `replay_adam_and_readback()` method. Locate the section that checks the previous readback event:

```rust
if let Some(ref event) = self.readback_event {
    if !event.is_complete() {
        event.synchronize()...
    }
}
```

Replace with a non-blocking check that reads stale values if the event isn't complete:

```rust
// Non-blocking: if previous readback isn't ready, use stale values.
// Training guard reads loss/grad_norm from GPU buffers directly —
// these CPU scalars are for logging only.
if let Some(ref event) = self.readback_event {
    if event.is_complete() {
        // Safe to read fresh values from pinned buffer
        self.readback_pending = false;
    }
    // If not complete, we'll read the previous values (still valid)
}
```

Ensure the CPU scalar reads (loss, grad_norm) happen unconditionally from the pinned buffer — they'll be from the previous-or-current step depending on event completion.

- [ ] **Step 2: Async diversity loss readback**

In `fused_training.rs`, find `readback_diversity_loss()`. Replace:
```rust
cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream());
cudarc::driver::sys::cuMemcpyDtoH_v2(host.as_mut_ptr().cast(), ptr, ...);
```

With async DtoH to pinned buffer + double-buffered return:
```rust
unsafe {
    cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
        self.trainer.readback_pinned_offset(9).cast(),
        ptr,
        std::mem::size_of::<f32>(),
        self.stream.cu_stream(),
    );
}
let prev = unsafe { *self.trainer.readback_pinned_offset(9) };
prev
```

Add `pub fn readback_pinned_offset(&self, idx: usize) -> *mut f32` accessor on `GpuDqnTrainer`:
```rust
pub fn readback_pinned_offset(&self, idx: usize) -> *mut f32 {
    unsafe { self.readback_pinned.add(idx) }
}
```

- [ ] **Step 3: Build, test, commit**

Run smoke tests. Commit: "perf: non-blocking readback + async diversity loss — zero sync in hot loop"

---

### Task 3: Spectral Norm Weight View Refactor (eliminates D2D copies)

**Context:** Spectral norm operates on separate per-layer `CudaSlice` weight tensors, then 10× `memcpy_dtod_async` copies them back to the flattened `params_buf`. ~25MB/step wasted.

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

- [ ] **Step 1: Compute weight offsets within params_buf**

The flattened `params_buf` stores all weights contiguously. Each weight matrix has a known offset computed from the parameter layout. Find `compute_param_sizes()` and the flattening logic. Create a `WeightViewOffsets` struct that stores `(offset, rows, cols)` for each weight matrix (W_s1, W_s2, b_s1, b_s2, W_v, W_a, etc.).

- [ ] **Step 2: Replace per-layer CudaSlice with pointer+offset views**

Instead of allocating separate `spec_u_s1`, `spec_v_s1` etc. tensors and copying weights into them, compute raw device pointers as `params_buf_ptr + byte_offset` for each weight matrix. The spectral norm kernels already take raw pointers — just pass the correct offset.

- [ ] **Step 3: Remove all D2D sync copies**

Delete the `sync_w!` macro and all 10 `memcpy_dtod_async` calls in `apply_spectral_norm`. Since spectral norm now operates directly on `params_buf`, no copy is needed.

- [ ] **Step 4: Build, test, commit**

Run smoke tests. Commit: "perf: weight views into params_buf — eliminate 25MB/step D2D copies"

---

### Task 4: Fused Batched Spectral Norm Kernel

**Context:** After Task 3 eliminates D2D copies, spectral norm still launches 6-12 single-block kernels (one per weight matrix). Fuse into 1-2 multi-block launches.

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu`
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`

- [ ] **Step 1: Write batched spectral norm kernel**

Add a kernel `spectral_norm_batched` that takes an array of weight matrix descriptors (pointer, rows, cols, u_vec_ptr, v_vec_ptr) and processes one matrix per block:

```cuda
extern "C" __global__ void spectral_norm_batched(
    const float* __restrict__ params_buf,      // flattened params
    float* __restrict__ u_vecs,                // [num_matrices, max_dim] u vectors
    float* __restrict__ v_vecs,                // [num_matrices, max_dim] v vectors
    const int* __restrict__ offsets,            // [num_matrices] byte offsets into params_buf
    const int* __restrict__ rows,              // [num_matrices]
    const int* __restrict__ cols,              // [num_matrices]
    int num_matrices,
    int power_iterations)
{
    int mat_idx = blockIdx.x;
    if (mat_idx >= num_matrices) return;
    // ... power iteration for this matrix using shared memory
}
```

Launch with `grid=(num_weight_matrices, 1, 1), block=(256, 1, 1)`.

- [ ] **Step 2: Replace per-matrix launches with single batched launch**

In `apply_spectral_norm`, replace 6-12 individual kernel launches with one `spectral_norm_batched` launch. Pre-upload the descriptor arrays (offsets, rows, cols) once during construction.

- [ ] **Step 3: Build, test, commit**

Run smoke tests. Commit: "perf: batched spectral norm — 12 launches → 1"

---

### Task 5: Single-Block Kernel Fusion

**Context:** 9 kernels launch with `grid_dim=(1,1,1)` including `grad_norm_finalize` with literally 1 thread.

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu`
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`

- [ ] **Step 1: Fuse grad_norm_finalize into grad_norm kernel**

`grad_norm_finalize` does `grad_norm_bf16 = __float2bfloat16(sqrtf(grad_norm_f32))` — a single scalar operation launched as a kernel. Fuse this into the last block of the `compute_grad_norm` kernel using `if (blockIdx.x == 0 && threadIdx.x == 0)` after the final reduction.

- [ ] **Step 2: Eliminate remaining single-block kernels**

For each remaining grid=(1,1,1) kernel in the per-step hot path:
- If it's a scalar operation (finalize, cast), fuse into the preceding reduction kernel
- If it's a small reduction, ensure block_dim >= 256 and use proper warp-level primitives
- If it can't be fused, at minimum ensure it's not in the per-step path

- [ ] **Step 3: Build, test, commit**

Run smoke tests. Commit: "perf: fuse single-block finalizer kernels"

---

### Task 6: Eliminate Unnecessary memset_zeros

**Context:** 5 `memset_zeros` calls per step on buffers fully overwritten by the next kernel.

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

- [ ] **Step 1: Audit each memset_zeros**

For each memset in the per-step path:
- `total_loss_buf` (8 bytes) — check if loss kernel initializes or accumulates
- `mse_loss_buf` (8 bytes) — same
- `grad_buf` (~1MB) — backward kernels use atomicAdd, so memset IS needed. Keep.
- `d_value_logits_buf`, `d_adv_logits_buf` — check if gradient kernels overwrite fully or accumulate

- [ ] **Step 2: Remove unnecessary memsets**

For buffers that are fully overwritten (not accumulated via atomicAdd), remove the preceding `memset_zeros`. For scalar buffers (8 bytes), the memset overhead exceeds the data size — fuse initialization into the writing kernel.

- [ ] **Step 3: Build, test, commit**

Run smoke tests. Commit: "perf: eliminate unnecessary per-step memset_zeros"

---

### Task 7: TF32 for Forward GEMMs

**Context:** All cuBLAS GEMMs use `CUBLAS_COMPUTE_32F`. On H100, `CUBLAS_COMPUTE_32F_FAST_TF32` provides 2-3× throughput for forward passes where ~10-bit mantissa precision is sufficient.

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs`
- Modify: `crates/ml/src/cuda_pipeline/batched_backward.rs`

- [ ] **Step 1: Change forward GEMM compute type to TF32**

In `batched_forward.rs`, find all `CUBLAS_COMPUTE_32F` references. Change to `CUBLAS_COMPUTE_32F_FAST_TF32` for the forward pass GEMMs. This affects:
- `CublasForward::forward_online_raw` — online network forward
- `CublasForward::forward_target_raw` — target network forward (if separate)

- [ ] **Step 2: Keep F32 for backward GEMMs**

In `batched_backward.rs`, verify `CUBLAS_COMPUTE_32F` is used for all backward GEMMs. Do NOT change these — gradient computation needs full precision.

- [ ] **Step 3: Build, test, commit**

Run smoke tests. Commit: "perf: TF32 tensor cores for forward GEMMs (2-3× throughput)"

---

### Task 8: Dual-Stream DDQN Parallelization

**Context:** `graph_forward_ddqn` replays on the main stream despite having a separate `double_dqn_stream`. The online forward and DDQN forward could overlap.

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

- [ ] **Step 1: Launch DDQN graph on separate stream**

In `replay_forward()`, the DDQN graph is already captured on `double_dqn_stream`. Change the replay to also use `double_dqn_stream`:

Find where `graph_forward_ddqn` is launched and ensure it uses `double_dqn_stream.cu_stream()` instead of `self.stream.cu_stream()`.

- [ ] **Step 2: Add event synchronization before graph_adam**

After both forward graphs complete, graph_adam needs both results. Record an event on `double_dqn_stream` after DDQN forward, and have the main stream wait on it before launching graph_adam:

```rust
// Record completion of DDQN forward on ddqn stream
let ddqn_done = self.double_dqn_stream.record_event(None)?;
// Main stream waits for DDQN to finish before Adam
self.stream.wait_for_event(&ddqn_done)?;
```

- [ ] **Step 3: Build, test, commit**

Run smoke tests. Commit: "perf: parallel DDQN forward on separate CUDA stream"

---

### Task 9: CUDA Graph Expansion for Spectral Norm

**Context:** After Tasks 3-4 reduce spectral norm to 1-2 kernel launches, capture it as a CUDA graph for zero-overhead replay.

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`

- [ ] **Step 1: Capture spectral norm as `graph_spectral`**

After the first step, capture the spectral norm operation as a CUDA graph. Since spectral norm uses fixed weight pointers (views into params_buf from Task 3), the graph topology is constant across steps.

Add a `graph_spectral: Option<RawCudaGraph>` field to `FusedTrainingContext`. Capture after step 1, replay thereafter.

- [ ] **Step 2: Replay graph on subsequent steps**

In `run_full_step`, replace the direct spectral norm call with graph replay:
```rust
if let Some(ref graph) = self.graph_spectral {
    graph.launch(self.stream.cu_stream())?;
} else {
    // First step: run ungraphed, then capture
    self.trainer.apply_spectral_norm(...)?;
    self.graph_spectral = Some(capture_spectral_graph(...)?);
}
```

- [ ] **Step 3: Build, test, commit**

Run smoke tests. Commit: "perf: CUDA graph for spectral norm — zero launch overhead"

---

### Task 10: Deploy and Validate on H100

**Files:** None (testing only)

- [ ] **Step 1: Run full smoke test suite**

Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep 'test result'`
Expected: `test result: ok. 19 passed; 0 failed;`

- [ ] **Step 2: Push and deploy**

```bash
git push origin main
./scripts/argo-train.sh dqn --baseline --watch
```

- [ ] **Step 3: Monitor H100 training**

Verify:
- `H100_HANG4: step 0 — PER update done` (seg_tree fix works)
- `H100_LOOP: step 1` appears (no hang)
- `Epoch 1` completes
- Training step time in epoch summary (should show improvement)
- Q-stats still reported in epoch logs

- [ ] **Step 4: Compare epoch timing**

Previous best epoch time (before optimizations) vs. current. Expected improvement: ~19% faster per epoch from combined fixes.
