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:

```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<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**

```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<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**

```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<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**

```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::<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**

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