foxhunt/docs/superpowers/plans/2026-05-25-cuda-performance-p1-p2-p3.md

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

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

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

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

# 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

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

• Step 5: Commit

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:

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

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

SQLX_OFFLINE=true cargo check -p ml-alpha

Expected: 0 errors.

• Step 3: Run local smoke

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

• Step 4: Commit

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

    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:

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

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

• Step 5: Commit

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

    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:

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

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

• Step 6: Commit

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

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

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

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

"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

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×+