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

KVConnector

KVConnector

class max.pipelines.kv_cache.connectors.KVConnector(*args, **kwargs)

source

Bases: Protocol

Protocol for KV cache connectors managing external (non-device) tiers.

The manager owns device tensors, block allocation, and device-side prefix cache. Connectors handle external tier operations (such as host memory) via load/offload methods.

All block hashes crossing this Protocol are in canonical bytes form: 8 big-endian bytes for ahash64, 32 bytes for full SHA-256 digests. The block hasher produces this canonical form directly, so callers pass the hashes through unchanged; a connector that needs a narrower wire encoding (such as dKV’s 64-bit key) validates and converts at its own boundary.

Required call ordering per inference step:

connector.load() # post this step’s onloads 2. connector.offload() # kick off this step’s offloads 3. [model executes]

No barrier orders any of it. Each call hands back a KVTransfer, and what depends on the copy waits on that: the manager defers committing an onloaded prefix and holds the request out of the batch, and a connector publishes an offloaded block only once its bytes are written. So the model in step 3 never reads KV that has not landed, and the host is free to build the next batch while a copy is in flight.

abandon_lookup()

abandon_lookup(lookup)

source

Hands back a start_lookup() the manager will not load from.

Returns at once and never raises into the caller. Safe on a handle that has not answered yet, that was already loaded from, or that found nothing. No-op by default, since a lookup that holds nothing has nothing to hand back.

Parameters:

lookup (KVLookup)

Return type:

None

defers_lookup()

defers_lookup(replica_idx=0)

source

Returns whether start_lookup() can return before it answers.

Only such a lookup holds a request out of the batch, so only it counts against the scheduler’s bound on unanswered lookups, and a connector that answers in place is asked in place about the gap a short device hit leaves. False by default, matching the default start_lookup().

Parameters:

replica_idx (int)

Return type:

bool

disk_byte_count

property disk_byte_count: ByteCount

source

Disk tier occupancy in bytes. Empty (0 of 0) if not applicable.

host_byte_count

property host_byte_count: ByteCount

source

Host tier occupancy in bytes. Empty (0 of 0) if not applicable.

leaves

property leaves: Mapping[str, KVCacheGroupId]

source

Returns a mapping from leaf_id to group_id serviced by the connector.

load()

load(block_ids, block_hashes, replica_idx=0, hint=None)

source

Loads exactly the blocks it is told to, leaf by leaf.

block_hashes[leaf_id][i] is read into block_ids[leaf_id][i], so the two are the same length on every leaf. Which hashes those are – the whole prefix for a full-attention leaf, a window’s worth for a windowed one – the caller settled from a lookup(). The slots below a leaf’s share are the null block, which the caller fills in; this never sees them.

Must follow a lookup() with the same replica_idx and hint, at most once, since a connector may be reading out of what that lookup found.

Parameters:

• block_ids (Mapping[str, Sequence[int]]) – Device block IDs to load into, per leaf. Leaf counts differ by design: a full-attention leaf takes one block per hash of the hit, a windowed leaf only its window’s worth.
• block_hashes (Mapping[str, Sequence[bytes]]) – The hashes to read, per leaf, in the order that leaf’s row holds them. Canonical bytes form (8 big-endian bytes for ahash64-family, 32 bytes for SHA-256).
• replica_idx (int) – DP replica whose device buffers receive the loaded blocks. The external tier itself is replica-agnostic (keyed by hash); this only selects the H2D destination.
• hint (bytes | None) – The request’s dkv_cache_hint as raw JSON bytes, or None when it carried none. Opaque to the manager: only the dKV connector reads it, to route blocks to the peer that holds them. It never affects correctness, since an unusable hint costs a cache miss and nothing else.

Returns:

A KVTransfer for the H2D copy, which the manager polls before reading the loaded KV. A load that moved nothing returns a CompletedTransfer.

Raises:

• KVLoadRefused – If it cannot move every block it was asked for. Loading less is not an option – see KVLoadRefused.
• KVLoadFailed – If posting the copy failed. A copy that fails after it was posted raises from the returned transfer instead.

Return type:

KVTransfer

load_lookup()

load_lookup(lookup, block_ids, block_hashes, replica_idx=0, hint=None)

source

Loads as load() does, reading out of what lookup found.

lookup is complete and was started with the same replica_idx and hint, and the hashes are a contiguous run of the ones it was asked about, starting wherever the device hit ends. Consumes the handle however it exits, so no abandon_lookup() is owed afterwards. By default the handle carries no state and this is a plain load().

Parameters:

• lookup (KVLookup)
• block_ids (Mapping[str, Sequence[int]])
• block_hashes (Mapping[str, Sequence[bytes]])
• replica_idx (int)
• hint (bytes | None)

Return type:

KVTransfer

lookup()

lookup(block_hashes, replica_idx=0, hint=None)

source

Reports which of block_hashes each leaf holds. Presence only.

What that presence is worth is the manager’s to decide, since it depends on how far each leaf’s attention reads back. The manager reconciles the masks with max._kv_core’s prefix-hit rules, the same rules it runs over the device pools, so a connector needs no notion of attention shape, window widths or null blocks.

Advisory, and the caller owes it nothing: it may look up and then not load. A connector that keeps state between the two calls – dKV holds a lease its load() reads out of – must bound that itself and stay correct when no load follows, or when two lookups run back to back.

