feat(rl): CUDA Graph capture for pre-snapshot kernel pipeline (Graph A)

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>
This commit is contained in:
jgrusewski
2026-05-25 22:16:53 +02:00
parent 0b7895a4fb
commit 4c94366f63

View File

@@ -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;