feat(ml-alpha): from-scratch Mamba2 block — foundation (Phase 1d.1, session 1)

GPU-pure stateful encoder skeleton for the snapshot-stream falsification.
This session lands the build infrastructure + weight allocation + kernel
loading; forward/backward + training loop in follow-up sessions.

- build.rs compiles `../ml/src/cuda_pipeline/mamba2_temporal_kernel.cu` to
  `mamba2_temporal_kernel.cubin` in OUT_DIR (rerun-if-env-changed=CUDA_COMPUTE_CAP
  per the L40S/H100 cubin-staleness pattern). Zero header dependencies → single
  nvcc invocation; no NVRTC.
- `Mamba2Block` holds all parameters on GPU (`OwnedGpuLinear` from ml-core
  for the projection layers, raw `CudaSlice<f32>` for `W_c` which the kernel
  reads directly). Xavier init via ml-core, which uses pinned host buffers
  for the seed transfer.
- Both `mamba2_scan_projected_fwd` and `mamba2_scan_projected_bwd` kernel
  symbols resolve at construction; forward and backward paths in follow-up.
- State dim hardcoded at ≤16 in the kernel; config validation rejects >16.

Tests (3 passing on real GPU):
- Reject state_dim > 16
- Reject zero dims
- Constructs + loads both kernels + correct param count (8417 for 81×64×16×1)

Aligns with project memories:
- feedback_no_nvrtc: pre-compiled cubin via build.rs
- feedback_no_htod_htoh_only_mapped_pinned: pinned via ml-core init helpers
- ml-alpha invariant: no `ml`/`ml-supervised` dep (only the .cu source file)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

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<PathBuf> {
+    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,
+    }
+}

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};

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<f32>`).
+//! - 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<CudaStream>,
+
+    // ── 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<f32>` (no bias) because the kernel reads it directly.
+    pub w_c: CudaSlice<f32>,
+    /// Output projection: `hidden_dim → 1`. Reduces the enriched hidden
+    /// state to the binary direction logit.
+    pub w_out: OwnedGpuLinear,
+
+    // ── Kernel handles ────────────────────────────────────────────────
+    _module: Arc<CudaModule>,
+    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<CudaStream>) -> Result<Self> {
+        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<Arc<CudaStream>> {
+        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);
+    }
+}