[Core][Hardware][NVIDIA] Add custom all-reduce suspend hooks

csrc/libtorch_stable/custom_all_reduce.cu

@@ -119,6 +119,11 @@ void dispose(fptr_t _fa) {
   delete reinterpret_cast<vllm::CustomAllreduce*>(_fa);
 }
 
+void custom_ar_close(fptr_t _fa) {
+  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
+  fa->close();
+}
+
 int64_t meta_size() { return sizeof(vllm::Signal); }
 
 void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs) {
@@ -140,6 +145,14 @@ get_graph_buffer_ipc_meta(fptr_t _fa) {
   return std::make_tuple(bytes, offsets);
 }
 
+std::tuple<std::vector<int64_t>, std::vector<int64_t>>
+get_graph_buffer_ipc_meta_for_reinit(fptr_t _fa) {
+  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
+  auto [handle, offsets] = fa->get_graph_buffer_ipc_meta_for_reinit();
+  std::vector<int64_t> bytes(handle.begin(), handle.end());
+  return std::make_tuple(bytes, offsets);
+}
+
 // Use vector<int64_t> to represent byte data for python binding compatibility.
 void register_graph_buffers(fptr_t _fa,
                             const std::vector<std::vector<int64_t>>& handles,
@@ -154,6 +167,33 @@ void register_graph_buffers(fptr_t _fa,
   fa->register_graph_buffers(bytes, offsets);
 }
 
+void custom_ar_prepare_for_suspend(fptr_t _fa) {
+  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
+  fa->prepare_for_suspend();
+}
+
+void custom_ar_reinit_after_resume(
+    fptr_t _fa, const std::vector<std::vector<int64_t>>& handles,
+    const std::vector<std::vector<int64_t>>& offsets,
+    const std::vector<fptr_t>& fake_signal_ptrs,
+    const std::vector<fptr_t>& fake_buffer_ptrs) {
+  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
+  STD_TORCH_CHECK(fake_signal_ptrs.size() == fa->world_size_);
+  STD_TORCH_CHECK(fake_buffer_ptrs.size() == fa->world_size_);
+  std::vector<std::string> bytes;
+  bytes.reserve(handles.size());
+  for (const auto& handle : handles) {
+    bytes.emplace_back(handle.begin(), handle.end());
+  }
+  vllm::Signal* signal_ptrs[8];
+  void* buffer_ptrs[8];
+  for (int i = 0; i < fa->world_size_; i++) {
+    signal_ptrs[i] = reinterpret_cast<vllm::Signal*>(fake_signal_ptrs[i]);
+    buffer_ptrs[i] = reinterpret_cast<void*>(fake_buffer_ptrs[i]);
+  }
+  fa->reinit_after_resume(bytes, offsets, signal_ptrs, buffer_ptrs);
+}
+
 std::tuple<fptr_t, torch::stable::Tensor> allocate_shared_buffer_and_handle(
     int64_t size) {
   int device_index;
@@ -197,6 +237,21 @@ fptr_t open_mem_handle(torch::stable::Tensor& mem_handle) {
   return reinterpret_cast<fptr_t>(ipc_ptr);
 }
 
+torch::stable::Tensor get_mem_handle(fptr_t buffer) {
+  auto handle = torch::stable::empty(
+      {static_cast<int64_t>(sizeof(cudaIpcMemHandle_t))},
+      torch::headeronly::ScalarType::Byte, std::nullopt,
+      torch::stable::Device(torch::stable::DeviceType::CPU));
+  STD_CUDA_CHECK(cudaIpcGetMemHandle(
+      reinterpret_cast<cudaIpcMemHandle_t*>(handle.mutable_data_ptr()),
+      reinterpret_cast<void*>(buffer)));
+  return handle;
+}
+
+void close_mem_handle(fptr_t ptr) {
+  STD_CUDA_CHECK(cudaIpcCloseMemHandle(reinterpret_cast<void*>(ptr)));
+}
+
 void free_shared_buffer(fptr_t buffer) {
   STD_CUDA_CHECK(cudaFree(reinterpret_cast<void*>(buffer)));
 }

csrc/libtorch_stable/ops.h

@@ -374,16 +374,27 @@ void all_reduce(fptr_t _fa, torch::stable::Tensor& inp,
                 torch::stable::Tensor& out, fptr_t reg_buffer,
                 int64_t reg_buffer_sz_bytes);
 void dispose(fptr_t _fa);
+void custom_ar_close(fptr_t _fa);
 int64_t meta_size();
 void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
 std::tuple<std::vector<int64_t>, std::vector<int64_t>>
 get_graph_buffer_ipc_meta(fptr_t _fa);
+std::tuple<std::vector<int64_t>, std::vector<int64_t>>
+get_graph_buffer_ipc_meta_for_reinit(fptr_t _fa);
 void register_graph_buffers(fptr_t _fa,
                             const std::vector<std::vector<int64_t>>& handles,
                             const std::vector<std::vector<int64_t>>& offsets);
+void custom_ar_prepare_for_suspend(fptr_t _fa);
+void custom_ar_reinit_after_resume(
+    fptr_t _fa, const std::vector<std::vector<int64_t>>& handles,
+    const std::vector<std::vector<int64_t>>& offsets,
+    const std::vector<int64_t>& fake_signal_ptrs,
+    const std::vector<int64_t>& fake_buffer_ptrs);
 std::tuple<int64_t, torch::stable::Tensor> allocate_shared_buffer_and_handle(
     int64_t size);
 int64_t open_mem_handle(torch::stable::Tensor& mem_handle);
+torch::stable::Tensor get_mem_handle(int64_t buffer);
+void close_mem_handle(int64_t ptr);
 void free_shared_buffer(int64_t buffer);
 
 // Activation kernels (shared CUDA/ROCm)

csrc/libtorch_stable/torch_bindings.cpp

@@ -888,13 +888,22 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ar) {
       "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, "
       "int reg_buffer_sz_bytes) -> ()");
   custom_ar.def("dispose(int fa) -> ()");
+  custom_ar.def("custom_ar_close(int fa) -> ()");
   custom_ar.def("meta_size() -> int");
   custom_ar.def("register_buffer(int fa, int[] ipc_tensors) -> ()");
   custom_ar.def("get_graph_buffer_ipc_meta(int fa) -> (int[], int[])");
+  custom_ar.def(
+      "get_graph_buffer_ipc_meta_for_reinit(int fa) -> (int[], int[])");
   custom_ar.def(
       "register_graph_buffers(int fa, int[][] handles, int[][] offsets) -> ()");
+  custom_ar.def("custom_ar_prepare_for_suspend(int fa) -> ()");
+  custom_ar.def(
+      "custom_ar_reinit_after_resume(int fa, int[][] handles, "
+      "int[][] offsets, int[] signal_ptrs, int[] buffer_ptrs) -> ()");
   custom_ar.def("allocate_shared_buffer_and_handle(int size) -> (int, Tensor)");
   custom_ar.def("open_mem_handle(Tensor mem_handle) -> int");
+  custom_ar.def("get_mem_handle(int ptr) -> Tensor");
+  custom_ar.def("close_mem_handle(int ptr) -> ()");
   custom_ar.def("free_shared_buffer(int ptr) -> ()");
 }
 
@@ -909,13 +918,22 @@ STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CPU, custom_ar) {
 
 STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CompositeExplicitAutograd, custom_ar) {
   custom_ar.impl("dispose", TORCH_BOX(&dispose));
+  custom_ar.impl("custom_ar_close", TORCH_BOX(&custom_ar_close));
   custom_ar.impl("meta_size", TORCH_BOX(&meta_size));
   custom_ar.impl("register_buffer", TORCH_BOX(&register_buffer));
   custom_ar.impl("get_graph_buffer_ipc_meta",
                  TORCH_BOX(&get_graph_buffer_ipc_meta));
+  custom_ar.impl("get_graph_buffer_ipc_meta_for_reinit",
+                 TORCH_BOX(&get_graph_buffer_ipc_meta_for_reinit));
   custom_ar.impl("register_graph_buffers", TORCH_BOX(&register_graph_buffers));
+  custom_ar.impl("custom_ar_prepare_for_suspend",
+                 TORCH_BOX(&custom_ar_prepare_for_suspend));
+  custom_ar.impl("custom_ar_reinit_after_resume",
+                 TORCH_BOX(&custom_ar_reinit_after_resume));
   custom_ar.impl("allocate_shared_buffer_and_handle",
                  TORCH_BOX(&allocate_shared_buffer_and_handle));
+  custom_ar.impl("get_mem_handle", TORCH_BOX(&get_mem_handle));
+  custom_ar.impl("close_mem_handle", TORCH_BOX(&close_mem_handle));
   custom_ar.impl("free_shared_buffer", TORCH_BOX(&free_shared_buffer));
 }
 

