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

KVCacheParams

KVCacheParams​

class max.nn.kv_cache.KVCacheParams(dtype, head_dim, num_layers, devices, enable_prefix_caching=False, enable_dp_cross_replica_prefix_copy=True, per_layer_buffers=False, kv_hash_algo='ahash64', kv_hash_seed=None, kv_connector_config=<factory>, page_size=128, slots_per_page=None, data_parallel_degree=1, kvcache_quant_config=None, speculative_method=None, num_draft_tokens=0, window_size=None, stride_alignment=None, scratch=False)

source

Bases: KVCacheParamInterface

Configuration parameters for key-value cache management in transformer models.

This class encapsulates all configuration options for managing KV caches during inference, including parallelism settings, and memory management.

Parameters:

allocate_buffers()​

allocate_buffers(total_num_pages, _prefix='')

source

Allocates the buffers for the KV cache.

Parameters:

  • total_num_pages (int)
  • _prefix (str)

Return type:

list[KVCacheBuffer]

build_runtime_inputs()​

build_runtime_inputs(assignments, buffers, _prefix='', *, staging=None)

source

Builds the runtime KV-cache leaf spanning all replicas.

assignments and buffers are indexed by data-parallel replica. The returned KVCacheInputs lists one KVCacheInputsPerDevice per (replica, TP shard), in the same replica-major order as get_symbolic_inputs().

Parameters:

Return type:

tuple[KVCacheInputsPerDevice[Buffer, Buffer], …]

bytes_per_block​

property bytes_per_block: int

source

Returns the number of bytes per cache block.

When TP>1, each block is sharded across the devices in the tensor parallel group. This method returns the total memory needed to store a block across these devices. Includes memory needed for scales if quantization is enabled.

Returns:

The number of bytes per cache block.

bytes_per_scale_block​

property bytes_per_scale_block: int

source

Returns the number of bytes per scale block.

bytes_per_value_block​

property bytes_per_value_block: int

source

Returns the number of bytes per value block.

data_parallel_degree​

data_parallel_degree: int = 1

source

Degree of data parallelism. Devices are grouped replica-major, with n_devices // data_parallel_degree TP shards per replica.

devices​

devices: Sequence[DeviceRef]

source

Devices to use for the KV cache.

devices_per_replica​

property devices_per_replica: Sequence[Sequence[DeviceRef]]

source

Returns the devices per replica.

dtype​

dtype: DType

source

Data type for storing key and value tensors in the cache.

dtype_shorthand​

property dtype_shorthand: str

source

Returns a shorthand textual representation of the data type.

Returns:

“bf16” for bfloat16 dtype, “f32” otherwise.

enable_dp_cross_replica_prefix_copy​

enable_dp_cross_replica_prefix_copy: bool = True

source

Whether a prefix-cache block resident on another data-parallel (DP) replica’s device may be materialized locally via a device-to-device copy to serve a cache hit. When False, cross-replica reuse is only served from the shared external tier via the KV connector (or recomputed). Only relevant when data_parallel_degree > 1 and prefix caching is enabled.

enable_prefix_caching​

enable_prefix_caching: bool = False

source

Whether to enable prefix caching for efficient reuse of common prompt prefixes.

get_symbolic_inputs()​

get_symbolic_inputs(namespace='')

source

Computes the symbolic inputs for the KV cache.

Parameters:

namespace (str) – Prefix disambiguating this cache’s per-pool page-count dim from sibling caches in a multi-group tree (empty for a single-group cache).

Returns:

The symbolic inputs for the KV cache.

Return type:

tuple[KVCacheInputsPerDevice[TensorType, BufferType], …]

group_id​

property group_id: KVCacheGroupId

source

Returns the group id this cache pools under: scratch for a scratch cache, sliding_window with the window size when one is set, else full.

head_dim​

head_dim: int

source

Dimensionality of each attention head.

is_fp8_kv_dtype​

property is_fp8_kv_dtype: bool

source

Whether the KV cache stores FP8 data, for dispatch resolution.

