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)
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 (e.g. {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)
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.
build_runtime_inputs()
build_runtime_inputs(assignments, buffers, _prefix='')
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.
bytes_per_block
property bytes_per_block: int
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
Cache parameter sets to aggregate. Values may be leaf
KVCacheParams or RecurrentStateParams instances, or
nested MultiKVCacheParams trees.
data_parallel_degree
data_parallel_degree: int
devices
enable_dp_cross_replica_prefix_copy
property enable_dp_cross_replica_prefix_copy: bool
Whether DP cross-replica prefix copies are enabled (shared across all caches).
enable_prefix_caching
property enable_prefix_caching: bool
Whether prefix caching is enabled (shared across all caches).
from_params()
classmethod from_params(params)
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 (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, CacheLeafParamInterface]) – Named mapping of
CacheLeafParamInterfaceinstances 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
MultiKVCacheParamsaggregating all provided params. -
Raises:
-
ValueError – If no params are provided, or if none of them is a cache an attention op reads.
-
Return type:
get_symbolic_inputs()
get_symbolic_inputs(namespace='')
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)
Returns the union of probe cache lengths across all child caches.
kv_connector_config
kv_connector_config: KVConnectorConfigInterface
kv_hash_algo
property kv_hash_algo: Literal['ahash64', 'sha256', 'sha256_64']
Hash algorithm used for KV-cache block identity.
kv_hash_seed
Resolved 32-byte cluster seed for sha256/sha256_64. None for ahash64.
leaves()
leaves(_prefix='')
Returns the leaves of every child, prefixed by the child’s name.
n_devices
property n_devices: int
Returns the number of devices.
num_draft_tokens
num_draft_tokens: int = 0
page_size
page_size: int
replicates_kv_across_tp
property replicates_kv_across_tp: bool
Whether every device holds identical KV state.
resolve_attn_key()
resolve_attn_key(batch_size, max_prompt_length, max_cache_valid_length)
Resolves the dispatch shape tree mirroring the attention caches.
slab_to_bound_views()
slab_to_bound_views(slabs)
Returns whatever the children bind, in one mapping.
Keys come from the leaves, which are unique across the tree.
slab_to_buffer_views()
slab_to_buffer_views(buffers, padded_page_bytes=None, _prefix='')
Converts a slab of memory into a buffer view.
Only the attention children have pages to view.
speculative_method
speculative_method: Literal['eagle', 'mtp', 'dflash', 'dflash2'] | None = None
tensor_parallel_degree
property tensor_parallel_degree: int
Returns the tensor parallel degree.
unflatten_basic_kv_tree()
unflatten_basic_kv_tree(it)
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:
-
Return type:
-
tuple[list[KVCacheInputsPerDevice[TensorValue, BufferValue]], …]
unflatten_kv_inputs()
unflatten_kv_inputs(it)
Unflattens the KV cache inputs from a graph-input iterator.
-
Parameters:
-
Return type:
-
MultiKVCacheInputs[TensorValue, BufferValue]