# 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).