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

KVCacheParamInterface

KVCacheParamInterface

class max.nn.kv_cache.KVCacheParamInterface(*args, **kwargs)

source

Bases: Protocol

Interface for KV cache parameters.

allocate_buffers()

allocate_buffers(total_num_pages)

source

Allocates the buffers for the KV cache.

Parameters:

total_num_pages (int)

Return type:

Sequence[KVCacheBufferInterface]

build_runtime_inputs()

build_runtime_inputs(assignments, buffers, _prefix='')

source

Builds the runtime KV-cache inputs spanning all replicas.

assignments and buffers are indexed by data-parallel replica. Returns a single KVCacheInputs leaf (or a MultiKVCacheInputs tree) whose leaves each hold every (replica, TP shard) device’s inputs.

Parameters:

• assignments (Sequence[KVCacheAssignments])
• buffers (Sequence[KVCacheBufferInterface])
• _prefix (str)

Return type:

KVCacheInputsInterface[Buffer, Buffer]

bytes_per_block

property bytes_per_block: int

source

Number of bytes per cache block.

data_parallel_degree

data_parallel_degree: int

source

Degree of data parallelism.

devices

devices: Sequence[DeviceRef]

source

Devices to use for the cache.

enable_dp_cross_replica_prefix_copy

property enable_dp_cross_replica_prefix_copy: bool

source

Whether a prefix-cache hit resident on another data-parallel replica’s device may be served by a device-to-device copy.

enable_prefix_caching

property enable_prefix_caching: bool

source

Whether prefix caching is enabled.

flattened_kv_inputs()

flattened_kv_inputs()

source

Flattens the symbolic inputs for the KV cache.

Return type:

list[TensorType | BufferType]

get_symbolic_inputs()

get_symbolic_inputs(namespace='')

source

Returns the symbolic inputs for the KV cache.

Parameters:

namespace (str) – Prefix that disambiguates this cache’s page-pool symbolic dim from sibling caches in a multi-group tree. Empty for a single-group cache, leaving its names unchanged.

Return type:

KVCacheInputsInterface[TensorType, BufferType]

graph_capture_probe_cache_lengths()

graph_capture_probe_cache_lengths(max_cache_length, q_max_seq_len=1)

source

Returns the cache lengths to probe during decode graph capture.

Parameters:

• max_cache_length (int)
• q_max_seq_len (int)

Return type:

list[int]

kv_connector_config

kv_connector_config: KVConnectorConfigInterface

source

The KV connector’s type and settings.

kv_hash_algo

property kv_hash_algo: Literal['ahash64', 'sha256', 'sha256_64']

source

Hash algorithm used for KV-cache block identity.

kv_hash_seed

property kv_hash_seed: bytes | None

source

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

leaves()

leaves(_prefix='')

source

Returns the leaves of the KV cache.

Parameters:

_prefix (str)

Return type:

Mapping[str, KVLeafRegion]

n_devices

property n_devices: int

source

Returns the total number of devices.

num_draft_tokens

num_draft_tokens: int = 0

source

num_draft_tokens_per_step

property num_draft_tokens_per_step: int

source

Number of draft tokens written per draft forward.

Zero when speculative decoding is disabled; one for autoregressive drafts (eagle, mtp); equal to num_draft_tokens for block drafts (dflash, dflash2).

page_size

page_size: int

source

Number of tokens per page (block).

replicates_kv_across_tp

property replicates_kv_across_tp: bool

source

Whether every device holds identical KV state.

resolve_attn_key()

resolve_attn_key(batch_size, max_prompt_length, max_cache_valid_length)

source

Resolves the decode dispatch shape for the given shape.

Returns an AttnKeyInterface for a single cache, or a MultiAttnKey tree mirroring the cache tree.

Parameters:

• batch_size (int)
• max_prompt_length (int)
• max_cache_valid_length (int)

Return type:

AttnKeyInterface

slab_to_buffer_views()

slab_to_buffer_views(buffers)

source

Converts a slab of memory into a buffer view.

Parameters:

buffers (Sequence[Buffer])

Return type:

KVCacheBufferInterface

speculative_method

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

source

tensor_parallel_degree

property tensor_parallel_degree: int

source

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 symbolic inputs for the KV cache.

Parameters:

it (Iterator[Any])

Return type:

KVCacheInputsInterface[TensorValue, BufferValue]