feat: add data parallel rank to KVEventBatch #18925

PeaBrane · 2025-05-29T20:34:19Z

Overview:

Add data_parallel_rank to KvEventBatch, so it can be published along with the block storage / removal events
Perform a port offset based on the dp rank for EventPublisher so there is no port conflict
Allow the mock subscriber to subscribe to multiple ports for testing
Added a test for testing the behavior of the publisher under async multi-proc, dp=2 and tp=2

Resolves #18924

github-actions · 2025-05-29T20:34:27Z

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

🚀

njhill

Thanks @PeaBrane looks good to me, just a few minor comments.

vllm/v1/core/sched/scheduler.py

vllm/distributed/kv_events.py

tests/v1/engine/test_engine_core_client.py

vllm/config.py

PeaBrane · 2025-05-31T04:17:26Z

@njhill thanks for such a careful review! I accepted all your suggestions / edits. Changes should scheduler.py should be minimal now, and most of the changes moved into kv_events.py

njhill · 2025-06-01T18:14:37Z

vllm/config.py

+    data_parallel_rank: int = 0
+    """The DP rank where the event is emitted from
+    """


I guess we don't expect this to be set explicitly right? It should always be set to the value of data_parallel_rank from parallel_config. It's here just to pass in to the publisher factory rather than something to be provided as part of the "external" KVEventsConfig like the other values.

That should be reflected in the comment at least, and perhaps could rename it to _data_parallel_rank or even pass to the factory as a second arg rather than being part of this config class?

I agree, this is a similar pattern to the Metrics stuff where the only argument to the stat logger factory is a config and a dp rank.

StatLoggerFactory = Callable[[VllmConfig, int], "StatLoggerBase"]

I think it would make the most sense to remove from the config and add to the create() function

@classmethod def create(cls, config: Optional[KVEventsConfig], dp_rank: int) -> EventPublisher:

The only small problem here is with the StatLogger the user owns the Factory method, while here it is written in vLLM, so we have to decide if we want to always pass the dp_rank to the Publisher or how exactly to manage that.

can make it an optional kwarg, lemme try that

njhill

Thanks @PeaBrane @alec-flowers.

Just one small thing to fix I think.

vllm/distributed/kv_events.py

njhill · 2025-06-03T00:40:06Z

vllm/distributed/kv_events.py

+    # Create a modified config with port offsetting
+    modified_config = copy.deepcopy(kv_events_config)
+
+    # Apply port offsetting to the endpoint
+    original_endpoint = modified_config.endpoint
+    modified_config.endpoint = _offset_endpoint_port(original_endpoint,
+                                                     data_parallel_rank)
+    modified_config.data_parallel_rank = data_parallel_rank
+
+    # Apply port offsetting to the replay_endpoint if it exists
+    if modified_config.replay_endpoint:
+        original_replay_endpoint = modified_config.replay_endpoint
+        modified_config.replay_endpoint = _offset_endpoint_port(
+            original_replay_endpoint, data_parallel_rank)


I think this would be fine as a follow-on if we want to get this merged but I'm thinking it may be better to also move the port offsetting inside the publisher implementation since it's kind of ZMQ-specific and tied to the DP rank that we're now passing. WDYT?

Yes it shouldn't be outside. It would break any other publisher.

Signed-off-by: PeaBrane <yanrpei@gmail.com>

Co-authored-by: Nick Hill <nhill@redhat.com> Signed-off-by: Yan Ru Pei <yanrpei@gmail.com> Signed-off-by: alec-flowers <aflowers@nvidia.com>

Signed-off-by: PeaBrane <yanrpei@gmail.com>

Signed-off-by: alec-flowers <aflowers@nvidia.com>

Co-authored-by: Nick Hill <nhill@redhat.com> Signed-off-by: alec-flowers <aflowers@nvidia.com>

Signed-off-by: alec-flowers <aflowers@nvidia.com>

njhill

Thanks @PeaBrane @alec-flowers

Signed-off-by: PeaBrane <yanrpei@gmail.com>

tests/v1/engine/test_engine_core_client.py

Signed-off-by: Nick Hill <nhill@redhat.com>

njhill · 2025-06-03T20:29:53Z

Not sure why there's a bunch of apparently unrelated test failures now :-(

mergify bot added the v1 label May 29, 2025

PeaBrane force-pushed the rupei/kv-events-dp branch from 78ac551 to aedd6a3 Compare May 29, 2025 20:38

PeaBrane marked this pull request as ready for review May 30, 2025 02:44

PeaBrane requested review from WoosukKwon, robertgshaw2-redhat, njhill, ywang96, comaniac and alexm-redhat as code owners May 30, 2025 02:44

PeaBrane mentioned this pull request May 30, 2025

[FEATURE]: more general router ai-dynamo/dynamo#1289

Open

njhill reviewed May 30, 2025

View reviewed changes

njhill reviewed May 31, 2025

View reviewed changes

vllm/config.py Outdated Show resolved Hide resolved

PeaBrane force-pushed the rupei/kv-events-dp branch from b85d1a8 to 4752d38 Compare May 31, 2025 03:57

PeaBrane requested review from hmellor, DarkLight1337, simon-mo, zhuohan123 and youkaichao as code owners May 31, 2025 03:57

mergify bot added the documentation Improvements or additions to documentation label May 31, 2025

njhill reviewed Jun 1, 2025

View reviewed changes

njhill reviewed Jun 2, 2025

View reviewed changes

vllm/distributed/kv_events.py Outdated Show resolved Hide resolved

