mojolearn.mamba#

Mamba-1, Mamba-2 and Mamba-3 blocks in float32.

The block and state classes are also exported from mojolearn. Each block takes caller-supplied weights and provides forward(x, state=None), step(x, state), allocate_state(...) and backward(x, grad_output). The recurrent state is explicit and owned by the caller, so prefill, continuation and token-by-token decode give the same bits. Forward runs under numeric_mode ‘identical’ (the default) or ‘fast’. backward returns the gradients of a zero-state prefill with respect to x and every constructor weight (11 arrays for Mamba-1, 10 each for Mamba-2 and Mamba-3), recomputing the forward internally; gradients through a carried state or a decode step are not available. Inputs are read through the buffer protocol, so PyTorch and NumPy are not required.

Mamba1DecodeSession (from Mamba1Block.decode_session) keeps a Mamba-1 block and its state resident on the device for fast decoding.

Examples

With float32 weights, x and dy of the right shapes:

from mojolearn.mamba import Mamba2Block
block = Mamba2Block(weights, numeric_mode="identical")
y = block.forward(x)
gradients = block.backward(x, dy)
state = block.allocate_state(batch_size=x.shape[0])
y0 = block.step(x[:, :1], state)

Classes

Mamba1Block(weights)

One Mamba-1 block: RMSNorm, selective-scan mixer and residual.

Mamba1DecodeSession(block, state)

Token-by-token decoding on one Mamba1Block and one Mamba1State.

Mamba1State(conv_window, h, *[, profile_id, ...])

Recurrent state of a Mamba1Block, owned by the caller.

Mamba2Block(weights, *[, dt_limit])

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

Mamba2State(conv_window, h, buffer_xbc, ...)

Recurrent state of a Mamba2Block, owned by the caller.

Mamba3Block(weights)

One Mamba-3 (SISO) block: RMSNorm, mixer and residual.

Mamba3State(theta, h, buffer_qrot, ...[, ...])

Recurrent state of a Mamba3Block, owned by the caller.
class mojolearn.mamba.Mamba1Block(weights)[source]#

Bases: _MambaBase

One Mamba-1 block: RMSNorm, selective-scan mixer and residual.

Matches Hugging Face’s MambaBlock.forward. Takes (B, L, d_model) float32 and returns (B, L, d_model) float32, with the recurrent state explicit and caller-owned (Mamba1State).

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).

Fixed by the block: d_state 16, d_conv 4, expand 2, dt_rank ceil(d_model / 16), convolution bias on, projection bias off, SiLU activation.

Parameters:

weights (dict of str to float32 buffer) – Weights under the reference parameter names, with d_inner = 2 * d_model and dt_rank = ceil(d_model / 16):

  • norm.weight (d_model,)

  • in_proj.weight (2 * d_inner, d_model)

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

  • conv1d.bias (d_inner,)

  • x_proj.weight (dt_rank + 32, d_inner)

  • dt_proj.weight (d_inner, dt_rank)

  • dt_proj.bias (d_inner,)

  • A_log (d_inner, 16)

  • D (d_inner,)

  • out_proj.weight (d_model, d_inner)

Variables:

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

  • d_inner (int) – Inner dimension, 2 * d_model.

  • dt_rank (int) – Rank of the dt projection.

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

Examples

>>> import numpy as np, mojolearn
>>> block = mojolearn.Mamba1Block(weights)   # dict of float32 arrays, as above
>>> x = np.random.default_rng(0).random((2, 16, block.d_model), dtype=np.float32)
>>> y = block.forward(x)
>>> state = block.allocate_state(2)
>>> y = block.forward(x, state)
>>> y_next = block.step(x[:, -1], state)
allocate_state(batch_size)[source]#

Return a zero Mamba1State for batch_size sequences.

Parameters:

batch_size (int) – Number of sequences B.

Returns:

