[P/D][Bugfix]: Fix the metadata corruption issue in Nixl when TP > 1.

[chaunceyjiang]

FIX #19338

Introduced By #18833

[gemini-code-assist]

Hello @chaunceyjiang, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

Hello team, Gemini here to provide a quick summary of this pull request. Based on the title, this PR aims to fix an AssertionError that occurs when self.num_blocks is equal to nixl_agent_meta.num_blocks. The core change involves modifying an assertion in the vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py file to use a strict greater-than comparison (>) instead of greater-than-or-equal-to (>=). This suggests that the logic requires the local number of blocks to be strictly larger than the remote agent's number of blocks for the assertion to pass correctly in all intended scenarios.

Highlights

• Bugfix: Addresses an AssertionError related to block count comparison in the KV transfer connector.
• Assertion Logic Update: Modifies an assertion in nixl_connector.py from >= to > to correctly handle cases where block counts might be equal.

Changelog

• vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
  • Fixed an AssertionError on line 706 by changing the comparison operator from >= to > in the assertion assert self.num_blocks >= nixl_agent_meta.num_blocks.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Code Review

This pull request aims to fix an AssertionError related to self.num_blocks and nixl_agent_meta.num_blocks. However, the current change appears to make the assertion stricter in a way that would not fix the AssertionError if it was occurring on the specified line as implied by the PR title.

A more detailed explanation of the original bug and the rationale for this specific change would be very helpful. Currently, the PR description is empty, which makes it difficult to fully understand the context and intent of this bugfix.

Summary of Findings

• Effectiveness of the Fix: The primary concern is that the proposed change to the assertion (self.num_blocks > nixl_agent_meta.num_blocks) does not appear to fix the AssertionError if it was caused by self.num_blocks < nixl_agent_meta.num_blocks (as implied by the failure of the original assert self.num_blocks >= nixl_agent_meta.num_blocks). The new assertion is stricter and would also fail under the same condition.
• Clarity of Intent: The pull request lacks a description, making it difficult to understand the exact nature of the bug being fixed and the rationale behind making the assertion stricter. Clarification is needed to ensure the change is appropriate and correctly addresses the underlying issue.

Merge Readiness

This pull request should not be merged in its current state. The core change seems to not address the AssertionError as described in the PR title. Further clarification and potentially a different approach are needed to resolve the underlying issue. I am unable to approve this pull request; it requires further review and changes by the author.

[chaunceyjiang]

The root cause lies here:

vllm/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

Lines 470 to 473 in 9ef9173

	# Handshake with remote agent-rank0 first to get the tp_size of remote
	path = make_zmq_path("tcp", host, port)
	logger.debug("Querying master rank metadata on path: %s", path)
	metadata = handshake(path, 0)

when the process with rank 7 is retrieving the metadata, it incorrectly records the num_blocks.

vllm/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

Lines 708 to 712 in 9ef9173

	# Create dst descs and xfer side handles. TP workers have same #blocks.
	if engine_id in self.dst_num_blocks:
	assert self.dst_num_blocks[engine_id] == nixl_agent_meta.num_blocks
	else:
	self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks

Specifically, rank 7 of D records the num_blocks from rank 0 of P.

[NickLucche]

this won't work when local TP > remote TP

[chaunceyjiang]

indeed

[chaunceyjiang]

this won't work when local TP > remote TP

@NickLucche I have updated the calculation method. Please take a look again.

[chaunceyjiang]

/cc @NickLucche @njhill PTAL.

[NickLucche]

Thanks for reporting the issue! I think there's two overlapping matters here, let me clarify.

