@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / bd877162eb

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 19:00 UTC

[BugFix] Support online dense model DP without overhead (#30739)

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

Nick Hill committed Jan 2, 2026 at 07:36 UTC bd877162ebd7a2eb36961a8e49824422dd35bcdc

20 files changed +344 -145

tests/test_config.py

+48 -5

@@ -205,8 +205,8 @@ def test_default_pooling_type(model_id, default_pooling_type, pooling_type):

205 )

206 def test_moe_model_detection(model_id, expected_is_moe_model):

207 model_config = ModelConfig(model_id)

208 - # Just check that is_moe_model field exists and is a boolean

209 - assert model_config.is_model_moe() == expected_is_moe_model

208 + # Just check that is_moe field exists and is a boolean

209 + assert model_config.is_moe == expected_is_moe_model

210

211

212 @pytest.mark.parametrize(

@@ -224,7 +224,7 @@ def test_moe_model_detection(model_id, expected_is_moe_model):

224 def test_is_quantized(model_id, quantized):

225 model_config = ModelConfig(model_id)

226 # Just check that quantized field exists and is a boolean

227 - assert model_config.is_quantized() == quantized

227 + assert model_config.is_quantized == quantized

228

229

230 @pytest.mark.skipif(

@@ -925,7 +925,7 @@ def test_vllm_config_callable_defaults():

925 model_config=quantized_model, optimization_level=OptimizationLevel.O2

926 )

927 enable_if_quantized = lambda cfg: (

928 - cfg.model_config is not None and cfg.model_config.is_quantized()

928 + cfg.model_config is not None and cfg.model_config.is_quantized

929 )

930 assert enable_if_quantized(config_quantized) is True

931 assert enable_if_quantized(config_no_model) is False

@@ -936,7 +936,7 @@ def test_vllm_config_callable_defaults():

936 model_config=moe_model, optimization_level=OptimizationLevel.O2

937 )

938 enable_if_sequential = lambda cfg: (

939 - cfg.model_config is not None and not cfg.model_config.is_model_moe()

939 + cfg.model_config is not None and not cfg.model_config.is_moe

940 )

941 assert enable_if_sequential(config_moe) is False

942 assert enable_if_sequential(config_quantized) is True

@@ -1050,3 +1050,46 @@ def test_scheduler_config_init():

1050 with pytest.raises(AttributeError):

1051 # InitVar does not become an attribute

1052 print(SchedulerConfig.default_factory().max_model_len)

1053 +

1054 +

1055 +@pytest.mark.parametrize(

1056 + (

1057 + "model_id",

1058 + "data_parallel_size",

1059 + "external_lb",

1060 + "expected_needs_coordinator",

1061 + ),

1062 + [

1063 + # Non-MoE model with DP=1 should not need coordinator

1064 + ("facebook/opt-125m", 1, False, False),

1065 + # Non-MoE model with DP>1 internal LB should need coordinator

1066 + ("facebook/opt-125m", 2, False, True),

1067 + # Non-MoE model with DP>1 external LB should not need coordinator

1068 + ("facebook/opt-125m", 2, True, False),

1069 + # MoE model with DP=1 should not need coordinator

1070 + ("mistralai/Mixtral-8x7B-Instruct-v0.1", 1, False, False),

1071 + # MoE model with DP>1 internal LB should need both coordinator

1072 + # and wave coordination

1073 + ("mistralai/Mixtral-8x7B-Instruct-v0.1", 2, False, True),

1074 + # MoE model with DP>1 external LB needs coordinator for wave coordination

1075 + # (wave coordination runs in coordinator process)

1076 + ("mistralai/Mixtral-8x7B-Instruct-v0.1", 2, True, True),

1077 + ],

1078 +)

1079 +def test_needs_dp_coordination(

1080 + model_id,

1081 + data_parallel_size,

1082 + external_lb,

1083 + expected_needs_coordinator,

1084 +):

1085 + """Test that DP coordinator and wave coordination are configured correctly."""

1086 + from vllm.config import ParallelConfig

1087 +

1088 + model_config = ModelConfig(model_id)

1089 + parallel_config = ParallelConfig(

1090 + data_parallel_size=data_parallel_size,

1091 + data_parallel_external_lb=external_lb,

1092 + )

1093 + vllm_config = VllmConfig(model_config=model_config, parallel_config=parallel_config)

1094 +

1095 + assert vllm_config.needs_dp_coordinator == expected_needs_coordinator

tests/v1/engine/test_engine_core_client.py

@@ -133,6 +133,7 @@ def test_mp_client_uses_env_timeout(monkeypatch: pytest.MonkeyPatch):

133 parallel_config = SimpleNamespace(

134 data_parallel_size=1,

135 data_parallel_rank=0,

136 + data_parallel_index=0,

137 data_parallel_size_local=1,

138 data_parallel_rank_local=None,

139 data_parallel_hybrid_lb=False,

vllm/compilation/backends.py

+1 -1

@@ -630,7 +630,7 @@ class VllmBackend:

630 os.makedirs(cache_dir, exist_ok=True)

631 self.compilation_config.cache_dir = cache_dir

632 rank = vllm_config.parallel_config.rank

633 - dp_rank = vllm_config.parallel_config.data_parallel_rank

633 + dp_rank = vllm_config.parallel_config.data_parallel_index

634 local_cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}", self.prefix)

635 os.makedirs(local_cache_dir, exist_ok=True)

636 self.compilation_config.local_cache_dir = local_cache_dir

vllm/compilation/decorators.py

