[Experimental] Breakable Piecewise Cuda Graph (#22218)

Signed-off-by: Oasis-Git <ayw.sirius19@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

python/sglang/srt/layers/radix_attention.py

@@ -23,6 +23,12 @@ from torch import nn
 
 from sglang.srt.compilation.compilation_config import register_split_op
 from sglang.srt.compilation.piecewise_context_manager import get_forward_context
+from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import (
+    eager_on_graph,
+)
+from sglang.srt.model_executor.breakable_cuda_graph.context import (
+    is_in_breakable_cuda_graph,
+)
 from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:
@@ -119,9 +125,14 @@ class RadixAttention(nn.Module):
                 output = q.new_empty((q.shape[0], self.tp_q_head_num * self.v_head_dim))
             else:
                 output = torch.empty_like(q)
-            unified_attention_with_output(
-                q, k, v, output, save_kv_cache, self.layer_id, **kwargs
-            )
+            if is_in_breakable_cuda_graph():
+                bcg_unified_attention_with_output(
+                    q, k, v, output, save_kv_cache, self.layer_id, **kwargs
+                )
+            else:
+                unified_attention_with_output(
+                    q, k, v, output, save_kv_cache, self.layer_id, **kwargs
+                )
             return output
         else:
             return forward_batch.attn_backend.forward(
@@ -197,3 +208,6 @@ def unified_attention_with_output(
 
     output[:real_num_tokens].view(ret.shape).copy_(ret)
     return
+
+
+bcg_unified_attention_with_output = eager_on_graph(True)(unified_attention_with_output)

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

@@ -11,10 +11,21 @@
 # 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, NamedTuple
+from typing import Any, Callable
 
 import torch
 
@@ -43,17 +54,11 @@ def _check_cuda_bindings():
         )
 
 
-class GraphBreakInfo(NamedTuple):
-    # python function breaking the graph
-    func: Callable
-    # output of the function (must be a tensor so we keep them)
-    output: Any
-    # raw handle after capture or raw exec handle after instantiate
-    graph_handle: Any
-
-
-_captured_graphs_var: ContextVar[list[GraphBreakInfo] | None] = ContextVar(
-    "captured_graphs", default=None
+# 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
@@ -84,7 +89,10 @@ def _is_capturing(stream_ptr: int) -> bool:
     )
 
 
-# hook wait_stream to track forks/joins during breakable capture.
+# 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
@@ -106,9 +114,6 @@ def _hooked_wait_stream(self: torch.cuda.Stream, other: torch.cuda.Stream):
     is_other_cap = other is capturing or other.cuda_stream == cap_ptr
 
     if is_self_cap and not is_other_cap:
-        # Join: capturing_stream.wait_stream(other).
-        # other might not be part of the capture because we join it in the last segment
-        # skip the wait to avoid cuda error
         if (
             _capture_status(other.cuda_stream)
             != rt.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
@@ -117,7 +122,6 @@ def _hooked_wait_stream(self: torch.cuda.Stream, other: torch.cuda.Stream):
         _original_wait_stream(self, other)
         forked.discard(other)
     elif is_other_cap and not is_self_cap:
-        # Fork: other.wait_stream(capturing_stream).
         _original_wait_stream(self, other)
         forked.add(self)
     else:
@@ -143,49 +147,20 @@ def _uninstall_wait_stream_hook():
             _original_wait_stream = None
 
 
-def _end_capture_segment(stream: torch.cuda.Stream):
-    """End a capture segment, auto-joining any forked streams first."""
-    # Join forked streams that are still part of this capture.
-    forked = _forked_streams_var.get()
-    if forked:
-        assert _original_wait_stream is not None
-        for s in forked:
-            if _is_capturing(s.cuda_stream):
-                _original_wait_stream(stream, s)
-        forked.clear()
+def _weak_ref_if_tensor(x):
+    """Return a weak-ref tensor view (shared storage, no refcount) for tensors;
+    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.
 
-    graph = checkCudaErrors(rt.cudaStreamEndCapture(stream.cuda_stream))
-    assert graph is not None
-    return graph
+    ``weak_ref_tensors`` is imported lazily: the module hard-raises on
+    non-CUDA/NPU platforms, and we only reach this code during an active
+    BCG capture (which can't happen on CPU-only runners anyway)."""
+    if torch.is_tensor(x):
+        from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors
 
-
-def _begin_capture_segment(stream: torch.cuda.Stream):
-    checkCudaErrors(
-        rt.cudaStreamBeginCapture(
-            stream.cuda_stream,
-            rt.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal,
-        )
-    )
-
-
-def _instantiate_graph(graph_ptr: int) -> int:
-    graph_exec = checkCudaErrors(
-        rt.cudaGraphInstantiateWithFlags(
-            graph_ptr,
-            rt.cudaGraphInstantiateFlags.cudaGraphInstantiateFlagAutoFreeOnLaunch,
-        )
-    )
-    assert graph_exec is not None
-    checkCudaErrors(rt.cudaGraphDestroy(graph_ptr))
-    return graph_exec
-
-
-def _destroy_graph_exec(graph_exec_ptr: int) -> None:
-    checkCudaErrors(rt.cudaGraphExecDestroy(graph_exec_ptr))
-
-
-def _replay_graph(graph_exec_ptr: int, stream_ptr: int) -> None:
-    checkCudaErrors(rt.cudaGraphLaunch(graph_exec_ptr, stream_ptr))
+        return weak_ref_tensors(x)
+    return x
 
 
 def _copy_output(dst: Any, src: Any) -> Any:
@@ -199,7 +174,6 @@ def _copy_output(dst: Any, src: Any) -> Any:
         dst.copy_(src)
         return dst
 
-    # Handle objects with __dict__ (dataclasses, regular objects)
     if hasattr(dst, "__dict__") and hasattr(src, "__dict__"):
         for key, src_val in src.__dict__.items():
             dst_val = getattr(dst, key, None)
@@ -209,7 +183,6 @@ def _copy_output(dst: Any, src: Any) -> Any:
                 setattr(dst, key, src_val)
         return dst
 
-    # Handle dicts of tensors
     if isinstance(dst, dict) and isinstance(src, dict):
         for key, src_val in src.items():
             dst_val = dst.get(key)
@@ -228,32 +201,35 @@ def eager_on_graph(enable: bool):
             return inner
 
         def wrapper(*args, **kwargs):
-            stream = get_current_stream()
-            if not _is_capturing(stream.cuda_stream):
+            capture = _current_capture_var.get()
+            if capture is None:
                 return inner(*args, **kwargs)
-            last_graph = _end_capture_segment(stream)
-            logger.debug(f"Break graph due to function: {inner.__name__}")
-            # run the function once to allocate the output tensor captured by later graphs
+
+            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)
 
-            # Store the callable and its arguments so replay can re-invoke with
-            # the same argument *references* (which point to CUDA graph input
-            # buffers whose contents are updated before replay).
+            # Weak-ref the closure state. Storage lives with the segment
+            # CUDAGraphs' mempool pin; Python refs don't need to prevent
+            # pool reuse across layers.
             captured_inner = inner
-            captured_args = args
-            captured_kwargs = kwargs
-            captured_output = output
+            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()}
+            captured_output = _weak_ref_if_tensor(output)
 
             def replay_fn():
                 new_out = captured_inner(*captured_args, **captured_kwargs)
                 return _copy_output(captured_output, new_out)
 
-            captured_graphs = _captured_graphs_var.get()
-            assert (
-                captured_graphs is not None
-            ), "eager_on_graph wrapper called outside of BreakableCUDAGraphCapture"
-            captured_graphs.append(GraphBreakInfo(replay_fn, output, last_graph))
-            _begin_capture_segment(stream)
+            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
@@ -261,60 +237,37 @@ def eager_on_graph(enable: bool):
     return decorator
 
 
-class BreakableCUDAGraph(torch.cuda.CUDAGraph):
+class BreakableCUDAGraph:
+    """Container holding one ``torch.cuda.CUDAGraph`` per segment plus an
+    eager break function between consecutive segments."""
 
-    def __new__(cls) -> "BreakableCUDAGraph":
-        return super().__new__(cls, True)
+    def __init__(self) -> None:
+        self._segments: list[torch.cuda.CUDAGraph] = []
+        self._break_fns: list[Callable[[], Any]] = []
 
-    def capture_begin(self, pool=None, capture_error_mode: str = "global") -> None:
-        _check_cuda_bindings()
-        super().capture_begin(pool, capture_error_mode)
-        stream = get_current_stream()
-        # torch graph will not record any operation but only for compatibility
-        _end_capture_segment(stream)
-        _begin_capture_segment(stream)
-
-    def capture_end(self):
-        stream = get_current_stream()
-        self.last_graph = _end_capture_segment(stream)
-        self.last_graph_exec = _instantiate_graph(self.last_graph)
-        breaks = _captured_graphs_var.get()
-        self._exec = []
-        if breaks:
-            for replay_fn, output, handle in breaks:
-                graph_exec = _instantiate_graph(handle)
-                self._exec.append(GraphBreakInfo(replay_fn, output, graph_exec))
-
-        # start a dummy capture so torch's capture_end() can finalize
-        _begin_capture_segment(stream)
-        super().capture_end()
-
-    def replay(self):
+    def replay(self) -> None:
         stream = torch.cuda.current_stream()
         token = _current_stream_var.set(stream)
         try:
-            if not self._exec:
-                _replay_graph(self.last_graph_exec, stream.cuda_stream)
-                return
-            for func, _, handle in self._exec:
-                _replay_graph(handle, stream.cuda_stream)
-                func()
-            _replay_graph(self.last_graph_exec, stream.cuda_stream)
+            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 __del__(self):
-        try:
-            if hasattr(self, "_exec"):
-                for _, _, handle in self._exec:
-                    _destroy_graph_exec(handle)
-            if hasattr(self, "last_graph_exec"):
-                _destroy_graph_exec(self.last_graph_exec)
-        except Exception:
-            pass
 
+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.
+    """
 
-class BreakableCUDAGraphCapture(torch.cuda.graph):
     def __init__(
         self,
         cuda_graph: BreakableCUDAGraph,
@@ -322,31 +275,66 @@ class BreakableCUDAGraphCapture(torch.cuda.graph):
         stream: torch.cuda.Stream | None = None,
         capture_error_mode: str = "global",
     ):
-        super().__init__(
-            cuda_graph, pool=pool, stream=stream, capture_error_mode=capture_error_mode
-        )
-        self._stream = stream
         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
 
     def __enter__(self):
         _install_wait_stream_hook()
-        self._breaks_token = _captured_graphs_var.set([])
-        self._stream_token = _current_stream_var.set(self._stream)
-        self._forked_streams_token = _forked_streams_var.set(set())
-        return super().__enter__()
+        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):
-        super().__exit__(*args)
-        _current_stream_var.reset(self._stream_token)
-        _captured_graphs_var.reset(self._breaks_token)
-        _forked_streams_var.reset(self._forked_streams_token)
-        _uninstall_wait_stream_hook()
+        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:
+        graph = torch.cuda.CUDAGraph()
+        graph.capture_begin(
+            pool=self._pool, capture_error_mode=self._capture_error_mode
+        )
+        self.cuda_graph._segments.append(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_capturing(side.cuda_stream):
+                    _original_wait_stream(main_stream, side)
+            forked.clear()
+        self.cuda_graph._segments[-1].capture_end()
 
 
 @eager_on_graph(True)
-def break_graph():
+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/breakable_cuda_graph/context.py

@@ -0,0 +1,39 @@
+# 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 (BCG) runner.
+
+Kept intentionally separate from ``compilation/piecewise_context_manager.py``:
+BCG no longer inherits from the torch.compile-based PCG path, so its
+capture/replay lifecycle is managed on its own.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+
+_in_breakable_cuda_graph = False
+
+
+def is_in_breakable_cuda_graph() -> bool:
+    return _in_breakable_cuda_graph
+
+
+@contextmanager
+def enable_breakable_cuda_graph():
+    global _in_breakable_cuda_graph
+    _in_breakable_cuda_graph = True
+    try:
+        yield
+    finally:
+        _in_breakable_cuda_graph = False

python/sglang/srt/model_executor/breakable_cuda_graph_runner.py

@@ -0,0 +1,402 @@
+# 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 (BCG) runner.
+
+Captures the model forward as a sequence of ``torch.cuda.CUDAGraph`` segments
+split at attention layers. Functionally parallel to the torch.compile-based
+PCG runner but does not depend on torch.compile or FX graph splitting — graph
+breaks are inserted eagerly via :func:`eager_on_graph` decorated callables
+(radix attention for dense models, mamba for hybrid models).
+"""
+
+from __future__ import annotations
+
+import bisect
+import logging
+from typing import TYPE_CHECKING, Union
+
+import torch
+import tqdm
+
+from sglang.srt.compilation.piecewise_context_manager import set_forward_context
+from sglang.srt.distributed import get_tensor_model_parallel_rank
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    set_graph_pool_id,
+)
+from sglang.srt.distributed.parallel_state import graph_capture
+from sglang.srt.layers.dp_attention import (
+    set_dp_buffer_len,
+    set_is_extend_in_batch,
+)
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.layers.pooler import EmbeddingPoolerOutput
+from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import (
+    BreakableCUDAGraph,
+    BreakableCUDAGraphCapture,
+)
+from sglang.srt.model_executor.breakable_cuda_graph.context import (
+    enable_breakable_cuda_graph,
+)
+from sglang.srt.model_executor.cuda_graph_runner import (
+    get_global_graph_memory_pool,
+    set_global_graph_memory_pool,
+)
+from sglang.srt.model_executor.forward_batch_info import (
+    PPProxyTensors,
+)
+from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
+    PiecewiseCudaGraphRunner,
+    freeze_gc,
+)
+from sglang.srt.utils import get_available_gpu_memory, log_info_on_rank0
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+
+class BreakableCudaGraphRunner:
+    """Breakable CUDA graph runner.
+
+    Captures the model forward as a series of ``torch.cuda.CUDAGraph`` segments
+    with graph breaks at attention layers. Simpler than the torch.compile-based
+    PCG runner: no FX tracing, no compiled-kernel fusion — just segment-level
+    graph capture of the eager kernel stream.
+    """
+
+    # replay_prepare shares its buffer-population logic with the PCG runner —
+    # bind the method here without inheriting. __init__, capture, and replay
+    # diverge enough that inheritance would obscure more than it saves.
+    replay_prepare = PiecewiseCudaGraphRunner.replay_prepare
+
+    def __init__(self, model_runner: ModelRunner):
+        self.model_runner = model_runner
+        self.device = model_runner.device
+        self.device_module = torch.get_device_module(self.device)
+        self.graphs = {}
+        self.output_buffers = {}
+
+        self.quant_config = getattr(model_runner.model, "quant_config", None)
+        self.is_multimodal = model_runner.is_multimodal
+        # Read by the shared replay_prepare (bound from PiecewiseCudaGraphRunner).
+        self.capture_return_pooled_hidden_states = not model_runner.is_generation
+
+        # Capture sizes
+        capture_tokens = model_runner.server_args.piecewise_cuda_graph_tokens
+        assert capture_tokens is not None
+        self.capture_num_tokens = sorted(capture_tokens)
+        self.max_num_tokens = (
+            max(self.capture_num_tokens) if self.capture_num_tokens else 8192
+        )
+        self.max_bs = model_runner.req_to_token_pool.size
+
+        log_info_on_rank0(
+            logger,
+            f"[BCG] Capture num tokens: {self.capture_num_tokens}",
+        )
+
+        self._init_buffers(model_runner)
+
+        self.attention_layers = model_runner.attention_layers
+        self.moe_layers = model_runner.moe_layers
+        self.moe_fusions = model_runner.moe_fusions
+
+        with torch.device(self.device):
+            self.static_seq_lens = torch.zeros((self.max_bs,), dtype=torch.int64)
+            self.static_extend_seq_lens = torch.zeros((self.max_bs,), dtype=torch.int64)
+            self.static_extend_prefix_lens = torch.zeros(
+                (self.max_bs,), dtype=torch.int64
+            )
+            self.static_extend_start_loc = torch.zeros(
+                (self.max_bs,), dtype=torch.int64
+            )
+            self.static_req_pool_indices = torch.zeros(
+                (self.max_bs,), dtype=torch.int64
+            )
+            self.static_orig_seq_lens = torch.zeros((self.max_bs,), dtype=torch.int64)
+
+        # Memory pool
+        if get_global_graph_memory_pool() is None:
+            set_global_graph_memory_pool(self.device_module.graph_pool_handle())
+        set_graph_pool_id(get_global_graph_memory_pool())
+
+        # Warmup then capture
+        self._warmup()
+        self.device_module.synchronize()
+        self.model_runner.tp_group.barrier()
+        self._capture_all()
+
+        self.raw_num_tokens = 0
+
+    def _init_buffers(self, model_runner):
+        """Initialize input buffers."""
+        from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
+            PrefillInputBuffers,
+        )
+        from sglang.srt.utils import is_npu
+
+        with torch.device(self.device):
+            input_ids = torch.zeros((self.max_num_tokens,), dtype=torch.int64)
+            out_cache_loc = torch.zeros(
+                (self.max_num_tokens,),
+                dtype=torch.int64 if not is_npu() else torch.int32,
+            )
+            out_cache_loc_swa = (
+                torch.zeros((self.max_num_tokens,), dtype=torch.int64)
+                if model_runner.is_hybrid_swa
+                else None
+            )
+            positions = torch.zeros((self.max_num_tokens,), dtype=torch.int64)
+            if self.is_multimodal:
+                input_embeds = torch.zeros(
+                    (self.max_num_tokens, model_runner.model_config.hidden_size),
+                    dtype=model_runner.dtype,
+                )
+                mrope_positions = torch.zeros(
+                    (3, self.max_num_tokens), dtype=torch.int64
+                )
+            else:
+                input_embeds = None
+                mrope_positions = None
+
+        self.buffers = PrefillInputBuffers(
+            input_ids=input_ids,
+            out_cache_loc=out_cache_loc,
+            out_cache_loc_swa=out_cache_loc_swa,
+            mamba_track_indices=None,
+            mamba_track_mask=None,
+            mamba_track_seqlens=None,
+            positions=positions,
+            input_embeds=input_embeds,
+            mrope_positions=mrope_positions,
+        )
+        self.buffers.share_buffers()
+
+    def _run_forward(self, forward_batch, num_tokens):
+        """Run model forward with proper context."""
+        forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
+        set_dp_buffer_len(None, num_tokens, forward_batch.dp_padding_mode.is_max_len())
+        set_is_extend_in_batch(False)
+
+        with set_forward_context(
+            forward_batch,
+            self.attention_layers,
+            self.quant_config,
+            self.moe_layers,
+            self.moe_fusions,
+        ):
+            output = self.model_runner.model.forward(
+                forward_batch.input_ids,
+                forward_batch.positions,
+                forward_batch,
+            )
+        return output
+
+    def _build_capture_forward_batch(self, num_tokens):
+        """Build a ForwardBatch for capture using static buffers for stable addresses."""
+        from sglang.srt.layers.dp_attention import DpPaddingMode
+        from sglang.srt.model_executor.forward_batch_info import (
+            CaptureHiddenMode,
+            ForwardBatch,
+            ForwardMode,
+        )
+
+        buffers = self.buffers
+        bs = 1
+        self.static_seq_lens[:bs].fill_(num_tokens)
+        self.static_extend_seq_lens[:bs].fill_(num_tokens)
+        self.static_extend_prefix_lens[:bs].zero_()
+        self.static_extend_start_loc[:bs].zero_()
+        self.static_req_pool_indices[:bs].copy_(torch.arange(bs, device=self.device))
+        self.static_orig_seq_lens[:bs].fill_(num_tokens)
+
+        return ForwardBatch(
+            forward_mode=ForwardMode.EXTEND,
+            batch_size=bs,
+            input_ids=buffers.input_ids[:num_tokens],
+            input_embeds=(
+                buffers.input_embeds[:num_tokens] if self.is_multimodal else None
+            ),
+            req_pool_indices=self.static_req_pool_indices[:bs],
+            seq_lens=self.static_seq_lens[:bs],
+            next_token_logits_buffer=None,
+            orig_seq_lens=self.static_orig_seq_lens[:bs],
+            seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
+            req_to_token_pool=self.model_runner.req_to_token_pool,
+            token_to_kv_pool=self.model_runner.token_to_kv_pool,
+            attn_backend=self.model_runner.attn_backend,
+            out_cache_loc=buffers.out_cache_loc[:num_tokens],
+            out_cache_loc_swa=(
+                buffers.out_cache_loc_swa[:num_tokens]
+                if buffers.out_cache_loc_swa is not None
+                else None
+            ),
+            seq_lens_sum=num_tokens,
+            mamba_track_indices=None,
+            mamba_track_mask=None,
+            mamba_track_seqlens=None,
+            encoder_lens=None,
+            return_logprob=False,
+            extend_num_tokens=num_tokens,
+            extend_seq_lens=self.static_extend_seq_lens[:bs],
+            extend_prefix_lens=self.static_extend_prefix_lens[:bs],
+            extend_start_loc=self.static_extend_start_loc[:bs],
+            extend_prefix_lens_cpu=torch.tensor([0], device="cpu"),
+            extend_seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
+            extend_logprob_start_lens_cpu=torch.tensor([num_tokens], device="cpu"),
+            positions=buffers.positions[:num_tokens],
+            global_num_tokens_gpu=None,
+            global_num_tokens_for_logprob_gpu=None,
+            dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
+            global_dp_buffer_len=None,
+            mrope_positions=(
+                buffers.mrope_positions[:, :num_tokens] if self.is_multimodal else None
+            ),
+            spec_algorithm=None,
+            spec_info=None,
+            capture_hidden_mode=CaptureHiddenMode.NULL,
+            num_token_non_padded=None,
+            global_forward_mode=ForwardMode.EXTEND,
+            lora_ids=None,
+        )
+
+    def _warmup(self):
+        """Warmup the model with a forward pass."""
+        num_tokens = self.capture_num_tokens[0]
+        forward_batch = self._build_capture_forward_batch(num_tokens)
+        self.model_runner.attn_backend.init_forward_metadata(forward_batch)
+        self._run_forward(forward_batch, num_tokens)
+
+    def _capture_all(self):
+        """Capture breakable CUDA graphs for all token sizes."""
+        with freeze_gc(
+            self.model_runner.server_args.enable_cudagraph_gc
+        ), graph_capture() as graph_capture_context, enable_breakable_cuda_graph():
+            stream = graph_capture_context.stream
+            pool = get_global_graph_memory_pool()
+
+            capture_range = (
+                tqdm.tqdm(list(reversed(self.capture_num_tokens)))
+                if get_tensor_model_parallel_rank() == 0
+                else reversed(self.capture_num_tokens)
+            )
+            for num_tokens in capture_range:
+                if get_tensor_model_parallel_rank() == 0:
+                    avail_mem = get_available_gpu_memory(
+                        self.model_runner.device,
+                        self.model_runner.gpu_id,
+                        empty_cache=False,
+                    )
+                    capture_range.set_description(
+                        f"[BCG] Capturing ({num_tokens=} {avail_mem=:.2f} GB)"
+                    )
+
+                graph, output = self._capture_one(num_tokens, pool, stream)
+                self.graphs[num_tokens] = graph
+                self.output_buffers[num_tokens] = output
+
+    def can_run(self, forward_batch: "ForwardBatch"):
+        # BCG graphs are captured with batch_size=1 (see _build_capture_forward_batch);
+        # the captured logits-gather / sampler path yields bs=1 outputs. Multi-req
+        # prefill would silently return wrong-shaped logits, corrupting downstream
+        # output_ids and breaking the subsequent decode step. Reject here so the
+        # caller falls back to the eager extend path.
+        if forward_batch.batch_size > 1:
+            return False
+        if forward_batch.input_embeds is not None:
+            return False
+        if forward_batch.replace_embeds is not None:
+            return False
+        num_tokens = len(forward_batch.input_ids)
+        if forward_batch.return_logprob:
+            for start_len, seq_len in zip(
+                forward_batch.extend_logprob_start_lens_cpu,
+                forward_batch.extend_seq_lens_cpu,
+            ):
+                if start_len is not None and start_len < seq_len:
+                    return False
+        return num_tokens <= self.max_num_tokens
+
+    def _capture_one(self, num_tokens, pool, stream):
+        """Capture a breakable CUDA graph for one token size."""
+        forward_batch = self._build_capture_forward_batch(num_tokens)
+        self.model_runner.attn_backend.init_forward_metadata(forward_batch)
+
+        def run_once():
+            return self._run_forward(forward_batch, num_tokens)
+
+        for _ in range(2):
+            self.device_module.synchronize()
+            self.model_runner.tp_group.barrier()
+            run_once()
+
+        graph = BreakableCUDAGraph()
+        with BreakableCUDAGraphCapture(cuda_graph=graph, pool=pool, stream=stream):
+            output = run_once()
+
+        return graph, output
+
+    def replay(
+        self,
+        forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
+        num_tokens = len(forward_batch.input_ids)
+        index = bisect.bisect_left(self.capture_num_tokens, num_tokens)
+        static_num_tokens = self.capture_num_tokens[index]
+
+        with enable_breakable_cuda_graph():
+            static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
+            bs = forward_batch.batch_size
+
+            # Update static buffers used by graph segments (esp. logits processor).
+            # The graph reads from these addresses — they must have serving-time values.
+            self.static_seq_lens[:bs].copy_(forward_batch.seq_lens)
+            self.static_extend_seq_lens[:bs].copy_(forward_batch.extend_seq_lens)
+            self.static_extend_prefix_lens[:bs].copy_(forward_batch.extend_prefix_lens)
+            self.static_extend_start_loc[:bs].copy_(forward_batch.extend_start_loc)
+            self.static_req_pool_indices[:bs].copy_(forward_batch.req_pool_indices)
+            if forward_batch.orig_seq_lens is not None:
+                self.static_orig_seq_lens[:bs].copy_(forward_batch.orig_seq_lens)
+
+            # Set forward context and replay
+            self.model_runner.attn_backend.init_forward_metadata(forward_batch)
+            with set_forward_context(
+                static_forward_batch,
+                self.attention_layers,
+                self.quant_config,
+                self.moe_layers,
+                self.moe_fusions,
+            ):
+                self.graphs[static_num_tokens].replay()
+
+        output = self.output_buffers[static_num_tokens]
+        if isinstance(output, LogitsProcessorOutput):
+            return LogitsProcessorOutput(
+                next_token_logits=output.next_token_logits[: self.raw_num_tokens],
+                hidden_states=(
+                    output.hidden_states[: self.raw_num_tokens]
+                    if output.hidden_states is not None
+                    else None
+                ),
+            )
+        elif isinstance(output, EmbeddingPoolerOutput):
+            return output
+        else:
+            assert isinstance(output, PPProxyTensors)
+            raise NotImplementedError(
+                "PPProxyTensors is not supported in BreakableCudaGraphRunner."
+            )

python/sglang/srt/model_executor/model_runner.py

@@ -124,6 +124,9 @@ from sglang.srt.lora.lora_registry import LoRARef
 from sglang.srt.managers.schedule_batch import sanity_check_mm_pad_shift_value
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.model_executor.breakable_cuda_graph_runner import (
+    BreakableCudaGraphRunner,
+)
 from sglang.srt.model_executor.cpu_graph_runner import CPUGraphRunner
 from sglang.srt.model_executor.cuda_graph_runner import (
     CudaGraphRunner,
@@ -2726,7 +2729,11 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             f"Capture piecewise CUDA graph begin. avail mem={before_mem:.2f} GB"
         )
 
-        self.piecewise_cuda_graph_runner = PiecewiseCudaGraphRunner(self)
+        if self.server_args.enable_breakable_cuda_graph:
+            # Experimental feature
+            self.piecewise_cuda_graph_runner = BreakableCudaGraphRunner(self)
+        else:
+            self.piecewise_cuda_graph_runner = PiecewiseCudaGraphRunner(self)
 
         after_mem = get_available_gpu_memory(self.device, self.gpu_id)
         mem_usage = before_mem - after_mem

python/sglang/srt/models/nemotron_h.py

@@ -59,6 +59,12 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import (
+    eager_on_graph,
+)
+from sglang.srt.model_executor.breakable_cuda_graph.context import (
+    is_in_breakable_cuda_graph,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
@@ -429,6 +435,11 @@ class NemotronHMambaDecoderLayer(nn.Module):
         else:
             hidden_states, residual = self.norm(hidden_states, residual)
 
+        if is_in_breakable_cuda_graph():
+            output = torch.empty_like(hidden_states)
+            breakable_nemotron_mamba2_with_output(hidden_states, output, self.layer_id)
+            return output, residual
+
         if is_in_piecewise_cuda_graph():
             output = torch.empty_like(hidden_states)
             nemotron_mamba2_with_output(hidden_states, output, self.layer_id)
@@ -912,4 +923,8 @@ def nemotron_mamba2_with_output(
 
     # Copy result back; output may be larger (padded) so only fill actual tokens
     output[:num_actual_tokens].view(ret.shape).copy_(ret)
-    return
+
+
+breakable_nemotron_mamba2_with_output = eager_on_graph(True)(
+    nemotron_mamba2_with_output
+)

python/sglang/srt/server_args.py

@@ -624,6 +624,7 @@ class ServerArgs:
     cuda_graph_bs: Optional[List[int]] = None
     disable_cuda_graph: bool = False
     disable_cuda_graph_padding: bool = False
+    enable_breakable_cuda_graph: bool = False
     enable_profile_cuda_graph: bool = False
     enable_cudagraph_gc: bool = False
     debug_cuda_graph: bool = False
@@ -5852,6 +5853,11 @@ class ServerArgs:
             action="store_true",
             help="Disable cuda graph when padding is needed. Still uses cuda graph when padding is not needed.",
         )
+        parser.add_argument(
+            "--enable-breakable-cuda-graph",
+            action="store_true",
+            help="Use breakable CUDA graph for piecewise capture instead of torch.compile-based splitting.",
+        )
         parser.add_argument(
             "--enable-profile-cuda-graph",
             action="store_true",

test/registered/breakable_cuda_graph/test_breakable_cuda_graph.py

@@ -1,19 +1,29 @@
-"""Unit tests for the breakable CUDA graph mechanism.
+"""Tests for the breakable CUDA graph (BCG) runner.
 
-Tests the core capture/replay logic with simple tensor operations,
-verifying that graph breaks work correctly and outputs are properly
-propagated across segments.
+Two test classes:
+- ``TestBreakableCUDAGraphBasic`` / ``TestCopyOutput`` / ``TestBreakGraphHelper``:
+  unit tests for the core capture / replay mechanism (simple tensor ops).
+- ``TestBreakableCudaGraph``: integration test — spin up Qwen3-8B with
+  ``--enable-breakable-cuda-graph`` and check mgsm_en accuracy.
 """
 
 import unittest
 
 import torch
 
+from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    SimpleNamespace,
+    popen_launch_server,
+)
 
-# CI Registration
-register_cuda_ci(est_time=30, suite="stage-b-test-1-gpu-small")
+# CI Registration — large suite to fit the integration test's server startup.
+register_cuda_ci(est_time=130, suite="stage-b-test-1-gpu-large")
 
 
 def _skip_if_no_cuda(test_func):
@@ -281,5 +291,41 @@ class TestBreakGraphHelper(CustomTestCase):
         self.assertTrue(torch.allclose(y, torch.full((4,), 13.0, device=self.device)))
 
 
+class TestBreakableCudaGraph(CustomTestCase):
+    """Integration: Qwen3-8B with --enable-breakable-cuda-graph on mgsm_en."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "Qwen/Qwen3-8B"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--enable-breakable-cuda-graph",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k_accuracy(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="mgsm_en",
+            num_examples=1319,
+            num_threads=1024,
+        )
+
+        metrics = run_eval(args)
+        score = metrics["score"]
+        print(f"mgsm_en accuracy with breakable CUDA graph: {score:.3f}")
+
+        self.assertGreaterEqual(score, 0.80)
+
+
 if __name__ == "__main__":
     unittest.main()