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>
91 lines
3.1 KiB
Rust
91 lines
3.1 KiB
Rust
//! 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());
|
|
}
|
|
}
|