[Memory] Reuse output storage across full prefill CUDA graphs (#38038)

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

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)

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.