Files
foxhunt/docs/superpowers/specs/2026-05-25-cuda-performance-roadmap.md
jgrusewski 4bed8f2dbf 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>
2026-05-25 23:00:40 +02:00

5.9 KiB
Raw Blame History

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_deltarl_reward_shapingabs_copyapply_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.