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

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 6, 07:00 UTC

elastic_ep: Fix issues with repeated scale up/down cycles (#37131)

Signed-off-by: Itay Alroy <ialroy@nvidia.com> Co-authored-by: Ron Tourgeman <rtourgeman@nvidia.com>

Itay Alroy committed Mar 21, 2026 at 01:13 UTC c57d38d603213a9acfd5e83f38d45f9d635124fb

10 files changed +129 -90

vllm/distributed/device_communicators/cuda_communicator.py

@@ -338,6 +338,7 @@ class CudaCommunicator(DeviceCommunicatorBase):

338

339 def destroy(self):

340 if self.pynccl_comm is not None:

341 + self.pynccl_comm.destroy()

342 self.pynccl_comm = None

343 if self.ca_comm is not None:

344 self.ca_comm = None

vllm/distributed/device_communicators/pynccl.py

@@ -145,6 +145,13 @@ class PyNcclCommunicator:

145 stream.synchronize()

146 del data

147

148 + def destroy(self):

149 + if self.available and not self.disabled:

150 + with torch.accelerator.device_index(self.device.index):

151 + self.nccl.ncclCommDestroy(self.comm)

152 + self.available = False

153 + self.disabled = True

154 +

155 def all_reduce(

156 self,

157 in_tensor: torch.Tensor,

vllm/distributed/elastic_ep/elastic_execute.py

+68 -34

@@ -145,11 +145,37 @@ class ElasticEPScalingExecutor:

145 raise ValueError(f"Unknown execute method: {execute_method}")

146 return method(*args, **kwargs)

147

148 + def _set_eplb_suppressed(self, suppressed: bool) -> None:

149 + self.worker.model_runner.eep_eplb_suppressed = suppressed

150 + ep_group = get_standby_ep_group() or get_ep_group()

151 + if ep_group.rank == 0:

152 + logger.info(

153 + "[Elastic EP] EPLB %s elastic scaling transition",

154 + "disabled during" if suppressed else "re-enabled after",

155 + )

156 +

157 + def load_model(self) -> None:

158 + (

159 + expanded_physical_to_logical,

160 + num_logical_experts,

161 + old_num_physical_experts,

162 + ) = self.receive_expert_mapping()

163 + num_physical_experts = expanded_physical_to_logical.shape[1]

164 + self.worker.parallel_config.eplb_config.num_redundant_experts = (

165 + num_physical_experts - num_logical_experts

166 + )

167 + self.worker.load_model(load_dummy_weights=True)

168 + self.worker.model_runner.setup_eplb_from_mapping(

169 + expanded_physical_to_logical, old_num_physical_experts

170 + )

171 + self._set_eplb_suppressed(True)

172 +

173 def create_standby_groups(

174 self, reconfig_request: ReconfigureDistributedRequest

175 ) -> None:

176 self.reconfig_request = reconfig_request

177 new_dp_size = reconfig_request.new_data_parallel_size

178 + old_dp_size = get_dp_group().world_size

179 world_size = self.worker.vllm_config.parallel_config.world_size

180 new_world_size_across_dp = world_size * new_dp_size

181 updated_config = copy.copy(self.worker.vllm_config)

@@ -165,11 +191,8 @@ class ElasticEPScalingExecutor:

191 coord_store_port=reconfig_request.coord_store_port,

192 enable_eplb=updated_config.parallel_config.enable_eplb,

193 )

168 - self.worker.model_runner.eep_eplb_suppressed = True

169 - standby_ep_group = get_standby_ep_group()

170 - assert standby_ep_group is not None

171 - if standby_ep_group.rank == 0:

172 - logger.info("[Elastic EP] EPLB disabled during elastic scaling transition")

194 + if new_dp_size > old_dp_size:

195 + self._set_eplb_suppressed(True)

196

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

198 standby_dp_group = get_standby_dp_group()

@@ -237,13 +260,31 @@ class ElasticEPScalingExecutor:

260 device=self.worker.device,

261 )

262

263 + def _release_cuda_graphs(self) -> None:

264 + if isinstance(self.worker.model_runner.model, CUDAGraphWrapper):

265 + wrapper = self.worker.model_runner.model

266 + wrapper.concrete_cudagraph_entries = {}

267 +

268 + elif isinstance(self.worker.model_runner.model, UBatchWrapper):

269 + raise RuntimeError("DBO is not yet supported in elastic EP")

270 +

271 + torch.compiler.reset()

272 + with set_current_vllm_config(self.worker.vllm_config):

273 + reset_compile_wrapper(self.worker.model_runner.get_model())

