From e9dd1db62b16d9ce5c55a925291fdb2420b0d3e4 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 30 Jun 2026 13:42:54 +0000 Subject: [PATCH 1/5] [Core] Release NCCL communicator memory in sleep mode Sleep mode offloads weights and discards the KV cache, but each PyNccl communicator keeps its dynamic GPU buffers resident the whole time (~958 MiB/GPU on an 8-GPU TP+EP worker). Release them with NCCL's native ncclCommSuspend(NCCL_SUSPEND_MEM) / ncclCommResume, which frees a communicator's allocations while keeping its topology/connection state, so resume is cheap (no re-init, no bootstrap rendezvous). Implemented along vLLM's existing NCCL conventions, no new module: - pynccl_wrapper.py: bind ncclCommSuspend/ncclCommResume in NCCLLibrary (skipped gracefully with a one-time warning on NCCL < 2.29.7). - pynccl.py: PyNcclCommunicator.suspend()/resume(), mirroring register_comm_window(); gated on disabled and nccl_version >= 22907. - parallel_state.py: suspend_nccl_comms()/resume_nccl_comms() iterate the _groups registry and act on each group's communicator. Collective and synchronous on return, so no extra device sync is needed. - gpu_worker.py: suspend after allocator.sleep(), resume after wake_up(). The symbols exist only in NCCL >= 2.29.7, so this is a no-op on torch 2.11.0 (NCCL 2.28.9) and activates automatically on torch >= 2.12.1 (NCCL 2.29.7), with no code change. Co-Authored-By: Claude Opus 4.8 Signed-off-by: aoshen02 --- .../device_communicators/pynccl.py | 21 +++++++++ .../device_communicators/pynccl_wrapper.py | 23 ++++++++++ vllm/distributed/parallel_state.py | 43 +++++++++++++++++++ vllm/v1/worker/gpu_worker.py | 11 +++++ 4 files changed, 98 insertions(+) diff --git a/vllm/distributed/device_communicators/pynccl.py b/vllm/distributed/device_communicators/pynccl.py index 9f305c718f9d..3a02cc82cb05 100644 --- a/vllm/distributed/device_communicators/pynccl.py +++ b/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 diff --git a/vllm/distributed/device_communicators/pynccl_wrapper.py b/vllm/distributed/device_communicators/pynccl_wrapper.py index 5ca8cc7c77f4..a8d5c723da44 100644 --- a/vllm/distributed/device_communicators/pynccl_wrapper.py +++ b/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"]()) diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 11b9e24e864c..035516e87ad4 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -127,6 +127,49 @@ 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) + freed = abs(free_before - torch.cuda.mem_get_info()[0]) + logger.info("NCCL %s: %d comms, %.1f MiB", label, len(comms), freed / 1024**2) + + +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]() diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 2947101d9778..b66eecbc8bba 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/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 From 3cc12e7cc735e2252e5e85c4e1cb7e3c0d3efccf Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Wed, 1 Jul 2026 00:39:52 +0000 Subject: [PATCH 2/5] [Core] Log freed/allocated direction for NCCL suspend/resume Address review feedback: replace abs() in the sleep-mode NCCL log with a signed delta plus an explicit freed/allocated direction, so a wrong-sign result (e.g. an accidental allocation during suspend) is surfaced instead of masked. Co-Authored-By: Claude Opus 4.8 Signed-off-by: aoshen02 --- vllm/distributed/parallel_state.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 035516e87ad4..82ab45a93ddb 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -152,8 +152,15 @@ def _apply_to_nccl_comms(label: str, action: Callable[[Any], None]) -> None: free_before = torch.cuda.mem_get_info()[0] for pynccl in comms: action(pynccl) - freed = abs(free_before - torch.cuda.mem_get_info()[0]) - logger.info("NCCL %s: %d comms, %.1f MiB", label, len(comms), freed / 1024**2) + 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: From b27115fbbe53df43f92107a5abd2276ec4a37dff Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 5 Jul 2026 11:17:45 +0000 Subject: [PATCH 3/5] [Core] Generalize sleep comm release to a device-communicator surface Per review on #46234: instead of parallel_state reaching into device_communicator.pynccl_comm, add suspend()/resume() to DeviceCommunicatorBase (default no-op) and override in CudaCommunicator to forward to its pynccl_comm. parallel_state now walks device communicators polymorphically (suspend_device_comms/resume_device_comms), so FlashInfer / all2all comms can add their own suspend later without touching parallel_state. Non-CUDA backends inherit the base no-op. The freed-bytes probe also moves to the torch.accelerator memory API. Co-Authored-By: Claude Opus 4.8 Signed-off-by: aoshen02 --- .../base_device_communicator.py | 10 ++++ .../device_communicators/cuda_communicator.py | 8 +++ vllm/distributed/parallel_state.py | 50 +++++++++---------- vllm/v1/worker/gpu_worker.py | 18 +++---- 4 files changed, 52 insertions(+), 34 deletions(-) diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index 6fd889daeb05..72645615ebd4 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -321,6 +321,16 @@ def broadcast(self, tensor: torch.Tensor, src: int = 0) -> torch.Tensor: def destroy(self): pass + def suspend(self) -> None: + """Release reclaimable communicator memory (default: no-op). + + Overridden by communicators that can free their device buffers while a + engine sleeps and restore them on ``resume``. Collective across ranks. + """ + + def resume(self) -> None: + """Restore memory released by ``suspend`` (default: no-op).""" + def prepare_communication_buffer_for_model(self, model: torch.nn.Module) -> None: """ Prepare the communication buffer for the model. diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index b92015b18808..6836cda4d331 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -516,6 +516,14 @@ def destroy(self): self.all2all_manager.destroy() self.all2all_manager = None # type: ignore[assignment] + def suspend(self) -> None: + if self.pynccl_comm is not None: + self.pynccl_comm.suspend() + + def resume(self) -> None: + if self.pynccl_comm is not None: + self.pynccl_comm.resume() + def all_gatherv( self, input_: torch.Tensor | list[torch.Tensor], diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 82ab45a93ddb..8a265a92776c 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -127,35 +127,34 @@ 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. +def _apply_to_device_comms(label: str, action: Callable[[Any], None]) -> None: + """Apply ``action`` to every group's device communicator, collectively. + + Walks the registered parallel groups and skips those without a device + communicator (absent at ``world_size == 1``). Each communicator's + ``suspend``/``resume`` is a no-op unless it holds releasable device memory + (see ``DeviceCommunicatorBase``). Collective across ranks 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: + dc = group.device_communicator + if dc is None: continue - comms.append(pynccl) + comms.append(dc) 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 + free_before = torch.accelerator.get_memory_info()[0] + for dc in comms: + action(dc) + delta = torch.accelerator.get_memory_info()[0] - free_before direction = "freed" if delta > 0 else "allocated" logger.info( - "NCCL %s: %d comms, %.1f MiB %s", + "device-comm %s: %d comms, %.1f MiB %s", label, len(comms), abs(delta) / 1024**2, @@ -163,18 +162,19 @@ def _apply_to_nccl_comms(label: str, action: Callable[[Any], None]) -> None: ) -def suspend_nccl_comms() -> None: - """Release idle NCCL communicator memory on every group (collective). +def suspend_device_comms() -> None: + """Release idle device 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. + Must run on every rank with communicators idle. Communicator suspend hooks + (e.g. NCCL's ncclCommSuspend) have internal cross-rank barriers; a no-op + where unsupported (older NCCL, world_size 1, non-CUDA). """ - _apply_to_nccl_comms("suspend", lambda c: c.suspend()) + _apply_to_device_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 resume_device_comms() -> None: + """Restore all suspended device communicators before reuse (collective).""" + _apply_to_device_comms("resume", lambda c: c.resume()) def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor: diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index d568207c0e34..98709d8ffd0f 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -43,8 +43,8 @@ Handle, get_pp_group, get_tp_group, - resume_nccl_comms, - suspend_nccl_comms, + resume_device_comms, + suspend_device_comms, ) from vllm.distributed.weight_transfer import ( WeightTransferEngine, @@ -199,10 +199,10 @@ def sleep(self, level: int = 1) -> None: self._get_sleep_mode_backend().suspend(level) - # 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() + # Release idle device communicator memory (e.g. NCCL >= 2.29.7). + # Collective across ranks; no-op where unsupported. Done after the + # cumem unmap so the reported freed bytes include the comm release. + suspend_device_comms() torch.accelerator.synchronize() deadline = time.monotonic() + (5.0 if current_platform.is_rocm() else 0) @@ -224,9 +224,9 @@ def sleep(self, level: int = 1) -> None: def wake_up(self, tags: list[str] | None = None) -> None: self._get_sleep_mode_backend().resume(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 device communicator memory before any collective runs again. + # Collective across ranks; no-op where unsupported. + resume_device_comms() # Restore the buffers after level 2 sleep if len(self._sleep_saved_buffers): From a3444d1ff3265da77228bbac2143440d2b116f3d Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 5 Jul 2026 19:48:44 +0800 Subject: [PATCH 4/5] Update base_device_communicator.py Signed-off-by: aoshen02 --- .../device_communicators/base_device_communicator.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index 72645615ebd4..3af5c2de636d 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -322,11 +322,7 @@ def destroy(self): pass def suspend(self) -> None: - """Release reclaimable communicator memory (default: no-op). - - Overridden by communicators that can free their device buffers while a - engine sleeps and restore them on ``resume``. Collective across ranks. - """ + """Release reclaimable communicator memory (default: no-op).""" def resume(self) -> None: """Restore memory released by ``suspend`` (default: no-op).""" From ad16cece3df99eec5a2c60d842da0646f62a6a16 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 5 Jul 2026 19:49:54 +0800 Subject: [PATCH 5/5] Update pynccl_wrapper.py Signed-off-by: aoshen02 --- vllm/distributed/device_communicators/pynccl_wrapper.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/vllm/distributed/device_communicators/pynccl_wrapper.py b/vllm/distributed/device_communicators/pynccl_wrapper.py index a8d5c723da44..622d59f5ffe2 100644 --- a/vllm/distributed/device_communicators/pynccl_wrapper.py +++ b/vllm/distributed/device_communicators/pynccl_wrapper.py @@ -296,9 +296,6 @@ 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);