A high-throughput and memory-efficient inference and serving engine for LLMs
[EPLB] Reject NCCL-based EPLB communicators with async EPLB (#44978)
Signed-off-by: Markov Ilya <markovilya197@gmail.com>
Ilya Markov committed
Jun 10, 2026 at 21:51 UTC
6471ec75bdcbe54200d4a35d1ca284b1baac6336
8 files changed
+37
-32
tests/distributed/test_elastic_ep.py
+4
@@ -78,6 +78,8 @@ def test_elastic_ep_scaling():
78
"--enable-eplb",
79
"--eplb-config.num_redundant_experts",
80
"0",
81
+ "--eplb-config.use_async",
82
+ "false",
83
"--data-parallel-backend",
84
"ray",
85
"--data-parallel-size",
@@ -151,6 +153,8 @@ def test_elastic_ep_scaling_uneven():
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",
tests/distributed/test_eplb_execute.py
+1
-3
@@ -644,9 +644,7 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:
644
(2, 2, 2, 3),
645
],
646
)
647
-@pytest.mark.parametrize(
648
- "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"]
649
-)
647
+@pytest.mark.parametrize("eplb_communicator", ["torch_gloo", "nixl"])
648
def test_async_transfer_layer_without_mtp(
649
world_size: int,
650
num_layers: int,
tests/kernels/moe/test_moe_layer.py
+3
@@ -1285,6 +1285,9 @@ def _test_body_eplb(
1285
expert_weights = [list(eplb_moe_layer.get_expert_weights())]
1286
1287
expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
1288
+ assert vllm_config.parallel_config.eplb_config.communicator is not None, (
1289
+ "EPLB communicator backend must be set by ParallelConfig"
1290
+ )
1291
communicator = create_eplb_communicator(
1292
group_coordinator=get_eplb_group(),
1293
backend=vllm_config.parallel_config.eplb_config.communicator,
vllm/config/parallel.py
+15
-1
@@ -94,13 +94,20 @@ class EPLBConfig:
94
- "torch_gloo": Use torch.distributed gloo with CPU staging
95
- "nixl": Use NIXL/ RIXL with staged send/recv buffers
96
- "pynccl": Use PyNccl send/recv
97
- - None: Auto-select backend ("torch_gloo" for async, "torch_nccl" for sync)
97
+ - None: Auto-select backend (prefers "nixl", falls back to "torch_gloo")
98
"""
99
100
@model_validator(mode="after")
101
def _validate_eplb_config(self) -> Self:
102
if self.use_async and self.policy != "default":
103
raise ValueError("Async EPLB is only supported with the default policy.")
104
+ if self.use_async and self.communicator in ("torch_nccl", "pynccl"):
105
+ raise ValueError(
106
+ f"{self.communicator} communicator is incompatible with "
107
+ "async EPLB due to NCCL multi-stream conflicts. Use "
108
+ "'torch_gloo' or 'nixl' instead, or leave communicator "
109
+ "unset for automatic selection."
110
+ )
111
if self.log_balancedness and self.log_balancedness_interval <= 0:
112
raise ValueError("log_balancedness_interval must be greater than 0.")
113
return self
@@ -787,6 +794,13 @@ class ParallelConfig:
794
if self.enable_elastic_ep:
795
if not self.enable_eplb:
796
raise ValueError("Elastic EP is only supported with enable_eplb=True.")
797
+ if self.eplb_config.use_async:
798
+ raise ValueError(
799
+ "Elastic EP requires the pynccl communicator, which is "
800
+ "incompatible with async EPLB due to NCCL multi-stream "
801
+ "conflicts. Disable async EPLB (eplb_config.use_async=False) "
802
+ "to use elastic EP."
803
+ )
804
if self.pipeline_parallel_size > 1:
805
raise ValueError(
806
"Elastic EP is not supported with pipeline parallelism "
vllm/distributed/elastic_ep/elastic_execute.py
+3
@@ -470,6 +470,9 @@ class ElasticEPScalingExecutor:
470
eplb_model_state.expert_buffer = [
471
torch.empty_like(w) for w in model.expert_weights[0]
472
]
473
+ assert parallel_config.eplb_config.communicator is not None, (
474
+ "EPLB communicator backend must be set by ParallelConfig"
475
+ )
476
eplb_model_state.communicator = create_eplb_communicator(
477
group_coordinator=get_eplb_group(),
478
backend=parallel_config.eplb_config.communicator,
vllm/distributed/eplb/eplb_communicator.py
+1
-5
@@ -617,7 +617,7 @@ class PyNcclEplbCommunicator(EplbCommunicator):
617
618
def create_eplb_communicator(
619
group_coordinator: GroupCoordinator,
620
- backend: str | None,
620
+ backend: str,
621
expert_weights: Sequence[Sequence[torch.Tensor]],
622
expert_buffer: Sequence[torch.Tensor],
623
) -> EplbCommunicator:
@@ -628,7 +628,6 @@ def create_eplb_communicator(
628
device and CPU communication groups.
629
backend: Communicator backend name (``"torch_nccl"``,
630
``"torch_gloo"``, ``"pynccl"``, or ``"nixl"``).
631
- Falls back to ``"torch_nccl"`` when *None*.
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,
@@ -641,9 +640,6 @@ def create_eplb_communicator(
640
expert_buffer: Pre-allocated receive buffer tensors (one per
641
weight tensor in a single layer).
642
"""
644
- if backend is None:
645
- backend = "torch_nccl"
646
-
643
first_layer = expert_weights[0] if expert_weights else []
644
tensor_device_type = first_layer[0].device.type if first_layer else "cpu"
645
torch_group = (
vllm/distributed/eplb/eplb_state.py
+3
@@ -447,6 +447,9 @@ class EplbState:
447
self._init_should_record_tensor(model)
448
expert_buffer = [torch.empty_like(w) for w in model.expert_weights[0]]
449
450
+ assert self.parallel_config.eplb_config.communicator is not None, (
451
+ "EPLB communicator backend must be set by ParallelConfig"
452
+ )
453
communicator = create_eplb_communicator(
454
group_coordinator=get_eplb_group(),
455
backend=self.parallel_config.eplb_config.communicator,
vllm/distributed/eplb/eplb_utils.py
+7
-23
@@ -74,8 +74,6 @@ def override_envs_for_eplb(
74
"""
75
is_data_parallel = parallel_config.data_parallel_size > 1
76
is_eplb_enabled = parallel_config.enable_eplb
77
- async_eplb = parallel_config.eplb_config.use_async
78
- is_deepep_ll = parallel_config.all2all_backend == "deepep_low_latency"
77
is_mega_moe = moe_backend == "deep_gemm_mega_moe"
78
is_nccl_based_eplb_communicator = parallel_config.eplb_config.communicator in (
79
"torch_nccl",
@@ -85,29 +83,16 @@ def override_envs_for_eplb(
83
# Override NCCL_MAX_CTAS to avoid hangs when EPLB's NCCL weight exchange
84
# contends with MoE backend's cooperative-launch on GPU SMs.
85
#
88
- # DeepEP low-latency:
89
- # The hang happens when two ranks interleave kernel launches differently
90
- # between NCCL collectives (used by async EPLB weight exchange) and DeepEP
91
- # low-latency (LL) kernels. DeepEP LL uses a cooperative launch and tries
92
- # to reserve a large fraction of the GPU's SMs; if those SMs are currently
93
- # occupied by NCCL, the DeepEP LL launch blocks until enough SMs are
94
- # freed.
95
- #
96
- # If rank A enters DeepEP LL in main thread while rank B is still executing
97
- # NCCL in async thread, rank A can block waiting for SMs, while rank B can
98
- # block inside NCCL waiting for rank A to participate in the collective.
99
- # This circular wait causes a deadlock.
100
- # Limiting NCCL occupancy via NCCL_MAX_CTAS leaves space for the DeepEP
101
- # cooperative kernel to launch and complete, breaking the deadlock.
102
- # See: https://github.com/deepseek-ai/DeepEP/issues/496
103
- #
104
- # DeepGEMM Mega MoE also uses cooperative launch and will cause hang even
105
- # with sync EPLB.
86
+ # DeepGEMM Mega MoE uses cooperative launch, which tries to reserve a
87
+ # large fraction of the GPU's SMs. If those SMs are occupied by NCCL,
88
+ # the cooperative launch blocks until enough SMs are freed, causing a
89
+ # deadlock. Limiting NCCL occupancy via NCCL_MAX_CTAS leaves space for
90
+ # the cooperative kernel to launch and complete.
91
if (
92
is_data_parallel
93
and is_eplb_enabled
94
and is_nccl_based_eplb_communicator
110
- and ((is_deepep_ll and async_eplb) or is_mega_moe)
95
+ and is_mega_moe
96
):
97
current_value_str = os.getenv("NCCL_MAX_CTAS")
98
@@ -116,10 +101,9 @@ def override_envs_for_eplb(
101
102
override_value = 8
103
os.environ["NCCL_MAX_CTAS"] = str(override_value)
119
- backend = "deepep_low_latency" if is_deepep_ll else "deep_gemm_mega_moe"
104
logger.info_once(
105
f"EPLB: Setting NCCL_MAX_CTAS={override_value} "
106
f"for expert parallel with NCCL-based EPLB communicator and "
123
- f"cooperative MoE backend ({backend})",
107
+ f"cooperative MoE backend (deep_gemm_mega_moe)",
108
scope="global",
109
)