# CUDA Performance P1-P3 Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Take training throughput from 1,840 transitions/sec to 70,000+ on L40S via batch size scaling (P1), complete CUDA Graph capture (P2), and fuse the reward pipeline (P3).

**Architecture:** P1 is a config change (b=16→256) + PER capacity scaling. P2 wraps remaining kernel sequences with cudarc's begin_capture/end_capture/launch three-state machine. P3 fuses 6 sequential reward kernels into one. All changes are in the RL trainer and CUDA kernel layer.

**Tech Stack:** Rust 1.85, cudarc 0.19, CUDA 12.4, pre-compiled cubins via build.rs

---

## File Map

| File | Responsibility | Tasks |
|------|---------------|-------|
| `infra/k8s/argo/alpha-rl-template.yaml` | Default n-backtests param | T1 |
| `crates/ml-alpha/src/trainer/integrated.rs` | Trainer: PER scaling, Graph B+C, borrow fixes | T2, T3, T4, T5 |
| `crates/ml-alpha/cuda/rl_fused_reward_pipeline.cu` | New: fused reward kernel | T6 |
| `crates/ml-alpha/build.rs` | Register new kernel | T6 |

---

## Task 1: P1 — Batch Size Scaling (Argo default + PER auto-scale)

**Files:**
- Modify: `infra/k8s/argo/alpha-rl-template.yaml:44-45`
- Modify: `crates/ml-alpha/src/trainer/integrated.rs` (PER capacity logic)

- [ ] **Step 1: Update Argo template default n-backtests from 16 to 128**

```yaml
      - name: n-backtests
        value: "128"
```

- [ ] **Step 2: Auto-scale PER capacity to 4× batch size minimum**

In `IntegratedTrainer::new()`, after `let b_size = cfg.perception.n_batch;`, replace the fixed PER capacity with:

```rust
let per_capacity = cfg.per_capacity.max(4 * b_size);
```

Find where `cfg.per_capacity` is passed to the replay buffer constructor (around line 1680 where `cfg.per_capacity` appears) and replace with `per_capacity`.

- [ ] **Step 3: Verify local smoke at b=64**

```bash
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml-alpha --test integrated_trainer_smoke -- --ignored --nocapture
```

The existing smoke uses `IntegratedTrainerConfig::default()` which has `n_batch=16`. Verify it still passes. Then manually test b=64 is allocable on RTX 3050 Ti (4GB):

```bash
# Quick alloc test — just construct the trainer, don't step
SQLX_OFFLINE=true cargo test -p ml-alpha --test integrated_trainer_smoke -- --ignored --nocapture
```

- [ ] **Step 4: Apply the Argo template to cluster**

```bash
kubectl apply -f infra/k8s/argo/alpha-rl-template.yaml -n foxhunt
```

- [ ] **Step 5: Commit**

```bash
git add infra/k8s/argo/alpha-rl-template.yaml crates/ml-alpha/src/trainer/integrated.rs
git commit -m "perf(rl): scale batch size to 128 default + auto-size PER capacity"
```

---

## Task 2: P2 — Fix Graph C Borrow Checker (reduce_axis0 pattern)

**Files:**
- Modify: `crates/ml-alpha/src/trainer/integrated.rs` (reduce_axis0_free calls)

The E0502 errors occur because `reduce_axis0_free` takes `&self.stream` + `&self.reduce_axis0_fn` (immutable borrow of self) AND `&mut self.ss_q_grad_w_d` (mutable borrow of self) in the same call. The fix is to extract the stream and fn references into locals before the call.

- [ ] **Step 1: Fix all reduce_axis0_free borrow conflicts**

The pattern for every call site is:

```rust
// BEFORE (fails borrow check):
reduce_axis0_free(&self.stream, &self.reduce_axis0_fn, &self.ss_q_grad_w_per_batch_d, b_size, k_dqn * HIDDEN_DIM, &mut self.ss_q_grad_w_d)?;

// AFTER (compiles):
let (stream, fn_ref) = (&self.stream, &self.reduce_axis0_fn);
reduce_axis0_free(stream, fn_ref, &self.ss_q_grad_w_per_batch_d, b_size, k_dqn * HIDDEN_DIM, &mut self.ss_q_grad_w_d)?;
```

Alternatively, extract `stream` and `reduce_axis0_fn` into locals at the TOP of `step_synthetic`:

```rust
let stream = &self.stream;
let reduce_fn = &self.reduce_axis0_fn;
```

Then use `stream` and `reduce_fn` throughout. This fixes all ~15 call sites at once.

- [ ] **Step 2: Verify compilation**

```bash
SQLX_OFFLINE=true cargo check -p ml-alpha
```

Expected: 0 errors.

- [ ] **Step 3: Run local smoke**

```bash
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml-alpha --test integrated_trainer_smoke -- --ignored --nocapture
```

- [ ] **Step 4: Commit**