Not a reservation either. A block reported here can be evicted before load() asks for it, which surfaces as KVLoadRefused.

Parameters:

• block_hashes (Sequence[bytes]) – Hashes to ask about, in prefix order and in canonical bytes form (see the class docstring).
• replica_idx (int) – DP replica asking. The external tier is replica-agnostic (keyed by hash); this only selects the client.
• hint (bytes | None) – As load(), and it must be the SAME hint, or a lookup would ask the co-located store about a peer-held prefix and report it absent.

Returns:

One mask per leaf, keyed as leaves keys them, each exactly len(block_hashes) long and positional. True means every TP shard of that leaf holds the block – not some shard, not a count. A weaker predicate pads null pages into a full-attention row while still claiming the whole prefix, which is silent wrong KV rather than a missed hit.

Return type:

Mapping[str, Sequence[bool]]

metrics

property metrics: KVCacheMetrics

source

Transfer metrics for this connector. Returns empty metrics by default.

name

property name: str

source

Connector name for logging/debugging.

offload()

offload(block_ids, block_hashes, replica_idx=0)

source

Offloads each leaf’s blocks under that leaf’s hashes.

block_ids[leaf_id][i] is written under the hash block_hashes[leaf_id][i]. Leaves may differ in length, since a recurrent leaf publishes only at its checkpoint boundaries.

Parameters:

• block_ids (Mapping[str, Sequence[int]]) – Device block IDs to offload, per leaf.
• block_hashes (Mapping[str, Sequence[bytes]]) – Hashes per leaf, in prefix order. Canonical bytes form (8 big-endian bytes for ahash64-family, 32 for SHA-256).
• replica_idx (int) – DP replica whose device buffers source the offloaded blocks. The external tier itself is replica-agnostic.

Returns:

A KVConnectorTransfer for the D2H copy; g0_blocks_per_leaf are the device source blocks the manager keeps pinned until it lands. An offload that moved nothing returns a CompletedTransfer and is pinned nowhere.

Return type:

KVConnectorTransfer

poll_transfers()

poll_transfers()

source

Let the connector reclaim what its settled transfers left behind.

Called every time the manager drains its in-flight transfers, before it polls them. A connector whose transfers settle themselves still has work no poll covers: a transfer dropped before it completed, and book-keeping a settled one recorded off to the side. Cheap, non-blocking, and never raises into the scheduler. No-op by default.

Return type:

None

probe()

probe(block_hashes, replica_idxs, hint=None)

source

Reports what each replica could load of block_hashes, for pricing.

One answer per entry of replica_idxs, each shaped as a lookup() answer, read without taking whatever a lookup takes. Called once per arriving request to weigh the replicas against each other, so a connector whose lookup leases or stages anything answers this from an index instead. An answer may cover fewer leaves than leaves names when one is representative, and the manager reconciles the leaves it covers. Answers may share one mapping, so a caller reads them and never changes them. None means the connector has no such answer, and the manager prices each replica with a lookup(). No-op by default.

Parameters:

• block_hashes (Sequence[bytes])
• replica_idxs (Sequence[int])
• hint (bytes | None)

Return type:

Sequence[Mapping[str, Sequence[bool]]] | None

reset_prefix_cache()

reset_prefix_cache()

source

Reset prefix cache. No-op by default.

Return type:

None

shutdown()

shutdown()

source

Clean shutdown of connector resources.

Return type:

None

start_lookup()

start_lookup(block_hashes, replica_idx=0, hint=None)

source

Starts a lookup() and returns at once.

Same arguments and same answer as lookup(), read off the handle once it completes. The manager holds the request out of the batch meanwhile, so a connector whose lookup is a round trip keeps it off the scheduler thread. A handle the manager does not load from goes back through abandon_lookup().

By default the lookup runs here and comes back already answered, which is right for a connector that answers from local memory.

Parameters:

• block_hashes (Sequence[bytes])
• replica_idx (int)
• hint (bytes | None)

Return type:

KVLookup

take_metrics()

take_metrics()

source

Reads and clears the per-batch transfer counters.

Return type:

KVCacheMetrics

touch()

touch(block_hashes, replica_idx=0)

source

Refresh the external tier’s recency for blocks served from device (G0).

Best-effort and fire-and-forget: returns immediately, processes asynchronously, ignores the result, and never raises into the caller. A block served from the on-device prefix cache issues no other external-tier traffic, so without this its external-tier LRU recency can freeze and the tier can evict a block that is still hot on device. There is no companion barrier; a missed touch costs at most a later refetch, never correctness. No-op by default.

Contract: pass the complete set in sequence order from the true root – the full sequence for a full-attention group, the full active window for a sliding-window group. Never a root-omitting slice: a partial touch reserves a later recency stamp and inverts eviction order (the omitted root ages below the touched subset and evicts first). Missing keys are tolerated, so it is always safe to pass the whole sequence.

Parameters:

• block_hashes (Sequence[bytes]) – Hashes of the device-served blocks, in canonical bytes form (8 big-endian bytes for ahash64-family, 32 bytes for SHA-256). Root-anchored and in sequence order (see the contract above).
• replica_idx (int) – DP replica that served the blocks. The external tier is replica-agnostic (keyed by hash); this only selects the client.

Return type:

None