foxhunt/docs/superpowers/plans/2026-04-06-per-prefix-sum.md

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:

// 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

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:

    // 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:

    // 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):

        // 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:

        // 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:

            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:

    seg_tree: CudaSlice<f32>,
    capacity_pow2: usize,

with:

    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:

    let cap_pow2 = cap.next_power_of_two();
    let seg = a32f(stream, 2 * cap_pow2, "seg_tree")?;

with:

    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

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

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:

    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:

    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

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:

        // 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

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:

        // 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:

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

with:

    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

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;

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

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.