IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/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. /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:

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]

Models

Cookbook

Puzzles

Blog

Community

Contact