vllm/distributed/kv_events.py Outdated Show resolved Hide resolved

alec-flowers force-pushed the rupei/kv-events-dp branch from 72a5e0c to c21d367 Compare June 2, 2025 19:42

njhill reviewed Jun 3, 2025

View reviewed changes

njhill added the ready ONLY add when PR is ready to merge/full CI is needed label Jun 3, 2025

PeaBrane added 4 commits June 2, 2025 18:23

publish dp name in kv event, and only publish if master in dp group

662c7dc

Signed-off-by: PeaBrane <yanrpei@gmail.com>

use dp and tp 2 test_kv_cache_events

591b26d

Signed-off-by: PeaBrane <yanrpei@gmail.com>

works for dp > 1 but not tp > 1

a79414b

Signed-off-by: PeaBrane <yanrpei@gmail.com>

cleanups

16e8650

Signed-off-by: PeaBrane <yanrpei@gmail.com>

PeaBrane added 6 commits June 2, 2025 18:23

restore dp and tp

d9de161

Signed-off-by: PeaBrane <yanrpei@gmail.com>

loop until fully done

29b2427

Signed-off-by: PeaBrane <yanrpei@gmail.com>

revert unneeded master checks

7f5474c

Signed-off-by: PeaBrane <yanrpei@gmail.com>

dp_rank does not have to be an attribute

34b118f

Signed-off-by: PeaBrane <yanrpei@gmail.com>

restore non-empty events check

68b9e65

Signed-off-by: PeaBrane <yanrpei@gmail.com>

let publisher handle dp rank identification

6f973a4

Signed-off-by: PeaBrane <yanrpei@gmail.com>

alec-flowers force-pushed the rupei/kv-events-dp branch 2 times, most recently from 9f03056 to ed77172 Compare June 3, 2025 01:25

PeaBrane and others added 7 commits June 2, 2025 18:34

no need for elif with return guard

7d92408

Co-authored-by: Nick Hill <nhill@redhat.com> Signed-off-by: Yan Ru Pei <yanrpei@gmail.com> Signed-off-by: alec-flowers <aflowers@nvidia.com>

same as above

f6f958e

Co-authored-by: Nick Hill <nhill@redhat.com> Signed-off-by: Yan Ru Pei <yanrpei@gmail.com> Signed-off-by: alec-flowers <aflowers@nvidia.com>

fix test loop logic

07aff59

Co-authored-by: Nick Hill <nhill@redhat.com> Signed-off-by: Yan Ru Pei <yanrpei@gmail.com> Signed-off-by: alec-flowers <aflowers@nvidia.com>

idiomatic async timeout, default dp_rank to 0

013f060

Signed-off-by: PeaBrane <yanrpei@gmail.com>

address reviewer comments

ec0b817

Signed-off-by: alec-flowers <aflowers@nvidia.com>

Update vllm/distributed/kv_events.py

eaf4bc4

Co-authored-by: Nick Hill <nhill@redhat.com> Signed-off-by: alec-flowers <aflowers@nvidia.com>

Update vllm/distributed/kv_events.py

7c94dc7

Co-authored-by: Nick Hill <nhill@redhat.com> Signed-off-by: alec-flowers <aflowers@nvidia.com>

alec-flowers force-pushed the rupei/kv-events-dp branch 2 times, most recently from 6ae7ef2 to 8bf0885 Compare June 3, 2025 01:43

refactor: dp port mainpulations moved to zmqpublisher class

9c40bc9

Signed-off-by: alec-flowers <aflowers@nvidia.com>

alec-flowers force-pushed the rupei/kv-events-dp branch from 8bf0885 to 9c40bc9 Compare June 3, 2025 03:28

alec-flowers mentioned this pull request Jun 3, 2025

feat: sglang integration of dynamo event manager ai-dynamo/dynamo#1323

Open

njhill approved these changes Jun 3, 2025

View reviewed changes

fix async timeout and multi gpu marker

3119617

Signed-off-by: PeaBrane <yanrpei@gmail.com>

njhill reviewed Jun 3, 2025

View reviewed changes

tests/v1/engine/test_engine_core_client.py Show resolved Hide resolved

run test_kv_cache_events_dp test with 4 GPUs

399005a

Signed-off-by: Nick Hill <nhill@redhat.com>

mergify bot added the ci/build label Jun 3, 2025

njhill enabled auto-merge (squash) June 3, 2025 18:05

simon-mo merged commit b712be9 into vllm-project:main Jun 4, 2025
84 of 91 checks passed

PeaBrane mentioned this pull request Jun 4, 2025

feat: router supporting intra-worker dp routing ai-dynamo/dynamo#1285

Draft

Uh oh!

feat: add data parallel rank to KVEventBatch #18925

feat: add data parallel rank to KVEventBatch #18925

Uh oh!

Conversation

PeaBrane commented May 29, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented May 29, 2025

Uh oh!

njhill left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

PeaBrane commented May 31, 2025

Uh oh!

njhill Jun 1, 2025

Choose a reason for hiding this comment

Uh oh!

alec-flowers Jun 2, 2025

Choose a reason for hiding this comment

Uh oh!

PeaBrane Jun 2, 2025

Choose a reason for hiding this comment

Uh oh!

njhill left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

njhill Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

alec-flowers Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

njhill left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

njhill commented Jun 3, 2025

Uh oh!

Uh oh!

Uh oh!

PeaBrane commented May 29, 2025 •

edited by github-actions bot

Loading