Mamba1State

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 (Mamba1State 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 (Mamba1State) – Required; updated in place.

Returns:

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

decode_session(state)[source]#

Open a Mamba1DecodeSession for fast token-by-token decoding.

On a GPU the session keeps the weights and the state on the device across steps; its step gives the same bytes as Mamba1Block.step. The session copies the weights and the state when it opens:

  • until close() or sync_state(), the state’s arrays are stale, and forward/step on this block refuse that state by name;

  • weight edits after opening are not seen; close and open a new session;

  • edits to the state arrays are not seen until load_state().

Parameters:

state (Mamba1State) – State to decode from.

Returns:

Mamba1DecodeSession

class mojolearn.mamba.Mamba1DecodeSession(block, state)[source]#

Bases: object

Token-by-token decoding on one Mamba1Block and one Mamba1State.

Made by Mamba1Block.decode_session(state). Every output is byte for byte what Mamba1Block.step gives on the same block and state. On a GPU the weights, the state and the decode buffers stay on the device across steps; on a CPU-only install the session holds its own copies and steps on the host. A session is not thread-safe. Use it as a context manager to close it automatically.

Variables:

  • state (Mamba1State) – The state this session decodes from.

  • is_open (bool) – Whether the session is open.

Examples

>>> with block.decode_session(state) as session:
...     for x_t in tokens:
...         y_t = session.step(x_t)
property state#

The Mamba1State this session decodes from.

property is_open#

Whether the session is open.

step(x)[source]#

Decode one token per row.

Parameters:

x (float32 buffer, shape (B, 1, d_model) or (B, d_model)) – B must equal the session’s batch size.

Returns:

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

sync_state()[source]#

Copy the session’s current state into the Mamba1State arrays and return the state.

load_state()[source]#

Reload the session from the Mamba1State arrays (after you edit them) and return the state.

close()[source]#

Copy the state back, release the session’s buffers and close it. Closing twice is a no-op.

class mojolearn.mamba.Mamba1State(conv_window, h, *, profile_id=1, absolute_position=0, affine_boundary=None, affine_a=None, affine_b=None)[source]#

Bases: object

Recurrent state of a Mamba1Block, owned by the caller.

Make one with Mamba1Block.allocate_state; it is all zeros before the first token. forward and step update the arrays in place. The arrays are mojolearn.Array when allocate_state made them, or any writable float32 buffer (a NumPy array included) you supply; their bytes can be saved and restored exactly.

Variables:

  • conv_window (float32 buffer, shape (B, d_inner, 4)) – The last four inputs to the convolution, oldest first.

  • h (float32 buffer, shape (B, d_inner, 16)) – SSM state.

class mojolearn.mamba.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)
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))

class mojolearn.mamba.Mamba2State(conv_window, h, buffer_xbc, buffer_dtraw, buffered_tokens=0)[source]#

Bases: object

Recurrent state of a Mamba2Block, owned by the caller.

With H = d_model / 32 heads and CD = 2 * d_model + 256 convolution channels. Make one with Mamba2Block.allocate_state; it is all zeros (and buffered_tokens == 0) before the first token. forward and step update it in place. Setting h nonzero on a fresh state supplies the reference’s initial_states.

Variables:

  • conv_window (float32 buffer, shape (B, CD, 4)) – The last four pre-convolution xBC inputs, oldest first.

  • h (float32 buffer, shape (B, H, 64, 128)) – SSM state as of the last completed 256-token chunk.

  • buffer_xbc (float32 buffer, shape (B, 256, CD)) – Post-convolution, post-SiLU xBC rows of the open chunk.

  • buffer_dtraw (float32 buffer, shape (B, 256, H)) – Raw dt rows of the open chunk.

  • buffered_tokens (int) – Number of valid rows in the open-chunk buffers, in [0, 256). Decoding resumes the open chunk from these rows.

class mojolearn.mamba.Mamba3Block(weights)[source]#

Bases: _MambaBase

One Mamba-3 (SISO) block: RMSNorm, mixer and residual.

Matches mamba_ssm’s Mamba3.forward (SISO) inside the non-fused block. Takes (B, L, d_model) float32 and returns (B, L, d_model) float32, with the recurrent state explicit and caller-owned (Mamba3State).

