foxhunt/docs/superpowers/plans/2026-03-22-gpu-segment-tree.md

GPU Segment Tree for PER Sampling

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 O(n) prefix sum PER sampling with O(log n) GPU segment tree, eliminating per-epoch training time growth on H100 (12ms→23ms/step).

Architecture: Binary segment tree as flat CudaSlice<f32>[2*capacity_pow2]. Two CUDA kernels: seg_tree_update (batch leaf write + propagate up) and seg_tree_sample (parallel root-to-leaf traversal with Philox RNG). Replaces prefix_sum + searchsorted pipeline entirely.

Tech Stack: Rust, cudarc 0.19, CUDA (NVRTC), Philox 4×32 PRNG

Spec: docs/superpowers/specs/2026-03-22-gpu-segment-tree-design.md

---

Files

Create

File	Responsibility
crates/ml-dqn/src/seg_tree_kernel.cu	seg_tree_update + seg_tree_sample CUDA kernels

Modify

File	Changes
crates/ml-dqn/src/gpu_replay_buffer.rs	Replace prefix sum with segment tree in struct, constructor, sample, update

Delete

File	Reason
crates/ml-dqn/src/prefix_sum_kernel.cu	Replaced by seg_tree_kernel.cu

---

Task 1: Create the segment tree CUDA kernels

Files:

• Create: crates/ml-dqn/src/seg_tree_kernel.cu

• Step 1: Write seg_tree_update kernel

// Batch update: write priorities^alpha to leaves, propagate sums to root.
// One thread per index. Each thread writes its leaf then walks up to root.
// Concurrent writes to shared ancestors use atomicExch for correctness.
extern "C" __global__ void seg_tree_update(
    float* __restrict__ tree,
    const unsigned int* __restrict__ indices,
    const float* __restrict__ td_errors,
    float alpha, float epsilon,
    int capacity, int batch_size)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    int idx = indices[i];
    if (idx < 0 || idx >= capacity) return;

    // Compute priority = (|td_error| + epsilon)^alpha
    float td = td_errors[i];
    float priority = powf(fabsf(td) + epsilon, alpha);

    // Write leaf
    int leaf = capacity + idx;
    tree[leaf] = priority;

    // Propagate up to root
    int node = leaf >> 1;
    while (node >= 1) {
        tree[node] = tree[2 * node] + tree[2 * node + 1];
        node >>= 1;
    }
}

// Batch insert: write raw priorities (already computed) to leaves + propagate.
extern "C" __global__ void seg_tree_insert(
    float* __restrict__ tree,
    const unsigned int* __restrict__ indices,
    const float* __restrict__ priorities,
    int capacity, int batch_size)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    int idx = indices[i];
    if (idx < 0 || idx >= capacity) return;

    int leaf = capacity + idx;
    tree[leaf] = priorities[i];

    int node = leaf >> 1;
    while (node >= 1) {
        tree[node] = tree[2 * node] + tree[2 * node + 1];
        node >>= 1;
    }
}

// Proportional sampling: parallel root-to-leaf traversal.
// Each thread generates a random threshold and traverses the tree.
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)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    // Philox RNG: generate uniform in [0, total_sum)
    unsigned int state = rng_states[i];
    unsigned int hi, lo;
    lo = state * 0xD2511F53u;
    hi = __umulhi(state, 0xD2511F53u);
    for (int r = 0; r < 10; r++) {
        unsigned int t = hi ^ state;
        hi = lo * 0xCD9E8D57u;
        lo = __umulhi(lo, 0xCD9E8D57u);
        lo ^= t;
        state += 0x9E3779B9u;
    }
    unsigned int bits = lo ^ hi;
    float u = (float)(bits >> 8) * (1.0f / 16777216.0f);
    float total_sum = tree[1];  // root = total priority sum
    float threshold = u * total_sum;

    // Update RNG state
    rng_states[i] = state;

    // Tree traversal: root → leaf
    int node = 1;
    while (node < capacity) {
        float left_sum = tree[2 * node];
        if (threshold <= left_sum) {
            node = 2 * node;
        } else {
            threshold -= left_sum;
            node = 2 * node + 1;
        }
    }

    // node is now a leaf: index = node - capacity
    int idx = node - capacity;
    // Clamp to valid range (tree may have empty leaves beyond buffer size)
    if (idx < 0) idx = 0;
    out_indices[i] = idx;
}

• Step 2: Commit

git add crates/ml-dqn/src/seg_tree_kernel.cu
git commit -m "feat: GPU segment tree CUDA kernels for O(log n) PER sampling"

---

Task 2: Replace prefix sum with segment tree in GpuReplayBuffer

Files:

• Modify: crates/ml-dqn/src/gpu_replay_buffer.rs

