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 6, 03:15 UTC
[EPLB] Remove main waits in case of slow EPLB (#36271)
Signed-off-by: ilmarkov <markovilya197@gmail.com>
Ilya Markov committed
Mar 24, 2026 at 12:50 UTC
14acf429ac08b6d538ca6feb3e06b6d13895804d
2 files changed
+3
-14
vllm/distributed/eplb/async_worker.py
+3
-3
@@ -160,9 +160,9 @@ async def transfer_run_periodically(
160
is_profile=is_profile,
161
cuda_stream=cuda_stream,
162
)
163
- event = torch.cuda.Event(blocking=False)
164
- cuda_stream.record_event(event)
165
- model_state.buffer_ready_event = event
163
+ # block the async thread until the transfer to
164
+ # the intermediate buffer is complete.
165
+ cuda_stream.synchronize()
166
model_state.ep_buffer_ready = 1
167
finally:
168
model_state.buffer_lock.release()
vllm/distributed/eplb/eplb_state.py
-11
@@ -176,11 +176,6 @@ class EplbModelState:
176
"""
177
The lock to protect the expert buffer.
178
"""
179
- buffer_ready_event: torch.cuda.Event | None
180
- """
181
- CUDA event recorded when the async worker finishes filling the buffer.
182
- The main thread waits on this before consuming the buffer.
183
- """
179
buffer_consumed_event: torch.cuda.Event | None
180
"""
181
CUDA event recorded after the main thread finishes consuming the buffer.
@@ -480,7 +475,6 @@ class EplbState:
475
model=model,
476
expert_buffer=expert_buffer,
477
buffer_lock=threading.Lock(),
483
- buffer_ready_event=None,
478
buffer_consumed_event=None,
479
window_ready_event=None,
480
ep_buffer_ready=0,
@@ -919,11 +913,6 @@ class EplbState:
913
)
914
try:
915
assert model_state.new_physical_to_logical_map is not None
922
- device_index = model_state.cuda_device_index or self.cuda_device_index
923
- if model_state.buffer_ready_event is not None and device_index is not None:
924
- stream = torch.cuda.current_stream(device=device_index)
925
- stream.wait_event(model_state.buffer_ready_event)
926
- model_state.buffer_ready_event = None
916
expert_weights = model_state.model.expert_weights[
917
model_state.layer_to_transfer
918
]