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
DKVConnector
DKVConnector
class max.pipelines.kv_cache.connectors.dkv.DKVConnector(leaves, replica_kv_memory, local_block_store_endpoint, devices, params)
Bases: KVConnector
KVConnector backed by the dkv_connector Rust client.
A single instance serves every DP replica. The underlying Rust client is
inherently per-endpoint (its load/offload reference block ids into
one registered device-buffer set and carry no replica/group key), so this
shim owns one client per leaf per DP replica in self._clients and routes
each call to the matching leaf client for the request’s processing replica.
Under backend dedup there is one client per DP replica. Each client
registers its replica’s FULL TP GPU set (so MLA keeps
device_buffers.len() == tp and NVLink broadcast stays engaged), and every
DP replica of a tenant handshakes the same per-tenant store identity
(kv_shard_id / replica_id zeroed), so the server keys ONE
region-sharded store per tenant_id. load() / offload()
therefore issue exactly one call, to the processing replica’s client. The
MHA/GQA-vs-MLA distinction is carried by the per-block BlockKey
tp_shard_id the Rust client builds, not by the store key: identical-KV
shards dedup, distinct head shards co-reside.
Constructs and admits one dKV Rust client per DP replica.
-
Parameters:
-
- leaves (Mapping[str, KVCacheGroupId]) – Mapping from leaf identity to attention group.
- replica_kv_memory (Sequence[Mapping[str, KVCacheMemory]]) – Per-replica offload-ready KV memory by leaf.
- local_block_store_endpoint (str) – Co-located dKV control-plane endpoint.
- devices (Sequence[Device]) – The pipeline’s flat, ordered device list across replicas.
- params (KVCacheParamInterface) – KV-cache parameters, folded into the tenant store’s layout
kv_config_hash.
-
Raises:
-
ValueError – If either operator-injected variable is missing —
MODULAR_NIXL_TRANSFER_BACKEND(dKV auto-selects no transport) orMODULAR_DKV_TENANT_ID(dKV has no default/legacy single-tenant path, so it fails model load rather than silently keying an unfenced shared store). Both are checked in one pass, so a pod missing both is told about both.
abandon_lookup()
abandon_lookup(lookup)
Hands back every leaf’s lease in lookup without waiting.
A no-op on an empty lookup and on a handle already loaded from. Guarded per leaf so one failure strands nothing after it. A lookup answered in place hands back the prepared slot lease only while it is still that lookup’s, so it never cancels a later request’s.
-
Parameters:
-
lookup (KVLookup)
-
Return type:
-
None
defers_lookup()
defers_lookup(replica_idx=0)
Returns whether replica_idx’s extension has the async entries.
Without them start_lookup() answers through the blocking
lookup() before it returns.
leaves
property leaves: Mapping[str, KVCacheGroupId]
Returns a mapping from leaf_id to group_id serviced by the connector.
load()
load(block_ids, block_hashes, replica_idx=0, hint=None)
Moves each leaf’s share out of the lease its lookup() took.
One load_prepared per leaf, reading out of the locations that
leaf’s own lookup already found. Which hashes each leaf gets – the
whole prefix for a full leaf, the window for a windowed one – the
cache manager settled from the lookup’s masks; this only has to find
where they sit in the leased candidate.
hint is ignored: it rode in on the lease, and honouring a different
one here would mean re-reading, which is what the lease exists to
avoid. The manager passes the same one either way.
Consumes the lease however this exits, raising included.
-
Returns:
-
A
_DkvBatchTransferover the posted READs. It polls complete once every leaf’s H2D has landed, so the manager pins the destination blocks and the scheduler holds the request out of the batch until then, rather than ordering the copies on a pre-forward barrier that stalls batch construction. -
Raises:
-
- ValueError – If the rows do not match the leaf tree one block per hash, if no lookup leased this replica, or if the rows disagree about where they end.
- KVLoadRefused – If a leaf delivered less than its own lease promised.
- KVLoadFailed – If posting a leaf’s copy failed, or a copy that a short delivery drained did. Every leaf that posted has been drained by then.
-
Parameters:
-
Return type:
-
KVTransfer
load_lookup()
load_lookup(lookup, block_ids, block_hashes, replica_idx=0, hint=None)
Moves each leaf’s share out of the lease lookup holds.
As load(), except that the lease comes from the handle its
start_lookup() returned rather than from the client’s prepared
slot. hint is ignored for the same reason. Consumes the handle
however this exits, raising included. Any other handle, which is
what an extension without the async entries answers with, goes to
load() and its prepared slot lease.
-
Raises:
-
- ValueError – As
load(), or iflookupis a handle this connector leased on another replica. - KVLoadRefused – As
load(), or if a later lookup on the replica superseded the prepared slot leaselookuptook. - KVLoadFailed – As
load().
- ValueError – As
-
Parameters:
-
Return type:
-
KVTransfer
lookup()
lookup(block_hashes, replica_idx=0, hint=None)
Reports which of block_hashes dkv holds, per leaf, under a lease.
One leased read per leaf, because the lease has to live on the client
that will post the transfer. That is what lets load() re-read
nothing, so a leaf falling short there is a fault rather than a race.
The lease is ours, not the caller’s: the protocol has no release. A
lookup no load follows leaves one candidate’s pins outstanding until
the next lookup abandons them, or the client’s plan_ttl (10 s)
expires them. Bounded, and never a correctness problem.
hint routes the read to the peer that may hold the blocks;
load() inherits it from the lease. See _usable_hint().
metrics
property metrics: KVCacheMetrics
Transfer metrics for this connector. Returns empty metrics by default.
name
property name: str
Connector name for logging/debugging.
offload()
offload(block_ids, block_hashes, replica_idx=0)
Offloads replica_idx’s device blocks to the dkv service by hash.
Each hash follows the same 8-or-32 byte contract as load()
(truncated to its first 8 bytes at the dkv boundary; see
_to_dkv_u64()).
The dKV store dedups by composite key (tp_shard_id, group, seq_hash) and does not chain blocks under a parent, so the Rust
client builds the keys (and the NUMA striping plan) from the hashes
alone.
Routes to the processing replica’s per-leaf clients, one call per leaf.
-
Returns:
-
A
_DkvBatchTransferover the posted WRITEs, reportingblock_idsas the device sources it pins. dKV marks a block readable only when its write settles, so the manager holds the source out of the eviction path until this reads complete rather than letting a D2H drain into a page that has been reused. -
Parameters:
-
Return type:
-
KVConnectorTransfer
poll_transfers()
poll_transfers()
Drains what each leaf’s settled transfers left behind.
One call per leaf client, because each owns its own connector: a transfer dropped before it settled defers its reader-pin release onto that client’s orphan queue, which nothing else in a MAX worker drains, and a settled one records its throughput sample off to the side for a live caller to land.
-
Return type:
-
None
probe()
probe(block_hashes, replica_idxs, hint=None)
Reports what each replica could load of block_hashes, for pricing.
One presence call per shard layout in the tree, for pricing a
request before it is bound. A client keys every group it is asked
about with its own layout, so each layout’s leaves are asked on a
client of that layout, which is one call for a uniform tree. Every
replica of a tenant shares one store, so one healthy replica’s
clients answer for every replica that can load, and a replica with
a degraded leaf client loads nothing until it reconnects, so it
answers all misses. Those answers share two masks mappings, one for
every replica that can load and one for every replica that cannot,
so a caller reads them and never changes them in place. Presence
pins nothing, stages no disk block and bumps no read tally, all
three of which a leased lookup does, so it leaves every
request’s outstanding lease alone.
Without presence on any layout’s client, its server or a hinted
peer, prices with one client lookup on the first full-attention
leaf, which pins and releases inline and stages disk blocks. It
takes no hint, so pricing never dials a peer on the scheduler
thread. None only for a tree without a full-attention leaf.
reset_prefix_cache()
reset_prefix_cache()
Reset prefix cache. No-op by default.
-
Return type:
-
None
shutdown()
shutdown()
Clean shutdown of connector resources.
-
Return type:
-
None
start_lookup()
start_lookup(block_hashes, replica_idx=0, hint=None)
Starts one leased read per leaf and returns before any answers.
The reads run concurrently because each leaf has its own client, and
the handle rather than a client slot carries their leases, so one
request’s lookup never cancels another’s. load_lookup() reads
out of the handle and abandon_lookup() hands it back, and a
handle the manager drops does the same on its own. hint routes
the reads as in lookup().
An extension without the async entries answers through the blocking
lookup() here instead. That lease lives in the prepared slot,
so a later lookup on the replica supersedes it and the request it
belonged to is then served as a miss.
A leaf whose read fails to start hands back the reads already started and answers a miss, the way a lookup that fails while it is polled does.
take_metrics()
take_metrics()
Reads and clears the per-batch transfer counters.
-
Return type:
touch()
touch(block_hashes, replica_idx=0)
Refreshes replica_idx’s dkv recency for blocks it served.
A block served from MAX’s on-device (G0) prefix cache issues no other
dkv traffic, so its dkv LRU recency can freeze and dkv can evict it
while it is still hot on device. This forwards the served blocks’
hashes to the Rust client’s touch, which the server treats as an
access that bumps recency.
Touch contract: the CALLER passes the full root-anchored sequence, and
this narrows it per leaf before it goes on the wire.
RegionLru::touch forward-touches exactly the keys it is handed and
states that obligation (region_lru.rs, GroupKind::Sliding), so:
- a full-attention leaf takes the whole sequence, which the server walks in REVERSE so the shared prefix ends most-recently-used, and
- a sliding leaf takes only its active window, which the server walks FORWARD so the window ends most-recently-used and slid-out blocks age toward eviction.
Handing a sliding leaf the whole chain would bump every slid-out block
back to MRU on each admission and erase the pressure
GROUP_RECENCY_SLIDING_WINDOW exists to create – roughly 780 blocks
touched to protect 8 for gemma-4 at 100k tokens with a 1k window. The
offload path leaves capacity to the server precisely because that
pressure exists, so the two have to agree.
Each block_hashes element follows the same 8-or-32 byte contract as
load() (truncated to its first 8 bytes at the dkv boundary; see
_to_dkv_u64()). Best-effort and fire-and-forget: the Rust client
queues the touch RPC behind the settles queued before it and returns
immediately, so this never blocks the caller and a missed touch costs
at most a later refetch, never correctness. A no-op when
MODULAR_DKV_DISABLE_G0_TOUCH is set (the kill-switch, read once at
construction).
Not chunked. shard_keys mints one key per TP shard, so a long
sequence on a full leaf is a large request (TP=8 over 780 blocks is
~6k keys). Touches run in the order they are sent, so chunks would
land in order, but splitting here costs a Python call and a queued
RPC per chunk; chunking belongs in Toucher::send, which queues the
whole walk (SERVOPT-1615). Windowing the sliding leaves removes the
pathological case (a 1k window is 8 blocks however long the sequence
is).
Routes to the processing replica’s per-leaf clients, one call per leaf.