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

PagedKVCacheManager

`PagedKVCacheManager`

class max.pipelines.kv_cache.PagedKVCacheManager(params, session, total_num_pages, enable_runtime_checks=False, *, max_batch_size)

source

Bases: PagedKVCacheManagerInterface

Paged KVCache manager with data and tensor parallelism support.

import numpy as np
from max.driver import CPU
from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef
from max.nn.kv_cache import MHAKVCacheParams
from max.pipelines.context import TextContext, TokenBuffer
from max.pipelines.kv_cache import PagedKVCacheManager
from max.pipelines.modeling.types import RequestID

params = MHAKVCacheParams(
    dtype=DType.float32,
    n_kv_heads=8,
    head_dim=128,
    num_layers=2,
    page_size=128,
    devices=[DeviceRef.CPU()],
)
kv_manager = PagedKVCacheManager(
    params=params,
    session=InferenceSession(devices=[CPU()]),
    total_num_pages=8,
    max_batch_size=4,
)

def make_context() -> TextContext:
    tokens = np.array([1, 2, 3, 4], dtype=np.int64)
    return TextContext(
        request_id=RequestID(),
        max_length=1000,
        tokens=TokenBuffer(tokens),
    )

ctx1 = make_context()
ctx2 = make_context()

kv_manager.claim(ctx1)
kv_manager.claim(ctx2)

# Allocate blocks for these requests
kv_manager.alloc(ctx1)
kv_manager.alloc(ctx2)

# Get KVCache inputs to feed to graph
kv_cache_inputs = kv_manager.runtime_inputs([[ctx1, ctx2]])

# Run model...
# Update requests with newly generated tokens
ctx1.update(42)
ctx2.update(42)

# Commit newly written blocks to prefix cache
kv_manager.step(ctx1)
kv_manager.step(ctx2)

# Release metadata and KV blocks for these requests
kv_manager.release(ctx1)
kv_manager.release(ctx2)

Initialize the multi-device paged KV cache manager.

Parameters:

params (KVCacheParamInterface) – KV cache parameters. Pass MultiKVCacheParams for models with more than one KV cache.
session (InferenceSession) – The MAX Engine inference session.
total_num_pages (int) – The total number of pages to allocate.
max_batch_size (int) – Maximum runtime batch size used to preallocate per-replica runtime lookup-table/cache-length row capacity.
enable_runtime_checks (bool) – Whether to enable runtime checks.

`alloc()`

alloc(ctx, *, defer_lookup=False)

source

Allocates blocks for a request.

When prefix caching is enabled, some of the allocated blocks may be retrieved from the prefix cache and the context’s active token window is advanced accordingly.

Parameters:

ctx (TextContext) – The text generation context for the request. The request must already be assigned to a replica via claim.
defer_lookup (bool) – Unused, since this manager has no connector to ask.

Returns:

this manager has no external tier to onload from.

Return type:

An already-complete CompletedTransfer

Raises:

InsufficientBlocksError – If there are insufficient free blocks to
satisfy the allocation. The request is left as it was before the –
call, so the caller can release or retry it. –

`alloc_dummy()`

alloc_dummy(ctx, replica_idx=0)

source

Claims a dummy request and maps it to the replica’s null block.

Parameters:

ctx (TextContext)
replica_idx (int)

Return type:

None

`block_count()`

block_count(replica_idx=0)

source

Returns the device KV cache block occupancy for the given replica.

Parameters:: replica_idx (int)
Return type:: BlockCount

`claim()`

claim(ctx, replica_idx=0)

source

Pins a request to one replica, which owns it until it is released.

Parameters:

ctx (TextContext)
replica_idx (int)

Return type:

None

`contains()`

contains(ctx)

source

Returns whether the request is claimed on any replica.

Parameters:: ctx (TextContext)
Return type:: bool

`disk_byte_count()`

disk_byte_count()

source

Returns an empty disk tier; this manager has none.

Return type:: ByteCount

`effective_max_seq_length`

property effective_max_seq_length: int | None

source

Returns the effective maximum sequence length that can be served by the block manager.

`get_device_buffer()`

get_device_buffer(replica_idx)

source

Returns the replica’s KV buffer (single leaf or tree).

HACK: this exists only for the transfer engine; callers flatten via KVCacheBufferInterface.all_buffers.

Parameters:: replica_idx (int)
Return type:: KVCacheBufferInterface

`get_metrics_aggregated()`

get_metrics_aggregated()

source

Returns aggregated metrics across all replicas.

Return type:: KVCacheMetrics

`get_prefix_cache_hit_counts()`

get_prefix_cache_hit_counts(ctx)

source

Counts each replica’s contiguous cached prefix for a request.

