oxidizr 0.1.0-beta1

// Mamba-1 (State Space Model) Implementation
// Ported from candle-examples/mamba-minimal with critical fixes for training stability
// Key: Use sequential scan (not parallel) to avoid numerical divergence in backward pass

use candle_core::{Result, Tensor, DType};
use candle_nn::{VarBuilder, Linear, linear_no_bias, Conv1d, Conv1dConfig, Module};
use crate::config::ModelConfig;

/// Maximum absolute value for hidden state clamping in Mamba SSM
///
/// Empirically chosen to prevent NaN in long sequences (seq_len=256+)
/// without being overly restrictive. Larger values (±50) caused OOM,
/// smaller values may be too restrictive for model expressiveness.
///
/// # Rationale
/// Over long sequences, the recurrent hidden state accumulates numerical errors
/// exponentially. Even with well-conditioned state transitions, floating-point
/// errors compound. This clamp prevents overflow while maintaining model capacity.
///
/// # Tested Values
/// - ±30: Works without NaN for seq_len up to 512, minimal memory overhead
/// - ±50: Also prevents NaN but caused OOM with batch_size=5, seq_len=128
/// - ±10: Untested (may be too restrictive)
///
/// # TODO
/// Systematically benchmark ±10, ±20, ±30, ±40, ±50 for quality vs stability tradeoffs
const HIDDEN_STATE_CLAMP_MAX: f32 = 30.0;

pub struct MambaBlock {
    in_proj: Linear,
    conv1d: Conv1d,
    x_proj: Linear,
    dt_proj: Linear,
    a_log: Tensor,  // Log of state matrix A
    d: Tensor,      // Skip connection parameter
    out_proj: Linear,

    #[allow(dead_code)]  // Stored for potential debugging/introspection
    d_inner: usize,
    d_state: usize,
    #[allow(dead_code)]  // Stored for potential debugging/introspection
    d_conv: usize,
}

impl MambaBlock {
    /// Create a new Mamba-1 block
    ///
    /// # Arguments
    /// * `cfg` - Model configuration containing hidden_size and layer parameters
    /// * `vb` - VarBuilder for parameter initialization
    ///
    /// # Returns
    /// Initialized MambaBlock with learned parameters (A, D) and projections
    ///
    /// # Architecture
    /// - Expansion factor: 2x (d_inner = 2 × hidden_size)
    /// - State dimension: N=16 (standard Mamba-1)
    /// - Conv kernel size: 4 (depthwise convolution)
    ///
    /// # Learned Parameters
    /// - `a_log`: Log of state transition matrix A [d_inner, d_state]
    /// - `d`: Skip connection parameter [d_inner]
    /// - Various projection weights (in_proj, x_proj, dt_proj, out_proj)
    pub fn new(cfg: &ModelConfig, vb: VarBuilder) -> Result<Self> {
        let d_model = cfg.hidden_size;
        let d_inner = d_model * 2;  // Expansion factor
        let d_state = 16;  // Mamba-1 uses N=16
        let d_conv = 4;    // Conv1D kernel size

        // Input projection: hidden → 2*hidden
        let in_proj = linear_no_bias(d_model, d_inner * 2, vb.pp("in_proj"))?;

        // Conv1D for local context
        let conv1d = candle_nn::conv1d(
            d_inner,
            d_inner,
            d_conv,
            Conv1dConfig {
                groups: d_inner,  // Depthwise convolution
                padding: d_conv - 1,
                ..Default::default()
            },
            vb.pp("conv1d")
        )?;

        // SSM projections
        let x_proj = linear_no_bias(d_inner, d_state + d_state + d_inner, vb.pp("x_proj"))?;
        let dt_proj = candle_nn::linear(d_inner, d_inner, vb.pp("dt_proj"))?;

        // State matrix A: initialized as log(1..d_state)
        let a_log = vb.get((d_inner, d_state), "a_log")?;

        // Skip connection parameter D
        let d = vb.get(d_inner, "d")?;

        // Output projection: 2*hidden → hidden
        let out_proj = linear_no_bias(d_inner, d_model, vb.pp("out_proj"))?;

        Ok(Self {
            in_proj,
            conv1d,
            x_proj,
            dt_proj,
            a_log,
            d,
            out_proj,
            d_inner,
            d_state,
            d_conv,
        })
    }

