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.

dispatch_metadata_spec()

dispatch_metadata_spec()

source

Returns the spec of the metadata resolve_attn_key() packs.

None when this cache resolves a tree of keys rather than one, so a caller cannot size a single input for it.

Return type:

DispatchMetadataSpec | None

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

The KV connector’s type and settings.

kv_hash_algo

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

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. 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, dspark).

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

speculative_method

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

source

staged_dispatch_metadata()

staged_dispatch_metadata(_prefix='')

source

The dispatch metadata this cache stages, keyed by name.

Empty for a kernel that reads it on the host. A draft key under speculative decoding is a second input of the same shape.

Parameters:

_prefix (str)

Return type:

Mapping[str, DispatchMetadataSpec]

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