diff --git a/docs/plans/2026-03-02-gpu-per-sumtree-implementation.md b/docs/plans/2026-03-02-gpu-per-sumtree-implementation.md new file mode 100644 index 000000000..a5c0c5a1d --- /dev/null +++ b/docs/plans/2026-03-02-gpu-per-sumtree-implementation.md @@ -0,0 +1,615 @@ +# GPU-Resident PER Sum-Tree Implementation Plan + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** Move Prioritized Experience Replay entirely to GPU — sampling, priority updates, experience storage — eliminating the last CPU bottleneck in the DQN training loop. + +**Architecture:** A new `GpuReplayBuffer` stores experiences as contiguous GPU tensors in a ring buffer. Sampling uses CUDA parallel prefix-sum + binary search instead of a CPU segment tree. Priority updates happen GPU-side from TD errors (no `to_vec1()` needed). Loss readback uses async DMA to eliminate per-step GPU stalls. + +**Tech Stack:** Rust, cudarc 0.17.3 (via candle-core/cuda), NVRTC for custom kernels, candle-core Tensors for integration. + +**Design doc:** `docs/plans/2026-03-02-gpu-per-sumtree-design.md` + +--- + +### Task 1: GpuBatch struct and BatchSample extension + +**Files:** +- Modify: `crates/ml/src/dqn/replay_buffer_type.rs:14-29` + +**Step 1: Add GpuBatch struct after BatchSample** + +After `BatchSample` (line 29), add: + +```rust +/// Pre-built GPU tensors for a training batch. +/// When present, compute_gradients() uses these directly — no CPU→GPU transfer. +#[cfg(feature = "cuda")] +#[derive(Debug, Clone)] +pub struct GpuBatch { + pub states: candle_core::Tensor, // [batch_size, state_dim] f32 on GPU + pub actions: candle_core::Tensor, // [batch_size] u32 on GPU + pub rewards: candle_core::Tensor, // [batch_size] f32 on GPU + pub next_states: candle_core::Tensor, // [batch_size, state_dim] f32 on GPU + pub dones: candle_core::Tensor, // [batch_size] f32 on GPU (0.0/1.0) + pub weights: candle_core::Tensor, // [batch_size] f32 on GPU (IS weights) + pub indices: candle_core::Tensor, // [batch_size] u32 on GPU (buffer indices) +} +``` + +**Step 2: Add gpu_batch field to BatchSample** + +Change `BatchSample` struct (lines 16-20) to: + +```rust +pub struct BatchSample { + pub experiences: Vec, + pub weights: Vec, + pub indices: Vec, + #[cfg(feature = "cuda")] + pub gpu_batch: Option, +} +``` + +**Step 3: Update BatchSample::uniform() constructor** + +Update `uniform()` (line 25-28) to include `gpu_batch: None`. + +**Step 4: Verify build** + +Run: `SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | head -20` +Expected: Fix any compilation errors from the new field in match arms across the codebase. + +**Step 5: Commit** + +```bash +git add crates/ml/src/dqn/replay_buffer_type.rs +git commit -m "feat(ml): add GpuBatch struct and gpu_batch field to BatchSample" +``` + +--- + +### Task 2: GpuReplayBuffer core struct + +**Files:** +- Create: `crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs` + +**Step 1: Write the failing test** + +At the bottom of the new file, add: + +```rust +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_gpu_replay_buffer_creation() { + // CPU fallback test — no CUDA required + let config = GpuReplayBufferConfig { + capacity: 1000, + state_dim: 51, + alpha: 0.6, + beta_start: 0.4, + beta_max: 1.0, + beta_annealing_steps: 100_000, + epsilon: 1e-6, + }; + let buf = GpuReplayBuffer::new(config, &candle_core::Device::Cpu); + assert!(buf.is_ok()); + let buf = buf.unwrap(); + assert_eq!(buf.len(), 0); + assert_eq!(buf.capacity(), 1000); + } +} +``` + +**Step 2: Implement GpuReplayBuffer struct** + +```rust +//! GPU-Resident Replay Buffer with parallel prefix-sum PER sampling +//! +//! Stores all experience data as contiguous GPU tensors in a ring buffer. +//! Sampling uses prefix-sum + binary search (proportional) or radix sort +//! (rank-based) — all on GPU with zero CPU involvement. + +use candle_core::{Device, Tensor, DType}; +use crate::MLError; + +/// Configuration for GPU replay buffer +#[derive(Debug, Clone)] +pub struct GpuReplayBufferConfig { + pub capacity: usize, + pub state_dim: usize, + pub alpha: f32, + pub beta_start: f32, + pub beta_max: f32, + pub beta_annealing_steps: usize, + pub epsilon: f32, // Small constant added to priorities +} + +/// GPU-resident ring buffer for experience replay +pub struct GpuReplayBuffer { + config: GpuReplayBufferConfig, + device: Device, + + // GPU tensor storage (pre-allocated, ring buffer) + states: Tensor, // [capacity, state_dim] f32 + next_states: Tensor, // [capacity, state_dim] f32 + actions: Tensor, // [capacity] u32 + rewards: Tensor, // [capacity] f32 + dones: Tensor, // [capacity] f32 + priorities: Tensor, // [capacity] f32 + + // Ring buffer state (CPU-side, just indices) + write_cursor: usize, + size: usize, + max_priority: f32, + + // Beta annealing + current_step: usize, +} +``` + +Implement `new()`, `len()`, `capacity()`, `is_empty()`, `can_sample()`. + +The `new()` method pre-allocates zero tensors for all storage on the given device. + +**Step 3: Run test** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib gpu_replay_buffer::tests::test_gpu_replay_buffer_creation` +Expected: PASS + +**Step 4: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs crates/ml/src/cuda_pipeline/mod.rs +git commit -m "feat(ml): GpuReplayBuffer core struct with ring buffer storage" +``` + +--- + +### Task 3: Experience insertion (GPU→GPU ring buffer) + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs` + +**Step 1: Write the failing test** + +```rust +#[test] +fn test_gpu_replay_buffer_insert_and_len() { + let config = GpuReplayBufferConfig { + capacity: 100, state_dim: 4, alpha: 0.6, + beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 1000, epsilon: 1e-6, + }; + let device = Device::Cpu; + let mut buf = GpuReplayBuffer::new(config, &device).unwrap(); + + // Insert a batch of 10 experiences + let states = Tensor::zeros(&[10, 4], DType::F32, &device).unwrap(); + let next_states = Tensor::zeros(&[10, 4], DType::F32, &device).unwrap(); + let actions = Tensor::zeros(&[10], DType::U32, &device).unwrap(); + let rewards = Tensor::ones(&[10], DType::F32, &device).unwrap(); + let dones = Tensor::zeros(&[10], DType::F32, &device).unwrap(); + + buf.insert_batch(&states, &next_states, &actions, &rewards, &dones).unwrap(); + assert_eq!(buf.len(), 10); + + // Insert more, verify ring buffer wrapping + for _ in 0..15 { + buf.insert_batch(&states, &next_states, &actions, &rewards, &dones).unwrap(); + } + assert_eq!(buf.len(), 100); // Capped at capacity +} +``` + +**Step 2: Implement insert_batch()** + +Method copies incoming tensors into the ring buffer using `Tensor::narrow()` + `Tensor::copy_()` (or slice assignment). New experiences get `max_priority`. Write cursor wraps at capacity. + +Key: When batch wraps around end of ring buffer, split into two copies (end segment + start segment). + +**Step 3: Run test, verify pass** + +**Step 4: Commit** + +```bash +git commit -am "feat(ml): GpuReplayBuffer::insert_batch with ring buffer wrapping" +``` + +--- + +### Task 4: Proportional sampling kernel (prefix-sum + binary search) + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs` + +**Step 1: Write the failing test** + +```rust +#[test] +fn test_gpu_replay_buffer_proportional_sample() { + let config = GpuReplayBufferConfig { + capacity: 100, state_dim: 4, alpha: 0.6, + beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 1000, epsilon: 1e-6, + }; + let device = Device::Cpu; + let mut buf = GpuReplayBuffer::new(config, &device).unwrap(); + + // Fill with 50 experiences (varying rewards for priority diversity) + let states = Tensor::randn(0.0_f32, 1.0, &[50, 4], &device).unwrap(); + let next_states = Tensor::randn(0.0_f32, 1.0, &[50, 4], &device).unwrap(); + let actions = Tensor::zeros(&[50], DType::U32, &device).unwrap(); + let rewards = Tensor::randn(0.0_f32, 1.0, &[50], &device).unwrap(); + let dones = Tensor::zeros(&[50], DType::F32, &device).unwrap(); + buf.insert_batch(&states, &next_states, &actions, &rewards, &dones).unwrap(); + + // Sample batch of 16 + let batch = buf.sample_proportional(16).unwrap(); + assert_eq!(batch.states.dim(0).unwrap(), 16); + assert_eq!(batch.states.dim(1).unwrap(), 4); + assert_eq!(batch.weights.dim(0).unwrap(), 16); + assert_eq!(batch.indices.dim(0).unwrap(), 16); +} +``` + +**Step 2: Implement sample_proportional()** + +For **CPU device** (test fallback): Use Candle tensor ops — `cumsum`, random uniform tensor, `searchsorted` equivalent via manual binary search in Rust. + +For **CUDA device**: Compile NVRTC kernel that does: +1. Generate random values via simple LCG (or use curand if available) +2. Compute cumulative sum of priorities (or maintain incrementally) +3. Binary search per thread: each of B threads finds its index +4. Gather experiences into output tensors +5. Compute IS weights: `w_i = (N * p_i / sum)^(-beta) / max_weight` + +Return `GpuBatch` with all tensors on GPU. + +**Step 3: Run test, verify pass** + +**Step 4: Commit** + +```bash +git commit -am "feat(ml): proportional PER sampling on GPU with prefix-sum" +``` + +--- + +### Task 5: Rank-based sampling kernel + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs` + +**Step 1: Write the failing test** + +```rust +#[test] +fn test_gpu_replay_buffer_rank_sample() { + // Same setup as proportional test + let config = GpuReplayBufferConfig { + capacity: 100, state_dim: 4, alpha: 0.6, + beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 1000, epsilon: 1e-6, + }; + let device = Device::Cpu; + let mut buf = GpuReplayBuffer::new(config, &device).unwrap(); + + let states = Tensor::randn(0.0_f32, 1.0, &[50, 4], &device).unwrap(); + let next_states = Tensor::randn(0.0_f32, 1.0, &[50, 4], &device).unwrap(); + let actions = Tensor::zeros(&[50], DType::U32, &device).unwrap(); + let rewards = Tensor::randn(0.0_f32, 1.0, &[50], &device).unwrap(); + let dones = Tensor::zeros(&[50], DType::F32, &device).unwrap(); + buf.insert_batch(&states, &next_states, &actions, &rewards, &dones).unwrap(); + + let batch = buf.sample_rank_based(16).unwrap(); + assert_eq!(batch.states.dim(0).unwrap(), 16); + assert_eq!(batch.weights.dim(0).unwrap(), 16); +} +``` + +**Step 2: Implement sample_rank_based()** + +For **CPU device**: Sort priorities descending, compute rank probabilities `1/rank^alpha`, normalize, cumulative sum, sample. + +For **CUDA device**: Use `Tensor::sort()` (candle wraps cub sort) → rank probability kernel → sample. + +**Step 3: Run test, verify pass** + +**Step 4: Commit** + +```bash +git commit -am "feat(ml): rank-based PER sampling on GPU with radix sort" +``` + +--- + +### Task 6: GPU-side priority update (no to_vec1) + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs` + +**Step 1: Write the failing test** + +```rust +#[test] +fn test_gpu_replay_buffer_priority_update() { + let config = GpuReplayBufferConfig { + capacity: 100, state_dim: 4, alpha: 0.6, + beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 1000, epsilon: 1e-6, + }; + let device = Device::Cpu; + let mut buf = GpuReplayBuffer::new(config, &device).unwrap(); + + let states = Tensor::randn(0.0_f32, 1.0, &[50, 4], &device).unwrap(); + let next_states = Tensor::randn(0.0_f32, 1.0, &[50, 4], &device).unwrap(); + let actions = Tensor::zeros(&[50], DType::U32, &device).unwrap(); + let rewards = Tensor::randn(0.0_f32, 1.0, &[50], &device).unwrap(); + let dones = Tensor::zeros(&[50], DType::F32, &device).unwrap(); + buf.insert_batch(&states, &next_states, &actions, &rewards, &dones).unwrap(); + + // Sample, then update priorities with fake TD errors + let batch = buf.sample_proportional(16).unwrap(); + let td_errors = Tensor::randn(0.0_f32, 1.0, &[16], &device).unwrap(); + buf.update_priorities_gpu(&batch.indices, &td_errors).unwrap(); + + // Priorities should have changed — sample again and verify different weights + let batch2 = buf.sample_proportional(16).unwrap(); + // Just verify it doesn't crash — statistical tests are in integration tests + assert_eq!(batch2.weights.dim(0).unwrap(), 16); +} +``` + +**Step 2: Implement update_priorities_gpu()** + +Takes `indices: &Tensor` (u32 on GPU) and `td_errors: &Tensor` (f32 on GPU). +Computes `new_priority = |td_error|^alpha + epsilon` via Candle ops (`.abs()`, `.powf()`, `.add()`). +Scatter-writes into `self.priorities` at given indices. +Updates `max_priority` (via `.max()` on the new priorities). + +Key: This is ALL tensor ops — no `to_vec1()`, no CPU involvement. + +**Step 3: Run test, verify pass** + +**Step 4: Commit** + +```bash +git commit -am "feat(ml): GPU-side priority scatter update from TD errors" +``` + +--- + +### Task 7: ReplayBufferType::GpuPrioritized variant + +**Files:** +- Modify: `crates/ml/src/dqn/replay_buffer_type.rs:36-41` + +**Step 1: Add the enum variant** + +Add to `ReplayBufferType` enum (after line 40): + +```rust +#[cfg(feature = "cuda")] +GpuPrioritized(Arc>), +``` + +**Step 2: Implement sample/update/add dispatch for the new variant** + +In `sample()` (line 84): Add `GpuPrioritized` match arm that calls `gpu_buf.lock().sample_proportional(batch_size)` and wraps in `BatchSample` with `gpu_batch: Some(...)`. + +In `update_priorities()` (line 145): Add arm that calls `gpu_buf.lock().update_priorities_gpu()` — but needs `Tensor` inputs. Add a new method `update_priorities_gpu()` on `ReplayBufferType` that takes tensors directly. + +In `add()` / `add_batch()`: For now, these are no-ops for GpuPrioritized — experiences are inserted via `insert_batch()` from the GPU experience collector directly. + +**Step 3: Update Debug impl and all match arms across the file** + +**Step 4: Verify build** + +Run: `SQLX_OFFLINE=true cargo check -p ml --lib` + +**Step 5: Commit** + +```bash +git commit -am "feat(ml): ReplayBufferType::GpuPrioritized variant with dispatch" +``` + +--- + +### Task 8: Wire GpuReplayBuffer into DQN trainer + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/trainer.rs` + +**Step 1: Buffer creation** + +Find where `ReplayBufferType::Prioritized` or `ReplayBufferType::Uniform` is created during trainer init. Add `#[cfg(feature = "cuda")]` block that creates `GpuPrioritized` when device is CUDA and `use_per` is true. + +**Step 2: Experience insertion** + +Find where GPU experience collector inserts experiences (search for `store_experience` or `add_batch` near the collector code). Change to call `gpu_buffer.insert_batch()` with the GPU tensors directly when `GpuPrioritized`. + +**Step 3: compute_gradients() — use GpuBatch** + +In `crates/ml/src/dqn/dqn.rs:2039` (`compute_gradients`), which calls `compute_loss_internal()`: +- Check if `batch.gpu_batch.is_some()` +- If yes: use GPU tensors directly for states, actions, rewards, next_states, dones, weights +- Skip the CPU→GPU tensor construction code (the `Tensor::from_vec` calls) +- Skip `diff.detach().to_vec1()` (line 1961) — TD errors stay on GPU +- Return `GradientResult` with `td_errors: Vec::new()` and add a new field `td_errors_gpu: Option` + +**Step 4: Priority update wiring** + +In `train_step_with_accumulation()` (trainer.rs ~line 3478), when `GpuPrioritized`: +- Collect `td_errors_gpu` tensors from `GradientResult` instead of `Vec` +- Call `buffer.update_priorities_gpu(&indices_tensor, &td_errors_tensor)` instead of CPU `update_priorities()` + +**Step 5: Verify build** + +Run: `SQLX_OFFLINE=true cargo check -p ml --lib` + +**Step 6: Commit** + +```bash +git commit -am "feat(ml): wire GpuReplayBuffer into DQN trainer hot path" +``` + +--- + +### Task 9: Async loss readback + +**Files:** +- Modify: `crates/ml/src/dqn/dqn.rs` (around line 1945-1948) +- Modify: `crates/ml/src/trainers/dqn/trainer.rs` (epoch logging) + +**Step 1: Defer loss.to_scalar()** + +In `compute_loss_internal()` (dqn.rs line 1945-1948), instead of: +```rust +let loss_value = loss_tensor.to_scalar::()?; +``` + +Return `loss_value: None` and keep `loss_tensor` alive. The caller accumulates loss tensors on GPU. + +**Step 2: Batch loss readback at accumulation boundary** + +In `train_step_with_accumulation()`, after the accumulation loop: +- Sum all loss tensors on GPU: `loss_sum = accumulated_losses.sum()?` +- Single `loss_sum.to_scalar::()` for logging +- This is 1 sync per accumulation group instead of N + +**Step 3: Test existing training tests still pass** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- dqn_trainer` +Expected: All pass (CPU fallback path unchanged) + +**Step 4: Commit** + +```bash +git commit -am "perf(ml): batch loss readback — 1 GPU sync per accumulation instead of N" +``` + +--- + +### Task 10: Integration test — full GPU PER training loop + +**Files:** +- Create: `crates/ml/tests/gpu_per_integration_test.rs` + +**Step 1: Write integration test** + +```rust +//! Integration test: GPU-resident PER training loop +//! Verifies that GpuReplayBuffer produces equivalent training dynamics +//! to the CPU PrioritizedReplayBuffer. + +#[tokio::test] +async fn test_gpu_per_training_loop() { + // Create DQN trainer with GpuReplayBuffer + // Run 3 epochs on synthetic data + // Verify: loss decreases, priorities update, no panics + // Compare: GPU PER loss trajectory vs CPU PER (within 10% tolerance) +} +``` + +**Step 2: Write KS distribution test** + +```rust +#[test] +fn test_gpu_per_sampling_distribution_matches_cpu() { + // Sample 10K indices from GPU PER and CPU PER with same priorities + // KS test: distributions should not differ significantly (p > 0.01) +} +``` + +**Step 3: Run tests** + +Run: `SQLX_OFFLINE=true cargo test -p ml --test gpu_per_integration_test` + +**Step 4: Commit** + +```bash +git commit -am "test(ml): GPU PER integration tests — training loop + distribution correctness" +``` + +--- + +### Task 11: Fallback and OOM handling + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs` +- Modify: `crates/ml/src/trainers/dqn/trainer.rs` + +**Step 1: OOM fallback in GpuReplayBuffer::new()** + +Wrap tensor allocations in a `match` that catches CUDA OOM. If allocation fails: +- Log warning: "GPU replay buffer allocation failed, falling back to CPU PER" +- Return `Err(MLError::GpuMemoryError(...))` + +**Step 2: Trainer fallback** + +In trainer init, wrap `GpuReplayBuffer::new()` in: +```rust +match GpuReplayBuffer::new(config, device) { + Ok(buf) => ReplayBufferType::GpuPrioritized(Arc::new(Mutex::new(buf))), + Err(e) => { + warn!("GPU PER failed ({}), falling back to CPU PER", e); + ReplayBufferType::Prioritized(...) + } +} +``` + +**Step 3: Test fallback** + +Write test that forces OOM (absurd capacity) and verifies CPU fallback. + +**Step 4: Commit** + +```bash +git commit -am "feat(ml): GPU PER graceful fallback to CPU on OOM" +``` + +--- + +### Task 12: Clippy, warnings, final verification + +**Files:** All modified files + +**Step 1: Full workspace check** + +```bash +SQLX_OFFLINE=true cargo check --workspace +SQLX_OFFLINE=true cargo clippy --workspace -- -D warnings +``` + +**Step 2: Run all ml tests** + +```bash +SQLX_OFFLINE=true cargo test -p ml --lib +``` +Expected: 2390+ tests pass, 0 warnings + +**Step 3: Commit any fixes** + +```bash +git commit -am "fix(ml): clippy and warning cleanup for GPU PER" +``` + +--- + +### Task 13: Update memory and backlog + +**Files:** +- Modify: `/home/jgrusewski/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/gpu-saturation-backlog.md` +- Modify: `/home/jgrusewski/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/MEMORY.md` + +**Step 1:** Mark "DQN PER CPU sum-tree" as resolved in the backlog. + +**Step 2:** Update MEMORY.md with GPU PER details. + +**Step 3: Push and trigger pipeline** + +```bash +git push origin main +``` + +Trigger `train-validate-rl` job to verify epoch time improvement (target: 31s → ~26s on L40S).