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

MHAKVCacheParams

MHAKVCacheParams

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

KV cache parameters for multi-head attention (MHA).

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)

allow_kv_head_replication

allow_kv_head_replication: bool = False

source

Allows TP wider than n_kv_heads. When set and n_devices is a multiple of n_kv_heads, each KV head is replicated across a group of devices (n_kv_heads_per_device == 1).

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

each slot holds a key and a value tensor.

Type:

Returns two

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

Returns the KV heads on one device, or 1 per device group when heads are replicated (allow_kv_head_replication).

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