@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 658c8131c7

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

Itay Alroy committed Sep 12, 2026 at 17:26 UTC 658c8131c73194f812cd80d493e0409475b94235

22 files changed +372 -270

requirements/kv_connectors.txt

+1 -1

@@ -2,7 +2,7 @@ lmcache >= 0.3.9

2 # CuPy 14.1.0 imports pytest from cupy.testing._random, which breaks runtime

3 # images. 14.1.1 fixes the root cause, so only 14.1.0 is excluded.

4 cupy-cuda13x != 14.1.0

5 -nixl == 1.3.2

5 +nixl == 1.4.1

6 # CUDA 12 build. On the CUDA 13 image install-kv-connectors.sh swaps this for

7 # the mooncake-transfer-engine-cuda13 variant of the same version.

8 mooncake-transfer-engine >= 0.3.12

tests/distributed/test_elastic_ep.py

+33 -65

@@ -27,7 +27,7 @@ MODEL_NAME = "deepseek-ai/DeepSeek-V2-Lite-Chat"

27 NUM_GSM8K_QUESTIONS = 256

28 EXPECTED_ACCURACY = 0.58

29 ACCURACY_TOL = 0.08

30 -MAX_NUM_SEQS = 32

30 +DEFAULT_ALL2ALL_BACKEND = "allgather_reducescatter"

33 def _send_scale_command(server: RemoteOpenAIServer, new_dp_size: int) -> bool:

@@ -161,7 +161,12 @@ def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float:

161 return accuracy

162

163

164 -def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str]:

164 +def _base_serve_args(

165 + dp_size: int = 2,

166 + max_dp_size: int = 4,

167 + enforce_eager: bool = False,

168 + all2all_backend: str = DEFAULT_ALL2ALL_BACKEND,

169 +) -> list[str]:

170 args = [

171 "--trust-remote-code",

172 "--tensor-parallel-size",

@@ -170,12 +175,12 @@ def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str]

175 "0.8",

176 "--max-model-len",

177 "4096",

173 - "--max-num-seqs",

174 - str(MAX_NUM_SEQS),

178 "--enable-expert-parallel",

179 "--all2all-backend",

177 - "allgather_reducescatter",

180 + all2all_backend,

181 "--enable-elastic-ep",

182 + "--elastic-ep-max-dp-size",

183 + str(max_dp_size),

184 "--enable-eplb",

185 "--eplb-config.num_redundant_experts",

186 "0",

@@ -205,28 +210,43 @@ def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str]

210

211

212 @pytest.mark.parametrize(

208 - ("enforce_eager", "traffic_mode"),

213 + ("enforce_eager", "traffic_mode", "all2all_backend"),

214 [

210 - pytest.param(True, "none", id="enforce_eager_none"),

211 - pytest.param(True, "light", id="enforce_eager_light"),

212 - pytest.param(True, "heavy", id="enforce_eager_heavy"),

213 - pytest.param(False, "heavy", id="cuda_graphs_heavy"),

215 + pytest.param(True, "none", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_none"),

216 + pytest.param(True, "light", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_light"),

217 + pytest.param(True, "heavy", DEFAULT_ALL2ALL_BACKEND, id="enforce_eager_heavy"),

218 + pytest.param(

219 + False,

220 + "heavy",

221 + DEFAULT_ALL2ALL_BACKEND,

222 + id="cuda_graphs_heavy",

223 + ),

224 + pytest.param(False, "heavy", "nixl_ep", id="cuda_graphs_heavy_nixl_ep"),

225 ],

226 )

227 @multi_gpu_test(num_gpus=4)

217 -def test_elastic_ep_scaling(enforce_eager: bool, traffic_mode: str):

228 +def test_elastic_ep_scaling(

229 + enforce_eager: bool, traffic_mode: str, all2all_backend: str

230 +):

231 from vllm.distributed.eplb.eplb_communicator import has_nixl

232 + from vllm.utils.import_utils import has_nixl_ep

233

234 if not has_nixl():

235 pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")

236 + if all2all_backend == "nixl_ep" and not has_nixl_ep():

237 + pytest.skip("NIXL EP is not installed")

238

239 initial_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_INITIAL_DP", "2"))

240 target_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_TARGET_DP", "4"))

241 assert target_dp_size > initial_dp_size

226 - vllm_serve_args = _base_serve_args(initial_dp_size, enforce_eager)

242 + vllm_serve_args = _base_serve_args(

243 + initial_dp_size, target_dp_size, enforce_eager, all2all_backend

244 + )

245

246 with RemoteOpenAIServer(

229 - MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200

247 + MODEL_NAME,

248 + vllm_serve_args,

249 + max_wait_seconds=1200,

250 ) as server:

251 initial_accuracy = _run_gsm8k_eval(server, "Initial")

252

@@ -255,55 +275,3 @@ def test_elastic_ep_scaling(enforce_eager: bool, traffic_mode: str):

275 f"(diff: {scale_down_accuracy - initial_accuracy:+.3f})"

276 )

277 print(f" Tolerance: {ACCURACY_TOL:.3f}")

258 -

259 -

260 -@multi_gpu_test(num_gpus=4)

261 -def test_elastic_ep_scaling_uneven():

262 - """Test scale up with uneven worker distribution.

263 -

264 - This tests the case where num_new_workers % old_dp_size != 0,

265 - specifically 2 -> 3 where remainder = 1 % 2 = 1.

266 - This exercises the remainder handling in sender-receiver pairing.

267 - """

268 - from vllm.distributed.eplb.eplb_communicator import has_nixl

269 -

270 - if not has_nixl():

271 - pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")

272 -

273 - vllm_serve_args = _base_serve_args()

274 -

275 - with RemoteOpenAIServer(

276 - MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200

277 - ) as server:

278 - initial_accuracy = _run_gsm8k_eval(server, "Initial (2 GPUs)")

279 -

280 - # Scale 2 -> 3: This has remainder = 1 % 2 = 1

281 - # Tests uneven sender-receiver pairing

282 - assert _send_scale_command(server, 3)

283 - scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (3 GPUs)")

284 -

285 - assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, (

286 - f"Scale up accuracy {scale_up_accuracy:.3f} dropped more than "

287 - f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}"

288 - )

289 -

290 - # Scale back down to 2

291 - assert _send_scale_command(server, 2)

292 - scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)")

293 -

294 - assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, (

295 - f"Scale down accuracy {scale_down_accuracy:.3f} dropped more than "

296 - f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}"

297 - )

298 -

299 - print("\nAccuracy Summary (Uneven Scaling):")

300 - print(f" Initial: {initial_accuracy:.3f}")

301 - print(

302 - f" Scale up: {scale_up_accuracy:.3f} "

303 - f"(diff: {scale_up_accuracy - initial_accuracy:+.3f})"

304 - )

305 - print(

306 - f" Scale down: {scale_down_accuracy:.3f} "

307 - f"(diff: {scale_down_accuracy - initial_accuracy:+.3f})"

308 - )

309 - print(f" Tolerance: {ACCURACY_TOL:.3f}")

tests/distributed/test_eplb_utils.py

+11 -4

@@ -16,10 +16,14 @@ def _make_model_state(

16 phy2log: torch.Tensor,

17 log2phy: torch.Tensor,

18 logcnt: torch.Tensor,

19 + phy2log_storage: torch.Tensor | None = None,

20 ) -> MagicMock:

20 - """Build a minimal EplbModelState mock with only the three map tensors."""

21 + """Build a minimal EplbModelState mock with its map tensors."""

22 state = MagicMock()

23 state.physical_to_logical_map = phy2log

24 + state.physical_to_logical_map_buffer = (

25 + phy2log if phy2log_storage is None else phy2log_storage

26 + )

