chore: import upstream snapshot with attribution

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

@@ -0,0 +1,28 @@
+"""Low-level primitives used by the CUDA graph backends.
+
+Subpackages:
+  - breakable_cuda_graph: BreakableCUDAGraph + capture context,
+    eager_on_graph decorator, is_in_breakable_cuda_graph flag.
+  - piecewise_cuda_graph: shared piecewise context manager
+    (set_tc_piecewise_forward_context, is_in_tc_piecewise_cuda_graph).
+
+Backends in cuda_graph_backend/ import from here. Runners do not.
+"""
+
+# Generic failure-message hint for decode-style CUDA graph capture paths
+# (Full backend used by decode + EAGLE draft runners).
+CUDA_GRAPH_CAPTURE_FAILED_MSG = (
+    "Possible solutions:\n"
+    "1. set --mem-fraction-static to a smaller value (e.g., 0.8 or 0.7)\n"
+    "2. set --cuda-graph-max-bs-decode to a smaller value (e.g., 16)\n"
+    "3. disable decode CUDA graph by --cuda-graph-backend-decode=disabled. "
+    "(Not recommended. Huge performance loss)\n"
+    "Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
+)
+
+PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG = (
+    "Fail when using backend: {backend} for prefill runner.\n"
+    "Possible suggestions:\n"
+    "{suggestions}"
+    "Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
+)

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/__init__.py

