spec: GPU-resident PER + async non-blocking diag streaming
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
225
docs/superpowers/specs/2026-05-25-gpu-per-async-diag.md
Normal file
225
docs/superpowers/specs/2026-05-25-gpu-per-async-diag.md
Normal file
@@ -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<f32> [capacity × HIDDEN_DIM] — circular buffer
|
||||
replay_h_tp1_d: CudaSlice<f32> [capacity × HIDDEN_DIM] — circular buffer
|
||||
replay_scalars_d: CudaSlice<f32> [capacity × 7] — (action, scaled_reward, raw_reward, done, log_pi_old, n_step_gamma, n_step_return_raw)
|
||||
priority_tree_d: CudaSlice<f32> [2 × capacity] — binary sum-tree (leaves at [capacity..2×capacity))
|
||||
write_head_d: CudaSlice<u32> [1] — circular write position
|
||||
replay_len_d: CudaSlice<u32> [1] — number of valid entries (saturates at capacity)
|
||||
max_priority_d: CudaSlice<f32> [1] — max TD-error seen (for new entries)
|
||||
```
|
||||
|
||||
**N-step ring buffer (device-resident):**
|
||||
|
||||
```
|
||||
nstep_ring_d: CudaSlice<f32> [B × N_STEP_MAX × (1+1+1)] — per-batch ring of (scaled_reward, raw_reward, done)
|
||||
nstep_h_t_ring_d: CudaSlice<f32> [B × N_STEP_MAX × HIDDEN_DIM] — per-batch ring of h_t entries
|
||||
nstep_write_idx_d: CudaSlice<u32> [B] — per-batch ring write index
|
||||
nstep_count_d: CudaSlice<u32> [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
|
||||
Reference in New Issue
Block a user