docs: Phase 1 implementation plan — eliminate PER CPU roundtrips
7 tasks: Philox RNG kernel, pre-allocated buffers, GPU-native sampling, adam_step cleanup, profiling validation, to_host deprecation. Target: 37s → 12-15s/epoch on H100. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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<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"
|
||||
```
|
||||
Reference in New Issue
Block a user