274 +

275 + gc.collect()

276 + torch.accelerator.synchronize()

277 + torch.accelerator.empty_cache()

278 +

279 def switch_and_remove(self) -> None:

280 + self._release_cuda_graphs()

281 _replace_active_groups(world=None, dp=None, ep=None, eplb=None, node_count=None)

282

283 def switch_and_prepare(self) -> None:

284 old_dp_size = get_dp_group().world_size

285 old_ep_size = get_ep_group().world_size

286

287 + self._release_cuda_graphs()

288 _replace_active_groups(**pop_standby_groups())

289

290 parallel_config = self.worker.vllm_config.parallel_config

@@ -384,13 +425,6 @@ class ElasticEPScalingExecutor:

425 compilation_counter.stock_torch_compile_count += 1

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

427

387 - # release all previously captured CUDA graphs

388 - if isinstance(self.worker.model_runner.model, CUDAGraphWrapper):

389 - wrapper = self.worker.model_runner.model

390 - wrapper.concrete_cudagraph_entries = {}

391 - elif isinstance(self.worker.model_runner.model, UBatchWrapper):

392 - raise RuntimeError("DBO is not yet supported in elastic EP")

393 -

428 multi_block_table = self.worker.model_runner.input_batch.block_table

429 saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = []

430 for bt in multi_block_table.block_tables:

@@ -399,14 +433,6 @@ class ElasticEPScalingExecutor:

433 )

434 multi_block_table.clear()

435

402 - # reset the compile wrapper

403 - torch.compiler.reset()

404 - with set_current_vllm_config(self.worker.vllm_config):

405 - reset_compile_wrapper(self.worker.model_runner.get_model())

406 -

407 - gc.collect()

408 - torch.accelerator.synchronize()

409 - torch.accelerator.empty_cache()

436 unlock_workspace()

437 self.worker.compile_or_warm_up_model()

438 lock_workspace()

@@ -416,8 +442,12 @@ class ElasticEPScalingExecutor:

442 ):

443 bt.block_table.gpu.copy_(saved_gpu)

444 bt.block_table.cpu.copy_(saved_cpu)

445 + if new_dp_size < old_dp_size:

446 + self._set_eplb_suppressed(False)

447

420 - def perform_eplb_reshuffle(self, new_dp_size: int | None = None) -> None:

448 + def _perform_eplb_reshuffle(

449 + self, rank_mapping: dict[int, int] | None = None

450 + ) -> None:

451 if get_ep_group().rank == 0:

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

453

@@ -428,20 +458,9 @@ class ElasticEPScalingExecutor:

458 eplb_model_state = eplb_state.model_states[model_config.compute_hash()]

459 is_async_enabled = eplb_state.is_async

460 eplb_state.is_async = False

431 - if new_dp_size is None:

461 + if rank_mapping is None:

462 eplb_state.rearrange()

463 else:

434 - # scale down

435 - parallel_config = self.worker.vllm_config.parallel_config

436 - tp_size = parallel_config.tensor_parallel_size

437 - old_ep_size = parallel_config.data_parallel_size * tp_size

438 - new_ep_size = new_dp_size * tp_size

439 -

440 - rank_mapping = {

441 - old_ep_rank: old_ep_rank if old_ep_rank < new_ep_size else -1

442 - for old_ep_rank in range(old_ep_size)

443 - }

444 -

464 eplb_state.rearrange(rank_mapping=rank_mapping)

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

466 torch.accelerator.synchronize()

@@ -451,10 +470,25 @@ class ElasticEPScalingExecutor:

470 eplb_model_state.physical_to_logical_map.shape[1]

471 )

472 eplb_state.is_async = is_async_enabled

454 - self.worker.model_runner.eep_eplb_suppressed = False

473 if get_ep_group().rank == 0:

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

475

476 + def perform_eplb_reshuffle(self) -> None:

477 + self._perform_eplb_reshuffle()

478 + self._set_eplb_suppressed(False)

479 +

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

481 + self._set_eplb_suppressed(True)

482 + parallel_config = self.worker.vllm_config.parallel_config

483 + tp_size = parallel_config.tensor_parallel_size

484 + old_ep_size = parallel_config.data_parallel_size * tp_size

485 + new_ep_size = new_dp_size * tp_size

486 + rank_mapping = {

487 + old_ep_rank: old_ep_rank if old_ep_rank < new_ep_size else -1

488 + for old_ep_rank in range(old_ep_size)

489 + }

490 + self._perform_eplb_reshuffle(rank_mapping=rank_mapping)

491 +

492 def receive_weights(self) -> None:

