diff --git a/crates/ml-alpha/src/cfc/trunk.rs b/crates/ml-alpha/src/cfc/trunk.rs index 7a845e0cc..73d71b5e3 100644 --- a/crates/ml-alpha/src/cfc/trunk.rs +++ b/crates/ml-alpha/src/cfc/trunk.rs @@ -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, probs_d: CudaSlice, proj_out_d: CudaSlice, + + // 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, } impl CfcTrunk { @@ -116,6 +134,15 @@ impl CfcTrunk { let proj_g: Vec = vec![1.0; PROJ_DIM]; let proj_n: Vec = 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, + host: &[f32], + staging: &MappedF32Buffer, + dst: &mut CudaSlice, + ) -> 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::(); + 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( diff --git a/crates/ml-alpha/tests/graph_a_replay.rs b/crates/ml-alpha/tests/graph_a_replay.rs new file mode 100644 index 000000000..c98af60b8 --- /dev/null +++ b/crates/ml-alpha/tests/graph_a_replay.rs @@ -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()); + } +}