[BugFix][Nixl][PD] Fix heterogenous TP

[NickLucche]

In the attempt to deprecate V0 here #21785, we introduced a subtle bug that would prevent a KVConnector from forwarding its KV cache layout preference.
This would result in the following log:

(EngineCore_0 pid=1204777) INFO 08-11 16:32:10 [utils.py:113] Connectors do not specify a kv cache layout, defaulting to NHD.

which, in the case of NixlConnector for D TP != P TP means breaking one of the single core assumption about the layout. This results in a garbled output.

This PR brings things back to

(EngineCore_0 pid=1306559) INFO 08-11 17:17:38 [nixl_connector.py:151] NixlConnector setting KV cache layout to HND for better xfer performance.

Also added @njhill fix for KVConnectorOutput or it would break even earlier than observing the correctness regression.

[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.

🚀

[NickLucche]

Thanks @wseaton for spotting the correctness issue 🙏🏻

[gemini-code-assist]

Code Review

This pull request effectively addresses a critical bug that caused garbled output in heterogeneous tensor parallelism setups. The fix correctly determines the KV cache layout by refactoring the KVConnectorFactory to expose a get_connector_class method, which is a clean and effective solution. Additionally, the inclusion of a null check in KVOutputAggregator is a good defensive measure that prevents potential runtime errors. The changes are well-implemented, improving both the correctness and maintainability of the code.

[njhill]

Thanks @NickLucche

[DarkLight1337]

Can you merge from main to fix CI?

[NickLucche]

@DarkLight1337 rebased

[DarkLight1337]

Please fix pre-commit

[NickLucche]

Why isn't the pre-commit CI job showing the diff of the changes?

I think yapf and isort are conflicting , yapf changes the file and isort brings it back to how it was lol so no changes are left but pre-commit still fails

(tmp) ➜  vllm git:(pd-fix-hetero) ✗ pre-commit run yapf --files vllm/distributed/kv_transfer/kv_connector/factory.py  
yapf.....................................................................Failed
- hook id: yapf
- files were modified by this hook

Reformatting vllm/distributed/kv_transfer/kv_connector/factory.py

(tmp) ➜  vllm git:(pd-fix-hetero) ✗ pre-commit run isort --files vllm/distributed/kv_transfer/kv_connector/factory.py 
isort....................................................................Failed
- hook id: isort
- files were modified by this hook

Fixing /home/nicolo/llmd/vllm/vllm/distributed/kv_transfer/kv_connector/factory.py

(tmp) ➜  vllm git:(pd-fix-hetero) ✗ gd vllm/distributed/kv_transfer/kv_connector/factory.py | cat                                            

@hmellor what do you think?

Also I don't like that isort behaves differently based on whether the imports are installed or not in your venv, which is causing the CI to not match my local env..

[DarkLight1337]

If they are conflicting, you can disable yapf for those lines.

[NickLucche]

ok done, still it would be nice to sort this out with configs

[DarkLight1337]

We are slowly migrating to using ruff to do both so there will be no more conflicts. Hopefully the process will speed up once V0 is removed

[NickLucche]

This is failing on main

v1/kv_connector/unit/test_remote_decode_lifecycle.py::test_short_prompt_lifecycle

but it really shouldn't, is someone looking into it already?

[NickLucche]

I put up another PR for that failing test.. #22727 @DarkLight1337

[hmellor]

Hopefully the process will speed up once V0 is removed

That's the plan! While so much is changing in tests/ and vllm/ it'd cause too many conflicts. Everything outside of these directorie should already be using Ruff though