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 module

max.nn.kv_cache

Cache configuration

`KVCacheBuffer`	A collection of KVCache buffers for one data-parallel replica.
`KVCacheParamInterface`	Interface for KV cache parameters.
`KVCacheParams`	Configuration parameters for key-value cache management in transformer models.
`MHAKVCacheParams`
`MLAKVCacheParams`
`MSAKVCacheParams`
`KVCacheQuantizationConfig`	Configuration for KVCache quantization.
`KVConnectorType`	Identifies which off-device backing store the KV cache uses.
`KVCacheMemory`	A single KV cache shard as a 2-D `uint8` view.
`MultiKVCacheParams`	Aggregates multiple KV cache parameter sets into a recursive tree.
`ReplicatedKVCacheMemory`	A replicated KV cache unit (rank-0 shard plus its TP peers).

Cache inputs

`KVCacheInputs`	Symbolic graph input types for a leaf KV cache.
`KVCacheInputsPerDevice`	Symbolic graph input types for a single device's paged KV cache.
`BatchCharacteristics`	Upper-bound batch shape used to prepare decode attention metadata.
`PagedCacheValues`	alias of `KVCacheInputsPerDevice`[`TensorValue`, `BufferValue`]

Attention dispatch

`AttnKey`	A resolved decode-attention dispatch shape.
`MHAAttnKey`	Decode dispatch metadata for multi-head attention (MHA).
`MLAAttnKey`	Decode dispatch metadata for multi-latent attention (MLA).
`MSAAttnKey`	Decode dispatch metadata for multi-step attention (MSA).

Metrics

`KVCacheMetrics`	Metrics for the KV cache.

Functions

`build_max_lengths_tensors`	Builds two `[1]` uint32 scalar buffers of maximum lengths.
`compute_max_seq_len_fitting_in_cache`	Computes the maximum sequence length that can fit in the available memory.
`compute_num_device_blocks`	Computes the number of blocks that can be allocated based on the available cache memory.
`compute_num_host_blocks`	Computes the number of blocks that can be allocated on the host.
`estimated_memory_size`	Computes the estimated memory size of the KV cache used by all replicas.
`spec_decode_cache_slack`	Computes the extra KV positions a request may occupy past `max_seq_len`.

Cache configuration
Cache inputs
Attention dispatch
Metrics
Functions