docs: GPU segment tree spec + plan for O(log n) PER sampling
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
314
docs/superpowers/plans/2026-03-22-gpu-segment-tree.md
Normal file
314
docs/superpowers/plans/2026-03-22-gpu-segment-tree.md
Normal file
@@ -0,0 +1,314 @@
|
||||
# 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
|
||||
```
|
||||
109
docs/superpowers/specs/2026-03-22-gpu-segment-tree-design.md
Normal file
109
docs/superpowers/specs/2026-03-22-gpu-segment-tree-design.md
Normal file
@@ -0,0 +1,109 @@
|
||||
# GPU Segment Tree for PER Sampling
|
||||
|
||||
## Goal
|
||||
|
||||
Replace O(n) prefix sum in PER sampling with O(log n) GPU segment tree. Eliminates the per-epoch training time growth on H100 (12ms→23ms/step as buffer fills from 128K to 500K).
|
||||
|
||||
## Problem
|
||||
|
||||
`sample_proportional()` recomputes `prefix_sum(priorities[0..n])` every call — O(n) work. As the replay buffer fills, this grows linearly. On H100 with 500K buffer: sampling takes 5ms at epoch 2, growing to 23ms at epoch 5. This is 40% of training step time.
|
||||
|
||||
## Architecture
|
||||
|
||||
Binary segment tree stored as a flat `CudaSlice<f32>` of size `2 * capacity`. Internal nodes store sums of children. Leaves store priorities. Two CUDA kernels: `seg_tree_update` (batch leaf update + propagate) and `seg_tree_sample` (parallel root-to-leaf traversal).
|
||||
|
||||
### Tree Layout
|
||||
|
||||
```
|
||||
Array: tree[0..2*capacity] (index 0 unused, root at index 1)
|
||||
|
||||
[1] root = total priority sum
|
||||
/ \
|
||||
[2] [3]
|
||||
/ \ / \
|
||||
[4] [5] [6] [7]
|
||||
/ \ / \ / \ / \
|
||||
[8][9][10][11][12][13][14][15] ← leaves = priorities[0..capacity]
|
||||
|
||||
Leaf for priority[i] = tree[capacity + i]
|
||||
Parent of node j = j / 2
|
||||
Left child of j = 2*j, right child = 2*j + 1
|
||||
```
|
||||
|
||||
Capacity is rounded up to next power of 2 for balanced tree.
|
||||
|
||||
### Operations
|
||||
|
||||
**Update** — O(log₂(capacity) × batch_size), 1 kernel launch:
|
||||
- Write new priority to leaf: `tree[capacity + idx] = priority`
|
||||
- Propagate up: `tree[parent] = tree[left] + tree[right]` until root
|
||||
- Each thread handles one index, all threads independent
|
||||
|
||||
**Sample** — O(log₂(capacity) × batch_size), 1 kernel launch:
|
||||
- Generate random threshold in [0, tree[1]) via Philox RNG
|
||||
- Traverse root → leaf: go left if threshold ≤ left child sum, else subtract left and go right
|
||||
- Each thread handles one sample, all threads independent
|
||||
|
||||
### CUDA Kernels
|
||||
|
||||
```cuda
|
||||
// Batch update: write new priorities to leaves, propagate sums to root.
|
||||
// Grid: (batch_size, 1, 1), Block: (1, 1, 1) — 1 thread per update.
|
||||
extern "C" __global__ void seg_tree_update(
|
||||
float* __restrict__ tree,
|
||||
const unsigned int* __restrict__ indices,
|
||||
const float* __restrict__ new_priorities,
|
||||
int capacity, int batch_size)
|
||||
|
||||
// Proportional sampling: parallel root-to-leaf traversal.
|
||||
// Grid: ceil(batch_size/256), Block: (256, 1, 1).
|
||||
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)
|
||||
```
|
||||
|
||||
### Memory
|
||||
|
||||
- Segment tree: `2 * capacity_pow2 * sizeof(f32)` bytes
|
||||
- 500K capacity → pow2 = 524288 → 4MB
|
||||
- 100K capacity → pow2 = 131072 → 1MB
|
||||
- Replaces: `pa_buf` + `cs_buf` + `block_sums` ≈ same total
|
||||
|
||||
### Integration with GpuReplayBuffer
|
||||
|
||||
**Replace:**
|
||||
- `pa_buf: CudaSlice<f32>` → removed (no power-alpha buffer needed)
|
||||
- `cs_buf: CudaSlice<f32>` → removed (no cumulative sum buffer)
|
||||
- `block_sums: CudaSlice<f32>` → removed
|
||||
- `pfx_sum()` method → removed
|
||||
- `prefix_sum_local`, `prefix_sum_block_sums`, `prefix_sum_propagate` kernels → removed
|
||||
- `pow_alpha_f32` kernel → integrated into `seg_tree_update`
|
||||
- `searchsorted` kernel → replaced by `seg_tree_sample` traversal
|
||||
|
||||
**Add:**
|
||||
- `seg_tree: CudaSlice<f32>` — `[2 * capacity_pow2]`
|
||||
- `capacity_pow2: usize` — capacity rounded to next power of 2
|
||||
- `seg_tree_update_kernel: CudaFunction`
|
||||
- `seg_tree_sample_kernel: CudaFunction`
|
||||
|
||||
**Modified methods:**
|
||||
- `insert_batch()` → after writing priorities, call `seg_tree_update` on inserted indices
|
||||
- `sample_proportional()` → call `seg_tree_sample` instead of prefix_sum + searchsorted
|
||||
- `update_priorities_gpu()` → call `seg_tree_update` instead of scatter + pfx rebuild
|
||||
- `total_priority()` → read `tree[1]` (root = total sum, O(1))
|
||||
|
||||
### Expected Performance
|
||||
|
||||
| Operation | Current (O(n)) | Segment tree (O(log n)) |
|
||||
|-----------|----------------|------------------------|
|
||||
| Sample 1024 from 500K | 5-23ms (growing) | ~0.1ms (constant) |
|
||||
| Update 1024 priorities | 0.5ms + rebuild | ~0.1ms |
|
||||
| Total per training step | 12-28ms | ~5ms (constant) |
|
||||
|
||||
### Non-Goals
|
||||
|
||||
- Rank-based PER (separate code path, not affected)
|
||||
- CPU replay buffer (separate implementation)
|
||||
- Changing the PER algorithm itself (alpha, beta, IS weights)
|
||||
Reference in New Issue
Block a user