# PER Prefix Sum Replacement Implementation Plan

> **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 the PER segment tree (hangs on H100) with a decoupled lookback prefix sum scan + binary search sampling — the NVIDIA-standard O(N) parallel scan algorithm.

**Architecture:** Delete `seg_tree_kernel.cu` and all segment tree code. Create `per_kernels.cu` with 4 kernels: `per_update_pa` (leaf write, accepts bf16 directly), `per_prefix_scan` (Merrill-Garland decoupled lookback), `per_sample` (binary search + gather), `per_insert_pa` (insert-time leaf write). Delete the `CastKernels` / `OnceLock` / `cast_kernels.cubin` infrastructure (source of H100 hang). Delete `update_td_f32` scratch buffer. Two kernel launches for the entire PER cycle instead of 25+.

**Tech Stack:** Rust 1.85, CUDA 12.4 (SM 9.0), cudarc 0.19, half (bf16)

---

### Task 1: Write `per_kernels.cu` with All 4 Kernels

**Files:**
- Create: `crates/ml-dqn/src/per_kernels.cu`

- [ ] **Step 1: Write the `per_update_pa` kernel**

Create `crates/ml-dqn/src/per_kernels.cu`:

```cuda
// PER (Prioritized Experience Replay) kernels — prefix sum architecture.
//
// Replaces the segment tree with a flat prefix sum array using NVIDIA's
// decoupled lookback scan algorithm (Merrill & Garland, 2016).
//
// Four kernels:
//   1. per_update_pa      — priority update: compute pa from bf16 TD errors,
//                           write to priorities_pa[] and priorities[].
//   2. per_insert_pa      — insert: write priority^alpha to priorities_pa[].
//   3. per_prefix_scan    — decoupled lookback inclusive scan of priorities_pa.
//   4. per_sample         — proportional sampling via binary search on prefix sum.

#include <cuda_bf16.h>

// ═══════════════════════════════════════════════════════════════════════
// 1. per_update_pa — Priority leaf update from bf16 TD errors
//
// Reads bf16 td_errors directly (no separate bf16→f32 cast kernel).
// Writes both raw priorities and priority^alpha values.
// Launch: grid=(ceil(batch_size/256), 1, 1), block=(256, 1, 1)
// ═══════════════════════════════════════════════════════════════════════

extern "C" __global__ void per_update_pa(
    float* __restrict__ priorities_pa,
    float* __restrict__ priorities,
    float* __restrict__ batch_max,
    const unsigned int* __restrict__ indices,
    const __nv_bfloat16* __restrict__ td_errors_bf16,
    float alpha, float epsilon,
    int capacity, int batch_size)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    unsigned int idx = indices[i];
    if (idx >= (unsigned int)capacity) return;

    float td = fabsf(__bfloat162float(td_errors_bf16[i]));
    float new_prio = powf(td, alpha) + epsilon;
    if (new_prio < epsilon) new_prio = epsilon;
    if (new_prio > 1e6f) new_prio = 1e6f;

    priorities[idx] = new_prio;

    int ival = __float_as_int(new_prio);
    atomicMax((int*)batch_max, ival);

    float pa = powf(new_prio, alpha);
    priorities_pa[idx] = pa;
}

// ═══════════════════════════════════════════════════════════════════════
// 2. per_insert_pa — Insert priorities at max_priority
//
// Called during insert_batch. Takes raw priority values (max_priority fill),
// computes priority^alpha, writes to priorities_pa[].
// Launch: grid=(ceil(batch_size/256), 1, 1), block=(256, 1, 1)
// ═══════════════════════════════════════════════════════════════════════

extern "C" __global__ void per_insert_pa(
    float* __restrict__ priorities_pa,
    const unsigned int* __restrict__ indices,
    const float* __restrict__ raw_priorities,
    float alpha,
    int capacity, int batch_size)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    unsigned int idx = indices[i];
    if (idx >= (unsigned int)capacity) return;

    priorities_pa[idx] = powf(raw_priorities[i], alpha);
}

// ═══════════════════════════════════════════════════════════════════════
// 3. per_prefix_scan — Decoupled lookback inclusive prefix sum
//
// Single-pass O(N) parallel scan using the Merrill-Garland algorithm.
// Each block processes TILE_SIZE elements, publishes its aggregate,
// then looks back through preceding blocks to determine its global prefix.
//
// Status flags per tile (packed in tile_descriptors):
//   bits [31:30] = state: 0=INVALID, 1=AGGREGATE, 2=PREFIX
//   bits [29:0]  = unused (value stored separately for f32 precision)
//
// Launch: grid=(num_tiles, 1, 1), block=(TILE_SIZE, 1, 1)
//   where num_tiles = ceil(size / TILE_SIZE), TILE_SIZE = 256
// ═══════════════════════════════════════════════════════════════════════

#define PER_TILE_SIZE 256
#define STATE_INVALID    0u
#define STATE_AGGREGATE  1u
#define STATE_PREFIX     2u

extern "C" __global__ void per_prefix_scan(
    const float* __restrict__ input,
    float* __restrict__ output,
    volatile unsigned int* __restrict__ tile_state,
    volatile float* __restrict__ tile_aggregate,
    volatile float* __restrict__ tile_prefix,
    int size)
{
    int tile_idx = blockIdx.x;
    int tid = threadIdx.x;
    int global_idx = tile_idx * PER_TILE_SIZE + tid;

    // ── Phase 1: Load tile into shared memory ──
    __shared__ float sdata[PER_TILE_SIZE];
    float val = (global_idx < size) ? input[global_idx] : 0.0f;
    sdata[tid] = val;
    __syncthreads();

    // ── Phase 2: Intra-tile inclusive scan (Hillis-Steele) ──
    for (int offset = 1; offset < PER_TILE_SIZE; offset <<= 1) {
        float addend = (tid >= offset) ? sdata[tid - offset] : 0.0f;
        __syncthreads();
        sdata[tid] += addend;
        __syncthreads();
    }

    // Tile aggregate = last element of scanned tile
    float tile_agg = sdata[PER_TILE_SIZE - 1];

    // ── Phase 3: Publish aggregate and determine prefix ──
    __shared__ float s_prefix;

    if (tid == 0) {
        // Publish this tile's aggregate
        tile_aggregate[tile_idx] = tile_agg;
        __threadfence();
        tile_state[tile_idx] = STATE_AGGREGATE;
        __threadfence();

        if (tile_idx == 0) {
            // First tile: prefix is 0 (identity for addition)
            tile_prefix[tile_idx] = 0.0f;
            __threadfence();
            tile_state[tile_idx] = STATE_PREFIX;
            s_prefix = 0.0f;
        } else {
            // Decoupled lookback: scan backwards through preceding tiles
            float running_prefix = 0.0f;
            int lookback = tile_idx - 1;

            while (true) {
                unsigned int state = tile_state[lookback];
                if (state == STATE_PREFIX) {
                    running_prefix += tile_prefix[lookback] + tile_aggregate[lookback];
                    break;
                } else if (state == STATE_AGGREGATE) {
                    running_prefix += tile_aggregate[lookback];
                    lookback--;
                }
                // STATE_INVALID: spin (tile hasn't started yet)
            }

            tile_prefix[tile_idx] = running_prefix;
            __threadfence();
            tile_state[tile_idx] = STATE_PREFIX;
            s_prefix = running_prefix;
        }
    }
    __syncthreads();

    // ── Phase 4: Add prefix to all elements and write output ──
    if (global_idx < size) {
        output[global_idx] = sdata[tid] + s_prefix;
    }
}

// ═══════════════════════════════════════════════════════════════════════
// 4. per_sample — Proportional sampling via binary search
//
// Each thread generates a random threshold in [0, total_sum) and performs
// binary search on the prefix sum array to find the corresponding index.
// Also gathers priorities_pa[index] for IS weight computation.
//
// Launch: grid=(ceil(batch_size/256), 1, 1), block=(256, 1, 1)
// ═══════════════════════════════════════════════════════════════════════

// Philox single-round helper (counter-based PRNG)
__device__ __forceinline__
unsigned int philox_single(unsigned int counter, unsigned int key) {
    unsigned int hi, lo;
    lo = counter * 0xD2511F53u;
    hi = __umulhi(counter, 0xD2511F53u);
    for (int r = 0; r < 10; r++) {
        unsigned int t = hi ^ key;
        hi = lo * 0xCD9E8D57u;
        lo = __umulhi(lo, 0xCD9E8D57u);
        lo ^= t;
        key += 0x9E3779B9u;
    }
    return lo ^ hi;
}

extern "C" __global__ void per_sample(
    const float* __restrict__ prefix_sum,
    const float* __restrict__ priorities_pa,
    long long* __restrict__ out_indices,
    float* __restrict__ out_priorities,
    float* __restrict__ total_sum_out,
    unsigned int seed,
    int size, int batch_size)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    // Thread 0 writes total_sum for IS weight computation
    float total_sum = prefix_sum[size - 1];
    if (i == 0) {
        total_sum_out[0] = total_sum;
    }

    // Generate uniform random threshold in [0, total_sum)
    unsigned int bits = philox_single((unsigned int)i, seed);
    float u = (float)(bits >> 8) * (1.0f / 16777216.0f);
    float threshold = u * total_sum;

    // Binary search: find smallest idx where prefix_sum[idx] >= threshold
    int lo = 0, hi = size - 1;
    while (lo < hi) {
        int mid = (lo + hi) >> 1;
        if (prefix_sum[mid] < threshold) {
            lo = mid + 1;
        } else {
            hi = mid;
        }
    }

    // Clamp to valid range
    if (lo >= size) lo = size - 1;
    if (lo < 0) lo = 0;

    out_indices[i] = (long long)lo;
    out_priorities[i] = priorities_pa[lo];
}
```

