gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ml-alpha): CfcTrunk forward — sequential perception dispatch

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-16 21:55:00 +02:00
parent 9ddde632b2
commit 1ed81cf6c0
3 changed files with 414 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
crates/ml-alpha/src/cfc/mod.rs
View File

@@ -5,3 +5,6 @@
pub mod snap_features;
pub mod step;
pub mod trunk;
pub use trunk::{CfcConfig, CfcTrunk};

305
crates/ml-alpha/src/cfc/trunk.rs Normal file
View File

@@ -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<CudaStream>` 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<CudaStream>,
// Cubin modules (kept alive for the function refs below)
_snap_module: Arc<CudaModule>,
_step_module: Arc<CudaModule>,
_heads_module: Arc<CudaModule>,
_proj_module: Arc<CudaModule>,
snap_fn: CudaFunction,
step_fn: CudaFunction,
heads_fn: CudaFunction,
proj_fn: CudaFunction,
// Device-resident weights
w_in_d: CudaSlice<f32>,
w_rec_d: CudaSlice<f32>,
b_d: CudaSlice<f32>,
tau_d: CudaSlice<f32>,
heads_w_d: CudaSlice<f32>,
heads_b_d: CudaSlice<f32>,
proj_w_d: CudaSlice<f32>,
proj_b_d: CudaSlice<f32>,
proj_g_d: CudaSlice<f32>,
proj_n_d: CudaSlice<f32>,
// Ping-pong hidden state
h_ping: CudaSlice<f32>,
h_pong: CudaSlice<f32>,
// Pre-allocated per-step scratch (raw input + features + outputs)
bid_px_d: CudaSlice<f32>,
bid_sz_d: CudaSlice<f32>,
ask_px_d: CudaSlice<f32>,
ask_sz_d: CudaSlice<f32>,
prev_bid_sz_d: CudaSlice<f32>,
prev_ask_sz_d: CudaSlice<f32>,
snap_feat_d: CudaSlice<f32>,
probs_d: CudaSlice<f32>,
proj_out_d: CudaSlice<f32>,
}
impl CfcTrunk {
pub fn new_random(dev: &MlDevice, cfg: &CfcConfig, seed: u64) -> Result<Self> {
let stream: Arc<CudaStream> = 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<f32> = (0..cfg.n_hid * cfg.n_in).map(|_| r.gen_range(-scale_in..scale_in)).collect();
let w_rec: Vec<f32> = (0..cfg.n_hid * cfg.n_hid).map(|_| r.gen_range(-scale_rec..scale_rec)).collect();
let b: Vec<f32> = 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<f32> = (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<f32> = (0..N_HORIZONS * cfg.n_hid).map(|_| r.gen_range(-head_scale..head_scale)).collect();
let heads_b: Vec<f32> = vec![0.0; N_HORIZONS];
let proj_w: Vec<f32> = (0..PROJ_DIM * cfg.n_hid).map(|_| r.gen_range(-head_scale..head_scale)).collect();
let proj_b: Vec<f32> = vec![0.0; PROJ_DIM];
let proj_g: Vec<f32> = vec![1.0; PROJ_DIM];
let proj_n: Vec<f32> = 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::<f32>(cfg.n_hid).context("h_ping alloc")?,
h_pong: stream.alloc_zeros::<f32>(cfg.n_hid).context("h_pong alloc")?,
bid_px_d: stream.alloc_zeros::<f32>(10).context("bid_px alloc")?,
bid_sz_d: stream.alloc_zeros::<f32>(10).context("bid_sz alloc")?,
ask_px_d: stream.alloc_zeros::<f32>(10).context("ask_px alloc")?,
ask_sz_d: stream.alloc_zeros::<f32>(10).context("ask_sz alloc")?,
prev_bid_sz_d: stream.alloc_zeros::<f32>(10).context("prev_bid_sz alloc")?,
prev_ask_sz_d: stream.alloc_zeros::<f32>(10).context("prev_ask_sz alloc")?,
snap_feat_d: stream.alloc_zeros::<f32>(FEATURE_DIM).context("snap_feat alloc")?,
probs_d: stream.alloc_zeros::<f32>(N_HORIZONS).context("probs alloc")?,
proj_out_d: stream.alloc_zeros::<f32>(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<Vec<f32>> {
download(&self.stream, &self.h_ping)
}
pub fn config(&self) -> &CfcConfig { &self.cfg }
}
// Helpers ---------------------------------------------------------------
fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
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::<f32>(n).context("trunk upload alloc")?;
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("trunk upload DtoD")?;
}
}
Ok(dst)
}
fn upload_into(stream: &Arc<CudaStream>, host: &[f32], dst: &mut CudaSlice<f32>) -> 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::<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_into DtoD")?;
}
Ok(())
}
fn copy_dtod(stream: &Arc<CudaStream>, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>) -> Result<()> {
let n = src.len();
assert_eq!(n, dst.len());
let nbytes = n * std::mem::size_of::<f32>();
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<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
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::<f32>();
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())
}

106
crates/ml-alpha/tests/trunk_forward.rs Normal file
View File

@@ -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::<f32>() / 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");
}
}