@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a2257f95b7

[Elastic EP] Reduce eager-mode reconfiguration downtime (#51885)

Signed-off-by: Itay Alroy <ialroy@nvidia.com>

Itay Alroy committed Aug 19, 2026 at 05:13 UTC a2257f95b79287933defef5be40b74ce93053da3

12 files changed +194 -303

tests/distributed/test_eplb_execute.py

-122

@@ -782,125 +782,3 @@ def test_rearrange_expert_weights_profile_mode(world_size):

782 _test_rearrange_expert_weights_profile_mode,

783 world_size,

784 )

785 -

786 -

787 -def _test_nixl_deferred_init_worker(

788 - env,

789 - world_size: int,

790 - num_layers: int,

791 - num_local_experts: int,

792 - num_logical_experts: int,

793 -) -> None:

794 - """Exercise NixlEplbCommunicator with defer_remote_setup=True (elastic EP path)."""

795 - from vllm.distributed.eplb.eplb_communicator import NixlEplbCommunicator

796 -

797 - set_env_vars_and_device(env)

798 -

799 - vllm_config = VllmConfig()

800 - vllm_config.parallel_config.tensor_parallel_size = world_size

801 -

802 - with set_current_vllm_config(vllm_config):

803 - ensure_model_parallel_initialized(

804 - tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1

805 - )

806 -

807 - ep_group_coordinator = get_tp_group()

808 - ep_group = ep_group_coordinator.cpu_group

809 - ep_rank = torch.distributed.get_rank()

810 - device = torch.device(f"cuda:{ep_rank}")

811 -

812 - total_physical_experts = world_size * num_local_experts

813 - hidden_sizes = [32, 64]

814 -

815 - redundancy_config = create_redundancy_config(

816 - num_logical_experts, total_physical_experts

817 - )

818 - old_indices = create_expert_indices_with_redundancy(

819 - num_layers,

820 - num_logical_experts,

821 - total_physical_experts,

822 - redundancy_config,

823 - )

824 -

825 - new_redundancy_config = create_redundancy_config(

826 - num_logical_experts, total_physical_experts

827 - )

828 - new_indices = create_expert_indices_with_redundancy(

829 - num_layers,

830 - num_logical_experts,

831 - total_physical_experts,

832 - new_redundancy_config,

833 - )

834 -

835 - expert_weights = create_expert_weights(

836 - num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices

837 - )

838 -

839 - expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]

840 -

841 - communicator = NixlEplbCommunicator(

842 - cpu_group=ep_group_coordinator.cpu_group,

843 - all_expert_weights=expert_weights,

844 - expert_buffer=expert_buffer,

845 - defer_remote_setup=True,

846 - )

847 - assert not communicator._remote_state_initialized

848 -

849 - rearrange_expert_weights_inplace(

850 - old_indices,

851 - new_indices,

852 - expert_weights,

853 - expert_buffer,

854 - ep_group,

855 - communicator,

856 - )

857 -

858 - assert communicator._remote_state_initialized

859 -

860 - local_ok = verify_expert_weights_after_shuffle(

861 - expert_weights,

862 - new_indices,

863 - hidden_sizes,

864 - ep_rank,

865 - num_local_experts,

866 - )

867 -

868 - local_ok = (

869 - verify_redundant_experts_have_same_weights(

870 - expert_weights,

871 - new_indices,

872 - hidden_sizes,

873 - ep_rank,

874 - world_size,

875 - num_local_experts,

876 - )

877 - and local_ok

878 - )

879 - assert_verification_synced(

880 - local_ok,

881 - "Deferred NIXL init verification failed on at least one rank.",

882 - )

883 -

884 -

885 -@pytest.mark.skipif(not has_nixl(), reason="NIXL is not available")

886 -@pytest.mark.parametrize(

887 - "world_size,num_layers,num_local_experts,num_logical_experts",

888 - [(2, 2, 3, 4)],

889 -)

890 -def test_nixl_deferred_init(

891 - world_size,

892 - num_layers,

893 - num_local_experts,

894 - num_logical_experts,

895 -):

896 - """Test NixlEplbCommunicator with defer_remote_setup=True (elastic EP path)."""

897 -

898 - if torch.accelerator.device_count() < world_size:

899 - pytest.skip(f"Need at least {world_size} GPUs to run the test")

900 - distributed_run(

901 - _test_nixl_deferred_init_worker,

902 - world_size,

903 - num_layers,

904 - num_local_experts,

905 - num_logical_experts,

906 - )

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

+1 -2

@@ -159,13 +159,12 @@ def test_v2_setup_eplb_from_mapping_rebuilds_state(monkeypatch):

159

160 runner = _make_runner(model=SimpleNamespace(is_moe=True))

161 mapping = torch.tensor([[0, 1, 2, 3]], dtype=torch.int64)

162 - mrv2.GPUModelRunner.setup_eplb_from_mapping(runner, mapping, 2)

162 + mrv2.GPUModelRunner.setup_eplb_from_mapping(runner, mapping)

163

164 assert runner.eplb_state is not None

