110 lines
4.2 KiB
Markdown
110 lines
4.2 KiB
Markdown
# 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)
|