```bash
git add crates/ml-alpha/src/trainer/integrated.rs
git commit -m "fix(rl): resolve borrow-checker conflicts in pre-allocated gradient reduce"
```

---

## Task 3: P2 — Graph B Capture (post-fill reward pipeline)

**Files:**
- Modify: `crates/ml-alpha/src/trainer/integrated.rs`

**Prerequisite:** Task 2 (Graph C buffers compile clean).

Graph B covers ~20 kernels from `extract_realized_pnl_delta` through `launch_var_over_abs_mean` in `step_with_lobsim`. The `rl_write_u64` kernel and `ts_ns_d` device buffer are already wired. The `postfill_graph` field already exists but currently captures the pre-fill A2 section — rename and add a `reward_graph` field.

- [ ] **Step 1: Add `reward_graph: Option<CudaGraph>` field**

In the struct definition (near line 426):

```rust
    prefill_graph: Option<CudaGraph>,
    postfill_graph: Option<CudaGraph>,
    reward_graph: Option<CudaGraph>,
    graph_warmup_done: bool,
```

Initialize in constructor: `reward_graph: None,`

- [ ] **Step 2: Wrap post-fill section with three-state machine**

After `lobsim.step_fill_from_market_targets(...)` and before the `extract_realized_pnl_delta` block, add:

```rust
if self.reward_graph.is_some() {
    self.reward_graph.as_ref().unwrap().launch()
        .context("reward graph launch")?;
} else {
    let capturing_reward = self.postfill_graph.is_some();
    if capturing_reward {
        self.stream
            .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)
            .map_err(|e| anyhow::anyhow!("reward begin_capture: {e}"))?;
    }

    // ... existing post-fill kernels (extract_realized_pnl_delta through
    //     launch_var_over_abs_mean) stay here unchanged ...

    if capturing_reward {
        let graph = self.stream
            .end_capture(CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH)
            .context("reward end_capture")?
            .ok_or_else(|| anyhow::anyhow!("reward end_capture returned None"))?;
        self.reward_graph = Some(graph);
    }
} // end else (reward warmup / capture)
```

- [ ] **Step 3: Verify no host ops inside capture region**

Audit: no `alloc_zeros`, no `stream.synchronize()`, no `lobsim.pos_bytes()` HOST calls inside the region. The `lobsim.pos_d()`, `lobsim.bid_px_d()` calls return stable device pointers — safe.

The `launch_rl_fused_controllers(b_size)` and `launch_var_over_abs_mean(...)` are method calls — verify they contain only kernel launches (already confirmed: pure `launch_builder` + `launch`).

- [ ] **Step 4: Run local smoke + compute-sanitizer**

```bash
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml-alpha --test integrated_trainer_smoke -- --ignored --nocapture
```

- [ ] **Step 5: Commit**

```bash
git add crates/ml-alpha/src/trainer/integrated.rs
git commit -m "feat(rl): CUDA Graph B capture for post-fill reward/EMA/controller pipeline (~20 kernels)"
```

---

## Task 4: P2 — Graph C Capture (replay training step)

**Files:**
- Modify: `crates/ml-alpha/src/trainer/integrated.rs`

**Prerequisite:** Tasks 2+3 (buffers pre-allocated, compiles clean).

Graph C captures `step_synthetic`'s kernel sequence. The three-state machine is the same pattern. Add `training_graph: Option<CudaGraph>`.

- [ ] **Step 1: Add `training_graph: Option<CudaGraph>` field + warmup counter**

```rust
    training_graph: Option<CudaGraph>,
    training_graph_warmup_done: bool,
```

Initialize: `training_graph: None, training_graph_warmup_done: false,`

- [ ] **Step 2: Wrap step_synthetic with capture**

At the top of `step_synthetic`, after the `let b_size = ...` + `let k_dqn = ...` lines:

```rust
if self.training_graph.is_some() {
    self.training_graph.as_ref().unwrap().launch()
        .context("training graph launch")?;
    // Skip the rest — graph replays all kernels.
    self.stream.synchronize().context("training graph sync")?;
    // Read back losses from mapped-pinned...
    return Ok(step_synthetic_stats_from_device(self)?);
}

let capturing_training = self.training_graph_warmup_done;
if !self.training_graph_warmup_done {
    self.training_graph_warmup_done = true;
}
if capturing_training {
    self.stream
        .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)
        .map_err(|e| anyhow::anyhow!("training begin_capture: {e}"))?;
}

// ... existing step_synthetic body ...

if capturing_training {
    let graph = self.stream
        .end_capture(CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH)
        .context("training end_capture")?
        .ok_or_else(|| anyhow::anyhow!("training end_capture returned None"))?;
    self.training_graph = Some(graph);
}
```

- [ ] **Step 3: Extract loss readback to separate method**

