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 9, 10:00 UTC
[BugFix] Async Eplb fix potential race condition (#32881)
Signed-off-by: ilmarkov <markovilya197@gmail.com>
Ilya Markov committed
Jan 29, 2026 at 11:31 UTC
d09135fbd0b02046f3d75a6f0e8de28e70249e46
2 files changed
+18
vllm/distributed/eplb/async_worker.py
+6
@@ -86,6 +86,12 @@ async def transfer_run_periodically(
86
if model_state.layer_to_transfer >= current_num_layers:
87
break
88
89
+ # Wait for the main thread to finish consuming the buffer
90
+ # before overwriting it
91
+ if model_state.buffer_consumed_event is not None:
92
+ cuda_stream.wait_event(model_state.buffer_consumed_event)
93
+ model_state.buffer_consumed_event = None
94
+
95
(
96
model_state.is_unchanged,
97
model_state.is_received_locally,
vllm/distributed/eplb/eplb_state.py
+12
@@ -151,6 +151,11 @@ class EplbModelState:
151
CUDA event recorded when the async worker finishes filling the buffer.
152
The main thread waits on this before consuming the buffer.
153
"""
154
+ buffer_consumed_event: torch.cuda.Event | None
155
+ """
156
+ CUDA event recorded after the main thread finishes consuming the buffer.
157
+ The async worker waits on this before writing to the buffer again.
158
+ """
159
ep_buffer_ready: int
160
"""
161
The flag indicates whether the expert buffer is ready for transfer.
@@ -502,6 +507,7 @@ class EplbState:
507
expert_buffer=expert_buffer,
508
buffer_lock=threading.Lock(),
509
buffer_ready_event=None,
510
+ buffer_consumed_event=None,
511
ep_buffer_ready=0,
512
layer_to_transfer=0,
513
rebalanced=False,
@@ -1012,6 +1018,12 @@ class EplbState:
1018
new_indices=new_indices,
1019
ep_rank=ep_group.rank(),
1020
)
1021
+ # Record event after consuming buffer to signal async thread
1022
+ # that it's safe to overwrite the buffer
1023
+ consumed_event = torch.cuda.Event()
1024
+ consumed_event.record()
1025
+ model_state.buffer_consumed_event = consumed_event
1026
+
1027
transferred_layer = model_state.layer_to_transfer
1028
self._update_layer_mapping_from_new(model_state, transferred_layer)
1029
# After the main thread consumes, advance layer_to_transfer