Unlike quantized_kv_cache (which also requires valid scale config), this checks only the storage dtype—matching the compile-time detection in the MLA decode kernel.

TODO(SERVOPT-1094): Once SnapMLA uses a valid scale_dtype, this can be replaced by quantized_kv_cache.

is_nvfp4_kv_cache​

property is_nvfp4_kv_cache: bool

source

Whether the cache stores packed NVFP4: uint8 values holding two E2M1 elements each, with E4M3 scales.

kv_cache_scale_dtype​

property kv_cache_scale_dtype: DType

source

Returns the dtype of the KV cache scales.

Returns:

The dtype of the KV cache scales.

kv_connector_config​

kv_connector_config: KVConnectorConfigInterface

source

Holds the connector type and its settings. The default is a null connector (no external caching).

kv_dim​

property kv_dim: int

source

Returns the number of key/value tensors each cache slot holds.

kv_hash_algo​

kv_hash_algo: Literal['ahash64', 'sha256'] = 'ahash64'

source

Hash algorithm used for KV-cache block identity.

kv_hash_seed​

kv_hash_seed: bytes | None = None

source

Resolved 32-byte cluster seed for sha256. None for ahash64.

Set by KVCacheConfig.to_params via resolve_kv_hash_seed.

kvcache_quant_config​

kvcache_quant_config: KVCacheQuantizationConfig | None = None

source

KVCache quantization config. Currently only FP8 quantization supported.

leaf_id()​

leaf_id(_prefix='')

source

Returns the id this cache’s pages are pooled under.

Parameters:

_prefix (str)

Return type:

str

leaf_kind​

leaf_kind: ClassVar[CacheLeafKind] = 'attention'

source

leaves()​

leaves(_prefix='')

source

Returns the leaves of the KV cache.

A leaf reports what one device holds, the unit a slab tiles in, where bytes_per_block counts a block across the whole replica. The division is exact: that figure is this one times the degree.

Parameters:

_prefix (str)

Return type:

Mapping[str, KVLeafRegion]

n_devices​

property n_devices: int

source

Returns the number of devices.

Returns:

The number of devices.

n_kv_heads_per_device​

property n_kv_heads_per_device: int

source

Returns the number of KV attention heads on one device.

num_draft_tokens​

num_draft_tokens: int = 0

source

Total draft tokens generated per speculative iteration.

Zero when no speculative decoding is configured.

num_layers​

num_layers: int

source

Number of layers in the model.

page_size​

page_size: int = 128

source

Number of tokens per page (block).

This value is expressed in tokens, not bytes. The byte footprint of a page is derived from pipeline configuration.

Current constraints: the page size must be a multiple of 128 and at least 128.

per_layer_buffers​

per_layer_buffers: bool = False

source

When True, allocate one standalone single-layer buffer per layer instead of one [..., num_layers, ...] multi-layer buffer.

Each attention dispatch then binds only its own per-layer buffer, so the pool total can exceed a per-allocation size cap (such as a device’s maximum single allocation) while every individual buffer stays under it. Defaults to False (one multi-layer buffer), keeping all other backends and models byte-identical.

quantized_kv_cache​

property quantized_kv_cache: bool

source

Returns whether KV cache quantization is enabled.

replicates_kv_across_tp​

property replicates_kv_across_tp: bool

source

Whether every device holds identical KV state.

row_bytes​

property row_bytes: int

source

Returns one value row, num_heads * head_size * dtype_size.

scale_row_bytes​

property scale_row_bytes: int

source

Returns one scale row; see row_bytes.

scratch​

scratch: bool = False

source

Whether the cache is one block a request, kept until release and never published. Excludes window_size.

shape_per_block​

property shape_per_block: list[int]

source

Returns the shape of each cache block.

The slot dimension is slots_per_page, which equals page_size unless the architecture stores a page’s tokens compressed.

Returns:

The shape of the cache block.

shape_per_layer_block​

property shape_per_layer_block: list[int]

source