- [ ] **Step 2: Add to build.rs**

In `crates/ml-dqn/build.rs`, add `"per_kernels.cu"` to the CUDA source files list (where `"seg_tree_kernel.cu"` currently is):

Replace:
```
        "seg_tree_kernel.cu",
```
with:
```
        "per_kernels.cu",
```

- [ ] **Step 3: Verify kernel compilation**

Run: `SQLX_OFFLINE=true cargo check -p ml-dqn --lib 2>&1 | tail -5`
Expected: The cubin compiles (build.rs runs nvcc). Rust compilation may fail because kernel names changed — that's fixed in Task 2.

- [ ] **Step 4: Commit**

```bash
git add crates/ml-dqn/src/per_kernels.cu crates/ml-dqn/build.rs
git commit -m "feat: per_kernels.cu — decoupled lookback scan + binary search sampling

Four PER kernels replacing the segment tree:
1. per_update_pa — priority leaf update from bf16 TD errors directly
2. per_insert_pa — insert priorities at max_priority
3. per_prefix_scan — Merrill-Garland decoupled lookback inclusive scan
4. per_sample — proportional sampling via binary search on prefix sum"
```

---

### Task 2: Replace Segment Tree Data Structures in `GpuReplayBuffer`

**Files:**
- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs`
- Delete: `crates/ml-dqn/src/seg_tree_kernel.cu`

- [ ] **Step 1: Replace `ReplayKernels` segment tree fields with PER kernel fields**

In `gpu_replay_buffer.rs`, replace lines 90-95:

```rust
    // Segment tree kernels (replace prefix_sum + searchsorted + pow_alpha pipeline)
    seg_tree_update_leaves: CudaFunction,
    seg_tree_rebuild_level: CudaFunction,
    seg_tree_insert: CudaFunction,
    seg_tree_sample: CudaFunction,
    seg_tree_gather_prios: CudaFunction,
