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, slots_per_page=None, data_parallel_degree=1, kvcache_quant_config=None, speculative_method=None, num_draft_tokens=0, window_size=None, stride_alignment=None, scratch=False, *, n_kv_heads, allow_kv_head_replication=False)
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'])
- kv_hash_seed (bytes | None)
- kv_connector_config (KVConnectorConfigInterface)
- page_size (int)
- slots_per_page (int | None)
- data_parallel_degree (int)
- kvcache_quant_config (KVCacheQuantizationConfig | None)
- speculative_method (Literal['eagle', 'mtp', 'dflash', 'dflash2', 'dspark'] | None)
- num_draft_tokens (int)
- window_size (int | None)
- stride_alignment (int | None)
- scratch (bool)
- n_kv_heads (int)
- allow_kv_head_replication (bool)
dispatch_metadata_spec()
dispatch_metadata_spec()
Returns the spec MSAAttnKey packs into.
-
Return type:
-
DispatchMetadataSpec
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.
-
Parameters:
-
Return type: