From 611c82b4edcf59049f31ee788afdfedf69fd10cc Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 16 May 2026 21:46:58 +0200 Subject: [PATCH] feat(ml-alpha): snap_feature_assemble kernel + CPU oracle Per-snapshot 32-dim feature vector (mid log-return, spread, depth, OFI, trade-flow, dt). Single-block single-thread kernel; uploads via MappedF32Buffer DtoD into CudaSlice per the addendum Pattern 3. Bit-equiv tested CPU vs GPU at eps<=1e-5 over (synthetic input, reserved-slots-are-zero, zero-prev-mid edge case). Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/cuda/snap_feature_assemble.cu | 57 ++++++++- crates/ml-alpha/src/cfc/mod.rs | 7 ++ crates/ml-alpha/src/cfc/oracle.rs | 28 +++++ crates/ml-alpha/src/cfc/snap_features.rs | 119 ++++++++++++++++++ crates/ml-alpha/src/lib.rs | 1 + .../ml-alpha/tests/snap_feature_bit_equiv.rs | 69 ++++++++++ 6 files changed, 279 insertions(+), 2 deletions(-) create mode 100644 crates/ml-alpha/src/cfc/mod.rs create mode 100644 crates/ml-alpha/src/cfc/oracle.rs create mode 100644 crates/ml-alpha/src/cfc/snap_features.rs create mode 100644 crates/ml-alpha/tests/snap_feature_bit_equiv.rs diff --git a/crates/ml-alpha/cuda/snap_feature_assemble.cu b/crates/ml-alpha/cuda/snap_feature_assemble.cu index fb98190cb..6fac2d08b 100644 --- a/crates/ml-alpha/cuda/snap_feature_assemble.cu +++ b/crates/ml-alpha/cuda/snap_feature_assemble.cu @@ -1,2 +1,55 @@ -// placeholder — real implementation in the task adding this kernel -extern "C" __global__ void snap_feature_assemble_stub() {} +// snap_feature_assemble.cu +// +// Reads one MBP-10 snapshot in struct-of-arrays layout, emits the 32-dim +// feature vector documented in spec Section 2. +// +// Layout enforced bit-for-bit by tests: +// out[0] = mid log-return since prev snapshot +// out[1] = L1 spread in ticks +// out[2..7] = bid-side per-level log-size (L1..L5) +// out[7..12] = ask-side per-level log-size (L1..L5) +// out[12..17]= per-level OFI (bid_delta - ask_delta) +// out[17] = log(trade_count + 1) +// out[18] = trade_signed_vol +// out[19] = dt_ms (snapshot inter-arrival time) +// out[20..32]= reserved, zero +// +// Single-thread per-snapshot kernel; the call site launches with one +// block of one thread per snapshot. + +extern "C" __global__ void snap_feature_assemble( + const float* __restrict__ bid_px, // [10] + const float* __restrict__ bid_sz, // [10] + const float* __restrict__ ask_px, // [10] + const float* __restrict__ ask_sz, // [10] + const float* __restrict__ prev_bid_sz, // [10] for OFI baseline + const float* __restrict__ prev_ask_sz, // [10] + float prev_mid, + float trade_signed_vol, + int trade_count, + long ts_ns, + long prev_ts_ns, + float tick_size, + float* __restrict__ out // [32] +) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + + const float mid = 0.5f * (bid_px[0] + ask_px[0]); + out[0] = (mid > 0.0f && prev_mid > 0.0f) ? logf(mid / prev_mid) : 0.0f; + out[1] = (ask_px[0] - bid_px[0]) / tick_size; + + for (int i = 0; i < 5; ++i) { + out[2 + i] = log1pf(bid_sz[i]); + out[7 + i] = log1pf(ask_sz[i]); + const float bid_delta = bid_sz[i] - prev_bid_sz[i]; + const float ask_delta = ask_sz[i] - prev_ask_sz[i]; + out[12 + i] = bid_delta - ask_delta; + } + + out[17] = log1pf((float) trade_count); + out[18] = trade_signed_vol; + const float dt_ms = (float)(ts_ns - prev_ts_ns) * 1e-6f; + out[19] = dt_ms; + + for (int i = 20; i < 32; ++i) out[i] = 0.0f; +} diff --git a/crates/ml-alpha/src/cfc/mod.rs b/crates/ml-alpha/src/cfc/mod.rs new file mode 100644 index 000000000..dcf741df6 --- /dev/null +++ b/crates/ml-alpha/src/cfc/mod.rs @@ -0,0 +1,7 @@ +//! CfC trunk module — snapshot-level Closed-form Continuous-time +//! recurrent network. Per spec Section 2 + 2026-05-16 stacked +//! amendment: this module owns the CfC layer that sits on top of a +//! frozen Mamba2 sequence encoder. + +pub mod snap_features; +pub mod oracle; diff --git a/crates/ml-alpha/src/cfc/oracle.rs b/crates/ml-alpha/src/cfc/oracle.rs new file mode 100644 index 000000000..93b16b28f --- /dev/null +++ b/crates/ml-alpha/src/cfc/oracle.rs @@ -0,0 +1,28 @@ +//! CPU oracles for bit-equiv tests of `cfc` kernels. + +use super::snap_features::{Mbp10RawInput, ES_TICK_SIZE, FEATURE_DIM}; + +/// CPU mirror of `snap_feature_assemble.cu`. Indices match the spec +/// Section 2 feature table. `prev_bid_sz` and `prev_ask_sz` are assumed +/// zero (matches the one-shot GPU helper's initialization). +pub fn snap_feature_assemble_cpu(input: &Mbp10RawInput) -> [f32; FEATURE_DIM] { + let mut out = [0f32; FEATURE_DIM]; + let mid = 0.5 * (input.bid_px[0] + input.ask_px[0]); + out[0] = if mid > 0.0 && input.prev_mid > 0.0 { + (mid / input.prev_mid).ln() + } else { + 0.0 + }; + out[1] = (input.ask_px[0] - input.bid_px[0]) / ES_TICK_SIZE; + for i in 0..5 { + out[2 + i] = (input.bid_sz[i] + 1.0).ln(); + out[7 + i] = (input.ask_sz[i] + 1.0).ln(); + // OFI per level — prev_*_sz initialized to zero in the GPU helper, + // so bid_delta = bid_sz[i], ask_delta = ask_sz[i]. + out[12 + i] = input.bid_sz[i] - input.ask_sz[i]; + } + out[17] = (input.trade_count as f32 + 1.0).ln(); + out[18] = input.trade_signed_vol; + out[19] = (input.ts_ns - input.prev_ts_ns) as f32 * 1e-6; + out +} diff --git a/crates/ml-alpha/src/cfc/snap_features.rs b/crates/ml-alpha/src/cfc/snap_features.rs new file mode 100644 index 000000000..1da586213 --- /dev/null +++ b/crates/ml-alpha/src/cfc/snap_features.rs @@ -0,0 +1,119 @@ +//! `snap_feature_assemble` kernel binding. +//! +//! Per spec Section 2, builds a 32-dim feature vector from one MBP-10 +//! snapshot. Production code calls this inside the captured Graph A +//! (Task 11); the standalone helper exists for bit-equiv tests. + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg}; +use ml_core::device::MlDevice; + +use crate::pinned_mem::MappedF32Buffer; + +const KERNEL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/snap_feature_assemble.cubin")); + +pub const ES_TICK_SIZE: f32 = 0.25; +pub const FEATURE_DIM: usize = 32; + +/// One raw MBP-10 snapshot in struct-of-arrays form. Phase A loader +/// converts predecoded sidecar rows to this shape. +#[derive(Clone, Copy, Debug, Default)] +pub struct Mbp10RawInput { + pub bid_px: [f32; 10], + pub bid_sz: [f32; 10], + pub ask_px: [f32; 10], + pub ask_sz: [f32; 10], + pub prev_mid: f32, + pub trade_signed_vol: f32, + pub trade_count: u32, + pub ts_ns: u64, + pub prev_ts_ns: u64, +} + +/// Standalone GPU run for bit-equiv tests. Production code uses the +/// captured Graph A path in `cfc::trunk` (Task 11). +pub fn snap_feature_assemble_gpu(dev: &MlDevice, input: &Mbp10RawInput) -> Result<[f32; FEATURE_DIM]> { + let stream: &Arc = dev.cuda_stream().context("CUDA stream")?; + let ctx = dev.cuda_context().context("CUDA context")?; + let module = ctx.load_cubin(KERNEL_CUBIN.to_vec()).context("load snap_feature_assemble cubin")?; + let func = module.load_function("snap_feature_assemble").context("load function")?; + + // Stage raw input via mapped-pinned -> DtoD into CudaSlice. For one-shot + // tests this is fine; production wires the same data through the + // captured Graph A's mapped-pinned slot. + let bid_px = upload(stream, &input.bid_px)?; + let bid_sz = upload(stream, &input.bid_sz)?; + let ask_px = upload(stream, &input.ask_px)?; + let ask_sz = upload(stream, &input.ask_sz)?; + let prev_bid_sz = stream.alloc_zeros::(10).context("prev_bid_sz alloc")?; + let prev_ask_sz = stream.alloc_zeros::(10).context("prev_ask_sz alloc")?; + let mut out_d = stream.alloc_zeros::(FEATURE_DIM).context("out alloc")?; + + 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 cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }; + + let mut launch = stream.launch_builder(&func); + launch + .arg(&bid_px) + .arg(&bid_sz) + .arg(&ask_px) + .arg(&ask_sz) + .arg(&prev_bid_sz) + .arg(&prev_ask_sz) + .arg(&prev_mid) + .arg(&trade_signed_vol) + .arg(&trade_count) + .arg(&ts_ns) + .arg(&prev_ts_ns) + .arg(&tick_size) + .arg(&mut out_d); + unsafe { launch.launch(cfg).context("snap_feature_assemble launch")?; } + + let v = download(stream, &out_d)?; + let mut out = [0f32; FEATURE_DIM]; + out.copy_from_slice(&v); + Ok(out) +} + +fn upload(stream: &Arc, host: &[f32]) -> Result> { + let n = host.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("upload staging: {e}"))?; + staging.write_from_slice(host); + let mut dst = stream.alloc_zeros::(n).context("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("upload DtoD")?; + } + } + Ok(dst) +} + +fn download(stream: &Arc, src: &CudaSlice) -> Result> { + let n = src.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("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("download DtoD")?; + } + stream.synchronize().context("download sync")?; + Ok(staging.read_all()) +} diff --git a/crates/ml-alpha/src/lib.rs b/crates/ml-alpha/src/lib.rs index ed209a18f..3b53403df 100644 --- a/crates/ml-alpha/src/lib.rs +++ b/crates/ml-alpha/src/lib.rs @@ -25,6 +25,7 @@ // Task 10+: pub mod trainer; // Task 12: pub mod data; // Task 16+: pub mod gate; +pub mod cfc; pub mod isv; pub mod pinned; pub mod pinned_mem; diff --git a/crates/ml-alpha/tests/snap_feature_bit_equiv.rs b/crates/ml-alpha/tests/snap_feature_bit_equiv.rs new file mode 100644 index 000000000..59e8c175b --- /dev/null +++ b/crates/ml-alpha/tests/snap_feature_bit_equiv.rs @@ -0,0 +1,69 @@ +//! Bit-equiv test: CPU snap_feature oracle vs GPU kernel. + +use approx::assert_relative_eq; +use ml_alpha::cfc::oracle::snap_feature_assemble_cpu; +use ml_alpha::cfc::snap_features::{snap_feature_assemble_gpu, Mbp10RawInput, FEATURE_DIM}; +use ml_core::device::MlDevice; + +fn test_device() -> MlDevice { + MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests") +} + +fn synthetic_input() -> 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: 1_700_000_001_000_000_000, + prev_ts_ns: 1_700_000_000_500_000_000, + } +} + +#[test] +fn snap_feature_oracle_matches_gpu() { + let dev = test_device(); + let input = synthetic_input(); + let cpu = snap_feature_assemble_cpu(&input); + let gpu = snap_feature_assemble_gpu(&dev, &input).expect("gpu run"); + for i in 0..FEATURE_DIM { + assert_relative_eq!( + cpu[i], + gpu[i], + epsilon = 1e-5, + max_relative = 1e-5 + ); + } +} + +#[test] +fn snap_feature_reserved_slots_are_zero() { + let dev = test_device(); + let input = synthetic_input(); + let gpu = snap_feature_assemble_gpu(&dev, &input).expect("gpu run"); + for i in 20..32 { + assert_eq!(gpu[i], 0.0, "reserved slot {i} should be 0, got {}", gpu[i]); + } +} + +#[test] +fn snap_feature_zero_prev_mid_yields_zero_log_return() { + let dev = test_device(); + let mut input = synthetic_input(); + input.prev_mid = 0.0; + let gpu = snap_feature_assemble_gpu(&dev, &input).expect("gpu run"); + assert_eq!(gpu[0], 0.0, "log-return should be 0 when prev_mid is 0"); +}