[Bugfix][NIXL] Reject push handshakes between different P/D block sizes by lucamotz · Pull Request #59747 · vllm-project/vllm

lucamotz · 2026-10-02T07:04:56Z

Overview

NixlPushConnector now rejects P/D peers with different block sizes at the handshake, so their requests fail at once instead of stalling or crashing D.

Claims

A P/D pair that differs in kernel block size or kernel blocks per logical block is rejected at the handshake: each request fails at once and D keeps serving.
Pull and equal-block-size push are unchanged.

Validation

pytest tests/v1/kv_connector/unit/test_nixl_push_connector.py tests/v1/kv_connector/unit/test_nixl_desc_geometry.py

test_push_handshake_rejects_different_block_sizes covers the new check. Of the two tests of the removed check, test_layer_handshake_rejects_unsupported_geometry keeps only its TP case, and test_packed_mla_rejects_unequal_block_sizes_before_peer_registration is removed: the new check does not depend on packed MLA or PP and runs before any registration.

CPU, this PR at 8ab8d31c9b on main a70e48deba: the new test's 3 cases fail without this PR and pass with it. Other NIXL unit tests are unchanged.

GB10, Qwen3-0.6B, NixlPushConnector, P and D on one GPU, greedy, --enforce-eager, main 0cbac6cd13 without and with this PR at 5a908a9e72:

P 16 / D 64, FLASH_ATTN: on main, P's handshake fails with Local block size 16 is not divisible by remote block size 64 and D's requests are still pending after 150 s. With this PR, D's handshake raises NixlPushConnector requires identical P/D block sizes. and each request fails at once.
FLASHINFER, P 128 / D 64, both with 64-token kernel blocks: on main, D's EngineCore dies with IndexError: tuple index out of range in get_transfer_results. With this PR, D's handshake rejects it the same way.
In both rejected cases D keeps serving requests.
64/64: outputs and logprobs are bitwise equal to main's push.

Details

NixlPushConnector cannot convert or re-pair blocks when P and D differ in kernel block size or kernel blocks per logical block, and D registers P for notifications without checking either. On main, for such a pair:

D handshakes with P and sends PUSH_REG with its block ids.
P handshakes with D and WRITEs into those blocks. Depending on the sizes, P's handshake fails or P writes misaligned blocks.
D waits for a WRITE that never comes, or its receive fails on the block-size-ratio assert or on empty push block ids.

This PR rejects such peers in _validate_remote_parallel_config, which _nixl_handshake runs for P's and both of D's handshakes (registration and heartbeat) before registering the peer, so the request fails like any failed handshake. It replaces the PP-only block-size check in add_remote_agent and removes push's block_size_ratio > 1 branch in _xfer_blocks, which can no longer run.

No open PR rejects push peers with different block sizes.

Pull Request Checklist

I used vLLM's /pr-checklist skill. (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 Code) was used for the analysis, patch and tests. I reviewed every changed line.

🤖 Generated with Claude Code

mergify · 2026-10-09T08:07:40Z

Documentation preview: https://vllm--59747.org.readthedocs.build/en/59747/

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

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 GB200, 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 the smoke test's clean stage: 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>

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>

sawsa307

Some small nits, overall LGTM. cc @njhill and @NickLucche to take a final look, thanks!

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>

The push connector cannot convert or re-pair blocks for a peer whose kernel block size or kernel blocks per logical block differ, and D registers P for notifications without checking either. Depending on the sizes, D then waits for a WRITE that never comes (P rejects its handshake, or fails to build the WRITE, after D sent the registration), or P writes head-sliced or misaligned blocks and D's receive post-process fails on the block-size ratio assert or on the empty push remote block ids. Reject such peers in _validate_remote_parallel_config, which _nixl_handshake runs for P's and both of D's handshakes (registration and heartbeat) before registering the peer, so the request fails on D like any failed handshake. This subsumes the PP-only check in add_remote_agent and leaves push's block_size_ratio > 1 branch in _xfer_blocks unreachable, so it is removed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Luca Motz <luca.motz@icloud.com>

…orker Review follow-up: override _validate_remote_parallel_config in the push worker instead of branching on _TRANSFER_MODE in the base, include the local and remote sizes in the error, and note the push restriction in the hetero block size footnote. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Luca Motz <luca.motz@icloud.com>

aijanai

Hi,
Thanks for also removing the dead code for P > D.
LGTM

mergify Bot added bug Something isn't working kv-connector labels Oct 2, 2026

lucamotz mentioned this pull request Oct 5, 2026

[Core][NIXL] Transfer KV between different P/D KV cache group plans #60113

Draft

5 of 6 tasks

lucamotz force-pushed the nixl-push-block-size-guard branch from 5a908a9 to d64fd6a Compare October 9, 2026 08:07

mergify Bot added the documentation Improvements or additions to documentation label Oct 9, 2026

lucamotz force-pushed the nixl-push-block-size-guard branch from d64fd6a to 22086c1 Compare October 9, 2026 17:42

lucamotz marked this pull request as ready for review October 9, 2026 17:43

lucamotz requested review from ApostaC, NickLucche, ivanium, orozery, wzhao18 and xuechendi as code owners October 9, 2026 17:43

claude Bot reviewed Oct 9, 2026

View reviewed changes

This was referenced Oct 12, 2026

[KV Transfer] Connector-agnostic KV checksums for PD disaggregation (NIXL pull/push, Mooncake) sawsa307/vllm#4

Open

[KV Transfer] Connector-agnostic KV checksums for P/D disaggregation #61488

Closed

sawsa307 mentioned this pull request Oct 14, 2026

[KV Connector][3/N] KV checksums for NIXL push #61580

Closed

5 of 6 tasks

sawsa307 reviewed Oct 16, 2026

View reviewed changes

Comment thread vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py Resolved

Comment thread docs/features/nixl_connector_compatibility.md Resolved

mergify Bot added the nixl NIXL KV connector / multi-tiered KV offloading (KVCR) label Oct 29, 2026

lucamotz and others added 2 commits October 30, 2026 08:19

lucamotz force-pushed the nixl-push-block-size-guard branch from cb04c65 to 8ab8d31 Compare October 30, 2026 08:02

aijanai approved these changes Oct 30, 2026

View reviewed changes

Navigation Menu

Repository navigation

[Bugfix][NIXL] Reject push handshakes between different P/D block sizes - #59747

[Bugfix][NIXL] Reject push handshakes between different P/D block sizes#59747
lucamotz wants to merge 2 commits into
vllm-project:mainfrom
lucamotz:nixl-push-block-size-guard

lucamotz commented Oct 2, 2026 •

edited

Loading

Uh oh!

mergify Bot commented Oct 9, 2026

Uh oh!

claude Bot left a comment

Uh oh!

sawsa307 left a comment

Uh oh!

Uh oh!

Uh oh!

aijanai left a comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

lucamotz commented Oct 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Claims

Validation

Details

Uh oh!

mergify Bot commented Oct 9, 2026

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

sawsa307 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

aijanai left a comment

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

lucamotz commented Oct 2, 2026 •

edited

Loading