std1.0.0-beta-7 SafeC docs

std :: ml::ssm

Struct std::Mamba

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

struct Mamba {
    &FloatTensor inProj;
    &FloatTensor convW;
    &FloatTensor convB;
    &FloatTensor xProj;
    &FloatTensor dtProj;
    &FloatTensor dtBias;
    &FloatTensor Alog;
    &FloatTensor Dskip;
    &FloatTensor outProj;
    unsigned long dModel;
    unsigned long dInner;
    unsigned long dState;
    unsigned long dConv;
    unsigned long dtRank;
}

One Mamba block (the mixer of Mamba / Mamba-1 language models, without the surrounding RMSNorm and residual add):

(x, z)  = split(X @ inProj)                     dInner each
x       = silu(causal_depthwise_conv1d(x))      kernel dConv
(dt, B, C) = split(x @ xProj)                   dtRank, dState, dState
delta   = softplus(dt @ dtProj + dtBias)
y       = selective_scan(x, delta, -exp(Alog), B, C, Dskip) * silu(z)
out     = y @ outProj

Weights are stored [in, out] like the rest of std::ml (transpose a PyTorch nn.Linear's weight to load it); convW is [dInner, dConv], the PyTorch Conv1d weight [dInner, 1, dConv] with the middle axis dropped.

Fields

§&FloatTensor inProjsource

[dModel, 2 * dInner]

§&FloatTensor convWsource

[dInner, dConv]

§&FloatTensor convBsource

[dInner]

§&FloatTensor xProjsource

[dInner, dtRank + 2 * dState]

§&FloatTensor dtProjsource

[dtRank, dInner]

§&FloatTensor dtBiassource

[dInner]

§&FloatTensor Alogsource

[dInner, dState]; A = -exp(Alog)

§&FloatTensor Dskipsource

[dInner]

§&FloatTensor outProjsource

[dInner, dModel]

§unsigned long dModelsource
§unsigned long dInnersource

expand * dModel

§unsigned long dStatesource

N

§unsigned long dConvsource
§unsigned long dtRanksource

ceil(dModel / 16)

Methods

§&FloatTensor forward(const &FloatTensor X) constsource · definition

X: [T, dModel] -> [T, dModel], the whole sequence at once (prefill).

§void collect_params(&Vec<&FloatTensor> out) constsource · definition

Parameterized (module.h): every weight tensor.

§void collect_state(&Vec<&FloatTensor> out) constsource · definition
§void drop()source · definition
§void free() consumesource · definition

= drop()

§&FloatTensor step(&MambaState state, const &FloatTensor x) constsource · definition

x: [1, dModel] (or [dModel]) -> [1, dModel]; advances 'state'.

§static struct Mamba new( unsigned long dModel, unsigned long dState, unsigned long dConv, unsigned long expand, )source · definition

Deterministic initialization following the paper (and HF transformers' MambaMixer): Alog[d][n] = log(n + 1), Dskip = 1, dtBias chosen so softplus(dtBias) spans [0.001, 0.1] log-uniformly, other weights small and pseudo-random. Typical: dState 16, dConv 4, expand 2.