[Refactor] Share CUDA graph memory pool across prefill and decode (#28973)

Co-authored-by: cctry <cctry@fb.com>

python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py

@@ -36,6 +36,9 @@ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import
     eager_on_graph,
     enable_breakable_cuda_graph,
 )
+from sglang.srt.model_executor.runner_utils.pool import (
+    get_or_create_global_graph_memory_pool,
+)
 from sglang.srt.utils import get_bool_env_var
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
 
@@ -82,7 +85,7 @@ class BreakableCudaGraphBackend(BaseCudaGraphBackend):
     @contextmanager
     def capture_session(self, stream: torch.cuda.Stream):
         if self._pool is None:
-            self._pool = self._device_module.graph_pool_handle()
+            self._pool = get_or_create_global_graph_memory_pool(self._device_module)
         set_graph_pool_id(self._pool)
         self._capture_stream = stream
         self._shared_output_buffer = None

python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py

@@ -30,6 +30,9 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
 from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
     BaseCudaGraphBackend,
 )
+from sglang.srt.model_executor.runner_utils.pool import (
+    get_or_create_global_graph_memory_pool,
+)
 from sglang.srt.utils import get_bool_env_var
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
 
@@ -66,7 +69,7 @@ class FullCudaGraphBackend(BaseCudaGraphBackend):
     @contextmanager
     def capture_session(self, stream: torch.cuda.Stream):
         if self._pool is None:
-            self._pool = self._device_module.graph_pool_handle()
+            self._pool = get_or_create_global_graph_memory_pool(self._device_module)
         set_graph_pool_id(self._pool)
         self._capture_stream = stream
         try:

python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py

@@ -46,6 +46,9 @@ from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     enable_tc_piecewise_cuda_graph,
 )
+from sglang.srt.model_executor.runner_utils.pool import (
+    get_or_create_global_graph_memory_pool,
+)
 from sglang.srt.utils import is_hip
 
 if TYPE_CHECKING:
@@ -158,7 +161,9 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
                 )
 
                 if self._pool is None:
-                    self._pool = self._device_module.graph_pool_handle()
+                    self._pool = get_or_create_global_graph_memory_pool(
+                        self._device_module
+                    )
                 set_graph_pool_id(self._pool)
 
                 self.install_compile(

python/sglang/srt/model_executor/runner_utils/pool.py

@@ -11,11 +11,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Shared graph memory pool used by the speculative-draft cuda graph
-runners. The new DecodeCudaGraphRunner and PrefillCudaGraphRunner
-backends each own their pool internally; this global is retained for the
-EAGLE / multi-step draft runners that haven't been folded into the new
-backend interface.
+"""Process-wide CUDA graph memory pool shared across the prefill and
+decode graph backends. The two phases never replay concurrently, so
+sharing one pool reserves only the larger phase's capture footprint.
 """
 
 from __future__ import annotations
@@ -32,3 +30,12 @@ def get_global_graph_memory_pool() -> Optional[Any]:
 def set_global_graph_memory_pool(val: Any) -> None:
     global _global_graph_memory_pool
     _global_graph_memory_pool = val
+
+
+def get_or_create_global_graph_memory_pool(device_module: Any) -> Any:
+    """Return the shared graph memory pool, creating it on first use so
+    later backends reuse the same handle."""
+    global _global_graph_memory_pool
+    if _global_graph_memory_pool is None:
+        _global_graph_memory_pool = device_module.graph_pool_handle()
+    return _global_graph_memory_pool