mamba-rs 0.7.3

/// Scan mode for GPU SSM recurrence.
///
/// Controls whether the SSM recurrence uses sequential O(T) scan
/// or parallel prefix scan O(T) work / O(log T) depth.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ScanMode {
    /// Sequential O(T) scan — zero overhead, optimal for short sequences.
    Sequential,
    /// Parallel prefix scan — warp shuffle + shared memory, optimal for long T.
    Parallel,
    /// Auto-select: Sequential for T <= 64, Parallel for T > 64
    /// (PARALLEL_SCAN_THRESHOLD, measured on RTX 6000 Ada).
    #[default]
    Auto,
}

impl ScanMode {
    /// Auto threshold: sequential at or below, parallel above. Single source
    /// of truth for the GPU dispatch (re-exported as
    /// `mamba_ssm::gpu::forward::PARALLEL_SCAN_THRESHOLD`).
    pub const PARALLEL_SCAN_THRESHOLD: usize = 64;

    /// Resolve the effective scan choice for a sequence length and state dim.
    ///
    /// `d_state > 64` ALWAYS takes the parallel path regardless of the
    /// requested mode. Historically the sequential kernels were hard-capped
    /// at 64; they are capacity-sized and guarded now, but the route stays —
    /// it is part of the numeric identity of existing runs, and the
    /// parallel scan is the faster path at large state sizes anyway.
    pub fn use_parallel(self, seq_len: usize, d_state: usize) -> bool {
        if d_state > 64 {
            return true;
        }
        match self {
            Self::Sequential => false,
            Self::Parallel => true,
            Self::Auto => seq_len > Self::PARALLEL_SCAN_THRESHOLD,
        }
    }

    /// Resolve Auto to a concrete mode based on sequence length.
    /// Resolve `Auto` for a shape. Delegates to [`Self::use_parallel`] so the
    /// `d_state > 64` sequential override can never diverge between the two
    /// entry points (this used to duplicate the
    /// threshold WITHOUT the d_state override - a safety-divergent copy).
    pub fn resolve(self, seq_len: usize, d_state: usize) -> Self {
        if self.use_parallel(seq_len, d_state) {
            Self::Parallel
        } else {
            Self::Sequential
        }
    }
}

/// Configuration for a Mamba SSM model.
///
/// Source: Gu & Dao (2023) "Mamba: Linear-Time Sequence Modeling with Selective State Spaces"
/// Paper: <https://arxiv.org/abs/2312.00752>
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct MambaConfig {
    /// Model dimension — input features projected to this size.
    pub d_model: usize,

    /// SSM state dimension. Controls memory capacity per channel.
    /// Paper default: 16.
    pub d_state: usize,

    /// Local convolution width before SSM. Paper default: 4.
    pub d_conv: usize,

    /// Expansion factor. `d_inner = expand * d_model`.
    /// Paper default: 2.
    pub expand: usize,

    /// Number of stacked Mamba layers. Paper default: varies by model size.
    pub n_layers: usize,

    /// GPU SSM scan mode. Default: Auto (see [`ScanMode`]).
    /// Only affects GPU training forward/backward. CPU always uses sequential.
    pub scan_mode: ScanMode,

    /// RMSNorm epsilon for every layer norm and the final norm_f.
    /// Default 1e-5 (state-spaces/mamba). FalconMamba checkpoints use 1e-6 —
    /// the HF loader plumbs the checkpoint's value through here.
    pub rms_norm_eps: f32,
}

impl MambaConfig {
    /// Inner dimension (expanded). Used for in_proj, SSM, conv1d.
    pub fn d_inner(&self) -> usize {
        self.expand * self.d_model
    }

    /// dt_rank = ceil(d_model / 16). Controls delta projection bottleneck.
    pub fn dt_rank(&self) -> usize {
        self.d_model.div_ceil(16)
    }

    /// x_proj output dimension: dt_rank + 2 * d_state (for delta, B, C).
    pub fn xdbl_dim(&self) -> usize {
        self.dt_rank() + 2 * self.d_state
    }
}

impl MambaConfig {
    /// Validate configuration constraints.
    ///
    /// Returns `Err` if any dimension would cause kernel failures
    /// (e.g., d_inner not divisible by 4 for vectorized CUDA kernels).
    pub fn validate(&self) -> Result<(), String> {
        if self.d_model == 0 {
            return Err("d_model must be > 0".into());
        }
        if self.d_state == 0 {
            return Err("d_state must be > 0".into());
        }
        if self.d_conv == 0 {
            return Err("d_conv must be > 0".into());
        }
        if self.expand == 0 {
            return Err("expand must be > 0".into());
        }
        if self.n_layers == 0 {
            return Err("n_layers must be > 0".into());
        }
        // Upper bounds. Configs can come from untrusted sources (safetensors
        // metadata, HF config.json): keep dimension products from wrapping
        // and keep grid.y (= d_inner for the CUDA parallel scan) under the
        // hardware limit of 65535.
        if self.d_model > (1 << 20) {
            return Err(format!("d_model ({}) must be <= 2^20", self.d_model));
        }
        if self.n_layers > 4096 {
            return Err(format!("n_layers ({}) must be <= 4096", self.n_layers));
        }
        let d_inner = self
            .expand
            .checked_mul(self.d_model)
            .ok_or("expand * d_model overflows usize")?;
        if d_inner > 65535 {
            return Err(format!(
                "d_inner ({d_inner}) must be <= 65535 (CUDA grid.y limit for the parallel scan)"
            ));
        }
        // 256 is the reference implementations' own maximum. The
        // sequential SSM kernels size their per-thread state arrays at
        // JIT time from the model config (past ~128 the compiler spills
        // to local memory — correct, slower); the training dispatch
        // still routes d_state > 64 to the parallel scan as its faster
        // numeric route.
        if self.d_state > 256 {
            return Err(format!(
                "d_state ({}) must be <= 256 (CUDA parallel scan MAX_DSTATE limit)",
                self.d_state
            ));
        }
        // CUDA conv1d kernel unrolls d_conv — max 8
        if self.d_conv > 8 {
            return Err(format!(
                "d_conv ({}) must be <= 8 (CUDA kernel limit)",
                self.d_conv
            ));
        }
        // CUDA float4 kernels require d_inner divisible by 4
        if !self.d_inner().is_multiple_of(4) {
            return Err(format!(
                "d_inner ({}) must be divisible by 4 (d_model={} * expand={})",
                self.d_inner(),
                self.d_model,
                self.expand
            ));
        }
        Ok(())
    }
}

/// Paper defaults: d_model=128, d_state=16, d_conv=4, expand=2, n_layers=3.
impl Default for MambaConfig {
    fn default() -> Self {
        Self {
            d_model: 128,
            d_state: 16,
            d_conv: 4,
            expand: 2,
            n_layers: 3,
            scan_mode: ScanMode::Auto,
            rms_norm_eps: crate::ops::fast_math::RMS_NORM_EPS,
        }
    }
}