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

PagedKVCacheManagerInterface

PagedKVCacheManagerInterface

class max.pipelines.kv_cache.PagedKVCacheManagerInterface

source

Bases: ABC

Interface for the KVCache manager.

This is largely a temporarily class to help phase in the JengaKVCacheManager. We want to delete the legacy PagedKVCacheManager and rename JengaKVCacheManager to PagedKVCacheManager soon.

alloc()

abstract alloc(ctx, *, defer_lookup=False)

source

Allocates blocks for a request.

defer_lookup says the caller holds the request out of the batch while the returned transfer is incomplete and calls alloc again once it completes. Only then may a connector lookup run off the calling thread, since the request then holds no pages until that next call. A caller that reads the request’s blocks right after this returns leaves it off.

Parameters:

• ctx (TextContext)
• defer_lookup (bool)

Return type:

KVTransfer

alloc_dummy()

abstract 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

awaits_lookup()

awaits_lookup(ctx)

source

Returns whether the request’s last alloc left a lookup out.

Read after a deferred alloc, so a caller counts only the lookups that actually hold a request out, re-asks included. A manager that never defers a lookup reads False.

Parameters:

ctx (TextContext)

Return type:

bool

block_count()

abstract 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 token multiple a context-encoding chunk ends on.

Zero when no group of this cache needs the alignment.

claim()

abstract 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()

abstract contains(ctx)

source

Returns whether the request is claimed on any replica.

Parameters:

ctx (TextContext)

Return type:

bool

disk_byte_count()

abstract disk_byte_count()

source

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

Return type:

ByteCount

effective_max_seq_length

abstract 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()

abstract get_device_buffer(replica_idx)

source

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

Parameters:

replica_idx (int)

Return type:

KVCacheBufferInterface

get_metrics_aggregated()

abstract get_metrics_aggregated()

source

Returns aggregated metrics across all replicas.

Return type:

KVCacheMetrics

get_prefix_cache_hit_counts()

abstract get_prefix_cache_hit_counts(ctx)

source

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

Parameters:

ctx (TextContext)

Return type:

list[PrefixCacheHits]

get_req_blocks_per_leaf()

abstract get_req_blocks_per_leaf(ctx)

source

Returns the block IDs the request holds, keyed by leaf id.

Covers the replica the request was claimed on, with one entry per leaf that KVCacheParamInterface.leaves() names. A manager whose leaves share one block space repeats the same row under every leaf.

Parameters:

ctx (TextContext)

Return type:

dict[str, list[int]]

hand_back_lookup()

hand_back_lookup(ctx)

source

Hands back the request’s lookup without loading from it.

For a request whose answered lookup would otherwise wait in a queue past its lease. Its next deferred alloc starts a fresh one. A manager may keep the lookup instead, for one that already gave up, a request handed back too often, or an answer still fresh enough to load. A manager that never defers a lookup has nothing to hand back.

Returns False when the manager kept an answer a later call may still hand back, and True when there is nothing left to hand back.

Parameters:

ctx (TextContext)

Return type:

bool

host_byte_count()

abstract host_byte_count()

source

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

Return type:

ByteCount

num_req_blocks()

abstract num_req_blocks(ctx)

source

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

Parameters:

ctx (TextContext)

Return type:

int

params

params: KVCacheParamInterface

source

pending_transfers_exist()

abstract pending_transfers_exist(replica_idx=0)

source

Returns whether any async KV transfer is in flight on the replica.

Parameters:

replica_idx (int)

Return type:

bool

poll_transfers()

abstract poll_transfers()

source

Drains completed async KV transfers (onloads and offloads).

Return type:

None

pressure_pct()

pressure_pct(replica_idx=0)

source

Returns how close the replica is to refusing work, in [0, 100].

Defaults to block occupancy. A manager whose blocks are coarser than what a request allocates overrides this.

Parameters:

replica_idx (int)

Return type:

float

release()

abstract 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()

abstract reset_prefix_cache()

source

Resets the device prefix caches and every connector’s tiers.

Return type:

None

runtime_inputs()

abstract 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()

abstract 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()

abstract shutdown()

source

Releases the KV connector’s external resources.

Return type:

None

start_lookup()

start_lookup(ctx)

source

Starts the lookup a deferred alloc would start, taking no pages.

For a request that has to wait its turn anyway, so the lookup runs while it waits. Asked only when starts_lookup reads True, so a manager that never defers a lookup is never asked.

Parameters:

ctx (TextContext)

Return type:

KVTransfer

starts_lookup()

starts_lookup(ctx)

source

Returns whether a deferred alloc would start a connector lookup.

Takes nothing, so a caller that bounds how many lookups it has out can hold the request back before one is started. A manager that never defers a lookup reads False.

Parameters:

ctx (TextContext)

Return type:

bool

step()

abstract step(ctx)

source

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

Parameters:

ctx (TextContext)

Return type:

None

take_metrics_aggregated()

abstract take_metrics_aggregated()

source

Reads and clears aggregated metrics across all replicas.

Return type:

KVCacheMetrics