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

JengaKVCacheManager

JengaKVCacheManager

class max.pipelines.kv_cache.JengaKVCacheManager(*, params, leaf_infos, pools, kv_buffers, replica_kv_memory, max_batch_size, max_seq_len, max_num_input_tokens=None, connector=None, groups, slabs=())

source

Bases: JengaBlockManager, PagedKVCacheManagerInterface

Paged KV cache manager backed by a single fungible memory slab.

The slab is divided up based on the Jenga two-level huge/little page allocation strategy to allow multipple KV types to share the same physical memory.

Assigns blocks out of pools, one per replica.

groups decides which blocks a request holds; leaves says what each one costs and how the graph reaches it. block_size is in tokens.

Parameters:

• params (KVCacheParamInterface)
• leaf_infos (Mapping[str, KVLeafInfo])
• pools (Sequence[JengaBlockPool])
• kv_buffers (Sequence[KVCacheBufferInterface])
• replica_kv_memory (Sequence[Mapping[str, KVCacheMemory]])
• max_batch_size (int)
• max_seq_len (int)
• max_num_input_tokens (int | None)
• connector (KVConnector | None)
• groups (Mapping[str, KVGroupCoordinatorInterface])
• slabs (Sequence[Buffer])

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

chunk_alignment_tokens

property chunk_alignment_tokens: int

source

Returns the page size if any group needs the alignment, else zero.

A recurrent checkpoint fires only on an exact page boundary. With prefix caching off there are no checkpoints.

create()

classmethod create(*, params, available_bytes, max_batch_size, max_seq_len, max_num_input_tokens=None)

source

Creates a JengaKVCacheManager.

available_bytes is the KV budget across all devices from memory estimation (same contract as compute_num_device_blocks). Each device slab is sized from available_bytes // len(params.devices).

Parameters:

• params (KVCacheParamInterface)
• available_bytes (int)
• max_batch_size (int)
• max_seq_len (int)
• max_num_input_tokens (int | None)

Return type:

JengaKVCacheManager

disk_byte_count()

disk_byte_count()

source

Returns the disk KV tier occupancy in bytes, shared by every replica.

Return type:

ByteCount

get_device_buffer()

get_device_buffer(replica_idx)

source

Returns the device buffer for the given replica.

Parameters:

replica_idx (int)

Return type:

KVCacheBufferInterface

get_metrics_aggregated()

get_metrics_aggregated()

source

Returns aggregated metrics across all replicas.

Return type:

KVCacheMetrics

host_byte_count()

host_byte_count()

source

Returns the host KV tier occupancy in bytes, shared by every replica.

Return type:

ByteCount

max_seq_len_fitting_in_cache()

classmethod max_seq_len_fitting_in_cache(params, available_bytes)

source

Returns the longest request a slab of available_bytes holds one of.

Plans the geometry create() allocates for the same budget, so the answer is a length create accepts. Net of the speculative-decode slack a request may occupy past its cap.

Returns:

The longest sequence, or None when no length exhausts the slab.

Parameters:

• params (KVCacheParamInterface)
• available_bytes (int)

Return type:

int | None

memory_size()

classmethod memory_size(params, available_bytes, max_batch_size, max_seq_len)

source

Returns the slab bytes a full batch at max_seq_len needs.

Every leaf, a state or scratch leaf included, is priced in the whole huge blocks an empty pool would carve for it, the same test the pool admits a request by. The answer is checked against the geometry create() plans for it, since a smaller budget can choose a different huge block, and is never more than available_bytes.

Parameters:

• params (KVCacheParamInterface) – The cache the slab holds.
• available_bytes (int) – The KV budget across all devices.
• max_batch_size (int) – Requests held at once.
• max_seq_len (int) – The longest request, before speculative slack.

Returns:

The bytes across all devices.

Raises:

RuntimeError – If available_bytes cannot hold a single block.

Return type:

int

pressure_pct()

pressure_pct(replica_idx=0)

source

Returns the fullest leaf’s page occupancy, in [0, 100].

Reads 100 once no leaf can allocate another page.

Parameters:

replica_idx (int)

Return type:

float

runtime_inputs()

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

source

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

Parameters:

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

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.

Parameters:

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

Return type:

tuple[KVCacheInputsPerDevice[Buffer, Buffer], …]

shutdown()

shutdown()

source

Releases the connector’s external resources.

Drains in-flight host/disk transfers and frees the shared pinned host buffer; for the tiered connector this also removes its offload directory. One connector backs every replica, so this shuts it down once. A no-op for the null connector.

Return type:

None

step()

step(ctx)

source

Checkpoints the state onto its successor, then records the forward.

The checkpoint runs before super().step so the copy reads a block the commit has not published and advance has not freed, and on the stream the forward wrote it on.

Checkpoints are skipped when prefix caching is off, since nothing can publish them.

Parameters:

ctx (TextContext)

Return type:

None

take_metrics_aggregated()

take_metrics_aggregated()

source

Reads and clears aggregated metrics across all replicas.

Return type:

KVCacheMetrics