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.
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
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)
Reports what each replica could load of block_hashes.
One lookup of the first full-attention leaf, for pricing a request
before it is bound: every replica of a tenant shares one store, so one
healthy replica’s client answers for every replica that can load. It
costs two RPCs, a read_blocks that pins and a release_blocks
before it returns, and the read bumps that leaf’s read tally on the
server. A replica with a degraded leaf client loads nothing until it
reconnects, so it answers all misses. A tree with no full-attention
leaf raises ValueError: a sliding leaf’s mask cannot price a run. Holds no pin past the call,
takes no hint, which would dial peers, and leaves every replica’s
outstanding lease alone.
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
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.