```

with:

```rust
    // PER prefix sum kernels (decoupled lookback scan + binary search)
    per_update_pa: CudaFunction,
    per_insert_pa: CudaFunction,
    per_prefix_scan: CudaFunction,
    per_sample: CudaFunction,
```

- [ ] **Step 2: Replace kernel loading in `ReplayKernels::compile`**

Replace lines 110-135 (the ST_CUBIN loading + 5 function loads):

```rust
        // Load precompiled segment tree kernels cubin
        static ST_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/seg_tree_kernel.cubin"));
        let st_mod = ctx.load_cubin(ST_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("seg_tree cubin load: {e}")))?;
        // ... 5 load_function calls ...
```

with:

```rust
        // Load precompiled PER prefix sum kernels cubin
        static PER_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/per_kernels.cubin"));
        let per_mod = ctx.load_cubin(PER_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("per cubin load: {e}")))?;
        let per_ld = |n: &str| -> Result<CudaFunction, MLError> {
            per_mod.load_function(n).map_err(|e| MLError::ModelError(format!("per {n}: {e}")))
        };
```

And update the struct init to use:
```rust
            per_update_pa: per_ld("per_update_pa")?,
            per_insert_pa: per_ld("per_insert_pa")?,
            per_prefix_scan: per_ld("per_prefix_scan")?,
            per_sample: per_ld("per_sample")?,