165 assert runner.eplb_state.built_from_mapping is True

166 assert FakeEplbState.from_mapping_kwargs is not None

167 assert FakeEplbState.from_mapping_kwargs["expanded_physical_to_logical"] is mapping

168 - assert FakeEplbState.from_mapping_kwargs["num_valid_physical_experts"] == 2

168

169

170 def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):

vllm/config/parallel.py

+1 -1

@@ -973,7 +973,7 @@ class ParallelConfig:

973

974 if self.enable_eplb and self.eplb_config.communicator is None:

975 # Prefer NIXL when available: zero-copy RDMA reads, compatible

976 - # with both async EPLB and elastic EP (deferred remote setup).

976 + # with both async EPLB and elastic EP.

977 # Fallbacks: pynccl for elastic EP (stateless groups need it),

978 # torch_gloo for static EP. torch_nccl is avoided because NCCL

979 # is incompatible with async EPLB (multi-stream conflicts) and

vllm/distributed/elastic_ep/elastic_execute.py

+108 -54

@@ -29,12 +29,15 @@ from vllm.distributed.elastic_ep.standby_state import (

29 create_standby_groups,

30 get_standby_dp_group,

31 get_standby_ep_group,

32 + get_standby_eplb_group,

33 pop_standby_groups,

34 )

34 -from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator

35 +from vllm.distributed.eplb.eplb_communicator import (

36 + EplbCommunicator,

37 +)

38 from vllm.distributed.parallel_state import (

39 + GroupCoordinator,

40 _replace_active_groups,

37 - get_eplb_group,

41 )

42 from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator

43 from vllm.logger import init_logger

@@ -149,10 +152,12 @@ class ElasticEPScalingExecutor:

152 self.worker_ref = weakref.ref(worker)

153 self.reconfig_request = None

154 self._staged_moe_quant_methods: dict[nn.Module, FusedMoEMethodBase] = {}

155 + self._prepared_eplb_communicator: EplbCommunicator | None = None

156 self._async_executor = ThreadPoolExecutor(

157 max_workers=1, thread_name_prefix="ElasticEPAsync"

158 )

159 self._async_future: Future[None] | None = None

160 + self._group_cleanup_future: Future[None] | None = None

161

162 @property

163 def worker(self):

@@ -207,12 +212,40 @@ class ElasticEPScalingExecutor:

212 self._async_future = None

213 future.result()

214

215 + def _destroy_retired_groups(

216 + self, groups: tuple[GroupCoordinator | None, ...]

217 + ) -> None:

218 + from vllm.platforms import current_platform

219 +

220 + current_platform.set_device(self.worker.device)

221 + for group in groups:

222 + if group is not None:

223 + group.destroy()

224 +

225 + def _start_group_cleanup(self, groups: tuple[GroupCoordinator | None, ...]) -> None:

226 + assert self._group_cleanup_future is None

227 + self._group_cleanup_future = self._async_executor.submit(

228 + self._destroy_retired_groups, groups

229 + )

230 +

231 + def _wait_for_group_cleanup(self) -> None:

232 + if (future := self._group_cleanup_future) is not None:

233 + self._group_cleanup_future = None

234 + future.result()

235 +

236 + def shutdown(self) -> None:

237 + try:

238 + self._wait_for_group_cleanup()

239 + finally:

240 + self._async_executor.shutdown()

241 +

242 def load_model(self) -> None:

243 self.worker.load_model(load_dummy_weights=True)

244

213 - def create_standby_groups(

245 + def prepare_reconfiguration(

246 self, reconfig_request: ReconfigureDistributedRequest, use_all2all: bool

247 ) -> None:

248 + self._wait_for_group_cleanup()

249 self.reconfig_request = reconfig_request

250 new_dp_size = reconfig_request.new_data_parallel_size

251 old_dp_size = get_dp_group().world_size

@@ -227,8 +260,20 @@ class ElasticEPScalingExecutor:

260 use_all2all=use_all2all,

261 enable_eplb=parallel_config.enable_eplb,

262 )

230 - if new_dp_size < old_dp_size:

231 - self.stage_standby_moe_quant_methods()

263 + self.stage_standby_moe_quant_methods()

264 + self._prepare_eplb_communicator(get_standby_eplb_group())

265 + if new_dp_size > old_dp_size:

266 + self.transfer_weights(old_dp_size, new_dp_size)

267 + self._warm_target_groups(get_standby_dp_group(), get_standby_ep_group())

268 +

269 + def _prepare_eplb_communicator(self, eplb_group) -> None:

270 + assert eplb_group is not None

271 + model_runner = self.worker.model_runner

272 + eplb_state = model_runner.eplb_state

273 + assert eplb_state is not None

274 + self._prepared_eplb_communicator = eplb_state.create_communicator(

275 + model_runner.model_config, eplb_group

276 + )

277

278 def transfer_weights(self, old_dp_size: int, new_dp_size: int) -> None:

279 standby_dp_group = get_standby_dp_group()

@@ -276,6 +321,15 @@ class ElasticEPScalingExecutor:

