share pool id amongst cuda graphs (#73075)

lixinqi · web-flow · commit 35d7f4c39556 · 2025-06-04T20:14:48.000+08:00
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
@@ -1265,6 +1265,19 @@ PYBIND11_MODULE(libpaddle, m) {
                     platform::BeginCUDAGraphCapture(
                         place, static_cast<paddle::gpuStreamCaptureMode>(mode));
                   })
+      .def_static(
+          "begin_capture_with_pool_id",
+          [](phi::GPUPlace place, int mode, std::optional<int64_t> pool_id) {
+            if (pool_id.has_value()) {
+              platform::BeginCUDAGraphCapture(
+                  place,
+                  static_cast<paddle::gpuStreamCaptureMode>(mode),
+                  pool_id.value());
+            } else {
+              platform::BeginCUDAGraphCapture(
+                  place, static_cast<paddle::gpuStreamCaptureMode>(mode));
+            }
+          })
       .def_static("end_capture", &platform::EndCUDAGraphCapture)
       .def_static("gen_new_memory_pool_id",
                   &phi::backends::gpu::CUDAGraph::UniqueMemoryPoolID)
diff --git a/python/paddle/device/cuda/graphs.py b/python/paddle/device/cuda/graphs.py
@@ -41,7 +41,7 @@ def is_cuda_graph_supported():
 
 
 class CUDAGraph:
-    def __init__(self, place=None, mode="thread_local"):
+    def __init__(self, place=None, mode="thread_local", pool_id=None):
         assert (
             CoreCUDAGraph is not None
         ), "CUDA Graph is only supported on PaddlePaddle compiled with NVIDIA GPU."
@@ -53,9 +53,12 @@ def __init__(self, place=None, mode="thread_local"):
         self._place = place
         assert mode in ALL_MODES
         self._mode = ALL_MODES.index(mode)
+        self._pool_id = pool_id
 
     def capture_begin(self):
-        CoreCUDAGraph.begin_capture(self._place, self._mode)
+        CoreCUDAGraph.begin_capture_with_pool_id(
+            self._place, self._mode, self._pool_id
+        )
 
     def capture_end(self):
         self._graph = CoreCUDAGraph.end_capture()