493 dp_group = get_dp_group()

494 assert isinstance(dp_group, StatelessGroupCoordinator)

vllm/distributed/elastic_ep/elastic_state.py

+29 -7

@@ -43,9 +43,10 @@ class ScaleUpExistingEngineState(enum.IntEnum):

45 class ScaleUpNewEngineState(enum.IntEnum):

46 - PREPARE = 0

47 - EPLB_RESHUFFLE = 1

48 - COMPLETE = 2

46 + PRE_KV_INIT = 0

47 + PREPARE = 1

48 + EPLB_RESHUFFLE = 2

49 + COMPLETE = 3

52 class ScaleDownRemainingEngineState(enum.IntEnum):

@@ -104,7 +105,7 @@ class ElasticEPScalingState:

105 self.state: EngineState

106 if scale_type == "scale_up":

107 self.state = (

107 - ScaleUpNewEngineState.PREPARE

108 + ScaleUpNewEngineState.PRE_KV_INIT

109 if worker_type == "new"

110 else ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT

111 )

@@ -142,6 +143,12 @@ class ElasticEPScalingState:

143 else self._progress_remaining_engine()

144 )

145

146 + def run_pre_kv_init_states(self) -> None:

147 + assert self.scale_type == "scale_up" and self.worker_type == "new"

148 + assert self.state == ScaleUpNewEngineState.PRE_KV_INIT

149 + assert self.progress()

150 + assert self.state == ScaleUpNewEngineState.PREPARE

151 +

152 def _execute_tcp_store_barrier(

153 self, dp_store, group_rank, group_size, barrier_id, timeout=None

154 ):

@@ -303,7 +310,23 @@ class ElasticEPScalingState:

310 state = self.state

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

312

306 - if state == ScaleUpNewEngineState.PREPARE:

313 + if state == ScaleUpNewEngineState.PRE_KV_INIT:

314 + self.engine_core._eep_send_engine_core_notification(

315 + EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY

316 + )

317 + self.model_executor.collective_rpc(

318 + "elastic_ep_execute", args=("receive_weights",)

319 + )

320 + self.engine_core.available_gpu_memory_for_kv_cache = (

321 + ParallelConfig.sync_kv_cache_memory_size(self.new_dp_group, -1)

322 + )

323 + self.model_executor.collective_rpc(

324 + "elastic_ep_execute", args=("prepare_new_worker",)

325 + )

326 + self.state = ScaleUpNewEngineState.PREPARE

327 + return True

328 +

329 + elif state == ScaleUpNewEngineState.PREPARE:

330 tensor = torch.tensor([0, 0, 0], dtype=torch.int32, device="cpu")

331 torch.distributed.all_reduce(

332 tensor,

@@ -403,7 +426,6 @@ class ElasticEPScalingState:

426 self.engine_core._eep_send_engine_core_notification(

427 EEPNotificationType.SHUTDOWN_COMPLETE

428 )

406 - self.engine_core.shutdown()

429 return True

430

431 else:

@@ -525,7 +547,7 @@ class ElasticEPScalingState:

547 self.model_executor.collective_rpc(

548 "elastic_ep_execute",

549 args=(

528 - "perform_eplb_reshuffle",

550 + "perform_scale_down_eplb_reshuffle",

551 self.reconfig_request.new_data_parallel_size,

552 ),

553 )

vllm/v1/engine/core.py

+3 -14

@@ -1694,6 +1694,8 @@ class DPEngineCoreProc(EngineCoreProc):

1694 if self.eep_scaling_state is not None:

1695 _ = self.eep_scaling_state.progress()

1696 if self.eep_scaling_state.is_complete():

1697 + if self.eep_scaling_state.worker_type == "removing":

1698 + raise SystemExit

1699 self.process_input_queue_block = True

1700 self.eep_scaling_state = None

1701

@@ -1857,20 +1859,7 @@ class DPEngineCoreProc(EngineCoreProc):

1859 scale_type="scale_up",

1860 reconfig_request=None,

1861 )

1860 - self.model_executor.collective_rpc("init_device")

1861 - self.model_executor.collective_rpc("load_model")

1862 - self._eep_send_engine_core_notification(

1863 - EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY

1864 - )

1865 - self.model_executor.collective_rpc(

1866 - "elastic_ep_execute", args=("receive_weights",)

1867 - )

1868 - self.available_gpu_memory_for_kv_cache = (

1869 - ParallelConfig.sync_kv_cache_memory_size(self.dp_group, -1)

1870 - )

