# PER Prefix Sum — Replace Segment Tree with Decoupled Lookback Scan

## Problem

The PER segment tree (`seg_tree: CudaSlice<f32>`, 67M floats = 256MB) hangs on H100 during priority updates. The tree structure requires either:
- atomicAdd propagation from leaves to root (8192 threads contending on root = minutes of L2 cache serialization)
- Full tree rebuild (33M nodes × 25 levels = seconds of wasted compute for 8192 changed leaves)

Both approaches fail at H100 scale because the tree was designed for CPU-era single-element updates, not massively parallel GPU workloads.

## Solution

Replace the segment tree with a flat prefix sum array using NVIDIA's decoupled lookback scan algorithm (Merrill & Garland, 2016). Single-pass O(N) parallel scan + O(B log N) binary search sampling in 2 kernel launches. ~42µs per PER cycle on H100.

## Architecture

### Data Structures

**Delete:**
- `seg_tree: CudaSlice<f32>` — 67M floats (256MB). Gone.
- `capacity_pow2: usize` — only needed for tree indexing. Gone.

**Keep:**
- `priorities: CudaSlice<f32>` — `[capacity]` raw priorities `|td|^alpha + eps`

**Add:**
- `priorities_pa: CudaSlice<f32>` — `[capacity]` priority^alpha values (replaces tree leaves)
- `prefix_sum: CudaSlice<f32>` — `[capacity]` inclusive prefix sum of `priorities_pa[0..size]`
- `scan_status: CudaSlice<u32>` — `[num_tiles]` decoupled lookback state flags (zeroed before each scan)
- `scan_values: CudaSlice<f32>` — `[2 × num_tiles]` aggregate + prefix values (packed)

Where `num_tiles = ceil(capacity / TILE_SIZE)`, `TILE_SIZE = 1024`.

### Kernels

**Delete:** `seg_tree_kernel.cu` — all 5 kernels (`seg_tree_update_leaves`, `seg_tree_rebuild_level`, `seg_tree_insert`, `seg_tree_sample`, `seg_tree_gather_prios`).

**Add** in new file `per_kernels.cu`:

#### 1. `per_update_pa` — Priority leaf update (O(B), no propagation)

```cuda
extern "C" __global__ void per_update_pa(
    float* __restrict__ priorities_pa,        // [capacity] — write priority^alpha
    float* __restrict__ priorities,           // [capacity] — write raw priority
    float* __restrict__ batch_max,            // [1] atomicMax accumulator
    const unsigned int* __restrict__ indices, // [batch_size] buffer indices
    const float* __restrict__ td_errors,      // [batch_size] f32 TD errors
    float alpha, float epsilon,
    int capacity, int batch_size)
```

Per thread: compute `p = |td|^alpha + eps`, write `priorities[idx] = p`, compute `pa = p^alpha`, write `priorities_pa[idx] = pa`, atomicMax batch_max. No tree, no propagation, O(1) per thread.

Launch: `grid=(ceil(B/256), 1, 1), block=(256, 1, 1)`.

#### 2. `per_prefix_scan` — Decoupled lookback inclusive scan (O(N), single pass)

```cuda
extern "C" __global__ void per_prefix_scan(
    const float* __restrict__ priorities_pa,  // [capacity] input
    float* __restrict__ prefix_sum,           // [capacity] output (inclusive scan)
    volatile unsigned int* __restrict__ tile_status,  // [num_tiles] state flags
    volatile float* __restrict__ tile_values, // [2*num_tiles] [agg, prefix] pairs
    int size)                                 // actual valid entries
```

Algorithm:
1. Each block loads a tile of TILE_SIZE elements from `priorities_pa`
2. Intra-block inclusive scan via warp shuffles + shared memory (Blelloch/Hillis-Steele)
3. Block 0: prefix = 0, publish tile aggregate + prefix, set status = PREFIX_AVAILABLE
4. Block i>0: spin on `tile_status[i-1]`, look back through chain until finding PREFIX_AVAILABLE, accumulate aggregates, compute own prefix, publish, set status = PREFIX_AVAILABLE
5. Add prefix to all scanned elements, write to `prefix_sum`

Launch: `grid=(ceil(size/TILE_SIZE), 1, 1), block=(TILE_SIZE, 1, 1)`, where `TILE_SIZE = 1024`.

For N=20M: 19532 blocks, each processing 1024 elements. At H100's 3TB/s bandwidth: 20M × 4B × 2 (read+write) = 160MB → ~53µs. Actual: ~20-30µs (L2 cache + SM occupancy).

#### 3. `per_sample` — Proportional sampling via binary search (O(B log N))

```cuda
extern "C" __global__ void per_sample(
    const float* __restrict__ prefix_sum,     // [size] inclusive prefix sum
    const float* __restrict__ priorities_pa,  // [capacity] for pa gathering
    long long* __restrict__ out_indices,      // [batch_size] output indices (i64)
    float* __restrict__ out_priorities,       // [batch_size] sampled pa values
    float* __restrict__ total_sum_out,        // [1] = prefix_sum[size-1]
    unsigned int seed,                        // Philox RNG seed
    int size, int batch_size)
```

Per thread:
1. Thread 0 writes `total_sum_out[0] = prefix_sum[size-1]`
2. Generate `u = Philox(tid, seed) * prefix_sum[size-1]` (uniform in [0, total_sum))
3. Binary search `prefix_sum[0..size]` for smallest `i` where `prefix_sum[i] >= u`
4. Write `out_indices[tid] = i`
5. Write `out_priorities[tid] = priorities_pa[i]` (gather for IS weights)