321 )

322 torch.accelerator.synchronize()

323

324 + def _warm_target_groups(self, dp_group, ep_group) -> None:

325 + assert dp_group is not None and ep_group is not None

326 + stream = torch.Stream(device=dp_group.device)

327 + with stream:

328 + tensor = torch.zeros(1, dtype=torch.int32, device=dp_group.device)

329 + for group in (dp_group, ep_group):

330 + torch.distributed.all_reduce(tensor, group=group.device_group)

331 + stream.synchronize()

332 +

333 def broadcast_expert_mapping(self) -> None:

334 standby_dp_group = get_standby_dp_group()

335 assert standby_dp_group is not None

@@ -362,15 +416,22 @@ class ElasticEPScalingExecutor:

416 torch.accelerator.empty_cache()

417

418 def switch_and_remove(self) -> None:

419 + # Removing ranks skipped preparation, so wait for prior cleanup here.

420 + self._wait_for_group_cleanup()

421 self._release_cuda_graphs()

366 - _replace_active_groups(world=None, dp=None, ep=None, eplb=None, node_count=None)

422 + retired_groups = _replace_active_groups(

423 + world=None, dp=None, ep=None, eplb=None, node_count=None

424 + )

425 + self._start_group_cleanup(retired_groups)

426 + # Finish collective cleanup before this worker is shut down.

427 + self._wait_for_group_cleanup()

428

368 - def switch_and_prepare(self) -> None:

429 + def switch_and_prepare(self) -> tuple[GroupCoordinator | None, ...]:

430 old_dp_size = get_dp_group().world_size

431 old_ep_size = get_ep_group().world_size

432

433 self._release_cuda_graphs()

373 - _replace_active_groups(**pop_standby_groups())

434 + retired_groups = _replace_active_groups(**pop_standby_groups())

435

436 parallel_config = self.worker.vllm_config.parallel_config

437 reconfig_request = self.reconfig_request

@@ -453,7 +514,6 @@ class ElasticEPScalingExecutor:

514 )

515 eplb_model_state.expert_load_pass = expanded_expert_load_pass

516 eplb_model_state.expert_load_window = expanded_expert_load_window

456 - eplb_state.num_valid_physical_experts = old_num_physical_experts

517 else:

518 assert pad_size < 0

