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: CacheLeafParamInterface, Protocol

A cache leaf a model reads through an attention op.

It resolves a dispatch shape, names the cache lengths worth probing at graph capture, and hands out the paged buffers the op indexes. It also defines the pool: its page size, external tier, and block hash.

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.

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

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.

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

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 a 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, padded_page_bytes=None, _prefix='')

source

Converts a slab of memory into a buffer view.

Parameters:

• buffers (Sequence[Buffer])
• padded_page_bytes (Mapping[str, int] | None)
• _prefix (str)

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