Launch: `grid=(ceil(B/256), 1, 1), block=(256, 1, 1)`.

For B=8192, N=20M: 8192 threads × 25 binary search steps × 1 global memory read each = 200K reads. Cache-friendly (binary search converges to same regions for similar thresholds).

#### 4. `per_insert_pa` — Insert priorities at max_priority (O(B), for insert_batch)

```cuda
extern "C" __global__ void per_insert_pa(
    float* __restrict__ priorities_pa,        // [capacity]
    const unsigned int* __restrict__ indices, // [batch_size]
    const float* __restrict__ priorities,     // [batch_size] raw priority values (max_priority fill)
    float alpha,
    int capacity, int batch_size)
```

Per thread: `priorities_pa[indices[i]] = powf(priorities[i], alpha)`. No propagation — caller runs `per_prefix_scan` after.

Launch: `grid=(ceil(B/256), 1, 1), block=(256, 1, 1)`.

### Per-Step Training Flow

```
1. per_update_pa       — 8192 leaf writes               ~5µs
2. per_prefix_scan     — scan priorities_pa[0..size]     ~25µs
3. per_sample          — binary search + gather pa       ~10µs
   (replaces seg_tree_sample + seg_tree_gather_prios + DtoD root copy)
4. i64_to_u32          — existing kernel                 ~2µs
5. gather states/actions/rewards/dones/episodes          existing kernels
6. is_weights_f32      — existing kernel, total_sum_out  ~3µs
7. reduce_max + normalize_weights                        existing kernels
Total PER: ~45µs (was: infinite hang or 10+ seconds)
```

### Insert Flow

```
1. fill priorities_pa slots with max_priority^alpha      (per_insert_pa)
2. fill priorities slots with max_priority               (existing scatter_insert_f32)
3. per_prefix_scan to rebuild prefix_sum                 (full scan)
```

No `rebuild_tree`, no 25 kernel launches. One scan kernel.

### `sample_proportional` Changes

Replace steps 1, 4, 5 (tree sample, gather prios, DtoD root copy) with:

```rust
// Step 0: Zero scan scratch
stream.memset_zeros(&mut self.scan_status)?;

// Step 1: Prefix scan
unsafe {
    stream.launch_builder(&self.kernels.per_prefix_scan)
        .arg(&self.priorities_pa).arg(&mut self.prefix_sum)
        .arg(&self.scan_status).arg(&self.scan_values)
        .arg(&(self.size as i32))
        .launch(scan_config)?;
}

// Step 2: Sample + gather
unsafe {
    stream.launch_builder(&self.kernels.per_sample)
        .arg(&self.prefix_sum).arg(&self.priorities_pa)
        .arg(&mut self.sample_indices_i64).arg(&mut self.sample_priorities)
        .arg(&mut self.total_sum_buf)
        .arg(&seed).arg(&(self.size as i32)).arg(&bsi)
        .launch(sample_config)?;
}
```

Steps 3-8 (gather, IS weights, normalize) remain unchanged.

### `update_priorities_gpu` Changes

Replace `seg_tree_update_leaves` + `rebuild_tree/atomicAdd` with:
1. bf16→f32 cast (existing, uses pre-allocated scratch)
2. `per_update_pa` kernel (leaf write only, O(B))
3. Batch max merge (existing)

No tree propagation, no rebuild. The prefix sum is rebuilt in the next `sample_proportional` call.

### Dead Code Removal

| Item | Action |
|------|--------|
| `seg_tree_kernel.cu` | DELETE file |
| `seg_tree: CudaSlice<f32>` field | DELETE |
| `capacity_pow2: usize` field | DELETE |
| `seg_tree_update_leaves` kernel | DELETE |
| `seg_tree_rebuild_level` kernel | DELETE |
| `seg_tree_insert` kernel | DELETE |
| `seg_tree_sample` kernel | DELETE |
| `seg_tree_gather_prios` kernel | DELETE |
| `rebuild_tree()` method | DELETE |
| `ReplayKernels::seg_tree_*` fields | DELETE (replace with `per_*` fields) |
| `ST_CUBIN` static | DELETE (replace with `PER_CUBIN`) |
| All `seg_tree` references in build.rs | Replace with `per_kernels.cu` |
| All `capacity_pow2` references | Replace with `capacity` |

### VRAM Savings

| Before | After | Savings |
|--------|-------|---------|
| `seg_tree`: 2 × 33M × 4B = 256MB | `priorities_pa`: 20M × 4B = 80MB | |
| | `prefix_sum`: 20M × 4B = 80MB | |
| | `scan_scratch`: ~80KB | |
| **Total: 256MB** | **Total: 160MB** | **96MB freed** |

### Testing

- All 19 smoke tests must pass
- Verify correct PER sampling distribution: high-priority experiences sampled more frequently
- Verify IS weights are correctly computed (sum to ~batch_size after normalization)
- H100 deployment: training progresses past step 0, epochs complete

### Success Criteria

- Zero segment tree code in codebase
- PER update + sample in 2-3 kernel launches (not 25+)
- Total PER cycle < 50µs on H100
- No hangs, no contention, no cross-stream races
- Correct proportional sampling (validated by smoke tests)