+1 -1

@@ -403,7 +403,7 @@ def _support_torch_compile(

403 )

404

405 rank = self.vllm_config.parallel_config.rank

406 - dp_rank = self.vllm_config.parallel_config.data_parallel_rank

406 + dp_rank = self.vllm_config.parallel_config.data_parallel_index

407 cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}")

408 aot_compilation_path = os.path.join(cache_dir, "model")

409 try:

vllm/config/model.py

+6 -7

@@ -642,7 +642,7 @@ class ModelConfig:

642 cls = "Transformers"

643 # If 'hf_config != hf_text_config' it's a nested config, i.e. multimodal

644 cls += "MultiModal" if self.hf_config != self.hf_text_config else ""

645 - cls += "MoE" if self.get_num_experts() > 1 else ""

645 + cls += "MoE" if self.is_moe else ""

646 # Check if the architecture we're wrapping has defaults

647 runner = None

648 task = None

@@ -1001,8 +1001,7 @@ class ModelConfig:

1001 self.enforce_eager = True

1002

1003 def _verify_with_expert_parallelism(self) -> None:

1004 - num_experts = self.get_num_experts()

1005 - if num_experts < 1:

1004 + if not self.is_moe:

1005 raise ValueError(

1006 "Number of experts in the model must be greater than 0 "

1007 "when expert parallelism is enabled."

@@ -1797,11 +1796,11 @@ class ModelConfig:

1796 logger.debug("Generative models support prefix caching.")

1797 return True

1798

1800 - def is_model_moe(

1801 - self,

1802 - ) -> bool:

1803 - return self.get_num_experts() > 1

1799 + @property

1800 + def is_moe(self) -> bool:

1801 + return self.get_num_experts() > 0

1802

1803 + @property

1804 def is_quantized(self) -> bool:

1805 return getattr(self.hf_config, "quantization_config", None) is not None

1806

vllm/config/parallel.py

+15

@@ -119,6 +119,8 @@ class ParallelConfig:

119 between local data parallel ranks, but an external LB balances

120 between vLLM nodes/replicas. Set explicitly in conjunction with

121 --data-parallel-start-rank."""

122 + is_moe_model: bool | None = None

123 + """Whether the deployed model is MoE (if known)."""

124 enable_expert_parallel: bool = False

125 """Use expert parallelism instead of tensor parallelism for MoE layers."""

126 enable_eplb: bool = False

@@ -255,6 +257,10 @@ class ParallelConfig:

257 Block_size should be divisible by cp_kv_cache_interleave_size.

258 """

259

260 + data_parallel_index: int = Field(init=False)

261 + """Equal to the data parallel rank but not used for torch process groups

262 + and not overridden for dense models."""

263 +

264 _api_process_count: int = Field(default=1, gt=0)

265 """

266 The number of API processes initialized.

@@ -466,6 +472,7 @@ class ParallelConfig:

472 "data_parallel_rank",

473 "data_parallel_rank_local",

474 "data_parallel_size_local",

475 + "data_parallel_index",

476 "data_parallel_backend",

477 "data_parallel_external_lb",

478 "data_parallel_hybrid_lb",

@@ -546,6 +553,14 @@ class ParallelConfig:

553 self.data_parallel_master_ip = envs.VLLM_DP_MASTER_IP

554 self.data_parallel_master_port = envs.VLLM_DP_MASTER_PORT

555

556 + if self.data_parallel_size > 1 and self.is_moe_model is False:

557 + raise ValueError(

558 + "Offline data parallel mode is not supported/useful"

559 + " for dense models."

560 + )

561 +

562 + self.data_parallel_index = self.data_parallel_rank

563 +

564 if self.distributed_executor_backend == "external_launcher":

565 os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"

566 logger.info("Disabling V1 multiprocessing for external launcher.")

vllm/config/vllm.py

+33 -8

@@ -343,6 +343,29 @@ class VllmConfig:

343 # i.e., batch_size <= self.compilation_config.max_cudagraph_capture_size

344 return self.compilation_config.bs_to_padded_graph_size[batch_size]

345

346 + @property

347 + def needs_dp_coordinator(self) -> bool:

348 + """

349 + Determine if the DPCoordinator process is needed.

350 +

351 + The DPCoordinator is needed in two cases:

352 + 1. For MoE models with DP > 1: to handle wave coordination

353 + (even in external LB mode, since wave coordination runs in the coordinator)

354 + 2. For non-MoE models in internal/hybrid LB mode: to collect and publish

355 + queue stats for load balancing across DP ranks

356 +

357 + Returns:

358 + True if DPCoordinator process is needed, False otherwise.

359 + """

360 +

361 + # For non-MoE models, only need coordinator in internal/hybrid LB mode

362 + # (for stats collection).

363 + return self.parallel_config.data_parallel_size > 1 and (

364 + self.model_config is None

365 + or self.model_config.is_moe

366 + or not self.parallel_config.data_parallel_external_lb

367 + )

368 +

369 def enable_trace_function_call_for_thread(self) -> None:

370 """

371 Set up function tracing for the current thread,

@@ -522,6 +545,8 @@ class VllmConfig:

545 self.model_config.verify_with_parallel_config(self.parallel_config)

546 self.model_config.verify_dual_chunk_attention_config(self.load_config)

547

548 + self.parallel_config.is_moe_model = self.model_config.is_moe

549 +

550 self.cache_config.verify_with_parallel_config(self.parallel_config)

551

552 if self.lora_config is not None:

@@ -827,9 +852,14 @@ class VllmConfig:

852 )

