iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Elastic EP] Reuse CUDA graphs across reconfiguration (#54985)

Browse Source
This commit is contained in:
Itay Alroy
2026-09-12 10:26:39 -04:00
committed by GitHub
parent 13e221f830
commit 658c8131c7
22 changed files with 372 additions and 270 deletions
Download Patch File Download Diff File Expand all files Collapse all files
requirements/kv_connectors.txt
+1 -1
View File
@@ -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
tests/distributed/test_elastic_ep.py
+33 -65
View File
@@ -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}")
tests/distributed/test_eplb_utils.py
+11 -4
View File
@@ -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)
vllm/config/parallel.py
+17 -1
View File
@@ -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":
vllm/distributed/device_communicators/all2all.py
+42 -22
View File
@@ -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]
vllm/distributed/device_communicators/base_device_communicator.py
+12
View File
@@ -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:
"""
vllm/distributed/elastic_ep/elastic_execute.py
+77 -116
View File
@@ -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()
vllm/distributed/elastic_ep/elastic_state.py
+1 -1
View File
@@ -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(
vllm/distributed/eplb/eplb_state.py
+66 -12
View File
@@ -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)
vllm/engine/arg_utils.py
+6
View File
@@ -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,
vllm/envs.py
-5
View File
@@ -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"))
vllm/model_executor/layers/fused_moe/all2all_utils.py
+15 -18
View File
@@ -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,
vllm/model_executor/layers/fused_moe/config.py
+1
View File
@@ -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
vllm/model_executor/layers/fused_moe/eep_reconfigure.py
+3 -1
View File
@@ -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)
vllm/model_executor/layers/fused_moe/layer.py
+1
View File
@@ -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,
vllm/model_executor/layers/fused_moe/modular_kernel.py
-7
View File
@@ -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):
vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py
+3 -14
View File
@@ -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,
vllm/model_executor/layers/fused_moe/utils.py
+2 -1
View File
@@ -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)
vllm/model_executor/warmup/flashinfer_autotune_cache.py
+39
View File
@@ -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(
vllm/v1/engine/core_client.py
+6
View File
@@ -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
vllm/v1/worker/dp_utils.py
+31 -1
View File
@@ -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,
vllm/v1/worker/gpu/dp_utils.py
+5 -1
View File
@@ -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]