minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[BCG] Share output buffers across capture sizes + typed ShapeKey (#27857)

Browse Source
This commit is contained in:
Yuwei An
2026-06-11 11:58:05 -07:00
committed by GitHub
parent 7f57b344c9
commit 880e6f66fc
12 changed files with 133 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py
+5 -4
View File
@@ -24,6 +24,7 @@ from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
set_graph_pool_id,
)
from sglang.srt.model_executor.runner.shape_key import ShapeKey
from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
BaseCudaGraphBackend,
)
@@ -75,7 +76,7 @@ class NPUCudaGraphBackend(BaseCudaGraphBackend):
def capture_one(
self,
shape_key: Any,
shape_key: ShapeKey,
forward_fn: Callable[[], Any],
dummies: Optional[Any] = None,
post_warmup_hook: Optional[Callable[[], None]] = None,
@@ -121,7 +122,7 @@ class NPUCudaGraphBackend(BaseCudaGraphBackend):
self._graphs[shape_key] = graph
self._outputs[shape_key] = out
def can_run(self, forward_batch: ForwardBatch, shape_key: Any) -> bool:
def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool:
return shape_key in self._graphs
@contextmanager
@@ -130,7 +131,7 @@ class NPUCudaGraphBackend(BaseCudaGraphBackend):
def replay(
self,
shape_key: Any,
shape_key: ShapeKey,
static_forward_batch: ForwardBatch,
**kwargs,
) -> Any:
@@ -139,7 +140,7 @@ class NPUCudaGraphBackend(BaseCudaGraphBackend):
def replay_with_input_update(
self,
shape_key: Any,
shape_key: ShapeKey,
seq_lens: Any,
attr_name: str = None,
attr_type: Any = None,
python/sglang/srt/model_executor/runner/__init__.py
+1 -3
View File
@@ -25,12 +25,10 @@ from sglang.srt.model_executor.runner.base_cuda_graph_runner import ( # noqa: F
from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
DecodeCudaGraphRunner,
)
from sglang.srt.model_executor.runner.decode_cuda_graph_runner import ( # noqa: F401
_make_graph_key as _default_make_graph_key,
)
from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import ( # noqa: F401
PrefillCudaGraphRunner,
)
from sglang.srt.model_executor.runner.shape_key import ShapeKey # noqa: F401
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( # noqa: F401
TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG,
)
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+6 -13
View File
@@ -73,6 +73,7 @@ from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
freeze_gc,
get_batch_sizes_to_capture,
)
from sglang.srt.model_executor.runner.shape_key import ShapeKey
from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import (
BreakableCudaGraphBackend,
)
@@ -114,18 +115,6 @@ if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
def _make_graph_key(bs, stream_idx=None, variant_label=None):
"""Build a graph dict key from batch size, stream index, and lora variant.
Standalone function so speculative runners (which don't subclass
DecodeCudaGraphRunner) can use the same key encoding.
"""
key = bs if stream_idx is None else f"{stream_idx}_{bs}"
if variant_label is not None:
key = f"{variant_label}_{key}"
return key
def build_replay_fb_view(
forward_batch: ForwardBatch,
buffers: DecodeInputBuffers,
@@ -499,7 +488,11 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
return torch.int64
def _make_graph_key(self, bs, stream_idx=None, variant_label=None):
return _make_graph_key(bs, stream_idx, variant_label)
return ShapeKey(
size=bs,
stream_idx=stream_idx,
variant_label=variant_label,
)
def _resolve_lora_variant(self, forward_batch: ForwardBatch):
if not getattr(self, "record_nolora_graph", False):
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+3 -2
View File
@@ -58,6 +58,7 @@ from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
BaseCudaGraphRunner,
freeze_gc,
)
from sglang.srt.model_executor.runner.shape_key import ShapeKey
from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import (
BreakableCudaGraphBackend,
)
@@ -658,7 +659,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
else:
post_warmup_hook = getattr(attn_backend, "on_after_cuda_graph_warmup", None)
self.backend.capture_one(
num_tokens,
ShapeKey(size=num_tokens),
run_once,
dummies=None,
post_warmup_hook=post_warmup_hook,
@@ -820,7 +821,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# model.forward eagerly with the live multi-req
# static_forward_batch. The outer's logits_processor /
# pooler then runs on top with live multi-req metadata.
shape_key = self._static_num_tokens
shape_key = ShapeKey(size=self._static_num_tokens)
def replay_layer_forward(*args, **layer_kwargs):
return self.backend.replay(
python/sglang/srt/model_executor/runner/shape_key.py
+23
View File
@@ -0,0 +1,23 @@
"""ShapeKey — typed identifier for one captured CUDA-graph shape."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
@dataclass(frozen=True)
class ShapeKey:
"""Identifies one captured CUDA-graph shape across all runners.
size: the per-phase capture size — what the runner iterates over.
- prefill: num_tokens
- decode: bs
stream_idx: pdmux stream index, or None for single-stream runners.
variant_label: LoRA-variant label ("lora" / "nolora"), or None
for runners that don't record per-variant graphs.
"""
size: int
stream_idx: Optional[int] = None
variant_label: Optional[str] = None
python/sglang/srt/model_executor/runner_backend/base_cuda_graph_backend.py
+5 -3
View File
@@ -7,6 +7,8 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, Optional
import torch
from sglang.srt.model_executor.runner.shape_key import ShapeKey
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -43,14 +45,14 @@ class BaseCudaGraphBackend(ABC):
@abstractmethod
def capture_one(
self,
shape_key: Any,
shape_key: ShapeKey,
forward_fn,
dummies: Optional[Any] = None,
post_warmup_hook: Optional[Callable[[], None]] = None,
) -> None: ...
@abstractmethod
def can_run(self, forward_batch: ForwardBatch, shape_key: Any) -> bool: ...
def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool: ...
@abstractmethod
def replay_session(self) -> Iterator[None]: ...
@@ -58,7 +60,7 @@ class BaseCudaGraphBackend(ABC):
@abstractmethod
def replay(
self,
shape_key: Any,
shape_key: ShapeKey,
static_forward_batch: ForwardBatch,
**kwargs,
) -> Any: ...
python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py
+75 -5
View File
@@ -13,6 +13,8 @@ import torch
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
set_graph_pool_id,
)
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
from sglang.srt.model_executor.runner.shape_key import ShapeKey
from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
BaseCudaGraphBackend,
)
@@ -53,6 +55,7 @@ class BreakableCudaGraphBackend(BaseCudaGraphBackend):
self._tp_group = cuda_graph_runner.model_runner.tp_group
self._capture_stream: Optional[torch.cuda.Stream] = None
self._debug_eager = debug_eager
self._shared_output_buffer: Optional[Any] = None
self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
enable=enable_memory_saver
and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
@@ -71,6 +74,7 @@ class BreakableCudaGraphBackend(BaseCudaGraphBackend):
self._pool = self._device_module.graph_pool_handle()
set_graph_pool_id(self._pool)
self._capture_stream = stream
self._shared_output_buffer = None
try:
with self.replay_session():
yield
@@ -79,7 +83,7 @@ class BreakableCudaGraphBackend(BaseCudaGraphBackend):
def capture_one(
self,
shape_key: Any,
shape_key: ShapeKey,
forward_fn: Callable[[], Any],
dummies: Optional[Any] = None,
post_warmup_hook: Optional[Callable[[], None]] = None,
@@ -95,16 +99,81 @@ class BreakableCudaGraphBackend(BaseCudaGraphBackend):
captured_fn = (
eager_on_graph(True)(forward_fn) if self._debug_eager else forward_fn
)
size = shape_key.size
with BreakableCUDAGraphCapture(
cuda_graph=graph,
pool=self._pool,
stream=self._capture_stream,
):
out = captured_fn()
self._graphs[shape_key] = graph
self._outputs[shape_key] = out
if self._shared_output_buffer is not None:
self._copy_output_to_buffer(out, self._shared_output_buffer, size)
def can_run(self, forward_batch: ForwardBatch, shape_key: Any) -> bool:
if self._shared_output_buffer is None:
self._shared_output_buffer = out
stored = self._slice_output(out, size)
else:
stored = self._slice_output(self._shared_output_buffer, size)
self._graphs[shape_key] = graph
self._outputs[shape_key] = stored
def _slice_output(self, output: Any, num_tokens: int) -> Any:
if output is None:
return None
if torch.is_tensor(output):
return output[:num_tokens]
if isinstance(output, PPProxyTensors):
return output[:num_tokens]
if isinstance(output, tuple):
return tuple(self._slice_output(item, num_tokens) for item in output)
if isinstance(output, list):
return [self._slice_output(item, num_tokens) for item in output]
raise TypeError(f"Unsupported BCG output type: {type(output)}")
def _copy_output_to_buffer(
self, output: Any, output_buffer: Any, num_tokens: int
) -> None:
if output is None or output_buffer is None:
if output is None and output_buffer is None:
return
raise ValueError(
"BCG output structure changed between capture sizes: "
f"{type(output)} vs {type(output_buffer)}"
)
if torch.is_tensor(output) and torch.is_tensor(output_buffer):
output_buffer[:num_tokens].copy_(output[:num_tokens])
return
if isinstance(output, PPProxyTensors) and isinstance(
output_buffer, PPProxyTensors
):
if output.tensors.keys() != output_buffer.tensors.keys():
raise ValueError(
"BCG output proxy structure changed between capture sizes: "
f"{output.tensors.keys()} != {output_buffer.tensors.keys()}"
)
for key, tensor in output.tensors.items():
self._copy_output_to_buffer(
tensor, output_buffer.tensors[key], num_tokens
)
return
if isinstance(output, (list, tuple)) and isinstance(
output_buffer, type(output)
):
if len(output) != len(output_buffer):
raise ValueError(
"BCG output sequence structure changed between capture sizes: "
f"{len(output)} != {len(output_buffer)}"
)
for item, buffer in zip(output, output_buffer):
self._copy_output_to_buffer(item, buffer, num_tokens)
return
raise TypeError(
"Unsupported BCG output buffer pair: "
f"{type(output)} vs {type(output_buffer)}"
)
def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool:
return shape_key in self._graphs
@contextmanager
@@ -114,7 +183,7 @@ class BreakableCudaGraphBackend(BaseCudaGraphBackend):
def replay(
self,
shape_key: Any,
shape_key: ShapeKey,
static_forward_batch: ForwardBatch,
**kwargs,
) -> Any:
@@ -125,3 +194,4 @@ class BreakableCudaGraphBackend(BaseCudaGraphBackend):
self._graphs.clear()
self._outputs.clear()
self._pool = None
self._shared_output_buffer = None
python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py
+4 -3
View File
@@ -14,6 +14,7 @@ from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
set_graph_pool_id,
)
from sglang.srt.model_executor.runner.shape_key import ShapeKey
from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
BaseCudaGraphBackend,
)
@@ -62,7 +63,7 @@ class FullCudaGraphBackend(BaseCudaGraphBackend):
def capture_one(
self,
shape_key: Any,
shape_key: ShapeKey,
forward_fn: Callable[[], Any],
dummies: Optional[Any] = None,
post_warmup_hook: Optional[Callable[[], None]] = None,
@@ -96,7 +97,7 @@ class FullCudaGraphBackend(BaseCudaGraphBackend):
self._graphs[shape_key] = graph
self._outputs[shape_key] = out
def can_run(self, forward_batch: ForwardBatch, shape_key: Any) -> bool:
def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool:
return shape_key in self._graphs
@contextmanager
@@ -105,7 +106,7 @@ class FullCudaGraphBackend(BaseCudaGraphBackend):
def replay(
self,
shape_key: Any,
shape_key: ShapeKey,
static_forward_batch: ForwardBatch,
**kwargs,
) -> Any:
python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py
+4 -3
View File
@@ -27,6 +27,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
)
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
from sglang.srt.layers.utils import MultiPlatformOp
from sglang.srt.model_executor.runner.shape_key import ShapeKey
from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
BaseCudaGraphBackend,
)
@@ -193,7 +194,7 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
def capture_one(
self,
shape_key: Any,
shape_key: ShapeKey,
forward_fn: Callable[[], Any],
dummies: Optional[Any] = None,
post_warmup_hook: Optional[Callable[[], None]] = None,
@@ -207,7 +208,7 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
if post_warmup_hook is not None:
post_warmup_hook()
def can_run(self, forward_batch: ForwardBatch, shape_key: Any) -> bool:
def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool:
# torch.compile manages its per-shape cache internally.
# _run_compile_pass warms every shape in capture_num_tokens at __init__.
return True
@@ -219,7 +220,7 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
def replay(
self,
shape_key: Any,
shape_key: ShapeKey,
static_forward_batch: ForwardBatch,
**kwargs,
) -> Any:
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
+3 -2
View File
@@ -23,6 +23,7 @@ from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
from sglang.srt.model_executor.runner import (
DecodeCudaGraphRunner,
DeepEPCudaGraphRunnerAdapter,
ShapeKey,
get_batch_sizes_to_capture,
model_capture_mode,
)
@@ -247,8 +248,8 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
return torch.int64
def _make_graph_key(self, bs, stream_idx=None, variant_label=None):
# EAGLE doesn't use stream_idx / lora variants; key is just bs.
return bs
# EAGLE doesn't use stream_idx / lora variants.
return ShapeKey(size=bs)
# -----------------------------------------------------------------
# can_run
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
+2 -1
View File
@@ -23,6 +23,7 @@ from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
from sglang.srt.model_executor.runner import (
DecodeCudaGraphRunner,
DeepEPCudaGraphRunnerAdapter,
ShapeKey,
get_batch_sizes_to_capture,
model_capture_mode,
)
@@ -259,7 +260,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
return torch.int64
def _make_graph_key(self, bs, stream_idx=None, variant_label=None):
return bs
return ShapeKey(size=bs)
def can_run(self, forward_batch: ForwardBatch):
if self.require_mlp_tp_gather:
python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
+2 -1
View File
@@ -47,6 +47,7 @@ from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
from sglang.srt.model_executor.runner import (
DecodeCudaGraphRunner,
DeepEPCudaGraphRunnerAdapter,
ShapeKey,
get_batch_sizes_to_capture,
model_capture_mode,
)
@@ -286,7 +287,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
return self.backend.replay(shape_key, forward_batch)
def _make_graph_key(self, bs, stream_idx=None, variant_label=None):
return bs
return ShapeKey(size=bs)
def can_run(self, forward_batch: ForwardBatch):
if self.require_mlp_tp_gather: