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.legacy.kv_cache.PagedKVLeafRegion(leaf_id, group_id, bytes_per_page, page_size, *, stride_alignment=None)

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)
• stride_alignment (int | None)

blocks_to_reserve()

blocks_to_reserve(num_blocks, *, enable_prefix_caching)

source

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

A sliding-window leaf only needs the pages holding the tokens its window looks back at, plus the page it is writing now. With prefix caching on, it keeps one extra page, so a replay of a prompt that ends exactly on a page boundary can resume one page earlier.

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