27 state.logical_to_physical_map = log2phy

28 state.logical_replica_count = logcnt

29 return state

@@ -27,15 +31,17 @@ def _make_model_state(

32 def test_commit_eplb_maps_shape_change():

33 """

30 - The normal path copies the physical_to_logical map in-place. When the number of

31 - physical experts changes, the old map should be replaced entirely.

34 + When the number of physical experts changes, resize the active map within its

35 + preallocated storage.

36 """

37 num_layers, num_logical, num_physical = 2, 4, 6

38 max_replicas = 3

40 # Build current state tensors

41 + storage = torch.full((num_layers, num_physical + 2), -1, dtype=torch.long)

42 model_state = _make_model_state(

38 - phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),

43 + phy2log=storage[:, :num_physical],

44 + phy2log_storage=storage,

45 log2phy=torch.full(

46 (num_layers, num_logical, max_replicas), -1, dtype=torch.long

47 ),

@@ -54,6 +60,7 @@ def test_commit_eplb_maps_shape_change():

60 # Check that the number of physical experts has been updated and that the values

61 # match

62 assert model_state.physical_to_logical_map.shape[1] == num_physical + 2

63 + assert model_state.physical_to_logical_map.data_ptr() == storage.data_ptr()

64 assert torch.equal(model_state.physical_to_logical_map, new_phy2log_larger)

vllm/config/parallel.py

+17 -1

@@ -216,6 +216,8 @@ class ParallelConfig:

216

217 enable_elastic_ep: bool = False

218 """Enable elastic expert parallelism with stateless NCCL groups for DP/EP."""

219 + elastic_ep_max_dp_size: int = Field(default=None, ge=1) # type: ignore[assignment]

220 + """Maximum data parallel size supported by elastic expert parallelism."""

221

222 enable_dbo: bool = False

223 """Enable dual batch overlap for the model executor."""

@@ -435,7 +437,9 @@ class ParallelConfig:

437 )

438 """The configurations for fault tolerance."""

439

438 - @field_validator("disable_nccl_for_dp_synchronization", mode="wrap")

440 + @field_validator(

441 + "disable_nccl_for_dp_synchronization", "elastic_ep_max_dp_size", mode="wrap"

442 + )

443 @classmethod

444 def _skip_none_validation(cls, value: Any, handler: Callable) -> Any:

445 """Skip validation if the value is `None` when initialisation is delayed."""

@@ -941,6 +945,18 @@ class ParallelConfig:

945 " for dense models."

946 )

947

948 + max_dp_size = self.elastic_ep_max_dp_size

949 + self.elastic_ep_max_dp_size = (

950 + max_dp_size

951 + if self.enable_elastic_ep and max_dp_size is not None

952 + else self.data_parallel_size

953 + )

954 + if self.elastic_ep_max_dp_size < self.data_parallel_size:

955 + raise ValueError(

956 + "--elastic-ep-max-dp-size must be greater than or equal to "

957 + f"the initial data_parallel_size ({self.data_parallel_size})."

958 + )

959 +

960 self.data_parallel_index = self.data_parallel_rank

961

962 if self.distributed_executor_backend == "external_launcher":

vllm/distributed/device_communicators/all2all.py

+42 -22

@@ -1,6 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import threading

4 +from collections.abc import Iterator

5 +from contextlib import contextmanager

6 from dataclasses import dataclass

7 from typing import Any

@@ -377,6 +379,7 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):

379 @dataclass

380 class _NixlEPBufferState:

381 buffer: Any

382 + max_num_ep_ranks: int

383 connected_ep_size: int

384 active_ep_size: int

385

@@ -402,7 +405,10 @@ class NixlEPAll2AllManager(All2AllManagerBase):

405 super().__init__(cpu_group, tcp_store_group)

406 self.support_fault_tolerance = True

407

405 - self.max_num_ep_ranks = envs.VLLM_NIXL_EP_MAX_NUM_RANKS

408 + @property

409 + def max_num_ep_ranks(self) -> int:

410 + assert NixlEPAll2AllManager._buffer is not None

411 + return NixlEPAll2AllManager._buffer.max_num_ep_ranks

412

413 def _init_buffer(

414 self,

@@ -412,11 +418,19 @@ class NixlEPAll2AllManager(All2AllManagerBase):

418 ) -> None:

419 from nixl_ep import Buffer # type: ignore[import-not-found]

420

415 - max_num_global_experts = self.max_num_ep_ranks * num_experts_per_rank

421 + parallel_config = get_current_vllm_config().parallel_config

422 + max_num_ep_ranks = (

423 + self.world_size

424 + * parallel_config.elastic_ep_max_dp_size

425 + // parallel_config.data_parallel_size

426 + if parallel_config.enable_elastic_ep

427 + else self.world_size

428 + )

429 + max_num_global_experts = max_num_ep_ranks * num_experts_per_rank

430 num_rdma_bytes = Buffer.get_rdma_size_hint(

431 num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank,

432 hidden=token_hidden_size,

419 - num_ranks=self.max_num_ep_ranks,

433 + num_ranks=max_num_ep_ranks,

434 num_experts=max_num_global_experts,

435 )

436 assert NixlEPAll2AllManager._buffer is None, (

@@ -427,7 +441,7 @@ class NixlEPAll2AllManager(All2AllManagerBase):

441 tcp_store_group=self.tcp_store_group.store,

442 )

443 buffer.update_memory_buffers(

430 - num_ranks=self.max_num_ep_ranks,

444 + num_ranks=max_num_ep_ranks,

445 num_experts_per_rank=num_experts_per_rank,

446 num_rdma_bytes=num_rdma_bytes,

447 )

@@ -435,6 +449,7 @@ class NixlEPAll2AllManager(All2AllManagerBase):

449 buffer.connect_ranks(ranks_to_connect)

450 NixlEPAll2AllManager._buffer = _NixlEPBufferState(

451 buffer=buffer,

452 + max_num_ep_ranks=max_num_ep_ranks,

453 connected_ep_size=self.world_size,

454 active_ep_size=self.world_size,

455 )

@@ -474,6 +489,7 @@ class NixlEPAll2AllManager(All2AllManagerBase):

489

490 for rank in range(state.active_ep_size, target_ep_size):

491 state.buffer.update_mask_buffer(rank, mask=False)

492 + torch.accelerator.synchronize()

493 state.active_ep_size = target_ep_size

494

495 def _stage_ep_size(self) -> None:

@@ -486,8 +502,12 @@ class NixlEPAll2AllManager(All2AllManagerBase):

502 if target_ep_size > state.connected_ep_size:

503 self._connect_to_ep_size(target_ep_size, make_active=False)

504

489 - def commit_staged_state(self) -> None:

490 - """Commit staged NIXL EP state to the active communication set."""

505 + def stage_ep_size(self) -> None:

506 + with NixlEPAll2AllManager._lock:

507 + self._stage_ep_size()

508 +

509 + def commit_ep_size(self) -> None:

510 + """Commit the staged EP size to the active communication set."""

511 with NixlEPAll2AllManager._lock:

512 assert NixlEPAll2AllManager._buffer is not None

513 state = NixlEPAll2AllManager._buffer

@@ -500,31 +520,31 @@ class NixlEPAll2AllManager(All2AllManagerBase):

520

521 self._unmask_connected_ranks(target_ep_size)

522

503 - def _ensure_ep_size(self, *, stage: bool) -> None:

504 - if stage:

505 - self._stage_ep_size()

506 - else:

507 - self.commit_staged_state()

523 + @contextmanager

524 + def mask_remote_ranks(self) -> Iterator[None]:

525 + peers = [rank for rank in range(self.world_size) if rank != self.rank]

526 + assert NixlEPAll2AllManager._buffer is not None