    /// Forward pass through the Mamba block
    ///
    /// # Arguments
    /// * `x` - Input tensor of shape [batch_size, seq_len, d_model]
    ///
    /// # Returns
    /// Output tensor of shape [batch_size, seq_len, d_model]
    ///
    /// # Numerical Stability
    /// Uses sequential scan with hidden state clamping (±30) to prevent NaN
    /// in long sequences (seq_len > 200). See `selective_scan_sequential`
    /// for implementation details.
    ///
    /// # Errors
    /// Returns error if tensor operations fail or shapes are incompatible
    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let (_b_sz, seq_len, _d_model) = x.dims3()?;

        // 1. Input projection and split
        let x_and_res = self.in_proj.forward(x)?;  // [B, Seq, 2*d_inner]
        let d_inner_2 = x_and_res.dim(2)? / 2;
        let x = x_and_res.narrow(2, 0, d_inner_2)?.contiguous()?;  // [B, Seq, d_inner]
        let res = x_and_res.narrow(2, d_inner_2, d_inner_2)?.contiguous()?;  // [B, Seq, d_inner]

        // 2. Conv1D for local context
        let x_t = x.transpose(1, 2)?.contiguous()?;  // [B, d_inner, Seq]
        let x_conv = self.conv1d.forward(&x_t)?;
        let x_conv = x_conv.narrow(2, 0, seq_len)?;  // Remove padding
        let x_conv = x_conv.transpose(1, 2)?.contiguous()?;  // [B, Seq, d_inner]

        // 3. Activation
        let x_conv = candle_nn::ops::silu(&x_conv)?;

        // 4. SSM: Selective scan
        let y = self.ssm(&x_conv)?;

        // 5. Gating with residual branch
        let y = (y * candle_nn::ops::silu(&res)?)?;

