gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

spec: update CUDA graph spec with all session findings and implementation plan

Browse Source 
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-25 21:41:10 +02:00
parent a7c1d763d8
commit 3be9996689
1 changed files with 199 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files

327
docs/superpowers/specs/2026-05-25-cuda-graph-step-pipeline.md
View File

@@ -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. > **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 ## 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`.
- **No pointer swaps inside capture.** All device buffers must have stable addresses.
- **No host branches inside captured region.** Per `pearl_no_host_branches_in_captured_graph`. All kernel launches must be unconditional. - **No scalar arg changes inside capture.** All args are ISV-driven or constant.
- **No pointer swaps inside capture.** All device buffers must have stable addresses across replays. - **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 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-side malloc during capture.** All buffers pre-allocated at init. - **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. rl_increment_step → ISV[548] += 1
1. forward_encoder (ALREADY GRAPHED in perception trainer) 2. dqn_head.forward → q_logits_d
2. dqn_head.forward → q_logits_d 3. iqn_head.forward (tau sample) → iqn_q_values_d
3. iqn_head.forward → iqn_q_values_d 4. iqn_head.expected_q → iqn_expected_q_d
4. iqn_head.expected_q → iqn_expected_q_d 5. rl_ensemble_action_value → ensemble_q_d
5. ensemble_action_value → ensemble_q_d 6. noisy_exploration.resample+fwd → noise added to ensemble_q_d
6. value_head.forward → v_d 7. value_head.forward → v_d
7. q_pi_agree_b → ISV[407] 8. rl_q_pi_agree_b → ISV[407]
8. q_pi_distill_grad → pi_grad_logits 9. rl_q_pi_distill_grad → pi_grad_logits
9. pi_action_kernel → actions_d 10. rl_pi_action_kernel → actions_d
10. argmax_expected_q → next_actions_d 11. argmax_expected_q → next_actions_d
11. log_pi_at_action → log_pi_old_d 12. log_pi_at_action → log_pi_old_d
12. confidence_gate → actions_d (override) 13. rl_confidence_gate → actions_d (override)
13. frd_gate → actions_d (override) 14. rl_frd_gate → actions_d (override)
14. session_risk_check → actions_d (override) 15. rl_session_risk_check → actions_d (override)
15. min_hold_check → actions_d (override) 16. rl_min_hold_check → actions_d (override)
16. asymmetric_trail_decay → unit_trail_distance_d 17. rl_asymmetric_trail_decay → unit_trail_distance_d
17. trail_mutate → unit_trail_distance_d 18. rl_trail_mutate → unit_trail_distance_d
18. trail_stop_check → actions_d (override) 19. rl_trail_stop_check → actions_d (override)
19. heat_cap_check → actions_d (override) 20. rl_position_heat_check → actions_d (override)
20. actions_to_market_targets → market_targets_d 21. 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
``` ```
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) ### Uncapturable: lobsim fill
- 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)
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) This is the LobSim simulator with internal state. Cannot be captured.
- Input: h_t (from encoder, already graphed separately)
- Output: market_targets_d
**Graph A2: "post_fill"** — kernels 21-36 (extract PnL through controllers) ### Graph B: Post-fill (PnL extraction → controllers)
- Input: pos_state_d (from lobsim fill)
- Output: rewards_d, dones_d, ISV updates
**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 22. extract_realized_pnl_delta → rewards_d, dones_d
- Output: weight updates, grad_h_t 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 ### Uncapturable: PER push + sample
- **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)
### 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: Host-side PER operations with variable indices. Cannot be captured.
- `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
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. ✓ **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.
### 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).
## Implementation Plan ## 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: Replace host-side RNG for IQN tau and noisy exploration with device-side xorshift32 (matching `rl_pi_action_kernel`'s pattern).
- `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
### 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 **New kernel: `rl_sample_noise.cu`**
2. Warmup: run one step without capture ```cuda
3. Capture: `cudaStreamBeginCapture` → launch kernels 2-20 → `cudaStreamEndCapture` // Generate factored noise for NoisyLinear.
4. Instantiate `cudaGraphExec_t` // Uses device-side PRNG → f(x) = sign(x) × sqrt(|x|).
5. Each step: `cudaGraphLaunch(graphExec, stream)` 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<CudaGraphExec>,
### 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 ## Expected Performance
| Metric | Current | With Graphs | | Metric | Current | With Graphs |
|--------|---------|-------------| |--------|---------|-------------|
| Launches/step | ~50 | ~5 (3 graphs + lobsim + PER) | | Launches/step | ~55 | ~8 (3 graphs + encoder + lobsim + PER + sync) |
| Launch overhead/step | ~300μs | ~30μs | | Launch overhead/step | ~300μs | ~50μs |
| Step time (b=16) | ~850μs | ~550μs | | Step time (b=16) | ~170μs/step (5.8k steps/sec) | ~110μs/step (~9k steps/sec) |
| Steps/sec | ~1,180 | ~1,800 | | 1M run time | ~2.9h | **~1.9h** |
| 1M run time | ~14 min | **~9 min** |
## Kill Criteria ## Kill Criteria
- Graph capture succeeds without `STREAM_CAPTURE_ISOLATION` errors - Graph capture succeeds without `STREAM_CAPTURE_ISOLATION` errors
- Replay produces identical results to non-graphed execution (bit-exact ISV values after 100 steps) - Replay produces identical Q-loss trajectory to non-graphed execution (within fp noise)
- Throughput improves ≥ 1.5× at b=16 - 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.