foxhunt/docs/superpowers/plans/2026-03-21-h100-epoch-optimization-phase1.md

# H100 Epoch Optimization Phase 1: Eliminate CPU Roundtrips

> **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:** Eliminate all CPU roundtrips in the PER sampling hot path, reducing DQN training epoch time from 37s to ~12-15s on H100.

**Architecture:** Replace CPU-side random threshold generation + DtoH total_sum readback with a GPU Philox RNG kernel. Pre-allocate all PER sampling buffers to eliminate per-step `cuMemAlloc`. Keep the total_sum on GPU as a buffer pointer arg instead of a scalar.

**Tech Stack:** Rust, cudarc 0.19, CUDA (nvcc), Philox 4×32 PRNG

**Spec:** `docs/superpowers/specs/2026-03-21-h100-epoch-optimization-design.md`

---

### Task 1: Add Philox threshold generation kernel to replay buffer kernels

**Files:**
- Create: `crates/ml-dqn/src/per_threshold_kernel.cu`
- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs:195-240` (add kernel to PERKernels struct)

- [ ] **Step 1: Write the CUDA kernel**

Create `crates/ml-dqn/src/per_threshold_kernel.cu`:

```cuda
// Philox 4×32-10 PRNG for PER threshold generation.
// Generates batch_size uniform random values in [0, total_sum).
// Zero CPU involvement — total_sum read from GPU buffer.

__device__ __forceinline__
unsigned int philox_single(unsigned int counter, unsigned int key) {
    unsigned int hi, lo;
    lo = counter * 0xD2511F53u;
    hi = __umulhi(counter, 0xD2511F53u);
    for (int r = 0; r < 10; r++) {
        unsigned int t = hi ^ key;
        hi = lo * 0xCD9E8D57u;
        lo = __umulhi(lo, 0xCD9E8D57u);
        lo ^= t;
        key += 0x9E3779B9u;
    }
    return lo ^ hi;
}

extern "C" __global__
void per_generate_thresholds(
    float* __restrict__ thresholds,
    const float* __restrict__ total_sum_buf,  // [1] GPU scalar
    unsigned int seed,
    int batch_size)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;
    float total_sum = total_sum_buf[0];
    unsigned int bits = philox_single((unsigned int)i, seed);
    float u = (float)(bits >> 8) * (1.0f / 16777216.0f);
    thresholds[i] = u * total_sum;
}
```

- [ ] **Step 2: Verify kernel file exists**

Run: `ls -la crates/ml-dqn/src/per_threshold_kernel.cu`

- [ ] **Step 3: Commit**

```bash
git add crates/ml-dqn/src/per_threshold_kernel.cu
git commit -m "feat: Philox PRNG kernel for GPU-native PER threshold generation"
```

---

### Task 2: Add kernel to PERKernels struct and compile at init

**Files:**
- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs:185-240` (PERKernels struct + compilation)

- [ ] **Step 1: Add field to PERKernels**

In `gpu_replay_buffer.rs`, add to the `PERKernels` struct (after line ~208):
```rust
    per_gen_thresholds: CudaFunction,
```

- [ ] **Step 2: Compile and load in PERKernels::new()**

After the existing kernel loads (line ~237), add:
```rust
let per_thresh_ptx = crate::compile_ptx_for_device(
    include_str!("per_threshold_kernel.cu"), ctx,
)?;
let per_thresh_mod = ctx.load_module(per_thresh_ptx)?;
let per_gen_thresholds = per_thresh_mod.load_function("per_generate_thresholds")?;
```

And add `per_gen_thresholds` to the struct initialization.

- [ ] **Step 3: Verify compilation**

Run: `SQLX_OFFLINE=true cargo check -p ml-dqn 2>&1 | tail -3`
Expected: no errors

- [ ] **Step 4: Commit**

```bash
git add crates/ml-dqn/src/gpu_replay_buffer.rs
git commit -m "feat: compile Philox threshold kernel in PERKernels init"
```

---

### Task 3: Pre-allocate PER sampling buffers

