@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ac614587f5

[EPLB] Enable nixl eplb communicator for elastic ep (#45013)

Signed-off-by: Markov Ilya <markovilya197@gmail.com> Signed-off-by: Markov Ilya <markovilya19@gmail.com> Co-authored-by: Markov Ilya <markovilya19@gmail.com>

Ilya Markov committed Jun 22, 2026 at 19:54 UTC ac614587f514e9032bff75eb04c86ddeb3c9dbef

7 files changed +329 -97

tests/distributed/test_elastic_ep.py

+34 -35

@@ -59,9 +59,8 @@ def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float:

59 return accuracy

62 -@multi_gpu_test(num_gpus=4)

63 -def test_elastic_ep_scaling():

64 - vllm_serve_args = [

62 +def _base_serve_args(use_async_eplb: bool = False) -> list[str]:

63 + args = [

64 "--trust-remote-code",

65 "--tensor-parallel-size",

66 "1",

@@ -79,7 +78,11 @@ def test_elastic_ep_scaling():

78 "--eplb-config.num_redundant_experts",

79 "0",

80 "--eplb-config.use_async",

82 - "false",

81 + "true" if use_async_eplb else "false",

82 + "--eplb-config.step_interval",

83 + "10",

84 + "--eplb-config.window_size",

85 + "5",

86 "--data-parallel-backend",

87 "ray",

88 "--data-parallel-size",

@@ -90,7 +93,23 @@ def test_elastic_ep_scaling():

94 leader_address = os.environ.get("LEADER_ADDRESS")

95 if leader_address:

93 - vllm_serve_args.extend(["--data-parallel-address", leader_address])

96 + args.extend(["--data-parallel-address", leader_address])

97 +

98 + return args

99 +

100 +

101 +@pytest.mark.parametrize(

102 + "use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"]

103 +)

104 +@multi_gpu_test(num_gpus=4)

105 +def test_elastic_ep_scaling(use_async_eplb: bool):

106 + if use_async_eplb:

107 + from vllm.distributed.eplb.eplb_communicator import has_nixl

108 +

109 + if not has_nixl():

110 + pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")

111 +

112 + vllm_serve_args = _base_serve_args(use_async_eplb)

113

114 with RemoteOpenAIServer(

115 MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200

@@ -128,44 +147,24 @@ def test_elastic_ep_scaling():

147 print(f" Tolerance: {ACCURACY_TOL:.3f}")

148

149

150 +@pytest.mark.parametrize(

151 + "use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"]

152 +)

153 @multi_gpu_test(num_gpus=4)

132 -def test_elastic_ep_scaling_uneven():

154 +def test_elastic_ep_scaling_uneven(use_async_eplb: bool):

155 """Test scale up with uneven worker distribution.

156

157 This tests the case where num_new_workers % old_dp_size != 0,

158 specifically 2 -> 3 where remainder = 1 % 2 = 1.

159 This exercises the remainder handling in sender-receiver pairing.

160 """

139 - vllm_serve_args = [

140 - "--trust-remote-code",

141 - "--tensor-parallel-size",

142 - "1",

143 - "--gpu-memory-utilization",

144 - "0.8",

145 - "--max-model-len",

146 - "4096",

147 - "--max-num-seqs",

148 - str(MAX_NUM_SEQS),

149 - "--enable-expert-parallel",

150 - "--all2all-backend",

151 - "allgather_reducescatter",

152 - "--enable-elastic-ep",

153 - "--enable-eplb",

154 - "--eplb-config.num_redundant_experts",

155 - "0",

156 - "--eplb-config.use_async",

157 - "false",

158 - "--data-parallel-backend",

159 - "ray",

160 - "--data-parallel-size",

161 - "2",

162 - "--api-server-count",

163 - "1",

164 - ]

161 + if use_async_eplb:

162 + from vllm.distributed.eplb.eplb_communicator import has_nixl

163

166 - leader_address = os.environ.get("LEADER_ADDRESS")

167 - if leader_address:

168 - vllm_serve_args.extend(["--data-parallel-address", leader_address])

164 + if not has_nixl():

165 + pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")

166 +

167 + vllm_serve_args = _base_serve_args(use_async_eplb)

168

169 with RemoteOpenAIServer(

170 MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200

tests/distributed/test_eplb_execute.py

+122

@@ -782,3 +782,125 @@ def test_rearrange_expert_weights_profile_mode(world_size):

782 _test_rearrange_expert_weights_profile_mode,

783 world_size,

784 )

785 +

786 +

787 +def _test_nixl_deferred_init_worker(

788 + env,

789 + world_size: int,

790 + num_layers: int,

791 + num_local_experts: int,

792 + num_logical_experts: int,

793 +) -> None:

794 + """Exercise NixlEplbCommunicator with defer_remote_setup=True (elastic EP path)."""

795 + from vllm.distributed.eplb.eplb_communicator import NixlEplbCommunicator

796 +

797 + set_env_vars_and_device(env)

798 +

799 + vllm_config = VllmConfig()

800 + vllm_config.parallel_config.tensor_parallel_size = world_size

801 +

802 + with set_current_vllm_config(vllm_config):

803 + ensure_model_parallel_initialized(

804 + tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1

805 + )

806 +

807 + ep_group_coordinator = get_tp_group()

808 + ep_group = ep_group_coordinator.cpu_group

809 + ep_rank = torch.distributed.get_rank()

810 + device = torch.device(f"cuda:{ep_rank}")

811 +

812 + total_physical_experts = world_size * num_local_experts

813 + hidden_sizes = [32, 64]

814 +

815 + redundancy_config = create_redundancy_config(

816 + num_logical_experts, total_physical_experts

817 + )

818 + old_indices = create_expert_indices_with_redundancy(

819 + num_layers,

820 + num_logical_experts,

821 + total_physical_experts,

822 + redundancy_config,

823 + )

824 +

825 + new_redundancy_config = create_redundancy_config(

826 + num_logical_experts, total_physical_experts

827 + )

828 + new_indices = create_expert_indices_with_redundancy(

829 + num_layers,

830 + num_logical_experts,

831 + total_physical_experts,

832 + new_redundancy_config,

833 + )

834 +

835 + expert_weights = create_expert_weights(

836 + num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices

837 + )

838 +

839 + expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]

840 +

841 + communicator = NixlEplbCommunicator(

842 + cpu_group=ep_group_coordinator.cpu_group,

843 + all_expert_weights=expert_weights,

844 + expert_buffer=expert_buffer,

845 + defer_remote_setup=True,

846 + )

847 + assert not communicator._remote_state_initialized

848 +

849 + rearrange_expert_weights_inplace(

850 + old_indices,

851 + new_indices,

852 + expert_weights,

853 + expert_buffer,

854 + ep_group,

855 + communicator,

856 + )

857 +

858 + assert communicator._remote_state_initialized

859 +

860 + local_ok = verify_expert_weights_after_shuffle(

861 + expert_weights,

862 + new_indices,

863 + hidden_sizes,

864 + ep_rank,

865 + num_local_experts,

866 + )

867 +

868 + local_ok = (

869 + verify_redundant_experts_have_same_weights(

870 + expert_weights,

871 + new_indices,

872 + hidden_sizes,

873 + ep_rank,

874 + world_size,

875 + num_local_experts,

876 + )

877 + and local_ok

878 + )

879 + assert_verification_synced(

880 + local_ok,

881 + "Deferred NIXL init verification failed on at least one rank.",

882 + )

883 +

884 +

885 +@pytest.mark.skipif(not has_nixl(), reason="NIXL is not available")

886 +@pytest.mark.parametrize(

887 + "world_size,num_layers,num_local_experts,num_logical_experts",

888 + [(2, 2, 3, 4)],

889 +)

890 +def test_nixl_deferred_init(

891 + world_size,

892 + num_layers,

893 + num_local_experts,

894 + num_logical_experts,

895 +):

896 + """Test NixlEplbCommunicator with defer_remote_setup=True (elastic EP path)."""

897 +

898 + if torch.accelerator.device_count() < world_size:

899 + pytest.skip(f"Need at least {world_size} GPUs to run the test")

900 + distributed_run(

901 + _test_nixl_deferred_init_worker,

902 + world_size,

903 + num_layers,

904 + num_local_experts,

905 + num_logical_experts,

906 + )

vllm/config/parallel.py

+22 -22

@@ -803,13 +803,6 @@ class ParallelConfig:

803 if self.enable_elastic_ep:

804 if not self.enable_eplb:

805 raise ValueError("Elastic EP is only supported with enable_eplb=True.")

806 - if self.eplb_config.use_async:

807 - raise ValueError(

808 - "Elastic EP requires the pynccl communicator, which is "

809 - "incompatible with async EPLB due to NCCL multi-stream "

810 - "conflicts. Disable async EPLB (eplb_config.use_async=False) "

811 - "to use elastic EP."

812 - )

806 if self.pipeline_parallel_size > 1:

807 raise ValueError(

808 "Elastic EP is not supported with pipeline parallelism "

@@ -821,6 +814,15 @@ class ParallelConfig:

814 "or data_parallel_hybrid_lb. Elastic EP relies on a single API "

815 "server and core client to coordinate scale up/down."

816 )

817 + if self.eplb_config.use_async:

818 + from vllm.distributed.nixl_utils import is_nixl_available

819 +

820 + if not is_nixl_available():

821 + raise ValueError(

822 + "Elastic EP with async EPLB requires the NIXL "

823 + "package. Either install NIXL or set "

824 + "--eplb-config.use_async=false."

825 + )

826

827 if self.data_parallel_size > 1 or self.data_parallel_size_local == 0:

828 # Data parallel was specified in the engine args.

@@ -929,23 +931,21 @@ class ParallelConfig:

931 )

