# GPU Segment Tree for PER Sampling

> **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:** Replace O(n) prefix sum PER sampling with O(log n) GPU segment tree, eliminating per-epoch training time growth on H100 (12ms→23ms/step).

**Architecture:** Binary segment tree as flat `CudaSlice<f32>[2*capacity_pow2]`. Two CUDA kernels: `seg_tree_update` (batch leaf write + propagate up) and `seg_tree_sample` (parallel root-to-leaf traversal with Philox RNG). Replaces prefix_sum + searchsorted pipeline entirely.

**Tech Stack:** Rust, cudarc 0.19, CUDA (NVRTC), Philox 4×32 PRNG

**Spec:** `docs/superpowers/specs/2026-03-22-gpu-segment-tree-design.md`

---

## Files

### Create
| File | Responsibility |
|------|---------------|
| `crates/ml-dqn/src/seg_tree_kernel.cu` | `seg_tree_update` + `seg_tree_sample` CUDA kernels |

### Modify
| File | Changes |
|------|---------|
| `crates/ml-dqn/src/gpu_replay_buffer.rs` | Replace prefix sum with segment tree in struct, constructor, sample, update |

### Delete
| File | Reason |
|------|--------|
| `crates/ml-dqn/src/prefix_sum_kernel.cu` | Replaced by seg_tree_kernel.cu |

---

### Task 1: Create the segment tree CUDA kernels

**Files:**
- Create: `crates/ml-dqn/src/seg_tree_kernel.cu`

- [ ] **Step 1: Write `seg_tree_update` kernel**

```cuda
// Batch update: write priorities^alpha to leaves, propagate sums to root.
// One thread per index. Each thread writes its leaf then walks up to root.
// Concurrent writes to shared ancestors use atomicExch for correctness.
extern "C" __global__ void seg_tree_update(
    float* __restrict__ tree,
    const unsigned int* __restrict__ indices,
    const float* __restrict__ td_errors,
    float alpha, float epsilon,
    int capacity, int batch_size)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    int idx = indices[i];
    if (idx < 0 || idx >= capacity) return;

    // Compute priority = (|td_error| + epsilon)^alpha
    float td = td_errors[i];
    float priority = powf(fabsf(td) + epsilon, alpha);

    // Write leaf
    int leaf = capacity + idx;
    tree[leaf] = priority;

    // Propagate up to root
    int node = leaf >> 1;
    while (node >= 1) {
        tree[node] = tree[2 * node] + tree[2 * node + 1];
        node >>= 1;
    }
}

// Batch insert: write raw priorities (already computed) to leaves + propagate.
extern "C" __global__ void seg_tree_insert(
    float* __restrict__ tree,
    const unsigned int* __restrict__ indices,
    const float* __restrict__ priorities,
    int capacity, int batch_size)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    int idx = indices[i];
    if (idx < 0 || idx >= capacity) return;

    int leaf = capacity + idx;
    tree[leaf] = priorities[i];

    int node = leaf >> 1;
    while (node >= 1) {
        tree[node] = tree[2 * node] + tree[2 * node + 1];
        node >>= 1;
    }
}

// Proportional sampling: parallel root-to-leaf traversal.
// Each thread generates a random threshold and traverses the tree.
extern "C" __global__ void seg_tree_sample(
    const float* __restrict__ tree,
    int* __restrict__ out_indices,
    unsigned int* __restrict__ rng_states,
    int capacity, int batch_size)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    // Philox RNG: generate uniform in [0, total_sum)
    unsigned int state = rng_states[i];
    unsigned int hi, lo;
    lo = state * 0xD2511F53u;
    hi = __umulhi(state, 0xD2511F53u);
    for (int r = 0; r < 10; r++) {
        unsigned int t = hi ^ state;
        hi = lo * 0xCD9E8D57u;
        lo = __umulhi(lo, 0xCD9E8D57u);
        lo ^= t;
        state += 0x9E3779B9u;
    }
    unsigned int bits = lo ^ hi;
    float u = (float)(bits >> 8) * (1.0f / 16777216.0f);
    float total_sum = tree[1];  // root = total priority sum
    float threshold = u * total_sum;

    // Update RNG state
    rng_states[i] = state;

    // Tree traversal: root → leaf
    int node = 1;
    while (node < capacity) {
        float left_sum = tree[2 * node];
        if (threshold <= left_sum) {
            node = 2 * node;
        } else {
            threshold -= left_sum;
            node = 2 * node + 1;
        }
    }

    // node is now a leaf: index = node - capacity
    int idx = node - capacity;
    // Clamp to valid range (tree may have empty leaves beyond buffer size)
    if (idx < 0) idx = 0;
    out_indices[i] = idx;
}
```

- [ ] **Step 2: Commit**

```bash
git add crates/ml-dqn/src/seg_tree_kernel.cu
git commit -m "feat: GPU segment tree CUDA kernels for O(log n) PER sampling"
```

---

### Task 2: Replace prefix sum with segment tree in GpuReplayBuffer