**Files:**
- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs` (GpuReplayBuffer struct + new())

Currently `sample_proportional` allocates ~10 buffers per call via `a32f()`, `a16()`, `a32u()`. These must be pre-allocated at construction and reused.

- [ ] **Step 1: Add pre-allocated fields to GpuReplayBuffer**

Add to the struct (after existing fields):
```rust
    // Pre-allocated PER sampling buffers (reused across sample() calls)
    sample_thresholds: CudaSlice<f32>,     // [max_batch_size]
    sample_indices_i64: CudaSlice<i64>,    // [max_batch_size]
    sample_indices_u32: CudaSlice<u32>,    // [max_batch_size]
    sample_states: CudaSlice<u16>,         // [max_batch_size * state_dim]
    sample_next_states: CudaSlice<u16>,    // [max_batch_size * state_dim]
    sample_actions: CudaSlice<u32>,        // [max_batch_size]
    sample_rewards: CudaSlice<f32>,        // [max_batch_size]
    sample_dones: CudaSlice<f32>,          // [max_batch_size]
    sample_priorities: CudaSlice<f32>,     // [max_batch_size]
    sample_weights: CudaSlice<f32>,        // [max_batch_size]
    sample_max_weight: CudaSlice<f32>,     // [1]
    total_sum_buf: CudaSlice<f32>,         // [1] — holds cs_buf[n-1] via DtoD
    rng_step: u32,                          // counter for Philox seed
```

- [ ] **Step 2: Allocate in constructor**

In `GpuReplayBuffer::new()`, after existing allocations:
```rust
let max_bs = config.max_batch_size.unwrap_or(1024);
let sd = config.state_dim;
let sample_thresholds = stream.alloc_zeros::<f32>(max_bs)?;
let sample_indices_i64 = stream.alloc_zeros::<i64>(max_bs)?;
// ... (all fields above)
let total_sum_buf = stream.alloc_zeros::<f32>(1)?;
let rng_step = 0_u32;
```

- [ ] **Step 3: Verify compilation**

Run: `SQLX_OFFLINE=true cargo check -p ml-dqn 2>&1 | tail -3`

- [ ] **Step 4: Commit**

```bash
git add crates/ml-dqn/src/gpu_replay_buffer.rs
git commit -m "feat: pre-allocate all PER sampling buffers (zero cuMemAlloc per step)"
```

---

### Task 4: Rewrite sample_proportional to use GPU RNG + pre-allocated buffers

**Files:**
- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs:414-514` (sample_proportional)

This is the core change. Replace:
1. `cs_total()` DtoH readback → DtoD copy of `cs_buf[n-1]` to `total_sum_buf`
2. CPU `rand::thread_rng()` → `per_generate_thresholds` kernel launch
3. Per-call `a32f()`/`a16()`/`a32u()` allocs → pre-allocated buffers

- [ ] **Step 1: Replace cs_total with DtoD copy**

Replace lines 429-430:
```rust
// OLD: let ts = self.cs_total(n)?;
// NEW: DtoD copy of cs_buf[n-1] to total_sum_buf (zero CPU)
let cs_last = self.cs_buf.slice((n-1)..n);
self.stream.memcpy_dtod(&cs_last, &mut self.total_sum_buf)?;
```

- [ ] **Step 2: Replace CPU RNG with GPU kernel**

Replace lines 431-435:
```rust
// OLD: CPU rand + HtoD upload
// NEW: GPU Philox kernel
self.rng_step = self.rng_step.wrapping_add(1);
let seed = self.rng_step;
let blocks = ((batch_size + 255) / 256) as u32;
unsafe {
    self.stream.launch_builder(&self.kernels.per_gen_thresholds)
        .arg(&mut self.sample_thresholds)
        .arg(&self.total_sum_buf)
        .arg(&seed)
        .arg(&bsi)
        .launch(LaunchConfig { grid_dim: (blocks,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 })?;
}
```

- [ ] **Step 3: Replace all per-call allocations with pre-allocated buffers**

Replace `a32f(&self.stream, batch_size, "t")?` etc. with `self.sample_thresholds`, `self.sample_indices_i64`, etc. throughout the method. Update all kernel `.arg()` calls to reference `self.sample_*` fields.

- [ ] **Step 4: Update is_weights to read total_sum from GPU buffer**

The `is_weights_f32` kernel currently takes `ts: f32` as scalar arg. Change to `total_sum_buf: *const f32` pointer arg. Update kernel source to read `total_sum_buf[0]`.

- [ ] **Step 5: Remove cs_total() method**

Delete the `fn cs_total()` method entirely — it was the DtoH readback.

- [ ] **Step 6: Verify compilation**

