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, postfill_graph: Option, + reward_graph: Option, 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