[ROCm][Bugfix] Fix elastic EP scaling deadlock (#56610)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Matt committed
Sep 12, 2026 at 16:09 UTC
72d4d8315749c146cf19f047fecfec491911770d
4 files changed
+106
-38
vllm/distributed/device_communicators/pynccl.py
+31
@@ -3,6 +3,7 @@
3
4
5
# ===================== import region =====================
6
+import contextlib
7
import threading
8
9
import torch
@@ -26,6 +27,33 @@ from vllm.utils.torch_utils import current_stream
27
28
logger = init_logger(__name__)
29
30
+_warmup_tls = threading.local()
31
+
32
+
33
+def _warmup_deferred() -> bool:
34
+ return getattr(_warmup_tls, "deferred", False)
35
+
36
+
37
+@contextlib.contextmanager
38
+def defer_comm_warmup_on_rocm():
39
+ """Build communicators on this thread without their warm-up on ROCm, where
40
+ the warm-up can deadlock if built next to a live engine whose existing
41
+ communicators still have collectives in flight; instead the warm-up is
42
+ deferred to when the engine is quiet (during commit).
43
+ """
44
+ from vllm.platforms import current_platform
45
+
46
+ if not current_platform.is_rocm():
47
+ yield
48
+ return
49
+ prev = _warmup_deferred()
50
+ _warmup_tls.deferred = True
51
+ try:
52
+ yield
53
+ finally:
54
+ _warmup_tls.deferred = prev
55
+
56
+
57
_NCCL_SYMM_OPS_REGISTERED = False
58
59
_NCCL_SUSPEND_MEM = 0x01
@@ -153,6 +181,9 @@ class PyNcclCommunicator:
181
self.world_size, self.unique_id, self.rank
182
)
183
184
+ if _warmup_deferred():
185
+ return
186
+
187
stream = current_stream()
188
# A small all_reduce for warmup.
189
data = torch.zeros(1, device=device)
vllm/distributed/elastic_ep/elastic_execute.py
+13
-2
@@ -50,6 +50,7 @@ from vllm.model_executor.warmup.flashinfer_autotune_cache import (
50
sync_flashinfer_autotune_cache,
51
)
52
from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
53
+from vllm.platforms import current_platform
54
from vllm.utils import is_moe_layer
55
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
56
from vllm.v1.worker.dp_utils import skip_dp_coordination
@@ -283,7 +284,9 @@ class ElasticEPScalingExecutor:
284
self._prepare_eplb_communicator(get_standby_eplb_group())
285
if new_dp_size > old_dp_size:
286
self.transfer_weights(old_dp_size, new_dp_size)
286
- self._warm_target_groups(get_standby_dp_group(), standby_ep_group)
287
+ # On ROCm this deadlocks while serving; warm_and_capture warms these at commit.
288
+ if not current_platform.is_rocm():
289
+ self._warm_target_groups(get_standby_dp_group(), standby_ep_group)
290
if new_dp_size > old_dp_size and self._can_reuse_fused_moe_kernel():
291
target_world_group = get_standby_world_group()
292
assert target_world_group is not None
@@ -577,7 +580,11 @@ class ElasticEPScalingExecutor:
580
mapping = self.receive_expert_mapping()
581
self.worker.model_runner.setup_eplb_from_mapping(mapping)
582
if not self._can_reuse_fused_moe_kernel():
583
+ # warm_and_capture's dummy run warms the deferred dp/ep groups.
584
self.warm_and_capture()
585
+ elif current_platform.is_rocm():
586
+ # Reuse path skips warm_and_capture; warm the deferred groups here.
587
+ self._warm_target_groups(get_dp_group(), get_ep_group())
588
self._perform_eplb_reshuffle(async_op=True)
589
if is_existing_worker:
590
self._start_group_cleanup(retired_groups)
@@ -590,6 +597,9 @@ class ElasticEPScalingExecutor:
597
retired_groups = self.switch_and_prepare()
598
if not self._can_reuse_fused_moe_kernel():
599
self.warm_and_capture()
600
+ elif current_platform.is_rocm():
601
+ # Reuse path skips warm_and_capture; warm the deferred groups here.
602
+ self._warm_target_groups(get_dp_group(), get_ep_group())
603
self._start_group_cleanup(retired_groups)
604
605
def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None:
@@ -646,7 +656,8 @@ class ElasticEPScalingExecutor:
656
expert_weights=expert_weights,
657
)
658
torch.accelerator.synchronize()
649
- self._warm_target_groups(get_dp_group(), get_ep_group())
659
+ if not current_platform.is_rocm():
660
+ self._warm_target_groups(get_dp_group(), get_ep_group())
661
if self._can_reuse_fused_moe_kernel():
662
sync_flashinfer_autotune_cache(self.worker.model_runner, get_world_group())
663
vllm/distributed/elastic_ep/standby_state.py
+23
-14
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import torch
4
5
+from vllm.distributed.device_communicators.pynccl import defer_comm_warmup_on_rocm
6
from vllm.distributed.parallel_state import (
7
_init_stateless_group,
8
_node_count,
@@ -78,27 +79,35 @@ def create_standby_groups(
79
)
80
standby_dp_ranks = all_ranks.transpose(1, 3).reshape(-1, new_dp_size).unbind(0)
81
standby_dp_ranks = [x.tolist() for x in standby_dp_ranks]
81
- _STANDBY_DP = _init_stateless_group(
82
- standby_dp_ranks, "dp", master_ip, backend, coord_store=coord_store
83
- )
82
85
- standby_ep_ranks = (
86
- all_ranks.transpose(1, 2).reshape(-1, new_dp_size * tp_size).unbind(0)
87
- )
88
- standby_ep_ranks = [x.tolist() for x in standby_ep_ranks]
89
- _STANDBY_EP = _init_stateless_group(
90
- standby_ep_ranks, "ep", master_ip, backend, coord_store, use_all2all=use_all2all
91
- )
83
+ # Deferred to commit so the warm-up runs while the engine is paused.
84
+ with defer_comm_warmup_on_rocm():
85
+ _STANDBY_DP = _init_stateless_group(
86
+ standby_dp_ranks, "dp", master_ip, backend, coord_store=coord_store
87
+ )
88
93
- if enable_eplb:
94
- _STANDBY_EPLB = _init_stateless_group(
89
+ standby_ep_ranks = (
90
+ all_ranks.transpose(1, 2).reshape(-1, new_dp_size * tp_size).unbind(0)
91
+ )
92
+ standby_ep_ranks = [x.tolist() for x in standby_ep_ranks]
93
+ _STANDBY_EP = _init_stateless_group(
94
standby_ep_ranks,
96
- "eplb",
95
+ "ep",
96
master_ip,
97
backend,
99
- coord_store=coord_store,
98
+ coord_store,
99
+ use_all2all=use_all2all,
100
)
101
102
+ if enable_eplb:
103
+ _STANDBY_EPLB = _init_stateless_group(
104
+ standby_ep_ranks,
105
+ "eplb",
106
+ master_ip,
107
+ backend,
108
+ coord_store=coord_store,
109
+ )
110
+
111
112
def pop_standby_groups() -> dict:
113
"""Return all standby groups and clear the standby state."""
vllm/distributed/parallel_state.py
+39
-22
@@ -1919,6 +1919,20 @@ def init_distributed_environment(
1919
_INNER_DP_WORLD = _WORLD
1920
1921
1922
+def _elastic_join_warmup_ctx() -> AbstractContextManager[Any]:
1923
+ """Defer communicator warm-up when joining a live elastic-EP cluster on ROCm.
1924
+
1925
+ The existing ranks defer their matching warm-up in create_standby_groups(),
1926
+ so a joining worker must too or the warm-up collective has no peers. Both
1927
+ sides warm up at commit.
1928
+ """
1929
+ from vllm.distributed.device_communicators.pynccl import defer_comm_warmup_on_rocm
1930
+
1931
+ if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
1932
+ return defer_comm_warmup_on_rocm()
1933
+ return nullcontext()
1934
+
1935
+
1936
def initialize_model_parallel(
1937
tensor_model_parallel_size: int = 1,
1938
pipeline_model_parallel_size: int = 1,
@@ -2101,13 +2115,14 @@ def initialize_model_parallel(
2115
group_ranks = all_ranks.transpose(1, 4).reshape(-1, data_parallel_size).unbind(0)
2116
group_ranks = [x.tolist() for x in group_ranks]
2117
if enable_elastic_ep:
2104
- _DP = _init_stateless_group(
2105
- group_ranks,
2106
- "dp",
2107
- parallel_config.data_parallel_master_ip,
2108
- backend,
2109
- coord_store=coord_store,
2110
- )
2118
+ with _elastic_join_warmup_ctx():
2119
+ _DP = _init_stateless_group(
2120
+ group_ranks,
2121
+ "dp",
2122
+ parallel_config.data_parallel_master_ip,
2123
+ backend,
2124
+ coord_store=coord_store,
2125
+ )
2126
else:
2127
_DP = init_model_parallel_group(
2128
group_ranks, get_world_group().local_rank, backend, group_name="dp"
@@ -2130,14 +2145,15 @@ def initialize_model_parallel(
2145
group_ranks = [x.tolist() for x in group_ranks]
2146
use_all2all = parallel_config.use_all2all
2147
if enable_elastic_ep:
2133
- _EP = _init_stateless_group(
2134
- group_ranks,
2135
- "ep",
2136
- parallel_config.data_parallel_master_ip,
2137
- backend,
2138
- coord_store=coord_store,
2139
- use_all2all=use_all2all,
2140
- )
2148
+ with _elastic_join_warmup_ctx():
2149
+ _EP = _init_stateless_group(
2150
+ group_ranks,
2151
+ "ep",
2152
+ parallel_config.data_parallel_master_ip,
2153
+ backend,
2154
+ coord_store=coord_store,
2155
+ use_all2all=use_all2all,
2156
+ )
2157
else:
2158
_EP = init_model_parallel_group(
2159
group_ranks,
@@ -2155,13 +2171,14 @@ def initialize_model_parallel(
2171
assert _EPLB is None, "EPLB group is already initialized"
2172
if config.parallel_config.enable_eplb:
2173
if enable_elastic_ep:
2158
- _EPLB = _init_stateless_group(
2159
- group_ranks,
2160
- "eplb",
2161
- parallel_config.data_parallel_master_ip,
2162
- backend,
2163
- coord_store=coord_store,
2164
- )
2174
+ with _elastic_join_warmup_ctx():
2175
+ _EPLB = _init_stateless_group(
2176
+ group_ranks,
2177
+ "eplb",
2178
+ parallel_config.data_parallel_master_ip,
2179
+ backend,
2180
+ coord_store=coord_store,
2181
+ )
2182
else:
2183
_EPLB = init_model_parallel_group(
2184
group_ranks,