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

MSAKVCacheParams

MSAKVCacheParams

class max.nn.kv_cache.MSAKVCacheParams(dtype, head_dim, num_layers, devices, enable_prefix_caching=False, enable_dp_cross_replica_prefix_copy=True, per_layer_buffers=False, kv_hash_algo='ahash64', kv_hash_seed=None, kv_connector_config=<factory>, page_size=128, data_parallel_degree=1, kvcache_quant_config=None, speculative_method=None, num_draft_tokens=0, window_size=None, *, n_kv_heads, allow_kv_head_replication=False)

source

Bases: MHAKVCacheParams

KV cache parameters for multi-step attention (MSA).

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_hash_algo (Literal['ahash64', 'sha256', 'sha256_64'])
• kv_hash_seed (bytes | None)
• kv_connector_config (KVConnectorConfigInterface)
• page_size (int)
• data_parallel_degree (int)
• kvcache_quant_config (KVCacheQuantizationConfig | None)
• speculative_method (Literal['eagle', 'mtp', 'dflash', 'dflash2'] | None)
• num_draft_tokens (int)
• window_size (int | None)
• n_kv_heads (int)
• allow_kv_head_replication (bool)

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]

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)
• max_prompt_length (int)
• max_cache_valid_length (int)

Return type:

AttnKeyInterface