[DP] Support external DP Load Balancer mode #19790

njhill · 2025-06-18T05:57:13Z

See https://docs.google.com/document/d/1mSYsWQEbp4Oq50ghFWUun7OgagbuUJivAC8Ds-pu-xU/edit?tab=t.0#heading=h.ggq72ssewhm3

This also includes some de-duplication of the core engine and dp coordinator processorchestration logic, consolidated into a launch_core_engines context manager util method.

This is called:

from the run_multi_api_server in serve.py when there are multiple API server processes, in which case RPC connection details are passed into the engine core client running within each API server..
within the MPClient EngineCoreClient superclass otherwise (encapsulated behind the engine core client).

Some utility classes/functions related to engine process management moved from v1/utils.py to v1/engine/utils.py to avoid circular imports.

github-actions · 2025-06-18T05:57:22Z

👋 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

Summary of Changes

Hello @njhill, 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!

This pull request introduces comprehensive support for an external data parallel load balancer mode, allowing vLLM deployments to integrate more seamlessly with external orchestration systems. It refines the internal communication patterns, configures instances for direct local engine interaction, and adjusts stats reporting to align with an externally managed load balancing strategy.

Highlights

External Data Parallel Load Balancer Mode: Introduced a new data_parallel_external_lb configuration option and a --data-parallel-rank CLI argument to enable an external load balancer mode for data parallelism. When enabled, vLLM instances are expected to be managed by an external system for request distribution.
Decoupled Coordinator and Engine Communication: In external LB mode, the data parallel coordinator (responsible for internal load balancing and stats) now only runs on the rank 0 instance. Non-rank-0 instances will communicate directly with their local engine, streamlining the setup for external orchestration.
Refined Engine Handshake Process: The engine startup handshake mechanism has been updated to support the new external LB mode. This includes a two-stage handshake for non-rank-0 instances to retrieve both global coordinator information (from rank 0) and local client I/O addresses.
Adjusted Stats Reporting: Engines operating in external LB mode will no longer publish detailed request queue statistics to the coordinator. Instead, they will only report high-level wave completion and running state changes, reducing overhead when an external system handles load metrics.
Client-Side Adaptations for DP Modes: The MPClient hierarchy has been refactored to differentiate between internal (load-balancing) and external (direct-to-local-engine) data parallel modes. A new DPLBAsyncMPClient class was introduced to specifically handle the internal load-balancing logic.

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.

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

gemini-code-assist

Code Review

This pull request adds support for an external Data Parallel (DP) Load Balancer mode. The changes are comprehensive, touching configuration, argument parsing, server entrypoints, and core engine components. The introduction of a dual handshake mechanism for non-rank-0 external LB instances and the refactoring of client types (DPAsyncMPClient vs. DPLBAsyncMPClient) are notable improvements for handling different DP scenarios. The overall logic appears sound. One minor issue with a docstring in vllm/config.py needs addressing, and a check for data_parallel_rank in EngineArgs.

vllm/config.py

vllm/engine/arg_utils.py

vllm/v1/engine/core.py

vllm/v1/engine/core_client.py

mergify · 2025-06-24T19:53:01Z

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

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

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

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

tlrmchlsmth · 2025-07-01T14:39:23Z

vllm/v1/engine/utils.py

Anything new in here, or is it all moved over from vllm/utils.py?

Most is moved but the launch_core_engines function is new. Sorry, I know moving stuff is crap for reviews. Let me double-check diffs of the parts other than that function.

tlrmchlsmth · 2025-07-01T16:10:20Z

vllm/entrypoints/cli/serve.py

+        # For dp ranks > 0 in external DP LB mode, we must delay the
+        # start of the API servers until the local engine is started
+        # (after the launcher context manager exits),
+        # since we get the front-end stats update address from the coordinator
+        # via the handshake with the local engine.
+        if dp_rank == 0 or not external_dp_lb:
+            # Start API servers using the manager.
+            api_server_manager = APIServerProcessManager(
+                **api_server_manager_kwargs)


Just wondering, why can't we also delay for the case that dp_rank == 0? Do all the stats get published through the 0th API server?

We can delay for rank == 0 case too, I guess that might be cleaner. We don't need to on rank 0 because it's where the coordinator is started and so we know in advance the zmq address that the API servers use to connect to the coordinator.

tlrmchlsmth · 2025-07-01T16:12:59Z

vllm/v1/engine/coordinator.py

+        # When in external LB mode, load stats aren't published, only changes
+        # to request wave / running state, so we don't need to rate-limit the
+        # updates to the front-end proc(s).
+        min_stats_update_interval_ms = 0 if external_lb else 100


Are we going to be able to remove the every 200ms metrics logging spam in llm-d with this change?

Unfortunately no, these stats are the ones for the internal load-balancing, independent of prometheus.

There may actually be a lot more such logs on aggregate since llm-d will need to hit the metrics endpoint on each of the ranks/pods instead of just one ... !

tlrmchlsmth · 2025-07-01T17:34:48Z

vllm/entrypoints/cli/serve.py

@@ -148,7 +144,7 @@ def signal_handler(signum, frame):
        start_index=args.data_parallel_start_rank,


Does data_parallel_start_rank need to be set at all with this approach?

