chore: import upstream snapshot with attribution

python/sglang/srt/model_executor/runner_backend/__init__.py

@@ -0,0 +1,31 @@
+"""Capture-mechanism backends for CUDA graphs.
+
+A backend owns *how* a captured artifact is produced and replayed for
+one shape; it is phase-agnostic. Runners (cuda_graph_runner/) own
+*what* data flows in and out.
+
+Public API:
+  - BaseCudaGraphBackend — abstract interface.
+  - FullCudaGraphBackend — single torch.cuda.CUDAGraph per shape.
+  - BreakableCudaGraphBackend — segmented capture with eager break
+    markers; no torch.compile.
+  - TcPiecewiseCudaGraphBackend — torch.compile-driven piecewise
+    capture; FX-splits the model at attention layers.
+"""
+
+from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (  # noqa: F401
+    BaseCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import (  # noqa: F401
+    BreakableCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import (  # noqa: F401
+    FullCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.tc_piecewise_cuda_graph_backend import (  # noqa: F401
+    TcPiecewiseCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.utils import (  # noqa: F401
+    resolve_decode_backend,
+    resolve_prefill_backend,
+)

python/sglang/srt/model_executor/runner_backend/base_cuda_graph_backend.py

@@ -0,0 +1,81 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Backend interface for CUDA graph capture/replay."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Callable, Iterator, Optional
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.runner.shape_key import ShapeKey
+
+
+class BaseCudaGraphBackend(ABC):
+    """Pure ABC: no state, no defaults. Each implementation owns its
+    per-backend state and binds the handles it needs from the
+    cuda_graph_runner passed to its __init__.
+
+    Methods:
+      - capture_session(stream) — context wrapping the runner's outer
+        capture loop; backends bind stream / pool and open per-backend
+        capture flags here.
+      - capture_one(shape_key, forward_fn, dummies, post_warmup_hook)
+        — record the replayable artifact for shape_key; one call per
+        shape inside capture_session.
+      - can_run(forward_batch, shape_key) — can this backend replay
+        for the given batch at the given shape.
+      - replay_session() — context wrapping replay-time model code;
+        backends open the "we are replaying" flag here when they have
+        one.
+      - replay(shape_key, static_forward_batch, **kwargs) — invoke
+        the captured artifact.
+      - cleanup() — release pool and drop captured artifacts.
+
+    Notes:
+      - The outer capture loop is runner-specific; it lives on the
+        runner, not here.
+    """
+
+    @abstractmethod
+    def capture_session(self, stream: torch.cuda.Stream) -> Iterator[None]: ...
+
+    @abstractmethod
+    def capture_one(
+        self,
+        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: ShapeKey) -> bool: ...
+
+    @abstractmethod
+    def replay_session(self) -> Iterator[None]: ...
+
+    @abstractmethod
+    def replay(
+        self,
+        shape_key: ShapeKey,
+        static_forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any: ...
+
+    @abstractmethod
+    def cleanup(self) -> None: ...

python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py

@@ -0,0 +1,252 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""BreakableCudaGraphBackend — segment-captured graphs with eager break
+markers (eager_on_graph decorators on attention / mamba layers).
+No torch.compile.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+
+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_backend.base_cuda_graph_backend import (
+    BaseCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.cuda_graph_dedup_mixin import (
+    DedupedCudaGraphMixin,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
+    BreakableCUDAGraph,
+    BreakableCUDAGraphCapture,
+    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
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
+        BaseCudaGraphRunner,
+    )
+    from sglang.srt.model_executor.runner.shape_key import ShapeKey
+
+
+class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
+    """Segmented capture: graphs break at attention / mamba boundaries;
+    attention metadata is recomputed at replay outside captured segments.
+    """
+
+    def __init__(
+        self,
+        cuda_graph_runner: BaseCudaGraphRunner,
+        *,
+        enable_memory_saver: bool = False,
+        debug_eager: bool = False,
+    ) -> None:
+        self._model_runner = cuda_graph_runner.model_runner
+        self._graphs: Dict[Any, BreakableCUDAGraph] = {}
+        self._outputs: Dict[Any, Any] = {}
+        self._pool = None
+        self._device_module = cuda_graph_runner.device_module
+        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")
+        )
+        if (
+            self._memory_saver_adapter is not None
+            and self._memory_saver_adapter.enabled
+        ):
+            raise NotImplementedError(
+                "Breakable CUDA graph is not compatible with memory saver mode"
+            )
+
+    @contextmanager
+    def capture_session(self, stream: torch.cuda.Stream):
+        if self._pool is None:
+            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
+        self.begin_cuda_graph_capture()
+        try:
+            with self.replay_session():
+                yield
+        finally:
+            try:
+                self.end_cuda_graph_capture()
+            finally:
+                self._capture_stream = None
+
+    def capture_one(
+        self,
+        shape_key: ShapeKey,
+        forward_fn: Callable[[], Any],
+        dummies: Optional[Any] = None,
+        post_warmup_hook: Optional[Callable[[], None]] = None,
+    ) -> None:
+        warmup_out = None
+        for _ in range(2):
+            self._device_module.synchronize()
+            self._tp_group.barrier()
+            warmup_out = forward_fn()
+            if post_warmup_hook is not None:
+                post_warmup_hook()
+
+        graph = BreakableCUDAGraph()
+        captured_fn = (
+            eager_on_graph(True)(forward_fn) if self._debug_eager else forward_fn
+        )
+        size = shape_key.size
+        if self._shared_output_buffer is None:
+            self._shared_output_buffer = self._alloc_full_buffer(warmup_out, size)
+        with BreakableCUDAGraphCapture(
+            cuda_graph=graph,
+            pool=self._pool,
+            stream=self._capture_stream,
+        ):
+            out = captured_fn()
+            out_rows = self._output_rows(out, size)
+            self._copy_output_to_buffer(out, self._shared_output_buffer, out_rows)
+
+        stored = self._slice_output(self._shared_output_buffer, out_rows)
+        self._graphs[shape_key] = graph
+        self._outputs[shape_key] = stored
+
+    def _output_rows(self, output: Any, cap: int) -> int:
+        """Leading-dim row count actually produced by the body, clamped to ``cap``.
+
+        A body that shards or prunes its output along dim 0 returns fewer than
+        ``cap`` rows; everything else returns exactly ``cap``.
+        """
+        if torch.is_tensor(output):
+            return min(cap, output.shape[0])
+        if isinstance(output, PPProxyTensors):
+            rows = [t.shape[0] for t in output.tensors.values()]
+            return min([cap, *rows])
+        if isinstance(output, (list, tuple)) and output:
+            return min(self._output_rows(o, cap) for o in output if o is not None)
+        return cap
+
+    def _alloc_full_buffer(self, output: Any, size: int) -> Any:
+        """A same-structure buffer as ``output`` but with ``size`` leading rows."""
+        if output is None:
+            return None
+        if torch.is_tensor(output):
+            return output.new_empty((size, *output.shape[1:]))
+        if isinstance(output, PPProxyTensors):
+            return PPProxyTensors(
+                {
+                    key: t.new_empty((size, *t.shape[1:]))
+                    for key, t in output.tensors.items()
+                }
+            )
+        if isinstance(output, tuple):
+            return tuple(self._alloc_full_buffer(o, size) for o in output)
+        if isinstance(output, list):
+            return [self._alloc_full_buffer(o, size) for o in output]
+        raise TypeError(f"Unsupported BCG output type: {type(output)}")
+
+    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
+    def replay_session(self):
+        with enable_breakable_cuda_graph():
+            yield
+
+    def replay(
+        self,
+        shape_key: ShapeKey,
+        static_forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any:
+        self._graphs[shape_key].replay()
+        return self._outputs[shape_key]
+
+    def cleanup(self) -> None:
+        self.close()
+        self._graphs.clear()
+        self._outputs.clear()
+        self._pool = None
+        self._shared_output_buffer = None

python/sglang/srt/model_executor/runner_backend/cuda_graph_dedup_mixin.py

@@ -0,0 +1,375 @@
+"""Shared CUDA graph executable-dedup plumbing for CUDA graph backends."""
+
+from __future__ import annotations
+
+import heapq
+import logging
+from dataclasses import dataclass, field
+
+import torch
+
+try:
+    from cuda.bindings import driver as cuda_drv
+    from cuda.bindings import runtime as cuda_rt
+except ImportError:
+    cuda_drv = None
+    cuda_rt = None
+
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.cuda_utils import (
+    checkCudaErrors,
+)
+from sglang.srt.utils import get_bool_env_var
+
+logger = logging.getLogger(__name__)
+
+
+def dedup_update(graph_exec: int, raw_graph: int) -> tuple[bool, str]:
+    assert cuda_rt is not None
+    err, info = cuda_rt.cudaGraphExecUpdate(graph_exec, raw_graph)
+    if info is None:
+        return False, f"err={int(err)}"
+    result = info.result
+    ok = (
+        err == cuda_rt.cudaError_t.cudaSuccess
+        and result == cuda_rt.cudaGraphExecUpdateResult.cudaGraphExecUpdateSuccess
+    )
+    return ok, "" if ok else f"err={int(err)} result={result}"
+
+
+def maybe_cuda_result(result):
+    return None if int(result[0]) != 0 else checkCudaErrors(result)
+
+
+def kernel_name(params) -> str:
+    assert cuda_drv is not None
+    for handle, getter in (
+        (getattr(params, "kern", None), cuda_drv.cuKernelGetName),
+        (getattr(params, "func", None), cuda_drv.cuFuncGetName),
+    ):
+        if handle is None or int(handle) == 0:
+            continue
+        name = maybe_cuda_result(getter(handle))
+        if name is not None:
+            return name.decode("utf-8", "replace")
+    return f"func:{int(getattr(params, 'func', 0))}"
+
+
+def kernel_attrs(node) -> tuple[tuple[str, object], ...]:
+    assert cuda_drv is not None
+    attrs = []
+    for name, attr_name, get_value in (
+        (
+            "cooperative",
+            "CU_LAUNCH_ATTRIBUTE_COOPERATIVE",
+            lambda v: int(v.cooperative),
+        ),
+        (
+            "clusterDim",
+            "CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION",
+            lambda v: (
+                int(v.clusterDim.x),
+                int(v.clusterDim.y),
+                int(v.clusterDim.z),
+            ),
+        ),
+        (
+            "clusterSchedulingPolicyPreference",
+            "CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE",
+            lambda v: int(v.clusterSchedulingPolicyPreference),
+        ),
+        (
+            "preferredClusterDim",
+            "CU_LAUNCH_ATTRIBUTE_PREFERRED_CLUSTER_DIMENSION",
+            lambda v: (
+                int(v.preferredClusterDim.x),
+                int(v.preferredClusterDim.y),
+                int(v.preferredClusterDim.z),
+            ),
+        ),
+        (
+            "sharedMemCarveout",
+            "CU_LAUNCH_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT",
+            lambda v: int(v.sharedMemCarveout),
+        ),
+    ):
+        attr = getattr(cuda_drv.CUkernelNodeAttrID, attr_name, None)
+        if attr is None:
+            continue
+        value = maybe_cuda_result(cuda_drv.cuGraphKernelNodeGetAttribute(node, attr))
+        if value is not None:
+            attrs.append((name, get_value(value)))
+    return tuple(attrs)
+
+
+def kernel_node_payload(node):
+    assert cuda_drv is not None
+    params = checkCudaErrors(cuda_drv.cuGraphKernelNodeGetParams(node))
+    return (
+        kernel_name(params),
+        (int(params.gridDimX), int(params.gridDimY), int(params.gridDimZ)),
+        (int(params.blockDimX), int(params.blockDimY), int(params.blockDimZ)),
+        int(params.sharedMemBytes),
+        kernel_attrs(node),
+    )
+
+
+def graph_node_payload(node):
+    assert cuda_drv is not None
+    node_type = checkCudaErrors(cuda_drv.cuGraphNodeGetType(node))
+    match node_type:
+        case cuda_drv.CUgraphNodeType.CU_GRAPH_NODE_TYPE_KERNEL:
+            payload = kernel_node_payload(node)
+        case cuda_drv.CUgraphNodeType.CU_GRAPH_NODE_TYPE_MEMCPY:
+            params = checkCudaErrors(cuda_drv.cuGraphMemcpyNodeGetParams(node))
+            payload = (int(params.srcMemoryType), int(params.dstMemoryType))
+        case cuda_drv.CUgraphNodeType.CU_GRAPH_NODE_TYPE_MEMSET:
+            params = checkCudaErrors(cuda_drv.cuGraphMemsetNodeGetParams(node))
+            payload = (int(params.elementSize),)
+        case cuda_drv.CUgraphNodeType.CU_GRAPH_NODE_TYPE_GRAPH:
+            child_graph = checkCudaErrors(cuda_drv.cuGraphChildGraphNodeGetGraph(node))
+            payload = graph_signature(child_graph)
+        case cuda_drv.CUgraphNodeType.CU_GRAPH_NODE_TYPE_EMPTY:
+            payload = ()
+        case _:
+            payload = ()
+    return (node_type.name, payload)
+
+
+def graph_signature(raw_graph: int):
+    assert cuda_drv is not None
+    _, num_nodes = checkCudaErrors(cuda_drv.cuGraphGetNodes(raw_graph, 0))
+    nodes, _ = checkCudaErrors(cuda_drv.cuGraphGetNodes(raw_graph, num_nodes))
+    node_indices = {int(node): i for i, node in enumerate(nodes)}
+
+    _, _, _, num_edges = checkCudaErrors(cuda_drv.cuGraphGetEdges(raw_graph, 0))
+    from_nodes, to_nodes, _, _ = checkCudaErrors(
+        cuda_drv.cuGraphGetEdges(raw_graph, num_edges)
+    )
+    edges = [
+        (node_indices[int(src)], node_indices[int(dst)])
+        for src, dst in zip(from_nodes, to_nodes)
+    ]
+
+    children = [[] for _ in nodes]
+    indegree = [0] * len(nodes)
+    for src, dst in edges:
+        children[src].append(dst)
+        indegree[dst] += 1
+
+    ready = [i for i, degree in enumerate(indegree) if degree == 0]
+    heapq.heapify(ready)
+    order = []
+    while ready:
+        node_idx = heapq.heappop(ready)
+        order.append(node_idx)
+        for child_idx in sorted(children[node_idx]):
+            indegree[child_idx] -= 1
+            if indegree[child_idx] == 0:
+                heapq.heappush(ready, child_idx)
+    assert len(order) == len(nodes), "CUDA graph contains a dependency cycle"
+
+    topo_indices = {node_idx: i for i, node_idx in enumerate(order)}
+    topo_edges = tuple(
+        sorted((topo_indices[src], topo_indices[dst]) for src, dst in edges)
+    )
+    return (
+        tuple(graph_node_payload(nodes[node_idx]) for node_idx in order),
+        topo_edges,
+    )
+
+
+@dataclass(slots=True)
+class GraphExecGroup:
+    graph_exec: int
+    current_raw_graph: int
+    compat_exec: int | None
+    graphs: list[DedupedCudaGraph] = field(default_factory=list)
+
+
+@dataclass(eq=False, slots=True)
+class DedupedCudaGraph:
+    raw_graph: int
+    original_graph: object | None
+    registry: DedupedCudaGraphRegistry
+    group: GraphExecGroup | None = None
+
+    def replay(self, stream: int | None = None) -> None:
+        if stream is None:
+            stream = torch.cuda.current_stream().cuda_stream
+        self.registry.replay(self, stream)
+
+
+class DedupedCudaGraphRegistry:
+    def __init__(self):
+        self.groups: dict[tuple, GraphExecGroup] = {}
+        self.sealed = False
+
+    def instantiate(self, raw_graph: int) -> int:
+        assert cuda_rt is not None
+        graph_exec = checkCudaErrors(
+            cuda_rt.cudaGraphInstantiateWithFlags(raw_graph, 0)
+        )
+        return graph_exec
+
+    def destroy_exec(self, graph_exec: int) -> None:
+        assert cuda_rt is not None
+        checkCudaErrors(cuda_rt.cudaGraphExecDestroy(graph_exec))
+
+    def register(self, captured_graph) -> DedupedCudaGraph:
+        assert not self.sealed
+        raw_graph = captured_graph.raw_cuda_graph()
+        signature = graph_signature(raw_graph)
+        graph = DedupedCudaGraph(raw_graph, captured_graph, self)
+
+        group = self.groups.get(signature)
+        if group is not None:
+            assert group.compat_exec is not None
+            ok, detail = dedup_update(group.compat_exec, graph.raw_graph)
+            assert ok, f"CUDA graph dedup register update failed ({detail})"
+            graph.group = group
+            group.graphs.append(graph)
+            return graph
+
+        group = GraphExecGroup(
+            graph_exec=self.instantiate(graph.raw_graph),
+            current_raw_graph=graph.raw_graph,
+            compat_exec=self.instantiate(graph.raw_graph),
+            graphs=[graph],
+        )
+        graph.group = group
+        self.groups[signature] = group
+        return graph
+
+    def seal(self) -> None:
+        if self.sealed:
+            return
+        self.sealed = True
+        for group in self.groups.values():
+            if group.compat_exec is not None:
+                self.destroy_exec(group.compat_exec)
+                group.compat_exec = None
+
+    def stats(self) -> tuple[int, int]:
+        return sum(len(group.graphs) for group in self.groups.values()), len(
+            self.groups
+        )
+
+    def replay(self, graph: DedupedCudaGraph, stream: int) -> None:
+        assert cuda_rt is not None
+        group = graph.group
+        assert (
+            group is not None
+        ), "captured CUDA graph does not belong to this dedup state"
+
+        raw_graph = graph.raw_graph
+        graph_exec = group.graph_exec
+        if group.current_raw_graph != raw_graph:
+            ok, detail = dedup_update(graph_exec, raw_graph)
+            assert ok, (
+                "CUDA graph dedup replay update failed "
+                f"({detail}); captured graph is not compatible with its dedup group"
+            )
+            group.current_raw_graph = raw_graph
+
+        checkCudaErrors(cuda_rt.cudaGraphLaunch(graph_exec, stream))
+
+    def close(self) -> None:
+        self.sealed = True
+
+        for group in self.groups.values():
+            if group.compat_exec is not None:
+                self.destroy_exec(group.compat_exec)
+                group.compat_exec = None
+            self.destroy_exec(group.graph_exec)
+            for graph in group.graphs:
+                if graph.original_graph is not None:
+                    graph.original_graph.reset()
+                graph.original_graph = None
+                graph.group = None
+            group.graphs.clear()
+
+        self.groups.clear()
+
+
+class DedupedCudaGraphMixin:
+    deduped_cuda_graph: DedupedCudaGraphRegistry | None = None
+
+    def _dedup_registries(self) -> list[DedupedCudaGraphRegistry]:
+        registries = getattr(self, "_deduped_cuda_graph_registries", None)
+        if registries is None:
+            registries = []
+            self._deduped_cuda_graph_registries = registries
+        return registries
+
+    def _memory_saver_cuda_graph_enabled(self) -> bool:
+        adapter = getattr(self, "_memory_saver_adapter", None)
+        if adapter is not None and getattr(adapter, "enabled", False):
+            return True
+
+        model_runner = getattr(self, "model_runner", None)
+        if model_runner is None:
+            model_runner = getattr(self, "_model_runner", None)
+        server_args = getattr(model_runner, "server_args", None)
+        return bool(
+            server_args is not None
+            and getattr(server_args, "enable_memory_saver", False)
+            and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
+        )
+
+    def build_deduped_cuda_graph(self):
+        if not envs.SGLANG_ENABLE_CUDA_GRAPH_DEDUP.get():
+            return None
+        if cuda_drv is None or cuda_rt is None:
+            return None
+        try:
+            graph = torch.cuda.CUDAGraph(keep_graph=True)
+            if not hasattr(graph, "raw_cuda_graph"):
+                return None
+            return DedupedCudaGraphRegistry()
+        except TypeError:
+            return None
+        except Exception as e:
+            logger.warning(
+                "[CudaGraph][dedup] %s init failed (%s); using plain executables.",
+                type(self).__name__,
+                e,
+            )
+            return None
+
+    def begin_cuda_graph_capture(self) -> None:
+        if self.deduped_cuda_graph is not None:
+            self.end_cuda_graph_capture()
+
+        if self._memory_saver_cuda_graph_enabled():
+            self.deduped_cuda_graph = None
+            return
+
+        self.deduped_cuda_graph = self.build_deduped_cuda_graph()
+        if self.deduped_cuda_graph is not None:
+            self._dedup_registries().append(self.deduped_cuda_graph)
+
+    def end_cuda_graph_capture(self) -> None:
+        dedup = self.deduped_cuda_graph
+        self.deduped_cuda_graph = None
+        if dedup is not None:
+            captured, execs = dedup.stats()
+            dedup.seal()
+            logger.info("captured %d CUDA graphs, deduped to %d execs", captured, execs)
+
+    def close(self) -> None:
+        registries = self._dedup_registries()
+        seen: set[int] = set()
+        for registry in [self.deduped_cuda_graph, *registries]:
+            if registry is None or id(registry) in seen:
+                continue
+            seen.add(id(registry))
+            registry.close()
+        registries.clear()
+        self.deduped_cuda_graph = None
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            pass

python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py

@@ -0,0 +1,135 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""FullCudaGraphBackend — captures the entire model forward as one
+torch.cuda.CUDAGraph per shape.
+"""
+
+from __future__ import annotations
+
+from contextlib import AbstractContextManager, contextmanager
+from functools import partial
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+
+import torch
+
+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_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
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
+        BaseCudaGraphRunner,
+    )
+    from sglang.srt.model_executor.runner.shape_key import ShapeKey
+
+
+class FullCudaGraphBackend(BaseCudaGraphBackend):
+    """One torch.cuda.CUDAGraph per shape; attention metadata is
+    captured inside the graph. Memory-saver-aware.
+    """
+
+    def __init__(
+        self,
+        cuda_graph_runner: BaseCudaGraphRunner,
+        *,
+        enable_memory_saver: bool = False,
+    ) -> None:
+        self._graphs: Dict[Any, torch.cuda.CUDAGraph] = {}
+        self._outputs: Dict[Any, Any] = {}
+        self._pool = None
+        self._device_module = cuda_graph_runner.device_module
+        self._tp_group = cuda_graph_runner.model_runner.tp_group
+        self._capture_stream: Optional[torch.cuda.Stream] = None
+        self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
+            enable=enable_memory_saver
+            and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
+        )
+
+    @contextmanager
+    def capture_session(self, stream: torch.cuda.Stream):
+        if self._pool is None:
+            self._pool = get_or_create_global_graph_memory_pool(self._device_module)
+        set_graph_pool_id(self._pool)
+        self._capture_stream = stream
+        try:
+            yield
+        finally:
+            self._capture_stream = None
+
+    def capture_one(
+        self,
+        shape_key: ShapeKey,
+        forward_fn: Callable[[], Any],
+        dummies: Optional[Any] = None,
+        post_warmup_hook: Optional[Callable[[], None]] = None,
+    ) -> None:
+        # 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):
+            self._device_module.synchronize()
+            self._tp_group.barrier()
+            forward_fn()
+            if post_warmup_hook is not None:
+                post_warmup_hook()
+
+        graph = torch.cuda.CUDAGraph()
+
+        graph_ctx: Callable[..., AbstractContextManager]
+        if (
+            self._memory_saver_adapter is not None
+            and self._memory_saver_adapter.enabled
+        ):
+            graph_ctx = partial(
+                self._memory_saver_adapter.cuda_graph,
+                tag=GPU_MEMORY_TYPE_CUDA_GRAPH,
+            )
+        else:
+            graph_ctx = self._device_module.graph
+
+        with graph_ctx(cuda_graph=graph, pool=self._pool, stream=self._capture_stream):
+            out = forward_fn()
+
+        self._graphs[shape_key] = graph
+        self._outputs[shape_key] = out
+
+    def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool:
+        return shape_key in self._graphs
+
+    @contextmanager
+    def replay_session(self):
+        yield
+
+    def replay(
+        self,
+        shape_key: ShapeKey,
+        static_forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any:
+        self._graphs[shape_key].replay()
+        return self._outputs[shape_key]
+
+    def cleanup(self) -> None:
+        self._graphs.clear()
+        self._outputs.clear()
+        self._pool = None

python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py

@@ -0,0 +1,259 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""TcPiecewiseCudaGraphBackend — torch.compile-driven piecewise CUDA graph.
+
+FX-splits the model forward at attention layers; per-shape compiled
+callables internally capture sub-graphs via
+compilation/cuda_piecewise_backend. torch.compile owns the per-shape
+cache so this backend has no _graphs table — only a single
+_compiled_fn reused for every shape.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+import torch
+import tqdm
+
+from sglang.srt.compilation.compilation_config import CompilationConfig
+from sglang.srt.compilation.compile import install_torch_compiled
+from sglang.srt.compilation.compile_phase import (
+    enable_torch_compile_warmup,
+    set_pcg_capture_stream,
+)
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    set_graph_pool_id,
+)
+from sglang.srt.layers.moe.utils import get_moe_a2a_backend
+from sglang.srt.layers.utils import MultiPlatformOp
+from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
+    BaseCudaGraphBackend,
+)
+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.runtime_context import get_parallel
+from sglang.srt.utils import is_hip
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
+        BaseCudaGraphRunner,
+    )
+    from sglang.srt.model_executor.runner.shape_key import ShapeKey
+    from sglang.srt.server_args import ServerArgs
+
+
+_VALID_COMPILERS = ("eager", "inductor")
+
+
+def _toggle_multi_platform_ops(
+    model: torch.nn.Module, *, reverse: bool, num_tokens: int
+) -> None:
+    """Recursively flip MultiPlatformOp submodules into / out of
+    torch.compile mode."""
+    for sub in model._modules.values():
+        if isinstance(sub, MultiPlatformOp):
+            if reverse:
+                sub.leave_torch_compile()
+            else:
+                sub.enter_torch_compile(num_tokens=num_tokens)
+        if isinstance(sub, torch.nn.Module):
+            _toggle_multi_platform_ops(sub, reverse=reverse, num_tokens=num_tokens)
+
+
+class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
+    """torch.compile-driven piecewise capture; attention metadata
+    recomputed at replay outside the compiled callable's sub-graphs.
+    """
+
+    def __init__(self, cuda_graph_runner: BaseCudaGraphRunner) -> None:
+        model_runner = cuda_graph_runner.model_runner
+        self._pool = None
+        self._device_module = cuda_graph_runner.device_module
+        self._tp_group = model_runner.tp_group
+        self._capture_stream: Optional[torch.cuda.Stream] = None
+        self._compile_config: CompilationConfig = self.build_compilation_config(
+            model_runner.server_args
+        )
+        self._language_model: torch.nn.Module = getattr(
+            model_runner.model, "language_model", model_runner.model
+        )
+        self._run_compile_pass(cuda_graph_runner)
+        # model_runner.model.forward is the wrapper that builds LogitsProcessorOutput.
+        # The compiled trampoline is dispatched internally by it.
+        self._compiled_fn: Callable = model_runner.model.forward
+
+    @staticmethod
+    def build_compilation_config(server_args: ServerArgs) -> CompilationConfig:
+        """Construct a CompilationConfig from ServerArgs and
+        register the MoE A2A split-op when DeepEP / Mooncake is in use."""
+        prefill = server_args.cuda_graph_config.prefill
+        num_tokens = prefill.bs
+        compiler = prefill.tc_compiler
+        assert num_tokens is not None, "cuda_graph_config[prefill].bs is not set"
+        assert compiler in _VALID_COMPILERS, (
+            f"By now, only {_VALID_COMPILERS} are supported for the "
+            "tc_piecewise prefill compiler."
+        )
+
+        config = CompilationConfig(
+            num_tokens,
+            compiler,
+            server_args.enable_torch_compile_debug_mode,
+        )
+
+        if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
+            config.add_split_op("sglang.moe_forward_piecewise_cuda_graph_impl")
+
+        return config
+
+    @staticmethod
+    def install_compile(
+        language_model: Any,
+        *,
+        compile_config: CompilationConfig,
+        graph_pool: Any,
+        fullgraph: bool = True,
+        dynamic_arg_dims: Optional[Any] = None,
+    ) -> None:
+        """Wrap language_model.model.forward with torch.compile."""
+        install_torch_compiled(
+            language_model,
+            fullgraph=fullgraph,
+            dynamic_arg_dims=dynamic_arg_dims,
+            compile_config=compile_config,
+            graph_pool=graph_pool,
+        )
+
+    def _run_compile_pass(self, cuda_graph_runner: BaseCudaGraphRunner) -> None:
+        """JIT-activate kernels at the smallest shape, install
+        torch.compile, then run one forward per shape inside
+        enable_torch_compile_warmup to drive FX / inductor through
+        every shape without capturing cuda graphs yet."""
+        language_model = self._language_model
+        # Some multimodal models (e.g. Gemma4) store the inner transformer
+        # directly as `language_model` rather than wrapping it in a
+        # ForCausalLM that has a `.model` child.  Fall back to the module
+        # itself when `.model` is absent.
+        inner_model = getattr(language_model, "model", language_model)
+        compiler = self._compile_config.compiler
+        with enable_tc_piecewise_cuda_graph():
+            try:
+                if compiler != "eager":
+                    _toggle_multi_platform_ops(
+                        inner_model, reverse=False, num_tokens=16
+                    )
+
+                cuda_graph_runner._run_dummy_forward(
+                    num_tokens=cuda_graph_runner.capture_num_tokens[0]
+                )
+
+                if self._pool is None:
+                    self._pool = get_or_create_global_graph_memory_pool(
+                        self._device_module
+                    )
+                set_graph_pool_id(self._pool)
+
+                self.install_compile(
+                    inner_model,
+                    compile_config=self._compile_config,
+                    graph_pool=self._pool,
+                )
+
+                with enable_torch_compile_warmup():
+                    if is_hip():
+                        # AMD: single Dynamo trace is sufficient; the capture
+                        # phase does per-shape JIT kernel warmup before each
+                        # CUDA graph recording.  The N-iteration loop is
+                        # redundant and extremely slow on ROCm (~30 min).
+                        cuda_graph_runner._run_dummy_forward(
+                            num_tokens=cuda_graph_runner.capture_num_tokens[-1]
+                        )
+                    else:
+                        compile_range = (
+                            tqdm.tqdm(
+                                list(reversed(cuda_graph_runner.capture_num_tokens))
+                            )
+                            if get_parallel().tp_rank == 0
+                            else reversed(cuda_graph_runner.capture_num_tokens)
+                        )
+                        for num_tokens in compile_range:
+                            if get_parallel().tp_rank == 0:
+                                compile_range.set_description(
+                                    f"Compiling num tokens ({num_tokens=})"
+                                )
+                            cuda_graph_runner._run_dummy_forward(num_tokens=num_tokens)
+            finally:
+                _toggle_multi_platform_ops(inner_model, reverse=True, num_tokens=16)
+
+    @contextmanager
+    def capture_session(self, stream: torch.cuda.Stream):
+        self._capture_stream = stream
+        try:
+            with self.replay_session():
+                with set_pcg_capture_stream(stream):
+                    yield
+        finally:
+            self._capture_stream = None
+
+    def capture_one(
+        self,
+        shape_key: ShapeKey,
+        forward_fn: Callable[[], Any],
+        dummies: Optional[Any] = None,
+        post_warmup_hook: Optional[Callable[[], None]] = None,
+    ) -> None:
+        # Call 1 warms FX state; call 2 captures the cuda graph inside capture_session.
+        # See cuda_piecewise_backend.py for the FX backend that drives the capture.
+        for _ in range(2):
+            self._device_module.synchronize()
+            self._tp_group.barrier()
+            forward_fn()
+            if post_warmup_hook is not None:
+                post_warmup_hook()
+
+    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
+
+    @contextmanager
+    def replay_session(self):
+        with enable_tc_piecewise_cuda_graph():
+            yield
+
+    def replay(
+        self,
+        shape_key: ShapeKey,
+        static_forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any:
+        return self._compiled_fn(
+            static_forward_batch.input_ids,
+            static_forward_batch.positions,
+            static_forward_batch,
+            **kwargs,
+        )
+
+    def cleanup(self) -> None:
+        self._compiled_fn = None
+        self._compile_config = None
+        self._language_model = None
+        self._pool = None

python/sglang/srt/model_executor/runner_backend/utils.py

@@ -0,0 +1,124 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""runner_backend utilities — phase → BaseCudaGraphBackend resolution.
+
+Centralizes per-phase backend resolution so platform overrides (NPU,
+out-of-tree) and future backend additions can plug in without
+modifying the runner files. Phase / backend identifiers used here
+live in :mod:`.cuda_graph_config`.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from sglang.srt.model_executor.cuda_graph_config import Backend
+from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
+    BaseCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import (
+    BreakableCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import (
+    FullCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.tc_piecewise_cuda_graph_backend import (
+    TcPiecewiseCudaGraphBackend,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
+        BaseCudaGraphRunner,
+    )
+
+logger = logging.getLogger(__name__)
+
+# Track first occurrence of each fallback warning to avoid log spam.
+_TC_PIECEWISE_DECODE_FALLBACK_LOGGED = False
+
+
+def resolve_decode_backend(
+    cuda_graph_runner: BaseCudaGraphRunner,
+) -> BaseCudaGraphBackend:
+    """Pick a backend instance from cuda_graph_config['decode']['backend'].
+
+    NPU device returns NPUCudaGraphBackend regardless of mode (only
+    the Full-style backend is wired for NPU today).
+    """
+    model_runner = cuda_graph_runner.model_runner
+    cfg = model_runner.server_args.cuda_graph_config
+    backend_name = cfg.decode.backend if cfg is not None else Backend.FULL
+
+    enable_memory_saver = model_runner.server_args.enable_memory_saver
+
+    if model_runner.device == "npu":
+        from sglang.srt.hardware_backend.npu.graph_runner.npu_cudagraph_backend import (
+            NPUCudaGraphBackend,
+        )
+
+        return NPUCudaGraphBackend(
+            cuda_graph_runner, enable_memory_saver=enable_memory_saver
+        )
+    elif model_runner.device == "xpu":
+        if backend_name not in (Backend.FULL, Backend.DISABLED):
+            raise ValueError(
+                f"XPU only supports cuda_graph_config decode backend 'full', got '{backend_name}'"
+            )
+        from sglang.srt.hardware_backend.xpu.graph_runner.xpu_full_graph_backend import (
+            FullXPUGraphBackend,
+        )
+
+        return FullXPUGraphBackend(cuda_graph_runner)
+
+    if backend_name == Backend.BREAKABLE:
+        return BreakableCudaGraphBackend(
+            cuda_graph_runner,
+            enable_memory_saver=enable_memory_saver,
+            debug_eager=model_runner.server_args.debug_cuda_graph,
+        )
+    if backend_name == Backend.TC_PIECEWISE:
+        global _TC_PIECEWISE_DECODE_FALLBACK_LOGGED
+        if not _TC_PIECEWISE_DECODE_FALLBACK_LOGGED:
+            logger.warning(
+                "cuda_graph_config decode='tc_piecewise' is not yet implemented; "
+                "falling back to 'full'."
+            )
+            _TC_PIECEWISE_DECODE_FALLBACK_LOGGED = True
+    return FullCudaGraphBackend(
+        cuda_graph_runner, enable_memory_saver=enable_memory_saver
+    )
+
+
+def resolve_prefill_backend(
+    cuda_graph_runner: BaseCudaGraphRunner,
+) -> BaseCudaGraphBackend:
+    """Pick a backend instance from cuda_graph_config['prefill']['backend']."""
+    model_runner = cuda_graph_runner.model_runner
+    cfg = model_runner.server_args.cuda_graph_config
+    backend_name = cfg.prefill.backend if cfg is not None else Backend.TC_PIECEWISE
+
+    if backend_name == Backend.BREAKABLE:
+        return BreakableCudaGraphBackend(
+            cuda_graph_runner,
+            enable_memory_saver=model_runner.server_args.enable_memory_saver,
+            debug_eager=model_runner.server_args.debug_cuda_graph,
+        )
+    if backend_name == Backend.FULL:
+        return FullCudaGraphBackend(
+            cuda_graph_runner,
+            enable_memory_saver=model_runner.server_args.enable_memory_saver,
+        )
+    # Default: tc_piecewise.
+    return TcPiecewiseCudaGraphBackend(cuda_graph_runner)