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, 17:30 UTC
[EPLB] Remove asyncio infrastructure from Async EPLB (#40730)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Apr 23, 2026 at 17:21 UTC
62b1bbe470ea820b728cf2b6eee4f5f8b9655a57
3 files changed
+16
-26
tests/distributed/test_eplb_execute.py
+8
-11
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import asyncio
4
import random
5
6
import pytest
@@ -361,16 +360,14 @@ def _test_async_transfer_layer_without_mtp_worker(
360
communicator.set_stream(cuda_stream)
361
362
for layer_idx in range(num_layers):
364
- transfer_metadata = asyncio.run(
365
- transfer_layer(
366
- old_layer_indices=old_indices_cpu[layer_idx],
367
- new_layer_indices=new_indices_cpu[layer_idx],
368
- expert_weights=expert_weights[layer_idx],
369
- expert_weights_buffer=expert_buffer,
370
- ep_group=ep_group,
371
- communicator=communicator,
372
- cuda_stream=cuda_stream,
373
- )
363
+ transfer_metadata = transfer_layer(
364
+ old_layer_indices=old_indices_cpu[layer_idx],
365
+ new_layer_indices=new_indices_cpu[layer_idx],
366
+ expert_weights=expert_weights[layer_idx],
367
+ expert_weights_buffer=expert_buffer,
368
+ ep_group=ep_group,
369
+ communicator=communicator,
370
+ cuda_stream=cuda_stream,
371
)
372
cuda_stream.synchronize()
373
move_from_buffer(
vllm/distributed/eplb/async_worker.py
+7
-14
@@ -4,7 +4,6 @@
4
The async worker that transfers experts in the background.
5
"""
6
7
-import asyncio
7
import threading
8
from typing import TYPE_CHECKING
9
@@ -36,21 +35,15 @@ def start_async_worker(
35
assert device_index is not None
36
torch.accelerator.set_device_index(device_index)
37
cuda_stream = torch.cuda.Stream(device=device_index)
39
- loop = asyncio.new_event_loop()
40
- asyncio.set_event_loop(loop)
38
try:
42
- loop.run_until_complete(
43
- transfer_run_periodically(
44
- state=state,
45
- eplb_group=eplb_group,
46
- cuda_stream=cuda_stream,
47
- is_profile=is_profile,
48
- )
39
+ transfer_run_periodically(
40
+ state=state,
41
+ eplb_group=eplb_group,
42
+ cuda_stream=cuda_stream,
43
+ is_profile=is_profile,
44
)
45
except Exception as exc: # pragma: no cover - diagnostic path
46
logger.exception("async loop error (Rank %d): %s", rank, str(exc))
52
- finally:
53
- loop.close()
47
48
thread = threading.Thread(target=thread_target, daemon=True)
49
thread.start()
@@ -83,7 +76,7 @@ def run_rebalance_experts(
76
return new_physical_to_logical_map
77
78
86
-async def transfer_run_periodically(
79
+def transfer_run_periodically(
80
state: "EplbState",
81
eplb_group: ProcessGroup,
82
cuda_stream: torch.cuda.Stream,
@@ -118,7 +111,7 @@ async def transfer_run_periodically(
111
# model_state.expert_buffer, which will be consumed by the main thread in
112
# move_to_workspace
113
while model_state.rebalanced and layer_idx < num_layers:
121
- transfer_metadata = await transfer_layer(
114
+ transfer_metadata = transfer_layer(
115
old_layer_indices=physical_to_logical_map_cpu[layer_idx],
116
new_layer_indices=new_physical_to_logical_map[layer_idx],
117
expert_weights=model_state.model.expert_weights[layer_idx],
vllm/distributed/eplb/rebalance_execute.py
+1
-1
@@ -418,7 +418,7 @@ def move_from_buffer(
418
w[dst].copy_(w[src], non_blocking=True)
419
420
421
-async def transfer_layer(
421
+def transfer_layer(
422
old_layer_indices: torch.Tensor,
423
new_layer_indices: torch.Tensor,
424
expert_weights: Sequence[torch.Tensor],