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 10, 05:15 UTC
[EPLB][Cleanup] Remove `is_async_enabled` from `EplbModelState` (#32050)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Jan 13, 2026 at 10:19 UTC
6beef12b9b8a412c4f05566903bf53d0ccbe44f9
2 files changed
+8
-20
vllm/distributed/eplb/async_worker.py
+4
-4
@@ -30,6 +30,7 @@ def start_async_worker(
30
ep_group = get_ep_group().device_group
31
rank = ep_group.rank()
32
device_index = state.cuda_device_index
33
+ assert state.is_async
34
35
def thread_target() -> None:
36
assert device_index is not None
@@ -42,9 +43,9 @@ def start_async_worker(
43
transfer_run_periodically(
44
state=state,
45
ep_group=ep_group,
46
+ cuda_stream=cuda_stream,
47
is_profile=is_profile,
48
rank_mapping=rank_mapping,
47
- cuda_stream=cuda_stream,
49
)
50
)
51
except Exception as exc: # pragma: no cover - diagnostic path
@@ -60,17 +61,16 @@ def start_async_worker(
61
async def transfer_run_periodically(
62
state: "EplbState",
63
ep_group: ProcessGroup,
64
+ cuda_stream: torch.cuda.Stream,
65
is_profile: bool = False,
66
rank_mapping: dict[int, int] | None = None,
65
- cuda_stream: torch.cuda.Stream = None,
67
) -> None:
68
while True:
69
await asyncio.to_thread(state.rearrange_event.wait)
70
logger.info("async worker woke up for EPLB transfer")
71
72
+ assert state.is_async
73
for model_state in state.model_states.values():
72
- if not model_state.is_async_enabled:
73
- continue
74
current_num_layers = model_state.model.num_moe_layers
75
while (
76
model_state.rebalanced
vllm/distributed/eplb/eplb_state.py
+4
-16
@@ -182,10 +182,6 @@ class EplbModelState:
182
"""
183
intermediate variable between `move_to_buffer` and `move_to_workspace`.
184
"""
185
- is_async_enabled: bool
186
- """
187
- The flag indicates whether the EPLB is running in async mode.
188
- """
185
cuda_device_index: int | None
186
"""
187
CUDA device index for the async EPLB worker thread.
@@ -518,7 +514,6 @@ class EplbState:
514
recv_expert_ids=np.array([]),
515
recv_dst_rows=np.array([]),
516
),
521
- is_async_enabled=self.is_async,
517
cuda_device_index=self.cuda_device_index,
518
new_physical_to_logical_map=new_physical_to_logical_map,
519
new_logical_to_physical_map=new_logical_to_physical_map,
@@ -630,19 +625,12 @@ class EplbState:
625
626
if self.is_async:
627
for eplb_model_state in self.model_states.values():
633
- if not eplb_model_state.is_async_enabled:
634
- continue
635
-
628
all_ranks_buffer_ready = False
629
if eplb_model_state.pending_global_ready_check:
630
all_ranks_buffer_ready = self._all_ranks_buffer_ready(
631
eplb_model_state
632
)
641
- if (
642
- eplb_model_state.is_async_enabled
643
- and eplb_model_state.ep_buffer_ready
644
- and all_ranks_buffer_ready
645
- ):
633
+ if eplb_model_state.ep_buffer_ready and all_ranks_buffer_ready:
634
self.move_to_workspace(
635
model_state=eplb_model_state,
636
ep_group=ep_group,
@@ -664,8 +652,8 @@ class EplbState:
652
)
653
654
if self.expert_rearrangement_step >= self.expert_rearrangement_step_interval:
667
- if any(
668
- eplb_model_state.is_async_enabled and eplb_model_state.rebalanced
655
+ if self.is_async and any(
656
+ eplb_model_state.rebalanced
657
for eplb_model_state in self.model_states.values()
658
):
659
# Still performing asynchronous rearrangement
@@ -822,7 +810,7 @@ class EplbState:
810
eplb_model_state.physical_to_logical_map,
811
)
812
825
- if not eplb_model_state.is_async_enabled or is_profile:
813
+ if not self.is_async or is_profile:
814
# Update expert weights
815
rearrange_expert_weights_inplace(
816
eplb_model_state.physical_to_logical_map,