Enum Mamba2SsdPath 

pub enum Mamba2SsdPath {
    Minimal(Option<usize>),
    Serial(Option<usize>),
    SerialRecalculated(Option<usize>),
}

Algorithm selection for the Mamba-2 chunkwise SSD.

Each variant carries an optional chunk length. Larger values increase the intra-chunk GEMM work and reduce the inter-chunk scan length; the optimal value is approximately √(state_rank · per_head_dim) (see Self::optimal_chunk_len). None falls back to that optimal value.

Variants

Minimal(Option<usize>)

Minimal SSD: mostly batched matmuls; backward via autodiff.

See Mamba2SsdInput::ssd_minimal. For training, prefer Self::SerialRecalculated.

Based on /mamba_ssm/modules/ssd_minimal.py from the state-spaces/mamba github reference.

Serial(Option<usize>)

(Hybrid) serial SSD: a serial loop over the chunks plus batched matmuls; backward via autodiff.

See Mamba2SsdInput::ssd_serial. For a memory-saving custom backward, see Self::SerialRecalculated.

Based on 5 kernels under /mamba_ssm/ops/triton/ from the state-spaces/mamba github reference:

• ssd_chunk_state.py (K1, K3).
• ssd_bmm.py (K2).
• ssd_state_passing.py (K4).
• ssd_chunk_scan.py (K5).

SerialRecalculated(Option<usize>)

(Hybrid) serial SSD with a custom, memory-efficient backward that recomputes the forward intermediates instead of storing them.

See Mamba2SsdInput::ssd_serial_recalculated. For a plain autodiff backward, see Self::Serial.

Based on the combined kernel /mamba_ssm/ops/triton/ssd_combined.py from the state-spaces/mamba github reference.

Implementations

impl Mamba2SsdPath

pub fn optimal_chunk_len(state_rank: usize, per_head_dim: usize) -> usize

Optimal chunk length, approximately √(state_rank · per_head_dim), rounded up to a multiple of 32 and capped at 512.

pub fn chunk_len(&self) -> Option<usize>

The chunk length carried by this variant, if any.

pub fn chunk_len_or_optimal( &self, state_rank: usize, per_head_dim: usize, ) -> usize

The chunk length carried by this variant, or Self::optimal_chunk_len when unset.

pub fn default_optimal_from_block<B: Backend>(block: &Mamba2) -> Self

The recommended default path for a given block: Self::SerialRecalculated with Self::optimal_chunk_len for the block’s dimensions.

Trait Implementations

impl Clone for Mamba2SsdPath

fn clone(&self) -> Mamba2SsdPath

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for Mamba2SsdPath

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for Mamba2SsdPath

fn default() -> Mamba2SsdPath

Returns the “default value” for a type.