spec: CUDA Graph capture for RL step pipeline (2-3× throughput)
Captures 50+ kernel launches as 3 CUDA Graphs (pre-fill, post-fill, replay-step) for single-launch replay. Eliminates ~300μs/step of CPU launch overhead at b=16. Key challenges: device-resident step counter (no scalar arg changes in graph), lobsim fill breaks the graph (split into pre/post), ISV write ordering. Estimated 1.5-3× throughput improvement. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
183
docs/superpowers/specs/2026-05-25-cuda-graph-step-pipeline.md
Normal file
183
docs/superpowers/specs/2026-05-25-cuda-graph-step-pipeline.md
Normal file
@@ -0,0 +1,183 @@
|
||||
# CUDA Graph Capture for the RL Step Pipeline
|
||||
|
||||
> **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.
|
||||
|
||||
**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`.
|
||||
|
||||
**Tech Stack:** cudarc 0.19 CUDA Graph API, Rust, pre-allocated device buffers
|
||||
|
||||
---
|
||||
|
||||
## 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-side malloc during capture.** All buffers pre-allocated at init.
|
||||
|
||||
---
|
||||
|
||||
## Current Pipeline (40+ launches per step)
|
||||
|
||||
```
|
||||
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
|
||||
```
|
||||
|
||||
Total: **50+ kernel launches per step** (more with K-loop replay iterations).
|
||||
|
||||
## Graph Capture Strategy
|
||||
|
||||
### What to capture
|
||||
|
||||
Two separate graphs:
|
||||
|
||||
**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)
|
||||
|
||||
Actually, the lobsim fill in the middle breaks the graph. Split into:
|
||||
|
||||
**Graph A1: "pre_fill"** — kernels 2-20 (Q forward through market_targets)
|
||||
- Input: h_t (from encoder, already graphed separately)
|
||||
- Output: market_targets_d
|
||||
|
||||
**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: "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
|
||||
|
||||
### What NOT to capture
|
||||
|
||||
- **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)
|
||||
|
||||
### Conditional kernels problem
|
||||
|
||||
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
|
||||
|
||||
These use `current_step` which changes each step — **this is a scalar arg change** which violates graph capture rules.
|
||||
|
||||
**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.
|
||||
|
||||
### Pointer stability
|
||||
|
||||
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).
|
||||
|
||||
## Implementation Plan
|
||||
|
||||
### Phase 1: Device-resident step counter
|
||||
|
||||
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
|
||||
|
||||
### Phase 2: Graph A1 (pre-fill)
|
||||
|
||||
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)`
|
||||
|
||||
### Phase 3: Graph A2 (post-fill)
|
||||
|
||||
Same pattern for kernels 21-36.
|
||||
|
||||
### Phase 4: Graph B (replay step)
|
||||
|
||||
Capture one iteration of the replay training loop. Replay K times per step via `cudaGraphLaunch`.
|
||||
|
||||
### Phase 5: Fuse controllers
|
||||
|
||||
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.
|
||||
|
||||
## 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** |
|
||||
|
||||
## 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
|
||||
Reference in New Issue
Block a user