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

KVLeafRegion

KVLeafRegion

class max.nn.kv_cache.KVLeafRegion(leaf_id, group_id, bytes_per_page, *, stride_alignment=None)

source

Bases: object

One addressable region of the cache pool.

Parameters:

• leaf_id (str)
• group_id (KVCacheGroupId)
• bytes_per_page (int)
• stride_alignment (int | None)

blocks_to_reserve()

blocks_to_reserve(num_blocks, *, enable_prefix_caching)

source

Returns how many blocks one request draws to fill num_blocks slots.

Fewer than num_blocks when some slots hold the null block: a sliding window keeps only its span, and a state keeps only its live block and one checkpoint.

Parameters:

• num_blocks (int)
• enable_prefix_caching (bool)

Return type:

int

bound_row_copies()

bound_row_copies(src, dst)

source

Returns the rows to copy between, per bound input.

Empty for a leaf the graph reaches through a per-forward table, whose blocks are never overwritten in place.

Parameters:

• src (int) – The block to copy.
• dst (int) – The block to copy it to.

Return type:

Mapping[str, tuple[range, range]]

bound_row_span()

bound_row_span(block)

source

Returns the rows one block occupies, per bound input.

Empty for the same leaves bound_row_copies() is empty for.

Parameters:

block (int) – The block whose rows to name.

Return type:

Mapping[str, range]

bytes_per_page

bytes_per_page: int

source

cacheable

property cacheable: bool

source

Whether this leaf’s blocks are addressed by content.

False makes the leaf invisible to everything keyed on a hash: prefix hits, commits, and both directions of an external tier. It still draws from the pool and still counts against admission.

group_id

group_id: KVCacheGroupId

source

leaf_id

leaf_id: str

source

staged_input_shapes()

staged_input_shapes(batch_size, num_blocks)

source

Returns the shape and dtype of each input a forward stages.

Called once at the widest bounds a forward can reach, to allocate the tensors the graph binds, and again per forward to size host staging.

Parameters:

• batch_size (int) – Requests the forward runs.
• num_blocks (int) – Blocks the deepest row of the batch reaches.

Return type:

Mapping[str, tuple[tuple[int, …], DType]]

stride_alignment

stride_alignment: int | None = None

source

The bytes a padded page stride of this leaf must be a multiple of, or None if some reader needs packed pages.

Setting it promises every kernel reading the leaf takes a page stride, so the pool may pad the page to tile a smaller huge block. A leaf read through PagedKVCache must declare a multiple of its row.

write_staged_inputs()

write_staged_inputs(plans, into)

source

Writes one forward’s blocks into the arrays declared for them.

Parameters:

• plans (Sequence[Sequence[int]]) – The blocks each row touches, in batch-row order.
• into (Mapping[str, ndarray]) – The array to fill, per staged key.

Return type:

None