The graph replay skips the kernel dispatches but still needs to read loss scalars back. Extract the mapped-pinned loss reads at the end of step_synthetic into a helper that both the graph-replay path and the capture path call AFTER sync.

- [ ] **Step 4: Verify no host ops in capture region**

Check: no `alloc_zeros`, no `synchronize`, no host-side branching. The `self.launch_*` helper methods must be pure kernel launches.

- [ ] **Step 5: Run local smoke**

```bash
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml-alpha --test integrated_trainer_smoke -- --ignored --nocapture
```

- [ ] **Step 6: Commit**

```bash
git add crates/ml-alpha/src/trainer/integrated.rs
git commit -m "feat(rl): CUDA Graph C capture for replay training step (~25 kernels)"
```

---

## Task 5: P1 — Benchmark b=128 on L40S

**Files:**
- None (infrastructure test)

- [ ] **Step 1: Push and submit b=128 run**

```bash
git push origin ml-alpha-phase-a
argo submit -n foxhunt --from=wftmpl/alpha-rl \
  -p git-branch=ml-alpha-phase-a \
  -p n-steps=5000 \
  -p n-backtests=128
```

- [ ] **Step 2: Measure throughput**

Wait for completion, check `elapsed` in logs. Calculate steps/sec and transitions/sec.

Expected: ~700 steps/sec × 128 = ~89,600 transitions/sec (48× improvement over baseline).

- [ ] **Step 3: If passes, submit b=256**

```bash
argo submit -n foxhunt --from=wftmpl/alpha-rl \
  -p git-branch=ml-alpha-phase-a \
  -p n-steps=5000 \
  -p n-backtests=256
```

---

## Task 6: P3 — Fused Reward Pipeline Kernel

**Files:**
- Create: `crates/ml-alpha/cuda/rl_fused_reward_pipeline.cu`
- Modify: `crates/ml-alpha/build.rs`
- Modify: `crates/ml-alpha/src/trainer/integrated.rs`

- [ ] **Step 1: Write the fused kernel**

```cuda
// rl_fused_reward_pipeline.cu — Fuses 6 sequential post-fill kernels:
// extract_realized_pnl_delta + rl_reward_shaping + abs_copy +
// ema_update_on_done(duration) + ema_update_on_done(abs_pnl) +
// ema_update_on_done(outcome) + rl_recent_outcome_update.
//
// Single launch replaces 7 kernels. Eliminates 6 L2 cache round-trips
// on rewards_d and 6 kernel launch overheads.
//
// Grid=(1,1,1), Block=(b_size, 1, 1). Each thread handles one batch.
// Shared memory for tree-reduce of EMAs.
```

The kernel:
1. Reads `pos_d[b]`, `prev_realized_pnl_d[b]` → computes reward, done
2. Applies reward shaping (entry cost, hold bonus, min-hold penalty)
3. Writes `raw_rewards_d[b]` (snapshot before scale)
4. Computes `|reward|` → `reward_abs_d[b]`
5. Updates `steps_since_done_d[b]`, emits `trade_duration_emit_d[b]`
6. Tree-reduce: 3× EMA updates to ISV (mean_trade_duration, mean_abs_pnl, outcome)
7. Writes `outcome_ema_d[b]` (per-batch)
8. Writes final `rewards_d[b]`, `dones_d[b]`

- [ ] **Step 2: Register in build.rs**

```rust
"rl_fused_reward_pipeline",  // P3: 7→1 fused reward extraction + shaping + EMA update
```

- [ ] **Step 3: Wire into trainer — replace 7 kernel launches**

Replace the section from `extract_realized_pnl_delta` through `rl_recent_outcome_update` with a single launch of the fused kernel.

- [ ] **Step 4: Run local smoke + oracle test**

Verify loss trajectory matches the unfused version (within fp noise).

- [ ] **Step 5: Commit**

```bash
git add crates/ml-alpha/cuda/rl_fused_reward_pipeline.cu crates/ml-alpha/build.rs crates/ml-alpha/src/trainer/integrated.rs
git commit -m "perf(rl): fuse 7 post-fill reward kernels into one launch (P3)"
```

---

## Kill Criteria

- P1: b=128 completes without OOM on L40S and throughput ≥ 5× vs b=16
- P2 Graph B: local smoke passes after capture (step 1 warmup, step 2 capture, step 3+ replay)
- P2 Graph C: same three-state machine, local smoke passes
- P3: loss trajectory at step 1000 matches within 1% of unfused version

## Expected Final Throughput

| Config | Transitions/sec | vs Baseline |
|--------|----------------|-------------|
| Baseline (b=16, no graphs) | 1,840 | 1× |
| P1 only (b=128) | ~89,600 | 48× |
| P1 + P2 (b=128, all graphs) | ~92,000 | 50× |
| P1 + P2 + P3 (b=128, fused) | ~95,000 | 52× |
| P1 at b=256 | ~150,000+ | 80×+ |
