IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

MHAKVCacheParams

MHAKVCacheParams​

class max.nn.kv_cache.MHAKVCacheParams(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)

source

Bases: KVCacheParams

Parameters:

allow_kv_head_replication​

allow_kv_head_replication: bool = False

source

when set and n_devices is a multiple of n_kv_heads, replicate each KV head across a group of devices (n_kv_heads_per_device == 1).

Type:

Allow TP wider than n_kv_heads

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​

n_kv_heads: int

source

Total number of key-value attention heads across all devices.

n_kv_heads_per_device​

property n_kv_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

Models

Cookbook

Puzzles

Blog

Community

Contact