mojolearn.TransformerBlock#

class mojolearn.TransformerBlock(weights, *, n_heads, n_kv_heads=None, head_dim=None, window=0, rope_theta=10000.0, rope_scaling=None, rope_dim=None, max_positions=8192, qkv_bias=False, o_bias=False, norm='rmsnorm', norm_eps=1e-06, norm_bias=False, mlp='swiglu', mlp_bias=False, qk_norm=False, attn_softcap=None, numeric_profile=None)[source]#

Bases: NumericModeMixin

One Llama-shaped transformer decoder block, float32.

Computes Hugging Face’s LlamaDecoderLayer.forward: input RMSNorm, causal self-attention with RoPE, grouped-query attention and a KV cache, residual, post-attention RMSNorm, SiLU-gated MLP, residual. Options cover common Llama-family variants (LayerNorm, Gemma’s 1 + w RMSNorm, GELU/GeGLU MLPs, Qwen2 q/k/v biases, Qwen3 q/k norms, Gemma2 attention softcap, sliding-window attention, linear and Llama 3 RoPE scaling, partial rotary). Attention is the eager formulation; there is no FlashAttention, SDPA, paged attention, attention dropout or mask other than causal / sliding window.

There is no initializer: you pass the weights, as a dict keyed by Hugging Face parameter names. With QW = n_heads * head_dim (= d_model), KW = n_kv_heads * head_dim and IT = intermediate size (read from gate_proj.weight), the required weights are, in torch’s [out, in] layout:

`input_layernorm.weight`	(d_model,)
`post_attention_layernorm.weight`	(d_model,)
`q_proj.weight`	(QW, d_model)
`k_proj.weight`	(KW, d_model)
`v_proj.weight`	(KW, d_model)
`o_proj.weight`	(d_model, QW)
`gate_proj.weight`	(IT, d_model)
`up_proj.weight`	(IT, d_model)
`down_proj.weight`	(d_model, IT)

Ungated MLPs (mlp='gelu' or 'gelu_tanh') have no gate_proj.weight. Optional tensors, when their option is on: q_proj.bias, k_proj.bias, v_proj.bias (qkv_bias), o_proj.bias (o_bias), input_layernorm.bias and post_attention_layernorm.bias (norm_bias), up_proj.bias, down_proj.bias and, when gated, gate_proj.bias (mlp_bias), q_norm.weight and k_norm.weight of shape (head_dim,) (qk_norm). Weights are converted to float32 once, and each shape is checked.

Shapes must satisfy d_model == n_heads * head_dim, n_heads % n_kv_heads == 0 and an even head_dim; absolute positions must stay below max_positions. Inputs must be float32: bfloat16, float16 and float64 raise errors, and non-finite weights, inputs or cache values raise an error before any computation.

The KV cache is explicit (TransformerState): allocate_state makes a zero cache, and forward and step update it. In the identical mode, decoding token by token gives the same bits as one prefill over the whole sequence.

Parameters:

weights (dict) – float32 buffers (NumPy arrays, array.array, memoryviews, …) keyed by the parameter names above.
n_heads (int) – Number of attention heads. Required, keyword-only.
n_kv_heads (int or None, default None) – Number of key/value heads (grouped-query attention); None means n_heads.
head_dim (int or None, default None) – Head dimension; None means d_model // n_heads.
window (int, default 0) – 0 is full causal attention; W > 0 is sliding-window attention, where position p attends to positions [max(0, p - W + 1), p] and the KV cache is a ring of W slots.
rope_theta (float, default 10000.0) – RoPE base.
rope_scaling (dict or None, default None) – None, {"type": "linear", "factor": f} or {"type": "llama3", "factor": ..., "low_freq_factor": ..., "high_freq_factor": ..., "original_max_position_embeddings": ...}.
rope_dim (int or None, default None) – Number of rotary dimensions (partial rotary); None means head_dim.
max_positions (int, default 8192) – Ceiling on absolute positions.
qkv_bias (bool, default False) – Use q/k/v projection biases (Qwen2).
o_bias (bool, default False) – Use an output projection bias.
norm ({‘rmsnorm’, ‘layernorm’, ‘rmsnorm_offset’}, default ‘rmsnorm’) – Normalization; ‘rmsnorm_offset’ scales by 1 + w (Gemma).
norm_eps (float, default 1e-6) – Normalization epsilon.
norm_bias (bool, default False) – Use LayerNorm biases (with norm='layernorm').
mlp ({‘swiglu’, ‘gelu’, ‘gelu_tanh’, ‘geglu’, ‘geglu_tanh’}, default ‘swiglu’) – MLP type.
mlp_bias (bool, default False) – Use MLP biases.
qk_norm (bool, default False) – RMS-normalize queries and keys over head_dim before RoPE (Qwen3).
attn_softcap (float or None, default None) – Tanh softcap of the attention scores (Gemma2).
numeric_profile ({‘fp32_v1’, ‘fixed15_v1’} or None, default None) – Arithmetic profile. None and ‘fp32_v1’ are float32. ‘fixed15_v1’ quantizes the projection weights once to 15-bit fixed point and is inference only.
numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Examples

