diff --git a/crates/ml-alpha/cuda/cfc_step.cu b/crates/ml-alpha/cuda/cfc_step.cu index 11508d365..bf0c95325 100644 --- a/crates/ml-alpha/cuda/cfc_step.cu +++ b/crates/ml-alpha/cuda/cfc_step.cu @@ -1,2 +1,36 @@ -// placeholder — real implementation in the task adding this kernel -extern "C" __global__ void cfc_step_stub() {} +// cfc_step.cu — one CfC time step (Hasani 2022 closed-form recurrence). +// +// pre[i] = sum_k w_in[i,k] * x[k] + sum_k w_rec[i,k] * h_old[k] + b[i] +// decay[i] = exp(-dt_s / max(tau[i], 1e-6)) +// h_new[i] = h_old[i] * decay[i] + (1 - decay[i]) * tanh(pre[i]) +// +// Launched as one block of n_hid threads (or grid of blocks if n_hid > 128). +// Each thread computes one hidden unit; reductions are intra-thread (each +// thread does its own dot product). No atomicAdd. Per +// `feedback_no_atomicadd.md` and the GPU/CPU contract. + +extern "C" __global__ void cfc_step( + const float* __restrict__ w_in, // [n_hid, n_in], row-major + const float* __restrict__ w_rec, // [n_hid, n_hid], row-major + const float* __restrict__ b, // [n_hid] + const float* __restrict__ tau, // [n_hid] + const float* __restrict__ x, // [n_in] + const float* __restrict__ h_old, // [n_hid] + float dt_s, + int n_in, + int n_hid, + float* __restrict__ h_new // [n_hid] +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n_hid) return; + + float pre = b[i]; + for (int k = 0; k < n_in; ++k) { + pre += w_in[i * n_in + k] * x[k]; + } + for (int k = 0; k < n_hid; ++k) { + pre += w_rec[i * n_hid + k] * h_old[k]; + } + const float decay = expf(-dt_s / fmaxf(tau[i], 1e-6f)); + h_new[i] = h_old[i] * decay + (1.0f - decay) * tanhf(pre); +} diff --git a/crates/ml-alpha/src/cfc/mod.rs b/crates/ml-alpha/src/cfc/mod.rs index dcf741df6..38bf1ab33 100644 --- a/crates/ml-alpha/src/cfc/mod.rs +++ b/crates/ml-alpha/src/cfc/mod.rs @@ -4,4 +4,4 @@ //! frozen Mamba2 sequence encoder. pub mod snap_features; -pub mod oracle; +pub mod step; diff --git a/crates/ml-alpha/src/cfc/oracle.rs b/crates/ml-alpha/src/cfc/oracle.rs deleted file mode 100644 index 93b16b28f..000000000 --- a/crates/ml-alpha/src/cfc/oracle.rs +++ /dev/null @@ -1,28 +0,0 @@ -//! 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/step.rs b/crates/ml-alpha/src/cfc/step.rs new file mode 100644 index 000000000..582676eea --- /dev/null +++ b/crates/ml-alpha/src/cfc/step.rs @@ -0,0 +1,104 @@ +//! `cfc_step` kernel binding — one CfC time step. + +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"), "/cfc_step.cubin")); + +#[derive(Clone, Debug)] +pub struct CfcWeights { + pub w_in: Vec, // [n_hid, n_in] + pub w_rec: Vec, // [n_hid, n_hid] + pub b: Vec, // [n_hid] + pub tau: Vec, // [n_hid] + pub n_in: usize, + pub n_hid: usize, +} + +pub fn cfc_step_gpu( + dev: &MlDevice, + w: &CfcWeights, + x: &[f32], + h_old: &[f32], + dt_s: f32, +) -> Result> { + assert_eq!(x.len(), w.n_in); + assert_eq!(h_old.len(), w.n_hid); + + let stream: &Arc = dev.cuda_stream().context("cfc_step stream")?; + let ctx = dev.cuda_context().context("cfc_step ctx")?; + let module = ctx.load_cubin(KERNEL_CUBIN.to_vec()).context("load cfc_step cubin")?; + let func = module.load_function("cfc_step").context("load function")?; + + let w_in_d = upload(stream, &w.w_in)?; + let w_rec_d = upload(stream, &w.w_rec)?; + let b_d = upload(stream, &w.b)?; + let tau_d = upload(stream, &w.tau)?; + let x_d = upload(stream, x)?; + let h_old_d = upload(stream, h_old)?; + let mut h_new_d = stream + .alloc_zeros::(w.n_hid) + .context("h_new alloc")?; + + let n_in_i = w.n_in as i32; + let n_hid_i = w.n_hid as i32; + let block_dim = 128.min(w.n_hid as u32); + let grid_dim = ((w.n_hid as u32) + block_dim - 1) / block_dim; + let cfg = LaunchConfig { + grid_dim: (grid_dim, 1, 1), + block_dim: (block_dim, 1, 1), + shared_mem_bytes: 0, + }; + + let mut launch = stream.launch_builder(&func); + launch + .arg(&w_in_d) + .arg(&w_rec_d) + .arg(&b_d) + .arg(&tau_d) + .arg(&x_d) + .arg(&h_old_d) + .arg(&dt_s) + .arg(&n_in_i) + .arg(&n_hid_i) + .arg(&mut h_new_d); + unsafe { launch.launch(cfg).context("cfc_step launch")?; } + + download(stream, &h_new_d) +} + +fn upload(stream: &Arc, host: &[f32]) -> Result> { + let n = host.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("cfc_step upload staging: {e}"))?; + staging.write_from_slice(host); + let mut dst = stream.alloc_zeros::(n).context("cfc_step 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("cfc_step 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!("cfc_step 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("cfc_step download DtoD")?; + } + stream.synchronize().context("cfc_step download sync")?; + Ok(staging.read_all()) +} diff --git a/crates/ml-alpha/tests/cfc_step_bit_equiv.rs b/crates/ml-alpha/tests/cfc_step_bit_equiv.rs new file mode 100644 index 000000000..65f4a393b --- /dev/null +++ b/crates/ml-alpha/tests/cfc_step_bit_equiv.rs @@ -0,0 +1,135 @@ +//! CfC step GPU invariants. +//! +//! Per `feedback_no_cpu_test_fallbacks.md`: no CPU oracle. Validation is +//! via analytically-known invariants of the closed-form recurrence: +//! +//! h_new[i] = h_old[i] * decay[i] + (1 - decay[i]) * tanh(W_in·x + W_rec·h_old + b) +//! decay[i] = exp(-dt_s / tau[i]) +//! +//! Key invariants tested: +//! - dt = 0 → decay = 1 → h_new = h_old (identity) +//! - tau very large with dt small → decay ≈ 1 → h_new ≈ h_old +//! - Zero weights + zero bias → pre = 0 → tanh(0) = 0 → h_new = h_old · decay +//! - Output is bounded by max(|h_old|, 1) since tanh ∈ [-1, 1] + +use approx::assert_relative_eq; +use ml_alpha::cfc::step::{cfc_step_gpu, CfcWeights}; +use ml_core::device::MlDevice; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +fn test_device() -> MlDevice { + MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests") +} + +fn rand_weights(seed: u64, n_in: usize, n_hid: usize) -> CfcWeights { + let mut r = ChaCha8Rng::seed_from_u64(seed); + let w_in: Vec = (0..n_hid * n_in).map(|_| r.gen_range(-0.1..0.1)).collect(); + let w_rec: Vec = (0..n_hid * n_hid).map(|_| r.gen_range(-0.05..0.05)).collect(); + let b: Vec = (0..n_hid).map(|_| r.gen_range(-0.01..0.01)).collect(); + let tau: Vec = (0..n_hid).map(|_| r.gen_range(0.05..2.0)).collect(); + CfcWeights { + w_in, + w_rec, + b, + tau, + n_in, + n_hid, + } +} + +#[test] +fn cfc_step_zero_dt_returns_h_old() { + let dev = test_device(); + let w = rand_weights(0xDEAD_BEEF, 32, 128); + let h_old: Vec = (0..w.n_hid).map(|i| 0.01 * i as f32).collect(); + let x = vec![0.0; w.n_in]; + let h_gpu = cfc_step_gpu(&dev, &w, &x, &h_old, 0.0).expect("gpu"); + for i in 0..w.n_hid { + assert_relative_eq!(h_gpu[i], h_old[i], epsilon = 1e-6); + } +} + +#[test] +fn cfc_step_zero_weights_decays_h_old() { + // Zero W_in, W_rec, b → pre = 0 → tanh(0) = 0 → h_new = h_old · decay. + let dev = test_device(); + let n_in = 32; + let n_hid = 64; + let w = CfcWeights { + w_in: vec![0.0; n_hid * n_in], + w_rec: vec![0.0; n_hid * n_hid], + b: vec![0.0; n_hid], + tau: vec![1.0; n_hid], + n_in, + n_hid, + }; + let h_old: Vec = (0..n_hid).map(|i| 0.5 + 0.01 * i as f32).collect(); + let dt_s = 0.5_f32; + let expected_decay = (-dt_s / 1.0_f32).exp(); + let h_gpu = cfc_step_gpu(&dev, &w, &vec![0.0; n_in], &h_old, dt_s).expect("gpu"); + for i in 0..n_hid { + let expected = h_old[i] * expected_decay; + assert_relative_eq!(h_gpu[i], expected, epsilon = 1e-5); + } +} + +#[test] +fn cfc_step_large_tau_preserves_h_old() { + // tau >> dt → decay ≈ 1 → h_new ≈ h_old. + let dev = test_device(); + let n_in = 16; + let n_hid = 32; + let mut r = ChaCha8Rng::seed_from_u64(7); + let w = CfcWeights { + w_in: (0..n_hid * n_in).map(|_| r.gen_range(-0.1..0.1)).collect(), + w_rec: (0..n_hid * n_hid).map(|_| r.gen_range(-0.05..0.05)).collect(), + b: vec![0.0; n_hid], + tau: vec![10_000.0; n_hid], // huge + n_in, + n_hid, + }; + let h_old: Vec = (0..n_hid).map(|_| r.gen_range(-0.5..0.5)).collect(); + let x: Vec = (0..n_in).map(|_| r.gen_range(-0.5..0.5)).collect(); + let h_gpu = cfc_step_gpu(&dev, &w, &x, &h_old, 0.001).expect("gpu"); + for i in 0..n_hid { + // decay = exp(-0.001/10000) ≈ 1 - 1e-7 ≈ 1 + assert_relative_eq!(h_gpu[i], h_old[i], epsilon = 1e-3); + } +} + +#[test] +fn cfc_step_output_bounded_by_one_plus_h_old_norm() { + // tanh ∈ [-1, 1], decay ∈ [0, 1], so |h_new[i]| ≤ |h_old[i]| · 1 + 1 · 1 + let dev = test_device(); + let w = rand_weights(0xCFC_5EED, 32, 128); + let mut r = ChaCha8Rng::seed_from_u64(42); + let x: Vec = (0..w.n_in).map(|_| r.gen_range(-1.0..1.0)).collect(); + let h_old: Vec = (0..w.n_hid).map(|_| r.gen_range(-1.0..1.0)).collect(); + let h_gpu = cfc_step_gpu(&dev, &w, &x, &h_old, 0.02).expect("gpu"); + for i in 0..w.n_hid { + let bound = h_old[i].abs() + 1.0; + assert!( + h_gpu[i].abs() <= bound + 1e-5, + "h_new[{i}]={} exceeds |h_old|+1 bound {}", + h_gpu[i], + bound + ); + } +} + +#[test] +fn cfc_step_runs_on_small_hidden() { + // n_hid < 128 — exercises the grid-dim ceiling. Just assert the call + // succeeds and outputs are finite. + let dev = test_device(); + let w = rand_weights(0x1234, 8, 16); + let mut r = ChaCha8Rng::seed_from_u64(99); + let x: Vec = (0..w.n_in).map(|_| r.gen_range(-1.0..1.0)).collect(); + let h_old: Vec = (0..w.n_hid).map(|_| r.gen_range(-1.0..1.0)).collect(); + let h_gpu = cfc_step_gpu(&dev, &w, &x, &h_old, 0.01).expect("gpu"); + assert_eq!(h_gpu.len(), w.n_hid); + for v in &h_gpu { + assert!(v.is_finite(), "h_new must be finite, got {}", v); + } +} diff --git a/crates/ml-alpha/tests/snap_feature_bit_equiv.rs b/crates/ml-alpha/tests/snap_feature_bit_equiv.rs index 59e8c175b..c885f7464 100644 --- a/crates/ml-alpha/tests/snap_feature_bit_equiv.rs +++ b/crates/ml-alpha/tests/snap_feature_bit_equiv.rs @@ -1,8 +1,10 @@ -//! Bit-equiv test: CPU snap_feature oracle vs GPU kernel. +//! `snap_feature_assemble` GPU invariants. +//! +//! Per `feedback_no_cpu_test_fallbacks.md`: no CPU oracle. Validation is +//! via analytically-known synthetic inputs + property assertions. 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_alpha::cfc::snap_features::{snap_feature_assemble_gpu, Mbp10RawInput, ES_TICK_SIZE, FEATURE_DIM}; use ml_core::device::MlDevice; fn test_device() -> MlDevice { @@ -34,36 +36,77 @@ fn synthetic_input() -> Mbp10RawInput { } #[test] -fn snap_feature_oracle_matches_gpu() { +fn spread_in_ticks_matches_synthetic() { 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 - ); + let gpu = snap_feature_assemble_gpu(&dev, &input).expect("gpu"); + // ask[0]=5500.25, bid[0]=5500.00, tick_size=0.25 → 1 tick. + assert_relative_eq!(gpu[1], (input.ask_px[0] - input.bid_px[0]) / ES_TICK_SIZE, epsilon = 1e-6); + assert_relative_eq!(gpu[1], 1.0, epsilon = 1e-6); +} + +#[test] +fn mid_log_return_sign_tracks_input() { + let dev = test_device(); + // mid=5500.125; prev_mid=5499.875 → positive log-return + let input_up = synthetic_input(); + let gpu_up = snap_feature_assemble_gpu(&dev, &input_up).expect("gpu"); + assert!(gpu_up[0] > 0.0, "mid > prev_mid must give positive log-return, got {}", gpu_up[0]); + + let mut input_down = synthetic_input(); + input_down.prev_mid = 5501.0; + let gpu_down = snap_feature_assemble_gpu(&dev, &input_down).expect("gpu"); + assert!(gpu_down[0] < 0.0, "mid < prev_mid must give negative log-return, got {}", gpu_down[0]); +} + +#[test] +fn ofi_sign_tracks_bid_minus_ask_size() { + let dev = test_device(); + let input = synthetic_input(); + let gpu = snap_feature_assemble_gpu(&dev, &input).expect("gpu"); + // prev sizes init to 0, so OFI[i] = bid_sz[i] - ask_sz[i]. + // bid_sz[0]=12, ask_sz[0]=10 → OFI[0] = 2. + assert_relative_eq!(gpu[12], input.bid_sz[0] - input.ask_sz[0], epsilon = 1e-6); + // Level 4: bid_sz[4]=16, ask_sz[4]=16 → OFI = 0. + assert_relative_eq!(gpu[16], input.bid_sz[4] - input.ask_sz[4], epsilon = 1e-6); +} + +#[test] +fn log_size_features_strictly_positive() { + let dev = test_device(); + let input = synthetic_input(); + let gpu = snap_feature_assemble_gpu(&dev, &input).expect("gpu"); + // bid log-sizes at indices 2..7 and ask log-sizes at 7..12 are log1pf of + // positive sizes — must all be strictly positive. + for i in 2..12 { + assert!(gpu[i] > 0.0, "log-size at {i} must be > 0 for positive size input, got {}", gpu[i]); } } #[test] -fn snap_feature_reserved_slots_are_zero() { +fn 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]); + let gpu = snap_feature_assemble_gpu(&dev, &input).expect("gpu"); + for i in 20..FEATURE_DIM { + assert_eq!(gpu[i], 0.0, "reserved slot {i} must be 0, got {}", gpu[i]); } } #[test] -fn snap_feature_zero_prev_mid_yields_zero_log_return() { +fn 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"); + let gpu = snap_feature_assemble_gpu(&dev, &input).expect("gpu"); + assert_eq!(gpu[0], 0.0, "log-return must be 0 when prev_mid is 0"); +} + +#[test] +fn dt_in_ms_matches_input_nanoseconds() { + let dev = test_device(); + let input = synthetic_input(); + let gpu = snap_feature_assemble_gpu(&dev, &input).expect("gpu"); + let expected_ms = (input.ts_ns - input.prev_ts_ns) as f32 * 1e-6; + assert_relative_eq!(gpu[19], expected_ms, epsilon = 1e-3); }