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

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

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