feat(ml-alpha): CUDA Graph A capture for perception forward
Captures snap_feature_assemble -> cfc_step -> heads -> projection into
a single replayable graph. Scalars (dt_s, ts_ns, prev_mid, ...) are
frozen at capture time per cudarc 0.19 semantics; the trunk
re-captures when those change. A follow-up task moves scalars into a
device-resident buffer for cross-step replay stability.
Key learning: cudarc's default event-tracking creates cross-stream
dependencies that begin_capture rejects with
CUDA_ERROR_STREAM_CAPTURE_ISOLATION. Pattern (from crates/ml/.../
fused_training.rs): bracket begin/end_capture with
context.disable_event_tracking() / enable_event_tracking(). Mode
remains CU_STREAM_CAPTURE_MODE_RELAXED. Pre-allocate MappedF32Buffer
staging slots as struct fields (host-malloc during/around capture is
also a trigger).
The captured forward writes h_pong directly (no ping-pong swap inside
the captured region — the swap mutates pointer identity which would
invalidate captured kernel args). Heads and projection both read
h_pong.
Tests (3/3 on sm_86):
- graph_a_replay_matches_sequential: captured replay output equals
sequential dispatch on same input at eps<=1e-5 (probs) / 1e-4 (proj)
- graph_a_replay_is_deterministic: 3 consecutive replays produce
bit-identical output
- graph_a_replay_outputs_finite: probs in [0,1], proj finite
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -12,7 +12,8 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
|
||||
use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
|
||||
use cudarc::driver::{CudaFunction, CudaGraph, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
|
||||
use ml_core::device::MlDevice;
|
||||
use rand::{Rng, SeedableRng};
|
||||
use rand_chacha::ChaCha8Rng;
|
||||
@@ -78,6 +79,23 @@ pub struct CfcTrunk {
|
||||
snap_feat_d: CudaSlice<f32>,
|
||||
probs_d: CudaSlice<f32>,
|
||||
proj_out_d: CudaSlice<f32>,
|
||||
|
||||
// Pre-allocated mapped-pinned staging slots for raw input. Reused
|
||||
// across upload_into calls so no host-malloc happens during or
|
||||
// around CUDA Graph capture (allocations on the capturing thread
|
||||
// cause CUDA_ERROR_STREAM_CAPTURE_ISOLATION).
|
||||
stg_bid_px: MappedF32Buffer,
|
||||
stg_bid_sz: MappedF32Buffer,
|
||||
stg_ask_px: MappedF32Buffer,
|
||||
stg_ask_sz: MappedF32Buffer,
|
||||
|
||||
// CUDA Graph A — captured perception forward. Scalars (dt_s, ts_ns,
|
||||
// prev_mid, ...) are frozen at capture time per cudarc 0.19 semantics
|
||||
// (`launch_builder.arg(&scalar)` records the value, not the pointer).
|
||||
// The trunk re-captures via `capture_graph_a` whenever those scalars
|
||||
// would change meaningfully. A follow-up task moves scalars into a
|
||||
// device-resident buffer so the graph stays valid across changes.
|
||||
graph_a: Option<CudaGraph>,
|
||||
}
|
||||
|
||||
impl CfcTrunk {
|
||||
@@ -116,6 +134,15 @@ impl CfcTrunk {
|
||||
let proj_g: Vec<f32> = vec![1.0; PROJ_DIM];
|
||||
let proj_n: Vec<f32> = vec![0.0; PROJ_DIM];
|
||||
|
||||
let stg_bid_px = unsafe { MappedF32Buffer::new(10) }
|
||||
.map_err(|e| anyhow::anyhow!("stg bid_px alloc: {e}"))?;
|
||||
let stg_bid_sz = unsafe { MappedF32Buffer::new(10) }
|
||||
.map_err(|e| anyhow::anyhow!("stg bid_sz alloc: {e}"))?;
|
||||
let stg_ask_px = unsafe { MappedF32Buffer::new(10) }
|
||||
.map_err(|e| anyhow::anyhow!("stg ask_px alloc: {e}"))?;
|
||||
let stg_ask_sz = unsafe { MappedF32Buffer::new(10) }
|
||||
.map_err(|e| anyhow::anyhow!("stg ask_sz alloc: {e}"))?;
|
||||
|
||||
Ok(Self {
|
||||
cfg: cfg.clone(),
|
||||
w_in_d: upload(&stream, &w_in)?,
|
||||
@@ -148,9 +175,176 @@ impl CfcTrunk {
|
||||
step_fn,
|
||||
heads_fn,
|
||||
proj_fn,
|
||||
stg_bid_px,
|
||||
stg_bid_sz,
|
||||
stg_ask_px,
|
||||
stg_ask_sz,
|
||||
graph_a: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn upload_pre_allocated(
|
||||
stream: &Arc<CudaStream>,
|
||||
host: &[f32],
|
||||
staging: &MappedF32Buffer,
|
||||
dst: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
let n = host.len();
|
||||
assert_eq!(n, dst.len());
|
||||
assert_eq!(n, staging.len);
|
||||
staging.write_from_slice(host);
|
||||
if n > 0 {
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst_ptr,
|
||||
staging.dev_ptr,
|
||||
nbytes,
|
||||
stream.cu_stream(),
|
||||
)
|
||||
.context("upload_pre_allocated DtoD")?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Capture the four perception kernels into CUDA Graph A. The
|
||||
/// scalars (dt_s, ts_ns, prev_mid, trade flow, trade count) at the
|
||||
/// time of capture are baked in; production usage must re-capture
|
||||
/// when those would change. For Phase A training the scalars come
|
||||
/// from the trainer's per-batch context and capture happens fresh
|
||||
/// each call.
|
||||
pub fn capture_graph_a(&mut self, input: &Mbp10RawInput) -> Result<()> {
|
||||
// Stage raw input into device buffers BEFORE capture using
|
||||
// pre-allocated mapped-pinned staging (host-malloc during
|
||||
// capture trips CUDA_ERROR_STREAM_CAPTURE_ISOLATION).
|
||||
Self::upload_pre_allocated(&self.stream, &input.bid_px, &self.stg_bid_px, &mut self.bid_px_d)?;
|
||||
Self::upload_pre_allocated(&self.stream, &input.bid_sz, &self.stg_bid_sz, &mut self.bid_sz_d)?;
|
||||
Self::upload_pre_allocated(&self.stream, &input.ask_px, &self.stg_ask_px, &mut self.ask_px_d)?;
|
||||
Self::upload_pre_allocated(&self.stream, &input.ask_sz, &self.stg_ask_sz, &mut self.ask_sz_d)?;
|
||||
self.stream.synchronize().context("graph A: pre-capture sync")?;
|
||||
|
||||
// cudarc's default event tracking creates cross-stream deps that
|
||||
// the capture engine rejects with STREAM_CAPTURE_ISOLATION.
|
||||
// Mirror the working pattern in crates/ml/.../fused_training.rs.
|
||||
let ctx = self.stream.context().clone();
|
||||
let _ = ctx.check_err();
|
||||
unsafe { ctx.disable_event_tracking(); }
|
||||
|
||||
let begin_result = self
|
||||
.stream
|
||||
.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED);
|
||||
if let Err(e) = begin_result {
|
||||
unsafe { ctx.enable_event_tracking(); }
|
||||
let _ = ctx.check_err();
|
||||
return Err(anyhow::anyhow!("graph A begin_capture: {e}"));
|
||||
}
|
||||
|
||||
let dispatch_result = self.dispatch_perception(input);
|
||||
|
||||
let graph_result = self
|
||||
.stream
|
||||
.end_capture(CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
|
||||
|
||||
unsafe { ctx.enable_event_tracking(); }
|
||||
let _ = ctx.check_err();
|
||||
dispatch_result?;
|
||||
|
||||
let graph = graph_result
|
||||
.context("graph A end_capture")?
|
||||
.ok_or_else(|| anyhow::anyhow!("end_capture returned None — no work was captured"))?;
|
||||
self.graph_a = Some(graph);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Replay the captured graph. Returns probs[5] + projection[8].
|
||||
/// Caller is responsible for updating the trunk's pre-allocated
|
||||
/// input buffers (via `update_input_buffers`) before calling this.
|
||||
pub fn perception_forward_captured(&mut self) -> Result<([f32; N_HORIZONS], [f32; PROJ_DIM])> {
|
||||
let graph = self
|
||||
.graph_a
|
||||
.as_ref()
|
||||
.context("Graph A not captured — call capture_graph_a first")?;
|
||||
graph.launch().context("graph A launch")?;
|
||||
self.stream.synchronize().context("graph A sync")?;
|
||||
|
||||
let probs_vec = download(&self.stream, &self.probs_d)?;
|
||||
let proj_vec = download(&self.stream, &self.proj_out_d)?;
|
||||
let mut probs = [0f32; N_HORIZONS];
|
||||
probs.copy_from_slice(&probs_vec);
|
||||
let mut proj = [0f32; PROJ_DIM];
|
||||
proj.copy_from_slice(&proj_vec);
|
||||
Ok((probs, proj))
|
||||
}
|
||||
|
||||
fn dispatch_perception(&mut self, input: &Mbp10RawInput) -> Result<()> {
|
||||
let prev_mid = input.prev_mid;
|
||||
let trade_signed_vol = input.trade_signed_vol;
|
||||
let trade_count = input.trade_count as i32;
|
||||
let ts_ns = input.ts_ns as i64;
|
||||
let prev_ts_ns = input.prev_ts_ns as i64;
|
||||
let tick_size = ES_TICK_SIZE;
|
||||
let cfg1 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
|
||||
{
|
||||
let mut launch = self.stream.launch_builder(&self.snap_fn);
|
||||
launch
|
||||
.arg(&self.bid_px_d).arg(&self.bid_sz_d).arg(&self.ask_px_d).arg(&self.ask_sz_d)
|
||||
.arg(&self.prev_bid_sz_d).arg(&self.prev_ask_sz_d)
|
||||
.arg(&prev_mid).arg(&trade_signed_vol).arg(&trade_count)
|
||||
.arg(&ts_ns).arg(&prev_ts_ns).arg(&tick_size)
|
||||
.arg(&mut self.snap_feat_d);
|
||||
unsafe { launch.launch(cfg1).context("snap launch")?; }
|
||||
}
|
||||
|
||||
let dt_s = input.ts_ns.saturating_sub(input.prev_ts_ns) as f32 * 1e-9;
|
||||
let n_in_i = self.cfg.n_in as i32;
|
||||
let n_hid_i = self.cfg.n_hid as i32;
|
||||
let block_dim = 128.min(self.cfg.n_hid as u32);
|
||||
let grid_dim = ((self.cfg.n_hid as u32) + block_dim - 1) / block_dim;
|
||||
let cfg2 = LaunchConfig { grid_dim: (grid_dim, 1, 1), block_dim: (block_dim, 1, 1), shared_mem_bytes: 0 };
|
||||
{
|
||||
let mut launch = self.stream.launch_builder(&self.step_fn);
|
||||
launch
|
||||
.arg(&self.w_in_d).arg(&self.w_rec_d).arg(&self.b_d).arg(&self.tau_d)
|
||||
.arg(&self.snap_feat_d).arg(&self.h_ping)
|
||||
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i)
|
||||
.arg(&mut self.h_pong);
|
||||
unsafe { launch.launch(cfg2).context("step launch")?; }
|
||||
}
|
||||
// NOTE: ping-pong swap intentionally OMITTED inside the captured
|
||||
// path. Inside Graph A we always write to h_pong from h_ping; the
|
||||
// swap happens via a DtoD copy outside the captured region after
|
||||
// each replay. Doing the swap as `std::mem::swap` would change
|
||||
// pointer identity, which breaks captured kernel args.
|
||||
let cfg3 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (N_HORIZONS as u32, 1, 1), shared_mem_bytes: 0 };
|
||||
{
|
||||
let mut launch = self.stream.launch_builder(&self.heads_fn);
|
||||
launch.arg(&self.heads_w_d).arg(&self.heads_b_d).arg(&self.h_pong).arg(&mut self.probs_d);
|
||||
unsafe { launch.launch(cfg3).context("heads launch")?; }
|
||||
}
|
||||
let cfg4 = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (PROJ_DIM as u32, 1, 1), shared_mem_bytes: 0 };
|
||||
{
|
||||
let mut launch = self.stream.launch_builder(&self.proj_fn);
|
||||
launch
|
||||
.arg(&self.proj_w_d).arg(&self.proj_b_d)
|
||||
.arg(&self.proj_g_d).arg(&self.proj_n_d)
|
||||
.arg(&self.h_pong).arg(&mut self.proj_out_d);
|
||||
unsafe { launch.launch(cfg4).context("proj launch")?; }
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Mirror snapshot input into the pre-allocated device buffers
|
||||
/// without dispatching the kernels (used between Graph A replays).
|
||||
pub fn update_input_buffers(&mut self, input: &Mbp10RawInput) -> Result<()> {
|
||||
Self::upload_pre_allocated(&self.stream, &input.bid_px, &self.stg_bid_px, &mut self.bid_px_d)?;
|
||||
Self::upload_pre_allocated(&self.stream, &input.bid_sz, &self.stg_bid_sz, &mut self.bid_sz_d)?;
|
||||
Self::upload_pre_allocated(&self.stream, &input.ask_px, &self.stg_ask_px, &mut self.ask_px_d)?;
|
||||
Self::upload_pre_allocated(&self.stream, &input.ask_sz, &self.stg_ask_sz, &mut self.ask_sz_d)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Sequential dispatch: snap_features -> cfc_step -> heads -> projection.
|
||||
/// Returns (probs[5], projection[8]). Updates h_ping/h_pong in place.
|
||||
pub fn forward_snapshot(
|
||||
|
||||
90
crates/ml-alpha/tests/graph_a_replay.rs
Normal file
90
crates/ml-alpha/tests/graph_a_replay.rs
Normal file
@@ -0,0 +1,90 @@
|
||||
//! CUDA Graph A capture + replay.
|
||||
//!
|
||||
//! Validates that the captured graph produces identical output to the
|
||||
//! sequential dispatch path for the SAME input (scalars are frozen at
|
||||
//! capture time per cudarc 0.19 semantics — see the trunk struct doc).
|
||||
|
||||
use approx::assert_relative_eq;
|
||||
use ml_alpha::cfc::snap_features::Mbp10RawInput;
|
||||
use ml_alpha::cfc::{CfcConfig, CfcTrunk};
|
||||
use ml_alpha::heads::{N_HORIZONS, PROJ_DIM};
|
||||
use ml_core::device::MlDevice;
|
||||
|
||||
fn test_device() -> MlDevice {
|
||||
MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
|
||||
}
|
||||
|
||||
fn synthetic_input(dt_ns: u64) -> Mbp10RawInput {
|
||||
let mut bid_px = [0.0f32; 10];
|
||||
let mut bid_sz = [0.0f32; 10];
|
||||
let mut ask_px = [0.0f32; 10];
|
||||
let mut ask_sz = [0.0f32; 10];
|
||||
for i in 0..10 {
|
||||
bid_px[i] = 5500.00 - 0.25 * i as f32;
|
||||
ask_px[i] = 5500.25 + 0.25 * i as f32;
|
||||
bid_sz[i] = 12.0 + i as f32;
|
||||
ask_sz[i] = 10.0 + i as f32 * 1.5;
|
||||
}
|
||||
Mbp10RawInput {
|
||||
bid_px, bid_sz, ask_px, ask_sz,
|
||||
prev_mid: 5499.875,
|
||||
trade_signed_vol: 4.0,
|
||||
trade_count: 7,
|
||||
ts_ns: dt_ns,
|
||||
prev_ts_ns: 0,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn graph_a_replay_matches_sequential() {
|
||||
let dev = test_device();
|
||||
let input = synthetic_input(20_000_000);
|
||||
|
||||
// Trunk A: sequential dispatch.
|
||||
let mut trunk_seq = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0xCAFE).expect("seq init");
|
||||
let (probs_seq, proj_seq) = trunk_seq.forward_snapshot(&input).expect("seq forward");
|
||||
|
||||
// Trunk B: same seed, capture + replay.
|
||||
let mut trunk_g = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0xCAFE).expect("graph init");
|
||||
trunk_g.capture_graph_a(&input).expect("capture");
|
||||
let (probs_g, proj_g) = trunk_g.perception_forward_captured().expect("replay");
|
||||
|
||||
for k in 0..N_HORIZONS {
|
||||
assert_relative_eq!(probs_seq[k], probs_g[k], epsilon = 1e-5, max_relative = 1e-5);
|
||||
}
|
||||
for j in 0..PROJ_DIM {
|
||||
assert_relative_eq!(proj_seq[j], proj_g[j], epsilon = 1e-4, max_relative = 1e-4);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn graph_a_replay_is_deterministic() {
|
||||
let dev = test_device();
|
||||
let input = synthetic_input(15_000_000);
|
||||
let mut trunk = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0x5EED).expect("init");
|
||||
trunk.capture_graph_a(&input).expect("capture");
|
||||
|
||||
let (p1, _) = trunk.perception_forward_captured().expect("replay 1");
|
||||
let (p2, _) = trunk.perception_forward_captured().expect("replay 2");
|
||||
let (p3, _) = trunk.perception_forward_captured().expect("replay 3");
|
||||
|
||||
for k in 0..N_HORIZONS {
|
||||
assert_eq!(p1[k], p2[k], "replay {k} differs between calls 1 and 2");
|
||||
assert_eq!(p2[k], p3[k], "replay {k} differs between calls 2 and 3");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn graph_a_replay_outputs_finite() {
|
||||
let dev = test_device();
|
||||
let input = synthetic_input(40_000_000);
|
||||
let mut trunk = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0xBEEF).expect("init");
|
||||
trunk.capture_graph_a(&input).expect("capture");
|
||||
let (probs, proj) = trunk.perception_forward_captured().expect("replay");
|
||||
for &p in &probs {
|
||||
assert!(p.is_finite() && (0.0..=1.0).contains(&p));
|
||||
}
|
||||
for &v in &proj {
|
||||
assert!(v.is_finite());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user