From 62b1bbe470ea820b728cf2b6eee4f5f8b9655a57 Mon Sep 17 00:00:00 2001
From: Sage Moore <sage@neuralmagic.com>
Date: Thu, 23 Apr 2026 17:21:15 -0700
Subject: [PATCH] [EPLB] Remove asyncio infrastructure from Async EPLB (#40730)

Signed-off-by: Sage Moore <sage@neuralmagic.com>
---
 tests/distributed/test_eplb_execute.py     | 19 ++++++++-----------
 vllm/distributed/eplb/async_worker.py      | 21 +++++++--------------
 vllm/distributed/eplb/rebalance_execute.py |  2 +-
 3 files changed, 16 insertions(+), 26 deletions(-)

diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py
index 7f8895cd2c1..d9e6a739b01 100644
--- a/tests/distributed/test_eplb_execute.py
+++ b/tests/distributed/test_eplb_execute.py
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import asyncio
 import random
 
 import pytest
@@ -361,16 +360,14 @@ def _test_async_transfer_layer_without_mtp_worker(
         communicator.set_stream(cuda_stream)
 
         for layer_idx in range(num_layers):
-            transfer_metadata = asyncio.run(
-                transfer_layer(
-                    old_layer_indices=old_indices_cpu[layer_idx],
-                    new_layer_indices=new_indices_cpu[layer_idx],
-                    expert_weights=expert_weights[layer_idx],
-                    expert_weights_buffer=expert_buffer,
-                    ep_group=ep_group,
-                    communicator=communicator,
-                    cuda_stream=cuda_stream,
-                )
+            transfer_metadata = transfer_layer(
+                old_layer_indices=old_indices_cpu[layer_idx],
+                new_layer_indices=new_indices_cpu[layer_idx],
+                expert_weights=expert_weights[layer_idx],
+                expert_weights_buffer=expert_buffer,
+                ep_group=ep_group,
+                communicator=communicator,
+                cuda_stream=cuda_stream,
             )
             cuda_stream.synchronize()
             move_from_buffer(
diff --git a/vllm/distributed/eplb/async_worker.py b/vllm/distributed/eplb/async_worker.py
index a47b5ce29c2..542606fe741 100644
--- a/vllm/distributed/eplb/async_worker.py
+++ b/vllm/distributed/eplb/async_worker.py
@@ -4,7 +4,6 @@
 The async worker that transfers experts in the background.
 """
 
-import asyncio
 import threading
 from typing import TYPE_CHECKING
 
@@ -36,21 +35,15 @@ def start_async_worker(
         assert device_index is not None
         torch.accelerator.set_device_index(device_index)
         cuda_stream = torch.cuda.Stream(device=device_index)
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
         try:
-            loop.run_until_complete(
-                transfer_run_periodically(
-                    state=state,
-                    eplb_group=eplb_group,
-                    cuda_stream=cuda_stream,
-                    is_profile=is_profile,
-                )
+            transfer_run_periodically(
+                state=state,
+                eplb_group=eplb_group,
+                cuda_stream=cuda_stream,
+                is_profile=is_profile,
             )
         except Exception as exc:  # pragma: no cover - diagnostic path
             logger.exception("async loop error (Rank %d): %s", rank, str(exc))
-        finally:
-            loop.close()
 
     thread = threading.Thread(target=thread_target, daemon=True)
     thread.start()
@@ -83,7 +76,7 @@ def run_rebalance_experts(
     return new_physical_to_logical_map
 
 
-async def transfer_run_periodically(
+def transfer_run_periodically(
     state: "EplbState",
     eplb_group: ProcessGroup,
     cuda_stream: torch.cuda.Stream,
@@ -118,7 +111,7 @@ async def transfer_run_periodically(
             # model_state.expert_buffer, which will be consumed by the main thread in
             # move_to_workspace
             while model_state.rebalanced and layer_idx < num_layers:
-                transfer_metadata = await transfer_layer(
+                transfer_metadata = transfer_layer(
                     old_layer_indices=physical_to_logical_map_cpu[layer_idx],
                     new_layer_indices=new_physical_to_logical_map[layer_idx],
                     expert_weights=model_state.model.expert_weights[layer_idx],
diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py
index a68fbda86cc..f348521c00e 100644
--- a/vllm/distributed/eplb/rebalance_execute.py
+++ b/vllm/distributed/eplb/rebalance_execute.py
@@ -418,7 +418,7 @@ def move_from_buffer(
             w[dst].copy_(w[src], non_blocking=True)
 
 
-async def transfer_layer(
+def transfer_layer(
     old_layer_indices: torch.Tensor,
     new_layer_indices: torch.Tensor,
     expert_weights: Sequence[torch.Tensor],