diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index f60314d59..bfb6c15e2 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -4166,9 +4166,11 @@ impl IntegratedTrainer { .launch() .context("prefill graph launch")?; } else { + let capturing_prefill = self.graph_warmup_done; if !self.graph_warmup_done { self.graph_warmup_done = true; - } else { + } + if capturing_prefill { self.stream .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED) .map_err(|e| anyhow::anyhow!("prefill begin_capture: {e}"))?; @@ -4494,9 +4496,7 @@ impl IntegratedTrainer { } } - // End of prefill kernel sequence. If we were capturing, - // finalize the graph. - if self.prefill_graph.is_none() && self.graph_warmup_done { + if capturing_prefill { let graph = self .stream .end_capture( @@ -4535,6 +4535,24 @@ impl IntegratedTrainer { let pos_bytes_i = lobsim.pos_bytes() as i32; let b_size_i = b_size as i32; + // ── Graph A2: post-snapshot / pre-fill kernel pipeline ───────── + // 7 kernels from session_risk through actions_to_market_targets. + // Same three-state machine as prefill — captures after warmup, + // replays on third+ step. Shares the warmup counter with Graph A. + if self.postfill_graph.is_some() { + self.postfill_graph + .as_ref() + .unwrap() + .launch() + .context("postfill graph launch")?; + } else { + let capturing_postfill = self.prefill_graph.is_some(); + if capturing_postfill { + self.stream + .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED) + .map_err(|e| anyhow::anyhow!("postfill begin_capture: {e}"))?; + } + // Session risk circuit breaker — block opening actions when // cumulative PnL EMA drops below threshold. { @@ -4728,6 +4746,18 @@ impl IntegratedTrainer { } } + if capturing_postfill { + let graph = self + .stream + .end_capture( + CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ) + .context("postfill end_capture")? + .ok_or_else(|| anyhow::anyhow!("postfill end_capture returned None"))?; + self.postfill_graph = Some(graph); + } + } // end else (postfill warmup / capture dispatch) + // Run the fill + pnl_track on whatever's now in // market_targets_d. lobsim