[P/D] Update output of `get_finished` to newly defined class `KVConnectorFinishOutput` by lk-chen · Pull Request #21790 · vllm-project/vllm

lk-chen · 2025-07-28T22:26:23Z

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.

This PR is on top of #21785

Purpose

We need to pipe more information from kv_connector implementations to model_runner (e.g. finished_loading_dict in #21534 needed by #19329 , and metrics to be added in #21784 ). It's not sustainable to append ad-hoc fields along with existing ones everywhere (ModelRunnerOutput, IntermediateTensors etc.). This PR aggregates all these info into a new class.

Test Plan

CI unit tests

Test Result

(Optional) Documentation Update

Signed-off-by: Linkun <github@lkchen.net>

Signed-off-by: Linkun Chen <github@lkchen.net>

github-actions · 2025-07-28T22:26:34Z

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

🚀

mergify · 2025-07-28T22:27:04Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @lk-chen.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

gemini-code-assist

Code Review

This pull request introduces a new KVConnectorFinishOutput class to standardize the output from kv_connector implementations, which is a great step towards improving code clarity and maintainability. The removal of the V0 connector code is also a welcome cleanup.

I've found a few critical issues that need to be addressed. The __init__ method of the new KVConnectorFinishOutput class has an incorrect return type hint. Additionally, several connectors are not providing all the required arguments when instantiating this new class, which would lead to runtime errors. I've provided specific suggestions for these issues in the comments. Addressing these will ensure the refactoring is robust and correct.

gemini-code-assist · 2025-07-28T22:28:07Z

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

+            finished_sending: set[str],
+            finished_recving: set[str],
+            finished_loading_num_tokens: dict[str, int],
+        ) -> "KVConnectorBase_V1.KVConnectorFinishOutput":


The __init__ method in Python should always return None. Specifying a return type hint other than None is incorrect and will cause a TypeError at runtime.

Suggested change

) -> "KVConnectorBase_V1.KVConnectorFinishOutput":

) -> None:

gemini-code-assist · 2025-07-28T22:28:07Z

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

+        finished_sending, finished_recving = self._lmcache_engine.get_finished(
+            finished_req_ids)
+        return KVConnectorBase_V1.KVConnectorFinishOutput(
+            finished_sending=finished_sending,
+            finished_recving=finished_recving)


The KVConnectorFinishOutput constructor is missing the required keyword argument finished_loading_num_tokens. This will cause a TypeError at runtime.

return KVConnectorBase_V1.KVConnectorFinishOutput( finished_sending=finished_sending, finished_recving=finished_recving, finished_loading_num_tokens={})

gemini-code-assist · 2025-07-28T22:28:07Z

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

        finished_sending: set[str] = set()
        finished_recving: set[str] = set()
        for c in self._connectors:


With the change of the return type of get_finished to KVConnectorFinishOutput, the unpacking on line 114 (sending, recving = c.get_finished(finished_req_ids)) will raise a TypeError because get_finished no longer returns a tuple. You need to update the loop body to handle the new return type.

For example:

output = c.get_finished(finished_req_ids) sending = output.finished_sending recving = output.finished_recving

gemini-code-assist · 2025-07-28T22:28:08Z

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

+        return KVConnectorBase_V1.KVConnectorFinishOutput(
+            finished_sending=finished_sending,
+            finished_recving=finished_recving)


The KVConnectorFinishOutput constructor is missing the required keyword argument finished_loading_num_tokens. This will cause a TypeError at runtime.

Suggested change

return KVConnectorBase_V1.KVConnectorFinishOutput(

finished_sending=finished_sending,

finished_recving=finished_recving)

return KVConnectorBase_V1.KVConnectorFinishOutput(

finished_sending=finished_sending,

finished_recving=finished_recving,

finished_loading_num_tokens={})

gemini-code-assist · 2025-07-28T22:28:08Z

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

+        return KVConnectorBase_V1.KVConnectorFinishOutput(
+            finished_sending=done_sending, finished_recving=done_recving)


The KVConnectorFinishOutput constructor is missing the required keyword argument finished_loading_num_tokens. This will cause a TypeError at runtime.

Suggested change

return KVConnectorBase_V1.KVConnectorFinishOutput(

finished_sending=done_sending, finished_recving=done_recving)

return KVConnectorBase_V1.KVConnectorFinishOutput(

finished_sending=done_sending, finished_recving=done_recving, finished_loading_num_tokens={})

gemini-code-assist · 2025-07-28T22:28:08Z

vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py

+        finished_sending, finished_recving = self.p2p_nccl_engine.get_finished(
+            finished_req_ids, no_compile_layers)
+        return KVConnectorBase_V1.KVConnectorFinishOutput(
+            finished_sending=finished_sending or set(),
+            finished_recving=finished_recving or set(),
+        )


The KVConnectorFinishOutput constructor is missing the required keyword argument finished_loading_num_tokens. This will cause a TypeError at runtime.

return KVConnectorBase_V1.KVConnectorFinishOutput( finished_sending=finished_sending or set(), finished_recving=finished_recving or set(), finished_loading_num_tokens={}, )

Signed-off-by: Linkun Chen <github@lkchen.net>

sdavidbd · 2025-07-29T17:25:02Z

Hi @lk-chen, I completely agree on the need to encapsulate KV connector output in a structured way. I introduced a similar concept in #19330 — see KVConnectorOutput.

In that PR, I didn’t yet propagate this class through ModelRunnerOutput or IntermediateTensors, as doing so would require broader refactoring that goes beyond the scope of the current change. There are also ongoing efforts to generalize and unify the structure of these fields, which are being discussed in #19330 and #19555. Would be great to align these directions going forward.

