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
MSAKVCacheParams
MSAKVCacheParams
class max.nn.kv_cache.MSAKVCacheParams(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, *, n_kv_heads: 'int', allow_kv_head_replication: 'bool' = False)
Bases: MHAKVCacheParams
-
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)
- 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)
Returns cache lengths to probe for distinct num_partitions.
resolve_attn_key()
resolve_attn_key(batch_size, max_prompt_length, max_cache_valid_length)
Resolves the decode attention dispatch shape for the given shape.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!