853

854 # Do this after all the updates to compilation_config.mode

855 + effective_dp_size = (

856 + self.parallel_config.data_parallel_size

857 + if self.model_config is None or self.model_config.is_moe

858 + else 1

859 + )

860 self.compilation_config.set_splitting_ops_for_v1(

861 all2all_backend=self.parallel_config.all2all_backend,

832 - data_parallel_size=self.parallel_config.data_parallel_size,

862 + data_parallel_size=effective_dp_size,

863 )

864

865 if self.compilation_config.pass_config.enable_sp:

@@ -1297,13 +1327,8 @@ class VllmConfig:

1327 if self.compilation_config.debug_dump_path is None:

1328 return None

1329 tp_rank = self.parallel_config.rank

1300 - dp_rank = self.parallel_config.data_parallel_rank

1301 - data_parallel_size = self.parallel_config.data_parallel_size

1302 - append_path = (

1303 - f"rank_{tp_rank}"

1304 - if data_parallel_size == 1

1305 - else f"rank_{tp_rank}_dp_{dp_rank}"

1306 - )

1330 + dp_rank = self.parallel_config.data_parallel_index

1331 + append_path = f"rank_{tp_rank}_dp_{dp_rank}"

1332 path = self.compilation_config.debug_dump_path / append_path

1333 return path

1334

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

+1 -1

@@ -915,6 +915,6 @@ def get_mooncake_side_channel_port(vllm_config: VllmConfig) -> int:

915 # This logic is now centralized

916 return (

917 envs.VLLM_MOONCAKE_BOOTSTRAP_PORT

918 - + vllm_config.parallel_config.data_parallel_rank

918 + + vllm_config.parallel_config.data_parallel_index

919 * vllm_config.parallel_config.tensor_parallel_size

920 )

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

+1 -1

@@ -471,7 +471,7 @@ class NixlConnectorScheduler:

471 self.side_channel_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST

472 self.side_channel_port = (

473 envs.VLLM_NIXL_SIDE_CHANNEL_PORT

474 - + vllm_config.parallel_config.data_parallel_rank

474 + + vllm_config.parallel_config.data_parallel_index

475 )

476 assert vllm_config.kv_transfer_config is not None

477 if current_platform.device_type == "cpu":

vllm/distributed/parallel_state.py

+22 -15

@@ -1115,7 +1115,11 @@ _EP: GroupCoordinator | None = None

1115

1116

1117 def get_ep_group() -> GroupCoordinator:

1118 - assert _EP is not None, "expert parallel group is not initialized"

1118 + assert _EP is not None, (

1119 + "expert parallel group is not initialized. "

1120 + "EP group is only created for MoE models with num_experts > 0. "

1121 + "This function should only be called for MoE models."

1122 + )

1123 return _EP

1124

1125