1871 - self.model_executor.collective_rpc(

1872 - "elastic_ep_execute", args=("prepare_new_worker",)

1873 - )

1862 + self.eep_scaling_state.run_pre_kv_init_states()

1863 self.process_input_queue_block = False

1864

1865

vllm/v1/executor/multiproc_executor.py

+8 -7

@@ -602,13 +602,14 @@ class WorkerProc:

602 )

603

604 # Load model

605 - is_eep_new_worker = envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH

606 - if not is_eep_new_worker:

607 - self.worker.init_device()

608 - # Update process title now that parallel groups are initialized

609 - self.setup_proc_title_and_log_prefix(

610 - enable_ep=vllm_config.parallel_config.enable_expert_parallel

611 - )

605 + self.worker.init_device()

606 + # Update process title now that parallel groups are initialized

607 + self.setup_proc_title_and_log_prefix(

608 + enable_ep=vllm_config.parallel_config.enable_expert_parallel

609 + )

610 + if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:

611 + self.worker.elastic_ep_execute("load_model")

612 + else:

613 self.worker.load_model()

614

615 scheduler_config = vllm_config.scheduler_config

vllm/v1/executor/ray_executor.py

+4 -3

@@ -382,9 +382,10 @@ class RayDistributedExecutor(Executor):

382 all_kwargs.append(kwargs)

383 self.collective_rpc("init_worker", args=(all_kwargs,))

384

385 - is_eep_new_worker = envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH

386 - if not is_eep_new_worker:

387 - self.collective_rpc("init_device")

385 + self.collective_rpc("init_device")

386 + if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:

387 + self.collective_rpc("elastic_ep_execute", args=("load_model",))

388 + else:

389 self.collective_rpc("load_model")

390

391 def _update_block_size(worker):

vllm/v1/executor/uniproc_executor.py

+6 -4

@@ -43,12 +43,14 @@ class UniProcExecutor(Executor):

43 max_workers=1, thread_name_prefix="WorkerAsyncOutput"

44 )

46 - is_eep_new_worker = envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH

46 self.driver_worker.init_worker(all_kwargs=[kwargs])

48 - if not is_eep_new_worker:

49 - self.driver_worker.init_device()

47 + self.driver_worker.init_device()

48 +

49 + if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:

50 + self.driver_worker.elastic_ep_execute("load_model")

51 + else:

52 self.driver_worker.load_model()

51 - current_platform.update_block_size_for_backend(self.vllm_config)

53 + current_platform.update_block_size_for_backend(self.vllm_config)

55 def _distributed_args(self) -> tuple[str, int, int]:

56 """Return (distributed_init_method, rank, local_rank)."""

vllm/v1/worker/gpu_worker.py

+2 -20

@@ -315,30 +315,12 @@ class Worker(WorkerBase):

315

316 # FIXME(youkaichao & ywang96): Use TorchDispatchMode instead of memory pool

317 # to hijack tensor allocation.

318 - def load_model(self) -> None:

319 - dummy_weights = os.environ.get("VLLM_ELASTIC_EP_SCALE_UP_LAUNCH") == "1"

320 - if dummy_weights:

321 - (

322 - expanded_physical_to_logical,

323 - num_logical_experts,

324 - old_num_physical_experts,

325 - ) = self.elastic_ep_executor.receive_expert_mapping()

326 - num_physical_experts = expanded_physical_to_logical.shape[1]

327 - self.parallel_config.eplb_config.num_redundant_experts = (

328 - num_physical_experts - num_logical_experts

329 - )

330 -

318 + def load_model(self, *, load_dummy_weights: bool = False) -> None:

319 with (

320 self._maybe_get_memory_pool_context(tag="weights"),

321 set_current_vllm_config(self.vllm_config),

322 ):

335 - self.model_runner.load_model(load_dummy_weights=dummy_weights)

336 -

337 - if dummy_weights:

338 - self.model_runner.setup_eplb_from_mapping(

339 - expanded_physical_to_logical, old_num_physical_experts

340 - )

341 - self.model_runner.eep_eplb_suppressed = True

323 + self.model_runner.load_model(load_dummy_weights=load_dummy_weights)

324

325 def update_config(self, overrides: dict[str, Any]) -> None:

326 self.model_runner.update_config(overrides)

vllm/v1/worker/worker_base.py

+1 -1

@@ -122,7 +122,7 @@ class WorkerBase:

122

123 return format_model_inspection(self.get_model())

124

125 - def load_model(self) -> None:

125 + def load_model(self, *, load_dummy_weights: bool = False) -> None:

126 """Load model onto target device."""

127 raise NotImplementedError

128