@@ -0,0 +1,21 @@
+"""Breakable primitives — segmented CUDA graph capture with eager break points.
+
+Public API (also reachable via the deeper module paths):
+  - BreakableCUDAGraph, BreakableCUDAGraphCapture — capture/replay
+  - eager_on_graph — decorator that marks a callable as a graph break
+  - break_graph — helper that inserts a bare graph break
+  - enable_breakable_cuda_graph — context that flips the Breakable runtime flag
+  - is_in_breakable_cuda_graph — runtime flag getter
+
+"""
+
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (  # noqa: F401
+    BreakableCUDAGraph,
+    BreakableCUDAGraphCapture,
+    break_graph,
+    eager_on_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (  # noqa: F401
+    enable_breakable_cuda_graph,
+    is_in_breakable_cuda_graph,
+)

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py

@@ -0,0 +1,391 @@
+# 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.
+# ==============================================================================
+"""Breakable CUDA Graph: capture a region as a sequence of
+torch.cuda.CUDAGraph segments separated by eager break points.
+
+Each segment is a real torch.cuda.CUDAGraph. Its destructor calls
+releasePool on the shared mempool, so the pool's use_count tracks how
+many segments are alive; the pool stays pinned as long as any segment graph
+is alive. This lets weak_ref_tensor views of intermediate pool-allocated
+tensors remain valid across replays — we don't need Python-managed bridge
+buffers to keep break-point tensors at stable addresses.
+"""
+
+import logging
+import threading
+from contextvars import ContextVar
+from typing import Any, Callable
+
+import torch
+
+try:
+    from cuda.bindings import runtime as rt
+except ImportError:
+    rt = None
+
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.cuda_utils import (
+    checkCudaErrors,
+)
+from sglang.srt.utils import is_hip
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "eager_on_graph",
+    "BreakableCUDAGraph",
+    "BreakableCUDAGraphCapture",
+    "break_graph",
+]
+
+
+def _check_cuda_bindings():
+    if rt is None:
+        raise ImportError(
+            "Breakable CUDA graph on NVIDIA requires the 'cuda-python' package. "
+            "Install it with: pip install cuda-python"
+        )
+
+
+# Active BreakableCUDAGraphCapture context for the currently-capturing thread.
+# eager_on_graph's wrapper uses this to split the current torch.cuda.CUDAGraph
+# at break points.
+_current_capture_var: ContextVar["BreakableCUDAGraphCapture | None"] = ContextVar(
+    "current_capture", default=None
+)
+_current_stream_var: ContextVar[torch.cuda.Stream | None] = ContextVar(
+    "current_stream", default=None
+)
+_forked_streams_var: ContextVar[set[torch.cuda.Stream] | None] = ContextVar(
+    "forked_streams", default=None
+)
+
+
+def get_current_stream(device: torch.device | None = None) -> torch.cuda.Stream:
+    stream = _current_stream_var.get()
+    if stream is None:
+        return torch.cuda.current_stream(device)
+    return stream
+
+
+def _capture_status(stream_ptr: int) -> "rt.cudaStreamCaptureStatus":
+    _check_cuda_bindings()
+    status, *_ = checkCudaErrors(rt.cudaStreamGetCaptureInfo(stream_ptr))
+    return status
+
+
+def _is_stream_capturing(stream: torch.cuda.Stream) -> bool:
+    # On ROCm/HIP, cuda-python is unavailable, so use the portable torch API
+    # (which maps to the HIP runtime). On NVIDIA, keep querying the CUDA runtime
+    # directly via cuda-python: torch.cuda.is_current_stream_capturing() has
+    # proven unreliable there, so we preserve the original behavior.
+    if is_hip():
+        with torch.cuda.stream(stream):
+            return torch.cuda.is_current_stream_capturing()
+    return (
+        _capture_status(stream.cuda_stream)
+        == rt.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
+    )
+
+
+# Hook torch.cuda.Stream.wait_stream to track side-stream forks/joins that happen
+# during breakable capture. We need this because capture_end() on a torch
+# CUDAGraph fails if there are still side streams participating in the capture
+# — so before ending each segment we auto-join any forked-but-not-rejoined streams.
+_original_wait_stream: Callable | None = None
+_hook_lock = threading.Lock()
+_hook_refcount = 0
+
+
+def _hooked_wait_stream(self: torch.cuda.Stream, other: torch.cuda.Stream):
+    assert _original_wait_stream is not None
+    forked = _forked_streams_var.get()
+    if forked is None:
+        _original_wait_stream(self, other)
+        return
+    capturing = _current_stream_var.get()
+    if capturing is None:
+        _original_wait_stream(self, other)
+        return
+
+    cap_ptr = capturing.cuda_stream
+    is_self_cap = self is capturing or self.cuda_stream == cap_ptr
+    is_other_cap = other is capturing or other.cuda_stream == cap_ptr
+
+    if is_self_cap and not is_other_cap:
+        if not _is_stream_capturing(other):
+            return
+        _original_wait_stream(self, other)
+        forked.discard(other)
+    elif is_other_cap and not is_self_cap:
+        _original_wait_stream(self, other)
+        forked.add(self)
+    else:
+        _original_wait_stream(self, other)
+
+
+def _install_wait_stream_hook():
+    global _original_wait_stream, _hook_refcount
+    with _hook_lock:
+        if _hook_refcount == 0:
+            _original_wait_stream = torch.cuda.Stream.wait_stream
+            torch.cuda.Stream.wait_stream = _hooked_wait_stream  # type: ignore[assignment]
+        _hook_refcount += 1
+
+
+def _uninstall_wait_stream_hook():
+    global _original_wait_stream, _hook_refcount
+    with _hook_lock:
+        _hook_refcount -= 1
+        if _hook_refcount == 0:
+            assert _original_wait_stream is not None, "wait_stream hook not installed"
+            torch.cuda.Stream.wait_stream = _original_wait_stream  # type: ignore[assignment]
+            _original_wait_stream = None
+
+
+def _weak_ref_if_tensor(x):
+    """Return a weak-ref tensor view (shared storage, no refcount) for tensors;
+    recurse into tuples/lists; pass-through for non-tensors. Weak-ref'ing
+    captured args lets the shared mempool reclaim per-layer intermediates
+    between segments — storage stays alive for each segment CUDAGraph's
+    lifetime via its pool use_count.
+
+    weak_ref_tensors is imported lazily because it hard-raises on
+    platforms without a CUDA/HIP/NPU backend; we only reach this code during
+    an active Breakable capture, which runs only on those backends."""
+    if torch.is_tensor(x):
+        from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors
+
+        return weak_ref_tensors(x)
+    if isinstance(x, tuple):
+        return tuple(_weak_ref_if_tensor(e) for e in x)
+    if isinstance(x, list):
+        return [_weak_ref_if_tensor(e) for e in x]
+    return x
+
+
+def _copy_output(dst: Any, src: Any) -> Any:
+    """Copy src output into dst in-place where possible.
+
+    Handles plain tensors, tuples/lists of tensors, dataclass/object with
+    tensor attributes, and dicts of tensors. Returns dst if in-place copy
+    succeeded, otherwise returns src.
+    """
+    if torch.is_tensor(dst) and torch.is_tensor(src):
+        dst.copy_(src)
+        return dst
+
+    if (
+        isinstance(dst, (tuple, list))
+        and isinstance(src, (tuple, list))
+        and len(dst) == len(src)
+    ):
+        copied = [_copy_output(d, s) for d, s in zip(dst, src)]
+        return tuple(copied) if isinstance(dst, tuple) else copied
+
+    if hasattr(dst, "__dict__") and hasattr(src, "__dict__"):
+        for key, src_val in src.__dict__.items():
+            dst_val = getattr(dst, key, None)
+            if torch.is_tensor(dst_val) and torch.is_tensor(src_val):
+                dst_val.copy_(src_val)
+            else:
+                setattr(dst, key, src_val)
+        return dst
+
+    if isinstance(dst, dict) and isinstance(src, dict):
+        for key, src_val in src.items():
+            dst_val = dst.get(key)
+            if torch.is_tensor(dst_val) and torch.is_tensor(src_val):
+                dst_val.copy_(src_val)
+            else:
+                dst[key] = src_val
+        return dst
+
+    return src
+
+
+def eager_on_graph(enable: bool):
+    def decorator(inner: Callable):
+        if not enable:
+            return inner
+
+        def wrapper(*args, **kwargs):
+            capture = _current_capture_var.get()
+            if capture is None:
+                return inner(*args, **kwargs)
+
+            logger.debug("Break graph due to function: %s", inner.__name__)
+
+            # End the segment that captured up to this break point.
+            capture._end_current_segment()
+
+            # Run the eager function once so it allocates its outputs and
+            # writes real data into them.
+            output = inner(*args, **kwargs)
+
+            # Weak-ref captured inputs produced by graph segments. Their storage
+            # is pinned by the segment CUDAGraphs' mempool use-count, so Python
+            # refs do not need to keep every intermediate alive.
+            captured_inner = inner
+            captured_args = tuple(_weak_ref_if_tensor(a) for a in args)
+            captured_kwargs = {k: _weak_ref_if_tensor(v) for k, v in kwargs.items()}
+            # The eager break output is different: it is allocated between graph
+            # captures and is the static input address consumed by the next
+            # captured segment. Keep a strong reference so replay can safely
+            # copy fresh eager output into that bridge buffer.
+            captured_output = output
+
+            def replay_fn():
+                new_out = captured_inner(*captured_args, **captured_kwargs)
+                return _copy_output(captured_output, new_out)
+
+            capture.cuda_graph._break_fns.append(replay_fn)
+
+            # Start a fresh CUDAGraph segment for the remainder of the forward.
+            capture._begin_new_segment()
+            return output
+
+        return wrapper
+
+    return decorator
+
+
+class BreakableCUDAGraph:
+    """Container holding one torch.cuda.CUDAGraph per segment plus an
+    eager break function between consecutive segments."""
+
+    def __init__(self, deduped_cuda_graph=None) -> None:
+        self._segments: list[Any] = []
+        self._break_fns: list[Callable[[], Any]] = []
+        self._deduped_cuda_graph = deduped_cuda_graph
+
+    def replay(self) -> None:
+        stream = torch.cuda.current_stream()
+        token = _current_stream_var.set(stream)
+        try:
+            for i, seg in enumerate(self._segments):
+                seg.replay()
+                if i < len(self._break_fns):
+                    self._break_fns[i]()
+        finally:
+            _current_stream_var.reset(token)
+
+    def _append_segment(
+        self, graph: torch.cuda.CUDAGraph, needs_instantiate: bool
+    ) -> None:
+        if self._deduped_cuda_graph is not None:
+            self._segments.append(self._deduped_cuda_graph.register(graph))
+            return
+        if needs_instantiate:
+            graph.instantiate()
+        self._segments.append(graph)
+
+
+class BreakableCUDAGraphCapture:
+    """Context manager that captures the enclosed code as one or more
+    torch.cuda.CUDAGraph segments separated by eager break points.
+
+    Each segment shares the supplied pool (MempoolId_t tuple) so
+    pool-allocated intermediates can be reused across segments. While any
+    segment is alive, its beginAllocateToPool call keeps the mempool's
+    use_count > 0, which makes weak_ref_tensor of segment-allocated
+    tensors safe across subsequent replays.
+    """
+
+    def __init__(
+        self,
+        cuda_graph: BreakableCUDAGraph,
+        pool=None,
+        stream: torch.cuda.Stream | None = None,
+        capture_error_mode: str = "global",
+    ):
+        assert isinstance(
+            cuda_graph, BreakableCUDAGraph
+        ), "cuda_graph must be a BreakableCUDAGraph"
+        self.cuda_graph = cuda_graph
+        self._pool = pool if pool is not None else (0, 0)
+        self._stream = stream
+        self._capture_error_mode = capture_error_mode
+        self._stream_ctx = None
+        self._capture_token = None
+        self._stream_token = None
+        self._forked_token = None
+        self._current_graph: torch.cuda.CUDAGraph | None = None
+        self._current_graph_needs_instantiate = False
+
+    def __enter__(self):
+        _install_wait_stream_hook()
+        if self._stream is not None:
+            self._stream_ctx = torch.cuda.stream(self._stream)
+            self._stream_ctx.__enter__()
+        self._capture_token = _current_capture_var.set(self)
+        self._stream_token = _current_stream_var.set(
+            self._stream or torch.cuda.current_stream()
+        )
+        self._forked_token = _forked_streams_var.set(set())
+        self._begin_new_segment()
+        return self
+
+    def __exit__(self, *args: object):
+        try:
+            self._end_current_segment()
+        finally:
+            _forked_streams_var.reset(self._forked_token)
+            _current_stream_var.reset(self._stream_token)
+            _current_capture_var.reset(self._capture_token)
+            if self._stream_ctx is not None:
+                self._stream_ctx.__exit__(*args)
+                self._stream_ctx = None
+            _uninstall_wait_stream_hook()
+        return False
+
+    def _begin_new_segment(self) -> None:
+        # keep_graph retains the raw graph for dedup; skip it on the plain path.
+        if self.cuda_graph._deduped_cuda_graph is not None:
+            try:
+                graph = torch.cuda.CUDAGraph(keep_graph=True)
+                self._current_graph_needs_instantiate = True
+            except TypeError:
+                graph = torch.cuda.CUDAGraph()
+                self._current_graph_needs_instantiate = False
+        else:
+            graph = torch.cuda.CUDAGraph()
+            self._current_graph_needs_instantiate = False
+        graph.capture_begin(
+            pool=self._pool, capture_error_mode=self._capture_error_mode
+        )
+        self._current_graph = graph
+
+    def _end_current_segment(self) -> None:
+        # Auto-join any side streams forked during this segment but not joined.
+        main_stream = get_current_stream()
+        forked = _forked_streams_var.get()
+        if forked:
+            assert _original_wait_stream is not None
+            for side in list(forked):
+                if _is_stream_capturing(side):
+                    _original_wait_stream(main_stream, side)
+            forked.clear()
+        graph = self._current_graph
+        assert graph is not None
+        graph.capture_end()
+        self.cuda_graph._append_segment(graph, self._current_graph_needs_instantiate)
+        self._current_graph = None
+        self._current_graph_needs_instantiate = False
+
+
+@eager_on_graph(True)
+def break_graph() -> None:
+    """Insert a graph break. The @eager_on_graph decorator does the actual
+    segment split; this function body intentionally does nothing."""
+    pass

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/context.py