932

933 if self.enable_eplb and self.eplb_config.communicator is None:

932 - if self.enable_elastic_ep:

933 - # Elastic EP requires stateless mode

934 - # (torch.distributed.batch_isend_irecv doesn't

935 - # support stateless mode), so we use PyNCCL backend

934 + # Prefer NIXL when available: zero-copy RDMA reads, compatible

935 + # with both async EPLB and elastic EP (deferred remote setup).

936 + # Fallbacks: pynccl for elastic EP (stateless groups need it),

937 + # torch_gloo for static EP. torch_nccl is avoided because NCCL

938 + # is incompatible with async EPLB (multi-stream conflicts) and

939 + # batched isend/irecv hangs under high load.

940 + # See https://github.com/pytorch/pytorch/issues/174288

941 + from vllm.distributed.nixl_utils import is_nixl_available

942 +

943 + if is_nixl_available():

944 + self.eplb_config.communicator = "nixl"

945 + elif self.enable_elastic_ep:

946 self.eplb_config.communicator = "pynccl"

947 else:

938 - # Avoid torch_nccl: NCCL is fundamentally incompatible

939 - # with async EPLB due to multi-stream conflicts, and

940 - # batched isend/irecv hangs under high load.

941 - # See https://github.com/pytorch/pytorch/issues/174288