527 + buffer = NixlEPAll2AllManager._buffer.buffer

528 + for rank in peers:

529 + buffer.update_mask_buffer(rank, mask=True)

530 + torch.accelerator.synchronize()

531 + try:

532 + yield

533 + finally:

534 + for rank in peers:

535 + buffer.update_mask_buffer(rank, mask=False)

536 + torch.accelerator.synchronize()

537

538 def get_handle(self, kwargs):

539 with NixlEPAll2AllManager._lock:

511 - stage = bool(kwargs.get("stage", False))

540 state = NixlEPAll2AllManager._buffer

541 if state is None:

514 - assert not stage, (

515 - "NIXL EP staged initialization requires an existing buffer"

516 - )

542 max_num_tokens_per_dp_rank = kwargs["max_num_tokens_per_dp_rank"]

518 - num_experts_per_rank = (

519 - kwargs["num_global_experts"] // kwargs["num_ep_ranks"]

520 - )

543 self._init_buffer(

544 max_num_tokens_per_dp_rank=max_num_tokens_per_dp_rank,

545 token_hidden_size=kwargs["token_hidden_size"],

524 - num_experts_per_rank=num_experts_per_rank,

546 + num_experts_per_rank=kwargs["num_local_experts"],

547 )

526 - else:

527 - self._ensure_ep_size(stage=stage)

548

549 assert NixlEPAll2AllManager._buffer is not None

550 handle = NixlEPAll2AllManager._buffer.buffer

@@ -564,7 +584,7 @@ class NixlEPAll2AllManager(All2AllManagerBase):

584 assert state is not None

585 if NixlEPAll2AllManager._mask is None:

586 NixlEPAll2AllManager._mask = torch.zeros(

567 - self.max_num_ep_ranks, device="cuda", dtype=torch.int32

587 + state.max_num_ep_ranks, device="cuda", dtype=torch.int32

588 )

589 state.buffer.query_mask_buffer(NixlEPAll2AllManager._mask)

590 return NixlEPAll2AllManager._mask[: state.active_ep_size]

vllm/distributed/device_communicators/base_device_communicator.py

+12

@@ -1,6 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import threading

4 +from collections.abc import Iterator

5 +from contextlib import contextmanager

6 from weakref import WeakValueDictionary

8 import torch

@@ -157,6 +159,16 @@ class All2AllManagerBase:

159 def destroy(self):

160 pass

161

162 + def stage_ep_size(self) -> None:

163 + pass

164 +

165 + def commit_ep_size(self) -> None:

166 + pass

167 +

168 + @contextmanager

169 + def mask_remote_ranks(self) -> Iterator[None]:

170 + yield

171 +

172

173 class DeviceCommunicatorBase:

174 """

vllm/distributed/elastic_ep/elastic_execute.py

+77 -116

@@ -2,34 +2,32 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import gc

4 import weakref

5 -from collections.abc import Iterable, Sequence

5 +from collections.abc import Iterable, Iterator, Sequence

6 from concurrent.futures import Future, ThreadPoolExecutor

7 +from contextlib import contextmanager, nullcontext

8 from dataclasses import replace

9 from typing import TYPE_CHECKING

11 import torch

12 import torch.nn as nn

12 -import torch.nn.functional as F

13 from torch.distributed import P2POp

15 -from vllm.compilation.counter import compilation_counter

15 from vllm.compilation.cuda_graph import CUDAGraphWrapper

16 from vllm.compilation.wrapper import reset_compile_wrapper

18 -from vllm.config import (

19 - CompilationMode,

20 - set_current_vllm_config,

21 -)

17 +from vllm.config import ParallelConfig, set_current_vllm_config

18 from vllm.distributed import (

19 get_dp_group,

20 get_ep_group,

21 get_pcp_group,

22 get_tp_group,

23 + get_world_group,

24 )

25 from vllm.distributed.elastic_ep.standby_state import (

26 create_standby_groups,

27 get_standby_dp_group,

28 get_standby_ep_group,

29 get_standby_eplb_group,

30 + get_standby_world_group,

31 pop_standby_groups,

32 )

33 from vllm.distributed.eplb.eplb_communicator import (

@@ -41,13 +39,20 @@ from vllm.distributed.parallel_state import (

39 )

40 from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator

41 from vllm.logger import init_logger

42 +from vllm.model_executor.layers.fused_moe.all2all_utils import (

43 + get_ep_all2all_manager,

44 +)

45 from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig

46 from vllm.model_executor.layers.fused_moe.eep_reconfigure import (

47 make_eep_staged_quant_method,

48 )

49 +from vllm.model_executor.warmup.flashinfer_autotune_cache import (

50 + sync_flashinfer_autotune_cache,

51 +)

52 from vllm.model_executor.warmup.kernel_warmup import kernel_warmup

53 from vllm.utils import is_moe_layer

54 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType

55 +from vllm.v1.worker.dp_utils import skip_dp_coordination

56 from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper

57 from vllm.v1.worker.workspace import lock_workspace, unlock_workspace

@@ -59,6 +64,15 @@ if TYPE_CHECKING:

64 )

67 +def can_reuse_fused_moe_kernel(parallel_config: ParallelConfig) -> bool:

68 + """Whether fused MoE kernels can be reused across EP rank changes.

69 +

70 + The EP backend must preserve its CUDA graph entries when ranks are connected,

71 + disconnected, masked, or unmasked.

