[Refactor] Cuda Graph Runner/Backend Refactor (#23906)

Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>

python/sglang/auto_benchmark_lib.py

@@ -58,7 +58,7 @@ PROGRESS_FLAG_KEYS = (
     "num_continuous_decode_steps",
     "stream_interval",
     "page_size",
-    "cuda_graph_max_bs",
+    "cuda_graph_max_bs_decode",
     "speculative_num_steps",
     "speculative_eagle_topk",
     "speculative_num_draft_tokens",
@@ -84,7 +84,7 @@ PROGRESS_FLAG_ALIASES = {
     "num_continuous_decode_steps": "decode_steps",
     "stream_interval": "stream",
     "page_size": "page",
-    "cuda_graph_max_bs": "cg_bs",
+    "cuda_graph_max_bs_decode": "cg_bs",
     "speculative_num_steps": "spec_steps",
     "speculative_eagle_topk": "eagle_topk",
     "speculative_num_draft_tokens": "draft_tok",

python/sglang/bench_one_batch.py

@@ -77,6 +77,7 @@ from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.managers.scheduler_components.dp_attn import prepare_mlp_sync_batch_raw
 from sglang.srt.mem_cache.base_prefix_cache import EvictParams
+from sglang.srt.model_executor.cuda_graph_config import Phase
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.sampling.sampling_params import SamplingParams
@@ -944,7 +945,10 @@ def latency_test(
 
 
 def main(server_args, bench_args):
-    server_args.cuda_graph_max_bs = max(bench_args.batch_size)
+    # Post-init write to the legacy cuda_graph_max_bs_decode field would
+    # not propagate to cuda_graph_config; update the decode phase directly.
+    if server_args.cuda_graph_config is not None:
+        server_args.cuda_graph_config[Phase.DECODE].max_bs = max(bench_args.batch_size)
 
     _set_envs_and_config(server_args)
 

python/sglang/compile_deep_gemm.py

@@ -23,6 +23,7 @@ from sglang.srt.entrypoints.warmup import warmup
 from sglang.srt.environ import envs
 from sglang.srt.managers.io_struct import GenerateReqInput
 from sglang.srt.managers.tokenizer_manager import TokenizerManager
+from sglang.srt.model_executor.cuda_graph_config import Backend, Phase
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import kill_process_tree
 
@@ -169,8 +170,11 @@ def launch_server_process_and_send_one_request(
 
 
 def refine_server_args(server_args: ServerArgs, compile_args: CompileArgs):
-    # Disable cuda graph and torch compile to save time
-    server_args.disable_cuda_graph = True
+    # Disable cuda graph and torch compile to save time. Writes after
+    # ServerArgs.__post_init__ don't propagate to cuda_graph_config via the
+    # legacy disable_cuda_graph field, so flip both phases directly.
+    server_args.cuda_graph_config[Phase.DECODE].backend = Backend.DISABLED
+    server_args.cuda_graph_config[Phase.PREFILL].backend = Backend.DISABLED
     server_args.enable_torch_compile = False
     print(f"Disable CUDA Graph and Torch Compile to save time...")
 

python/sglang/srt/arg_groups/argparse_actions.py

@@ -67,6 +67,34 @@ class DeprecatedStoreTrueAction(argparse.Action):
         setattr(namespace, self.dest, True)
 
 
+class DeprecatedStoreConstAction(argparse.Action):
+    """Deprecated boolean flag that stores a fixed string/value into ``dest``
+    and prints a warning. Used to translate a legacy boolean flag into a
+    setting on the new per-phase config dict (e.g.
+    ``--disable-piecewise-cuda-graph`` -> ``cuda_graph_backend_prefill="disabled"``)."""
+
+    def __init__(
+        self,
+        option_strings,
+        dest,
+        new_flag=None,
+        const_value=None,
+        nargs=0,
+        default=None,
+        **kwargs,
+    ):
+        self.new_flag = new_flag
+        self.const_value = const_value
+        super().__init__(option_strings, dest, nargs=nargs, default=default, **kwargs)
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        replacement = f" Use '{self.new_flag}' instead." if self.new_flag else ""
+        print_deprecated_warning(
+            f"'{option_string}' is deprecated and will be removed in a future release.{replacement}"
+        )
+        setattr(namespace, self.dest, self.const_value)
+
+
 class DeprecatedAliasStoreAction(argparse.Action):
     """Deprecated alias that stores its value and prints a warning."""
 

python/sglang/srt/compilation/compile.py

@@ -9,7 +9,9 @@ from typing import Any, Callable, Optional, Union
 import torch
 
 from sglang.srt.compilation.compilation_config import CompilationConfig
-from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -186,7 +188,7 @@ def install_torch_compiled(
         state["compiled_callable"] = compiled_callable
 
     def trampoline(self, *args, **kwargs):
-        use_compiled = is_in_piecewise_cuda_graph()
+        use_compiled = is_in_tc_piecewise_cuda_graph()
         if use_compiled:
             if not state["compiled"]:
                 _ensure_compiled(self, *args, **kwargs)

python/sglang/srt/compilation/compile_phase.py

@@ -0,0 +1,57 @@
+"""torch.compile-internal phase markers used by the tc_piecewise backend.
+
+Two pieces of state, both private to the torch.compile path (the
+``cuda_piecewise_backend`` FX backend and the runner that drives it):
+
+* ``_in_torch_compile_warmup`` — true during the warmup-compile loop
+  where we run the compiled callable to trigger inductor compilation
+  but explicitly do **not** capture into a CUDA graph yet.
+  ``cuda_piecewise_backend`` reads this to short-circuit the capture
+  branch.
+* ``_pcg_capture_stream`` — the CUDA stream on which the runner is
+  performing capture, surfaced so the FX backend can use the same
+  stream for its own ``torch.cuda.graph(...)`` calls.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+
+import torch
+
+_in_torch_compile_warmup = False
+_pcg_capture_stream: "torch.cuda.Stream | None" = None
+
+
+def is_in_torch_compile_warmup() -> bool:
+    """True while inside the tc_piecewise warmup-compile pass. Strict subset of
+    ``torch.compiler.is_compiling()``.
+    """
+    return _in_torch_compile_warmup
+
+
+@contextmanager
+def enable_torch_compile_warmup():
+    """Mark the enclosed scope as the tc_piecewise warmup-compile pass. The FX
+    piecewise backend uses this to skip CUDA graph capture during warmup.
+    """
+    global _in_torch_compile_warmup
+    _in_torch_compile_warmup = True
+    try:
+        yield
+    finally:
+        _in_torch_compile_warmup = False
+
+
+def get_pcg_capture_stream() -> "torch.cuda.Stream | None":
+    return _pcg_capture_stream
+
+
+@contextmanager
+def set_pcg_capture_stream(stream: torch.cuda.Stream):
+    global _pcg_capture_stream
+    _pcg_capture_stream = stream
+    try:
+        yield
+    finally:
+        _pcg_capture_stream = None

python/sglang/srt/compilation/cuda_piecewise_backend.py

@@ -13,9 +13,9 @@ import torch.fx as fx
 
 from sglang.srt.compilation.compilation_config import CompilationConfig
 from sglang.srt.compilation.compilation_counter import compilation_counter
-from sglang.srt.compilation.piecewise_context_manager import (
+from sglang.srt.compilation.compile_phase import (
     get_pcg_capture_stream,
-    is_in_pcg_torch_compile,
+    is_in_torch_compile_warmup,
 )
 from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors
 from sglang.srt.utils import is_hip
@@ -145,7 +145,7 @@ class CUDAPiecewiseBackend:
             if self.is_last_graph and not self.to_be_compiled_sizes:
                 self.check_for_ending_compilation()
 
-        if is_in_pcg_torch_compile():
+        if is_in_torch_compile_warmup():
             return entry.runnable(*args)
 
         if entry.cudagraph is None:

python/sglang/srt/compilation/piecewise_context_manager.py

@@ -1,138 +0,0 @@
-from __future__ import annotations
-
-import logging
-from contextlib import contextmanager
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, List, Optional
-
-import torch
-
-logger = logging.getLogger(__name__)
-
-
-if TYPE_CHECKING:
-    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
-
-_in_piecewise_cuda_graph = False
-_in_pcg_torch_compile = False
-_pcg_capture_stream = None
-
-
-def is_in_piecewise_cuda_graph():
-    return _in_piecewise_cuda_graph
-
-
-def is_in_pcg_torch_compile():
-    return _in_pcg_torch_compile
-
-
-def get_pcg_capture_stream():
-    return _pcg_capture_stream
-
-
-@contextmanager
-def enable_piecewise_cuda_graph_compile():
-    global _in_pcg_torch_compile
-    _in_pcg_torch_compile = True
-    yield
-    _in_pcg_torch_compile = False
-
-
-@contextmanager
-def enable_piecewise_cuda_graph():
-    global _in_piecewise_cuda_graph
-    _in_piecewise_cuda_graph = True
-    try:
-        yield
-    except Exception as e:
-        logger.error(
-            "Piecewise CUDA Graph failed with error: %s\n%s",
-            e,
-            PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG,
-        )
-        raise
-    finally:
-        _in_piecewise_cuda_graph = False
-
-
-@contextmanager
-def set_pcg_capture_stream(stream: torch.cuda.Stream):
-    global _pcg_capture_stream
-    _pcg_capture_stream = stream
-    yield
-    _pcg_capture_stream = None
-
-
-@dataclass
-class ForwardContext:
-    def __init__(self):
-        self.forward_batch = None
-        self.attention_layers = None
-        self.quant_config = None
-        self.moe_layers = None
-        self.moe_fusions = None
-        self.dsa_indexers = None
-        self.num_tokens: Optional[int] = None
-        self.raw_num_tokens: Optional[int] = None
-
-    def set_forward_batch(self, forward_batch: ForwardBatch):
-        self.forward_batch = forward_batch
-
-    def set_attention_layers(self, layers: List[Any]):
-        self.attention_layers = layers
-
-    def set_quant_config(self, quant_config: Any):
-        self.quant_config = quant_config
-
-    def set_moe_layers(self, layers: List[Any]):
-        self.moe_layers = layers
-
-    def set_moe_fusions(self, fusions: List[Any]):
-        self.moe_fusions = fusions
-
-    def set_dsa_indexers(self, indexers: List[Any]):
-        self.dsa_indexers = indexers
-
-
-_forward_context: Optional[ForwardContext] = None
-
-
-def get_forward_context() -> Optional[ForwardContext]:
-    if _forward_context is None:
-        return None
-    return _forward_context
-
-
-@contextmanager
-def set_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 _forward_context
-    _forward_context = ForwardContext()
-    _forward_context.set_forward_batch(forward_batch)
-    _forward_context.set_attention_layers(attention_layers)
-    _forward_context.set_quant_config(quant_config)
-    _forward_context.set_moe_layers(moe_layers)
-    _forward_context.set_moe_fusions(moe_fusions)
-    if dsa_indexers is not None:
-        _forward_context.set_dsa_indexers(dsa_indexers)
-    _forward_context.num_tokens = num_tokens
-    _forward_context.raw_num_tokens = raw_num_tokens
-    try:
-        yield
-    finally:
-        _forward_context = None
-
-
-PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG = (
-    "Piecewise CUDA Graph is enabled by default as an experimental feature.\n"
-    "To work around this error, add --disable-piecewise-cuda-graph to your launch command.\n"
-    "Please report this issue at https://github.com/sgl-project/sglang/issues/new/choose"
-)

python/sglang/srt/compilation/torch_compile_decoration.py

@@ -0,0 +1,83 @@
+"""torch.compile decoration helpers used by the decode-Full path under
+``--enable-torch-compile``.
+
+``patch_model`` wraps the model forward with ``torch.compile`` for batch
+sizes that fall in the compile bucket list and returns the raw forward
+otherwise. ``set_torch_compile_config`` flips the inductor/dynamo config
+flags expected by that path.
+
+Note: the prefill-tc_piecewise path (``TcPiecewiseCudaGraphBackend``) does NOT
+use ``patch_model`` — it goes through ``compilation/compile.py``'s
+``install_torch_compiled``. ``_to_torch`` here is duplicated by
+tc_piecewise's local ``_toggle_multi_platform_ops``; the duplication is kept
+because the two paths have different lifecycle requirements.
+"""
+
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+
+import torch
+
+from sglang.srt.distributed.parallel_state import GroupCoordinator
+from sglang.srt.layers.utils import MultiPlatformOp
+from sglang.srt.utils import get_bool_env_var, is_hip
+from sglang.srt.utils.patch_torch import monkey_patch_torch_compile
+
+_is_hip = is_hip()
+
+
+def _to_torch(model: torch.nn.Module, reverse: bool, num_tokens: int) -> None:
+    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):
+            _to_torch(sub, reverse, num_tokens)
+
+
+@contextmanager
+def patch_model(
+    model: torch.nn.Module,
+    enable_compile: bool,
+    num_tokens: int,
+    tp_group: GroupCoordinator,
+):
+    """Patch the model to make it compatible with torch.compile."""
+    backup_ca_comm = None
+
+    try:
+        if enable_compile:
+            _to_torch(model, reverse=False, num_tokens=num_tokens)
+            backup_ca_comm = tp_group.ca_comm
+            yield torch.compile(
+                torch.no_grad()(model.forward),
+                mode=os.environ.get(
+                    "SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs"
+                ),
+                dynamic=_is_hip and get_bool_env_var("SGLANG_TORCH_DYNAMIC_SHAPE"),
+            )
+        else:
+            yield model.forward
+    finally:
+        if enable_compile:
+            _to_torch(model, reverse=True, num_tokens=num_tokens)
+            tp_group.ca_comm = backup_ca_comm
+
+
+def set_torch_compile_config() -> None:
+    import torch._dynamo.config
+    import torch._inductor.config
+
+    torch._inductor.config.coordinate_descent_tuning = True
+    torch._inductor.config.triton.unique_kernel_names = True
+    torch._inductor.config.fx_graph_cache = True
+
+    torch._dynamo.config.accumulated_cache_size_limit = 1024
+    if hasattr(torch._dynamo.config, "cache_size_limit"):
+        torch._dynamo.config.cache_size_limit = 1024
+
+    monkey_patch_torch_compile()

python/sglang/srt/debug_utils/pr_fix_toggle.py

@@ -24,7 +24,7 @@ patches:
         replacement: |
           hidden_states = logits_output.hidden_states
 
-  - target: sglang.srt.speculative.eagle_draft_cuda_graph_runner.EAGLEDraftCudaGraphRunner.capture_one_batch_size
+  - target: sglang.srt.speculative.eagle_draft_cuda_graph_runner.EAGLEDraftCudaGraphRunner.capture_one_shape
     edits:
       - match: |
           forward_batch.spec_info.hidden_states = hidden_states_backup

python/sglang/srt/distributed/device_communicators/custom_all_reduce.py

@@ -13,13 +13,15 @@ import torch.distributed as dist
 from torch.distributed import ProcessGroup
 
 import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops
-from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
 from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
 from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
     can_use_custom_all_reduce_with_nvlink,
     is_weak_contiguous,
 )
 from sglang.srt.environ import envs
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.utils import (
     get_bool_env_var,
     is_cuda,
@@ -316,7 +318,7 @@ class CustomAllreduce:
                 # Could be warmup OR piecewise cuda graph split op execution.
                 # In piecewise cuda graph, split ops run eagerly outside the graph
                 # but _IS_CAPTURING is still True. We need to do real all-reduce.
-                if is_in_piecewise_cuda_graph():
+                if is_in_tc_piecewise_cuda_graph():
                     # Split op execution - do real all-reduce
                     return self._all_reduce_impl(input, registered=False)
                 else:
@@ -350,8 +352,8 @@ def dispatch_custom_allreduce(
 
     On CUDA, the JIT-compiled v2 implementation is used by default.
     Set SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2=0 to fall back to the legacy CustomAllreduce.
-    Note: ``ServerArgs._handle_environment_variables`` forces this env to "0" when
-    ``nnodes > 1`` since custom AR is intra-node only.
+    Note: ServerArgs._handle_environment_variables forces this env to "0" when
+    nnodes > 1 since custom AR is intra-node only.
     """
     if _is_cuda and envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get():
         from .custom_all_reduce_v2 import (

python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py

@@ -8,11 +8,13 @@ import torch.distributed as dist
 from torch.distributed import ProcessGroup
 
 from sglang.jit_kernel.all_reduce import AllReduceAlgo, get_custom_all_reduce_cls
-from sglang.srt.distributed import is_in_piecewise_cuda_graph
 from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
     can_use_custom_all_reduce_with_nvlink,
     is_weak_contiguous,
 )
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.utils import is_sm100_supported, log_info_on_rank0
 
 logger = logging.getLogger(__name__)
@@ -121,7 +123,7 @@ class CustomAllReduceV2:
         return inp_size <= self.max_size
 
     def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor:
-        if is_in_piecewise_cuda_graph():  # disable inplace optimization
+        if is_in_tc_piecewise_cuda_graph():  # disable inplace optimization
             try:
                 self.obj.set_cuda_graph_capture(False)
                 return self._all_reduce(input)

python/sglang/srt/distributed/device_communicators/pymscclpp.py

@@ -7,10 +7,12 @@ import torch
 import torch.distributed as dist
 from torch.distributed import ProcessGroup, ReduceOp
 
-from sglang.srt.compilation.piecewise_context_manager import (
+from sglang.srt.compilation.compile_phase import (
     get_pcg_capture_stream,
-    is_in_pcg_torch_compile,
-    is_in_piecewise_cuda_graph,
+    is_in_torch_compile_warmup,
+)
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
 )
 from sglang.srt.server_args import get_global_server_args
 
@@ -343,8 +345,8 @@ class PyMscclppCommunicator:
         # (compile, capture, or replay) as it changes the allreduce dispatch
         # path and triggers recompilation.
         if (
-            is_in_piecewise_cuda_graph()
-            or is_in_pcg_torch_compile()
+            is_in_tc_piecewise_cuda_graph()
+            or is_in_torch_compile_warmup()
             or get_pcg_capture_stream() is not None
         ):
             return False

python/sglang/srt/distributed/parallel_state.py

@@ -43,9 +43,11 @@ import torch.distributed
 from torch.distributed import Backend, ProcessGroup
 
 from sglang.srt.compilation.compilation_config import register_split_op
-from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
 from sglang.srt.distributed.utils import set_global_tcp_store
 from sglang.srt.environ import envs
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.utils import (
     get_current_device_stream_fast,
     get_int_env_var,
@@ -641,7 +643,7 @@ class GroupCoordinator:
             and self.torch_symm_mem_comm.should_torch_symm_mem_allreduce(input_)
         ):
             outplace_all_reduce_method = "torch_symm_mem"
-        elif is_in_piecewise_cuda_graph() and self.pynccl_comm is not None:
+        elif is_in_tc_piecewise_cuda_graph() and self.pynccl_comm is not None:
             # For piecewise cuda graph, we use pynccl outplace allreduce
             outplace_all_reduce_method = "pynccl"
         if outplace_all_reduce_method is not None:
@@ -708,7 +710,7 @@ class GroupCoordinator:
         if (
             getattr(ca_comm, "_IS_CAPTURING", False)
             and not torch.cuda.is_current_stream_capturing()
-            and is_in_piecewise_cuda_graph()
+            and is_in_tc_piecewise_cuda_graph()
         ):
             if not hasattr(ca_comm, "fused_ar_rms"):
                 return None
@@ -863,7 +865,7 @@ class GroupCoordinator:
             if getattr(ca_comm, "_IS_CAPTURING", False):
                 if torch.cuda.is_current_stream_capturing():
                     ca_comm.all_gather_reg(input, out=output, dim=0)
-                elif is_in_piecewise_cuda_graph():
+                elif is_in_tc_piecewise_cuda_graph():
                     ca_comm.all_gather_unreg(input, out=output, dim=0)
                 else:
                     # True CUDA graph warmup: avoid a different host collective.

python/sglang/srt/hardware_backend/mlx/model_runner_stub.py

@@ -186,7 +186,7 @@ class MlxModelRunnerStub(ModelRunner):
         )
 
         # No CUDA graphs, no attention backend
-        self.graph_runner = None
+        self.decode_cuda_graph_runner = None
         self.graph_mem_usage = 0
         self.attn_backend = None
 

python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py

@@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, List, Optional
 
 import torch
 
+from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
     MultiLayerEagleDraftExtendCudaGraphRunner,
@@ -82,7 +83,7 @@ class MultiLayerEagleMultiStepDraftExtendNpuGraphRunner(
         super().__init__(eagle_worker)
 
     def _init_and_capture(self):
-        if self.eagle_worker.server_args.disable_cuda_graph:
+        if cuda_graph_fully_disabled():
             self.runners = [None] * self.speculative_num_steps
             return
 

python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py

@@ -0,0 +1,166 @@
+"""NPUCudaGraphBackend — Ascend NPU full-graph capture (torch.npu.NPUGraph).
+
+Mirrors FullCudaGraphBackend with two differences:
+  - Captures via torch.npu.graph(...) into torch.npu.NPUGraph.
+  - replay_with_input_update(shape_key, seq_lens, attr_name) rebinds
+    the recorded graph's input bindings for variable seq_lens at replay
+    time (NPU's NPUGraph.update(...) API).
+
+torch.npu is imported lazily inside methods so the module loads on
+non-NPU hosts.
+"""
+
+from __future__ import annotations
+
+import threading
+from contextlib import AbstractContextManager, contextmanager
+from functools import partial
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+
+import numpy as np
+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.utils import empty_context, 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,
+    )
+
+
+class NPUCudaGraphBackend(BaseCudaGraphBackend):
+    """One torch.npu.NPUGraph per shape; attention metadata captured
+    inside the graph. replay_with_input_update substitutes fresh
+    seq_lens without re-recording."""
+
+    def __init__(
+        self,
+        cuda_graph_runner: BaseCudaGraphRunner,
+        *,
+        enable_memory_saver: bool = False,
+    ) -> None:
+        self._graphs: Dict[Any, Any] = {}
+        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 = None
+        self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
+            enable=enable_memory_saver
+            and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
+        )
+        self._enable_torch_compile = getattr(
+            cuda_graph_runner, "enable_torch_compile", False
+        )
+
+    @contextmanager
+    def capture_session(self, stream):
+        if self._pool is None:
+            self._pool = self._device_module.graph_pool_handle()
+        set_graph_pool_id(self._pool)
+        self._capture_stream = stream
+        try:
+            yield
+        finally:
+            self._capture_stream = None
+
+    def capture_one(
+        self,
+        shape_key: Any,
+        forward_fn: Callable[[], Any],
+        dummies: Optional[Any] = None,
+        post_warmup_hook: Optional[Callable[[], None]] = None,
+    ) -> None:
+        import torch_npu  # noqa: F401  (verifies NPU availability)
+
+        # 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.npu.NPUGraph()
+
+        if self._enable_torch_compile:
+            skip_guard_context = torch.compiler.set_stance(skip_guard_eval_unsafe=True)
+        else:
+            skip_guard_context = empty_context()
+
+        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 = torch.npu.graph
+
+        with skip_guard_context, graph_ctx(
+            graph,
+            pool=self._pool,
+            stream=self._capture_stream,
+            auto_dispatch_capture=True,
+        ):
+            out = forward_fn()
+
+        self._graphs[shape_key] = graph
+        self._outputs[shape_key] = out
+
+    def can_run(self, forward_batch: ForwardBatch, shape_key: Any) -> bool:
+        return shape_key in self._graphs
+
+    @contextmanager
+    def replay_session(self):
+        yield
+
+    def replay(
+        self,
+        shape_key: Any,
+        static_forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any:
+        self._graphs[shape_key].replay()
+        return self._outputs[shape_key]
+
+    def replay_with_input_update(
+        self,
+        shape_key: Any,
+        seq_lens: list,
+        attr_name: str,
+        attr_type: Any,
+    ) -> Any:
+        """Rebind seq_lens on the recorded NPU graph in a background
+        thread, then replay. Used when the model is not deepseek-nsa."""
+        if isinstance(attr_type, torch.Tensor):
+            seq_lens = torch.from_numpy(np.array(seq_lens).astype(np.int32))
+
+        graph = self._graphs[shape_key]
+
+        def _update():
+            graph.update(cpu_update_input=[{attr_name: seq_lens}])
+
+        thread = threading.Thread(target=_update)
+        thread.start()
+        graph.replay()
+        thread.join()
+        return self._outputs[shape_key]
+
+    def cleanup(self) -> None:
+        self._graphs.clear()
+        self._outputs.clear()
+        self._pool = None

python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py

@@ -11,13 +11,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Run the model with npu graph and torch.compile."""
+"""Run the model with NPU graph and torch.compile.
+
+NPUGraphRunner is a thin subclass of DecodeCudaGraphRunner: the
+factory returns NPUCudaGraphBackend for NPU devices, so all
+capture/replay mechanics live in the backend. This class adds:
+  - NPU-specific patch_model monkey-patch for the decode-Full +
+    torch.compile path.
+  - Profile context override (NPU profiler emits to disk, not in-mem).
+  - Replay override that issues an async NPUGraph.update for
+    seq_lens before replay (skipped for deepseek-nsa).
+  - Smaller cache_loc dtype (int32 instead of int64).
+"""
 
 from __future__ import annotations
 
 import logging
 import os
-import threading
 from contextlib import contextmanager
 from pathlib import Path
 from typing import TYPE_CHECKING, Dict, Optional, Union
@@ -25,11 +35,10 @@ from typing import TYPE_CHECKING, Dict, Optional, Union
 import numpy as np
 import torch
 
-import sglang
 from sglang.srt.configs.model_config import AttentionArch, is_deepseek_dsa
 from sglang.srt.distributed.parallel_state import GroupCoordinator
 from sglang.srt.environ import envs
-from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
+from sglang.srt.model_executor.runner import DecodeCudaGraphRunner
 from sglang.srt.utils import (
     empty_context,
     get_bool_env_var,
@@ -71,8 +80,8 @@ def patch_model_npu(
         yield model.forward
 
 
-class NPUGraphRunner(CudaGraphRunner):
-    """A NPUGraphRunner runs the forward pass of a model with npu graph and torch.compile."""
+class NPUGraphRunner(DecodeCudaGraphRunner):
+    """A NPUGraphRunner runs the forward pass of a model with NPU graph and torch.compile."""
 
     def __init__(
         self,
@@ -82,7 +91,11 @@ class NPUGraphRunner(CudaGraphRunner):
         speculative_num_steps: Optional[int] = None,
         speculative_num_draft_tokens: Optional[int] = None,
     ):
-        sglang.srt.model_executor.cuda_graph_runner.patch_model = patch_model_npu
+        # NPU patch_model override: monkey-patch torch_compile_decoration's
+        # patch_model with the NPU-specific version.
+        from sglang.srt.compilation import torch_compile_decoration
+
+        torch_compile_decoration.patch_model = patch_model_npu
         super().__init__(
             model_runner,
             attn_backend=attn_backend,
@@ -215,9 +228,8 @@ class NPUGraphRunner(CudaGraphRunner):
                     forward_batch.mrope_positions
                 )
 
-        self.update_attr_name = self._get_update_attr_name()
-        self.update_attr_type = self._get_update_attr_type()
-        # Replay
+        graph_key = self._make_graph_key(self.bs)
+
         if not is_deepseek_dsa(self.model_runner.model_config.hf_config):
             if forward_batch.forward_mode.is_target_verify():
                 seq_lens_cpu = forward_batch.seq_lens.cpu() + self.num_tokens_per_bs
@@ -226,14 +238,15 @@ class NPUGraphRunner(CudaGraphRunner):
                 seq_lens = forward_batch.seq_lens.cpu().tolist() + [0] * (
                     self.bs - self.raw_bs
                 )
-            thread = threading.Thread(target=self._update_inputs, args=(seq_lens,))
-            thread.start()
-            self.graphs[self.bs].replay()
-            thread.join()
+            output = self.backend.replay_with_input_update(
+                graph_key,
+                seq_lens=seq_lens,
+                attr_name=self.attr_name[AttentionArch.MLA],
+                attr_type=self.attr_type[AttentionArch.MLA],
+            )
         else:
-            self.graphs[self.bs].replay()
+            output = self.backend.replay(graph_key, forward_batch)
 
-        output = self.output_buffers[self.bs]
         if isinstance(output, LogitsProcessorOutput):
             if self.is_dllm:
                 next_token_logits = None

python/sglang/srt/hardware_backend/npu/utils.py

@@ -56,27 +56,28 @@ def set_default_server_args(args: "ServerArgs"):
         args.page_size = 128
 
     # NPU memory settings
+    decode = args.cuda_graph_config.decode
     npu_mem = get_npu_memory_capacity()
     if npu_mem <= 32 * 1024:
         # Ascend 910B4,910B4_1
-        # (chunked_prefill_size 4k, cuda_graph_max_bs 16 if tp < 4 else 64)
+        # (chunked_prefill_size 4k, max_bs 16 if tp < 4 else 64)
         if args.chunked_prefill_size is None:
             args.chunked_prefill_size = 4 * 1024
-        if args.cuda_graph_max_bs is None:
+        if decode.max_bs is None:
             if args.tp_size < 4:
-                args.cuda_graph_max_bs = 16
+                decode.max_bs = 16
             else:
-                args.cuda_graph_max_bs = 64
+                decode.max_bs = 64
     elif npu_mem <= 64 * 1024:
         # Ascend 910B1,910B2,910B2C,910B3,910_9391,910_9392,910_9381,910_9382,910_9372,910_9362
-        # (chunked_prefill_size 8k, cuda_graph_max_bs 64 if tp < 4 else 256)
+        # (chunked_prefill_size 8k, max_bs 64 if tp < 4 else 256)
         if args.chunked_prefill_size is None:
             args.chunked_prefill_size = 8 * 1024
-        if args.cuda_graph_max_bs is None:
+        if decode.max_bs is None:
             if args.tp_size < 4:
-                args.cuda_graph_max_bs = 64
+                decode.max_bs = 64
             else:
-                args.cuda_graph_max_bs = 256
+                decode.max_bs = 256
 
     # NPU does not support CustomAllReduce
     args.disable_custom_all_reduce = True
@@ -216,7 +217,7 @@ def init_zbal(world_size, gpu_id, world_rank, do_check=True):
         gva_is_inited = True
 
         if do_check and not ret:
-            logger.error(f"[ZBAL] zbal init failed!")
+            logger.error("[ZBAL] zbal init failed!")
             sys.exit(-1)
 
         return ret
@@ -271,7 +272,7 @@ def lazy_init_zbal_gva_mem(
 
     gva_is_inited = True
     if do_check and not res:
-        logger.error(f"[ZBAL] zbal lazy init failed!")
+        logger.error("[ZBAL] zbal lazy init failed!")
         sys.exit(-1)
     return res
 

python/sglang/srt/kv_canary/api.py

@@ -12,6 +12,11 @@ from sglang.srt.kv_canary.pool_patcher.api import attach_canary_buffers
 from sglang.srt.kv_canary.pool_patcher.utils import wrap_method
 from sglang.srt.kv_canary.runner.canary_manager import CanaryManager
 from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
 if TYPE_CHECKING:
@@ -32,10 +37,11 @@ def install_canary(
     if config.mode is CanaryMode.NONE:
         return None
 
-    assert server_args.disable_piecewise_cuda_graph, (
+    assert not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE), (
         "kv-canary: piecewise cuda graph is not supported by the current "
-        "SingleForwardManager design; pass --disable-piecewise-cuda-graph "
-        "when canary is enabled"
+        "SingleForwardManager design; set "
+        "--cuda-graph-backend-prefill=disabled (or =breakable) when canary "
+        "is enabled"
     )
 
     perturb_config = PerturbConfig.from_env()

python/sglang/srt/kv_canary/capacities.py

@@ -63,7 +63,10 @@ class CanaryLaunchCapacities:
                 f"kv-canary: pool_slot_count must be positive, got {pool_slot_count}"
             )
 
-        cuda_graph_max_bs = server_args.cuda_graph_max_bs or 0
+        cuda_graph_config = server_args.cuda_graph_config
+        cuda_graph_max_bs = (
+            cuda_graph_config.decode.max_bs if cuda_graph_config is not None else 0
+        ) or 0
         if cuda_graph_max_bs < 0:
             raise ValueError(
                 f"kv-canary: cuda_graph_max_bs must be non-negative, got {cuda_graph_max_bs}"

python/sglang/srt/layers/activation.py

@@ -30,6 +30,11 @@ from sglang.srt.distributed import (
 from sglang.srt.environ import envs
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.utils import MultiPlatformOp
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     cpu_has_amx_support,
@@ -128,6 +133,9 @@ class SiluAndMul(MultiPlatformOp):
         return out
 
     def forward_musa(self, x: torch.Tensor) -> torch.Tensor:
+        if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE):
+            return self.forward_native(x)
+
         if not hasattr(self, "_musa_swish_glu"):
             # XXX (MUSA): nn.SwishGLU seems to have better performance than silu_and_mul on MUSA, we can switch to it for now. We can consider implementing a silu_and_mul kernel for MUSA in the future if needed.
             self._musa_swish_glu = nn.SwishGLU()

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -12,10 +12,6 @@ from sglang.jit_kernel.fused_store_index_cache import (
     can_use_dsa_fused_store,
     fused_store_index_k_cache,
 )
-from sglang.srt.compilation.piecewise_context_manager import (
-    get_forward_context,
-    is_in_piecewise_cuda_graph,
-)
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.dsa.utils import (
     aiter_can_use_preshuffle_paged_mqa,
@@ -26,6 +22,10 @@ from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor
 from sglang.srt.layers.layernorm import LayerNorm
 from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz
 from sglang.srt.layers.utils import MultiPlatformOp
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    get_tc_piecewise_forward_context,
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.state_capturer.indexer_topk import (
     maybe_capture_indexer_topk,
 )
@@ -83,13 +83,13 @@ from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.rotary_embedding import get_rope_wrapper
 from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_output
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.forward_context import (
     get_attn_backend,
     get_req_to_token_pool,
     get_token_to_kv_pool,
 )
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.server_args import get_global_server_args
 
 _use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get()
@@ -118,8 +118,8 @@ if _is_cuda:
         ), "Internal error: piecewise CUDA graph is only supported on CUDA"
         from sglang.srt.layers.attention.dsa.triton_kernel import act_quant
 
-        forward_batch = get_forward_context().forward_batch
-        indexer = get_forward_context().dsa_indexers[layer_id]
+        forward_batch = get_tc_piecewise_forward_context().forward_batch
+        indexer = get_tc_piecewise_forward_context().dsa_indexers[layer_id]
         metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch)
 
         # slice off padding from piecewise CUDA graph
@@ -198,7 +198,7 @@ def _broadcast_indexer_topk_from_rank0(
     if topk_indices is None or not envs.SGLANG_DSA_TOPK_BROADCAST.get():
         return topk_indices
 
-    if is_in_piecewise_cuda_graph():
+    if is_in_tc_piecewise_cuda_graph():
         broadcast_indexer_topk_from_rank0_(topk_indices)
     else:
         _broadcast_indexer_topk_from_rank0_impl(topk_indices)
@@ -999,7 +999,7 @@ class Indexer(MultiPlatformOp):
         cp_index: List[Tuple[int, int, int]] = None,
     ) -> torch.Tensor:
         assert (
-            not is_in_piecewise_cuda_graph()
+            not is_in_tc_piecewise_cuda_graph()
         ), "DSA context parallel (_get_topk_ragged_with_cp) not supported under piecewise CUDA graph"
         if TYPE_CHECKING:
             assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool)
@@ -1148,7 +1148,7 @@ class Indexer(MultiPlatformOp):
         layer_id: int,
     ) -> Optional[torch.Tensor]:
         assert (
-            not is_in_piecewise_cuda_graph()
+            not is_in_tc_piecewise_cuda_graph()
         ), "DSA forward_indexer (non-CUDA loop path) not supported under piecewise CUDA graph"
         if not _is_npu:
             from sglang.srt.layers.attention.dsa.tilelang_kernel import fp8_index
@@ -1339,10 +1339,10 @@ class Indexer(MultiPlatformOp):
         # a tuple like (x_fp8, x_scale[, y]). Use `x_meta` for shape/device queries.
         x_meta = x[0] if isinstance(x, tuple) else x
 
-        # In piecewise CUDA graph mode, metadata is fetched inside custom ops via get_forward_context() to
+        # In piecewise CUDA graph mode, metadata is fetched inside custom ops via get_tc_piecewise_forward_context() to
         # prevent Dynamo from guarding on forward_metadata identity (which changes each
         # replay when init_forward_metadata creates a new ForwardMetadata object).
-        if not is_in_piecewise_cuda_graph():
+        if not is_in_tc_piecewise_cuda_graph():
             metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch)
             if metadata is None:
                 return None
@@ -1360,7 +1360,7 @@ class Indexer(MultiPlatformOp):
         # We can only skip the logits computation if cuda graph is not involved
         skip_logits_computation = False
         if (
-            not is_in_piecewise_cuda_graph()
+            not is_in_tc_piecewise_cuda_graph()
             and forward_batch.forward_mode.is_extend_without_speculative()
         ):
             if forward_batch.seq_lens_cpu is not None:
@@ -1417,7 +1417,7 @@ class Indexer(MultiPlatformOp):
                         act_quant=act_quant,
                     )
                 current_stream.wait_stream(self.alt_stream)
-            elif not is_in_piecewise_cuda_graph():
+            elif not is_in_tc_piecewise_cuda_graph():
                 q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
                 self._store_index_k_cache(
                     forward_batch=forward_batch,
@@ -1471,7 +1471,7 @@ class Indexer(MultiPlatformOp):
             else:
                 x_for_gate = x
 
-            if is_in_piecewise_cuda_graph():
+            if is_in_tc_piecewise_cuda_graph():
                 weights = logits_head_gate_pcg(
                     x_for_gate,
                     self.weights_proj.weight,
@@ -1485,7 +1485,7 @@ class Indexer(MultiPlatformOp):
         if _is_cuda or _is_hip:
             # In piecewise CUDA graph, any access to seq_lens_cpu creates a Dynamo shape guard.
             # Piecewise CUDA graph never has empty batches.
-            if not is_in_piecewise_cuda_graph():
+            if not is_in_tc_piecewise_cuda_graph():
                 assert forward_batch.seq_lens_cpu is not None
                 if len(forward_batch.seq_lens_cpu) == 0:
                     # this seems b/c max-pad, no worries?
@@ -1556,7 +1556,7 @@ class Indexer(MultiPlatformOp):
                     topk_result = torch.cat([topk_result_prev, topk_result_next], dim=0)
                     topk_result = _broadcast_indexer_topk_from_rank0(topk_result)
                     return maybe_capture_indexer_topk(layer_id, topk_result)
-                elif is_in_piecewise_cuda_graph():
+                elif is_in_tc_piecewise_cuda_graph():
                     assert (
                         not enable_dual_stream
                     ), "Internal error: piecewise CUDA graph should not be enabled with dual stream"

python/sglang/srt/layers/attention/dsa_backend.py

@@ -2252,13 +2252,13 @@ class DeepseekSparseAttnBackend(
         """
         Decide all attention prefill dispatch strategies for this batch.
         """
-        from sglang.srt.compilation.piecewise_context_manager import (
-            is_in_piecewise_cuda_graph,
+        from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+            is_in_tc_piecewise_cuda_graph,
         )
         from sglang.srt.utils import get_device_sm, is_blackwell
 
         # Decide MHA vs MLA
-        if is_in_piecewise_cuda_graph():
+        if is_in_tc_piecewise_cuda_graph():
             # Can't branch on seq_lens_cpu in PCG, force mha off to guarantee correctness.
             self.use_mha = False
         elif (

python/sglang/srt/layers/attention/fla/layernorm_gated.py

@@ -15,7 +15,11 @@ import triton.language as tl
 from einops import rearrange
 
 from sglang.jit_kernel.utils import is_arch_support_pdl
-from sglang.srt.server_args import get_global_server_args
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
 from sglang.srt.utils import (
     cdiv,
     cpu_has_amx_support,
@@ -190,12 +194,8 @@ def _get_sm_count(device: torch.device) -> int:
 def calc_rows_per_block(M: int, device: torch.device) -> int:
     # When piecewise cuda graph is enabled, use a constant value to avoid
     # torch.compile creating guards on the dynamic batch dimension.
-    try:
-        if not get_global_server_args().disable_piecewise_cuda_graph:
-            return MAX_ROWS_PER_BLOCK
-    except ValueError:
-        # Global server args not initialized (e.g., in unit tests)
-        pass
+    if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE):
+        return MAX_ROWS_PER_BLOCK
     sm_count = _get_sm_count(device)
     rows_per_block = next_power_of_2(cdiv(M, 2 * sm_count))
     rows_per_block = min(rows_per_block, MAX_ROWS_PER_BLOCK)

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -34,6 +34,7 @@ from sglang.jit_kernel.flash_attention import (
     flash_attn_varlen_func,
     flash_attn_with_kvcache,
 )
+from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
 
 
 @dataclass
@@ -212,10 +213,7 @@ class FlashAttentionBackend(AttentionBackend):
         self.num_splits = (
             1
             if model_runner.server_args.enable_deterministic_inference
-            or (
-                self.fa_impl_ver == 4
-                and not model_runner.server_args.disable_cuda_graph
-            )
+            or (self.fa_impl_ver == 4 and not cuda_graph_fully_disabled())
             else 0
         )
 
@@ -2028,7 +2026,7 @@ class FlashAttentionBackend(AttentionBackend):
         """Shared capture+replay body for the cuda-graph init path.
 
         Public entry: :py:meth:`init_forward_metadata_out_graph`. This helper
-        formerly lived as the legacy ``init_forward_metadata_replay_cuda_graph``;
+        formerly lived as the legacy init_forward_metadata_replay_cuda_graph;
         the capture path used to wrap it. Both legacy method overrides
         are gone.
         """

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -17,7 +17,6 @@ from typing import TYPE_CHECKING, Callable, List, Optional, Union
 import torch
 
 from sglang.kernel_api_logging import debug_kernel_api
-from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
 from sglang.srt.dllm.config import DllmConfig
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
@@ -28,7 +27,15 @@ from sglang.srt.layers.attention.utils import (
 from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.layers.radix_attention import AttentionType
 from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.speculative.spec_info import SpecInput
 from sglang.srt.speculative.spec_utils import (
     draft_kv_indices_buffer_width,
@@ -281,7 +288,9 @@ class FlashInferAttnBackend(AttentionBackend):
 
         fmha_backend = "auto"
         if is_sm100_supported():
-            if not model_runner.server_args.disable_piecewise_cuda_graph:
+            # Disable CUTLASS backend when piecewise cuda graph is enabled
+            # due to TMA descriptor initialization issues on B200
+            if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE):
                 logger.info(
                     "CUTLASS backend is disabled when piecewise cuda graph is enabled "
                     "due to TMA descriptor initialization issues on SM100 GPUs. "
@@ -592,7 +601,7 @@ class FlashInferAttnBackend(AttentionBackend):
             else:
                 use_ragged = (
                     not self.enable_deterministic
-                    and not is_in_piecewise_cuda_graph()
+                    and not is_in_tc_piecewise_cuda_graph()
                     and not self.use_paged
                 )
                 extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu)

python/sglang/srt/layers/attention/flashinfer_mla_backend.py

@@ -15,7 +15,6 @@ from typing import TYPE_CHECKING, Callable, Optional, Union
 
 import torch
 
-from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.attention.flashinfer_backend import (
@@ -24,6 +23,9 @@ from sglang.srt.layers.attention.flashinfer_backend import (
 from sglang.srt.layers.attention.utils import assert_buffer_fits
 from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.spec_info import SpecInput
 from sglang.srt.speculative.spec_utils import (
@@ -408,7 +410,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
                 not get_global_server_args().flashinfer_mla_disable_ragged
                 and extend_no_prefix
                 # Piecewise cuda graph should use paged prefill to be compatible with prefix cache
-                and not is_in_piecewise_cuda_graph()
+                and not is_in_tc_piecewise_cuda_graph()
             )
 
             self.indices_updater_prefill.update(

python/sglang/srt/layers/attention/tokenspeed_mla_backend.py

@@ -222,7 +222,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
         k_pe: torch.Tensor,
         positions: torch.Tensor,
         layer: "DeepseekV2AttentionMLA",
-        forward_batch: "ForwardBatch",
+        forward_batch: ForwardBatch,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Build FP8 (Q, K, V) for the FMHA kernel and write FP8 KV cache."""
         kv = layer.kv_b_proj(kv_a)[0]

python/sglang/srt/layers/attention/triton_backend.py

@@ -15,6 +15,7 @@ from sglang.srt.layers.attention.triton_ops.metadata import get_num_kv_splits_tr
 from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.layers.radix_attention import AttentionType
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
+from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.speculative.spec_utils import (
     draft_kv_indices_buffer_width,
@@ -184,8 +185,8 @@ class TritonAttnBackend(AttentionBackend):
             self.use_pdl = False
 
         self.allow_bidirectional_attention_in_extend = (
-            model_runner.server_args.disable_cuda_graph
-            and (model_runner.server_args.chunked_prefill_size == -1)
+            cuda_graph_fully_disabled()
+            and model_runner.server_args.chunked_prefill_size == -1
         )
 
         # Decide whether enable deterministic inference with batch-invariant operations
@@ -332,8 +333,8 @@ class TritonAttnBackend(AttentionBackend):
     ):
         """Fill KV (and SWA) cuda-graph buffers for decode/idle mode.
 
-        Returns ``(kv_indptr, window_kv_indptr, window_kv_lens)`` where
-        ``window_kv_lens`` is ``None`` when sliding-window is disabled.
+        Returns (kv_indptr, window_kv_indptr, window_kv_lens) where
+        window_kv_lens is None when sliding-window is disabled.
         """
         seq_lens = seq_lens[:bs]
         req_pool_indices = req_pool_indices[:bs]
@@ -431,7 +432,7 @@ class TritonAttnBackend(AttentionBackend):
     ):
         """Fill QO + KV cuda-graph buffers for draft_extend mode.
 
-        Returns ``(qo_indptr, kv_indptr, num_tokens_per_bs)``.
+        Returns (qo_indptr, kv_indptr, num_tokens_per_bs).
         """
         seq_lens = seq_lens[:bs]
         # V2 draft-extend fills num_draft_tokens per req (the cuda-graph runner's
@@ -848,7 +849,7 @@ class TritonAttnBackend(AttentionBackend):
 
         Called by capture after the buffer-update helpers have already run
         (either via replay or directly).  All fields reference the same
-        ``self.cuda_graph_*`` tensors that the captured graph kernels will
+        self.cuda_graph_* tensors that the captured graph kernels will
         read — the Python object is rebuilt each capture, but the underlying
         GPU memory addresses are stable.
         """
@@ -1514,9 +1515,9 @@ def update_sliding_window_buffer(
 ):
     """Fill window KV buffers for sliding-window attention.
 
-    Pass ``window_kv_indices`` to write into a pre-allocated buffer (CUDA-graph
-    path); omit it (or pass ``None``) to allocate a fresh tensor (eager path,
-    requires ``device``).
+    Pass window_kv_indices to write into a pre-allocated buffer (CUDA-graph
+    path); omit it (or pass None) to allocate a fresh tensor (eager path,
+    requires device).
     """
     window_kv_lens = torch.minimum(
         seq_lens,

python/sglang/srt/layers/attention/trtllm_mla_backend.py

@@ -14,7 +14,6 @@ import triton
 import triton.language as tl
 
 from sglang.jit_kernel.fixup_zero_kv import fixup_zero_kv_rows
-from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.flashinfer_mla_backend import (
     FlashInferMLAAttnBackend,
@@ -30,6 +29,9 @@ from sglang.srt.layers.attention.utils import (
 from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import is_flashinfer_available, is_float4_e2m1fn_x2
 
@@ -562,11 +564,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         """Get the fill value for sequence lengths in CUDA graph."""
         return 1
 
-    def init_mha_chunk_metadata(self, forward_batch: "ForwardBatch") -> None:
+    def init_mha_chunk_metadata(self, forward_batch: ForwardBatch) -> None:
         has_prefix = any(forward_batch.extend_prefix_lens_cpu)
         fallback_to_flashinfer_impl = (
             self.disable_chunked_prefix_cache and has_prefix
-        ) or is_in_piecewise_cuda_graph()
+        ) or is_in_tc_piecewise_cuda_graph()
         if fallback_to_flashinfer_impl:
             super().init_mha_chunk_metadata(
                 forward_batch, disable_flashinfer_ragged=True
@@ -627,7 +629,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
             has_prefix = any(forward_batch.extend_prefix_lens_cpu)
             fallback_to_flashinfer_impl = (
                 self.disable_chunked_prefix_cache and has_prefix
-            ) or is_in_piecewise_cuda_graph()
+            ) or is_in_tc_piecewise_cuda_graph()
             if fallback_to_flashinfer_impl:
                 super().init_forward_metadata(forward_batch)
 
@@ -789,7 +791,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         return output[:total_tokens, :, :]
 
     def _compute_decode_bmm1_scale(self, layer: RadixAttention) -> float:
-        """BMM1 scale ``q_scale * k_scale * softmax_scale``. k_scale only
+        """BMM1 scale q_scale * k_scale * softmax_scale. k_scale only
         applies when the KV cache stores FP8."""
         q_scale = 1.0
         if self.data_type == torch.float8_e4m3fn:
@@ -867,7 +869,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
     ):
         """Hook for subclasses to swap the ragged prefill kernel. Q/K/V arrive
         in model-native dtype; subclasses do any kernel-specific quantization.
-        Returns the output tensor or ``(output, lse)`` if ``return_lse``."""
+        Returns the output tensor or (output, lse) if return_lse."""
         q_scale = k_scale = v_scale = 1.0
         if self.data_type == torch.float8_e4m3fn:
             q, k, v, k_scale, v_scale = _quantize_fp8_qkv(q, k, v, layer)

python/sglang/srt/layers/communicator.py

@@ -70,6 +70,11 @@ from sglang.srt.layers.utils.cp_utils import (
     is_mla_prefill_cp_enabled,
     mla_use_prefill_cp,
 )
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -274,7 +279,7 @@ class AttnTpContext:
             and not is_dp_attention_enabled()
             and get_moe_a2a_backend().is_none()
             and not enable_moe_dense_fully_dp()
-            and get_global_server_args().disable_piecewise_cuda_graph
+            and not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
             and get_global_server_args().speculative_algorithm != "EAGLE3"
         )
         if get_global_server_args().enable_attn_tp_input_scattered:

python/sglang/srt/layers/layernorm.py

@@ -26,6 +26,11 @@ from sglang.srt.batch_invariant_ops import (
 )
 from sglang.srt.environ import envs
 from sglang.srt.layers.utils import MultiPlatformOp
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     cpu_has_amx_support,
@@ -431,6 +436,9 @@ class RMSNorm(MultiPlatformOp):
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE):
+            return self.forward_native(x, residual, post_residual_addition)
+
         if not x.is_contiguous():
             x = x.contiguous()
 

python/sglang/srt/layers/moe/ep_moe/layer.py

@@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, Any, Dict, Optional
 
 import torch
 
-from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
 from sglang.srt.environ import envs
 from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.moe import (
@@ -26,6 +25,9 @@ from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.quantization.fp8 import Fp8Config
 from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
 from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config, W4AFp8MoEMethod
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.utils import get_bool_env_var, is_hip, is_npu
 
 if TYPE_CHECKING:
@@ -150,7 +152,7 @@ class DeepEPMoE(FusedMoE):
         hidden_states: torch.Tensor,
         topk_output: TopKOutput,
     ):
-        if is_in_piecewise_cuda_graph():
+        if is_in_tc_piecewise_cuda_graph():
             assert TopKOutputChecker.format_is_standard(
                 topk_output
             ), "Only standard topk output is supported for piecewise cuda graph"

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -12,10 +12,6 @@ from torch.nn.parameter import UninitializedParameter
 
 from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs
 from sglang.srt.batch_overlap.two_batch_overlap import MaybeTboDeepEPDispatcher
-from sglang.srt.compilation.piecewise_context_manager import (
-    get_forward_context,
-    is_in_piecewise_cuda_graph,
-)
 from sglang.srt.distributed import (
     get_moe_expert_parallel_rank,
     get_moe_expert_parallel_world_size,
@@ -63,6 +59,10 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import (
 from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod
 from sglang.srt.layers.quantization.modelopt_quant import ModelOptNvFp4FusedMoEMethod
 from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    get_tc_piecewise_forward_context,
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.model_loader.weight_utils import narrow_padded_param_and_loaded_weight
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
@@ -1071,7 +1071,7 @@ class FusedMoE(torch.nn.Module):
             from sglang.srt.hardware_backend.npu.moe.fuseep import forward_fuseep
 
             return forward_fuseep(self, hidden_states, topk_output)
-        if is_in_piecewise_cuda_graph():
+        if is_in_tc_piecewise_cuda_graph():
             if TopKOutputChecker.format_is_standard(topk_output):
                 return moe_forward_piecewise_cuda_graph_impl(
                     hidden_states,
@@ -1308,7 +1308,7 @@ def moe_forward_piecewise_cuda_graph_impl(
     topk_output = StandardTopKOutput(
         topk_weights=topk_weights, topk_ids=topk_ids, router_logits=router_logits
     )
-    forward_context = get_forward_context()
+    forward_context = get_tc_piecewise_forward_context()
     moe_layer = forward_context.moe_layers[layer_id]
     return moe_layer.forward_impl(hidden_states, topk_output)
 
@@ -1335,6 +1335,6 @@ def fused_moe_bypassed_piecewise_cuda_graph_impl(
             renormalize=renormalize,
         ),
     )
-    forward_context = get_forward_context()
+    forward_context = get_tc_piecewise_forward_context()
     moe_layer = forward_context.moe_layers[layer_id]
     return moe_layer.forward_impl(hidden_states, topk_output)

python/sglang/srt/layers/moe/mega_moe.py

@@ -26,7 +26,7 @@ from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
 from sglang.srt.layers.dp_attention import get_dp_global_num_tokens
 from sglang.srt.layers.moe.utils import get_moe_a2a_backend
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
+from sglang.srt.model_executor.runner import get_is_capture_mode
 
 if TYPE_CHECKING:
     from deep_gemm import SymmBuffer
@@ -114,7 +114,7 @@ def should_use_mega_moe(moe: "DeepseekV2MoE", hidden_states: torch.Tensor) -> bo
 def forward_mega_moe(
     moe: "DeepseekV2MoE",
     hidden_states: torch.Tensor,
-    forward_batch: Optional["ForwardBatch"] = None,
+    forward_batch: Optional[ForwardBatch] = None,
     input_ids_global: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     num_tokens = hidden_states.shape[0]
@@ -151,7 +151,7 @@ def forward_mega_moe(
 def _run_mega_routed(
     moe: "DeepseekV2MoE",
     hidden_states: torch.Tensor,
-    forward_batch: Optional["ForwardBatch"],
+    forward_batch: Optional[ForwardBatch],
     input_ids_global: Optional[torch.Tensor],
     num_tokens: int,
 ) -> torch.Tensor:

python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py

@@ -11,6 +11,7 @@ from sglang.srt.layers.moe.moe_runner.base import (
     MoeRunnerConfig,
     register_fused_func,
 )
+from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
 from sglang.srt.utils.common import log_info_on_rank0, print_warning_once
 
 if TYPE_CHECKING:
@@ -251,7 +252,9 @@ def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None:
     )
 
     server_args = get_global_server_args()
-    use_cuda_graph = not server_args.disable_cuda_graph
+    # CuteDSL wrapper preallocates CG buffers used by any captured graph
+    # that routes through this MoE — decode and prefill alike.
+    use_cuda_graph = not cuda_graph_fully_disabled()
 
     # Size the wrapper's CUDA-graph buffers for the largest number of tokens a
     # single forward can route through this layer.
@@ -301,17 +304,17 @@ class CuteDslFp4MoeQuantInfo(MoeQuantInfo):
 
     Shared by the two CuteDSL runner entries:
 
-    * "v2" standard path (a2a=``none``/``flashinfer``): consumed by the
-      ``@register_fused_func("none", "flashinfer_cutedsl")`` entry, which
-      drives ``CuteDslMoEWrapper.run``. Weights are ``[Up, Gate]``
-      interleaved with MMA-layout blockscales. ``wrapper`` is set;
-      ``w*_scale`` are scalarized.
+    * "v2" standard path (a2a=none/flashinfer): consumed by the
+      @register_fused_func("none", "flashinfer_cutedsl") entry, which
+      drives CuteDslMoEWrapper.run. Weights are [Up, Gate]
+      interleaved with MMA-layout blockscales. wrapper is set;
+      w*_scale are scalarized.
 
-    * "v1" DeepEP low-latency path (a2a=``deepep``): consumed by the
-      ``@register_fused_func("deepep", "flashinfer_cutedsl")`` entry,
-      which drives ``flashinfer_cutedsl_moe_masked``. Weights are
-      ``[Gate, Up]`` non-interleaved with swizzled blockscales.
-      ``wrapper`` is ``None``; ``w*_scale`` are per-expert.
+    * "v1" DeepEP low-latency path (a2a=deepep): consumed by the
+      @register_fused_func("deepep", "flashinfer_cutedsl") entry,
+      which drives flashinfer_cutedsl_moe_masked. Weights are
+      [Gate, Up] non-interleaved with swizzled blockscales.
+      wrapper is None; w*_scale are per-expert.
     """
 
     # FP4 packed weights (uint8)
@@ -332,10 +335,10 @@ class CuteDslFp4MoeQuantInfo(MoeQuantInfo):
     a1_scale: torch.Tensor
     a2_scale: torch.Tensor
 
-    # v2 only: lazily-created CuteDslMoEWrapper (``None`` on the v1 path).
+    # v2 only: lazily-created CuteDslMoEWrapper (None on the v1 path).
     wrapper: Optional[Any] = None
 
-    # v1 only: ``True`` when DeepEP pre-quantizes activations to NVFP4.
+    # v1 only: True when DeepEP pre-quantizes activations to NVFP4.
     use_nvfp4_dispatch: bool = False
 
     # v1 only: SBO down-GEMM overlap args.

python/sglang/srt/layers/quantization/fp8_utils.py

@@ -1502,12 +1502,14 @@ def apply_fp8_linear(
         # eliminating a separate kernel launch per linear layer.
         # weight_scale shape does not matter here -- it is only used in the
         # GEMM epilogue, not in the activation quant fusion. Only activates when
-        # piecewise_cuda_graph_compiler=inductor; eager PCG and decode both
-        # use the faster custom kernel.
+        # cuda_graph_config[prefill].tc_compiler=inductor; eager PCG and
+        # decode both use the faster custom kernel.
+
         if (
             input_scale is not None
             and input_scale.numel() == 1
-            and get_global_server_args().piecewise_cuda_graph_compiler == "inductor"
+            and get_global_server_args().cuda_graph_config.prefill.tc_compiler
+            == "inductor"
         ):
             qinput = (
                 (input_2d * input_scale.reciprocal())

python/sglang/srt/layers/quantization/marlin_utils.py

@@ -33,7 +33,9 @@ if TYPE_CHECKING:
     from sglang.srt.layers.linear import LinearBase
     from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
 
-from sglang.srt.compilation.piecewise_context_manager import get_forward_context
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    get_tc_piecewise_forward_context,
+)
 
 try:
     from vllm import _custom_ops as ops
@@ -499,7 +501,7 @@ def apply_gptq_marlin_linear(
         dtype=input.dtype,
     )
 
-    forward_context = get_forward_context()
+    forward_context = get_tc_piecewise_forward_context()
     if forward_context is None:
         output = gptq_marlin_gemm(
             reshaped_x,
@@ -569,7 +571,7 @@ def apply_awq_marlin_linear(
         dtype=input.dtype,
     )
 
-    forward_context = get_forward_context()
+    forward_context = get_tc_piecewise_forward_context()
     if forward_context is None:
         output = gptq_marlin_gemm(
             reshaped_x,
@@ -906,7 +908,7 @@ def unified_apply_gptq_marlin_gemm(
     use_fp32_reduce: bool,
     is_zp_float: bool,
 ) -> torch.Tensor:
-    quant_config = get_forward_context().quant_config
+    quant_config = get_tc_piecewise_forward_context().quant_config
     quant_type = quant_config.quant_type
     return gptq_marlin_gemm(
         input,

python/sglang/srt/layers/radix_attention.py

@@ -22,14 +22,14 @@ import torch
 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 (
+from sglang.srt.model_executor.forward_context import get_attn_backend
+from sglang.srt.model_executor.runner_backend_utils.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_context import get_attn_backend
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    get_tc_piecewise_forward_context,
+)
 from sglang.srt.utils import is_hip
 from sglang.srt.utils.custom_op import register_custom_op
 
@@ -124,13 +124,16 @@ class RadixAttention(nn.Module):
             else:
                 k = k.view(-1, self.tp_k_head_num, self.v_head_dim)
 
-        if forward_batch.forward_mode.is_extend() and get_forward_context() is not None:
+        if (
+            forward_batch.forward_mode.is_extend()
+            and get_tc_piecewise_forward_context() is not None
+        ):
             if self.qk_head_dim != self.v_head_dim:
                 output = q.new_empty((q.shape[0], self.tp_q_head_num * self.v_head_dim))
             else:
                 output = torch.empty_like(q)
             if is_in_breakable_cuda_graph():
-                bcg_unified_attention_with_output(
+                breakable_unified_attention_with_output(
                     q, k, v, output, save_kv_cache, self.layer_id, **kwargs
                 )
             else:
@@ -170,7 +173,7 @@ def unified_attention_with_output(
     llama_4_scaling: Optional[torch.Tensor] = None,
     topk_indices: Optional[torch.Tensor] = None,
 ) -> None:
-    context = get_forward_context()
+    context = get_tc_piecewise_forward_context()
     forward_batch = context.forward_batch
     attention_layers = context.attention_layers
     attention_layer = attention_layers[layer_id]
@@ -250,4 +253,6 @@ def unified_attention_with_output(
     return
 
 
-bcg_unified_attention_with_output = eager_on_graph(True)(unified_attention_with_output)
+breakable_unified_attention_with_output = eager_on_graph(True)(
+    unified_attention_with_output
+)

python/sglang/srt/layers/radix_linear_attention.py

@@ -21,14 +21,14 @@ import torch
 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 (
+from sglang.srt.model_executor.forward_context import get_attn_backend
+from sglang.srt.model_executor.runner_backend_utils.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_context import get_attn_backend
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    get_tc_piecewise_forward_context,
+)
 from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:
@@ -82,7 +82,10 @@ class RadixLinearAttention(nn.Module):
         a: torch.Tensor,
         b: torch.Tensor,
     ) -> torch.Tensor:
-        if forward_batch.forward_mode.is_extend() and get_forward_context() is not None:
+        if (
+            forward_batch.forward_mode.is_extend()
+            and get_tc_piecewise_forward_context() is not None
+        ):
             # Output shape from linear attention: (1, seq_len, num_v_heads, head_v_dim)
             seq_len = mixed_qkv.shape[0]
             output = torch.empty(
@@ -129,7 +132,7 @@ def unified_linear_attention_with_output(
     """
     Custom op wrapper for linear attention computation only.
     """
-    context = get_forward_context()
+    context = get_tc_piecewise_forward_context()
     forward_batch = context.forward_batch
     attention_layers = context.attention_layers
     attention_layer = attention_layers[layer_id]

python/sglang/srt/lora/lora_moe_runners.py

@@ -29,7 +29,7 @@ from typing import Callable
 
 import torch
 
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.utils import is_cuda, is_hip, is_xpu, next_power_of_2
 
 _is_cuda = is_cuda()

python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py

@@ -66,7 +66,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora(
     from sglang.srt.lora.trtllm_lora_temp.triton_ops import (
         merged_experts_fused_moe_lora_add,
     )
-    from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
+    from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode
 
     assert runner_config.activation == "silu" and runner_config.is_gated, (
         "experimental_sgl_trtllm LoRA currently supports the gated SwiGLU FP8 "
@@ -327,7 +327,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora(
     from sglang.srt.lora.trtllm_lora_temp.triton_ops import (
         merged_experts_fused_moe_lora_add,
     )
-    from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
+    from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode
 
     assert (
         runner_config.activation == "silu" and runner_config.is_gated

python/sglang/srt/managers/overlap_utils.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Sequence, Union
 import torch
 
 from sglang.srt.environ import envs
+from sglang.srt.model_executor.cuda_graph_config import Backend
 from sglang.srt.speculative.spec_utils import spec_need_hidden_states
 from sglang.srt.speculative.triton_ops.gather_spec_extras import gather_spec_extras
 from sglang.srt.utils import is_cuda, is_hip, is_npu
@@ -30,7 +31,11 @@ def decide_needs_cpu_seq_lens(
     if server_args.enable_two_batch_overlap:
         # FIXME: support TBO without seq lens cpu value
         return True
-    if not server_args.disable_piecewise_cuda_graph:
+    cuda_graph_config = server_args.cuda_graph_config
+    if (
+        cuda_graph_config is not None
+        and cuda_graph_config.prefill.backend == Backend.TC_PIECEWISE
+    ):
         # FIXME: support PCG without seq lens cpu value
         return True
     # Skip unset slots (e.g. draft_extend_attn_backend on some spec configs);

python/sglang/srt/managers/scheduler_components/dp_attn.py

@@ -14,6 +14,7 @@ from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.observability.metrics_collector import DPCooperationInfo
 from sglang.srt.server_args import ServerArgs
@@ -276,7 +277,7 @@ class SchedulerDPAttnAdapter:
             attn_cp_size=self.ps.attn_cp_size,
             tp_group=self.tp_group,
             get_idle_batch=self.get_idle_batch,
-            disable_cuda_graph=self.server_args.disable_cuda_graph,
+            disable_cuda_graph=cuda_graph_fully_disabled(),
             require_mlp_tp_gather=require_mlp_tp_gather(self.server_args),
             disable_overlap_schedule=self.server_args.disable_overlap_schedule,
             offload_tags=self.offload_tags,

python/sglang/srt/mem_cache/memory_pool.py

@@ -129,7 +129,7 @@ def _set_kv_buffer_impl(
             row_dim,
         )
 
-    from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
+    from sglang.srt.model_executor.runner import get_is_capture_mode
 
     if get_is_capture_mode() and alt_stream is not None:
         current_stream = device_module.current_stream()
@@ -1531,7 +1531,7 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool):
         layer_id_override: Optional[int] = None,
     ):
         maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MHA-FP4)")
-        from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
+        from sglang.srt.model_executor.runner import get_is_capture_mode
 
         if layer_id_override is not None:
             layer_id = layer_id_override

python/sglang/srt/model_executor/breakable_cuda_graph_runner.py

@@ -1,541 +0,0 @@
-# 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 inspect
-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 (
-    CaptureHiddenMode,
-    PPProxyTensors,
-)
-from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
-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.
-    """
-
-    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
-
-        self.capture_hidden_mode = CaptureHiddenMode.NULL
-        if model_runner.server_args.enable_return_hidden_states:
-            self.capture_hidden_mode = CaptureHiddenMode.FULL
-        if (
-            model_runner.spec_algorithm is not None
-            and model_runner.spec_algorithm.is_eagle()
-        ):
-            if model_runner.is_draft_worker:
-                self.capture_hidden_mode = CaptureHiddenMode.LAST
-            else:
-                self.capture_hidden_mode = CaptureHiddenMode.FULL
-
-        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
-        self.use_captured_attn_metadata = (
-            model_runner.attn_backend.use_captured_forward_metadata_for_breakable_cuda_graph
-        )
-        self.attn_metadata_buffers = {} if self.use_captured_attn_metadata else None
-
-        # Resolve the inner transformer-stack module (the same boundary PCG draws
-        # via patch_model). At replay we monkey-patch this module's forward with
-        # a closure that replays the captured CUDAGraph and returns the captured
-        # hidden_states; the outer model.forward then runs logits_processor /
-        # pooler eagerly with the live (multi-req) forward_batch.
-        language_model = getattr(
-            model_runner.model, "language_model", model_runner.model
-        )
-        if hasattr(language_model, "model") and hasattr(language_model.model, "layers"):
-            self.layer_model = language_model.model
-        else:
-            # If we can't find the inner layer_model, disable BCG.
-            self.layer_model = None
-            logger.warning(
-                "[BCG] Could not resolve inner layer_model on %s. BCG is "
-                "disabled for this model; prefill will fall back to eager.",
-                type(language_model).__name__,
-            )
-            return
-        self.use_input_embeds = self.is_multimodal
-        if self.use_input_embeds:
-            sig = inspect.signature(self.layer_model.forward)
-            params = list(sig.parameters)
-            if "input_embeds" not in params:
-                raise ValueError(
-                    f"layer_model.forward must accept 'input_embeds' for "
-                    f"multimodal BCG, got params: {params}"
-                )
-            self._input_embeds_arg_idx = params.index("input_embeds")
-
-        # 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 _has_inactive_dp_rank(self, forward_batch: "ForwardBatch") -> bool:
-        global_num_tokens = forward_batch.global_num_tokens_cpu
-        if global_num_tokens is None:
-            return False
-
-        # DSV4 DP attention / DeepEP collectives need every DP rank to enter
-        # the same replay path. Sparse-DP batches fall back to eager to avoid
-        # hanging ranks that have zero local tokens.
-        return len(global_num_tokens) > 1 and any(
-            int(num_tokens) == 0 for num_tokens in global_num_tokens
-        )
-
-    def _init_buffers(self, model_runner):
-        """Initialize input buffers."""
-        from sglang.srt.model_executor.cuda_graph_buffer_registry import (
-            build_prefill_registry,
-        )
-        from sglang.srt.utils import is_npu
-
-        cache_loc_dtype = torch.int64 if not is_npu() else torch.int32
-        if model_runner.is_draft_worker:
-            from sglang.srt.speculative.eagle_utils import get_draft_hidden_dim
-
-            hidden_dim = get_draft_hidden_dim(model_runner)
-            self.static_draft_hidden_states = torch.zeros(
-                (self.max_num_tokens, hidden_dim),
-                dtype=model_runner.dtype,
-                device=self.device,
-            )
-
-        # Registry owns (allocates + pools) the token-axis input buffers.
-        self.buffer_registry = build_prefill_registry(
-            device=self.device,
-            max_bs=1,
-            max_num_token=self.max_num_tokens,
-            cache_loc_dtype=cache_loc_dtype,
-            is_multimodal=self.is_multimodal,
-            hidden_size=model_runner.model_config.hidden_size,
-            embed_dtype=model_runner.dtype,
-            enable_mamba_track=False,
-            share_pool=not is_npu(),
-            source=None,
-        )
-
-    @torch.no_grad()
-    def _run_forward(self, forward_batch, num_tokens):
-        """Run layer-stack forward with proper context.
-
-        Captures only the inner transformer stack (layer_model). The outer
-        model.forward's tail (logits_processor / pooler) is intentionally
-        excluded — it has bs-shaped kernels that would bake batch_size=1
-        into the captured graph.
-
-        ``@torch.no_grad`` mirrors the decorator on the outer ``*ForCausalLM.forward``
-        (e.g. qwen3.py:507). Calling ``layer_model.forward`` directly skips that
-        decorator, so we apply it here — without it some MoE @torch.compile
-        kernels (``torch.sum(out=...)``) fail dynamo with "out= doesn't support
-        autograd", and mamba state ops can spuriously track gradients.
-        """
-        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.layer_model.forward(
-                forward_batch.input_ids,
-                forward_batch.positions,
-                forward_batch,
-                input_embeds=forward_batch.input_embeds,
-            )
-        return output
-
-    def _build_capture_forward_batch(self, num_tokens):
-        """Build a bs=1 placeholder ForwardBatch for capture.
-
-        bs=1 here is only a placeholder for attention/mamba breaks' metadata
-        shapes; replay supplies live multi-req metadata via replay_prepare.
-        Captured kernels run only on the token-major layer stack and are
-        bs-invariant.
-        """
-        from sglang.srt.layers.dp_attention import DpPaddingMode
-        from sglang.srt.model_executor.forward_batch_info import (
-            ForwardBatch,
-            ForwardMode,
-        )
-
-        spec_info = None
-        if self.model_runner.is_draft_worker:
-            from sglang.srt.speculative.eagle_info import EagleDraftInput
-
-            spec_info = EagleDraftInput(
-                hidden_states=self.static_draft_hidden_states[:num_tokens],
-            )
-
-        registry = self.buffer_registry
-        bs = 1
-
-        def _slot(name):
-            return registry.get_slot(name).slice_for(bs, num_tokens)
-
-        with torch.device(self.device):
-            seq_lens = torch.full((bs,), num_tokens, dtype=torch.int64)
-            extend_seq_lens = torch.full((bs,), num_tokens, dtype=torch.int64)
-            extend_prefix_lens = torch.zeros((bs,), dtype=torch.int64)
-            extend_start_loc = torch.zeros((bs,), dtype=torch.int64)
-            req_pool_indices = torch.arange(bs, dtype=torch.int64)
-            orig_seq_lens = torch.full((bs,), num_tokens, dtype=torch.int64)
-
-        return ForwardBatch(
-            forward_mode=ForwardMode.EXTEND,
-            batch_size=bs,
-            input_ids=_slot("input_ids"),
-            input_embeds=(
-                _slot("input_embeds") if registry.has_slot("input_embeds") else None
-            ),
-            req_pool_indices=req_pool_indices,
-            seq_lens=seq_lens,
-            next_token_logits_buffer=None,
-            orig_seq_lens=orig_seq_lens,
-            seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
-            out_cache_loc=_slot("out_cache_loc"),
-            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=extend_seq_lens,
-            extend_prefix_lens=extend_prefix_lens,
-            extend_start_loc=extend_start_loc,
-            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=_slot("positions"),
-            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=(
-                _slot("mrope_positions")
-                if registry.has_slot("mrope_positions")
-                else None
-            ),
-            spec_algorithm=None,
-            spec_info=spec_info,
-            capture_hidden_mode=self.capture_hidden_mode,
-            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)
-        with (
-            forward_context(
-                ForwardContext(attn_backend=self.model_runner.attn_backend)
-            ),
-            set_forward_context(
-                forward_batch,
-                self.attention_layers,
-                self.quant_config,
-                self.moe_layers,
-                self.moe_fusions,
-            ),
-        ):
-            self._init_forward_metadata_for_capture(forward_batch, num_tokens)
-            self._run_forward(forward_batch, num_tokens)
-
-    def _init_forward_metadata_for_capture(self, forward_batch, num_tokens):
-        attn_backend = self.model_runner.attn_backend
-        if not self.use_captured_attn_metadata:
-            attn_backend.init_forward_metadata(forward_batch)
-            return
-        metadata = attn_backend.init_forward_metadata_for_breakable_cuda_graph_capture(
-            forward_batch
-        )
-        assert self.attn_metadata_buffers is not None
-        self.attn_metadata_buffers[num_tokens] = metadata
-
-    def _prepare_forward_metadata_for_replay(
-        self, forward_batch, static_forward_batch, num_tokens
-    ):
-        attn_backend = self.model_runner.attn_backend
-        if not self.use_captured_attn_metadata:
-            attn_backend.init_forward_metadata(forward_batch)
-            return
-        assert self.attn_metadata_buffers is not None
-        metadata = self.attn_metadata_buffers[num_tokens]
-        attn_backend.prepare_forward_metadata_for_breakable_cuda_graph_replay(
-            metadata,
-            forward_batch,
-            static_forward_batch=static_forward_batch,
-        )
-
-    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"):
-        if self.layer_model is None:
-            return False
-        if forward_batch.forward_mode.is_target_verify():
-            return False
-        if forward_batch.capture_hidden_mode != self.capture_hidden_mode:
-            return False
-        if forward_batch.input_embeds is not None:
-            return False
-        if forward_batch.replace_embeds is not None:
-            return False
-        if self._has_inactive_dp_rank(forward_batch):
-            return False
-        if (
-            forward_batch.global_num_tokens_cpu is not None
-            and not forward_batch.can_run_dp_breakable_cuda_graph
-        ):
-            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._init_forward_metadata_for_capture(forward_batch, num_tokens)
-
-        def run_once():
-            return self._run_forward(forward_batch, num_tokens)
-
-        with forward_context(
-            ForwardContext(attn_backend=self.model_runner.attn_backend)
-        ):
-            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_prepare(self, forward_batch, **kwargs):
-        # TODO: fix PiecewiseCudaGraphRunner to support draft workers as well.
-        static_forward_batch = PiecewiseCudaGraphRunner.replay_prepare(
-            self, forward_batch, **kwargs
-        )
-        if self.model_runner.is_draft_worker and forward_batch.spec_info is not None:
-            num_tokens = len(forward_batch.input_ids)
-            self.static_draft_hidden_states[:num_tokens].copy_(
-                forward_batch.spec_info.hidden_states
-            )
-        return static_forward_batch
-
-    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]
-
-        captured_graph = self.graphs[static_num_tokens]
-        captured_hidden = self.output_buffers[static_num_tokens]
-
-        # Closure replaces layer_model.forward for the duration of the outer
-        # model.forward call. Replays the captured CUDAGraph and hands the
-        # outer forward the captured hidden_states; logits_processor / pooler
-        # then runs eagerly on top with the live multi-req forward_batch.
-        def replay_layer_forward(*args, **layer_kwargs):
-            ie = layer_kwargs.get("input_embeds") or (
-                args[self._input_embeds_arg_idx]
-                if self.use_input_embeds and len(args) > self._input_embeds_arg_idx
-                else None
-            )
-            if self.use_input_embeds:
-                if ie is None:
-                    raise ValueError("BCG replay expects input_embeds but got None")
-                self.buffer_registry.get_slot("input_embeds").slice_for(
-                    1, static_num_tokens
-                ).copy_(ie[:static_num_tokens])
-            else:
-                if ie is not None:
-                    raise ValueError(
-                        "BCG replay got unexpected input_embeds on non-multimodal model"
-                    )
-            captured_graph.replay()
-            return captured_hidden
-
-        with enable_breakable_cuda_graph():
-            static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
-
-            original_layer_forward = self.layer_model.forward
-            self.layer_model.forward = replay_layer_forward
-            try:
-                self._prepare_forward_metadata_for_replay(
-                    forward_batch, static_forward_batch, static_num_tokens
-                )
-                with set_forward_context(
-                    static_forward_batch,
-                    self.attention_layers,
-                    self.quant_config,
-                    self.moe_layers,
-                    self.moe_fusions,
-                ):
-                    output = self.model_runner.model.forward(
-                        static_forward_batch.input_ids,
-                        static_forward_batch.positions,
-                        static_forward_batch,
-                        **kwargs,
-                    )
-            finally:
-                self.layer_model.forward = original_layer_forward
-        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/cpu_graph_runner.py

@@ -29,7 +29,6 @@ import tqdm
 from sglang.srt.distributed import get_tensor_model_parallel_rank
 from sglang.srt.distributed.parallel_state import GroupCoordinator
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
-from sglang.srt.model_executor.cuda_graph_runner import model_capture_mode
 from sglang.srt.model_executor.forward_batch_info import (
     CaptureHiddenMode,
     ForwardBatch,
@@ -38,6 +37,7 @@ from sglang.srt.model_executor.forward_batch_info import (
     enable_num_token_non_padded,
 )
 from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
+from sglang.srt.model_executor.runner_utils.capture_mode import model_capture_mode
 from sglang.srt.utils import (
     empty_context,
     log_info_on_rank0,
@@ -126,10 +126,10 @@ def set_torch_compile_config():
 def get_batch_sizes_to_capture(model_runner: ModelRunner):
     # torch compile speeds up decoding by reducing python overhead on CPU
     server_args = model_runner.server_args
-    # Note that we reuse server_args.cuda_graph_bs here.
+    # Reuse cuda_graph_config[decode].bs here.
     # Users can customize the batch sizes supported by cpu_graph, such as:
-    # --cuda-graph-bs 1 2 4 8 16
-    capture_bs = server_args.cuda_graph_bs
+    # --cuda-graph-bs-decode 1 2 4 8 16
+    capture_bs = server_args.cuda_graph_config.decode.bs
     assert (
         max(capture_bs) <= server_args.torch_compile_max_bs
     ), f"{capture_bs=}, {server_args.torch_compile_max_bs=}"

python/sglang/srt/model_executor/cuda_graph_config.py

@@ -0,0 +1,211 @@
+"""Phase / backend identifiers, the canonical default for
+cuda_graph_config, and the --cuda-graph-config JSON CLI parser.
+
+Module-level imports are pure stdlib — no torch / sglang.srt deps — so
+ServerArgs can import everything here without pulling in backend
+classes. check_cuda_graph_backend lazy-imports get_global_server_args
+inside the function body to preserve that invariant.
+"""
+
+import argparse
+import dataclasses
+import json
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+
+class Phase:
+    """The two phases of model forward."""
+
+    DECODE = "decode"
+    PREFILL = "prefill"
+    ALL = (DECODE, PREFILL)
+
+
+class Backend:
+    """CUDA graph capture backends a phase can use."""
+
+    FULL = "full"
+    BREAKABLE = "breakable"
+    TC_PIECEWISE = "tc_piecewise"
+    DISABLED = "disabled"
+    ALL = (FULL, BREAKABLE, TC_PIECEWISE, DISABLED)
+
+
+ALLOWED_BACKENDS_PER_PHASE = {
+    Phase.DECODE: (
+        Backend.FULL,
+        Backend.BREAKABLE,
+        Backend.TC_PIECEWISE,
+        Backend.DISABLED,
+    ),
+    # full is rejected for prefill — full CUDA graph capture only
+    # fits fixed-shape and prefill is variable-shape. Use breakable
+    # or tc_piecewise for prefill.
+    Phase.PREFILL: (Backend.BREAKABLE, Backend.TC_PIECEWISE, Backend.DISABLED),
+}
+
+# Per-phase settings schema. Keys other than backend are runner-level
+# (read by any backend in that phase); tc_compiler is the lone
+# backend-specific knob (only meaningful when backend == tc_piecewise).
+# For prefill, bs carries the captured shape size (token count for
+# tc_piecewise, request count for breakable) — one shape knob per phase.
+ALLOWED_KEYS_PER_PHASE = {
+    Phase.DECODE: ("backend", "max_bs", "bs", "tc_compiler"),
+    Phase.PREFILL: ("backend", "max_bs", "bs", "tc_compiler"),
+}
+
+
+@dataclass
+class PhaseConfig:
+    """Per-phase CUDA graph settings."""
+
+    backend: str = Backend.DISABLED
+    max_bs: Optional[int] = None
+    bs: Optional[List[int]] = None
+    # Only meaningful when backend == tc_piecewise; ignored otherwise.
+    tc_compiler: str = "eager"
+
+
+@dataclass
+class CudaGraphConfig:
+    """Top-level CUDA graph config: one PhaseConfig per phase."""
+
+    decode: PhaseConfig = field(
+        default_factory=lambda: PhaseConfig(backend=Backend.FULL)
+    )
+    prefill: PhaseConfig = field(
+        default_factory=lambda: PhaseConfig(backend=Backend.TC_PIECEWISE)
+    )
+
+    def __getitem__(self, phase: str) -> PhaseConfig:
+        """Phase-string lookup; kept for migration ergonomics."""
+        if phase not in Phase.ALL:
+            raise KeyError(phase)
+        return getattr(self, phase)
+
+    def to_dict(self) -> Dict[str, Dict[str, Any]]:
+        # Diff-only, not asdict: the parser locks every (phase, key) it sees,
+        # so emitting defaults would lock fields the caller never set.
+        baseline = default_cuda_graph_config()
+        return {
+            Phase.DECODE: _diff_phase(self.decode, baseline.decode),
+            Phase.PREFILL: _diff_phase(self.prefill, baseline.prefill),
+        }
+
+    @classmethod
+    def from_dict(cls, raw: Optional[Dict[str, Dict[str, Any]]]) -> "CudaGraphConfig":
+        """Build from a (partial) dict of overrides, defaults fill the rest.
+        Unknown phases / keys are silently dropped — the JSON-input
+        validator (parse_cuda_graph_config_arg) rejects them upstream."""
+        cfg = cls()
+        if not raw:
+            return cfg
+        for phase, phase_settings in raw.items():
+            if phase not in Phase.ALL or not isinstance(phase_settings, dict):
+                continue
+            phase_cfg = getattr(cfg, phase)
+            allowed = ALLOWED_KEYS_PER_PHASE[phase]
+            for key, value in phase_settings.items():
+                if key in allowed:
+                    setattr(phase_cfg, key, value)
+        return cfg
+
+
+def default_cuda_graph_config() -> CudaGraphConfig:
+    """Fresh CudaGraphConfig populated with canonical defaults."""
+    return CudaGraphConfig()
+
+
+def _diff_phase(actual: PhaseConfig, baseline: PhaseConfig) -> Dict[str, Any]:
+    """Return only fields whose value differs from the per-phase default."""
+    return {
+        f.name: getattr(actual, f.name)
+        for f in dataclasses.fields(actual)
+        if getattr(actual, f.name) != getattr(baseline, f.name)
+    }
+
+
+def check_cuda_graph_backend(phase: str, backend: str) -> bool:
+    """True if cuda_graph_config[phase].backend == backend on the
+    global server args. Returns False if the global server args have not
+    been initialized yet (e.g. unit tests, early startup)."""
+    from sglang.srt.server_args import get_global_server_args
+
+    try:
+        server_args = get_global_server_args()
+    except ValueError:
+        return False
+    cfg = server_args.cuda_graph_config
+    if cfg is None or phase not in Phase.ALL:
+        return False
+    return getattr(cfg, phase).backend == backend
+
+
+def cuda_graph_fully_disabled() -> bool:
+    """True iff cuda_graph_config has Backend.DISABLED on every phase.
+
+    Use at sites that ask the legacy server_args.disable_cuda_graph
+    question ("no CG anywhere globally") — e.g., preallocating buffers
+    that any captured graph would otherwise reuse, or one-shot init
+    that's a no-op when CG is completely off.
+    """
+    return check_cuda_graph_backend(
+        Phase.DECODE, Backend.DISABLED
+    ) and check_cuda_graph_backend(Phase.PREFILL, Backend.DISABLED)
+
+
+def parse_cuda_graph_config_arg(raw: str) -> Dict[str, Dict[str, Any]]:
+    """argparse type for --cuda-graph-config: parse JSON dict of
+    phase → settings dict. Each phase's settings dict is itself validated
+    against ALLOWED_KEYS_PER_PHASE. Returns a plain dict — the
+    precedence pipeline in ServerArgs converts to CudaGraphConfig
+    after merging."""
+    try:
+        parsed = json.loads(raw)
+    except json.JSONDecodeError as e:
+        raise argparse.ArgumentTypeError(f"--cuda-graph-config must be JSON: {e}")
+    if not isinstance(parsed, dict):
+        raise argparse.ArgumentTypeError(
+            f"--cuda-graph-config must be a JSON object, got {type(parsed).__name__}"
+        )
+
+    result: Dict[str, Dict[str, Any]] = {}
+    for phase, phase_settings in parsed.items():
+        phase = str(phase)
+        if phase not in Phase.ALL:
+            raise argparse.ArgumentTypeError(
+                f"--cuda-graph-config: unknown phase '{phase}', expected one of {Phase.ALL}"
+            )
+        if not isinstance(phase_settings, dict):
+            raise argparse.ArgumentTypeError(
+                f"--cuda-graph-config['{phase}'] must be a JSON object, got "
+                f"{type(phase_settings).__name__}"
+            )
+        allowed = ALLOWED_KEYS_PER_PHASE[phase]
+        result[phase] = {}
+        for key, value in phase_settings.items():
+            if key not in allowed:
+                raise argparse.ArgumentTypeError(
+                    f"--cuda-graph-config['{phase}']: unknown key '{key}', expected one of {allowed}"
+                )
+            result[phase][key] = value
+    return result
+
+
+def explicit_keys_in(
+    settings: Optional[Dict[str, Dict[str, Any]]],
+) -> set:
+    """Return the set of (phase, key) tuples present in settings
+    (the raw dict form, as it arrives from CLI/SDK). Used by ServerArgs
+    to track keys the user explicitly set so the auto-disable cascade can
+    skip them."""
+    out: set = set()
+    if not settings:
+        return out
+    for phase, phase_settings in settings.items():
+        if not isinstance(phase_settings, dict):
+            continue
+        for key in phase_settings.keys():
+            out.add((phase, key))
+    return out

python/sglang/srt/model_executor/forward_context.py

@@ -1,23 +1,23 @@
 """Per-forward-call control context.
 
-Owns ``ForwardContext`` — a frozen dataclass holding control configs the model
-layer reads at depth via ``get_forward_context()``. The only mandatory field
-today is ``attn_backend``; pool refs are derived from ``attn_backend.*``
-(every backend caches them at ``__init__``), so a published ``ForwardContext``
+Owns ForwardContext — a frozen dataclass holding control configs the model
+layer reads at depth via get_forward_context(). The only mandatory field
+today is attn_backend; pool refs are derived from attn_backend.*
+(every backend caches them at __init__), so a published ForwardContext
 is enough to resolve the active pools without a separate global.
 
-``ModelRunner._forward_raw`` publishes a fresh ``ForwardContext`` for the
+ModelRunner._forward_raw publishes a fresh ForwardContext for the
 duration of each forward; callers that need a per-call override (PDmux
 per-stream backend, frozen-KV MTP draft loop, TBO per-child dispatch) use
-``dataclasses.replace`` and wrap the override scope with ``forward_context()``.
+dataclasses.replace and wrap the override scope with forward_context().
 
-Distinct from ``sglang.srt.compilation.piecewise_context_manager.ForwardContext``,
+Distinct from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph.TcPiecewiseForwardContext,
 which collects compilation-time refs for the piecewise CUDA graph backend.
 
-Concurrency: ``_current`` is a plain module-level global, not thread-local.
-This matches the ``global_server_args`` precedent and is safe because each
+Concurrency: _current is a plain module-level global, not thread-local.
+This matches the global_server_args precedent and is safe because each
 forward runs synchronously on a single Python thread per worker process. If
-worker threads ever share a process, migrate to ``contextvars.ContextVar``.
+worker threads ever share a process, migrate to contextvars.ContextVar.
 """
 
 from __future__ import annotations
@@ -33,9 +33,9 @@ if TYPE_CHECKING:
 
 @dataclass(frozen=True, slots=True)
 class ForwardContext:
-    """Per-forward-call control configs. Read via ``get_forward_context()``;
+    """Per-forward-call control configs. Read via get_forward_context();
     extend by adding fields here. Frozen so accidental mutation raises at
-    write time — use ``dataclasses.replace`` for per-call overrides."""
+    write time — use dataclasses.replace for per-call overrides."""
 
     attn_backend: AttentionBackend
 
@@ -45,7 +45,7 @@ _current: Optional[ForwardContext] = None
 
 def set_forward_context(ctx: Optional[ForwardContext]) -> Optional[ForwardContext]:
     """Set the active context; return the previous one for explicit
-    save/restore. Prefer the ``forward_context()`` context manager."""
+    save/restore. Prefer the forward_context() context manager."""
     global _current
     prev, _current = _current, ctx
     return prev

python/sglang/srt/model_executor/model_runner.py

@@ -35,10 +35,7 @@ import torch.distributed as dist
 from torch import nn
 
 from sglang.jit_kernel.ngram_embedding import update_token_table_decode
-from sglang.srt.compilation.piecewise_context_manager import (
-    enable_piecewise_cuda_graph,
-    set_forward_context,
-)
+from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config
 from sglang.srt.configs import (
     BailingHybridConfig,
     FalconH1Config,
@@ -140,19 +137,17 @@ 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_buffer_registry import (
     CudaGraphBufferRegistry,
     build_decode_registry,
     build_prefill_registry,
 )
-from sglang.srt.model_executor.cuda_graph_runner import (
-    CudaGraphRunner,
-    _allocate_decode_buffers,
-    set_torch_compile_config,
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+    cuda_graph_fully_disabled,
 )
 from sglang.srt.model_executor.forward_batch_info import (
     CaptureHiddenMode,
@@ -169,10 +164,13 @@ from sglang.srt.model_executor.hook_manager import register_forward_hooks
 from sglang.srt.model_executor.model_runner_kv_cache_mixin import (
     ModelRunnerKVCacheMixin,
 )
-from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
-    PiecewiseCudaGraphRunner,
-)
 from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig
+from sglang.srt.model_executor.runner import (
+    PrefillCudaGraphRunner,
+)
+from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
+    _allocate_decode_buffers,
+)
 from sglang.srt.model_loader.loader import DefaultModelLoader, get_model_loader
 from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
     RemoteInstanceWeightLoaderBackend,
@@ -763,12 +761,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         # Init lora
         if server_args.enable_lora:
             self.init_lora_manager()
-            if not server_args.disable_cuda_graph:
+            if not cuda_graph_fully_disabled():
                 # Phase 1 of LoRA CUDA graph init: pre-allocate large MoE
                 # intermediate buffers before init_memory_pool() so memory
-                # profiling accounts for them.  Phase 2 (dense LoRA batch
-                # metadata) is handled in CudaGraphRunner.__init__() via
-                # lora_manager.init_cuda_graph_batch_info().
+                # profiling accounts for them. The buffers are reused by
+                # any captured graph (decode today; widen here so any
+                # future prefill capture path also picks them up).
                 self._init_lora_cuda_graph_moe_buffers()
 
         # Enable batch invariant mode
@@ -801,7 +799,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         self.init_indexer_capturer()
 
         # TODO: Refactor device-specific init branches into platform interface (separate PR).
-        # Must be called BEFORE init_device_graphs() so CUDA graph capture
+        # Must be called BEFORE init_decode_cuda_graph() so CUDA graph capture
         # runs with aux hidden state capture enabled.
         self.init_aux_hidden_state_capture()
 
@@ -831,13 +829,13 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             self.init_attention_backend()
             self.kernel_warmup()
             self._pre_initialize_flashinfer_allreduce_workspace()
-            self.init_device_graphs()
+            self.init_decode_cuda_graph()
         elif self.device == "cpu":
             self.init_attention_backend()
-            self.init_device_graphs()
+            self.init_decode_cuda_graph()
         elif self.device == "npu":
             self.init_attention_backend()
-            # lazy init for zbal with mix mode(before graph capture when enable_cuda_graph)
+            # lazy init for zbal with mix mode (before graph capture when enable_cuda_graph)
             if envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() > 0 and not self.is_draft_worker:
                 from sglang.srt.hardware_backend.npu.utils import lazy_init_zbal_gva_mem
 
@@ -848,16 +846,16 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                     get_world_group().world_size,
                     get_world_group().cpu_group,
                 )
-            self.init_device_graphs()
+            self.init_decode_cuda_graph()
         elif current_platform.is_out_of_tree():
             self.init_attention_backend()
             if current_platform.support_cuda_graph():
-                self.init_device_graphs()
+                self.init_decode_cuda_graph()
             else:
-                self.graph_runner = None
+                self.decode_cuda_graph_runner = None
                 self.graph_mem_usage = 0
         else:
-            self.graph_runner = None
+            self.decode_cuda_graph_runner = None
             self.graph_mem_usage = 0
             self.init_attention_backend()
 
@@ -865,7 +863,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             register_forward_hooks(self.model, server_args.forward_hooks)
 
         # Initialize piecewise CUDA graph
-        self.init_piecewise_cuda_graphs()
+        self.init_prefill_cuda_graph()
 
         self.prealloc_symmetric_memory_pool()
 
@@ -1720,7 +1718,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 and current_platform.support_cuda_graph()
             )
         ):
-            self.init_device_graphs()
+            self.init_decode_cuda_graph()
 
         logger.info("Update weights end.")
         return True, "Succeeded to update model weights."
@@ -2082,7 +2080,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         """
         from sglang.srt.lora.layers import FusedMoEWithLoRA
 
-        max_bs = self.server_args.cuda_graph_max_bs
+        max_bs = self.server_args.cuda_graph_config.decode.max_bs
         max_loras = self.server_args.max_loras_per_batch
         for module in self.model.modules():
             if isinstance(module, FusedMoEWithLoRA):
@@ -2546,8 +2544,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
     ):
         """Run a dummy forward pass for warmup/profiling.
 
-        ``forward_mode_override`` forces EXTEND/DECODE regardless of
-        ``is_generation`` (used by the PP-parallel DeepGEMM warmup).
+        forward_mode_override forces EXTEND/DECODE regardless of
+        is_generation (used by the PP-parallel DeepGEMM warmup).
         """
         if forward_mode_override is not None:
             capture_forward_mode = forward_mode_override
@@ -2845,7 +2843,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
     def maybe_update_ngram_token_table(
         self,
         next_token_ids: torch.Tensor,
-        forward_batch: "ForwardBatch",
+        forward_batch: ForwardBatch,
     ):
         """Update the ngram embedding token table after sampling."""
         ngram_embedding_info = forward_batch.ngram_embedding_info
@@ -2862,9 +2860,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             column_starts=ngram_embedding_info.out_column_starts,
         )
 
-    def init_device_graphs(self):
+    def init_decode_cuda_graph(self):
         """Capture device graphs."""
-        self.graph_runner = None
+        self.decode_cuda_graph_runner = None
         self.graph_mem_usage = 0
 
         if not self.is_generation:
@@ -2874,7 +2872,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         if self.server_args.model_impl.lower() == ModelImpl.MINDSPORE:
             return
 
-        if self.device != "cpu" and self.server_args.disable_cuda_graph:
+        if self.device != "cpu" and check_cuda_graph_backend(
+            Phase.DECODE, Backend.DISABLED
+        ):
             return
 
         if self.device == "cpu" and not self.server_args.enable_torch_compile:
@@ -2896,16 +2896,20 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         )
         if current_platform.is_out_of_tree():
             GraphRunnerCls = current_platform.get_graph_runner_cls()
-            self.graph_runner = GraphRunnerCls(self)
+            self.decode_cuda_graph_runner = GraphRunnerCls(self)
         else:
+            from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
+                DecodeCudaGraphRunner,
+            )
+
             graph_runners = defaultdict(
-                lambda: CudaGraphRunner,
+                lambda: DecodeCudaGraphRunner,
                 {
                     "cpu": CPUGraphRunner,
                     "npu": NPUGraphRunner,
                 },
             )
-            self.graph_runner = graph_runners[self.device](self)
+            self.decode_cuda_graph_runner = graph_runners[self.device](self)
 
         after_mem = get_available_gpu_memory(self.device, self.gpu_id)
         self.graph_mem_usage = before_mem - after_mem
@@ -2914,13 +2918,15 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             f"mem usage={self.graph_mem_usage:.2f} GB. avail mem={after_mem:.2f} GB."
         )
 
-    def init_piecewise_cuda_graphs(self, force_for_draft_worker: bool = False):
+    def init_prefill_cuda_graph(self, force_for_draft_worker: bool = False):
         """Initialize piecewise CUDA graph runner."""
-        self.piecewise_cuda_graph_runner = None
+        self.prefill_cuda_graph_runner = None
 
-        if self.server_args.disable_piecewise_cuda_graph:
+        if check_cuda_graph_backend(Phase.PREFILL, Backend.DISABLED):
             logger.info(
-                "Disable piecewise CUDA graph because --disable-piecewise-cuda-graph is set"
+                "Disable prefill CUDA graph because cuda_graph_config "
+                "resolved prefill.backend='disabled' (e.g. via "
+                "--cuda-graph-backend-prefill=disabled or auto-disable rules)."
             )
             return
 
@@ -2938,7 +2944,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             return
 
         # Disable piecewise CUDA graph for non capture size
-        if not self.server_args.piecewise_cuda_graph_tokens:
+        if not self.server_args.cuda_graph_config.prefill.bs:
             logger.warning(
                 "Disable piecewise CUDA graph because the capture size is not set"
             )
@@ -2971,8 +2977,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 elif hasattr(layer.self_attn, "attn_mqa"):
                     # For DeepSeek model
                     attn_layer = layer.self_attn.attn_mqa
-                    if _is_hip and hasattr(layer.self_attn, "attn_mha"):
-                        attn_layer._pcg_mha_companion = layer.self_attn.attn_mha
             # For hybrid model
             elif hasattr(layer, "attn"):
                 attn_layer = layer.attn
@@ -3037,11 +3041,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             f"Capture piecewise CUDA graph begin. avail mem={before_mem:.2f} GB"
         )
 
-        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)
+        self.prefill_cuda_graph_runner = PrefillCudaGraphRunner(self)
 
         after_mem = get_available_gpu_memory(self.device, self.gpu_id)
         mem_usage = before_mem - after_mem
@@ -3278,12 +3278,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
 
         # Check piecewies cuda graph
         can_run_graph = (
-            self.piecewise_cuda_graph_runner is not None
-            and self.piecewise_cuda_graph_runner.can_run(forward_batch)
+            self.prefill_cuda_graph_runner is not None
+            and self.prefill_cuda_graph_runner.can_run(forward_batch)
         )
         if can_run_graph:
             # TODO: device_timer.wrap is too broad here — it also includes
-            # replay_prepare time. Move timing into the piecewise cuda graph
+            # replay_prepare time. Move timing into the prefill cuda graph
             # runner to capture only the model.forward part.
             ctx = (
                 self.device_timer.wrap(metadata={"category": "extend"})
@@ -3291,7 +3291,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 else contextlib.nullcontext()
             )
             with ctx:
-                ret = self.piecewise_cuda_graph_runner.replay(forward_batch, **kwargs)
+                ret = self.prefill_cuda_graph_runner.replay(forward_batch, **kwargs)
             return (ret, can_run_graph)
 
         if not self.server_args.enable_pdmux and self.device == "cuda":
@@ -3311,37 +3311,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             else contextlib.nullcontext()
         )
         with ctx:
-            if _is_hip and self.piecewise_cuda_graph_runner is not None:
-                # AMD/HIP: when PCG is enabled but the batch exceeds max captured
-                # size, run eagerly under enable_piecewise_cuda_graph() and
-                # set_forward_context() so that (a) Dynamo guards on
-                # _in_piecewise_cuda_graph stay consistent with the PCG-traced
-                # graph (preventing runtime recompilation) and (b) PCG-specific
-                # code paths (MoE, attention) can access their layer objects.
-                with (
-                    enable_piecewise_cuda_graph(),
-                    set_forward_context(
-                        forward_batch,
-                        self.attention_layers,
-                        getattr(self.model, "quant_config", None),
-                        self.moe_layers,
-                        self.moe_fusions,
-                        dsa_indexers=self.dsa_indexers,
-                    ),
-                ):
-                    ret = self.model.forward(
-                        forward_batch.input_ids,
-                        forward_batch.positions,
-                        forward_batch,
-                        **kwargs,
-                    )
-            else:
-                ret = self.model.forward(
-                    forward_batch.input_ids,
-                    forward_batch.positions,
-                    forward_batch,
-                    **kwargs,
-                )
+            ret = self.model.forward(
+                forward_batch.input_ids,
+                forward_batch.positions,
+                forward_batch,
+                **kwargs,
+            )
         return (ret, can_run_graph)
 
     def forward_idle(
@@ -3472,7 +3447,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             output.routed_experts_output = experts_capturer.on_forward_end(
                 forward_batch=forward_batch,
                 can_run_graph=output.can_run_graph,
-                cuda_graph_batch=getattr(self.graph_runner, "bs", None),
+                cuda_graph_batch=getattr(self.decode_cuda_graph_runner, "bs", None),
                 no_copy_to_cpu=no_copy_to_cpu,
             )
 
@@ -3480,7 +3455,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             output.indexer_topk_output = indexer_capturer.on_forward_end(
                 forward_batch=forward_batch,
                 can_run_graph=output.can_run_graph,
-                cuda_graph_batch=getattr(self.graph_runner, "bs", None),
+                cuda_graph_batch=getattr(self.decode_cuda_graph_runner, "bs", None),
                 no_copy_to_cpu=no_copy_to_cpu,
             )
 
@@ -3506,9 +3481,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         reinit_attn_backend: bool = False,
         split_forward_count: int = 1,
     ) -> ModelRunnerOutput:
-        # Honor an outer-published context (spec workers wrap each per-step
-        # draft forward with the i-th child backend); otherwise publish this
-        # runner's own attn_backend for the forward.
         if has_forward_context():
             ctx_mgr = contextlib.nullcontext()
         else:
@@ -3521,21 +3493,21 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             )
             can_run_graph = bool(
                 mode_check()
-                and self.graph_runner
-                and self.graph_runner.can_run(forward_batch)
+                and self.decode_cuda_graph_runner
+                and self.decode_cuda_graph_runner.can_run(forward_batch)
             )
 
-            # Hisparse coordinator — backends now read it from self.model_runner.
             if (
                 forward_batch.forward_mode.is_decode()
                 and self.hisparse_coordinator is not None
             ):
+                forward_batch.hisparse_coordinator = self.hisparse_coordinator
                 self.hisparse_coordinator.wait_for_pending_backup()
                 self.hisparse_coordinator.num_real_reqs.fill_(forward_batch.batch_size)
 
             # Replay cuda graph if applicable
             if can_run_graph:
-                ret = self.graph_runner.replay(
+                ret = self.decode_cuda_graph_runner.replay(
                     forward_batch,
                     pp_proxy_tensors=pp_proxy_tensors,
                 )

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -1,860 +0,0 @@
-# Copyright 2023-2024 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.
-# ==============================================================================
-"""Run the model with cuda graph and torch.compile."""
-
-from __future__ import annotations
-
-import bisect
-import gc
-import logging
-import warnings
-from contextlib import contextmanager
-from typing import TYPE_CHECKING, Union
-
-import torch
-import tqdm
-
-from sglang.srt.batch_overlap.two_batch_overlap import TboCudaGraphRunnerPlugin
-from sglang.srt.compilation.compilation_config import CompilationConfig
-from sglang.srt.compilation.compile import install_torch_compiled
-from sglang.srt.compilation.piecewise_context_manager import (
-    enable_piecewise_cuda_graph,
-    enable_piecewise_cuda_graph_compile,
-    set_forward_context,
-    set_pcg_capture_stream,
-)
-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 (
-    DpPaddingMode,
-    get_attention_cp_size,
-    get_attention_tp_rank,
-    get_attention_tp_size,
-    set_dp_buffer_len,
-    set_is_extend_in_batch,
-)
-from sglang.srt.layers.logits_processor import LogitsProcessorOutput
-from sglang.srt.layers.moe.utils import get_moe_a2a_backend
-from sglang.srt.layers.pooler import EmbeddingPoolerOutput
-from sglang.srt.layers.utils import MultiPlatformOp
-from sglang.srt.model_executor.cuda_graph_buffer_registry import build_prefill_registry
-from sglang.srt.model_executor.forward_batch_info import (
-    CaptureHiddenMode,
-    ForwardBatch,
-    ForwardMode,
-    PPProxyTensors,
-)
-from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
-from sglang.srt.utils import (
-    get_available_gpu_memory,
-    get_bool_env_var,
-    is_hip,
-    is_musa,
-    is_npu,
-    log_info_on_rank0,
-    require_gathered_buffer,
-)
-
-_is_hip = is_hip()
-_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
-# Suppress Dynamo warning about tracing through lru_cache-wrapped functions (e.g., is_arch_support_pdl).
-warnings.filterwarnings("ignore", message=".*lru_cache.*", module="torch._dynamo")
-logger = logging.getLogger(__name__)
-
-if TYPE_CHECKING:
-    from sglang.srt.model_executor.model_runner import ModelRunner
-
-_is_musa = is_musa()
-
-
-@contextmanager
-def freeze_gc(enable_cudagraph_gc: bool):
-    """
-    Optimize garbage collection during CUDA graph capture.
-    Clean up, then freeze all remaining objects from being included
-    in future collections if GC is disabled during capture.
-    """
-    gc.collect()
-    should_freeze = not enable_cudagraph_gc
-    if should_freeze:
-        gc.freeze()
-    try:
-        yield
-    finally:
-        if should_freeze:
-            gc.unfreeze()
-
-
-def _to_torch(model: torch.nn.Module, reverse: bool, num_tokens: int):
-    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):
-            _to_torch(sub, reverse, num_tokens)
-
-
-@contextmanager
-def patch_model(model: torch.nn.Module, compiler: str):
-    try:
-        if compiler != "eager":
-            _to_torch(model, reverse=False, num_tokens=16)
-        yield model
-    finally:
-        _to_torch(model, reverse=True, num_tokens=16)
-
-
-# Reuse this memory pool across all cuda graph runners.
-global_graph_memory_pool = None
-
-
-def get_global_graph_memory_pool():
-    return global_graph_memory_pool
-
-
-def set_global_graph_memory_pool(val):
-    global global_graph_memory_pool
-    global_graph_memory_pool = val
-
-
-def set_torch_compile_config():
-    import torch._dynamo.config
-
-    # Resolve torch._dynamo.exc.FailOnRecompileLimitHit
-    torch._dynamo.config.accumulated_cache_size_limit = 1024
-    if hasattr(torch._dynamo.config, "cache_size_limit"):
-        torch._dynamo.config.cache_size_limit = 1024
-
-    if _is_musa:
-        from sglang.srt.hardware_backend.musa.utils.patch_torch import (
-            patch_fx_custom_device,
-        )
-
-        patch_fx_custom_device()
-
-
-class PiecewiseCudaGraphRunner:
-    """A PiecewiseCudaGraphRunner runs the forward pass of a model with cuda graph and torch.compile."""
-
-    def is_mamba_track_enabled(self):
-        return (
-            self.model_runner.server_args.enable_mamba_extra_buffer()
-            and not self.model_runner.server_args.disable_radix_cache
-            and self.model_runner.spec_algorithm.is_none()
-        )
-
-    def __init__(self, model_runner: ModelRunner):
-        # Parse args
-        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.tp_size = model_runner.server_args.tp_size
-        self.dp_size = model_runner.server_args.dp_size
-        self.pp_size = model_runner.server_args.pp_size
-
-        self.attn_tp_size = get_attention_tp_size()
-        self.attn_tp_rank = get_attention_tp_rank()
-
-        set_torch_compile_config()
-
-        assert (
-            self.model_runner.server_args.piecewise_cuda_graph_tokens is not None
-        ), "piecewise_cuda_graph_tokens is not set"
-        assert self.model_runner.server_args.piecewise_cuda_graph_compiler in [
-            "eager",
-            "inductor",
-        ], "By now, only eager and inductor are supported for piecewise cuda graph compiler."
-        self.compile_config = CompilationConfig(
-            self.model_runner.server_args.piecewise_cuda_graph_tokens,
-            self.model_runner.server_args.piecewise_cuda_graph_compiler,
-            self.model_runner.server_args.enable_torch_compile_debug_mode,
-        )
-        if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
-            self.compile_config.add_split_op(
-                "sglang.moe_forward_piecewise_cuda_graph_impl"
-            )
-
-        self.quant_config = getattr(self.model_runner.model, "quant_config", None)
-
-        # Batch sizes to capture
-        self.capture_num_tokens = self.compile_config.get_capture_sizes()
-        # When the layer communicator scatters/gathers across the attention TP
-        # group (e.g. with --moe-dense-tp-size 1), the model's reduce_scatter
-        # requires the token count to be divisible by attn_tp_size * attn_cp_size.
-        # Drop captures that would violate this (mirrors the filter used by
-        # the regular CUDA graph runner in get_batch_sizes_to_capture).
-        if require_gathered_buffer(self.model_runner.server_args):
-            mul_base = self.attn_tp_size
-            attn_cp_size = get_attention_cp_size()
-            if mul_base % attn_cp_size != 0:
-                mul_base *= attn_cp_size
-            filtered = [n for n in self.capture_num_tokens if n % mul_base == 0]
-            assert (
-                len(filtered) > 0
-            ), f"No piecewise CUDA graph capture sizes are multiples of {mul_base}"
-            self.capture_num_tokens = filtered
-        log_info_on_rank0(
-            logger, f"Capture cuda graph num tokens {self.capture_num_tokens}"
-        )
-        self.capture_forward_mode = ForwardMode.EXTEND
-        self.capture_hidden_mode = CaptureHiddenMode.NULL
-
-        # If returning hidden states is enabled, or if speculative prefill needs
-        # aux hidden states (DFLASH), capture the FULL variant up front.
-        if (
-            model_runner.server_args.enable_return_hidden_states
-            or model_runner.spec_algorithm.is_dflash()
-        ):
-            self.capture_hidden_mode = CaptureHiddenMode.FULL
-
-        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
-
-        self.is_multimodal = model_runner.is_multimodal
-        self.mamba_track_enabled = self.is_mamba_track_enabled()
-        # Classification/reward forwards branch on return_pooled_hidden_states; piecewise
-        # CUDA graph capture must use the same flag value as replay for those models.
-        self.capture_return_pooled_hidden_states = not model_runner.is_generation
-
-        with torch.device(self.device):
-            self.tbo_plugin = TboCudaGraphRunnerPlugin()
-
-        # Registry owns (allocates + pools) the token-axis input buffers.
-        self.buffer_registry = build_prefill_registry(
-            device=self.device,
-            max_bs=self.max_bs,
-            max_num_token=self.max_num_tokens,
-            cache_loc_dtype=self._cache_loc_dtype(),
-            is_multimodal=self.is_multimodal,
-            hidden_size=self.model_runner.model_config.hidden_size,
-            embed_dtype=self.model_runner.dtype,
-            enable_mamba_track=self.mamba_track_enabled,
-            share_pool=not is_npu(),
-            source=None,
-        )
-
-        self.attention_layers = self.model_runner.attention_layers
-        self.moe_layers = self.model_runner.moe_layers
-        self.moe_fusions = self.model_runner.moe_fusions
-        self.dsa_indexers = getattr(self.model_runner, "dsa_indexers", None)
-
-        if get_global_graph_memory_pool() is None:
-            set_global_graph_memory_pool(self.device_module.graph_pool_handle())
-        # Set graph pool id globally to be able to use symmetric memory
-        set_graph_pool_id(get_global_graph_memory_pool())
-
-        with enable_piecewise_cuda_graph():
-            language_model = getattr(
-                self.model_runner.model, "language_model", self.model_runner.model
-            )
-            layer_model = (
-                language_model.model
-                if hasattr(language_model, "model")
-                and hasattr(language_model.model, "layers")
-                else language_model
-            )
-            with patch_model(
-                layer_model, self.compile_config.compiler
-            ) as patched_model:
-
-                # Dummy warmup for jit kernel
-                self.warmup_compile(num_tokens=self.capture_num_tokens[0])
-
-                install_torch_compiled(
-                    patched_model,
-                    fullgraph=True,
-                    dynamic_arg_dims=None,
-                    compile_config=self.compile_config,
-                    graph_pool=get_global_graph_memory_pool(),
-                )
-
-                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).
-                    with enable_piecewise_cuda_graph_compile():
-                        self.warmup_compile(num_tokens=self.capture_num_tokens[-1])
-                else:
-                    with enable_piecewise_cuda_graph_compile():
-                        compile_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 enumerate(compile_range):
-                            if get_tensor_model_parallel_rank() == 0:
-                                compile_range.set_description(
-                                    f"Compiling num tokens ({num_tokens=})"
-                                )
-                            self.warmup_compile(num_tokens=num_tokens)
-
-                set_global_graph_memory_pool(self.device_module.graph_pool_handle())
-                set_graph_pool_id(get_global_graph_memory_pool())
-
-                if _use_aiter:
-                    self._pre_warm_aiter_chip_info()
-
-                self.device_module.synchronize()
-                self.model_runner.tp_group.barrier()
-                # Capture
-                self.capture()
-
-        self.raw_num_tokens = 0
-
-    _aiter_chip_info_cached = False
-
-    @classmethod
-    def _pre_warm_aiter_chip_info(cls):
-        """Pre-populate aiter chip info env vars before CUDA graph capture.
-
-        aiter's get_cu_num_custom_op and get_gfx_custom_op call
-        subprocess.run(rocminfo) to query GPU info. During CUDA graph capture
-        the GPU context is locked, so rocminfo hangs indefinitely. Pre-calling
-        them here caches the results as environment variables so the subprocess
-        is never invoked during capture. Only runs once per process.
-        """
-        if cls._aiter_chip_info_cached:
-            return
-        cls._aiter_chip_info_cached = True
-
-        import os
-
-        try:
-            from aiter.jit.utils.chip_info import get_cu_num, get_gfx
-
-            if not os.environ.get("CU_NUM"):
-                cu_num = get_cu_num()
-                os.environ["CU_NUM"] = str(cu_num)
-                logger.info(f"Pre-warmed aiter CU_NUM={cu_num}")
-
-            if not os.environ.get("GPU_ARCHS"):
-                gfx = get_gfx()
-                os.environ["GPU_ARCHS"] = gfx
-                logger.info(f"Pre-warmed aiter GPU_ARCHS={gfx}")
-        except ImportError:
-            pass
-        except Exception as e:
-            logger.warning(f"Failed to pre-warm aiter chip info: {e}")
-
-    def warmup_compile(self, num_tokens: int):
-        """Warmup the model with a simple forward pass before CUDA graph capture."""
-        registry = self.buffer_registry
-        bs = 1
-
-        def _slot(name):
-            return registry.get_slot(name).slice_for(bs, num_tokens)
-
-        input_ids = _slot("input_ids")
-        positions = _slot("positions")
-        out_cache_loc = _slot("out_cache_loc")
-        input_embeds = (
-            _slot("input_embeds") if registry.has_slot("input_embeds") else None
-        )
-        mrope_positions = (
-            _slot("mrope_positions") if registry.has_slot("mrope_positions") else None
-        )
-        mamba_track_indices = (
-            _slot("mamba_track_indices")
-            if registry.has_slot("mamba_track_indices")
-            else None
-        )
-        mamba_track_mask = (
-            _slot("mamba_track_mask") if registry.has_slot("mamba_track_mask") else None
-        )
-        mamba_track_seqlens = (
-            _slot("mamba_track_seqlens")
-            if registry.has_slot("mamba_track_seqlens")
-            else None
-        )
-        with torch.device(self.device):
-            forward_batch = ForwardBatch(
-                forward_mode=ForwardMode.EXTEND,
-                batch_size=1,
-                input_ids=input_ids,
-                input_embeds=input_embeds,
-                req_pool_indices=torch.arange(1, device=self.device),
-                seq_lens=torch.tensor([num_tokens], device=self.device),
-                next_token_logits_buffer=None,
-                orig_seq_lens=torch.tensor([num_tokens], device=self.device),
-                seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
-                out_cache_loc=out_cache_loc,
-                seq_lens_sum=num_tokens,
-                mamba_track_indices=mamba_track_indices,
-                mamba_track_mask=mamba_track_mask,
-                mamba_track_seqlens=mamba_track_seqlens,
-                encoder_lens=None,
-                return_logprob=False,
-                extend_num_tokens=num_tokens,
-                extend_seq_lens=torch.tensor([num_tokens], device=self.device),
-                extend_prefix_lens=torch.tensor([0], device=self.device),
-                extend_start_loc=torch.tensor([0], device=self.device),
-                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=positions,
-                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=mrope_positions,
-                spec_algorithm=None,
-                spec_info=None,
-                capture_hidden_mode=self.capture_hidden_mode,
-                num_token_non_padded=None,
-                num_token_non_padded_cpu=num_tokens,
-                global_forward_mode=ForwardMode.EXTEND,
-                lora_ids=None,
-                return_pooled_hidden_states=self.capture_return_pooled_hidden_states,
-            )
-
-        self.model_runner.attn_backend.init_forward_metadata(forward_batch)
-        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 forward_context(
-            ForwardContext(attn_backend=self.model_runner.attn_backend)
-        ):
-            with set_forward_context(
-                forward_batch,
-                self.attention_layers,
-                self.quant_config,
-                self.moe_layers,
-                self.moe_fusions,
-                dsa_indexers=self.dsa_indexers,
-            ):
-                _ = self.model_runner.model.forward(
-                    forward_batch.input_ids,
-                    forward_batch.positions,
-                    forward_batch,
-                )
-
-    def _cache_loc_dtype(self):
-        return torch.int64 if not is_npu() else torch.int32
-
-    def can_run(self, forward_batch: ForwardBatch):
-        # Disable piecewise cuda graph for input embeddings
-        # TODO(yuwei): fix it
-        if forward_batch.input_embeds is not None:
-            return False
-        # PCG graphs are captured with ForwardMode.EXTEND and spec_info=None.
-        # TARGET_VERIFY has different spec_info and capture_hidden_mode,
-        # so it must not use PCG-captured graphs.
-        if forward_batch.forward_mode.is_target_verify():
-            return False
-        # PCG graphs are captured with the runner's capture_hidden_mode.
-        # If the batch needs a different mode (e.g. FULL for speculative
-        # decoding), PCG replay would return wrong/missing hidden_states.
-        if forward_batch.capture_hidden_mode != self.capture_hidden_mode:
-            return False
-        # Disable for token embedding overrides (dynamic per-request)
-        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
-        if num_tokens <= self.max_num_tokens:
-            return True
-        return False
-
-    def capture(self) -> None:
-        # Trigger CUDA graph capture for specific shapes.
-        # Capture the large shapes first so that the smaller shapes
-        # can reuse the memory pool allocated for the large shapes.
-        with (
-            freeze_gc(self.model_runner.server_args.enable_cudagraph_gc),
-            graph_capture() as graph_capture_context,
-        ):
-            stream = graph_capture_context.stream
-            with set_pcg_capture_stream(stream):
-                avail_mem = get_available_gpu_memory(
-                    self.model_runner.device,
-                    self.model_runner.gpu_id,
-                    empty_cache=False,
-                )
-                # Reverse the order to enable better memory sharing across cuda graphs.
-                capture_range = (
-                    tqdm.tqdm(list(reversed(self.capture_num_tokens)))
-                    if get_tensor_model_parallel_rank() == 0
-                    else reversed(self.capture_num_tokens)
-                )
-                for i, num_tokens in enumerate(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"Capturing num tokens ({num_tokens=} {avail_mem=:.2f} GB)"
-                        )
-
-                    self.capture_one_batch_size(num_tokens)
-
-    def capture_one_batch_size(self, num_tokens: int):
-        registry = self.buffer_registry
-        bs = 1
-
-        # Graph inputs — views into the registry's (adopted) graph-resident
-        # slots; capture burns these addresses into the graph.
-        def _slot(name):
-            return registry.get_slot(name).slice_for(bs, num_tokens)
-
-        input_ids = _slot("input_ids")
-        positions = _slot("positions")
-        out_cache_loc = _slot("out_cache_loc")
-        input_embeds = (
-            _slot("input_embeds") if registry.has_slot("input_embeds") else None
-        )
-        mrope_positions = (
-            _slot("mrope_positions") if registry.has_slot("mrope_positions") else None
-        )
-        mamba_track_indices = (
-            _slot("mamba_track_indices")
-            if registry.has_slot("mamba_track_indices")
-            else None
-        )
-        mamba_track_mask = (
-            _slot("mamba_track_mask") if registry.has_slot("mamba_track_mask") else None
-        )
-        mamba_track_seqlens = (
-            _slot("mamba_track_seqlens")
-            if registry.has_slot("mamba_track_seqlens")
-            else None
-        )
-
-        global_dp_buffer_len = None
-        global_num_tokens_cpu = None
-
-        if self.model_runner.server_args.enable_lora:
-            # It is safe to capture CUDA graph using empty LoRA id, as the LoRA kernels will always be launched whenever
-            # `--enable-lora` is set to True (and return immediately if the LoRA id is empty for perf optimization).
-            lora_ids = [None] * bs
-        else:
-            lora_ids = None
-
-        with torch.device(self.device):
-            forward_batch = ForwardBatch(
-                forward_mode=ForwardMode.EXTEND,
-                batch_size=bs,
-                input_ids=input_ids,
-                input_embeds=input_embeds,
-                req_pool_indices=torch.arange(bs, device=self.device),
-                seq_lens=torch.tensor([num_tokens], device=self.device),
-                next_token_logits_buffer=None,
-                orig_seq_lens=torch.tensor([num_tokens], device=self.device),
-                seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
-                out_cache_loc=out_cache_loc,
-                seq_lens_sum=num_tokens,
-                mamba_track_indices=mamba_track_indices,
-                mamba_track_mask=mamba_track_mask,
-                mamba_track_seqlens=mamba_track_seqlens,
-                encoder_lens=None,
-                return_logprob=False,
-                extend_num_tokens=num_tokens,
-                extend_seq_lens=torch.tensor([num_tokens], device=self.device),
-                extend_prefix_lens=torch.tensor([0], device=self.device),
-                extend_start_loc=torch.tensor([0], device=self.device),
-                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=positions,
-                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=mrope_positions,
-                spec_algorithm=None,
-                spec_info=None,
-                capture_hidden_mode=self.capture_hidden_mode,
-                num_token_non_padded=None,
-                num_token_non_padded_cpu=num_tokens,
-                global_forward_mode=ForwardMode.EXTEND,
-                lora_ids=None,
-                return_pooled_hidden_states=self.capture_return_pooled_hidden_states,
-            )
-        # Setup hooks below read get_attn_backend() and must run inside the
-        # same ForwardContext as the warmup/capture forward.
-        with forward_context(
-            ForwardContext(attn_backend=self.model_runner.attn_backend)
-        ):
-            self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
-
-            if lora_ids is not None:
-                self.model_runner.lora_manager.prepare_lora_batch(forward_batch)
-
-            self.model_runner.attn_backend.init_forward_metadata(forward_batch)
-
-            # Run and capture
-            def run_once():
-                # Clean intermediate result cache for DP attention
-                forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = (
-                    None
-                )
-                set_dp_buffer_len(
-                    global_dp_buffer_len,
-                    num_tokens,
-                    forward_batch.dp_padding_mode.is_max_len(),
-                    global_num_tokens_cpu,
-                )
-                # FIXME: the implementation is hacky. `is_extend_in_batch`` is for determining the deepep mode.
-                # It is True in this context but we need to set it to use low latency deepep mode.
-                set_is_extend_in_batch(False)
-
-                kwargs = {}
-                with set_forward_context(
-                    forward_batch,
-                    self.attention_layers,
-                    self.quant_config,
-                    self.moe_layers,
-                    self.moe_fusions,
-                    dsa_indexers=self.dsa_indexers,
-                ):
-                    self.model_runner.model.forward(
-                        forward_batch.input_ids,
-                        forward_batch.positions,
-                        forward_batch,
-                        **kwargs,
-                    )
-                return
-
-            # run twice for warmup at the first time and cuda graph capture at the second time
-            # detail lies in sglang/python/sglang/srt/compilation/cuda_piecewise_backend.py
-            for _ in range(2):
-                self.device_module.synchronize()
-                self.model_runner.tp_group.barrier()
-                run_once()
-
-        return
-
-    def replay_prepare(
-        self,
-        forward_batch: ForwardBatch,
-        **kwargs,
-    ):
-        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]
-        self.raw_num_tokens = num_tokens
-        bs = forward_batch.batch_size
-        registry = self.buffer_registry
-        # Reset the padded token tail (ZERO) + copy the [:num_tokens] head for
-        # every graph-resident slot in one grouped pass. input_embeds is
-        # reset-only (the model writes embeds into it inside the graph).
-        registry.fill_from(
-            forward_batch,
-            raw_bs=bs,
-            padded_bs=bs,
-            raw_num_tokens=num_tokens,
-            padded_num_tokens=static_num_tokens,
-        )
-
-        def _slot(name):
-            return registry.get_slot(name).slice_for(bs, static_num_tokens)
-
-        input_ids = _slot("input_ids")
-        positions = _slot("positions")
-        out_cache_loc = _slot("out_cache_loc")
-        mamba_track_indices = (
-            _slot("mamba_track_indices")
-            if registry.has_slot("mamba_track_indices")
-            else None
-        )
-        mamba_track_mask = (
-            _slot("mamba_track_mask") if registry.has_slot("mamba_track_mask") else None
-        )
-        mamba_track_seqlens = (
-            _slot("mamba_track_seqlens")
-            if registry.has_slot("mamba_track_seqlens")
-            else None
-        )
-        input_embeds = (
-            _slot("input_embeds") if registry.has_slot("input_embeds") else None
-        )
-        mrope_positions = (
-            _slot("mrope_positions")
-            if (
-                registry.has_slot("mrope_positions")
-                and forward_batch.mrope_positions is not None
-            )
-            else None
-        )
-
-        next_token_logits_buffer = None
-
-        # Normalize MIXED→EXTEND so dynamo's guard (captured with EXTEND=1) doesn't fail on MIXED=3.
-        pcg_forward_mode = (
-            ForwardMode.EXTEND
-            if forward_batch.forward_mode == ForwardMode.MIXED
-            else forward_batch.forward_mode
-        )
-        pcg_global_forward_mode = (
-            ForwardMode.EXTEND
-            if forward_batch.global_forward_mode == ForwardMode.MIXED
-            else forward_batch.global_forward_mode
-        )
-
-        static_forward_batch = ForwardBatch(
-            forward_mode=pcg_forward_mode,
-            batch_size=bs,
-            input_ids=input_ids,
-            input_embeds=input_embeds,
-            req_pool_indices=forward_batch.req_pool_indices,
-            seq_lens=forward_batch.seq_lens,
-            next_token_logits_buffer=next_token_logits_buffer,
-            orig_seq_lens=forward_batch.orig_seq_lens,
-            seq_lens_cpu=forward_batch.seq_lens_cpu,
-            out_cache_loc=out_cache_loc,
-            seq_lens_sum=forward_batch.seq_lens_sum,
-            mamba_track_indices=mamba_track_indices,
-            mamba_track_mask=mamba_track_mask,
-            mamba_track_seqlens=mamba_track_seqlens,
-            encoder_lens=forward_batch.encoder_lens,
-            return_logprob=False,
-            extend_seq_lens=forward_batch.extend_seq_lens,
-            extend_prefix_lens=forward_batch.extend_prefix_lens,
-            extend_start_loc=forward_batch.extend_start_loc,
-            extend_prefix_lens_cpu=forward_batch.extend_prefix_lens_cpu,
-            extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
-            extend_logprob_start_lens_cpu=forward_batch.extend_logprob_start_lens_cpu,
-            extend_num_tokens=forward_batch.extend_num_tokens,
-            extend_input_logprob_token_ids_gpu=forward_batch.extend_input_logprob_token_ids_gpu,
-            positions=positions,
-            global_num_tokens_gpu=forward_batch.global_num_tokens_gpu,
-            global_num_tokens_for_logprob_gpu=forward_batch.global_num_tokens_for_logprob_gpu,
-            dp_padding_mode=forward_batch.dp_padding_mode,
-            global_dp_buffer_len=forward_batch.global_dp_buffer_len,
-            mrope_positions=mrope_positions,
-            spec_algorithm=forward_batch.spec_algorithm,
-            spec_info=forward_batch.spec_info,
-            capture_hidden_mode=forward_batch.capture_hidden_mode,
-            num_token_non_padded=forward_batch.num_token_non_padded,
-            num_token_non_padded_cpu=forward_batch.num_token_non_padded_cpu,
-            global_forward_mode=pcg_global_forward_mode,
-            lora_ids=forward_batch.lora_ids,
-            sampling_info=forward_batch.sampling_info,
-            mm_inputs=forward_batch.mm_inputs,
-            temperature=forward_batch.temperature,
-            top_p=forward_batch.top_p,
-            dimensions=forward_batch.dimensions,
-            return_pooled_hidden_states=(
-                self.capture_return_pooled_hidden_states
-                or forward_batch.return_pooled_hidden_states
-            ),
-        )
-
-        return static_forward_batch
-
-    def replay(
-        self,
-        forward_batch: ForwardBatch,
-        **kwargs,
-    ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
-        with enable_piecewise_cuda_graph():
-            static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
-            static_num_tokens = len(static_forward_batch.input_ids)
-            raw_num_tokens = self.raw_num_tokens
-            # Replay
-            with set_forward_context(
-                static_forward_batch,
-                self.attention_layers,
-                self.quant_config,
-                self.moe_layers,
-                self.moe_fusions,
-                dsa_indexers=self.dsa_indexers,
-                num_tokens=static_num_tokens,
-                raw_num_tokens=raw_num_tokens,
-            ):
-                self.model_runner.attn_backend.init_forward_metadata(forward_batch)
-                output = self.model_runner.model.forward(
-                    static_forward_batch.input_ids,
-                    static_forward_batch.positions,
-                    static_forward_batch,
-                    **kwargs,
-                )
-                if isinstance(output, LogitsProcessorOutput):
-                    # Preserve mm_input_embeds when speculative decoding is
-                    # enabled. The speculative draft's prefill path
-                    # (eagle_worker_v2._draft_extend_for_prefill) reads
-                    # mm_input_embeds off this LogitsProcessorOutput to reuse
-                    # the target's encoder embeddings instead of re-embedding
-                    # multimodal placeholder token ids.
-                    mm_input_embeds = None
-                    if (
-                        self.model_runner.spec_algorithm.is_speculative()
-                        and output.mm_input_embeds is not None
-                    ):
-                        mm_input_embeds = output.mm_input_embeds[: self.raw_num_tokens]
-                    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
-                        ),
-                        mm_input_embeds=mm_input_embeds,
-                    )
-                elif isinstance(output, EmbeddingPoolerOutput):
-                    return output
-                else:
-                    assert isinstance(output, PPProxyTensors)
-                    # TODO(Yuwei): support PP Support
-                    raise NotImplementedError(
-                        "PPProxyTensors is not supported in PiecewiseCudaGraphRunner yet."
-                    )
-
-    def get_spec_info(self, num_tokens: int):
-        spec_info = None
-        if (
-            self.model_runner.spec_algorithm.is_eagle()
-            or self.model_runner.spec_algorithm.is_standalone()
-        ):
-            from sglang.srt.speculative.eagle_utils import EagleVerifyInput
-
-            if self.model_runner.is_draft_worker:
-                raise RuntimeError("This should not happen.")
-            else:
-                spec_info = EagleVerifyInput(
-                    draft_token=None,
-                    custom_mask=self.custom_mask,
-                    positions=None,
-                    retrieve_index=None,
-                    retrieve_next_token=None,
-                    retrieve_next_sibling=None,
-                    retrieve_cum_len=None,
-                    spec_steps=self.model_runner.server_args.speculative_num_steps,
-                    topk=self.model_runner.server_args.speculative_eagle_topk,
-                    draft_token_num=self.model_runner.server_args.speculative_num_draft_tokens,
-                    capture_hidden_mode=CaptureHiddenMode.FULL,
-                    seq_lens_sum=None,
-                    seq_lens_cpu=None,
-                )
-
-        return spec_info

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

@@ -0,0 +1,49 @@
+"""Phase-aware CUDA graph runners.
+
+One concrete runner per phase. Each runner owns its phase-specific
+shape semantics (decode → batch size; prefill → token count) and
+delegates capture/replay mechanics to a pluggable
+BaseCudaGraphBackend chosen via cuda_graph_config.
+
+Public API:
+  - BaseCudaGraphRunner — abstract base; shared init + bucket
+    padding + capture-loop scaffolding.
+  - DecodeCudaGraphRunner — concrete decode-phase runner.
+  - PrefillCudaGraphRunner — concrete prefill-phase runner.
+  - Buffer dataclasses, capture-mode flags, the global memory pool,
+    and the DeepEP adapter live in
+    sglang.srt.model_executor.runner_utils; they are
+    re-exported here for the EAGLE / multi-step draft cuda graph
+    runners that were authored against the legacy public surface.
+"""
+
+from sglang.srt.model_executor.runner.base_cuda_graph_runner import (  # noqa: F401
+    BaseCudaGraphRunner,
+    freeze_gc,
+    get_batch_sizes_to_capture,
+)
+from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
+    DecodeCudaGraphRunner,
+)
+from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (  # noqa: F401
+    _make_graph_key as _default_make_graph_key,
+)
+from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (  # noqa: F401
+    PrefillCudaGraphRunner,
+)
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (  # noqa: F401
+    TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG,
+)
+from sglang.srt.model_executor.runner_utils import (  # noqa: F401
+    DecodeInputBuffers,
+    DeepEPCudaGraphRunnerAdapter,
+    PrefillInputBuffers,
+    _grouped_foreach_copy_,
+    _set_capture_lora_variant,
+    compile_in_capture_mode,
+    get_capture_lora_variant,
+    get_global_graph_memory_pool,
+    get_is_capture_mode,
+    model_capture_mode,
+    set_global_graph_memory_pool,
+)

python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py

@@ -0,0 +1,177 @@
+"""Shared scaffolding for the prefill and decode CUDA graph runners."""
+
+from __future__ import annotations
+
+import bisect
+import gc
+import logging
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, List, Sequence, Tuple
+
+import torch
+
+from sglang.srt.batch_overlap.two_batch_overlap import TboCudaGraphRunnerPlugin
+from sglang.srt.layers.dp_attention import (
+    get_attention_cp_size,
+    get_attention_tp_rank,
+    get_attention_tp_size,
+)
+from sglang.srt.utils import require_gathered_buffer
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
+        BaseCudaGraphBackend,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def freeze_gc(enable_cudagraph_gc: bool):
+    """Optimize garbage collection during CUDA graph capture.
+
+    Clean up first, then freeze remaining objects from being included in
+    future collections if GC is disabled during capture.
+    """
+    gc.collect()
+    should_freeze = not enable_cudagraph_gc
+    if should_freeze:
+        gc.freeze()
+    try:
+        yield
+    finally:
+        if should_freeze:
+            gc.unfreeze()
+            gc.collect()
+
+
+def get_batch_sizes_to_capture(
+    model_runner: ModelRunner, num_tokens_per_bs: int = 1
+) -> Tuple[List[int], List[int]]:
+    """Build the (capture_bs, compile_bs) lists for the decode runner.
+
+    Filters cuda_graph_config[decode].bs by attention-tp/cp alignment
+    constraints and clamps to req_to_token_pool.size.
+    """
+
+    server_args = model_runner.server_args
+    capture_bs = server_args.cuda_graph_config.decode.bs
+    num_max_requests = model_runner.req_to_token_pool.size
+
+    mul_base = 1
+    if server_args.enable_two_batch_overlap:
+        mul_base *= 2
+        num_tokens_per_bs = 1
+
+    if require_gathered_buffer(server_args):
+        mul_base *= get_attention_tp_size()
+
+    if mul_base % get_attention_cp_size() != 0:
+        mul_base *= get_attention_cp_size()
+
+    num_max_requests = (num_max_requests + mul_base - 1) // mul_base * mul_base
+    if max(capture_bs) > num_max_requests:
+        capture_bs += [num_max_requests]
+
+    capture_bs = [bs for bs in capture_bs if bs * num_tokens_per_bs % mul_base == 0]
+    capture_bs = [bs for bs in capture_bs if bs <= num_max_requests]
+    capture_bs = list(sorted(set(capture_bs)))
+
+    assert len(capture_bs) > 0 and capture_bs[0] > 0, f"{capture_bs=}"
+    compile_bs = (
+        [bs for bs in capture_bs if bs <= server_args.torch_compile_max_bs]
+        if server_args.enable_torch_compile
+        else []
+    )
+    return capture_bs, compile_bs
+
+
+class BaseCudaGraphRunner(ABC):
+    """Abstract base for phase-specific cuda-graph runners.
+
+    A subclass (DecodeCudaGraphRunner / PrefillCudaGraphRunner) owns one
+    phase and plugs in a BaseCudaGraphBackend that handles the
+    capture / replay mechanics. The runner orchestrates bucket
+    selection, static buffer population, attention metadata init,
+    replay dispatch, and output slicing.
+
+    Methods:
+      - can_run(forward_batch) — should forward_batch go through cuda
+        graph replay (vs eager fallback)?
+      - capture_prepare(size, ...) — build the dummy ForwardBatch and
+        per-capture local state needed by capture_one_shape.
+      - capture() — outer capture loop; iterates over shapes and calls
+        capture_one_shape for each.
+      - capture_one_shape(size, ...) — drive one model forward at this
+        shape into the backend's captured artifact.
+      - replay_prepare(forward_batch, ...) — pad to the nearest captured
+        bucket, populate static input buffers, init attention metadata.
+      - replay(forward_batch, ...) — dispatch one batch through cuda
+        graph replay.
+
+    Notes:
+      - buffers and backend are populated by the subclass before
+        capture(); the base only declares them.
+    """
+
+    # Subclasses populate before calling capture().
+    buffers: ForwardInputBuffers
+    backend: BaseCudaGraphBackend
+
+    def __init__(self, model_runner: ModelRunner) -> None:
+        self.model_runner = model_runner
+        self.device = model_runner.device
+        self.device_module = torch.get_device_module(self.device)
+        self.tp_size = model_runner.server_args.tp_size
+        self.dp_size = model_runner.server_args.dp_size
+        self.pp_size = model_runner.server_args.pp_size
+        self.attn_tp_size = get_attention_tp_size()
+        self.attn_tp_rank = get_attention_tp_rank()
+        self.tbo_plugin = TboCudaGraphRunnerPlugin()
+
+    @staticmethod
+    def _pad_to_bucket(raw_size: int, buckets: Sequence[int]) -> int:
+        """Return the smallest buckets[i] >= raw_size.
+
+        Caller's can_run must reject raw_size > max(buckets) before
+        reaching replay_prepare; this assertion makes the contract
+        explicit (bisect_left returns len(buckets) when the value
+        exceeds all buckets, which would otherwise IndexError below
+        with no diagnostic).
+        """
+        assert raw_size <= buckets[-1], (
+            f"size {raw_size} exceeds max captured bucket {buckets[-1]}; "
+            f"can_run should have rejected this batch"
+        )
+        index = bisect.bisect_left(buckets, raw_size)
+        return buckets[index]
+
+    @abstractmethod
+    def can_run(self, forward_batch: ForwardBatch) -> bool: ...
+
+    @abstractmethod
+    def capture_prepare(self, size: int, *args, **kwargs) -> Any: ...
+
+    @abstractmethod
+    def capture(self) -> None: ...
+
+    @abstractmethod
+    def capture_one_shape(self, size: int, *args, **kwargs) -> Any: ...
+
+    @abstractmethod
+    def replay_prepare(
+        self,
+        forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any: ...
+
+    @abstractmethod
+    def replay(
+        self,
+        forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any: ...

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -0,0 +1,846 @@
+# 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.
+# ==============================================================================
+"""PrefillCudaGraphRunner — runs the EXTEND phase under a pluggable backend.
+
+Backend selection comes from cuda_graph_config.prefill:
+  - "tc_piecewise"     — default, TcPiecewiseCudaGraphBackend: torch.compile
+                      wraps the model; per-shape graphs live in
+                      torch.compile's internal cache. Multi-batch supported.
+  - "breakable" — BreakableCudaGraphBackend: segmented capture (no
+                      torch.compile). Captures with bs=1; rejects multi-req
+                      prefill in can_run.
+  - "full"      — rejected at config validation; not supported for prefill.
+  - "disabled"  — handled at the model_runner level — runner not
+                      constructed.
+"""
+
+from __future__ import annotations
+
+import logging
+import warnings
+from typing import TYPE_CHECKING, Dict, Optional, Union
+
+import torch
+import tqdm
+
+from sglang.srt.distributed import get_tensor_model_parallel_rank
+from sglang.srt.distributed.parallel_state import graph_capture
+from sglang.srt.layers.dp_attention import (
+    DpPaddingMode,
+    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.cuda_graph_buffer_registry import (
+    CudaGraphBufferRegistry,
+    build_prefill_registry,
+)
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+    PPProxyTensors,
+)
+from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
+from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
+    BaseCudaGraphRunner,
+    freeze_gc,
+)
+from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import (
+    BreakableCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.utils import (
+    resolve_prefill_backend,
+)
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    set_tc_piecewise_forward_context,
+)
+from sglang.srt.model_executor.runner_utils.buffers import (
+    PrefillInputBuffers,
+)
+from sglang.srt.utils import (
+    get_available_gpu_memory,
+    is_npu,
+    log_info_on_rank0,
+    require_attn_tp_gather,
+    require_mlp_tp_gather,
+)
+
+# Suppress Dynamo warning about tracing through lru_cache-wrapped functions.
+warnings.filterwarnings("ignore", message=".*lru_cache.*", module="torch._dynamo")
+logger = logging.getLogger(__name__)
+
+# Names of the static prefill input tensors a Breakable-backed prefill
+# runner owns. Each is a 1-D int64 tensor of length max_bs; captured
+# Breakable segments read from these stable addresses.
+_PREFILL_STATIC_FIELDS = (
+    "seq_lens",
+    "extend_seq_lens",
+    "extend_prefix_lens",
+    "extend_start_loc",
+    "req_pool_indices",
+    "orig_seq_lens",
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+
+class PrefillCudaGraphRunner(BaseCudaGraphRunner):
+    """Prefill-phase CUDA graph runner.
+
+    Owns: PrefillInputBuffers, capture-num-tokens list, attention layers
+    snapshot, and the pluggable self.backend. The backend handles capture
+    + replay mechanics; this runner handles dummy ForwardBatch construction,
+    buffer population, attention metadata init, and output slicing.
+    """
+
+    def __init__(self, model_runner: ModelRunner):
+        super().__init__(model_runner)
+        # --- core state ------------------------------------------------
+        self.quant_config = getattr(self.model_runner.model, "quant_config", None)
+        self.is_multimodal = model_runner.is_multimodal
+        # Classification/reward forwards branch on return_pooled_hidden_states;
+        # capture must use the same flag value as replay for those models.
+        self.capture_return_pooled_hidden_states = not model_runner.is_generation
+
+        # --- bucket sizes ---------------------------------------------
+        # bs in prefill carries the captured shape (token count for
+        # tc_piecewise) — one shape knob per phase.
+        capture_tokens = model_runner.server_args.cuda_graph_config.prefill.bs
+        assert capture_tokens is not None, "cuda_graph_config[prefill].bs is not set"
+        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"Capture cuda graph num tokens {self.capture_num_tokens}"
+        )
+
+        self.capture_forward_mode = ForwardMode.EXTEND
+        self.capture_hidden_mode = CaptureHiddenMode.NULL
+        # If returning hidden states is enabled, or if speculative prefill
+        # needs aux hidden states (DFLASH), capture the FULL variant up front.
+        # Ported from main #27468.
+        if (
+            model_runner.server_args.enable_return_hidden_states
+            or model_runner.spec_algorithm.is_dflash()
+        ):
+            self.capture_hidden_mode = CaptureHiddenMode.FULL
+
+        self.mamba_track_enabled = self._is_mamba_track_enabled()
+
+        # --- buffers ---------------------------------------------------
+        self.buffers: PrefillInputBuffers = PrefillInputBuffers.create(
+            device=self.device,
+            max_bs=self.max_bs,
+            max_num_tokens=self.max_num_tokens,
+            cache_loc_dtype=self._cache_loc_dtype(),
+            is_hybrid_swa=model_runner.is_hybrid_swa,
+            is_multimodal=self.is_multimodal,
+            hidden_size=self.model_runner.model_config.hidden_size,
+            dtype=self.model_runner.dtype,
+            enable_mamba_track=self.mamba_track_enabled,
+        )
+        self.buffers.share_buffers()
+        # Token-axis FB-shared slot registry adopting PrefillInputBuffers
+        # storage; same physical tensors, stable data_ptr for capture vs
+        # replay. Replaces populate_from_forward_batch on capture/replay paths.
+        self.buffer_registry: CudaGraphBufferRegistry = build_prefill_registry(
+            device=self.device,
+            max_bs=self.max_bs,
+            max_num_token=self.max_num_tokens,
+            cache_loc_dtype=self._cache_loc_dtype(),
+            is_multimodal=self.is_multimodal,
+            hidden_size=self.model_runner.model_config.hidden_size,
+            embed_dtype=self.model_runner.dtype,
+            enable_mamba_track=self.mamba_track_enabled,
+            source=self.buffers,
+        )
+
+        self.attention_layers = self.model_runner.attention_layers
+        self.moe_layers = self.model_runner.moe_layers
+        self.moe_fusions = self.model_runner.moe_fusions
+        self.dsa_indexers = getattr(self.model_runner, "dsa_indexers", None)
+
+        self.dp_size = model_runner.server_args.dp_size
+        self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args)
+        self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args)
+
+        # --- backend ---------------------------------------------------
+        # When the backend is Breakable, captured segments need stable
+        # tensor addresses, so we own a set of static int64 buffers here
+        # and rebind them into capture-time dummy inputs / replay-time
+        # serving inputs below. Other backends don't need this.
+        # Initialize the slot to None BEFORE constructing the backend:
+        # TcPiecewise runs its compile pass during __init__ which calls
+        # _run_dummy_forward -> capture_prepare, and capture_prepare reads
+        # self._prefill_static_buffers. self.layer_model has the same
+        # ordering requirement: _run_forward checks `self.layer_model is
+        # not None` to decide whether to call the inner stack or outer
+        # model.forward, and that check fires inside TcPiecewise's
+        # _run_compile_pass before backend resolution returns.
+        self._prefill_static_buffers: Optional[Dict[str, torch.Tensor]] = None
+        self.layer_model = None
+        self.backend = resolve_prefill_backend(self)
+        if isinstance(self.backend, BreakableCudaGraphBackend):
+            with torch.device(self.device):
+                self._prefill_static_buffers = {
+                    name: torch.zeros((self.max_bs,), dtype=torch.int64)
+                    for name in _PREFILL_STATIC_FIELDS
+                }
+
+        # Some attention backends (e.g. DSV4) opt into a captured-metadata
+        # contract under BCG: capture-time builds a per-bucket metadata
+        # object the backend then refreshes in place at replay. We honor
+        # the contract only when the backend is Breakable; FullCG and
+        # TC_PIECEWISE use the eager init_forward_metadata path.
+        if isinstance(self.backend, BreakableCudaGraphBackend):
+            self.use_captured_attn_metadata = (
+                model_runner.attn_backend.use_captured_forward_metadata_for_breakable_cuda_graph
+            )
+        else:
+            self.use_captured_attn_metadata = False
+        self.attn_metadata_buffers: Optional[Dict[int, object]] = (
+            {} if self.use_captured_attn_metadata else None
+        )
+
+        # --- BCG: resolve inner layer_model for capture/replay --------
+        # BCG captures only the inner transformer stack (layer_model.forward)
+        # — not the outer model.forward. The outer's tail (logits_processor /
+        # pooler) has bs-shaped kernels that would bake bs=1 into the captured
+        # graph and break multi-req replay. At replay, we monkey-patch
+        # layer_model.forward to replay the captured graph and return the
+        # captured hidden states; the outer model.forward then runs
+        # logits_processor eagerly on top with the live multi-req metadata.
+        # Mirrors main's BreakableCudaGraphRunner. (Slot pre-init lives
+        # above next to _prefill_static_buffers — TcPiecewise's compile
+        # pass runs during backend construction and reads self.layer_model.)
+        if isinstance(self.backend, BreakableCudaGraphBackend):
+            language_model = getattr(
+                self.model_runner.model, "language_model", self.model_runner.model
+            )
+            if hasattr(language_model, "model") and hasattr(
+                language_model.model, "layers"
+            ):
+                self.layer_model = language_model.model
+            else:
+                raise RuntimeError(
+                    f"BCG could not resolve inner layer_model on "
+                    f"{type(language_model).__name__}; BCG is unsupported for "
+                    f"this model architecture."
+                )
+
+        # --- capture --------------------------------------------------
+        self.device_module.synchronize()
+        self.model_runner.tp_group.barrier()
+        self.capture()
+
+        self.raw_num_tokens = 0
+
+    # -----------------------------------------------------------------
+    # Helpers
+    # -----------------------------------------------------------------
+    def _is_mamba_track_enabled(self) -> bool:
+        return (
+            self.model_runner.server_args.enable_mamba_extra_buffer()
+            and not self.model_runner.server_args.disable_radix_cache
+            and self.model_runner.spec_algorithm.is_none()
+        )
+
+    def _cache_loc_dtype(self):
+        return torch.int64 if not is_npu() else torch.int32
+
+    @torch.no_grad()
+    def _run_forward(self, forward_batch: ForwardBatch, num_tokens: int):
+        """Run forward inside the prefill set_tc_piecewise_forward_context.
+
+        BCG path: captures only the inner layer_model.forward (transformer
+        stack), excluding the outer model.forward tail (logits_processor /
+        pooler). The captured output is bs=1 hidden states; replay then runs
+        the outer tail eagerly with live multi-req metadata.
+
+        TC_PIECEWISE path: captures the outer model.forward; torch.compile
+        FX-traces produce bs-invariant kernels.
+
+        ``@torch.no_grad`` mirrors the decorator on the outer
+        ``*ForCausalLM.forward``. For BCG, calling ``layer_model.forward``
+        directly skips that decorator, so we apply it here — without it
+        some MoE ``@torch.compile`` kernels (``torch.sum(out=...)``) fail
+        dynamo with "out= doesn't support autograd", and mamba state ops
+        can spuriously track gradients.
+        """
+        forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
+        set_dp_buffer_len(
+            forward_batch.global_dp_buffer_len,
+            num_tokens,
+            forward_batch.dp_padding_mode.is_max_len(),
+            forward_batch.global_num_tokens_cpu,
+        )
+        set_is_extend_in_batch(False)
+
+        with forward_context(
+            ForwardContext(attn_backend=self.model_runner.attn_backend)
+        ), set_tc_piecewise_forward_context(
+            forward_batch,
+            self.attention_layers,
+            self.quant_config,
+            self.moe_layers,
+            self.moe_fusions,
+            dsa_indexers=self.dsa_indexers,
+        ):
+            if self.layer_model is not None:
+                return self.layer_model.forward(
+                    forward_batch.input_ids,
+                    forward_batch.positions,
+                    forward_batch,
+                    forward_batch.input_embeds,
+                )
+            return self.model_runner.model.forward(
+                forward_batch.input_ids,
+                forward_batch.positions,
+                forward_batch,
+            )
+
+    def _run_dummy_forward(self, num_tokens: int) -> None:
+        """Build a dummy ForwardBatch at this shape, init attn metadata,
+        run forward once. Used by TcPiecewiseCudaGraphBackend.prepare
+        for both the JIT-activate forward (single shape, before
+        torch.compile install) and the compile-loop pass (every shape,
+        inside enable_torch_compile_warmup).
+        """
+        fb, attn_backend = self.capture_prepare(num_tokens)
+        attn_backend.init_forward_metadata(fb)
+        self._run_forward(fb, num_tokens)
+
+    def _has_inactive_dp_rank(self, forward_batch: ForwardBatch) -> bool:
+        # DSV4 DP attention / DeepEP collectives need every DP rank to enter
+        # the same replay path. Sparse-DP batches (one or more ranks with
+        # zero local tokens) fall back to eager to avoid hanging ranks.
+        global_num_tokens = forward_batch.global_num_tokens_cpu
+        if global_num_tokens is None:
+            return False
+        return len(global_num_tokens) > 1 and any(
+            int(num_tokens) == 0 for num_tokens in global_num_tokens
+        )
+
+    def _init_forward_metadata_for_capture(
+        self, forward_batch: ForwardBatch, num_tokens: int
+    ) -> None:
+        """Capture-time metadata init for the BCG-with-captured-metadata
+        contract. For opt-in backends (DSV4), call the BCG-specific entry
+        and stash the returned per-bucket metadata object; otherwise fall
+        back to the generic eager init that BCG/TC_PIECEWISE use today."""
+        attn_backend = self.model_runner.attn_backend
+        if not self.use_captured_attn_metadata:
+            attn_backend.init_forward_metadata(forward_batch)
+            return
+        metadata = attn_backend.init_forward_metadata_for_breakable_cuda_graph_capture(
+            forward_batch
+        )
+        assert self.attn_metadata_buffers is not None
+        self.attn_metadata_buffers[num_tokens] = metadata
+
+    def _prepare_forward_metadata_for_replay(
+        self,
+        forward_batch: ForwardBatch,
+        static_forward_batch: ForwardBatch,
+        num_tokens: int,
+    ) -> None:
+        """Replay-time metadata refresh for the BCG-with-captured-metadata
+        contract. For opt-in backends, refresh the stashed per-bucket
+        metadata in place against the current batch; otherwise fall back
+        to the generic eager init."""
+        attn_backend = self.model_runner.attn_backend
+        if not self.use_captured_attn_metadata:
+            attn_backend.init_forward_metadata(forward_batch)
+            return
+        assert self.attn_metadata_buffers is not None
+        metadata = self.attn_metadata_buffers[num_tokens]
+        attn_backend.prepare_forward_metadata_for_breakable_cuda_graph_replay(
+            metadata,
+            forward_batch,
+            static_forward_batch=static_forward_batch,
+        )
+
+    # -----------------------------------------------------------------
+    # can_run
+    # -----------------------------------------------------------------
+    def can_run(self, forward_batch: ForwardBatch) -> bool:
+        if forward_batch.input_embeds is not None:
+            return False
+        if forward_batch.replace_embeds is not None:
+            return False
+        # tc_piecewise captures with ForwardMode.EXTEND and spec_info=None.
+        if forward_batch.forward_mode.is_target_verify():
+            return False
+        if forward_batch.capture_hidden_mode != self.capture_hidden_mode:
+            return False
+        # BCG-with-captured-metadata under DP attention: every rank must
+        # have local tokens, and the batch must declare itself replayable.
+        # These gates are no-ops for non-DP / non-opt-in paths because
+        # global_num_tokens_cpu stays None.
+        if self._has_inactive_dp_rank(forward_batch):
+            return False
+        if (
+            forward_batch.global_num_tokens_cpu is not None
+            and not forward_batch.can_run_dp_breakable_cuda_graph
+        ):
+            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
+        if num_tokens > self.max_num_tokens:
+            return False
+        # No backend-level shape check here: replay_prepare bucket-pads
+        # num_tokens up to the nearest captured shape, so eligibility is
+        # bounded by num_tokens <= self.max_num_tokens (already
+        # checked above), not by exact shape membership.
+        #
+        # Multi-req replay is supported by BCG via the layer_model.forward
+        # monkey-patch in replay(): the captured bs=1 graph runs the
+        # transformer stack, then the outer model.forward runs
+        # logits_processor eagerly on top with live multi-req metadata.
+        return True
+
+    # -----------------------------------------------------------------
+    # capture_prepare
+    # -----------------------------------------------------------------
+    def capture_prepare(
+        self, num_tokens: int
+    ) -> tuple[ForwardBatch, "AttentionBackend"]:
+        """Build a dummy prefill ForwardBatch for capture/warmup at this shape.
+
+        Default tensor inputs are fresh literals; under a Breakable
+        backend, we swap in slices of our static buffers so captured
+        segments read from stable addresses.
+
+        Returns ``(forward_batch, attn_backend)`` to mirror decode's
+        capture_prepare signature.
+        """
+        buffers = self.buffers
+        bs = 1
+
+        with torch.device(self.device):
+            shape_inputs = {
+                "req_pool_indices": torch.arange(bs, device=self.device),
+                "seq_lens": torch.tensor([num_tokens], device=self.device),
+                "orig_seq_lens": torch.tensor([num_tokens], device=self.device),
+                "extend_seq_lens": torch.tensor([num_tokens], device=self.device),
+                "extend_prefix_lens": torch.tensor([0], device=self.device),
+                "extend_start_loc": torch.tensor([0], device=self.device),
+            }
+        if self._prefill_static_buffers is not None:
+            s = self._prefill_static_buffers
+            s["seq_lens"][:bs].fill_(num_tokens)
+            s["extend_seq_lens"][:bs].fill_(num_tokens)
+            s["extend_prefix_lens"][:bs].zero_()
+            s["extend_start_loc"][:bs].zero_()
+            s["req_pool_indices"][:bs].copy_(
+                torch.arange(bs, device=s["req_pool_indices"].device)
+            )
+            s["orig_seq_lens"][:bs].fill_(num_tokens)
+            for name in _PREFILL_STATIC_FIELDS:
+                shape_inputs[name] = s[name][:bs]
+
+        registry = self.buffer_registry
+
+        def _slot(name):
+            return registry.get_slot(name).slice_for(bs, num_tokens)
+
+        if self.require_mlp_tp_gather:
+            global_num_tokens_cpu = [num_tokens] * self.dp_size
+        elif self.require_attn_tp_gather:
+            global_num_tokens_cpu = [num_tokens]
+        else:
+            global_num_tokens_cpu = None
+
+        if global_num_tokens_cpu is not None:
+            global_dp_buffer_len = sum(global_num_tokens_cpu)
+            num_tokens_tensor = torch.tensor(
+                global_num_tokens_cpu, dtype=torch.int32, device=self.device
+            )
+            global_num_tokens_gpu = num_tokens_tensor
+            global_num_tokens_for_logprob_gpu = num_tokens_tensor
+        else:
+            global_dp_buffer_len = None
+            global_num_tokens_gpu = None
+            global_num_tokens_for_logprob_gpu = None
+
+        with torch.device(self.device):
+            forward_batch = ForwardBatch(
+                forward_mode=ForwardMode.EXTEND,
+                batch_size=bs,
+                input_ids=_slot("input_ids"),
+                input_embeds=(
+                    _slot("input_embeds") if registry.has_slot("input_embeds") else None
+                ),
+                req_pool_indices=shape_inputs["req_pool_indices"],
+                seq_lens=shape_inputs["seq_lens"],
+                next_token_logits_buffer=None,
+                orig_seq_lens=shape_inputs["orig_seq_lens"],
+                seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
+                out_cache_loc=_slot("out_cache_loc"),
+                seq_lens_sum=num_tokens,
+                mamba_track_indices=(
+                    _slot("mamba_track_indices")
+                    if registry.has_slot("mamba_track_indices")
+                    else None
+                ),
+                mamba_track_mask=(
+                    _slot("mamba_track_mask")
+                    if registry.has_slot("mamba_track_mask")
+                    else None
+                ),
+                mamba_track_seqlens=(
+                    _slot("mamba_track_seqlens")
+                    if registry.has_slot("mamba_track_seqlens")
+                    else None
+                ),
+                encoder_lens=None,
+                return_logprob=False,
+                extend_num_tokens=num_tokens,
+                extend_seq_lens=shape_inputs["extend_seq_lens"],
+                extend_prefix_lens=shape_inputs["extend_prefix_lens"],
+                extend_start_loc=shape_inputs["extend_start_loc"],
+                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=_slot("positions"),
+                global_num_tokens_gpu=global_num_tokens_gpu,
+                global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
+                global_num_tokens_cpu=global_num_tokens_cpu,
+                dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
+                global_dp_buffer_len=global_dp_buffer_len,
+                mrope_positions=(
+                    _slot("mrope_positions")
+                    if registry.has_slot("mrope_positions")
+                    else None
+                ),
+                spec_algorithm=None,
+                spec_info=None,
+                # Use self.capture_hidden_mode so dflash spec (which needs
+                # FULL aux hidden states) captures with the right mode.
+                # Ported from main #27468.
+                capture_hidden_mode=self.capture_hidden_mode,
+                num_token_non_padded=None,
+                num_token_non_padded_cpu=num_tokens,
+                global_forward_mode=ForwardMode.EXTEND,
+                lora_ids=None,
+                return_pooled_hidden_states=self.capture_return_pooled_hidden_states,
+            )
+            self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
+        return forward_batch, self.model_runner.attn_backend
+
+    # -----------------------------------------------------------------
+    # capture
+    # -----------------------------------------------------------------
+    def capture(self) -> None:
+        with freeze_gc(self.model_runner.server_args.enable_cudagraph_gc):
+            with graph_capture() as graph_capture_context:
+                self.stream = graph_capture_context.stream
+                with self.backend.capture_session(self.stream):
+                    self._capture_one_stream()
+
+    def _capture_one_stream(self) -> None:
+        avail_mem = get_available_gpu_memory(
+            self.model_runner.device,
+            self.model_runner.gpu_id,
+            empty_cache=False,
+        )
+        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"Capturing num tokens ({num_tokens=} {avail_mem=:.2f} GB)"
+                )
+            self.capture_one_shape(num_tokens)
+
+    # -----------------------------------------------------------------
+    # capture_one_shape
+    # -----------------------------------------------------------------
+    def capture_one_shape(self, size: int) -> None:
+        """Per-shape capture: build dummy ForwardBatch + run_once,
+        delegate to backend. size is the prefill token count.
+        """
+        num_tokens = size
+        forward_batch, attn_backend = self.capture_prepare(num_tokens)
+        self._init_forward_metadata_for_capture(forward_batch, num_tokens)
+
+        def run_once():
+            return self._run_forward(forward_batch, num_tokens)
+
+        # Main's monolithic BCG runner never invokes
+        # on_after_cuda_graph_warmup between warmup iterations — the BCG
+        # contract is to keep warmup state untouched and let
+        # init_forward_metadata_in_graph (recorded inside the captured
+        # forward) do any raw->full upgrade. cg-refactor's runner_backend
+        # abstraction exposes a post_warmup_hook for backends that need
+        # workspace cleanup between iterations; suppress it for BCG so
+        # DSV4's hook (which restores forward_metadata to a stale
+        # _current_capture_raw left over from decode CG capture) doesn't
+        # corrupt warmup iter 2's metadata read.
+        if isinstance(self.backend, BreakableCudaGraphBackend):
+            post_warmup_hook = None
+        else:
+            post_warmup_hook = getattr(attn_backend, "on_after_cuda_graph_warmup", None)
+        self.backend.capture_one(
+            num_tokens,
+            run_once,
+            dummies=None,
+            post_warmup_hook=post_warmup_hook,
+        )
+
+    # -----------------------------------------------------------------
+    # replay_prepare
+    # -----------------------------------------------------------------
+    def replay_prepare(self, forward_batch: ForwardBatch, **kwargs) -> ForwardBatch:
+        """Pad, populate static buffers, and build the static_forward_batch
+        the model code reads during replay.
+        """
+        buffers = self.buffers
+        num_tokens = len(forward_batch.input_ids)
+        static_num_tokens = self._pad_to_bucket(num_tokens, self.capture_num_tokens)
+        self.raw_num_tokens = num_tokens
+
+        bs = forward_batch.batch_size
+
+        self.buffer_registry.fill_from(
+            forward_batch,
+            raw_bs=bs,
+            padded_bs=bs,
+            raw_num_tokens=num_tokens,
+            padded_num_tokens=static_num_tokens,
+        )
+
+        registry = self.buffer_registry
+
+        def _slot(name):
+            return registry.get_slot(name).slice_for(bs, static_num_tokens)
+
+        mamba_track_indices = (
+            _slot("mamba_track_indices")
+            if registry.has_slot("mamba_track_indices")
+            else None
+        )
+        mamba_track_mask = (
+            _slot("mamba_track_mask") if registry.has_slot("mamba_track_mask") else None
+        )
+        mamba_track_seqlens = (
+            _slot("mamba_track_seqlens")
+            if registry.has_slot("mamba_track_seqlens")
+            else None
+        )
+
+        input_ids = _slot("input_ids")
+        input_embeds = (
+            _slot("input_embeds") if registry.has_slot("input_embeds") else None
+        )
+        positions = _slot("positions")
+        out_cache_loc = _slot("out_cache_loc")
+        mrope_positions = (
+            _slot("mrope_positions")
+            if registry.has_slot("mrope_positions")
+            and forward_batch.mrope_positions is not None
+            else None
+        )
+
+        # Normalize MIXED→EXTEND so dynamo's guard (captured with EXTEND=1)
+        # doesn't fail on MIXED=3.
+        pcg_forward_mode = (
+            ForwardMode.EXTEND
+            if forward_batch.forward_mode == ForwardMode.MIXED
+            else forward_batch.forward_mode
+        )
+        pcg_global_forward_mode = (
+            ForwardMode.EXTEND
+            if forward_batch.global_forward_mode == ForwardMode.MIXED
+            else forward_batch.global_forward_mode
+        )
+
+        static_forward_batch = ForwardBatch(
+            forward_mode=pcg_forward_mode,
+            batch_size=bs,
+            input_ids=input_ids,
+            input_embeds=input_embeds,
+            req_pool_indices=forward_batch.req_pool_indices,
+            seq_lens=forward_batch.seq_lens,
+            next_token_logits_buffer=None,
+            orig_seq_lens=forward_batch.orig_seq_lens,
+            seq_lens_cpu=forward_batch.seq_lens_cpu,
+            out_cache_loc=out_cache_loc,
+            seq_lens_sum=forward_batch.seq_lens_sum,
+            mamba_track_indices=mamba_track_indices,
+            mamba_track_mask=mamba_track_mask,
+            mamba_track_seqlens=mamba_track_seqlens,
+            encoder_lens=forward_batch.encoder_lens,
+            return_logprob=False,
+            extend_seq_lens=forward_batch.extend_seq_lens,
+            extend_prefix_lens=forward_batch.extend_prefix_lens,
+            extend_start_loc=forward_batch.extend_start_loc,
+            extend_prefix_lens_cpu=forward_batch.extend_prefix_lens_cpu,
+            extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
+            extend_logprob_start_lens_cpu=forward_batch.extend_logprob_start_lens_cpu,
+            extend_num_tokens=forward_batch.extend_num_tokens,
+            extend_input_logprob_token_ids_gpu=forward_batch.extend_input_logprob_token_ids_gpu,
+            positions=positions,
+            global_num_tokens_gpu=forward_batch.global_num_tokens_gpu,
+            global_num_tokens_for_logprob_gpu=forward_batch.global_num_tokens_for_logprob_gpu,
+            dp_padding_mode=forward_batch.dp_padding_mode,
+            global_dp_buffer_len=forward_batch.global_dp_buffer_len,
+            mrope_positions=mrope_positions,
+            spec_algorithm=forward_batch.spec_algorithm,
+            spec_info=forward_batch.spec_info,
+            capture_hidden_mode=forward_batch.capture_hidden_mode,
+            num_token_non_padded=forward_batch.num_token_non_padded,
+            num_token_non_padded_cpu=forward_batch.num_token_non_padded_cpu,
+            global_forward_mode=pcg_global_forward_mode,
+            lora_ids=forward_batch.lora_ids,
+            sampling_info=forward_batch.sampling_info,
+            mm_inputs=forward_batch.mm_inputs,
+            temperature=forward_batch.temperature,
+            top_p=forward_batch.top_p,
+            dimensions=forward_batch.dimensions,
+            return_pooled_hidden_states=(
+                self.capture_return_pooled_hidden_states
+                or forward_batch.return_pooled_hidden_states
+            ),
+        )
+
+        # Under Breakable, copy serving-time values into the static
+        # buffers so the addresses captured segments hold stay live with
+        # current data.
+        if self._prefill_static_buffers is not None:
+            bs = forward_batch.batch_size
+            s = self._prefill_static_buffers
+            s["seq_lens"][:bs].copy_(forward_batch.seq_lens)
+            s["extend_seq_lens"][:bs].copy_(forward_batch.extend_seq_lens)
+            s["extend_prefix_lens"][:bs].copy_(forward_batch.extend_prefix_lens)
+            s["extend_start_loc"][:bs].copy_(forward_batch.extend_start_loc)
+            s["req_pool_indices"][:bs].copy_(forward_batch.req_pool_indices)
+            if forward_batch.orig_seq_lens is not None:
+                s["orig_seq_lens"][:bs].copy_(forward_batch.orig_seq_lens)
+
+        self._prepare_forward_metadata_for_replay(
+            forward_batch, static_forward_batch, static_num_tokens
+        )
+
+        self._static_num_tokens = static_num_tokens
+        return static_forward_batch
+
+    # -----------------------------------------------------------------
+    # replay
+    # -----------------------------------------------------------------
+    def replay(
+        self, forward_batch: ForwardBatch, **kwargs
+    ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
+        with self.backend.replay_session():
+            static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
+
+            if self.layer_model is not None:
+                # BCG path. The captured graph is a bs=1 replay of
+                # layer_model.forward. Monkey-patch layer_model.forward to
+                # call backend.replay (which fires the captured graph and
+                # returns the captured hidden_states), then drive the outer
+                # model.forward eagerly with the live multi-req
+                # static_forward_batch. The outer's logits_processor /
+                # pooler then runs on top with live multi-req metadata.
+                shape_key = self._static_num_tokens
+
+                def replay_layer_forward(*args, **layer_kwargs):
+                    return self.backend.replay(
+                        shape_key, static_forward_batch, **kwargs
+                    )
+
+                original_layer_forward = self.layer_model.forward
+                self.layer_model.forward = replay_layer_forward
+                try:
+                    with forward_context(
+                        ForwardContext(attn_backend=self.model_runner.attn_backend)
+                    ), set_tc_piecewise_forward_context(
+                        static_forward_batch,
+                        self.attention_layers,
+                        self.quant_config,
+                        self.moe_layers,
+                        self.moe_fusions,
+                        dsa_indexers=self.dsa_indexers,
+                    ):
+                        output = self.model_runner.model.forward(
+                            static_forward_batch.input_ids,
+                            static_forward_batch.positions,
+                            static_forward_batch,
+                            **kwargs,
+                        )
+                finally:
+                    self.layer_model.forward = original_layer_forward
+            else:
+                # TC_PIECEWISE path. backend.replay calls the compiled
+                # outer model.forward directly (torch.compile handles
+                # multi-req via bs-invariant FX-traced kernels).
+                with forward_context(
+                    ForwardContext(attn_backend=self.model_runner.attn_backend)
+                ), set_tc_piecewise_forward_context(
+                    static_forward_batch,
+                    self.attention_layers,
+                    self.quant_config,
+                    self.moe_layers,
+                    self.moe_fusions,
+                    dsa_indexers=self.dsa_indexers,
+                ):
+                    output = self.backend.replay(
+                        self._static_num_tokens, static_forward_batch, **kwargs
+                    )
+
+            if isinstance(output, LogitsProcessorOutput):
+                # Preserve mm_input_embeds for speculative decoding.
+                mm_input_embeds = None
+                if (
+                    self.model_runner.spec_algorithm.is_speculative()
+                    and output.mm_input_embeds is not None
+                ):
+                    mm_input_embeds = output.mm_input_embeds[: self.raw_num_tokens]
+                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
+                    ),
+                    mm_input_embeds=mm_input_embeds,
+                )
+            elif isinstance(output, EmbeddingPoolerOutput):
+                return output
+            else:
+                assert isinstance(output, PPProxyTensors)
+                raise NotImplementedError(
+                    "PPProxyTensors is not supported in PrefillCudaGraphRunner yet."
+                )

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,67 @@
+"""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
+
+
+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: Any,
+        forward_fn,
+        dummies: Optional[Any] = None,
+        post_warmup_hook: Optional[Callable[[], None]] = None,
+    ) -> None: ...
+
+    @abstractmethod
+    def can_run(self, forward_batch: ForwardBatch, shape_key: Any) -> bool: ...
+
+    @abstractmethod
+    def replay_session(self) -> Iterator[None]: ...
+
+    @abstractmethod
+    def replay(
+        self,
+        shape_key: Any,
+        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,127 @@
+"""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.runner_backend.base_cuda_graph_backend import (
+    BaseCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
+    BreakableCUDAGraph,
+    BreakableCUDAGraphCapture,
+    eager_on_graph,
+    enable_breakable_cuda_graph,
+)
+from sglang.srt.utils import get_bool_env_var, is_hip
+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,
+    )
+
+
+class BreakableCudaGraphBackend(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:
+        if is_hip():
+            raise RuntimeError("Breakable CUDA graph is not supported on ROCm/HIP")
+        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._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 = self._device_module.graph_pool_handle()
+        set_graph_pool_id(self._pool)
+        self._capture_stream = stream
+        try:
+            with self.replay_session():
+                yield
+        finally:
+            self._capture_stream = None
+
+    def capture_one(
+        self,
+        shape_key: Any,
+        forward_fn: Callable[[], Any],
+        dummies: Optional[Any] = None,
+        post_warmup_hook: Optional[Callable[[], None]] = None,
+    ) -> None:
+        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 = BreakableCUDAGraph()
+        captured_fn = (
+            eager_on_graph(True)(forward_fn) if self._debug_eager else forward_fn
+        )
+        with BreakableCUDAGraphCapture(
+            cuda_graph=graph,
+            pool=self._pool,
+            stream=self._capture_stream,
+        ):
+            out = captured_fn()
+        self._graphs[shape_key] = graph
+        self._outputs[shape_key] = out
+
+    def can_run(self, forward_batch: ForwardBatch, shape_key: Any) -> bool:
+        return shape_key in self._graphs
+
+    @contextmanager
+    def replay_session(self):
+        with enable_breakable_cuda_graph():
+            yield
+
+    def replay(
+        self,
+        shape_key: Any,
+        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/full_cuda_graph_backend.py

@@ -0,0 +1,118 @@
+"""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.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,
+    )
+
+
+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 = self._device_module.graph_pool_handle()
+        set_graph_pool_id(self._pool)
+        self._capture_stream = stream
+        try:
+            yield
+        finally:
+            self._capture_stream = None
+
+    def capture_one(
+        self,
+        shape_key: Any,
+        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: Any) -> bool:
+        return shape_key in self._graphs
+
+    @contextmanager
+    def replay_session(self):
+        yield
+
+    def replay(
+        self,
+        shape_key: Any,
+        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,225 @@
+"""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 import get_tensor_model_parallel_rank
+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,
+)
+
+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.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
+        compiler = self._compile_config.compiler
+        with enable_tc_piecewise_cuda_graph():
+            try:
+                if compiler != "eager":
+                    _toggle_multi_platform_ops(
+                        language_model.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 = self._device_module.graph_pool_handle()
+                set_graph_pool_id(self._pool)
+
+                self.install_compile(
+                    language_model.model,
+                    compile_config=self._compile_config,
+                    graph_pool=self._pool,
+                )
+
+                with enable_torch_compile_warmup():
+                    compile_range = (
+                        tqdm.tqdm(list(reversed(cuda_graph_runner.capture_num_tokens)))
+                        if get_tensor_model_parallel_rank() == 0
+                        else reversed(cuda_graph_runner.capture_num_tokens)
+                    )
+                    for num_tokens in compile_range:
+                        if get_tensor_model_parallel_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(
+                    language_model.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: Any,
+        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: Any) -> 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: Any,
+        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,96 @@
+"""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
+        )
+
+    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,
+        )
+    # Default: tc_piecewise. (prefill, full) is rejected at config validation.
+    return TcPiecewiseCudaGraphBackend(cuda_graph_runner)

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

@@ -0,0 +1,22 @@
+"""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 non-piecewise CUDA graph capture
+# paths (Full backend used by decode + EAGLE draft runners). The
+# piecewise-specific variant lives in
+# piecewise_cuda_graph.context_manager and points users at
+# --disable-piecewise-cuda-graph, which doesn't apply here.
+CUDA_GRAPH_CAPTURE_FAILED_MSG = (
+    "CUDA graph capture failed.\n"
+    "To work around this error, add --disable-cuda-graph to your launch command\n"
+    "(or use --disable-decode-cuda-graph to disable only the decode phase).\n"
+    "Please report this issue at https://github.com/sgl-project/sglang/issues/new/choose"
+)

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
+  - enable_breakable_cuda_graph — context that flips the Breakable runtime flag
+  - is_in_breakable_cuda_graph — runtime flag getter
+
+The legacy model_executor/breakable_cuda_graph/ package is a
+backwards-compat shim that re-exports from here.
+"""
+
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (  # noqa: F401
+    BreakableCUDAGraph,
+    BreakableCUDAGraphCapture,
+    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

@@ -12,12 +12,12 @@
 # limitations under the License.
 # ==============================================================================
 """Breakable CUDA Graph: capture a region as a sequence of
-``torch.cuda.CUDAGraph`` segments separated by eager break points.
+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
+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
+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.
 """
@@ -34,7 +34,9 @@ try:
 except ImportError:
     rt = None
 
-from sglang.srt.model_executor.breakable_cuda_graph.cuda_utils import checkCudaErrors
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.cuda_utils import (
+    checkCudaErrors,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -153,9 +155,9 @@ def _weak_ref_if_tensor(x):
     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: the module hard-raises on
+    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)."""
+    Breakable 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
 
@@ -238,7 +240,7 @@ def eager_on_graph(enable: bool):
 
 
 class BreakableCUDAGraph:
-    """Container holding one ``torch.cuda.CUDAGraph`` per segment plus an
+    """Container holding one torch.cuda.CUDAGraph per segment plus an
     eager break function between consecutive segments."""
 
     def __init__(self) -> None:
@@ -259,12 +261,12 @@ class BreakableCUDAGraph:
 
 class BreakableCUDAGraphCapture:
     """Context manager that captures the enclosed code as one or more
-    ``torch.cuda.CUDAGraph`` segments separated by eager break points.
+    torch.cuda.CUDAGraph segments separated by eager break points.
 
-    Each segment shares the supplied ``pool`` (``MempoolId_t`` tuple) so
+    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
+    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.
     """
 

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

@@ -11,12 +11,7 @@
 # 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.
-"""
+"""Runtime state for the breakable CUDA graph runner."""
 
 from __future__ import annotations
 

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

@@ -0,0 +1,21 @@
+"""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.
+  - TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG.
+
+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
+    TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG,
+    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,112 @@
+"""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
+
+logger = logging.getLogger(__name__)
+
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+_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 "we are inside a piecewise CUDA graph
+    capture/replay". Sets _in_tc_piecewise_cuda_graph true for the duration.
+
+    Errors during capture surface a hint that lets users disable the
+    feature while filing a bug.
+    """
+    global _in_tc_piecewise_cuda_graph
+    _in_tc_piecewise_cuda_graph = True
+    try:
+        yield
+    except Exception as e:
+        logger.error(
+            "Piecewise CUDA Graph failed with error: %s\n%s",
+            e,
+            TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_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)
+
+
+_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,
+):
+    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,
+    )
+    try:
+        yield
+    finally:
+        _tc_piecewise_forward_context = None
+
+
+TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG = (
+    "Piecewise CUDA Graph is enabled by default as an experimental feature.\n"
+    "To work around this error, add --disable-piecewise-cuda-graph to your launch command.\n"
+    "Please report this issue at https://github.com/sgl-project/sglang/issues/new/choose"
+)

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

@@ -0,0 +1,28 @@
+"""Low-level utilities used by the CUDA graph runners.
+
+Mirror of cuda_graph_backend_utils/ for runner-side state — buffer
+dataclasses, process-global capture flags, the speculative-shared
+graph memory pool, and the DeepEP capture/replay adapter. Runners in
+cuda_graph_runner/ import from here; nothing here should import
+back into cuda_graph_runner/.
+"""
+
+from sglang.srt.model_executor.runner_utils.buffers import (  # noqa: F401
+    DecodeInputBuffers,
+    PrefillInputBuffers,
+    _grouped_foreach_copy_,
+)
+from sglang.srt.model_executor.runner_utils.capture_mode import (  # noqa: F401
+    _set_capture_lora_variant,
+    compile_in_capture_mode,
+    get_capture_lora_variant,
+    get_is_capture_mode,
+    model_capture_mode,
+)
+from sglang.srt.model_executor.runner_utils.deepep_adapter import (  # noqa: F401
+    DeepEPCudaGraphRunnerAdapter,
+)
+from sglang.srt.model_executor.runner_utils.pool import (  # noqa: F401
+    get_global_graph_memory_pool,
+    set_global_graph_memory_pool,
+)

python/sglang/srt/model_executor/runner_utils/buffers.py

@@ -0,0 +1,442 @@
+"""Static-buffer dataclasses used by the CUDA graph runners.
+
+DecodeInputBuffers backs the decode-phase capture/replay path.
+PrefillInputBuffers backs the prefill-phase capture/replay path.
+
+Both subclass ForwardInputBuffers so that buffer-pool sharing works
+the same way as for non-cuda-graph forward paths.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.forward_batch_info import (
+    ForwardBatch,
+    NgramEmbeddingInfo,
+    PPProxyTensors,
+    compute_local_num_token_non_padded,
+)
+from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
+
+_has_foreach_copy = hasattr(torch, "_foreach_copy_")
+
+
+def _grouped_foreach_copy_(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:
+    """Call torch._foreach_copy_ grouped by (dst_dtype, src_dtype) pairs."""
+
+    def foreach_copy(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:
+        if _has_foreach_copy:
+            torch._foreach_copy_(dsts, srcs)
+        else:
+            for dst, src in zip(dsts, srcs):
+                dst.copy_(src)
+
+    groups: Dict[Tuple[torch.dtype, torch.dtype], Tuple[List, List]] = {}
+    for dst, src in zip(dsts, srcs):
+        key = (dst.dtype, src.dtype)
+        if key not in groups:
+            groups[key] = ([], [])
+        groups[key][0].append(dst)
+        groups[key][1].append(src)
+    for group_dsts, group_srcs in groups.values():
+        foreach_copy(group_dsts, group_srcs)
+
+
+@dataclass
+class DecodeInputBuffers(ForwardInputBuffers):
+
+    input_ids: torch.Tensor
+    input_embeds: torch.Tensor
+    req_pool_indices: torch.Tensor
+    seq_lens: torch.Tensor
+    seq_lens_cpu: torch.Tensor
+    out_cache_loc: torch.Tensor
+    out_cache_loc_swa: Optional[torch.Tensor]
+    positions: torch.Tensor
+    mrope_positions: torch.Tensor
+    num_token_non_padded: torch.Tensor
+    custom_mask: torch.Tensor
+    next_token_logits_buffer: torch.Tensor
+    mamba_track_indices: Optional[torch.Tensor]
+    mamba_track_mask: Optional[torch.Tensor]
+    global_num_tokens_gpu: torch.Tensor
+    global_num_tokens_for_logprob_gpu: torch.Tensor
+    encoder_lens: Optional[torch.Tensor]
+    pp_proxy_tensors: Optional[Dict[str, torch.Tensor]]
+    ngram_embedding_info: Optional["NgramEmbeddingInfo"]
+    rids_int: Optional[torch.Tensor]
+    bootstrap_room_ids_int: Optional[torch.Tensor]
+
+    @classmethod
+    def create(
+        cls,
+        *,
+        device: torch.device,
+        max_bs: int,
+        max_num_token: int,
+        hidden_size: int,
+        vocab_size: int,
+        dtype: torch.dtype,
+        dp_size: int,
+        pp_size: int,
+        is_encoder_decoder: bool,
+        require_mlp_tp_gather: bool,
+        seq_len_fill_value: int,
+        encoder_len_fill_value: int,
+        num_tokens_per_bs: int,
+        cache_loc_dtype: torch.dtype,
+        enable_mamba_track: bool,
+        ne_token_table: Optional[torch.Tensor] = None,
+        is_hybrid_swa: bool = False,
+        hc_hidden_size: Optional[int] = None,
+    ) -> "DecodeInputBuffers":
+        with torch.device(device):
+            input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
+            input_embeds = torch.zeros((max_num_token, hidden_size), dtype=dtype)
+            req_pool_indices = torch.zeros((max_bs,), dtype=torch.int64)
+            seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int32)
+            out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype)
+            out_cache_loc_swa = (
+                torch.zeros((max_num_token,), dtype=torch.int64)
+                if is_hybrid_swa
+                else None
+            )
+            positions = torch.zeros((max_num_token,), dtype=torch.int64)
+            mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64)
+            num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
+            custom_mask = torch.ones(
+                (max_bs * seq_len_fill_value + max_num_token) * num_tokens_per_bs,
+                dtype=torch.bool,
+            )
+            next_token_logits_buffer = torch.zeros(
+                (max_num_token, vocab_size),
+                dtype=torch.float,
+            )
+            mamba_track_indices = (
+                torch.zeros((max_bs,), dtype=torch.int64)
+                if enable_mamba_track
+                else None
+            )
+            mamba_track_mask = (
+                torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None
+            )
+
+            if pp_size > 1:
+                is_mhc = hc_hidden_size is not None
+                hs = hc_hidden_size if is_mhc else hidden_size
+                pp_proxy_tensors = {
+                    "hidden_states": torch.zeros((max_bs, hs), dtype=dtype),
+                }
+                if not is_mhc:
+                    pp_proxy_tensors["residual"] = torch.zeros(
+                        (max_bs, hidden_size), dtype=dtype
+                    )
+            else:
+                pp_proxy_tensors = None
+
+            if is_encoder_decoder:
+                encoder_lens = torch.full(
+                    (max_bs,), encoder_len_fill_value, dtype=torch.int32
+                )
+            else:
+                encoder_lens = None
+
+            if require_mlp_tp_gather:
+                global_num_tokens_gpu = torch.zeros((dp_size,), dtype=torch.int32)
+                global_num_tokens_for_logprob_gpu = torch.zeros(
+                    (dp_size,), dtype=torch.int32
+                )
+            else:
+                global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
+                global_num_tokens_for_logprob_gpu = torch.zeros((1,), dtype=torch.int32)
+
+            ngram_embedding_info = (
+                NgramEmbeddingInfo(
+                    token_table=ne_token_table,
+                    column_starts=torch.zeros([max_bs], dtype=torch.int32),
+                    req_lens=torch.ones([max_bs], dtype=torch.int32),
+                    out_column_starts=torch.zeros([max_bs], dtype=torch.int32),
+                    out_req_lens=torch.ones([max_bs], dtype=torch.int32),
+                )
+                if ne_token_table is not None
+                else None
+            )
+
+            if envs.SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE.get():
+                rids_int = torch.zeros((max_bs,), dtype=torch.int64)
+                bootstrap_room_ids_int = torch.full((max_bs,), -1, dtype=torch.int64)
+            else:
+                rids_int = None
+                bootstrap_room_ids_int = None
+
+        seq_lens_cpu = torch.full(
+            (max_bs,),
+            seq_len_fill_value,
+            dtype=torch.int32,
+            device="cpu",
+        )
+
+        return cls(
+            input_ids=input_ids,
+            input_embeds=input_embeds,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens_cpu,
+            out_cache_loc=out_cache_loc,
+            out_cache_loc_swa=out_cache_loc_swa,
+            positions=positions,
+            mrope_positions=mrope_positions,
+            num_token_non_padded=num_token_non_padded,
+            custom_mask=custom_mask,
+            next_token_logits_buffer=next_token_logits_buffer,
+            mamba_track_indices=mamba_track_indices,
+            mamba_track_mask=mamba_track_mask,
+            encoder_lens=encoder_lens,
+            global_num_tokens_gpu=global_num_tokens_gpu,
+            global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
+            pp_proxy_tensors=pp_proxy_tensors,
+            ngram_embedding_info=ngram_embedding_info,
+            rids_int=rids_int,
+            bootstrap_room_ids_int=bootstrap_room_ids_int,
+        )
+
+    def populate_from_forward_batch(
+        self,
+        *,
+        forward_batch: ForwardBatch,
+        raw_bs: int,
+        raw_num_token: int,
+        bs: int,
+        seq_len_fill_value: int,
+        require_gathered_buffer: bool,
+        num_tokens_per_bs: int,
+        dsa_enable_prefill_cp: bool,
+        enable_num_token_non_padded_flag: bool,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
+    ):
+        if bs != raw_bs:
+            self.seq_lens.fill_(seq_len_fill_value)
+            self.out_cache_loc.zero_()
+            if self.mamba_track_indices is not None:
+                self.mamba_track_indices.zero_()
+            if self.mamba_track_mask is not None:
+                self.mamba_track_mask.fill_(False)
+
+        # Build batched copy lists for all GPU tensors.
+        dsts = [
+            self.input_ids[:raw_num_token],
+            self.req_pool_indices[:raw_bs],
+            self.seq_lens[:raw_bs],
+            self.out_cache_loc[:raw_num_token],
+            self.positions[:raw_num_token],
+        ]
+        srcs = [
+            forward_batch.input_ids,
+            forward_batch.req_pool_indices,
+            forward_batch.seq_lens,
+            forward_batch.out_cache_loc,
+            forward_batch.positions,
+        ]
+
+        if self.ngram_embedding_info is not None:
+            ngram_embedding_info = forward_batch.ngram_embedding_info
+            self.ngram_embedding_info.column_starts[:raw_bs].copy_(
+                ngram_embedding_info.column_starts
+            )
+            self.ngram_embedding_info.req_lens[:raw_bs].copy_(
+                ngram_embedding_info.req_lens
+            )
+
+        if (
+            self.mamba_track_indices is not None
+            and forward_batch.mamba_track_indices is not None
+        ):
+            dsts.append(self.mamba_track_indices[:raw_bs])
+            srcs.append(forward_batch.mamba_track_indices)
+        if (
+            self.mamba_track_mask is not None
+            and forward_batch.mamba_track_mask is not None
+        ):
+            dsts.append(self.mamba_track_mask[:raw_bs])
+            srcs.append(forward_batch.mamba_track_mask)
+
+        if self.encoder_lens is not None and forward_batch.encoder_lens is not None:
+            dsts.append(self.encoder_lens[:raw_bs])
+            srcs.append(forward_batch.encoder_lens)
+
+        if forward_batch.mrope_positions is not None:
+            dsts.append(self.mrope_positions[:, :raw_num_token])
+            srcs.append(forward_batch.mrope_positions)
+
+        if self.rids_int is not None and forward_batch.rids_int is not None:
+            dsts.append(self.rids_int[:raw_bs])
+            srcs.append(forward_batch.rids_int)
+        if (
+            self.bootstrap_room_ids_int is not None
+            and forward_batch.bootstrap_room_ids_int is not None
+        ):
+            dsts.append(self.bootstrap_room_ids_int[:raw_bs])
+            srcs.append(forward_batch.bootstrap_room_ids_int)
+
+        if require_gathered_buffer:
+            self.global_num_tokens_gpu.fill_(bs * num_tokens_per_bs)
+            self.global_num_tokens_for_logprob_gpu.fill_(bs * num_tokens_per_bs)
+
+        if enable_num_token_non_padded_flag:
+            if require_gathered_buffer and not dsa_enable_prefill_cp:
+                num_tokens_per_dp = bs * num_tokens_per_bs
+                local = compute_local_num_token_non_padded(
+                    global_num_token_non_padded=forward_batch.num_token_non_padded,
+                    num_tokens_per_dp=num_tokens_per_dp,
+                )
+                dsts.append(self.num_token_non_padded)
+                srcs.append(local)
+            else:
+                dsts.append(self.num_token_non_padded)
+                srcs.append(forward_batch.num_token_non_padded)
+
+        # Pipeline-parallel proxy tensors.
+        if pp_proxy_tensors is not None and self.pp_proxy_tensors is not None:
+            for key, buf in self.pp_proxy_tensors.items():
+                src = pp_proxy_tensors.tensors[key]
+                dim = src.shape[0]
+                dsts.append(buf[:dim])
+                srcs.append(src)
+
+        # SWA cache location (int32, separate from the int64 batch above).
+        if (
+            self.out_cache_loc_swa is not None
+            and forward_batch.out_cache_loc_swa is not None
+        ):
+            dsts.append(self.out_cache_loc_swa[:raw_num_token])
+            srcs.append(forward_batch.out_cache_loc_swa[:raw_num_token])
+
+        # Batch all GPU copies, grouped by dtype pair.
+        _grouped_foreach_copy_(dsts, srcs)
+
+        if forward_batch.seq_lens_cpu is not None:
+            if bs != raw_bs:
+                self.seq_lens_cpu.fill_(seq_len_fill_value)
+            self.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
+
+
+@dataclass
+class PrefillInputBuffers(ForwardInputBuffers):
+    input_ids: torch.Tensor
+    out_cache_loc: torch.Tensor
+    out_cache_loc_swa: Optional[torch.Tensor]
+    mamba_track_indices: Optional[torch.Tensor]
+    mamba_track_mask: Optional[torch.Tensor]
+    mamba_track_seqlens: Optional[torch.Tensor]
+    positions: torch.Tensor
+    input_embeds: Optional[torch.Tensor]
+    mrope_positions: Optional[torch.Tensor]
+
+    @classmethod
+    def create(
+        cls,
+        *,
+        device: torch.device,
+        max_bs: int,
+        max_num_tokens: int,
+        cache_loc_dtype: torch.dtype,
+        is_hybrid_swa: bool,
+        is_multimodal: bool,
+        hidden_size: int,
+        dtype: torch.dtype,
+        enable_mamba_track: bool,
+    ) -> "PrefillInputBuffers":
+        with torch.device(device):
+            input_ids = torch.zeros((max_num_tokens,), dtype=torch.int64)
+            out_cache_loc = torch.zeros((max_num_tokens,), dtype=cache_loc_dtype)
+            out_cache_loc_swa = (
+                torch.zeros((max_num_tokens,), dtype=torch.int64)
+                if is_hybrid_swa
+                else None
+            )
+            mamba_track_indices = (
+                torch.zeros((max_bs,), dtype=torch.int64)
+                if enable_mamba_track
+                else None
+            )
+            mamba_track_mask = (
+                torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None
+            )
+            mamba_track_seqlens = (
+                torch.zeros((max_bs,), dtype=torch.int32)
+                if enable_mamba_track
+                else None
+            )
+            positions = torch.zeros((max_num_tokens,), dtype=torch.int64)
+
+            if is_multimodal:
+                input_embeds = torch.zeros((max_num_tokens, hidden_size), dtype=dtype)
+                mrope_positions = torch.zeros((3, max_num_tokens), dtype=torch.int64)
+            else:
+                input_embeds = None
+                mrope_positions = None
+
+        return cls(
+            input_ids=input_ids,
+            out_cache_loc=out_cache_loc,
+            out_cache_loc_swa=out_cache_loc_swa,
+            mamba_track_indices=mamba_track_indices,
+            mamba_track_mask=mamba_track_mask,
+            mamba_track_seqlens=mamba_track_seqlens,
+            positions=positions,
+            input_embeds=input_embeds,
+            mrope_positions=mrope_positions,
+        )
+
+    def populate_from_forward_batch(
+        self,
+        *,
+        forward_batch: ForwardBatch,
+        raw_num_tokens: int,
+        static_num_tokens: int,
+        is_multimodal: bool,
+    ) -> None:
+        """Copy serving-batch values into static buffers and zero out
+        the padding region between raw_num_tokens and
+        static_num_tokens.
+        """
+        if static_num_tokens != raw_num_tokens:
+            self.out_cache_loc.zero_()
+            self.input_ids[raw_num_tokens:static_num_tokens].zero_()
+            self.positions[raw_num_tokens:static_num_tokens].zero_()
+            if is_multimodal:
+                self.input_embeds[raw_num_tokens:static_num_tokens].zero_()
+            if forward_batch.mrope_positions is not None:
+                self.mrope_positions[:, raw_num_tokens:static_num_tokens].zero_()
+
+        bs = forward_batch.batch_size
+
+        self.input_ids[:raw_num_tokens].copy_(forward_batch.input_ids)
+        self.positions[:raw_num_tokens].copy_(forward_batch.positions)
+        self.out_cache_loc[:raw_num_tokens].copy_(forward_batch.out_cache_loc)
+
+        if (
+            self.mamba_track_indices is not None
+            and forward_batch.mamba_track_indices is not None
+        ):
+            self.mamba_track_indices[:bs].copy_(forward_batch.mamba_track_indices)
+        if (
+            self.mamba_track_mask is not None
+            and forward_batch.mamba_track_mask is not None
+        ):
+            self.mamba_track_mask[:bs].copy_(forward_batch.mamba_track_mask)
+        if (
+            self.mamba_track_seqlens is not None
+            and forward_batch.mamba_track_seqlens is not None
+        ):
+            self.mamba_track_seqlens[:bs].copy_(forward_batch.mamba_track_seqlens)
+
+        if forward_batch.mrope_positions is not None:
+            self.mrope_positions[:, :raw_num_tokens].copy_(
+                forward_batch.mrope_positions
+            )

python/sglang/srt/model_executor/runner_utils/capture_mode.py

@@ -0,0 +1,56 @@
+"""Process-global capture-mode flags shared by the decode runner and the
+speculative-draft runners. Read by model code that needs to take a
+capture-time branch (e.g. lora dual-graph capture decides per-batch
+which variant to use).
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Optional
+
+import torch
+
+# Detect whether the current forward pass is in capture mode.
+is_capture_mode = False
+
+# When capturing dual MoE backends, tracks which variant is being captured.
+# None = not dual, "lora" = capturing lora variant, "nolora" = capturing nolora variant.
+_capture_lora_variant: Optional[str] = None
+
+
+def get_is_capture_mode() -> bool:
+    return is_capture_mode
+
+
+def compile_in_capture_mode(func):
+    """Decorator: wrap func with torch.compile only when defined
+    inside model capture mode; passthrough otherwise.
+
+    Used by model code (e.g. DeepSeek-V4) to opt nested helpers into
+    torch.compile during cuda-graph capture without paying the
+    compilation cost in the eager forward path.
+    """
+    if is_capture_mode:
+        return torch.compile(func)
+    return func
+
+
+def get_capture_lora_variant() -> Optional[str]:
+    """Return the lora variant being captured, or None if not in dual capture."""
+    return _capture_lora_variant
+
+
+def _set_capture_lora_variant(variant: Optional[str]) -> None:
+    global _capture_lora_variant
+    _capture_lora_variant = variant
+
+
+@contextmanager
+def model_capture_mode():
+    global is_capture_mode
+    is_capture_mode = True
+    try:
+        yield
+    finally:
+        is_capture_mode = False

python/sglang/srt/model_executor/runner_utils/deepep_adapter.py

@@ -0,0 +1,29 @@
+"""DeepEP capture/replay adapter — records the dispatch mode used during
+capture and re-applies it during replay so DeepEP all-to-all has
+consistent expert routing across the captured graph.
+"""
+
+from __future__ import annotations
+
+from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPBuffer
+from sglang.srt.layers.moe.utils import get_deepep_mode, get_moe_a2a_backend
+
+
+class DeepEPCudaGraphRunnerAdapter:
+    def __init__(self) -> None:
+        # Record DeepEP mode used during capture to ensure replay consistency.
+        self._captured_deepep_mode = None
+
+    def capture(self, is_extend_in_batch: bool) -> None:
+        if not get_moe_a2a_backend().is_deepep():
+            return
+        self._captured_deepep_mode = get_deepep_mode().resolve(
+            is_extend_in_batch=is_extend_in_batch
+        )
+        DeepEPBuffer.set_dispatch_mode(self._captured_deepep_mode)
+
+    def replay(self) -> None:
+        if not get_moe_a2a_backend().is_deepep():
+            return
+        assert self._captured_deepep_mode is not None
+        DeepEPBuffer.set_dispatch_mode(self._captured_deepep_mode)

python/sglang/srt/model_executor/runner_utils/pool.py

@@ -0,0 +1,21 @@
+"""Shared graph memory pool used by the speculative-draft cuda graph
+runners. The new DecodeCudaGraphRunner and PrefillCudaGraphRunner
+backends each own their pool internally; this global is retained for the
+EAGLE / multi-step draft runners that haven't been folded into the new
+backend interface.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+_global_graph_memory_pool: Optional[Any] = None
+
+
+def get_global_graph_memory_pool() -> Optional[Any]:
+    return _global_graph_memory_pool
+
+
+def set_global_graph_memory_pool(val: Any) -> None:
+    global _global_graph_memory_pool
+    _global_graph_memory_pool = val

python/sglang/srt/models/bailing_moe.py

@@ -73,8 +73,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.utils import (
     apply_qk_norm,

python/sglang/srt/models/bailing_moe_linear.py

@@ -57,8 +57,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA, DeepseekV2MLP, _is_hip
 from sglang.srt.models.utils import WeightsMapper

python/sglang/srt/models/cohere2_moe.py

@@ -29,8 +29,8 @@ from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.layers.rotary_embedding import get_rope
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.utils import add_prefix, get_compiler_backend, is_cuda, make_layers
 

python/sglang/srt/models/deepseek_common/attention_backend_handler.py

@@ -1,7 +1,9 @@
-from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
 from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
 from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
 from sglang.srt.model_executor.forward_context import get_attn_backend
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.models.deepseek_common.attention_forward_methods.forward_methods import (
     AttnForwardMethod,
 )
@@ -72,7 +74,7 @@ def _support_mha_one_shot(attn, forward_batch, backend_name):
 
 
 def _handle_attention_backend(attn, forward_batch, backend_name):
-    if is_in_piecewise_cuda_graph():
+    if is_in_tc_piecewise_cuda_graph():
         return AttnForwardMethod.MLA
 
     # MLA prefill CP forces absorbed MLA regardless of prefix length: the
@@ -130,7 +132,7 @@ def handle_attention_fa4(attn, forward_batch):
 
 
 def handle_attention_trtllm_mla(attn, forward_batch):
-    if is_in_piecewise_cuda_graph():
+    if is_in_tc_piecewise_cuda_graph():
         return AttnForwardMethod.MLA
 
     sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch)
@@ -170,7 +172,7 @@ def handle_attention_dsa(attn, forward_batch):
 
 
 def handle_attention_triton(attn, forward_batch):
-    if is_in_piecewise_cuda_graph():
+    if is_in_tc_piecewise_cuda_graph():
         return AttnForwardMethod.MLA
 
     # when deterministic inference is enabled, use MLA

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -4,7 +4,6 @@ from typing import TYPE_CHECKING, Optional
 
 import torch
 
-from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
 from sglang.srt.environ import envs
 from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
@@ -29,6 +28,9 @@ from sglang.srt.model_executor.forward_context import (
     get_attn_backend,
     get_token_to_kv_pool,
 )
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.models.deepseek_common.utils import (
     FORWARD_ABSORB_CORE_ATTENTION_BACKENDS,
     _is_cpu,
@@ -147,7 +149,7 @@ class DeepseekMLAForwardMixin:
         llama_4_scaling: Optional[torch.Tensor] = None,
         prev_topk_indices: Optional[torch.Tensor] = None,
     ):
-        from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
+        from sglang.srt.model_executor.runner import get_is_capture_mode
 
         q_lora = None
         topk_indices = None
@@ -709,7 +711,7 @@ class DeepseekMLAForwardMixin:
             )
             attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
         else:
-            if is_in_piecewise_cuda_graph():
+            if is_in_tc_piecewise_cuda_graph():
                 # torch dynamo requires out= op was called where output tensor was non-contiguous
                 attn_bmm_output = (
                     torch.bmm(attn_output.transpose(0, 1), self.w_vc)

python/sglang/srt/models/deepseek_v2.py

@@ -135,8 +135,13 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
     get_embedding_tp_kwargs,
 )
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.models.deepseek_common.attention_backend_handler import (
     AttentionBackendRegistry,
 )
@@ -724,7 +729,7 @@ class DeepseekV2MoE(nn.Module):
                     ModelOptFp4LinearMethod,
                 )
                 and fc1_n % 128 == 0
-                and get_global_server_args().disable_piecewise_cuda_graph
+                and not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
             ):
                 self.shared_experts.gate_up_proj._interleave_for_swiglu_fusion = True
                 self.shared_experts._enable_nvfp4_gemm_swiglu_fusion = True
@@ -2439,7 +2444,7 @@ class DeepseekV2Model(nn.Module):
             # NOTE: torch dynamo does not support graph break in context manager
             ctx = (
                 nullcontext()
-                if not get_global_server_args().disable_piecewise_cuda_graph
+                if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
                 else get_global_expert_distribution_recorder().with_current_layer(i)
             )
             with ctx:

python/sglang/srt/models/deepseek_v4.py

@@ -28,7 +28,6 @@ from sglang.jit_kernel.dsv4 import (
     fused_rope_inplace,
 )
 from sglang.srt.compilation.compilation_config import register_split_op
-from sglang.srt.compilation.piecewise_context_manager import get_forward_context
 from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
 from sglang.srt.distributed import (
     get_pp_group,
@@ -84,21 +83,29 @@ from sglang.srt.layers.utils.cp_utils import (
 )
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from sglang.srt.mem_cache.memory_pool import RadixAttention
-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.cuda_graph_runner import (
-    compile_in_capture_mode,
-    get_is_capture_mode,
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
 )
 from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
 from sglang.srt.model_executor.forward_context import (
     get_attn_backend,
     get_token_to_kv_pool,
 )
+from sglang.srt.model_executor.runner import (
+    compile_in_capture_mode,
+    get_is_capture_mode,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
+    eager_on_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+    is_in_breakable_cuda_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    get_tc_piecewise_forward_context,
+)
 from sglang.srt.model_loader.utils import maybe_executor_submit, should_async_load
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.dbrx import ReplicatedLinear
@@ -211,7 +218,7 @@ def deepseek_v4_attention_with_output(
     attn_sink: torch.Tensor,
     save_kv_cache: bool,
 ) -> None:
-    context = get_forward_context()
+    context = get_tc_piecewise_forward_context()
     forward_batch = context.forward_batch
     attention_layers = context.attention_layers
     attention_layer = attention_layers[layer_id]
@@ -1678,7 +1685,7 @@ class DeepseekV4Model(nn.Module):
             last_layer = layer
             ctx = (
                 nullcontext()
-                if not get_global_server_args().disable_piecewise_cuda_graph
+                if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
                 else get_global_expert_distribution_recorder().with_current_layer(i)
             )
             with ctx:

python/sglang/srt/models/exaone_moe.py

@@ -63,8 +63,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_layers

python/sglang/srt/models/glm4_moe.py

@@ -82,8 +82,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
 from sglang.srt.models.utils import apply_qk_norm

python/sglang/srt/models/glm4_moe_lite.py

@@ -68,8 +68,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
     DeepseekV2WeightLoaderMixin,

python/sglang/srt/models/gpt_oss.py

@@ -26,10 +26,6 @@ import torch
 from torch import nn
 from transformers import PretrainedConfig
 
-from sglang.srt.compilation.piecewise_context_manager import (
-    get_forward_context,
-    is_in_piecewise_cuda_graph,
-)
 from sglang.srt.distributed import (
     get_moe_expert_parallel_rank,
     get_moe_expert_parallel_world_size,
@@ -70,6 +66,10 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    get_tc_piecewise_forward_context,
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.utils import (
     create_fused_set_kv_buffer_arg,
@@ -299,7 +299,7 @@ class GptOssSparseMoeBlock(nn.Module):
         else:
             router_input = hidden_states
 
-        if is_in_piecewise_cuda_graph():
+        if is_in_tc_piecewise_cuda_graph():
             final_hidden_states = moe_impl(self.layer_id, hidden_states)
         else:
             router_logits, _ = self.router(router_input)
@@ -326,7 +326,7 @@ class GptOssSparseMoeBlock(nn.Module):
 
 @register_custom_op(out_shape="hidden_states")
 def moe_impl(layer_id: int, hidden_states: torch.Tensor) -> torch.Tensor:
-    forward_context = get_forward_context()
+    forward_context = get_tc_piecewise_forward_context()
     moe_fusion = forward_context.moe_fusions[layer_id]
     router_logits, _ = moe_fusion.router(hidden_states)
     topk_output = moe_fusion.topk(hidden_states, router_logits)

python/sglang/srt/models/grok.py

@@ -56,8 +56,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.loader import DefaultModelLoader
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.utils import add_prefix, is_npu

python/sglang/srt/models/hunyuan_v3.py

@@ -45,7 +45,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
 from sglang.srt.managers.schedule_batch import ForwardBatch
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.utils import is_cuda
 from sglang.srt.utils.hf_transformers_utils import get_rope_config

python/sglang/srt/models/kimi_linear.py

@@ -39,8 +39,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,

python/sglang/srt/models/llada2.py

@@ -72,8 +72,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.utils import (
     apply_qk_norm,

python/sglang/srt/models/minimax_m2.py

@@ -74,6 +74,11 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
@@ -633,7 +638,7 @@ class MiniMaxM2MoE(nn.Module):
         if router_logits is not None:
             ctx = (
                 nullcontext()
-                if not get_global_server_args().disable_piecewise_cuda_graph
+                if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
                 else get_global_expert_distribution_recorder().with_current_layer(
                     self.layer_id
                 )
@@ -671,7 +676,7 @@ class MiniMaxM2MoE(nn.Module):
         if self.ep_size > 1:
             ctx = (
                 nullcontext()
-                if not get_global_server_args().disable_piecewise_cuda_graph
+                if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
                 else get_global_expert_distribution_recorder().with_current_layer(
                     self.layer_id
                 )
@@ -1171,7 +1176,7 @@ class MiniMaxM2Model(nn.Module):
             for i in range(self.start_layer, self.end_layer):
                 ctx = (
                     nullcontext()
-                    if not get_global_server_args().disable_piecewise_cuda_graph
+                    if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
                     else get_global_expert_distribution_recorder().with_current_layer(i)
                 )
                 with ctx:

python/sglang/srt/models/mllama.py

@@ -957,7 +957,7 @@ class MllamaForConditionalGeneration(nn.Module):
         positions: torch.Tensor,
         forward_batch: ForwardBatch,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
-        from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
+        from sglang.srt.model_executor.runner import get_is_capture_mode
 
         batched_images, batched_ar_ids, batched_ar_mask, encoder_lens_need = (
             self._batch_image_inputs(forward_batch)

python/sglang/srt/models/moss_vl.py

@@ -44,8 +44,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
 from sglang.srt.managers.schedule_batch import MultimodalInputs
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix

python/sglang/srt/models/nemotron_h.py

@@ -24,10 +24,6 @@ import torch
 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,
-    is_in_piecewise_cuda_graph,
-)
 from sglang.srt.configs import NemotronHConfig
 from sglang.srt.configs.nemotron_h import ATTENTION, MAMBA, MLP, MOE
 from sglang.srt.distributed import (
@@ -62,14 +58,16 @@ 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_executor.forward_context import get_attn_backend
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
+    eager_on_graph,
+    is_in_breakable_cuda_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    get_tc_piecewise_forward_context,
+    is_in_tc_piecewise_cuda_graph,
+)
 from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
@@ -232,9 +230,10 @@ class NemotronHMoE(nn.Module):
         self,
         hidden_states: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        # torch.compile cannot trace CUDA streams, so use the non-overlapping
-        # path when inside piecewise CUDA graph compilation.
-        if _is_cuda and not is_in_piecewise_cuda_graph():
+        # torch.compile cannot trace CUDA streams. Take the
+        # non-overlapping path only during dynamo tracing; replay can
+        # use the overlapping fast path since dynamo is no longer active.
+        if _is_cuda and not torch.compiler.is_compiling():
             return self._forward_core_shared_routed_overlap(hidden_states)
         else:
             return self._forward_core_normal(hidden_states)
@@ -447,7 +446,7 @@ class NemotronHMambaDecoderLayer(nn.Module):
             breakable_nemotron_mamba2_with_output(hidden_states, output, self.layer_id)
             return output, residual
 
-        if is_in_piecewise_cuda_graph():
+        if is_in_tc_piecewise_cuda_graph():
             output = torch.empty_like(hidden_states)
             nemotron_mamba2_with_output(hidden_states, output, self.layer_id)
             return output, residual
@@ -1022,7 +1021,7 @@ def nemotron_mamba2_with_output(
     layer_id: int,
 ) -> None:
     """Split op for Mamba2 forward in piecewise CUDA graph mode."""
-    context = get_forward_context()
+    context = get_tc_piecewise_forward_context()
     forward_batch = context.forward_batch
     attention_layers = context.attention_layers
     mamba_layer = attention_layers[layer_id]

python/sglang/srt/models/olmo2.py

@@ -46,8 +46,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.utils import add_prefix, is_cuda, make_layers
 
@@ -63,7 +63,7 @@ def get_attention_sliding_window_size(config):
 class Olmo2Attention(nn.Module):
     """
     This is the attention block where the output is computed as
-    ``Attention(LN(x))`` in ``MLP(LN(x + Attention(LN(x))))``
+    Attention(LN(x)) in MLP(LN(x + Attention(LN(x))))
     (plus another skip connection).
     """
 
@@ -216,7 +216,7 @@ class Olmo2Attention(nn.Module):
 class Olmo2MLP(nn.Module):
     """
     This is the MLP block where the output is computed as
-    ``MLP(x)`` in ``LN(MLP(x + LN(Attention(x))))``
+    MLP(x) in LN(MLP(x + LN(Attention(x))))
     (plus another skip connection).
     """
 
@@ -265,7 +265,7 @@ class Olmo2MLP(nn.Module):
 class Olmo2DecoderLayer(nn.Module):
     """
     This is a typical transformer block where the output is
-    computed as ``MLP(LN(x + Attention(LN(x))))``
+    computed as MLP(LN(x + Attention(LN(x))))
     (plus another skip connection).
     """
 

python/sglang/srt/models/qwen2_moe.py

@@ -87,8 +87,13 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
@@ -129,7 +134,7 @@ def can_fuse_shared_expert(
 ) -> bool:
     """Whether the shared expert may be fused as an extra MoE expert (Qwen3.5 + Aiter).
 
-    Caller must still gate on ``support_shared_expert_fusion`` and ``_use_aiter``.
+    Caller must still gate on support_shared_expert_fusion and _use_aiter.
     """
     if (
         get_global_server_args().disable_shared_experts_fusion is True
@@ -868,7 +873,7 @@ class Qwen2MoeModel(nn.Module):
             for i in range(self.start_layer, self.end_layer):
                 ctx = (
                     nullcontext()
-                    if not get_global_server_args().disable_piecewise_cuda_graph
+                    if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
                     else get_global_expert_distribution_recorder().with_current_layer(i)
                 )
                 with ctx:

python/sglang/srt/models/qwen3.py

@@ -22,6 +22,11 @@ from sglang.srt.layers.rotary_embedding import get_rope
 from sglang.srt.layers.rotary_embedding.mrope import MRotaryEmbedding
 from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
 from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
 from sglang.srt.model_loader.weight_utils import (
@@ -416,7 +421,7 @@ class Qwen3DecoderLayer(nn.Module):
             cache=(
                 [self.mlp.gate_up_proj.weight, self.mlp.down_proj.weight]
                 if _is_npu
-                and not get_global_server_args().disable_piecewise_cuda_graph
+                and check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
                 and (
                     hasattr(self.mlp.gate_up_proj, "weight")
                     and hasattr(self.mlp.down_proj, "weight")

python/sglang/srt/models/qwen3_5.py

@@ -69,8 +69,13 @@ from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
 from sglang.srt.layers.rotary_embedding import get_rope
 from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     sharded_weight_loader,
@@ -448,7 +453,7 @@ class Qwen3_5GatedDeltaNet(nn.Module):
         if (
             _is_cpu
             or _is_npu
-            or not get_global_server_args().disable_piecewise_cuda_graph
+            or check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
         ):
             DUAL_STREAM_TOKEN_THRESHOLD = 0
         else:

python/sglang/srt/models/qwen3_next.py

@@ -35,8 +35,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     sharded_weight_loader,
@@ -58,6 +58,11 @@ logger = logging.getLogger(__name__)
 
 from sglang.jit_kernel.triton.gdn_fused_proj import fused_qkvzba_split_reshape_cat
 from sglang.srt.layers.attention.fla.fused_norm_gate import FusedRMSNormGated
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    Phase,
+    check_cuda_graph_backend,
+)
 
 _is_cuda = is_cuda()
 _is_npu = is_npu()
@@ -196,7 +201,7 @@ class Qwen3GatedDeltaNet(nn.Module):
                     else {}
                 ),
             )
-            if not get_global_server_args().disable_piecewise_cuda_graph
+            if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
             else FusedRMSNormGated(
                 self.head_v_dim,
                 eps=self.layer_norm_epsilon,
@@ -372,7 +377,7 @@ class Qwen3GatedDeltaNet(nn.Module):
         if (
             _is_cpu
             or _is_npu
-            or not get_global_server_args().disable_piecewise_cuda_graph
+            or check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
         ):
             DUAL_STREAM_TOKEN_THRESHOLD = 0
         else:

python/sglang/srt/models/sarvam_moe.py

@@ -52,12 +52,12 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.model_executor.forward_context import (
     get_attn_backend,
     get_token_to_kv_pool,
 )
+from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.bailing_moe import BailingMoEForCausalLM
 from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mha import (