IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /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. /get-started.md).

Python class

MultiKVCacheParams

MultiKVCacheParams​

class max.nn.kv_cache.MultiKVCacheParams(children, page_size, data_parallel_degree, devices, kv_connector_config, enable_prefix_caching=False, enable_dp_cross_replica_prefix_copy=True, kv_hash_algo='ahash64', kv_hash_seed=None, speculative_method=None, num_draft_tokens=0)

source

Bases: KVCacheParamInterface

Aggregates multiple cache parameter sets into a recursive tree.

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

A RecurrentStateParams is a child like any other, but answers fewer questions, so attention-only aggregates run over _attention_children.

Parameters:

allocate_buffers()​

allocate_buffers(total_num_pages, _prefix='')

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 that allocates one.

Parameters:

  • total_num_pages (int)
  • _prefix (str)

Return type:

list[KVCacheBufferInterface]

build_runtime_inputs()​

build_runtime_inputs(assignments, buffers, _prefix='', *, staging=None)

source

Builds the runtime KV-cache tree spanning all replicas.

Each child builds itself from every replica’s assignment plus that replica’s child buffer, if it allocated one; the per-replica assignment (cache lengths / lookup table / dispatch shape / state rows) is shared across child caches since they all map the same sequence. The tree comes out in the order the graph declared its inputs.

Parameters:

Return type:

KVCacheInputs[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.

child()​

child(key, kind)

source

Returns the child at key as the kind it must be.

Raises:

  • KeyError – If the tree has no child at key.
  • TypeError – If the child at key is not a kind.

Parameters:

Return type:

_Child

children​

children: dict[str, CacheLeafParamInterface]

source

Cache parameter sets to aggregate. Values may be leaf KVCacheParams or RecurrentStateParams instances, 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​

enable_dp_cross_replica_prefix_copy: bool = True

source

Whether a DP cross-replica prefix copy may serve a hit.

enable_prefix_caching​

enable_prefix_caching: bool = False

source

Whether prefix caching is enabled, a value every child must share.

from_params()​

classmethod from_params(params)

source

Creates a MultiKVCacheParams from one or more param sets.

Children may be leaf KVCacheParams instances, one RecurrentStateParams, or nested MultiKVCacheParams trees, enabling arbitrarily deep KV cache hierarchies (such as {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, CacheLeafParamInterface]) – Named mapping of CacheLeafParamInterface instances to aggregate. At least one must be a cache an attention op reads, since the pool’s configuration is read off one.

Returns:

A new MultiKVCacheParams aggregating all provided params.

Raises:

ValueError – If no params are provided, or if none of them is a cache an attention op reads.

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:

dict[str, KVCacheInputs[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​

kv_hash_algo: Literal['ahash64', 'sha256'] = 'ahash64'

source

Hash algorithm used for block identity.

kv_hash_seed​

kv_hash_seed: bytes | None = None

source

Resolved cluster seed for sha256.

leaf_kind​

property leaf_kind: CacheLeafKind

source

Attention where the subtree holds any, else recurrent.

leaves()​

leaves(_prefix='')

source

Returns the leaves of every child, prefixed by the child’s name.

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.

A leaf’s own answer, not the tree’s: an MHA leaf shards its heads where an MLA one replicates its latent, so siblings can differ.

resolve_attn_key()​

resolve_attn_key(batch_size, max_prompt_length, max_cache_valid_length)

source

Resolves the dispatch shape tree mirroring the attention caches.

Parameters:

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

Return type:

AttnKeyInterface

slab_to_bound_views()​

slab_to_bound_views(slabs)

source

Returns whatever the children bind, in one mapping.

Keys come from the leaves, which are unique across the tree.

Parameters:

slabs (Sequence[Buffer])

Return type:

Mapping[str, list[Buffer]]

slab_to_buffer_views()​

slab_to_buffer_views(buffers, padded_page_bytes=None, _prefix='')

source

Converts a slab of memory into every child’s view of it.

Parameters:

Return type:

KVCacheBufferInterface

speculative_method​

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

source

Speculative decoding method propagated from SpeculativeConfig.

staged_dispatch_metadata()​

staged_dispatch_metadata(_prefix='')

source

The metadata every child stages, under the child’s own prefix.

Parameters:

_prefix (str)

Return type:

Mapping[str, DispatchMetadataSpec]

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.

Returns one entry per attention child, in declaration order.

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:

dict[str, KVCacheInputs[TensorValue, BufferValue]]

Models

Cookbook

Puzzles

Blog

Community