mojolearn.Mamba2Block#

class mojolearn.Mamba2Block(weights, *, dt_limit=(0.0, inf))[source]#

Bases: _MambaBase

One Mamba-2 (SSD) block: RMSNorm, mixer and residual.

Matches Hugging Face’s Mamba2Block.forward (the non-memory-efficient mixer). Takes (B, L, d_model) float32 and returns (B, L, d_model) float32, with the recurrent state explicit and caller-owned (Mamba2State).

Constructor keyword numeric_mode selects ‘identical’ (the default) or ‘fast’. Inputs, weights and state must be float32: bfloat16, float16 and float64 are refused by name rather than converted, and non-finite values are refused by name. Arrays are read through the buffer protocol, so NumPy is not required; arrays the block allocates are mojolearn.Array, zero-copy under numpy.asarray. There is no initializer: weights are passed in, so a cross-check can give both sides the same values. Projection weights may also be packed with mojolearn.lowbit; they are expanded exactly to float32 (see weight_format).

d_model must be a multiple of 32. Fixed by the block: d_state 128, head dimension 64, one group, chunk size 256, expand 2, RMSNorm on with the gate applied before the norm. Decoding resumes the open chunk, so step and a forward with a carried state continue a prefill bit for bit.

After every call, h_last_ holds the SSM state after the final padded chunk. This is a reporting value, not the resumption state; feeding it back as initial_states matches resumption only on a chunk boundary.

Parameters:

  • weights (dict of str to float32 buffer) – With d_inner = 2 * d_model, H = d_inner / 64 heads and CD = d_inner + 256:

    • block_norm.weight (d_model,), the block norm

    • in_proj.weight (2 * d_inner + 256 + H, d_model), column order z | xBC | dt

    • conv1d.weight (CD, 1, 4) or (CD, 4)

    • conv1d.bias (CD,)

    • dt_bias (H,)

    • A_log (H,), one value per head

    • D (H,)

    • norm.weight (d_inner,), the gated norm

    • out_proj.weight (d_model, d_inner)

  • dt_limit (tuple of float, default (0.0, inf)) – Clamp (lo, hi) applied to dt.

Variables:

  • d_model (int) – Model dimension, read from block_norm.weight.

  • weight_format (str) – Format the weights were given in.

Examples

>>> block = mojolearn.Mamba2Block(weights, numeric_mode="identical")
>>> y = block.forward(x)
>>> grads = block.backward(x, dy)
>>> state = block.allocate_state(batch_size=x.shape[0])
>>> y0 = block.step(x[:, :1], state)
__init__(weights, *, dt_limit=(0.0, inf))[source]#

Methods

allocate_state(batch_size)

Return a zero Mamba2State for batch_size sequences.

backward(x, grad_output)

Gradients of a zero-state prefill (a vector-Jacobian product).

forward(x[, state, lengths])

Run the block over a sequence.

get_params([deep])

numeric_mode_used()

The numeric mode this estimator runs in.

set_params(**params)

Set constructor parameters.

step(x, state)

Run one decode token.

vendor_used()

The GPU backend this estimator runs on.
allocate_state(batch_size)[source]#

Return a zero Mamba2State for batch_size sequences.

Set its h nonzero before the first call to supply initial_states.

Parameters:

batch_size (int) – Number of sequences B.

Returns:

Mamba2State

forward(x, state=None, *, lengths=None)[source]#

Run the block over a sequence.

Parameters:

  • x (float32 buffer, shape (B, L, d_model)) – Input activations: a NumPy array, array.array('f'), mojolearn.Array or any float32 buffer. Other dtypes raise TypeError.

  • state (Mamba2State or None, default None) – None runs a self-contained prefill from a zero state and discards the final state. Pass a state from allocate_state to carry it: it is read at entry and updated in place, so a later forward or step continues the sequence exactly as if the whole sequence had run in one call.

  • lengths (sequence of int or None, default None) – Makes the batch ragged: B integers in [1, L]. Row i is real at positions [0, lengths[i]) and padding after. Each real position’s output equals the row run alone at its own length, and every padding position’s output is exactly +0.0. Not allowed together with state.

Returns:

y (mojolearn.Array, shape (B, L, d_model)) – Block output, residual included. numpy.asarray(y) is zero-copy.

backward(x, grad_output)[source]#

Gradients of a zero-state prefill (a vector-Jacobian product).

Recomputes the forward with the current weights on every call; no earlier forward or cached state is used or changed. Runs under numeric_mode ‘identical’ or ‘fast’; other modes raise NotImplementedError. Gradients with respect to a carried state, decode backward and final-state cotangents are not available.

Parameters:

  • x (float32 buffer, shape (B, L, d_model)) – Forward input.

  • grad_output (float32 buffer, shape (B, L, d_model)) – Gradient of the loss with respect to the block output. Must be finite.

Returns:

grads (dict of str to mojolearn.Array) – 'x' plus one entry per constructor weight, under the same names.

step(x, state)[source]#

Run one decode token.

The same computation as forward at L = 1 with the state carried, so decoding token by token gives the same bits as prefilling the whole sequence.

Parameters:

  • x (float32 buffer, shape (B, 1, d_model) or (B, d_model)) – One token per row.

  • state (Mamba2State) – Required; updated in place.

Returns:

y (mojolearn.Array, shape (B, 1, d_model))