942 - # Prefer nixl when available; fall back to torch_gloo.

943 - from vllm.distributed.nixl_utils import is_nixl_available

944 -

945 - if is_nixl_available():

946 - self.eplb_config.communicator = "nixl"

947 - else:

948 - self.eplb_config.communicator = "torch_gloo"

948 + self.eplb_config.communicator = "torch_gloo"

949

950 @property

951 def use_ray(self) -> bool:

vllm/distributed/elastic_ep/elastic_execute.py

+11

@@ -207,6 +207,9 @@ class ElasticEPScalingExecutor:

207 )

208 if new_dp_size > old_dp_size:

209 self._set_eplb_suppressed(True)

210 + eplb_state = self.worker.model_runner.eplb_state

211 + if eplb_state is not None:

212 + eplb_state.drain_async()

213 elif new_dp_size < old_dp_size:

214 self._stage_standby_moe_quant_methods()

215

@@ -540,6 +543,11 @@ class ElasticEPScalingExecutor:

543 eplb_model_state.physical_to_logical_map.shape[1]

544 )

545 eplb_state.is_async = is_async_enabled

546 + # Start the async worker thread if it doesn't exist yet (idempotent).

547 + # This is needed for new workers after scale-up: they create EplbState

548 + # in setup_eplb_from_mapping() but don't start the thread there because

549 + # groups aren't ready yet.

550 + eplb_state.start_async_loop()

551 if get_ep_group().rank == 0:

552 logger.info("[Elastic EP] Expert resharding completed")

553

@@ -549,6 +557,9 @@ class ElasticEPScalingExecutor:

557

558 def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None:

559 self._set_eplb_suppressed(True)

560 + eplb_state = self.worker.model_runner.eplb_state

561 + if eplb_state is not None:

562 + eplb_state.drain_async()

563 parallel_config = self.worker.vllm_config.parallel_config

564 tp_size = parallel_config.tensor_parallel_size

565 old_ep_size = parallel_config.data_parallel_size * tp_size

vllm/distributed/eplb/async_worker.py

+27 -13

@@ -8,7 +8,6 @@ import threading

8 from typing import TYPE_CHECKING

10 import torch

11 -from torch.distributed import ProcessGroup

12 from vllm.distributed.parallel_state import get_eplb_group

13 from vllm.logger import init_logger

