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.legacy.kv_cache.KVCacheParams(dtype, head_dim, num_layers, devices, enable_prefix_caching=False, enable_dp_cross_replica_prefix_copy=True, 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, v_head_dim=None)

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.

That is scratch for a scratch cache, sliding_window with the window size when one is set, else full.

has_split_head_dims​

property has_split_head_dims: bool

source

Whether the key and value heads have different dims.

When they differ, one packed row holds both (see v_head_dim).

head_dim​

head_dim: int

source

Dimensionality of each attention head; the key (and query) head dim when v_head_dim differs.

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.

Each uint8 value holds two E2M1 elements, 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.

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_scale_block​

property shape_per_scale_block: list[int]

source

Returns the shape of each quantization scale block.

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), and for split head dims, whose row holds the key and then the value.

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], …]

v_head_dim​

v_head_dim: int | None = None

source

Dimensionality of each value head.

None (the default) means head_dim, and construction turns a value equal to head_dim into None. It stays unresolved so that dataclasses.replace(params, head_dim=...) on an unsplit cache leaves it unsplit. When it differs (has_split_head_dims), each cache slot holds one packed row per KV head, the key in its first head_dim columns and the value in the next v_head_dim, so neither is padded to the other’s width.

window_size​

window_size: int | None = None

source

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

Models

Cookbook

Puzzles

Blog

Community