foxhunt/crates/ml-alpha/src/mamba2_block.rs

//! 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::cublas::CudaBlas;
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use ml_core::cuda_autograd::gpu_tensor::GpuTensor;
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,
    pub kernel_reduce_d_w_c: CudaFunction,

    /// cuBLAS handle for the four linear projections (input / A / B / output).
    pub cublas: CudaBlas,
}

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_alpha_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_alpha_scan_fwd")
            .map_err(|e| anyhow!("Mamba2Block: forward kernel symbol resolve: {e}"))?;
        let kernel_bwd = module
            .load_function("mamba2_alpha_scan_bwd")
            .map_err(|e| anyhow!("Mamba2Block: backward kernel symbol resolve: {e}"))?;
        let kernel_reduce_d_w_c = module
            .load_function("mamba2_alpha_reduce_d_w_c")
            .map_err(|e| anyhow!("Mamba2Block: d_w_c reduction kernel resolve: {e}"))?;

        // cuBLAS handle for the four GEMM projections.
        let cublas = CudaBlas::new(Arc::clone(&stream))
            .map_err(|e| anyhow!("Mamba2Block: cuBLAS init failed: {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,
            kernel_reduce_d_w_c,
            cublas,
        })
    }

    // ── Forward pass ───────────────────────────────────────────────────

    /// GPU-native forward inference. `input` is a flat `[n_batch × seq_len × in_dim]`
    /// `GpuTensor` (row-major); returns a flat `[n_batch]` tensor of raw logits
    /// (one per sequence, predicted from the final-position state).
    ///
    /// Steps:
    ///   1. `x = input @ Wᵢₙ.T + bᵢₙ`         (cuBLAS sgemm + bias add)
    ///   2. `a = x @ Wₐ.T + bₐ`                 (cuBLAS sgemm + bias add)
    ///   3. `b = x @ Wᵦ.T + bᵦ`                 (cuBLAS sgemm + bias add)
    ///   4. zero-init `h_s2[B, hidden_dim]`, scratch `h_enriched[B, hidden_dim]`
    ///   5. launch `mamba2_alpha_scan_fwd(a, b, w_c, h_s2, h_enriched, B, K, H, S)`
    ///   6. `logit = h_enriched @ Wₒᵤₜ.T + bₒᵤₜ`  (cuBLAS sgemm + bias add)
    ///
    /// All intermediate tensors live on GPU; no host roundtrip. The caller is
    /// responsible for shuttling input/output between host and GPU at the
    /// pipeline boundary (pinned-memory transfers handled by ml-core helpers
    /// elsewhere in the stack).
    pub fn forward(&self, input: &GpuTensor) -> Result<GpuTensor> {
        let c = &self.config;
        let n_batch = match input.shape() {
            [b, k, d] if *k == c.seq_len && *d == c.in_dim => *b,
            shape => {
                return Err(anyhow!(
                    "Mamba2Block::forward: expected input shape [B, {}, {}], got {:?}",
                    c.seq_len, c.in_dim, shape
                ));
            }
        };
        let n_rows = n_batch * c.seq_len; // flattened over time for the per-step projections

        // Reshape input to [N*K, in_dim] (logical reshape — same storage).
        let input_2d = GpuTensor::new(
            input.cuda_data().clone(),
            vec![n_rows, c.in_dim],
        )
        .map_err(|e| anyhow!("reshape input → 2D: {e}"))?;

        // (1) input projection: x = input @ W_in.T + b_in → [N*K, hidden_dim]
        let (x, _act_in) = self.w_in.inner.forward_with_slices(
            &input_2d, &self.w_in.weight, &self.w_in.bias, &self.cublas, &self.stream,
        ).map_err(|e| anyhow!("w_in forward: {e}"))?;

        // (2) a_proj = x @ W_a.T + b_a → [N*K, state_dim]
        let (a_proj, _act_a) = self.w_a.inner.forward_with_slices(
            &x, &self.w_a.weight, &self.w_a.bias, &self.cublas, &self.stream,
        ).map_err(|e| anyhow!("w_a forward: {e}"))?;

        // (3) b_proj = x @ W_b.T + b_b → [N*K, state_dim]
        let (b_proj, _act_b) = self.w_b.inner.forward_with_slices(
            &x, &self.w_b.weight, &self.w_b.bias, &self.cublas, &self.stream,
        ).map_err(|e| anyhow!("w_b forward: {e}"))?;

        // (4) initial hidden h_s2 and output buffer h_enriched, both [N, hidden_dim].
        // h_s2 is the residual the kernel adds to the scan context; for the
        // pure-supervised pipeline we start with zero residual (no carry from
        // a previous chunk).
        let h_s2 = GpuTensor::zeros(&[n_batch, c.hidden_dim], &self.stream)
            .map_err(|e| anyhow!("alloc h_s2: {e}"))?;
        let mut h_enriched = GpuTensor::zeros(&[n_batch, c.hidden_dim], &self.stream)
            .map_err(|e| anyhow!("alloc h_enriched: {e}"))?;

        // (5) launch the scan kernel. Grid: (N, ceil(H/32)), Block: 32.
        let block_threads: u32 = 32;
        let grid_x: u32 = n_batch as u32;
        let grid_y: u32 = ((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
        let cfg = LaunchConfig {
            grid_dim: (grid_x, grid_y, 1),
            block_dim: (block_threads, 1, 1),
            shared_mem_bytes: 0,
        };
        let n_i32 = n_batch as i32;
        let k_i32 = c.seq_len as i32;
        let sh2_i32 = c.hidden_dim as i32;
        let st_i32 = c.state_dim as i32;
        unsafe {
            self.stream
                .launch_builder(&self.kernel_fwd)
                .arg(a_proj.cuda_data())
                .arg(b_proj.cuda_data())
                .arg(&self.w_c)
                .arg(h_s2.cuda_data())
                .arg(h_enriched.data_mut())
                .arg(&n_i32)
                .arg(&k_i32)
                .arg(&sh2_i32)
                .arg(&st_i32)
                .launch(cfg)
                .map_err(|e| anyhow!("mamba2_alpha_scan_fwd launch: {e}"))?;
        }

        // (6) output projection: logit = h_enriched @ W_out.T + b_out → [N, 1]
        let (logit, _act_out) = self.w_out.inner.forward_with_slices(
            &h_enriched, &self.w_out.weight, &self.w_out.bias, &self.cublas, &self.stream,
        ).map_err(|e| anyhow!("w_out forward: {e}"))?;

        Ok(logit)
    }

    /// 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_forward_shape_and_finite() {
        let stream = match cuda_stream_or_skip() {
            Some(s) => s,
            None => {
                eprintln!("CUDA unavailable; skipping forward smoke");
                return;
            }
        };
        let cfg = Mamba2BlockConfig {
            in_dim: 81, hidden_dim: 32, state_dim: 8, seq_len: 16,
        };
        let n_batch = 4;
        let block = Mamba2Block::new(cfg.clone(), Arc::clone(&stream)).expect("Mamba2Block::new");

        // Build deterministic synthetic input [N, K, in_dim] = [4, 16, 81] = 5184 f32.
        let n_input = n_batch * cfg.seq_len * cfg.in_dim;
        let host_input: Vec<f32> = (0..n_input).map(|i| ((i as f32) * 1e-3).sin()).collect();
        let dev_input = stream.clone_htod(&host_input).expect("htod");
        let input_tensor = GpuTensor::new(dev_input, vec![n_batch, cfg.seq_len, cfg.in_dim])
            .expect("input tensor");

        let logit = block.forward(&input_tensor).expect("forward");
        // Output shape must be [N, 1].
        assert_eq!(logit.shape(), &[n_batch, 1]);
        // All logits must be finite (no NaN / Inf — sanity check on the
        // selective scan + projections).
        let host_logit = logit.to_host(&stream).expect("dtoh");
        assert_eq!(host_logit.len(), n_batch);
        for (i, &z) in host_logit.iter().enumerate() {
            assert!(z.is_finite(), "logit[{i}] non-finite: {z}");
        }
    }

    #[test]
    fn test_mamba2_block_forward_rejects_wrong_shape() {
        let stream = match cuda_stream_or_skip() {
            Some(s) => s,
            None => return,
        };
        let cfg = Mamba2BlockConfig {
            in_dim: 81, hidden_dim: 32, state_dim: 8, seq_len: 16,
        };
        let block = Mamba2Block::new(cfg.clone(), Arc::clone(&stream)).expect("init");
        // Wrong in_dim.
        let bad = GpuTensor::zeros(&[2, 16, 99], &stream).expect("bad alloc");
        assert!(block.forward(&bad).is_err());
        // Wrong seq_len.
        let bad2 = GpuTensor::zeros(&[2, 8, 81], &stream).expect("bad alloc 2");
        assert!(block.forward(&bad2).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);
    }
}