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

ReplicatedKVCacheMemory

ReplicatedKVCacheMemory

class max.nn.kv_cache.ReplicatedKVCacheMemory(buffer, peers)

source

Bases: KVCacheMemory

A replicated KV cache unit (rank-0 shard plus its TP peers).

All shards hold identical data (MLA); D2H reads from buffer (rank-0) and H2D broadcasts back to buffer and every entry in peers. Each buffer has shape [num_pages, bytes_per_page] with dtype uint8.

Parameters:

all_buffers

property all_buffers: list[Buffer]

source

Returns the rank-0 shard buffer followed by every peer buffer.

For replicated caches (MLA) all shards hold identical data; the inherited KVCacheMemory.copy_block_to() zips these against the destination’s shards (strict=True), so each shard is fanned out with an independent point-to-point copy (no broadcast collective).

peers

peers: list[Buffer]

source

Was this page helpful?