For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

RecurrentStateRegion

RecurrentStateRegion

class max.nn.kv_cache.RecurrentStateRegion(leaf_id, num_layers, row_shape, dtype, scratch=False)

source

Bases: object

Shape and dtype of one kind of recurrent state, for one pool leaf.

Parameters:

• leaf_id (str)
• num_layers (int)
• row_shape (tuple[int, ...])
• dtype (DType)
• scratch (bool)

bytes_per_page

property bytes_per_page: int

source

Bytes one page of this leaf holds on one device.

dtype

dtype: DType

source

leaf_id

leaf_id: str

source

num_layers

num_layers: int

source

pool_key

property pool_key: str

source

Key the leaf’s flat pool view is staged under.

row_elements

property row_elements: int

source

Elements in one layer’s state.

row_shape

row_shape: tuple[int, ...]

source

Shape of one layer’s state, per device.

rows_dim

property rows_dim: str

source

Symbolic dim naming this leaf’s row count.

rows_of()

rows_of(page)

source

Returns the rows a page’s layers occupy, layer l at index l.

Parameters:

page (int)

Return type:

range

scratch

scratch: bool = False

source

Whether the region is per-request scratch rather than a checkpoint.

A scratch region is drawn once per request and never published, so it is invisible to prefix hits, to eviction order and to an external tier, and its block does not rotate at a page boundary. Its rows are addressed exactly like a published state’s, so a kernel reads it the same way.