@@ -1400,20 +1404,23 @@ def initialize_model_parallel(

1404

1405 global _EP

1406 assert _EP is None, "expert parallel group is already initialized"

1403 - group_ranks = (

1404 - all_ranks.transpose(1, 2)

1405 - .reshape(

1406 - -1,

1407 - data_parallel_size

1408 - * prefill_context_model_parallel_size

1409 - * tensor_model_parallel_size,

1407 + # Don't create EP group for dense models.

1408 + if config is None or config.model_config is None or config.model_config.is_moe:

1409 + group_ranks = (

1410 + all_ranks.transpose(1, 2)

1411 + .reshape(

1412 + -1,

1413 + data_parallel_size

1414 + * prefill_context_model_parallel_size

1415 + * tensor_model_parallel_size,

1416 + )

1417 + .unbind(0)

1418 )

1411 - .unbind(0)

1412 - )

1413 - group_ranks = [x.tolist() for x in group_ranks]

1414 - _EP = init_model_parallel_group(

1415 - group_ranks, get_world_group().local_rank, backend, group_name="ep"

1416 - )

1419 + group_ranks = [x.tolist() for x in group_ranks]

1420 + _EP = init_model_parallel_group(

1421 + group_ranks, get_world_group().local_rank, backend, group_name="ep"

1422 + )

1423 + # If no EP group needed, _EP remains None

1424

1425 logger.info_once(

1426 "rank %s in world size %s is assigned as "

@@ -1425,7 +1432,7 @@ def initialize_model_parallel(

1432 _PP.rank_in_group,

1433 _PCP.rank_in_group,

1434 _TP.rank_in_group,

1428 - _EP.rank_in_group,

1435 + _EP.rank_in_group if _EP is not None else "N/A",

1436 )

1437

1438

vllm/engine/arg_utils.py

@@ -1575,6 +1575,7 @@ class EngineArgs:

1575 data_parallel_rpc_port=data_parallel_rpc_port,

1576 data_parallel_backend=self.data_parallel_backend,

1577 data_parallel_hybrid_lb=self.data_parallel_hybrid_lb,

1578 + is_moe_model=model_config.is_moe,

1579 enable_expert_parallel=self.enable_expert_parallel,

1580 all2all_backend=self.all2all_backend,

1581 enable_dbo=self.enable_dbo,

vllm/forward_context.py

@@ -102,6 +102,7 @@ class DPMetadata:

102 ) -> "DPMetadata":

103 assert num_tokens_across_dp_cpu is not None

104 assert parallel_config.data_parallel_size > 1

105 + assert parallel_config.is_moe_model is not False

106 dp_rank = parallel_config.data_parallel_rank

107 batchsize = num_tokens

108

vllm/v1/core/sched/scheduler.py

+1 -1

@@ -127,7 +127,7 @@ class Scheduler(SchedulerInterface):

127

128 self.kv_event_publisher = EventPublisherFactory.create(

129 self.kv_events_config,

130 - self.parallel_config.data_parallel_rank,

130 + self.parallel_config.data_parallel_index,

131 )

132 self.ec_connector = None

133 if self.vllm_config.ec_transfer_config is not None:

vllm/v1/engine/coordinator.py

+62 -43

@@ -55,7 +55,9 @@ class DPCoordinator:

55 request wave / running state changes.

56 """

58 - def __init__(self, parallel_config: ParallelConfig):

58 + def __init__(

59 + self, parallel_config: ParallelConfig, enable_wave_coordination: bool = True

60 + ):

61 dp_size = parallel_config.data_parallel_size

62 assert dp_size > 1, "Coordinator only used for data parallel"

@@ -83,6 +85,7 @@ class DPCoordinator:

85 "front_publish_address": front_publish_address,

86 "back_output_address": back_output_address,

87 "back_publish_address": back_publish_address,

88 + "enable_wave_coordination": enable_wave_coordination,

89 },

90 daemon=True,

91 )

@@ -110,13 +113,19 @@ class EngineState:

113

114

115 class DPCoordinatorProc:

113 - def __init__(self, engine_count: int, min_stats_update_interval_ms: int = 100):

116 + def __init__(

117 + self,

118 + engine_count: int,

119 + min_stats_update_interval_ms: int = 100,

120 + enable_wave_coordination: bool = True,

121 + ):

122 set_process_title("DPCoordinator")

123 self.ctx = zmq.Context()

124

125 self.engines = [EngineState() for _ in range(engine_count)]

126

127 self.stats_update_interval_ms = min_stats_update_interval_ms

128 + self.enable_wave_coordination = enable_wave_coordination

129

130 @staticmethod

131 def run_coordinator(

@@ -125,10 +134,12 @@ class DPCoordinatorProc:

134 back_output_address: str,

135 back_publish_address: str,

136 min_stats_update_interval_ms: int = 100,

137 + enable_wave_coordination: bool = True,

138 ):

139 coordinator = DPCoordinatorProc(

140 engine_count=engine_count,

141 min_stats_update_interval_ms=min_stats_update_interval_ms,

142 + enable_wave_coordination=enable_wave_coordination,

143 )

144 try:

145 coordinator.process_input_socket(

@@ -265,22 +276,25 @@ class DPCoordinatorProc:

276 )

277 continue # Skip normal engine notification processing

278

268 - # We received a message on the front-end XPUB socket,

269 - # from an API server sending a new request while the

270 - # engines are paused, so that we can wake the other

271 - # engines.

272 - engine_to_exclude, wave = decoded

273 - if not engines_running:

274 - if wave < current_wave:

275 - # If the wave number is stale, ensure the message

276 - # is handled by all the engines.

277 - engine_to_exclude = None

278 -

279 - engines_running = True

280 - wave_state_changed = True

281 - self._send_start_wave(

282 - publish_back, current_wave, engine_to_exclude

283 - )

279 + # Wave coordination: handle new-request messages from front-end.

280 + # Only process these when wave coordination is enabled

281 + if self.enable_wave_coordination:

282 + # We received a message on the front-end XPUB socket,

283 + # from an API server sending a new request while the

284 + # engines are paused, so that we can wake the other

285 + # engines.

286 + engine_to_exclude, wave = decoded

287 + if not engines_running:

288 + if wave < current_wave:

289 + # If the wave number is stale, ensure the message

290 + # is handled by all the engines.

291 + engine_to_exclude = None

292 +

293 + engines_running = True

294 + wave_state_changed = True

295 + self._send_start_wave(

296 + publish_back, current_wave, engine_to_exclude

297 + )

298

299 if output_back in events:

300 # We received a message from one of the engines.

@@ -325,34 +339,39 @@ class DPCoordinatorProc:

339 stats[1] = scheduler_stats.num_running_reqs

340 stats_changed = True

341

328 - if (wave := outputs.wave_complete) is not None:

329 - # 2. Notification from rank 0 engine that we've

330 - # moved into the global paused state

331 - # (engines_running==False).

332 - if current_wave <= wave:

333 - new_wave = wave + 1

342 + # Wave coordination: handle wave completion and start notifications

343 + # Only process these when wave coordination is enabled

344 + if self.enable_wave_coordination:

345 + if (wave := outputs.wave_complete) is not None:

346 + # 2. Notification from rank 0 engine that we've

347 + # moved into the global paused state

348 + # (engines_running==False).

349 + if current_wave <= wave:

350 + new_wave = wave + 1

351 + logger.debug(

352 + "Moving DP wave from %d to %d.",

353 + current_wave,

354 + new_wave,

355 + )

356 + current_wave = new_wave

357 + engines_running = False

358 + wave_state_changed = True

359 + elif (wave := outputs.start_wave) is not None and (

360 + wave > current_wave

361 + or (wave == current_wave and not engines_running)

362 + ):

363 + # 3. The engine received request for a non-current wave

364 + # so we must ensure that other engines progress to the

365 + # next wave (race condition handling).

366 logger.debug(

335 - "Moving DP wave from %d to %d.", current_wave, new_wave

367 + "Starting wave %d after notification of "

368 + "stale wave request from engine.",

369 + wave,

370 )

337 - current_wave = new_wave

338 - engines_running = False

371 + current_wave = wave

372 + engines_running = True

373 wave_state_changed = True

340 - elif (wave := outputs.start_wave) is not None and (

341 - wave > current_wave

342 - or (wave == current_wave and not engines_running)

343 - ):

344 - # 3. The engine received request for a non-current wave

345 - # so we must ensure that other engines progress to the

346 - # next wave (race condition handling).

347 - logger.debug(

348 - "Starting wave %d after notification of "

349 - "stale wave request from engine.",

350 - wave,

351 - )

352 - current_wave = wave

353 - engines_running = True

354 - wave_state_changed = True

355 - self._send_start_wave(publish_back, wave, eng_index)

374 + self._send_start_wave(publish_back, wave, eng_index)

375

376 if wave_state_changed:

377 message = (None, current_wave, engines_running)

vllm/v1/engine/core.py

+94 -23

@@ -84,6 +84,7 @@ class EngineCore:

84 executor_class: type[Executor],

85 log_stats: bool,

86 executor_fail_callback: Callable | None = None,

87 + include_finished_set: bool = False,

88 ):

89 # plugins need to be loaded at the engine/scheduler level too

90 from vllm.plugins import load_general_plugins

@@ -91,7 +92,7 @@ class EngineCore:

92 load_general_plugins()

94 self.vllm_config = vllm_config

94 - if vllm_config.parallel_config.data_parallel_rank == 0:

95 + if not vllm_config.parallel_config.data_parallel_rank_local:

96 logger.info(

97 "Initializing a V1 LLM engine (v%s) with config: %s",

98 VLLM_VERSION,

@@ -138,7 +139,7 @@ class EngineCore:

139 vllm_config=vllm_config,

140 kv_cache_config=kv_cache_config,

141 structured_output_manager=self.structured_output_manager,

141 - include_finished_set=vllm_config.parallel_config.data_parallel_size > 1,

142 + include_finished_set=include_finished_set,

143 log_stats=self.log_stats,

144 block_size=scheduler_block_size,

145 )

@@ -605,6 +606,7 @@ class EngineCoreProc(EngineCore):

606 executor_class: type[Executor],

607 log_stats: bool,

608 client_handshake_address: str | None = None,

609 + *,

610 engine_index: int = 0,

611 ):

612 self.input_queue = queue.Queue[tuple[EngineCoreRequestType, Any]]()

@@ -636,17 +638,22 @@ class EngineCoreProc(EngineCore):

638 self.has_coordinator,

639 self.frontend_stats_publish_address,

640 )

639 - # Only publish request queue stats to coordinator for "internal"

640 - # and "hybrid" LB modes .

641 - self.publish_dp_lb_stats = (

641 + internal_dp_balancing = (

642 self.has_coordinator

643 and not vllm_config.parallel_config.data_parallel_external_lb

644 )

645 + # Only publish request queue stats to coordinator for "internal"

646 + # and "hybrid" LB modes.

647 + self.publish_dp_lb_stats = internal_dp_balancing

648

649 self._init_data_parallel(vllm_config)

650

651 super().__init__(

649 - vllm_config, executor_class, log_stats, executor_fail_callback

652 + vllm_config,

653 + executor_class,

654 + log_stats,

655 + executor_fail_callback,

656 + internal_dp_balancing,

657 )

658

659 # Background Threads and Queues for IO. These enable us to

@@ -854,18 +861,29 @@ class EngineCoreProc(EngineCore):

861

862 engine_core: EngineCoreProc | None = None

863 try:

857 - parallel_config: ParallelConfig = kwargs["vllm_config"].parallel_config

858 - if parallel_config.data_parallel_size > 1 or dp_rank > 0:

864 + vllm_config: VllmConfig = kwargs["vllm_config"]

865 + parallel_config: ParallelConfig = vllm_config.parallel_config

866 + data_parallel = parallel_config.data_parallel_size > 1 or dp_rank > 0

867 + if data_parallel:

868 + parallel_config.data_parallel_rank_local = local_dp_rank

869 set_process_title("EngineCore", f"DP{dp_rank}")

860 - decorate_logs()

870 + else:

871 + set_process_title("EngineCore")

872 + decorate_logs()

873 +

874 + parallel_config.data_parallel_index = dp_rank

875 + if data_parallel and vllm_config.model_config.is_moe:

876 # Set data parallel rank for this engine process.

877 parallel_config.data_parallel_rank = dp_rank

863 - parallel_config.data_parallel_rank_local = local_dp_rank

878 engine_core = DPEngineCoreProc(*args, **kwargs)

879 else:

866 - set_process_title("EngineCore")

867 - decorate_logs()

868 - engine_core = EngineCoreProc(*args, **kwargs)

880 + # Non-MoE DP ranks are completely independent, so treat like DP=1.

881 + # Note that parallel_config.data_parallel_index will still reflect

882 + # the original DP rank.

883 + parallel_config.data_parallel_size = 1

884 + parallel_config.data_parallel_size_local = 1

885 + parallel_config.data_parallel_rank = 0

886 + engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs)

887

888 engine_core.run_busy_loop()

889

@@ -1195,6 +1213,10 @@ class DPEngineCoreProc(EngineCoreProc):

1213 log_stats: bool,

1214 client_handshake_address: str | None = None,

1215 ):

1216 + assert vllm_config.model_config.is_moe, (

1217 + "DPEngineCoreProc should only be used for MoE models"

1218 + )

1219 +

1220 # Counts forward-passes of the model so that we can synchronize

1221 # finished with DP peers every N steps.

1222 self.step_counter = 0

@@ -1210,7 +1232,7 @@ class DPEngineCoreProc(EngineCoreProc):

1232 executor_class,

1233 log_stats,

1234 client_handshake_address,

1213 - dp_rank,

1235 + engine_index=dp_rank,

1236 )

1237

1238 def _init_data_parallel(self, vllm_config: VllmConfig):

@@ -1391,7 +1413,7 @@ class DPEngineCoreProc(EngineCoreProc):

1413 )

1414

1415

1394 -class DPEngineCoreActor(DPEngineCoreProc):

1416 +class EngineCoreActorMixin:

1417 """

1418 Ray actor for running EngineCore in a data parallel context

1419 """

@@ -1399,15 +1421,12 @@ class DPEngineCoreActor(DPEngineCoreProc):

1421 def __init__(

1422 self,

1423 vllm_config: VllmConfig,

1402 - local_client: bool,

1424 addresses: EngineZmqAddresses,

1404 - executor_class: type[Executor],

1405 - log_stats: bool,

1425 dp_rank: int = 0,

1426 local_dp_rank: int = 0,

1427 ):

1428 self.addresses = addresses

1410 - vllm_config.parallel_config.data_parallel_rank = dp_rank

1429 + vllm_config.parallel_config.data_parallel_index = dp_rank

1430 vllm_config.parallel_config.data_parallel_rank_local = local_dp_rank

1431

1432 # Set CUDA_VISIBLE_DEVICES as early as possible in actor life cycle

@@ -1429,8 +1448,6 @@ class DPEngineCoreActor(DPEngineCoreProc):

1448 # of ray.

1449 self._set_visible_devices(vllm_config, local_dp_rank)

1450

1432 - super().__init__(vllm_config, local_client, "", executor_class, log_stats)

1433 -

1451 def _set_visible_devices(self, vllm_config: VllmConfig, local_dp_rank: int):

1452 from vllm.platforms import current_platform

1453

@@ -1491,7 +1508,7 @@ class DPEngineCoreActor(DPEngineCoreProc):

1508 Run the engine core busy loop.

1509 """

1510 try:

1494 - self.run_busy_loop()

1511 + self.run_busy_loop() # type: ignore[attr-defined]

1512 except SystemExit:

1513 logger.debug("EngineCore exiting.")

1514 raise

@@ -1499,4 +1516,58 @@ class DPEngineCoreActor(DPEngineCoreProc):

1516 logger.exception("EngineCore encountered a fatal error.")

1517 raise

1518 finally:

1502 - self.shutdown()

1519 + self.shutdown() # type: ignore[attr-defined]

1520 +

1521 +

1522 +class DPMoEEngineCoreActor(EngineCoreActorMixin, DPEngineCoreProc):

1523 + """Used for MoE model data parallel cases."""

1524 +

1525 + def __init__(

1526 + self,

1527 + vllm_config: VllmConfig,

1528 + local_client: bool,

1529 + addresses: EngineZmqAddresses,

1530 + executor_class: type[Executor],

1531 + log_stats: bool,

1532 + dp_rank: int = 0,

1533 + local_dp_rank: int = 0,

1534 + ):

1535 + vllm_config.parallel_config.data_parallel_rank = dp_rank

1536 +

1537 + EngineCoreActorMixin.__init__(

1538 + self, vllm_config, addresses, dp_rank, local_dp_rank

1539 + )

1540 + DPEngineCoreProc.__init__(

1541 + self, vllm_config, local_client, "", executor_class, log_stats

1542 + )

1543 +

1544 +

1545 +class EngineCoreActor(EngineCoreActorMixin, EngineCoreProc):

1546 + """Used for non-MoE and/or non-DP cases."""

1547 +

1548 + def __init__(

1549 + self,

1550 + vllm_config: VllmConfig,

1551 + local_client: bool,

1552 + addresses: EngineZmqAddresses,

1553 + executor_class: type[Executor],

1554 + log_stats: bool,

1555 + dp_rank: int = 0,

1556 + local_dp_rank: int = 0,

1557 + ):

1558 + vllm_config.parallel_config.data_parallel_size = 1

1559 + vllm_config.parallel_config.data_parallel_size_local = 1

1560 + vllm_config.parallel_config.data_parallel_rank = 0

1561 +

1562 + EngineCoreActorMixin.__init__(

1563 + self, vllm_config, addresses, dp_rank, local_dp_rank

1564 + )

1565 + EngineCoreProc.__init__(

1566 + self,

1567 + vllm_config,

1568 + local_client,

1569 + "",

1570 + executor_class,

1571 + log_stats,

1572 + engine_index=dp_rank,

1573 + )

vllm/v1/engine/core_client.py

+1 -1

@@ -502,7 +502,7 @@ class MPClient(EngineCoreClient):

502

503 parallel_config = vllm_config.parallel_config

504 dp_size = parallel_config.data_parallel_size

505 - dp_rank = parallel_config.data_parallel_rank

505 + dp_rank = parallel_config.data_parallel_index

506 dp_local_size = parallel_config.data_parallel_size_local

507 offline_mode = parallel_config.data_parallel_rank_local is not None

508 # Client manages local+remote EngineCores in pure internal LB case.

vllm/v1/engine/llm_engine.py

+2 -1

@@ -65,8 +65,9 @@ class LLMEngine:

66 self.log_stats = log_stats

68 - executor_backend = self.vllm_config.parallel_config.distributed_executor_backend

68 parallel_config = vllm_config.parallel_config

69 + executor_backend = parallel_config.distributed_executor_backend

70 +

71 self.external_launcher_dp = (

72 parallel_config.data_parallel_size > 1

73 and executor_backend == "external_launcher"

vllm/v1/engine/utils.py

+37 -19

@@ -75,7 +75,6 @@ class EngineHandshakeMetadata:

76 addresses: EngineZmqAddresses

77 parallel_config: dict[str, int | str | list[int]]

78 - parallel_config_hash: str | None = None

80 class CoreEngineProcManager:

@@ -249,12 +248,19 @@ class CoreEngineActorManager:

248 from ray.runtime_env import RuntimeEnv

249 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

250

252 - from vllm.v1.engine.core import DPEngineCoreActor

251 + from vllm.v1.engine.core import DPMoEEngineCoreActor, EngineCoreActor

252 +

253 + dp_size = vllm_config.parallel_config.data_parallel_size

254 + actor_class = (

255 + DPMoEEngineCoreActor

256 + if dp_size > 1 and vllm_config.model_config.is_moe

257 + else EngineCoreActor

258 + )

259

260 self.local_engine_actors: list[ray.ActorHandle] = []

261 self.remote_engine_actors: list[ray.ActorHandle] = []

262

257 - env_vars_list = get_env_vars_to_copy(destination="DPEngineCoreActor")

263 + env_vars_list = get_env_vars_to_copy(destination=actor_class.__name__)

264 self.env_vars_dict = {

265 name: os.environ[name] for name in env_vars_list if name in os.environ

266 }

@@ -263,7 +269,6 @@ class CoreEngineActorManager:

269 self.addresses = addresses

270 self.executor_class = executor_class

271 self.log_stats = log_stats

266 - dp_size = vllm_config.parallel_config.data_parallel_size

272 local_engine_count = vllm_config.parallel_config.data_parallel_size_local

273 world_size = vllm_config.parallel_config.world_size

274

@@ -314,7 +319,7 @@ class CoreEngineActorManager:

319 runtime_env = RuntimeEnv(env_vars=actor_env_vars)

320

321 actor = (

317 - ray.remote(DPEngineCoreActor)

322 + ray.remote(actor_class)

323 .options(

324 scheduling_strategy=PlacementGroupSchedulingStrategy(

325 placement_group=pg,

@@ -624,7 +629,13 @@ class CoreEngineActorManager:

629 from ray.runtime_env import RuntimeEnv

630 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

631

627 - from vllm.v1.engine.core import DPEngineCoreActor

632 + from vllm.v1.engine.core import DPMoEEngineCoreActor, EngineCoreActor

633 +

634 + actor_class = (

635 + DPMoEEngineCoreActor

636 + if cur_vllm_config.model_config.is_moe

637 + else EngineCoreActor

638 + )

639

640 cur_data_parallel_size = len(self.local_engine_actors) + len(

641 self.remote_engine_actors

@@ -667,7 +678,7 @@ class CoreEngineActorManager:

678 )

679

680 actor = (

670 - ray.remote(DPEngineCoreActor)

681 + ray.remote(actor_class)

682 .options(

683 scheduling_strategy=PlacementGroupSchedulingStrategy(

684 placement_group=pg,

@@ -804,12 +815,19 @@ def launch_core_engines(

815 ],

816 )

817

807 - # Run the DP Coordinator process with rank 0 when in

808 - # online DP mode.

809 - run_coordinator = dp_size > 1 and not offline_mode and dp_rank == 0

818 + # Run the DP Coordinator process with rank 0 when in online DP mode.

819 + # The coordinator is needed for:

820 + # 1. Internal/hybrid LB: collecting and publishing queue stats for load balancing

821 + # 2. MoE models: wave coordination in addition to stats

822 + run_coordinator = (

823 + vllm_config.needs_dp_coordinator and not offline_mode and dp_rank == 0

824 + )

825

826 if run_coordinator:

812 - coordinator = DPCoordinator(parallel_config)

827 + coordinator = DPCoordinator(

828 + parallel_config,

829 + enable_wave_coordination=vllm_config.model_config.is_moe,

830 + )

831

832 addresses.coordinator_input, addresses.coordinator_output = (

833 coordinator.get_engine_socket_addresses()

@@ -905,6 +923,7 @@ def launch_core_engines(

923 addresses,

924 engines_to_handshake,

925 parallel_config,

926 + dp_size > 1 and vllm_config.model_config.is_moe,

927 vllm_config.cache_config,

928 local_engine_manager,

929 coordinator.proc if coordinator else None,

@@ -916,6 +935,7 @@ def wait_for_engine_startup(

935 addresses: EngineZmqAddresses,

936 core_engines: list[CoreEngine],

937 parallel_config: ParallelConfig,

938 + coordinated_dp: bool,

939 cache_config: CacheConfig,

940 proc_manager: CoreEngineProcManager | None,

941 coord_process: Process | None,

@@ -997,8 +1017,7 @@ def wait_for_engine_startup(

1017 )

1018

1019 if status == "HELLO" and engine.state == CoreEngineState.NEW:

1000 - # Send init message with DP config info and config hash.

1001 - # The config hash ensures all DP workers have compatible configs.

1020 + # Send init message with DP config info.

1021 init_message = msgspec.msgpack.encode(

1022 EngineHandshakeMetadata(

1023 addresses=addresses,

@@ -1010,10 +1029,9 @@ def wait_for_engine_startup(

1029 "_data_parallel_master_port_list",

1030 "data_parallel_size",

1031 )

1013 - },

1014 - parallel_config_hash=parallel_config.compute_hash()

1015 - if parallel_config.data_parallel_size > 1

1016 - else None,

1032 + }

1033 + if coordinated_dp

1034 + else {},

1035 )

1036 )

1037 handshake_socket.send_multipart((eng_identity, init_message), copy=False)

@@ -1034,8 +1052,8 @@ def wait_for_engine_startup(

1052 if addresses.frontend_stats_publish_address is None:

1053 addresses.frontend_stats_publish_address = msg.get("dp_stats_address")

1054

1037 - # Validate config hash consistency across DP workers

1038 - if parallel_config.data_parallel_size > 1:

1055 + # Validate config hash consistency across DP workers for MoE models.

1056 + if coordinated_dp:

1057 worker_config_hash = msg.get("parallel_config_hash")

1058 expected_hash = parallel_config.compute_hash()

1059 if worker_config_hash != expected_hash:

vllm/v1/worker/gpu/model_runner.py

+9 -9

@@ -98,9 +98,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

98 self.max_num_reqs = self.scheduler_config.max_num_seqs

99 self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()

100

101 - self.dp_size = self.parallel_config.data_parallel_size

102 - self.dp_rank = self.parallel_config.data_parallel_rank

103 -

101 self.use_async_scheduling = self.scheduler_config.async_scheduling

102 self.output_copy_stream = torch.cuda.Stream(self.device)

103 self.output_copy_event = torch.cuda.Event()

@@ -268,7 +265,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

265 if not skip_attn:

266 self.prepare_dummy_attn_metadata(input_batch)

267

271 - num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, num_tokens)

268 + dp_size = self.parallel_config.data_parallel_size

269 + num_tokens_across_dp = make_num_tokens_across_dp(dp_size, num_tokens)

270 num_sampled_tokens = np.ones(input_batch.num_reqs, dtype=np.int32)

271 with (

272 self.maybe_dummy_run_with_lora(

@@ -312,7 +310,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

310 self._dummy_sampler_run(sample_hidden_states)

311 if self.do_spec_decode:

312 num_tokens_across_dp = make_num_tokens_across_dp(

315 - self.dp_size, self.max_num_tokens

313 + self.parallel_config.data_parallel_size, self.max_num_tokens

314 )

315 self.speculator.run_model(

316 self.max_num_tokens,

@@ -807,7 +805,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

805 scheduler_output: SchedulerOutput,

806 ) -> tuple[CUDAGraphMode, int, torch.Tensor | None]:

807 total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

810 - if self.dp_size == 1:

808 + dp_size = self.parallel_config.data_parallel_size

809 + if dp_size == 1:

810 # No DP. Only consider CUDA graphs.

811 if total_num_scheduled_tokens == 0:

812 # Special case: no tokens to run.

@@ -835,11 +834,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

834 cudagraph_size_before_dp = -1

835

836 assert cudagraph_size_before_dp is not None

837 + dp_rank = self.parallel_config.data_parallel_rank

838 num_tokens_across_dp, cudagraph_size_across_dp = get_batch_metadata_across_dp(

839 total_num_scheduled_tokens,

840 cudagraph_size_before_dp,

841 - self.dp_size,

842 - self.dp_rank,

841 + dp_size,

842 + dp_rank,

843 )

844 if all(cudagraph_size_across_dp >= 0):

845 # If all ranks can use CUDA graph, pad to the maximum number of tokens

@@ -850,7 +850,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

850 # If any of the ranks cannot use CUDA graph, use eager mode for all ranks.

851 # No padding is needed except for ranks that have no tokens to run.

852 num_tokens_across_dp = torch.clamp(num_tokens_across_dp, min=1)

853 - num_tokens_after_padding = num_tokens_across_dp[self.dp_rank]

853 + num_tokens_after_padding = num_tokens_across_dp[dp_rank]

854 cudagraph_mode = CUDAGraphMode.NONE

855 return cudagraph_mode, num_tokens_after_padding, num_tokens_across_dp

856

vllm/v1/worker/gpu_worker.py

+7 -9

@@ -179,22 +179,20 @@ class Worker(WorkerBase):

179 self.cache_config.num_cpu_blocks = num_cpu_blocks

180

181 def init_device(self):

182 - device = self.device_config.device

183 - if isinstance(device, torch.device) and device.type == "cuda":

182 + if self.device_config.device_type == "cuda":

183 # This env var set by Ray causes exceptions with graph building.

184 os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)

185 + parallel_config = self.parallel_config

186 if (

187 - self.parallel_config.data_parallel_size > 1

188 - and self.parallel_config.data_parallel_size_local > 0

189 - and self.parallel_config.distributed_executor_backend

190 - not in ["ray", "external_launcher"]

191 - and self.vllm_config.parallel_config.data_parallel_backend != "ray"

192 - and self.vllm_config.parallel_config.nnodes_within_dp == 1

187 + parallel_config.distributed_executor_backend

188 + not in ("ray", "external_launcher")

189 + and parallel_config.data_parallel_backend != "ray"

190 + and parallel_config.nnodes_within_dp == 1

191 ):

192 # Use local DP rank if available, otherwise use global DP rank.

193 dp_local_rank = self.parallel_config.data_parallel_rank_local

194 if dp_local_rank is None:

197 - dp_local_rank = self.parallel_config.data_parallel_rank

195 + dp_local_rank = self.parallel_config.data_parallel_index

196

197 tp_pp_world_size = (

198 self.parallel_config.pipeline_parallel_size