IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Python class

MultiKVCacheParams

MultiKVCacheParams​

class max.nn.kv_cache.MultiKVCacheParams(children, page_size, data_parallel_degree, devices, kv_connector, host_kvcache_swap_space_gb, speculative_method=None, num_draft_tokens=0)

source

Bases: KVCacheParamInterface

Aggregates multiple KV cache parameter sets into a recursive tree.

Children may be leaf KVCacheParams instances or nested MultiKVCacheParams subtrees, so arbitrarily deep hierarchies are supported (e.g. {target: {sliding, mla}, draft: mha}). The whole tree is consumed through the KVCacheParamInterface — callers never need to know the depth.

Parameters:

allocate_buffers()​

allocate_buffers(total_num_pages)

source

Allocates per-replica buffers for every cache in the tree.

Returns one MultiKVCacheBuffer per data-parallel replica, each holding that replica’s KVCacheBuffer for every child cache.

Parameters:

total_num_pages (int)

Return type:

list[KVCacheBufferInterface]

build_runtime_inputs()​

build_runtime_inputs(assignments, buffers)

source

Builds the runtime KV-cache tree spanning all replicas.

Each child leaf is built from every replica’s assignment plus that replica’s child buffer; the per-replica assignment (cache lengths / lookup table / dispatch shape) is shared across child caches since they all map the same sequence.

Parameters:

  • assignments (Sequence[KVCacheAssignments])
  • buffers (Sequence[KVCacheBufferInterface])

Return type:

KVCacheInputsInterface[Buffer, Buffer]

bytes_per_block​

property bytes_per_block: int

source

Total bytes per block across all KV caches.

Since all caches allocate memory for the same sequence, the total memory cost per block is the sum across all param sets.

children​

children: dict[str, KVCacheParamInterface]

source

KV cache parameter sets to aggregate. Values may be leaf KVCacheParams or nested MultiKVCacheParams trees.

data_parallel_degree​

data_parallel_degree: int

source

devices​

devices: Sequence[DeviceRef]

source

enable_dp_cross_replica_prefix_copy​

property enable_dp_cross_replica_prefix_copy: bool

source

Whether DP cross-replica prefix copies are enabled (shared across all caches).

enable_prefix_caching​

property enable_prefix_caching: bool

source

Whether prefix caching is enabled (shared across all caches).

from_params()​

classmethod from_params(params)

source

Creates a MultiKVCacheParams from one or more param sets.

Children may be leaf KVCacheParams instances or nested MultiKVCacheParams trees, enabling arbitrarily deep KV cache hierarchies (e.g. {target: {sliding, mla}, draft: mha}). All children must share the same page_size, data_parallel_degree, n_devices, kv_connector, and host_kvcache_swap_space_gb values.

Parameters:

params (Mapping[str, KVCacheParamInterface]) – Named mapping of KVCacheParamInterface instances to aggregate.

Returns:

A new MultiKVCacheParams aggregating all provided params.

Raises:

ValueError – If no params are provided.

Return type:

MultiKVCacheParams

get_symbolic_inputs()​

get_symbolic_inputs(namespace='')

source

Returns the symbolic inputs for the KV cache tree.

Each child inherits a distinct namespace so sibling groups’ page-pool dims stay independent; nested subtrees compose the prefix.

Parameters:

namespace (str)

Return type:

MultiKVCacheInputs[TensorType, BufferType]

graph_capture_probe_cache_lengths()​

graph_capture_probe_cache_lengths(max_cache_length, q_max_seq_len=1)

source

Returns the union of probe cache lengths across all child caches.

Parameters:

  • max_cache_length (int)
  • q_max_seq_len (int)

Return type:

list[int]

host_kvcache_swap_space_gb​

host_kvcache_swap_space_gb: float | None

source

kv_connector​

kv_connector: KVConnectorType | None

source

kv_connector_config​

property kv_connector_config: Any

source

Connector config (shared across all caches).

kv_hash_algo​

property kv_hash_algo: Literal['ahash64', 'sha256', 'sha256_64']

source

Hash algorithm used for KV-cache block identity.

kv_hash_seed​

property kv_hash_seed: bytes | None

source

Resolved 32-byte cluster seed for sha256/sha256_64. None for ahash64.

n_devices​

property n_devices: int

source

Returns the number of devices.

num_draft_tokens​

num_draft_tokens: int = 0

source

page_size​

page_size: int

source

replicates_kv_across_tp​

property replicates_kv_across_tp: bool

source

Whether every device holds identical KV state.

resolve_attn_key()​

resolve_attn_key(batch_size, max_prompt_length, max_cache_valid_length)

source

Resolves the dispatch shape tree mirroring the cache tree.

Parameters:

  • batch_size (int)
  • max_prompt_length (int)
  • max_cache_valid_length (int)

Return type:

AttnKeyInterface

speculative_method​

speculative_method: Literal['eagle', 'mtp', 'dflash'] | None = None

source

tensor_parallel_degree​

property tensor_parallel_degree: int

source

Returns the tensor parallel degree.

unflatten_basic_kv_tree()​

unflatten_basic_kv_tree(it)

source

Unflattens a basic KV tree from a graph-input iterator.

Requires that the model is a basic height-1 tree. This method does not work on nested trees.

Parameters:

it (Iterator[Any])

Return type:

tuple[list[KVCacheInputsPerDevice[TensorValue, BufferValue]], …]

unflatten_kv_inputs()​

unflatten_kv_inputs(it)

source

Unflattens the KV cache inputs from a graph-input iterator.

Parameters:

it (Iterator[Any])

Return type:

MultiKVCacheInputs[TensorValue, BufferValue]

Models

Cookbook

Puzzles

Blog

Community

Contact