[Core] Release NCCL communicator memory in sleep mode

vllm/distributed/device_communicators/pynccl.py

@@ -27,6 +27,11 @@
 
 _NCCL_SYMM_OPS_REGISTERED = False
 
+# ncclCommSuspend flag: release dynamic GPU memory, keep topology/connection.
+_NCCL_SUSPEND_MEM = 0x01
+# ncclCommSuspend / ncclCommResume were introduced in NCCL 2.29.7 (22907).
+_NCCL_SUSPEND_MIN_VERSION = 22907
+
 
 def register_nccl_symmetric_ops(pynccl_comm):
     from vllm.distributed.device_communicators.pynccl_allocator import (
@@ -419,6 +424,22 @@ def register_comm_window_raw(self, ptr: int, size: int):
     def deregister_comm_window(self, window):
         return self.nccl.ncclCommWindowDeregister(self.comm, window)
 
+    def suspend(self):
+        """Release dynamic GPU memory, keeping topology/connection state.
+
+        Collective across the group's ranks (ncclCommSuspend has an internal
+        cross-rank barrier). No-op on NCCL < 2.29.7.
+        """
+        if self.disabled or self.nccl_version < _NCCL_SUSPEND_MIN_VERSION:
+            return
+        self.nccl.ncclCommSuspend(self.comm, _NCCL_SUSPEND_MEM)
+
+    def resume(self):
+        """Restore a suspended communicator (collective). No-op on NCCL < 2.29.7."""
+        if self.disabled or self.nccl_version < _NCCL_SUSPEND_MIN_VERSION:
+            return
+        self.nccl.ncclCommResume(self.comm)
+
     def batch_isend_irecv(self, p2p_ops: list, stream=None):
         if self.disabled:
             return

vllm/distributed/device_communicators/pynccl_wrapper.py

@@ -296,6 +296,13 @@ class NCCLLibrary:
         # shutdown when peer ranks may already be gone.
         # ncclResult_t  ncclCommAbort(ncclComm_t comm);
         Function("ncclCommAbort", ncclResult_t, [ncclComm_t]),
+        # Release a communicator's dynamic GPU memory while keeping its
+        # topology/connection state, so resume is cheap. NCCL >= 2.29.7 only;
+        # bound optionally (see __init__) and absent on older NCCL.
+        # ncclResult_t ncclCommSuspend(ncclComm_t comm, int flags);
+        Function("ncclCommSuspend", ncclResult_t, [ncclComm_t, ctypes.c_int]),
+        # ncclResult_t ncclCommResume(ncclComm_t comm);
+        Function("ncclCommResume", ncclResult_t, [ncclComm_t]),
         # ncclResult_t ncclGroupStart();
         Function("ncclGroupStart", ncclResult_t, []),
         # ncclResult_t ncclGroupEnd();
@@ -358,6 +365,16 @@ def __init__(self, so_file: str | None = None):
                     f.argtypes = func.argtypes
                     _funcs[func.name] = f
                 except AttributeError:
+                    if func.name in ("ncclCommSuspend", "ncclCommResume"):
+                        # Only present in NCCL >= 2.29.7. Skip on older NCCL;
+                        # sleep-mode NCCL memory release degrades to a no-op.
+                        logger.warning_once(
+                            "NCCL communicator suspend/resume not found in %s "
+                            "(needs NCCL >= 2.29.7); sleep-mode NCCL memory "
+                            "release is disabled.",
+                            so_file,
+                        )
+                        continue
                     if func.name in [
                         "ncclCommWindowRegister",
                         "ncclCommWindowDeregister",
@@ -557,6 +574,12 @@ def ncclCommDestroy(self, comm: ncclComm_t) -> None:
     def ncclCommAbort(self, comm: ncclComm_t) -> None:
         self.NCCL_CHECK(self._funcs["ncclCommAbort"](comm))
 
+    def ncclCommSuspend(self, comm: ncclComm_t, flags: int) -> None:
+        self.NCCL_CHECK(self._funcs["ncclCommSuspend"](comm, flags))
+
+    def ncclCommResume(self, comm: ncclComm_t) -> None:
+        self.NCCL_CHECK(self._funcs["ncclCommResume"](comm))
+
     def ncclGroupStart(self) -> None:
         self.NCCL_CHECK(self._funcs["ncclGroupStart"]())
 

vllm/distributed/parallel_state.py

@@ -127,6 +127,56 @@ def _register_group(group: "GroupCoordinator") -> None:
     _groups[group.unique_name] = weakref.ref(group)
 
 
+def _apply_to_nccl_comms(label: str, action: Callable[[Any], None]) -> None:
+    """Apply ``action`` to every group's NCCL communicator, collectively.
+
+    Walks the registered parallel groups and skips those without a CUDA NCCL
+    communicator (``pynccl_comm`` is CUDA-only; also absent at ``world_size
+    == 1`` or when disabled). ncclCommSuspend/Resume are collective (internal
+    cross-rank barrier) and synchronous on return, so no extra sync is needed.
+    """
+    comms = []
+    for group_ref in _groups.values():
+        group = group_ref()
+        if group is None:
+            continue
+        # pynccl_comm exists only on CudaCommunicator (not CPU/XPU/Ray, nor on
+        # groups with no device communicator such as world_size == 1).
+        pynccl = getattr(group.device_communicator, "pynccl_comm", None)
+        if pynccl is None or pynccl.disabled:
+            continue
+        comms.append(pynccl)
+    if not comms:
+        return
+
+    free_before = torch.cuda.mem_get_info()[0]
+    for pynccl in comms:
+        action(pynccl)
+    delta = torch.cuda.mem_get_info()[0] - free_before
+    direction = "freed" if delta > 0 else "allocated"
+    logger.info(
+        "NCCL %s: %d comms, %.1f MiB %s",
+        label,
+        len(comms),
+        abs(delta) / 1024**2,
+        direction,
+    )
+
+
+def suspend_nccl_comms() -> None:
+    """Release idle NCCL communicator memory on every group (collective).
+
+    Must run on every rank with the communicators idle; ncclCommSuspend has an
+    internal cross-rank barrier. No-op on NCCL < 2.29.7 or at world_size 1.
+    """
+    _apply_to_nccl_comms("suspend", lambda c: c.suspend())
+
+
+def resume_nccl_comms() -> None:
+    """Restore all suspended NCCL communicators before reuse (collective)."""
+    _apply_to_nccl_comms("resume", lambda c: c.resume())
+
+
 def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
     assert group_name in _groups, f"Group {group_name} is not found."
     group = _groups[group_name]()

vllm/v1/worker/gpu_worker.py

@@ -43,6 +43,8 @@
     Handle,
     get_pp_group,
     get_tp_group,
+    resume_nccl_comms,
+    suspend_nccl_comms,
 )
 from vllm.distributed.weight_transfer import (
     WeightTransferEngine,
@@ -184,6 +186,11 @@ def sleep(self, level: int = 1) -> None:
         allocator = get_mem_allocator_instance()
         allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
 
+        # Release idle NCCL communicator memory (NCCL >= 2.29.7). Collective
+        # across ranks; no-op at world_size 1 or on older NCCL. Done after the
+        # cumem unmap so the reported freed bytes include the NCCL release.
+        suspend_nccl_comms()
+
         torch.accelerator.synchronize()
         deadline = time.monotonic() + (5.0 if current_platform.is_rocm() else 0)
         while True:
@@ -205,6 +212,10 @@ def wake_up(self, tags: list[str] | None = None) -> None:
         allocator = get_mem_allocator_instance()
         allocator.wake_up(tags)
 
+        # Restore NCCL communicator memory before any collective runs again.
+        # Collective across ranks; no-op at world_size 1 or on older NCCL.
+        resume_nccl_comms()
+
         # Restore the buffers after level 2 sleep
         if len(self._sleep_saved_buffers):
             model = self.model_runner.model