```

- [ ] **Step 3: Replace struct fields in `GpuReplayBuffer`**

Replace:
```rust
    seg_tree: CudaSlice<f32>,
    capacity_pow2: usize,
```

with:
```rust
    priorities_pa: CudaSlice<f32>,
    prefix_sum: CudaSlice<f32>,
    scan_tile_state: CudaSlice<u32>,
    scan_tile_aggregate: CudaSlice<f32>,
    scan_tile_prefix: CudaSlice<f32>,
    num_scan_tiles: usize,
```

- [ ] **Step 4: Delete cast kernel infrastructure**

Delete entirely from `gpu_replay_buffer.rs`:
- `struct CastKernels` (line 48-50)
- `static CAST_KERNELS: OnceLock<...>` (line 52)
- `fn get_cast_kernels(...)` (lines 54-71)
- `update_td_f32: CudaSlice<f32>` field from struct
- The pre-init `get_cast_kernels(stream)?` call in constructor
- `use std::sync::OnceLock;` import if now unused

- [ ] **Step 5: Update constructor to allocate new buffers**

Replace the segment tree allocation:
```rust
    let cap_pow2 = cap.next_power_of_two();
    let seg = a32f(stream, 2 * cap_pow2, "seg_tree")?;
```

with:
```rust
    let priorities_pa = a32f(stream, cap, "priorities_pa")?;
    let prefix_sum = a32f(stream, cap, "prefix_sum")?;
    let tile_size = 256_usize; // PER_TILE_SIZE in CUDA kernel
    let num_scan_tiles = cap.div_ceil(tile_size);
    let scan_tile_state = stream.alloc_zeros::<u32>(num_scan_tiles)
        .map_err(|e| MLError::ModelError(format!("alloc scan_tile_state: {e}")))?;
    let scan_tile_aggregate = a32f(stream, num_scan_tiles, "scan_tile_agg")?;
    let scan_tile_prefix = a32f(stream, num_scan_tiles, "scan_tile_pfx")?;
```

Remove `update_td_f32` allocation. Update struct initializer to use new fields.

- [ ] **Step 6: Delete `seg_tree_kernel.cu`**

```bash
rm crates/ml-dqn/src/seg_tree_kernel.cu
```

- [ ] **Step 7: Build check**

Run: `SQLX_OFFLINE=true cargo check -p ml-dqn --lib 2>&1 | tail -10`
Expected: Errors in methods that still reference old fields — fixed in Task 3.

- [ ] **Step 8: Commit**

```bash
git add -u crates/ml-dqn/src/
git commit -m "refactor: replace segment tree with prefix sum buffers

