[Distributed] Use NVSHMEM symm-mem backend across nodes

Select NVSHMEM before the first symmetric-memory allocation for process groups spanning multiple nodes, since CUDA IPC handles are node-local. Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-07-16 20:44:32 +00:00
1 changed files with 19 additions and 0 deletions
@@ -87,6 +87,25 @@ class SymmMemCommunicator:
            self.max_size = SYMM_MEM_ALL_REDUCE_MAX_SIZES[self.device_capability][
                self.world_size
            ]
+        from vllm.distributed.parallel_state import in_the_same_node_as
+
+        if not all(in_the_same_node_as(self.group, source_rank=0)):
+            # Cross-node group (multi-node NVLink): the default CUDA backend
+            # exchanges CUDA IPC handles, which are same-node only. Select
+            # the NVSHMEM backend before the process's first symm-mem
+            # allocation so rendezvous maps fabric-reachable peer VAs.
+            try:
+                torch_symm_mem.set_backend("NVSHMEM")
+                logger.info_once(
+                    "SymmMemCommunicator: multi-node group, using the "
+                    "NVSHMEM symmetric-memory backend"
+                )
+            except RuntimeError as e:
+                logger.warning_once(
+                    "SymmMemCommunicator: could not select the NVSHMEM "
+                    "backend for a multi-node group: %s",
+                    str(e),
+                )
        try:
            self.buffer = torch_symm_mem.empty(
                self.max_size // self.dtype.itemsize,