@@ -0,0 +1,60 @@
+# 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.
+# ==============================================================================
+"""Runtime state for the breakable CUDA graph runner."""
+
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+
+from sglang.srt.model_executor.cuda_graph_config import Backend
+from sglang.srt.model_executor.runner_backend_utils import (
+    PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG,
+)
+
+logger = logging.getLogger(__name__)
+
+_in_breakable_cuda_graph = False
+
+
+def is_in_breakable_cuda_graph() -> bool:
+    return _in_breakable_cuda_graph
+
+
+@contextmanager
+def enable_breakable_cuda_graph():
+    """Mark the enclosed scope as inside a BCG capture/replay. Any exception
+    raised inside is logged with the BCG-specific failure hint, then re-raised
+    for the caller to handle."""
+    global _in_breakable_cuda_graph
+    _in_breakable_cuda_graph = True
+    try:
+        yield
+    except Exception as exc:
+        msg = PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG.format(
+            backend=Backend.BREAKABLE, suggestions=BCG_FAILURE_HINT
+        )
+        logger.error(f"{type(exc).__name__}: {exc}\n{msg}")
+        raise
+    finally:
+        _in_breakable_cuda_graph = False
+
+
+BCG_FAILURE_HINT = (
+    "1. change to tc_piecewise by --cuda-graph-backend-prefill=tc_piecewise\n"
+    "2. disable the prefill CUDA graph by --cuda-graph-backend-prefill=disabled\n"
+    "3. if it is an OOM problem, set --mem-fraction-static to a smaller value "
+    "(e.g., 0.8 or 0.7) or set --cuda-graph-max-bs-prefill to a smaller value "
+    "(e.g., 2048)\n"
+)

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py

