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

MultiKVCacheParams

MultiKVCacheParams

class max.nn.kv_cache.MultiKVCacheParams(children, page_size, data_parallel_degree, devices, kv_connector_config, 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 num_blocks.

Parameters:

• children (dict[str, KVCacheParamInterface])
• page_size (int)
• data_parallel_degree (int)
• devices (Sequence[DeviceRef])
• kv_connector_config (KVConnectorConfigInterface)
• speculative_method (Literal['eagle', 'mtp', 'dflash', 'dflash2'] | None)
• num_draft_tokens (int)

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, _prefix='')

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])
• _prefix (str)

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

Degree of data parallelism, a value every child cache must share.

devices

devices: Sequence[DeviceRef]

source

Devices to use for the KV caches.

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, and kv_connector_config 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]

kv_connector_config

kv_connector_config: KVConnectorConfigInterface

source

The KV connector’s type and settings, a value every child must share.

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.

leaves()

leaves(_prefix='')

source

Returns the leaves of the KV cache.

Parameters:

_prefix (str)

Return type:

Mapping[str, KVLeafRegion]

n_devices

property n_devices: int

source

Returns the number of devices.

num_draft_tokens

num_draft_tokens: int = 0

source

Total draft tokens generated per speculative iteration.

page_size

page_size: int

source

Number of tokens per page, a value every child cache must share.

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

slab_to_buffer_views()

slab_to_buffer_views(buffers)

source

Converts a slab of memory into a buffer view.

Parameters:

buffers (Sequence[Buffer])

Return type:

KVCacheBufferInterface

speculative_method

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

source

Speculative decoding method propagated from SpeculativeConfig.

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]