[ROCm][Bugfix] Use TCP store when AITER custom all-reduce is enabled (#51635)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>
vllmellm committed
Aug 10, 2026 at 19:07 UTC
436be94e135fd07fc94e4c165a1b6870af72bf31
3 files changed
+34
-3
vllm/utils/network_utils.py
+11
@@ -136,6 +136,17 @@ def get_file_store_init_method() -> str:
136
return f"file://{tempfile.gettempdir()}/vllm_dist_{uuid4().hex}"
137
138
139
+def aiter_requires_tcp_store() -> bool:
140
+ """AITER custom all-reduce requires a pure-TCP default store (its IPC
141
+ metadata exchange asserts on ``TCPStore``); the file:// rendezvous yields a
142
+ ``FileStore`` and trips that assertion. Prefer the TCP rendezvous
143
+ (pre-#50999) for ROCm + AITER custom AR until AITER accepts FileStore.
144
+ """
145
+ from vllm._aiter_ops import rocm_aiter_ops
146
+
147
+ return rocm_aiter_ops.is_custom_all_reduce_enabled()
148
+
149
+
150
def get_tcp_uri(ip: str, port: int) -> str:
151
if is_valid_ipv6_address(ip):
152
return f"tcp://[{ip}]:{port}"
vllm/v1/executor/multiproc_executor.py
+10
-1
@@ -46,8 +46,12 @@ from vllm.platforms import current_platform
46
from vllm.tracing import instrument, maybe_init_worker_tracer
47
from vllm.utils import numa_utils
48
from vllm.utils.network_utils import (
49
+ aiter_requires_tcp_store,
50
+ get_distributed_init_method,
51
get_file_store_init_method,
52
get_ip,
53
+ get_loopback_ip,
54
+ get_open_port,
55
)
56
from vllm.utils.ompmultiprocessing import OMPProcessManager
57
from vllm.utils.system_utils import (
@@ -127,7 +131,12 @@ class MultiprocExecutor(Executor):
131
132
set_multiprocessing_worker_envs(self.local_world_size)
133
130
- distributed_init_method = get_file_store_init_method()
134
+ if aiter_requires_tcp_store():
135
+ distributed_init_method = get_distributed_init_method(
136
+ get_loopback_ip(), get_open_port()
137
+ )
138
+ else:
139
+ distributed_init_method = get_file_store_init_method()
140
self.rpc_broadcast_mq: MessageQueue | None = None
141
scheduler_output_handle: Handle | None = None
142
# Initialize worker and set up message queues for SchedulerOutputs
vllm/v1/executor/uniproc_executor.py
+13
-2
@@ -12,7 +12,13 @@ import torch.distributed as dist
12
import vllm.envs as envs
13
from vllm.logger import init_logger
14
from vllm.platforms import current_platform
15
-from vllm.utils.network_utils import get_file_store_init_method
15
+from vllm.utils.network_utils import (
16
+ aiter_requires_tcp_store,
17
+ get_distributed_init_method,
18
+ get_file_store_init_method,
19
+ get_ip,
20
+ get_open_port,
21
+)
22
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
23
from vllm.v1.executor.abstract import Executor
24
from vllm.v1.executor.vllm_net_devices import set_worker_net_device
@@ -70,7 +76,12 @@ class UniProcExecutor(Executor):
76
77
def _distributed_args(self) -> tuple[str, int, int]:
78
"""Return (distributed_init_method, rank, local_rank)."""
73
- distributed_init_method = get_file_store_init_method()
79
+ if aiter_requires_tcp_store():
80
+ distributed_init_method = get_distributed_init_method(
81
+ get_ip(), get_open_port()
82
+ )
83
+ else:
84
+ distributed_init_method = get_file_store_init_method()
85
# set local rank as the device index if specified
86
device_info = self.vllm_config.device_config.device.__str__().split(":")
87
local_rank = int(device_info[1]) if len(device_info) > 1 else 0