mojolearn.TransformerState#

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

__init__(batch_size, n_kv_heads, head_dim, max_tokens, k_cache, v_cache, cached_tokens=0, window=0)[source]#

Methods

keys()

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

values()

The used part of the value cache, in the shape returned by keys() and with the same view/copy rule.

Attributes

capacity

max_tokens, or window under a sliding window.
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.