[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)