mojolearn.transformer#

Llama-style transformer decoder block.

TransformerBlock is one Llama-shaped decoder layer (input RMSNorm, causal self-attention with RoPE, grouped-query attention and a KV cache, residual, post-attention RMSNorm, SiLU-gated MLP, residual), float32 in and out. It is a building block, not an estimator: it has no fit. You pass in the weights, and the KV cache is an explicit, caller-owned TransformerState. Prefill, continuation from a carried cache and single-token decode all run through the same computation. TransformerDecodeSession keeps the weights and cache on the device across decode steps.

NumPy is not required: inputs are read through the buffer protocol and outputs are mojolearn.Array objects, which numpy.asarray wraps without a copy. All three names are also exported from mojolearn.

Examples

>>> from mojolearn.transformer import TransformerBlock
>>> blk = TransformerBlock(weights, n_heads=2)      # dict of float32 buffers
>>> y = blk.forward(x)                              # (B, L, d_model) -> same
>>> st = blk.allocate_state(batch_size=1, max_tokens=64)
>>> y = blk.forward(x, st)                          # prefill, cache carried
>>> y_t = blk.step(x_t, st)                         # one decode token

Classes

`TransformerBlock`(weights, *, n_heads[, ...])	One Llama-shaped transformer decoder block, float32.
`TransformerDecodeSession`(block, state)	Decode session on one `TransformerBlock` and one `TransformerState`.
`TransformerState`(batch_size, n_kv_heads, ...)	KV cache of a `TransformerBlock`, owned by the caller.

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

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

class mojolearn.transformer.TransformerDecodeSession(block, state)[source]#

Bases: object

Decode session on one TransformerBlock and one TransformerState.

Made by TransformerBlock.decode_session(state); see there for ownership rules. On a GPU the weights and KV cache stay on the device across steps. Every output is byte for byte the block’s own step or forward on the same state. A session is not thread-safe and refuses re-entrant use. It is a context manager: leaving the with block calls close().

Examples

>>> st = blk.allocate_state(batch_size=1, max_tokens=128)
>>> with blk.decode_session(st) as sess:
...     y = sess.forward(prompt)        # (1, L, d_model)
...     y_t = sess.step(x_t)            # (1, 1, d_model)

property state#: The TransformerState this session decodes.

property is_open#: Whether the session is still open.

step(x)[source]#

Decode one token per row.

x is (B, 1, d_model) or (B, d_model) float32; returns (B, 1, d_model) and advances state.cached_tokens. Raises ValueError if the session is closed or B differs from the state’s batch size.

forward(x)[source]#

Run L tokens per row on the session’s cache (a prefill or a chunked continuation).

x is (B, L, d_model) float32; returns the same shape. Raises ValueError if the session is closed or B differs from the state’s batch size.

sync_state()[source]#: Copy the session’s cache into the state’s buffers; the session stays open and keeps owning the state.

load_state()[source]#: Reload the session’s cache from the state’s buffers and cached_tokens, after you edit them.

close()[source]#

Copy the cache back into the state, release the session’s buffers and hand the state back.

Calling it on a closed session does nothing.

discard()[source]#

Close without copying the cache back into the state.

Use it when you throw the state away, or after a failed call.

class mojolearn.transformer.TransformerState(batch_size, n_kv_heads, head_dim, max_tokens, k_cache, v_cache, cached_tokens=0, window=0)[source]#

Bases: object

KV cache of a TransformerBlock, owned by the caller.

Made by TransformerBlock.allocate_state. With B = batch_size, KV = n_kv_heads, SMAX = max_tokens and HD = head_dim it holds:

k_cache: flat float32 buffer of B * KV * SMAX * HD values. The first B * KV * cached_tokens * HD values are the [B, KV, cached_tokens, HD] key cache; the rest is spare capacity. It is flat because the packing stride is the used length, not SMAX.
v_cache: the value cache, same shape and packing.
cached_tokens: number of positions the cache holds, in [0, max_tokens]. Slot j is absolute position j; the next call continues at position cached_tokens.

A new state is all zeros. forward and step read both buffers and write them back whole, and update cached_tokens. The buffers are plain float32 data, so you can serialize them any way you like and they round-trip exactly.

With a sliding window (TransformerBlock(window=W), W > 0) each buffer is a ring of W slots per (batch, kv head): B * KV * W * HD floats, slot = position % W, and cached_tokens keeps counting past W.

Parameters:

batch_size (int) – Number of sequences.
n_kv_heads (int) – Number of key/value heads.
head_dim (int) – Head dimension.
max_tokens (int) – Maximum total positions per sequence (prefill plus decode).
k_cache, v_cache (buffer of float32) – Key and value cache buffers.
cached_tokens (int, default 0) – Positions already held.
window (int, default 0) – Sliding-window size; 0 means full causal attention.

property capacity#

max_tokens, or window under a sliding window.

Type:: Slots per (batch, kv head) in each buffer

keys()[source]#

The used part of the key cache as a (B, n_kv_heads, cached_tokens, head_dim) array.

Without a window this is a view into k_cache, valid until the next call updates the state. With a window it is a copy of the min(cached_tokens, window) most recent positions, in ascending position order.

values()[source]#: The used part of the value cache, in the shape returned by keys() and with the same view/copy rule.