diff --git a/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py index 19ee9a2d8..5f495e349 100644 --- a/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py +++ b/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 diff --git a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py index 2e9292f17..e475ed9ca 100644 --- a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py +++ b/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: diff --git a/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py index d07354452..0054f0f2a 100644 --- a/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py +++ b/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( diff --git a/python/sglang/srt/model_executor/runner_utils/pool.py b/python/sglang/srt/model_executor/runner_utils/pool.py index 3268e6956..21fc49e65 100644 --- a/python/sglang/srt/model_executor/runner_utils/pool.py +++ b/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