For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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_connector=None, kv_hash_algo='ahash64', kv_hash_seed=None, kv_connector_config=None, host_kvcache_swap_space_gb=None, page_size=128, data_parallel_degree=1, kvcache_quant_config=None, speculative_method=None, num_draft_tokens=0)

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:

dtype (DType)
head_dim (int)
num_layers (int)
devices (Sequence[DeviceRef])
enable_prefix_caching (bool)
enable_dp_cross_replica_prefix_copy (bool)
per_layer_buffers (bool)
kv_connector (KVConnectorType | None)
kv_hash_algo (Literal['ahash64', 'sha256', 'sha256_64'])
kv_hash_seed (bytes | None)
kv_connector_config (Any)
host_kvcache_swap_space_gb (float | None)
page_size (int)
data_parallel_degree (int)
kvcache_quant_config (KVCacheQuantizationConfig | None)
speculative_method (Literal['eagle', 'mtp', 'dflash'] | None)
num_draft_tokens (int)

`allocate_buffers()`

allocate_buffers(total_num_pages)

source

Allocates the buffers for the KV cache.

Parameters:: total_num_pages (int)
Return type:: list[KVCacheBuffer]

`build_runtime_inputs()`

build_runtime_inputs(assignments, buffers)

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:

assignments (Sequence[KVCacheAssignments])
buffers (Sequence[KVCacheBufferInterface])

Return type:

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

`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:: KVCacheInputs[TensorType, BufferType]

`head_dim`

head_dim: int

source

Dimensionality of each attention head.

`host_kvcache_swap_space_gb`

host_kvcache_swap_space_gb: float | None = None

source

Amount of host memory (in GB) to reserve for KV cache swapping. Required when local or tiered connector is used.

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

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

kv_connector: KVConnectorType | None = None

source

Type of KV cache connector to use (null, local, tiered, dkv).

`kv_connector_config`

kv_connector_config: Any = None

source

Connector-specific configuration (KVConnectorConfig from the pipelines layer).

`kv_dim`

property kv_dim: int

source

`kv_hash_algo`

kv_hash_algo: Literal['ahash64', 'sha256', 'sha256_64'] = '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/sha256_64. 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.

`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

`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 (e.g. 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 FP8 KV cache quantization is enabled.

Returns:: True when the cache dtype is float8_e4m3fn or float8_e4m3fnuz and a valid quantization scale dtype is configured; False otherwise.

`replicates_kv_across_tp`

property replicates_kv_across_tp: bool

source

Whether every device holds identical KV state.

`shape_per_block`

property shape_per_block: list[int]

source

Returns the shape of each cache block.

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

`speculative_method`

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

source

Speculative decoding method propagated from SpeculativeConfig

`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:: KVCacheInputs[TensorValue, BufferValue]

KVCacheParams​

allocate_buffers()​

build_runtime_inputs()​

bytes_per_block​

data_parallel_degree​

devices​

devices_per_replica​

dtype​

dtype_shorthand​

enable_dp_cross_replica_prefix_copy​

enable_prefix_caching​

get_symbolic_inputs()​

head_dim​

host_kvcache_swap_space_gb​

is_fp8_kv_dtype​

kv_cache_scale_dtype​

kv_connector​

kv_connector_config​

kv_dim​

kv_hash_algo​

kv_hash_seed​

kvcache_quant_config​

n_devices​

n_kv_heads_per_device​

num_draft_tokens​

num_layers​

page_size​

per_layer_buffers​

quantized_kv_cache​

replicates_kv_across_tp​

shape_per_block​

shape_per_layer_block​

shape_per_scale_block​

speculative_method​

tensor_parallel_degree​

unflatten_basic_kv_tree()​

unflatten_kv_inputs()​