mojolearn.TransformerBlockInference#

class mojolearn.TransformerBlockInference(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)#

Forward-only TransformerBlock on the CPU, from weights trained on a GPU.

Takes the same constructor arguments as TransformerBlock (the nine named weights, n_heads, n_kv_heads, head_dim, window) and applies the same weight checks. Supports forward(x, lengths=None) from a zero state, allocate_state(B, max_tokens), forward(x, state) and step(x, state); decoding one token at a time with a carried state gives the same bits as one forward pass over the whole sequence. A block has no output head, so it returns hidden states, not logits. Runs on the host binding shipped with mojolearn, loaded on first use.

Examples

>>> blk = mojolearn.TransformerBlockInference(weights, n_heads=4)
>>> y = blk.forward(x)

See also

User guide: Neural networks

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)

Not available: inference is forward only.

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.
backward(x, grad_output)#

Not available: inference is forward only. Always raises NotImplementedError; train with the GPU block class.

allocate_state(batch_size, max_tokens)#

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

decode_session(state)#

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

forward(x, state=None, *, lengths=None)#

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)

numeric_mode_used()#

The numeric mode this estimator runs in.

Returns:

str – 'identical' or 'fast', read from the compiled binary the estimator uses, not from the argument it was given.

set_params(**params)#

Set constructor parameters. The estimator is rebuilt through its constructor, so every constructor check runs again and any fitted state is dropped, as a new setting invalidates it.

step(x, state)#

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)

vendor_used()#

The GPU backend this estimator runs on.

Returns:

str – 'metal', 'cuda' or 'hip', read from the compiled binary the estimator uses.