519 eplb_model_state.expert_load_pass = eplb_model_state.expert_load_pass[

@@ -462,7 +522,6 @@ class ElasticEPScalingExecutor:

522 eplb_model_state.expert_load_window = eplb_model_state.expert_load_window[

523 :, :, :num_physical_experts

524 ]

465 - eplb_state.num_valid_physical_experts = num_physical_experts

525

526 model = self.worker.model_runner.get_model()

527 model.expert_weights = []

@@ -485,18 +544,9 @@ class ElasticEPScalingExecutor:

544 if getattr(module._quant_method, "wraps_legacy_quant_method", False):

545 module._replace_quant_method(module._quant_method.old_quant_method)

546

488 - eplb_model_state.expert_buffer = [

489 - torch.empty_like(w) for w in model.expert_weights[0]

490 - ]

491 - assert parallel_config.eplb_config.communicator is not None, (

492 - "EPLB communicator backend must be set by ParallelConfig"

493 - )

494 - eplb_model_state.communicator = create_eplb_communicator(

495 - group_coordinator=get_eplb_group(),

496 - backend=parallel_config.eplb_config.communicator,

497 - expert_weights=model.expert_weights,

498 - expert_buffer=eplb_model_state.expert_buffer,

499 - )

547 + assert self._prepared_eplb_communicator is not None

548 + eplb_state.update_communicator(model_config, self._prepared_eplb_communicator)

549 + self._prepared_eplb_communicator = None

550

551 if (

552 self.worker.vllm_config.compilation_config.mode

@@ -513,8 +563,12 @@ class ElasticEPScalingExecutor:

563 compilation_counter.stock_torch_compile_count += 1

564 self.worker.model_runner.model.compile(fullgraph=True, backend=backend)

565

566 + return retired_groups

567 +

568 def _perform_eplb_reshuffle(

517 - self, rank_mapping: dict[int, int] | None = None

569 + self,

570 + rank_mapping: dict[int, int] | None = None,

571 + async_op: bool = False,

572 ) -> None:

573 if get_ep_group().rank == 0:

574 logger.info("[Elastic EP] Starting expert resharding...")

@@ -522,21 +576,19 @@ class ElasticEPScalingExecutor:

576 eplb_state = self.worker.model_runner.eplb_state

577 assert eplb_state is not None

578

525 - model_config = self.worker.model_runner.model_config

526 - eplb_model_state = eplb_state.model_states[model_config.compute_hash()]

579 is_async_enabled = eplb_state.is_async

528 - eplb_state.is_async = False

580 + run_async = async_op and is_async_enabled

581 + eplb_state.is_async = run_async

582 if rank_mapping is None:

583 eplb_state.rearrange()

584 else:

585 eplb_state.rearrange(rank_mapping=rank_mapping)

533 - # NOTE(yongji): check whether we need to synchronize here

534 - torch.accelerator.synchronize()

586 + if not run_async:

587 + # Wait for non-blocking expert resharding copies before continuing

588 + # the Elastic EP reconfiguration.

589 + torch.accelerator.synchronize()

590 # reset expert_rearrangement_step to ensure all ranks are synchronized

591 eplb_state.expert_rearrangement_step = 0

537 - eplb_state.num_valid_physical_experts = (

538 - eplb_model_state.physical_to_logical_map.shape[1]

539 - )

592 eplb_state.is_async = is_async_enabled

593 # Start the async worker thread if it doesn't exist yet (idempotent).

594 # This is needed for new workers after scale-up: they create EplbState

@@ -544,25 +596,31 @@ class ElasticEPScalingExecutor:

596 # groups aren't ready yet.

597 eplb_state.start_async_loop()

598 if get_ep_group().rank == 0:

547 - logger.info("[Elastic EP] Expert resharding completed")

599 + logger.info(

600 + "[Elastic EP] Expert resharding %s",

601 + "scheduled" if run_async else "completed",

602 + )

603

604 def commit_scale_up(self, is_existing_worker: bool) -> None:

605 if is_existing_worker:

606 self.broadcast_expert_mapping()

552 - self.switch_and_prepare()

607 + retired_groups = self.switch_and_prepare()

608 else:

554 - mapping, _, num_valid_experts = self.receive_expert_mapping()

555 - self.worker.model_runner.setup_eplb_from_mapping(mapping, num_valid_experts)

556 - self._perform_eplb_reshuffle()

609 + mapping = self.receive_expert_mapping()

610 + self.worker.model_runner.setup_eplb_from_mapping(mapping)

611 self.warm_and_capture()

612 + self._perform_eplb_reshuffle(async_op=True)

613 + if is_existing_worker:

614 + self._start_group_cleanup(retired_groups)

615

616 def commit_scale_down(self, new_dp_size: int, removing: bool) -> None:

617 self.perform_scale_down_eplb_reshuffle(new_dp_size)

618 if removing:

619 self.switch_and_remove()

620 else:

564 - self.switch_and_prepare()

621 + retired_groups = self.switch_and_prepare()

622 self.warm_and_capture()

623 + self._start_group_cleanup(retired_groups)

624

625 def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None:

626 eplb_state = self.worker.model_runner.eplb_state

@@ -578,7 +636,7 @@ class ElasticEPScalingExecutor:

636 }

637 self._perform_eplb_reshuffle(rank_mapping=rank_mapping)

638

581 - def receive_weights(self) -> None:

639 + def prepare_new_worker(self) -> None:

640 dp_group = get_dp_group()

641 assert isinstance(dp_group, StatelessGroupCoordinator)

642 new_dp_size = dp_group.world_size

@@ -618,19 +676,18 @@ class ElasticEPScalingExecutor:

676 expert_weights=expert_weights,

677 )

678 torch.accelerator.synchronize()

679 + self._warm_target_groups(get_dp_group(), get_ep_group())

680

622 - def receive_expert_mapping(self) -> tuple[torch.Tensor, int, int]:

681 + def receive_expert_mapping(self) -> torch.Tensor:

682 dp_group = get_dp_group()

683 assert isinstance(dp_group, StatelessGroupCoordinator)

625 - physical_to_logical, num_local_physical_experts, num_logical_experts = (

626 - broadcast_expert_mapping(

627 - physical_to_logical=None,

628 - num_local_physical_experts=None,

629 - num_logical_experts=None,

630 - dp_group=dp_group,

631 - src_rank=0,

632 - device=self.worker.device,

633 - )

684 + physical_to_logical, num_local_physical_experts, _ = broadcast_expert_mapping(

685 + physical_to_logical=None,

686 + num_local_physical_experts=None,

687 + num_logical_experts=None,

688 + dp_group=dp_group,

689 + src_rank=0,

690 + device=self.worker.device,

691 )

692 num_moe_layers = physical_to_logical.shape[0]

693 new_dp_size = get_dp_group().world_size

@@ -642,13 +699,10 @@ class ElasticEPScalingExecutor:

699 dtype=physical_to_logical.dtype,

700 device=physical_to_logical.device,

701 )

645 - old_num_physical_experts = physical_to_logical.shape[1]

646 - expanded_physical_to_logical[:, :old_num_physical_experts] = physical_to_logical

