Repository navigation
Conversation
sawsa307
marked this pull request as ready for review
October 13, 2026 19:19
sawsa307
requested review from
ApostaC,
WoosukKwon,
ivanium,
orozery,
robertgshaw2-redhat and
wzhao18
as code owners
October 13, 2026 19:19
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
sawsa307
requested review from
NickLucche,
ProExpertProg,
ZJY0516,
alexm-redhat,
heheda12345,
houseroad,
mgoin,
njhill,
tdoublep,
tlrmchlsmth,
tomeras91,
vadiklyutiy,
xuechendi,
yewentao256,
youkaichao and
ywang96
as code owners
October 13, 2026 19:19
mergify
Bot
added
performance
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>
sawsa307
force-pushed
the
kv-checksum-upstream
branch
from
October 13, 2026 19:51
62ba0cd to
ef3d386
Compare
sawsa307
marked this pull request as draft
October 13, 2026 20:03
njhill
added
the
verified
6 tasks done
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
njhill
reviewed
Oct 14, 2026
Member
There was a problem hiding this comment.
We can remove the MRV1 files (including this and kv_connector_model_runner_mixin.py), and include this in the MRV1-unsupported startup check.
Comment on lines
+800
to
+802
| init_kv_checksum_worker( | ||
| self.vllm_config, self.kv_cache_config, kv_caches_dict | ||
| ) |
Member
There was a problem hiding this comment.
Can we just initialize the kv checksum worker inside get_kv_connector and have it be a field in there?
| return host.to(device, non_blocking=True) | ||
|
|
||
|
|
||
| _KV_CHECKSUM_WORKER: KVChecksumWorker | None = None |
Member
There was a problem hiding this comment.
Could we avoid this global state? Just have a function to create/return the checksum worker (or None if not enabled)
Comment on lines
+668
to
+672
| self._copy_stream = ( | ||
| torch.cuda.Stream(device) | ||
| if device is not None and device.type == "cuda" | ||
| else None | ||
| ) |
Member
There was a problem hiding this comment.
We can keep this targeted at GPU; this will still work with ROCm, XPU, CPU
Suggested change
| self._copy_stream = ( | |
| torch.cuda.Stream(device) | |
| if device is not None and device.type == "cuda" | |
| else None | |
| ) | |
| self._copy_stream = torch.cuda.Stream(device) |
This was referenced Oct 14, 2026
ZJY0516
added
ready
This was referenced Oct 17, 2026
sawsa307
added a commit
to sawsa307/vllm
that referenced
this pull request
Oct 19, 2026
…e replicated Problem: with heterogeneous TP between prefill and decode, NixlConnector silently transfers the wrong Mamba2 conv state when TP does not divide n_groups (e.g. granite-4.0-h with n_groups=1), and decode output is wrong without any error (vllm-project#61526). Cause: when n_groups % tp != 0, mamba2_state_shape pads the groups to the TP size, so every rank holds a full replica of B and C. The conv transfer derived the remote layout from the local shapes alone and assumed B and C are sharded across TP like x. With D TP > P TP, odd D ranks read P's C in place of B, and every D rank read C from past the end of P's conv state. With P TP > D TP, D read the first part of the same B/C replica from every P rank. Fix: Mamba2 layers declare how TP splits their states (MambaStateSplit, MambaSpec.state_splits, MambaMixer2.get_state_splits, built by MambaStateShapeCalculator.mamba2_state_splits from the unpadded n_groups). derive_mamba_conv_split records the global dims, from which MambaConvSplitInfo.proj_layout gives the (columns, replicas) layout of any TP size. Each D rank then reads, for every segment, the global columns it holds from the P rank that holds them, and checks that P's conv state size matches the layout derived for P. When no segment is replicated the offsets are unchanged. Under heterogeneous TP, a Mamba2 layer without a declared split now fails the handshake instead of transferring corrupt state. Mamba1 and GDN are unchanged. Push mode writes one chunk of the producer's page to each consumer rank, which cannot fill B/C that a larger consumer TP replicates, so NixlPushConnector now rejects P TP < D TP in that case at handshake. The TP-split declaration (MambaStateSplit, MambaSpec.state_splits, Mamba2 get_state_splits) is also in the KV checksum PRs (vllm-project#61488, vllm-project#62211), which use it; whichever lands second drops the duplicate. No handshake change: both sides derive the layout from the same model config and their own TP size, which the handshake already carries; the existing ssm_sizes field is used to check the result. Tests: a CPU test builds the global [x|B|C] conv state, slices each rank's page for n_groups in {1, 8} and (P, D) TP in {(1,2), (2,1), (2,4), (4,2)}, takes the source ranks from compute_tp_mapping and the local destination slices from _build_local_splits_from_plan, simulates every D rank's reads and checks D's page equals its expected page; it fails for n_groups=1 with the previous offsets. Two more tests cover the handshake rejections in pull and push mode. The TP8 MLA-hybrid geometry test now uses KDA dims that 8 ranks can split. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
sawsa307
added a commit
to sawsa307/vllm
that referenced
this pull request
Oct 19, 2026
…e replicated Problem: with heterogeneous TP between prefill and decode, NixlConnector silently transfers the wrong Mamba2 conv state when TP does not divide n_groups (e.g. granite-4.0-h with n_groups=1), and decode output is wrong without any error (vllm-project#61526). Cause: when n_groups % tp != 0, mamba2_state_shape pads the groups to the TP size, so every rank holds a full replica of B and C. The conv transfer derived the remote layout from the local shapes alone and assumed B and C are sharded across TP like x. With D TP > P TP, odd D ranks read P's C in place of B, and every D rank read C from past the end of P's conv state. With P TP > D TP, D read the first part of the same B/C replica from every P rank. Fix: Mamba2 layers declare how TP splits their states (MambaStateSplit, MambaSpec.state_splits, MambaMixer2.get_state_splits, built by MambaStateShapeCalculator.mamba2_state_splits from the unpadded n_groups). derive_mamba_conv_split records the global dims, from which MambaConvSplitInfo.proj_layout gives the (columns, replicas) layout of any TP size. Each D rank then reads, for every segment, the global columns it holds from the P rank that holds them, and checks that P's conv state size matches the layout derived for P. When no segment is replicated the offsets are unchanged. Under heterogeneous TP, a Mamba2 layer without a declared split now fails the handshake instead of transferring corrupt state. Mamba1 and GDN are unchanged. Push mode writes one chunk of the producer's page to each consumer rank, which cannot fill B/C that a larger consumer TP replicates, so NixlPushConnector now rejects P TP < D TP in that case at handshake. The TP-split declaration (MambaStateSplit, MambaSpec.state_splits, Mamba2 get_state_splits) is also in the KV checksum PRs (vllm-project#61488, vllm-project#62211), which use it; whichever lands second drops the duplicate. No handshake change: both sides derive the layout from the same model config and their own TP size, which the handshake already carries; the existing ssm_sizes field is used to check the result. Tests: a CPU test builds the global [x|B|C] conv state, slices each rank's page for n_groups in {1, 8} and (P, D) TP in {(1,2), (2,1), (2,4), (4,2)}, takes the source ranks from compute_tp_mapping and the local destination slices from _build_local_splits_from_plan, simulates every D rank's reads and checks D's page equals its expected page; it fails for n_groups=1 with the previous offsets. Two more tests cover the handshake rejections in pull and push mode. The TP8 MLA-hybrid geometry test now uses KDA dims that 8 ranks can split. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: David Cheung <d7cheung@gmail.com>
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
Connector-agnostic KV checksums for P/D disaggregation. Some transfers complete without error but deliver the wrong KV: a read after the producer has already reused its blocks, or KV placed at the wrong position because producer and consumer disagree on the layout. With
kv_transfer_config.enable_kv_checksumon both sides, the consumer verifies every KV load against checksums computed by the producer. Underkv_checksum_fail_closed, a mismatched or unverifiable load goes through the existing KV load failure path (kv_load_failure_policy:fail, orrecomputelocally).Related, not duplicated: #57912 and RFC #54363 add integrity checks for the filesystem KV offload tier; this PR covers producer-to-consumer transfers. It also references #58860, #59746 / #59747 and RFC #47812 / #48216 (see the commit messages).
Claims
Validation
Unit tests (CPU, plus CUDA tests where a GPU is available):
595 passed and 12 skipped (the CUDA tests, on a CPU host). One test failed:
test_multi_example_connector_consistencyneeds a gated model and fails the same way on the base commit. On GPU,test_kv_checksum_cuda.pypasses, and the Triton kernel matches the PyTorch path bit for bit. pre-commit and mypy pass.Mutation checks: each of the following edits makes the listed tests fail.
End to end on GPU (NIXL,
VLLM_GPU_SYNC_CHECK=error, with a test-only fault-injecting connector that is not part of this PR):Outputs unchanged (determinism test). Setup: Qwen3-8B, P/D TP1 -> TP1 with NIXL pull,
VLLM_BATCH_INVARIANT=1, FlashInfer autotune off, temperature 0, concurrency 1, the first 200 gsm8k questions, max_tokens 512. Comparison is on full output token ids:Three runs on fresh servers: two with checksums off and one with checksums on (consumer fail-closed). Each pair of runs is compared on all 200 outputs:
Accuracy (MiniMax-M3, P/D TP4 -> TP4, NIXL pull). The comparison is checksums on (fail-closed with recompute) against checksums off and against a single instance. Differences are paired by problem; the 95% band treats the problem set as fixed.
gpqa_diamond: P/D on − single instance = −1.11 pp (±1.5).
In the first aime25 run with checksums on, the consumer stalled for about 80 s. Its NIXL heartbeats stopped, and the producer's 30 s lease expired for 31 requests. The consumer later read those blocks after the producer had reused them. The checksums flagged exactly those 31 loads, matching the producer's lease-expiry log one to one. The consumer recomputed them locally, and every request completed with a full answer. This is the failure reported in #61229; #61549 makes the NIXL pull consumer check the lease itself.
Heterogeneous TP (MiniMax-M3, P TP4 -> D DEP8, i.e. DP8 x TP1 with EP8 on two nodes). Each decode rank pulls its KV heads from all 4 producer ranks.
Data-parallel producer (MiniMax-M3, P DEP2 -> D DEP8, i.e. P DP2 x TP4 with EP8 on two nodes). Each request is prefilled by one of the two producer engines, and the consumer checks the checksums from whichever engine prefilled it.
Performance. Qwen3-8B, P/D TP2 -> TP2 with NIXL pull on GB200, prefix caching on, consumer fail-closed with recompute,
vllm bench serveon random prompts. Checksums off and on were launched 3 times each, interleaved; values are the mean of the 3 launches, with 0 failed requests in every run.The cost at short prompts is mostly the fixed cost of one kernel launch per layer in the steps that finish a prefill or a load. Fusing the layers of one layout into one launch reduces the concurrency-1 cost to +1.5 ms and the concurrency-32 cost to -1.5%; that change is being evaluated and is not in this PR yet.
Details
The series is split by component. Each commit message gives the design, the alternatives that were rejected, known limitations and validation.
kv_checksum_carrier. NIXL pull carries the checksums inkv_transfer_params.NIXL_CONNECTOR_VERSION(once for the series).benchmarks/kernels/benchmark_kv_checksum.py.MambaSpec.state_splits, for Mamba1, Mamba2 and GDN), so state can be compared across different TP sizes.max_tp_shardsis treated as replicated, which matches how NIXL classifies it.Limitations:
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.
AI assistance (Claude) was used to write and review this change. I reviewed every changed line and ran the validation above.
🤖 Generated with Claude Code