# 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