tests/distributed/test_custom_all_reduce.py

@@ -8,7 +8,9 @@
 import torch
 import torch.distributed as dist
 
+from vllm import _custom_ops as ops
 from vllm.distributed.communication_op import tensor_model_parallel_all_reduce  # noqa
+from vllm.distributed.device_communicators.custom_all_reduce import _LifecycleState
 from vllm.distributed.parallel_state import get_tp_group, graph_capture
 
 from ..utils import (
@@ -104,22 +106,135 @@ def eager_allreduce(
         num_communication = rank // tp_size + 1
         sz = 1024
         fa = get_tp_group().device_communicator.ca_comm
+        assert fa is not None
+        with pytest.raises(RuntimeError, match="exactly one eager registered buffer"):
+            ops.register_buffer(fa._ptr, fa.buffer_ptrs)
         inp = torch.ones(sz, dtype=torch.float32, device=device)
         out = inp
         for _ in range(num_communication):
             out = fa.all_reduce(out, registered=False)
         torch.testing.assert_close(out, inp * (tp_size**num_communication))
 
+        if torch.version.hip is None:
+            expected_peer_mappings = 2 * (fa.world_size - 1)
+            assert len(fa._open_peer_ptrs) == expected_peer_mappings
+            fa.prepare_for_suspend()
+            assert not fa._open_peer_ptrs
+            fa.prepare_for_suspend()
+            with pytest.raises(RuntimeError, match="detached"):
+                tensor_model_parallel_all_reduce(inp)
+            fa.reinit_after_resume()
+            assert len(fa._open_peer_ptrs) == expected_peer_mappings
+            fa.reinit_after_resume()
+            eager_out = tensor_model_parallel_all_reduce(inp)
+            torch.testing.assert_close(eager_out, inp * tp_size)
+
         inp = torch.ones(sz * 4, dtype=torch.bfloat16, device=device)
         out = inp
         for _ in range(num_communication):
             out = fa.all_reduce(out, registered=False)
         torch.testing.assert_close(out, inp * (tp_size**num_communication))
+        if torch.version.hip is None:
+            fa.prepare_for_suspend()
+            fa.close()
+            fa.close()
+
+
+@ray.remote(num_gpus=1, max_calls=1)
+def graph_allreduce_suspend_resume(
+    monkeypatch: pytest.MonkeyPatch,
+    tp_size,
+    pp_size,
+    rank,
+    distributed_init_port,
+):
+    with monkeypatch.context() as m:
+        m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+        m.delenv("HIP_VISIBLE_DEVICES", raising=False)
+        device = torch.device(f"cuda:{rank}")
+        torch.accelerator.set_device_index(device)
+        init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
+        ensure_model_parallel_initialized(tp_size, pp_size)
+        tp_group = get_tp_group()
+        device_communicator = tp_group.device_communicator
+        assert device_communicator is not None
+        fa = device_communicator.ca_comm
+        assert fa is not None and not fa.disabled
+
+        selected_calls = 0
+        original_custom_all_reduce = fa.custom_all_reduce
+
+        def tracked_custom_all_reduce(input_: torch.Tensor) -> torch.Tensor | None:
+            nonlocal selected_calls
+            selected_calls += 1
+            return original_custom_all_reduce(input_)
+
+        m.setattr(fa, "custom_all_reduce", tracked_custom_all_reduce)
+        inp1 = torch.full((1024,), fa.rank + 1, dtype=torch.float32, device=device)
+        inp2 = torch.full((2048,), fa.rank + 2, dtype=torch.float32, device=device)
+        with graph_capture(device=device) as graph_capture_context:
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, stream=graph_capture_context.stream):
+                out1 = tensor_model_parallel_all_reduce(inp1)
+                out1_repeat = tensor_model_parallel_all_reduce(inp1)
+                out2 = tensor_model_parallel_all_reduce(inp2)
+        assert selected_calls == 3
+        assert fa._lifecycle_state == _LifecycleState.ACTIVE
+
+        expected_peer_mappings = 2 * (fa.world_size - 1)
+        assert len(fa._open_peer_ptrs) == expected_peer_mappings
+        fa.prepare_for_suspend()
+        assert not fa._open_peer_ptrs
+        fa.prepare_for_suspend()
+        with pytest.raises(RuntimeError, match="detached"):
+            tensor_model_parallel_all_reduce(inp1)
+        fa.reinit_after_resume()
+        assert len(fa._open_peer_ptrs) == expected_peer_mappings
+        fa.reinit_after_resume()
+
+        for replay_idx in range(2):
+            inp1.fill_(fa.rank + 3 + replay_idx)
+            inp2.fill_(fa.rank + 4 + replay_idx)
+            graph.replay()
+            torch.accelerator.synchronize()
+            torch.testing.assert_close(
+                out1,
+                torch.full_like(
+                    out1, sum(range(3 + replay_idx, tp_size + 3 + replay_idx))
+                ),
+            )
+            torch.testing.assert_close(out1_repeat, out1)
+            torch.testing.assert_close(
+                out2,
+                torch.full_like(
+                    out2, sum(range(4 + replay_idx, tp_size + 4 + replay_idx))
+                ),
+            )
+
+        selected_before_eager = selected_calls
+        eager_out = tensor_model_parallel_all_reduce(inp1)
+        assert selected_calls == selected_before_eager + 1
+        torch.testing.assert_close(eager_out, out1)
+        fa.close()
+        fa.close()
 
 
 @pytest.mark.parametrize("tp_size", [2])
 @pytest.mark.parametrize("pipeline_parallel_size", [1, 2])
-@pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce])
+@pytest.mark.parametrize(
+    "test_target",
+    [
+        eager_allreduce,
+        graph_allreduce,
+        pytest.param(
+            graph_allreduce_suspend_resume,
+            marks=pytest.mark.skipif(
+                torch.version.hip is not None,
+                reason="Custom allreduce suspend/resume requires CUDA.",
+            ),
+        ),
+    ],
+)
 def test_custom_allreduce(
     monkeypatch: pytest.MonkeyPatch,
     tp_size,

vllm/_custom_ops.py

@@ -2907,6 +2907,10 @@ def dispose(fa: int) -> None:
     torch.ops._C_custom_ar.dispose(fa)
 
 
+def custom_ar_close(fa: int) -> None:
+    torch.ops._C_custom_ar.custom_ar_close(fa)
+
+
 def meta_size() -> int:
     return torch.ops._C_custom_ar.meta_size()
 
@@ -2919,12 +2923,34 @@ def get_graph_buffer_ipc_meta(fa: int) -> tuple[list[int], list[int]]:
     return torch.ops._C_custom_ar.get_graph_buffer_ipc_meta(fa)
 
 
+def get_graph_buffer_ipc_meta_for_reinit(
+    fa: int,
+) -> tuple[list[int], list[int]]:
+    return torch.ops._C_custom_ar.get_graph_buffer_ipc_meta_for_reinit(fa)
+
+
 def register_graph_buffers(
     fa: int, handles: list[list[int]], offsets: list[list[int]]
 ) -> None:
     torch.ops._C_custom_ar.register_graph_buffers(fa, handles, offsets)
 
 
+def custom_ar_prepare_for_suspend(fa: int) -> None:
+    torch.ops._C_custom_ar.custom_ar_prepare_for_suspend(fa)
+
+
+def custom_ar_reinit_after_resume(
+    fa: int,
+    handles: list[list[int]],
+    offsets: list[list[int]],
+    signal_ptrs: list[int],
+    buffer_ptrs: list[int],
+) -> None:
+    torch.ops._C_custom_ar.custom_ar_reinit_after_resume(
+        fa, handles, offsets, signal_ptrs, buffer_ptrs
+    )
+
+
 def allocate_shared_buffer_and_handle(size: int) -> tuple[int, torch.Tensor]:
     return torch.ops._C_custom_ar.allocate_shared_buffer_and_handle(size)
 
@@ -2933,6 +2959,14 @@ def open_mem_handle(mem_handle: torch.Tensor):
     return torch.ops._C_custom_ar.open_mem_handle(mem_handle)
 
 
+def get_mem_handle(ptr: int) -> torch.Tensor:
+    return torch.ops._C_custom_ar.get_mem_handle(ptr)
+
+
+def close_mem_handle(ptr: int) -> None:
+    torch.ops._C_custom_ar.close_mem_handle(ptr)
+
+
 def free_shared_buffer(ptr: int) -> None:
     torch.ops._C_custom_ar.free_shared_buffer(ptr)
 

vllm/distributed/device_communicators/cuda_communicator.py

@@ -508,6 +508,12 @@ def destroy(self):
             self.pynccl_comm.destroy()
             self.pynccl_comm = None
         if self.ca_comm is not None:
+            try:
+                self.ca_comm.close()
+            except Exception:
+                logger.warning(
+                    "Failed to close custom allreduce communicator.", exc_info=True
+                )
             self.ca_comm = None
         if self.fi_ar_comm is not None:
             self.fi_ar_comm.destroy()