Run: `SQLX_OFFLINE=true cargo check -p ml-dqn 2>&1 | tail -3`

- [ ] **Step 7: Run tests**

Run: `SQLX_OFFLINE=true cargo test -p ml-dqn --lib -- --test-threads=2 2>&1 | tail -3`
Expected: 359 passed, 0 failed

- [ ] **Step 8: Commit**

```bash
git add crates/ml-dqn/src/gpu_replay_buffer.rs
git commit -m "perf: GPU-native PER sampling — zero CPU roundtrips per training step"
```

---

### Task 5: Remove HtoD for adam_step counter

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:790-793,855-856`

- [ ] **Step 1: Add GPU-side adam_step increment kernel**

Add a tiny kernel or use `cuMemsetD32` to increment the step counter on GPU instead of `memcpy_htod(&[self.adam_step], &mut self.t_buf)`.

Alternative (simpler): keep the HtoD but document it as 4-byte cold-path. The adam_step is 1 i32 per training step — negligible vs the PER fix. Mark as `// gpu-entry: 4B adam_step (acceptable — 1 scalar/step)`.

- [ ] **Step 2: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "perf: document adam_step HtoD as acceptable 4B entry point"
```

---

### Task 6: Profile and validate

**Files:**
- Read: training_loop.rs phase breakdown output

- [ ] **Step 1: Run on RTX 3050 with phase profiling**

```bash
SQLX_OFFLINE=true cargo run --release --example train_baseline_rl -p ml -- \
  --model dqn --data-dir test_data/futures-baseline --symbol ES.FUT \
  --epochs 3 --max-steps-per-epoch 50 --batch-size 64 \
  --train-months 3 --val-months 1 --test-months 1 --step-months 3 \
  2>&1 | grep "phase breakdown"
```

Expected: training phase time reduced by 2-3× compared to baseline 14,858ms.

- [ ] **Step 2: Run full test suite**

```bash
SQLX_OFFLINE=true cargo test -p ml-core --lib && \
SQLX_OFFLINE=true cargo test -p ml-dqn --lib && \
SQLX_OFFLINE=true cargo test -p ml --lib && \
SQLX_OFFLINE=true cargo test -p ml --test smoke_test_real_data -- --test-threads=1
```

Expected: all pass, 0 failures.

- [ ] **Step 3: Push and run H100 CI**

```bash
git push origin main
argo submit --from workflowtemplate/gpu-test-pipeline -p revision=main -n foxhunt
```

Expected: phase breakdown shows training < 15s (down from 37.56s).

- [ ] **Step 4: Commit profiling results to spec**

Update `docs/superpowers/specs/2026-03-21-h100-epoch-optimization-design.md` with actual measured speedup.

---

### Task 7: Deprecate to_host() and add guard

**Files:**
- Modify: `crates/ml-core/src/cuda_autograd/gpu_tensor.rs` (deprecate to_host)
- Modify: `crates/ml-core/src/cuda_autograd/stream_ops.rs` (deprecate to_vec)
- Modify: `scripts/gpu-hotpath-guard.sh` (add to_host pattern)

- [ ] **Step 1: Rename to_host → to_host_checkpoint, add #[deprecated]**

```rust
#[deprecated(note = "Use to_host_checkpoint() for cold-path exports only. Hot-path DtoH is forbidden.")]
pub fn to_host(&self, stream: &Arc<CudaStream>) -> Result<Vec<f32>, MLError> {
    self.to_host_checkpoint(stream)
}

/// Cold-path only: download tensor to host for checkpoint export.
/// NOT for use in training hot paths.
pub fn to_host_checkpoint(&self, stream: &Arc<CudaStream>) -> Result<Vec<f32>, MLError> {
    // ... existing implementation
}
```

- [ ] **Step 2: Add to_host to gpu-hotpath-guard.sh**

Add `\.to_host(` to the HOT_PATH patterns in `scripts/gpu-hotpath-guard.sh`.

- [ ] **Step 3: Fix all existing callers to use to_host_checkpoint**

Update all legitimate callers (checkpoint save, test assertions, epoch-boundary metrics) to use `to_host_checkpoint()`.

- [ ] **Step 4: Run full test suite**

- [ ] **Step 5: Commit**

```bash
git commit -m "perf: deprecate to_host(), add to_host_checkpoint() for cold paths only"
```