While you're correct the num_blocks being recorded is that of rank 0, the main point though is that the assumption here is that all tp ranks in a vllm instance have the same number of blocks, hence the assert self.dst_num_blocks[engine_id] == nixl_agent_meta.num_blocks.
When this assumption does not hold, the actual fix here is to record the minimum num_blocks in the tp group (similar to https://github.com/vllm-project/vllm/blob/87360308/vllm/v1/core/kv_cache_utils.py#L771), which can also be achieved by just sending the value recorded in the config.
Alternatively, one could extend self.dst_num_blocks to a map [str, dict[int,int] ], where the value is the number of blocks per tp rank, but I feel like that would be overkill.

The second issue you reported assert self.num_blocks >= nixl_agent_meta.num_blocks is a separate condition I was already discussing with @ptarasiewiczNV . The rationale is that the consumer should have enough space to fit all the blocks from the producer. We can easily relax that one with some other lower bound as long as the kv cache is not being transferred in its entirity with nixl.

[NickLucche]

this won't work when local TP > remote TP

[chaunceyjiang]

the main point though is that the assumption here is that all tp ranks in a vllm instance have the same number of blocks,

In fact, they are not the same. I've provided some logs here.

(VllmWorker rank=0 pid=2271553) INFO 06-09 07:25:02 [nixl_connector.py:527] Registering KV_Caches: use_mla: False, num_blocks: 111680, block_shape: torch.Size([16, 1, 128]), per_layer_kv_cache_shape: torch.Size([2, 111680, 16, 1, 128])
(VllmWorker rank=1 pid=2271554) INFO 06-09 07:25:02 [nixl_connector.py:527] Registering KV_Caches: use_mla: False, num_blocks: 111488, block_shape: torch.Size([16, 1, 128]), per_layer_kv_cache_shape: torch.Size([2, 111488, 16, 1, 128])
(VllmWorker rank=2 pid=2271555) INFO 06-09 07:25:02 [nixl_connector.py:527] Registering KV_Caches: use_mla: False, num_blocks: 111488, block_shape: torch.Size([16, 1, 128]), per_layer_kv_cache_shape: torch.Size([2, 111488, 16, 1, 128])
(VllmWorker rank=3 pid=2271556) INFO 06-09 07:25:02 [nixl_connector.py:527] Registering KV_Caches: use_mla: False, num_blocks: 111488, block_shape: torch.Size([16, 1, 128]), per_layer_kv_cache_shape: torch.Size([2, 111488, 16, 1, 128])
(VllmWorker rank=4 pid=2271557) INFO 06-09 07:25:02 [nixl_connector.py:527] Registering KV_Caches: use_mla: False, num_blocks: 111488, block_shape: torch.Size([16, 1, 128]), per_layer_kv_cache_shape: torch.Size([2, 111488, 16, 1, 128])
(VllmWorker rank=5 pid=2271558) INFO 06-09 07:25:02 [nixl_connector.py:527] Registering KV_Caches: use_mla: False, num_blocks: 111488, block_shape: torch.Size([16, 1, 128]), per_layer_kv_cache_shape: torch.Size([2, 111488, 16, 1, 128])
(VllmWorker rank=6 pid=2271559) INFO 06-09 07:25:02 [nixl_connector.py:527] Registering KV_Caches: use_mla: False, num_blocks: 112448, block_shape: torch.Size([16, 1, 128]), per_layer_kv_cache_shape: torch.Size([2, 112448, 16, 1, 128])
(VllmWorker rank=7 pid=2271560) INFO 06-09 07:25:02 [nixl_connector.py:527] Registering KV_Caches: use_mla: False, num_blocks: 111488, block_shape: torch.Size([16, 1, 128]), per_layer_kv_cache_shape: torch.Size([2, 111488, 16, 1, 128])

[NickLucche]

As I commented earlier, I'd fix the num_blocks issue rather than skipping the checks.
Can you either send the num_blocks off of kv_cache config here https://github.com/vllm-project/vllm/blob/main/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py#L62 (which is already the min in the group) or simply store the min during add_remote_agent?

[NickLucche]

rather than an early exit here, which is equivalent to skipping descs at L727, I'd prefer fixing the num_blocks issue for once and still execute the other asserts (those are extra correctness checks that won't hurt to have).

[chaunceyjiang]

Got it, I'll look into the solution you proposed.

Just to add/clarify:

Here, it only skips the first handshake execution. During the second handshake execution, it will not be skipped.

vllm/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

Lines 470 to 473 in 9ef9173

	# Handshake with remote agent-rank0 first to get the tp_size of remote
	path = make_zmq_path("tcp", host, port)
	logger.debug("Querying master rank metadata on path: %s", path)
	metadata = handshake(path, 0)

The first handshake execution is solely for obtaining the remote tp_size, so it can be safely skipped.

vllm/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

Lines 480 to 483 in 9ef9173

	path = make_zmq_path("tcp", host, port + p_remote_rank)
	logger.debug("Querying metadata on path: %s at remote rank %s",
	path, p_remote_rank)
	_ = handshake(path, p_remote_rank)

The second handshake execution, however, will not be skipped—all checks will be performed as expected

[NickLucche]

yeah absolutely, this was also the original intent: handshake with rank0 is to get tp_size.
But skipping a bunch of extra asserts on top won't get us to solving the core issue.

[GuanLuo]

The second issue you reported assert self.num_blocks >= nixl_agent_meta.num_blocks is a separate condition I was already discussing with @ptarasiewiczNV . The rationale is that the consumer should have enough space to fit all the blocks from the producer. We can easily relax that one with some other lower bound as long as the kv cache is not being transferred in its entirity with nixl.

@NickLucche I am hitting the same issue as well, I also think that this assertion is too restricted. The number of blocks allocated depends on the GPU usage at the time and can easily fail this assertion on the same TP configs.. In my case it is off by 8 blocks over 5800

logger.info(f"[DEBUG] {self.num_blocks} v.s. {nixl_agent_meta.num_blocks}")
assert self.num_blocks >= nixl_agent_meta.num_blocks
# log
(RayWorkerWrapper pid=3776752) INFO 06-11 15:53:59 [nixl_connector.py:740] [DEBUG] 5800 v.s. 5808

[NickLucche]

#19532