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

KVCacheBuffer

KVCacheBuffer​

class max.nn.kv_cache.KVCacheBuffer(replicates_kv_across_tp, values, values_packed=None, scales=None, scales_packed=None, values_padded=None, scales_padded=None, values_per_layer=None, scales_per_layer=None, *, leaf_id)

source

Bases: KVCacheBufferInterface

A collection of KVCache buffers for one data-parallel replica.

Two buffer kinds are supported: values and (optionally, for FP8 quantization) scales. The length of each list corresponds to the tensor-parallel degree, with one buffer per TP shard.

replicates_kv_across_tp is True when the KV data is replicated identically across TP shards and False when it is sharded. The data is replicated in certain cases like TP + MLA, TP + MiniMaxM3IndexerAttn, and so on.

Parameters:

all_buffers​

property all_buffers: list[Buffer]

source

Returns all value and scale buffers in a single flat list.

Returns:

A list containing every value buffer followed by every scale buffer (if scales are present).

leaf_id​

leaf_id: str

source

The pool leaf these pages belong to, as leaves() names it.

replicates_kv_across_tp​

replicates_kv_across_tp: bool

source

scales​

scales: list[Buffer] | None = None

source

Per-TP-shard scale buffers for a quantized cache; None when unquantized.

scales_packed​

scales_packed: list[Buffer] | None = None

source

Contiguous aliases of scales; see values_packed.

scales_padded​

scales_padded: list[Buffer] | None = None

source

The same as values_padded, for scales.

scales_per_layer​

scales_per_layer: list[list[Buffer]] | None = None

source

Per-TP-shard, per-layer scale buffers for a quantized KV cache backed by per_layer_buffers (mirrors values_per_layer). scales[shard] aliases scales_per_layer[shard][0]. None for a single multi-layer scale buffer or an unquantized cache.

to_memory()​

to_memory()

source

Converts to offload-ready memory units, keyed by pool leaf id.

Every buffer is re-viewed as 2-D uint8 pages so consumers can treat all caches uniformly regardless of dtype or shape. A padded leaf hands over its padded grid instead, so a page’s pad travels with it.

Per-layer buffers are deliberately not enumerated – only each shard’s layer-0 alias – which is why allocate_buffers rejects per_layer_buffers alongside off-device connectors and DP > 1.

Returns:

This cache’s values leaf, plus its scales leaf when quantized.

Return type:

Mapping[str, KVCacheMemory]

values​

values: list[Buffer]

source

Page views, strided when the pool padded this leaf. The canonical form: the page stride is read off these.

values_packed​

values_packed: list[Buffer] | None = None

source

Contiguous aliases of values, for binding as graph inputs.

Model execution rejects a non-contiguous buffer, so a padded leaf cannot be bound as its strided view. These cover the same allocation packed, and are deliberately shorter than the span their stride reaches – see contiguous_page_view_and_stride. None when nothing was padded, in which case values is already contiguous.

values_padded​

values_padded: list[Buffer] | None = None

source

Contiguous [num_pages, page_stride] uint8 under values, pad included, for to_memory(). None when the pool did not pad this leaf.

values_per_layer​

values_per_layer: list[list[Buffer]] | None = None

source

Per-TP-shard, per-layer value buffers when the pool uses per_layer_buffers.

values_per_layer[shard] is the list of single-layer buffers for that shard, and values[shard] aliases values_per_layer[shard][0] so the single-buffer values invariants (and consumers) stay valid. None for a normal single multi-layer buffer.