# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import gc import weakref from collections.abc import Iterable, Iterator, Sequence from concurrent.futures import Future, ThreadPoolExecutor from contextlib import contextmanager, nullcontext from dataclasses import replace from typing import TYPE_CHECKING import torch import torch.nn as nn from torch.distributed import P2POp from vllm.compilation.cuda_graph import CUDAGraphWrapper from vllm.compilation.wrapper import reset_compile_wrapper from vllm.config import ParallelConfig, set_current_vllm_config from vllm.distributed import ( get_dp_group, get_ep_group, get_pcp_group, get_tp_group, get_world_group, ) from vllm.distributed.elastic_ep.standby_state import ( create_standby_groups, get_standby_dp_group, get_standby_ep_group, get_standby_eplb_group, get_standby_world_group, pop_standby_groups, ) from vllm.distributed.eplb.eplb_communicator import ( EplbCommunicator, ) from vllm.distributed.parallel_state import ( GroupCoordinator, _replace_active_groups, ) from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.all2all_utils import ( get_ep_all2all_manager, ) from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig from vllm.model_executor.layers.fused_moe.eep_reconfigure import ( make_eep_staged_quant_method, ) from vllm.model_executor.warmup.flashinfer_autotune_cache import ( sync_flashinfer_autotune_cache, ) from vllm.model_executor.warmup.kernel_warmup import kernel_warmup from vllm.platforms import current_platform from vllm.utils import is_moe_layer from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType from vllm.v1.worker.dp_utils import skip_dp_coordination from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper from vllm.v1.worker.workspace import lock_workspace, unlock_workspace logger = init_logger(__name__) if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( FusedMoEMethodBase, ) def can_reuse_fused_moe_kernel(parallel_config: ParallelConfig) -> bool: """Whether fused MoE kernels can be reused across EP rank changes. The EP backend must preserve its CUDA graph entries when ranks are connected, disconnected, masked, or unmasked. """ return parallel_config.all2all_backend == "nixl_ep" def batch_transfer_weights( model: nn.Module, is_sender: bool, peer_rank: int, dp_group: StatelessGroupCoordinator, expert_weights: Sequence[Iterable[torch.Tensor]], ) -> None: device_comm = dp_group.device_communicator if device_comm is None: raise ValueError("No device communicator found") expert_weights_set = set() for weight_group in expert_weights: for weight in weight_group: expert_weights_set.add(weight.data_ptr()) state_dict = model.state_dict() all_params = [] for name, param in state_dict.items(): if name.endswith("expert_map") or name.find("._shared_experts") != -1: continue if param.data_ptr() not in expert_weights_set: all_params.append(param.data) assert len(all_params) > 0 p2p_ops = [] for param in all_params: # PyNccl transfers flat memory and does not honor tensor strides. transfer_param = param.contiguous() op = object.__new__(P2POp) op.op = torch.distributed.isend if is_sender else torch.distributed.irecv op.tensor = transfer_param op.group_peer = peer_rank p2p_ops.append(op) if transfer_param is not param: device_comm.batch_isend_irecv(p2p_ops) p2p_ops.clear() if not is_sender: param.copy_(transfer_param) if p2p_ops: device_comm.batch_isend_irecv(p2p_ops) def broadcast_expert_mapping( physical_to_logical: torch.Tensor | None, dp_group: StatelessGroupCoordinator, device: torch.device, src_rank: int = 0, ) -> torch.Tensor: if dp_group.rank_in_group == src_rank: assert physical_to_logical is not None assert physical_to_logical.dtype == torch.int64 shape_tensor = torch.tensor( list(physical_to_logical.shape), dtype=torch.int64, device="cpu" ) else: shape_tensor = torch.empty(2, dtype=torch.int64, device="cpu") shape_tensor = dp_group.tcp_store_group.broadcast(shape_tensor, src_rank) if dp_group.rank_in_group != src_rank: assert device is not None physical_to_logical = torch.empty( tuple(shape_tensor.tolist()), dtype=torch.int64, device=device, ) assert physical_to_logical is not None return dp_group.broadcast(physical_to_logical, src_rank) class ElasticEPScalingExecutor: def __init__(self, worker): self.worker_ref = weakref.ref(worker) self.reconfig_request = None self._staged_moe_quant_methods: dict[nn.Module, FusedMoEMethodBase] = {} self._prepared_eplb_communicator: EplbCommunicator | None = None self._async_executor = ThreadPoolExecutor( max_workers=1, thread_name_prefix="ElasticEPAsync" ) self._async_future: Future[None] | None = None self._group_cleanup_future: Future[None] | None = None @property def worker(self): worker = self.worker_ref() if worker is None: raise RuntimeError("Worker has been garbage collected") return worker def _can_reuse_fused_moe_kernel(self) -> bool: return can_reuse_fused_moe_kernel( self.worker.vllm_config.parallel_config, ) @contextmanager def _disable_flashinfer_autotune(self) -> Iterator[None]: kernel_config = self.worker.vllm_config.kernel_config enabled = kernel_config.enable_flashinfer_autotune kernel_config.enable_flashinfer_autotune = False try: yield finally: kernel_config.enable_flashinfer_autotune = enabled def execute(self, execute_method: str, *args, **kwargs): method = getattr(self, execute_method, None) if method is None: raise ValueError(f"Unknown execute method: {execute_method}") return method(*args, **kwargs) def start_async(self, execute_method: str, *args, **kwargs) -> str: if self._async_future is not None: raise RuntimeError("Another Elastic EP async method is active") if args and isinstance(args[0], ReconfigureDistributedRequest): self.reconfig_request = args[0] dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank done_key = f"eep_async/{execute_method}/{dp_rank}/{self.worker.rank}" self._async_future = self._async_executor.submit( self._run_async, execute_method, *args, **kwargs ) self._async_future.add_done_callback(lambda _: self._mark_async_done(done_key)) return done_key def _run_async(self, execute_method: str, *args, **kwargs) -> None: from vllm.platforms import current_platform self.worker.vllm_config.enable_trace_function_call_for_thread() assert hasattr(self.worker, "device") current_platform.set_device(self.worker.device) with set_current_vllm_config(self.worker.vllm_config): self.execute(execute_method, *args, **kwargs) def _mark_async_done(self, done_key: str) -> None: from vllm.distributed.utils import get_cached_tcp_store_client assert self.reconfig_request is not None get_cached_tcp_store_client( self.reconfig_request.new_data_parallel_master_ip, self.reconfig_request.coord_store_port, ).set(done_key, b"1") def clear_async(self) -> None: future = self._async_future if future is None: raise RuntimeError("No Elastic EP async method is active") if not future.done(): raise RuntimeError("Elastic EP async method is not done") self._async_future = None future.result() def _destroy_retired_groups( self, groups: tuple[GroupCoordinator | None, ...] ) -> None: from vllm.platforms import current_platform current_platform.set_device(self.worker.device) for group in groups: if group is not None: group.destroy() def _start_group_cleanup(self, groups: tuple[GroupCoordinator | None, ...]) -> None: assert self._group_cleanup_future is None self._group_cleanup_future = self._async_executor.submit( self._destroy_retired_groups, groups ) def _wait_for_group_cleanup(self) -> None: if (future := self._group_cleanup_future) is not None: self._group_cleanup_future = None future.result() def shutdown(self) -> None: try: self._wait_for_group_cleanup() finally: self._async_executor.shutdown() def load_model(self) -> None: self.worker.load_model(load_dummy_weights=True) def prepare_reconfiguration( self, reconfig_request: ReconfigureDistributedRequest, use_all2all: bool ) -> None: self._wait_for_group_cleanup() self.reconfig_request = reconfig_request new_dp_size = reconfig_request.new_data_parallel_size old_dp_size = get_dp_group().world_size parallel_config = self.worker.vllm_config.parallel_config world_size = parallel_config.world_size new_world_size_across_dp = world_size * new_dp_size create_standby_groups( new_dp_size=new_dp_size, new_world_size_across_dp=new_world_size_across_dp, master_ip=reconfig_request.new_data_parallel_master_ip, coord_store_port=reconfig_request.coord_store_port, use_all2all=use_all2all, enable_eplb=parallel_config.enable_eplb, ) standby_ep_group = get_standby_ep_group() assert standby_ep_group is not None all2all_manager = get_ep_all2all_manager(standby_ep_group) all2all_manager.stage_ep_size() if not self._can_reuse_fused_moe_kernel(): self.stage_standby_moe_quant_methods(all2all_manager) self._prepare_eplb_communicator(get_standby_eplb_group()) if new_dp_size > old_dp_size: self.transfer_weights(old_dp_size, new_dp_size) # On ROCm this deadlocks while serving; warm_and_capture warms these at commit. if not current_platform.is_rocm(): self._warm_target_groups(get_standby_dp_group(), standby_ep_group) if new_dp_size > old_dp_size and self._can_reuse_fused_moe_kernel(): target_world_group = get_standby_world_group() assert target_world_group is not None sync_flashinfer_autotune_cache(self.worker.model_runner, target_world_group) def _prepare_eplb_communicator(self, eplb_group) -> None: assert eplb_group is not None model_runner = self.worker.model_runner eplb_state = model_runner.eplb_state assert eplb_state is not None self._prepared_eplb_communicator = eplb_state.create_communicator( model_runner.model_config, eplb_group ) def transfer_weights(self, old_dp_size: int, new_dp_size: int) -> None: standby_dp_group = get_standby_dp_group() assert standby_dp_group is not None # Broadcast old_dp_size to all workers in standby group if standby_dp_group.rank_in_group < old_dp_size: old_dp_size_tensor = torch.tensor( [old_dp_size], dtype=torch.int64, device="cpu" ) else: old_dp_size_tensor = torch.empty(1, dtype=torch.int64, device="cpu") old_dp_size_tensor = standby_dp_group.tcp_store_group.broadcast( old_dp_size_tensor, 0 ) num_new_workers = new_dp_size - old_dp_size dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank # Sender-receiver pairing: the first new_workers % old_dp_size # senders get (k+1) contiguous receivers, the rest get k # receivers. num_dst_per_sender = num_new_workers // old_dp_size remainder = num_new_workers % old_dp_size if dp_rank < remainder: recv_begin = dp_rank * (num_dst_per_sender + 1) recv_end = recv_begin + num_dst_per_sender + 1 else: recv_begin = ( remainder * (num_dst_per_sender + 1) + (dp_rank - remainder) * num_dst_per_sender ) recv_end = recv_begin + num_dst_per_sender ranks_to_send = list(range(old_dp_size + recv_begin, old_dp_size + recv_end)) model = self.worker.model_runner.get_model() for new_worker_rank in sorted(ranks_to_send): batch_transfer_weights( model=model, is_sender=True, peer_rank=new_worker_rank, dp_group=standby_dp_group, expert_weights=model.expert_weights, ) torch.accelerator.synchronize() def _warm_target_groups(self, dp_group, ep_group) -> None: assert dp_group is not None and ep_group is not None stream = torch.Stream(device=dp_group.device) with stream: tensor = torch.zeros(1, dtype=torch.int32, device=dp_group.device) for group in (dp_group, ep_group): torch.distributed.all_reduce(tensor, group=group.device_group) stream.synchronize() def broadcast_expert_mapping(self) -> None: standby_dp_group = get_standby_dp_group() assert standby_dp_group is not None model_config = self.worker.model_runner.model_config eplb_state = self.worker.model_runner.eplb_state assert eplb_state is not None eplb_state.drain_async() eplb_model_state = eplb_state.model_states[model_config.compute_hash()] physical_to_logical = eplb_model_state.physical_to_logical_map_buffer broadcast_expert_mapping( physical_to_logical=physical_to_logical, dp_group=standby_dp_group, src_rank=0, device=self.worker.device, ) def _make_eep_moe_config(self, module, dp_group, ep_group): parallel_config = self.worker.vllm_config.parallel_config tp_size = get_tp_group().world_size sp_size = tp_size if parallel_config.use_sequence_parallel_moe else 1 moe_parallel_config = FusedMoEParallelConfig.make( tp_size_=tp_size, pcp_size_=get_pcp_group().world_size, dp_size_=dp_group.world_size, sp_size_=sp_size, vllm_parallel_config=parallel_config, ) return replace( module.moe_config, num_experts=module.moe_config.num_local_experts * ep_group.world_size, moe_parallel_config=moe_parallel_config, ) def stage_standby_moe_quant_methods(self, all2all_manager) -> None: standby_dp_group = get_standby_dp_group() standby_ep_group = get_standby_ep_group() model = self.worker.model_runner.get_model() moe_modules = [module for module in model.modules() if is_moe_layer(module)] self._staged_moe_quant_methods.clear() with set_current_vllm_config(self.worker.vllm_config): for module in moe_modules: staged_quant_method = make_eep_staged_quant_method( module, self._make_eep_moe_config( module, standby_dp_group, standby_ep_group, ), all2all_manager=all2all_manager, ) if staged_quant_method is not None: self._staged_moe_quant_methods[module] = staged_quant_method def _commit_staged_moe_quant_methods(self) -> None: model = self.worker.model_runner.get_model() moe_modules = [module for module in model.modules() if is_moe_layer(module)] for module in moe_modules: staged_quant_method = self._staged_moe_quant_methods.pop(module, None) if staged_quant_method is None: continue assert staged_quant_method.moe_kernel is not None module._replace_quant_method(staged_quant_method) self._staged_moe_quant_methods.clear() def _release_cuda_graphs(self) -> None: if isinstance(self.worker.model_runner.model, CUDAGraphWrapper): wrapper = self.worker.model_runner.model wrapper.concrete_cudagraph_entries = {} elif isinstance(self.worker.model_runner.model, UBatchWrapper): self.worker.model_runner.model.clear_graphs() torch.compiler.reset() with set_current_vllm_config(self.worker.vllm_config): reset_compile_wrapper(self.worker.model_runner.get_model()) gc.collect() torch.accelerator.synchronize() torch.accelerator.empty_cache() def switch_and_remove(self) -> None: # Removing ranks skipped preparation, so wait for prior cleanup here. self._wait_for_group_cleanup() self._release_cuda_graphs() retired_groups = _replace_active_groups( world=None, dp=None, ep=None, eplb=None, node_count=None ) self._start_group_cleanup(retired_groups) # Finish collective cleanup before this worker is shut down. self._wait_for_group_cleanup() def switch_and_prepare(self) -> tuple[GroupCoordinator | None, ...]: if not self._can_reuse_fused_moe_kernel(): self._release_cuda_graphs() retired_groups = _replace_active_groups(**pop_standby_groups()) parallel_config = self.worker.vllm_config.parallel_config reconfig_request = self.reconfig_request assert reconfig_request is not None new_dp_size = reconfig_request.new_data_parallel_size new_ep_size = get_ep_group().world_size parallel_config.data_parallel_size = new_dp_size if ( reconfig_request.new_data_parallel_rank != ReconfigureRankType.KEEP_CURRENT_RANK ): parallel_config.data_parallel_rank = reconfig_request.new_data_parallel_rank if ( reconfig_request.new_data_parallel_rank_local != ReconfigureRankType.KEEP_CURRENT_RANK ): parallel_config.data_parallel_rank_local = ( reconfig_request.new_data_parallel_rank_local ) parallel_config.data_parallel_master_ip = ( reconfig_request.new_data_parallel_master_ip ) parallel_config.data_parallel_master_port = ( reconfig_request.new_data_parallel_master_port ) # Reconfigure MoE modules with new EP size moe_modules = [ module for module in self.worker.model_runner.model.modules() if is_moe_layer(module) ] num_local_experts = moe_modules[0].moe_config.num_local_experts assert all( module.moe_config.num_local_experts == num_local_experts for module in moe_modules ), "All MoE modules must have the same number of experts" dp_group = get_dp_group() ep_group = get_ep_group() for module in moe_modules: new_moe_config = self._make_eep_moe_config(module, dp_group, ep_group) module._set_moe_config(new_moe_config) # Update EPLB state eplb_state = self.worker.model_runner.eplb_state assert eplb_state is not None model_config = self.worker.model_runner.model_config eplb_model_state = eplb_state.model_states[model_config.compute_hash()] num_physical_experts = num_local_experts * new_ep_size num_logical_experts = eplb_model_state.logical_replica_count.shape[1] parallel_config.eplb_config.num_redundant_experts = ( num_physical_experts - num_logical_experts ) eplb_state.reconfigure_physical_expert_slots(model_config, num_physical_experts) model = self.worker.model_runner.get_model() model.expert_weights = [] with set_current_vllm_config(self.worker.vllm_config): model.set_eplb_state( eplb_model_state.expert_load_pass_buffer, eplb_model_state.logical_to_physical_map, eplb_model_state.logical_replica_count, ) eplb_state._propagate_shared_tensors( model, eplb_model_state.num_unpadded_tokens_tensors ) model.update_physical_experts_metadata( num_physical_experts=num_physical_experts, num_local_physical_experts=num_local_experts, ) self._commit_staged_moe_quant_methods() get_ep_all2all_manager().commit_ep_size() # Legacy modular methods need to be recreated for the new EP size. for module in moe_modules: if getattr(module._quant_method, "wraps_legacy_quant_method", False): module._replace_quant_method(module._quant_method.old_quant_method) assert self._prepared_eplb_communicator is not None eplb_state.update_communicator(model_config, self._prepared_eplb_communicator) self._prepared_eplb_communicator = None return retired_groups def _perform_eplb_reshuffle( self, rank_mapping: dict[int, int] | None = None, async_op: bool = False, ) -> None: if get_ep_group().rank == 0: logger.info("[Elastic EP] Starting expert resharding...") eplb_state = self.worker.model_runner.eplb_state assert eplb_state is not None is_async_enabled = eplb_state.is_async run_async = async_op and is_async_enabled eplb_state.is_async = run_async if rank_mapping is None: eplb_state.rearrange() else: eplb_state.rearrange(rank_mapping=rank_mapping) if not run_async: # Wait for non-blocking expert resharding copies before continuing # the Elastic EP reconfiguration. torch.accelerator.synchronize() # reset expert_rearrangement_step to ensure all ranks are synchronized eplb_state.expert_rearrangement_step = 0 eplb_state.is_async = is_async_enabled # Start the async worker thread if it doesn't exist yet (idempotent). # This is needed for new workers after scale-up: they create EplbState # in setup_eplb_from_mapping() but don't start the thread there because # groups aren't ready yet. eplb_state.start_async_loop() if get_ep_group().rank == 0: logger.info( "[Elastic EP] Expert resharding %s", "scheduled" if run_async else "completed", ) def commit_scale_up(self, is_existing_worker: bool) -> None: if is_existing_worker: self.broadcast_expert_mapping() retired_groups = self.switch_and_prepare() else: mapping = self.receive_expert_mapping() self.worker.model_runner.setup_eplb_from_mapping(mapping) if not self._can_reuse_fused_moe_kernel(): # warm_and_capture's dummy run warms the deferred dp/ep groups. self.warm_and_capture() elif current_platform.is_rocm(): # Reuse path skips warm_and_capture; warm the deferred groups here. self._warm_target_groups(get_dp_group(), get_ep_group()) self._perform_eplb_reshuffle(async_op=True) if is_existing_worker: self._start_group_cleanup(retired_groups) def commit_scale_down(self, new_dp_size: int, removing: bool) -> None: self.perform_scale_down_eplb_reshuffle(new_dp_size) if removing: self.switch_and_remove() else: retired_groups = self.switch_and_prepare() if not self._can_reuse_fused_moe_kernel(): self.warm_and_capture() elif current_platform.is_rocm(): # Reuse path skips warm_and_capture; warm the deferred groups here. self._warm_target_groups(get_dp_group(), get_ep_group()) self._start_group_cleanup(retired_groups) def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None: eplb_state = self.worker.model_runner.eplb_state if eplb_state is not None: eplb_state.drain_async() parallel_config = self.worker.vllm_config.parallel_config tp_size = parallel_config.tensor_parallel_size old_ep_size = parallel_config.data_parallel_size * tp_size new_ep_size = new_dp_size * tp_size rank_mapping = { old_ep_rank: old_ep_rank if old_ep_rank < new_ep_size else -1 for old_ep_rank in range(old_ep_size) } self._perform_eplb_reshuffle(rank_mapping=rank_mapping) def prepare_new_worker(self) -> None: dp_group = get_dp_group() assert isinstance(dp_group, StatelessGroupCoordinator) new_dp_size = dp_group.world_size dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank # Receive old_dp_size broadcasted during transfer_weights old_dp_size_tensor = torch.empty(1, dtype=torch.int64, device="cpu") old_dp_size_tensor = dp_group.tcp_store_group.broadcast(old_dp_size_tensor, 0) old_dp_size = int(old_dp_size_tensor[0].item()) # Calculate which existing worker will send to this new worker num_new_workers = new_dp_size - old_dp_size new_worker_idx = dp_rank - old_dp_size num_dst_per_sender = num_new_workers // old_dp_size remainder = num_new_workers % old_dp_size if new_worker_idx < remainder * (num_dst_per_sender + 1): sender_rank = new_worker_idx // (num_dst_per_sender + 1) else: sender_rank = ( remainder + (new_worker_idx - remainder * (num_dst_per_sender + 1)) // num_dst_per_sender ) model = self.worker.model_runner.get_model() expert_weights = [ module.get_expert_weights() for module in model.modules() if is_moe_layer(module) ] batch_transfer_weights( model=model, is_sender=False, peer_rank=sender_rank, dp_group=dp_group, expert_weights=expert_weights, ) torch.accelerator.synchronize() if not current_platform.is_rocm(): self._warm_target_groups(get_dp_group(), get_ep_group()) if self._can_reuse_fused_moe_kernel(): sync_flashinfer_autotune_cache(self.worker.model_runner, get_world_group()) def receive_expert_mapping(self) -> torch.Tensor: dp_group = get_dp_group() assert isinstance(dp_group, StatelessGroupCoordinator) return broadcast_expert_mapping( physical_to_logical=None, dp_group=dp_group, src_rank=0, device=self.worker.device, ) def warmup_new_worker(self) -> None: with set_current_vllm_config(self.worker.vllm_config): kernel_warmup(self.worker, process_local_only=True) if self._can_reuse_fused_moe_kernel(): self.warm_and_capture() def warm_and_capture(self) -> None: # Save and clear block tables so the dummy MoE forward doesn't # write dummy slot mappings into real KV-cache blocks. multi_block_table = self.worker.model_runner.input_batch.block_table saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = [] for bt in multi_block_table.block_tables: saved_block_tables.append( (bt.block_table.gpu.clone(), bt.block_table.cpu.clone()) ) multi_block_table.clear() # _ensure_workspace_size allocates a fresh tensor on grow, leaving # any captured CUDA graph with a stale data pointer; drop graphs # before re-warm so captures realign with the resized buffer. self._release_cuda_graphs() unlock_workspace() # Grow the MoE workspace at max_num_tokens. compile_or_warm_up_model # alone only exercises cudagraph-capture sizes and can leave the # workspace too small for post-reshuffle routing. Use _dummy_run # directly with skip_eplb=True so dummy routing doesn't pollute the # just-rebalanced EPLB stats. runner = self.worker.model_runner all2all_manager = get_ep_all2all_manager() reuse_kernel = self._can_reuse_fused_moe_kernel() with ( skip_dp_coordination() if reuse_kernel else nullcontext(), all2all_manager.mask_remote_ranks() if reuse_kernel else nullcontext(), self._disable_flashinfer_autotune() if reuse_kernel else nullcontext(), ): runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True) self.worker.compile_or_warm_up_model() lock_workspace() for bt, (saved_gpu, saved_cpu) in zip( multi_block_table.block_tables, saved_block_tables ): bt.block_table.gpu.copy_(saved_gpu) bt.block_table.cpu.copy_(saved_cpu)