Repository navigation
Conversation
Add worker-side computation of KV cache checksums that do not depend on
the KV connector, the TP or PP sizes, or the KV cache memory layout. The
scheduler side, which asks for checksums and compares them, follows in
the next commit. Enable it with kv_transfer_config.enable_kv_checksum.
Design
------
- Connector-agnostic. Workers compute the checksums from the KV cache
itself, the scheduler sums and compares them, and connectors only
carry opaque bytes from producer to consumer. A connector needs no
checksum code of its own, and every connector is verified by the
same definition. Computing them inside each connector was rejected:
every connector would reimplement it, and the values would depend on
the connector's view of TP and the cache layout.
- One value per (KV cache group, block position), independent of TP,
PP and memory layout, so a producer and a consumer with different
TP/PP sizes or attention backends compare the same numbers.
- No host sync on the engine step. Results go to host memory through a
side copy stream and an event; see below.
What a checksum is
------------------
A checksum covers one block position of one KV cache group, over the
unsharded model: all layers of the group and all KV heads.
1. A slice is the KV of one layer and one KV head in the block, that is,
the K and V of the block's tokens for that head.
2. Each slice gets its own checksum. Lay the slice out token by token,
split it into little-endian 8-byte pieces x_0, x_1, ..., and compute
s = sum_j x_j * (2 * j + 1) mod 2**64
3. The block's checksum adds up the slices of every layer and head, each
multiplied by its own odd factor:
c = sum_{layer, head} m * s mod 2**64
m = 2 * (layer * H + head) + 1
H = 2**16 (MAX_KV_HEADS_PER_LAYER), a fixed bound on any layer's KV
heads; groups with more are skipped. Layers are numbered by the layer
number in their names, and heads globally across TP. m therefore
depends on neither the model's head counts nor its TP or PP split,
so every rank and PP stage of the producer and the consumer agrees
on it.
Example: 2 layers, 2 KV heads, 2 tokens per block, 16-byte cells (one
token's K and V for one head). A block has 4 slices of 2 tokens x 2
pieces each:
slice (layer, head) m s
(0, 0) 1 1*x0 + 3*x1 + 5*x2 + 7*x3
(0, 1) 3 ... (token 0: x0 x1, token 1: x2 x3)
(1, 0) 131073 ...
(1, 1) 131075 ...
c = 1*s(0,0) + 3*s(0,1) + 131073*s(1,0) + 131075*s(1,1) mod 2**64
Why this definition:
- Factors and weights are odd, hence invertible mod 2**64. A change to
any single piece changes c by an odd multiple of the change, so it is
always detected.
- Within a slice, distinct weights catch tokens in the wrong place.
Across slices, distinct factors catch swapped heads or layers.
- Random corruption of several pieces goes undetected with probability
about 2**-64.
- It is not a cryptographic hash and does not defend against deliberate
tampering.
Each worker only computes the slices it holds, and they add up to c:
- TP splits KV heads, never a head, so every slice lies whole on one TP
rank. Under TP2 in the example, rank 0 holds the head 0 slices and
rank 1 holds the head 1 slices.
- PP splits layers. Under PP2, stage 0 holds the layer 0 slices and
stage 1 holds the layer 1 slices.
- Each slice is laid out by token within its head, whatever the memory
layout. Head-major or token-major (NHD) caches, and kernel-block
splits, give the same s.
Granularity
-----------
There is one checksum per (KV cache group, block position). Every other
axis (layers, KV heads, TP ranks, PP stages) is summed into it, so its
size does not grow with the model or the parallelism.
The two remaining axes cannot be merged:
- Block positions. A consumer with a local prefix-cache hit loads only
a suffix of the request's blocks, starting wherever its own cache
ends. The producer cannot know where that is, so it must provide one
value per block, and the consumer compares only the blocks it loaded.
Coarser units, say one value per k blocks, fail whenever the hit does
not end on a multiple of k.
- KV cache groups. Each group has its own block table and covers
different tokens. For example, a sliding-window group keeps only the
window. The same block position therefore holds unrelated data in
different groups.
Cost: 8 bytes per block per group, for example 16 KB for a 32k-token
prompt with block size 16 and one group. A mismatch is localized to a
(group, block position), not to a layer or a rank.
Replication
-----------
When several TP ranks hold the same KV heads, those copies must not be
summed twice. A partial checksum is therefore keyed by the
(replicas, replica_index) of its layers:
- Specs that declare max_tp_shards (MLA) give the replica count through
layer_tp_replicas, made public for this.
- Regular attention shards KV heads across TP and replicates one head on
consecutive ranks once TP exceeds them. More than one head per rank
therefore means no replicas. One head per rank is ambiguous, and is
resolved with the model's KV head count only when every layer shares
it: no per-layer override, no separate global head count, and a draft
model with the same count. Otherwise the group is skipped.
- A producer only needs replica 0. A consumer reports every replica, so
each copy can be checked on its own.
The last block of a request only covers its first num_tokens slots. This
excludes slots that later decode or draft tokens may write, so the
producer's value does not depend on when it is taken.
Skipped groups
--------------
Groups that cannot be checksummed are logged as a warning at startup and
skipped:
- layers without per-token KV blocks (Mamba state; covered later in
this series);
- compressed caches (tokens_per_state != 1), e.g. the compressed MLA
and indexer caches of DeepSeek V4.x. A group mixing them with
ordinary layers is skipped as a whole;
- KV heads packed into head slots (num_head_slots);
- host-resident (HiSparse) groups;
- caches that do not hold num_blocks blocks, whose cells or strides are
not whole 8-byte pieces, or whose kernel blocks do not tile the block
size;
- unknown head replication, or more than 2**16 KV heads per layer;
- layer names without a single layer number, or two caches with the same
layer number.
Each record lists its worker's skipped groups (skipped_groups), so the
scheduler can tell a partial sum from a group with no layers on that
worker. It also carries a fingerprint per group: the sum of a hash of
each of the worker's layers in it (name, spec type, block size, dtype).
Summed over PP stages it identifies the group's layers whatever the PP
layout, so a producer and a consumer can check they checksum the same
layers.
How it is computed
------------------
The scheduler sets the requests to checksum each step in
SchedulerOutput.kv_checksum_scheduled (KVChecksumScheduled), as
KVChecksumBlocks per request id:
- reqs_to_send are checksummed in post_forward of that step, once
the forward and the drafter have written the KV.
- reqs_to_recv are remembered, and checksummed in the step whose
transfer results report their load finished. Their kernels are queued
before any later forward can write those blocks. A request that ends,
or fails to load, is dropped.
Every request a worker processes appears in its record, with no groups
if none of them has layers on that worker, so the scheduler can count
which workers reported it.
KVChecksumWorker builds, for every layer of every transfer group, a
view of its [B, H, N, C] cache as 8-byte pieces, with each piece's
weight m * (2 * j + 1). KV-sharing aliases are counted once, under the
lowest layer number. Gathers are chunked (64 MB) to bound transient
device memory.
No host sync on the engine step
-------------------------------
This follows the model runner's sampled-token copy (AsyncOutput):
- Block ids reach the device through pinned memory.
- Partial sums accumulate on the device.
- In the same post_forward, the worker's own copy stream waits for the
main stream, queues a non-blocking copy to pinned host memory, and
records a blocking event.
- The device values stay referenced until the copy is done.
Because the copy is ordered when the values are produced, it does not
matter which thread later reads them. KVChecksumOutput.finalize, called
when the output is pickled for the scheduler or on first use, only waits
for the event. On devices without CUDA streams the copy is synchronous.
AsyncOutput itself is not reused: it exists only for steps that run a
forward and sample on the last PP rank. A consumer's load usually
finishes in a step without a forward, and non-last PP ranks hold KV but
do not sample, so reusing it would need a second path for those cases.
The cost of the separate path is one extra CUDA stream per worker.
The results ride KVConnectorOutput.kv_checksums, the per-worker channel
that KVOutputAggregator already merges across TP and PP workers.
KVChecksumOutput.aggregate concatenates the per-worker records, and
is_empty counts the field so PP>1 stages do not drop it.
On the v1 runner with a drafter, the drafter writes KV after the
connector step (defer_finalize). post_forward then leaves the sends,
and their copy, to finalize_sends, which finalize_kv_connector calls
after the drafter. Only the last PP rank defers: it is the only one
that runs the drafter and calls finalize_kv_connector.
kv_checksum_enabled disables KV checksums, with a warning, for context
parallelism (DCP/PCP > 1), and for the Ray executor with PP > 1 unless
VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1: its non-last stages' connector
outputs never reach the scheduler.
Cost
----
The checksums are computed only in the step where a request finishes
its prefill (producer) or its load (consumer). The PyTorch
implementation in this commit reads the gathered KV about five times per
layer (gather, mask, multiply, reduce) and launches a few kernels per
layer and chunk. For a 32k-token prompt on an 8B GQA model (about 4 GB
of KV) that is roughly 25 GB of device memory traffic, plus
kernel-launch time on the engine thread. The Triton kernel later in
this series reads each piece once and replaces this path on CUDA; its
commit has the measured numbers.
Tests
-----
tests/v1/kv_connector/unit/test_kv_checksum.py checks every result
against an independent NumPy implementation of the definition:
- head-major and token-major layouts, and kernel-block splits (the
value must not depend on the memory layout);
- TP2 shards and PP2 stages summing to the whole model's value (the
parallel split must not change it);
- GQA and MLA replicas, per replica (a copy is neither double-counted
nor left unchecked);
- swapped tokens and swapped layers changing the value (misplaced data
is detected);
- slots past num_tokens (later decode or draft writes do not change the
producer's value);
- KV-sharing counted once, and chunked gathers;
- every skipped group kind, next to a checksummed group;
- consumer lifecycle (finish, failure, abort);
- deferred sends covering KV written before finalize_sends;
- group fingerprints independent of the PP split;
- serialization and aggregation across workers;
- on CUDA, a finalize from another thread while the main stream is busy
(the copy is ordered after the checksums for any reader).
Validation
----------
End-to-end GPU validation (real backends and NIXL transfers, with
VLLM_GPU_SYNC_CHECK=error) used a fault-injecting connector that is not
part of this series; the results are in the pull request description.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: David Cheung <d7cheung@gmail.com>
…verification The scheduler side of KV checksums: which blocks the workers checksum, how the producer's checksums reach the consumer, and how the consumer compares them before a loaded request can run. With this commit NIXL pull mode verifies transferred KV end to end. Enable it with kv_transfer_config.enable_kv_checksum on both sides; kv_checksum_fail_closed turns failed checks into KV load failures. Design ------ - The scheduler is the only place that sees every worker's record, so it sums the TP and PP parts there, on both sides. The payload is one value per (group, block position) of the whole model, so the consumer needs no knowledge of the producer's TP/PP layout. - The connector only carries bytes. A connector declares how through kv_checksum_carrier; this commit adds the generic kv_transfer_params carrier, which needs no connector code beyond that declaration. - A failed check is a KV load failure. It takes the path connector-reported failures take, so kv_load_failure_policy decides what happens next and no new failure handling is added. - The producer checksums the whole prompt once, in the step that finishes it, with or without chunked prefill. Workers keep no state across steps. Per-chunk partial sums were rejected: they would need per-rank partial state in the scheduler, cleanup on preemption, and handling of the block spanning two chunks. The cost is that the whole prompt's checksum work lands in its last prefill step. Life of a request ----------------- Producer (P): 1. In the step that finishes the prompt's KV (scheduled tokens cross num_prompt_tokens), Scheduler._update_after_schedule calls KVChecksumScheduler.add_send. It puts the request's blocks, per transfer group, into SchedulerOutput.kv_checksum_scheduled .reqs_to_send for the workers (build_checksum_meta, which also resets the step's state, like the connectors' build_connector_meta). 2. At the top of update_from_output, before the loop that finishes requests, update_from_output collects every worker's record. 3. When the request finishes, Scheduler._free_request calls request_finished, which sums the workers' parts and adds them to the kv_transfer_params the connector returns, as "kv_checksums". The proxy forwards them to the consumer with the request. Consumer (D): 1. When a request starts an async load, schedule() calls add_recv. It decodes the expected checksums from the request's kv_transfer_params and puts the blocks the load writes into reqs_to_recv, in the same step, so workers know them before the load starts. 2. Workers finish loading in different steps; update_from_output keeps each worker's record until the aggregated finished_recving reports the request. 3. It then compares, after the connector's own failed_recving and invalid_block_ids are handled (loads already known to be bad are not compared), and before the request can be scheduled again. Which blocks ------------ - Producer: per transfer group, the blocks holding tokens [0, num_computed_tokens). - Consumer: the blocks the load writes, tokens [num_local_computed_tokens, num_computed_tokens), i.e. after a local prefix hit. - Both: leading null blocks (out of a sliding window) are left out, a group with a null block after a non-null one is left out entirely, and sliding-window groups keep their last cdiv(window, block_size) + 1 blocks, the clip both NIXL and Mooncake apply to what they transfer. A connector that carries checksums must load exactly these blocks. - Only positions both sides checksummed are compared, and the last block only if both sides hold the same number of tokens in it. - Block sizes are per group, from the group's spec. Verdicts -------- Each group is compared on its own: - Mismatch: some compared position differs. Logged as an error with the groups and block positions, even if other groups are unverified. - Unverified: the group cannot be compared: it is checksummed by one side only, its fingerprint differs (other layers, block size or dtype; e.g. different block sizes on P and D), a worker checksummed other blocks than scheduled, a replica is missing, or no position is in common. A load is also unverified as a whole when not every worker (TP * PP) reported it, or the producer sent no (valid) checksums. Unverified loads log a warning once per fixed reason, and the details per request at debug level. Group fingerprints come from the workers' records (see the previous commit), summed over PP stages, so producers and consumers with different PP layouts compare the same layers. The scheduler's own KVCacheConfig only lists the first worker's layers. Groups a worker cannot checksum by configuration (SSM state, packed head slots, ...; see the previous commit) are reported in skipped_groups. A group skipped on both sides is left out of the verdict, so hybrid models verify their attention groups under fail-closed. A group skipped on one side only is unverified. Failure semantics ----------------- - Fail-open (default): mismatched and unverified loads are only logged; the loaded blocks are then used and prefix-cached as usual. This lets checksums be rolled out, and observed, without failing requests. - Fail-closed (kv_checksum_fail_closed): mismatched and unverified loads fail. They go through _handle_failed_recving, so kv_load_failure_policy applies: "fail" (default) ends the request with an error, "recompute" sets num_computed_tokens to 0 and recomputes the whole prompt locally on the consumer. - The KV is never transferred again: the connectors clear do_remote_prefill on the first load. Recovery is per request, not per worker or block. Recomputing only from the first bad block is future work, building on group-aware load-failure recovery (RFC vllm-project#47812, PR vllm-project#48216); the comparison already knows the bad positions per group. - A connector that does not implement kv_checksum_carrier can carry no checksums, so no load could be verified: enabling KV checksums with it fails at startup when fail-closed, and disables them with a warning when fail-open. Replicas -------- The consumer sums its workers per replica key. Check k adds copy k % R of the layers replicated R times, so each check covers every layer once, every copy is in some check, and each check must equal the producer's value (which sums replica 0 only). Carrier ------- KVConnectorBase_V1.kv_checksum_carrier(request, receiving) says how a connector carries a request's checksums; the default, None, means the request is not checksummed (e.g. offloading connectors). - KVChecksumCarrier.KV_TRANSFER_PARAMS: generic code adds and reads the "kv_checksums" key; the connector needs no other change. - NixlPullConnector returns it for producer requests (do_remote_decode) and for regular loads. A bidirectional load (a load with do_remote_decode set) reads blocks whose source sent no checksums, so it returns None. - MultiConnector: for a load, the sub-connector chosen to load it; for a producer request, the first sub-connector with a carrier. The NIXL push and Mooncake commits add a connector-owned carrier. Wire format: base64 of msgpack (version, num_tokens, [(group_id, fingerprint, start, little-endian uint64 values)]). A producer whose workers did not all report a request (or reported other blocks) sends no checksums for it, with a warning. For a 128k-token prompt with one group and 16-token blocks this is about 87 KB of base64. Proxies must forward unknown kv_transfer_params keys. Rollout and limits ------------------ - Enable on producers before consumers; otherwise consumers see loads without checksums, unverified (failed when fail-closed). - A producer request preempted in its final prefill step and re-prefilled before that step's output arrives can mix the two steps' values; the consumer then reports a mismatch or unverified, never a false match. - NIXL's bidirectional transfer (a prefiller reloading a decoder's KV from an earlier turn) is not checksummed: its carrier is None, so such loads are neither verified nor counted as unverified. - A group skipped on both sides leaves the load verified by its other groups only. For example, groups with tokens_per_state > 1 (the compressed MLA and indexer caches of DeepSeek V4.x) are skipped, and the load passes on its remaining groups, also when fail-closed. - Not covered: SSM state (later in this series); DCP/PCP and the Ray executor with PP (disabled in the previous commit); NIXL push and Mooncake (next commits). Tests ----- tests/v1/kv_connector/unit/test_kv_checksum_scheduler.py runs a producer's and a consumer's real schedulers (NIXL pull) with real checksum workers, copying the producer's blocks into the consumer's: - correct loads verify for (TP, PP) pairs (1, 1)/(1, 1), (2, 1)/(1, 1), (1, 1)/(2, 1), (1, 1)/(4, 1) (replicated heads on the consumer), (1, 2)/(1, 1) and (1, 1)/(2, 2), each scheduler holding only its first worker's layers (no false mismatch across layouts); - with a local prefix hit only the loaded blocks are compared; - a corrupted byte (logged as a mismatch) and missing checksums: only logged when fail-open, an error with policy "fail", a local recompute with "recompute"; - a group skipped on both sides is ignored, on one side only is unverified; - a connector without a carrier is refused at startup when fail-closed and disables checksums when fail-open. Disabling the comparison makes the fail-closed corruption tests fail, and counting one PP stage's fingerprint makes the PP pairs fail. test_scheduler.py and test_async_scheduler.py set the new attribute on their hand-built schedulers. Validation ---------- End-to-end GPU results for NIXL pull, including a producer with TP2 and a consumer with TP1, are in the pull request description. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
In NIXL push mode the producer WRITEs the KV into the consumer's blocks. The consumer's kv_transfer_params are built before the producer finishes, and nothing returns to the consumer through them afterwards, so KV checksums travel with the transfer itself: on the notif of the WRITE that delivers the KV. This adds the connector-owned checksum carrier and implements it for NixlPushConnector. Connector-owned carrier ----------------------- KVChecksumCarrier.CONNECTOR: the connector moves the checksums in its own messages. The generic code and the connector exchange only bytes; the comparison is unchanged. - Producer: after request_finished, KVChecksumScheduler hands the summed checksums (encoded KVChecksumPayload bytes) to the new KVConnectorBase_V1.send_kv_checksums(request, checksums). - Consumer: the worker connector returns them in the new KVConnectorTransferResults.kv_checksums, keyed by request, by the time it reports the load in finished_recving. Both model runners pass them to KVChecksumWorker.post_forward, which adds them to its record (KVChecksumRecord.received), so they reach the scheduler on the existing per-worker channel. The scheduler takes the first copy for a pending load, before it compares the step's finished loads. - MultiConnector forwards send_kv_checksums to the sub-connector that carries the request's checksums, and merges kv_checksums of its sub-connectors' transfer results. A load whose checksums never arrive is unverified. For both carriers, checksums are only sent for requests that finished normally (stop or length), not aborted ones. NIXL push --------- - NixlBaseConnector.kv_checksum_carrier is shared by pull and push; each sets its carrier (_KV_CHECKSUM_CARRIER). - P scheduler: send_kv_checksums keeps the checksums of a request whose finished blocks are about to be pushed, and build_connector_meta ships them with the blocks (NixlConnectorMetadata .push_finished_kv_checksums). Checksums of requests with nothing to push (e.g. aborted) are dropped, so nothing is left to clean up. - P worker: start_load_kv hands them to the writer thread with the finished blocks, in the same queue item; the writer keeps them until the request is evicted. _completion_notif builds the WRITE's notif: KV_CHECKSUM_NOTIF_PREFIX + msgpack((remote_request_id, tp_size, checksums)), instead of "remote_request_id:tp_size" when there are checksums. - D worker: _get_new_notifs parses both forms with the same completion accounting, keeps the first checksums received for a request it is loading, and get_transfer_results reports them with the request's finished_recving. - NIXL_CONNECTOR_VERSION 14 -> 15 (new notif form), so mixed versions fail the handshake instead of misreading the notif. Without checksums (disabled, or not summed on the producer) the notif keeps the plain form, so the default path is unchanged. Design choices -------------- - Why the next step's metadata: only the P worker's writer thread can send NIXL messages, and the scheduler reaches workers only through the next SchedulerOutput. The finished blocks to push already take that path (request_finished records them, the next build_connector_meta ships them), and the checksums ride along in the same metadata and queue item, so they add no latency to the WRITE. - Why the completion notif: it is already the consumer's signal that the load is done, and NIXL delivers it after the WRITE's data. The checksums therefore cannot arrive after the consumer reports the load finished, so the consumer never waits for them. A separate notif would have to be ordered against, and survive the loss of, the completion notif. - Why every writing producer rank carries them: which producer ranks write to a given consumer rank depends on the TP ratio, PP and MLA. A single designated carrier would need a proof that it always writes, otherwise the checksums would not arrive and the load would be unverified. Every rank that writes sends the same bytes (summed in the producer's scheduler), and the consumer keeps the first copy, without merging. The cost is duplicated payloads when several producer ranks write to one consumer rank. - The NIXL version is bumped once for the whole series, which lands as one PR. The checksum payload inside the notif is opaque to NIXL and versioned on its own; later commits in this series extend it without changing the notif form. Known limitations ----------------- - NIXL's bidirectional transfer (a prefiller reloading a decoder's KV from an earlier turn) is not verified: the decoder computes no checksums for the KV it exports. Its carrier is None, so such loads are not checksummed rather than unverified. - If a producer's lease expires while its WRITE waits for a deferred handshake, the writer drops the checksums with the request and the late WRITE (a pre-existing hazard: its blocks may have been freed) goes out without them. The consumer reports the load unverified, which fails it under fail-closed, rather than a mismatch. - A consumer load that never completes (some notifs lost) keeps its received checksums, as it already keeps its recv metadata. - The notif grows by 8 bytes per block and group (about 64 KB for a 128k-token prompt with 16-token blocks), sent once per WRITE. PUSH_REG carries block-id lists of similar size. - A checksum notif that fails to decode is logged and dropped: like a lost notif, the consumer's load then never completes and the request waits until aborted. Tests ----- - test_kv_checksum_scheduler.py runs every end-to-end case for both NixlConnector (pull) and NixlPushConnector (push): the producer's push metadata carries the checksums, and the consumer's workers report them with the finished load. Ignoring received checksums makes the push cases fail. - test_nixl_push_connector.py: the P scheduler ships checksums with the finished blocks and drops those of an aborted request; start_load_kv queues them for the writer; a completion notif built by a producer worker, with and without checksums, completes the consumer's load and reports its checksums. Validation ---------- On GPU, end-to-end runs with a fault-injecting connector (not part of this series) met their expected results for push with TP1 -> TP1, producer TP2 -> consumer TP1 and producer TP1 -> consumer TP2: no false mismatch, every injected corruption detected, no host sync. These runs used Qwen/Qwen3-0.6B; no MLA model was run. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
NIXL transfers KV between producers and consumers with different block sizes (block_size_ratio), but KV checksums could not compare them: a piece's weight depended on its position within its block, so the checksums of two 16-token blocks did not add up to that of the 32-token block holding the same tokens. Such loads were unverified (the group fingerprint included the block size), which fails them all under kv_checksum_fail_closed. Definition ---------- Pieces are now numbered from the request's first token rather than the block's: piece p of the token at position t in the request is j = t * P + p (P pieces per token), weighing m * (2 * j + 1) as before. Block checksums then add up over consecutive blocks to the checksum of the tokens they hold, whatever the block size. Swapping two blocks also changes the sum of the chunk holding them, which it did not before. The worker adds the block's start to each weight: a block starting at token f sums x * w + f * x * (2 * P * m), with w the weights of a block at the request's start (_LayerView.weights, .offset_weights), so the precomputed weights stay per layer rather than per block position. The cost is a second weighted sum per row on the PyTorch path; the Triton kernel later in this series computes the weight in place instead. Comparison ---------- The producer only adds each group's block size to the payload; the consumer's scheduler does the regrouping. With different block sizes, it sums the smaller blocks' checksums into chunks of the larger block size on each side and compares chunks; a chunk is compared only where both sides cover it whole and hold the same number of tokens in it, blocks past the last token counting as zero. With equal block sizes a chunk is a block, as before. Mismatches are reported as the consumer's block positions, so with different block sizes a mismatch is localized to one larger block. The group fingerprint no longer includes the block size. Unverified: - block sizes that are not multiples of each other; - a consumer with smaller blocks whose first loaded block starts inside a producer block (a local prefix hit or window clip not aligned to the producer's blocks). That block's chunk also holds tokens the consumer did not load, so it cannot be compared. Skipping the chunk instead would leave the first loaded block unchecked while reporting the load verified. Adding the consumer's locally cached tokens to the chunk would report false mismatches: KV a consumer computed while decoding is in general not bit-identical to the producer's prefill of the same tokens. The payload version stays 1: the payload is new in this series, which lands as one PR. The worker and scheduler flows are unchanged: which blocks are checksummed, when, and how checksums are carried. Known limitations ----------------- - The unaligned case above is unverified, so it fails under fail-closed and is only logged under fail-open. It arises in push mode with larger producer blocks. NIXL itself mispairs blocks across different block sizes when the consumer has a partial prefix hit; upstream vllm-project#59746 (pull: pair blocks by the request's last token) and vllm-project#59747 (push: reject different P/D block sizes at handshake) are open. Once both land, a consumer's first loaded block can no longer start inside a producer block. Checksumming producer blocks at a finer granularity would make the case comparable; it is not done while those PRs are pending. Tests ----- - test_kv_checksum.py: the reference numbers pieces from the request's start; two 16-token blocks add up to the 32-token block with the same tokens, including a partial last block. Dropping the block-start term makes the definition tests fail. - test_kv_checksum_scheduler.py: producer and consumer block sizes 16/32 and 32/16, NIXL pull and push: a correct load verifies and a corrupted byte fails it under fail-closed. Comparing unsummed blocks makes the 16/32 cases fail. A consumer with 16-token blocks after a 16-token local prefix hit, loading from a producer with 32-token blocks (push), is unverified; without the alignment check it verifies without its first loaded block. Validation ---------- On GPU, NIXL pull with producer block size 16 and consumer block size 32, on this series' base, reported a mismatch for 4 of 4 requests in a run without injected faults: NIXL on that base corrupts this transfer (fixed upstream by vllm-project#58860, which is not in the base). With recompute, the consumer's outputs matched the reference. A clean run with vllm-project#58860 applied, and push with producer block size 32 and consumer block size 16, have not been run on GPU. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
Mooncake P2P verifies transferred KV with KV checksums. Its router sends a request to the producer and the consumer at the same time, and the producer's request_finished returns no kv_transfer_params, so nothing reaches the consumer through them. The only producer-to-consumer message after the KV is written is the MooncakeXferResponse that reports the transfer done, so the checksums ride on it (KVChecksumCarrier.CONNECTOR, added for NIXL push earlier in this series). The checksum code itself is unchanged; only the connector carries the bytes. - MooncakeConnector.kv_checksum_carrier: CONNECTOR for producer requests (do_remote_decode with a transfer_id) and for loads. - Producer scheduler: send_kv_checksums keeps the checksums of a request whose blocks are about to be sent, and build_connector_meta ships them with them (MooncakeConnectorMetadata.kv_checksums). Those of requests with nothing to send are dropped. As in NIXL push, the blocks already reach the workers through the next step's metadata, so the checksums add no latency. - Producer worker: record_send_reqs stores them on the request's SendBlockMeta before marking it ready, and send_kv_to_decode adds them, for the requests it sent successfully, to the response (MooncakeXferResponse.kv_checksums, by consumer request id). Error responses carry none. - Consumer worker: process_pulling_result queues the checksums of each request on the receiver loop before it finishes the load (with heterogeneous TP several producer workers send the same ones), and get_transfer_results drains the queue on the worker thread, without waiting on the receiver loop, keeping the first copy until it reports them with the request's finished_recving. - Timing: the producer sends only after send_meta.ready, which is set by the metadata that follows request_finished, so the checksums are always there when the response goes out. On the consumer, the checksums are queued before the load is marked finished, both in one synchronous step of the receiver loop, and get_transfer_results takes the finished set before draining the queue, so a finished load's checksums are always there. - Compatibility: the new response field is optional (omit_defaults), and msgspec ignores unknown fields, so mixed versions keep working without checksums (their loads are unverified). No protocol version bump is needed, unlike the NIXL notif form. Known limitations ----------------- - Only Mooncake P2P. The Mooncake store connector does not carry checksums. - Mooncake sends KV only for producer requests that finished with length (FINISHED_LENGTH_CAPPED); other requests have nothing sent, so no checksums either. - Like NIXL push, every producer worker's response carries the whole payload, so a consumer worker loading from several producer workers receives duplicate copies. Tests ----- - test_kv_checksum_scheduler.py runs every end-to-end case for MooncakeConnector too, with the real producer and consumer schedulers and checksum workers: the consumer's params are built as the router would, and the checksums the producer's scheduler ships in its next metadata are handed to the consumer's workers as the connector would report them. Not shipping them makes the Mooncake cases that expect checksums fail. - test_mooncake_connector.py covers each hop of the connector: the producer scheduler ships checksums with the finished blocks and drops an aborted request's; record_send_reqs attaches them; the producer's response carries them for its ok requests; the consumer worker's get_transfer_results reports them with the finished load. Validation ---------- Mooncake has not been run end to end on GPU with KV checksums. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
The PyTorch path of _LayerView.accumulate reads each block's KV about five times per layer (gather, mask, multiply, reduce) and launches a few kernels per layer and gather chunk. On CUDA with Triton, one kernel per layer now computes the same sums, reading each 8-byte piece once and computing the weights in place instead of reading a weight table. Kernel ------ kv_checksum/kernels.py: _layer_checksum_kernel runs one program per block row and chunk of 16K pieces, so a block of attention KV is usually one program while a multi-megabyte state row spreads over many SMs. A program walks its chunk in one flat loop over [kernel block, head, token, piece], in tiles of up to 1024, so Triton can pipeline the loads. For each piece it computes the weight in place, m * (2 * j + 1) with j counted from the request's first token (the row's first token position is an argument), masks token slots at or past the row's valid token count, and accumulates in uint64, whose wrap-around is the definition's mod 2**64 (unsigned, so Triton's debug-mode overflow checks do not fire on it). Address offsets are int64. Each program adds its sum to the row's output column with an int64 atomic add: integer addition is exact, so the result does not depend on the order, and the kernel matches the PyTorch path bit for bit. The kernel takes the cache's strides, so head-major, token-major and kernel-split layouts need no copy. The block shape (kernel blocks, heads, tokens, pieces) is a compile-time constant, so the index arithmetic compiles to multiplies and shifts; the layer's slice offset and the output column's pointer and stride are not specialized, so one variant per layer shape serves every layer and column. Integration ----------- _LayerView uses the kernel when its cache is on CUDA and Triton is available, and keeps the PyTorch path otherwise (CPU, tests); the PyTorch weights and gather chunking are only set up for that path, which also serves as the reference implementation. The call signature is the same, so _compute is unchanged. KVChecksumWorker warms the kernel up per layer at startup, with an all-masked row, so its compilation does not land in the first checksummed step. Cost: a second implementation to keep bit-identical with the PyTorch path (asserted by the CUDA tests and the benchmark), and one compiled variant per distinct layer shape at startup. Tests ----- - tests/v1/kv_connector/unit/test_kv_checksum_cuda.py (new, not marked cpu_test, so the GPU job of v1/kv_connector/unit runs it): the kernel's checksums equal the independent NumPy implementation of the definition for head- and token-major layouts and kernel-block splits with partial and empty blocks, including a layer with a large slice factor; for one-head blocks of 32 tokens x 36 pieces (several tiles, not a power of two, so tiling and tail masking are exercised); and a layer's sums land in one column of a wider matrix, leaving the others untouched. - The existing CUDA test of the copy's ordering moves here from test_kv_checksum.py, whose cpu_test mark kept it out of the GPU job. The test_kv_checksum.py helpers take head and cell sizes. - benchmarks/kernels/benchmark_kv_checksum.py times one layer's checksums of N blocks with both paths, reports the KV read rate, and checks the two paths agree. Validation ---------- On GPU, the CUDA tests pass and the benchmark asserts that both paths agree bit for bit. The kernel takes 0.039 ms against 0.314 ms for the PyTorch path on 2048 blocks x 8 heads x 16 tokens x 128 head size, and 0.38 ms against 5.0 ms on rows of about 1 MB, the size of Mamba states. ROCm has not been run. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
Hybrid models' Mamba (SSM) state groups were skipped on both sides, so
only their attention KV was verified. This checksums the Mamba state a
request transfers too, also between producers and consumers of
different TP sizes for Mamba1, Mamba2 and gated delta net (GDN) layers.
What is checksummed
-------------------
A Mamba layer's cache is an int8 view [blocks, 1, 1, state bytes] of one
page per block, packing its states from the page's start: conv, ssm,
then any extra states (ReplaySSM ring, KDA record); padding up to the
page size lies outside the view. The checksum covers the first two
states (or the only one): NIXL transfers exactly those, Mooncake the
whole page, and the extra states may never be transferred.
TP-independent definition
-------------------------
Attention checksums are TP-independent because every slice lies whole
on one rank. A Mamba layer as one slice per TP shard is not: the conv
state concatenates several sub-projections (x, B, C for Mamba2; Q, K, V
for GDN), each sharded by TP on its own, so a rank's page layout
depends on the TP size. A layer whose TP split is known therefore uses
its states' global rows as slices:
- A state's rows run along the dimension TP shards: conv channels and
temporal heads (channels for Mamba1). In the global (TP=1) order the
segments follow each other:
Mamba1 conv [x] temporal [channels]
Mamba2 conv [x, B, C] temporal [heads]
GDN conv [Q, K, V] temporal [value heads]
- Each global row is one slice, numbered across the layer's states in
that order: m = 2 * (layer * H + row) + 1. A rank holds, of each
segment, the rows of its shard: the segment's first global row +
shard * local rows, onward.
- Within a row, j numbers its elements, the pieces of these slices,
read as unsigned 2- or 4-byte integers: a conv row holds only a few
elements (e.g. 3 bf16, 6 bytes), not whole 8-byte pieces. Reading
each element at its own width needs no copy in either conv layout.
States of other element sizes leave the group unchecksummed.
A row lies whole on one rank, so the ranks' parts sum to the TP=1
checksum whatever the TP size or the conv layout (dim-first or not),
exactly as attention heads do. Attention keeps its 8-byte pieces and is
unchanged.
Replicated sub-projections: Mamba2 with n_groups == 1 and TP > 1 pads
its groups so every rank holds the same B and C. Those segments have TP
replicas and use the existing replica keys, as replicated attention
heads do: the producer sends replica 0 and the consumer checks every
replica. A group may mix rows with one replica (x, temporal) and with
TP replicas (B, C); the consumer's per-replica sums already handle
that. Upstream NIXL's heterogeneous-TP Mamba transfer reads B and C at
the offsets of sharded groups (derive_mamba_conv_split), which is wrong
for n_groups == 1; checking every replica lets the checksums report it.
Layout from the layers
----------------------
The segment layout comes from the layer that defines the state:
- MambaStateSplit (kv_cache_interface): the local dimension TP shards
(axis), and the run of (local size, replicas) segments along it.
- MambaSpec.state_splits, filled by MambaBase.get_kv_cache_spec from
the new MambaBase.get_state_splits (default None).
- MambaStateShapeCalculator.mamba1/mamba2/gated_delta_net_state_splits,
next to the matching shape functions and computed from the same
arguments, used by MambaMixer, MambaMixer2 and the Qwen and Olmo GDN
layers. The conv axis follows is_conv_state_dim_first.
- The calculators return None, rather than raise, for sizes the TP size
does not divide: get_kv_cache_spec runs whether or not checksums are
enabled.
Deriving the segments from MambaSpec.shapes alone was rejected: shapes
cannot tell sharded from replicated B and C. NIXL's heterogeneous-TP
Mamba transfer derives the conv split that way
(derive_mamba_conv_split), assuming B is sharded. Reading n_groups from
the HF config would need per-model field names. The cost of asking the
layers is a new optional MambaSpec field and three layer classes
touched outside the checksum code.
Worker and kernel
-----------------
A layer with state_splits gets one _LayerView per segment, viewing the
segment as [blocks, 1, rows, 1, elements]: the rows take the place of
heads, so the existing per-head slice numbering (_LayerCache's new
first_head plus shard * heads) yields the global rows, and nothing
token-based (positions, valid-token masks) applies. Views that would
need a copy, or splits that do not match the state shapes, skip the
group rather than fall back. A layer with more than H rows is skipped
like a layer with more than H KV heads.
The Triton kernel and the PyTorch path take the piece width
(PIECE_BITS, a compile-time constant) and zero-extend narrower pieces
to 64 bits; attention views stay int64, so their kernel is unchanged.
A Mamba layer in a group of mixed specs (UniformTypeKVCacheSpecs) is
skipped: the scheduler picks a Mamba group's state block by its group
spec, which it would not recognize. Layers without state_splits (KDA,
ShortConv, linear attention, ...) keep one slice per TP shard over
their covered bytes as 8-byte pieces, comparable only at equal TP
sizes; see below.
Fingerprint
-----------
The group fingerprint of a Mamba layer covers its checksummed states:
for a layer with state_splits, each state's global segment sizes, its
other dimensions and its dtype, which are the same for every TP size;
otherwise the local shapes and dtypes. Producers and consumers with
different speculative token counts have different conv state shapes,
so their Mamba groups are unverified rather than mismatched.
Which block
-----------
Mamba state is not per token. The scheduler takes the request's
running-state slot: its non-null blocks clipped with
clip_ssm_state_blocks, the rule NIXL transfers with, for
mamba_cache_mode "none" ([running, speculative...]) and "align"
(position-indexed, running state in the last non-speculative block).
It is recorded at position 0 and compared there, without the attention
groups' last-block token count check. Both sides truncate the prompt by
the same backoff, so they hold the same state, h(N - 1). The comparison
takes a Mamba group as one state per block whatever the block sizes in
the payload: without prefix caching the block spans the whole sequence,
so producer and consumer may differ in it for the same state.
Timing: the producer checksums it in the step that finishes the prompt,
after the forward wrote it; the consumer in the step its load finishes,
before any forward of the request (an align-mode copy to the next block
runs in that forward, after the checksum kernels on the same stream).
Different TP sizes
------------------
The payload now carries the producer's TP size. A Mamba group whose
layers have state_splits is compared across TP sizes like any other
group. One without them (and not tp_replicated) is left out when the
TP sizes differ, like a group neither side checksums, with a warning
once: a hybrid model's attention groups still verify the load, also
when fail-closed (a pure-SSM model's load is then unverified). Treating
it as unverified would fail every such hybrid load with heterogeneous
TP under fail-closed, which NIXL otherwise supports.
The payload is now (version, body), each decoded with its type, so a
payload of another version is reported as such, and a malformed one (as
a client can send in kv_transfer_params) only ever raises ValueError,
which makes the load unverified. The payload version stays 1: the
payload is new in this series.
Known limitations
-----------------
- Mamba state of layers without state_splits is not verified between
different TP sizes.
- Extra states (ReplaySSM ring, KDA record) are not checksummed.
- Different speculative token counts on the two sides leave the Mamba
groups unverified.
Tests
-----
- test_kv_checksum.py: a Mamba layer's checksum equals the definition
over its covered states only (changing an extra state or padding
does not change it, flipping a state bit does); TP shards sum to the
whole state; replicated state is sent as replica 0 and checked per
replica; fingerprints follow the state shapes; state that is not
whole 8-byte pieces is skipped; on the model runner's views of a
cache from allocate_kv_cache (two Mamba layers interleaved per block,
padded pages, a nonzero offset), each layer's checksum equals the
definition over its state bytes located from the tensor's offset and
strides. For a TP-split Mamba2 layer, the producer's sums at TP 1, 2
and 4, sharded or with one group, dim-first or not, all equal the
definition over the global state, with bf16 or fp32 conv state;
corrupting one rank's copy of B
fails that replica's check only; the fingerprint is the same across
TP sizes and conv layouts and differs for another conv width; and the
Mamba1, Mamba2 and GDN splits partition their state shapes into
segments of the same global sizes at TP 1, 2 and 4, with B and C
replicated for one group. Numbering rows per rank, reading pieces
signed, or ignoring B and C's replicas (in the worker or the
calculator) makes these cases fail.
- test_kv_checksum_scheduler.py: with a hybrid config (prefix caching,
align mode, one speculative scratch block per request), for NIXL
pull, NIXL push and Mooncake, a correct load verifies and a flipped
state bit is a mismatch. Without a TP split, between TP 1 and TP 2
the state is left out and the load still verifies fail-closed. With
a TP-split Mamba2 layer of one group, producer TP 1 -> consumer TP 2
and TP 2 -> TP 1 verify, and a flipped bit in the last consumer
rank's B is a mismatch; leaving such layers out across TP sizes makes
the mismatch cases fail. The test takes the running-state block as
the row's last block before its speculative scratch blocks, and
checks the workers checksum exactly that block. Making the worker
skip Mamba groups makes the mismatch cases fail, and taking the
row's last block (the scratch block) instead of the running state
makes every Mamba case fail. Malformed checksums (deeply nested
msgpack) leave the load unverified for every connector, without an
exception escaping the scheduler. The test engine now samples nothing
for an unfinished prefill, as the model runner does, since align mode
splits prefills at block boundaries.
- test_kv_checksum_cuda.py: on CUDA, a Mamba state of 24K pieces
spanning two kernel chunks equals the definition, and the kernel's
sums of TP-split rows (2- and 4-byte pieces) equal the definition.
Validation
----------
On upstream NIXL, heterogeneous-TP transfer of Mamba2 state with
n_groups == 1 (granite-4.0-h-micro, producer TP2 -> consumer TP1 and
TP1 -> TP2) gives wrong output without any error, while Nemotron-H-8B
(n_groups 8) and Qwen3.5-0.8B (GDN) are correct. Checking every B and
C replica against the producer's is what lets the checksums report
that case. On GPU, granite producer TP2 -> consumer TP1 and TP1 -> TP2,
with NIXL pull and push, reported a mismatch for every request, only in
the Mamba groups at the running-state block, and recompute matched the
single-instance outputs; granite at equal TP, and Nemotron-H and
Qwen3.5 across TP, verified with no mismatch.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: David Cheung <d7cheung@gmail.com>
In push mode, the only notifs D sends to P are PUSH_REG, which the writer thread consumes, and HB, which _get_new_notifs handles before the counting. Only pull_worker.py fills _pending_recv_notifs, so no code path sends P a plain "<req_id>:<n>" notif. The producer-side branch of NixlPushConnectorWorker._get_new_notifs was carried over from pull mode and never runs: P reports finished sends from its WRITE handles in get_transfer_results. Remove the branch, its guard, the transfer_topo assert it needed and the counter cleanup for finished pushes. _get_new_notifs now counts only consumer-side completion notifs and always returns an empty set. A completion notif for a request that is also being sent (one request both loading and saving on a kv_both instance) is now counted as a load instead of as a send. Tests ----- test_nixl_push_connector.py no longer sets transfer_topo or checks the return value of _get_new_notifs; the consumer-side assertions are unchanged. Validation ---------- - With a tripwire in the removed branch, the 259 tests of test_nixl_push_connector.py, test_nixl_heartbeat.py and test_nixl_connector.py pass without reaching it (the two abort-timeout cases fail locally with and without it). - On GPU, the tripwire was not reached in end-to-end push runs with TP1 -> TP1, producer TP2 -> consumer TP1 and producer TP1 -> consumer TP2, and those runs met their expected results. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
MiniMax-M3's lightning indexer writes one index-key vector per token, from index_k, a single head that MinimaxM3QKVParallelLinearWithIndexer replicates to every TP rank. Its cache spec did not say so: max_tp_shards was unset, which AttentionSpec documents as "unknown, treated as not replicated". Declare max_tp_shards=1, as DeepSeek V3.2 declares for its indexer cache of the same shape. Behavior today is unchanged: max_tp_shards is read through layer_tp_replicas only. kv_tp_replicas, the gcd of the layers' replica counts, stays 1 because the GQA layers report 1, and the offloading connector reads it only for MLA models. NIXL already transfers the cache whole to every rank, classifying it by its MLA spec type. KV checksums take a cache's replica count from max_tp_shards when set. Without it they inferred it from the model's KV head count, which does not describe this one-head cache: unsupported at TP 2, and too few replicas at TP 4 and 8. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
A KV cache group may hold several distinct caches of one layer:
MiniMax-M3's block-sparse attention keeps its GQA KV and its lightning
indexer's side cache (one index-key vector per token) in one group,
both under their decoder layer's module. Their slices were numbered
from the same head, so the worker skipped any such group as ambiguous
("two of its KV caches share a layer number"), leaving every
MiniMax-M3 load unverified: unchecked under fail-open, failed under
fail-closed.
The caches of one layer module (the layer name up to its layer number,
e.g. "model.layers.3") now number their heads consecutively, in layer
name order: the second takes the heads after the first's global heads
(TP shards times heads per rank, or global rows for TP-split Mamba
state). A layer module lies whole on one PP stage, so producers and
consumers of any TP or PP sizes agree on the numbering; a side with a
cache the other lacks has a different group fingerprint, so the load
is unverified rather than mismatched. KV-sharing aliases are still
counted once, under the lowest (layer number, name) using the cache.
Caches of different modules with one layer number (e.g. a draft
model's "draft_model.model.layers.0" next to "model.layers.0") still
skip the group: under PP they can sit on different stages, which
would number them apart. A layer whose caches together exceed
MAX_KV_HEADS_PER_LAYER heads is skipped, as one cache with that many
heads is.
Swapping data between a layer's caches now changes the checksum too,
since their slices have distinct factors.
Tests
-----
test_kv_checksum.py:
- Two caches of one layer module are both covered, the second in name
order taking the heads after the first's (the definition over their
heads concatenated).
- MiniMax-M3's layout: attention KV of two heads, sharded at TP2 and
replicated at TP4, and a one-head indexer cache held whole by every
rank (max_tp_shards=1), in a group of mixed specs. The producer's
shards sum to the definition with each head counted once, and every
consumer replica check equals it.
- One layer number in two modules ("model.layers.0" and
"mtp.layers.0") still skips the group.
Numbering both caches from the same head makes the first two fail.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: David Cheung <d7cheung@gmail.com>
A KV checksum counts a layer's TP replicas, to send replica 0 and check every consumer replica. A spec without max_tp_shards and with one KV head per rank was resolved with the model's KV head count, the HF num_key_value_heads. That count describes the model's attention heads, not a one-head MLA-type cache: it gave too few replicas for TP sizes it divides and no answer for the others, so such a cache was skipped or checksummed per TP shard, and producers and consumers of different TP sizes compared different numberings. An MLA-type cache (MLAAttentionSpec or SlidingWindowMLASpec) without max_tp_shards is now taken as held whole by every rank, as NIXL classifies and transfers it (it reads such regions whole, by spec type). Specs that declare max_tp_shards keep using it. In the tree this changes: - Kimi-K3's MLA cache, whose config's num_key_value_heads is above 1: previously skipped or numbered per shard, now checksummed as replicated. - HiddenStateCacheSpec (extract_hidden_states), an MLAAttentionSpec of num_hidden_states heads identical on every rank: previously counted as one replica per rank, now as replicated. The other MLA-type specs without max_tp_shards are compressed caches, which the checksums skip. A later head-sharded MLA-type cache would be transferred wrong by NIXL as well; its checksums would then report a mismatch. Tests ----- test_kv_checksum.py: - An MLA cache of one head is replicated on every rank whether its spec declares max_tp_shards=1 or leaves it unset, with a model KV head count of 8 as in Kimi-K3's config. - The MiniMax-M3 layout test (attention KV sharded at TP2 and replicated at TP4 next to a one-head indexer cache) also runs with the indexer's spec leaving max_tp_shards unset. Without the rule the unset cases fail. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
Each step, KVChecksumWorker._compute expands every request's blocks
into per-row block ids, first tokens and valid token counts. It did
this with Python loops, one element per block, and copied the three
lists to the GPU separately for every KV cache group: list to tensor,
pin_memory(), then a non-blocking copy, three times per group.
Now each group builds its rows with numpy, vectorized across requests,
and all groups' rows are copied with one pinned host-to-device copy
per call. Each group starts at an even column, so every slice the
kernel reads is 16-byte aligned, like the slices warm_up compiles for;
without the padding, a group with an odd number of rows would misalign
the slices after it, and Triton would compile another variant while
serving.
A reused pinned staging buffer (vLLM's UvaBufferPool) was considered
and not used: pin_memory() is backed by PyTorch's caching host
allocator, which already reuses pinned blocks and records an event on
the copy, so a pool would save only allocator bookkeeping while adding
reuse-lifetime constraints.
Host time to build the rows (Apple M-series CPU, before -> after):
block size prompt x requests rows before after
16 128K x 1 8192 1516 us 110 us
16 16K x 32 32768 5875 us 534 us
128 128K x 1 1024 193 us 17 us
128 4K x 128 4096 757 us 360 us
128 1K x 1 8 3 us 3 us
GPU-side effects (one copy instead of 3 per group) are measured
separately.
Tests: loads of two requests and two groups finishing in one step,
starting past the first block, ending in partial blocks, with an odd
number of rows in the first group, match the definition on CPU
(test_kv_checksum.py) and with the Triton kernel
(test_kv_checksum_cuda.py). The CPU test fails if the per-request
start offset is dropped. Alignment itself is not asserted.
.venv/bin/python -m pytest tests/v1/kv_connector/unit/test_kv_checksum.py \
tests/v1/kv_connector/unit/test_kv_checksum_scheduler.py # 143 passed
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: David Cheung <d7cheung@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Address review on the KV checksum framework:
- Drop the Model Runner V1 hooks (gpu_model_runner.py and
kv_connector_model_runner_mixin.py). V1 now rejects
enable_kv_checksum at startup ("Model Runner V1 does not support:
KV checksums"). With them go post_forward's defer_sends and
finalize_sends: they existed for V1's drafter, which writes KV after
the connector step. V2 calls post_forward after the speculator has
proposed, so the sends already cover the drafter's KV.
- ActiveKVConnector creates the checksum worker and holds it as a
field, instead of a module-level global set by
init_kv_checksum_worker. create_kv_checksum_worker returns the
worker, or None when checksums are off. get_kv_connector takes the
KV cache config for it.
- The copy to host always runs on a torch.cuda.Stream. ROCm and XPU
provide it, and the CPU worker stubs streams and events, so the
synchronous fallback for devices without streams is gone. The
connector unit tests stub them the same way where CUDA is
unavailable (fixture cuda_stream_stubs).
Ported from vllm-project#61578 (1a2920b) onto the combined branch: kept
post_forward's received checksums and replica_zero_only, and the
layer view warm-up.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: David Cheung <d7cheung@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
…tput vllm-project#56430 made get_kv_connector_output abstract on AsyncModelRunnerOutput, so AsyncKVChecksumOutput could no longer be instantiated. Return the wrapped output's KV connector output, without checksums: a load it reports as finished then fails verification and is recomputed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
… objects Tests that build Scheduler, ActiveKVConnector, VllmConfig or a Mamba layer without __init__ now set kv_checksum, kv_checksum_worker, kv_transfer_config and get_state_splits, which this change reads. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
… init GPUModelRunner.main_stream caches the current stream on first access. Passing it to get_kv_connector in initialize_kv_cache read it before vLLM switches to its dedicated stream, so the cached stream could be the default stream. Output copies then waited on the wrong stream and could read sampled token ids before the sampler wrote them. The connector now takes the current stream in post_forward, when it starts the checksum copy. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
A connector that offloads KV and loads it back can now offload each block's checksum with the KV, and have the loads checked against it, the same way P/D loads are checked against the producer's checksums. The contract between the checksum worker and a connector: - init_block_checksums(group_ids) returns a (num_blocks,) int64 device tensor, one entry per block: the sum of this worker's partial checksums of the block. One rank checks its own KV, so no cross-rank sum is needed. - save_block_checksums(runs) fills the entries of full blocks, on the current stream, after the forward that wrote them and before the connector reads them to offload. A run carries the blocks' hashes, and each entry is XORed with a tag of its block hash and its group's fingerprint (block_checksum_tags). - The connector moves a block's entry wherever it moves the block's KV, and writes it back before it reports the load finished. Where the entry is kept (with the KV, in object metadata, ...) is up to it. Connector API: set_kv_checksum_worker() hands the checksum worker to the connector before register_kv_caches. get_kv_checksum_sources(request) names the cache tiers whose loads of the request bring the checksums back. on_kv_checksum_failure(request, mismatches) lets the connector decide what happens to offloaded KV that failed its check (drop the bad entries, its whole cache, ...): without that, a retry of the request, or another request with the same prefix, may load the same KV again. MultiConnector forwards the worker to every connector, but routes the other two to the connector serving the request, as another one may report the same tier (e.g. SimpleCPUOffloadConnector reports HOST). Scheduler: an asynchronous load from those tiers that is not a P/D load is scheduled as reqs_to_load. Once it finishes, each worker reports its checksums of the loaded blocks together with the offloaded ones, and the scheduler compares them per block position. A failed check takes the existing KV load failure path and calls on_kv_checksum_failure. The scheduler removes the tags of the blocks the request expects at those positions, from its own block hashes, before it compares. KV loaded at another position, or from another TP or PP shard, already checksums differently, since each worker weighs pieces by their token position, head and layer. But a stored checksum is kept with its KV, so without the tag an intact value stored for the same position of another prefix (another request's key, a reused or late-written value) would match its own checksum; with it, such a value fails like corrupted KV, and so does one stored by a model with other layers or dtypes. P/D loads need no tag: their reference comes from the producer of the same request, apart from the KV. A load of groups no checksums are offloaded for is not checked: which groups are covered is fixed at startup. P/D loads are unchanged. Synchronous loads (load_kv_async=False, e.g. LMCacheConnectorV1, MoRIIO READ mode, Mooncake store or UMBP with load_async=False) fit the same contract. They skip WAITING_FOR_REMOTE_KVS and finished_recving, and the step that loads them also uses them, so a bad load is found after the fact. They can be scheduled for checking in the step that loads them, checksummed in that step's post_forward (the connector writes the offloaded entries back during the load, before the forward), and a mismatch reported as invalid blocks rather than a failed load. The scheduler's existing handling of invalid blocks for synchronous loads then rolls the requests back: it drops their outputs of that step, truncates num_computed_tokens to the valid prefix, and evicts or recomputes the blocks. That handling supports a single KV cache group. Tested on GB200: tests/v1/kv_connector/unit/test_kv_checksum*.py, test_multi_connector.py and test_mooncake_store_*.py, 557 passed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
With KV checksums enabled, MooncakeStoreConnector stores each block's checksum with the block's KV, checks the KV it loads back against it, and removes the blocks that fail from the store. - When: KV checksums are enabled, keys are rank-local (no store TP sharding, which lets ranks of another TP size read a shard that a rank checksummed as a whole), every group is full attention, every key holds one whole block (hash block size equal to the block size, no partial-hash hits), and the KV cache is not in a custom_mem_pool (NVLINK/BAREX), which the checksum tensor is not allocated from. Values with checksums get their own key namespace (@kv_checksums), so they never mix with values stored without one. - Store: the block checksum tensor (init_block_checksums) is registered with the store and appended as one more 8-byte region of every value, so each put and get carries a block's entry with its KV. wait_for_save queues save_block_checksums for the blocks each save covers, with their hashes from the request, before it records the CUDA event the sending thread waits on. - Load: loads are asynchronous, so the generic check applies: the entry comes back with the KV, and the scheduler compares it per block. - Failure: on_kv_checksum_failure maps the mismatching block positions to their hashes and sends a new REMOVE_MSG over the LookupKey channel. Worker rank 0 calls batch_remove(keys, force=True) for every rank's key of those blocks. Force is needed because the load that found the corruption still holds a read lease on them. A remove that fails, e.g. for a key being written, is logged. - Otherwise, with KV checksums enabled, the scheduler logs at startup that loads from the store are not checked, and the docs list the requirements. Tested on GB200: tests/v1/kv_connector/unit/test_kv_checksum*.py, test_multi_connector.py and test_mooncake_store_*.py, 557 passed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
Contributor
|
Documentation preview: https://vllm--64126.org.readthedocs.build/en/64126/ |
mergify
Bot
added
documentation
mergify
Bot
added
kv-connector
mooncake
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
mergify
Bot
added
needs-rebase
nixl
9 of 48 tasks
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
Check KV loaded back from KV offloading against checksums stored with it, and implement it for
MooncakeStoreConnector. Each stored checksum is bound to the block it was stored for, so a value that is intact but belongs to another prefix fails the check like corrupted KV.Stacked on #62211 (connector-agnostic KV checksums for P/D): the first commit adds the offloading contract to it, the second implements it for the Mooncake store. Review only the last two commits until #62211 merges.
Claims
enable_kv_checksum,MooncakeStoreConnectorstores each block's checksum with the block's KV, and every load from the store is checked against it, per block, with the same kernel and scheduler path as P/D loads.kv_load_failure_policy), and the failing blocks are removed from the store on every rank, so later requests do not load them again.custom_mem_pool, HiSparse) are reported at startup instead of left silently unchecked.Out of scope: KV that never leaves the GPU (e.g. a kernel reading past a request's own blocks). These checksums only check KV that is loaded from elsewhere.
Validation
Unit tests:
.venv/bin/python -m pytest -q \ tests/v1/kv_connector/unit/test_kv_checksum.py \ tests/v1/kv_connector/unit/test_kv_checksum_scheduler.py \ tests/v1/kv_connector/unit/test_multi_connector.py \ tests/v1/kv_connector/unit/test_mooncake_store_*.pytest_multi_example_connector_consistencydeselected: it needs the gatedmeta-llama/Llama-3.2-1B-Instruct).test_offloaded_kv_load_is_checked_against_its_checksums[True-other_prefix]loads intact KV whose checksums were stored for the same positions of another request. Without the tag (the save call reverted to the old signature) the load passes the check, so this case fails (assert set() == {'id-1'}; the other 5 cases pass); with the tag both positions fail and all 6 pass (GB200).End to end on GB200, Qwen3-8B,
MooncakeStoreConnector(kv_both) withenable_kv_checksum, TP1. Instance A stores 16 4k-token prompts; instance B, with an empty GPU cache and the samemooncake_master, serves them from the store; C (no connector) is a reference.On detection B logs
Removing 1 Mooncake store blocks of request ... that failed their KV checksums., the key is gone right after, and the request is recomputed. In the uninjected run and the misplacement run, A and B ran withVLLM_GPU_SYNC_CHECK=error: no sync violations.Earlier runs of the same setup before the tag was added (same storage layout, the stored value differs only by the XOR):
TP1 and TP2: B's external prefix cache hit rate 99.6% (every full block; the partial last block is recomputed), 0 mismatches, 0 load failures. Outputs vs. A and vs. C differ in the usual cross-instance way (TP2: B==A 5/8, B==C 6/8, A==C 6/8).
Overhead on store hits, checksums on vs. off, median of 3 rounds:
No measurable overhead. These loads go over TCP and dominate TTFT, so this does not bound the overhead relative to a faster (RDMA) load; the checksum kernels themselves are measured in [KV Transfer] Connector-agnostic KV checksums for P/D disaggregation #62211. The tag adds one 8-byte-per-block host-to-device copy per save and a NumPy XOR per checked block on the scheduler.
Setup:
mooncake-transfer-engine0.3.13.post1,protocol: tcp, both instances withMC_TCP_MAX_QUEUED_TRANSFERS_PER_PEER=65535 MC_TCP_MAX_PENDING_ADMISSIONS_PER_PEER=65535 MC_TCP_ADMISSION_TIMEOUT_MS=60000. Without them every cross-instance store read fails (-800 after a 60 s timeout), with or without KV checksums: one 4k-token load is 240 keys x 36 regions, far over the default per-peer queue of 1024 transfers.CI: the new tests are in existing files that CI already runs; the Triton-kernel ones need a GPU.
Details
Contract (first commit). The checksum worker gives a connector a
(num_blocks,)int64 device tensor (init_block_checksums) and fills the entries of full blocks before the connector offloads them (save_block_checksums). The connector passes each block's hash with it, moves a block's entry wherever it moves the block's KV, and writes it back before it reports the load finished.get_kv_checksum_sources(request)names the cache tiers whose loads bring the checksums back; such an asynchronous load is checksummed on its finishing step and compared per block by the scheduler.on_kv_checksum_failurelets the connector drop the bad KV.Why the tag. A P/D load is checked against a reference that comes from the producer of the same request, separately from the KV, so KV from the wrong place cannot match it. An offloaded checksum travels with its KV, so on its own it only shows that a value is internally consistent. The checksum already depends on the token position, head and layer, so KV loaded at the wrong position or into another shard fails; what it cannot tell is whether the value at the right position belongs to this prefix. The tag (
block_checksum_tags:block_hash[:8] ^ group fingerprint) binds the stored checksum to what the reader expects: the block hash covers the prefix tokens, LoRA, multimodal inputs and cache salt, and the fingerprint (from #62211) covers the group's layer names, types and dtypes. The scheduler derives the expected tag fromrequest.block_hashes, not from the store. Leaving the rank out of the tag keeps replicas' stored checksums equal. P/D is unchanged. The tag adds no new requirement: cross-instance hits already need reproducible block hashes, since they are the store keys.Mooncake store (second commit). The checksum tensor is registered with the store and appended as one more 8-byte region of every value, so put and get carry a block's entry with its KV and need no extra round trip.
wait_for_savequeues the checksums, with the blocks' hashes from the request, before the CUDA event the sending thread waits on. Keys are rank-local because each rank checks its own shard; values with checksums use a separate key namespace (@kv_checksums), so they never mix with values stored without one, at the cost of not reusing blocks stored before checksums were enabled. On failure, rank 0 runsbatch_remove(keys, force=True)over the LookupKey channel (newREMOVE_MSG); force is needed because the failing load still holds a read lease.Size. About 330 non-test lines in the first commit and 260 in the second, most of it the per-block plumbing in the checksum worker and scheduler and the Mooncake register/put/get/remove paths; the tag itself is about 30.
Limitations. Synchronous loads (
load_async=False) are not checked; Mooncake store asserts asynchronous loads. This detects misplaced values; it does not fix their cause (kvcache-ai/Mooncake#4491 should be fixed in Mooncake).Not a duplicate. #57912 adds CRC32C for the filesystem offload tier; this reuses #62211's GPU checksums and its scheduler path, and covers the Mooncake store.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
🤖 Generated with Claude Code