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
KVTransferEngine
KVTransferEngine
class max.pipelines.kv_cache.KVTransferEngine(name, memory)
Bases: TransferEngine
KVCache Transfer Engine with support for Data Parallelism (DP) and Tensor Parallelism (TP).
The engine accepts per-replica producer-authored NIXL groups
(KVCacheMemoryGroup). The outer list
is indexed by DP replica; the inner list is that replica’s group list — one
group per logical (child, kind) tensor, from to_memory_groups().
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’replicatedfield (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[KVCacheMemoryGroup]]) – Per-replica group lists as
[replica][group]. Each entry is aKVCacheMemoryGroup— one logical(child, kind)tensor carrying every TP-shard view, as returned byKVCacheBuffer.to_memory_groups(). All replicas must have the same group count and consistent replication kind. The page count (including the null block) is read from the groups themselves, so every group must agree ontotal_num_pages.
from_paged_kv_cache()
classmethod from_paged_kv_cache(name, kv_cache)
Construct an engine wired to a PagedKVCacheManager.
Calls KVCacheBuffer.to_memory_groups() on each replica’s device
buffer to obtain the producer-authored NIXL groups, then passes them to
the constructor, which carries each group’s replicated field as
replicated_per_group.
For models with multiple KV caches (e.g., speculative decoding with a separate target and draft KV), each child cache contributes its own group(s) so that heterogeneous buffer shapes (e.g., 61-layer MLA target vs. 1-layer Eagle draft) are registered as independent NIXL groups.
Quantized caches (values + scales): to_memory_groups() 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.
-
Parameters:
-
- name (str)
- kv_cache (PagedKVCacheManager)
-
Return type:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!