From 1ed81cf6c073bb64101e6026645472cc35016a36 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 16 May 2026 21:55:00 +0200 Subject: [PATCH] =?UTF-8?q?feat(ml-alpha):=20CfcTrunk=20forward=20?= =?UTF-8?q?=E2=80=94=20sequential=20perception=20dispatch?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CfcTrunk owns weights, ping-pong hidden buffers, and pre-allocated per-step scratch (snap features, probs, projection output). Modules and CudaFunction handles cached at new_random so the hot path avoids reload. forward_snapshot dispatches snap_feature_assemble -> cfc_step -> heads -> projection sequentially; Graph A capture (Task 11) will fold these into a single launch. The Mamba2 prefix (per 2026-05-16 spec amendment) is added in a follow-up task before Graph A capture. Tests (5/5 on sm_86): - probs in [0,1] across all 5 horizons - hidden state changes after forward - probs + proj are finite - layer-norm proj has near-zero mean - 50-step run leaves hidden finite Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/src/cfc/mod.rs | 3 + crates/ml-alpha/src/cfc/trunk.rs | 305 +++++++++++++++++++++++++ crates/ml-alpha/tests/trunk_forward.rs | 106 +++++++++ 3 files changed, 414 insertions(+) create mode 100644 crates/ml-alpha/src/cfc/trunk.rs create mode 100644 crates/ml-alpha/tests/trunk_forward.rs diff --git a/crates/ml-alpha/src/cfc/mod.rs b/crates/ml-alpha/src/cfc/mod.rs index 38bf1ab33..9baaf1f60 100644 --- a/crates/ml-alpha/src/cfc/mod.rs +++ b/crates/ml-alpha/src/cfc/mod.rs @@ -5,3 +5,6 @@ pub mod snap_features; pub mod step; +pub mod trunk; + +pub use trunk::{CfcConfig, CfcTrunk}; diff --git a/crates/ml-alpha/src/cfc/trunk.rs b/crates/ml-alpha/src/cfc/trunk.rs new file mode 100644 index 000000000..7a845e0cc --- /dev/null +++ b/crates/ml-alpha/src/cfc/trunk.rs @@ -0,0 +1,305 @@ +//! `CfcTrunk` — owns CfC weights, multi-horizon heads, projection, and a +//! ping-pong hidden buffer. `forward_snapshot` dispatches the four +//! perception kernels sequentially. The Mamba2 prefix (per the +//! 2026-05-16 spec amendment) lands in a follow-up task before Graph A +//! capture; this task validates the CfC + heads + projection forward +//! path standalone. +//! +//! Per the addendum API patterns: cache `Arc` and the +//! `CudaFunction` handles in the struct so the hot path avoids module +//! reload. + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg}; +use ml_core::device::MlDevice; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +use crate::cfc::snap_features::{Mbp10RawInput, ES_TICK_SIZE, FEATURE_DIM}; +use crate::heads::{HIDDEN_DIM, N_HORIZONS, PROJ_DIM}; +use crate::pinned_mem::MappedF32Buffer; + +const SNAP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/snap_feature_assemble.cubin")); +const STEP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cfc_step.cubin")); +const HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horizon_heads.cubin")); +const PROJ_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/projection.cubin")); + +#[derive(Clone, Debug)] +pub struct CfcConfig { + pub n_in: usize, + pub n_hid: usize, +} + +impl Default for CfcConfig { + fn default() -> Self { + Self { n_in: FEATURE_DIM, n_hid: HIDDEN_DIM } + } +} + +pub struct CfcTrunk { + cfg: CfcConfig, + stream: Arc, + + // Cubin modules (kept alive for the function refs below) + _snap_module: Arc, + _step_module: Arc, + _heads_module: Arc, + _proj_module: Arc, + snap_fn: CudaFunction, + step_fn: CudaFunction, + heads_fn: CudaFunction, + proj_fn: CudaFunction, + + // Device-resident weights + w_in_d: CudaSlice, + w_rec_d: CudaSlice, + b_d: CudaSlice, + tau_d: CudaSlice, + heads_w_d: CudaSlice, + heads_b_d: CudaSlice, + proj_w_d: CudaSlice, + proj_b_d: CudaSlice, + proj_g_d: CudaSlice, + proj_n_d: CudaSlice, + + // Ping-pong hidden state + h_ping: CudaSlice, + h_pong: CudaSlice, + + // Pre-allocated per-step scratch (raw input + features + outputs) + bid_px_d: CudaSlice, + bid_sz_d: CudaSlice, + ask_px_d: CudaSlice, + ask_sz_d: CudaSlice, + prev_bid_sz_d: CudaSlice, + prev_ask_sz_d: CudaSlice, + snap_feat_d: CudaSlice, + probs_d: CudaSlice, + proj_out_d: CudaSlice, +} + +impl CfcTrunk { + pub fn new_random(dev: &MlDevice, cfg: &CfcConfig, seed: u64) -> Result { + let stream: Arc = dev.cuda_stream().context("trunk stream")?.clone(); + let ctx = dev.cuda_context().context("trunk ctx")?; + + let snap_module = ctx.load_cubin(SNAP_CUBIN.to_vec()).context("load snap cubin")?; + let step_module = ctx.load_cubin(STEP_CUBIN.to_vec()).context("load step cubin")?; + let heads_module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("load heads cubin")?; + let proj_module = ctx.load_cubin(PROJ_CUBIN.to_vec()).context("load proj cubin")?; + let snap_fn = snap_module.load_function("snap_feature_assemble").context("snap_fn")?; + let step_fn = step_module.load_function("cfc_step").context("step_fn")?; + let heads_fn = heads_module.load_function("multi_horizon_heads").context("heads_fn")?; + let proj_fn = proj_module.load_function("projection_kernel").context("proj_fn")?; + + // Random init. + let mut r = ChaCha8Rng::seed_from_u64(seed); + let scale_in = (1.0_f32 / cfg.n_in as f32).sqrt(); + let scale_rec = (1.0_f32 / cfg.n_hid as f32).sqrt(); + let w_in: Vec = (0..cfg.n_hid * cfg.n_in).map(|_| r.gen_range(-scale_in..scale_in)).collect(); + let w_rec: Vec = (0..cfg.n_hid * cfg.n_hid).map(|_| r.gen_range(-scale_rec..scale_rec)).collect(); + let b: Vec = vec![0.0; cfg.n_hid]; + // tau log-uniform over [10ms, 1000s] per spec Section 2 (ln(0.01)=-4.6, ln(1000)=6.9). + let tau: Vec = (0..cfg.n_hid) + .map(|_| { + let u: f32 = r.gen_range(0.0..1.0); + (-4.6 + u * (4.6 + 6.9)).exp() + }) + .collect(); + let head_scale = scale_rec; + let heads_w: Vec = (0..N_HORIZONS * cfg.n_hid).map(|_| r.gen_range(-head_scale..head_scale)).collect(); + let heads_b: Vec = vec![0.0; N_HORIZONS]; + let proj_w: Vec = (0..PROJ_DIM * cfg.n_hid).map(|_| r.gen_range(-head_scale..head_scale)).collect(); + let proj_b: Vec = vec![0.0; PROJ_DIM]; + let proj_g: Vec = vec![1.0; PROJ_DIM]; + let proj_n: Vec = vec![0.0; PROJ_DIM]; + + Ok(Self { + cfg: cfg.clone(), + w_in_d: upload(&stream, &w_in)?, + w_rec_d: upload(&stream, &w_rec)?, + b_d: upload(&stream, &b)?, + tau_d: upload(&stream, &tau)?, + heads_w_d: upload(&stream, &heads_w)?, + heads_b_d: upload(&stream, &heads_b)?, + proj_w_d: upload(&stream, &proj_w)?, + proj_b_d: upload(&stream, &proj_b)?, + proj_g_d: upload(&stream, &proj_g)?, + proj_n_d: upload(&stream, &proj_n)?, + h_ping: stream.alloc_zeros::(cfg.n_hid).context("h_ping alloc")?, + h_pong: stream.alloc_zeros::(cfg.n_hid).context("h_pong alloc")?, + bid_px_d: stream.alloc_zeros::(10).context("bid_px alloc")?, + bid_sz_d: stream.alloc_zeros::(10).context("bid_sz alloc")?, + ask_px_d: stream.alloc_zeros::(10).context("ask_px alloc")?, + ask_sz_d: stream.alloc_zeros::(10).context("ask_sz alloc")?, + prev_bid_sz_d: stream.alloc_zeros::(10).context("prev_bid_sz alloc")?, + prev_ask_sz_d: stream.alloc_zeros::(10).context("prev_ask_sz alloc")?, + snap_feat_d: stream.alloc_zeros::(FEATURE_DIM).context("snap_feat alloc")?, + probs_d: stream.alloc_zeros::(N_HORIZONS).context("probs alloc")?, + proj_out_d: stream.alloc_zeros::(PROJ_DIM).context("proj_out alloc")?, + stream, + _snap_module: snap_module, + _step_module: step_module, + _heads_module: heads_module, + _proj_module: proj_module, + snap_fn, + step_fn, + heads_fn, + proj_fn, + }) + } + + /// Sequential dispatch: snap_features -> cfc_step -> heads -> projection. + /// Returns (probs[5], projection[8]). Updates h_ping/h_pong in place. + pub fn forward_snapshot( + &mut self, + input: &Mbp10RawInput, + ) -> Result<([f32; N_HORIZONS], [f32; PROJ_DIM])> { + // Stage raw input into pre-allocated device buffers via mapped-pinned. + upload_into(&self.stream, &input.bid_px, &mut self.bid_px_d)?; + upload_into(&self.stream, &input.bid_sz, &mut self.bid_sz_d)?; + upload_into(&self.stream, &input.ask_px, &mut self.ask_px_d)?; + upload_into(&self.stream, &input.ask_sz, &mut self.ask_sz_d)?; + + // Step 1: snap_feature_assemble. + 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")?; } + } + + // Step 2: cfc_step (h_ping -> h_pong). + 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")?; } + } + std::mem::swap(&mut self.h_ping, &mut self.h_pong); + + // Step 3: heads. + 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_ping).arg(&mut self.probs_d); + unsafe { launch.launch(cfg3).context("heads launch")?; } + } + + // Step 4: projection. + 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_ping).arg(&mut self.proj_out_d); + unsafe { launch.launch(cfg4).context("proj launch")?; } + } + + // Copy current bid_sz/ask_sz into prev_*_sz for next-step OFI baseline (DtoD). + copy_dtod(&self.stream, &self.bid_sz_d, &mut self.prev_bid_sz_d)?; + copy_dtod(&self.stream, &self.ask_sz_d, &mut self.prev_ask_sz_d)?; + + self.stream.synchronize().context("trunk forward sync")?; + + // Readback (slow-path for now; Task 11 captures into Graph A and + // keeps everything on-device). + 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)) + } + + pub fn snapshot_hidden(&self) -> Result> { + download(&self.stream, &self.h_ping) + } + + pub fn config(&self) -> &CfcConfig { &self.cfg } +} + +// Helpers --------------------------------------------------------------- + +fn upload(stream: &Arc, host: &[f32]) -> Result> { + let n = host.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("trunk upload staging: {e}"))?; + staging.write_from_slice(host); + let mut dst = stream.alloc_zeros::(n).context("trunk upload alloc")?; + 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("trunk upload DtoD")?; + } + } + Ok(dst) +} + +fn upload_into(stream: &Arc, host: &[f32], dst: &mut CudaSlice) -> Result<()> { + let n = host.len(); + assert_eq!(n, dst.len()); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("upload_into staging: {e}"))?; + staging.write_from_slice(host); + 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_into DtoD")?; + } + Ok(()) +} + +fn copy_dtod(stream: &Arc, src: &CudaSlice, dst: &mut CudaSlice) -> Result<()> { + let n = src.len(); + assert_eq!(n, dst.len()); + let nbytes = n * std::mem::size_of::(); + unsafe { + let (src_ptr, _g1) = src.device_ptr(stream); + let (dst_ptr, _g2) = dst.device_ptr_mut(stream); + cudarc::driver::result::memcpy_dtod_async(dst_ptr, src_ptr, nbytes, stream.cu_stream()) + .context("copy_dtod")?; + } + Ok(()) +} + +fn download(stream: &Arc, src: &CudaSlice) -> Result> { + let n = src.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("trunk download staging: {e}"))?; + let nbytes = n * std::mem::size_of::(); + unsafe { + let (src_ptr, _g) = src.device_ptr(stream); + cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream()) + .context("trunk download DtoD")?; + } + stream.synchronize().context("trunk download sync")?; + Ok(staging.read_all()) +} diff --git a/crates/ml-alpha/tests/trunk_forward.rs b/crates/ml-alpha/tests/trunk_forward.rs new file mode 100644 index 000000000..803732244 --- /dev/null +++ b/crates/ml-alpha/tests/trunk_forward.rs @@ -0,0 +1,106 @@ +//! CfcTrunk forward — invariants on the sequentially-dispatched perception path. + +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 forward_returns_probs_in_unit_interval() { + let dev = test_device(); + let mut trunk = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0xABCD).expect("init"); + let (probs, proj) = trunk.forward_snapshot(&synthetic_input(1_000_000)).expect("forward"); + assert_eq!(probs.len(), N_HORIZONS); + assert_eq!(proj.len(), PROJ_DIM); + for &p in &probs { + assert!((0.0..=1.0).contains(&p), "head prob {p} out of [0,1]"); + } +} + +#[test] +fn forward_updates_hidden_state() { + let dev = test_device(); + let mut trunk = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0xCAFE).expect("init"); + let h0 = trunk.snapshot_hidden().expect("h0"); + trunk + .forward_snapshot(&synthetic_input(20_000_000)) + .expect("forward"); + let h1 = trunk.snapshot_hidden().expect("h1"); + let diff: f32 = h0.iter().zip(&h1).map(|(a, b)| (a - b).abs()).sum(); + assert!( + diff > 1e-6, + "hidden state must change after a forward step (got total |Δh| = {diff})" + ); +} + +#[test] +fn forward_outputs_are_finite() { + let dev = test_device(); + let mut trunk = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0x1234).expect("init"); + let (probs, proj) = trunk.forward_snapshot(&synthetic_input(50_000_000)).expect("forward"); + for &p in &probs { + assert!(p.is_finite(), "prob {p} not finite"); + } + for &v in &proj { + assert!(v.is_finite(), "proj {v} not finite"); + } +} + +#[test] +fn forward_layer_norm_proj_has_near_zero_mean() { + let dev = test_device(); + let mut trunk = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0x5EED).expect("init"); + let (_, proj) = trunk.forward_snapshot(&synthetic_input(20_000_000)).expect("forward"); + let mean: f32 = proj.iter().sum::() / PROJ_DIM as f32; + assert_relative_eq!(mean, 0.0, epsilon = 1e-3); +} + +#[test] +fn forward_multi_step_advances_state_consistently() { + // After many forward steps, the hidden state should still be finite and + // within the bounded range expected by the CfC formula (tanh of pre). + let dev = test_device(); + let mut trunk = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0xBEEF).expect("init"); + let mut last_ts = 0u64; + for k in 0..50 { + last_ts += 20_000_000; + let mut input = synthetic_input(last_ts); + input.prev_ts_ns = last_ts - 20_000_000; + // Slowly mutate trade_signed_vol so the input changes each step. + input.trade_signed_vol = k as f32 * 0.1; + trunk.forward_snapshot(&input).expect("forward"); + } + let h = trunk.snapshot_hidden().expect("snapshot"); + for (i, v) in h.iter().enumerate() { + assert!(v.is_finite(), "hidden[{i}] = {v} not finite after 50 steps"); + } +}