4ac452ad98 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/distributed/test_pynccl.py

+39

@@ -472,3 +472,42 @@ def test_ncclGetUniqueId():

472 # 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

473 # as long as the function doesn't raise an exception, we're good

474 assert unique_id is not None

475 +

476 +

477 +def test_pynccl_suspend_resume_idempotent():

478 + """Repeated suspend/resume (e.g. staged wake-ups) reach NCCL once each."""

479 + from unittest.mock import Mock

480 +

481 + comm = object.__new__(PyNcclCommunicator)

482 + comm.disabled = False

483 + comm._suspended = False

484 + comm.comm = object()

485 + comm.nccl = Mock()

486 +

487 + comm.suspend()

488 + comm.suspend()

489 + comm.resume()

490 + comm.resume()

491 +

492 + assert comm.nccl.ncclCommSuspend.call_count == 1

493 + assert comm.nccl.ncclCommResume.call_count == 1

494 +

495 +

496 +def test_pynccl_suspend_noop_without_symbol():

497 + """RCCL / NCCL < 2.29.7 lack ncclCommSuspend; suspend() must no-op, not

498 + crash, and resume() must then no-op too since nothing was suspended."""

499 + from unittest.mock import Mock

500 +

501 + comm = object.__new__(PyNcclCommunicator)

502 + comm.disabled = False

503 + comm._suspended = False

504 + comm.comm = object()

505 + comm.nccl = Mock()

506 + comm.nccl.has_symbol.return_value = False

507 +

508 + comm.suspend()

509 + comm.resume()

510 +

511 + comm.nccl.ncclCommSuspend.assert_not_called()

512 + comm.nccl.ncclCommResume.assert_not_called()

513 + assert comm._suspended is False

tests/v1/worker/test_kv_cache_allocation_scope.py

+4 -1

@@ -113,11 +113,14 @@ def test_kv_wake_does_not_run_model_runner_recovery() -> None:

113 worker = cast(

114 Worker,

115 SimpleNamespace(

116 - _get_sleep_mode_backend=lambda: SimpleNamespace(resume=lambda tags: None),

116 + sleep_mode_backend=SimpleNamespace(resume=lambda tags: None),

117 _sleep_saved_buffers={},

118 _sleep_saved_draft_buffers={},

119 model_runner=runner,

120 synchronize_device=lambda: None,

121 + vllm_config=SimpleNamespace(

122 + model_config=SimpleNamespace(enable_nccl_comm_suspend=False)

123 + ),

124 ),

125 )

126 layout_tensors = runner.layout_tensors

tests/v1/worker/test_sleep_mode_backend.py

+49

@@ -38,6 +38,55 @@ def test_new_backend_starts_in_running_state():

38 assert CuMemBackend().state() == "RUNNING"

39

40

41 +@pytest.mark.parametrize("enable_nccl_comm_suspend", [True, False])

42 +def test_worker_drives_communicator_suspension(monkeypatch, enable_nccl_comm_suspend):

43 + """Comm walkers run around sleep/wake only when explicitly enabled."""

44 + from types import SimpleNamespace

45 +

46 + from vllm.v1.worker.gpu_worker import Worker

47 +

48 + calls: list[tuple[str, object]] = []

49 +

50 + class Backend:

51 + def suspend(self, level: int = 1) -> None:

52 + calls.append(("backend.suspend", level))

53 +

54 + def resume(self, tags: list[str] | None = None) -> None:

55 + calls.append(("backend.resume", tuple(tags) if tags else None))

56 +

57 + worker = object.__new__(Worker)

58 + worker._sleep_mode_backend = Backend()

59 + worker._sleep_saved_buffers = {}

60 + worker._sleep_saved_draft_buffers = {}

61 + worker.vllm_config = SimpleNamespace(

62 + model_config=SimpleNamespace(enable_nccl_comm_suspend=enable_nccl_comm_suspend)

63 + )

64 +

65 + monkeypatch.setattr("torch.accelerator.synchronize", lambda: None)

66 + monkeypatch.setattr("torch.accelerator.get_memory_info", lambda: (0, 0))

