mojolearn.Mamba3Block#

class mojolearn.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)
__init__(weights)[source]#

Methods

allocate_state(batch_size)

Return a zero Mamba3State 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.

session_info()

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

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