**Files:**
- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs`

This is the main task. Read the full file first, then make these changes:

- [ ] **Step 1: Replace struct fields**

Remove:
```rust
pa_buf: CudaSlice<f32>,
cs_buf: CudaSlice<f32>,
block_sums: CudaSlice<f32>,
```

Add:
```rust
seg_tree: CudaSlice<f32>,      // [2 * capacity_pow2] segment tree
capacity_pow2: usize,           // capacity rounded up to power of 2
```

- [ ] **Step 2: Replace kernel fields in ReplayKernels**

Remove from `ReplayKernels`:
```rust
pow_alpha_f32: CudaFunction,
prefix_sum: CudaFunction,
prefix_sum_local: CudaFunction,
prefix_sum_block_sums: CudaFunction,
prefix_sum_propagate: CudaFunction,
searchsorted: CudaFunction,
```

Add:
```rust
seg_tree_update: CudaFunction,
seg_tree_insert: CudaFunction,
seg_tree_sample: CudaFunction,
```

- [ ] **Step 3: Update kernel compilation**

In `ReplayKernels::new()`, compile `seg_tree_kernel.cu` instead of `prefix_sum_kernel.cu`. Load the 3 new functions. Remove the old prefix sum + pow_alpha + searchsorted loads.

- [ ] **Step 4: Update constructor**

In `GpuReplayBuffer::new()`:
- Compute `capacity_pow2 = config.capacity.next_power_of_two()`
- Allocate `seg_tree = stream.alloc_zeros::<f32>(2 * capacity_pow2)`
- Remove `pa_buf`, `cs_buf`, `block_sums` allocations
- Initialize `rng_states` buffer for Philox (reuse existing `rng_step` or add per-sample RNG states)

- [ ] **Step 5: Rewrite `sample_proportional()`**

Replace the entire method body:
```rust
pub fn sample_proportional(&mut self, batch_size: usize) -> Result<GpuBatchSlices, MLError> {
    // 1. Launch seg_tree_sample kernel (root-to-leaf traversal)
    // 2. Use sampled indices to gather states, actions, rewards, dones, priorities
    // 3. Compute IS weights from priorities + total_sum (tree[1])
    // Return GpuBatchSlices
}
```

The key difference: no prefix sum, no searchsorted. Just 1 kernel launch for sampling + gather kernels.

- [ ] **Step 6: Rewrite `update_priorities_gpu()`**

Replace prefix scatter + rebuild with:
```rust
pub fn update_priorities_gpu(&mut self, indices, td_errors, bs) -> Result<(), MLError> {
    // Launch seg_tree_update kernel: writes new priorities to leaves, propagates up
    // Batch max accumulation stays the same (atomicMax or separate kernel)
}
```

- [ ] **Step 7: Update `insert_batch()`**

After writing priorities to the ring buffer, also update the segment tree leaves:
```rust
// After scatter_insert of priorities:
// Launch seg_tree_insert kernel on the inserted indices
```

- [ ] **Step 8: Remove `pfx_sum()` method**

Delete the entire `pfx_sum()` method and all prefix sum infrastructure.

- [ ] **Step 9: Update `total_priority()` / `cs_total()`**

Total priority sum is now `tree[1]` (root node). Replace any remaining `cs_total()` or total_sum readback with a DtoD copy from `seg_tree[1]`.

- [ ] **Step 10: Verify compilation + tests**

Run: `SQLX_OFFLINE=true cargo check -p ml-dqn`
Run: `SQLX_OFFLINE=true cargo test -p ml-dqn --lib`
Expected: 359 passed, 0 failed

- [ ] **Step 11: Commit**

```bash
git add crates/ml-dqn/src/gpu_replay_buffer.rs
git commit -m "perf: replace O(n) prefix sum with O(log n) GPU segment tree for PER"
```

---

### Task 3: Delete old prefix sum kernel + validate

**Files:**
- Delete: `crates/ml-dqn/src/prefix_sum_kernel.cu`

- [ ] **Step 1: Verify no remaining references**

```bash
grep -rn "prefix_sum_kernel\|prefix_sum_local\|prefix_sum_block_sums\|prefix_sum_propagate" crates/ml-dqn/src/
```

- [ ] **Step 2: Delete**

```bash
rm crates/ml-dqn/src/prefix_sum_kernel.cu
```

- [ ] **Step 3: Run full test suite**

```bash
SQLX_OFFLINE=true cargo test -p ml-dqn --lib
SQLX_OFFLINE=true cargo test -p ml --lib
```

- [ ] **Step 4: Run GPU smoke test**

```bash
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline \
  cargo test -p ml --lib -- test_train_step_produces_finite_metrics --ignored --test-threads=1
```

- [ ] **Step 5: Profile locally**

```bash
SQLX_OFFLINE=true cargo run --release --example train_baseline_rl -p ml -- \
  --model dqn --data-dir test_data/futures-baseline --symbol ES.FUT \
  --training-profile dqn-production --epochs 5 --max-steps-per-epoch 50 \
  --train-months 3 --val-months 1 --test-months 1 --step-months 3 \
  2>&1 | grep "step breakdown"
```

Expected: `sample` time should be CONSTANT across epochs (not growing).

- [ ] **Step 6: Commit + push**

```bash
git add -A crates/ml-dqn/src/
git commit -m "refactor: delete prefix_sum_kernel.cu (replaced by segment tree)"
git push origin main
```