foxhunt/docs/plans/2026-03-02-gpu-per-sumtree-implementation.md

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:

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

pub struct BatchSample {
    pub experiences: Vec<Experience>,
    pub weights: Vec<f32>,
    pub indices: Vec<usize>,
    #[cfg(feature = "cuda")]
    pub gpu_batch: Option<GpuBatch>,
}

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

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:

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

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

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

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

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

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

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

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

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

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

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

#[cfg(feature = "cuda")]
GpuPrioritized(Arc<parking_lot::Mutex<crate::cuda_pipeline::gpu_replay_buffer::GpuReplayBuffer>>),

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

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

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<f32>
• 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

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:

let loss_value = loss_tensor.to_scalar::<f32>()?;

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

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

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

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

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:

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

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

SQLX_OFFLINE=true cargo check --workspace
SQLX_OFFLINE=true cargo clippy --workspace -- -D warnings

Step 2: Run all ml tests

SQLX_OFFLINE=true cargo test -p ml --lib

Expected: 2390+ tests pass, 0 warnings

Step 3: Commit any fixes

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

git push origin main

Trigger train-validate-rl job to verify epoch time improvement (target: 31s → ~26s on L40S).