From 4c94366f631e0e0718064a5b9e0b29126be65480 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Mon, 25 May 2026 22:16:53 +0200
Subject: [PATCH] feat(rl): CUDA Graph capture for pre-snapshot kernel pipeline
 (Graph A)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Three-state machine (warmup → capture → replay) wraps the 20 pre-
snapshot kernels in step_with_lobsim: FRD/DQN/IQN/V/π forwards,
ensemble action-value, noisy exploration, Q→π agreement, π-driven
action selection, argmax Bellman, log_pi, confidence gate, FRD gate.

On first step, kernels dispatch eagerly (warmup). On second step,
stream capture records the kernel topology. On third+ steps, a single
cudaGraphLaunch replays all 20 kernels — ~200μs launch overhead → ~10μs.

Host-side step_counter increment moved outside the capture region
(host ops are invisible to graph capture).

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

diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs
index 37127718a..f60314d59 100644
--- a/crates/ml-alpha/src/trainer/integrated.rs
+++ b/crates/ml-alpha/src/trainer/integrated.rs
@@ -78,9 +78,10 @@ use std::sync::Arc;
 
 use anyhow::{Context, Result};
 use cudarc::driver::{
-    CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig,
-    PushKernelArg,
+    CudaFunction, CudaGraph, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut,
+    LaunchConfig, PushKernelArg,
 };
+use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
 use ml_core::device::MlDevice;
 
 use crate::cfc::snap_features::Mbp10RawInput;
@@ -416,6 +417,14 @@ pub struct IntegratedTrainer {
 
     stream: Arc<CudaStream>,
 
+    // ── CUDA Graph capture for the RL step pipeline ─────────────────
+    // Three-state machine: first step = warmup (eager), second =
+    // capture, third+ = replay. Split around apply_snapshot (HtoD
+    // barrier that can't be captured).
+    prefill_graph: Option<CudaGraph>,
+    postfill_graph: Option<CudaGraph>,
+    graph_warmup_done: bool,
+
     // ── Kernel handles for grad combine + cross-batch reduce ──────────
     _grad_h_module: Arc<CudaModule>,
     grad_h_accumulate_fn: CudaFunction,
@@ -1661,6 +1670,9 @@ impl IntegratedTrainer {
             isv_d,
             isv_host,
             stream,
+            prefill_graph: None,
+            postfill_graph: None,
+            graph_warmup_done: false,
             _grad_h_module: grad_h_module,
             grad_h_accumulate_fn,
             _reduce_axis0_module: reduce_axis0_module,
@@ -4138,6 +4150,30 @@ impl IntegratedTrainer {
             .forward_encoder(snapshots)
             .context("step_with_lobsim: forward_encoder(snapshots)")?;
 
+        // Host-side step counter — must stay outside graph capture region
+        // (host ops aren't captured, would silently stop incrementing).
+        self.step_counter = self.step_counter.wrapping_add(1);
+
+        // ── Graph A: pre-snapshot kernel pipeline ──────────────────────
+        // 20 kernels from FRD/DQN/IQN/V/π forwards + action selection +
+        // gates. Three-state machine: warmup (first) → capture (second)
+        // → replay (third+). All device pointers are stable (pre-allocated
+        // at init), so the captured graph replays with fresh data.
+        if self.prefill_graph.is_some() {
+            self.prefill_graph
+                .as_ref()
+                .unwrap()
+                .launch()
+                .context("prefill graph launch")?;
+        } else {
+            if !self.graph_warmup_done {
+                self.graph_warmup_done = true;
+            } else {
+                self.stream
+                    .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)
+                    .map_err(|e| anyhow::anyhow!("prefill begin_capture: {e}"))?;
+            }
+
         // ── Step 2: Q + V forwards on h_t for action sampling. ────────
         let h_t_borrow: &CudaSlice<f32> = self.perception.h_t_view();
         debug_assert_eq!(h_t_borrow.len(), b_size * HIDDEN_DIM);
@@ -4318,10 +4354,6 @@ impl IntegratedTrainer {
         //
         // The host Thompson + host argmax + host log-softmax loops the
         // flawed Phase F+G shipped are GONE. The xorshift32 PRNG state
-        // is per-batch device-resident; the host-side step_counter
-        // increment that drove the old ChaCha8 host RNG is also gone.
-        self.step_counter = self.step_counter.wrapping_add(1);
-
         // ── π-driven action selection (Option B fix per
         // `pearl_q_thompson_actor_makes_pi_dead_weight`). Replaces the
         // prior Q-Thompson selector. π is now the actor: actions come
@@ -4462,6 +4494,20 @@ impl IntegratedTrainer {
             }
         }
 
+            // End of prefill kernel sequence. If we were capturing,
+            // finalize the graph.
+            if self.prefill_graph.is_none() && self.graph_warmup_done {
+                let graph = self
+                    .stream
+                    .end_capture(
+                        CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
+                    )
+                    .context("prefill end_capture")?
+                    .ok_or_else(|| anyhow::anyhow!("prefill end_capture returned None"))?;
+                self.prefill_graph = Some(graph);
+            }
+        } // end else (warmup / capture dispatch)
+
         // ── Step 5 (Phase R6): GPU-pure env step. ─────────────────────
         // No per-batch host loop — actions land in lobsim's
         // market_targets_d via the `actions_to_market_targets` kernel
@@ -5356,7 +5402,6 @@ impl IntegratedTrainer {
             }
 
             let buf = &self.n_step_buffer[b];
-            let buf_len = buf.len();
 
             // Compute n-step discounted return: R_n = Σ γᵏ rₖ
             let mut r_n_scaled = 0.0_f32;