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:

// 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

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):

    per_gen_thresholds: CudaFunction,

• Step 2: Compile and load in PERKernels::new()

After the existing kernel loads (line ~237), add:

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

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):

    // 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:

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

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:

// 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:

// 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

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

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

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

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

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]

#[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

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