diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt index 1760ed6517..9d9d127881 100644 --- a/requirements/kv_connectors.txt +++ b/requirements/kv_connectors.txt @@ -2,7 +2,7 @@ lmcache >= 0.3.9 # CuPy 14.1.0 imports pytest from cupy.testing._random, which breaks runtime # images. 14.1.1 fixes the root cause, so only 14.1.0 is excluded. cupy-cuda13x != 14.1.0 -nixl == 1.3.2 +nixl == 1.4.1 # CUDA 12 build. On the CUDA 13 image install-kv-connectors.sh swaps this for # the mooncake-transfer-engine-cuda13 variant of the same version. mooncake-transfer-engine >= 0.3.12 diff --git a/tests/distributed/test_elastic_ep.py b/tests/distributed/test_elastic_ep.py index 01c254c2d0..8ccea86283 100644 --- a/tests/distributed/test_elastic_ep.py +++ b/tests/distributed/test_elastic_ep.py @@ -27,7 +27,7 @@ MODEL_NAME = "deepseek-ai/DeepSeek-V2-Lite-Chat" NUM_GSM8K_QUESTIONS = 256 EXPECTED_ACCURACY = 0.58 ACCURACY_TOL = 0.08 -MAX_NUM_SEQS = 32 +DEFAULT_ALL2ALL_BACKEND = "allgather_reducescatter" def _send_scale_command(server: RemoteOpenAIServer, new_dp_size: int) -> bool: @@ -161,7 +161,12 @@ def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float: return accuracy -def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str]: +def _base_serve_args( + dp_size: int = 2, + max_dp_size: int = 4, + enforce_eager: bool = False, + all2all_backend: str = DEFAULT_ALL2ALL_BACKEND, +) -> list[str]: args = [ "--trust-remote-code", "--tensor-parallel-size", @@ -170,12 +175,12 @@ def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str] "0.8", "--max-model-len", "4096", - "--max-num-seqs", - str(MAX_NUM_SEQS), "--enable-expert-parallel", "--all2all-backend", - "allgather_reducescatter", + all2all_backend, "--enable-elastic-ep", + "--elastic-ep-max-dp-size", + str(max_dp_size), "--enable-eplb", "--eplb-config.num_redundant_experts", "0", @@ -205,28 +210,43 @@ def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str] @pytest.mark.parametrize( - ("enforce_eager", "traffic_mode"), + ("enforce_eager", "traffic_mode", "all2all_backend"), [ - pytest.param(True, "none", id="enforce_eager_none"), - pytest.param(True, "light", id="enforce_eager_light"), - pytest.param(True, "heavy", id="enforce_eager_heavy"), - pytest.param(False, "heavy", id="cuda_graphs_heavy"), + pytest.param(True, "none", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_none"), + pytest.param(True, "light", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_light"), + pytest.param(True, "heavy", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_heavy"), + pytest.param( + False, + "heavy", + DEFAULT_ALL2ALL_BACKEND, + id="cuda_graphs_heavy", + ), + pytest.param(False, "heavy", "nixl_ep", id="cuda_graphs_heavy_nixl_ep"), ], ) @multi_gpu_test(num_gpus=4) -def test_elastic_ep_scaling(enforce_eager: bool, traffic_mode: str): +def test_elastic_ep_scaling( + enforce_eager: bool, traffic_mode: str, all2all_backend: str +): from vllm.distributed.eplb.eplb_communicator import has_nixl + from vllm.utils.import_utils import has_nixl_ep if not has_nixl(): pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") + if all2all_backend == "nixl_ep" and not has_nixl_ep(): + pytest.skip("NIXL EP is not installed") initial_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_INITIAL_DP", "2")) target_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_TARGET_DP", "4")) assert target_dp_size > initial_dp_size - vllm_serve_args = _base_serve_args(initial_dp_size, enforce_eager) + vllm_serve_args = _base_serve_args( + initial_dp_size, target_dp_size, enforce_eager, all2all_backend + ) with RemoteOpenAIServer( - MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200 + MODEL_NAME, + vllm_serve_args, + max_wait_seconds=1200, ) as server: initial_accuracy = _run_gsm8k_eval(server, "Initial") @@ -255,55 +275,3 @@ def test_elastic_ep_scaling(enforce_eager: bool, traffic_mode: str): f"(diff: {scale_down_accuracy - initial_accuracy:+.3f})" ) print(f" Tolerance: {ACCURACY_TOL:.3f}") - - -@multi_gpu_test(num_gpus=4) -def test_elastic_ep_scaling_uneven(): - """Test scale up with uneven worker distribution. - - This tests the case where num_new_workers % old_dp_size != 0, - specifically 2 -> 3 where remainder = 1 % 2 = 1. - This exercises the remainder handling in sender-receiver pairing. - """ - from vllm.distributed.eplb.eplb_communicator import has_nixl - - if not has_nixl(): - pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") - - vllm_serve_args = _base_serve_args() - - with RemoteOpenAIServer( - MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200 - ) as server: - initial_accuracy = _run_gsm8k_eval(server, "Initial (2 GPUs)") - - # Scale 2 -> 3: This has remainder = 1 % 2 = 1 - # Tests uneven sender-receiver pairing - assert _send_scale_command(server, 3) - scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (3 GPUs)") - - assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, ( - f"Scale up accuracy {scale_up_accuracy:.3f} dropped more than " - f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}" - ) - - # Scale back down to 2 - assert _send_scale_command(server, 2) - scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)") - - assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, ( - f"Scale down accuracy {scale_down_accuracy:.3f} dropped more than " - f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}" - ) - - print("\nAccuracy Summary (Uneven Scaling):") - print(f" Initial: {initial_accuracy:.3f}") - print( - f" Scale up: {scale_up_accuracy:.3f} " - f"(diff: {scale_up_accuracy - initial_accuracy:+.3f})" - ) - print( - f" Scale down: {scale_down_accuracy:.3f} " - f"(diff: {scale_down_accuracy - initial_accuracy:+.3f})" - ) - print(f" Tolerance: {ACCURACY_TOL:.3f}") diff --git a/tests/distributed/test_eplb_utils.py b/tests/distributed/test_eplb_utils.py index 53a4ce21af..a6cfb0d29d 100644 --- a/tests/distributed/test_eplb_utils.py +++ b/tests/distributed/test_eplb_utils.py @@ -16,10 +16,14 @@ def _make_model_state( phy2log: torch.Tensor, log2phy: torch.Tensor, logcnt: torch.Tensor, + phy2log_storage: torch.Tensor | None = None, ) -> MagicMock: - """Build a minimal EplbModelState mock with only the three map tensors.""" + """Build a minimal EplbModelState mock with its map tensors.""" state = MagicMock() state.physical_to_logical_map = phy2log + state.physical_to_logical_map_buffer = ( + phy2log if phy2log_storage is None else phy2log_storage + ) state.logical_to_physical_map = log2phy state.logical_replica_count = logcnt return state @@ -27,15 +31,17 @@ def _make_model_state( def test_commit_eplb_maps_shape_change(): """ - The normal path copies the physical_to_logical map in-place. When the number of - physical experts changes, the old map should be replaced entirely. + When the number of physical experts changes, resize the active map within its + preallocated storage. """ num_layers, num_logical, num_physical = 2, 4, 6 max_replicas = 3 # Build current state tensors + storage = torch.full((num_layers, num_physical + 2), -1, dtype=torch.long) model_state = _make_model_state( - phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long), + phy2log=storage[:, :num_physical], + phy2log_storage=storage, log2phy=torch.full( (num_layers, num_logical, max_replicas), -1, dtype=torch.long ), @@ -54,6 +60,7 @@ def test_commit_eplb_maps_shape_change(): # Check that the number of physical experts has been updated and that the values # match assert model_state.physical_to_logical_map.shape[1] == num_physical + 2 + assert model_state.physical_to_logical_map.data_ptr() == storage.data_ptr() assert torch.equal(model_state.physical_to_logical_map, new_phy2log_larger) diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index bd2532b1cc..0213e06283 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -216,6 +216,8 @@ class ParallelConfig: enable_elastic_ep: bool = False """Enable elastic expert parallelism with stateless NCCL groups for DP/EP.""" + elastic_ep_max_dp_size: int = Field(default=None, ge=1) # type: ignore[assignment] + """Maximum data parallel size supported by elastic expert parallelism.""" enable_dbo: bool = False """Enable dual batch overlap for the model executor.""" @@ -435,7 +437,9 @@ class ParallelConfig: ) """The configurations for fault tolerance.""" - @field_validator("disable_nccl_for_dp_synchronization", mode="wrap") + @field_validator( + "disable_nccl_for_dp_synchronization", "elastic_ep_max_dp_size", mode="wrap" + ) @classmethod def _skip_none_validation(cls, value: Any, handler: Callable) -> Any: """Skip validation if the value is `None` when initialisation is delayed.""" @@ -941,6 +945,18 @@ class ParallelConfig: " for dense models." ) + max_dp_size = self.elastic_ep_max_dp_size + self.elastic_ep_max_dp_size = ( + max_dp_size + if self.enable_elastic_ep and max_dp_size is not None + else self.data_parallel_size + ) + if self.elastic_ep_max_dp_size < self.data_parallel_size: + raise ValueError( + "--elastic-ep-max-dp-size must be greater than or equal to " + f"the initial data_parallel_size ({self.data_parallel_size})." + ) + self.data_parallel_index = self.data_parallel_rank if self.distributed_executor_backend == "external_launcher": diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index eb187a99f4..080e7feeff 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import threading +from collections.abc import Iterator +from contextlib import contextmanager from dataclasses import dataclass from typing import Any @@ -377,6 +379,7 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase): @dataclass class _NixlEPBufferState: buffer: Any + max_num_ep_ranks: int connected_ep_size: int active_ep_size: int @@ -402,7 +405,10 @@ class NixlEPAll2AllManager(All2AllManagerBase): super().__init__(cpu_group, tcp_store_group) self.support_fault_tolerance = True - self.max_num_ep_ranks = envs.VLLM_NIXL_EP_MAX_NUM_RANKS + @property + def max_num_ep_ranks(self) -> int: + assert NixlEPAll2AllManager._buffer is not None + return NixlEPAll2AllManager._buffer.max_num_ep_ranks def _init_buffer( self, @@ -412,11 +418,19 @@ class NixlEPAll2AllManager(All2AllManagerBase): ) -> None: from nixl_ep import Buffer # type: ignore[import-not-found] - max_num_global_experts = self.max_num_ep_ranks * num_experts_per_rank + parallel_config = get_current_vllm_config().parallel_config + max_num_ep_ranks = ( + self.world_size + * parallel_config.elastic_ep_max_dp_size + // parallel_config.data_parallel_size + if parallel_config.enable_elastic_ep + else self.world_size + ) + max_num_global_experts = max_num_ep_ranks * num_experts_per_rank num_rdma_bytes = Buffer.get_rdma_size_hint( num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank, hidden=token_hidden_size, - num_ranks=self.max_num_ep_ranks, + num_ranks=max_num_ep_ranks, num_experts=max_num_global_experts, ) assert NixlEPAll2AllManager._buffer is None, ( @@ -427,7 +441,7 @@ class NixlEPAll2AllManager(All2AllManagerBase): tcp_store_group=self.tcp_store_group.store, ) buffer.update_memory_buffers( - num_ranks=self.max_num_ep_ranks, + num_ranks=max_num_ep_ranks, num_experts_per_rank=num_experts_per_rank, num_rdma_bytes=num_rdma_bytes, ) @@ -435,6 +449,7 @@ class NixlEPAll2AllManager(All2AllManagerBase): buffer.connect_ranks(ranks_to_connect) NixlEPAll2AllManager._buffer = _NixlEPBufferState( buffer=buffer, + max_num_ep_ranks=max_num_ep_ranks, connected_ep_size=self.world_size, active_ep_size=self.world_size, ) @@ -474,6 +489,7 @@ class NixlEPAll2AllManager(All2AllManagerBase): for rank in range(state.active_ep_size, target_ep_size): state.buffer.update_mask_buffer(rank, mask=False) + torch.accelerator.synchronize() state.active_ep_size = target_ep_size def _stage_ep_size(self) -> None: @@ -486,8 +502,12 @@ class NixlEPAll2AllManager(All2AllManagerBase): if target_ep_size > state.connected_ep_size: self._connect_to_ep_size(target_ep_size, make_active=False) - def commit_staged_state(self) -> None: - """Commit staged NIXL EP state to the active communication set.""" + def stage_ep_size(self) -> None: + with NixlEPAll2AllManager._lock: + self._stage_ep_size() + + def commit_ep_size(self) -> None: + """Commit the staged EP size to the active communication set.""" with NixlEPAll2AllManager._lock: assert NixlEPAll2AllManager._buffer is not None state = NixlEPAll2AllManager._buffer @@ -500,31 +520,31 @@ class NixlEPAll2AllManager(All2AllManagerBase): self._unmask_connected_ranks(target_ep_size) - def _ensure_ep_size(self, *, stage: bool) -> None: - if stage: - self._stage_ep_size() - else: - self.commit_staged_state() + @contextmanager + def mask_remote_ranks(self) -> Iterator[None]: + peers = [rank for rank in range(self.world_size) if rank != self.rank] + assert NixlEPAll2AllManager._buffer is not None + buffer = NixlEPAll2AllManager._buffer.buffer + for rank in peers: + buffer.update_mask_buffer(rank, mask=True) + torch.accelerator.synchronize() + try: + yield + finally: + for rank in peers: + buffer.update_mask_buffer(rank, mask=False) + torch.accelerator.synchronize() def get_handle(self, kwargs): with NixlEPAll2AllManager._lock: - stage = bool(kwargs.get("stage", False)) state = NixlEPAll2AllManager._buffer if state is None: - assert not stage, ( - "NIXL EP staged initialization requires an existing buffer" - ) max_num_tokens_per_dp_rank = kwargs["max_num_tokens_per_dp_rank"] - num_experts_per_rank = ( - kwargs["num_global_experts"] // kwargs["num_ep_ranks"] - ) self._init_buffer( max_num_tokens_per_dp_rank=max_num_tokens_per_dp_rank, token_hidden_size=kwargs["token_hidden_size"], - num_experts_per_rank=num_experts_per_rank, + num_experts_per_rank=kwargs["num_local_experts"], ) - else: - self._ensure_ep_size(stage=stage) assert NixlEPAll2AllManager._buffer is not None handle = NixlEPAll2AllManager._buffer.buffer @@ -564,7 +584,7 @@ class NixlEPAll2AllManager(All2AllManagerBase): assert state is not None if NixlEPAll2AllManager._mask is None: NixlEPAll2AllManager._mask = torch.zeros( - self.max_num_ep_ranks, device="cuda", dtype=torch.int32 + state.max_num_ep_ranks, device="cuda", dtype=torch.int32 ) state.buffer.query_mask_buffer(NixlEPAll2AllManager._mask) return NixlEPAll2AllManager._mask[: state.active_ep_size] diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index 83ae25ff44..4c1900590b 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import threading +from collections.abc import Iterator +from contextlib import contextmanager from weakref import WeakValueDictionary import torch @@ -157,6 +159,16 @@ class All2AllManagerBase: def destroy(self): pass + def stage_ep_size(self) -> None: + pass + + def commit_ep_size(self) -> None: + pass + + @contextmanager + def mask_remote_ranks(self) -> Iterator[None]: + yield + class DeviceCommunicatorBase: """ diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index 6124d22b44..343edd171f 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -2,34 +2,32 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import gc import weakref -from collections.abc import Iterable, Sequence +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 -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.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 ( @@ -41,13 +39,20 @@ from vllm.distributed.parallel_state import ( ) 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.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 @@ -59,6 +64,15 @@ if TYPE_CHECKING: ) +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, @@ -105,31 +119,20 @@ def batch_transfer_weights( 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]: +) -> torch.Tensor: 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 @@ -140,11 +143,7 @@ def broadcast_expert_mapping( ) 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 + return dp_group.broadcast(physical_to_logical, src_rank) class ElasticEPScalingExecutor: @@ -166,6 +165,21 @@ class ElasticEPScalingExecutor: 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: @@ -260,11 +274,20 @@ class ElasticEPScalingExecutor: use_all2all=use_all2all, enable_eplb=parallel_config.enable_eplb, ) - self.stage_standby_moe_quant_methods() + 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) - self._warm_target_groups(get_standby_dp_group(), get_standby_ep_group()) + 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 @@ -338,14 +361,9 @@ class ElasticEPScalingExecutor: 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] + physical_to_logical = eplb_model_state.physical_to_logical_map_buffer 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, @@ -368,7 +386,7 @@ class ElasticEPScalingExecutor: moe_parallel_config=moe_parallel_config, ) - def stage_standby_moe_quant_methods(self) -> None: + 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() @@ -383,6 +401,7 @@ class ElasticEPScalingExecutor: 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 @@ -396,7 +415,6 @@ class ElasticEPScalingExecutor: 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: @@ -405,7 +423,7 @@ class ElasticEPScalingExecutor: wrapper.concrete_cudagraph_entries = {} elif isinstance(self.worker.model_runner.model, UBatchWrapper): - raise RuntimeError("DBO is not yet supported in elastic EP") + self.worker.model_runner.model.clear_graphs() torch.compiler.reset() with set_current_vllm_config(self.worker.vllm_config): @@ -427,10 +445,8 @@ class ElasticEPScalingExecutor: 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() + 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 @@ -487,47 +503,13 @@ class ElasticEPScalingExecutor: 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 - ] + 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, + eplb_model_state.expert_load_pass_buffer, eplb_model_state.logical_to_physical_map, eplb_model_state.logical_replica_count, ) @@ -539,6 +521,7 @@ class ElasticEPScalingExecutor: 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): @@ -548,21 +531,6 @@ class ElasticEPScalingExecutor: 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( @@ -608,7 +576,8 @@ class ElasticEPScalingExecutor: else: mapping = self.receive_expert_mapping() self.worker.model_runner.setup_eplb_from_mapping(mapping) - self.warm_and_capture() + if not self._can_reuse_fused_moe_kernel(): + self.warm_and_capture() self._perform_eplb_reshuffle(async_op=True) if is_existing_worker: self._start_group_cleanup(retired_groups) @@ -619,7 +588,8 @@ class ElasticEPScalingExecutor: self.switch_and_remove() else: retired_groups = self.switch_and_prepare() - self.warm_and_capture() + if not self._can_reuse_fused_moe_kernel(): + self.warm_and_capture() self._start_group_cleanup(retired_groups) def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None: @@ -677,42 +647,26 @@ class ElasticEPScalingExecutor: ) torch.accelerator.synchronize() 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) - physical_to_logical, num_local_physical_experts, _ = broadcast_expert_mapping( + return 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: + 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: - # 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 @@ -735,8 +689,15 @@ class ElasticEPScalingExecutor: # 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() + 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() diff --git a/vllm/distributed/elastic_ep/elastic_state.py b/vllm/distributed/elastic_ep/elastic_state.py index 6718cd8ae8..5ba1891b7a 100644 --- a/vllm/distributed/elastic_ep/elastic_state.py +++ b/vllm/distributed/elastic_ep/elastic_state.py @@ -204,7 +204,7 @@ class ElasticEPScalingState: return True elif state == ScaleUpNewEngineState.PREPARE: - self._collective_rpc("elastic_ep_execute", args=("warmup_local_kernels",)) + self._collective_rpc("elastic_ep_execute", args=("warmup_new_worker",)) self._mark_ready_for_switch() tensor = torch.tensor([0, 0, 0], dtype=torch.int32, device="cpu") torch.distributed.all_reduce( diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 580477ba23..bc3aa36bb2 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -121,6 +121,8 @@ class EplbModelState: [0, 2, 0, 1, 0, 3]] ``` """ + physical_to_logical_map_buffer: torch.Tensor + """Maximum-capacity buffer backing ``physical_to_logical_map``.""" logical_to_physical_map: torch.Tensor """ Mapping from logical experts to physical experts. @@ -168,6 +170,8 @@ class EplbModelState: Shape: (num_moe_layers, num_physical_experts) """ + expert_load_pass_buffer: torch.Tensor + """Maximum-capacity buffer backing ``expert_load_pass``.""" expert_load_window: torch.Tensor """ A sliding window of expert load. @@ -373,6 +377,19 @@ class EplbState: physical_to_logical_map_list, device=self.device, ) + physical_expert_capacity = model.num_physical_experts + if self.parallel_config.enable_elastic_ep: + physical_expert_capacity = ( + model.num_local_physical_experts + * self.parallel_config.elastic_ep_max_dp_size + * get_ep_group().world_size + // self.parallel_config.data_parallel_size + ) + physical_to_logical_map = torch.nn.functional.pad( + physical_to_logical_map, + (0, physical_expert_capacity - model.num_physical_experts), + value=-1, + ) # Assuming 8 GPUs per node, this supports up to # (1023 + 1) / 8 = 128 nodes for now. # TODO(rui): make this configurable @@ -407,6 +424,10 @@ class EplbState: ) .contiguous() ) + physical_to_logical_map_buffer = physical_to_logical_map + physical_to_logical_map = physical_to_logical_map_buffer[ + :, : model.num_physical_experts + ] logical_to_physical_map = ( logical_to_physical_map.unsqueeze(0) .expand( @@ -425,11 +446,12 @@ class EplbState: .contiguous() ) - expert_load_pass = torch.zeros( - (model.num_moe_layers, model.num_physical_experts), + expert_load_pass_buffer = torch.zeros( + (model.num_moe_layers, physical_expert_capacity), dtype=torch.int32, device=self.device, ) + expert_load_pass = expert_load_pass_buffer[:, : model.num_physical_experts] self.expert_load_window_size = self.parallel_config.eplb_config.window_size expert_load_window = torch.zeros( ( @@ -460,7 +482,7 @@ class EplbState: ] model.set_eplb_state( - expert_load_pass, + expert_load_pass_buffer, logical_to_physical_map, logical_replica_count, ) @@ -479,9 +501,11 @@ class EplbState: model_state = EplbModelState( physical_to_logical_map=physical_to_logical_map, + physical_to_logical_map_buffer=physical_to_logical_map_buffer, logical_to_physical_map=logical_to_physical_map, logical_replica_count=logical_replica_count, expert_load_pass=expert_load_pass, + expert_load_pass_buffer=expert_load_pass_buffer, expert_load_window=expert_load_window, model_name=model_config.model, model=model, @@ -1075,10 +1099,14 @@ class EplbState: expanded_physical_to_logical: torch.Tensor, ) -> None: eplb_model_state = self.model_states[model_config.compute_hash()] - eplb_model_state.physical_to_logical_map.copy_(expanded_physical_to_logical) + eplb_model_state.physical_to_logical_map_buffer.copy_( + expanded_physical_to_logical + ) + eplb_model_state.expert_load_pass.zero_() + physical_to_logical_map = eplb_model_state.physical_to_logical_map (logical_to_physical_map_cpu, logical_replica_count_cpu) = compute_logical_maps( - expanded_physical_to_logical.cpu(), + physical_to_logical_map.cpu(), eplb_model_state.model.num_logical_experts, ) @@ -1097,6 +1125,33 @@ class EplbState: eplb_model_state.logical_to_physical_map.copy_(logical_to_physical_map) eplb_model_state.logical_replica_count.copy_(logical_replica_count) + def reconfigure_physical_expert_slots( + self, + model_config: ModelConfig, + num_physical_experts: int, + ) -> None: + """Replace physical_to_logical_map and expert_load_pass with views + covering the new active size. + """ + model_state = self.model_states[model_config.compute_hash()] + old_num_physical_experts = model_state.model.num_physical_experts + first_slot, last_slot = sorted((old_num_physical_experts, num_physical_experts)) + physical_to_logical_map_buffer = model_state.physical_to_logical_map_buffer + expert_load_pass_buffer = model_state.expert_load_pass_buffer + assert last_slot <= physical_to_logical_map_buffer.shape[1] + expert_slots = slice(first_slot, last_slot) + physical_to_logical_map_buffer[:, expert_slots].fill_(-1) + expert_load_pass_buffer[:, expert_slots].zero_() + model_state.physical_to_logical_map = physical_to_logical_map_buffer[ + :, :num_physical_experts + ] + model_state.expert_load_pass = expert_load_pass_buffer[:, :num_physical_experts] + + pad_size = num_physical_experts - model_state.expert_load_window.shape[-1] + model_state.expert_load_window = torch.nn.functional.pad( + model_state.expert_load_window, (0, pad_size) + ) + def create_communicator( self, model_config: ModelConfig, group_coordinator: GroupCoordinator ) -> EplbCommunicator: @@ -1333,14 +1388,13 @@ def _commit_eplb_maps( if PIN_MEMORY and src.is_cpu: src = src.new_empty(src.shape, pin_memory=True).copy_(src) - # Rare Case: When the number of physical experts has changed, discard the old - # physical to logical expert map and use the new one. This only happens when the - # number of GPUs available to vLLM changes while vLLM is running. Otherwise copy the - # new map into the old one. + # When the number of physical experts changes, refresh the active map view to + # cover the new size before copying it. if src.shape[1] != dst.shape[1]: - model_state.physical_to_logical_map = src.to(dst.device, non_blocking=True) - else: - dst.copy_(src, non_blocking=True) + dst = model_state.physical_to_logical_map_buffer[:, : src.shape[1]] + assert dst.shape == src.shape + model_state.physical_to_logical_map = dst + dst.copy_(src, non_blocking=True) num_logical_experts = model_state.logical_to_physical_map.shape[1] new_logical, new_replica_count = compute_logical_maps(src, num_logical_experts) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index db398438d0..ae6f14b1df 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -507,6 +507,7 @@ class EngineArgs: linear_backend: LinearBackend = KernelConfig.linear_backend all2all_backend: All2AllBackend = ParallelConfig.all2all_backend enable_elastic_ep: bool = ParallelConfig.enable_elastic_ep + elastic_ep_max_dp_size: int = ParallelConfig.elastic_ep_max_dp_size enable_dbo: bool = ParallelConfig.enable_dbo ubatch_size: int = ParallelConfig.ubatch_size dbo_decode_token_threshold: int = ParallelConfig.dbo_decode_token_threshold @@ -1201,6 +1202,10 @@ class EngineArgs: parallel_group.add_argument( "--enable-elastic-ep", **parallel_kwargs["enable_elastic_ep"] ) + parallel_group.add_argument( + "--elastic-ep-max-dp-size", + **parallel_kwargs["elastic_ep_max_dp_size"], + ) parallel_group.add_argument( "--dbo-decode-token-threshold", **parallel_kwargs["dbo_decode_token_threshold"], @@ -2345,6 +2350,7 @@ class EngineArgs: enable_ep_weight_filter=self.enable_ep_weight_filter, all2all_backend=self.all2all_backend, enable_elastic_ep=self.enable_elastic_ep, + elastic_ep_max_dp_size=self.elastic_ep_max_dp_size, enable_dbo=self.enable_dbo, ubatch_size=self.ubatch_size, dbo_decode_token_threshold=self.dbo_decode_token_threshold, diff --git a/vllm/envs.py b/vllm/envs.py index 1c22466398..aa32fc38ba 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -316,7 +316,6 @@ if TYPE_CHECKING: VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: bool = False VLLM_ELASTIC_EP_DRAIN_REQUESTS: bool = False VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = True - VLLM_NIXL_EP_MAX_NUM_RANKS: int = 32 VLLM_XPU_ENABLE_XPU_GRAPH: bool = False VLLM_XPU_FORCE_N_CONTIG_WEIGHT: bool = False VLLM_XPU_USE_SAMPLER_KERNEL: bool = True @@ -2129,10 +2128,6 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS": lambda: bool( int(os.getenv("VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS", "1")) ), - # NIXL EP environment variables - "VLLM_NIXL_EP_MAX_NUM_RANKS": lambda: int( - os.getenv("VLLM_NIXL_EP_MAX_NUM_RANKS", "32") - ), # Whether enable XPU graph on Intel GPU "VLLM_XPU_ENABLE_XPU_GRAPH": lambda: bool( int(os.getenv("VLLM_XPU_ENABLE_XPU_GRAPH", "0")) diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 7861345453..5019867d53 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -8,6 +8,7 @@ import torch from vllm.config import get_current_vllm_config from vllm.distributed import ( + GroupCoordinator, get_ep_group, ) from vllm.logger import init_logger @@ -104,16 +105,12 @@ if current_platform.is_cuda_alike(): ) -def get_ep_all2all_manager(eep_stage: bool = False) -> Any: - if eep_stage: - from vllm.distributed.elastic_ep.standby_state import get_standby_ep_group - - ep_group = get_standby_ep_group() - assert ep_group is not None - device_communicator = ep_group.device_communicator - else: - device_communicator = get_ep_group().device_communicator - +def get_ep_all2all_manager( + ep_group: GroupCoordinator | None = None, +) -> Any: + if ep_group is None: + ep_group = get_ep_group() + device_communicator = ep_group.device_communicator assert device_communicator is not None all2all_manager = device_communicator.all2all_manager assert all2all_manager is not None @@ -168,7 +165,7 @@ def maybe_make_prepare_finalize( routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, allow_new_interface: bool = False, use_monolithic: bool = False, - eep_stage: bool = False, + all2all_manager: Any | None = None, ) -> FusedMoEPrepareAndFinalize | None: if not moe.moe_parallel_config.use_all2all_kernels: if not allow_new_interface: @@ -190,7 +187,8 @@ def maybe_make_prepare_finalize( "Detected DP deployment with no --enable-expert-parallel. " "Falling back to AllGather+ReduceScatter dispatch/combine." ) - all2all_manager = get_ep_all2all_manager(eep_stage) + if all2all_manager is None: + all2all_manager = get_ep_all2all_manager() return make_moe_prepare_and_finalize_naive_dp_ep( is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel, num_dispatchers=all2all_manager.world_size, @@ -199,7 +197,8 @@ def maybe_make_prepare_finalize( else: return make_moe_prepare_and_finalize_no_dp_ep(use_monolithic) - all2all_manager = get_ep_all2all_manager(eep_stage) + if all2all_manager is None: + all2all_manager = get_ep_all2all_manager() prepare_finalize: FusedMoEPrepareAndFinalize | None = None @@ -361,10 +360,7 @@ def maybe_make_prepare_finalize( all_to_all_args = dict( max_num_tokens_per_dp_rank=moe.max_num_tokens, token_hidden_size=moe.hidden_dim, - num_ep_ranks=all2all_manager.world_size, - num_global_experts=moe.num_experts, - num_local_experts=moe.num_experts // all2all_manager.world_size, - stage=eep_stage, + num_local_experts=moe.num_local_experts, ) handle = all2all_manager.get_handle(all_to_all_args) @@ -378,7 +374,8 @@ def maybe_make_prepare_finalize( prepare_finalize = NixlEPPrepareAndFinalize( handle, max_tokens_per_rank=moe.max_num_tokens, - num_dispatchers=all2all_manager.world_size, + num_dispatchers=all2all_manager.max_num_ep_ranks, + expert_capacity=(moe.num_local_experts * all2all_manager.max_num_ep_ranks), use_fp8_dispatch=use_fp8_dispatch, global_to_physical=global_to_physical, physical_to_global=physical_to_global, diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 8736b1393d..4c7760d007 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1306,6 +1306,7 @@ class FusedMoEConfig: moe_backend: MoEBackend = "auto" max_num_tokens: int = SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP + elastic_ep_max_dp_size: int | None = None has_bias: bool = False is_lora_enabled: bool = False diff --git a/vllm/model_executor/layers/fused_moe/eep_reconfigure.py b/vllm/model_executor/layers/fused_moe/eep_reconfigure.py index 2c6c3f6aa4..9b6a1da0f9 100644 --- a/vllm/model_executor/layers/fused_moe/eep_reconfigure.py +++ b/vllm/model_executor/layers/fused_moe/eep_reconfigure.py @@ -60,6 +60,8 @@ def _make_eep_experts( def make_eep_staged_quant_method( module: "MoERunner", moe_config: FusedMoEConfig, + *, + all2all_manager: Any, ) -> FusedMoEMethodBase | None: quant_method = module._quant_method if not quant_method.supports_internal_mk: @@ -93,7 +95,7 @@ def make_eep_staged_quant_method( routing_tables=None, allow_new_interface=True, use_monolithic=quant_method.is_monolithic, - eep_stage=True, + all2all_manager=all2all_manager, ) assert prepare_finalize is not None assert isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeModular) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 4b633dd99b..a9c1a66d87 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -348,6 +348,7 @@ def FusedMoEFactory( moe_backend=vllm_config.kernel_config.moe_backend, router_logits_dtype=router_logits_dtype, max_num_tokens=max_num_batched_tokens, + elastic_ep_max_dp_size=vllm_config.parallel_config.elastic_ep_max_dp_size, has_bias=has_bias, is_lora_enabled=vllm_config.lora_config is not None, activation=moe_activation, diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 778b95bf47..120c9e9857 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -252,13 +252,6 @@ class FusedMoEPrepareAndFinalize(ABC): """ return False - def on_commit(self) -> None: - """ - Runs after this prepare/finalize has been committed to the active - MoE kernel. - """ - return - # TODO: pass FusedMoEParallelConfig in as ctor parameter? class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize): diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py index fd7d45655d..b2c85ce0c0 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py @@ -6,8 +6,6 @@ import nixl_ep import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm.distributed import get_ep_group -from vllm.distributed.device_communicators.all2all import NixlEPAll2AllManager from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( @@ -79,6 +77,7 @@ class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): buffer: nixl_ep.Buffer, max_tokens_per_rank: int, num_dispatchers: int, + expert_capacity: int, use_fp8_dispatch: bool = False, global_to_physical: torch.Tensor | None = None, physical_to_global: torch.Tensor | None = None, @@ -94,6 +93,7 @@ class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): # combine function. self.handles: list[tuple | None] = [None, None] self.num_dispatchers_ = num_dispatchers + self.expert_capacity = expert_capacity topk_indices_dtype = self.topk_indices_dtype() @@ -141,17 +141,6 @@ class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): def max_num_tokens_per_rank(self) -> int | None: return self.max_tokens_per_rank - def on_commit(self) -> None: - device_communicator = get_ep_group().device_communicator - assert device_communicator is not None - all2all_manager = device_communicator.all2all_manager - assert isinstance(all2all_manager, NixlEPAll2AllManager) - # maybe_make_prepare_finalize(..., eep_stage=True) initializes self.buffer - # with get_handle(..., stage=True), which stages global NIXL state for the - # new config but leaves it inactive while the old config remains active. - # When EEP commit switches to this P/F, this P/F needs to commit that state. - all2all_manager.commit_staged_state() - def topk_indices_dtype(self) -> torch.dtype | None: return NIXL_EP_TOPK_INDICES_DTYPE @@ -267,7 +256,7 @@ class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): a1, dispatch_topk_ids, self.max_tokens_per_rank, - num_experts, + self.expert_capacity, use_fp8=self.use_fp8_dispatch, # round_scale needs to be set to dispatch in ue8m0 round_scale=self.use_ue8m0_dispatch, diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index 48e217b1a9..cfc3be3e8e 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -500,7 +500,8 @@ def fi_moe_largest_bucket(moe_config: "FusedMoEConfig") -> int: For a detailed explanation, see: `docs/serving/data_parallel_deployment.md` """ - return max(moe_config.max_num_tokens * moe_config.dp_size, 8192) + dp_size = moe_config.elastic_ep_max_dp_size or moe_config.dp_size + return max(moe_config.max_num_tokens * dp_size, 8192) @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) diff --git a/vllm/model_executor/warmup/flashinfer_autotune_cache.py b/vllm/model_executor/warmup/flashinfer_autotune_cache.py index 37ea4812fc..afc3ef3ac0 100644 --- a/vllm/model_executor/warmup/flashinfer_autotune_cache.py +++ b/vllm/model_executor/warmup/flashinfer_autotune_cache.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING import vllm.envs as envs if TYPE_CHECKING: + from vllm.distributed.parallel_state import GroupCoordinator from vllm.v1.worker.gpu_model_runner import GPUModelRunner @@ -40,6 +41,44 @@ def resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path: return output_dir / "autotune_configs.json" +def sync_flashinfer_autotune_cache( + runner: "GPUModelRunner", + group: "GroupCoordinator", +) -> None: + cache: bytes | str | None = None + if ( + group.rank_in_group == 0 + and runner.vllm_config.kernel_config.enable_flashinfer_autotune + ): + try: + from vllm.platforms import current_platform + from vllm.utils.flashinfer import has_flashinfer + + if has_flashinfer() and current_platform.has_device_capability(90): + from flashinfer.autotuner import AutoTuner + + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "autotune_configs.json" + AutoTuner.get().save_configs(str(path)) + cache = path.read_bytes() + except Exception as exc: + cache = f"{type(exc).__name__}: {exc}" + + cache = group.broadcast_object(cache) + if isinstance(cache, str): + raise RuntimeError(f"Failed to serialize FlashInfer autotune state: {cache}") + if cache is None or group.rank_in_group == 0: + return + + from flashinfer.autotuner import AutoTuner + + with tempfile.NamedTemporaryFile() as f: + f.write(cache) + f.flush() + if not AutoTuner.get().load_configs(f.name): + raise RuntimeError("FlashInfer autotune cache is incompatible") + + def write_flashinfer_autotune_cache(cache_path: Path, contents: bytes) -> None: cache_path.parent.mkdir(parents=True, exist_ok=True) fd, tmp_path = tempfile.mkstemp( diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index ffd10ab557..c354ba699c 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -1711,6 +1711,12 @@ class DPLBAsyncMPClient(DPAsyncMPClient): "Only ray DP backend supports scaling elastic EP" ) parallel_config = self.vllm_config.parallel_config + if new_data_parallel_size > parallel_config.elastic_ep_max_dp_size: + raise ValueError( + f"Cannot scale to data_parallel_size {new_data_parallel_size}; " + "--elastic-ep-max-dp-size is " + f"{parallel_config.elastic_ep_max_dp_size}." + ) num_experts = self.vllm_config.model_config.get_num_experts() num_physical_experts = ( num_experts + parallel_config.eplb_config.num_redundant_experts diff --git a/vllm/v1/worker/dp_utils.py b/vllm/v1/worker/dp_utils.py index 07bda70483..fe82ec1427 100644 --- a/vllm/v1/worker/dp_utils.py +++ b/vllm/v1/worker/dp_utils.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext +from contextvars import ContextVar import torch import torch.distributed as dist @@ -18,6 +19,22 @@ from vllm.v1.worker.ubatch_utils import ( logger = init_logger(__name__) +_skip_dp_coordination = ContextVar("skip_dp_coordination", default=False) + + +@contextmanager +def skip_dp_coordination(): + """Run without coordinating DP metadata with other DP ranks.""" + token = _skip_dp_coordination.set(True) + try: + yield + finally: + _skip_dp_coordination.reset(token) + + +def should_skip_dp_coordination() -> bool: + return _skip_dp_coordination.get() + def _get_device_and_group(parallel_config: ParallelConfig): # Use the actual device assigned to the DP group, not just the device type @@ -223,6 +240,19 @@ def coordinate_batch_across_dp( if num_tokens_padded is None: num_tokens_padded = num_tokens_unpadded + if should_skip_dp_coordination(): + should_ubatch = should_attempt_ubatching and not is_last_ubatch_empty( + num_tokens_unpadded, + num_tokens_padded, + parallel_config.num_ubatches, + ) + num_tokens_after_padding = torch.full( + (parallel_config.data_parallel_size,), + num_tokens_padded, + dtype=torch.int32, + ) + return should_ubatch, num_tokens_after_padding, cudagraph_mode + (should_ubatch, num_tokens_after_padding, synced_cudagraph_mode) = ( _synchronize_dp_ranks( num_tokens_unpadded, diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index dddf630cc3..b7c73789de 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -10,6 +10,7 @@ import torch.distributed as dist from vllm.config import ParallelConfig from vllm.config.compilation import CUDAGraphMode from vllm.distributed.parallel_state import get_dp_group +from vllm.v1.worker.dp_utils import should_skip_dp_coordination from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, CudaGraphManager, @@ -71,7 +72,10 @@ def sync_cudagraph_and_dp_padding( tensor[3][dp_rank] = max_query_len or -1 # (-1 means None) tensor[4][dp_rank] = int(allow_ubatching) tensor[5][dp_rank] = num_reqs - dist.all_reduce(tensor, group=group) + if should_skip_dp_coordination(): + tensor[:] = tensor[:, dp_rank, None].clone() + else: + dist.all_reduce(tensor, group=group) num_tokens_across_dp = tensor[0] cg_mode_across_dp = tensor[1]