# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import gc import weakref from collections.abc import Iterable, Sequence from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import replace from typing import TYPE_CHECKING import torch import torch.nn as nn import torch.nn.functional as F from torch.distributed import P2POp from vllm.compilation.counter import compilation_counter from vllm.compilation.cuda_graph import CUDAGraphWrapper from vllm.compilation.wrapper import reset_compile_wrapper from vllm.config import ( CompilationMode, set_current_vllm_config, ) from vllm.distributed import ( get_dp_group, get_ep_group, get_pcp_group, get_tp_group, ) from vllm.distributed.elastic_ep.standby_state import ( create_standby_groups, get_standby_dp_group, get_standby_ep_group, get_standby_eplb_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.config import FusedMoEParallelConfig from vllm.model_executor.layers.fused_moe.eep_reconfigure import ( make_eep_staged_quant_method, ) from vllm.model_executor.warmup.kernel_warmup import kernel_warmup from vllm.utils import is_moe_layer from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType 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 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, num_local_physical_experts: int | None, num_logical_experts: int | None, dp_group: StatelessGroupCoordinator, device: torch.device, src_rank: int = 0, ) -> tuple[torch.Tensor, int, int]: if dp_group.rank_in_group == src_rank: assert physical_to_logical is not None assert num_local_physical_experts is not None assert num_logical_experts is not None assert physical_to_logical.dtype == torch.int64 shape_tensor = torch.tensor( list(physical_to_logical.shape), dtype=torch.int64, device="cpu" ) metadata_tensor = torch.tensor( [num_local_physical_experts, num_logical_experts], dtype=torch.int64, device="cpu", ) else: shape_tensor = torch.empty(2, dtype=torch.int64, device="cpu") metadata_tensor = torch.empty(2, dtype=torch.int64, device="cpu") shape_tensor = dp_group.tcp_store_group.broadcast(shape_tensor, src_rank) metadata_tensor = dp_group.tcp_store_group.broadcast(metadata_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 physical_to_logical = dp_group.broadcast(physical_to_logical, src_rank) num_local_physical_experts = int(metadata_tensor[0].item()) num_logical_experts = int(metadata_tensor[1].item()) return physical_to_logical, num_local_physical_experts, num_logical_experts 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 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, ) self.stage_standby_moe_quant_methods() self._prepare_eplb_communicator(get_standby_eplb_group()) if new_dp_size > old_dp_size: self.transfer_weights(old_dp_size, new_dp_size) self._warm_target_groups(get_standby_dp_group(), get_standby_ep_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 num_physical_experts = physical_to_logical.shape[1] num_local_physical_experts = num_physical_experts // get_ep_group().world_size num_logical_experts = eplb_model_state.logical_replica_count.shape[1] broadcast_expert_mapping( physical_to_logical=physical_to_logical, num_local_physical_experts=num_local_physical_experts, num_logical_experts=num_logical_experts, 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) -> 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, ), ) 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) staged_quant_method.moe_kernel.prepare_finalize.on_commit() 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): raise RuntimeError("DBO is not yet supported in elastic EP") 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, ...]: old_dp_size = get_dp_group().world_size old_ep_size = get_ep_group().world_size 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 ) old_physical_to_logical = eplb_model_state.physical_to_logical_map num_moe_layers = old_physical_to_logical.shape[0] num_local_experts = eplb_model_state.expert_load_pass.shape[1] // old_ep_size if new_dp_size > old_dp_size: expanded_physical_to_logical = torch.full( (num_moe_layers, num_local_experts * new_ep_size), -1, dtype=old_physical_to_logical.dtype, device=old_physical_to_logical.device, ) expanded_physical_to_logical[:, : num_local_experts * old_ep_size] = ( old_physical_to_logical ) eplb_model_state.physical_to_logical_map = expanded_physical_to_logical old_num_physical_experts = eplb_model_state.expert_load_pass.shape[1] pad_size = num_physical_experts - old_num_physical_experts if new_dp_size > old_dp_size: assert pad_size > 0 expanded_expert_load_pass = F.pad( eplb_model_state.expert_load_pass, (0, pad_size), value=0 ) expanded_expert_load_window = F.pad( eplb_model_state.expert_load_window, (0, pad_size), value=0 ) eplb_model_state.expert_load_pass = expanded_expert_load_pass eplb_model_state.expert_load_window = expanded_expert_load_window else: assert pad_size < 0 eplb_model_state.expert_load_pass = eplb_model_state.expert_load_pass[ :, :num_physical_experts ] eplb_model_state.expert_load_window = eplb_model_state.expert_load_window[ :, :, :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, 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() # 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 if ( self.worker.vllm_config.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE ): # NOTE(yongji): when using stock torch.compile, # torch.compile is triggered during GPUModelRunner's load_model() # TODO(yongji):check do we need to re-trigger torch.compile here? # any changes to the tensor shapes in execution should already # be handled internally by torch.compile. backend = self.worker.vllm_config.compilation_config.init_backend( self.worker.vllm_config ) compilation_counter.stock_torch_compile_count += 1 self.worker.model_runner.model.compile(fullgraph=True, backend=backend) 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) self.warm_and_capture() 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() self.warm_and_capture() 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() self._warm_target_groups(get_dp_group(), get_ep_group()) def receive_expert_mapping(self) -> torch.Tensor: dp_group = get_dp_group() assert isinstance(dp_group, StatelessGroupCoordinator) physical_to_logical, num_local_physical_experts, _ = broadcast_expert_mapping( physical_to_logical=None, num_local_physical_experts=None, num_logical_experts=None, dp_group=dp_group, src_rank=0, device=self.worker.device, ) num_moe_layers = physical_to_logical.shape[0] new_dp_size = get_dp_group().world_size tp_size = self.worker.vllm_config.parallel_config.tensor_parallel_size new_ep_size = new_dp_size * tp_size expanded_physical_to_logical = torch.full( (num_moe_layers, num_local_physical_experts * new_ep_size), -1, dtype=physical_to_logical.dtype, device=physical_to_logical.device, ) expanded_physical_to_logical[:, : physical_to_logical.shape[1]] = ( physical_to_logical ) return expanded_physical_to_logical def warmup_local_kernels(self) -> None: with set_current_vllm_config(self.worker.vllm_config): kernel_warmup(self.worker, process_local_only=True) def warm_and_capture(self) -> None: # Must run on every DP sibling in lockstep: _dummy_run calls # coordinate_batch_across_dp whenever data_parallel_size > 1 # (gpu_model_runner.py:3663), which deadlocks if any rank skips it. # 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 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)