[Distributed] Use NVSHMEM symm-mem backend across nodes

[WoosukKwon]

Summary

• Select the PyTorch NVSHMEM symmetric-memory backend before the first allocation when a process group spans multiple nodes.
• Keep the current backend for same-node groups.
• Preserve the existing communicator-disable fallback if backend selection or allocation is unavailable.

The default CUDA backend exchanges node-local CUDA IPC handles. Multi-node NVLink process groups need NVSHMEM-backed symmetric virtual addresses.

Duplicate-work check

I searched open PRs for symm_mem NVSHMEM backend multi-node and found no open PR implementing this backend selection.

Validation

• .venv/bin/python -m pytest tests/distributed/test_symm_mem_allreduce.py::test_symm_mem_allreduce -v: 1 passed
• .venv/bin/pre-commit run --files vllm/distributed/device_communicators/symm_mem.py: passed

Model evaluation

Not applicable; this changes distributed symmetric-memory backend selection and does not change model computations or output.

AI assistance

AI assistance was used to extract and validate this change. The human submitter will review every changed line and is responsible for understanding and defending the change end-to-end.

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[ErenAta16]

The node-spanning detection reads right, and not all(in_the_same_node_as(self.group, source_rank=0)) is the correct predicate for "this group crosses a node boundary".

My concern is that set_backend is process-global while this call site is per-group. PyTorch documents both properties:

This is a global setting and affects all subsequent calls to torch._distributed._symmetric_memory.empty(). Note that the backend cannot be changed once a symmetric memory tensor has been allocated.

SymmMemCommunicator is constructed from cuda_communicator.py, and a single vLLM process builds one per parallel group, so TP, PP, DP and EP groups can each own an instance. That makes two orderings reachable, and neither matches the summary's "keep the current backend for same-node groups":

Cross-node group constructed first. It calls set_backend("NVSHMEM"), which is global, so every same-node group created afterwards also allocates through NVSHMEM. Nothing resets it, and those groups never execute the branch that would have opted them in. The same-node path is not preserved, it just is not the one that decided.

A same-node group allocates first. The backend is now frozen, so the later cross-node set_backend("NVSHMEM") raises. The except RuntimeError catches it and warns, and the cross-node group then proceeds on the CUDA IPC backend, which is exactly the configuration this PR exists to avoid. The failure surfaces later at allocation or rendezvous rather than at the point the decision was made, and the warning text ("could not select the NVSHMEM backend") reads like a capability problem rather than an ordering one.

The second is the one I would want addressed before merge, because the fallback is not neutral: continuing with a same-node backend for a multi-node group is a known-broken state, not a degraded one. Disabling the communicator there, which the description says is the intended behaviour when selection is unavailable, seems more consistent than warning and carrying on.

More broadly, a global switch driven from a per-group constructor means the outcome depends on construction order across unrelated groups. Deciding once during distributed init, based on whether any group in the process spans nodes, would remove the ordering dependency and would also make the "same-node groups keep the current backend" claim expressible: today it can only hold when there is exactly one group.

I do not have a multi-node setup to confirm the second ordering empirically; the above is from the PyTorch semantics and the construction sites rather than from a run. Happy to be told the group creation order makes one of these unreachable in practice.