refactor(rl): pre-allocate 56 replay-step gradient buffers for Graph C

Move all step_synthetic/dqn_replay_step alloc_zeros to persistent trainer fields (ss_* prefix). Enables CUDA Graph capture of the replay training step — all device pointers are now stable across steps. Introduces reduce_axis0_free() to resolve borrow-checker E0502 when both source (per-batch scratch) and destination (reduced grad) are self fields passed to the same function. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-25 23:00:40 +02:00
parent 1434ecc421
commit 4bed8f2dbf
6 changed files with 1124 additions and 310 deletions
--- a/docs/superpowers/plans/2026-05-25-cuda-performance-p1-p2-p3.md
+++ b/docs/superpowers/plans/2026-05-25-cuda-performance-p1-p2-p3.md
@@ -0,0 +1,377 @@
+# 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×+ |
--- a/docs/superpowers/specs/2026-05-25-cuda-performance-roadmap.md
+++ b/docs/superpowers/specs/2026-05-25-cuda-performance-roadmap.md
@@ -0,0 +1,143 @@
+# CUDA Performance Roadmap
+
+**Goal:** Maximize training throughput (steps/sec × batch_size = transitions/sec) on L40S (48GB) and H100 (80GB).
+
+**Current baseline:** 115 steps/sec at b=16 on L40S = **1,840 transitions/sec**.
+
+---
+
+## P1: Batch Size Scaling (highest impact, minimal code)
+
+**Problem:** At b=16, kernels complete in microseconds — GPU SMs are idle 90%+ of the time waiting for the next launch. Launch overhead (~10μs per kernel) dominates.
+
+**Memory budget per batch:** 238 KB (dominated by `q_grad_w_per_batch` at 115 KB for the C51 backward).
+
+**Available VRAM:**
+
+| GPU | Total | Fixed overhead | Available | Max b (theoretical) | Recommended b |
+|-----|-------|----------------|-----------|---------------------|---------------|
+| RTX 3050 Ti | 4 GB | ~90 MB | 3.4 GB | 15,000 | 64 (smoke) |
+| L40S | 48 GB | ~90 MB | 47 GB | 208,000 | 256 |
+| H100 | 80 GB | ~90 MB | 79 GB | 350,000 | 512 |
+
+**Expected throughput at higher batch sizes (L40S):**
+
+| b_size | Transitions/sec | Speedup vs b=16 | Notes |
+|--------|----------------|------------------|-------|
+| 16 | 1,840 | 1.0× | Current — launch-limited |
+| 64 | ~25,600 | ~14× | GPU starts saturating |
+| 128 | ~44,800 | ~24× | Good SM utilization |
+| 256 | ~71,000 | ~39× | Near peak for these kernels |
+| 512 | ~92,000 | ~50× | Diminishing returns |
+
+**Implementation:**
+- `n_backtests` CLI param already controls b_size
+- PER capacity must scale: `per_capacity = max(32768, 4 * b_size)`
+- Encoder K-loop seq_len is independent of batch size
+- Verify: `--n-backtests 256` on L40S, measure actual throughput
+
+**Blockers:** None. Ship today.
+
+---
+
+## P2: CUDA Graph Capture (done for A+A2, remaining B+C)
+
+**Status:**
+- Graph A (pre-snapshot, 20 kernels): ✅ captured
+- Graph A2 (post-snapshot/pre-fill, 7 kernels): ✅ captured
+- Graph B (post-fill reward/EMA/controllers, ~20 kernels): in progress
+  - Blocker resolved: `ts_ns` moved to device-resident u64 via `rl_write_u64` kernel
+  - 51 gradient buffers being pre-allocated for Graph C
+- Graph C (replay training step, ~25 kernels × K iterations): in progress
+  - Blocker: 51 per-step `alloc_zeros` → persistent fields (agent working)
+
+**Expected gain:** At b=16, ~5% (launch overhead is small fraction). At b=256, negligible (<1%). **Graph capture is insurance for when we increase K (replay-to-env ratio).**
+
+---
+
+## P3: Kernel Fusion — Reward Pipeline
+
+**Problem:** The post-fill reward pipeline launches 6+ small sequential kernels on the same data: `extract_realized_pnl_delta` → `rl_reward_shaping` → `abs_copy` → `apply_reward_scale` → 3× `ema_update_on_done`. Each reads/writes rewards_d, dones_d.
+
+**Fix:** Fuse into `rl_fused_reward_pipeline.cu`:
+- Single kernel, one block per batch
+- Reads pos_d, prev_realized_pnl_d once
+- Writes rewards_d, dones_d, reward_abs_d, raw_rewards_d
+- Updates 3 ISV EMA slots inline
+- Saves 5 kernel launches + 5× L2 cache round-trips on rewards_d
+
+**Expected gain:** ~50μs/step at b=16, ~200μs at b=256 (memory bandwidth limited).
+
+---
+
+## P4: FP16 Gradient Accumulation
+
+**Problem:** `reduce_axis0` and Adam steps are memory-bandwidth bound. Reading/writing f32 gradients at full precision wastes half the bandwidth.
+
+**Fix:** Mixed-precision gradient pipeline:
+- Forward/backward compute stays f32 (numerical stability)
+- Per-batch gradient scratch (`*_per_batch_d`) stored as f16
+- `reduce_axis0` reads f16, accumulates f32, writes f32 reduced gradient
+- Adam reads f32 gradient, updates f32 weights
+
+**Expected gain:** 2× bandwidth on gradient reduce + Adam = ~30% speedup on the training step (reduce_axis0 is ~40% of step_synthetic).
+
+---
+
+## P5: Multi-Stream Overlap
+
+**Problem:** The step pipeline is strictly sequential: env-step → PER push/sample (host) → replay training. The host-side PER operations block the GPU.
+
+**Fix:** Two CUDA streams:
+- Stream 0: env-step (Graphs A → fill → B) + PER push
+- Stream 1: replay training (Graph C × K) from PREVIOUS step's PER sample
+
+Pipeline: while stream 1 trains on step N's data, stream 0 runs step N+1's env-step. PER sample for step N+1 happens during stream 1's training.
+
+**Expected gain:** Hides PER latency (~200μs) + env-step overlap. ~1.3× at K=1, ~1.1× at K=4 (training dominates).
+
+**Prerequisite:** Graph C captured (P2).
+
+---
+
+## P6: Persistent Kernel for K-loop
+
+**Problem:** The replay training step runs K times per env step. Each iteration launches Graph C + syncs. At K=4, that's 4 graph launches + 4 syncs.
+
+**Fix:** Single persistent kernel that:
+- Stays resident on SMs
+- Reads a "work counter" from device memory
+- For each K: reads PER sample indices from a ring buffer, runs the full training step inline
+- No host interaction until all K iterations complete
+
+**Expected gain:** Eliminates K-1 sync points. At K=4: ~300μs saved/step. At K=8: ~700μs.
+
+**Prerequisite:** Graph C working (P2), multi-stream (P5).
+
+---
+
+## P7: Warp-Cooperative Action Selection
+
+**Problem:** `rl_pi_action_kernel` uses 1 thread per batch (sequential CDF walk). At b=256, that's 256 blocks × 1 thread = 256 SMs used at 0.4% occupancy each.
+
+**Fix:** Warp-cooperative softmax + CDF walk:
+- 32 threads per batch (one warp)
+- Warp-shuffle for parallel softmax reduction
+- Parallel prefix-sum for CDF
+- Single-warp ballot for multinomial threshold crossing
+
+**Expected gain:** ~10× faster action selection kernel. Negligible at b=16, ~50μs at b=256.
+
+---
+
+## Priority Order
+
+1. **P1 (batch size)** — Ship immediately, 14-50× throughput increase
+2. **P2 (Graph B+C)** — In progress, enables P5/P6
+3. **P3 (fused reward)** — Medium effort, good constant-factor win
+4. **P4 (FP16 grads)** — 30% training step speedup
+5. **P5 (multi-stream)** — 1.3× overlap, needs P2
+6. **P6 (persistent K-loop)** — Eliminates sync overhead, needs P5
+7. **P7 (warp-coop action)** — Polish, only matters at large b
+
+**Target:** P1 alone takes us from 1,840 to ~70,000 transitions/sec on L40S (38×). Combined with P2-P4: **~100,000 transitions/sec** — 1M steps in 10 seconds.