67 + monkeypatch.setattr(

68 + "vllm.v1.worker.gpu_worker.suspend_device_comms",

69 + lambda: calls.append(("comms.suspend", None)),

70 + )

71 + monkeypatch.setattr(

72 + "vllm.v1.worker.gpu_worker.resume_device_comms",

73 + lambda: calls.append(("comms.resume", None)),

74 + )

75 +

76 + worker.sleep(level=1)

77 + worker.wake_up(tags=["weights"])

78 +

79 + expected = [

80 + ("backend.suspend", 1),

81 + ("comms.suspend", None),

82 + ("backend.resume", ("weights",)),

83 + ("comms.resume", None),

84 + ]

85 + if not enable_nccl_comm_suspend:

86 + expected = [c for c in expected if not c[0].startswith("comms.")]

87 + assert calls == expected

88 +

89 +

90 def test_unknown_backend_raises():

91 with pytest.raises(ValueError, match="Unsupported sleep-mode backend"):

92 SleepModeBackendFactory.get_backend_class("does-not-exist")

vllm/config/model.py

+4

@@ -343,6 +343,10 @@ class ModelConfig:

343 (default) uses the built-in ``CuMemAllocator`` and is behavior-compatible

344 with prior releases. Additional backends (CUDA checkpoint, CRIU, durable

345 snapshot) may be registered in-tree or by plugins (RFC #34303)."""

346 + enable_nccl_comm_suspend: bool = False

347 + """Enable releasing NCCL communicator memory during sleep mode

348 + (``ncclCommSuspend``/``ncclCommResume``). Experimental; when disabled

349 + (the default) sleep still releases weights/KV-cache memory as before."""

350 enable_cumem_allocator: bool = False

351 """Enable the custom cumem allocator to leverage advanced GPU memory

352 allocation features such as multi-node NVLink support.

vllm/device_allocator/sleep_mode_backend.py

+1 -2

@@ -134,8 +134,7 @@ class CuMemBackend(SleepModeBackend):

134

135 @classmethod

136 def preserves_communicators(cls) -> bool:

137 - # Communicator buffers (e.g. NCCL) live outside CuMemAllocator's pool, so

138 - # an allocator-level sleep leaves them intact (no reinit needed on resume).

137 + # Comm identity/topology survive memory suspension; no reinit needed.

138 return True

139

140

vllm/distributed/device_communicators/base_device_communicator.py

+6

@@ -226,6 +226,12 @@ class DeviceCommunicatorBase:

226 def checkpoint_restore(self) -> None:

227 """Restore communicator state after checkpoint (default: no-op)."""

228

229 + def suspend(self) -> None:

230 + """Release reclaimable communicator memory (default: no-op)."""

231 +

232 + def resume(self) -> None:

233 + """Restore memory released by ``suspend`` (default: no-op)."""

234 +

235 def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:

236 if dim < 0:

237 # Convert negative dim to positive.

vllm/distributed/device_communicators/cuda_communicator.py

+8

@@ -631,6 +631,14 @@ class CudaCommunicator(DeviceCommunicatorBase):

631 self.all2all_manager.destroy()

632 self.all2all_manager = None # type: ignore[assignment]

633

634 + def suspend(self) -> None:

635 + if self.pynccl_comm is not None:

636 + self.pynccl_comm.suspend()

637 +

638 + def resume(self) -> None:

639 + if self.pynccl_comm is not None:

640 + self.pynccl_comm.resume()

641 +

642 def checkpoint_prepare(self) -> None:

643 # Only FlashInfer all-reduce and FlashInfer all2all are supported for now.

644 from .flashinfer_all_reduce import checkpoint_prepare_fi_ar_workspaces

vllm/distributed/device_communicators/pynccl.py

+24

@@ -27,6 +27,8 @@ logger = init_logger(__name__)

27

28 _NCCL_SYMM_OPS_REGISTERED = False

29

30 +_NCCL_SUSPEND_MEM = 0x01

31 +

32

33 def register_nccl_symmetric_ops(pynccl_comm):

34 from vllm.distributed.device_communicators.pynccl_allocator import (

@@ -105,6 +107,7 @@ class PyNcclCommunicator:

107

108 self.available = True

109 self.disabled = False

110 + self._suspended = False

111

112 self.nccl_version = self.nccl.ncclGetRawVersion()

113 if self.rank == 0:

@@ -479,6 +482,27 @@ class PyNcclCommunicator:

482 def deregister_comm_window(self, window):

483 return self.nccl.ncclCommWindowDeregister(self.comm, window)

484

485 + def suspend(self):

486 + """Release comm GPU memory (collective, idempotent); keeps topology."""

487 + if self.disabled or self._suspended:

488 + return

489 + if not self.nccl.has_symbol("ncclCommSuspend"):

490 + logger.warning_once(

491 + "ncclCommSuspend is not available in the loaded NCCL/RCCL "

492 + "library (requires NCCL >= 2.29.7); skipping communicator "

493 + "memory suspension."

494 + )

495 + return

496 + self.nccl.ncclCommSuspend(self.comm, _NCCL_SUSPEND_MEM)

497 + self._suspended = True

498 +

499 + def resume(self):

500 + """Restore a suspended comm (collective); no-op unless suspended."""

501 + if self.disabled or not self._suspended:

502 + return

503 + self.nccl.ncclCommResume(self.comm)

504 + self._suspended = False

505 +

506 def batch_isend_irecv(self, p2p_ops: list, stream=None):

507 if self.disabled:

508 return

vllm/distributed/device_communicators/pynccl_wrapper.py

+21

@@ -329,6 +329,10 @@ class NCCLLibrary:

329 # shutdown when peer ranks may already be gone.

330 # ncclResult_t ncclCommAbort(ncclComm_t comm);

331 Function("ncclCommAbort", ncclResult_t, [ncclComm_t]),

332 + # ncclResult_t ncclCommSuspend(ncclComm_t comm, int flags);

333 + Function("ncclCommSuspend", ncclResult_t, [ncclComm_t, ctypes.c_int]),

334 + # ncclResult_t ncclCommResume(ncclComm_t comm);

335 + Function("ncclCommResume", ncclResult_t, [ncclComm_t]),

336 # ncclResult_t ncclGroupStart();

337 Function("ncclGroupStart", ncclResult_t, []),

338 # ncclResult_t ncclGroupEnd();

@@ -417,10 +421,21 @@ class NCCLLibrary:

421 elif func.name == "ncclCommQueryProperties":

422 # Optional on NCCL versions older than 2.29.

423 continue

424 + elif func.name in ("ncclCommSuspend", "ncclCommResume"):

425 + # RCCL doesn't export these; NCCL >= 2.29.7 does, and

426 + # vLLM's CUDA path already requires that version.

427 + # PyNcclCommunicator checks has_symbol() before

428 + # calling either.

429 + if current_platform.is_rocm():

430 + continue

431 raise

432 NCCLLibrary.path_to_dict_mapping[so_file] = _funcs

433 self._funcs = NCCLLibrary.path_to_dict_mapping[so_file]

434

435 + def has_symbol(self, name: str) -> bool:

436 + """Whether the loaded NCCL/RCCL library exports the given symbol."""

437 + return name in self._funcs

438 +

439 def ncclGetErrorString(self, result: ncclResult_t) -> str:

440 return self._funcs["ncclGetErrorString"](result).decode("utf-8")

441

@@ -599,6 +614,12 @@ class NCCLLibrary:

614 def ncclCommAbort(self, comm: ncclComm_t) -> None:

615 self.NCCL_CHECK(self._funcs["ncclCommAbort"](comm))

616

617 + def ncclCommSuspend(self, comm: ncclComm_t, flags: int) -> None:

618 + self.NCCL_CHECK(self._funcs["ncclCommSuspend"](comm, flags))

619 +

620 + def ncclCommResume(self, comm: ncclComm_t) -> None:

621 + self.NCCL_CHECK(self._funcs["ncclCommResume"](comm))

622 +

623 def ncclGroupStart(self) -> None:

624 self.NCCL_CHECK(self._funcs["ncclGroupStart"]())

625

vllm/distributed/parallel_state.py

+10

@@ -180,6 +180,16 @@ def _apply_to_device_comms(

180 action(dc)

181

182

183 +def suspend_device_comms() -> None:

184 + """Release device communicator memory (collective; comms must be idle)."""

185 + _apply_to_device_comms(lambda comm: comm.suspend())

186 +

187 +

188 +def resume_device_comms() -> None:

189 + """Restore suspended device communicators (collective)."""

190 + _apply_to_device_comms(lambda comm: comm.resume())

191 +

192 +

193 def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor:

194 assert group_name in _groups, f"Group {group_name} is not found."

195 group = _groups[group_name]()

vllm/engine/arg_utils.py

+6

@@ -710,6 +710,7 @@ class EngineArgs:

710 generation_config: str = ModelConfig.generation_config

711 enable_sleep_mode: bool = ModelConfig.enable_sleep_mode

712 enable_cumem_allocator: bool = ModelConfig.enable_cumem_allocator

713 + enable_nccl_comm_suspend: bool = ModelConfig.enable_nccl_comm_suspend

714 override_generation_config: dict[str, Any] = get_field(

715 ModelConfig, "override_generation_config"

716 )

@@ -932,6 +933,10 @@ class EngineArgs:

933 model_group.add_argument(

934 "--enable-cumem-allocator", **model_kwargs["enable_cumem_allocator"]

935 )

936 + model_group.add_argument(

937 + "--enable-nccl-comm-suspend",

938 + **model_kwargs["enable_nccl_comm_suspend"],

939 + )

940 model_group.add_argument("--model-impl", **model_kwargs["model_impl"])

941 model_group.add_argument(

942 "--logits-processors", **model_kwargs["logits_processors"]

@@ -1823,6 +1828,7 @@ class EngineArgs:

1828 override_generation_config=self.override_generation_config,

1829 enable_sleep_mode=self.enable_sleep_mode,

1830 enable_cumem_allocator=self.enable_cumem_allocator,

1831 + enable_nccl_comm_suspend=self.enable_nccl_comm_suspend,

1832 model_impl=self.model_impl,

1833 logits_processors=self.logits_processors,

1834 video_pruning_rate=self.video_pruning_rate,

vllm/v1/worker/gpu_worker.py

+10 -3

@@ -45,6 +45,8 @@ from vllm.distributed.parallel_state import (

45 checkpoint_restore_distributed_state,

46 get_pp_group,

47 get_tp_group,

48 + resume_device_comms,

49 + suspend_device_comms,

50 )

51 from vllm.distributed.weight_transfer import (

52 WeightTransferEngine,

@@ -220,7 +222,8 @@ class Worker(WorkerBase):

222 # Resolved lazily on first sleep/wake; persists worker-process state.

223 self._sleep_mode_backend: SleepModeBackend | None = None

224

223 - def _get_sleep_mode_backend(self) -> "SleepModeBackend":

225 + @property

226 + def sleep_mode_backend(self) -> "SleepModeBackend":

227 if self._sleep_mode_backend is None:

228 from vllm.device_allocator.sleep_mode_backend import (

229 SleepModeBackendFactory,

@@ -247,7 +250,9 @@ class Worker(WorkerBase):

250 name: buffer.cpu().clone() for name, buffer in draft.named_buffers()

251 }

252

250 - self._get_sleep_mode_backend().suspend(level)

253 + self.sleep_mode_backend.suspend(level)

254 + if self.vllm_config.model_config.enable_nccl_comm_suspend:

255 + suspend_device_comms()

256

257 torch.accelerator.synchronize()

258 deadline = time.monotonic() + (5.0 if current_platform.is_rocm() else 0)

@@ -267,7 +272,9 @@ class Worker(WorkerBase):

272 )

273

274 def wake_up(self, tags: list[str] | None = None) -> None:

270 - self._get_sleep_mode_backend().resume(tags)

275 + self.sleep_mode_backend.resume(tags)

276 + if self.vllm_config.model_config.enable_nccl_comm_suspend:

277 + resume_device_comms()

278

279 # Restore the buffers after level 2 sleep

280 wake_weights = tags is None or "weights" in tags

[Core] Release NCCL communicator memory in sleep mode (#51485)