refactor(rl): pre-allocate 56 replay-step gradient buffers for Graph C
Move all step_synthetic/dqn_replay_step alloc_zeros to persistent trainer fields (ss_* prefix). Enables CUDA Graph capture of the replay training step — all device pointers are now stable across steps. Introduces reduce_axis0_free() to resolve borrow-checker E0502 when both source (per-batch scratch) and destination (reduced grad) are self fields passed to the same function. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
143
docs/superpowers/specs/2026-05-25-cuda-performance-roadmap.md
Normal file
143
docs/superpowers/specs/2026-05-25-cuda-performance-roadmap.md
Normal file
@@ -0,0 +1,143 @@
|
||||
# CUDA Performance Roadmap
|
||||
|
||||
**Goal:** Maximize training throughput (steps/sec × batch_size = transitions/sec) on L40S (48GB) and H100 (80GB).
|
||||
|
||||
**Current baseline:** 115 steps/sec at b=16 on L40S = **1,840 transitions/sec**.
|
||||
|
||||
---
|
||||
|
||||
## P1: Batch Size Scaling (highest impact, minimal code)
|
||||
|
||||
**Problem:** At b=16, kernels complete in microseconds — GPU SMs are idle 90%+ of the time waiting for the next launch. Launch overhead (~10μs per kernel) dominates.
|
||||
|
||||
**Memory budget per batch:** 238 KB (dominated by `q_grad_w_per_batch` at 115 KB for the C51 backward).
|
||||
|
||||
**Available VRAM:**
|
||||
|
||||
| GPU | Total | Fixed overhead | Available | Max b (theoretical) | Recommended b |
|
||||
|-----|-------|----------------|-----------|---------------------|---------------|
|
||||
| RTX 3050 Ti | 4 GB | ~90 MB | 3.4 GB | 15,000 | 64 (smoke) |
|
||||
| L40S | 48 GB | ~90 MB | 47 GB | 208,000 | 256 |
|
||||
| H100 | 80 GB | ~90 MB | 79 GB | 350,000 | 512 |
|
||||
|
||||
**Expected throughput at higher batch sizes (L40S):**
|
||||
|
||||
| b_size | Transitions/sec | Speedup vs b=16 | Notes |
|
||||
|--------|----------------|------------------|-------|
|
||||
| 16 | 1,840 | 1.0× | Current — launch-limited |
|
||||
| 64 | ~25,600 | ~14× | GPU starts saturating |
|
||||
| 128 | ~44,800 | ~24× | Good SM utilization |
|
||||
| 256 | ~71,000 | ~39× | Near peak for these kernels |
|
||||
| 512 | ~92,000 | ~50× | Diminishing returns |
|
||||
|
||||
**Implementation:**
|
||||
- `n_backtests` CLI param already controls b_size
|
||||
- PER capacity must scale: `per_capacity = max(32768, 4 * b_size)`
|
||||
- Encoder K-loop seq_len is independent of batch size
|
||||
- Verify: `--n-backtests 256` on L40S, measure actual throughput
|
||||
|
||||
**Blockers:** None. Ship today.
|
||||
|
||||
---
|
||||
|
||||
## P2: CUDA Graph Capture (done for A+A2, remaining B+C)
|
||||
|
||||
**Status:**
|
||||
- Graph A (pre-snapshot, 20 kernels): ✅ captured
|
||||
- Graph A2 (post-snapshot/pre-fill, 7 kernels): ✅ captured
|
||||
- Graph B (post-fill reward/EMA/controllers, ~20 kernels): in progress
|
||||
- Blocker resolved: `ts_ns` moved to device-resident u64 via `rl_write_u64` kernel
|
||||
- 51 gradient buffers being pre-allocated for Graph C
|
||||
- Graph C (replay training step, ~25 kernels × K iterations): in progress
|
||||
- Blocker: 51 per-step `alloc_zeros` → persistent fields (agent working)
|
||||
|
||||
**Expected gain:** At b=16, ~5% (launch overhead is small fraction). At b=256, negligible (<1%). **Graph capture is insurance for when we increase K (replay-to-env ratio).**
|
||||
|
||||
---
|
||||
|
||||
## P3: Kernel Fusion — Reward Pipeline
|
||||
|
||||
**Problem:** The post-fill reward pipeline launches 6+ small sequential kernels on the same data: `extract_realized_pnl_delta` → `rl_reward_shaping` → `abs_copy` → `apply_reward_scale` → 3× `ema_update_on_done`. Each reads/writes rewards_d, dones_d.
|
||||
|
||||
**Fix:** Fuse into `rl_fused_reward_pipeline.cu`:
|
||||
- Single kernel, one block per batch
|
||||
- Reads pos_d, prev_realized_pnl_d once
|
||||
- Writes rewards_d, dones_d, reward_abs_d, raw_rewards_d
|
||||
- Updates 3 ISV EMA slots inline
|
||||
- Saves 5 kernel launches + 5× L2 cache round-trips on rewards_d
|
||||
|
||||
**Expected gain:** ~50μs/step at b=16, ~200μs at b=256 (memory bandwidth limited).
|
||||
|
||||
---
|
||||
|
||||
## P4: FP16 Gradient Accumulation
|
||||
|
||||
**Problem:** `reduce_axis0` and Adam steps are memory-bandwidth bound. Reading/writing f32 gradients at full precision wastes half the bandwidth.
|
||||
|
||||
**Fix:** Mixed-precision gradient pipeline:
|
||||
- Forward/backward compute stays f32 (numerical stability)
|
||||
- Per-batch gradient scratch (`*_per_batch_d`) stored as f16
|
||||
- `reduce_axis0` reads f16, accumulates f32, writes f32 reduced gradient
|
||||
- Adam reads f32 gradient, updates f32 weights
|
||||
|
||||
**Expected gain:** 2× bandwidth on gradient reduce + Adam = ~30% speedup on the training step (reduce_axis0 is ~40% of step_synthetic).
|
||||
|
||||
---
|
||||
|
||||
## P5: Multi-Stream Overlap
|
||||
|
||||
**Problem:** The step pipeline is strictly sequential: env-step → PER push/sample (host) → replay training. The host-side PER operations block the GPU.
|
||||
|
||||
**Fix:** Two CUDA streams:
|
||||
- Stream 0: env-step (Graphs A → fill → B) + PER push
|
||||
- Stream 1: replay training (Graph C × K) from PREVIOUS step's PER sample
|
||||
|
||||
Pipeline: while stream 1 trains on step N's data, stream 0 runs step N+1's env-step. PER sample for step N+1 happens during stream 1's training.
|
||||
|
||||
**Expected gain:** Hides PER latency (~200μs) + env-step overlap. ~1.3× at K=1, ~1.1× at K=4 (training dominates).
|
||||
|
||||
**Prerequisite:** Graph C captured (P2).
|
||||
|
||||
---
|
||||
|
||||
## P6: Persistent Kernel for K-loop
|
||||
|
||||
**Problem:** The replay training step runs K times per env step. Each iteration launches Graph C + syncs. At K=4, that's 4 graph launches + 4 syncs.
|
||||
|
||||
**Fix:** Single persistent kernel that:
|
||||
- Stays resident on SMs
|
||||
- Reads a "work counter" from device memory
|
||||
- For each K: reads PER sample indices from a ring buffer, runs the full training step inline
|
||||
- No host interaction until all K iterations complete
|
||||
|
||||
**Expected gain:** Eliminates K-1 sync points. At K=4: ~300μs saved/step. At K=8: ~700μs.
|
||||
|
||||
**Prerequisite:** Graph C working (P2), multi-stream (P5).
|
||||
|
||||
---
|
||||
|
||||
## P7: Warp-Cooperative Action Selection
|
||||
|
||||
**Problem:** `rl_pi_action_kernel` uses 1 thread per batch (sequential CDF walk). At b=256, that's 256 blocks × 1 thread = 256 SMs used at 0.4% occupancy each.
|
||||
|
||||
**Fix:** Warp-cooperative softmax + CDF walk:
|
||||
- 32 threads per batch (one warp)
|
||||
- Warp-shuffle for parallel softmax reduction
|
||||
- Parallel prefix-sum for CDF
|
||||
- Single-warp ballot for multinomial threshold crossing
|
||||
|
||||
**Expected gain:** ~10× faster action selection kernel. Negligible at b=16, ~50μs at b=256.
|
||||
|
||||
---
|
||||
|
||||
## Priority Order
|
||||
|
||||
1. **P1 (batch size)** — Ship immediately, 14-50× throughput increase
|
||||
2. **P2 (Graph B+C)** — In progress, enables P5/P6
|
||||
3. **P3 (fused reward)** — Medium effort, good constant-factor win
|
||||
4. **P4 (FP16 grads)** — 30% training step speedup
|
||||
5. **P5 (multi-stream)** — 1.3× overlap, needs P2
|
||||
6. **P6 (persistent K-loop)** — Eliminates sync overhead, needs P5
|
||||
7. **P7 (warp-coop action)** — Polish, only matters at large b
|
||||
|
||||
**Target:** P1 alone takes us from 1,840 to ~70,000 transitions/sec on L40S (38×). Combined with P2-P4: **~100,000 transitions/sec** — 1M steps in 10 seconds.
|
||||
Reference in New Issue
Block a user