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>
This commit is contained in:
jgrusewski
377
docs/superpowers/plans/2026-05-25-cuda-performance-p1-p2-p3.md
Normal file
377
docs/superpowers/plans/2026-05-25-cuda-performance-p1-p2-p3.md
Normal file
@@ -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×+ |
|
||||
Reference in New Issue
Block a user