Delete seg_tree (67M floats = 256MB), capacity_pow2, CastKernels,
OnceLock, get_cast_kernels, update_td_f32 scratch buffer.
Add priorities_pa, prefix_sum, scan scratch buffers.
Delete seg_tree_kernel.cu entirely."
```

---

### Task 3: Rewrite `update_priorities_gpu` — No Cast, No Tree

**Files:**
- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs`

- [ ] **Step 1: Delete `rebuild_tree` method**

Delete the entire `fn rebuild_tree(...)` method.

- [ ] **Step 2: Rewrite `update_priorities_gpu`**

Replace the entire method body with:

```rust
    pub fn update_priorities_gpu(
        &mut self,
        indices: &CudaSlice<u32>,
        td_errors: &CudaSlice<half::bf16>,
        bs: usize,
        ext_stream: Option<&Arc<CudaStream>>,
    ) -> Result<(), MLError> {
        if bs == 0 { return Ok(()); }
        let stream_owned = ext_stream.cloned().unwrap_or_else(|| Arc::clone(&self.stream));
        let stream = &stream_owned;

        let (al, ep, bsi) = (self.config.alpha, self.config.epsilon, bs as i32);
        let cap_i = self.config.capacity as i32;

        stream.memset_zeros(&mut self.update_batch_max)
            .map_err(|e| MLError::ModelError(format!("zero batch_max: {e}")))?;

        // per_update_pa: reads bf16 td_errors directly, writes priorities + priorities_pa
        unsafe {
            stream.launch_builder(&self.kernels.per_update_pa)
                .arg(&self.priorities_pa)
                .arg(&self.priorities)
                .arg(&self.update_batch_max)
                .arg(indices)
                .arg(td_errors)
                .arg(&al).arg(&ep).arg(&cap_i).arg(&bsi)
                .launch(lcfg(bs))
                .map_err(|e| MLError::ModelError(format!("per_update_pa: {e}")))?;
        }

        // Merge batch max into pending (existing logic, unchanged)
        match self.pending_max_priority.take() {
            Some(prev) => {
                unsafe {
                    stream.launch_builder(&self.kernels.max_of_two_f32)
                        .arg(&self.update_max_merge).arg(&prev).arg(&self.update_batch_max)
                        .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })
                        .map_err(|e| MLError::ModelError(format!("max_of_two: {e}")))?;
                }
                self.pending_max_priority = Some(std::mem::replace(&mut self.update_max_merge, prev));
            }
            None => {
                let fresh = a32f(stream, 1, "bm_epoch")?;
                self.pending_max_priority = Some(std::mem::replace(&mut self.update_batch_max, fresh));
            }
        }
        Ok(())
    }
```

- [ ] **Step 3: Update `update_priorities_gpu_raw` to pass bf16 directly**

The raw variant no longer needs the bf16→f32 conversion:

```rust
    pub fn update_priorities_gpu_raw(
        &mut self,
        _indices_ptr: u64,
        td_errors: &CudaSlice<half::bf16>,
        bs: usize,
        ext_stream: Option<&Arc<CudaStream>>,
    ) -> Result<(), MLError> {
        let idx_slice = unsafe { &*(&self.sample_indices_u32 as *const CudaSlice<u32>) };
        self.update_priorities_gpu(idx_slice, td_errors, bs, ext_stream)
    }
```

- [ ] **Step 4: Build check**

Run: `SQLX_OFFLINE=true cargo check -p ml-dqn --lib 2>&1 | tail -10`
Expected: Errors in `sample_proportional` and `insert_batch` — fixed in Tasks 4-5.

- [ ] **Step 5: Commit**

```bash
git add crates/ml-dqn/src/gpu_replay_buffer.rs
git commit -m "refactor: per_update_pa accepts bf16 directly, no cast kernel

Eliminates the bf16→f32 cast kernel launch, OnceLock infrastructure,
and update_td_f32 scratch buffer. The per_update_pa CUDA kernel reads
bf16 td_errors via __bfloat162float() inline."
```

---

