burn_mamba/mamba2/

cache.rs

//! # Mamba-2 Inference Caches
//!
//! This module defines the state that must be preserved between calls during
//! autoregressive (token-by-token) generation.  During *training* or *prefill*
//! the full sequence is available at once and the chunked SSD algorithm is used
//! (see [`Mamba2::forward`]).  During *decoding* the model
//! processes one token per step and the SSM operates in its pure recurrent
//! form (see [`Mamba2::step`]):
//!
//! ```text
//!   hₜ = Āₜ hₜ₋₁ + B̄ₜ xₜ        (state update)
//!   yₜ = Cₜᵀ hₜ + D xₜ            (output)
//! ```
//!
//! Two pieces of state are required per layer:
//!
//! 1. **Convolution cache** — the last `conv_kernel` inputs to the depthwise
//!    Conv1d, kept so that every decoding step can apply the causal filter
//!    without re-processing previous tokens.
//!
//! 2. **SSM hidden state** — the matrix `hₜ ∈ ℝ^{per_head_dim×state_rank}` (per head), which
//!    compresses the entire past context into a fixed-size representation
//!    regardless of how many tokens have been generated.  This is the key
//!    memory-efficiency advantage of SSMs over attention: the KV-cache of a
//!    Transformer grows as O(sequence·state_rank) with sequence length, whereas the SSM state
//!    is always O(per_head_dim·state_rank).

use crate::mamba2::prelude::*;
use crate::modules::sanity as san;
use burn::module::Module;
use burn::prelude::*;

// ---------------------------------------------------------------------------
// Mamba2Caches  (one cache entry per layer)
// ---------------------------------------------------------------------------

/// A collection of per-layer caches for a complete Mamba-2 network.
///
/// During autoregressive decoding, a [`Mamba2Caches`] instance is threaded
/// through every layer-stack `step` call (the family-generic
/// [`crate::generic::Layers`]).  Each element of `caches` corresponds to one
/// (virtual) layer in the network.
#[derive(Module, Debug)]
pub struct Mamba2Caches {
    /// Per-layer caches.
    ///
    /// Length: `n_real_caches` (the number of *virtual* layers, which may
    /// exceed the number of *real* weight layers when weight-sharing / layer
    /// scheduling is in use).
    pub caches: Vec<Mamba2Cache>,
}

/// Configuration / factory for [`Mamba2Caches`].
#[derive(Config, Debug)]
pub struct Mamba2CachesConfig {
    /// Number of cache slots.  Equals the number of virtual layers in the
    /// network (one cache per layer, even when layers share weights).
    pub n_real_caches: usize,

    /// Shared configuration that determines the shape of each individual
    /// cache tensor.
    pub cache: Mamba2CacheConfig,
}

impl Mamba2CachesConfig {
    /// Convenience constructor that derives cache shapes directly from a
    /// [`Mamba2Config`] block configuration.
    pub fn new_from_block_config(
        n_real_caches: usize,
        batch: usize,
        block_config: Mamba2Config,
    ) -> Self {
        Self {
            n_real_caches,
            cache: Mamba2CacheConfig::new_from_block_config(batch, block_config),
        }
    }

    /// Allocate all cache tensors (zero-initialised) on `device`.
    pub fn init(&self, device: &Device) -> Mamba2Caches {
        let caches = (0..self.n_real_caches)
            .map(|_| self.cache.clone().init(device))
            .collect();
        Mamba2Caches { caches }
    }
}

// ---------------------------------------------------------------------------
// Mamba2Cache  (state for a single layer)
// ---------------------------------------------------------------------------

/// The mutable state carried between decoding steps for a **single** Mamba-2
/// layer.
///
/// Both tensors are updated in-place (via Burn's functional clone) at every
/// call to [`Mamba2::step`].
#[derive(Module, Debug)]
pub struct Mamba2Cache {
    /// **Convolution rolling window.**
    ///
    /// Stores the last `conv_kernel` pre-activation feature vectors fed into
    /// the depthwise Conv1d.  At each step, the oldest column is discarded and
    /// the new token's projection is appended (a left-shift followed by an
    /// insert into the rightmost column), maintaining strict causality.
    ///
    /// Shape: `[batch, conv_dim, conv_kernel]`
    ///   - `conv_dim  = d_inner + 2 · ngroups · state_rank`
    ///   - `conv_kernel` is typically 4
    pub conv_bvk: Tensor<3>,

