@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2601f18a82
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 12, 16:00 UTC

[EPLB] Optimize EPLB for Async Rearrange Experts (#22179)

Signed-off-by: David Chen <530634352@qq.com> Co-authored-by: SunChenxiang123 <1291824390@qq.com>

WeiQing Chen committed Nov 24, 2025 at 22:08 UTC 2601f18a8216b8e0b8743cded669b8b4d0e3e980
7 files changed +778 -77
tests/distributed/test_eplb_execute.py
+126 -1
@@ -1,13 +1,18 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import asyncio
5 import random
6
7 import pytest
8 import torch
9 import torch.distributed
10
10 -from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace
11 +from vllm.distributed.eplb.rebalance_execute import (
12 + move_from_buffer,
13 + rearrange_expert_weights_inplace,
14 + transfer_layer,
15 +)
16 from vllm.distributed.parallel_state import (
17 ensure_model_parallel_initialized,
18 get_tp_group,
@@ -231,6 +236,100 @@ def verify_redundant_experts_have_same_weights(
236 )
237
238
239 +def _test_async_transfer_layer_without_mtp_worker(
240 + env,
241 + world_size: int,
242 + num_layers: int,
243 + num_local_experts: int,
244 + num_logical_experts: int,
245 +) -> None:
246 + set_env_vars_and_device(env)
247 + ensure_model_parallel_initialized(
248 + tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
249 + )
250 +
251 + tp_group = get_tp_group()
252 + ep_group = tp_group.device_group
253 + ep_rank = torch.distributed.get_rank()
254 + device = torch.device(f"cuda:{ep_rank}")
255 +
256 + total_physical_experts = world_size * num_local_experts
257 + hidden_sizes = [16, 32]
258 +
259 + redundancy_config = create_redundancy_config(
260 + num_logical_experts,
261 + total_physical_experts,
262 + )
263 + old_indices = create_expert_indices_with_redundancy(
264 + num_layers,
265 + num_logical_experts,
266 + total_physical_experts,
267 + redundancy_config,
268 + )
269 +
270 + new_redundancy_config = create_redundancy_config(
271 + num_logical_experts,
272 + total_physical_experts,
273 + )
274 + new_indices = create_expert_indices_with_redundancy(
275 + num_layers,
276 + num_logical_experts,
277 + total_physical_experts,
278 + new_redundancy_config,
279 + )
280 +
281 + expert_weights = create_expert_weights(
282 + num_layers,
283 + num_local_experts,
284 + hidden_sizes,
285 + ep_rank,
286 + device,
287 + old_indices,
288 + )
289 +
290 + expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
291 + cuda_stream = torch.cuda.Stream(device=device)
292 +
293 + for layer_idx in range(num_layers):
294 + is_unchanged, is_received_locally, experts_recv_loc = asyncio.run(
295 + transfer_layer(
296 + old_global_expert_indices=old_indices,
297 + new_global_expert_indices=new_indices,
298 + expert_weights=expert_weights,
299 + expert_weights_buffer=expert_buffer,
300 + ep_group=ep_group,
301 + layer=layer_idx,
302 + cuda_stream=cuda_stream,
303 + )
304 + )
305 +
306 + cuda_stream.synchronize()
307 + move_from_buffer(
308 + expert_weights=expert_weights[layer_idx],
309 + expert_weights_buffer=expert_buffer,
310 + is_unchanged=is_unchanged,
311 + is_received_locally=is_received_locally,
312 + experts_recv_loc=experts_recv_loc,
313 + new_indices=new_indices[layer_idx].tolist(),
314 + ep_group=ep_group,
315 + )
316 +
317 + verify_expert_weights_after_shuffle(
318 + expert_weights,
319 + new_indices,
320 + hidden_sizes,
321 + ep_rank,
322 + num_local_experts,
323 + )
324 + verify_redundant_experts_have_same_weights(
325 + expert_weights,
326 + new_indices,
327 + hidden_sizes,
328 + world_size,
329 + num_local_experts,
330 + )
331 +
332 +
333 def _test_rearrange_expert_weights_with_redundancy(
334 env, world_size, num_layers, num_local_experts, num_logical_experts
335 ) -> None:
@@ -399,6 +498,32 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:
498 )
499
500
501 +@pytest.mark.parametrize(
502 + "world_size,num_layers,num_local_experts,num_logical_experts",
503 + [
504 + (2, 2, 2, 3),
505 + ],
506 +)
507 +def test_async_transfer_layer_without_mtp(
508 + world_size: int,
509 + num_layers: int,
510 + num_local_experts: int,
511 + num_logical_experts: int,
512 +):
513 + """Exercise async EPLB transfer path without MTP/spec decode."""
514 +
515 + if torch.cuda.device_count() < world_size:
516 + pytest.skip(f"Need at least {world_size} GPUs to run the test")
517 +
518 + distributed_run(
519 + _test_async_transfer_layer_without_mtp_worker,
520 + world_size,
521 + num_layers,
522 + num_local_experts,
523 + num_logical_experts,
524 + )
525 +
526 +
527 @pytest.mark.parametrize("world_size", [2, 4])
528 def test_rearrange_expert_weights_no_change(world_size):
529 """
tests/distributed/test_eplb_spec_decode.py
+38 -1
@@ -10,10 +10,11 @@ from tests.utils import large_gpu_mark
10
11 def get_model_args(
12 model_name: str,
13 - spec_model_name: str,
13 + spec_model_name: str | None,
14 spec_method: str,
15 tp_size: int,
16 model_max_len: int,
17 + use_async: bool = False,
18 ) -> dict:
19 speculative_config = {
20 "method": spec_method,
@@ -37,6 +38,8 @@ def get_model_args(
38 "enable_eplb": True,
39 "max_model_len": model_max_len,
40 }
41 + if use_async:
42 + model_args["eplb_config"] = {"use_async": True}
43 return model_args
44
45
@@ -94,3 +97,37 @@ def test_eplb_spec_decode(
97 measured_value - RTOL < expected_gsm8k_value
98 and measured_value + RTOL > expected_gsm8k_value
99 ), f"Expected: {expected_gsm8k_value} | Measured: {measured_value}"
100 +
101 +
102 +@large_gpu_mark(min_gb=80)
103 +def test_eplb_spec_decode_qwen3_next_mtp_async() -> None:
104 + """
105 + Ensure async EPLB works with MTP speculative decoding for Qwen3-Next.
106 + """
107 +
108 + TASK = "gsm8k"
109 + FILTER = "exact_match,strict-match"
110 + RTOL = 0.03
111 + expected_gsm8k_value = 0.86
112 +
113 + model_args = get_model_args(
114 + model_name="Qwen/Qwen3-Next-80B-A3B-Instruct",
115 + spec_model_name=None,
116 + spec_method="mtp",
117 + tp_size=4,
118 + model_max_len=4096,
119 + use_async=True,
120 + )
121 +
122 + results = lm_eval.simple_evaluate(
123 + model="vllm",
124 + model_args=model_args,
125 + tasks=TASK,
126 + batch_size=64,
127 + num_fewshot=8,
128 + )
129 + measured_value = results["results"][TASK][FILTER]
130 + assert (
131 + measured_value - RTOL < expected_gsm8k_value
132 + and measured_value + RTOL > expected_gsm8k_value
133 + ), f"Expected: {expected_gsm8k_value} | Measured: {measured_value}"
vllm/config/parallel.py
+4
@@ -60,6 +60,10 @@ class EPLBConfig:
60 Log the balancedness each step of expert parallelism.
61 This is turned off by default since it will cause communication overhead.
62 """
63 + use_async: bool = False
64 + """
65 + Whether to use non-blocking EPLB.
66 + """
67
68
69 @config
vllm/distributed/eplb/async_worker.py new
+115
@@ -0,0 +1,115 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +The async worker that transfers experts in the background.
5 +"""
6 +
7 +import asyncio
8 +import threading
9 +from typing import TYPE_CHECKING
10 +
11 +import torch
12 +from torch.distributed import ProcessGroup
13 +
14 +from vllm.distributed.parallel_state import get_ep_group
15 +from vllm.logger import init_logger
16 +
17 +from .rebalance_execute import transfer_layer
18 +
19 +if TYPE_CHECKING:
20 + from .eplb_state import EplbState
21 +
22 +logger = init_logger(__name__)
23 +
24 +
25 +def start_async_worker(
26 + state: "EplbState",
27 + rank_mapping: dict[int, int] | None = None,
28 + is_profile: bool = False,
29 +) -> threading.Thread:
30 + ep_group = get_ep_group().device_group
31 + rank = ep_group.rank()
32 + device_index = state.cuda_device_index
33 +
34 + def thread_target() -> None:
35 + assert device_index is not None
36 + torch.cuda.set_device(device_index)
37 + cuda_stream = torch.cuda.Stream(device=device_index)
38 + loop = asyncio.new_event_loop()
39 + asyncio.set_event_loop(loop)
40 + try:
41 + loop.run_until_complete(
42 + transfer_run_periodically(
43 + state=state,
44 + ep_group=ep_group,
45 + is_profile=is_profile,
46 + rank_mapping=rank_mapping,
47 + cuda_stream=cuda_stream,
48 + )
49 + )
50 + except Exception as exc: # pragma: no cover - diagnostic path
51 + logger.exception("async loop error (Rank %d): %s", rank, str(exc))
52 + finally:
53 + loop.close()
54 +
55 + thread = threading.Thread(target=thread_target, daemon=True)
56 + thread.start()
57 + return thread
58 +
59 +
60 +async def transfer_run_periodically(
61 + state: "EplbState",
62 + ep_group: ProcessGroup,
63 + is_profile: bool = False,
64 + rank_mapping: dict[int, int] | None = None,
65 + cuda_stream: torch.cuda.Stream = None,
66 +) -> None:
67 + while True:
68 + await asyncio.to_thread(state.rearrange_event.wait)
69 + logger.info("async worker woke up for EPLB transfer")
70 +
71 + for model_state in state.model_states.values():
72 + if not model_state.is_async_enabled:
73 + continue
74 + current_num_layers = model_state.model.num_moe_layers
75 + while (
76 + model_state.rebalanced
77 + and model_state.layer_to_transfer < current_num_layers
78 + ):
79 + if (
80 + not model_state.ep_buffer_ready
81 + and model_state.rebalanced
82 + and model_state.new_physical_to_logical_map is not None
83 + ):
84 + await asyncio.to_thread(model_state.buffer_lock.acquire)
85 + try:
86 + if model_state.layer_to_transfer >= current_num_layers:
87 + break
88 +
89 + (
90 + model_state.is_unchanged,
91 + model_state.is_received_locally,
92 + model_state.experts_recv_loc,
93 + ) = await transfer_layer(
94 + old_global_expert_indices=model_state.physical_to_logical_map,
95 + new_global_expert_indices=model_state.new_physical_to_logical_map,
96 + expert_weights=model_state.model.expert_weights,
97 + expert_weights_buffer=model_state.expert_buffer,
98 + ep_group=ep_group,
99 + is_profile=is_profile,
100 + layer=model_state.layer_to_transfer,
101 + cuda_stream=cuda_stream,
102 + rank_mapping=rank_mapping,
103 + )
104 + event = torch.cuda.Event(blocking=False)
105 + cuda_stream.record_event(event)
106 + model_state.buffer_ready_event = event
107 + model_state.ep_buffer_ready = 1
108 + finally:
109 + model_state.buffer_lock.release()
110 + else:
111 + if not model_state.rebalanced:
112 + break
113 + await asyncio.sleep(0.001)
114 +
115 + state.rearrange_event.clear()
vllm/distributed/eplb/eplb_state.py
+374 -57
@@ -26,6 +26,7 @@ MoE layer. If we have 32 EP ranks, then each GPU will hold 288 / 32 = 9 local
26 physical experts.
27 """
28
29 +import threading
30 import time
31 from collections.abc import Sequence
32 from dataclasses import dataclass
@@ -43,8 +44,9 @@ from vllm.distributed.utils import StatelessProcessGroup
44 from vllm.logger import init_logger
45 from vllm.model_executor.models.interfaces import MixtureOfExperts
46
47 +from .async_worker import start_async_worker
48 from .rebalance_algo import rebalance_experts
47 -from .rebalance_execute import rearrange_expert_weights_inplace
49 +from .rebalance_execute import move_from_buffer, rearrange_expert_weights_inplace
50
51 logger = init_logger(__name__)
52
@@ -132,6 +134,74 @@ class EplbModelState:
134 """
135 model_name: str
136 model: MixtureOfExperts
137 + expert_buffer: list[torch.Tensor]
138 + """
139 + The buffer to store the expert weights during transfer.
140 + """
141 + buffer_lock: threading.Lock
142 + """
143 + The lock to protect the expert buffer.
144 + """
145 + buffer_ready_event: torch.cuda.Event | None
146 + """
147 + CUDA event recorded when the async worker finishes filling the buffer.
148 + The main thread waits on this before consuming the buffer.
149 + """
150 + ep_buffer_ready: int
151 + """
152 + The flag indicates whether the expert buffer is ready for transfer.
153 + 0 or 1.
154 + """
155 + layer_to_transfer: int
156 + """
157 + The layer index to transfer in async mode.
158 + """
159 + rebalanced: bool
160 + """
161 + The flag indicates whether the experts rebalance have been computed.
162 + """
163 + pending_global_ready_check: bool
164 + """
165 + Whether the async EPLB needs to poll peers for buffer readiness.
166 + """
167 + is_unchanged: list[bool]
168 + """
169 + intermediate variable between `move_to_buffer` and `move_to_workspace`.
170 + The size is same as the num of physical experts in the current layer.
171 + """
172 + is_received_locally: list[bool]
173 + """
174 + intermediate variable between `move_to_buffer` and `move_to_workspace`.
175 + The size is same as the num of physical experts in the current layer.
176 + """
177 + experts_recv_loc: dict[int, int]
178 + """
179 + intermediate variable between `move_to_buffer` and `move_to_workspace`.
180 + The size is same as the num of physical experts in the current layer.
181 + """
182 + is_async_enabled: bool
183 + """
184 + The flag indicates whether the EPLB is running in async mode.
185 + """
186 + cuda_device_index: int | None
187 + """
188 + CUDA device index for the async EPLB worker thread.
189 + """
190 + new_physical_to_logical_map: torch.Tensor | None = None
191 + """
192 + intermediate variable between `move_to_buffer` and `move_to_workspace`.
193 + the size is same as physical_to_logical_map
194 + """
195 + new_logical_to_physical_map: torch.Tensor | None = None
196 + """
197 + intermediate variable between `move_to_buffer` and `move_to_workspace`.
198 + the size is same as logical_to_physical_map
199 + """
200 + new_logical_replica_count: torch.Tensor | None = None
201 + """
202 + intermediate variable between `move_to_buffer` and `move_to_workspace`.
203 + the size is same as logical_replica_count
204 + """
205
206
207 class EplbState:
@@ -164,12 +234,31 @@ class EplbState:
234 Otherwise, the rearrangement will hang at collective
235 communication calls.
236 """
167 - self.expert_rearrangement_step: int = 0
237 + self.expert_rearrangement_step_interval: int = 0
238 """
239 Interval for expert rearrangement steps.
240 This is a constant and is taken from the config.
241 """
172 - self.expert_rearrangement_step_interval: int = 0
242 + self.is_async: bool = False
243 + """
244 + The flag indicates whether the EPLB is running in async mode.
245 + """
246 + self.rearrange_event = threading.Event()
247 + """
248 + Event to signal when a new rearrangement is needed for the async thread.
249 + """
250 + self.async_worker: threading.Thread | None = None
251 + """
252 + Background thread handling async transfers.
253 + """
254 + self.cuda_device_index: int | None = None
255 + """
256 + CUDA device index for the async EPLB worker thread.
257 + """
258 + if self.device.type == "cuda":
259 + self.cuda_device_index = self.device.index
260 + if self.cuda_device_index is None and torch.cuda.is_available():
261 + self.cuda_device_index = torch.cuda.current_device()
262
263 @staticmethod
264 def build_initial_global_physical_to_logical_map(
@@ -239,6 +328,8 @@ class EplbState:
328 Build the initial EPLB state.
329 """
330 self.validate_ep_configuration(model)
331 + self.is_async = self.parallel_config.eplb_config.use_async
332 +
333 physical_to_logical_map_list = (
334 EplbState.build_initial_global_physical_to_logical_map(
335 model.num_routed_experts,
@@ -368,7 +459,12 @@ class EplbState:
459 physical_to_logical_map = new_physical_to_logical_map.to(self.device)
460 logical_to_physical_map.copy_(new_logical_to_physical_map)
461 logical_replica_count.copy_(new_logical_replica_count)
462 + else:
463 + new_physical_to_logical_map = None
464 +
465 + new_logical_to_physical_map = None
466
467 + new_logical_replica_count = None
468 model.set_eplb_state(
469 expert_load_pass,
470 logical_to_physical_map,
@@ -385,15 +481,33 @@ class EplbState:
481 )
482 self.expert_rearrangement_step = 0
483
388 - self.model_states[model_config.compute_hash()] = EplbModelState(
389 - physical_to_logical_map,
390 - logical_to_physical_map,
391 - logical_replica_count,
392 - expert_load_pass,
393 - expert_load_window,
394 - model_config.model,
395 - model,
484 + expert_buffer = [torch.empty_like(w) for w in model.expert_weights[0]]
485 +
486 + model_state = EplbModelState(
487 + physical_to_logical_map=physical_to_logical_map,
488 + logical_to_physical_map=logical_to_physical_map,
489 + logical_replica_count=logical_replica_count,
490 + expert_load_pass=expert_load_pass,
491 + expert_load_window=expert_load_window,
492 + model_name=model_config.model,
493 + model=model,
494 + expert_buffer=expert_buffer,
495 + buffer_lock=threading.Lock(),
496 + buffer_ready_event=None,
497 + ep_buffer_ready=0,
498 + layer_to_transfer=0,
499 + rebalanced=False,
500 + pending_global_ready_check=False,
501 + is_unchanged=[],
502 + is_received_locally=[],
503 + experts_recv_loc={},
504 + is_async_enabled=self.is_async,
505 + cuda_device_index=self.cuda_device_index,
506 + new_physical_to_logical_map=new_physical_to_logical_map,
507 + new_logical_to_physical_map=new_logical_to_physical_map,
508 + new_logical_replica_count=new_logical_replica_count,
509 )
510 + self.model_states[model_config.compute_hash()] = model_state
511
512 def step(
513 self,
@@ -420,7 +534,7 @@ class EplbState:
534 - `max_tokens`: The maximum load across ranks.
535 - `balancedness`: The ratio of average load to maximum load.
536 """
423 -
537 + ep_group = get_ep_group().device_group
538 if is_profile:
539 self.rearrange(is_profile=True)
540 return
@@ -488,7 +602,49 @@ class EplbState:
602 # rearrangement step and perform rearrangement to ensure all ranks are
603 # performing collective communication.
604 self.expert_rearrangement_step += 1
605 +
606 + if self.is_async:
607 + for eplb_model_state in self.model_states.values():
608 + if not eplb_model_state.is_async_enabled:
609 + continue
610 +
611 + all_ranks_buffer_ready = False
612 + if eplb_model_state.pending_global_ready_check:
613 + all_ranks_buffer_ready = self._all_ranks_buffer_ready(
614 + eplb_model_state
615 + )
616 + if (
617 + eplb_model_state.is_async_enabled
618 + and eplb_model_state.ep_buffer_ready
619 + and all_ranks_buffer_ready
620 + ):
621 + self.move_to_workspace(
622 + model_state=eplb_model_state,
623 + ep_group=ep_group,
624 + is_profile=is_profile,
625 + )
626 + if (
627 + eplb_model_state.layer_to_transfer
628 + >= eplb_model_state.model.num_moe_layers
629 + ):
630 + self.post_eplb(eplb_model_state, is_profile)
631 + eplb_model_state.rebalanced = False
632 + eplb_model_state.layer_to_transfer = 0
633 + eplb_model_state.pending_global_ready_check = False
634 + logger.info(
635 + "finish async transfer for model %s rank %d layer %d",
636 + eplb_model_state.model_name,
637 + ep_group.rank(),
638 + eplb_model_state.model.num_moe_layers,
639 + )
640 +
641 if self.expert_rearrangement_step >= self.expert_rearrangement_step_interval:
642 + if any(
643 + eplb_model_state.is_async_enabled and eplb_model_state.rebalanced
644 + for eplb_model_state in self.model_states.values()
645 + ):
646 + # Still performing asynchronous rearrangement
647 + return
648 self.expert_rearrangement_step = 0
649 self.rearrange()
650
@@ -524,7 +680,11 @@ class EplbState:
680 if is_main_rank:
681 torch.cuda.synchronize()
682 time_start = time.perf_counter()
527 - logger.info("Rearranging experts %s...", "(profile)" if is_profile else "")
683 + logger.info(
684 + "Rearranging experts %s %s...",
685 + "(async mode)" if self.is_async else "sync mode",
686 + "(profile)" if is_profile else "",
687 + )
688
689 if global_expert_loads is None:
690 # Map the physical expert load to global logical experts
@@ -593,6 +753,7 @@ class EplbState:
753 model = eplb_model_state.model
754 num_replicas = model.num_physical_experts
755 num_groups = model.num_expert_groups
756 +
757 if rank_mapping is not None and len(rank_mapping) == ep_group.size():
758 # NOTE(yongji): scale down, we need to rebalance the experts on
759 # remaining GPUs, transfer the experts while we haven't shutdown
@@ -608,7 +769,7 @@ class EplbState:
769 num_gpus = ep_group.size()
770
771 if num_gpus % num_nodes != 0:
611 - self.num_nodes = 1
772 + num_nodes = 1
773 logger.warning_once(
774 f"num_gpus % num_nodes != 0, "
775 "not using hierarchical rearrangement algorithm.\n"
@@ -631,59 +792,215 @@ class EplbState:
792 num_gpus,
793 )
794
634 - # Update expert weights
635 - rearrange_expert_weights_inplace(
636 - eplb_model_state.physical_to_logical_map,
637 - new_physical_to_logical_map,
638 - eplb_model_state.model.expert_weights,
639 - ep_group,
640 - is_profile,
641 - rank_mapping,
642 - )
795 + if not eplb_model_state.is_async_enabled or is_profile:
796 + # Update expert weights
797 + rearrange_expert_weights_inplace(
798 + eplb_model_state.physical_to_logical_map,
799 + new_physical_to_logical_map,
800 + eplb_model_state.model.expert_weights,
801 + ep_group,
802 + is_profile,
803 + rank_mapping,
804 + )
805
644 - if not is_profile:
645 - if (
646 - eplb_model_state.physical_to_logical_map.shape[1]
647 - != new_physical_to_logical_map.shape[1]
648 - ):
649 - eplb_model_state.physical_to_logical_map = (
650 - new_physical_to_logical_map.to(
651 - eplb_model_state.physical_to_logical_map.device
806 + if not is_profile:
807 + if (
808 + eplb_model_state.physical_to_logical_map.shape[1]
809 + != new_physical_to_logical_map.shape[1]
810 + ):
811 + eplb_model_state.physical_to_logical_map = (
812 + new_physical_to_logical_map.to(
813 + eplb_model_state.physical_to_logical_map.device
814 + )
815 )
816 + else:
817 + eplb_model_state.physical_to_logical_map.copy_(
818 + new_physical_to_logical_map
819 + )
820 + max_physical_slots = new_logical_to_physical_map.shape[-1]
821 + assert (
822 + max_physical_slots
823 + <= eplb_model_state.logical_to_physical_map.shape[-1]
824 )
654 - else:
655 - eplb_model_state.physical_to_logical_map.copy_(
656 - new_physical_to_logical_map
825 + new_logical_to_physical_map = torch.nn.functional.pad(
826 + new_logical_to_physical_map,
827 + (
828 + 0,
829 + eplb_model_state.logical_to_physical_map.shape[-1]
830 + - max_physical_slots,
831 + ),
832 + value=-1,
833 )
658 - max_physical_slots = new_logical_to_physical_map.shape[-1]
659 - assert (
660 - max_physical_slots
661 - <= eplb_model_state.logical_to_physical_map.shape[-1]
662 - )
663 - new_logical_to_physical_map = torch.nn.functional.pad(
834 + eplb_model_state.logical_to_physical_map.copy_(
835 + new_logical_to_physical_map
836 + )
837 + eplb_model_state.logical_replica_count.copy_(
838 + new_logical_replica_count
839 + )
840 + if is_main_rank:
841 + assert time_start is not None
842 + torch.cuda.synchronize()
843 + time_end = time.perf_counter()
844 + logger.info(
845 + "Rearranged experts%sin %.2f seconds.",
846 + " (profile) " if is_profile else " ",
847 + time_end - time_start,
848 + )
849 + else:
850 + device = eplb_model_state.physical_to_logical_map.device
851 + new_physical = new_physical_to_logical_map.to(device)
852 + max_slots = eplb_model_state.logical_to_physical_map.shape[-1]
853 + padded_logical = torch.nn.functional.pad(
854 new_logical_to_physical_map,
665 - (
666 - 0,
667 - eplb_model_state.logical_to_physical_map.shape[-1]
668 - - max_physical_slots,
669 - ),
855 + (0, max(0, max_slots - new_logical_to_physical_map.shape[-1])),
856 value=-1,
857 + ).to(eplb_model_state.logical_to_physical_map.device)
858 + new_replica = new_logical_replica_count.to(
859 + eplb_model_state.logical_replica_count.device
860 )
672 - eplb_model_state.logical_to_physical_map.copy_(
673 - new_logical_to_physical_map
674 - )
675 - eplb_model_state.logical_replica_count.copy_(new_logical_replica_count)
861
677 - if is_main_rank:
678 - assert time_start is not None
679 - torch.cuda.synchronize()
680 - time_end = time.perf_counter()
862 + eplb_model_state.new_physical_to_logical_map = new_physical
863 + eplb_model_state.new_logical_to_physical_map = padded_logical
864 + eplb_model_state.new_logical_replica_count = new_replica
865 +
866 + eplb_model_state.rebalanced = True
867 + eplb_model_state.layer_to_transfer = 0
868 + eplb_model_state.pending_global_ready_check = True
869 +
870 + # Signal async thread to start transferring layers
871 + if self.is_async and (not is_profile):
872 + self.rearrange_event.set()
873 + return None
874 +
875 + def start_async_loop(
876 + self,
877 + rank_mapping: dict[int, int] | None = None,
878 + is_profile: bool = False,
879 + ):
880 + if not self.is_async:
881 + return
882 + if self.async_worker is None:
883 + self.async_worker = start_async_worker(
884 + self,
885 + rank_mapping=rank_mapping,
886 + is_profile=is_profile,
887 + )
888 +
889 + def _update_layer_mapping_from_new(
890 + self, model_state: EplbModelState, layer: int
891 + ) -> None:
892 + if (
893 + model_state.new_physical_to_logical_map is None
894 + or model_state.new_logical_to_physical_map is None
895 + or model_state.new_logical_replica_count is None
896 + ):
897 + return
898 +
899 + target_device = model_state.physical_to_logical_map.device
900 + new_physical = model_state.new_physical_to_logical_map
901 + if model_state.physical_to_logical_map.shape[1] != new_physical.shape[1]:
902 + model_state.physical_to_logical_map = new_physical.to(target_device)
903 + else:
904 + model_state.physical_to_logical_map[layer].copy_(
905 + new_physical[layer].to(target_device)
906 + )
907 +
908 + logical_device = model_state.logical_to_physical_map.device
909 + new_logical = model_state.new_logical_to_physical_map[layer].to(logical_device)
910 + max_slots = model_state.logical_to_physical_map.shape[-1]
911 + slot_delta = max_slots - new_logical.shape[-1]
912 + if slot_delta > 0:
913 + new_logical = torch.nn.functional.pad(
914 + new_logical, (0, slot_delta), value=-1
915 + )
916 + model_state.logical_to_physical_map[layer].copy_(new_logical)
917 +
918 + replica_device = model_state.logical_replica_count.device
919 + model_state.logical_replica_count[layer].copy_(
920 + model_state.new_logical_replica_count[layer].to(replica_device)
921 + )
922 +
923 + def _all_ranks_buffer_ready(self, model_state: EplbModelState) -> bool:
924 + parallel_state = get_ep_group()
925 + cpu_group = getattr(parallel_state, "cpu_group", None)
926 + if cpu_group is not None and cpu_group.size() > 1:
927 + flag = torch.tensor(
928 + (int(model_state.ep_buffer_ready),), dtype=torch.int32, device="cpu"
929 + )
930 + all_reduce(flag, group=cpu_group)
931 + return int(flag.item()) == cpu_group.size()
932 +
933 + device_group = parallel_state.device_group
934 + if device_group.size() <= 1:
935 + return bool(model_state.ep_buffer_ready)
936 +
937 + device = getattr(
938 + parallel_state, "device", model_state.physical_to_logical_map.device
939 + )
940 + flag = torch.tensor(
941 + (int(model_state.ep_buffer_ready),), dtype=torch.int32, device=device
942 + )
943 + all_reduce(flag, group=device_group)
944 + return int(flag.item()) == device_group.size()
945 +
946 + def move_to_workspace(
947 + self,
948 + model_state: EplbModelState,
949 + ep_group: ProcessGroup,
950 + is_profile: bool = False,
951 + ):
952 + if not model_state.buffer_lock.acquire(blocking=False):
953 + return
954 + try:
955 + assert model_state.new_physical_to_logical_map is not None
956 + device_index = model_state.cuda_device_index or self.cuda_device_index
957 + if model_state.buffer_ready_event is not None and device_index is not None:
958 + stream = torch.cuda.current_stream(device=device_index)
959 + stream.wait_event(model_state.buffer_ready_event)
960 + model_state.buffer_ready_event = None
961 + move_from_buffer(
962 + expert_weights=model_state.model.expert_weights[
963 + model_state.layer_to_transfer
964 + ],
965 + expert_weights_buffer=model_state.expert_buffer,
966 + is_unchanged=model_state.is_unchanged,
967 + is_received_locally=model_state.is_received_locally,
968 + experts_recv_loc=model_state.experts_recv_loc,
969 + new_indices=model_state.new_physical_to_logical_map[
970 + model_state.layer_to_transfer
971 + ].tolist(),
972 + ep_group=ep_group,
973 + )
974 + transferred_layer = model_state.layer_to_transfer
975 + self._update_layer_mapping_from_new(model_state, transferred_layer)
976 + # After the main thread consumes, advance layer_to_transfer
977 + model_state.layer_to_transfer += 1
978 + model_state.ep_buffer_ready = 0
979 logger.info(
682 - "Rearranged experts%sin %.2f seconds.",
683 - " (profile) " if is_profile else " ",
684 - time_end - time_start,
980 + "model %s successfully move_to_workspace layer %d",
981 + model_state.model_name,
982 + transferred_layer,
983 )
686 - return None
984 + finally:
985 + try:
986 + model_state.buffer_lock.release()
987 + except Exception as e:
988 + logger.error(
989 + "Rank %d: buffer_lock release failed in move_to_workspace: %s",
990 + ep_group.rank(),
991 + str(e),
992 + )
993 +
994 + def post_eplb(self, model_state: EplbModelState, is_profile: bool = False) -> None:
995 + assert model_state.new_physical_to_logical_map is not None
996 + assert model_state.new_logical_to_physical_map is not None
997 + assert model_state.new_logical_replica_count is not None
998 + if not is_profile:
999 + for layer_idx in range(model_state.physical_to_logical_map.shape[0]):
1000 + self._update_layer_mapping_from_new(model_state, layer_idx)
1001 + model_state.new_physical_to_logical_map = None
1002 + model_state.new_logical_to_physical_map = None
1003 + model_state.new_logical_replica_count = None
1004
1005 @staticmethod
1006 def recv_state() -> tuple[list[torch.Tensor], list[torch.Tensor]]:
vllm/distributed/eplb/rebalance_execute.py
+119 -18
@@ -100,18 +100,19 @@ def get_ep_ranks_with_expert(
100 return ranks_to_send, ranks_to_recv_actual
101
102
103 -def shuffle_layer(
103 +def move_to_buffer(
104 num_local_experts: int,
105 - ep_rank: int,
105 old_indices: Sequence[int],
106 new_indices: Sequence[int],
107 expert_weights: Iterable[torch.Tensor],
108 expert_weights_buffer: Sequence[torch.Tensor],
109 + cuda_stream: torch.cuda.Stream | None,
110 ep_group: ProcessGroup,
111 -) -> None:
111 +) -> tuple[list[bool], list[bool], dict[int, int]]:
112 """
113 Perform expert weights rearrangement of one layer.
114 """
115 + ep_rank = ep_group.rank()
116 local2global = partial(
117 idx_local_to_global,
118 local_cnt=num_local_experts,
@@ -137,7 +138,8 @@ def shuffle_layer(
138 if old_indices[src_global] == new_indices[dst_global]:
139 is_received_locally[dst] = True
140 for weight, buffer in zip(expert_weights, expert_weights_buffer):
140 - buffer[dst].copy_(weight[src])
141 + with torch.cuda.stream(cuda_stream):
142 + buffer[dst].copy_(weight[src], non_blocking=True)
143
144 p2p_ops: list[P2POp] = []
145
@@ -225,25 +227,115 @@ def shuffle_layer(
227 ]
228
229 # 4. Execute the P2P operations. The real communication happens here.
228 - if p2p_ops:
230 + if p2p_ops and cuda_stream is not None:
231 + with torch.cuda.stream(cuda_stream):
232 + reqs = batch_isend_irecv(p2p_ops)
233 + for req in reqs:
234 + req.wait()
235 + elif p2p_ops:
236 reqs = batch_isend_irecv(p2p_ops)
237 for req in reqs:
238 req.wait()
239 + # wait for the communication to finish
240 + return is_unchanged, is_received_locally, experts_recv_loc
241 +
242 +
243 +def move_from_buffer(
244 + expert_weights: Iterable[torch.Tensor],
245 + expert_weights_buffer: list[torch.Tensor],
246 + is_unchanged: list[bool],
247 + is_received_locally: list[bool],
248 + experts_recv_loc: dict[int, int],
249 + new_indices: Sequence[int],
250 + ep_group: ProcessGroup,
251 +) -> None:
252 + ep_rank = ep_group.rank()
253 + num_local_experts = len(is_unchanged)
254 +
255 + local2global = partial(
256 + idx_local_to_global, local_cnt=num_local_experts, ep_rank=ep_rank
257 + )
258
233 - # 5. Copy the weights from the buffer back to the original weights.
259 for dst in range(num_local_experts):
260 if is_unchanged[dst]:
261 continue
262 if is_received_locally[dst]:
263 for weight, buffer in zip(expert_weights, expert_weights_buffer):
239 - weight[dst].copy_(buffer[dst])
264 + weight[dst].copy_(buffer[dst], non_blocking=True)
265 else:
266 expert = new_indices[local2global(dst)]
267 if expert == -1:
268 continue
269 src = experts_recv_loc[expert]
270 for weight, buffer in zip(expert_weights, expert_weights_buffer):
246 - weight[dst].copy_(buffer[src])
271 + weight[dst].copy_(buffer[src], non_blocking=True)
272 +
273 +
274 +async def transfer_layer(
275 + old_global_expert_indices: torch.Tensor,
276 + new_global_expert_indices: torch.Tensor,
277 + expert_weights: Sequence[Iterable[torch.Tensor]],
278 + expert_weights_buffer: Sequence[torch.Tensor],
279 + ep_group: ProcessGroup,
280 + is_profile: bool = False,
281 + layer: int = 0,
282 + cuda_stream: torch.cuda.Stream | None = None,
283 + rank_mapping: dict[int, int] | None = None,
284 +) -> tuple[list[bool], list[bool], dict[int, int]]:
285 + """
286 + Rearranges the expert weights in place according to the new expert indices.
287 +
288 + The value of the indices arguments are logical indices of the experts,
289 + while keys are physical.
290 +
291 + Args:
292 + old_global_expert_indices: Shape (num_moe_layers, num_physical_experts).
293 + new_global_expert_indices: Shape (num_moe_layers, num_physical_experts).
294 + expert_weights: A sequence of shape (num_moe_layers)(weight_count)
295 + of tensors of shape (num_local_physical_experts, hidden_size_i).
296 + For example, a linear layer may have up and down projection,
297 + so weight_count = 2. Each weight's hidden size can be different.
298 + ep_group: The device process group for expert parallelism.
299 + is_profile (bool): If `True`, do not perform any actual weight copy.
300 + This is used during profile run, where we only perform dummy
301 + communications to reserve enough memory for the buffers.
302 + """
303 + ep_size = ep_group.size()
304 + if rank_mapping is not None:
305 + if len(rank_mapping) == ep_group.size():
306 + # scale down
307 + new_global_expert_indices = _map_new_expert_indices_with_rank_mapping(
308 + new_global_expert_indices,
309 + rank_mapping,
310 + )
311 + else:
312 + # scale up
313 + old_global_expert_indices = _map_old_expert_indices_with_rank_mapping(
314 + old_global_expert_indices,
315 + rank_mapping,
316 + ep_group.size(),
317 + )
318 +
319 + assert old_global_expert_indices.shape[1] == new_global_expert_indices.shape[1]
320 + num_moe_layers, num_physical_experts = old_global_expert_indices.shape
321 + assert len(expert_weights) == num_moe_layers
322 + num_local_physical_experts = next(iter(expert_weights[0])).shape[0]
323 + assert new_global_expert_indices.shape == (num_moe_layers, num_physical_experts)
324 + assert num_physical_experts == ep_size * num_local_physical_experts
325 + # A buffer to hold the expert weights in one layer during the exchange.
326 + # NOTE: Currently we assume the same weights across different layers
327 + # have the same shape.
328 +
329 + is_unchanged, is_received_locally, experts_recv_loc = move_to_buffer(
330 + num_local_experts=num_local_physical_experts,
331 + old_indices=old_global_expert_indices[layer].tolist(),
332 + new_indices=new_global_expert_indices[layer].tolist(),
333 + expert_weights=expert_weights[layer],
334 + expert_weights_buffer=expert_weights_buffer,
335 + cuda_stream=cuda_stream,
336 + ep_group=ep_group,
337 + )
338 + return is_unchanged, is_received_locally, experts_recv_loc
339
340
341 def rearrange_expert_weights_inplace(
@@ -296,7 +388,6 @@ def rearrange_expert_weights_inplace(
388 num_local_physical_experts = next(iter(expert_weights[0])).shape[0]
389 assert new_global_expert_indices.shape == (num_moe_layers, num_physical_experts)
390
299 - ep_rank = ep_group.rank()
391 ep_size = ep_group.size()
392 assert num_physical_experts == ep_size * num_local_physical_experts
393
@@ -329,14 +420,24 @@ def rearrange_expert_weights_inplace(
420 torch.cuda.synchronize()
421
422 for layer in range(num_moe_layers):
332 - shuffle_layer(
333 - num_local_physical_experts,
334 - ep_rank,
335 - old_global_expert_indices_cpu[layer].tolist(),
336 - new_global_expert_indices_cpu[layer].tolist(),
337 - expert_weights[layer],
338 - expert_weights_buffer,
339 - ep_group,
423 + is_unchanged, is_received_locally, experts_recv_loc = move_to_buffer(
424 + num_local_experts=num_local_physical_experts,
425 + old_indices=old_global_expert_indices_cpu[layer].tolist(),
426 + new_indices=new_global_expert_indices_cpu[layer].tolist(),
427 + expert_weights=expert_weights[layer],
428 + expert_weights_buffer=expert_weights_buffer,
429 + cuda_stream=None,
430 + ep_group=ep_group,
431 + )
432 +
433 + move_from_buffer(
434 + expert_weights=expert_weights[layer],
435 + expert_weights_buffer=expert_weights_buffer,
436 + is_unchanged=is_unchanged,
437 + is_received_locally=is_received_locally,
438 + experts_recv_loc=experts_recv_loc,
439 + new_indices=new_global_expert_indices[layer].tolist(),
440 + ep_group=ep_group,
441 )
442
443
@@ -428,4 +529,4 @@ def _map_new_expert_indices_with_rank_mapping(
529 return mapped_expert_indices
530
531
431 -__all__ = ["rearrange_expert_weights_inplace"]
532 +__all__ = ["transfer_layer", "move_from_buffer"]
vllm/v1/worker/gpu_model_runner.py
+2
@@ -3370,6 +3370,8 @@ class GPUModelRunner(
3370 old_global_expert_indices,
3371 rank_mapping,
3372 )
3373 + if self.eplb_state.is_async:
3374 + self.eplb_state.start_async_loop(rank_mapping=rank_mapping)
3375
3376 if (
3377 self.vllm_config.compilation_config.mode