Signed-off-by: Linkun Chen <github@lkchen.net>

mergify · 2025-08-05T02:44:06Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @lk-chen.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

lk-chen · 2025-08-05T17:50:55Z

KVConnectorOutput introduced in #21980

lk-chen added 4 commits July 25, 2025 13:01

wip

c6f7fe5

Signed-off-by: Linkun <github@lkchen.net>

rm legacy code

c4104b1

Signed-off-by: Linkun Chen <github@lkchen.net>

Merge branch 'nixl_deprecate_v0' into nixl_xfer_log

6c7e861

Signed-off-by: Linkun Chen <github@lkchen.net>

fix test

4c776fd

Signed-off-by: Linkun Chen <github@lkchen.net>

lk-chen requested review from WoosukKwon, alexm-redhat, comaniac, njhill, robertgshaw2-redhat and ywang96 as code owners July 28, 2025 22:26

lk-chen removed request for WoosukKwon, comaniac and ywang96 July 28, 2025 22:27

mergify bot added the v1 label Jul 28, 2025

mergify bot added tpu Related to Google TPUs needs-rebase labels Jul 28, 2025

gemini-code-assist bot reviewed Jul 28, 2025

View reviewed changes

Merge branch 'main' into nixl_xfer_log

eb1ff49

Signed-off-by: Linkun Chen <github@lkchen.net>

mergify bot removed the needs-rebase label Jul 28, 2025

clean up

c3c8cc8

Signed-off-by: Linkun Chen <github@lkchen.net>

lk-chen added the ready ONLY add when PR is ready to merge/full CI is needed label Jul 29, 2025

lk-chen added 3 commits July 29, 2025 10:18

clean up

a7e5afc

Signed-off-by: Linkun Chen <github@lkchen.net>

not mention v0/v1

c1dd7ee

Signed-off-by: Linkun Chen <github@lkchen.net>

not mention v0/v1

9587e2c

Signed-off-by: Linkun Chen <github@lkchen.net>

lk-chen added 3 commits July 29, 2025 12:15

Merge branch 'main' into nixl_deprecate_v0

9d558d2

Merge branch 'nixl_deprecate_v0' into nixl_xfer_log

d42a20c

remove todo

4bb3330

Signed-off-by: Linkun Chen <github@lkchen.net>

lk-chen added 4 commits July 29, 2025 13:05

yapf

c1dae5a

Signed-off-by: Linkun Chen <github@lkchen.net>

Merge branch 'nixl_deprecate_v0' into nixl_xfer_log

164e98c

remove v0 test

b557f22

Signed-off-by: Linkun Chen <github@lkchen.net>

Merge branch 'nixl_deprecate_v0' into nixl_xfer_log

b1dcd37

Signed-off-by: Linkun Chen <github@lkchen.net>

lk-chen force-pushed the nixl_xfer_log branch from 64d2cc7 to b1dcd37 Compare July 30, 2025 00:12

lk-chen added 4 commits July 30, 2025 15:43

disable disagg test, this is PyNccl only

aa997cd

Signed-off-by: Linkun Chen <github@lkchen.net>

rename

26bc5b3

Signed-off-by: Linkun Chen <github@lkchen.net>

Merge branch 'main' into nixl_deprecate_v0

5b6a9aa

Signed-off-by: Linkun Chen <github@lkchen.net>

Merge branch 'nixl_deprecate_v0' into nixl_xfer_log

8374dfe

Signed-off-by: Linkun Chen <github@lkchen.net>

mergify bot added the ci/build label Jul 30, 2025

mergify bot added the needs-rebase label Aug 5, 2025

lk-chen closed this Aug 5, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[P/D] Update output of `get_finished` to newly defined class `KVConnectorFinishOutput`#21790

[P/D] Update output of `get_finished` to newly defined class `KVConnectorFinishOutput`#21790
lk-chen wants to merge 20 commits intovllm-project:mainfrom
lk-chen:nixl_xfer_log

lk-chen commented Jul 28, 2025 •

edited by github-actions bot

Loading

Uh oh!

github-actions bot commented Jul 28, 2025

Uh oh!

mergify bot commented Jul 28, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jul 28, 2025

Uh oh!

gemini-code-assist bot Jul 28, 2025

Uh oh!

gemini-code-assist bot Jul 28, 2025

Uh oh!

gemini-code-assist bot Jul 28, 2025

Uh oh!

gemini-code-assist bot Jul 28, 2025

Uh oh!

gemini-code-assist bot Jul 28, 2025

Uh oh!

sdavidbd commented Jul 29, 2025

Uh oh!

mergify bot commented Aug 5, 2025

Uh oh!

lk-chen commented Aug 5, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	) -> "KVConnectorBase_V1.KVConnectorFinishOutput":
	) -> None:

		return KVConnectorBase_V1.KVConnectorFinishOutput(
		finished_sending=done_sending, finished_recving=done_recving)

Uh oh!

Conversation

lk-chen commented Jul 28, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Essential Elements of an Effective PR Description Checklist

Purpose

Test Plan

Test Result

(Optional) Documentation Update

Uh oh!

github-actions bot commented Jul 28, 2025

Uh oh!

mergify bot commented Jul 28, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jul 28, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 28, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 28, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 28, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 28, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 28, 2025

Choose a reason for hiding this comment

Uh oh!

sdavidbd commented Jul 29, 2025

Uh oh!

mergify bot commented Aug 5, 2025

Uh oh!

lk-chen commented Aug 5, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

lk-chen commented Jul 28, 2025 •

edited by github-actions bot

Loading