@@ -0,0 +1,48 @@
+# 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.
+# ==============================================================================
+"""CUDA runtime binding utilities."""
+
+try:
+    from cuda.bindings import runtime as rt
+except ImportError:
+    rt = None
+
+
+def _cudaGetErrorString(error):
+    if rt is None:
+        return "<cuda.bindings not available>"
+    err, msg = rt.cudaGetErrorString(error)
+    if err != rt.cudaError_t.cudaSuccess:
+        return "<unknown>"
+    if isinstance(msg, bytes):
+        return msg.decode("utf-8", "replace")
+    return str(msg)
+
+
+def checkCudaErrors(result):
+    if rt is None:
+        raise RuntimeError(
+            "cuda.bindings is not available. "
+            "Install it with: pip install cuda-python"
+        )
+    if result[0] != rt.cudaError_t.cudaSuccess:
+        raise RuntimeError(
+            f"CUDA error {int(result[0])}({_cudaGetErrorString(result[0])})"
+        )
+    if len(result) == 1:
+        return None
+    elif len(result) == 2:
+        return result[1]
+    else:
+        return result[1:]

python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/__init__.py

@@ -0,0 +1,22 @@
+"""Piecewise CUDA graph utilities — shared between Breakable and tc_piecewise backends.
+
+Public API:
+  - is_in_tc_piecewise_cuda_graph() — true while inside any piecewise capture.
+  - enable_tc_piecewise_cuda_graph() — context manager that toggles the flag.
+  - TcPiecewiseForwardContext + set_tc_piecewise_forward_context + get_tc_piecewise_forward_context.
+  - TCPCG_FAILURE_HINT — backend-switch suggestion plugged into
+    PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG by the prefill runner.
+
+The torch.compile-warmup flag (is_in_torch_compile_warmup) lives in
+sglang.srt.compilation.compile_phase — it is torch.compile-internal,
+not piecewise-shared.
+"""
+
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph.context_manager import (  # noqa: F401
+    TCPCG_FAILURE_HINT,
+    TcPiecewiseForwardContext,
+    enable_tc_piecewise_cuda_graph,
+    get_tc_piecewise_forward_context,
+    is_in_tc_piecewise_cuda_graph,
+    set_tc_piecewise_forward_context,
+)

