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

KVTransferEngine

KVTransferEngine

class max.pipelines.kv_cache.KVTransferEngine(name, memory, leaf_ids=None)

source

Bases: TransferEngine

KVCache Transfer Engine with support for Data Parallelism (DP) and Tensor Parallelism (TP).

The engine accepts per-replica producer-authored NIXL groups (KVCacheMemory). The outer list is indexed by DP replica; the inner list is that replica’s group list — one group per pool leaf, in the order the cache declares its leaves.

KVTransferEngine is a thin layer on top of TransferEngine that adds:

• Validation of tensor shapes and device types
• NIXL group construction from the authored groups
• Per-group replication (replicated_per_group) authored from the groups’ replicated field (no caller plumbing needed)
• from_paged_kv_cache() convenience constructor

All NIXL transport operations are delegated to TransferEngine.

The TransferEngine communicates with other TransferEngines in other threads or processes. However, individual TransferEngines themselves are not thread-safe. It is intended to be used by MAX’s single-threaded scheduler.

Initialize the transfer engine from producer-authored NIXL groups.

Parameters:

• name (str) – Unique name for this engine.
• memory (Sequence[Sequence[KVCacheMemory]]) – Per-replica group lists as [replica][group]. Each entry is a KVCacheMemory — one pool leaf’s unit, carrying every TP-shard view, ordered by the caller. All replicas must have the same group count and consistent replication kind. Each group’s page count (including the null block) is read from the group itself, and every replica of a group must agree on it.
• leaf_ids (list[str]) – Id of each group, parallel to a replica’s group list. Transfers address pages by these ids. Defaults to each group’s position, "0", "1", and so on.

from_paged_kv_cache()

classmethod from_paged_kv_cache(name, kv_cache)

source

Construct an engine wired to a PagedKVCacheManager.

Calls to_memory() on each replica’s device buffer to obtain the producer-authored NIXL groups, orders them by the leaves the cache declares, and passes them to the constructor, which carries each group’s replicated field as replicated_per_group.

For models with multiple KV caches (such as speculative decoding with a separate target and draft KV), each child cache contributes its own group(s) so that heterogeneous buffer shapes (such as 61-layer MLA target vs. 1-layer Eagle draft) are registered as independent NIXL groups.

Quantized caches (values + scales): to_memory() authors a separate group for values and for scales (one group per child x kind). For non-quantized caches this collapses to one group per child, which is byte-identical to the previous all_buffers path.

Raises:

NotImplementedError – If the cache keeps recurrent state.

Parameters:

• name (str)
• kv_cache (PagedKVCacheManagerInterface)

Return type:

KVTransferEngine