This is the main task. Read the full file first, then make these changes:

• Step 1: Replace struct fields

Remove:

pa_buf: CudaSlice<f32>,
cs_buf: CudaSlice<f32>,
block_sums: CudaSlice<f32>,

Add:

seg_tree: CudaSlice<f32>,      // [2 * capacity_pow2] segment tree
capacity_pow2: usize,           // capacity rounded up to power of 2

• Step 2: Replace kernel fields in ReplayKernels

Remove from ReplayKernels:

pow_alpha_f32: CudaFunction,
prefix_sum: CudaFunction,
prefix_sum_local: CudaFunction,
prefix_sum_block_sums: CudaFunction,
prefix_sum_propagate: CudaFunction,
searchsorted: CudaFunction,

Add:

seg_tree_update: CudaFunction,
seg_tree_insert: CudaFunction,
seg_tree_sample: CudaFunction,

• Step 3: Update kernel compilation

In ReplayKernels::new(), compile seg_tree_kernel.cu instead of prefix_sum_kernel.cu. Load the 3 new functions. Remove the old prefix sum + pow_alpha + searchsorted loads.

• Step 4: Update constructor

In GpuReplayBuffer::new():

• Compute capacity_pow2 = config.capacity.next_power_of_two()

• Allocate seg_tree = stream.alloc_zeros::<f32>(2 * capacity_pow2)

• Remove pa_buf, cs_buf, block_sums allocations

• Initialize rng_states buffer for Philox (reuse existing rng_step or add per-sample RNG states)

• Step 5: Rewrite sample_proportional()

Replace the entire method body:

pub fn sample_proportional(&mut self, batch_size: usize) -> Result<GpuBatchSlices, MLError> {
    // 1. Launch seg_tree_sample kernel (root-to-leaf traversal)
    // 2. Use sampled indices to gather states, actions, rewards, dones, priorities
    // 3. Compute IS weights from priorities + total_sum (tree[1])
    // Return GpuBatchSlices
}

The key difference: no prefix sum, no searchsorted. Just 1 kernel launch for sampling + gather kernels.

• Step 6: Rewrite update_priorities_gpu()

Replace prefix scatter + rebuild with:

pub fn update_priorities_gpu(&mut self, indices, td_errors, bs) -> Result<(), MLError> {
    // Launch seg_tree_update kernel: writes new priorities to leaves, propagates up
    // Batch max accumulation stays the same (atomicMax or separate kernel)
}

• Step 7: Update insert_batch()

After writing priorities to the ring buffer, also update the segment tree leaves:

// After scatter_insert of priorities:
// Launch seg_tree_insert kernel on the inserted indices

• Step 8: Remove pfx_sum() method

Delete the entire pfx_sum() method and all prefix sum infrastructure.

• Step 9: Update total_priority() / cs_total()

Total priority sum is now tree[1] (root node). Replace any remaining cs_total() or total_sum readback with a DtoD copy from seg_tree[1].

• Step 10: Verify compilation + tests

Run: SQLX_OFFLINE=true cargo check -p ml-dqn Run: SQLX_OFFLINE=true cargo test -p ml-dqn --lib Expected: 359 passed, 0 failed

• Step 11: Commit

git add crates/ml-dqn/src/gpu_replay_buffer.rs
git commit -m "perf: replace O(n) prefix sum with O(log n) GPU segment tree for PER"

---

Task 3: Delete old prefix sum kernel + validate

Files:

• Delete: crates/ml-dqn/src/prefix_sum_kernel.cu

• Step 1: Verify no remaining references

grep -rn "prefix_sum_kernel\|prefix_sum_local\|prefix_sum_block_sums\|prefix_sum_propagate" crates/ml-dqn/src/

• Step 2: Delete

rm crates/ml-dqn/src/prefix_sum_kernel.cu

• Step 3: Run full test suite

SQLX_OFFLINE=true cargo test -p ml-dqn --lib
SQLX_OFFLINE=true cargo test -p ml --lib

• Step 4: Run GPU smoke test

SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline \
  cargo test -p ml --lib -- test_train_step_produces_finite_metrics --ignored --test-threads=1

• Step 5: Profile locally

SQLX_OFFLINE=true cargo run --release --example train_baseline_rl -p ml -- \
  --model dqn --data-dir test_data/futures-baseline --symbol ES.FUT \
  --training-profile dqn-production --epochs 5 --max-steps-per-epoch 50 \
  --train-months 3 --val-months 1 --test-months 1 --step-months 3 \
  2>&1 | grep "step breakdown"

Expected: sample time should be CONSTANT across epochs (not growing).

• Step 6: Commit + push

git add -A crates/ml-dqn/src/
git commit -m "refactor: delete prefix_sum_kernel.cu (replaced by segment tree)"
git push origin main