72 + """

73 + return parallel_config.all2all_backend == "nixl_ep"

74 +

75 +

76 def batch_transfer_weights(

77 model: nn.Module,

78 is_sender: bool,

@@ -105,31 +119,20 @@ def batch_transfer_weights(

119

120 def broadcast_expert_mapping(

121 physical_to_logical: torch.Tensor | None,

108 - num_local_physical_experts: int | None,

109 - num_logical_experts: int | None,

122 dp_group: StatelessGroupCoordinator,

123 device: torch.device,

124 src_rank: int = 0,

113 -) -> tuple[torch.Tensor, int, int]:

125 +) -> torch.Tensor:

126 if dp_group.rank_in_group == src_rank:

127 assert physical_to_logical is not None

116 - assert num_local_physical_experts is not None

117 - assert num_logical_experts is not None

128 assert physical_to_logical.dtype == torch.int64

129 shape_tensor = torch.tensor(

130 list(physical_to_logical.shape), dtype=torch.int64, device="cpu"

131 )

122 - metadata_tensor = torch.tensor(

123 - [num_local_physical_experts, num_logical_experts],

124 - dtype=torch.int64,

125 - device="cpu",

126 - )

132 else:

133 shape_tensor = torch.empty(2, dtype=torch.int64, device="cpu")

129 - metadata_tensor = torch.empty(2, dtype=torch.int64, device="cpu")

134

135 shape_tensor = dp_group.tcp_store_group.broadcast(shape_tensor, src_rank)

132 - metadata_tensor = dp_group.tcp_store_group.broadcast(metadata_tensor, src_rank)

136

137 if dp_group.rank_in_group != src_rank:

138 assert device is not None

@@ -140,11 +143,7 @@ def broadcast_expert_mapping(

143 )

144

145 assert physical_to_logical is not None

143 - physical_to_logical = dp_group.broadcast(physical_to_logical, src_rank)

144 - num_local_physical_experts = int(metadata_tensor[0].item())

145 - num_logical_experts = int(metadata_tensor[1].item())

146 -

147 - return physical_to_logical, num_local_physical_experts, num_logical_experts

146 + return dp_group.broadcast(physical_to_logical, src_rank)

147

148

149 class ElasticEPScalingExecutor:

@@ -166,6 +165,21 @@ class ElasticEPScalingExecutor:

165 raise RuntimeError("Worker has been garbage collected")

166 return worker

167

168 + def _can_reuse_fused_moe_kernel(self) -> bool:

169 + return can_reuse_fused_moe_kernel(

170 + self.worker.vllm_config.parallel_config,

171 + )

172 +

173 + @contextmanager

174 + def _disable_flashinfer_autotune(self) -> Iterator[None]:

175 + kernel_config = self.worker.vllm_config.kernel_config

176 + enabled = kernel_config.enable_flashinfer_autotune

177 + kernel_config.enable_flashinfer_autotune = False

178 + try:

179 + yield

180 + finally:

181 + kernel_config.enable_flashinfer_autotune = enabled

182 +

183 def execute(self, execute_method: str, *args, **kwargs):

184 method = getattr(self, execute_method, None)

185 if method is None:

@@ -260,11 +274,20 @@ class ElasticEPScalingExecutor:

274 use_all2all=use_all2all,

275 enable_eplb=parallel_config.enable_eplb,

276 )

263 - self.stage_standby_moe_quant_methods()

277 + standby_ep_group = get_standby_ep_group()

278 + assert standby_ep_group is not None

279 + all2all_manager = get_ep_all2all_manager(standby_ep_group)

280 + all2all_manager.stage_ep_size()

281 + if not self._can_reuse_fused_moe_kernel():

282 + self.stage_standby_moe_quant_methods(all2all_manager)

283 self._prepare_eplb_communicator(get_standby_eplb_group())

284 if new_dp_size > old_dp_size:

285 self.transfer_weights(old_dp_size, new_dp_size)

267 - self._warm_target_groups(get_standby_dp_group(), get_standby_ep_group())

286 + self._warm_target_groups(get_standby_dp_group(), standby_ep_group)

287 + if new_dp_size > old_dp_size and self._can_reuse_fused_moe_kernel():

288 + target_world_group = get_standby_world_group()

289 + assert target_world_group is not None

290 + sync_flashinfer_autotune_cache(self.worker.model_runner, target_world_group)

291

292 def _prepare_eplb_communicator(self, eplb_group) -> None:

293 assert eplb_group is not None

@@ -338,14 +361,9 @@ class ElasticEPScalingExecutor:

361 assert eplb_state is not None

362 eplb_state.drain_async()

363 eplb_model_state = eplb_state.model_states[model_config.compute_hash()]

341 - physical_to_logical = eplb_model_state.physical_to_logical_map

342 - num_physical_experts = physical_to_logical.shape[1]

343 - num_local_physical_experts = num_physical_experts // get_ep_group().world_size

344 - num_logical_experts = eplb_model_state.logical_replica_count.shape[1]

364 + physical_to_logical = eplb_model_state.physical_to_logical_map_buffer

365 broadcast_expert_mapping(

366 physical_to_logical=physical_to_logical,

347 - num_local_physical_experts=num_local_physical_experts,

348 - num_logical_experts=num_logical_experts,

367 dp_group=standby_dp_group,

368 src_rank=0,

369 device=self.worker.device,

@@ -368,7 +386,7 @@ class ElasticEPScalingExecutor:

386 moe_parallel_config=moe_parallel_config,

387 )

388

371 - def stage_standby_moe_quant_methods(self) -> None:

389 + def stage_standby_moe_quant_methods(self, all2all_manager) -> None:

390 standby_dp_group = get_standby_dp_group()

391 standby_ep_group = get_standby_ep_group()

392 model = self.worker.model_runner.get_model()

@@ -383,6 +401,7 @@ class ElasticEPScalingExecutor:

401 standby_dp_group,

402 standby_ep_group,

403 ),

404 + all2all_manager=all2all_manager,

405 )

406 if staged_quant_method is not None:

407 self._staged_moe_quant_methods[module] = staged_quant_method

@@ -396,7 +415,6 @@ class ElasticEPScalingExecutor:

415 continue

416 assert staged_quant_method.moe_kernel is not None

417 module._replace_quant_method(staged_quant_method)

399 - staged_quant_method.moe_kernel.prepare_finalize.on_commit()

418 self._staged_moe_quant_methods.clear()

419

420 def _release_cuda_graphs(self) -> None:

@@ -405,7 +423,7 @@ class ElasticEPScalingExecutor:

423 wrapper.concrete_cudagraph_entries = {}

424

425 elif isinstance(self.worker.model_runner.model, UBatchWrapper):

408 - raise RuntimeError("DBO is not yet supported in elastic EP")

426 + self.worker.model_runner.model.clear_graphs()

427

428 torch.compiler.reset()

429 with set_current_vllm_config(self.worker.vllm_config):

@@ -427,10 +445,8 @@ class ElasticEPScalingExecutor:

445 self._wait_for_group_cleanup()

446

447 def switch_and_prepare(self) -> tuple[GroupCoordinator | None, ...]:

430 - old_dp_size = get_dp_group().world_size

431 - old_ep_size = get_ep_group().world_size

432 -

433 - self._release_cuda_graphs()

448 + if not self._can_reuse_fused_moe_kernel():

449 + self._release_cuda_graphs()

450 retired_groups = _replace_active_groups(**pop_standby_groups())

451

452 parallel_config = self.worker.vllm_config.parallel_config

@@ -487,47 +503,13 @@ class ElasticEPScalingExecutor:

503 parallel_config.eplb_config.num_redundant_experts = (

504 num_physical_experts - num_logical_experts

505 )

490 - old_physical_to_logical = eplb_model_state.physical_to_logical_map

491 - num_moe_layers = old_physical_to_logical.shape[0]

492 - num_local_experts = eplb_model_state.expert_load_pass.shape[1] // old_ep_size

493 - if new_dp_size > old_dp_size:

494 - expanded_physical_to_logical = torch.full(

495 - (num_moe_layers, num_local_experts * new_ep_size),

496 - -1,

497 - dtype=old_physical_to_logical.dtype,

498 - device=old_physical_to_logical.device,

499 - )

500 - expanded_physical_to_logical[:, : num_local_experts * old_ep_size] = (

501 - old_physical_to_logical

502 - )

503 - eplb_model_state.physical_to_logical_map = expanded_physical_to_logical

504 -

505 - old_num_physical_experts = eplb_model_state.expert_load_pass.shape[1]

506 - pad_size = num_physical_experts - old_num_physical_experts

507 - if new_dp_size > old_dp_size:

508 - assert pad_size > 0

509 - expanded_expert_load_pass = F.pad(

510 - eplb_model_state.expert_load_pass, (0, pad_size), value=0

511 - )

512 - expanded_expert_load_window = F.pad(

513 - eplb_model_state.expert_load_window, (0, pad_size), value=0

514 - )

515 - eplb_model_state.expert_load_pass = expanded_expert_load_pass

516 - eplb_model_state.expert_load_window = expanded_expert_load_window

517 - else:

518 - assert pad_size < 0

519 - eplb_model_state.expert_load_pass = eplb_model_state.expert_load_pass[

520 - :, :num_physical_experts

521 - ]

522 - eplb_model_state.expert_load_window = eplb_model_state.expert_load_window[

523 - :, :, :num_physical_experts

524 - ]

506 + eplb_state.reconfigure_physical_expert_slots(model_config, num_physical_experts)

507

508 model = self.worker.model_runner.get_model()

509 model.expert_weights = []

510 with set_current_vllm_config(self.worker.vllm_config):

511 model.set_eplb_state(

530 - eplb_model_state.expert_load_pass,

512 + eplb_model_state.expert_load_pass_buffer,

513 eplb_model_state.logical_to_physical_map,

514 eplb_model_state.logical_replica_count,

515 )

@@ -539,6 +521,7 @@ class ElasticEPScalingExecutor:

521 num_local_physical_experts=num_local_experts,

522 )

523 self._commit_staged_moe_quant_methods()

524 + get_ep_all2all_manager().commit_ep_size()

525 # Legacy modular methods need to be recreated for the new EP size.

526 for module in moe_modules:

527 if getattr(module._quant_method, "wraps_legacy_quant_method", False):

@@ -548,21 +531,6 @@ class ElasticEPScalingExecutor:

531 eplb_state.update_communicator(model_config, self._prepared_eplb_communicator)

532 self._prepared_eplb_communicator = None

533

551 - if (

552 - self.worker.vllm_config.compilation_config.mode

553 - == CompilationMode.STOCK_TORCH_COMPILE

554 - ):

555 - # NOTE(yongji): when using stock torch.compile,

556 - # torch.compile is triggered during GPUModelRunner's load_model()

557 - # TODO(yongji):check do we need to re-trigger torch.compile here?

558 - # any changes to the tensor shapes in execution should already

559 - # be handled internally by torch.compile.

560 - backend = self.worker.vllm_config.compilation_config.init_backend(

561 - self.worker.vllm_config

562 - )

563 - compilation_counter.stock_torch_compile_count += 1

564 - self.worker.model_runner.model.compile(fullgraph=True, backend=backend)

565 -

534 return retired_groups

535

536 def _perform_eplb_reshuffle(

@@ -608,7 +576,8 @@ class ElasticEPScalingExecutor:

576 else:

577 mapping = self.receive_expert_mapping()

578 self.worker.model_runner.setup_eplb_from_mapping(mapping)

611 - self.warm_and_capture()

579 + if not self._can_reuse_fused_moe_kernel():

580 + self.warm_and_capture()

581 self._perform_eplb_reshuffle(async_op=True)

582 if is_existing_worker:

583 self._start_group_cleanup(retired_groups)

@@ -619,7 +588,8 @@ class ElasticEPScalingExecutor:

588 self.switch_and_remove()

589 else:

590 retired_groups = self.switch_and_prepare()

622 - self.warm_and_capture()

591 + if not self._can_reuse_fused_moe_kernel():

592 + self.warm_and_capture()

593 self._start_group_cleanup(retired_groups)

594

595 def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None:

@@ -677,42 +647,26 @@ class ElasticEPScalingExecutor:

647 )

648 torch.accelerator.synchronize()

649 self._warm_target_groups(get_dp_group(), get_ep_group())

650 + if self._can_reuse_fused_moe_kernel():

651 + sync_flashinfer_autotune_cache(self.worker.model_runner, get_world_group())

652

653 def receive_expert_mapping(self) -> torch.Tensor:

654 dp_group = get_dp_group()

655 assert isinstance(dp_group, StatelessGroupCoordinator)

684 - physical_to_logical, num_local_physical_experts, _ = broadcast_expert_mapping(

656 + return broadcast_expert_mapping(

657 physical_to_logical=None,

686 - num_local_physical_experts=None,

687 - num_logical_experts=None,

658 dp_group=dp_group,

659 src_rank=0,

660 device=self.worker.device,

661 )

692 - num_moe_layers = physical_to_logical.shape[0]

693 - new_dp_size = get_dp_group().world_size

694 - tp_size = self.worker.vllm_config.parallel_config.tensor_parallel_size

695 - new_ep_size = new_dp_size * tp_size

696 - expanded_physical_to_logical = torch.full(

697 - (num_moe_layers, num_local_physical_experts * new_ep_size),

698 - -1,

699 - dtype=physical_to_logical.dtype,

700 - device=physical_to_logical.device,

701 - )

702 - expanded_physical_to_logical[:, : physical_to_logical.shape[1]] = (

703 - physical_to_logical

704 - )

705 - return expanded_physical_to_logical

662

707 - def warmup_local_kernels(self) -> None:

663 + def warmup_new_worker(self) -> None:

664 with set_current_vllm_config(self.worker.vllm_config):

665 kernel_warmup(self.worker, process_local_only=True)

666 + if self._can_reuse_fused_moe_kernel():

667 + self.warm_and_capture()

668

669 def warm_and_capture(self) -> None:

712 - # Must run on every DP sibling in lockstep: _dummy_run calls

713 - # coordinate_batch_across_dp whenever data_parallel_size > 1

714 - # (gpu_model_runner.py:3663), which deadlocks if any rank skips it.

715 -

670 # Save and clear block tables so the dummy MoE forward doesn't

671 # write dummy slot mappings into real KV-cache blocks.

672 multi_block_table = self.worker.model_runner.input_batch.block_table

@@ -735,8 +689,15 @@ class ElasticEPScalingExecutor:

689 # directly with skip_eplb=True so dummy routing doesn't pollute the

690 # just-rebalanced EPLB stats.

691 runner = self.worker.model_runner

738 - runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True)

739 - self.worker.compile_or_warm_up_model()

692 + all2all_manager = get_ep_all2all_manager()

693 + reuse_kernel = self._can_reuse_fused_moe_kernel()

694 + with (

695 + skip_dp_coordination() if reuse_kernel else nullcontext(),

696 + all2all_manager.mask_remote_ranks() if reuse_kernel else nullcontext(),

697 + self._disable_flashinfer_autotune() if reuse_kernel else nullcontext(),

698 + ):

699 + runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True)

700 + self.worker.compile_or_warm_up_model()

701

702 lock_workspace()

703

vllm/distributed/elastic_ep/elastic_state.py

+1 -1

@@ -204,7 +204,7 @@ class ElasticEPScalingState:

204 return True

205

206 elif state == ScaleUpNewEngineState.PREPARE:

207 - self._collective_rpc("elastic_ep_execute", args=("warmup_local_kernels",))

207 + self._collective_rpc("elastic_ep_execute", args=("warmup_new_worker",))

208 self._mark_ready_for_switch()

209 tensor = torch.tensor([0, 0, 0], dtype=torch.int32, device="cpu")

210 torch.distributed.all_reduce(

vllm/distributed/eplb/eplb_state.py

+66 -12

@@ -121,6 +121,8 @@ class EplbModelState:

121 [0, 2, 0, 1, 0, 3]]

122 ```

123 """

