From 78e6067cf7643b3e07ebcaa7b1e339a6dd5c4da9 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 16 Jul 2026 20:44:32 +0000 Subject: [PATCH] [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 Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: Isotr0py Co-authored-by: Isotr0py Co-authored-by: Jee Jee Li Co-authored-by: Roger Wang Co-authored-by: Yifan Qiao Co-authored-by: Claude Fable 5 Co-authored-by: OpenAI Codex Signed-off-by: Woosuk Kwon --- .../device_communicators/symm_mem.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/vllm/distributed/device_communicators/symm_mem.py b/vllm/distributed/device_communicators/symm_mem.py index 8c174602c3c..5621514020e 100644 --- a/vllm/distributed/device_communicators/symm_mem.py +++ b/vllm/distributed/device_communicators/symm_mem.py @@ -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,