diff --git a/docs/superpowers/plans/2026-03-22-gpu-segment-tree.md b/docs/superpowers/plans/2026-03-22-gpu-segment-tree.md new file mode 100644 index 000000000..733cdce51 --- /dev/null +++ b/docs/superpowers/plans/2026-03-22-gpu-segment-tree.md @@ -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[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, +cs_buf: CudaSlice, +block_sums: CudaSlice, +``` + +Add: +```rust +seg_tree: CudaSlice, // [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::(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 { + // 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 +``` diff --git a/docs/superpowers/specs/2026-03-22-gpu-segment-tree-design.md b/docs/superpowers/specs/2026-03-22-gpu-segment-tree-design.md new file mode 100644 index 000000000..fda0f2504 --- /dev/null +++ b/docs/superpowers/specs/2026-03-22-gpu-segment-tree-design.md @@ -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` 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` → removed (no power-alpha buffer needed) +- `cs_buf: CudaSlice` → removed (no cumulative sum buffer) +- `block_sums: CudaSlice` → 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` — `[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)