# 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 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:** 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 (`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

- **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

### 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:

```
 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
```

**21 launches → 1 graph launch.**

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

**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).

**Blocker: noisy_exploration.resample_noise().** Same issue — host RNG + upload. Same fix: device-side PRNG for noise generation.

### Uncapturable: lobsim fill

```
lobsim.apply_snapshot(last_snap)
lobsim.step_fill_from_market_targets(ts_ns)
```

This is the LobSim simulator with internal state. Cannot be captured.

### Graph B: Post-fill (PnL extraction → controllers)

```
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
```

**15 launches → 1 graph launch.**

### Uncapturable: PER push + sample

```
push_to_replay (host-side n-step ring buffer + PER)
sample_and_gather (host-side PER sampling + DtoD copies)
```

Host-side PER operations with variable indices. Cannot be captured.

### Graph C: Replay training step

```
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
```

**~20 launches → 1 graph launch.** K-loop iterations replay this graph K times.

**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-side tau + noise generation

Replace host-side RNG for IQN tau and noisy exploration with device-side xorshift32 (matching `rl_pi_action_kernel`'s pattern).

**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
)
```

**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 2: Graph A capture (pre-fill)

```rust
// In IntegratedTrainer:
graph_a_exec: Option<CudaGraphExec>,

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(())
}
```

### Phase 3: Graph B capture (post-fill)

Same pattern for the 15 post-fill kernels.

### 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 | ~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 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.