124 + physical_to_logical_map_buffer: torch.Tensor

125 + """Maximum-capacity buffer backing ``physical_to_logical_map``."""

126 logical_to_physical_map: torch.Tensor

127 """

128 Mapping from logical experts to physical experts.

@@ -168,6 +170,8 @@ class EplbModelState:

170

171 Shape: (num_moe_layers, num_physical_experts)

172 """

173 + expert_load_pass_buffer: torch.Tensor

174 + """Maximum-capacity buffer backing ``expert_load_pass``."""

175 expert_load_window: torch.Tensor

176 """

177 A sliding window of expert load.

@@ -373,6 +377,19 @@ class EplbState:

377 physical_to_logical_map_list,

378 device=self.device,

379 )

380 + physical_expert_capacity = model.num_physical_experts

381 + if self.parallel_config.enable_elastic_ep:

382 + physical_expert_capacity = (

383 + model.num_local_physical_experts

384 + * self.parallel_config.elastic_ep_max_dp_size

385 + * get_ep_group().world_size

386 + // self.parallel_config.data_parallel_size

387 + )

388 + physical_to_logical_map = torch.nn.functional.pad(

389 + physical_to_logical_map,

390 + (0, physical_expert_capacity - model.num_physical_experts),

391 + value=-1,

392 + )

393 # Assuming 8 GPUs per node, this supports up to

394 # (1023 + 1) / 8 = 128 nodes for now.

395 # TODO(rui): make this configurable

@@ -407,6 +424,10 @@ class EplbState:

424 )

425 .contiguous()

426 )

427 + physical_to_logical_map_buffer = physical_to_logical_map

428 + physical_to_logical_map = physical_to_logical_map_buffer[

429 + :, : model.num_physical_experts

430 + ]

431 logical_to_physical_map = (

432 logical_to_physical_map.unsqueeze(0)

433 .expand(

@@ -425,11 +446,12 @@ class EplbState:

446 .contiguous()

447 )

448

428 - expert_load_pass = torch.zeros(

429 - (model.num_moe_layers, model.num_physical_experts),

449 + expert_load_pass_buffer = torch.zeros(

450 + (model.num_moe_layers, physical_expert_capacity),

451 dtype=torch.int32,

452 device=self.device,

453 )

454 + expert_load_pass = expert_load_pass_buffer[:, : model.num_physical_experts]

455 self.expert_load_window_size = self.parallel_config.eplb_config.window_size

456 expert_load_window = torch.zeros(

457 (

@@ -460,7 +482,7 @@ class EplbState:

482 ]

483

484 model.set_eplb_state(

463 - expert_load_pass,

485 + expert_load_pass_buffer,

486 logical_to_physical_map,

487 logical_replica_count,

488 )

@@ -479,9 +501,11 @@ class EplbState:

501

502 model_state = EplbModelState(

503 physical_to_logical_map=physical_to_logical_map,

504 + physical_to_logical_map_buffer=physical_to_logical_map_buffer,

505 logical_to_physical_map=logical_to_physical_map,

506 logical_replica_count=logical_replica_count,

507 expert_load_pass=expert_load_pass,

508 + expert_load_pass_buffer=expert_load_pass_buffer,

509 expert_load_window=expert_load_window,

510 model_name=model_config.model,

511 model=model,

@@ -1075,10 +1099,14 @@ class EplbState:

1099 expanded_physical_to_logical: torch.Tensor,

1100 ) -> None:

1101 eplb_model_state = self.model_states[model_config.compute_hash()]

1078 - eplb_model_state.physical_to_logical_map.copy_(expanded_physical_to_logical)

1102 + eplb_model_state.physical_to_logical_map_buffer.copy_(

1103 + expanded_physical_to_logical

1104 + )

1105 + eplb_model_state.expert_load_pass.zero_()

1106 + physical_to_logical_map = eplb_model_state.physical_to_logical_map

1107

1108 (logical_to_physical_map_cpu, logical_replica_count_cpu) = compute_logical_maps(

1081 - expanded_physical_to_logical.cpu(),

1109 + physical_to_logical_map.cpu(),

1110 eplb_model_state.model.num_logical_experts,

1111 )

1112

@@ -1097,6 +1125,33 @@ class EplbState:

1125 eplb_model_state.logical_to_physical_map.copy_(logical_to_physical_map)

1126 eplb_model_state.logical_replica_count.copy_(logical_replica_count)

1127

1128 + def reconfigure_physical_expert_slots(

1129 + self,

1130 + model_config: ModelConfig,

1131 + num_physical_experts: int,

1132 + ) -> None:

1133 + """Replace physical_to_logical_map and expert_load_pass with views

