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 4, 22:30 UTC
[EPLB] Refactor Async EPLB synchronization logic (#37601)
Signed-off-by: Sage Moore <sage@neuralmagic.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>
Sage Moore committed
Apr 20, 2026 at 10:05 UTC
3461c8b0277f2d1df6c7ea1ec789881c1d01650b
6 files changed
+317
-271
tests/distributed/test_eplb_events.py
new
+98
@@ -0,0 +1,98 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import threading
5
+import time
6
+
7
+import torch
8
+
9
+from vllm.distributed.eplb.eplb_utils import CpuGpuEvent
10
+
11
+
12
+def test_wait_blocks_until_record():
13
+ event = CpuGpuEvent()
14
+ record_stream = torch.cuda.Stream()
15
+ wait_stream = torch.cuda.Stream()
16
+ wait_returned = threading.Event()
17
+
18
+ def waiter():
19
+ event.wait(stream=wait_stream)
20
+ wait_returned.set()
21
+
22
+ t = threading.Thread(target=waiter)
23
+ t.start()
24
+
25
+ time.sleep(0.05)
26
+ assert not wait_returned.is_set(), "wait() returned before record() was called"
27
+
28
+ event.record(stream=record_stream)
29
+ t.join(timeout=5.0)
30
+
31
+ assert not event._recorded.is_set()
32
+
33
+
34
+def test_reuse_across_multiple_cycles():
35
+ wrapper = CpuGpuEvent()
36
+ record_stream = torch.cuda.Stream()
37
+ wait_stream = torch.cuda.Stream()
38
+ NUM_CYCLES = 8
39
+ completed_cycles = []
40
+ barriers = [threading.Barrier(2) for _ in range(NUM_CYCLES)]
41
+
42
+ def waiter():
43
+ for i in range(NUM_CYCLES):
44
+ wrapper.wait(stream=wait_stream)
45
+ completed_cycles.append(True)
46
+ barriers[i].wait()
47
+
48
+ t = threading.Thread(target=waiter)
49
+ t.start()
50
+
51
+ for i in range(NUM_CYCLES):
52
+ wrapper.record(stream=record_stream)
53
+ barriers[i].wait()
54
+
55
+ t.join(timeout=10.0)
56
+ assert len(completed_cycles) == NUM_CYCLES
57
+
58
+
59
+def test_producer_consumer():
60
+ """
61
+ This test uses the CpuGpuEvent to synchronize reads and writes to/from a shared GPU
62
+ tensor on multiple CPU threads.
63
+ """
64
+ worker_stream = torch.cuda.Stream()
65
+ # Create a single element counter that will be shared between two threads
66
+ buf = torch.zeros(1, device="cuda")
67
+ NUM_ROUNDS = 5
68
+
69
+ ready_cpu = [threading.Event() for _ in range(NUM_ROUNDS)]
70
+ events = [CpuGpuEvent() for _ in range(NUM_ROUNDS)]
71
+ errors: list[str] = []
72
+
73
+ # For each round, the worker thread (writer) sets the counter in buf and waits for
74
+ # the main thread to read it.
75
+ def worker():
76
+ for i in range(NUM_ROUNDS):
77
+ if i > 0:
78
+ events[i - 1].wait(stream=worker_stream)
79
+
80
+ with torch.cuda.stream(worker_stream):
81
+ buf.fill_(float(i + 1))
82
+
83
+ worker_stream.synchronize()
84
+ ready_cpu[i].set()
85
+
86
+ t = threading.Thread(target=worker)
87
+ t.start()
88
+
89
+ for i in range(NUM_ROUNDS):
90
+ ready_cpu[i].wait()
91
+ snapshot = buf.clone()
92
+ events[i].record()
93
+ val = snapshot.item()
94
+ if val != float(i + 1):
95
+ errors.append(f"round {i}: expected {i + 1:.1f}, got {val:.1f}")
96
+
97
+ t.join(timeout=10.0)
98
+ assert not errors, f"Buffer ordering errors: {errors}"
tests/distributed/test_eplb_utils.py
+2
-2
@@ -80,7 +80,7 @@ def test_commit_eplb_maps_for_layer_logical_padding():
80
.contiguous()
81
)
82
layer = 0
83
- _commit_eplb_maps_for_layer(model_state, new_phy2log, layer)
83
+ _commit_eplb_maps_for_layer(model_state, new_phy2log[layer], layer)
84
85
assert torch.all(model_state.logical_to_physical_map[layer, :, 2] == -1)
86
@@ -143,7 +143,7 @@ def test_commit_eplb_maps_for_layer():
143
)
144
new_logcnt = torch.tensor([[2, 1, 1], [1, 2, 1]], dtype=torch.long)
145
146
- _commit_eplb_maps_for_layer(model_state, new_phy2log, layer=0)
146
+ _commit_eplb_maps_for_layer(model_state, new_phy2log[0], layer=0)
147
148
# Layer 0 updated
149
assert torch.equal(model_state.physical_to_logical_map[0], new_phy2log[0])
vllm/distributed/eplb/async_worker.py
+69
-86
@@ -14,7 +14,8 @@ from torch.distributed import ProcessGroup
14
from vllm.distributed.parallel_state import get_eplb_group
15
from vllm.logger import init_logger
16
17
-from .rebalance_execute import transfer_layer
17
+from .eplb_utils import CpuGpuEvent
18
+from .rebalance_execute import AsyncEplbLayerResult, transfer_layer
19
20
if TYPE_CHECKING:
21
from .eplb_state import EplbModelState, EplbState
@@ -60,18 +61,14 @@ def run_rebalance_experts(
61
model_state: "EplbModelState",
62
eplb_state: "EplbState",
63
physical_to_logical_map_cpu: torch.Tensor,
63
-) -> None:
64
+ cuda_stream: torch.cuda.Stream,
65
+) -> torch.Tensor:
66
assert model_state.eplb_stats is not None
67
eplb_stats = model_state.eplb_stats
68
67
- # Wait for the main thread's all-reduce and clone to complete before
68
- # accessing the global_expert_load_window tensor.
69
- assert model_state.window_ready_event is not None
70
- model_state.window_ready_event.wait()
71
- model_state.window_ready_event = None
72
-
69
# Move the global expert load window to CPU for computation.
74
- global_expert_load_window = eplb_stats.global_expert_load_window.cpu()
70
+ with torch.cuda.stream(cuda_stream):
71
+ global_expert_load_window = eplb_stats.global_expert_load_window.cpu()
72
# Compute new expert mappings for the model
73
new_physical_to_logical_map = eplb_state.policy.rebalance_experts(
74
global_expert_load_window,
@@ -83,7 +80,7 @@ def run_rebalance_experts(
80
)
81
assert new_physical_to_logical_map.device == torch.device("cpu")
82
86
- model_state.new_physical_to_logical_map = new_physical_to_logical_map
83
+ return new_physical_to_logical_map
84
85
86
async def transfer_run_periodically(
@@ -93,85 +90,71 @@ async def transfer_run_periodically(
90
is_profile: bool = False,
91
) -> None:
92
while True:
96
- await asyncio.to_thread(state.rearrange_event.wait)
93
+ state.rearrange_event.wait(stream=cuda_stream)
94
logger.info("async worker woke up for EPLB transfer")
95
96
assert state.is_async
97
for model_state in state.model_states.values():
98
+ layer_idx = 0
99
# Set the async worker's CUDA stream on the communicator
100
model_state.communicator.set_stream(cuda_stream)
103
- rebalancing_algorithm_executed = False
104
- physical_to_logical_map_cpu = None
105
- current_num_layers = model_state.model.num_moe_layers
106
- while (
107
- model_state.rebalanced
108
- and model_state.layer_to_transfer < current_num_layers
109
- ):
110
- if not model_state.ep_buffer_ready and model_state.rebalanced:
111
- # Polling the lock directly in the async thread avoids
112
- # the thread switch overhead of asyncio.to_thread.
113
- # This is typically faster than offloading to a worker thread.
114
- while not model_state.buffer_lock.acquire(blocking=False):
115
- await asyncio.sleep(0)
116
- try:
117
- if model_state.layer_to_transfer >= current_num_layers:
118
- break
119
- if (
120
- not rebalancing_algorithm_executed
121
- or model_state.new_physical_to_logical_map is None
122
- ):
123
- # Move the physical_to_logical_map to CPU
124
- # for rebalancing and transfer_layer.
125
- physical_to_logical_map_cpu = (
126
- model_state.physical_to_logical_map.cpu()
127
- )
128
- run_rebalance_experts(
129
- model_state, state, physical_to_logical_map_cpu
130
- )
131
- rebalancing_algorithm_executed = True
132
- logger.info(
133
- "Async worker computed new indices for model %s",
134
- model_state.model_name,
135
- )
136
-
137
- assert model_state.new_physical_to_logical_map is not None
138
- assert physical_to_logical_map_cpu is not None
139
-
140
- layer_idx = model_state.layer_to_transfer
141
- old_layer_indices = physical_to_logical_map_cpu[layer_idx]
142
- new_layer_indices = model_state.new_physical_to_logical_map[
143
- layer_idx
144
- ]
145
-
146
- # Wait for the main thread to finish consuming the buffer
147
- # before initiating an EPLB transfer on another layer.
148
- if model_state.buffer_consumed_event is not None:
149
- cuda_stream.wait_event(model_state.buffer_consumed_event)
150
- model_state.buffer_consumed_event = None
151
-
152
- (
153
- model_state.is_unchanged,
154
- model_state.is_received_locally,
155
- model_state.recv_metadata,
156
- ) = await transfer_layer(
157
- old_layer_indices=old_layer_indices,
158
- new_layer_indices=new_layer_indices,
159
- expert_weights=model_state.model.expert_weights[layer_idx],
160
- expert_weights_buffer=model_state.expert_buffer,
161
- ep_group=eplb_group,
162
- communicator=model_state.communicator,
163
- is_profile=is_profile,
164
- cuda_stream=cuda_stream,
165
- )
166
- # block the async thread until the transfer to
167
- # the intermediate buffer is complete.
168
- cuda_stream.synchronize()
169
- model_state.ep_buffer_ready = 1
170
- finally:
171
- model_state.buffer_lock.release()
172
- else:
173
- if not model_state.rebalanced:
174
- break
175
- await asyncio.sleep(0.001)
176
-
177
- state.rearrange_event.clear()
101
+ num_layers = model_state.model.num_moe_layers
102
+
103
+ # Snapshot the physical_to_logical_map (synchronized with
104
+ # rearrange_event) and copy it to CPU
105
+ with torch.cuda.stream(cuda_stream):
106
+ physical_to_logical_map_cpu = model_state.physical_to_logical_map.cpu()
107
+
108
+ new_physical_to_logical_map = run_rebalance_experts(
109
+ model_state, state, physical_to_logical_map_cpu, cuda_stream
110
+ )
111
+ logger.info(
112
+ "Async worker computed new indices for model %s",
113
+ model_state.model_name,
114
+ )
115
+
116
+ # Execute one EPLB layer transfer per model forward pass. Each iteration
117
+ # of this loop will copy the new set of expert weights into
118
+ # model_state.expert_buffer, which will be consumed by the main thread in
119
+ # move_to_workspace
120
+ while model_state.rebalanced and layer_idx < num_layers:
121
+ (
122
+ is_unchanged,
123
+ is_received_locally,
124
+ recv_metadata,
125
+ ) = await transfer_layer(
126
+ old_layer_indices=physical_to_logical_map_cpu[layer_idx],
127
+ new_layer_indices=new_physical_to_logical_map[layer_idx],
128
+ expert_weights=model_state.model.expert_weights[layer_idx],
129
+ expert_weights_buffer=model_state.expert_buffer,
130
+ communicator=model_state.communicator,
131
+ ep_group=eplb_group,
132
+ is_profile=is_profile,
133
+ cuda_stream=cuda_stream,
134
+ )
135
+
136
+ # Wait until all writes to expert_buffer have finished before making the
137
+ # AsyncEplbLayerResult visible to the main thread.
138
+ cuda_stream.synchronize()
139
+
140
+ # This event guarantees that expert_buffer will not be overwritten by
141
+ # subsequent iterations of this loop until the main thread has consumed
142
+ # it. Record is called by the main thread after move_from_buffer().
143
+ consumed_event = CpuGpuEvent()
144
+
145
+ model_state.pending_result = AsyncEplbLayerResult(
146
+ layer_idx=layer_idx,
147
+ new_physical_to_logical_map=new_physical_to_logical_map[layer_idx],
148
+ is_unchanged=is_unchanged,
149
+ is_received_locally=is_received_locally,
150
+ recv_metadata=recv_metadata,
151
+ consumed_event=consumed_event,
152
+ )
153
+
154
+ # Block this thread until the main thread and main stream
155
+ # finish copying model_state.expert_buffer into
156
+ # model_state.model.expert_weights[layer_idx]
157
+ consumed_event.wait(stream=cuda_stream)
158
+ logger.debug("Layer %d transfer complete", layer_idx)
159
+ assert model_state.pending_result is None
160
+ layer_idx += 1
vllm/distributed/eplb/eplb_state.py
+64
-182
@@ -30,7 +30,6 @@ import threading
30
from collections.abc import Sequence
31
from dataclasses import dataclass
32
33
-import numpy as np
33
import torch
34
from torch.distributed import ProcessGroup, all_reduce
35
@@ -48,9 +47,10 @@ from vllm.model_executor.models.interfaces import MixtureOfExperts
47
48
from .async_worker import start_async_worker
49
from .eplb_communicator import EplbCommunicator, create_eplb_communicator
50
+from .eplb_utils import CpuGpuEvent
51
from .policy import EPLB_POLICIES, AbstractEplbPolicy, DefaultEplbPolicy
52
from .rebalance_execute import (
53
- RecvMetadata,
53
+ AsyncEplbLayerResult,
54
move_from_buffer,
55
rearrange_expert_weights_inplace,
56
)
@@ -174,55 +174,20 @@ class EplbModelState:
174
"""
175
The buffer to store the expert weights during transfer.
176
"""
177
- buffer_lock: threading.Lock
178
- """
179
- The lock to protect the expert buffer.
180
- """
181
- buffer_consumed_event: torch.cuda.Event | None
182
- """
183
- CUDA event recorded after the main thread finishes consuming the buffer.
184
- The async worker waits on this before writing to the buffer again.
185
- """
186
- window_ready_event: torch.cuda.Event | None
187
- """
188
- CUDA event recorded after all-reduce and clone on the main thread.
189
- The async worker waits on this before accessing global_expert_load_window.
190
- """
191
- ep_buffer_ready: int
192
- """
193
- The flag indicates whether the expert buffer is ready for transfer.
194
- 0 or 1.
195
- """
196
- layer_to_transfer: int
197
- """
198
- The layer index to transfer in async mode.
199
- """
177
rebalanced: bool
178
"""
202
- The flag indicates whether the experts rebalance have been computed.
203
- """
204
- pending_global_ready_check: bool
205
- """
206
- Whether the async EPLB needs to poll peers for buffer readiness.
179
+ This flag is only used when running Async EPLB. It is set to True by the main thread
180
+ after the new expert maps have been computed. This indicates that the async worker
181
+ should start transferring weights. move_to_workspace sets this flag to False when
182
+ all weights have been transferred and the new map has been successfully committed.
183
+
184
+ rebalanced relies on the GIL to synchronize access between the main thread and
185
+ the async worker.
186
"""
187
eplb_stats: EplbStats | None
188
"""
189
EPLB stats for the model.
190
"""
212
- is_unchanged: np.ndarray
213
- """
214
- intermediate variable between `move_to_buffer` and `move_to_workspace`.
215
- The size is same as the num of physical experts in the current layer.
216
- """
217
- is_received_locally: np.ndarray
218
- """
219
- intermediate variable between `move_to_buffer` and `move_to_workspace`.
220
- The size is same as the num of physical experts in the current layer.
221
- """
222
- recv_metadata: RecvMetadata
223
- """
224
- intermediate variable between `move_to_buffer` and `move_to_workspace`.
225
- """
191
cuda_device_index: int | None
192
"""
193
CUDA device index for the async EPLB worker thread.
@@ -231,10 +196,14 @@ class EplbModelState:
196
"""
197
The communicator for expert weight transfers.
198
"""
234
- new_physical_to_logical_map: torch.Tensor | None = None
199
+ pending_result: AsyncEplbLayerResult | None = None
200
"""
236
- intermediate variable between `move_to_buffer` and `move_to_workspace`.
237
- the size is same as physical_to_logical_map
201
+ Set by the async worker after all writes to expert_buffer are done. Consumed
202
+ and reset to None by the main thread in move_to_workspace() after the contents of
203
+ expert_buffer have been transferred out. At most one result is pending at a time.
204
+
205
+ pending_result relies on the GIL to synchronize access between the main thread and
206
+ the async worker.
207
"""
208
209
@@ -289,7 +258,7 @@ class EplbState:
258
"""
259
The flag indicates whether the EPLB is running in async mode.
260
"""
292
- self.rearrange_event = threading.Event()
261
+ self.rearrange_event: CpuGpuEvent = CpuGpuEvent()
262
"""
263
Event to signal when a new rearrangement is needed for the async thread.
264
"""
@@ -493,25 +462,10 @@ class EplbState:
462
model_name=model_config.model,
463
model=model,
464
expert_buffer=expert_buffer,
496
- buffer_lock=threading.Lock(),
497
- buffer_consumed_event=None,
498
- window_ready_event=None,
499
- ep_buffer_ready=0,
500
- layer_to_transfer=0,
465
rebalanced=False,
502
- pending_global_ready_check=False,
466
eplb_stats=None,
504
- is_unchanged=np.array([]),
505
- is_received_locally=np.array([]),
506
- recv_metadata=RecvMetadata(
507
- recv_primary_mask=np.array([]),
508
- recv_count=0,
509
- recv_expert_ids=np.array([]),
510
- recv_dst_rows=np.array([]),
511
- ),
467
cuda_device_index=self.cuda_device_index,
468
communicator=communicator,
514
- new_physical_to_logical_map=None,
469
)
470
self.model_states[model_config.compute_hash()] = model_state
471
self.num_valid_physical_experts = model.num_physical_experts
@@ -622,17 +576,17 @@ class EplbState:
576
self.expert_rearrangement_step += 1
577
578
if self.is_async:
579
+ # Run _move_to_workspace if all ranks have finished transferring the
580
+ # new weights to the intermediate buffer
581
for eplb_model_state in self.model_states.values():
626
- all_ranks_buffer_ready = False
627
- if eplb_model_state.pending_global_ready_check:
628
- all_ranks_buffer_ready = self._all_ranks_buffer_ready(
629
- eplb_model_state
630
- )
631
- if eplb_model_state.ep_buffer_ready and all_ranks_buffer_ready:
632
- self.move_to_workspace(
582
+ # rebalanced must remain consistent amongst all ranks otherwise the
583
+ # all_reduce in _all_ranks_result_ready will hang
584
+ if eplb_model_state.rebalanced and self._all_ranks_result_ready(
585
+ eplb_model_state
586
+ ):
587
+ _move_to_workspace(
588
model_state=eplb_model_state,
634
- ep_group=ep_group,
635
- is_profile=is_profile,
589
+ ep_rank=ep_group.rank(),
590
)
591
592
if self.expert_rearrangement_step >= self.expert_rearrangement_step_interval:
@@ -846,18 +800,10 @@ class EplbState:
800
num_nodes=num_nodes,
801
num_gpus=num_gpus,
802
)
849
- # Record event after clone to signal async worker
850
- # that load stats data is ready
851
- sync_event = torch.cuda.Event()
852
- sync_event.record()
853
- eplb_model_state.window_ready_event = sync_event
854
-
803
eplb_model_state.rebalanced = True
856
- eplb_model_state.layer_to_transfer = 0
857
- eplb_model_state.pending_global_ready_check = True
804
# Signal async thread to start transferring layers
805
if self.is_async and (not is_profile):
860
- self.rearrange_event.set()
806
+ self.rearrange_event.record()
807
return None
808
809
def start_async_loop(
@@ -873,121 +819,27 @@ class EplbState:
819
is_profile=is_profile,
820
)
821
876
- def _all_ranks_buffer_ready(self, model_state: EplbModelState) -> bool:
822
+ def _all_ranks_result_ready(self, model_state: EplbModelState) -> bool:
823
parallel_state = get_ep_group()
824
+ has_result = int(model_state.pending_result is not None)
825
+
826
cpu_group = getattr(parallel_state, "cpu_group", None)
827
if cpu_group is not None and cpu_group.size() > 1:
880
- flag = torch.tensor(
881
- (int(model_state.ep_buffer_ready),), dtype=torch.int32, device="cpu"
882
- )
828
+ flag = torch.tensor((has_result,), dtype=torch.int32, device="cpu")
829
all_reduce(flag, group=cpu_group)
830
return int(flag.item()) == cpu_group.size()
831
832
device_group = parallel_state.device_group
833
if device_group.size() <= 1:
888
- return bool(model_state.ep_buffer_ready)
834
+ return bool(has_result)
835
836
device = getattr(
837
parallel_state, "device", model_state.physical_to_logical_map.device
838
)
893
- flag = torch.tensor(
894
- (int(model_state.ep_buffer_ready),), dtype=torch.int32, device=device
895
- )
839
+ flag = torch.tensor((has_result,), dtype=torch.int32, device=device)
840
all_reduce(flag, group=device_group)
841
return int(flag.item()) == device_group.size()
842
899
- def move_to_workspace(
900
- self,
901
- model_state: EplbModelState,
902
- ep_group: ProcessGroup,
903
- is_profile: bool = False,
904
- ):
905
- # We call move_to_workspace only when ep_buffer_ready is 1.
906
- # It means we only need to wait for the lock for a short time.
907
- max_retries = 6 # 1 minute max
908
- retries = 0
909
- while not model_state.buffer_lock.acquire(blocking=True, timeout=10.0):
910
- retries += 1
911
- if retries >= max_retries:
912
- raise RuntimeError(
913
- f"Rank {ep_group.rank()}: buffer_lock timeout after "
914
- "{max_retries * 10}s"
915
- )
916
- logger.warning(
917
- "Rank %d: EPLB buffer_lock acquire failed, retrying (%d/%d)",
918
- ep_group.rank(),
919
- retries,
920
- max_retries,
921
- )
922
- try:
923
- assert model_state.new_physical_to_logical_map is not None
924
- expert_weights = model_state.model.expert_weights[
925
- model_state.layer_to_transfer
926
- ]
927
- expert_weights_buffer = model_state.expert_buffer
928
- new_indices = model_state.new_physical_to_logical_map[
929
- model_state.layer_to_transfer
930
- ].numpy()
931
- move_from_buffer(
932
- expert_weights=expert_weights,
933
- expert_weights_buffers=expert_weights_buffer,
934
- is_unchanged=model_state.is_unchanged,
935
- is_received_locally=model_state.is_received_locally,
936
- recv_metadata=model_state.recv_metadata,
937
- new_indices=new_indices,
938
- ep_rank=ep_group.rank(),
939
- )
940
-
941
- transferred_layer = model_state.layer_to_transfer
942
-
943
- transferred_layer = model_state.layer_to_transfer
944
- assert model_state.new_physical_to_logical_map is not None
945
- _commit_eplb_maps_for_layer(
946
- model_state,
947
- new_physical_to_logical_map=model_state.new_physical_to_logical_map,
948
- layer=transferred_layer,
949
- )
950
-
951
- # Record event after consuming buffer to signal async thread
952
- # that it's safe to overwrite the intermediate buffer
953
- consumed_event = torch.cuda.Event()
954
- consumed_event.record()
955
- model_state.buffer_consumed_event = consumed_event
956
-
957
- # After the main thread consumes, advance layer_to_transfer
958
- model_state.layer_to_transfer += 1
959
- model_state.ep_buffer_ready = 0
960
- logger.debug(
961
- "model %s successfully move_to_workspace layer %d",
962
- model_state.model_name,
963
- transferred_layer,
964
- )
965
- if model_state.layer_to_transfer >= model_state.model.num_moe_layers:
966
- self.post_eplb(model_state)
967
- model_state.rebalanced = False
968
- model_state.layer_to_transfer = 0
969
- model_state.pending_global_ready_check = False
970
- logger.info(
971
- "finish async transfer for model %s rank %d layer %d",
972
- model_state.model_name,
973
- ep_group.rank(),
974
- model_state.model.num_moe_layers,
975
- )
976
-
977
- finally:
978
- try:
979
- model_state.buffer_lock.release()
980
- except Exception as e:
981
- logger.error(
982
- "Rank %d: buffer_lock release failed in move_to_workspace: %s",
983
- ep_group.rank(),
984
- str(e),
985
- )
986
-
987
- def post_eplb(self, model_state: EplbModelState) -> None:
988
- assert model_state.new_physical_to_logical_map is not None
989
- model_state.new_physical_to_logical_map = None
990
-
843
def _allreduce_list(self, tensor_list: list[torch.Tensor]) -> list[torch.Tensor]:
844
"""
845
All-reduce a list of tensors.
@@ -1225,7 +1077,7 @@ def _commit_eplb_maps_for_layer(
1077
"""
1078
1079
# Commit physical_to_logical_map
1228
- src = new_physical_to_logical_map[layer]
1080
+ src = new_physical_to_logical_map
1081
dst = model_state.physical_to_logical_map[layer]
1082
assert src.shape == dst.shape, (
1083
"The number of physical experts must stay the same while running Async EPLB. "
@@ -1284,3 +1136,33 @@ def _commit_eplb_maps(
1136
src = new_replica_count
1137
dst = model_state.logical_replica_count
1138
dst.copy_(src, non_blocking=True)
1139
+
1140
+
1141
+def _move_to_workspace(
1142
+ model_state: EplbModelState,
1143
+ ep_rank: int,
1144
+) -> None:
1145
+ result = model_state.pending_result
1146
+ assert result is not None
1147
+ move_from_buffer(
1148
+ expert_weights=model_state.model.expert_weights[result.layer_idx],
1149
+ expert_weights_buffers=model_state.expert_buffer,
1150
+ is_unchanged=result.is_unchanged,
1151
+ is_received_locally=result.is_received_locally,
1152
+ recv_metadata=result.recv_metadata,
1153
+ new_indices=result.new_physical_to_logical_map.numpy(),
1154
+ ep_rank=ep_rank,
1155
+ )
1156
+
1157
+ _commit_eplb_maps_for_layer(
1158
+ model_state,
1159
+ new_physical_to_logical_map=result.new_physical_to_logical_map,
1160
+ layer=result.layer_idx,
1161
+ )
1162
+
1163
+ if result.layer_idx == model_state.model.num_moe_layers - 1:
1164
+ model_state.rebalanced = False
1165
+
1166
+ # Reset pending_result before unblocking the async worker
1167
+ model_state.pending_result = None
1168
+ result.consumed_event.record()
vllm/distributed/eplb/eplb_utils.py
+51
@@ -3,6 +3,9 @@
3
"""Utility functions for EPLB (Expert Parallel Load Balancing)."""
4
5
import os
6
+import threading
7
+
8
+import torch
9
10
from vllm.config import ParallelConfig
11
from vllm.logger import init_logger
@@ -10,6 +13,54 @@ from vllm.logger import init_logger
13
logger = init_logger(__name__)
14
15
16
+class CpuGpuEvent:
17
+ """
18
+ Combines a CUDA event with a CPU threading event to enforce record->wait
19
+ ordering across two threads.
20
+
21
+ This class is designed for exactly two threads: one producer that calls
22
+ record() and one consumer that calls wait(). Using it with more than two
23
+ threads is not supported and will produce undefined behavior.
24
+
25
+ CUDA events alone are insufficient for cross-thread synchronization because
26
+ waiting on an unrecorded CUDA event is a no-op. The wait will return
27
+ immediately instead of blocking. This class adds a threading.Event so
28
+ that the waiting thread blocks on the CPU side until record() is called, at
29
+ which point the CUDA event is guaranteed to be in-flight and event.wait() will
30
+ correctly synchronize the GPU stream.
31
+ """
32
+
33
+ def __init__(self):
34
+ self._event = torch.cuda.Event()
35
+ self._recorded = threading.Event()
36
+
37
+ def wait(self, stream: torch.cuda.Stream | None = None):
38
+ """
39
+ Blocks the calling thread until record finishes. Used to guarantee that the
40
+ record kernel is called before wait.
41
+
42
+ Should only be called by the Async Eplb thread.
43
+ """
44
+ self._recorded.wait()
45
+ self._event.wait(stream)
46
+ self._recorded.clear()
47
+
48
+ def record(self, stream: torch.cuda.Stream | None = None):
49
+ """
50
+ Unblocks the waiting thread after calling event.record().
51
+
52
+ Should only be called by the main thread.
53
+ """
54
+ if self._recorded.is_set():
55
+ raise RuntimeError(
56
+ "CpuGpuEvent.record() called before the previous event was "
57
+ "consumed by wait()"
58
+ )
59
+ self._event = torch.cuda.Event()
60
+ self._event.record(stream)
61
+ self._recorded.set()
62
+
63
+
64
def override_envs_for_eplb(parallel_config: ParallelConfig) -> None:
65
"""
66
Override environment variables for EPLB when specific conditions are met.
vllm/distributed/eplb/rebalance_execute.py
+33
-1
@@ -13,7 +13,11 @@ import numpy as np
13
import torch
14
from torch.distributed import ProcessGroup, all_gather
15
16
-from .eplb_communicator import EplbCommunicator
16
+from vllm.distributed.eplb.eplb_communicator import EplbCommunicator
17
+from vllm.distributed.eplb.eplb_utils import CpuGpuEvent
18
+from vllm.logger import init_logger
19
+
20
+logger = init_logger(__name__)
21
22
23
@dataclass
@@ -34,6 +38,34 @@ class RecvMetadata:
38
MoveToBufferResult = tuple[np.ndarray, np.ndarray, RecvMetadata]
39
40
41
+@dataclass
42
+class AsyncEplbLayerResult:
43
+ """
44
+ The result of one completed async EPLB layer transfer.
45
+ """
46
+
47
+ layer_idx: int
48
+ """Index of the MoE layer that was transferred."""
49
+ new_physical_to_logical_map: torch.Tensor
50
+ """
51
+ New physical→logical mapping for layers_idx, on CPU.
52
+ Shape: (num_physical_experts)
53
+ """
54
+ is_unchanged: np.ndarray
55
+ """Per-physical-expert flag: weight was not moved during transfer."""
56
+ is_received_locally: np.ndarray
57
+ """Per-physical-expert flag: weight was received on this rank."""
58
+ recv_metadata: RecvMetadata
59
+ """Metadata describing what was received during transfer_layer."""
60
+ consumed_event: CpuGpuEvent
61
+ """
62
+ Event used to synchronize access to the intermediate buffer. The async worker calls
63
+ wait() after it finishes transferring weights to the intermediate buffer. The main
64
+ thread calls record() after it finishes transferring weights out of the intermediate
65
+ buffer in _move_to_workspace()
66
+ """
67
+
68
+
69
def get_ep_ranks_with_experts_batch(
70
expert_ids: np.ndarray,
71
num_local_experts: int,