315 lines
9.2 KiB
Markdown
315 lines
9.2 KiB
Markdown
# 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
|
||
```
|