    /// **SSM hidden state** `hₜ`.
    ///
    /// This is the O(per_head_dim·state_rank) compressed summary of all tokens seen so far.
    /// Updated via `hₜ = Āₜ hₜ₋₁ + B̄ₜ xₜ` at each decoding step.
    ///
    /// The tensor is indexed as `[batch, nheads, per_head_dim, state_rank]`
    /// (i.e. `[batch, nheads, per_head_dim, state_rank]` in the paper's notation), which is the transpose
    /// of the mathematical `hₜ ∈ ℝ^{state_rank×per_head_dim}` but equivalent in content.
    ///
    /// Shape: `[batch, nheads, per_head_dim, state_rank]`
    pub ssm_bhpr: Tensor<4>,
}

impl Mamba2Cache {
    /// Run the [`NaN`/`Inf` guards](crate::utils::sanity) on every cached tensor.
    pub fn sanity(&self) {
        san(&self.conv_bvk);
        san(&self.ssm_bhpr);
    }
}

/// Configuration / factory for a single [`Mamba2Cache`].
#[derive(Config, Debug)]
pub struct Mamba2CacheConfig {
    /// Batch size.
    pub batch: usize,

    /// `state_rank` — the number of latent dimensions in the SSM hidden
    /// state.  Corresponds to `state_rank` in [`Mamba2Config`].
    #[config(default = 128)]
    pub state_rank: usize,

    /// Causal convolution window length.  Corresponds to `conv_kernel` in
    /// [`Mamba2Config`].
    #[config(default = 4)]
    pub conv_kernel: usize,

    /// Number of channels entering (and leaving) the depthwise convolution.
    /// Equal to `d_inner + 2 · ngroups · state_rank`.
    pub conv_dim: usize,

    /// Head dimension `per_head_dim`.  Corresponds to `per_head_dim` in [`Mamba2Config`].
    #[config(default = 64)]
    pub per_head_dim: usize,

    /// Number of SSM heads `nheads`.
    pub nheads: usize,
}

impl Mamba2CacheConfig {
    /// Derive cache shapes from a Mamba-2 block configuration plus a batch
    /// size.
    pub fn new_from_block_config(batch: usize, block_config: Mamba2Config) -> Self {
        Self {
            batch,
            state_rank: block_config.state_rank,
            conv_kernel: block_config.conv_kernel,
            conv_dim: block_config.conv_dim(),
            per_head_dim: block_config.per_head_dim,
            nheads: block_config.nheads(),
        }
    }

    /// Allocate zero-initialised cache tensors on `device`.
    ///
    /// Zero initialisation is correct because:
    /// - The convolution cache represents "no previous tokens" (identity padding).
    /// - The SSM state represents `h₀ = 0` (zero initial condition), which is
    ///   the standard default.  Learnable initial state (if configured) are
    ///   added on top of this inside [`Mamba2::forward`] /
    ///   [`Mamba2::step`].
    pub fn init(&self, device: &Device) -> Mamba2Cache {
        let conv_bvk = Tensor::zeros(
            Shape::new([self.batch, self.conv_dim, self.conv_kernel]),
            device,
        );
        let ssm_bhpr = Tensor::zeros(
            Shape::new([self.batch, self.nheads, self.per_head_dim, self.state_rank]),
            device,
        );
        Mamba2Cache { conv_bvk, ssm_bhpr }
    }
}

impl Mamba2Caches {
    /// Number of per-layer caches.
    pub fn caches_len(&self) -> usize {
        self.caches.len()
    }

    /// Wrap a vector of per-layer caches.
    pub fn from_vec(vec: Vec<Mamba2Cache>) -> Self {
        Self { caches: vec }
    }

    /// Wrap each per-layer cache in `Some` so the layer loop can `take` it
    /// without cloning (Burn tensors are reference-counted).
    pub fn into_options(self) -> Vec<Option<Mamba2Cache>> {
        self.caches.into_iter().map(Some).collect()
    }

    /// Inverse of [`Self::into_options`]: unwrap each slot and re-bundle.
    pub fn from_options(options: Vec<Option<Mamba2Cache>>) -> Self {
        let caches = options.into_iter().map(Option::unwrap).collect();
        Self::from_vec(caches)
    }
}