Computes the request’s block hashes once and counts every replica’s device hits, probing the shared external tiers once for all of them, without claiming the request or mutating any per-request state. Intended for prefix-aware data-parallel routing: callers can compare replicas’ hit num_blocks (across the device, host, and disk tiers) before deciding which replica should serve the request.

Parameters:: ctx (TextContext) – The request context to count cached prefix blocks for.
Returns:: One PrefixCacheHits per replica, indexed by replica.
Return type:: list[PrefixCacheHits]

`get_req_blocks_per_leaf()`

get_req_blocks_per_leaf(ctx)

source

Returns the request’s block IDs, repeated under every leaf.

Every leaf of this cache shares one block space.

Parameters:: ctx (TextContext)
Return type:: dict[str, list[int]]

`host_byte_count()`

host_byte_count()

source

Returns an empty host tier; this manager has none.

Return type:: ByteCount

`num_req_blocks()`

num_req_blocks(ctx)

source

Returns how many pages long the request is on its claimed replica.

Parameters:: ctx (TextContext)
Return type:: int

`pending_transfers_exist()`

pending_transfers_exist(replica_idx=0)

source

Returns False; this manager posts no async transfers.

Parameters:: replica_idx (int)
Return type:: bool

`poll_transfers()`

poll_transfers()

source

No-op; this manager posts no async transfers.

Return type:: None

`release()`

release(ctx)

source

Releases the blocks the request holds on the replica it was claimed on.

Parameters:: ctx (TextContext)
Return type:: None

`reset_prefix_cache()`

reset_prefix_cache()

source

Resets the device prefix caches.

Return type:: None

`runtime_inputs()`

runtime_inputs(batches, *, max_cache_length=None, batch_characteristics=None)

source

Gets the graph inputs for per-replica batches of requests.

Returns the KVCacheInputs pytree (a tuple of per-device leaves, or a dict of named subtrees for multi-cache models) whose leaves each hold one (DP replica, TP shard) device’s inputs.

This method will raise a RuntimeError if any request has insufficient blocks already allocated to it.

Parameters:

batches (Sequence[Sequence[TextContext]]) – Per-replica batches of requests
max_cache_length (int | None) – Optional explicit max cache length to size LUT views. If not provided, uses request-derived runtime length.
batch_characteristics (BatchCharacteristics | None) – Optional upper-bound batch shape applied uniformly across every replica when preparing attention dispatch metadata. When provided (such as graph-capture replay, where every DP replica must run the identical captured graph), the dispatch key is resolved once from these aligned values; the real per-replica values must not exceed them. When None, each replica prepares metadata from its own real values (which may differ per replica).

Return type:

KVCacheInputs[Buffer, Buffer]

`runtime_inputs_for_leaf()`

runtime_inputs_for_leaf(batches, *, max_cache_length=None, batch_characteristics=None)

source

Returns runtime_inputs() narrowed to a single leaf cache.

Convenience wrapper for single-cache (non-tree) models: it asserts the result is a per-device shard tuple and returns it. Raises AssertionError for multi-cache (dict) models.

Parameters:

batches (Sequence[Sequence[TextContext]])
max_cache_length (int | None)
batch_characteristics (BatchCharacteristics | None)

Return type:

tuple[KVCacheInputsPerDevice[Buffer, Buffer], …]

`shutdown()`

shutdown()

source

No-op; this manager holds no external resources.

Return type:: None

`step()`

step(ctx)

source

Commits the request’s newly written tokens into the prefix cache.

Parameters:: ctx (TextContext)
Return type:: None

`take_metrics_aggregated()`

take_metrics_aggregated()

source

Reads and clears aggregated metrics across all replicas.

Return type:: KVCacheMetrics

PagedKVCacheManager​

alloc()​

alloc_dummy()​

block_count()​

claim()​

contains()​

disk_byte_count()​

effective_max_seq_length​

get_device_buffer()​

get_metrics_aggregated()​

get_prefix_cache_hit_counts()​

get_req_blocks_per_leaf()​

host_byte_count()​

num_req_blocks()​

pending_transfers_exist()​

poll_transfers()​

release()​

reset_prefix_cache()​

runtime_inputs()​

runtime_inputs_for_leaf()​

shutdown()​

step()​

take_metrics_aggregated()​

`PagedKVCacheManager`

`alloc()`

`alloc_dummy()`

`block_count()`

`claim()`

`contains()`

`disk_byte_count()`

`effective_max_seq_length`

`get_device_buffer()`

`get_metrics_aggregated()`

`get_prefix_cache_hit_counts()`

`get_req_blocks_per_leaf()`

`host_byte_count()`

`num_req_blocks()`

`pending_transfers_exist()`

`poll_transfers()`

`release()`

`reset_prefix_cache()`

`runtime_inputs()`

`runtime_inputs_for_leaf()`

`shutdown()`

`step()`

`take_metrics_aggregated()`