### Task 4: Rewrite `sample_proportional` — Prefix Scan + Binary Search

**Files:**
- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs`

- [ ] **Step 1: Rewrite the sampling section**

In `sample_proportional`, replace the segment tree sampling steps (tree sample + gather prios + DtoD root copy) with prefix scan + binary search:

Replace steps 1, 4, 5 (seg_tree_sample, seg_tree_gather_prios, DtoD tree[1] copy) with:

```rust
        // Step 1: Zero scan scratch and run prefix scan
        self.stream.memset_zeros(&mut self.scan_tile_state)
            .map_err(|e| MLError::ModelError(format!("zero scan state: {e}")))?;

        let size_i = self.size as i32;
        let scan_blocks = self.num_scan_tiles as u32;
        unsafe {
            self.stream.launch_builder(&self.kernels.per_prefix_scan)
                .arg(&self.priorities_pa)
                .arg(&mut self.prefix_sum)
                .arg(&self.scan_tile_state)
                .arg(&self.scan_tile_aggregate)
                .arg(&self.scan_tile_prefix)
                .arg(&size_i)
                .launch(LaunchConfig {
                    grid_dim: (scan_blocks.max(1), 1, 1),
                    block_dim: (256, 1, 1), // PER_TILE_SIZE
                    shared_mem_bytes: 256 * 4, // PER_TILE_SIZE * sizeof(float)
                })
                .map_err(|e| MLError::ModelError(format!("per_prefix_scan: {e}")))?;
        }

        // Step 2: Sample + gather priorities_pa in one kernel
        self.rng_step = self.rng_step.wrapping_add(1);
        let seed = self.rng_step;
        unsafe {
            self.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(&size_i)
                .arg(&bsi)
                .launch(lcfg(batch_size))
                .map_err(|e| MLError::ModelError(format!("per_sample: {e}")))?;
        }
```

Delete the old steps:
- `seg_tree_sample` launch
- `seg_tree_gather_prios` launch
- The DtoD copy of `tree[1]` → `total_sum_buf`

Steps 2-8 (i64_to_u32, gather, IS weights, normalize) remain **unchanged**.

- [ ] **Step 2: Build check**

Run: `SQLX_OFFLINE=true cargo check -p ml-dqn --lib 2>&1 | tail -10`
Expected: Errors in `insert_batch` — fixed in Task 5.

- [ ] **Step 3: Commit**

```bash
git add crates/ml-dqn/src/gpu_replay_buffer.rs
git commit -m "feat: prefix scan + binary search sampling replaces segment tree

per_prefix_scan: Merrill-Garland decoupled lookback inclusive scan.
per_sample: binary search + priority gathering in one kernel.
2 launches instead of 3 (tree sample + gather + DtoD copy)."
```

---

### Task 5: Rewrite `insert_batch` / `insert_batch_bf16` + `clear()`

**Files:**
- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs`

- [ ] **Step 1: Replace `seg_tree_insert` + `rebuild_tree` in `insert_batch`**

Find the `seg_tree_insert` launch + `rebuild_tree` call in `insert_batch`. Replace with:

```rust
        // Write priority^alpha to priorities_pa (no tree, no propagation)
        unsafe {
            self.stream.launch_builder(&self.kernels.per_insert_pa)
                .arg(&self.priorities_pa).arg(&idx_buf).arg(&pt).arg(&al)
                .arg(&cap_i).arg(&bsi)
                .launch(lcfg(eff))
                .map_err(|e| MLError::ModelError(format!("per_insert_pa: {e}")))?;
        }
        // No rebuild needed — prefix_sum is rebuilt in sample_proportional
```

where `cap_i = self.config.capacity as i32`.

- [ ] **Step 2: Same replacement in `insert_batch_bf16`**

Identical change in the bf16 insert path.

- [ ] **Step 3: Update `clear()`**