Compared with Mamba2Block: no convolution (v is the raw x split), no dt clamp (so no dt_limit), A depends on the token and head with a floor clamp, the gate is applied inside the core (no gated output norm), and B and C have their own RMSNorms with per-head biases added after the norm.

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, rope fraction 0.5 (32 angles), A floor 1e-4, chunk size 64, expand 2. MIMO, output-projection norm, fused pre-gate options, more than one group and variable-length batches are not available. Decoding resumes the working chunk, so step and a forward with a carried state continue a prefill bit for bit.

After every call, h_last_, k_last_, v_last_ and theta_last_ hold reporting values, not the resumption state.

Parameters:

weights (dict of str to float32 buffer) – With d_inner = 2 * d_model, H = d_inner / 64 heads, N = 128 and d_in_proj = 2 * d_inner + 256 + 3 * H + 32:

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

  • in_proj.weight (d_in_proj, d_model), column order z | x | B | C | dd_dt | dd_A | trap | angle

  • dt_bias (H,)

  • B_norm.weight (N,), RMSNorm with eps 1e-5

  • C_norm.weight (N,)

  • B_bias (H, N)

  • C_bias (H, N)

  • D (H,)

  • out_proj.weight (d_model, d_inner)

Variables:

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

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

Examples

>>> block = mojolearn.Mamba3Block(weights)
>>> y = block.forward(x)
>>> state = block.allocate_state(x.shape[0])
>>> y = block.forward(x, state)
>>> y_next = block.step(x[:, -1], state)
allocate_state(batch_size)[source]#

Return a zero Mamba3State for batch_size sequences.

For a reference Input_States continuation, call Mamba3State.set_input_states on the fresh state.

Parameters:

batch_size (int) – Number of sequences B.

Returns:

Mamba3State

session_info()[source]#

Counters of the prefill session as a dict, or None before the first call.

Reports weight uploads, re-copies and reuses, backward stage reuses, forward recomputes and the number of stages held.

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 (Mamba3State 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 (Mamba3State) – Required; updated in place.

Returns:

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

class mojolearn.mamba.Mamba3State(theta, h, buffer_qrot, buffer_krot, buffer_v, buffer_dt, buffer_sig, buffer_adt, pending_k, pending_v, buffered_tokens=0, pending=False)[source]#

Bases: object

Recurrent state of a Mamba3Block, owned by the caller.

With H = d_model / 32 heads, N = 128 (QK head dimension), P = 64 (V head dimension), R = 32 rotary angles and Q = 64 (the Mamba-3 chunk size). Make one with Mamba3Block.allocate_state; it is all zeros (buffered_tokens == 0, pending == False) before the first token. forward and step update it in place.

Variables:

  • theta (float32 buffer, shape (B, H, R)) – Rotary angles, each in [0, 2*pi).

  • h (float32 buffer, shape (B, H, P, N)) – SSM state entering the last working chunk.

  • buffer_krot (buffer_qrot,) – Rotated, unscaled q and k rows of the working chunk.

  • buffer_v (float32 buffer, shape (B, Q, H, P)) – Raw v rows.

  • buffer_adt (buffer_dt, buffer_sig,) – dt, sigmoid(trap) and A*dt rows.

  • pending_k (float32 buffer, shape (B, H, N)) – Input-state k, used only when pending.

  • pending_v (float32 buffer, shape (B, H, P)) – Input-state v, used only when pending.

  • buffered_tokens (int) – Valid buffer rows, in [0, Q]. After the first call it is always in [1, Q].

  • pending (bool) – True when set_input_states supplied the state; consumed by the next call.

set_input_states(theta, h, k, v)[source]#

Start from a reference Input_States continuation.

Copies the given values into theta, h, pending_k and pending_v and marks them pending for the next call. The state must be fresh (no tokens yet); that is checked at the next call. Arrays must already be float32 with the shapes above; nothing is converted.

Parameters:

  • theta (float32 buffer, shape (B, H, R))

  • h (float32 buffer, shape (B, H, P, N))

  • k (float32 buffer, shape (B, H, N))

  • v (float32 buffer, shape (B, H, P))