1134 + covering the new active size.

1135 + """

1136 + model_state = self.model_states[model_config.compute_hash()]

1137 + old_num_physical_experts = model_state.model.num_physical_experts

1138 + first_slot, last_slot = sorted((old_num_physical_experts, num_physical_experts))

1139 + physical_to_logical_map_buffer = model_state.physical_to_logical_map_buffer

1140 + expert_load_pass_buffer = model_state.expert_load_pass_buffer

1141 + assert last_slot <= physical_to_logical_map_buffer.shape[1]

1142 + expert_slots = slice(first_slot, last_slot)

1143 + physical_to_logical_map_buffer[:, expert_slots].fill_(-1)

1144 + expert_load_pass_buffer[:, expert_slots].zero_()

1145 + model_state.physical_to_logical_map = physical_to_logical_map_buffer[

1146 + :, :num_physical_experts

1147 + ]

1148 + model_state.expert_load_pass = expert_load_pass_buffer[:, :num_physical_experts]

1149 +

1150 + pad_size = num_physical_experts - model_state.expert_load_window.shape[-1]

1151 + model_state.expert_load_window = torch.nn.functional.pad(

1152 + model_state.expert_load_window, (0, pad_size)

1153 + )

1154 +

1155 def create_communicator(

1156 self, model_config: ModelConfig, group_coordinator: GroupCoordinator

1157 ) -> EplbCommunicator:

@@ -1333,14 +1388,13 @@ def _commit_eplb_maps(

1388

1389 if PIN_MEMORY and src.is_cpu:

1390 src = src.new_empty(src.shape, pin_memory=True).copy_(src)

1336 - # Rare Case: When the number of physical experts has changed, discard the old

1337 - # physical to logical expert map and use the new one. This only happens when the

1338 - # number of GPUs available to vLLM changes while vLLM is running. Otherwise copy the

1339 - # new map into the old one.

1391 + # When the number of physical experts changes, refresh the active map view to

1392 + # cover the new size before copying it.

1393 if src.shape[1] != dst.shape[1]:

1341 - model_state.physical_to_logical_map = src.to(dst.device, non_blocking=True)

1342 - else:

1343 - dst.copy_(src, non_blocking=True)

1394 + dst = model_state.physical_to_logical_map_buffer[:, : src.shape[1]]

1395 + assert dst.shape == src.shape

1396 + model_state.physical_to_logical_map = dst

1397 + dst.copy_(src, non_blocking=True)

1398

1399 num_logical_experts = model_state.logical_to_physical_map.shape[1]

1400 new_logical, new_replica_count = compute_logical_maps(src, num_logical_experts)

vllm/engine/arg_utils.py

@@ -507,6 +507,7 @@ class EngineArgs:

507 linear_backend: LinearBackend = KernelConfig.linear_backend

508 all2all_backend: All2AllBackend = ParallelConfig.all2all_backend

509 enable_elastic_ep: bool = ParallelConfig.enable_elastic_ep

510 + elastic_ep_max_dp_size: int = ParallelConfig.elastic_ep_max_dp_size

511 enable_dbo: bool = ParallelConfig.enable_dbo

512 ubatch_size: int = ParallelConfig.ubatch_size

513 dbo_decode_token_threshold: int = ParallelConfig.dbo_decode_token_threshold

@@ -1201,6 +1202,10 @@ class EngineArgs:

1202 parallel_group.add_argument(

1203 "--enable-elastic-ep", **parallel_kwargs["enable_elastic_ep"]

1204 )

1205 + parallel_group.add_argument(

1206 + "--elastic-ep-max-dp-size",

1207 + **parallel_kwargs["elastic_ep_max_dp_size"],

1208 + )

1209 parallel_group.add_argument(

1210 "--dbo-decode-token-threshold",

1211 **parallel_kwargs["dbo_decode_token_threshold"],

@@ -2345,6 +2350,7 @@ class EngineArgs:

2350 enable_ep_weight_filter=self.enable_ep_weight_filter,

2351 all2all_backend=self.all2all_backend,

2352 enable_elastic_ep=self.enable_elastic_ep,

2353 + elastic_ep_max_dp_size=self.elastic_ep_max_dp_size,

2354 enable_dbo=self.enable_dbo,

2355 ubatch_size=self.ubatch_size,

2356 dbo_decode_token_threshold=self.dbo_decode_token_threshold,

vllm/envs.py

-5

@@ -316,7 +316,6 @@ if TYPE_CHECKING:

316 VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: bool = False

317 VLLM_ELASTIC_EP_DRAIN_REQUESTS: bool = False

318 VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = True

319 - VLLM_NIXL_EP_MAX_NUM_RANKS: int = 32

319 VLLM_XPU_ENABLE_XPU_GRAPH: bool = False

320 VLLM_XPU_FORCE_N_CONTIG_WEIGHT: bool = False

321 VLLM_XPU_USE_SAMPLER_KERNEL: bool = True

@@ -2129,10 +2128,6 @@ environment_variables: dict[str, Callable[[], Any]] = {

2128 "VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS": lambda: bool(

2129 int(os.getenv("VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS", "1"))

2130 ),

2132 - # NIXL EP environment variables

2133 - "VLLM_NIXL_EP_MAX_NUM_RANKS": lambda: int(

2134 - os.getenv("VLLM_NIXL_EP_MAX_NUM_RANKS", "32")

2135 - ),

2131 # Whether enable XPU graph on Intel GPU

2132 "VLLM_XPU_ENABLE_XPU_GRAPH": lambda: bool(

2133 int(os.getenv("VLLM_XPU_ENABLE_XPU_GRAPH", "0"))

vllm/model_executor/layers/fused_moe/all2all_utils.py

+15 -18

@@ -8,6 +8,7 @@ import torch

9 from vllm.config import get_current_vllm_config

10 from vllm.distributed import (

11 + GroupCoordinator,

12 get_ep_group,

13 )

14 from vllm.logger import init_logger

@@ -104,16 +105,12 @@ if current_platform.is_cuda_alike():

105 )

106

107

107 -def get_ep_all2all_manager(eep_stage: bool = False) -> Any:

108 - if eep_stage:

109 - from vllm.distributed.elastic_ep.standby_state import get_standby_ep_group

110 -

111 - ep_group = get_standby_ep_group()

112 - assert ep_group is not None

113 - device_communicator = ep_group.device_communicator

114 - else:

115 - device_communicator = get_ep_group().device_communicator

116 -

108 +def get_ep_all2all_manager(

109 + ep_group: GroupCoordinator | None = None,

110 +) -> Any:

111 + if ep_group is None:

112 + ep_group = get_ep_group()

113 + device_communicator = ep_group.device_communicator

114 assert device_communicator is not None

115 all2all_manager = device_communicator.all2all_manager

116 assert all2all_manager is not None

@@ -168,7 +165,7 @@ def maybe_make_prepare_finalize(

165 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

166 allow_new_interface: bool = False,

167 use_monolithic: bool = False,

171 - eep_stage: bool = False,

168 + all2all_manager: Any | None = None,

169 ) -> FusedMoEPrepareAndFinalize | None:

170 if not moe.moe_parallel_config.use_all2all_kernels:

171 if not allow_new_interface:

@@ -190,7 +187,8 @@ def maybe_make_prepare_finalize(

187 "Detected DP deployment with no --enable-expert-parallel. "

188 "Falling back to AllGather+ReduceScatter dispatch/combine."

189 )

193 - all2all_manager = get_ep_all2all_manager(eep_stage)

190 + if all2all_manager is None:

191 + all2all_manager = get_ep_all2all_manager()

192 return make_moe_prepare_and_finalize_naive_dp_ep(

193 is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel,

194 num_dispatchers=all2all_manager.world_size,

@@ -199,7 +197,8 @@ def maybe_make_prepare_finalize(

197 else:

198 return make_moe_prepare_and_finalize_no_dp_ep(use_monolithic)

199

202 - all2all_manager = get_ep_all2all_manager(eep_stage)

200 + if all2all_manager is None:

201 + all2all_manager = get_ep_all2all_manager()

202

203 prepare_finalize: FusedMoEPrepareAndFinalize | None = None

204

@@ -361,10 +360,7 @@ def maybe_make_prepare_finalize(

360 all_to_all_args = dict(

361 max_num_tokens_per_dp_rank=moe.max_num_tokens,

362 token_hidden_size=moe.hidden_dim,

364 - num_ep_ranks=all2all_manager.world_size,

365 - num_global_experts=moe.num_experts,

366 - num_local_experts=moe.num_experts // all2all_manager.world_size,

367 - stage=eep_stage,

363 + num_local_experts=moe.num_local_experts,

364 )

365 handle = all2all_manager.get_handle(all_to_all_args)

366

@@ -378,7 +374,8 @@ def maybe_make_prepare_finalize(

374 prepare_finalize = NixlEPPrepareAndFinalize(

375 handle,

376 max_tokens_per_rank=moe.max_num_tokens,

381 - num_dispatchers=all2all_manager.world_size,

377 + num_dispatchers=all2all_manager.max_num_ep_ranks,

378 + expert_capacity=(moe.num_local_experts * all2all_manager.max_num_ep_ranks),

379 use_fp8_dispatch=use_fp8_dispatch,

380 global_to_physical=global_to_physical,

381 physical_to_global=physical_to_global,

vllm/model_executor/layers/fused_moe/config.py

@@ -1306,6 +1306,7 @@ class FusedMoEConfig:

1306

1307 moe_backend: MoEBackend = "auto"

1308 max_num_tokens: int = SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP

1309 + elastic_ep_max_dp_size: int | None = None

1310 has_bias: bool = False

1311 is_lora_enabled: bool = False

1312

vllm/model_executor/layers/fused_moe/eep_reconfigure.py

+3 -1

@@ -60,6 +60,8 @@ def _make_eep_experts(

60 def make_eep_staged_quant_method(

61 module: "MoERunner",

62 moe_config: FusedMoEConfig,

63 + *,

64 + all2all_manager: Any,

65 ) -> FusedMoEMethodBase | None:

66 quant_method = module._quant_method

67 if not quant_method.supports_internal_mk:

@@ -93,7 +95,7 @@ def make_eep_staged_quant_method(

95 routing_tables=None,

96 allow_new_interface=True,

97 use_monolithic=quant_method.is_monolithic,

96 - eep_stage=True,

98 + all2all_manager=all2all_manager,

99 )

100 assert prepare_finalize is not None

101 assert isinstance(prepare_finalize, FusedMoEPrepareAndFinalizeModular)

vllm/model_executor/layers/fused_moe/layer.py

@@ -348,6 +348,7 @@ def FusedMoEFactory(

348 moe_backend=vllm_config.kernel_config.moe_backend,

349 router_logits_dtype=router_logits_dtype,

350 max_num_tokens=max_num_batched_tokens,

351 + elastic_ep_max_dp_size=vllm_config.parallel_config.elastic_ep_max_dp_size,

352 has_bias=has_bias,

353 is_lora_enabled=vllm_config.lora_config is not None,

354 activation=moe_activation,

vllm/model_executor/layers/fused_moe/modular_kernel.py

-7

@@ -252,13 +252,6 @@ class FusedMoEPrepareAndFinalize(ABC):

252 """