No it doesn't, this line though is in the headless path (in run_headless function), used for remote nodes in the internal LB mode.

Maybe there should be a warning if it's set just to let the user know it doesn't do anything

Added some validation to fail if invalid CLI arg combos are used.

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

tlrmchlsmth

LGTM, nice work!

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

Last-minute bug introduced in vllm-project#19790 Signed-off-by: Nick Hill <nhill@redhat.com>

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

hidva · 2025-07-16T03:27:11Z

vllm/v1/engine/core.py

@@ -799,10 +867,18 @@ def _init_data_parallel(self, vllm_config: VllmConfig):
        from vllm.platforms import current_platform
        device_control_env_var = current_platform.device_control_env_var
        world_size = vllm_config.parallel_config.world_size
-        os.environ[device_control_env_var] = ",".join(


local_dp_rank is always 0, this may break nixl connector.

# modify the engine_id and append the local_dp_rank to it to ensure # that the kv_transfer_config is unique for each DP rank. vllm_config.kv_transfer_config.engine_id = ( f"{vllm_config.kv_transfer_config.engine_id}_dp{local_dp_rank}" )

@njhill

Thanks @hidva. Actually in this case I don't think it's a problem since a separate vllm serve command / top-level process is run for each DP rank and thus different engine ids will be generated for each.

mergify bot added frontend v1 labels Jun 18, 2025

gemini-code-assist bot reviewed Jun 18, 2025

View reviewed changes

vllm/config.py Outdated Show resolved Hide resolved

vllm/engine/arg_utils.py Show resolved Hide resolved

vllm/v1/engine/core.py Show resolved Hide resolved

vllm/v1/engine/core_client.py Outdated Show resolved Hide resolved

vllm/v1/engine/core_client.py Outdated Show resolved Hide resolved

mergify bot added the needs-rebase label Jun 24, 2025

njhill force-pushed the ext-lb-dp branch from 6c9e48d to 88c0928 Compare June 24, 2025 20:07

mergify bot removed the needs-rebase label Jun 24, 2025

njhill marked this pull request as ready for review June 25, 2025 04:41

njhill requested review from WoosukKwon, robertgshaw2-redhat, ywang96, comaniac, alexm-redhat and aarnphm as code owners June 25, 2025 04:41

njhill force-pushed the ext-lb-dp branch from 7a0c86d to 833284a Compare June 25, 2025 05:27

njhill added the ready label Jun 25, 2025

njhill added 5 commits June 26, 2025 13:40

[DP] Support external DP Load Balancer mode

491f0aa

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

deduplicate/unify engine launching logic

666c197

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

move utils not directly related to engines back to v1/utils.py

85dfe7e

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

fix: ignore coordinator XPUB unsubscribe messages

a73c2c0

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

add CI test

a195cc5

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

njhill force-pushed the ext-lb-dp branch from 833284a to a195cc5 Compare June 26, 2025 20:44

njhill requested review from simon-mo, youkaichao, mgoin, tlrmchlsmth, houseroad and hmellor as code owners June 26, 2025 20:44

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

fix typing in test

04f2f97

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

tlrmchlsmth mentioned this pull request Jun 27, 2025

[RFC]: Data Parallel Attention and Expert Parallel MoEs #16037

Open

22 tasks

njhill added 6 commits June 26, 2025 21:13

fix test_engine_core_proc_instantiation_cuda_empty

2025b49

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

fix test formatting

13de014

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

Merge remote-tracking branch 'origin/main' into ext-lb-dp

120f9a2

Merge remote-tracking branch 'origin/main' into ext-lb-dp

05640de

test fixes

3cc7bd1

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

fix test import

410ee4d

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

tlrmchlsmth reviewed Jul 1, 2025

View reviewed changes

fix test properly hopefully

d93b903

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

tlrmchlsmth approved these changes Jul 2, 2025

View reviewed changes

validate mutually exclusive command line args

d0b5450

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

vllm-bot merged commit 657f2f3 into vllm-project:main Jul 2, 2025
101 of 103 checks passed

njhill deleted the ext-lb-dp branch July 2, 2025 18:12

njhill added a commit to njhill/vllm that referenced this pull request Jul 2, 2025

[BugFix] Fix DP headless mode arg validation

6714b87

Last-minute bug introduced in vllm-project#19790 Signed-off-by: Nick Hill <nhill@redhat.com>

njhill mentioned this pull request Jul 2, 2025

[BugFix] Fix DP headless mode arg validation #20398

Merged

huydhn pushed a commit to huydhn/vllm that referenced this pull request Jul 8, 2025

[DP] Support external DP Load Balancer mode (vllm-project#19790)

16ee279

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

hidva reviewed Jul 16, 2025

View reviewed changes

		@@ -148,7 +144,7 @@ def signal_handler(signum, frame):
		start_index=args.data_parallel_start_rank,

Uh oh!

[DP] Support external DP Load Balancer mode #19790

[DP] Support external DP Load Balancer mode #19790

Conversation

njhill commented Jun 18, 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 Jun 18, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mergify bot commented Jun 24, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tlrmchlsmth left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

hidva Jul 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

njhill commented Jun 18, 2025 •

edited by github-actions bot

Loading

hidva Jul 16, 2025 •

edited

Loading