| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import gc |
| 4 | import weakref |
| 5 | from collections.abc import Iterable, Iterator, Sequence |
| 6 | from concurrent.futures import Future, ThreadPoolExecutor |
| 7 | from contextlib import contextmanager, nullcontext |
| 8 | from dataclasses import replace |
| 9 | from typing import TYPE_CHECKING |
| 10 | |
| 11 | import torch |
| 12 | import torch.nn as nn |
| 13 | from torch.distributed import P2POp |
| 14 | |
| 15 | from vllm.compilation.cuda_graph import CUDAGraphWrapper |
| 16 | from vllm.compilation.wrapper import reset_compile_wrapper |
| 17 | from vllm.config import ParallelConfig, set_current_vllm_config |
| 18 | from vllm.distributed import ( |
| 19 | get_dp_group, |
| 20 | get_ep_group, |
| 21 | get_pcp_group, |
| 22 | get_tp_group, |
| 23 | get_world_group, |
| 24 | ) |
| 25 | from vllm.distributed.elastic_ep.standby_state import ( |
| 26 | create_standby_groups, |
| 27 | get_standby_dp_group, |
| 28 | get_standby_ep_group, |
| 29 | get_standby_eplb_group, |
| 30 | get_standby_world_group, |
| 31 | pop_standby_groups, |
| 32 | ) |
| 33 | from vllm.distributed.eplb.eplb_communicator import ( |
| 34 | EplbCommunicator, |
| 35 | ) |
| 36 | from vllm.distributed.parallel_state import ( |
| 37 | GroupCoordinator, |
| 38 | _replace_active_groups, |
| 39 | ) |
| 40 | from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator |
| 41 | from vllm.logger import init_logger |
| 42 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 43 | get_ep_all2all_manager, |
| 44 | ) |
| 45 | from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig |
| 46 | from vllm.model_executor.layers.fused_moe.eep_reconfigure import ( |
| 47 | make_eep_staged_quant_method, |
| 48 | ) |
| 49 | from vllm.model_executor.warmup.flashinfer_autotune_cache import ( |
| 50 | sync_flashinfer_autotune_cache, |
| 51 | ) |
| 52 | from vllm.model_executor.warmup.kernel_warmup import kernel_warmup |
| 53 | from vllm.platforms import current_platform |
| 54 | from vllm.utils import is_moe_layer |
| 55 | from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType |
| 56 | from vllm.v1.worker.dp_utils import skip_dp_coordination |
| 57 | from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager |
| 58 | from vllm.v1.worker.gpu.eplb_utils import preserve_serving_state |
| 59 | from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper |
| 60 | from vllm.v1.worker.workspace import lock_workspace, unlock_workspace |
| 61 | |
| 62 | logger = init_logger(__name__) |
| 63 | |
| 64 | if TYPE_CHECKING: |
| 65 | from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( |
| 66 | FusedMoEMethodBase, |
| 67 | ) |
| 68 | |
| 69 | |
| 70 | def can_reuse_fused_moe_kernel(parallel_config: ParallelConfig) -> bool: |
| 71 | """Whether fused MoE kernels can be reused across EP rank changes. |
| 72 | |
| 73 | The EP backend must preserve its CUDA graph entries when ranks are connected, |
| 74 | disconnected, masked, or unmasked. |
| 75 | """ |
| 76 | return parallel_config.all2all_backend == "nixl_ep" |
| 77 | |
| 78 | |
| 79 | def batch_transfer_weights( |
| 80 | model: nn.Module, |
| 81 | is_sender: bool, |
| 82 | peer_rank: int, |
| 83 | dp_group: StatelessGroupCoordinator, |
| 84 | expert_weights: Sequence[Iterable[torch.Tensor]], |
| 85 | ) -> None: |
| 86 | device_comm = dp_group.device_communicator |
| 87 | if device_comm is None: |
| 88 | raise ValueError("No device communicator found") |
| 89 | |
| 90 | expert_weights_set = set() |
| 91 | for weight_group in expert_weights: |
| 92 | for weight in weight_group: |
| 93 | expert_weights_set.add(weight.data_ptr()) |
| 94 | |
| 95 | state_dict = model.state_dict() |
| 96 | all_params = [] |
| 97 | |
| 98 | for name, param in state_dict.items(): |
| 99 | if name.endswith("expert_map") or name.find("._shared_experts") != -1: |
| 100 | continue |
| 101 | if param.data_ptr() not in expert_weights_set: |
| 102 | all_params.append(param.data) |
| 103 | |
| 104 | assert len(all_params) > 0 |
| 105 | p2p_ops = [] |
| 106 | for param in all_params: |
| 107 | # PyNccl transfers flat memory and does not honor tensor strides. |
| 108 | transfer_param = param.contiguous() |
| 109 | op = object.__new__(P2POp) |
| 110 | op.op = torch.distributed.isend if is_sender else torch.distributed.irecv |
| 111 | op.tensor = transfer_param |
| 112 | op.group_peer = peer_rank |
| 113 | p2p_ops.append(op) |
| 114 | if transfer_param is not param: |
| 115 | device_comm.batch_isend_irecv(p2p_ops) |
| 116 | p2p_ops.clear() |
| 117 | if not is_sender: |
| 118 | param.copy_(transfer_param) |
| 119 | if p2p_ops: |
| 120 | device_comm.batch_isend_irecv(p2p_ops) |
| 121 | |
| 122 | |
| 123 | def broadcast_expert_mapping( |
| 124 | physical_to_logical: torch.Tensor | None, |
| 125 | dp_group: StatelessGroupCoordinator, |
| 126 | device: torch.device, |
| 127 | src_rank: int = 0, |
| 128 | ) -> torch.Tensor: |
| 129 | if dp_group.rank_in_group == src_rank: |
| 130 | assert physical_to_logical is not None |
| 131 | assert physical_to_logical.dtype == torch.int64 |
| 132 | shape_tensor = torch.tensor( |
| 133 | list(physical_to_logical.shape), dtype=torch.int64, device="cpu" |
| 134 | ) |
| 135 | else: |
| 136 | shape_tensor = torch.empty(2, dtype=torch.int64, device="cpu") |
| 137 | |
| 138 | shape_tensor = dp_group.tcp_store_group.broadcast(shape_tensor, src_rank) |
| 139 | |
| 140 | if dp_group.rank_in_group != src_rank: |
| 141 | assert device is not None |
| 142 | physical_to_logical = torch.empty( |
| 143 | tuple(shape_tensor.tolist()), |
| 144 | dtype=torch.int64, |
| 145 | device=device, |
| 146 | ) |
| 147 | |
| 148 | assert physical_to_logical is not None |
| 149 | return dp_group.broadcast(physical_to_logical, src_rank) |
| 150 | |
| 151 | |
| 152 | class ElasticEPScalingExecutor: |
| 153 | def __init__(self, worker): |
| 154 | self.worker_ref = weakref.ref(worker) |
| 155 | self.reconfig_request = None |
| 156 | self._staged_moe_quant_methods: dict[nn.Module, FusedMoEMethodBase] = {} |
| 157 | self._prepared_eplb_communicator: EplbCommunicator | None = None |
| 158 | self._async_executor = ThreadPoolExecutor( |
| 159 | max_workers=1, thread_name_prefix="ElasticEPAsync" |
| 160 | ) |
| 161 | self._async_future: Future[None] | None = None |
| 162 | self._group_cleanup_future: Future[None] | None = None |
| 163 | |
| 164 | @property |
| 165 | def worker(self): |
| 166 | worker = self.worker_ref() |
| 167 | if worker is None: |
| 168 | raise RuntimeError("Worker has been garbage collected") |
| 169 | return worker |
| 170 | |
| 171 | def _can_reuse_fused_moe_kernel(self) -> bool: |
| 172 | return can_reuse_fused_moe_kernel( |
| 173 | self.worker.vllm_config.parallel_config, |
| 174 | ) |
| 175 | |
| 176 | @contextmanager |
| 177 | def _suppress_eplb(self) -> Iterator[None]: |
| 178 | # A rank warming up alone must not advance the EPLB step counter. |
| 179 | runner = self.worker.model_runner |
| 180 | was_suppressed = runner.eep_eplb_suppressed |
| 181 | runner.eep_eplb_suppressed = True |
| 182 | try: |
| 183 | yield |
| 184 | finally: |
| 185 | runner.eep_eplb_suppressed = was_suppressed |
| 186 | |
| 187 | @contextmanager |
| 188 | def _disable_flashinfer_autotune(self) -> Iterator[None]: |
| 189 | kernel_config = self.worker.vllm_config.kernel_config |
| 190 | enabled = kernel_config.enable_flashinfer_autotune |
| 191 | kernel_config.enable_flashinfer_autotune = False |
| 192 | try: |
| 193 | yield |
| 194 | finally: |
| 195 | kernel_config.enable_flashinfer_autotune = enabled |
| 196 | |
| 197 | def execute(self, execute_method: str, *args, **kwargs): |
| 198 | method = getattr(self, execute_method, None) |
| 199 | if method is None: |
| 200 | raise ValueError(f"Unknown execute method: {execute_method}") |
| 201 | return method(*args, **kwargs) |
| 202 | |
| 203 | def start_async(self, execute_method: str, *args, **kwargs) -> str: |
| 204 | if self._async_future is not None: |
| 205 | raise RuntimeError("Another Elastic EP async method is active") |
| 206 | if args and isinstance(args[0], ReconfigureDistributedRequest): |
| 207 | self.reconfig_request = args[0] |
| 208 | dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank |
| 209 | done_key = f"eep_async/{execute_method}/{dp_rank}/{self.worker.rank}" |
| 210 | self._async_future = self._async_executor.submit( |
| 211 | self._run_async, execute_method, *args, **kwargs |
| 212 | ) |
| 213 | self._async_future.add_done_callback(lambda _: self._mark_async_done(done_key)) |
| 214 | return done_key |
| 215 | |
| 216 | def _run_async(self, execute_method: str, *args, **kwargs) -> None: |
| 217 | from vllm.platforms import current_platform |
| 218 | |
| 219 | self.worker.vllm_config.enable_trace_function_call_for_thread() |
| 220 | assert hasattr(self.worker, "device") |
| 221 | current_platform.set_device(self.worker.device) |
| 222 | with set_current_vllm_config(self.worker.vllm_config): |
| 223 | self.execute(execute_method, *args, **kwargs) |
| 224 | |
| 225 | def _mark_async_done(self, done_key: str) -> None: |
| 226 | from vllm.distributed.utils import get_cached_tcp_store_client |
| 227 | |
| 228 | assert self.reconfig_request is not None |
| 229 | get_cached_tcp_store_client( |
| 230 | self.reconfig_request.new_data_parallel_master_ip, |
| 231 | self.reconfig_request.coord_store_port, |
| 232 | ).set(done_key, b"1") |
| 233 | |
| 234 | def clear_async(self) -> None: |
| 235 | future = self._async_future |
| 236 | if future is None: |
| 237 | raise RuntimeError("No Elastic EP async method is active") |
| 238 | if not future.done(): |
| 239 | raise RuntimeError("Elastic EP async method is not done") |
| 240 | self._async_future = None |
| 241 | future.result() |
| 242 | |
| 243 | def _destroy_retired_groups( |
| 244 | self, groups: tuple[GroupCoordinator | None, ...] |
| 245 | ) -> None: |
| 246 | from vllm.platforms import current_platform |
| 247 | |
| 248 | current_platform.set_device(self.worker.device) |
| 249 | for group in groups: |
| 250 | if group is not None: |
| 251 | group.destroy() |
| 252 | |
| 253 | def _start_group_cleanup(self, groups: tuple[GroupCoordinator | None, ...]) -> None: |
| 254 | assert self._group_cleanup_future is None |
| 255 | self._group_cleanup_future = self._async_executor.submit( |
| 256 | self._destroy_retired_groups, groups |
| 257 | ) |
| 258 | |
| 259 | def _wait_for_group_cleanup(self) -> None: |
| 260 | if (future := self._group_cleanup_future) is not None: |
| 261 | self._group_cleanup_future = None |
| 262 | future.result() |
| 263 | |
| 264 | def shutdown(self) -> None: |
| 265 | try: |
| 266 | self._wait_for_group_cleanup() |
| 267 | finally: |
| 268 | self._async_executor.shutdown() |
| 269 | |
| 270 | def load_model(self) -> None: |
| 271 | self.worker.load_model(load_dummy_weights=True) |
| 272 | |
| 273 | def prepare_reconfiguration( |
| 274 | self, reconfig_request: ReconfigureDistributedRequest, use_all2all: bool |
| 275 | ) -> None: |
| 276 | self._wait_for_group_cleanup() |
| 277 | self.reconfig_request = reconfig_request |
| 278 | new_dp_size = reconfig_request.new_data_parallel_size |
| 279 | old_dp_size = get_dp_group().world_size |
| 280 | parallel_config = self.worker.vllm_config.parallel_config |
| 281 | world_size = parallel_config.world_size |
| 282 | new_world_size_across_dp = world_size * new_dp_size |
| 283 | create_standby_groups( |
| 284 | new_dp_size=new_dp_size, |
| 285 | new_world_size_across_dp=new_world_size_across_dp, |
| 286 | master_ip=reconfig_request.new_data_parallel_master_ip, |
| 287 | coord_store_port=reconfig_request.coord_store_port, |
| 288 | use_all2all=use_all2all, |
| 289 | enable_eplb=parallel_config.enable_eplb, |
| 290 | ) |
| 291 | standby_ep_group = get_standby_ep_group() |
| 292 | assert standby_ep_group is not None |
| 293 | all2all_manager = get_ep_all2all_manager(standby_ep_group) |
| 294 | all2all_manager.stage_ep_size() |
| 295 | if not self._can_reuse_fused_moe_kernel(): |
| 296 | self.stage_standby_moe_quant_methods(all2all_manager) |
| 297 | self._prepare_eplb_communicator(get_standby_eplb_group()) |
| 298 | if new_dp_size > old_dp_size: |
| 299 | self.transfer_weights(old_dp_size, new_dp_size) |
| 300 | # On ROCm this deadlocks while serving; warm_and_capture warms these at commit. |
| 301 | if not current_platform.is_rocm(): |
| 302 | self._warm_target_groups(get_standby_dp_group(), standby_ep_group) |
| 303 | if new_dp_size > old_dp_size and self._can_reuse_fused_moe_kernel(): |
| 304 | target_world_group = get_standby_world_group() |
| 305 | assert target_world_group is not None |
| 306 | sync_flashinfer_autotune_cache(self.worker.model_runner, target_world_group) |
| 307 | |
| 308 | def _prepare_eplb_communicator(self, eplb_group) -> None: |
| 309 | assert eplb_group is not None |
| 310 | model_runner = self.worker.model_runner |
| 311 | eplb_state = model_runner.eplb_state |
| 312 | assert eplb_state is not None |
| 313 | self._prepared_eplb_communicator = eplb_state.create_communicator( |
| 314 | model_runner.model_config, eplb_group |
| 315 | ) |
| 316 | |
| 317 | def transfer_weights(self, old_dp_size: int, new_dp_size: int) -> None: |
| 318 | standby_dp_group = get_standby_dp_group() |
| 319 | assert standby_dp_group is not None |
| 320 | # Broadcast old_dp_size to all workers in standby group |
| 321 | if standby_dp_group.rank_in_group < old_dp_size: |
| 322 | old_dp_size_tensor = torch.tensor( |
| 323 | [old_dp_size], dtype=torch.int64, device="cpu" |
| 324 | ) |
| 325 | else: |
| 326 | old_dp_size_tensor = torch.empty(1, dtype=torch.int64, device="cpu") |
| 327 | old_dp_size_tensor = standby_dp_group.tcp_store_group.broadcast( |
| 328 | old_dp_size_tensor, 0 |
| 329 | ) |
| 330 | |
| 331 | num_new_workers = new_dp_size - old_dp_size |
| 332 | dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank |
| 333 | |
| 334 | # Sender-receiver pairing: the first new_workers % old_dp_size |
| 335 | # senders get (k+1) contiguous receivers, the rest get k |
| 336 | # receivers. |
| 337 | num_dst_per_sender = num_new_workers // old_dp_size |
| 338 | remainder = num_new_workers % old_dp_size |
| 339 | |
| 340 | if dp_rank < remainder: |
| 341 | recv_begin = dp_rank * (num_dst_per_sender + 1) |
| 342 | recv_end = recv_begin + num_dst_per_sender + 1 |
| 343 | else: |
| 344 | recv_begin = ( |
| 345 | remainder * (num_dst_per_sender + 1) |
| 346 | + (dp_rank - remainder) * num_dst_per_sender |
| 347 | ) |
| 348 | recv_end = recv_begin + num_dst_per_sender |
| 349 | |
| 350 | ranks_to_send = list(range(old_dp_size + recv_begin, old_dp_size + recv_end)) |
| 351 | |
| 352 | model = self.worker.model_runner.get_model() |
| 353 | for new_worker_rank in sorted(ranks_to_send): |
| 354 | batch_transfer_weights( |
| 355 | model=model, |
| 356 | is_sender=True, |
| 357 | peer_rank=new_worker_rank, |
| 358 | dp_group=standby_dp_group, |
| 359 | expert_weights=model.expert_weights, |
| 360 | ) |
| 361 | torch.accelerator.synchronize() |
| 362 | |
| 363 | def _warm_target_groups(self, dp_group, ep_group) -> None: |
| 364 | assert dp_group is not None and ep_group is not None |
| 365 | stream = torch.Stream(device=dp_group.device) |
| 366 | with stream: |
| 367 | tensor = torch.zeros(1, dtype=torch.int32, device=dp_group.device) |
| 368 | for group in (dp_group, ep_group): |
| 369 | torch.distributed.all_reduce(tensor, group=group.device_group) |
| 370 | stream.synchronize() |
| 371 | |
| 372 | def broadcast_expert_mapping(self) -> None: |
| 373 | standby_dp_group = get_standby_dp_group() |
| 374 | assert standby_dp_group is not None |
| 375 | model_config = self.worker.model_runner.model_config |
| 376 | eplb_state = self.worker.model_runner.eplb_state |
| 377 | assert eplb_state is not None |
| 378 | eplb_state.drain_async() |
| 379 | eplb_model_state = eplb_state.model_states[model_config.compute_hash()] |
| 380 | physical_to_logical = eplb_model_state.physical_to_logical_map_buffer |
| 381 | broadcast_expert_mapping( |
| 382 | physical_to_logical=physical_to_logical, |
| 383 | dp_group=standby_dp_group, |
| 384 | src_rank=0, |
| 385 | device=self.worker.device, |
| 386 | ) |
| 387 | |
| 388 | def _make_eep_moe_config(self, module, dp_group, ep_group): |
| 389 | parallel_config = self.worker.vllm_config.parallel_config |
| 390 | tp_size = get_tp_group().world_size |
| 391 | sp_size = tp_size if parallel_config.use_sequence_parallel_moe else 1 |
| 392 | moe_parallel_config = FusedMoEParallelConfig.make( |
| 393 | tp_size_=tp_size, |
| 394 | pcp_size_=get_pcp_group().world_size, |
| 395 | dp_size_=dp_group.world_size, |
| 396 | sp_size_=sp_size, |
| 397 | vllm_parallel_config=parallel_config, |
| 398 | ) |
| 399 | return replace( |
| 400 | module.moe_config, |
| 401 | num_experts=module.moe_config.num_local_experts * ep_group.world_size, |
| 402 | moe_parallel_config=moe_parallel_config, |
| 403 | ) |
| 404 | |
| 405 | def stage_standby_moe_quant_methods(self, all2all_manager) -> None: |
| 406 | standby_dp_group = get_standby_dp_group() |
| 407 | standby_ep_group = get_standby_ep_group() |
| 408 | model = self.worker.model_runner.get_model() |
| 409 | moe_modules = [module for module in model.modules() if is_moe_layer(module)] |
| 410 | self._staged_moe_quant_methods.clear() |
| 411 | with set_current_vllm_config(self.worker.vllm_config): |
| 412 | for module in moe_modules: |
| 413 | staged_quant_method = make_eep_staged_quant_method( |
| 414 | module, |
| 415 | self._make_eep_moe_config( |
| 416 | module, |
| 417 | standby_dp_group, |
| 418 | standby_ep_group, |
| 419 | ), |
| 420 | all2all_manager=all2all_manager, |
| 421 | ) |
| 422 | if staged_quant_method is not None: |
| 423 | self._staged_moe_quant_methods[module] = staged_quant_method |
| 424 | |
| 425 | def _commit_staged_moe_quant_methods(self) -> None: |
| 426 | model = self.worker.model_runner.get_model() |
| 427 | moe_modules = [module for module in model.modules() if is_moe_layer(module)] |
| 428 | for module in moe_modules: |
| 429 | staged_quant_method = self._staged_moe_quant_methods.pop(module, None) |
| 430 | if staged_quant_method is None: |
| 431 | continue |
| 432 | assert staged_quant_method.moe_kernel is not None |
| 433 | module._replace_quant_method(staged_quant_method) |
| 434 | self._staged_moe_quant_methods.clear() |
| 435 | |
| 436 | def _release_cuda_graphs(self) -> None: |
| 437 | manager = getattr(self.worker.model_runner, "cudagraph_manager", None) |
| 438 | speculator = getattr(self.worker.model_runner, "speculator", None) |
| 439 | if manager is not None: |
| 440 | # MRV2 captures through CudaGraphManager instead of wrapping the |
| 441 | # model, so neither wrapper branch below ever fires. |
| 442 | manager.release_graphs() |
| 443 | for attr in vars(speculator).values() if speculator else (): |
| 444 | if isinstance(attr, CudaGraphManager): |
| 445 | attr.release_graphs() |
| 446 | |
| 447 | elif isinstance(self.worker.model_runner.model, CUDAGraphWrapper): |
| 448 | wrapper = self.worker.model_runner.model |
| 449 | wrapper.concrete_cudagraph_entries = {} |
| 450 | |
| 451 | elif isinstance(self.worker.model_runner.model, UBatchWrapper): |
| 452 | self.worker.model_runner.model.clear_graphs() |
| 453 | |
| 454 | torch.compiler.reset() |
| 455 | with set_current_vllm_config(self.worker.vllm_config): |
| 456 | reset_compile_wrapper(self.worker.model_runner.get_model()) |
| 457 | if speculator is not None: |
| 458 | reset_compile_wrapper(speculator.model) |
| 459 | |
| 460 | gc.collect() |
| 461 | torch.accelerator.synchronize() |
| 462 | torch.accelerator.empty_cache() |
| 463 | |
| 464 | def switch_and_remove(self) -> None: |
| 465 | # Removing ranks skipped preparation, so wait for prior cleanup here. |
| 466 | self._wait_for_group_cleanup() |
| 467 | self._release_cuda_graphs() |
| 468 | retired_groups = _replace_active_groups( |
| 469 | world=None, dp=None, ep=None, eplb=None, node_count=None |
| 470 | ) |
| 471 | self._start_group_cleanup(retired_groups) |
| 472 | # Finish collective cleanup before this worker is shut down. |
| 473 | self._wait_for_group_cleanup() |
| 474 | |
| 475 | def switch_and_prepare(self) -> tuple[GroupCoordinator | None, ...]: |
| 476 | if not self._can_reuse_fused_moe_kernel(): |
| 477 | self._release_cuda_graphs() |
| 478 | retired_groups = _replace_active_groups(**pop_standby_groups()) |
| 479 | |
| 480 | parallel_config = self.worker.vllm_config.parallel_config |
| 481 | reconfig_request = self.reconfig_request |
| 482 | assert reconfig_request is not None |
| 483 | new_dp_size = reconfig_request.new_data_parallel_size |
| 484 | new_ep_size = get_ep_group().world_size |
| 485 | |
| 486 | parallel_config.data_parallel_size = new_dp_size |
| 487 | if ( |
| 488 | reconfig_request.new_data_parallel_rank |
| 489 | != ReconfigureRankType.KEEP_CURRENT_RANK |
| 490 | ): |
| 491 | parallel_config.data_parallel_rank = reconfig_request.new_data_parallel_rank |
| 492 | if ( |
| 493 | reconfig_request.new_data_parallel_rank_local |
| 494 | != ReconfigureRankType.KEEP_CURRENT_RANK |
| 495 | ): |
| 496 | parallel_config.data_parallel_rank_local = ( |
| 497 | reconfig_request.new_data_parallel_rank_local |
| 498 | ) |
| 499 | parallel_config.data_parallel_master_ip = ( |
| 500 | reconfig_request.new_data_parallel_master_ip |
| 501 | ) |
| 502 | parallel_config.data_parallel_master_port = ( |
| 503 | reconfig_request.new_data_parallel_master_port |
| 504 | ) |
| 505 | |
| 506 | # Reconfigure MoE modules with new EP size |
| 507 | moe_modules = [ |
| 508 | module |
| 509 | for module in self.worker.model_runner.model.modules() |
| 510 | if is_moe_layer(module) |
| 511 | ] |
| 512 | num_local_experts = moe_modules[0].moe_config.num_local_experts |
| 513 | assert all( |
| 514 | module.moe_config.num_local_experts == num_local_experts |
| 515 | for module in moe_modules |
| 516 | ), "All MoE modules must have the same number of experts" |
| 517 | dp_group = get_dp_group() |
| 518 | ep_group = get_ep_group() |
| 519 | for module in moe_modules: |
| 520 | new_moe_config = self._make_eep_moe_config(module, dp_group, ep_group) |
| 521 | module._set_moe_config(new_moe_config) |
| 522 | |
| 523 | # Update EPLB state |
| 524 | eplb_state = self.worker.model_runner.eplb_state |
| 525 | assert eplb_state is not None |
| 526 | model_config = self.worker.model_runner.model_config |
| 527 | eplb_model_state = eplb_state.model_states[model_config.compute_hash()] |
| 528 | |
| 529 | num_physical_experts = num_local_experts * new_ep_size |
| 530 | num_logical_experts = eplb_model_state.logical_replica_count.shape[1] |
| 531 | parallel_config.eplb_config.num_redundant_experts = ( |
| 532 | num_physical_experts - num_logical_experts |
| 533 | ) |
| 534 | eplb_state.reconfigure_physical_expert_slots(model_config, num_physical_experts) |
| 535 | |
| 536 | model = self.worker.model_runner.get_model() |
| 537 | model.expert_weights = [] |
| 538 | with set_current_vllm_config(self.worker.vllm_config): |
| 539 | model.set_eplb_state( |
| 540 | eplb_model_state.expert_load_pass_buffer, |
| 541 | eplb_model_state.logical_to_physical_map, |
| 542 | eplb_model_state.logical_replica_count, |
| 543 | ) |
| 544 | eplb_state._propagate_shared_tensors( |
| 545 | model, eplb_model_state.num_unpadded_tokens_tensors |
| 546 | ) |
| 547 | model.update_physical_experts_metadata( |
| 548 | num_physical_experts=num_physical_experts, |
| 549 | num_local_physical_experts=num_local_experts, |
| 550 | ) |
| 551 | self._commit_staged_moe_quant_methods() |
| 552 | get_ep_all2all_manager().commit_ep_size() |
| 553 | # Legacy modular methods need to be recreated for the new EP size. |
| 554 | for module in moe_modules: |
| 555 | if getattr(module._quant_method, "wraps_legacy_quant_method", False): |
| 556 | module._replace_quant_method(module._quant_method.old_quant_method) |
| 557 | |
| 558 | assert self._prepared_eplb_communicator is not None |
| 559 | eplb_state.update_communicator(model_config, self._prepared_eplb_communicator) |
| 560 | self._prepared_eplb_communicator = None |
| 561 | |
| 562 | return retired_groups |
| 563 | |
| 564 | def _perform_eplb_reshuffle( |
| 565 | self, |
| 566 | rank_mapping: dict[int, int] | None = None, |
| 567 | async_op: bool = False, |
| 568 | ) -> None: |
| 569 | if get_ep_group().rank == 0: |
| 570 | logger.info("[Elastic EP] Starting expert resharding...") |
| 571 | |
| 572 | eplb_state = self.worker.model_runner.eplb_state |
| 573 | assert eplb_state is not None |
| 574 | |
| 575 | is_async_enabled = eplb_state.is_async |
| 576 | run_async = async_op and is_async_enabled |
| 577 | eplb_state.is_async = run_async |
| 578 | eplb_state.rearrange(rank_mapping=rank_mapping, use_last_expert_load=True) |
| 579 | if not run_async: |
| 580 | # Wait for non-blocking expert resharding copies before continuing |
| 581 | # the Elastic EP reconfiguration. |
| 582 | torch.accelerator.synchronize() |
| 583 | # reset expert_rearrangement_step to ensure all ranks are synchronized |
| 584 | eplb_state.expert_rearrangement_step = 0 |
| 585 | eplb_state.is_async = is_async_enabled |
| 586 | # Start the async worker thread if it doesn't exist yet (idempotent). |
| 587 | # This is needed for new workers after scale-up: they create EplbState |
| 588 | # in setup_eplb_from_mapping() but don't start the thread there because |
| 589 | # groups aren't ready yet. |
| 590 | eplb_state.start_async_loop() |
| 591 | if get_ep_group().rank == 0: |
| 592 | logger.info( |
| 593 | "[Elastic EP] Expert resharding %s", |
| 594 | "scheduled" if run_async else "completed", |
| 595 | ) |
| 596 | |
| 597 | def commit_scale_up(self, is_existing_worker: bool) -> None: |
| 598 | if is_existing_worker: |
| 599 | self.broadcast_expert_mapping() |
| 600 | retired_groups = self.switch_and_prepare() |
| 601 | else: |
| 602 | mapping = self.receive_expert_mapping() |
| 603 | self.worker.model_runner.setup_eplb_from_mapping(mapping) |
| 604 | if not self._can_reuse_fused_moe_kernel(): |
| 605 | # warm_and_capture's dummy run warms the deferred dp/ep groups. |
| 606 | self.warm_and_capture() |
| 607 | elif current_platform.is_rocm(): |
| 608 | # Reuse path skips warm_and_capture; warm the deferred groups here. |
| 609 | self._warm_target_groups(get_dp_group(), get_ep_group()) |
| 610 | self._perform_eplb_reshuffle(async_op=True) |
| 611 | if is_existing_worker: |
| 612 | self._start_group_cleanup(retired_groups) |
| 613 | |
| 614 | def commit_scale_down(self, new_dp_size: int, removing: bool) -> None: |
| 615 | self.perform_scale_down_eplb_reshuffle(new_dp_size) |
| 616 | if removing: |
| 617 | self.switch_and_remove() |
| 618 | else: |
| 619 | retired_groups = self.switch_and_prepare() |
| 620 | if not self._can_reuse_fused_moe_kernel(): |
| 621 | self.warm_and_capture() |
| 622 | elif current_platform.is_rocm(): |
| 623 | # Reuse path skips warm_and_capture; warm the deferred groups here. |
| 624 | self._warm_target_groups(get_dp_group(), get_ep_group()) |
| 625 | self._start_group_cleanup(retired_groups) |
| 626 | |
| 627 | def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None: |
| 628 | eplb_state = self.worker.model_runner.eplb_state |
| 629 | if eplb_state is not None: |
| 630 | eplb_state.drain_async() |
| 631 | parallel_config = self.worker.vllm_config.parallel_config |
| 632 | tp_size = parallel_config.tensor_parallel_size |
| 633 | old_ep_size = parallel_config.data_parallel_size * tp_size |
| 634 | new_ep_size = new_dp_size * tp_size |
| 635 | rank_mapping = { |
| 636 | old_ep_rank: old_ep_rank if old_ep_rank < new_ep_size else -1 |
| 637 | for old_ep_rank in range(old_ep_size) |
| 638 | } |
| 639 | self._perform_eplb_reshuffle(rank_mapping=rank_mapping) |
| 640 | |
| 641 | def prepare_new_worker(self) -> None: |
| 642 | dp_group = get_dp_group() |
| 643 | assert isinstance(dp_group, StatelessGroupCoordinator) |
| 644 | new_dp_size = dp_group.world_size |
| 645 | dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank |
| 646 | |
| 647 | # Receive old_dp_size broadcasted during transfer_weights |
| 648 | old_dp_size_tensor = torch.empty(1, dtype=torch.int64, device="cpu") |
| 649 | old_dp_size_tensor = dp_group.tcp_store_group.broadcast(old_dp_size_tensor, 0) |
| 650 | old_dp_size = int(old_dp_size_tensor[0].item()) |
| 651 | |
| 652 | # Calculate which existing worker will send to this new worker |
| 653 | num_new_workers = new_dp_size - old_dp_size |
| 654 | new_worker_idx = dp_rank - old_dp_size |
| 655 | num_dst_per_sender = num_new_workers // old_dp_size |
| 656 | remainder = num_new_workers % old_dp_size |
| 657 | |
| 658 | if new_worker_idx < remainder * (num_dst_per_sender + 1): |
| 659 | sender_rank = new_worker_idx // (num_dst_per_sender + 1) |
| 660 | else: |
| 661 | sender_rank = ( |
| 662 | remainder |
| 663 | + (new_worker_idx - remainder * (num_dst_per_sender + 1)) |
| 664 | // num_dst_per_sender |
| 665 | ) |
| 666 | |
| 667 | model = self.worker.model_runner.get_model() |
| 668 | expert_weights = [ |
| 669 | module.get_expert_weights() |
| 670 | for module in model.modules() |
| 671 | if is_moe_layer(module) |
| 672 | ] |
| 673 | batch_transfer_weights( |
| 674 | model=model, |
| 675 | is_sender=False, |
| 676 | peer_rank=sender_rank, |
| 677 | dp_group=dp_group, |
| 678 | expert_weights=expert_weights, |
| 679 | ) |
| 680 | torch.accelerator.synchronize() |
| 681 | if not current_platform.is_rocm(): |
| 682 | self._warm_target_groups(get_dp_group(), get_ep_group()) |
| 683 | if self._can_reuse_fused_moe_kernel(): |
| 684 | sync_flashinfer_autotune_cache(self.worker.model_runner, get_world_group()) |
| 685 | |
| 686 | def receive_expert_mapping(self) -> torch.Tensor: |
| 687 | dp_group = get_dp_group() |
| 688 | assert isinstance(dp_group, StatelessGroupCoordinator) |
| 689 | return broadcast_expert_mapping( |
| 690 | physical_to_logical=None, |
| 691 | dp_group=dp_group, |
| 692 | src_rank=0, |
| 693 | device=self.worker.device, |
| 694 | ) |
| 695 | |
| 696 | def warmup_new_worker(self) -> None: |
| 697 | with set_current_vllm_config(self.worker.vllm_config): |
| 698 | kernel_warmup(self.worker, process_local_only=True) |
| 699 | if self._can_reuse_fused_moe_kernel(): |
| 700 | self.warm_and_capture() |
| 701 | |
| 702 | def warm_and_capture(self) -> None: |
| 703 | # _ensure_workspace_size allocates a fresh tensor on grow, leaving |
| 704 | # any captured CUDA graph with a stale data pointer; drop graphs |
| 705 | # before re-warm so captures realign with the resized buffer. |
| 706 | self._release_cuda_graphs() |
| 707 | unlock_workspace() |
| 708 | |
| 709 | runner = self.worker.model_runner |
| 710 | all2all_manager = get_ep_all2all_manager() |
| 711 | reuse_kernel = self._can_reuse_fused_moe_kernel() |
| 712 | serving_state = ( |
| 713 | preserve_serving_state(runner) |
| 714 | if self.worker.use_v2_model_runner |
| 715 | else runner.preserve_serving_state() |
| 716 | ) |
| 717 | with ( |
| 718 | skip_dp_coordination() if reuse_kernel else nullcontext(), |
| 719 | all2all_manager.mask_remote_ranks() if reuse_kernel else nullcontext(), |
| 720 | self._disable_flashinfer_autotune() if reuse_kernel else nullcontext(), |
| 721 | self._suppress_eplb(), |
| 722 | serving_state, |
| 723 | ): |
| 724 | runner._dummy_run( |
| 725 | runner.max_num_tokens, |
| 726 | is_profile=True, |
| 727 | skip_eplb=True, |
| 728 | randomize_inputs=self.worker.randomize_dummy_inputs, |
| 729 | ) |
| 730 | self.worker.compile_or_warm_up_model() |
| 731 | |
| 732 | lock_workspace() |