# 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)