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.legacy.kv_cache.KVCacheBuffer(replicates_kv_across_tp, values, values_packed=None, scales=None, scales_packed=None, values_padded=None, scales_padded=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:

• replicates_kv_across_tp (bool)
• values (list[Buffer])
• values_packed (list[Buffer] | None)
• scales (list[Buffer] | None)
• scales_packed (list[Buffer] | None)
• values_padded (list[Buffer] | None)
• scales_padded (list[Buffer] | None)
• leaf_id (str)

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.

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.

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.