python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py

@@ -0,0 +1,121 @@
+"""CUDA graph capture context manager + forward-context propagation.
+
+Owns two pieces of cross-cutting state used by *every* piecewise-style
+backend (currently breakable + tc_piecewise):
+
+* _in_tc_piecewise_cuda_graph — a process-global flag set true while we
+  are inside the capture or replay window of a piecewise CUDA graph.
+  Read by model code that needs to take the static-buffer / fixed-shape
+  branch. See refactor/plan.md §6.5 for the full semantics.
+* TcPiecewiseForwardContext — a dataclass propagated across attention/MoE
+  layers during capture and replay so that submodules can reach the
+  current ForwardBatch and per-layer metadata without threading
+  arguments through every call site. Named TcPiecewise… (matches
+  Backend.TC_PIECEWISE + enable_tc_piecewise_cuda_graph) to
+  disambiguate from the per-forward-call
+  sglang.srt.model_executor.forward_context.ForwardContext.
+
+This module deliberately does **not** own torch.compile-specific state
+(warmup flag, capture stream); those live in compilation/compile_phase.py.
+"""
+
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from sglang.srt.model_executor.cuda_graph_config import Backend
+from sglang.srt.model_executor.runner_backend_utils import (
+    PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+logger = logging.getLogger(__name__)
+
+_in_tc_piecewise_cuda_graph = False
+
+
+def is_in_tc_piecewise_cuda_graph() -> bool:
+    """True while inside tc_piecewise CUDA graph capture/replay."""
+    return _in_tc_piecewise_cuda_graph
+
+
+@contextmanager
+def enable_tc_piecewise_cuda_graph():
+    """Mark the enclosed scope as inside a tc_piecewise CUDA graph
+    capture/replay. Any exception raised inside is logged with the
+    PCG-specific failure hint, then re-raised for the caller to handle.
+    """
+    global _in_tc_piecewise_cuda_graph
+    _in_tc_piecewise_cuda_graph = True
+    try:
+        yield
+    except Exception as exc:
+        msg = PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG.format(
+            backend=Backend.TC_PIECEWISE, suggestions=TCPCG_FAILURE_HINT
+        )
+        logger.error(f"{type(exc).__name__}: {exc}\n{msg}")
+        raise
+    finally:
+        _in_tc_piecewise_cuda_graph = False
+
+
+@dataclass
+class TcPiecewiseForwardContext:
+    forward_batch: Optional[ForwardBatch] = None
+    attention_layers: Optional[List[Any]] = field(default=None)
+    quant_config: Any = None
+    moe_layers: Optional[List[Any]] = field(default=None)
+    moe_fusions: Optional[List[Any]] = field(default=None)
+    dsa_indexers: Optional[List[Any]] = field(default=None)
+    num_tokens: Optional[int] = None
+    raw_num_tokens: Optional[int] = None
+
+
+_tc_piecewise_forward_context: Optional[TcPiecewiseForwardContext] = None
+
+
+def get_tc_piecewise_forward_context() -> Optional[TcPiecewiseForwardContext]:
+    return _tc_piecewise_forward_context
+
+
+@contextmanager
+def set_tc_piecewise_forward_context(
+    forward_batch: ForwardBatch,
+    attention_layers: List[Any],
+    quant_config: Any,
+    moe_layers: List[Any],
+    moe_fusions: List[Any],
+    dsa_indexers: Optional[List[Any]] = None,
+    num_tokens: Optional[int] = None,
+    raw_num_tokens: Optional[int] = None,
+):
+    global _tc_piecewise_forward_context
+    _tc_piecewise_forward_context = TcPiecewiseForwardContext(
+        forward_batch=forward_batch,
+        attention_layers=attention_layers,
+        quant_config=quant_config,
+        moe_layers=moe_layers,
+        moe_fusions=moe_fusions,
+        dsa_indexers=dsa_indexers,
+        num_tokens=num_tokens,
+        raw_num_tokens=raw_num_tokens,
+    )
+    try:
+        yield
+    finally:
+        _tc_piecewise_forward_context = None
+
+
+TCPCG_FAILURE_HINT = (
+    "1. change to breakable by --cuda-graph-backend-prefill=breakable\n"
+    "2. disable the prefill CUDA graph by --cuda-graph-backend-prefill=disabled\n"
+    "3. if it is an OOM problem, set --mem-fraction-static to a smaller value "
+    "(e.g., 0.8 or 0.7) or set --cuda-graph-max-bs-prefill to a smaller value "
+    "(e.g., 2048)\n"
+)