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

PagedKVLeafRegion

PagedKVLeafRegion

class max.nn.kv_cache.PagedKVLeafRegion(leaf_id, group_id, bytes_per_page, page_size, *, row_bytes=1)

source

Bases: KVLeafRegion

A leaf the graph reaches through a per-forward page table.

Parameters:

• leaf_id (str)
• group_id (KVCacheGroupId)
• bytes_per_page (int)
• page_size (int)
• row_bytes (int)

blocks_to_reserve()

blocks_to_reserve(num_blocks, *, enable_prefix_caching)

source

Returns num_blocks, capped by the window when the leaf has one.

Parameters:

• num_blocks (int)
• enable_prefix_caching (bool)

Return type:

int

page_size

page_size: int

source

Number of tokens per page.

staged_input_shapes()

staged_input_shapes(batch_size, num_blocks)

source

Returns this leaf’s lookup table.

Parameters:

• batch_size (int)
• num_blocks (int)

Return type:

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

write_staged_inputs()

write_staged_inputs(plans, into)

source

Writes each row’s block ids, zeroing the slots it does not reach.

Parameters:

• plans (Sequence[Sequence[int]])
• into (Mapping[str, ndarray])

Return type:

None