From 6308be794e43b1208523b73c4bda1f2d421cbc72 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 25 May 2026 21:16:31 +0200 Subject: [PATCH] =?UTF-8?q?spec:=20CUDA=20Graph=20capture=20for=20RL=20ste?= =?UTF-8?q?p=20pipeline=20(2-3=C3=97=20throughput)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../2026-05-25-cuda-graph-step-pipeline.md | 183 ++++++++++++++++++ 1 file changed, 183 insertions(+) create mode 100644 docs/superpowers/specs/2026-05-25-cuda-graph-step-pipeline.md 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 new file mode 100644 index 000000000..91b54320d --- /dev/null +++ b/docs/superpowers/specs/2026-05-25-cuda-graph-step-pipeline.md @@ -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