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 module

max.nn.kv_cache

Cache configuration

KVCacheBuffer	A collection of KVCache buffers for one data-parallel replica.
KVCacheParamInterface	A cache leaf a model reads through an attention op.
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	One logical (child, kind) KV tensor as per-TP-shard uint8 views.
MultiKVCacheParams	Aggregates multiple cache parameter sets into a recursive tree.

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