        // 6. Output projection
        self.out_proj.forward(&y)
    }

    fn ssm(&self, x: &Tensor) -> Result<Tensor> {
        let (_b_sz, _seq_len, d_inner) = x.dims3()?;

        // 1. Project x to get B, C, and delta
        let x_dbl = self.x_proj.forward(x)?.contiguous()?;  // [B, Seq, d_state + d_state + d_inner]

        // Split into delta (time step), B (input proj), C (output proj)
        // CRITICAL: Must call contiguous() after narrow to fix CUDA memory padding
        let delta = x_dbl.narrow(2, 0, d_inner)?.contiguous()?;  // [B, Seq, d_inner]
        let b = x_dbl.narrow(2, d_inner, self.d_state)?.contiguous()?;  // [B, Seq, d_state]
        let c = x_dbl.narrow(2, d_inner + self.d_state, self.d_state)?.contiguous()?;  // [B, Seq, d_state]

        // 2. Transform delta (time step size)
        let delta = self.dt_proj.forward(&delta)?;  // [B, Seq, d_inner]
        // Softplus to ensure positive: softplus(x) = log(1 + exp(x))
        // NOTE: No clamping here - we rely on hidden state clamping downstream to prevent NaN
        // This minimizes memory overhead from intermediate tensors in the gradient graph
        let delta = (delta.exp()? + 1.0)?.log()?;  // Ensure positive

        // 3. Get state matrix A (state transition matrix)
        // NOTE: No clamping on a_log - learned parameters stay reasonable during training
        // Clamping here would create extra tensors and increase memory usage
        let a = self.a_log.neg()?.exp()?;  // [d_inner, d_state]

        // 4. Run selective scan (CRITICAL: Use sequential for training stability)
        let y = self.selective_scan_sequential(x, &delta, &a, &b, &c)?;

        // 5. Skip connection
        let y = (y + x.broadcast_mul(&self.d)?)?;

        Ok(y)
    }

    /// CRITICAL: Sequential scan for numerical stability in long sequences
    ///
    /// # Why Sequential vs Parallel
    /// - Parallel associative scan: Fast but causes NaN gradients during backprop
    /// - Sequential scan: Slower but numerically stable for training
    ///
    /// # Numerical Stability Strategy
    /// This function uses MINIMAL clamping for memory efficiency:
    /// - ONLY clamps hidden state h after each update (see `HIDDEN_STATE_CLAMP_MAX`)
    /// - Does NOT clamp: delta, a_log, delta_a, or intermediate computations
    /// - Rationale: Multiple clamps caused OOM, single clamp is sufficient
    ///
    /// # Tested Performance
    /// - seq_len=128, batch_size=5: Works without NaN, ~10GB VRAM
    /// - seq_len=256, batch_size=1: Works without NaN, ~5GB VRAM
    /// - Without hidden state clamp: Immediate NaN at seq_len=256
    ///
    /// # Parameters
    /// - u: Input sequence [B, Seq, d_inner]
    /// - delta: Time step sizes [B, Seq, d_inner]
    /// - a: State transition matrix [d_inner, d_state]
    /// - b: Input projection [B, Seq, d_state]
    /// - c: Output projection [B, Seq, d_state]
    fn selective_scan_sequential(
        &self,
        u: &Tensor,      // Input [B, Seq, d_inner]
        delta: &Tensor,  // Time step [B, Seq, d_inner]
        a: &Tensor,      // State matrix [d_inner, d_state]
        b: &Tensor,      // Input projection [B, Seq, d_state]
        c: &Tensor,      // Output projection [B, Seq, d_state]
    ) -> Result<Tensor> {
        let (b_sz, seq_len, d_inner) = u.dims3()?;
        let d_state = self.d_state;

        // Ensure f32 for recurrent state (even if model is bf16)
        let u = u.to_dtype(DType::F32)?;
        let delta = delta.to_dtype(DType::F32)?;
        let a = a.to_dtype(DType::F32)?;
        let b = b.to_dtype(DType::F32)?;
        let c = c.to_dtype(DType::F32)?;

        let mut outputs = Vec::new();

        // Process each batch independently
        for batch_idx in 0..b_sz {
            // Initialize hidden state for this batch
            let mut h = Tensor::zeros((d_inner, d_state), DType::F32, u.device())?;

            let mut batch_outputs = Vec::new();

            // Sequential scan over time
            for t in 0..seq_len {
                // Get inputs at time t - squeeze to remove batch and seq dims
                let u_t = u.narrow(0, batch_idx, 1)?.narrow(1, t, 1)?.squeeze(0)?.squeeze(0)?;
                let delta_t = delta.narrow(0, batch_idx, 1)?.narrow(1, t, 1)?.squeeze(0)?.squeeze(0)?;
                let b_t = b.narrow(0, batch_idx, 1)?.narrow(1, t, 1)?.squeeze(0)?.squeeze(0)?;
                let c_t = c.narrow(0, batch_idx, 1)?.narrow(1, t, 1)?.squeeze(0)?.squeeze(0)?;

                // Compute delta * A (broadcast delta across d_state)
                let delta_a = (delta_t.unsqueeze(1)?.broadcast_as((d_inner, d_state))? * &a)?;
                // NOTE: No clamping before exp() - saves memory by avoiding extra tensors
                // The hidden state clamp after state update is sufficient to prevent overflow
                let delta_a = delta_a.exp()?;

                // Compute delta * B * u (outer product-like operation)
                let delta_b_u = (delta_t.unsqueeze(1)? * u_t.unsqueeze(1)?)?
                    .broadcast_mul(&b_t.unsqueeze(0)?)?;

                // Update hidden state: h[t] = exp(delta*A) * h[t-1] + delta*B*u
                h = ((delta_a * &h)? + delta_b_u)?;

                // CRITICAL: Clamp hidden state to prevent NaN in long sequences
                // WHY THIS IS ESSENTIAL:
                // - Over long sequences (e.g., seq_len=256), the recurrent state accumulates errors
                // - Even with exp(delta*A) < 1, numerical errors compound exponentially
                // - Without this clamp: NaN appears immediately at seq_len=256
                // - With this clamp: Stable training up to seq_len=256+
                //
                // MEMORY OPTIMIZATION:
                // - This is the ONLY clamp in the entire Mamba forward pass
                // - Adding clamps on delta, a_log, delta_a caused OOM with batch_size=5, seq_len=128
                // - Single clamp here: same memory usage as no clamping, prevents NaN
                //
                // See HIDDEN_STATE_CLAMP_MAX constant for value rationale and TODO for systematic benchmarking
                h = h.clamp(-HIDDEN_STATE_CLAMP_MAX as f64, HIDDEN_STATE_CLAMP_MAX as f64)?;

                // Output: y[t] = h[t] · c
                let y_t = h.matmul(&c_t.unsqueeze(1)?)?;  // [d_inner, 1]
                let y_t = y_t.reshape((1, d_inner))?;

                batch_outputs.push(y_t);
            }

            // Stack outputs for this batch
            let batch_out = Tensor::cat(&batch_outputs, 0)?;  // [Seq, d_inner]
            outputs.push(batch_out.unsqueeze(0)?);
        }

        // Stack all batches
        let output = Tensor::cat(&outputs, 0)?;  // [B, Seq, d_inner]

        Ok(output)
    }
}