@@ -26,8 +25,7 @@ def start_async_worker(

25 state: "EplbState",

26 is_profile: bool = False,

27 ) -> threading.Thread:

29 - eplb_group = get_eplb_group().device_group

30 - rank = eplb_group.rank()

28 + rank = get_eplb_group().device_group.rank()

29 device_index = state.cuda_device_index

30 assert state.is_async

@@ -38,7 +36,6 @@ def start_async_worker(

36 try:

37 transfer_run_periodically(

38 state=state,

41 - eplb_group=eplb_group,

39 cuda_stream=cuda_stream,

40 is_profile=is_profile,

41 )

@@ -78,13 +75,15 @@ def run_rebalance_experts(

76 def transfer_run_periodically(

77 state: "EplbState",

81 - eplb_group: ProcessGroup,

78 cuda_stream: torch.cuda.Stream,

79 is_profile: bool = False,

80 ) -> None:

81 while True:

82 state.rearrange_event.wait(stream=cuda_stream)

87 - logger.info("async worker woke up for EPLB transfer")

83 +

84 + eplb_group = get_eplb_group().device_group

85 + eplb_cpu_group = get_eplb_group().cpu_group

86 + ep_rank = eplb_group.rank()

88 assert state.is_async

89 for model_state in state.model_states.values():

@@ -101,16 +100,32 @@ def transfer_run_periodically(

100 new_physical_to_logical_map = run_rebalance_experts(

101 model_state, state, physical_to_logical_map_cpu, cuda_stream

102 )

104 - logger.info(

105 - "Async worker computed new indices for model %s",

106 - model_state.model_name,

107 - )

103

104 # Execute one EPLB layer transfer per model forward pass. Each iteration

105 # of this loop will copy the new set of expert weights into

106 # 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:

107 + # move_to_workspace.

108 + # We sync the rebalanced flag across ranks before each iteration so

109 + # all ranks make a coordinated decision to continue or stop.

110 + while layer_idx < num_layers:

111 + flag = torch.tensor(

112 + [int(model_state.rebalanced)],

113 + dtype=torch.int32,

114 + device="cpu",

115 + )

116 + torch.distributed.all_reduce(flag, group=eplb_cpu_group)

117 + if int(flag.item()) != eplb_cpu_group.size():

118 + logger.warning(

119 + "async worker (rank=%d): layer %d coordinated stop "

120 + "(flag_sum=%d, group_size=%d)",

121 + ep_rank,

122 + layer_idx,

123 + int(flag.item()),

124 + eplb_cpu_group.size(),

125 + )

126 + model_state.rebalanced = False

127 + break

128 +

129 transfer_metadata = transfer_layer(

130 old_layer_indices=physical_to_logical_map_cpu[layer_idx],

131 new_layer_indices=new_physical_to_logical_map[layer_idx],

@@ -143,6 +158,5 @@ def transfer_run_periodically(

158 # finish copying model_state.expert_buffer into

159 # model_state.model.expert_weights[layer_idx]

160 consumed_event.wait(stream=cuda_stream)

146 - logger.debug("Layer %d transfer complete", layer_idx)

161 assert model_state.pending_result is None

162 layer_idx += 1

vllm/distributed/eplb/eplb_communicator.py

+71 -25

@@ -246,7 +246,20 @@ class NixlEplbCommunicator(EplbCommunicator):

246 cpu_group: ProcessGroup,

247 all_expert_weights: Sequence[Sequence[torch.Tensor]],

248 expert_buffer: Sequence[torch.Tensor],

249 + defer_remote_setup: bool = False,

250 ) -> None:

251 + """Create a NIXL-backed EPLB communicator.

252 +

253 + Args:

254 + cpu_group: CPU process group for metadata exchange.

255 + all_expert_weights: Expert weight tensors for all MoE layers.

256 + expert_buffer: Pre-allocated receive buffer tensors.

257 + defer_remote_setup: If True, postpone the collective

258 + all-gather of NIXL agent metadata until the first

259 + ``set_transfer_context`` call. Required for elastic EP

260 + where ranks join asynchronously and cannot participate

261 + in collectives at construction time.

262 + """

263 assert all_expert_weights, (

264 "NixlEplbCommunicator requires non-empty all_expert_weights."

265 )

@@ -302,10 +315,29 @@ class NixlEplbCommunicator(EplbCommunicator):

315 ] = {}

316

317 self._cuda_device_id = int(self._device.index or 0)

318 + self._remote_state_initialized = False

319 self._init_step("buffers", self._init_registered_buffers)

320 + if defer_remote_setup:

321 + logger.info_once("NIXL EPLB: deferring remote agent setup (elastic EP).")

322 + else:

323 + self._init_remote_state()

324 + self._log_initialized()

325 +

326 + def _init_remote_state(self) -> None:

327 + """Exchange NIXL agent metadata and RDMA pointer info with all peers.

328 +

329 + This is a collective operation (uses ``all_gather_object`` twice).

330 + Under elastic EP the call is deferred to the first

331 + ``set_transfer_context`` invocation, where all ranks are

332 + guaranteed to be synchronized.

333 + """

334 self._init_step("agents", self._init_remote_agents)

335 self._init_step("send meta", self._exchange_remote_send_meta)

308 - self._log_initialized()

336 + self._remote_state_initialized = True

337 +

338 + def _ensure_remote_state(self) -> None:

339 + if not self._remote_state_initialized:

340 + self._init_remote_state()

341

342 @property

343 def needs_profile_buffer_reservation(self) -> bool:

@@ -339,8 +371,7 @@ class NixlEplbCommunicator(EplbCommunicator):

371 pass

372

373 def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None:

342 - # Pre-compute expert_id -> src_row mapping for every rank so that

343 - # add_recv can immediately issue NIXL READs.

374 + self._ensure_remote_state()

375 assert not self._xfer_entries, (

376 f"set_transfer_context() called with {len(self._xfer_entries)} "

377 f"pending transfers from layer {self._layer_idx}; "

@@ -523,6 +554,21 @@ class NixlEplbCommunicator(EplbCommunicator):

554 )

555 return (local_handle, remote_handle, xfer_handle)

556

557 + def _post_read_barrier(self) -> None:

558 + """Correctness fence: prevents overwrite-while-remote-read race.

559 +

560 + We avoid ``torch.distributed.monitored_barrier`` because it

561 + calls ``get_backend(group)`` which fails for stateless groups

562 + (elastic EP). An async ``all_reduce`` + ``wait(timeout)``

563 + works with both regular and stateless groups and provides

564 + equivalent timeout detection.

565 + """

566 + _dummy = torch.zeros(1, dtype=torch.int32)

567 + work = torch.distributed.all_reduce(

568 + _dummy, group=self._cpu_group, async_op=True

569 + )

570 + work.wait(timeout=timedelta(minutes=5))

571 +

572 def execute(self) -> None:

573 assert self._layer_idx is not None or not self._xfer_entries, (

574 "set_transfer_context() must be called before execute() "

@@ -531,13 +577,7 @@ class NixlEplbCommunicator(EplbCommunicator):

577 try:

578 self._wait_for_all_transfers([x[2] for x in self._xfer_entries])

579

534 - # Post-READ barrier.

535 - # Correctness fence for zero-copy: prevents overwrite-while-

536 - # remote-read race.

537 - torch.distributed.monitored_barrier(

538 - group=self._cpu_group,

539 - timeout=timedelta(minutes=5),

540 - )

580 + self._post_read_barrier()

581 finally:

582 for local_h, remote_h, xfer_h in self._xfer_entries:

583 with contextlib.suppress(Exception):

@@ -628,11 +668,13 @@ def create_eplb_communicator(

668 device and CPU communication groups.

669 backend: Communicator backend name (``"torch_nccl"``,

670 ``"torch_gloo"``, ``"pynccl"``, or ``"nixl"``).

631 - Stateless (elastic EP) groups only support ``"torch_nccl"``

632 - and ``"pynccl"``; ``"torch_nccl"`` is silently promoted to

633 - ``"pynccl"`` in that case. When tensors reside on CPU,

634 - ``"torch_gloo"`` or ``"torch_nccl"`` are used via the CPU

635 - process group.

671 + Falls back to ``"torch_nccl"`` when *None*.

672 + Stateless (elastic EP) groups support ``"torch_nccl"``,

673 + ``"pynccl"``, and ``"nixl"``; ``"torch_nccl"`` is silently

674 + promoted to ``"pynccl"``. ``"nixl"`` uses deferred remote

675 + agent setup to avoid collective deadlocks during elastic

676 + scaling. When tensors reside on CPU, ``"torch_gloo"`` or

677 + ``"torch_nccl"`` are used via the CPU process group.

678 expert_weights: Expert weight tensors for *all* MoE layers.

679 Shape ``(num_layers)(num_tensors_per_layer)``.

680 NixlEplbCommunicator registers all layers with NIXL for

@@ -686,18 +728,21 @@ def create_eplb_communicator(

728

729 is_stateless = isinstance(group_coordinator, StatelessGroupCoordinator)

730 if is_stateless:

689 - if backend not in ("torch_nccl", "pynccl"):

731 + if backend == "nixl":

732 + pass # handled below with defer_remote_setup=True

733 + elif backend not in ("torch_nccl", "pynccl"):

734 raise ValueError(

691 - f"Elastic EP requires 'torch_nccl' or 'pynccl' EPLB communicator "

692 - f"(got '{backend}')."

735 + f"Elastic EP requires 'torch_nccl', 'pynccl', or 'nixl' "

736 + f"EPLB communicator (got '{backend}')."

737 )

694 - if backend == "torch_nccl":

695 - logger.warning(

696 - "Stateless elastic EP requires PyNCCL backend. "

697 - "Forcing EPLB communicator to 'pynccl'."

698 - )

699 - backend = "pynccl"

700 - return _create_pynccl()

738 + else:

739 + if backend == "torch_nccl":

740 + logger.warning(

741 + "Stateless elastic EP requires PyNCCL backend. "

742 + "Forcing EPLB communicator to 'pynccl'."

743 + )

744 + backend = "pynccl"

745 + return _create_pynccl()

746

747 if backend == "nixl":

748 if not has_nixl():

@@ -714,6 +759,7 @@ def create_eplb_communicator(

759 cpu_group=group_coordinator.cpu_group,

760 all_expert_weights=expert_weights,

761 expert_buffer=expert_buffer,

762 + defer_remote_setup=is_stateless,

763 )

764 except Exception as exc:

765 raise RuntimeError(

vllm/distributed/eplb/eplb_state.py

+42 -2

@@ -27,6 +27,7 @@ physical experts.

27 """

29 import threading

30 +import time

31 from collections.abc import Sequence

32 from dataclasses import dataclass

@@ -825,6 +826,45 @@ class EplbState:

826 is_profile=is_profile,

827 )

828

829 + def drain_async(self) -> None:

830 + """Drain in-flight async EPLB by consuming all remaining layer results.

831 +

832 + Each pending result is acknowledged (consumed_event recorded) so the

833 + async worker can proceed, but the transferred weights are intentionally

834 + NOT applied — a full synchronous rearrange is expected to follow.

835 +

836 + Ranks are kept in lockstep via _all_ranks_result_ready (all_reduce

837 + on the EP CPU group). The async worker's coordinated-stop collectives

838 + use the separate EPLB group, so the two sets of collectives do not

839 + interfere.

840 +

841 + No-op when no async cycle is in progress (rebalanced=False).

842 + """

843 + if not self.is_async:

844 + return

845 + for model_key, ms in self.model_states.items():

846 + needs_drain = ms.rebalanced

847 + if needs_drain:

848 + logger.info(

849 + "Draining async EPLB worker for model %s",

850 + model_key,

851 + )

852 + while ms.rebalanced:

853 + if self._all_ranks_result_ready(ms):

854 + result = ms.pending_result

855 + assert result is not None

856 + if result.layer_idx == ms.model.num_moe_layers - 1:

857 + ms.rebalanced = False

858 + ms.pending_result = None

859 + result.consumed_event.record()

860 + else:

861 + time.sleep(0.001)

862 + if needs_drain:

863 + logger.info(

864 + "Async EPLB worker drained for model %s",

865 + model_key,

866 + )

867 +

868 def _all_ranks_result_ready(self, model_state: EplbModelState) -> bool:

869 parallel_state = get_ep_group()

870 has_result = int(model_state.pending_result is not None)

@@ -850,8 +890,9 @@ class EplbState:

890 """

891 All-reduce a list of tensors.

892 """

893 + ep_group = get_ep_group().device_group

894 if len(tensor_list) == 1:

854 - all_reduce(tensor_list[0], group=get_ep_group().device_group)

895 + all_reduce(tensor_list[0], group=ep_group)

896 return tensor_list

897 assert all(t.dim() == 2 for t in tensor_list), "All tensors must be 2D."

898 assert all(t.shape[1] == tensor_list[0].shape[1] for t in tensor_list), (

@@ -863,7 +904,6 @@ class EplbState:

904 shapes = [t.shape for t in tensor_list]

905 concat_tensor = torch.cat(tensor_list, dim=0)

906

866 - ep_group = get_ep_group().device_group

907 all_reduce(concat_tensor, group=ep_group)

908

909 all_reduce_list = []