diff --git a/docs/superpowers/specs/2026-05-25-gpu-per-async-diag.md b/docs/superpowers/specs/2026-05-25-gpu-per-async-diag.md new file mode 100644 index 000000000..217ec4708 --- /dev/null +++ b/docs/superpowers/specs/2026-05-25-gpu-per-async-diag.md @@ -0,0 +1,225 @@ +# GPU-Resident PER + Async Non-Blocking Diag + +**Goal:** Eliminate the two host-side bottlenecks (PER + diag readback) that cap throughput at 34 sps @ b=256 on L40S. Target: 100+ sps @ b=256. + +**Architecture:** All replay buffer operations move to device (push, sample, gather, priority update). Diag readback uses a separate CUDA stream with double-buffered mapped-pinned memory — zero stalls on the training pipeline. + +**Tech Stack:** Rust, cudarc 0.19, CUDA 12.4, pre-compiled cubins + +--- + +## Part 1: Async Non-Blocking Diag Streaming + +### Problem + +The train binary (`alpha_rl_train.rs`) calls `read_slice_d` 10× per step to read ISV, rewards, actions, positions, etc. from device to host for the JSONL diag. Each `read_slice_d` involves `stream.synchronize()` which flushes the entire GPU pipeline — the GPU sits idle until the copy completes, then idle again waiting for the next kernel launch. + +### Solution: Separate diag stream + double-buffered mapped-pinned + +``` +Training stream: [kernels...] → never synced from diag path +Diag stream: [DtoD: device → mapped-pinned staging] → host reads without sync +``` + +**Components:** + +1. **`DiagStagingBuffer`** — pre-allocated mapped-pinned buffer large enough for ALL per-step diag data (ISV + rewards + actions + dones + positions + unit state). Two copies (double-buffer). + +2. **`diag_stream: CudaStream`** — separate CUDA stream for diag DtoD copies. Does NOT block the training stream. + +3. **Diag snapshot kernel: `rl_diag_snapshot`** — single kernel launched on `diag_stream` that copies all diag fields into the staging buffer in one pass: + - ISV[0..RL_SLOTS_END] → staging[0..] + - rewards_d[0..B] → staging[ISV_LEN..] + - dones_d[0..B] → staging[ISV_LEN+B..] + - actions_d[0..B] (as f32 cast) → staging[...] + - etc. + + Alternatively: a sequence of `memcpy_dtod_async` calls on `diag_stream` (simpler, same effect since they're all on the same stream and execute in order). + +4. **Host read path** — after launching the snapshot on `diag_stream`, the host reads the PREVIOUS step's staging buffer (which completed one step ago). Zero sync needed — the data was visible after the previous `diag_stream.synchronize()` which happened at the START of this step (overlapped with training). + +**Flow per step:** +``` +step N: + 1. diag_stream.synchronize() ← waits for step N-1's DtoD (overlaps with step N-1's training) + 2. Host reads staging_buf[N-1 % 2] ← previous step's data, zero stall + 3. Write JSONL from host data ← IO, overlapped with GPU + 4. Launch DtoD copies on diag_stream for step N's data into staging_buf[N % 2] + 5. Training stream: Graph A → fill → Graph B → PER → Graph C ← no stall +``` + +The `diag_stream.synchronize()` at step start waits for the PREVIOUS step's copies (which ran concurrently with the previous step's training). Since training takes ~30ms and the DtoD copies take <1ms, the sync is always instant. + +**No `--diag-every` flag.** Every step gets full diag at zero throughput cost. + +### Implementation location + +- New struct: `crates/ml-alpha/src/trainer/diag_staging.rs` +- Modified: `crates/ml-alpha/examples/alpha_rl_train.rs` (replace `read_slice_d` calls with staging buffer reads) +- Remove: the `--diag-every` flag (no longer needed) + +--- + +## Part 2: GPU-Resident Prioritized Experience Replay + +### Problem + +`push_to_replay` reads 5 device buffers to host, loops over b_size with per-element `alloc_zeros` + DtoD, computes n-step returns on CPU, pushes to a host-side VecDeque + priority tree. + +`sample_and_gather` does priority-weighted sampling via host-side tree traversal, then gathers per-element with DtoD copies. + +At b=256 this is ~256 host round-trips per step = ~15ms blocked. + +### Solution: All-device replay buffer + +**Device-resident data structures:** + +``` +replay_h_t_d: CudaSlice [capacity × HIDDEN_DIM] — circular buffer +replay_h_tp1_d: CudaSlice [capacity × HIDDEN_DIM] — circular buffer +replay_scalars_d: CudaSlice [capacity × 7] — (action, scaled_reward, raw_reward, done, log_pi_old, n_step_gamma, n_step_return_raw) +priority_tree_d: CudaSlice [2 × capacity] — binary sum-tree (leaves at [capacity..2×capacity)) +write_head_d: CudaSlice [1] — circular write position +replay_len_d: CudaSlice [1] — number of valid entries (saturates at capacity) +max_priority_d: CudaSlice [1] — max TD-error seen (for new entries) +``` + +**N-step ring buffer (device-resident):** + +``` +nstep_ring_d: CudaSlice [B × N_STEP_MAX × (1+1+1)] — per-batch ring of (scaled_reward, raw_reward, done) +nstep_h_t_ring_d: CudaSlice [B × N_STEP_MAX × HIDDEN_DIM] — per-batch ring of h_t entries +nstep_write_idx_d: CudaSlice [B] — per-batch ring write index +nstep_count_d: CudaSlice [B] — per-batch ring fill count +``` + +### Kernels + +**1. `rl_per_push.cu`** — Push transitions + compute n-step return + +``` +Grid=(b_size, 1, 1), Block=(1, 1, 1) +``` + +Per thread (one per batch element): +1. Write current (scaled_reward, raw_reward, done) to nstep ring at `nstep_write_idx[b]` +2. DtoD copy `h_t[b*HIDDEN_DIM..(b+1)*HIDDEN_DIM]` into nstep_h_t_ring +3. Increment nstep_count[b], advance nstep_write_idx[b] +4. If `nstep_count[b] >= N_STEP` OR `done`: + - Compute n-step return: `R_n = Σ γᵏ rₖ` over the ring + - Compute n_step_gamma: `γⁿ` (zero if any done in window) + - Read oldest h_t from ring → write to `replay_h_t_d[write_head]` + - Copy current h_tp1 → `replay_h_tp1_d[write_head]` + - Write scalars to `replay_scalars_d[write_head]` + - Set `priority_tree_d[capacity + write_head] = max_priority_d[0]` + - Advance `write_head` (mod capacity) + - Increment `replay_len` (clamped at capacity) + - Clear the ring for this batch element +5. If not flushing: do nothing (accumulate in ring) + +**2. `rl_per_sample.cu`** — Priority-weighted sampling + gather + +``` +Grid=(b_size, 1, 1), Block=(1, 1, 1) +``` + +Per thread: +1. Compute segment boundaries: `segment_size = total_priority / b_size`, `lo = b * segment_size`, `hi = lo + segment_size` +2. Draw `u ~ U(lo, hi)` from device PRNG (xorshift32, same pattern as rl_sample_tau) +3. Walk the sum-tree top-down: + ``` + idx = 1 // root + while idx < capacity: + left = 2*idx + if u <= tree[left]: + idx = left + else: + u -= tree[left] + idx = left + 1 + leaf = idx - capacity + ``` +4. Write `sample_indices[b] = leaf` +5. Gather: copy `replay_h_t_d[leaf*H..(leaf+1)*H]` → `sampled_h_t_d[b*H..(b+1)*H]` +6. Same for h_tp1, scalars → sampled buffers +7. Compute IS weight: `w = (replay_len × P(leaf))^(-β)`, normalize later + +**3. `rl_per_tree_rebuild.cu`** — Bottom-up tree rebuild (no atomics) + +``` +Grid=(capacity/256, 1, 1), Block=(256, 1, 1) +``` + +After training step updates priorities at the leaves: +1. Each thread handles one internal node (bottom-up sweep) +2. `tree[i] = tree[2i] + tree[2i+1]` +3. Multiple passes from leaves to root (log2(capacity) = 15 passes for 32768) +4. Each pass is independent — sync between passes via separate kernel launches (one launch per level) or a single kernel with `__syncthreads` if capacity fits in one block + +For capacity=32768, the tree has 15 levels. Rebuild all internal nodes: +- Level 14 (16384 nodes): Grid=(64, Block=256) — one pass +- Level 13 (8192): Grid=(32, Block=256) +- ...down to level 0 (1 node): Grid=(1, Block=1) + +Alternative (simpler): single kernel, one thread per leaf-pair at the bottom, then iterate up with `__syncthreads` at each level. At capacity=32768 this needs 16384 threads in a single block — too many. Use the multi-launch approach (15 tiny kernel launches, trivial overhead vs the training step). + +Actually simplest: **single kernel, grid-stride, multiple passes.** Launch once with Grid=(128), Block=(256). Each of the 15 levels is a for-loop iteration with a `__threadfence()` between levels (device-wide fence, no atomics). All 32768 threads are overkill — the tree has 32767 internal nodes, 256 threads handle it in 128 iterations per level. + +**4. `rl_per_update_priority.cu`** — Write new TD errors to leaves + +``` +Grid=(b_size, 1, 1), Block=(1, 1, 1) +``` + +Per thread: +1. Read `sample_indices[b]` (from the sampling step) +2. Compute new priority: `|TD_error[b]|^α + ε` +3. Write to `priority_tree_d[capacity + sample_indices[b]]` +4. Update `max_priority_d[0] = max(max_priority, new_priority)` — single writer (b==0 does the reduce) + +Then launch `rl_per_tree_rebuild` to propagate. + +### Integration + +Replace in `IntegratedTrainer`: +- `self.replay: ReplayBuffer` (host) → `GpuReplayBuffer` (device) +- `push_to_replay()` → launch `rl_per_push` +- `sample_and_gather()` → launch `rl_per_sample` +- After training step: launch `rl_per_update_priority` + `rl_per_tree_rebuild` + +The `sampled_h_t_d`, `sampled_h_tp1_d`, `sampled_rewards_d`, etc. are already pre-allocated persistent fields. The sample kernel writes directly into them. + +--- + +## Expected Performance + +| Component | Before (host) | After (GPU) | Savings/step | +|-----------|---------------|-------------|--------------| +| Diag readback | 10× sync ~5ms | 0 sync (async) | −5ms | +| PER push | 256 alloc + DtoD ~8ms | 1 kernel ~0.1ms | −8ms | +| PER sample | 256 tree walks ~4ms | 1 kernel ~0.1ms | −4ms | +| PER priority update | host loop ~1ms | 2 kernels ~0.05ms | −1ms | +| **Total saved** | | | **~18ms/step** | + +At b=256 current: ~30ms/step = 34 sps. +After: ~12ms/step = **~83 sps** (2.4×). +With Graph C replay: ~10ms/step = **~100 sps**. + +At b=256 × 100 sps = **25,600 transitions/sec** (14× original baseline). + +--- + +## Kill Criteria + +- Async diag: sps identical with/without JSONL writing (zero throughput cost) +- GPU PER push: no host readback in push path (`read_slice_d` gone) +- GPU PER sample: sampled data appears directly in `sampled_*_d` buffers +- Tree rebuild: `total_priority` (tree root) matches host reference within fp32 noise +- End-to-end: ≥80 sps @ b=256 on L40S with full diag + +## Constraints + +- Per `feedback_no_atomicadd`: tree rebuild uses bottom-up parallel scan, NOT atomicAdd +- Per `feedback_cpu_is_read_only`: no host compute in push/sample/update path +- Per `feedback_no_htod_htoh_only_mapped_pinned`: diag uses mapped-pinned staging only +- Per `feedback_no_nvrtc`: all kernels pre-compiled cubins +- Per `feedback_no_stubs`: every kernel wired and functional in same commit