253 return False

254

255 - def on_commit(self) -> None:

256 - """

257 - Runs after this prepare/finalize has been committed to the active

258 - MoE kernel.

259 - """

260 - return

261 -

255

256 # TODO: pass FusedMoEParallelConfig in as ctor parameter?

257 class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize):

vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py

+3 -14

@@ -6,8 +6,6 @@ import nixl_ep

6 import torch

8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 -from vllm.distributed import get_ep_group

10 -from vllm.distributed.device_communicators.all2all import NixlEPAll2AllManager

9 from vllm.logger import init_logger

10 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

11 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

@@ -79,6 +77,7 @@ class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):

77 buffer: nixl_ep.Buffer,

78 max_tokens_per_rank: int,

79 num_dispatchers: int,

80 + expert_capacity: int,

81 use_fp8_dispatch: bool = False,

82 global_to_physical: torch.Tensor | None = None,

83 physical_to_global: torch.Tensor | None = None,

@@ -94,6 +93,7 @@ class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):

93 # combine function.

94 self.handles: list[tuple | None] = [None, None]

95 self.num_dispatchers_ = num_dispatchers

96 + self.expert_capacity = expert_capacity

98 topk_indices_dtype = self.topk_indices_dtype()

@@ -141,17 +141,6 @@ class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):

