@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 35efdf6b34

[Elastic EP] Async preparation (#47288)

Signed-off-by: Itay Alroy <ialroy@nvidia.com>

Itay Alroy committed Jul 28, 2026 at 15:19 UTC 35efdf6b34f1ef76c23c8980f6e8a3b73cab50f1
21 files changed +701 -677
.buildkite/test_areas/expert_parallelism.yaml
+1
@@ -52,4 +52,5 @@ steps:
52 - vllm/compilation/
53 - tests/distributed/
54 commands:
55 + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
56 - pytest -v -s distributed/test_elastic_ep.py
tests/distributed/test_elastic_ep.py
+135 -31
@@ -3,7 +3,9 @@
3
4 import os
5 import subprocess
6 +import threading
7 import time
8 +from concurrent.futures import ThreadPoolExecutor
9
10 import pytest
11 import requests
@@ -40,6 +42,106 @@ def _send_scale_command(server: RemoteOpenAIServer, new_dp_size: int) -> bool:
42 return False
43
44
45 +def _traffic_loop(
46 + server: RemoteOpenAIServer,
47 + dp_rank: int | None,
48 + ready: threading.Barrier,
49 + stop: threading.Event,
50 + finished: threading.Event,
51 + is_probe: bool = False,
52 +) -> list[tuple[float, float, int | None]]:
53 + url = server.url_for("is_scaling_elastic_ep" if is_probe else "v1/completions")
54 + payload = {"model": MODEL_NAME, "prompt": "Hello", "max_tokens": 4}
55 + headers = None if dp_rank is None else {"X-data-parallel-rank": str(dp_rank)}
56 + request_payload = None if is_probe else payload
57 + responses = []
58 + is_ready = False
59 + while not stop.is_set():
60 + request_start = time.perf_counter()
61 + try:
62 + response = requests.post(
63 + url, json=request_payload, headers=headers, timeout=120
64 + )
65 + status_code = response.status_code
66 + except requests.exceptions.RequestException:
67 + status_code = None
68 + responses.append((request_start, time.perf_counter(), status_code))
69 + if status_code == 200:
70 + if not is_ready:
71 + ready.wait(timeout=120)
72 + is_ready = True
73 + if finished.is_set():
74 + return responses
75 + time.sleep(0.05)
76 + return responses
77 +
78 +
79 +def _downtime(responses: list[tuple[float, float, int | None]]) -> float:
80 + rejected = [end for _, end, status in responses if status == 503]
81 + if not rejected:
82 + return 0
83 + recovered = next(
84 + end for _, end, status in responses if status == 200 and end > rejected[-1]
85 + )
86 + return recovered - rejected[0]
87 +
88 +
89 +def _scale_with_traffic(
90 + server: RemoteOpenAIServer,
91 + source_dp_size: int,
92 + new_dp_size: int,
93 + traffic_mode: str,
94 +) -> None:
95 + traffic_clients: list[int | None] = []
96 + if traffic_mode == "light":
97 + traffic_clients = [0]
98 + elif traffic_mode == "heavy":
99 + traffic_clients = [None] * source_dp_size
100 + clients = [(None, True)] + [(rank, False) for rank in traffic_clients]
101 + ready = threading.Barrier(len(clients) + 1)
102 + stop = threading.Event()
103 + finished = threading.Event()
104 +
105 + with ThreadPoolExecutor(max_workers=len(clients)) as executor:
106 + futures = [
107 + executor.submit(
108 + _traffic_loop, server, rank, ready, stop, finished, is_probe
109 + )
110 + for rank, is_probe in clients
111 + ]
112 + try:
113 + ready.wait(timeout=120)
114 + start_time = time.perf_counter()
115 + assert _send_scale_command(server, new_dp_size)
116 + scale_seconds = time.perf_counter() - start_time
117 + finished.set()
118 + probe_result, *results = [future.result(timeout=120) for future in futures]
119 + finally:
120 + stop.set()
121 +
122 + bad_statuses = {
123 + status
124 + for responses in [probe_result, *results]
125 + for _, _, status in responses
126 + if status not in (200, 503)
127 + }
128 + assert not bad_statuses, f"traffic got unexpected statuses {bad_statuses}"
129 + probe_503 = [start for start, _, status in probe_result if status == 503]
130 + assert probe_503, "Scaling probe did not observe commit"
131 + assert not results or any(
132 + status == 200 and start_time <= request_start and request_end < probe_503[0]
133 + for responses in results
134 + for request_start, request_end, status in responses
135 + ), "No request completed successfully during preparation"
136 +
137 + print(
138 + f"[Elastic EP timing][{source_dp_size}->{new_dp_size}]"
139 + f"[traffic={traffic_mode}] "
140 + f"scale_seconds={scale_seconds:.3f} "
141 + f"downtime_seconds={_downtime(probe_result):.3f}"
142 + )
143 +
144 +
145 def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float:
146 assert server.port is not None
147 result = evaluate_gsm8k(
@@ -59,7 +161,7 @@ def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float:
161 return accuracy
162
163
62 -def _base_serve_args(use_async_eplb: bool = False) -> list[str]:
164 +def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str]:
165 args = [
166 "--trust-remote-code",
167 "--tensor-parallel-size",
@@ -78,57 +180,65 @@ def _base_serve_args(use_async_eplb: bool = False) -> list[str]:
180 "--eplb-config.num_redundant_experts",
181 "0",
182 "--eplb-config.use_async",
81 - "true" if use_async_eplb else "false",
183 + "true",
184 "--eplb-config.step_interval",
83 - "10",
185 + "300",
186 "--eplb-config.window_size",
187 "5",
188 "--data-parallel-backend",
189 "ray",
190 "--data-parallel-size",
89 - "2",
191 + str(dp_size),
192 "--api-server-count",
193 "1",
194 + "--disable-access-log-for-endpoints",
195 + "/is_scaling_elastic_ep",
196 ]
197
198 leader_address = os.environ.get("LEADER_ADDRESS")
199 if leader_address:
200 args.extend(["--data-parallel-address", leader_address])
201 + if enforce_eager:
202 + args.append("--enforce-eager")
203
204 return args
205
206
207 @pytest.mark.parametrize(
102 - "use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"]
208 + ("enforce_eager", "traffic_mode"),
209 + [
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"),
214 + ],
215 )
216 @multi_gpu_test(num_gpus=4)
105 -def test_elastic_ep_scaling(use_async_eplb: bool):
106 - if use_async_eplb:
107 - from vllm.distributed.eplb.eplb_communicator import has_nixl
217 +def test_elastic_ep_scaling(enforce_eager: bool, traffic_mode: str):
218 + from vllm.distributed.eplb.eplb_communicator import has_nixl
219
109 - if not has_nixl():
110 - pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")
220 + if not has_nixl():
221 + pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")
222
112 - vllm_serve_args = _base_serve_args(use_async_eplb)
223 + initial_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_INITIAL_DP", "2"))
224 + target_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_TARGET_DP", "4"))
225 + assert target_dp_size > initial_dp_size
226 + vllm_serve_args = _base_serve_args(initial_dp_size, enforce_eager)
227
228 with RemoteOpenAIServer(
229 MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200
230 ) as server:
117 - initial_accuracy = _run_gsm8k_eval(server, "Initial (2 GPUs)")
118 -
119 - assert _send_scale_command(server, 4)
120 - time.sleep(10)
121 - scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (4 GPUs)")
231 + initial_accuracy = _run_gsm8k_eval(server, "Initial")
232
233 + _scale_with_traffic(server, initial_dp_size, target_dp_size, traffic_mode)
234 + scale_up_accuracy = _run_gsm8k_eval(server, "After scale up")
235 assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, (
236 f"Scale up accuracy {scale_up_accuracy:.3f} dropped more than "
237 f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}"
238 )
239
128 - assert _send_scale_command(server, 2)
129 - time.sleep(5)
130 - scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)")
131 -
240 + _scale_with_traffic(server, target_dp_size, initial_dp_size, traffic_mode)
241 + scale_down_accuracy = _run_gsm8k_eval(server, "After scale down")
242 assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, (
243 f"Scale down accuracy {scale_down_accuracy:.3f} dropped more than "
244 f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}"
@@ -147,24 +257,20 @@ def test_elastic_ep_scaling(use_async_eplb: bool):
257 print(f" Tolerance: {ACCURACY_TOL:.3f}")
258
259
150 -@pytest.mark.parametrize(
151 - "use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"]
152 -)
260 @multi_gpu_test(num_gpus=4)
154 -def test_elastic_ep_scaling_uneven(use_async_eplb: bool):
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 """
161 - if use_async_eplb:
162 - from vllm.distributed.eplb.eplb_communicator import has_nixl
268 + from vllm.distributed.eplb.eplb_communicator import has_nixl
269
164 - if not has_nixl():
165 - pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")
270 + if not has_nixl():
271 + pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")
272
167 - vllm_serve_args = _base_serve_args(use_async_eplb)
273 + vllm_serve_args = _base_serve_args()
274
275 with RemoteOpenAIServer(
276 MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200
@@ -174,7 +280,6 @@ def test_elastic_ep_scaling_uneven(use_async_eplb: bool):
280 # Scale 2 -> 3: This has remainder = 1 % 2 = 1
281 # Tests uneven sender-receiver pairing
282 assert _send_scale_command(server, 3)
177 - time.sleep(10)
283 scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (3 GPUs)")
284
285 assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, (
@@ -184,7 +289,6 @@ def test_elastic_ep_scaling_uneven(use_async_eplb: bool):
289
290 # Scale back down to 2
291 assert _send_scale_command(server, 2)
187 - time.sleep(5)
292 scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)")
293
294 assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, (
vllm/config/parallel.py
+9
@@ -686,6 +686,14 @@ class ParallelConfig:
686 and self.data_parallel_size > 1
687 )
688
689 + @property
690 + def use_all2all(self) -> bool:
691 + return (
692 + self.data_parallel_size > 1
693 + or self.use_sequence_parallel_moe
694 + or (self.enable_expert_parallel and self.prefill_context_parallel_size > 1)
695 + )
696 +
697 @property
698 def use_batched_dp_moe(self) -> bool:
699 return (
@@ -786,6 +794,7 @@ class ParallelConfig:
794 "data_parallel_master_ip",
795 "data_parallel_master_port",
796 "_data_parallel_master_port_list",
797 + "_coord_store_port",
798 "data_parallel_rpc_port",
799 "rank",
800 "master_addr",
vllm/distributed/device_communicators/base_device_communicator.py
+3 -13
@@ -175,6 +175,7 @@ class DeviceCommunicatorBase:
175 unique_name: str = "",
176 global_ranks: list[int] | None = None,
177 global_world_size: int | None = None,
178 + use_all2all: bool = False,
179 ):
180 self.device = device or torch.device("cpu")
181 self.cpu_group = cpu_group
@@ -204,26 +205,15 @@ class DeviceCommunicatorBase:
205 self.global_world_size = dist.get_world_size()
206 self.rank_in_group = dist.get_group_rank(self.cpu_group, self.global_rank)
207
207 - use_ep = False
208 all2all_backend = None
209 from vllm.config import get_current_vllm_config_or_none
210
211 config = get_current_vllm_config_or_none()
212 if config is not None:
213 - # initialize the all2all manager for DP or sequence-parallel EP.
214 - parallel_config = config.parallel_config
215 - use_ep = (
216 - parallel_config.data_parallel_size > 1
217 - or parallel_config.use_sequence_parallel_moe
218 - or (
219 - parallel_config.enable_expert_parallel
220 - and parallel_config.prefill_context_parallel_size > 1
221 - )
222 - )
223 - all2all_backend = parallel_config.all2all_backend
213 + all2all_backend = config.parallel_config.all2all_backend
214
215 self.is_ep_communicator = unique_name.split(":")[0] == "ep"
226 - self.use_all2all = self.is_ep_communicator and use_ep
216 + self.use_all2all = self.is_ep_communicator and use_all2all
217 self.all2all_backend = all2all_backend
218 self.all2all_manager: All2AllManagerBase | None = None
219
vllm/distributed/device_communicators/cpu_communicator.py
+4 -1
@@ -24,8 +24,11 @@ class CpuCommunicator(DeviceCommunicatorBase):
24 device: torch.device | None = None,
25 device_group: ProcessGroup | None = None,
26 unique_name: str = "",
27 + use_all2all: bool = False,
28 ):
28 - super().__init__(cpu_group, device, device_group, unique_name)
29 + super().__init__(
30 + cpu_group, device, device_group, unique_name, use_all2all=use_all2all
31 + )
32 self.dist_module = torch.distributed
33
34 if (
vllm/distributed/device_communicators/cuda_communicator.py
+2
@@ -36,6 +36,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
36 global_ranks: list[int] | None = None,
37 global_world_size: int | None = None,
38 tcp_store_group: StatelessProcessGroup | None = None,
39 + use_all2all: bool = False,
40 ):
41 super().__init__(
42 cpu_group,
@@ -44,6 +45,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
45 unique_name,
46 global_ranks,
47 global_world_size,
48 + use_all2all=use_all2all,
49 )
50 if "tp" not in unique_name:
51 # custom allreduce or torch symm mem can be used only by tp
vllm/distributed/device_communicators/xpu_communicator.py
+4 -1
@@ -20,8 +20,11 @@ class XpuCommunicator(DeviceCommunicatorBase):
20 device: torch.device | None = None,
21 device_group: ProcessGroup | None = None,
22 unique_name: str = "",
23 + use_all2all: bool = False,
24 ):
24 - super().__init__(cpu_group, device, device_group, unique_name)
25 + super().__init__(
26 + cpu_group, device, device_group, unique_name, use_all2all=use_all2all
27 + )
28 self.ca_comm: None = None
29 if self.use_all2all:
30 if self.all2all_backend in ("naive", "allgather_reducescatter"):
vllm/distributed/elastic_ep/elastic_execute.py
+95 -86
@@ -1,9 +1,9 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import copy
3 import gc
4 import weakref
5 from collections.abc import Iterable, Sequence
6 +from concurrent.futures import Future, ThreadPoolExecutor
7 from dataclasses import replace
8 from typing import TYPE_CHECKING
9
@@ -43,6 +43,7 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
43 from vllm.model_executor.layers.fused_moe.eep_reconfigure import (
44 make_eep_staged_quant_method,
45 )
46 +from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
47 from vllm.utils import is_moe_layer
48 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
49 from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper
@@ -145,6 +146,10 @@ class ElasticEPScalingExecutor:
146 self.worker_ref = weakref.ref(worker)
147 self.reconfig_request = None
148 self._staged_moe_quant_methods: dict[nn.Module, FusedMoEMethodBase] = {}
149 + self._async_executor = ThreadPoolExecutor(
150 + max_workers=1, thread_name_prefix="ElasticEPAsync"
151 + )
152 + self._async_future: Future[None] | None = None
153
154 @property
155 def worker(self):
@@ -159,59 +164,68 @@ class ElasticEPScalingExecutor:
164 raise ValueError(f"Unknown execute method: {execute_method}")
165 return method(*args, **kwargs)
166
162 - def _set_eplb_suppressed(self, suppressed: bool) -> None:
163 - self.worker.model_runner.eep_eplb_suppressed = suppressed
164 - ep_group = get_standby_ep_group() or get_ep_group()
165 - if ep_group.rank == 0:
166 - logger.info(
167 - "[Elastic EP] EPLB %s elastic scaling transition",
168 - "disabled during" if suppressed else "re-enabled after",
169 - )
167 + def start_async(self, execute_method: str, *args, **kwargs) -> str:
168 + if self._async_future is not None:
169 + raise RuntimeError("Another Elastic EP async method is active")
170 + if args and isinstance(args[0], ReconfigureDistributedRequest):
171 + self.reconfig_request = args[0]
172 + dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank
173 + done_key = f"eep_async/{execute_method}/{dp_rank}/{self.worker.rank}"
174 + self._async_future = self._async_executor.submit(
175 + self._run_async, execute_method, *args, **kwargs
176 + )
177 + self._async_future.add_done_callback(lambda _: self._mark_async_done(done_key))
178 + return done_key
179 +
180 + def _run_async(self, execute_method: str, *args, **kwargs) -> None:
181 + from vllm.platforms import current_platform
182 +
183 + self.worker.vllm_config.enable_trace_function_call_for_thread()
184 + assert hasattr(self.worker, "device")
185 + current_platform.set_device(self.worker.device)
186 + with set_current_vllm_config(self.worker.vllm_config):
187 + self.execute(execute_method, *args, **kwargs)
188 +
189 + def _mark_async_done(self, done_key: str) -> None:
190 + from vllm.distributed.utils import get_cached_tcp_store_client
191 +
192 + assert self.reconfig_request is not None
193 + get_cached_tcp_store_client(
194 + self.reconfig_request.new_data_parallel_master_ip,
195 + self.reconfig_request.coord_store_port,
196 + ).set(done_key, b"1")
197 +
198 + def clear_async(self) -> None:
199 + future = self._async_future
200 + if future is None:
201 + raise RuntimeError("No Elastic EP async method is active")
202 + if not future.done():
203 + raise RuntimeError("Elastic EP async method is not done")
204 + self._async_future = None
205 + future.result()
206
207 def load_model(self) -> None:
172 - (
173 - expanded_physical_to_logical,
174 - num_logical_experts,
175 - old_num_physical_experts,
176 - ) = self.receive_expert_mapping()
177 - num_physical_experts = expanded_physical_to_logical.shape[1]
178 - self.worker.parallel_config.eplb_config.num_redundant_experts = (
179 - num_physical_experts - num_logical_experts
180 - )
208 self.worker.load_model(load_dummy_weights=True)
182 - self.worker.model_runner.setup_eplb_from_mapping(
183 - expanded_physical_to_logical, old_num_physical_experts
184 - )
185 - self._set_eplb_suppressed(True)
209
210 def create_standby_groups(
188 - self, reconfig_request: ReconfigureDistributedRequest
211 + self, reconfig_request: ReconfigureDistributedRequest, use_all2all: bool
212 ) -> None:
213 self.reconfig_request = reconfig_request
214 new_dp_size = reconfig_request.new_data_parallel_size
215 old_dp_size = get_dp_group().world_size
193 - world_size = self.worker.vllm_config.parallel_config.world_size
216 + parallel_config = self.worker.vllm_config.parallel_config
217 + world_size = parallel_config.world_size
218 new_world_size_across_dp = world_size * new_dp_size
195 - updated_config = copy.copy(self.worker.vllm_config)
196 - updated_config.parallel_config = copy.deepcopy(
197 - self.worker.vllm_config.parallel_config
219 + create_standby_groups(
220 + new_dp_size=new_dp_size,
221 + new_world_size_across_dp=new_world_size_across_dp,
222 + master_ip=reconfig_request.new_data_parallel_master_ip,
223 + coord_store_port=reconfig_request.coord_store_port,
224 + use_all2all=use_all2all,
225 + enable_eplb=parallel_config.enable_eplb,
226 )
199 - updated_config.parallel_config.data_parallel_size = new_dp_size
200 - with set_current_vllm_config(updated_config):
201 - create_standby_groups(
202 - new_dp_size=new_dp_size,
203 - new_world_size_across_dp=new_world_size_across_dp,
204 - master_ip=reconfig_request.new_data_parallel_master_ip,
205 - coord_store_port=reconfig_request.coord_store_port,
206 - enable_eplb=updated_config.parallel_config.enable_eplb,
207 - )
208 - if new_dp_size > old_dp_size:
209 - self._set_eplb_suppressed(True)
210 - eplb_state = self.worker.model_runner.eplb_state
211 - if eplb_state is not None:
212 - eplb_state.drain_async()
213 - elif new_dp_size < old_dp_size:
214 - self._stage_standby_moe_quant_methods()
227 + if new_dp_size < old_dp_size:
228 + self.stage_standby_moe_quant_methods()
229
230 def transfer_weights(self, old_dp_size: int, new_dp_size: int) -> None:
231 standby_dp_group = get_standby_dp_group()
@@ -265,6 +279,7 @@ class ElasticEPScalingExecutor:
279 model_config = self.worker.model_runner.model_config
280 eplb_state = self.worker.model_runner.eplb_state
281 assert eplb_state is not None
282 + eplb_state.drain_async()
283 eplb_model_state = eplb_state.model_states[model_config.compute_hash()]
284 physical_to_logical = eplb_model_state.physical_to_logical_map
285 num_physical_experts = physical_to_logical.shape[1]
@@ -278,10 +293,6 @@ class ElasticEPScalingExecutor:
293 src_rank=0,
294 device=self.worker.device,
295 )
281 - # New workers enter load_model after receiving the expert mapping.
282 - # Stage replacement MoE kernels before returning to the state machine
283 - # so existing ranks can participate in collective EP comm creation.
284 - self._stage_standby_moe_quant_methods()
296
297 def _make_eep_moe_config(self, module, dp_group, ep_group):
298 parallel_config = self.worker.vllm_config.parallel_config
@@ -300,7 +311,7 @@ class ElasticEPScalingExecutor:
311 moe_parallel_config=moe_parallel_config,
312 )
313
303 - def _stage_standby_moe_quant_methods(self) -> None:
314 + def stage_standby_moe_quant_methods(self) -> None:
315 standby_dp_group = get_standby_dp_group()
316 standby_ep_group = get_standby_ep_group()
317 model = self.worker.model_runner.get_model()
@@ -500,26 +511,6 @@ class ElasticEPScalingExecutor:
511 compilation_counter.stock_torch_compile_count += 1
512 self.worker.model_runner.model.compile(fullgraph=True, backend=backend)
513
503 - multi_block_table = self.worker.model_runner.input_batch.block_table
504 - saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = []
505 - for bt in multi_block_table.block_tables:
506 - saved_block_tables.append(
507 - (bt.block_table.gpu.clone(), bt.block_table.cpu.clone())
508 - )
509 - multi_block_table.clear()
510 -
511 - unlock_workspace()
512 - self.worker.compile_or_warm_up_model()
513 - lock_workspace()
514 -
515 - for bt, (saved_gpu, saved_cpu) in zip(
516 - multi_block_table.block_tables, saved_block_tables
517 - ):
518 - bt.block_table.gpu.copy_(saved_gpu)
519 - bt.block_table.cpu.copy_(saved_cpu)
520 - if new_dp_size < old_dp_size:
521 - self._set_eplb_suppressed(False)
522 -
514 def _perform_eplb_reshuffle(
515 self, rank_mapping: dict[int, int] | None = None
516 ) -> None:
@@ -553,12 +544,25 @@ class ElasticEPScalingExecutor:
544 if get_ep_group().rank == 0:
545 logger.info("[Elastic EP] Expert resharding completed")
546
556 - def perform_eplb_reshuffle(self) -> None:
547 + def commit_scale_up(self, is_existing_worker: bool) -> None:
548 + if is_existing_worker:
549 + self.broadcast_expert_mapping()
550 + self.switch_and_prepare()
551 + else:
552 + mapping, _, num_valid_experts = self.receive_expert_mapping()
553 + self.worker.model_runner.setup_eplb_from_mapping(mapping, num_valid_experts)
554 self._perform_eplb_reshuffle()
558 - self._set_eplb_suppressed(False)
555 + self.warm_and_capture()
556 +
557 + def commit_scale_down(self, new_dp_size: int, removing: bool) -> None:
558 + self.perform_scale_down_eplb_reshuffle(new_dp_size)
559 + if removing:
560 + self.switch_and_remove()
561 + else:
562 + self.switch_and_prepare()
563 + self.warm_and_capture()
564
565 def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None:
561 - self._set_eplb_suppressed(True)
566 eplb_state = self.worker.model_runner.eplb_state
567 if eplb_state is not None:
568 eplb_state.drain_async()
@@ -599,12 +603,17 @@ class ElasticEPScalingExecutor:
603 )
604
605 model = self.worker.model_runner.get_model()
606 + expert_weights = [
607 + module.get_expert_weights()
608 + for module in model.modules()
609 + if is_moe_layer(module)
610 + ]
611 batch_transfer_weights(
612 model=model,
613 is_sender=False,
614 peer_rank=sender_rank,
615 dp_group=dp_group,
607 - expert_weights=model.expert_weights,
616 + expert_weights=expert_weights,
617 )
618 torch.accelerator.synchronize()
619
@@ -643,14 +652,17 @@ class ElasticEPScalingExecutor:
652 with set_current_vllm_config(self.worker.vllm_config):
653 prepare_communication_buffer_for_model(self.worker.model_runner.get_model())
654
646 - def rewarm_workspace(self) -> None:
655 + def warmup_local_kernels(self) -> None:
656 + with set_current_vllm_config(self.worker.vllm_config):
657 + kernel_warmup(self.worker, process_local_only=True)
658 +
659 + def warm_and_capture(self) -> None:
660 # Must run on every DP sibling in lockstep: _dummy_run calls
661 # coordinate_batch_across_dp whenever data_parallel_size > 1
662 # (gpu_model_runner.py:3663), which deadlocks if any rank skips it.
663
651 - # Save and clear block tables so profile_run/compile_or_warm_up_model
652 - # don't write dummy slot mappings into real KV-cache blocks (mirrors
653 - # switch_and_prepare's pattern).
664 + # Save and clear block tables so the dummy MoE forward doesn't
665 + # write dummy slot mappings into real KV-cache blocks.
666 multi_block_table = self.worker.model_runner.input_batch.block_table
667 saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = []
668 for bt in multi_block_table.block_tables:
@@ -660,19 +672,16 @@ class ElasticEPScalingExecutor:
672 multi_block_table.clear()
673
674 # _ensure_workspace_size allocates a fresh tensor on grow, leaving
663 - # captured CUDA graphs with stale data pointers; drop graphs before
664 - # re-warm so captures realign with the resized buffer.
675 + # any captured CUDA graph with a stale data pointer; drop graphs
676 + # before re-warm so captures realign with the resized buffer.
677 self._release_cuda_graphs()
678 unlock_workspace()
679
668 - # Grow the MoE workspace at max_num_tokens.
669 - # compile_or_warm_up_model alone only exercises cudagraph-capture
670 - # sizes (≤64 tokens for this test) and leaves the workspace at
671 - # ~10-14 MB; the post-all-to-all per-rank token count under real
672 - # post-reshuffle routing needs hundreds of MB. Use _dummy_run
673 - # directly (rather than profile_run) with skip_eplb=True so dummy
674 - # routing doesn't pollute the just-rebalanced EPLB stats — same
675 - # convention compile_or_warm_up_model itself uses.
680 + # Grow the MoE workspace at max_num_tokens. compile_or_warm_up_model
681 + # alone only exercises cudagraph-capture sizes and can leave the
682 + # workspace too small for post-reshuffle routing. Use _dummy_run
683 + # directly with skip_eplb=True so dummy routing doesn't pollute the
684 + # just-rebalanced EPLB stats.
685 runner = self.worker.model_runner
686 runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True)
687 self.worker.compile_or_warm_up_model()
vllm/distributed/elastic_ep/elastic_state.py
+161 -297
@@ -1,18 +1,17 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import enum
4 -import time
4 import weakref
6 -from datetime import timedelta
7 -from typing import TYPE_CHECKING, Literal, TypeAlias
5 +from concurrent.futures import Future, ThreadPoolExecutor
6 +from typing import TYPE_CHECKING, Any, Literal, TypeAlias
7
8 import torch.distributed
9
10 from vllm.config import ParallelConfig
11 from vllm.distributed import (
13 - sched_yield,
12 stateless_destroy_torch_distributed_process_group,
13 )
14 +from vllm.distributed.utils import get_cached_tcp_store_client
15 from vllm.logger import init_logger
16 from vllm.v1.engine import (
17 EEPNotificationType,
@@ -31,35 +30,29 @@ WorkerType = Literal["existing", "new", "removing"]
30
31
32 class ScaleUpExistingEngineState(enum.IntEnum):
34 - WAIT_NEW_CORE_ENGINES_INIT = 0
35 - CREATE_STANDBY_GROUPS = 1
36 - TRANSFER_EXPERT_MAPPING = 2
37 - WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT = 3
38 - TRANSFER_WEIGHTS = 4
39 - SYNC_KV_CACHE_MEMORY_SIZE = 5
40 - SWITCH_AND_PREPARE = 6
41 - EPLB_RESHUFFLE = 7
42 - COMPLETE = 8
33 + CREATE_STANDBY_GROUPS = 0
34 + STAGE_QUANT_METHODS = 1
35 + TRANSFER_WEIGHTS = 2
36 + SYNC_KV_CACHE_MEMORY_SIZE = 3
37 + COMMIT_SCALE_UP = 4 # Blocks forward passes.
38 + COMPLETE = 5
39
40
41 class ScaleUpNewEngineState(enum.IntEnum):
42 PRE_KV_INIT = 0
43 PREPARE = 1
48 - EPLB_RESHUFFLE = 2
49 - COMPLETE = 3
44 + COMPLETE = 2
45
46
47 class ScaleDownRemainingEngineState(enum.IntEnum):
48 PREPARE = 0
54 - EPLB_RESHUFFLE = 1
55 - SWITCH_AND_PREPARE = 2
56 - COMPLETE = 3
49 + COMMIT_SCALE_DOWN = 1 # Blocks forward passes.
50 + COMPLETE = 2
51
52
53 class ScaleDownRemovingEngineState(enum.IntEnum):
54 PREPARE = 0
61 - EPLB_RESHUFFLE = 1
62 - COMPLETE = 2
55 + COMPLETE = 1
56
57
58 EngineState: TypeAlias = (
@@ -70,15 +63,6 @@ EngineState: TypeAlias = (
63 )
64
65
73 -class _BarrierTimeoutError(RuntimeError):
74 - """
75 - Exception raised for timeout
76 - in the first stage of our two-staged
77 - TCPStore based barrier to synchronize the
78 - execution of all engines in the DP group.
79 - """
80 -
81 -
66 class ElasticEPScalingState:
67 def __init__(
68 self,
@@ -94,20 +78,24 @@ class ElasticEPScalingState:
78 self.engine_core_ref = weakref.ref(engine_core)
79 self.vllm_config = vllm_config
80 self.old_dp_group = self.engine_core.dp_group if worker_type != "new" else None
97 - self.old_dp_store = self.engine_core.dp_store if worker_type != "new" else None
81 self.new_parallel_config: ParallelConfig = new_parallel_config
82 self.new_dp_group = self.engine_core.dp_group if worker_type == "new" else None
83 self.new_dp_store = self.engine_core.dp_store if worker_type == "new" else None
84 self.worker_type = worker_type
85 self.scale_type = scale_type
86 self.reconfig_request = reconfig_request
104 -
87 + self.commit_requested = False
88 + self._prepare_executor = ThreadPoolExecutor(
89 + max_workers=1, thread_name_prefix="ElasticEPPrepare"
90 + )
91 + self._prepare_future: Future[Any] | None = None
92 + self._new_dp_sync: tuple[object, Any] | None = None
93 self.state: EngineState
94 if scale_type == "scale_up":
95 self.state = (
96 ScaleUpNewEngineState.PRE_KV_INIT
97 if worker_type == "new"
110 - else ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT
98 + else ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS
99 )
100 else:
101 self.state = (
@@ -130,6 +118,31 @@ class ElasticEPScalingState:
118 raise RuntimeError("Engine core has been garbage collected")
119 return engine_core
120
121 + def _collective_rpc(self, *args, **kwargs):
122 + return self.model_executor.collective_rpc(*args, **kwargs)
123 +
124 + def _execute_async(self, execute_method: str, *args) -> bool:
125 + if self._prepare_future is None:
126 + done_keys = self._collective_rpc(
127 + "elastic_ep_execute",
128 + args=("start_async", execute_method, *args),
129 + )
130 + assert self.reconfig_request is not None
131 + coord_store = get_cached_tcp_store_client(
132 + self.reconfig_request.new_data_parallel_master_ip,
133 + self.reconfig_request.coord_store_port,
134 + )
135 + self._prepare_future = self._prepare_executor.submit(
136 + coord_store.wait, done_keys
137 + )
138 + if not self._prepare_future.done():
139 + return False
140 +
141 + self._prepare_future.result()
142 + self._collective_rpc("elastic_ep_execute", args=("clear_async",))
143 + self._prepare_future = None
144 + return True
145 +
146 def progress(self) -> bool:
147 if self.scale_type == "scale_up":
148 return (
@@ -149,157 +162,43 @@ class ElasticEPScalingState:
162 assert self.progress()
163 assert self.state == ScaleUpNewEngineState.PREPARE
164
152 - def _execute_tcp_store_barrier(
153 - self, dp_store, group_rank, group_size, barrier_id, timeout=None
154 - ):
155 - arrival_key = f"arrival_{barrier_id}_{group_rank}"
156 - dp_store.set(arrival_key, b"1")
157 -
158 - start_time = time.time()
159 - processes_arrived: set[int] = set()
160 -
161 - while len(processes_arrived) < group_size:
162 - if (
163 - timeout is not None
164 - and time.time() - start_time > timeout.total_seconds()
165 - ):
166 - raise _BarrierTimeoutError(
167 - f"Barrier timed out after {timeout.total_seconds()} seconds"
168 - )
169 -
170 - for i in range(group_size):
171 - if i in processes_arrived:
172 - continue
173 -
174 - key = f"arrival_{barrier_id}_{i}"
175 - present = dp_store.check([key])
176 - if present:
177 - processes_arrived.add(i)
178 -
179 - if len(processes_arrived) < group_size:
180 - sched_yield()
181 -
182 - def _staged_barrier(self, use_new_group: bool, barrier_name: str) -> bool:
183 - """
184 - Execute a two-staged barrier to synchronize all engines in the DP group.
185 -
186 - Some DP EngineCores may receive the reconfiguration notifications
187 - later than others, and already proceed to engine step (model forward)
188 - in the busy loop.
189 - In this case, EngineCores that already proceed to reconfiguration
190 - should skip reconfiguration and execute model forward for one more
191 - step, so in the next step, all EngineCores will be synchronized.
192 - We use a two-staged barrier to achieve this. The first time each
193 - EngineCore executes the barrier, if a timeout is reached before the
194 - barrier completes, that means some EngineCores have already entered
195 - engine step. The EngineCores that timed out will then proceed to
196 - engine step, and will synchronize with the other EngineCores in the
197 - next step with a barrier without timeout.
198 - """
199 - dp_group = self.new_dp_group if use_new_group else self.old_dp_group
200 - dp_store = self.new_dp_store if use_new_group else self.old_dp_store
201 - assert dp_group is not None and dp_store is not None
202 -
203 - group_rank = dp_group.rank()
204 - group_size = dp_group.size()
205 - barrier_id = f"eep_barrier_{barrier_name}"
206 - sync_key = f"{barrier_id}_sync"
207 -
208 - # TODO(yongji): figure out appropriate timeout for the barrier
209 - timeout = None if dp_store.check([sync_key]) else timedelta(seconds=5)
210 -
211 - try:
212 - self._execute_tcp_store_barrier(
213 - dp_store, group_rank, group_size, barrier_id, timeout=timeout
214 - )
215 - torch.distributed.barrier(dp_group)
216 - if group_rank == 0:
217 - dp_store.delete_key(sync_key)
218 - for i in range(group_size):
219 - dp_store.delete_key(f"arrival_{barrier_id}_{i}")
220 - return True
221 - except _BarrierTimeoutError as e:
222 - if timeout is None:
223 - raise RuntimeError("Unexpected timeout encountered") from e
224 - dp_store.compare_set(sync_key, "", b"1")
225 - return False
226 -
165 def _progress_existing_engine(self) -> bool:
166 state = self.state
229 - assert self.old_dp_group is not None and self.old_dp_store is not None
230 -
231 - if state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT:
232 - return False
167 + assert self.old_dp_group is not None
168
234 - elif state == ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS:
235 - # NOTE(yongji): wait for all existing workers to receive the request
236 - if (
237 - int(self.old_dp_store.get("eep_barrier_engine_count"))
238 - < self.old_dp_group.size()
239 - ):
240 - return False
241 - if not self._staged_barrier(
242 - use_new_group=False, barrier_name="create_standby_groups"
243 - ):
169 + if state == ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS:
170 + if not self._create_standby_groups():
171 return False
245 - if self.old_dp_group.rank() == 0:
246 - self.old_dp_store.delete_key("eep_barrier_engine_count")
247 - self._create_standby_groups()
248 - self.state = ScaleUpExistingEngineState.TRANSFER_EXPERT_MAPPING
172 + self.state = ScaleUpExistingEngineState.STAGE_QUANT_METHODS
173 return True
174
251 - elif state == ScaleUpExistingEngineState.TRANSFER_EXPERT_MAPPING:
252 - self._transfer_expert_mapping()
253 - self.state = ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT
175 + elif state == ScaleUpExistingEngineState.STAGE_QUANT_METHODS:
176 + if not self._execute_async("stage_standby_moe_quant_methods"):
177 + return False
178 + self.state = ScaleUpExistingEngineState.TRANSFER_WEIGHTS
179 return True
180
256 - elif state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT:
257 - return False
258 -
181 elif state == ScaleUpExistingEngineState.TRANSFER_WEIGHTS:
260 - if (
261 - int(self.old_dp_store.get("eep_barrier_engine_count"))
262 - < self.old_dp_group.size()
263 - ):
182 + if not self._transfer_weights():
183 return False
265 - if not self._staged_barrier(
266 - use_new_group=False, barrier_name="transfer_weights"
267 - ):
268 - return False
269 - if self.old_dp_group.rank() == 0:
270 - self.old_dp_store.delete_key("eep_barrier_engine_count")
271 - self._transfer_weights()
184 self.state = ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE
185 return True
186
187 elif state == ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE:
276 - self._sync_kv_cache_memory_size()
277 - self.state = ScaleUpExistingEngineState.SWITCH_AND_PREPARE
278 - return True
279 -
280 - elif state == ScaleUpExistingEngineState.SWITCH_AND_PREPARE:
281 - self._switch_and_prepare()
282 - self.state = ScaleUpExistingEngineState.EPLB_RESHUFFLE
283 - assert self.new_dp_store is not None
284 - self.new_dp_store.add("eep_barrier_engine_count", 1)
188 + if not self._sync_kv_cache_memory_size():
189 + return False
190 + self.state = ScaleUpExistingEngineState.COMMIT_SCALE_UP
191 + self._mark_ready_for_switch()
192 return True
193
287 - elif state == ScaleUpExistingEngineState.EPLB_RESHUFFLE:
288 - assert self.new_dp_group is not None and self.new_dp_store is not None
289 - if (
290 - int(self.new_dp_store.get("eep_barrier_engine_count"))
291 - < self.new_dp_group.size()
292 - ):
293 - return False
294 - if not self._staged_barrier(
295 - use_new_group=True, barrier_name="eplb_reshuffle"
296 - ):
194 + elif state == ScaleUpExistingEngineState.COMMIT_SCALE_UP:
195 + if not self.commit_requested:
196 return False
298 - if self.new_dp_group.rank() == 0:
299 - self.new_dp_store.delete_key("eep_barrier_engine_count")
300 - self._eplb_reshuffle()
197 + self._commit_new_dp_group()
198 + self._collective_rpc("elastic_ep_execute", args=("commit_scale_up", True))
199 self.state = ScaleUpExistingEngineState.COMPLETE
200 self._update_parallel_config()
201 + self._send_reconfigure_finished()
202 return True
203
204 else:
@@ -311,22 +210,17 @@ class ElasticEPScalingState:
210 assert self.new_dp_group is not None and self.new_dp_store is not None
211
212 if state == ScaleUpNewEngineState.PRE_KV_INIT:
314 - self.engine_core._eep_send_engine_core_notification(
315 - EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY
316 - )
317 - self.model_executor.collective_rpc(
318 - "elastic_ep_execute", args=("receive_weights",)
319 - )
213 + self._collective_rpc("elastic_ep_execute", args=("receive_weights",))
214 self.engine_core.available_gpu_memory_for_kv_cache = (
215 ParallelConfig.sync_kv_cache_memory_size(self.new_dp_group, -1)
216 )
323 - self.model_executor.collective_rpc(
324 - "elastic_ep_execute", args=("prepare_new_worker",)
325 - )
217 + self._collective_rpc("elastic_ep_execute", args=("prepare_new_worker",))
218 self.state = ScaleUpNewEngineState.PREPARE
219 return True
220
221 elif state == ScaleUpNewEngineState.PREPARE:
222 + self._collective_rpc("elastic_ep_execute", args=("warmup_local_kernels",))
223 + self._mark_ready_for_switch()
224 tensor = torch.tensor([0, 0, 0], dtype=torch.int32, device="cpu")
225 torch.distributed.all_reduce(
226 tensor,
@@ -337,22 +231,7 @@ class ElasticEPScalingState:
231 self.engine_core.engines_running = bool(data[0])
232 self.engine_core.current_wave = int(data[1])
233 self.engine_core.step_counter = int(data[2])
340 - self.state = ScaleUpNewEngineState.EPLB_RESHUFFLE
341 - self.new_dp_store.add("eep_barrier_engine_count", 1)
342 - return True
343 -
344 - elif state == ScaleUpNewEngineState.EPLB_RESHUFFLE:
345 - if (
346 - int(self.new_dp_store.get("eep_barrier_engine_count"))
347 - < self.new_dp_group.size()
348 - ):
349 - return False
350 - if not self._staged_barrier(
351 - use_new_group=True, barrier_name="eplb_reshuffle"
352 - ):
353 - return False
354 - assert self.new_dp_group.rank() > 0
355 - self._eplb_reshuffle()
234 + self._collective_rpc("elastic_ep_execute", args=("commit_scale_up", False))
235 self.state = ScaleUpNewEngineState.COMPLETE
236 return True
237
@@ -362,38 +241,23 @@ class ElasticEPScalingState:
241
242 def _progress_remaining_engine(self) -> bool:
243 state = self.state
365 - assert self.old_dp_group is not None and self.old_dp_store is not None
244 + assert self.old_dp_group is not None
245
246 if state == ScaleDownRemainingEngineState.PREPARE:
368 - self.state = ScaleDownRemainingEngineState.EPLB_RESHUFFLE
369 - self.old_dp_store.add("eep_barrier_engine_count", 1)
370 - return True
247 + if self._create_standby_groups():
248 + self.state = ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN
249 + self._mark_ready_for_switch()
250 + return True
251 + return False
252
372 - elif state == ScaleDownRemainingEngineState.EPLB_RESHUFFLE:
373 - if (
374 - int(self.old_dp_store.get("eep_barrier_engine_count"))
375 - < self.old_dp_group.size()
376 - ):
377 - return False
378 - if not self._staged_barrier(
379 - use_new_group=False, barrier_name="eplb_reshuffle"
380 - ):
253 + elif state == ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN:
254 + if not self.commit_requested:
255 return False
382 - if self.old_dp_group.rank() == 0:
383 - self.old_dp_store.delete_key("eep_barrier_engine_count")
384 - self._eplb_reshuffle_before_scale_down()
385 - self.state = ScaleDownRemainingEngineState.SWITCH_AND_PREPARE
386 - # NOTE(yongji): currently, after EPLB reshuffle
387 - # that redistributes experts to remaining workers, workers
388 - # to be removed will immediately initiate shutdown;
389 - # existing workers can no longer execute forward steps using
390 - # the old setup. In the future, we may keep
391 - # the removing workers alive a bit longer,
392 - # e.g., to drain in-batch requests.
393 - self._create_standby_groups()
394 - self._switch_and_prepare()
256 + self._commit_scale_down(removing=False)
257 + self._commit_new_dp_group()
258 self._update_parallel_config()
259 self.state = ScaleDownRemainingEngineState.COMPLETE
260 + self._send_reconfigure_finished()
261 return True
262
263 else:
@@ -402,26 +266,11 @@ class ElasticEPScalingState:
266
267 def _progress_removing_engine(self) -> bool:
268 state = self.state
405 - assert self.old_dp_group is not None and self.old_dp_store is not None
269 + assert self.old_dp_group is not None
270
271 if state == ScaleDownRemovingEngineState.PREPARE:
408 - self.state = ScaleDownRemovingEngineState.EPLB_RESHUFFLE
409 - self.old_dp_store.add("eep_barrier_engine_count", 1)
410 - return True
411 -
412 - if state == ScaleDownRemovingEngineState.EPLB_RESHUFFLE:
413 - if (
414 - int(self.old_dp_store.get("eep_barrier_engine_count"))
415 - < self.old_dp_group.size()
416 - ):
417 - return False
418 - if not self._staged_barrier(
419 - use_new_group=False, barrier_name="eplb_reshuffle"
420 - ):
421 - return False
272 assert self.old_dp_group.rank() > 0
423 - self._eplb_reshuffle_before_scale_down()
424 - self._switch_and_remove()
273 + self._commit_scale_down(removing=True)
274 self.state = ScaleDownRemovingEngineState.COMPLETE
275 self.engine_core._eep_send_engine_core_notification(
276 EEPNotificationType.SHUTDOWN_COMPLETE
@@ -432,22 +281,22 @@ class ElasticEPScalingState:
281 assert self.state == ScaleDownRemovingEngineState.COMPLETE
282 return True
283
435 - def handle_notification(self, notification_type: EEPNotificationType):
436 - assert self.worker_type != "new"
437 - assert self.old_dp_store is not None
438 - if (
439 - notification_type == EEPNotificationType.NEW_CORE_ENGINES_INIT_READY
440 - and self.state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT
441 - ):
442 - self.old_dp_store.add("eep_barrier_engine_count", 1)
443 - self.state = ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS
444 - elif (
445 - notification_type == EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY
446 - and self.state
447 - == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT
448 - ):
449 - self.old_dp_store.add("eep_barrier_engine_count", 1)
450 - self.state = ScaleUpExistingEngineState.TRANSFER_WEIGHTS
284 + def is_ready_for_switch(self) -> bool:
285 + return self.worker_type == "existing" and (
286 + self.state is ScaleUpExistingEngineState.COMMIT_SCALE_UP
287 + or self.state is ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN
288 + )
289 +
290 + @property
291 + def ready_key(self) -> str:
292 + return f"eep_ready/{self.engine_core.dp_rank}"
293 +
294 + def _mark_ready_for_switch(self) -> None:
295 + parallel_config = self.new_parallel_config
296 + get_cached_tcp_store_client(
297 + parallel_config.data_parallel_master_ip,
298 + parallel_config._coord_store_port,
299 + ).set(self.ready_key, b"1")
300
301 def is_complete(self) -> bool:
302 if self.scale_type == "scale_up":
@@ -462,50 +311,78 @@ class ElasticEPScalingState:
311 else self.state == ScaleDownRemainingEngineState.COMPLETE
312 )
313
465 - def _create_standby_groups(self):
314 + def _init_new_dp_group(self) -> tuple[Any, Any]:
315 + return self.new_parallel_config.stateless_init_dp_group(return_store=True)
316 +
317 + def _ensure_new_dp_group(self) -> bool:
318 + if self.new_dp_group is not None:
319 + return True
320 +
321 + if self._prepare_future is None:
322 + self._prepare_future = self._prepare_executor.submit(
323 + self._init_new_dp_group
324 + )
325 + if not self._prepare_future.done():
326 + return False
327 +
328 + self.new_dp_group, self.new_dp_store = self._prepare_future.result()
329 + self._prepare_future = None
330 + return True
331 +
332 + def _create_standby_groups(self) -> bool:
333 assert self.old_dp_group is not None
467 - self.new_dp_group, self.new_dp_store = (
468 - self.new_parallel_config.stateless_init_dp_group(return_store=True)
469 - )
470 - self.model_executor.collective_rpc(
471 - "elastic_ep_execute", args=("create_standby_groups", self.reconfig_request)
472 - )
334 + if not self._ensure_new_dp_group():
335 + return False
336 + if not self._execute_async(
337 + "create_standby_groups",
338 + self.reconfig_request,
339 + self.new_parallel_config.use_all2all,
340 + ):
341 + return False
342 if self.old_dp_group.rank() == 0:
343 logger.info("[Elastic EP] Created standby communication groups")
344 + return True
345
476 - def _transfer_weights(self):
346 + def _transfer_weights(self) -> bool:
347 assert self.reconfig_request is not None and self.old_dp_group is not None
348 old_dp_size = self.old_dp_group.size()
349 new_dp_size = self.reconfig_request.new_data_parallel_size
350
481 - self.model_executor.collective_rpc(
482 - "elastic_ep_execute", args=("transfer_weights", old_dp_size, new_dp_size)
483 - )
351 + if not self._execute_async("transfer_weights", old_dp_size, new_dp_size):
352 + return False
353 if self.old_dp_group.rank() == 0:
354 logger.info("[Elastic EP] Transferred weights to new workers")
355 + return True
356
487 - def _transfer_expert_mapping(self):
488 - assert self.old_dp_group is not None
489 - self.model_executor.collective_rpc(
490 - "elastic_ep_execute", args=("broadcast_expert_mapping",)
491 - )
492 - if self.old_dp_group.rank() == 0:
493 - logger.info("[Elastic EP] Broadcasted expert mapping to new workers")
494 -
495 - def _sync_kv_cache_memory_size(self):
357 + def _sync_kv_cache_memory_size(self) -> bool:
358 assert self.engine_core.available_gpu_memory_for_kv_cache > 0
359 assert self.new_dp_group is not None and self.old_dp_group is not None
498 - ParallelConfig.sync_kv_cache_memory_size(
499 - self.new_dp_group,
500 - self.engine_core.available_gpu_memory_for_kv_cache,
501 - )
360 +
361 + if self._new_dp_sync is None:
362 + tensor = torch.tensor(
363 + [self.engine_core.available_gpu_memory_for_kv_cache],
364 + dtype=torch.int64,
365 + device="cpu",
366 + )
367 + work = torch.distributed.all_reduce(
368 + tensor,
369 + op=torch.distributed.ReduceOp.MIN,
370 + group=self.new_dp_group,
371 + async_op=True,
372 + )
373 + self._new_dp_sync = (tensor, work)
374 + return False
375 +
376 + _, work = self._new_dp_sync
377 + if not work.is_completed():
378 + return False
379 + work.wait()
380 + self._new_dp_sync = None
381 if self.old_dp_group.rank() == 0:
382 logger.info("[Elastic EP] Synced KV cache memory size to new workers")
383 + return True
384
505 - def _switch_and_prepare(self):
506 - self.model_executor.collective_rpc(
507 - "elastic_ep_execute", args=("switch_and_prepare",)
508 - )
385 + def _commit_new_dp_group(self):
386 old_dp_group = self.old_dp_group
387 stateless_destroy_torch_distributed_process_group(old_dp_group)
388 assert self.new_dp_group is not None
@@ -529,41 +406,28 @@ class ElasticEPScalingState:
406 self.engine_core.current_wave = int(data[1])
407 self.engine_core.step_counter = int(data[2])
408 if new_dp_group.rank() == 0:
532 - self.engine_core._eep_send_engine_core_notification(
533 - EEPNotificationType.RECONFIGURE_FINISHED
534 - )
409 logger.info("[Elastic EP] Switched to new setup")
410
537 - def _eplb_reshuffle(self):
538 - self.model_executor.collective_rpc(
539 - "elastic_ep_execute", args=("perform_eplb_reshuffle",)
540 - )
541 - # Reshuffle changes per-rank token routing; the locked MoE workspace
542 - # may now be too small. Rewarm covers both new and existing engines.
543 - self.model_executor.collective_rpc(
544 - "elastic_ep_execute", args=("rewarm_workspace",)
545 - )
411 + def _send_reconfigure_finished(self):
412 assert self.new_dp_group is not None
413 if self.new_dp_group.rank() == 0:
548 - logger.info("[Elastic EP] EPLB reshuffle completed")
414 + self.engine_core._eep_send_engine_core_notification(
415 + EEPNotificationType.RECONFIGURE_FINISHED
416 + )
417
550 - def _eplb_reshuffle_before_scale_down(self):
418 + def _commit_scale_down(self, removing: bool):
419 assert self.reconfig_request is not None and self.old_dp_group is not None
552 - self.model_executor.collective_rpc(
420 + self._collective_rpc(
421 "elastic_ep_execute",
422 args=(
555 - "perform_scale_down_eplb_reshuffle",
423 + "commit_scale_down",
424 self.reconfig_request.new_data_parallel_size,
425 + removing,
426 ),
427 )
428 if self.old_dp_group.rank() == 0:
429 logger.info("[Elastic EP] EPLB reshuffle completed")
430
562 - def _switch_and_remove(self):
563 - self.model_executor.collective_rpc(
564 - "elastic_ep_execute", args=("switch_and_remove",)
565 - )
566 -
431 def _update_parallel_config(self):
432 assert self.reconfig_request is not None
433 reconfig_request = self.reconfig_request
vllm/distributed/elastic_ep/standby_state.py
+2 -1
@@ -39,6 +39,7 @@ def create_standby_groups(
39 new_world_size_across_dp: int,
40 master_ip: str,
41 coord_store_port: int,
42 + use_all2all: bool,
43 enable_eplb: bool = True,
44 backend: str | None = None,
45 ) -> None:
@@ -86,7 +87,7 @@ def create_standby_groups(
87 )
88 standby_ep_ranks = [x.tolist() for x in standby_ep_ranks]
89 _STANDBY_EP = _init_stateless_group(
89 - standby_ep_ranks, "ep", master_ip, backend, coord_store=coord_store
90 + standby_ep_ranks, "ep", master_ip, backend, coord_store, use_all2all=use_all2all
91 )
92
93 if enable_eplb:
vllm/distributed/parallel_state.py
+13 -1
@@ -414,6 +414,7 @@ class GroupCoordinator:
414 use_device_communicator: bool, # whether to use device communicator
415 use_message_queue_broadcaster: bool = False,
416 group_name: str | None = None,
417 + use_all2all: bool = False,
418 ):
419 group_name = group_name or "anonymous"
420 self.unique_name = _get_unique_name(group_name)
@@ -508,6 +509,7 @@ class GroupCoordinator:
509 device=self.device,
510 device_group=self.device_group,
511 unique_name=self.unique_name,
512 + use_all2all=use_all2all,
513 )
514
515 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
@@ -1321,6 +1323,7 @@ def init_model_parallel_group(
1323 use_message_queue_broadcaster: bool = False,
1324 group_name: str | None = None,
1325 use_device_communicator: bool = True,
1326 + use_all2all: bool = False,
1327 ) -> GroupCoordinator:
1328 return GroupCoordinator(
1329 group_ranks=group_ranks,
@@ -1329,6 +1332,7 @@ def init_model_parallel_group(
1332 use_device_communicator=use_device_communicator,
1333 use_message_queue_broadcaster=use_message_queue_broadcaster,
1334 group_name=group_name,
1335 + use_all2all=use_all2all,
1336 )
1337
1338
@@ -1339,6 +1343,7 @@ def _init_stateless_group(
1343 backend: str,
1344 coord_store: Store,
1345 use_device_communicator: bool = True,
1346 + use_all2all: bool = False,
1347 ) -> "StatelessGroupCoordinator":
1348 """Create a StatelessGroupCoordinator with the given parameters."""
1349 from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
@@ -1354,6 +1359,7 @@ def _init_stateless_group(
1359 coord_store=coord_store,
1360 global_rank=world.rank,
1361 global_world_size=world.world_size,
1362 + use_all2all=use_all2all,
1363 )
1364
1365
@@ -1924,6 +1930,7 @@ def initialize_model_parallel(
1930 .unbind(0)
1931 )
1932 group_ranks = [x.tolist() for x in group_ranks]
1933 + use_all2all = parallel_config.use_all2all
1934 if enable_elastic_ep:
1935 _EP = _init_stateless_group(
1936 group_ranks,
@@ -1931,10 +1938,15 @@ def initialize_model_parallel(
1938 parallel_config.data_parallel_master_ip,
1939 backend,
1940 coord_store=coord_store,
1941 + use_all2all=use_all2all,
1942 )
1943 else:
1944 _EP = init_model_parallel_group(
1937 - group_ranks, get_world_group().local_rank, backend, group_name="ep"
1945 + group_ranks,
1946 + get_world_group().local_rank,
1947 + backend,
1948 + group_name="ep",
1949 + use_all2all=use_all2all,
1950 )
1951
1952 # Create EPLB group with the same ranks as EP if EPLB is enabled.
vllm/distributed/stateless_coordinator.py
+2
@@ -79,6 +79,7 @@ class StatelessGroupCoordinator(GroupCoordinator):
79 host: str = "127.0.0.1",
80 global_rank: int = 0,
81 global_world_size: int = 1,
82 + use_all2all: bool = False,
83 ):
84 group_name = group_name or "anonymous"
85 self.unique_name = _get_unique_name(group_name)
@@ -191,6 +192,7 @@ class StatelessGroupCoordinator(GroupCoordinator):
192 global_ranks=self.ranks,
193 global_world_size=global_world_size,
194 tcp_store_group=self.tcp_store_group,
195 + use_all2all=use_all2all,
196 )
197
198 self.mq_broadcaster = None
vllm/entrypoints/serve/elastic_ep/api_router.py
+1 -8
@@ -12,10 +12,7 @@ from vllm.engine.protocol import EngineClient
12 from vllm.entrypoints.openai.engine.protocol import (
13 ErrorResponse,
14 )
15 -from vllm.entrypoints.serve.elastic_ep.middleware import (
16 - get_scaling_elastic_ep,
17 - set_scaling_elastic_ep,
18 -)
15 +from vllm.entrypoints.serve.elastic_ep.middleware import get_scaling_elastic_ep
16 from vllm.entrypoints.serve.utils.api_utils import validate_json_request
17 from vllm.logger import init_logger
18
@@ -64,8 +61,6 @@ async def scale_elastic_ep(raw_request: Request):
61 status_code=400, detail="drain_timeout must be a positive integer"
62 )
63
67 - # Set scaling flag to prevent new requests
68 - set_scaling_elastic_ep(True)
64 client = engine_client(raw_request)
65 try:
66 await client.scale_elastic_ep(new_data_parallel_size, drain_timeout)
@@ -83,8 +78,6 @@ async def scale_elastic_ep(raw_request: Request):
78 except Exception as e:
79 logger.error("Scale failed: %s", e)
80 raise HTTPException(status_code=500, detail="Scale failed") from e
86 - finally:
87 - set_scaling_elastic_ep(False)
81
82
83 @router.post("/is_scaling_elastic_ep")
vllm/model_executor/warmup/kernel_warmup.py
+17 -14
@@ -93,7 +93,7 @@ def _warmup_ll_bf16_router_gemm(model: torch.nn.Module) -> None:
93 )
94
95
96 -def kernel_warmup(worker: "Worker"):
96 +def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
97 from vllm.model_executor.warmup.minimax_m3_msa_warmup import (
98 minimax_m3_msa_warmup,
99 )
@@ -118,6 +118,22 @@ def kernel_warmup(worker: "Worker"):
118 )
119
120 # Run next so input-prep kernels JIT against pristine runner state.
121 + if worker.vllm_config.kernel_config.enable_jit_warmup:
122 + fa4_cutedsl_warmup(worker)
123 + sparse_mla_triton_warmup(worker)
124 +
125 + if current_platform.has_device_capability(90):
126 + _warmup_ll_bf16_router_gemm(worker.get_model())
127 +
128 + if worker.vllm_config.kernel_config.enable_cutedsl_warmup:
129 + # TODO(roberto): Remove after registered CuTeDSL warmups are migrated
130 + # to the shared JIT warmup infrastructure.
131 + # https://github.com/vllm-project/vllm/pull/47451
132 + cutedsl_warmup()
133 +
134 + if process_local_only:
135 + return
136 +
137 flashinfer_sparse_mla_decode_autotune_warmup(worker)
138 deepseek_v4_sparse_mla_attention_warmup(worker)
139
@@ -143,9 +159,6 @@ def kernel_warmup(worker: "Worker"):
159 elif has_flashinfer() and current_platform.has_device_capability(90):
160 flashinfer_autotune(worker.model_runner)
161
146 - if current_platform.has_device_capability(90):
147 - _warmup_ll_bf16_router_gemm(worker.get_model())
148 -
162 # FlashInfer attention warmup
163 # Only warmup if the model has FlashInfer attention groups
164 # and is not a pooling model
@@ -178,16 +191,6 @@ def kernel_warmup(worker: "Worker"):
191 create_mixed_batch=True,
192 )
193
181 - if worker.vllm_config.kernel_config.enable_cutedsl_warmup:
182 - # TODO(roberto): Remove after registered CuTeDSL warmups are migrated
183 - # to the shared JIT warmup infrastructure.
184 - # https://github.com/vllm-project/vllm/pull/47451
185 - cutedsl_warmup()
186 -
187 - if worker.vllm_config.kernel_config.enable_jit_warmup:
188 - fa4_cutedsl_warmup(worker)
189 - sparse_mla_triton_warmup(worker)
190 -
194
195 def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None:
196 if envs.VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS is not None:
vllm/v1/engine/__init__.py
-2
@@ -35,8 +35,6 @@ FT_STATUS_CALL_ID = -2
35
36
37 class EEPNotificationType(enum.Enum):
38 - NEW_CORE_ENGINES_INIT_READY = "NEW_CORE_ENGINES_INIT_READY"
39 - NEW_CORE_ENGINES_WEIGHTS_INIT_READY = "NEW_CORE_ENGINES_WEIGHTS_INIT_READY"
38 RECONFIGURE_FINISHED = "RECONFIGURE_FINISHED"
39 SHUTDOWN_COMPLETE = "SHUTDOWN_COMPLETE"
40
vllm/v1/engine/async_llm.py
+26 -19
@@ -109,6 +109,7 @@ class AsyncLLM(EngineClient):
109 maybe_register_config_serialize_by_value()
110
111 self.vllm_config = vllm_config
112 + self._elastic_ep_lock = asyncio.Lock()
113 self.model_config = vllm_config.model_config
114 self.observability_config = vllm_config.observability_config
115
@@ -998,17 +999,27 @@ class AsyncLLM(EngineClient):
999 "waiting for requests to drain."
1000 )
1001
1002 + async def _drain_requests_for_elastic_ep(self, drain_timeout: int) -> None:
1003 + try:
1004 + logger.info(
1005 + "VLLM_ELASTIC_EP_DRAIN_REQUESTS is set, "
1006 + "waiting for requests to drain before scaling"
1007 + )
1008 + await self.wait_for_requests_to_drain(drain_timeout)
1009 + except BaseException:
1010 + set_scaling_elastic_ep(False)
1011 + raise
1012 +
1013 async def scale_elastic_ep(
1014 self, new_data_parallel_size: int, drain_timeout: int = 300
1015 ):
1004 - """
1005 - Scale up or down the data parallel size by adding or removing
1006 - engine cores.
1007 - Args:
1008 - new_data_parallel_size: The new number of data parallel workers
1009 - drain_timeout:
1010 - Maximum time to wait for requests to drain (seconds)
1011 - """
1016 + """Scale the elastic EP data parallel size."""
1017 + async with self._elastic_ep_lock:
1018 + await self._scale_elastic_ep(new_data_parallel_size, drain_timeout)
1019 +
1020 + async def _scale_elastic_ep(
1021 + self, new_data_parallel_size: int, drain_timeout: int
1022 + ) -> None:
1023 old_data_parallel_size = self.vllm_config.parallel_config.data_parallel_size
1024 if old_data_parallel_size == new_data_parallel_size:
1025 logger.info(
@@ -1017,12 +1028,7 @@ class AsyncLLM(EngineClient):
1028 )
1029 return
1030
1020 - if envs.VLLM_ELASTIC_EP_DRAIN_REQUESTS:
1021 - logger.info(
1022 - "VLLM_ELASTIC_EP_DRAIN_REQUESTS is set, "
1023 - "waiting for requests to drain before scaling"
1024 - )
1025 - await self.wait_for_requests_to_drain(drain_timeout)
1031 + await self.engine_core.prepare_elastic_ep(new_data_parallel_size)
1032
1033 # recreate stat loggers
1034 if new_data_parallel_size > old_data_parallel_size and self.log_stats:
@@ -1042,11 +1048,12 @@ class AsyncLLM(EngineClient):
1048 self.logger_manager.log_engine_initialized()
1049
1050 set_scaling_elastic_ep(True)
1045 - try:
1046 - await self.engine_core.scale_elastic_ep(new_data_parallel_size)
1047 - self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
1048 - finally:
1049 - set_scaling_elastic_ep(False)
1051 + if envs.VLLM_ELASTIC_EP_DRAIN_REQUESTS:
1052 + await self._drain_requests_for_elastic_ep(drain_timeout)
1053 +
1054 + await self.engine_core.commit_elastic_ep()
1055 + self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
1056 + set_scaling_elastic_ep(False)
1057
1058 async def handle_fault(
1059 self, fault_tolerance_request: FaultToleranceRequest
vllm/v1/engine/coordinator.py
+3
@@ -376,6 +376,9 @@ class DPCoordinatorProc:
376 eng_index = outputs.engine_index
377 scheduler_stats = outputs.scheduler_stats
378 if scheduler_stats:
379 + # Elastic EP stats may arrive while the engine list changes.
380 + if eng_index >= len(self.engines):
381 + continue
382 # 1. Updated request load stats - update our local
383 # state with these.
384 stats = self.engines[eng_index].request_counts
vllm/v1/engine/core.py
+36 -39
@@ -327,8 +327,9 @@ class EngineCore:
327
328 vllm_config.validate_block_size()
329
330 - # Initialize kv cache and warmup the execution
330 self.model_executor.initialize_from_config(kv_cache_configs)
331 + if not envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
332 + self.model_executor.compile_or_warm_up_model()
333
334 elapsed = time.time() - start
335 compile_time = vllm_config.compilation_config.compilation_time
@@ -993,9 +994,7 @@ class EngineCore:
994 raise NotImplementedError
995
996 def _eep_send_engine_core_notification(
996 - self,
997 - notification_type: EEPNotificationType,
998 - vllm_config: VllmConfig | None = None,
997 + self, notification_type: EEPNotificationType
998 ):
999 raise NotImplementedError
1000
@@ -1070,11 +1069,6 @@ class EngineCoreProc(EngineCore):
1069
1070 self.addresses = addresses
1071 self.process_input_queue_block = True
1073 - if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
1074 - self._eep_send_engine_core_notification(
1075 - EEPNotificationType.NEW_CORE_ENGINES_INIT_READY,
1076 - vllm_config=vllm_config,
1077 - )
1072 self._init_data_parallel(vllm_config)
1073
1074 super().__init__(
@@ -2096,12 +2090,16 @@ class DPEngineCoreProc(EngineCoreProc):
2090 self._maybe_publish_request_counts()
2091
2092 if self.eep_scaling_state is not None:
2099 - _ = self.eep_scaling_state.progress()
2100 - if self.eep_scaling_state.is_complete():
2101 - if self.eep_scaling_state.worker_type == "removing":
2093 + state = self.eep_scaling_state
2094 + if state.commit_requested or not state.is_ready_for_switch():
2095 + state.progress()
2096 + if state.is_complete():
2097 + if state.worker_type == "removing":
2098 raise SystemExit
2099 self.process_input_queue_block = True
2100 self.eep_scaling_state = None
2101 + elif not state.commit_requested and state.is_ready_for_switch():
2102 + self.process_input_queue_block = True
2103
2104 executed = self._process_engine_step()
2105 self._maybe_publish_request_counts()
@@ -2171,7 +2169,7 @@ class DPEngineCoreProc(EngineCoreProc):
2169
2170 def reinitialize_distributed(
2171 self, reconfig_request: ReconfigureDistributedRequest
2174 - ) -> None:
2172 + ) -> str:
2173 from copy import deepcopy
2174
2175 from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState
@@ -2203,7 +2201,10 @@ class DPEngineCoreProc(EngineCoreProc):
2201 == ReconfigureRankType.SHUTDOWN_CURRENT_RANK
2202 )
2203
2206 - self.eep_scaling_state = ElasticEPScalingState(
2204 + if self.eep_scaling_state is not None:
2205 + raise RuntimeError("Elastic EP reconfiguration is already active")
2206 +
2207 + state = ElasticEPScalingState(
2208 model_executor=self.model_executor,
2209 engine_core=self,
2210 vllm_config=self.vllm_config,
@@ -2212,30 +2213,34 @@ class DPEngineCoreProc(EngineCoreProc):
2213 scale_type="scale_down" if is_scale_down else "scale_up",
2214 reconfig_request=reconfig_request,
2215 )
2216 + self.eep_scaling_state = state
2217 +
2218 self.process_input_queue_block = False
2219 logger.info(
2220 "[Elastic EP] Received reconfiguration request and starting scaling up/down"
2221 )
2222 + return state.ready_key
2223 +
2224 + def commit_prepared_elastic_ep(self) -> None:
2225 + state = self.eep_scaling_state
2226 + if state is None or state.commit_requested or not state.is_ready_for_switch():
2227 + raise RuntimeError("No prepared Elastic EP reconfiguration is ready")
2228 + state.commit_requested = True
2229 + self.process_input_queue_block = False
2230 + logger.info("[Elastic EP] Committing prepared reconfiguration")
2231
2232 def _eep_send_engine_core_notification(
2221 - self,
2222 - notification_type: EEPNotificationType,
2223 - vllm_config: VllmConfig | None = None,
2233 + self, notification_type: EEPNotificationType
2234 ):
2235 """
2236 Send notifications to EngineCoreClient, which can then forward
2237 the notifications to other engine core processes. It is used for:
2228 - 1) In scale up: new core engines to notify existing core engines
2229 - that they are ready;
2230 - 2) In scale down: removing core engines to notify EngineCoreClient
2238 + 1) In scale down: removing core engines to notify EngineCoreClient
2239 so EngineCoreClient can release their ray placement groups;
2232 - 3) Both scale up/down: to notify EngineCoreClient that existing
2240 + 2) Both scale up/down: to notify EngineCoreClient that existing
2241 core engines have already switched to the new parallel setup.
2242 """
2235 - if vllm_config is None:
2236 - dp_rank = self.vllm_config.parallel_config.data_parallel_rank
2237 - else:
2238 - dp_rank = vllm_config.parallel_config.data_parallel_rank
2243 + dp_rank = self.vllm_config.parallel_config.data_parallel_rank
2244 notification_data = (notification_type.value, dp_rank)
2245 outputs = EngineCoreOutputs(
2246 utility_output=UtilityOutput(
@@ -2257,22 +2262,11 @@ class DPEngineCoreProc(EngineCoreProc):
2262 ):
2263 socket.send_multipart(encoder.encode(outputs))
2264
2260 - def eep_handle_engine_core_notification(
2261 - self, notification_type: str | EEPNotificationType
2262 - ):
2263 - """
2264 - Handle notification received from EngineCoreClient
2265 - (forwarded from new core engines).
2266 - """
2267 - assert self.eep_scaling_state is not None
2268 - if isinstance(notification_type, str):
2269 - notification_type = EEPNotificationType(notification_type)
2270 - self.eep_scaling_state.handle_notification(notification_type)
2271 -
2265 def _eep_scale_up_before_kv_init(self):
2266 from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState
2267
2275 - self.eep_scaling_state = ElasticEPScalingState(
2268 + self.ignore_start_dp_wave = True
2269 + state = ElasticEPScalingState(
2270 model_executor=self.model_executor,
2271 engine_core=self,
2272 vllm_config=self.vllm_config,
@@ -2281,7 +2275,10 @@ class DPEngineCoreProc(EngineCoreProc):
2275 scale_type="scale_up",
2276 reconfig_request=None,
2277 )
2284 - self.eep_scaling_state.run_pre_kv_init_states()
2278 + if self.eep_scaling_state is not None:
2279 + raise RuntimeError("Elastic EP reconfiguration is already active")
2280 + self.eep_scaling_state = state
2281 + state.run_pre_kv_init_states()
2282 self.process_input_queue_block = False
2283
2284
vllm/v1/engine/core_client.py
+175 -137
@@ -225,7 +225,10 @@ class EngineCoreClient(ABC):
225 running state."""
226 raise NotImplementedError
227
228 - async def scale_elastic_ep(self, new_data_parallel_size: int) -> None:
228 + async def commit_elastic_ep(self) -> None:
229 + raise NotImplementedError
230 +
231 + async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None:
232 raise NotImplementedError
233
234 async def get_output_async(self) -> EngineCoreOutputs:
@@ -1490,6 +1493,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1493 )
1494
1495 assert len(self.core_engines) > 1
1496 + self._prepared_elastic_ep: tuple[int, int] | None = None
1497
1498 self.eng_start_index = (
1499 len(self.core_engines) * self.client_index
@@ -1603,31 +1607,16 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1607 if len(cache.pending_notifications[notification_type]) >= abs(
1608 cache.num_new_core_engines
1609 ):
1606 - if notification_type == EEPNotificationType.SHUTDOWN_COMPLETE:
1607 - assert isinstance(self.resources.engine_manager, CoreEngineActorManager)
1608 - assert cache.num_new_core_engines < 0
1609 - old_dp_size = len(cache.existing_core_engines)
1610 - new_dp_size = old_dp_size + cache.num_new_core_engines
1611 - self.resources.engine_manager.scale_down_elastic_ep(
1612 - old_dp_size, new_dp_size
1613 - )
1614 - else:
1615 - await asyncio.gather(
1616 - *[
1617 - self._call_utility_async(
1618 - "eep_handle_engine_core_notification",
1619 - notification_type,
1620 - engine=engine,
1621 - )
1622 - for engine in cache.existing_core_engines
1623 - ]
1624 - )
1625 - cache.pending_notifications[notification_type] = set()
1626 - if notification_type in [
1627 - EEPNotificationType.SHUTDOWN_COMPLETE,
1628 - EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY,
1629 - ]:
1630 - self.eep_scaling_cache = None
1610 + engine_manager = self.resources.engine_manager
1611 + assert isinstance(engine_manager, CoreEngineActorManager)
1612 + assert cache.num_new_core_engines < 0
1613 + old_dp_size = len(cache.existing_core_engines)
1614 + new_dp_size = old_dp_size + cache.num_new_core_engines
1615 + engine_manager.scale_down_elastic_ep(old_dp_size, new_dp_size)
1616 + self.vllm_config.parallel_config.data_parallel_size_local = len(
1617 + engine_manager.local_engine_actors
1618 + )
1619 + self.eep_scaling_cache = None
1620
1621 async def abort_requests_async(self, request_ids: list[str]) -> None:
1622 if not request_ids or self.resources.engine_dead:
@@ -1651,31 +1640,46 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1640 ) -> None:
1641 await self._send_input(EngineCoreRequestType.ABORT, request_ids, engine)
1642
1654 - async def scale_elastic_ep(self, new_data_parallel_size: int) -> None:
1655 - """Scale elastic EP data parallel size"""
1643 + async def commit_elastic_ep(self) -> None:
1644 + """Commit prepared elastic EP scaling."""
1645 + prepared = self._prepared_elastic_ep
1646 + if prepared is None:
1647 + raise RuntimeError("Elastic EP scaling has not been prepared")
1648 + new_data_parallel_size, num_redundant_experts = prepared
1649 cur_data_parallel_size = len(self.core_engines)
1657 -
1658 - assert new_data_parallel_size != cur_data_parallel_size, (
1659 - f"new_data_parallel_size {new_data_parallel_size} must be "
1660 - f"different from cur_data_parallel_size {cur_data_parallel_size}"
1650 + if new_data_parallel_size > cur_data_parallel_size:
1651 + await self._commit_scale_up_elastic_ep(new_data_parallel_size)
1652 + else:
1653 + await self._commit_scale_down_elastic_ep(new_data_parallel_size)
1654 + self.vllm_config.parallel_config.eplb_config.num_redundant_experts = (
1655 + num_redundant_experts
1656 )
1657 + self._prepared_elastic_ep = None
1658
1659 + async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None:
1660 + """Prepare elastic EP scaling without routing requests to new engines."""
1661 + if (prepared := self._prepared_elastic_ep) is not None:
1662 + if prepared[0] == new_data_parallel_size:
1663 + return
1664 + raise RuntimeError("Elastic EP scaling is already prepared")
1665 + cur_data_parallel_size = len(self.core_engines)
1666 assert self.vllm_config.parallel_config.data_parallel_backend == "ray", (
1667 "Only ray DP backend supports scaling elastic EP"
1668 )
1666 -
1667 - scale_up = new_data_parallel_size > cur_data_parallel_size
1668 -
1669 - if scale_up:
1670 - await self._scale_up_elastic_ep(
1671 - cur_data_parallel_size, new_data_parallel_size
1672 - )
1669 + parallel_config = self.vllm_config.parallel_config
1670 + num_experts = self.vllm_config.model_config.get_num_experts()
1671 + num_redundant_experts = (
1672 + num_experts + parallel_config.eplb_config.num_redundant_experts
1673 + ) * new_data_parallel_size // cur_data_parallel_size - num_experts
1674 + if new_data_parallel_size < cur_data_parallel_size:
1675 + await self._prepare_scale_down_elastic_ep(new_data_parallel_size)
1676 else:
1674 - await self._scale_down_elastic_ep(
1675 - cur_data_parallel_size, new_data_parallel_size
1677 + await self._prepare_scale_up_elastic_ep(
1678 + new_data_parallel_size, num_redundant_experts
1679 )
1680 + self._prepared_elastic_ep = new_data_parallel_size, num_redundant_experts
1681
1678 - async def _eep_wait_for_setup_switch_complete(self) -> None:
1682 + def _eep_wait_for_setup_switch_complete(self) -> asyncio.Future:
1683 """
1684 Wait for core engines to switch to the new setup.
1685
@@ -1687,9 +1691,26 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1691 future = asyncio.get_running_loop().create_future()
1692 self.utility_results[EEP_NOTIFICATION_CALL_ID] = future
1693 self._ensure_output_queue_task()
1690 - await future
1694 + return future
1695 +
1696 + def _wait_for_new_engine_ready(self, new_core_engines: list[bytes]) -> None:
1697 + new_engine_identities = set(new_core_engines)
1698 + sync_input_socket = zmq.Socket.shadow(self.input_socket)
1699 + while new_engine_identities:
1700 + if not sync_input_socket.poll(timeout=VLLM_ENGINE_READY_TIMEOUT_S * 1000):
1701 + raise TimeoutError(
1702 + f"Timed out waiting for new engine core processes to "
1703 + f"start. Waited "
1704 + f"{VLLM_ENGINE_READY_TIMEOUT_S}s (configured by "
1705 + f"VLLM_ENGINE_READY_TIMEOUT_S). To increase the "
1706 + f"timeout, set the environment variable: "
1707 + f"VLLM_ENGINE_READY_TIMEOUT_S=<seconds>"
1708 + )
1709 + identity, payload = sync_input_socket.recv_multipart()
1710 + new_engine_identities.discard(identity)
1711 + self._apply_ready_response(payload)
1712
1692 - def _setup_elastic_ep_reconfig_bootstrap(self) -> tuple[str, int]:
1713 + def _setup_elastic_ep_reconfig_bootstrap(self) -> None:
1714 from vllm.distributed.utils import create_tcp_store
1715 from vllm.utils.network_utils import get_open_ports_list
1716
@@ -1709,36 +1730,36 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1730 )
1731 parallel_config._coord_store_port = store.port
1732 self._coord_store = store
1712 - return ip, store.port
1733
1714 - async def _scale_up_elastic_ep(
1715 - self, cur_data_parallel_size: int, new_data_parallel_size: int
1716 - ) -> None:
1717 - """Scale up the data parallel size by creating new engine cores
1718 - and reconfiguring existing ones."""
1719 - cur_data_parallel_size = len(self.core_engines)
1720 -
1721 - self.eep_scaling_cache = ElasticScalingCache(
1722 - existing_core_engines=self.core_engines.copy(),
1723 - num_new_core_engines=new_data_parallel_size - cur_data_parallel_size,
1724 - pending_notifications=dict(),
1734 + def _make_reconfig_request(
1735 + self,
1736 + new_data_parallel_size: int,
1737 + rank_type: ReconfigureRankType = ReconfigureRankType.KEEP_CURRENT_RANK,
1738 + ) -> ReconfigureDistributedRequest:
1739 + parallel_config = self.vllm_config.parallel_config
1740 + return ReconfigureDistributedRequest(
1741 + new_data_parallel_size=new_data_parallel_size,
1742 + new_data_parallel_rank=rank_type,
1743 + new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK,
1744 + new_data_parallel_master_ip=parallel_config.data_parallel_master_ip,
1745 + new_data_parallel_master_port=parallel_config.data_parallel_master_port,
1746 + new_data_parallel_master_port_list=parallel_config._data_parallel_master_port_list,
1747 + coord_store_port=parallel_config._coord_store_port,
1748 )
1749
1727 - parallel_config = self.vllm_config.parallel_config
1728 - ip, coord_store_port = self._setup_elastic_ep_reconfig_bootstrap()
1750 + async def _prepare_scale_up_elastic_ep(
1751 + self,
1752 + new_data_parallel_size: int,
1753 + num_redundant_experts: int,
1754 + ) -> None:
1755 + """Prepare scale up by creating new engine cores and reconfiguring
1756 + existing ones."""
1757 + self._setup_elastic_ep_reconfig_bootstrap()
1758
1759 # Phase 1: Send reconfig messages to existing engines
1760 reconfig_futures = []
1761 for engine in self.core_engines:
1733 - reconfig_request = ReconfigureDistributedRequest(
1734 - new_data_parallel_size=new_data_parallel_size,
1735 - new_data_parallel_rank=ReconfigureRankType.KEEP_CURRENT_RANK,
1736 - new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK,
1737 - new_data_parallel_master_ip=ip,
1738 - new_data_parallel_master_port=parallel_config.data_parallel_master_port,
1739 - new_data_parallel_master_port_list=parallel_config._data_parallel_master_port_list,
1740 - coord_store_port=coord_store_port,
1741 - )
1762 + reconfig_request = self._make_reconfig_request(new_data_parallel_size)
1763 coro = self._call_utility_async(
1764 "reinitialize_distributed", reconfig_request, engine=engine
1765 )
@@ -1746,51 +1767,54 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1767
1768 # Phase 2: Create new engines
1769 assert isinstance(self.resources.engine_manager, CoreEngineActorManager)
1749 - parallel_config.eplb_config.num_redundant_experts = 0
1770 start_new_worker_future = asyncio.to_thread(
1771 self.resources.engine_manager.scale_up_elastic_ep,
1772 self.vllm_config,
1773 new_data_parallel_size,
1774 + num_redundant_experts,
1775 )
1755 - wait_future = self._eep_wait_for_setup_switch_complete()
1776
1777 # Phase 3: Wait for new engines to be created
1778 # and reconfig messages to be received
1779 await asyncio.gather(start_new_worker_future, *reconfig_futures)
1780 + ready_keys = [future.result() for future in reconfig_futures]
1781 + ready_keys.extend(
1782 + f"eep_ready/{rank}"
1783 + for rank in range(len(self.core_engines), new_data_parallel_size)
1784 + )
1785 + await asyncio.to_thread(self._coord_store.wait, ready_keys)
1786 logger.info("[Elastic EP] Successfully started new engines")
1787
1762 - # Create new CoreEngine objects for the new engines
1763 - new_engine_identities = set()
1764 - for i in range(cur_data_parallel_size, new_data_parallel_size):
1765 - new_engine = i.to_bytes(2, "little")
1766 - self.core_engines.append(new_engine)
1767 - # NOTE(yongji): we don't update lb_engines here,
1768 - # we let run_engine_stats_update_task to update it.
1769 - new_engine_identities.add(new_engine)
1770 -
1771 - # Wait for ready messages from new engines on the input socket
1772 - sync_input_socket = zmq.Socket.shadow(self.input_socket)
1773 - while new_engine_identities:
1774 - if not sync_input_socket.poll(
1775 - timeout=VLLM_ENGINE_READY_TIMEOUT_S * 1000 # convert to ms
1776 - ):
1777 - raise TimeoutError(
1778 - f"Timed out waiting for new engine core processes to "
1779 - f"start. Waited "
1780 - f"{VLLM_ENGINE_READY_TIMEOUT_S}s (configured by "
1781 - f"VLLM_ENGINE_READY_TIMEOUT_S). To increase the "
1782 - f"timeout, set the environment variable: "
1783 - f"VLLM_ENGINE_READY_TIMEOUT_S=<seconds>"
1784 - )
1785 - identity, payload = sync_input_socket.recv_multipart()
1786 - new_engine_identities.discard(identity)
1787 - self._apply_ready_response(payload)
1788 + async def _commit_scale_up_elastic_ep(self, new_data_parallel_size: int) -> None:
1789 + new_core_engines = [
1790 + rank.to_bytes(2, "little")
1791 + for rank in range(len(self.core_engines), new_data_parallel_size)
1792 + ]
1793
1789 - # NOTE(yongji): Before we schedule any requests on the new workers,
1790 - # we should wait for them to switch to the new setup.
1791 - await wait_future
1794 + await self.pause_scheduler_async(mode="keep", clear_cache=False)
1795 + wait_future = self._eep_wait_for_setup_switch_complete()
1796 + finish_futures = [
1797 + asyncio.create_task(
1798 + self._call_utility_async("commit_prepared_elastic_ep", engine=engine)
1799 + )
1800 + for engine in self.core_engines
1801 + ]
1802 + try:
1803 + await asyncio.gather(*finish_futures)
1804 + await wait_future
1805 + self._wait_for_new_engine_ready(new_core_engines)
1806 + except Exception:
1807 + wait_future.cancel()
1808 + raise
1809 +
1810 + self.core_engines.extend(new_core_engines)
1811 # Update the parallel config
1793 - self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
1812 + parallel_config = self.vllm_config.parallel_config
1813 + parallel_config.data_parallel_size = new_data_parallel_size
1814 + if isinstance(self.resources.engine_manager, CoreEngineActorManager):
1815 + parallel_config.data_parallel_size_local = len(
1816 + self.resources.engine_manager.local_engine_actors
1817 + )
1818 # Notify coordinator about scale up through existing
1819 # stats_update_task connection
1820 self._ensure_stats_update_task()
@@ -1803,10 +1827,23 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1827 "[Elastic EP] Scale up completed, new data parallel size: %s",
1828 new_data_parallel_size,
1829 )
1830 + await self.resume_scheduler_async()
1831
1807 - async def _scale_down_elastic_ep(
1808 - self, cur_data_parallel_size: int, new_data_parallel_size: int
1809 - ) -> None:
1832 + async def _prepare_scale_down_elastic_ep(self, new_data_parallel_size: int) -> None:
1833 + self._setup_elastic_ep_reconfig_bootstrap()
1834 +
1835 + reconfig_futures = []
1836 + for engine in self.core_engines[:new_data_parallel_size]:
1837 + reconfig_request = self._make_reconfig_request(new_data_parallel_size)
1838 + coro = self._call_utility_async(
1839 + "reinitialize_distributed", reconfig_request, engine=engine
1840 + )
1841 + reconfig_futures.append(asyncio.create_task(coro))
1842 +
1843 + ready_keys = await asyncio.gather(*reconfig_futures)
1844 + await asyncio.to_thread(self._coord_store.wait, ready_keys)
1845 +
1846 + async def _commit_scale_down_elastic_ep(self, new_data_parallel_size: int) -> None:
1847 """Scale down the data parallel size by shutting down and
1848 reconfiguring existing engine cores."""
1849 cur_data_parallel_size = len(self.core_engines)
@@ -1817,50 +1854,51 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1854 pending_notifications=dict(),
1855 )
1856
1820 - parallel_config = self.vllm_config.parallel_config
1821 - ip, coord_store_port = self._setup_elastic_ep_reconfig_bootstrap()
1822 -
1857 + old_core_engines = self.core_engines
1858 + # NOTE(yongji): Immediately stop sending requests to the removing engines.
1859 + self.core_engines = old_core_engines[:new_data_parallel_size]
1860 + self.lb_engines = self.lb_engines[:new_data_parallel_size]
1861 removed_dp_size = cur_data_parallel_size - new_data_parallel_size
1862 + pause_modes = ["keep"] * new_data_parallel_size + ["abort"] * removed_dp_size
1863 + pause_futures = [
1864 + self._call_utility_async("pause_scheduler", mode, False, engine=engine)
1865 + for mode, engine in zip(pause_modes, old_core_engines)
1866 + ]
1867 + await asyncio.gather(*pause_futures)
1868 assert isinstance(self.resources.engine_manager, CoreEngineActorManager)
1869 self.resources.engine_manager.remove_run_refs_for_scale_down(removed_dp_size)
1870 + wait_future = self._eep_wait_for_setup_switch_complete()
1871 reconfig_futures = []
1827 - for cur_dp_rank, engine in enumerate(self.core_engines):
1828 - reconfig_request = ReconfigureDistributedRequest(
1829 - new_data_parallel_size=new_data_parallel_size,
1830 - new_data_parallel_rank=ReconfigureRankType.KEEP_CURRENT_RANK,
1831 - new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK,
1832 - new_data_parallel_master_ip=ip,
1833 - new_data_parallel_master_port=parallel_config.data_parallel_master_port,
1834 - new_data_parallel_master_port_list=parallel_config._data_parallel_master_port_list,
1835 - coord_store_port=coord_store_port,
1836 - )
1837 - if cur_dp_rank >= new_data_parallel_size:
1838 - reconfig_request.new_data_parallel_rank = (
1839 - ReconfigureRankType.SHUTDOWN_CURRENT_RANK
1872 + for cur_dp_rank, engine in enumerate(old_core_engines):
1873 + if cur_dp_rank < new_data_parallel_size:
1874 + coro = self._call_utility_async(
1875 + "commit_prepared_elastic_ep", engine=engine
1876 + )
1877 + else:
1878 + reconfig_request = self._make_reconfig_request(
1879 + new_data_parallel_size,
1880 + ReconfigureRankType.SHUTDOWN_CURRENT_RANK,
1881 + )
1882 + coro = self._call_utility_async(
1883 + "reinitialize_distributed", reconfig_request, engine=engine
1884 )
1841 - coro = self._call_utility_async(
1842 - "reinitialize_distributed", reconfig_request, engine=engine
1843 - )
1885 reconfig_futures.append(asyncio.create_task(coro))
1886
1846 - # NOTE(yongji): Immediately stop sending requests to the removing engines.
1847 - self.core_engines = self.core_engines[:new_data_parallel_size]
1848 - self.lb_engines = self.lb_engines[:new_data_parallel_size]
1849 - wait_future = self._eep_wait_for_setup_switch_complete()
1850 -
1851 - await asyncio.gather(*reconfig_futures)
1887 + try:
1888 + await asyncio.gather(*reconfig_futures)
1889
1853 - self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
1854 - self._ensure_stats_update_task()
1855 - scale_down_marker = msgspec.msgpack.encode(
1856 - ("SCALE_ELASTIC_EP", new_data_parallel_size)
1857 - )
1858 - await self.first_req_send_socket.send(scale_down_marker)
1890 + self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
1891 + self._ensure_stats_update_task()
1892 + scale_down_marker = msgspec.msgpack.encode(
1893 + ("SCALE_ELASTIC_EP", new_data_parallel_size)
1894 + )
1895 + await self.first_req_send_socket.send(scale_down_marker)
1896 + await wait_future
1897 + await self.resume_scheduler_async()
1898 + except Exception:
1899 + wait_future.cancel()
1900 + raise
1901
1860 - # NOTE(yongji): Unlike scaling up,
1861 - # here we don't actually need to wait for the setup switch to complete.
1862 - # We may want to remove it in the future.
1863 - await wait_future
1902 logger.info(
1903 "[Elastic EP] Scale down completed, new data parallel size: %s",
1904 new_data_parallel_size,
vllm/v1/engine/utils.py
+8 -23
@@ -821,7 +821,10 @@ class CoreEngineActorManager:
821 return placement_groups, local_dp_ranks
822
823 def scale_up_elastic_ep(
824 - self, cur_vllm_config: VllmConfig, new_data_parallel_size: int
824 + self,
825 + cur_vllm_config: VllmConfig,
826 + new_data_parallel_size: int,
827 + num_redundant_experts: int,
828 ) -> None:
829 import copy
830
@@ -864,6 +867,9 @@ class CoreEngineActorManager:
867 if new_data_parallel_size > 1:
868 _apply_dp_identity_suffix(dp_vllm_config, rank)
869 dp_vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
870 + dp_vllm_config.parallel_config.eplb_config.num_redundant_experts = (
871 + num_redundant_experts
872 + )
873 dp_vllm_config.parallel_config.placement_group = pg
874
875 # Check if this placement group is on the head node
@@ -906,39 +912,18 @@ class CoreEngineActorManager:
912 self.created_placement_groups.append(pg)
913 self.placement_group_is_local.append(local_client)
914
909 - ray.get(
910 - [
911 - actor.wait_for_init.remote()
912 - for actor in (
913 - self.local_engine_actors[-new_local_engines:]
914 - if new_local_engines > 0
915 - else []
916 - )
917 - + self.remote_engine_actors[
918 - -(len(placement_groups) - new_local_engines) :
919 - ]
920 - ]
921 - )
922 -
915 actors = (
916 self.local_engine_actors[-new_local_engines:]
917 if new_local_engines > 0
918 else []
919 ) + self.remote_engine_actors[-(len(placement_groups) - new_local_engines) :]
920
921 + ray.get([actor.wait_for_init.remote() for actor in actors])
922 for actor in actors:
923 ref = actor.run.remote()
924 self.run_refs.append(ref)
925 self.actor_run_ref_dict[actor] = ref
926
934 - cur_vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
935 - # Update old_vllm_config with new data_parallel_size_local if any new
936 - # local engines were added
937 - if new_local_engines > 0:
938 - cur_vllm_config.parallel_config.data_parallel_size_local += (
939 - new_local_engines
940 - )
941 -
927 def scale_down_elastic_ep(
928 self, cur_data_parallel_size: int, new_data_parallel_size: int
929 ) -> None:
vllm/v1/executor/abstract.py
+4 -4
@@ -116,11 +116,11 @@ class Executor(ABC):
116 raise NotImplementedError
117
118 def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None:
119 - """
120 - Initialize the KV caches and begin the model execution loop of the
121 - underlying workers.
122 - """
119 + """Initialize the KV caches on the underlying workers."""
120 self.collective_rpc("initialize_from_config", args=(kv_cache_configs,))
121 +
122 + def compile_or_warm_up_model(self) -> None:
123 + """Compile/warm up the model and capture cudagraphs on workers."""
124 compilation_times: list[CompilationTimes] = self.collective_rpc(
125 "compile_or_warm_up_model"
126 )