Returns the block shape for a single-layer buffer.

Same as shape_per_block but with the layer dimension pinned to 1. Used when per_layer_buffers is set: the pool allocates num_layers such buffers per device instead of one multi-layer buffer. The attention kernel derives num_layers from this dim, so a single-layer buffer (num_layers == 1) with layer_idx == 0 is self-consistent.

shape_per_layer_scale_block​

property shape_per_layer_scale_block: list[int]

source

Scale-block shape for a single-layer buffer (layer dim pinned to 1).

The scale analog of shape_per_layer_block: used with per_layer_buffers on a quantized cache, where the pool allocates one single-layer scale buffer per layer instead of one multi-layer one.

shape_per_scale_block​

property shape_per_scale_block: list[int]

source

Returns the shape of each scale block used for KVCache quantization

Returns:

The shape of the KVCache quantization scales block.

slab_to_bound_views()​

slab_to_bound_views(slabs)

source

Returns nothing: these pages are addressed by the lookup table.

Parameters:

slabs (Sequence[Buffer])

Return type:

Mapping[str, list[Buffer]]

slab_to_buffer_views()​

slab_to_buffer_views(buffers, padded_page_bytes=None, _prefix='')

source

Converts a slab of memory into per-leaf page views.

Each view carries its own page-to-page distance in its dim-0 stride, so nothing downstream – the pool, the connector, the transfer engine – has to be told whether this pool padded its pages.

Parameters:

  • buffers (Sequence[Buffer]) – One raw arena slab per TP shard.
  • padded_page_bytes (Mapping[str, int] | None) – The pool’s padded page size per leaf, keyed as leaves() keys it. None means nothing was padded.
  • _prefix (str) – Leaf-id prefix identifying this node in a params tree.

Returns:

The buffer views for this cache.

Return type:

KVCacheBufferInterface

slots_per_page​

slots_per_page: int | None = None

source

Number of storage slots a page holds, which may be fewer than the page_size tokens it covers.

None (the default) means one slot per token; construction resolves it to page_size, so every later reader sees an int. An architecture whose attention compresses a run of tokens into a single cached entry – a pooled key, a compressed latent – sets this to the number of entries a page’s tokens compress into, and gets a page that many slots deep instead. It must divide page_size evenly, so a page never ends mid-slot.

Only the physical buffer shape and the kernels reading it are affected: block allocation, the page table and prefix caching stay in tokens, and a block still covers page_size of them.

speculative_method​

speculative_method: Literal['eagle', 'mtp', 'dflash', 'dflash2', 'dspark'] | None = None

source

Speculative decoding method propagated from SpeculativeConfig.

storage_head_dim​

property storage_head_dim: int

source

Returns the stored width of one head row in dtype elements.

Equals head_dim except for packed NVFP4, which stores two elements per byte, apart from its trailing FP8 elements (see KVCacheQuantizationConfig.fp8_tail_dim).

stride_alignment​

stride_alignment: int | None = None

source

The page stride alignment this cache’s leaves declare, or None to keep them packed. See KVLeafRegion.stride_alignment.

tensor_parallel_degree​

property tensor_parallel_degree: int

source

Returns the tensor parallel degree.

Returns:

The tensor parallel degree.

unflatten_basic_kv_tree()​

unflatten_basic_kv_tree(it)

source

Unflattens a basic KV tree from a graph-input iterator.

Requires that the model is a basic height-1 tree. This method does not work on nested trees.

Parameters:

it (Iterator[Any])

Return type:

tuple[list[KVCacheInputsPerDevice[TensorValue, BufferValue]], …]

unflatten_kv_inputs()​

unflatten_kv_inputs(it)

source

Unflattens the KV cache inputs from a graph-input iterator.

Parameters:

it (Iterator[Any])

Return type:

tuple[KVCacheInputsPerDevice[TensorValue, BufferValue], …]

window_size​

window_size: int | None = None

source

Window size for the sliding window attention. None for global attention.

Models

Cookbook

Puzzles

Blog

Community