141 def max_num_tokens_per_rank(self) -> int | None:

142 return self.max_tokens_per_rank

143

144 - def on_commit(self) -> None:

145 - device_communicator = get_ep_group().device_communicator

146 - assert device_communicator is not None

147 - all2all_manager = device_communicator.all2all_manager

148 - assert isinstance(all2all_manager, NixlEPAll2AllManager)

149 - # maybe_make_prepare_finalize(..., eep_stage=True) initializes self.buffer

150 - # with get_handle(..., stage=True), which stages global NIXL state for the

151 - # new config but leaves it inactive while the old config remains active.

152 - # When EEP commit switches to this P/F, this P/F needs to commit that state.

153 - all2all_manager.commit_staged_state()

154 -

144 def topk_indices_dtype(self) -> torch.dtype | None:

145 return NIXL_EP_TOPK_INDICES_DTYPE

146

@@ -267,7 +256,7 @@ class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):

256 a1,

257 dispatch_topk_ids,

258 self.max_tokens_per_rank,

270 - num_experts,

259 + self.expert_capacity,

260 use_fp8=self.use_fp8_dispatch,

261 # round_scale needs to be set to dispatch in ue8m0

262 round_scale=self.use_ue8m0_dispatch,

vllm/model_executor/layers/fused_moe/utils.py

+2 -1

@@ -500,7 +500,8 @@ def fi_moe_largest_bucket(moe_config: "FusedMoEConfig") -> int:

500

501 For a detailed explanation, see: `docs/serving/data_parallel_deployment.md`

502 """

503 - return max(moe_config.max_num_tokens * moe_config.dp_size, 8192)

503 + dp_size = moe_config.elastic_ep_max_dp_size or moe_config.dp_size

504 + return max(moe_config.max_num_tokens * dp_size, 8192)

505

506

507 @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)

vllm/model_executor/warmup/flashinfer_autotune_cache.py

+39

@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING

12 import vllm.envs as envs

14 if TYPE_CHECKING:

15 + from vllm.distributed.parallel_state import GroupCoordinator

16 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

@@ -40,6 +41,44 @@ def resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path:

41 return output_dir / "autotune_configs.json"

44 +def sync_flashinfer_autotune_cache(

45 + runner: "GPUModelRunner",

46 + group: "GroupCoordinator",

47 +) -> None:

48 + cache: bytes | str | None = None

49 + if (

50 + group.rank_in_group == 0

51 + and runner.vllm_config.kernel_config.enable_flashinfer_autotune

52 + ):

53 + try:

54 + from vllm.platforms import current_platform

55 + from vllm.utils.flashinfer import has_flashinfer

56 +

57 + if has_flashinfer() and current_platform.has_device_capability(90):

58 + from flashinfer.autotuner import AutoTuner

59 +

60 + with tempfile.TemporaryDirectory() as temp_dir:

61 + path = Path(temp_dir) / "autotune_configs.json"

62 + AutoTuner.get().save_configs(str(path))

63 + cache = path.read_bytes()

64 + except Exception as exc:

65 + cache = f"{type(exc).__name__}: {exc}"

66 +

67 + cache = group.broadcast_object(cache)

68 + if isinstance(cache, str):

69 + raise RuntimeError(f"Failed to serialize FlashInfer autotune state: {cache}")

70 + if cache is None or group.rank_in_group == 0:

71 + return

72 +

73 + from flashinfer.autotuner import AutoTuner

74 +

75 + with tempfile.NamedTemporaryFile() as f:

76 + f.write(cache)

77 + f.flush()

78 + if not AutoTuner.get().load_configs(f.name):

79 + raise RuntimeError("FlashInfer autotune cache is incompatible")

80 +

81 +

82 def write_flashinfer_autotune_cache(cache_path: Path, contents: bytes) -> None:

83 cache_path.parent.mkdir(parents=True, exist_ok=True)

84 fd, tmp_path = tempfile.mkstemp(

vllm/v1/engine/core_client.py

@@ -1711,6 +1711,12 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1711 "Only ray DP backend supports scaling elastic EP"

1712 )

1713 parallel_config = self.vllm_config.parallel_config

1714 + if new_data_parallel_size > parallel_config.elastic_ep_max_dp_size:

1715 + raise ValueError(

1716 + f"Cannot scale to data_parallel_size {new_data_parallel_size}; "

1717 + "--elastic-ep-max-dp-size is "

1718 + f"{parallel_config.elastic_ep_max_dp_size}."

1719 + )

1720 num_experts = self.vllm_config.model_config.get_num_experts()

1721 num_physical_experts = (

1722 num_experts + parallel_config.eplb_config.num_redundant_experts

vllm/v1/worker/dp_utils.py

+31 -1

@@ -1,7 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -from contextlib import nullcontext

4 +from contextlib import contextmanager, nullcontext

5 +from contextvars import ContextVar

7 import torch

8 import torch.distributed as dist

@@ -18,6 +19,22 @@ from vllm.v1.worker.ubatch_utils import (

20 logger = init_logger(__name__)

22 +_skip_dp_coordination = ContextVar("skip_dp_coordination", default=False)

23 +

24 +

25 +@contextmanager

26 +def skip_dp_coordination():

27 + """Run without coordinating DP metadata with other DP ranks."""

28 + token = _skip_dp_coordination.set(True)

29 + try:

30 + yield

31 + finally:

32 + _skip_dp_coordination.reset(token)

33 +

34 +

35 +def should_skip_dp_coordination() -> bool:

36 + return _skip_dp_coordination.get()

37 +

39 def _get_device_and_group(parallel_config: ParallelConfig):

40 # Use the actual device assigned to the DP group, not just the device type

@@ -223,6 +240,19 @@ def coordinate_batch_across_dp(

240 if num_tokens_padded is None:

241 num_tokens_padded = num_tokens_unpadded

242

243 + if should_skip_dp_coordination():

244 + should_ubatch = should_attempt_ubatching and not is_last_ubatch_empty(

245 + num_tokens_unpadded,

246 + num_tokens_padded,

247 + parallel_config.num_ubatches,

248 + )

249 + num_tokens_after_padding = torch.full(

250 + (parallel_config.data_parallel_size,),

251 + num_tokens_padded,

252 + dtype=torch.int32,

253 + )

254 + return should_ubatch, num_tokens_after_padding, cudagraph_mode

255 +

256 (should_ubatch, num_tokens_after_padding, synced_cudagraph_mode) = (

257 _synchronize_dp_ranks(

258 num_tokens_unpadded,

vllm/v1/worker/gpu/dp_utils.py

+5 -1

@@ -10,6 +10,7 @@ import torch.distributed as dist

10 from vllm.config import ParallelConfig

11 from vllm.config.compilation import CUDAGraphMode

12 from vllm.distributed.parallel_state import get_dp_group

13 +from vllm.v1.worker.dp_utils import should_skip_dp_coordination

14 from vllm.v1.worker.gpu.cudagraph_utils import (

15 BatchExecutionDescriptor,

16 CudaGraphManager,

@@ -71,7 +72,10 @@ def sync_cudagraph_and_dp_padding(

72 tensor[3][dp_rank] = max_query_len or -1 # (-1 means None)

73 tensor[4][dp_rank] = int(allow_ubatching)

74 tensor[5][dp_rank] = num_reqs

74 - dist.all_reduce(tensor, group=group)

75 + if should_skip_dp_coordination():

76 + tensor[:] = tensor[:, dp_rank, None].clone()

77 + else:

78 + dist.all_reduce(tensor, group=group)

80 num_tokens_across_dp = tensor[0]

81 cg_mode_across_dp = tensor[1]