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)
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)
- values_per_layer (list[list[Buffer]] | None)
- scales_per_layer (list[list[Buffer]] | None)
- leaf_id (str)
all_buffers
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
The pool leaf these pages belong to, as leaves() names it.
replicates_kv_across_tp
replicates_kv_across_tp: bool
scales
Per-TP-shard scale buffers for a quantized cache; None when
unquantized.
scales_packed
Contiguous aliases of scales; see values_packed.
scales_padded
The same as values_padded, for scales.
scales_per_layer
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()
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:
values
Page views, strided when the pool padded this leaf. The canonical form: the page stride is read off these.
values_packed
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
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
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.