diff --git a/docs/superpowers/plans/2026-03-21-h100-epoch-optimization-phase1.md b/docs/superpowers/plans/2026-03-21-h100-epoch-optimization-phase1.md new file mode 100644 index 000000000..ea47a2309 --- /dev/null +++ b/docs/superpowers/plans/2026-03-21-h100-epoch-optimization-phase1.md @@ -0,0 +1,334 @@ +# 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, // [max_batch_size] + sample_indices_i64: CudaSlice, // [max_batch_size] + sample_indices_u32: CudaSlice, // [max_batch_size] + sample_states: CudaSlice, // [max_batch_size * state_dim] + sample_next_states: CudaSlice, // [max_batch_size * state_dim] + sample_actions: CudaSlice, // [max_batch_size] + sample_rewards: CudaSlice, // [max_batch_size] + sample_dones: CudaSlice, // [max_batch_size] + sample_priorities: CudaSlice, // [max_batch_size] + sample_weights: CudaSlice, // [max_batch_size] + sample_max_weight: CudaSlice, // [1] + total_sum_buf: CudaSlice, // [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::(max_bs)?; +let sample_indices_i64 = stream.alloc_zeros::(max_bs)?; +// ... (all fields above) +let total_sum_buf = stream.alloc_zeros::(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) -> Result, 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) -> Result, 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" +```