[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>

vllm/utils/network_utils.py

@@ -136,6 +136,17 @@ def get_file_store_init_method() -> str:
     return f"file://{tempfile.gettempdir()}/vllm_dist_{uuid4().hex}"
 
 
+def aiter_requires_tcp_store() -> bool:
+    """AITER custom all-reduce requires a pure-TCP default store (its IPC
+    metadata exchange asserts on ``TCPStore``); the file:// rendezvous yields a
+    ``FileStore`` and trips that assertion. Prefer the TCP rendezvous
+    (pre-#50999) for ROCm + AITER custom AR until AITER accepts FileStore.
+    """
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    return rocm_aiter_ops.is_custom_all_reduce_enabled()
+
+
 def get_tcp_uri(ip: str, port: int) -> str:
     if is_valid_ipv6_address(ip):
         return f"tcp://[{ip}]:{port}"

vllm/v1/executor/multiproc_executor.py

@@ -46,8 +46,12 @@ from vllm.platforms import current_platform
 from vllm.tracing import instrument, maybe_init_worker_tracer
 from vllm.utils import numa_utils
 from vllm.utils.network_utils import (
+    aiter_requires_tcp_store,
+    get_distributed_init_method,
     get_file_store_init_method,
     get_ip,
+    get_loopback_ip,
+    get_open_port,
 )
 from vllm.utils.ompmultiprocessing import OMPProcessManager
 from vllm.utils.system_utils import (
@@ -127,7 +131,12 @@ class MultiprocExecutor(Executor):
 
         set_multiprocessing_worker_envs(self.local_world_size)
 
-        distributed_init_method = get_file_store_init_method()
+        if aiter_requires_tcp_store():
+            distributed_init_method = get_distributed_init_method(
+                get_loopback_ip(), get_open_port()
+            )
+        else:
+            distributed_init_method = get_file_store_init_method()
         self.rpc_broadcast_mq: MessageQueue | None = None
         scheduler_output_handle: Handle | None = None
         # Initialize worker and set up message queues for SchedulerOutputs

vllm/v1/executor/uniproc_executor.py

@@ -12,7 +12,13 @@ import torch.distributed as dist
 import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
-from vllm.utils.network_utils import get_file_store_init_method
+from vllm.utils.network_utils import (
+    aiter_requires_tcp_store,
+    get_distributed_init_method,
+    get_file_store_init_method,
+    get_ip,
+    get_open_port,
+)
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.executor.abstract import Executor
 from vllm.v1.executor.vllm_net_devices import set_worker_net_device
@@ -70,7 +76,12 @@ class UniProcExecutor(Executor):
 
     def _distributed_args(self) -> tuple[str, int, int]:
         """Return (distributed_init_method, rank, local_rank)."""
-        distributed_init_method = get_file_store_init_method()
+        if aiter_requires_tcp_store():
+            distributed_init_method = get_distributed_init_method(
+                get_ip(), get_open_port()
+            )
+        else:
+            distributed_init_method = get_file_store_init_method()
         # set local rank as the device index if specified
         device_info = self.vllm_config.device_config.device.__str__().split(":")
         local_rank = int(device_info[1]) if len(device_info) > 1 else 0