[BugFix] pp cannot run successfully under NixlConnector

[R2-Y]

[BugFix] pp cannot run successfully under NixlConnector

Purpose

I previously raised issue #22430 regarding NixlConnector failing with Pipeline Parallelism (PP) + Prefill-Decode (PD) configurations, but it hasn't received attention from the community yet. Since I need PP functionality for my project and my benchmarks show that NixlConnector delivers significantly better performance than LMCacheConnector, I'm motivated to contribute a fix.

My testing on Qwen 2.5 7B demonstrates NixlConnector's clear advantage:

Configuration	TTFT (ms)	Improvement
PP1 + TP2 + PD + LMCache	13,900.61	Baseline
PP1 + TP2 + PD + Nixl	10,333.13	25.7% faster

The error occurs during KV cache initialization when nixl_wrapper.register_memory() is called in a Pipeline Parallelism setup. After investigating the codebase, I discovered that the current implementation only considers Tensor Parallelism (TP) when initializing KV caches.

When registering cache data with NIXL, different PP stages that share the same TP rank are assigned to the same device ID. This causes device conflicts during PP memory registration, leading to the NIXL_ERR_BACKEND failure. Therefore, I modified the device ID assignment from using TP rank to automatically detecting the current device using torch.cuda.current_device().

Test Plan

Tested by running the model under PP + NixlConnector for both PD disaggregation and PD aggregation scenarios.

Test Result

The error no longer occurs.

(Optional) Documentation Update

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[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 issue where Pipeline Parallelism (PP) could not run successfully with the NixlConnector. The core of the change is to use the device ID from torch.cuda.current_device() instead of the tensor parallelism rank (tp_rank) to uniquely identify workers. This correctly resolves device and port conflicts for workers within a single instance when using PP. However, this change introduces a critical issue for cross-instance communication. The client-side handshake logic for connecting to a remote producer instance has not been updated and still relies on the tp_rank to determine the port, creating a mismatch that will lead to connection failures when the producer uses pipeline parallelism. This side effect needs to be addressed.

[R2-Y]

/gemini review

[gemini-code-assist]

Code Review

This pull request aims to fix an issue where Pipeline Parallelism (PP) fails with NixlConnector. The core of the fix is to use the correct device ID (torch.cuda.current_device()) instead of the tensor parallel rank for NIXL memory registration, which correctly resolves device conflicts in a PP setup. Additionally, the side-channel port calculation for worker processes is updated to include the pipeline parallel rank, ensuring unique listening ports for each worker and preventing port collisions. While these changes are necessary for PP support, the update to the worker's listening port is not reflected in the scheduler's logic, which will cause connection failures in Prefill-Decode disaggregation scenarios. This is a critical issue that needs to be addressed.

[R2-Y]

/gemini review

[gemini-code-assist]

Code Review

This pull request provides a crucial bugfix for running Pipeline Parallelism (PP) with the NixlConnector. The core issue, as you've correctly identified, was the use of the tensor parallelism (TP) rank as a device identifier for NIXL memory registration, which is not unique across different pipeline stages. Your solution to use torch.cuda.current_device() to get the actual device ID is the correct approach to resolve this conflict. Furthermore, you've updated the side-channel port calculation to properly account for data, tensor, and pipeline parallelism, ensuring each worker gets a unique port. The addition of the get_pipeline_model_parallel_rank helper function is a clean way to support this. The changes are well-targeted and appear to correctly solve the reported problem. I found no issues of high or critical severity in this pull request.

[R2-Y]

@WoosukKwon @robertgshaw2-redhat @lk-chen Hi could you help to review?

[lk-chen]

You may want to add some basic tests in test_nixl_connector.py, TestNixlHandshake has some good examples.

and cc @NickLucche