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

MLAKVCacheParams

MLAKVCacheParams

class max.nn.kv_cache.MLAKVCacheParams(dtype: 'DType', head_dim: 'int', num_layers: 'int', devices: 'Sequence[DeviceRef]', enable_prefix_caching: 'bool' = False, enable_dp_cross_replica_prefix_copy: 'bool' = True, per_layer_buffers: 'bool' = False, kv_connector: 'KVConnectorType | None' = None, kv_hash_algo: 'KVHashAlgo' = 'ahash64', kv_hash_seed: 'bytes | None' = None, kv_connector_config: 'Any' = None, host_kvcache_swap_space_gb: 'float | None' = None, page_size: 'int' = 128, data_parallel_degree: 'int' = 1, kvcache_quant_config: 'KVCacheQuantizationConfig | None' = None, speculative_method: 'SpeculativeMethod | None' = None, num_draft_tokens: 'int' = 0, *, num_q_heads: 'int')

source

Bases: KVCacheParams

Parameters:

• dtype (DType)
• head_dim (int)
• num_layers (int)
• devices (Sequence[DeviceRef])
• enable_prefix_caching (bool)
• enable_dp_cross_replica_prefix_copy (bool)
• per_layer_buffers (bool)
• kv_connector (KVConnectorType | None)
• kv_hash_algo (Literal['ahash64', 'sha256', 'sha256_64'])
• kv_hash_seed (bytes | None)
• kv_connector_config (Any)
• host_kvcache_swap_space_gb (float | None)
• page_size (int)
• data_parallel_degree (int)
• kvcache_quant_config (KVCacheQuantizationConfig | None)
• speculative_method (Literal['eagle', 'mtp', 'dflash'] | None)
• num_draft_tokens (int)
• num_q_heads (int)

graph_capture_probe_cache_lengths()

graph_capture_probe_cache_lengths(max_cache_length, q_max_seq_len=1)

source

Returns cache lengths to probe for distinct num_partitions.

Parameters:

• max_cache_length (int)
• q_max_seq_len (int)

Return type:

list[int]

kv_dim

property kv_dim: int

source

n_kv_heads_per_device

property n_kv_heads_per_device: int

source

num_q_heads

num_q_heads: int

source

Number of query attention heads, required so the MLA decode kernel can resolve its dispatch metadata.

num_q_heads_per_device

property num_q_heads_per_device: 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 attention dispatch shape for the given shape.

Parameters:

• batch_size (int) – Number of requests in the decode batch.
• max_prompt_length (int) – Per-step query width (1 for plain decode, 1 + num_spec_tokens for speculative verify).
• max_cache_valid_length (int) – Maximum valid cache length in the batch.

Returns:

The resolved AttnKeyInterface

Return type:

AttnKeyInterface