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 691d41cc3..3a047c8ff 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 @@ -47,6 +47,33 @@ if TYPE_CHECKING: from sglang.srt.model_executor.runner.shape_key import ShapeKey +def _allocate_output_buffer(output: Any) -> Optional[torch.Tensor]: + if not torch.is_tensor(output) or output.ndim == 0: + return None + return torch.empty_like(output) + + +def _output_fits_buffer(output: Any, output_buffer: torch.Tensor) -> bool: + return ( + torch.is_tensor(output) + and output.ndim == output_buffer.ndim + and output.shape[1:] == output_buffer.shape[1:] + and output.shape[0] <= output_buffer.shape[0] + and output.dtype == output_buffer.dtype + and output.device == output_buffer.device + ) + + +def _copy_output_to_buffer( + output: Any, output_buffer: torch.Tensor +) -> Optional[torch.Tensor]: + if not _output_fits_buffer(output, output_buffer): + return None + shared_output = output_buffer[: output.shape[0]] + shared_output.copy_(output) + return shared_output + + class FullCudaGraphBackend(BaseCudaGraphBackend): """One torch.cuda.CUDAGraph per shape; attention metadata is captured inside the graph. Memory-saver-aware. @@ -57,6 +84,7 @@ class FullCudaGraphBackend(BaseCudaGraphBackend): cuda_graph_runner: BaseCudaGraphRunner, *, enable_memory_saver: bool = False, + reuse_output_buffer: bool = False, ) -> None: self._graphs: Dict[Any, torch.cuda.CUDAGraph] = {} self._outputs: Dict[Any, Any] = {} @@ -66,6 +94,8 @@ class FullCudaGraphBackend(BaseCudaGraphBackend): self._tp_group = cuda_graph_runner.model_runner.tp_group self._capture_stream: Optional[torch.cuda.Stream] = None self._precarve = GraphPoolPrecarve() + self._reuse_output_buffer = reuse_output_buffer + self._output_buffer: Optional[torch.Tensor] = None self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create( enable=enable_memory_saver and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH") @@ -106,16 +136,27 @@ class FullCudaGraphBackend(BaseCudaGraphBackend): # Two warmups so kernels are loaded and one-time setup is paid before capture. # post_warmup_hook lets the attention backend reset state that warmup mutated. - for _ in range(2): + warmup_output = None + for warmup_step in range(2): self._device_module.synchronize() self._tp_group.barrier() with self._precarve.measure(): - forward_fn() + output = forward_fn() + if self._reuse_output_buffer and warmup_step == 1: + warmup_output = output + del output if profiler is not None: profiler.step() if post_warmup_hook is not None: post_warmup_hook() + if self._reuse_output_buffer and self._output_buffer is None: + # Prefill captures the largest shape first and replays one shape at + # a time, so all graphs can share this eager-tail input buffer. + self._output_buffer = _allocate_output_buffer(warmup_output) + self._reuse_output_buffer = self._output_buffer is not None + del warmup_output + graph = torch.cuda.CUDAGraph() graph_ctx: Callable[..., AbstractContextManager] @@ -136,6 +177,13 @@ class FullCudaGraphBackend(BaseCudaGraphBackend): ): self._precarve.mint() out = forward_fn() + if self._reuse_output_buffer: + output_buffer = self._output_buffer + assert output_buffer is not None + shared_output = _copy_output_to_buffer(out, output_buffer) + self._reuse_output_buffer = shared_output is not None + if shared_output is not None: + out = shared_output if profiler is not None: profiler.step() @@ -163,4 +211,5 @@ class FullCudaGraphBackend(BaseCudaGraphBackend): def cleanup(self) -> None: self._graphs.clear() self._outputs.clear() + self._output_buffer = None self._pool = None diff --git a/python/sglang/srt/model_executor/runner_backend/utils.py b/python/sglang/srt/model_executor/runner_backend/utils.py index 63c98b45c..1c68be8d6 100644 --- a/python/sglang/srt/model_executor/runner_backend/utils.py +++ b/python/sglang/srt/model_executor/runner_backend/utils.py @@ -122,6 +122,7 @@ def resolve_prefill_backend( return FullCudaGraphBackend( cuda_graph_runner, enable_memory_saver=get_exec().features.enable_memory_saver, + reuse_output_buffer=True, ) # Default: tc_piecewise. return TcPiecewiseCudaGraphBackend(cuda_graph_runner) diff --git a/test/registered/unit/model_executor/runner_backend/test_full_cuda_graph_backend.py b/test/registered/unit/model_executor/runner_backend/test_full_cuda_graph_backend.py index 32c767ff4..f653c8d9b 100644 --- a/test/registered/unit/model_executor/runner_backend/test_full_cuda_graph_backend.py +++ b/test/registered/unit/model_executor/runner_backend/test_full_cuda_graph_backend.py @@ -24,6 +24,8 @@ import unittest from types import SimpleNamespace from unittest import mock +import torch + from sglang.srt.model_executor.runner.shape_key import ShapeKey from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import ( FullCudaGraphBackend, @@ -59,6 +61,8 @@ def _make_backend(runner): backend._precarve = SimpleNamespace( measure=contextlib.nullcontext, mint=mock.Mock() ) + backend._reuse_output_buffer = False + backend._output_buffer = None backend._memory_saver_adapter = None backend._cuda_graph_runner = runner backend._device_module = runner.device_module @@ -107,6 +111,32 @@ class TestCaptureOneNoProfiling(CustomTestCase): self.assertEqual(backend._graphs[shape_key], "GRAPH") self.assertIs(backend._outputs[shape_key], sentinel_out) + def test_prefill_shapes_share_one_output_storage(self): + runner = _make_runner(enable_profile=False, profiler=None, mode_name="EXTEND") + backend = _make_backend(runner) + backend._reuse_output_buffer = True + + outputs = iter( + [ + torch.ones((4, 2)), + torch.ones((4, 2)), + torch.ones((4, 2)), + torch.ones((2, 2)), + torch.ones((2, 2)), + torch.ones((2, 2)), + ] + ) + with mock.patch("torch.cuda.CUDAGraph", side_effect=["GRAPH4", "GRAPH2"]): + backend.capture_one(ShapeKey(size=4), lambda: next(outputs)) + backend.capture_one(ShapeKey(size=2), lambda: next(outputs)) + + large = backend._outputs[ShapeKey(size=4)] + small = backend._outputs[ShapeKey(size=2)] + self.assertEqual(large.shape, (4, 2)) + self.assertEqual(small.shape, (2, 2)) + self.assertEqual(large.data_ptr(), small.data_ptr()) + self.assertEqual(backend._output_buffer.shape, (4, 2)) + def test_enable_flag_set_but_no_profiler_attr_does_not_step(self): # The runner advertises the flag but never created a profiler; the # getattr guard must keep capture_one on the non-profiling path.