>>> import numpy as np, mojolearn
>>> rng = np.random.default_rng(0)
>>> d, it = 64, 128
>>> shapes = {"input_layernorm.weight": (d,), "post_attention_layernorm.weight": (d,),
...           "q_proj.weight": (d, d), "k_proj.weight": (d, d), "v_proj.weight": (d, d),
...           "o_proj.weight": (d, d), "gate_proj.weight": (it, d), "up_proj.weight": (it, d),
...           "down_proj.weight": (d, it)}
>>> weights = {k: rng.standard_normal(s, dtype=np.float32) * 0.02 for k, s in shapes.items()}
>>> blk = mojolearn.TransformerBlock(weights, n_heads=4)
>>> x = rng.standard_normal((2, 16, d), dtype=np.float32)
>>> y = np.asarray(blk.forward(x))

__init__(weights, *, n_heads, n_kv_heads=None, head_dim=None, window=0, rope_theta=10000.0, rope_scaling=None, rope_dim=None, max_positions=8192, qkv_bias=False, o_bias=False, norm='rmsnorm', norm_eps=1e-06, norm_bias=False, mlp='swiglu', mlp_bias=False, qk_norm=False, attn_softcap=None, numeric_profile=None)[source]#

Methods

`allocate_state`(batch_size, max_tokens)	A zero `TransformerState` for `batch_size` sequences of up to `max_tokens` positions each.
`backward`(x, grad_output)	Gradients of a zero-state prefill with respect to the input and the weights.
`decode_session`(state)	Open a `TransformerDecodeSession` on this block and `state`.
`forward`(x[, state, lengths])	Run the block on `(B, L, d_model)` float32 input and return the output, same shape.
`get_params`([deep])
`numeric_mode_used`()	The numeric mode this estimator runs in.
`set_params`(**params)	Set constructor parameters.
`step`(x, state)	Decode one token per sequence, continuing the cache in `state`.
`vendor_used`()	The GPU backend this estimator runs on.

allocate_state(batch_size, max_tokens)[source]#

A zero TransformerState for batch_size sequences of up to max_tokens positions each.

max_tokens counts the prefill plus every later decode step; the capacity is fixed at allocation. Positions beyond max_positions are refused at the first call that reaches them. Under a sliding window the buffers are a ring of window slots and max_tokens only bounds positions.

Parameters:

batch_size (int) – Number of sequences.
max_tokens (int) – Maximum total positions per sequence.

Returns:

TransformerState

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

Run the block on (B, L, d_model) float32 input and return the output, same shape.

With state=None it runs a self-contained prefill from a zero cache and discards the final cache. Pass a TransformerState to carry the cache: it is read at entry and updated, so a later forward or step on that state continues the sequence. In the identical mode the result is bit-for-bit the same as one prefill over the whole sequence.

Parameters:

x (buffer of float32, shape (B, L, d_model)) – Input activations.
state (TransformerState or None, default None) – KV cache to read and update.
lengths (sequence of int or None, default None) – Makes the batch ragged: B lengths 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 padding positions output exactly +0.0. Not allowed together with state.

Returns:

mojolearn.Array of float32, shape (B, L, d_model)

step(x, state)[source]#

Decode one token per sequence, continuing the cache in state.

The same computation as forward with L = 1, at position state.cached_tokens.

Parameters:

x (buffer of float32, shape (B, 1, d_model) or (B, d_model)) – The new token’s activations.
state (TransformerState) – Required KV cache, updated in place.

Returns:

mojolearn.Array of float32, shape (B, 1, d_model)

backward(x, grad_output)[source]#

Gradients of a zero-state prefill with respect to the input and the weights.

Recomputes the forward from a zero cache at positions [0, L) (honoring window) and backpropagates grad_output. Each weight gradient is summed over the B * L tokens. Identical mode and the default block options only; other settings raise errors. There is no gradient for an incoming cache.

Parameters:

x (buffer of float32, shape (B, L, d_model)) – Input activations.
grad_output (buffer of float32, shape (B, L, d_model)) – Gradient of the loss with respect to the block output.

Returns:

dict – float32 gradients keyed "x" plus the nine weight names, each in its tensor’s shape.

decode_session(state)[source]#

Open a TransformerDecodeSession on this block and state.

On a GPU, the session keeps the weights, the KV cache and the decode buffers on the device between steps, so each token costs one input upload and one output download. Its outputs are byte for byte those of step and forward on the same block and state.

The session copies the weights and the cache when it opens. Until close() (or sync_state()), the state’s k_cache and v_cache are stale, and forward or step on this block refuse that state; cached_tokens stays current. Later edits to the weights are not seen: close and open a new session. Later edits to the cache buffers are not seen either: call load_state().

Parameters:: state (TransformerState) – The KV cache to decode from.
Returns:: TransformerDecodeSession