[NewFeature]custom_allreduce support cudagraph recapture (#4305)

* custom_allreduce support cudagraph recapture * add shut_down/restart default group
2025-10-06 09:07:10 +08:00 · 2025-09-29 15:56:54 +08:00
parent 7c919070f7
commit 81959c7d88
7 changed files with 32 additions and 3 deletions
--- a/custom_ops/gpu_ops/cpp_extensions.cc
+++ b/custom_ops/gpu_ops/cpp_extensions.cc
@@ -623,6 +623,8 @@ int64_t open_mem_handle(paddle::Tensor& mem_handle);

 void free_shared_buffer(int64_t buffer);

+void clear_ipc_handles(int64_t _fa);
+
 // speculative decoding Kernel
 std::vector<paddle::Tensor> SpeculateGetPaddingOffset(
    const paddle::Tensor& input_ids,
@@ -1229,6 +1231,8 @@ PYBIND11_MODULE(fastdeploy_ops, m) {

  m.def("free_shared_buffer", &free_shared_buffer, "free_shared_buffer");

+  m.def("clear_ipc_handles", &clear_ipc_handles, "clear_ipc_handles");
+
  m.def("open_mem_handle", &open_mem_handle, "open_mem_handle");

  m.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta, "get_graph_buffer_ipc_meta");