diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs new file mode 100644 index 000000000..2e256d150 --- /dev/null +++ b/crates/ml-alpha/build.rs @@ -0,0 +1,107 @@ +//! ml-alpha build.rs — precompile the Mamba2 temporal SSM scan kernel. +//! +//! Compiles `crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu` (the only +//! kernel ml-alpha needs from the parent ml crate's library, for Phase 1d.1's +//! GPU-resident stateful encoder) into an arch-specific cubin in OUT_DIR. +//! +//! Discipline: +//! - **rerun-if-env-changed=CUDA_COMPUTE_CAP** is mandatory: cargo target PVCs +//! are shared across L40S (sm_89) and H100 (sm_90) builds; without this, +//! stale cubins surface as `CUDA_ERROR_NO_BINARY_FOR_GPU` at runtime. +//! - The kernel source has zero `#include`s, so no header prepending is +//! needed (unlike ml-core's autograd kernels). +//! - On non-CUDA builds (`CARGO_FEATURE_CUDA` unset) we silently skip; the +//! library still compiles, only the Mamba2Block::new() runtime path fails +//! with a missing-cubin error. + +use std::path::{Path, PathBuf}; +use std::process::Command; + +fn main() { + println!("cargo:rerun-if-changed=build.rs"); + + if std::env::var("CARGO_FEATURE_CUDA").is_err() { + return; + } + + let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap()); + + // The kernel lives in the parent ml crate. We don't depend on that crate's + // library — only the .cu source file. Use a relative path from this + // crate's manifest directory; canonicalize so cargo emits a stable + // rerun-if-changed line. + let kernel_src_path = Path::new("../ml/src/cuda_pipeline/mamba2_temporal_kernel.cu") + .canonicalize() + .unwrap_or_else(|_| { + let workspace = std::env::var("CARGO_MANIFEST_DIR") + .map(PathBuf::from) + .unwrap_or_default(); + workspace + .join("../ml/src/cuda_pipeline/mamba2_temporal_kernel.cu") + .canonicalize() + .expect("Cannot locate mamba2_temporal_kernel.cu in ../ml/src/cuda_pipeline/") + }); + println!("cargo:rerun-if-changed={}", kernel_src_path.display()); + + println!("cargo:rerun-if-env-changed=CUDA_COMPUTE_CAP"); + let cuda_compute_cap = std::env::var("CUDA_COMPUTE_CAP").unwrap_or_else(|_| "80".to_string()); + let arch = format!("sm_{cuda_compute_cap}"); + + let nvcc = match find_nvcc() { + Some(p) => p, + None => { + eprintln!(" ml-alpha: nvcc not found, skipping mamba2 kernel precompilation"); + eprintln!(" Install CUDA toolkit or set CUDA_HOME for GPU builds"); + return; + } + }; + + let cubin_path = out_dir.join("mamba2_temporal_kernel.cubin"); + + let status = Command::new(&nvcc) + .args([ + "-cubin", + &format!("-arch={arch}"), + "-O3", + "--ftz=true", + "--fmad=true", + "--prec-div=true", + "--prec-sqrt=true", + "-o", + cubin_path.to_str().unwrap(), + kernel_src_path.to_str().unwrap(), + ]) + .status(); + + match status { + Ok(s) if s.success() => { + eprintln!( + " ml-alpha: Compiled mamba2_temporal_kernel.cu -> mamba2_temporal_kernel.cubin ({arch})" + ); + } + Ok(s) => panic!( + "ml-alpha: nvcc failed to compile mamba2_temporal_kernel.cu (exit={})", + s.code().unwrap_or(-1) + ), + Err(e) => panic!("ml-alpha: nvcc invocation error: {e}"), + } +} + +fn find_nvcc() -> Option { + if let Ok(home) = std::env::var("CUDA_HOME") { + let nvcc = PathBuf::from(home).join("bin/nvcc"); + if nvcc.exists() { + return Some(nvcc); + } + } + for path in &["/usr/local/cuda/bin/nvcc", "/usr/bin/nvcc"] { + let p = PathBuf::from(path); + if p.exists() { + return Some(p); + } + } + match Command::new("nvcc").arg("--version").output() { + Ok(output) if output.status.success() => Some(PathBuf::from("nvcc")), + _ => None, + } +} diff --git a/crates/ml-alpha/src/lib.rs b/crates/ml-alpha/src/lib.rs index d2137bec2..15f0d7bd4 100644 --- a/crates/ml-alpha/src/lib.rs +++ b/crates/ml-alpha/src/lib.rs @@ -52,6 +52,7 @@ pub mod training; pub mod eval; pub mod metrics_detail; pub mod calibration; +pub mod mamba2_block; pub use fxcache_reader::{FxCacheReader, FxCacheRecord, FxCacheMetadata}; pub use purged_split::{PurgedSplit, SplitIndices}; diff --git a/crates/ml-alpha/src/mamba2_block.rs b/crates/ml-alpha/src/mamba2_block.rs new file mode 100644 index 000000000..43dfa22ad --- /dev/null +++ b/crates/ml-alpha/src/mamba2_block.rs @@ -0,0 +1,237 @@ +//! Phase 1d.1 — From-scratch Mamba2 sequence block for ml-alpha. +//! +//! GPU-pure stateful encoder over snapshot streams. Architecture: +//! +//! ```text +//! input [B, K, in_dim] +//! │ cuBLAS GEMM (W_in) +//! ▼ +//! x [B, K, hidden_dim] +//! ├─ cuBLAS GEMM (W_a) ──▶ a_proj [B, K, state_dim] +//! ├─ cuBLAS GEMM (W_b) ──▶ b_proj [B, K, state_dim] +//! ▼ +//! mamba2_scan_projected_fwd kernel +//! (selective SSM scan over K timesteps; sigmoid gating; W_c mixes +//! state into per-position hidden output) +//! ▼ +//! h_enriched [B, hidden_dim] (residual added in kernel via h_s2; we +//! pass zeros for the SSM-residual term) +//! │ cuBLAS GEMM (W_out) +//! ▼ +//! logit [B, 1] +//! ``` +//! +//! Backward: analytical via `mamba2_scan_projected_bwd` kernel; cuBLAS +//! transposed-GEMMs for the projection backward passes. Lands in a follow-up +//! session — this module establishes the skeleton: weight allocation, +//! cubin loading, kernel-symbol resolution. +//! +//! Discipline (per project memory): +//! - All weights / state / activations live on GPU (`CudaSlice`). +//! - Weight init: `OwnedGpuLinear::xavier` from ml-core, which uses pinned +//! host buffers for the seed values. +//! - No `atomicAdd` (kernel uses block tree-reduce). +//! - State dim hardcoded at ≤ 16 in the kernel (`float x[16]`); we expose +//! it as a config field but the constructor rejects values > 16. + +use std::sync::Arc; + +use anyhow::{anyhow, Result}; +use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream}; +use ml_core::cuda_autograd::linear::OwnedGpuLinear; + +/// Mamba2 hidden state dim is hardcoded at 16 floats per (sample, batch) in +/// the kernel (`float x[16]`). Configurations exceeding this are rejected. +pub const MAMBA2_KERNEL_STATE_MAX: usize = 16; + +/// Configuration for a Mamba2 sequence block. +#[derive(Debug, Clone)] +pub struct Mamba2BlockConfig { + /// Per-snapshot feature dim (e.g. 81 for the snapshot pipeline). + pub in_dim: usize, + /// Hidden / model dim (=`sh2` in the kernel signature). + pub hidden_dim: usize, + /// SSM state dim (≤ 16). + pub state_dim: usize, + /// Sequence length per batch (K in kernel; number of timesteps). + pub seq_len: usize, +} + +impl Mamba2BlockConfig { + pub fn validate(&self) -> Result<()> { + if self.in_dim == 0 || self.hidden_dim == 0 || self.state_dim == 0 || self.seq_len < 2 { + return Err(anyhow!( + "Mamba2BlockConfig: invalid dims (in={}, hidden={}, state={}, seq_len={})", + self.in_dim, self.hidden_dim, self.state_dim, self.seq_len + )); + } + if self.state_dim > MAMBA2_KERNEL_STATE_MAX { + return Err(anyhow!( + "Mamba2BlockConfig: state_dim {} exceeds kernel max {} (mamba2_temporal_kernel.cu \ + hardcodes `float x[{}]` per thread)", + self.state_dim, MAMBA2_KERNEL_STATE_MAX, MAMBA2_KERNEL_STATE_MAX + )); + } + Ok(()) + } +} + +/// GPU-resident Mamba2 sequence block. +/// +/// Owns all parameter tensors (`W_in`, `W_a`, `W_b`, `W_c`, `W_out`) on the +/// CUDA device. The forward + backward kernels are loaded once at +/// construction from the precompiled cubin. +pub struct Mamba2Block { + pub config: Mamba2BlockConfig, + pub stream: Arc, + + // ── Parameters ──────────────────────────────────────────────────── + /// Input projection: `in_dim → hidden_dim`. Cuts the snapshot row down + /// to model space before the selective scan. + pub w_in: OwnedGpuLinear, + /// A-projection (gate): `hidden_dim → state_dim`. The kernel applies + /// sigmoid to produce per-state gates (`a_proj`). + pub w_a: OwnedGpuLinear, + /// B-projection (input contribution): `hidden_dim → state_dim`. Added + /// to the carried state each step (`b_proj`). + pub w_b: OwnedGpuLinear, + /// Context-mix weights `W_c` in the kernel signature, shape + /// `[hidden_dim, state_dim]`. Multiplies the final state to produce the + /// per-position context that's added to `h_s2`. Stored as a raw + /// `CudaSlice` (no bias) because the kernel reads it directly. + pub w_c: CudaSlice, + /// Output projection: `hidden_dim → 1`. Reduces the enriched hidden + /// state to the binary direction logit. + pub w_out: OwnedGpuLinear, + + // ── Kernel handles ──────────────────────────────────────────────── + _module: Arc, + pub kernel_fwd: CudaFunction, + pub kernel_bwd: CudaFunction, +} + +impl Mamba2Block { + /// Construct a fresh Mamba2 block with Xavier-initialised projections + /// and the SSM scan kernels loaded from the precompiled cubin. + pub fn new(config: Mamba2BlockConfig, stream: Arc) -> Result { + config.validate()?; + + // ── Load the precompiled cubin and resolve kernel symbols. ──── + // The cubin is produced by `build.rs` at compile time; no nvrtc. + static CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/mamba2_temporal_kernel.cubin")); + let module = stream + .context() + .load_cubin(CUBIN.to_vec()) + .map_err(|e| anyhow!("Mamba2Block: cubin load failed: {e}"))?; + let kernel_fwd = module + .load_function("mamba2_scan_projected_fwd") + .map_err(|e| anyhow!("Mamba2Block: forward kernel symbol resolve: {e}"))?; + let kernel_bwd = module + .load_function("mamba2_scan_projected_bwd") + .map_err(|e| anyhow!("Mamba2Block: backward kernel symbol resolve: {e}"))?; + + // ── Parameter allocation + initialisation. ────────────────────── + // All on GPU; init helpers in ml-core::cuda_autograd::init use + // pinned host buffers for the Glorot/Xavier seed transfer. + let w_in = OwnedGpuLinear::xavier("mamba2.w_in", config.in_dim, config.hidden_dim, &stream) + .map_err(|e| anyhow!("init w_in: {e}"))?; + let w_a = OwnedGpuLinear::xavier( + "mamba2.w_a", config.hidden_dim, config.state_dim, &stream, + ) + .map_err(|e| anyhow!("init w_a: {e}"))?; + let w_b = OwnedGpuLinear::xavier( + "mamba2.w_b", config.hidden_dim, config.state_dim, &stream, + ) + .map_err(|e| anyhow!("init w_b: {e}"))?; + + // W_c is a raw kernel-fed weight (no bias). Initialise via + // OwnedGpuLinear::xavier then keep only the weight slice. + let w_c_linear = OwnedGpuLinear::xavier( + "mamba2.w_c", config.state_dim, config.hidden_dim, &stream, + ) + .map_err(|e| anyhow!("init w_c: {e}"))?; + let w_c = w_c_linear.weight; // discard bias for kernel-direct use + + let w_out = OwnedGpuLinear::xavier("mamba2.w_out", config.hidden_dim, 1, &stream) + .map_err(|e| anyhow!("init w_out: {e}"))?; + + Ok(Self { + config, + stream, + w_in, + w_a, + w_b, + w_c, + w_out, + _module: module, + kernel_fwd, + kernel_bwd, + }) + } + + /// Total trainable parameter count (sum of all projections + W_c). + pub fn param_count(&self) -> usize { + let c = &self.config; + // W_in: in*hidden + hidden_bias + let n_w_in = c.in_dim * c.hidden_dim + c.hidden_dim; + // W_a / W_b: hidden*state + state_bias + let n_w_ab = (c.hidden_dim * c.state_dim + c.state_dim) * 2; + // W_c: hidden*state (no bias) + let n_w_c = c.hidden_dim * c.state_dim; + // W_out: hidden*1 + 1 + let n_w_out = c.hidden_dim + 1; + n_w_in + n_w_ab + n_w_c + n_w_out + } +} + +#[cfg(test)] +mod tests { + use super::*; + use cudarc::driver::CudaContext; + + fn cuda_stream_or_skip() -> Option> { + match CudaContext::new(0) { + Ok(ctx) => Some(ctx.default_stream()), + Err(_) => None, + } + } + + #[test] + fn test_mamba2_config_rejects_state_over_16() { + let cfg = Mamba2BlockConfig { + in_dim: 81, hidden_dim: 64, state_dim: 17, seq_len: 32, + }; + assert!(cfg.validate().is_err(), "state_dim > 16 must be rejected"); + } + + #[test] + fn test_mamba2_config_rejects_zero_dims() { + let cfg = Mamba2BlockConfig { + in_dim: 0, hidden_dim: 64, state_dim: 16, seq_len: 32, + }; + assert!(cfg.validate().is_err()); + } + + #[test] + fn test_mamba2_block_constructs_and_loads_kernels() { + let stream = match cuda_stream_or_skip() { + Some(s) => s, + None => { + eprintln!("CUDA unavailable; skipping GPU construction test"); + return; + } + }; + let cfg = Mamba2BlockConfig { + in_dim: 81, hidden_dim: 64, state_dim: 16, seq_len: 32, + }; + let block = Mamba2Block::new(cfg, stream).expect("Mamba2Block::new"); + // Both kernel handles must be resolved. + let _ = &block.kernel_fwd; + let _ = &block.kernel_bwd; + // Parameter count sanity: 81*64 + 64 + (64*16 + 16)*2 + 64*16 + 64 + 1 + // = 5184 + 64 + 2080 + 1024 + 65 = 8417 + let expected = 81 * 64 + 64 + (64 * 16 + 16) * 2 + 64 * 16 + 64 + 1; + assert_eq!(block.param_count(), expected); + } +}