647 - return (

648 - expanded_physical_to_logical,

649 - num_logical_experts,

650 - old_num_physical_experts,

702 + expanded_physical_to_logical[:, : physical_to_logical.shape[1]] = (

703 + physical_to_logical

704 )

705 + return expanded_physical_to_logical

706

707 def warmup_local_kernels(self) -> None:

708 with set_current_vllm_config(self.worker.vllm_config):

vllm/distributed/elastic_ep/elastic_state.py

+12 -37

@@ -30,12 +30,10 @@ WorkerType = Literal["existing", "new", "removing"]

32 class ScaleUpExistingEngineState(enum.IntEnum):

33 - CREATE_STANDBY_GROUPS = 0

34 - STAGE_QUANT_METHODS = 1

35 - TRANSFER_WEIGHTS = 2

36 - SYNC_KV_CACHE_MEMORY_SIZE = 3

37 - COMMIT_SCALE_UP = 4 # Blocks forward passes.

38 - COMPLETE = 5

33 + PREPARE = 0

34 + SYNC_KV_CACHE_MEMORY_SIZE = 1

35 + COMMIT_SCALE_UP = 2 # Blocks forward passes.

36 + COMPLETE = 3

39 class ScaleUpNewEngineState(enum.IntEnum):

@@ -95,7 +93,7 @@ class ElasticEPScalingState:

93 self.state = (

94 ScaleUpNewEngineState.PRE_KV_INIT

95 if worker_type == "new"

98 - else ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS

96 + else ScaleUpExistingEngineState.PREPARE

97 )

98 else:

99 self.state = (

@@ -166,20 +164,8 @@ class ElasticEPScalingState:

164 state = self.state

165 assert self.old_dp_group is not None

166

169 - if state == ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS:

170 - if not self._create_standby_groups():

171 - return False

172 - self.state = ScaleUpExistingEngineState.STAGE_QUANT_METHODS

173 - return True

174 -

175 - elif state == ScaleUpExistingEngineState.STAGE_QUANT_METHODS:

176 - if not self._execute_async("stage_standby_moe_quant_methods"):

177 - return False

178 - self.state = ScaleUpExistingEngineState.TRANSFER_WEIGHTS

179 - return True

180 -

181 - elif state == ScaleUpExistingEngineState.TRANSFER_WEIGHTS:

182 - if not self._transfer_weights():

167 + if state == ScaleUpExistingEngineState.PREPARE:

168 + if not self._prepare_workers():

169 return False

170 self.state = ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE

171 return True

@@ -210,7 +196,7 @@ class ElasticEPScalingState:

196 assert self.new_dp_group is not None and self.new_dp_store is not None

197

198 if state == ScaleUpNewEngineState.PRE_KV_INIT:

213 - self._collective_rpc("elastic_ep_execute", args=("receive_weights",))

199 + self._collective_rpc("elastic_ep_execute", args=("prepare_new_worker",))

200 self.engine_core.available_gpu_memory_for_kv_cache = (

201 ParallelConfig.sync_kv_cache_memory_size(self.new_dp_group, -1)

202 )

@@ -243,7 +229,7 @@ class ElasticEPScalingState:

229 assert self.old_dp_group is not None

230

231 if state == ScaleDownRemainingEngineState.PREPARE:

246 - if self._create_standby_groups():

232 + if self._prepare_workers():

233 self.state = ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN

234 self._mark_ready_for_switch()

235 return True

@@ -328,29 +314,18 @@ class ElasticEPScalingState:

314 self._prepare_future = None

315 return True

316

331 - def _create_standby_groups(self) -> bool:

317 + def _prepare_workers(self) -> bool:

318 assert self.old_dp_group is not None

319 if not self._ensure_new_dp_group():

320 return False

321 if not self._execute_async(

336 - "create_standby_groups",

322 + "prepare_reconfiguration",

323 self.reconfig_request,

324 self.new_parallel_config.use_all2all,

325 ):

326 return False

327 if self.old_dp_group.rank() == 0:

342 - logger.info("[Elastic EP] Created standby communication groups")

343 - return True

344 -

345 - def _transfer_weights(self) -> bool:

346 - assert self.reconfig_request is not None and self.old_dp_group is not None

347 - old_dp_size = self.old_dp_group.size()

348 - new_dp_size = self.reconfig_request.new_data_parallel_size

349 -

350 - if not self._execute_async("transfer_weights", old_dp_size, new_dp_size):

351 - return False

352 - if self.old_dp_group.rank() == 0:

353 - logger.info("[Elastic EP] Transferred weights to new workers")

328 + logger.info("[Elastic EP] Prepared reconfiguration")

329 return True

330

331 def _sync_kv_cache_memory_size(self) -> bool:

vllm/distributed/eplb/eplb_communicator.py

+12 -36

@@ -248,7 +248,6 @@ class NixlEplbCommunicator(EplbCommunicator):

248 cpu_group: ProcessGroup,

249 all_expert_weights: Sequence[Sequence[torch.Tensor]],

250 expert_buffer: Sequence[torch.Tensor],

251 - defer_remote_setup: bool = False,

251 ) -> None:

252 """Create a NIXL-backed EPLB communicator.

253

@@ -256,11 +255,6 @@ class NixlEplbCommunicator(EplbCommunicator):

255 cpu_group: CPU process group for metadata exchange.

256 all_expert_weights: Expert weight tensors for all MoE layers.

257 expert_buffer: Pre-allocated receive buffer tensors.

259 - defer_remote_setup: If True, postpone the collective

260 - all-gather of NIXL agent metadata until the first

261 - ``set_transfer_context`` call. Required for elastic EP

262 - where ranks join asynchronously and cannot participate

263 - in collectives at construction time.

258 """

259 assert all_expert_weights, (

260 "NixlEplbCommunicator requires non-empty all_expert_weights."

@@ -317,29 +311,17 @@ class NixlEplbCommunicator(EplbCommunicator):

311 ] = {}

312

313 self._cuda_device_id = int(self._device.index or 0)

320 - self._remote_state_initialized = False

314 self._init_step("buffers", self._init_registered_buffers)

322 - if defer_remote_setup:

323 - logger.info_once("NIXL EPLB: deferring remote agent setup (elastic EP).")

324 - else:

325 - self._init_remote_state()

315 + self._init_remote_state()

316 self._log_initialized()

317

318 def _init_remote_state(self) -> None:

319 """Exchange NIXL agent metadata and RDMA pointer info with all peers.

320

321 This is a collective operation (uses ``all_gather_object`` twice).

332 - Under elastic EP the call is deferred to the first

333 - ``set_transfer_context`` invocation, where all ranks are

334 - guaranteed to be synchronized.

322 """

323 self._init_step("agents", self._init_remote_agents)

324 self._init_step("send meta", self._exchange_remote_send_meta)

338 - self._remote_state_initialized = True

339 -

340 - def _ensure_remote_state(self) -> None:

341 - if not self._remote_state_initialized:

342 - self._init_remote_state()

325

326 @property

327 def needs_profile_buffer_reservation(self) -> bool:

@@ -373,7 +355,6 @@ class NixlEplbCommunicator(EplbCommunicator):

355 pass

356

357 def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None:

376 - self._ensure_remote_state()

358 assert not self._xfer_entries, (

359 f"set_transfer_context() called with {len(self._xfer_entries)} "

360 f"pending transfers from layer {self._layer_idx}; "

@@ -673,9 +654,8 @@ def create_eplb_communicator(

654 Falls back to ``"torch_nccl"`` when *None*.

655 Stateless (elastic EP) groups support ``"torch_nccl"``,

656 ``"pynccl"``, and ``"nixl"``; ``"torch_nccl"`` is silently

676 - promoted to ``"pynccl"``. ``"nixl"`` uses deferred remote

677 - agent setup to avoid collective deadlocks during elastic

678 - scaling. When tensors reside on CPU, ``"torch_gloo"`` or

657 + promoted to ``"pynccl"``. When tensors reside on CPU,

658 + ``"torch_gloo"`` or

659 ``"torch_nccl"`` are used via the CPU process group.

660 expert_weights: Expert weight tensors for *all* MoE layers.

661 Shape ``(num_layers)(num_tensors_per_layer)``.

@@ -729,22 +709,19 @@ def create_eplb_communicator(

709 ) from exc

710

711 is_stateless = isinstance(group_coordinator, StatelessGroupCoordinator)

732 - if is_stateless:

733 - if backend == "nixl":

734 - pass # handled below with defer_remote_setup=True

735 - elif backend not in ("torch_nccl", "pynccl"):

712 + if is_stateless and backend != "nixl":

713 + if backend not in ("torch_nccl", "pynccl"):

714 raise ValueError(

715 f"Elastic EP requires 'torch_nccl', 'pynccl', or 'nixl' "

716 f"EPLB communicator (got '{backend}')."

717 )

740 - else:

741 - if backend == "torch_nccl":

742 - logger.warning(

743 - "Stateless elastic EP requires PyNCCL backend. "

744 - "Forcing EPLB communicator to 'pynccl'."

745 - )

746 - backend = "pynccl"

747 - return _create_pynccl()

718 + if backend == "torch_nccl":

719 + logger.warning(

720 + "Stateless elastic EP requires PyNCCL backend. "

721 + "Forcing EPLB communicator to 'pynccl'."

722 + )

723 + backend = "pynccl"

724 + return _create_pynccl()

725

726 if backend == "nixl":

727 if not has_nixl():

@@ -761,7 +738,6 @@ def create_eplb_communicator(

738 cpu_group=group_coordinator.cpu_group,

739 all_expert_weights=expert_weights,

740 expert_buffer=expert_buffer,

764 - defer_remote_setup=is_stateless,

741 )

742 except Exception as exc:

743 raise RuntimeError(

vllm/distributed/eplb/eplb_state.py

+46 -25

@@ -37,6 +37,7 @@ from torch.distributed import ProcessGroup, all_reduce

37 from vllm.config import ModelConfig, ParallelConfig

38 from vllm.config.utils import compute_hash_cached

39 from vllm.distributed.parallel_state import (

40 + GroupCoordinator,

41 get_ep_group,

42 get_eplb_group,

43 get_node_count,

@@ -289,14 +290,6 @@ class EplbState:

290 """

291 CUDA device index for the async EPLB worker thread.

292 """

292 - self.num_valid_physical_experts: int = 0

293 - """

294 - Number of valid physical experts.

295 - This is the number of physical experts that are

296 - actually mapped to logical experts. In elastic EP,

297 - newly started EP ranks may not have physical experts

298 - mapped yet.

299 - """

293 if self.device.type == "cuda":

294 self.cuda_device_index = self.device.index

295 if self.cuda_device_index is None and torch.cuda.is_available():

@@ -499,7 +492,6 @@ class EplbState:

492 num_unpadded_tokens_tensors=num_unpadded_tokens_tensors,

493 )

494 self.model_states[model_config.compute_hash()] = model_state

502 - self.num_valid_physical_experts = model.num_physical_experts

495

496 def prepare_forward(

497 self,

@@ -766,28 +758,28 @@ class EplbState:

758 # Map the physical expert load to global logical experts

759 global_expert_load_windows = []

760 for eplb_model_state in self.model_states.values():

769 - expert_load_window = eplb_model_state.expert_load_window[

770 - :, :, : self.num_valid_physical_experts

771 - ]

761 + expert_load_window = eplb_model_state.expert_load_window

762 + physical_to_logical = eplb_model_state.physical_to_logical_map

763 + invalid_idx = eplb_model_state.model.num_logical_experts

764 logical_expert_load_window = torch.zeros(

765 self.expert_load_window_size,

766 eplb_model_state.model.num_moe_layers,

775 - eplb_model_state.model.num_logical_experts,

767 + invalid_idx + 1,

768 dtype=eplb_model_state.expert_load_window.dtype,

769 device=eplb_model_state.expert_load_window.device,

770 )

771 logical_expert_load_window.scatter_add_(

772 dim=-1,

781 - index=eplb_model_state.physical_to_logical_map[

782 - :, : self.num_valid_physical_experts

783 - ]

773 + index=physical_to_logical.masked_fill(

774 + physical_to_logical < 0, invalid_idx

775 + )

776 .unsqueeze(0)

777 .expand_as(expert_load_window)

778 .long(),

779 src=expert_load_window,

780 )

781

790 - global_expert_load_window = logical_expert_load_window.sum(dim=0)

782 + global_expert_load_window = logical_expert_load_window[..., :-1].sum(dim=0)

783 global_expert_load_windows.append(global_expert_load_window)

784 # Perform all-reduce to get the expert load across all ranks for each model

785 global_expert_load_windows = self._allreduce_list(global_expert_load_windows)

@@ -960,7 +952,7 @@ class EplbState:

952

953 Each pending result is acknowledged (consumed_event recorded) so the

954 async worker can proceed, but the transferred weights are intentionally

963 - NOT applied — a full synchronous rearrange is expected to follow.

955 + NOT applied — a full rearrange is expected to follow.

956

957 Ranks are kept in lockstep via _all_ranks_result_ready (all_reduce

958 on the EP CPU group). The async worker's coordinated-stop collectives

@@ -1060,7 +1052,6 @@ class EplbState:

1052 device: torch.device,

1053 parallel_config: ParallelConfig,

1054 expanded_physical_to_logical: torch.Tensor,

1063 - num_valid_physical_experts: int,

1055 ) -> "EplbState":

1056 eplb_state = cls(

1057 parallel_config=parallel_config,

@@ -1070,12 +1061,24 @@ class EplbState:

1061 model=model,

1062 model_config=model_config,

1063 )

1073 - eplb_state.num_valid_physical_experts = num_valid_physical_experts

1074 - eplb_model_state = eplb_state.model_states[model_config.compute_hash()]

1064 + eplb_state.update_mapping(

1065 + model_config,

1066 + expanded_physical_to_logical,

1067 + )

1068 +

1069 + return eplb_state

1070 +

1071 + def update_mapping(

1072 + self,

1073 + model_config: ModelConfig,

1074 + expanded_physical_to_logical: torch.Tensor,

1075 + ) -> None:

1076 + eplb_model_state = self.model_states[model_config.compute_hash()]

1077 eplb_model_state.physical_to_logical_map.copy_(expanded_physical_to_logical)

1078

1079 (logical_to_physical_map_cpu, logical_replica_count_cpu) = compute_logical_maps(

1078 - expanded_physical_to_logical.cpu(), model.num_logical_experts

1080 + expanded_physical_to_logical.cpu(),

1081 + eplb_model_state.model.num_logical_experts,

1082 )

1083

1084 max_num_replicas = eplb_model_state.logical_to_physical_map.shape[-1]

@@ -1087,13 +1090,31 @@ class EplbState:

1090 max_num_replicas - num_replicas,

1091 ),

1092 value=-1,

1090 - ).to(device)

1091 - logical_replica_count = logical_replica_count_cpu.to(device)

1093 + ).to(self.device)

1094 + logical_replica_count = logical_replica_count_cpu.to(self.device)

1095

1096 eplb_model_state.logical_to_physical_map.copy_(logical_to_physical_map)

1097 eplb_model_state.logical_replica_count.copy_(logical_replica_count)

1098

1096 - return eplb_state

1099 + def create_communicator(

1100 + self, model_config: ModelConfig, group_coordinator: GroupCoordinator

1101 + ) -> EplbCommunicator:

1102 + model_state = self.model_states[model_config.compute_hash()]

1103 + backend = self.parallel_config.eplb_config.communicator

1104 + assert backend is not None

1105 + return create_eplb_communicator(

1106 + group_coordinator,

1107 + backend,

1108 + model_state.model.expert_weights,

1109 + model_state.expert_buffer,

1110 + )

1111 +

1112 + def update_communicator(

1113 + self,

1114 + model_config: ModelConfig,

1115 + communicator: EplbCommunicator,

1116 + ) -> None:

1117 + self.model_states[model_config.compute_hash()].communicator = communicator

1118

1119

1120 @dataclass

vllm/distributed/parallel_state.py

+7 -7

@@ -1366,21 +1366,21 @@ def _replace_active_groups(

1366 ep: GroupCoordinator | None,

1367 eplb: GroupCoordinator | None,

1368 node_count: int | None,

1369 -) -> None:

1370 - """Destroy the current DP/EP/WORLD/EPLB groups and replace them.

1369 +) -> tuple[GroupCoordinator | None, ...]:

1370 + """Replace the active groups and return the groups they replaced.

1371

1372 - Destruction is collective — all ranks in the old groups must call this

1373 - function together. Pass all-``None`` to tear down without replacement.

1372 + The caller must destroy the returned DP, EP, WORLD, and EPLB groups

1373 + collectively and in that order. Pass all-``None`` to remove the active

1374 + groups without replacement.

1375 """

1376 global _WORLD, _DP, _EP, _EPLB, _NODE_COUNT

1376 - for group in (_DP, _EP, _WORLD, _EPLB):

1377 - if group is not None:

1378 - group.destroy()

1377 + old_groups = _DP, _EP, _WORLD, _EPLB

1378 _WORLD = world

1379 _DP = dp

1380 _EP = ep

1381 _EPLB = eplb

1382 _NODE_COUNT = node_count

1383 + return old_groups

1384

1385

1386 _TP: GroupCoordinator | None = None

vllm/v1/worker/gpu/eplb_utils.py

-2

@@ -144,7 +144,6 @@ class EPLBController:

144 model: nn.Module,

145 model_config: Any,

146 expanded_physical_to_logical: torch.Tensor,

147 - old_num_physical_experts: int,

147 ) -> None:

148 moe_model = get_mixture_of_experts_model(model)

149 assert moe_model is not None

@@ -155,6 +154,5 @@ class EPLBController:

154 device=self.device,

155 parallel_config=self.parallel_config,

156 expanded_physical_to_logical=expanded_physical_to_logical,

158 - num_valid_physical_experts=old_num_physical_experts,

157 )

158 self._has_registered_models = True

vllm/v1/worker/gpu/model_runner.py

-2

@@ -1962,13 +1962,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1962 def setup_eplb_from_mapping(

1963 self,

1964 expanded_physical_to_logical: torch.Tensor,

1965 - old_num_physical_experts: int,

1965 ) -> None:

1966 self.eplb.setup_from_mapping(

1967 self.model,

1968 self.model_config,

1969 expanded_physical_to_logical,

1971 - old_num_physical_experts,

1970 )

1971

1972 ########### EPLB methods end ###########

vllm/v1/worker/gpu_model_runner.py

+5 -15

@@ -3517,17 +3517,11 @@ class GPUModelRunner(

3517 def setup_eplb_from_mapping(

3518 self,

3519 expanded_physical_to_logical: torch.Tensor,

3520 - old_num_physical_experts: int,

3520 ) -> None:

3522 - assert self._moe_model is not None

3523 -

3524 - self.eplb_state = EplbState.from_mapping(

3525 - model=self._moe_model,

3526 - model_config=self.model_config,

3527 - device=self.device,

3528 - parallel_config=self.parallel_config,

3529 - expanded_physical_to_logical=expanded_physical_to_logical,

3530 - num_valid_physical_experts=old_num_physical_experts,

3521 + assert self.eplb_state is not None

3522 + self.eplb_state.update_mapping(

3523 + self.model_config,

3524 + expanded_physical_to_logical,

3525 )

3526

3527 def _pool(

@@ -5488,11 +5482,7 @@ class GPUModelRunner(

5482 # MixtureOfExperts themselves.

5483 self._moe_model = get_mixture_of_experts_model(self.model)

5484

5491 - if (

5492 - self.parallel_config.enable_eplb

5493 - and not load_dummy_weights

5494 - and self._moe_model is not None

5495 - ):

5485 + if self._moe_model is not None and self.parallel_config.enable_eplb:

5486 logger.info_once(

5487 "EPLB is enabled for MoE part of model %s.",

5488 self.model_config.model,

vllm/v1/worker/gpu_worker.py

@@ -1376,6 +1376,8 @@ class Worker(WorkerBase):

1376 if weight_transfer_engine := getattr(self, "weight_transfer_engine", None):

1377 weight_transfer_engine.shutdown()

1378

1379 + self.elastic_ep_executor.shutdown()

1380 +

1381 # Release GPU resources held by the model runner so that memory

1382 # can be reclaimed when running in-process

1383 if model_runner := getattr(self, "model_runner", None):