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:
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user