Replace:
```rust
    self.stream.memset_zeros(&mut self.seg_tree)
        .map_err(|e| MLError::ModelError(format!("seg_tree clear: {e}")))?;
```
with:
```rust
    self.stream.memset_zeros(&mut self.priorities_pa)
        .map_err(|e| MLError::ModelError(format!("priorities_pa clear: {e}")))?;
    self.stream.memset_zeros(&mut self.prefix_sum)
        .map_err(|e| MLError::ModelError(format!("prefix_sum clear: {e}")))?;
```

- [ ] **Step 4: Remove all remaining `seg_tree` / `capacity_pow2` references**

Search for any remaining references to `seg_tree`, `capacity_pow2`, `rebuild_tree`, `CastKernels`, `get_cast_kernels`, `update_td_f32` in `gpu_replay_buffer.rs` and delete them.

- [ ] **Step 5: Remove all diagnostic `eprintln` from PER code**

Search for `PER_DIAG` and delete all diagnostic eprintln calls in `gpu_replay_buffer.rs`.

- [ ] **Step 6: Build and verify**

Run: `SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | tail -5`
Expected: Clean compilation.

- [ ] **Step 7: Run smoke tests**

Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep 'test result'`
Expected: `test result: ok. 19 passed; 0 failed;`

- [ ] **Step 8: Commit**

```bash
git add -u crates/ml-dqn/src/
git commit -m "feat: complete PER prefix sum replacement — zero segment tree code

insert_batch uses per_insert_pa (leaf write only, no rebuild).
clear() zeroes priorities_pa + prefix_sum.
All seg_tree, capacity_pow2, CastKernels, rebuild_tree references deleted.
All PER_DIAG diagnostic eprintln removed."
```

---

### Task 6: Clean Up Integration Tests + External Callers

**Files:**
- Modify: `crates/ml/tests/gpu_per_integration_test.rs`
- Modify: `crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs`
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`

- [ ] **Step 1: Fix integration tests**

In `gpu_per_integration_test.rs`, update all calls to `update_priorities_gpu` and `update_priorities_gpu_raw` to match the new signatures (bf16 td_errors directly, ext_stream parameter).

- [ ] **Step 2: Fix gpu_residency smoke test**

Same signature update in `gpu_residency.rs`.

- [ ] **Step 3: Remove diagnostic sync in fused_training.rs**

In `fused_training.rs`, remove the debug `cuStreamSynchronize` + `check_err` after PER update (the `H100_HANG4: syncing stream to check errors` block).

- [ ] **Step 4: Remove all remaining H100 diagnostic eprintln**

Search for `H100_STEP`, `H100_LOOP`, `H100_HANG4`, `H100_DIAG`, `adam_readback`, `replay_forward`, `PER_DIAG` eprintln calls across the entire `crates/ml/` and `crates/ml-dqn/` and delete them all. Training is confirmed working once smoke tests pass — no more diagnostic noise.

- [ ] **Step 5: Build, full test, commit**

Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5` (full crate including tests)
Expected: Clean.

Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep 'test result'`
Expected: `test result: ok. 19 passed; 0 failed;`

```bash
git add -u
git commit -m "cleanup: fix tests, remove all H100 diagnostic eprintln

Integration tests updated for new PER API (bf16 td_errors, ext_stream).
All diagnostic eprintln removed (H100_STEP, PER_DIAG, adam_readback, etc).
Debug cuStreamSynchronize removed from fused training path."
```

---

### Task 7: Deploy to H100 and Validate

**Files:** None (testing only)

- [ ] **Step 1: Push and deploy**

```bash
git push origin main
./scripts/argo-train.sh dqn --baseline --watch
```

- [ ] **Step 2: Monitor training progression**

Verify:
- Training starts (no kernel loading hangs)
- Step 0 completes (PER update + sample works)
- Step 1 appears (training loop progresses)
- Epoch 1 completes with timing
- Q-stats and loss reported normally

- [ ] **Step 3: Compare epoch timing**

Check the `Training step breakdown` log for per-step timing. Expected: sample + fused + guard timings with no multi-second PER stalls.
