Module ml::ssm

#include <std/ml/ssm.h> source

state space models (Mamba's selective SSM).

A state space model mixes a sequence through a per-channel linear recurrence instead of attention: each of D channels carries an N-value hidden state h that every token updates and reads, h_t = exp(delta_t * A) * h_{t-1} + delta_t * B_t * x_t y_t = C_t . h_t + Dskip * x_t so a sequence costs O(T * D * N) time and decoding a token costs O(D * N) with a fixed-size state, however long the context — no KV cache. Mamba (Gu & Dao 2023, "S6") makes the step size delta and the B / C projections functions of the input ("selective"), which lets the model choose per token what to remember and what to forget.

Forward-only (inference), like the other layers in attention_advanced.h. The CPU scan runs the D channels in parallel on std/sync/parallel.h's pool.

Structs

Mamba	One Mamba block (the mixer of Mamba / Mamba-1 language models, without the surrounding RMSNorm and residual add):
MambaState	Recurrent decoding, one token at a time with a fixed-size state. Feeding tokens 0..T-1 through Mamba::step gives the same rows as Mamba::forward on the whole sequence.

Functions

ssm_selective_scan The selective scan on its own. x, delta: [T, D] (delta already positive, e.g. after softplus); A: [D, N] (negative for a decaying state); B, C: [T, N]; Dskip: [D] (the skip connection, D in the paper). Returns y: [T, D]. The state starts at zero.