From 3af2b40100f5441df250f7a6d191f34caad99142 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Mon, 25 May 2026 23:05:59 +0200
Subject: [PATCH] =?UTF-8?q?feat(rl):=20CUDA=20Graph=20B=20capture=20?=
 =?UTF-8?q?=E2=80=94=20post-fill=20reward/EMA/controller=20pipeline=20(~20?=
 =?UTF-8?q?=20kernels)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Three-state machine (warmup → capture → replay) wraps the post-fill
section: extract_realized_pnl_delta, unit_state_update, trade_context,
multires, step_counter, 3× EMA, abs_copy, reward_shaping, raw_rewards
snapshot, apply_reward_scale, reward_clamp, atom_support_update,
compute_advantage_return, var_over_abs_mean.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 crates/ml-alpha/src/trainer/integrated.rs | 32 +++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs
index 2e9742052..bd27056cb 100644
--- a/crates/ml-alpha/src/trainer/integrated.rs
+++ b/crates/ml-alpha/src/trainer/integrated.rs
@@ -425,6 +425,7 @@ pub struct IntegratedTrainer {
     // barrier that can't be captured).
     prefill_graph: Option<CudaGraph>,
     postfill_graph: Option<CudaGraph>,
+    reward_graph: Option<CudaGraph>,
     graph_warmup_done: bool,
 
     // ── Kernel handles for grad combine + cross-batch reduce ──────────
@@ -1902,6 +1903,7 @@ impl IntegratedTrainer {
             stream,
             prefill_graph: None,
             postfill_graph: None,
+            reward_graph: None,
             graph_warmup_done: false,
             _grad_h_module: grad_h_module,
             grad_h_accumulate_fn,
@@ -5042,6 +5044,24 @@ impl IntegratedTrainer {
             .step_fill_from_market_targets(last_snap.ts_ns)
             .context("step_with_lobsim: lobsim.step_fill_from_market_targets")?;
 
+        // ── Graph B: post-fill reward/EMA/controller pipeline ─────────
+        // ~20 kernels from extract_realized_pnl_delta through
+        // var_over_abs_mean. Same three-state machine as Graph A/A2 —
+        // captures after A2 is done, replays on subsequent steps.
+        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}"))?;
+            }
+
         // Extract per-batch (reward, done) deltas vs the trainer's
         // prev_realized_pnl_d snapshot. Updates prev_* in place for
         // the next step's delta.
@@ -5493,6 +5513,18 @@ impl IntegratedTrainer {
         self.launch_var_over_abs_mean(&self.advantages_d.clone(), b_size)
             .context("launch_var_over_abs_mean(advantages_d)")?;
 
+            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 7a (R7d): PER push + sample + gather. ────────────────
         // Push CURRENT step's b_size transitions (one per batch) to
         // ReplayBuffer; sample b_size indices (priority^α from