From 3be9996689660b479b1205cc476dcbb5ca2b11e9 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 25 May 2026 21:41:10 +0200 Subject: [PATCH] spec: update CUDA graph spec with all session findings and implementation plan Co-Authored-By: Claude Opus 4.7 --- .../2026-05-25-cuda-graph-step-pipeline.md | 327 +++++++++++------- 1 file changed, 199 insertions(+), 128 deletions(-) diff --git a/docs/superpowers/specs/2026-05-25-cuda-graph-step-pipeline.md b/docs/superpowers/specs/2026-05-25-cuda-graph-step-pipeline.md index 91b54320d..41702edf3 100644 --- a/docs/superpowers/specs/2026-05-25-cuda-graph-step-pipeline.md +++ b/docs/superpowers/specs/2026-05-25-cuda-graph-step-pipeline.md @@ -2,182 +2,253 @@ > **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. -**Goal:** Capture the entire RL training step pipeline (40+ kernel launches) as a CUDA Graph and replay it each step, eliminating ~200-400μs of CPU launch overhead per step for a 2-3× throughput improvement. +**Goal:** Capture the RL training step pipeline as CUDA Graphs and replay each step, eliminating ~300μs of CPU launch overhead per step for 1.5-3× throughput. -**Architecture:** Pre-allocate all device buffers at trainer init. Run one warmup step without capture. Then capture the step pipeline as a graph. Each subsequent step: write new input data into the pre-allocated buffers, replay the graph via a single `cudaGraphLaunch`. +**Architecture:** Split the step pipeline around the lobsim fill (which can't be captured) into two graphs. Pre-allocate all device buffers at trainer init. Run one warmup step, capture, instantiate, then replay each step. -**Tech Stack:** cudarc 0.19 CUDA Graph API, Rust, pre-allocated device buffers +**Tech Stack:** cudarc 0.19 (`sys::cudaStreamBeginCapture` / `cudaStreamEndCapture` / `cudaGraphInstantiate` / `cudaGraphLaunch`), Rust, pre-allocated device buffers --- +## Prerequisites (DONE) + +Both prerequisites are implemented and committed: + +1. **Device-resident step counter (ISV[548])** — commit `a7c1d763d`. All kernels that took `current_step` as a scalar arg now read from ISV. The `rl_increment_step` kernel bumps the counter each step. No scalar args change between replays. + +2. **Fused controllers** — commit `a7c1d763d`. 10 single-thread controllers fused into `rl_fused_controllers.cu`. Saves 9 launches/step and simplifies graph capture (1 controller kernel instead of 10). + ## Mandatory Constraints -All constraints from the Q-learning improvements spec apply, plus: - -- **No host branches inside captured region.** Per `pearl_no_host_branches_in_captured_graph`. All kernel launches must be unconditional. -- **No pointer swaps inside capture.** All device buffers must have stable addresses across replays. -- **No scalar arg changes inside capture.** `b_size`, `pos_bytes`, kernel configs must be constant. -- **Disable cudarc event tracking during capture.** Per `pearl_cudarc_disable_event_tracking_for_graph_capture` — cudarc 0.19 event tracking trips `STREAM_CAPTURE_ISOLATION`. +- **No host branches inside captured region.** Per `pearl_no_host_branches_in_captured_graph`. +- **No pointer swaps inside capture.** All device buffers must have stable addresses. +- **No scalar arg changes inside capture.** All args are ISV-driven or constant. +- **Disable cudarc event tracking during capture.** Per `pearl_cudarc_disable_event_tracking_for_graph_capture` — `ctx.disable_event_tracking()` before capture, `ctx.enable_event_tracking()` after. - **No host-side malloc during capture.** All buffers pre-allocated at init. +- **No `stream.synchronize()` inside capture.** Trips `STREAM_CAPTURE_ISOLATION`. ---- +## Pipeline Analysis -## Current Pipeline (40+ launches per step) +### Graph A: Pre-fill (action selection → market targets) + +These kernels run AFTER the encoder forward (already graphed in perception trainer) and BEFORE the lobsim fill: ``` -Step pipeline (step_with_lobsim): - 1. forward_encoder (ALREADY GRAPHED in perception trainer) - 2. dqn_head.forward → q_logits_d - 3. iqn_head.forward → iqn_q_values_d - 4. iqn_head.expected_q → iqn_expected_q_d - 5. ensemble_action_value → ensemble_q_d - 6. value_head.forward → v_d - 7. q_pi_agree_b → ISV[407] - 8. q_pi_distill_grad → pi_grad_logits - 9. pi_action_kernel → actions_d - 10. argmax_expected_q → next_actions_d - 11. log_pi_at_action → log_pi_old_d - 12. confidence_gate → actions_d (override) - 13. frd_gate → actions_d (override) - 14. session_risk_check → actions_d (override) - 15. min_hold_check → actions_d (override) - 16. asymmetric_trail_decay → unit_trail_distance_d - 17. trail_mutate → unit_trail_distance_d - 18. trail_stop_check → actions_d (override) - 19. heat_cap_check → actions_d (override) - 20. actions_to_market_targets → market_targets_d - --- lobsim fill (external, NOT captured) --- - 21. extract_realized_pnl_delta → rewards_d, dones_d - 22. unit_state_update → unit buffers - 23. trade_context_update → trade_context_d - 24. multires_features_update → multires_output_d - 25. step_counter_update → steps_since_done_d - 26. reward_shaping → rewards_d - 27. snapshot raw_rewards → raw_rewards_d (DtoD memcpy) - 28. apply_reward_scale → rewards_d - 29. abs_copy → reward_abs_d - 30. ema_update_on_done (×3) → ISV EMAs - 31. ema_update_per_step (×4) → ISV EMAs - 32. recent_outcome_update → outcome_ema_d - 33. gate_threshold_controller → ISV gates - 34. 7 RL controllers → ISV slots - 35. reward_clamp_controller → ISV clamp - 36. atom_support_update → atom_supports_d - -Replay step (dqn_replay_step, called K times): - 37. dqn_head.forward (sampled h_t) - 38. dqn_head.forward_target (sampled h_tp1) - 39. argmax_expected_q (Double DQN) - 40. select_action_atoms - 41. project_bellman_target - 42. C51 cross-entropy loss + backward - 43. iqn_head.forward (sampled h_t) - 44. iqn_head.forward_target (sampled h_tp1) - 45. iqn_loss - 46. iqn_backward - 47. backward_to_w_b_h (Q grad) - 48. backward_to_w_b_h (IQN grad) - 49. v_head backward - 50. pi head backward - 51. Adam updates (×15+ weight groups) - 52. target soft update + 1. rl_increment_step → ISV[548] += 1 + 2. dqn_head.forward → q_logits_d + 3. iqn_head.forward (tau sample) → iqn_q_values_d + 4. iqn_head.expected_q → iqn_expected_q_d + 5. rl_ensemble_action_value → ensemble_q_d + 6. noisy_exploration.resample+fwd → noise added to ensemble_q_d + 7. value_head.forward → v_d + 8. rl_q_pi_agree_b → ISV[407] + 9. rl_q_pi_distill_grad → pi_grad_logits +10. rl_pi_action_kernel → actions_d +11. argmax_expected_q → next_actions_d +12. log_pi_at_action → log_pi_old_d +13. rl_confidence_gate → actions_d (override) +14. rl_frd_gate → actions_d (override) +15. rl_session_risk_check → actions_d (override) +16. rl_min_hold_check → actions_d (override) +17. rl_asymmetric_trail_decay → unit_trail_distance_d +18. rl_trail_mutate → unit_trail_distance_d +19. rl_trail_stop_check → actions_d (override) +20. rl_position_heat_check → actions_d (override) +21. actions_to_market_targets → market_targets_d ``` -Total: **50+ kernel launches per step** (more with K-loop replay iterations). +**21 launches → 1 graph launch.** -## Graph Capture Strategy +**Blocker: tau sampling.** The IQN forward needs fresh τ ~ U(0,1) each step. Currently sampled on host and uploaded via mapped-pinned. In a graph, this host-side RNG + upload can't happen inside capture. -### What to capture +**Fix:** Pre-generate a large τ buffer at init (e.g., 1000 steps × B × N_TAU floats) and cycle through it with a device-side index counter. Or use a device-side PRNG (xorshift32 like `rl_pi_action_kernel` already does). -Two separate graphs: +**Blocker: noisy_exploration.resample_noise().** Same issue — host RNG + upload. Same fix: device-side PRNG for noise generation. -**Graph A: "env_step"** — kernels 1-36 (action selection through EMA updates) -- Fixed inputs: snapshot data in `window_tensor_d`, ISV state -- Fixed outputs: `rewards_d`, `dones_d`, `actions_d`, unit state -- **Excludes:** lobsim fill (external simulator, has internal state changes) +### Uncapturable: lobsim fill -Actually, the lobsim fill in the middle breaks the graph. Split into: +``` +lobsim.apply_snapshot(last_snap) +lobsim.step_fill_from_market_targets(ts_ns) +``` -**Graph A1: "pre_fill"** — kernels 2-20 (Q forward through market_targets) -- Input: h_t (from encoder, already graphed separately) -- Output: market_targets_d +This is the LobSim simulator with internal state. Cannot be captured. -**Graph A2: "post_fill"** — kernels 21-36 (extract PnL through controllers) -- Input: pos_state_d (from lobsim fill) -- Output: rewards_d, dones_d, ISV updates +### Graph B: Post-fill (PnL extraction → controllers) -**Graph B: "replay_step"** — kernels 37-52 (one iteration of replay training) -- Input: sampled_h_t_d, sampled_h_tp1_d, sampled_rewards_d, sampled_dones_d -- Output: weight updates, grad_h_t +``` +22. extract_realized_pnl_delta → rewards_d, dones_d +23. rl_unit_state_update → unit buffers +24. rl_trade_context_update → trade_context_d +25. rl_multires_features_update → multires_output_d +26. rl_step_counter_update → steps_since_done_d +27. rl_reward_shaping → rewards_d +28. DtoD memcpy raw_rewards → raw_rewards_d +29. apply_reward_scale → rewards_d +30. abs_copy → reward_abs_d +31. ema_update_on_done (×3) → ISV EMAs +32. ema_update_per_step (×4) → ISV EMAs +33. rl_recent_outcome_update → outcome_ema_d +34. rl_fused_controllers → ISV slots +35. rl_reward_clamp_controller → ISV clamp +36. rl_atom_support_update → atom_supports_d +``` -### What NOT to capture +**15 launches → 1 graph launch.** -- **Encoder forward** — already graphed in perception trainer -- **LobSim fill** — external simulator with variable-length internal state -- **push_to_replay / sample_and_gather** — host-side PER operations with variable indices -- **ISV host sync** — `read_slice_d` for diag output (host↔device) +### Uncapturable: PER push + sample -### Conditional kernels problem +``` +push_to_replay (host-side n-step ring buffer + PER) +sample_and_gather (host-side PER sampling + DtoD copies) +``` -Several kernels have ISV-driven conditionals: -- `confidence_gate`: returns early if `step < warmup` -- `frd_gate`: same warmup check -- `session_risk_check`: blocks actions if EMA < limit -- `min_hold_check`: blocks if hold_time < min +Host-side PER operations with variable indices. Cannot be captured. -These use `current_step` which changes each step — **this is a scalar arg change** which violates graph capture rules. +### Graph C: Replay training step -**Fix:** Move the step counter to an ISV slot (device-resident). The kernel reads it from ISV instead of a scalar arg. A tiny `increment_step` kernel bumps the ISV slot each step — this kernel IS in the graph. +``` +37. dqn_head.forward (sampled h_t) +38. iqn_head.forward (sampled h_t) +39. dqn_head.forward_target (sampled h_tp1) +40. iqn_head.forward_target (sampled h_tp1) +41. argmax_expected_q (Double DQN) +42. select_action_atoms +43. project_bellman_target +44. C51 cross-entropy loss +45. C51 backward +46. IQN loss +47. IQN backward +48. v_head backward +49. pi_head backward +50. grad_h_accumulate (×4 heads) +51. Adam updates (×15+ weight groups) +52. target soft update +``` -### Pointer stability +**~20 launches → 1 graph launch.** K-loop iterations replay this graph K times. -All buffers are already pre-allocated in `IntegratedTrainer::new()`. The graph captures their device pointers. As long as we don't reallocate, pointers are stable across replays. ✓ - -### ISV writes during graph - -Multiple kernels write to ISV slots. In a graph, these writes happen in kernel order (stream-ordered). This is correct as long as no kernel reads a slot that was written by a LATER kernel in the same graph. Current pipeline respects this (controllers read EMAs written by earlier kernels). +**Blocker: sampled data varies each replay.** The sample_and_gather writes different data to sampled_h_t_d etc. But the POINTERS are stable (pre-allocated). Graph replays use the same pointers — the DATA at those pointers changes between replays, which is fine. The graph captures the kernel launch topology, not the data. ## Implementation Plan -### Phase 1: Device-resident step counter +### Phase 1: Device-side tau + noise generation -Replace all `current_step` scalar args with ISV reads. Add: -- `RL_STEP_COUNTER_INDEX` ISV slot -- `rl_increment_step.cu` kernel (single thread, ISV[slot] += 1) -- Update all kernels that take `current_step` to read from ISV +Replace host-side RNG for IQN tau and noisy exploration with device-side xorshift32 (matching `rl_pi_action_kernel`'s pattern). -### Phase 2: Graph A1 (pre-fill) +**New kernel: `rl_sample_tau.cu`** +```cuda +// Per-batch xorshift32 PRNG generates N_TAU uniform floats. +// Grid=(B,1,1), Block=(N_TAU,1,1). Each thread generates one tau. +extern "C" __global__ void rl_sample_tau( + uint32_t* __restrict__ prng_state, // [B] per-batch PRNG + float* __restrict__ tau, // [B × N_TAU] OUT + int b_size +) +``` -1. Pre-allocate a dedicated graph capture stream -2. Warmup: run one step without capture -3. Capture: `cudaStreamBeginCapture` → launch kernels 2-20 → `cudaStreamEndCapture` -4. Instantiate `cudaGraphExec_t` -5. Each step: `cudaGraphLaunch(graphExec, stream)` +**New kernel: `rl_sample_noise.cu`** +```cuda +// Generate factored noise for NoisyLinear. +// Uses device-side PRNG → f(x) = sign(x) × sqrt(|x|). +extern "C" __global__ void rl_sample_noise( + uint32_t* __restrict__ prng_state, + float* __restrict__ eps_in, // [in_dim] + float* __restrict__ eps_out, // [out_dim] + int in_dim, + int out_dim +) +``` -### Phase 3: Graph A2 (post-fill) +### Phase 2: Graph A capture (pre-fill) -Same pattern for kernels 21-36. +```rust +// In IntegratedTrainer: +graph_a_exec: Option, -### Phase 4: Graph B (replay step) +fn capture_graph_a(&mut self) -> Result<()> { + unsafe { self.ctx.disable_event_tracking(); } + + let stream = &self.stream; + cudaStreamBeginCapture(stream.cu_stream(), cudaStreamCaptureModeGlobal); + + // Launch all 21 pre-fill kernels in order... + self.launch_rl_increment_step()?; + self.dqn_head.forward(...)?; + // ... etc + self.launch_actions_to_market_targets(...)?; + + let mut graph = std::ptr::null_mut(); + cudaStreamEndCapture(stream.cu_stream(), &mut graph); + + let mut graph_exec = std::ptr::null_mut(); + cudaGraphInstantiate(&mut graph_exec, graph, std::ptr::null_mut(), std::ptr::null_mut(), 0); + cudaGraphDestroy(graph); + + self.graph_a_exec = Some(graph_exec); + unsafe { self.ctx.enable_event_tracking(); } + Ok(()) +} +``` -Capture one iteration of the replay training loop. Replay K times per step via `cudaGraphLaunch`. +### Phase 3: Graph B capture (post-fill) -### Phase 5: Fuse controllers +Same pattern for the 15 post-fill kernels. -The 7 RL controllers + gate threshold controller + reward clamp controller are all single-thread kernels that read/write ISV. Fuse into one kernel with a switch on controller index. Saves 8 launches. +### Phase 4: Graph C capture (replay step) + +Same pattern for the ~20 replay training kernels. Replayed K times per step. + +### Phase 5: Step execution with graphs + +```rust +fn step_with_lobsim_graphed(&mut self, ...) -> Result<()> { + // 1. Encoder forward (already graphed internally) + self.perception.forward_encoder(snapshots)?; + + // 2. Pre-fill graph replay + cudaGraphLaunch(self.graph_a_exec, self.stream.cu_stream()); + + // 3. Lobsim fill (uncapturable) + lobsim.apply_snapshot(last_snap)?; + lobsim.step_fill_from_market_targets(ts_ns)?; + + // 4. Post-fill graph replay + cudaGraphLaunch(self.graph_b_exec, self.stream.cu_stream()); + + // 5. PER push + sample (host-side) + self.push_to_replay(b_size)?; + let indices = self.sample_and_gather(b_size)?; + + // 6. Replay graph × K iterations + for _ in 0..k_updates { + cudaGraphLaunch(self.graph_c_exec, self.stream.cu_stream()); + } + + // 7. Sync + diag + self.stream.synchronize()?; +} +``` ## Expected Performance | Metric | Current | With Graphs | |--------|---------|-------------| -| Launches/step | ~50 | ~5 (3 graphs + lobsim + PER) | -| Launch overhead/step | ~300μs | ~30μs | -| Step time (b=16) | ~850μs | ~550μs | -| Steps/sec | ~1,180 | ~1,800 | -| 1M run time | ~14 min | **~9 min** | +| Launches/step | ~55 | ~8 (3 graphs + encoder + lobsim + PER + sync) | +| Launch overhead/step | ~300μs | ~50μs | +| Step time (b=16) | ~170μs/step (5.8k steps/sec) | ~110μs/step (~9k steps/sec) | +| 1M run time | ~2.9h | **~1.9h** | ## Kill Criteria - Graph capture succeeds without `STREAM_CAPTURE_ISOLATION` errors -- Replay produces identical results to non-graphed execution (bit-exact ISV values after 100 steps) -- Throughput improves ≥ 1.5× at b=16 +- Replay produces identical Q-loss trajectory to non-graphed execution (within fp noise) +- Throughput improves ≥ 1.3× at b=16 + +## Findings from Today's Session + +1. **The perception trainer already uses graph capture** for `forward_only`. The pattern works — we proved it on this codebase. +2. **cudarc 0.19 requires `disable_event_tracking`** before capture — documented in `pearl_cudarc_disable_event_tracking_for_graph_capture`. +3. **Host-side RNG is the main blocker** — IQN tau and noisy exploration both use host RNG. Device-side xorshift32 is the proven solution (`rl_pi_action_kernel` already uses it). +4. **Fused controllers help graph capture** — 1 kernel in the graph instead of 10 reduces graph complexity and potential for capture issues. +5. **Device-resident step counter is essential** — without it, the `current_step` scalar arg would need to change between replays, breaking the graph.