For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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

MLAKVCacheParams​

graph_capture_probe_cache_lengths()​

kv_dim​

n_kv_heads_per_device​

num_q_heads​

num_q_heads_per_device​

replicates_kv_across_tp​

resolve_attn_key()​

`MLAKVCacheParams`

`graph_capture_probe_cache_lengths()`

`kv_dim`

`n_kv_heads_per_device`

`num_q_heads`

`num_q_heads_per_device`

`replicates_kv_across_tp`

`resolve_attn_key()`