diff --git a/docs_new/docs/advanced_features/server_arguments.mdx b/docs_new/docs/advanced_features/server_arguments.mdx index e1c1d63c6..a111d1fd2 100644 --- a/docs_new/docs/advanced_features/server_arguments.mdx +++ b/docs_new/docs/advanced_features/server_arguments.mdx @@ -1990,22 +1990,52 @@ Please consult the documentation below and [server_args.py](https://github.com/s bool flag (set to enable) - --cuda-graph-max-bs - Set the maximum batch size for cuda graph. It will extend the cuda graph capture batch size to this value. + --cuda-graph-config + Canonical per-phase CUDA graph settings as JSON, e.g. {`{"decode":{"backend":"full","max_bs":256},"prefill":{"backend":"tc_piecewise","tc_compiler":"eager"}}`}. JSON wins over the per-phase --cuda-graph-* convenience flags and over the legacy flags. Allowed backends: full, breakable, tc_piecewise, disabled (full is decode-only). + `None` + Type: JSON (dict-of-dicts) + + + --cuda-graph-backend-decode + Backend for the decode phase. Folds into cuda_graph_config[decode].backend. + `None` + full, breakable, tc_piecewise, disabled + + + --cuda-graph-backend-prefill + Backend for the prefill phase. Folds into cuda_graph_config[prefill].backend. + `None` + breakable, tc_piecewise, disabled + + + --cuda-graph-max-bs-decode + Maximum batch size captured for the decode CUDA graph. `None` Type: int - --cuda-graph-bs - Set the list of batch sizes for cuda graph. + --cuda-graph-max-bs-prefill + Maximum batch size captured for the prefill CUDA graph. + `None` + Type: int + + + --cuda-graph-bs-decode + Explicit list of batch sizes to capture for the decode CUDA graph. `None` List[int] - --disable-cuda-graph - Disable cuda graph. - False - bool flag (set to enable) + --cuda-graph-bs-prefill + Explicit list of batch sizes to capture for the prefill CUDA graph. + `None` + List[int] + + + --cuda-graph-tc-compiler + Compiler used by the tc_piecewise backend (only the prefill phase consumes it today). + `None` + eager, inductor --disable-cuda-graph-padding @@ -2019,6 +2049,12 @@ Please consult the documentation below and [server_args.py](https://github.com/s False bool flag (set to enable) + + --debug-cuda-graph + Eager-mode CUDA graph via the breakable backend: graph breaks let every op run eagerly while still going through the capture/replay path. Useful for debugging capture/replay issues. + False + bool flag (set to enable) + --enable-cudagraph-gc Enable garbage collection during CUDA graph capture. If disabled (default), GC is frozen during capture to speed up the process. @@ -2139,39 +2175,87 @@ Please consult the documentation below and [server_args.py](https://github.com/s False bool flag (set to enable) - - --disable-piecewise-cuda-graph - Disable piecewise cuda graph for extend/prefill. PCG is enabled by default. - False - bool flag (set to disable) - - - --enforce-piecewise-cuda-graph - Enforce piecewise cuda graph, skipping all auto-disable conditions. For testing only. - False - bool flag (set to enable) - - - --piecewise-cuda-graph-tokens - Set the list of tokens when using piecewise cuda graph. - `None` - Type: JSON list - - - --piecewise-cuda-graph-compiler - Set the compiler for piecewise cuda graph. Choices are: eager, inductor. - eager - eager, inductor - --torch-compile-max-bs Set the maximum batch size when using torch compile. 32 Type: int + + --cuda-graph-max-bs + Deprecated alias for --cuda-graph-max-bs-decode. + `None` + Type: int + + + --cuda-graph-bs + Deprecated alias for --cuda-graph-bs-decode. + `None` + List[int] + + + --disable-cuda-graph + Deprecated. Use --cuda-graph-backend-decode=disabled and/or --cuda-graph-backend-prefill=disabled. + False + bool flag (set to enable) + + + --enable-breakable-cuda-graph + Deprecated alias for --cuda-graph-backend-prefill=breakable. + False + bool flag (set to enable) + + + --prefill-cuda-graph-backend + Deprecated alias for --cuda-graph-backend-prefill. + `None` + breakable, tc_piecewise, disabled + + + --decode-cuda-graph-backend + Deprecated alias for --cuda-graph-backend-decode. + `None` + full, breakable, tc_piecewise, disabled + + + --disable-prefill-cuda-graph + Deprecated. Use --cuda-graph-backend-prefill=disabled. + False + bool flag (set to enable) + + + --disable-decode-cuda-graph + Deprecated. Use --cuda-graph-backend-decode=disabled. + False + bool flag (set to enable) + + + --disable-piecewise-cuda-graph + Deprecated alias for --cuda-graph-backend-prefill=disabled. + False + bool flag (set to enable) + + + --enforce-piecewise-cuda-graph + Deprecated alias for --cuda-graph-backend-prefill=tc_piecewise. Explicitly setting the prefill backend now skips the auto-disable cascade automatically. + False + bool flag (set to enable) + + + --piecewise-cuda-graph-tokens + Deprecated alias for --cuda-graph-bs-prefill. + `None` + List[int] + + + --piecewise-cuda-graph-compiler + Deprecated alias for --cuda-graph-tc-compiler. + eager + eager, inductor + --piecewise-cuda-graph-max-tokens - Set the maximum tokens when using piecewise cuda graph. + Deprecated alias for --cuda-graph-max-bs-prefill. 4096 Type: int diff --git a/python/sglang/auto_benchmark_lib.py b/python/sglang/auto_benchmark_lib.py index 1b82c9117..fe6d314ef 100644 --- a/python/sglang/auto_benchmark_lib.py +++ b/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", diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index 57f3eb3e3..1cfb78b56 100644 --- a/python/sglang/bench_one_batch.py +++ b/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) diff --git a/python/sglang/compile_deep_gemm.py b/python/sglang/compile_deep_gemm.py index 7abc6993b..d1ef48ae3 100644 --- a/python/sglang/compile_deep_gemm.py +++ b/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...") diff --git a/python/sglang/srt/arg_groups/argparse_actions.py b/python/sglang/srt/arg_groups/argparse_actions.py index 5540a10a7..3fdd39e97 100644 --- a/python/sglang/srt/arg_groups/argparse_actions.py +++ b/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.""" diff --git a/python/sglang/srt/compilation/compile.py b/python/sglang/srt/compilation/compile.py index 448c1beb0..1e1526fa5 100644 --- a/python/sglang/srt/compilation/compile.py +++ b/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) diff --git a/python/sglang/srt/compilation/compile_phase.py b/python/sglang/srt/compilation/compile_phase.py new file mode 100644 index 000000000..42ef0ce41 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/compilation/cuda_piecewise_backend.py b/python/sglang/srt/compilation/cuda_piecewise_backend.py index 100aa4762..a750eb967 100644 --- a/python/sglang/srt/compilation/cuda_piecewise_backend.py +++ b/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: diff --git a/python/sglang/srt/compilation/piecewise_context_manager.py b/python/sglang/srt/compilation/piecewise_context_manager.py deleted file mode 100644 index 620fd29e6..000000000 --- a/python/sglang/srt/compilation/piecewise_context_manager.py +++ /dev/null @@ -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" -) diff --git a/python/sglang/srt/compilation/torch_compile_decoration.py b/python/sglang/srt/compilation/torch_compile_decoration.py new file mode 100644 index 000000000..9c396b27c --- /dev/null +++ b/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() diff --git a/python/sglang/srt/debug_utils/pr_fix_toggle.py b/python/sglang/srt/debug_utils/pr_fix_toggle.py index 3829a84ee..aecdca3e1 100644 --- a/python/sglang/srt/debug_utils/pr_fix_toggle.py +++ b/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 diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index 3459165b7..e98584634 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/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 ( diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index 454090547..a55a5e1e5 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py +++ b/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) diff --git a/python/sglang/srt/distributed/device_communicators/pymscclpp.py b/python/sglang/srt/distributed/device_communicators/pymscclpp.py index ae62c4d90..45395876f 100644 --- a/python/sglang/srt/distributed/device_communicators/pymscclpp.py +++ b/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 diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index a4cc81ee2..95920bc1b 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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. diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py index adcba8a85..4dd39762c 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py +++ b/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 diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py index b0bbd995a..4103cebb6 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py new file mode 100644 index 000000000..db5bed56b --- /dev/null +++ b/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 diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py index 1e1df9598..46daed9cc 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/hardware_backend/npu/utils.py b/python/sglang/srt/hardware_backend/npu/utils.py index 7139351fc..5ab80ee38 100644 --- a/python/sglang/srt/hardware_backend/npu/utils.py +++ b/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 diff --git a/python/sglang/srt/kv_canary/api.py b/python/sglang/srt/kv_canary/api.py index 2800c4614..143563f01 100644 --- a/python/sglang/srt/kv_canary/api.py +++ b/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() diff --git a/python/sglang/srt/kv_canary/capacities.py b/python/sglang/srt/kv_canary/capacities.py index 5b5f3696e..62f17981c 100644 --- a/python/sglang/srt/kv_canary/capacities.py +++ b/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}" diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index 27f3c297c..ffd4896c1 100644 --- a/python/sglang/srt/layers/activation.py +++ b/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() diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 85fcd4b9e..8ba5ddcb9 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/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" diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index fcf6e1b1a..13eda3172 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/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 ( diff --git a/python/sglang/srt/layers/attention/fla/layernorm_gated.py b/python/sglang/srt/layers/attention/fla/layernorm_gated.py index 5a8fda41b..798033634 100644 --- a/python/sglang/srt/layers/attention/fla/layernorm_gated.py +++ b/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) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index a507b6270..93efdd098 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/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. """ diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index c79ef416a..8c0993516 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/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) diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 823450d5e..82ca1903a 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/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( diff --git a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py index af3da62c3..006588129 100644 --- a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py +++ b/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] diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index ebe59ce05..f9311b772 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 8b4817207..f661ece72 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/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) diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index b428adafc..78a482c4c 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/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: diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 89be4a16e..d416e937e 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/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() diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 691401230..46391d5bb 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/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" diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 640e2b504..73832ba49 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/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) diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py index 31eecf8a9..44c5ddf8a 100644 --- a/python/sglang/srt/layers/moe/mega_moe.py +++ b/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: diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py index b78da24c9..56f90cd6d 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py +++ b/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. diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 6c6827c7d..13efcf236 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/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()) diff --git a/python/sglang/srt/layers/quantization/marlin_utils.py b/python/sglang/srt/layers/quantization/marlin_utils.py index 0cfc0298d..bb2e1159f 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils.py +++ b/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, diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index d5f770fc1..7eaf30533 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/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 +) diff --git a/python/sglang/srt/layers/radix_linear_attention.py b/python/sglang/srt/layers/radix_linear_attention.py index 019b981b1..6696945d2 100644 --- a/python/sglang/srt/layers/radix_linear_attention.py +++ b/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] diff --git a/python/sglang/srt/lora/lora_moe_runners.py b/python/sglang/srt/lora/lora_moe_runners.py index 9e9e1b3f2..f44bc486a 100644 --- a/python/sglang/srt/lora/lora_moe_runners.py +++ b/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() diff --git a/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py b/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py index 88a5efac8..c093b9090 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py +++ b/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 diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index a2d0960c7..ec7ce2add 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/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); diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index 6e8c1d4c9..0ff6dafc4 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index dc57f9bb1..dd68367af 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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 diff --git a/python/sglang/srt/model_executor/breakable_cuda_graph/__init__.py b/python/sglang/srt/model_executor/breakable_cuda_graph/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/python/sglang/srt/model_executor/breakable_cuda_graph_runner.py b/python/sglang/srt/model_executor/breakable_cuda_graph_runner.py deleted file mode 100644 index d5efcd421..000000000 --- a/python/sglang/srt/model_executor/breakable_cuda_graph_runner.py +++ /dev/null @@ -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." - ) diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index c8e5a199f..5bd12c276 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/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=}" diff --git a/python/sglang/srt/model_executor/cuda_graph_config.py b/python/sglang/srt/model_executor/cuda_graph_config.py new file mode 100644 index 000000000..db685d3c6 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/model_executor/forward_context.py b/python/sglang/srt/model_executor/forward_context.py index 3a3a7e50f..eb84508df 100644 --- a/python/sglang/srt/model_executor/forward_context.py +++ b/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 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 3b30eb0e1..f590e2935 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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, ) diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py deleted file mode 100644 index 6f8a102fa..000000000 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ /dev/null @@ -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 diff --git a/python/sglang/srt/model_executor/runner/__init__.py b/python/sglang/srt/model_executor/runner/__init__.py new file mode 100644 index 000000000..c830e2434 --- /dev/null +++ b/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, +) diff --git a/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py new file mode 100644 index 000000000..86316ae91 --- /dev/null +++ b/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: ... diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py similarity index 67% rename from python/sglang/srt/model_executor/cuda_graph_runner.py rename to python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index ebbfc2b28..aa14563b1 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -1,4 +1,4 @@ -# Copyright 2023-2024 SGLang Team +# 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 @@ -11,18 +11,23 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Run the model with cuda graph and torch.compile.""" +"""DecodeCudaGraphRunner — runs DECODE / TARGET_VERIFY / DLLM_EXTEND under +a pluggable backend. + +Backend selection comes from cuda_graph_config.decode: + - "full" — default, FullCudaGraphBackend: one + torch.cuda.CUDAGraph per shape. + - "breakable" — experimental, BreakableCudaGraphBackend: + segmented capture (no torch.compile). + - "tc_piecewise" — not implemented for decode; logs a one-shot warning + and falls back to "full". +""" from __future__ import annotations -import bisect import contextlib -import gc import inspect import logging -import os -from contextlib import contextmanager -from functools import partial from types import SimpleNamespace from typing import TYPE_CHECKING, Callable, Optional, Union @@ -30,14 +35,12 @@ import torch import tqdm from torch.profiler import ProfilerActivity, profile -from sglang.srt.batch_overlap.two_batch_overlap import TboCudaGraphRunnerPlugin -from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH -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.compilation.torch_compile_decoration import ( + patch_model, + set_torch_compile_config, ) +from sglang.srt.distributed import get_tensor_model_parallel_rank from sglang.srt.distributed.parallel_state import ( - GroupCoordinator, graph_capture, set_pdmux_status, ) @@ -46,18 +49,17 @@ from sglang.srt.environ import envs from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp 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.token_dispatcher.deepep import DeepEPBuffer -from sglang.srt.layers.moe.utils import get_deepep_mode, get_moe_a2a_backend -from sglang.srt.layers.utils import MultiPlatformOp from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled -from sglang.srt.model_executor.cuda_graph_buffer_registry import build_decode_registry +from sglang.srt.model_executor.cuda_graph_buffer_registry import ( + CudaGraphBufferRegistry, + build_decode_registry, +) from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, @@ -68,21 +70,38 @@ 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.input_buffers import share_input_buffers_in +from sglang.srt.model_executor.runner.base_cuda_graph_runner import ( + BaseCudaGraphRunner, + freeze_gc, + get_batch_sizes_to_capture, +) +from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import ( + BreakableCudaGraphBackend, +) +from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend +from sglang.srt.model_executor.runner_backend_utils import ( + CUDA_GRAPH_CAPTURE_FAILED_MSG, +) +from sglang.srt.model_executor.runner_utils.buffers import ( + DecodeInputBuffers, +) +from sglang.srt.model_executor.runner_utils.capture_mode import ( + _set_capture_lora_variant, + model_capture_mode, +) +from sglang.srt.model_executor.runner_utils.deepep_adapter import ( + DeepEPCudaGraphRunnerAdapter, +) from sglang.srt.multiplex.pdmux_context import get_current_stream_idx, get_stream_groups from sglang.srt.utils import ( empty_context, get_available_gpu_memory, - get_bool_env_var, - is_hip, log_info_on_rank0, require_attn_tp_gather, require_gathered_buffer, require_mlp_sync, require_mlp_tp_gather, ) -from sglang.srt.utils.patch_torch import monkey_patch_torch_compile -from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter try: from kt_kernel import KTMoEWrapper @@ -91,52 +110,49 @@ try: except ImportError: KTRANSFORMERS_AVAILABLE = False -_is_hip = is_hip() - -if not _is_hip: - from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import ( - BreakableCUDAGraph, - BreakableCUDAGraphCapture, - eager_on_graph, - ) - logger = logging.getLogger(__name__) if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner +def _make_graph_key(bs, stream_idx=None, variant_label=None): + """Build a graph dict key from batch size, stream index, and lora variant. + + Standalone function so speculative runners (which don't subclass + DecodeCudaGraphRunner) can use the same key encoding. + """ + key = bs if stream_idx is None else f"{stream_idx}_{bs}" + if variant_label is not None: + key = f"{variant_label}_{key}" + return key + + def build_replay_fb_view( - forward_batch: "ForwardBatch", - buffers, + forward_batch: ForwardBatch, + buffers: DecodeInputBuffers, bs: int, raw_bs: int, num_tokens: int, seq_len_fill_value: int, - capture_forward_mode: "ForwardMode", + capture_forward_mode: ForwardMode, is_encoder_decoder: bool, ) -> SimpleNamespace: """Construct a ForwardBatch-like view for backend replay-side init. - Combines the original ``forward_batch`` (for unpadded / per-iter - fields like ``spec_info``, ``out_cache_loc``, and the runtime - ``actual_forward_mode``) with the padded capture-time buffers from - ``buffers`` (for ``req_pool_indices``, ``seq_lens``, ``seq_lens_cpu``, - ``encoder_lens``). + Combines the original forward_batch (for unpadded / per-iter + fields like spec_info, out_cache_loc, and the runtime + actual_forward_mode) with the padded capture-time buffers from + buffers (for req_pool_indices, seq_lens, seq_lens_cpu, + encoder_lens). - Field semantics: + forward_mode is the capture-time mode (used by backends for + bucket / dispatch decisions); actual_forward_mode is the + runtime mode (may be IDLE while the captured graph targets DECODE + — DSV4's replay metadata prep uses this for IDLE substitution). - - ``forward_mode``: the capture-time mode (``capture_forward_mode``), - used by backends for bucket / dispatch decisions (e.g. choosing - between decode / target-verify / draft-extend code paths). - - ``actual_forward_mode``: the original runtime ``forward_batch - .forward_mode``, which may be ``IDLE`` even when the captured - graph corresponds to ``DECODE``. DSV4's replay metadata prep - uses this for IDLE-batch substitution; other backends ignore it. - - This view subsumes the ``_replay_forward_batch`` side channel DSV4 - previously read out-of-band — step 04 swaps that mechanism for this - explicit fb_view field. + Subsumes the _replay_forward_batch side channel that DSV4 used to + read out-of-band before the init_forward_metadata 3-method ABC. """ return SimpleNamespace( batch_size=bs, @@ -283,161 +299,13 @@ def _allocate_decode_buffers( ) -# Detect whether the current forward pass is in capture mode -is_capture_mode = False +class DecodeCudaGraphRunner(BaseCudaGraphRunner): + """Decode-phase CUDA graph runner. - -def get_is_capture_mode(): - return is_capture_mode - - -def compile_in_capture_mode(func): - if get_is_capture_mode(): - return torch.compile(func) - return func - - -@contextmanager -def model_capture_mode(): - global is_capture_mode - is_capture_mode = True - - yield - - is_capture_mode = False - - -@contextmanager -def freeze_gc(enable_cudagraph_gc: bool): + Owns: static input buffers (DecodeInputBuffers), capture-bs list, + attention backend, two-batch-overlap plugin, DeepEP adapter, and the + pluggable self.backend that handles the actual capture/replay. """ - 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() - gc.collect() - - -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, - enable_compile: bool, - num_tokens: int, - tp_group: GroupCoordinator, -): - """Patch the model to make it compatible with 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 - # Use custom-allreduce here. - # We found the custom allreduce is much faster than the built-in allreduce in torch, - # even with ENABLE_INTRA_NODE_COMM=1. - # tp_group.ca_comm = None - 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(): - 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 # Experimental feature to reduce compilation times, will be on by default in future - - # FIXME: tmp workaround - 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() - - -def get_batch_sizes_to_capture(model_runner: ModelRunner, num_tokens_per_bs=1): - server_args = model_runner.server_args - capture_bs = server_args.cuda_graph_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 # tbo not test, set num_tokens_per_bs to 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() - - # pad `num_max_requests` to avoid being filtered out - num_max_requests = (num_max_requests + mul_base - 1) // mul_base * mul_base - if max(capture_bs) > num_max_requests: - # In some cases (e.g., with a small GPU or --max-running-requests), the #max-running-requests - # is very small. We add more values here to make sure we capture the maximum bs. - capture_bs += [num_max_requests] - - # Model input token count = bs * num_tokens_per_bs; must be a multiple of attn_tp_size. - 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 - - -# 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 - - -class CudaGraphRunner: - """A CudaGraphRunner runs the forward pass of a model with cuda graph and torch.compile.""" def __init__( self, @@ -447,12 +315,8 @@ class CudaGraphRunner: speculative_num_steps: Optional[int] = None, speculative_num_draft_tokens: Optional[int] = None, ): - # 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 = {} + super().__init__(model_runner) + # --- core state ------------------------------------------------ self.enable_torch_compile = model_runner.server_args.enable_torch_compile self.disable_padding = model_runner.server_args.disable_cuda_graph_padding self.is_encoder_decoder = model_runner.model_config.is_encoder_decoder @@ -472,9 +336,6 @@ class CudaGraphRunner: self.enable_profile_cuda_graph = ( model_runner.server_args.enable_profile_cuda_graph ) - 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.enable_pdmux = model_runner.server_args.enable_pdmux self.attn_tp_size = get_attention_tp_size() @@ -505,15 +366,13 @@ class CudaGraphRunner: else speculative_num_draft_tokens ) + # --- capture mode + tokens-per-bs ------------------------------ self.capture_forward_mode = ForwardMode.DECODE self.capture_hidden_mode = CaptureHiddenMode.NULL self.num_tokens_per_bs = 1 if model_runner.spec_algorithm.is_speculative(): if self.model_runner.is_draft_worker: - # Draft workers can use TARGET_VERIFY mode. - if ( - not self.model_runner.spec_algorithm.supports_target_verify_for_draft() - ): + if not self.model_runner.spec_algorithm.is_dflash(): raise RuntimeError("This should not happen") self.capture_forward_mode = ForwardMode.TARGET_VERIFY self.num_tokens_per_bs = ( @@ -525,7 +384,7 @@ class CudaGraphRunner: self.capture_forward_mode = ForwardMode.DLLM_EXTEND self.num_tokens_per_bs = self.dllm_config.block_size - # Batch sizes to capture + # --- bucket sizes --------------------------------------------- self.capture_bs, self.compile_bs = get_batch_sizes_to_capture( model_runner, self.num_tokens_per_bs ) @@ -533,16 +392,13 @@ class CudaGraphRunner: if KTRANSFORMERS_AVAILABLE: KTMoEWrapper.set_capture_batch_sizes(self.capture_bs) - # If returning hidden states is enabled, set initial capture hidden mode to full to avoid double-capture on startup if model_runner.server_args.enable_return_hidden_states: self.capture_hidden_mode = CaptureHiddenMode.FULL - # Attention backend self.max_bs = max(self.capture_bs) self.max_num_token = self.max_bs * self.num_tokens_per_bs self.attn_backend.init_cuda_graph_state(self.max_bs, self.max_num_token) - # Init PDMux if needed self.maybe_init_pdmux() self.seq_len_fill_value = ( self.attn_backend.get_cuda_graph_seq_len_fill_value() @@ -561,9 +417,6 @@ class CudaGraphRunner: set_torch_compile_config() if self.model_runner.server_args.enable_lora: - # Phase 2 of LoRA CUDA graph init: dense LoRA batch metadata. - # Phase 1 (MoE buffers) was handled earlier in ModelRunner via - # lora_manager.init_cuda_graph_moe_buffers(). self.model_runner.lora_manager.init_cuda_graph_batch_info( max_bs_in_cuda_graph=self.max_bs, num_tokens_per_bs=self.num_tokens_per_bs, @@ -576,7 +429,9 @@ class CudaGraphRunner: if self.require_gathered_buffer: assert self.require_mlp_tp_gather or self.require_attn_tp_gather - self.buffers = _allocate_decode_buffers( + + # --- buffers --------------------------------------------------- + self.buffers: DecodeInputBuffers = DecodeInputBuffers.create( device=self.device, max_bs=self.max_bs, max_num_token=self.max_num_token, @@ -599,9 +454,12 @@ class CudaGraphRunner: self.model_runner.model_config, "hc_hidden_size", None ), ) - share_input_buffers_in(self.buffers) - # The registry adopts these buffers (one data_ptr for capture + replay). - self.buffer_registry = build_decode_registry( + self.buffers.share_buffers() + # FB-shared slot registry adopting DecodeInputBuffers storage (same + # physical tensors, stable data_ptr for capture vs replay). Provides + # the unified fill_from / slot access surface, replacing + # populate_from_forward_batch on capture/replay paths. + self.buffer_registry: CudaGraphBufferRegistry = build_decode_registry( device=self.device, max_bs=self.max_bs, max_num_token=self.max_num_token, @@ -618,17 +476,21 @@ class CudaGraphRunner: source=self.buffers, ) - self.tbo_plugin = TboCudaGraphRunnerPlugin() + # --- backend --------------------------------------------------- + self.backend = resolve_decode_backend(self) - # Capture + # --- capture -------------------------------------------------- try: with model_capture_mode(): self.capture() except RuntimeError as e: raise Exception( - f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}" + f"Capture cuda graph failed: {e}\n" f"{CUDA_GRAPH_CAPTURE_FAILED_MSG}" ) + # ----------------------------------------------------------------- + # Helpers + # ----------------------------------------------------------------- def maybe_init_pdmux(self): if self.enable_pdmux: self.stream_groups = get_stream_groups() @@ -638,8 +500,22 @@ class CudaGraphRunner: def _cache_loc_dtype(self): return torch.int64 + def _make_graph_key(self, bs, stream_idx=None, variant_label=None): + return _make_graph_key(bs, stream_idx, variant_label) + + def _resolve_lora_variant(self, forward_batch: ForwardBatch): + if not getattr(self, "record_nolora_graph", False): + return None + if forward_batch.lora_ids is not None and any( + uid is not None for uid in forward_batch.lora_ids + ): + return "lora" + return "nolora" + + # ----------------------------------------------------------------- + # can_run + # ----------------------------------------------------------------- def can_run(self, forward_batch: ForwardBatch): - # Disable for token embedding overrides (dynamic per-request) if forward_batch.replace_embeds is not None: return False if self.require_mlp_tp_gather: @@ -658,7 +534,7 @@ class CudaGraphRunner: graph_key = f"{get_current_stream_idx()}_{cuda_graph_bs}" is_bs_supported = ( - graph_key in self.graphs + self.backend.can_run(forward_batch, graph_key) if self.disable_padding else cuda_graph_bs <= self.max_bs ) @@ -666,9 +542,6 @@ class CudaGraphRunner: if self.require_mlp_sync: is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph - # NOTE: cuda graph cannot handle mixed batch (encoder_len = 0) - # If mixed batch cannot be supported, then encoder_lens can be removed in cuda graph - # because the full_text_row_masked_out_mask tensor will always be ones is_encoder_lens_supported = ( torch.all(forward_batch.encoder_lens > 0) if self.is_encoder_decoder @@ -709,6 +582,9 @@ class CudaGraphRunner: and is_ngram_supported ) + # ----------------------------------------------------------------- + # Profiling helpers + # ----------------------------------------------------------------- def _init_profile_context_and_memory_record(self): profile_context = profile( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], @@ -718,7 +594,7 @@ class CudaGraphRunner: return profile_context def _post_process_after_profile(self, prof_context): - torch.cuda.memory._dump_snapshot(f"cuda_graph_runner_memory_usage.pickle") + torch.cuda.memory._dump_snapshot("cuda_graph_runner_memory_usage.pickle") torch.cuda.memory._record_memory_history(enabled=None) log_message = ( "Sorted by CUDA Time:\n" @@ -733,119 +609,28 @@ class CudaGraphRunner: ) logger.info(log_message) - def capture(self) -> None: - profile_context = empty_context() - if self.enable_profile_cuda_graph: - profile_context = self._init_profile_context_and_memory_record() - - def _capture_one_stream(stream_idx: Optional[int] = None): - 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_bs))) - if get_tensor_model_parallel_rank() == 0 - else reversed(self.capture_bs) - ) - for i, bs 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 batches ({bs=} {avail_mem=:.2f} GB)" - ) - - with patch_model( - self.model_runner.model, - bs in self.compile_bs, - num_tokens=bs * self.num_tokens_per_bs, - tp_group=self.model_runner.tp_group, - ) as forward: - ( - graph, - output_buffers, - ) = self.capture_one_batch_size(bs, forward, stream_idx) - # For pd_multiplexing, we need to save the graph and output buffers - key = bs if stream_idx is None else f"{stream_idx}_{bs}" - self.graphs[key] = graph - self.output_buffers[key] = output_buffers - - # 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): - if not self.enable_pdmux: - with graph_capture() as graph_capture_context, profile_context as prof: - self.stream = graph_capture_context.stream - _capture_one_stream() - else: - set_pdmux_status(False) - for i, sg in enumerate(self.stream_groups): - with ( - graph_capture(stream=sg[1]) as graph_capture_context, - profile_context as prof, - ): - self.stream = graph_capture_context.stream - _capture_one_stream(i) - - if self.enable_profile_cuda_graph: - self._post_process_after_profile(prof) - - def _capture_graph(self, graph, pool, stream, run_once_fn): - if self.model_runner.server_args.debug_cuda_graph: - assert ( - envs.SGLANG_USE_BREAKABLE_CUDA_GRAPH.get() - ), "Breakable CUDA graph is not enabled in debug mode" - - memory_saver_adapter = TorchMemorySaverAdapter.create( - enable=self.model_runner.server_args.enable_memory_saver - and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH") - ) - - if envs.SGLANG_USE_BREAKABLE_CUDA_GRAPH.get(): - if memory_saver_adapter.enabled: - raise NotImplementedError( - "Breakable CUDA graph is not compatible with memory saver mode" - ) - graph_ctx = BreakableCUDAGraphCapture - else: - graph_ctx = ( - partial(memory_saver_adapter.cuda_graph, tag=GPU_MEMORY_TYPE_CUDA_GRAPH) - if memory_saver_adapter.enabled - else self.device_module.graph - ) - - if self.model_runner.server_args.debug_cuda_graph: - captured_fn = eager_on_graph(True)(run_once_fn) - else: - captured_fn = run_once_fn - - with graph_ctx(cuda_graph=graph, pool=pool, stream=stream): - out = captured_fn() - return out - - def _create_device_graph(self): - if envs.SGLANG_USE_BREAKABLE_CUDA_GRAPH.get(): - if _is_hip: - raise RuntimeError("Breakable CUDA graph is not supported on ROCm/HIP") - return BreakableCUDAGraph() - return torch.cuda.CUDAGraph() - - def capture_one_batch_size( - self, bs: int, forward: Callable, stream_idx: Optional[int] = None + # ----------------------------------------------------------------- + # capture_prepare + # ----------------------------------------------------------------- + def capture_prepare( + self, + size: int, + stream_idx: Optional[int] = None, ): - buffers = self.buffers - graph = self._create_device_graph() - stream = self.stream + """Build the dummy decode ForwardBatch for capture at size (=bs), + populate static input buffers, choose the active attn backend, and + optionally build pp_proxy_tensors. + + Returns (forward_batch, attn_backend, pp_proxy_tensors); + pp_proxy_tensors is None unless pp_size > 1. + """ + bs = size + buffers: DecodeInputBuffers = self.buffers num_tokens = bs * self.num_tokens_per_bs - # Graph inputs: owned slots come from the registry; the rest off `buffers`. + # Registry-owned FB-shared slots come through the registry (which + # shares physical storage with self.buffers via source=...); the rest + # still come off buffers directly. registry = self.buffer_registry def _slot(name): @@ -869,8 +654,6 @@ class CudaGraphRunner: else None ) - # Adjust for attention TP if needed (matching replay path in - # populate_from_forward_batch). buffers.num_token_non_padded[...] = num_tokens if ( enable_num_token_non_padded() @@ -883,7 +666,7 @@ class CudaGraphRunner: ) buffers.num_token_non_padded.copy_(local) - # pipeline parallelism + pp_proxy_tensors = None if self.pp_size > 1: pp_proxy_tensors = PPProxyTensors( {k: v[:num_tokens] for k, v in buffers.pp_proxy_tensors.items()} @@ -913,13 +696,10 @@ class CudaGraphRunner: ) 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 - # mamba state tracking (registry-owned when enabled) mamba_track_indices = ( _slot("mamba_track_indices") if registry.has_slot("mamba_track_indices") @@ -948,7 +728,7 @@ class CudaGraphRunner: seq_lens_sum=seq_lens.sum().item(), mamba_track_indices=mamba_track_indices, mamba_track_mask=mamba_track_mask, - mamba_track_seqlens=None, # Prefill only + mamba_track_seqlens=None, encoder_lens=encoder_lens, return_logprob=False, positions=positions, @@ -956,6 +736,7 @@ class CudaGraphRunner: global_num_tokens_for_logprob_gpu=buffers.global_num_tokens_for_logprob_gpu, dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(), global_dp_buffer_len=global_dp_buffer_len, + global_num_tokens_cpu=global_num_tokens_cpu, mrope_positions=mrope_positions, spec_algorithm=self.model_runner.spec_algorithm, spec_info=spec_info, @@ -967,22 +748,111 @@ class CudaGraphRunner: bootstrap_room_ids_int=bootstrap_room_ids_int, ) - # Trip the coordinator so the hisparse code path is captured into the - # graph; backends read it from self.model_runner.hisparse_coordinator. - hisparse_coordinator = self.model_runner.hisparse_coordinator - if hisparse_coordinator is not None: - hisparse_coordinator.num_real_reqs.fill_(bs) + forward_batch.hisparse_coordinator = self.model_runner.hisparse_coordinator + if forward_batch.hisparse_coordinator is not None: + forward_batch.hisparse_coordinator.num_real_reqs.fill_(bs) if buffers.ngram_embedding_info is not None: forward_batch.ngram_embedding_info = buffers.ngram_embedding_info.slice(bs) + return forward_batch, attn_backend, pp_proxy_tensors + + # ----------------------------------------------------------------- + # capture + # ----------------------------------------------------------------- + def capture(self) -> None: + profile_context = empty_context() + if self.enable_profile_cuda_graph: + profile_context = self._init_profile_context_and_memory_record() + + with freeze_gc(self.model_runner.server_args.enable_cudagraph_gc): + if not self.enable_pdmux: + with graph_capture() as graph_capture_context, profile_context as prof: + self.stream = graph_capture_context.stream + with self.backend.capture_session(self.stream): + self._capture_one_stream() + else: + set_pdmux_status(False) + for i, sg in enumerate(self.stream_groups): + with ( + graph_capture(stream=sg[1]) as graph_capture_context, + profile_context as prof, + ): + self.stream = graph_capture_context.stream + with self.backend.capture_session(self.stream): + self._capture_one_stream(i) + + if self.enable_profile_cuda_graph: + self._post_process_after_profile(prof) + + def _capture_one_stream(self, stream_idx: Optional[int] = None) -> None: + avail_mem = get_available_gpu_memory( + self.model_runner.device, + self.model_runner.gpu_id, + empty_cache=False, + ) + # Reverse so cuda graphs share memory better. + capture_range = ( + tqdm.tqdm(list(reversed(self.capture_bs))) + if get_tensor_model_parallel_rank() == 0 + else reversed(self.capture_bs) + ) + lora_variants = ( + [("lora", True), ("nolora", False)] + if getattr(self, "record_nolora_graph", False) + else [(None, None)] + ) + for bs 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 batches ({bs=} {avail_mem=:.2f} GB)" + ) + + for variant_label, _variant_has_lora in lora_variants: + _set_capture_lora_variant(variant_label) + with patch_model( + self.model_runner.model, + bs in self.compile_bs, + num_tokens=bs * self.num_tokens_per_bs, + tp_group=self.model_runner.tp_group, + ) as forward: + self.capture_one_shape(bs, forward, stream_idx, variant_label) + + # ----------------------------------------------------------------- + # capture_one_shape + # ----------------------------------------------------------------- + def capture_one_shape( + self, + size: int, + forward: Callable, + stream_idx: Optional[int] = None, + variant_label: Optional[str] = None, + ): + bs = size + num_tokens = bs * self.num_tokens_per_bs + + # Sanity-check: --debug-cuda-graph requires breakable backend. + if self.model_runner.server_args.debug_cuda_graph: + assert isinstance( + self.backend, BreakableCudaGraphBackend + ), "Breakable CUDA graph is required for --debug-cuda-graph" + + forward_batch, attn_backend, pp_proxy_tensors = self.capture_prepare( + size, stream_idx=stream_idx + ) + # All setup hooks below read get_attn_backend() (TboForwardBatchPreparer, # DeepEP adapter, …) so they must run inside the same ForwardContext # that wraps the warmup/capture forward. with forward_context(ForwardContext(attn_backend=attn_backend)): self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens) - if lora_ids is not None: + if forward_batch.lora_ids is not None: self.model_runner.lora_manager.prepare_lora_batch(forward_batch) attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True) @@ -996,10 +866,10 @@ class CudaGraphRunner: None ) set_dp_buffer_len( - global_dp_buffer_len, + forward_batch.global_dp_buffer_len, num_tokens, forward_batch.dp_padding_mode.is_max_len(), - global_num_tokens_cpu, + forward_batch.global_num_tokens_cpu, ) set_is_extend_in_batch(False) @@ -1016,46 +886,38 @@ class CudaGraphRunner: and self.model_runner.is_draft_worker and "input_embeds" in inspect.signature(forward).parameters ): - kwargs["input_embeds"] = buffers.input_embeds[:num_tokens] + kwargs["input_embeds"] = self.buffers.input_embeds[:num_tokens] - logits_output_or_pp_proxy_tensors = forward( - input_ids, + return forward( + forward_batch.input_ids, forward_batch.positions, forward_batch, **kwargs, ) - return logits_output_or_pp_proxy_tensors self.deepep_adapter.capture(is_extend_in_batch=False) - canary_ctx = ( c.with_active_single_forward_manager(0) if (c := self.model_runner.canary_manager) is not None else contextlib.nullcontext() ) with canary_ctx: - for _ in range(2): - self.device_module.synchronize() - self.model_runner.tp_group.barrier() - run_once() - attn_backend.on_after_cuda_graph_warmup() - - 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()) - - out = self._capture_graph( - graph, get_global_graph_memory_pool(), stream, run_once + shape_key = self._make_graph_key(bs, stream_idx, variant_label) + self.backend.capture_one( + shape_key, + run_once, + dummies=None, + post_warmup_hook=getattr( + self.model_runner.attn_backend, + "on_after_cuda_graph_warmup", + None, + ), ) - return graph, out - + # ----------------------------------------------------------------- + # recapture + # ----------------------------------------------------------------- def recapture_if_needed(self, forward_batch: ForwardBatch): - - # If the required capture_hidden_mode changes, we need to recapture the graph - - # These are the different factors that can influence the capture_hidden_mode capture_hidden_mode_required_by_forward_batch = ( forward_batch.capture_hidden_mode ) @@ -1069,32 +931,51 @@ class CudaGraphRunner: else CaptureHiddenMode.NULL ) - # Determine the highest capture_hidden_mode required - # (If we have FULL, we can emulate LAST or NULL) - # (If we have LAST, we can emulate NULL) required_capture_hidden_mode = max( capture_hidden_mode_required_by_forward_batch, capture_hidden_mode_required_by_spec_info, capture_hidden_mode_required_for_returning_hidden_states, ) - # If the current hidden mode is no longer aligned with the required hidden mode, we need to set it to what is required and re-capture if self.capture_hidden_mode != required_capture_hidden_mode: self.capture_hidden_mode = required_capture_hidden_mode + self.backend.cleanup() self.capture() + # ----------------------------------------------------------------- + # replay_prepare + # ----------------------------------------------------------------- def replay_prepare( self, forward_batch: ForwardBatch, pp_proxy_tensors: Optional[PPProxyTensors] = None, ): + self.deepep_adapter.replay() + + if not forward_batch.needs_forward_metadata_init(): + self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids) + self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions) + if ( + self.model_runner.spec_algorithm.is_dflash() + and self.model_runner.is_draft_worker + and forward_batch.input_embeds is not None + ): + self.buffers.input_embeds[: self.raw_num_token].copy_( + forward_batch.input_embeds + ) + variant_label = self._resolve_lora_variant(forward_batch) + stream_idx = get_current_stream_idx() if self.enable_pdmux else None + self._replay_graph_key = self._make_graph_key( + self.bs, stream_idx, variant_label + ) + return + buffers = self.buffers self.recapture_if_needed(forward_batch) raw_bs = forward_batch.batch_size raw_num_token = raw_bs * self.num_tokens_per_bs - # Pad if self.require_mlp_tp_gather: max_num_tokens = max(forward_batch.global_num_tokens_cpu) max_batch_size = ( @@ -1104,10 +985,9 @@ class CudaGraphRunner: or self.model_runner.spec_algorithm.is_dflash() else max_num_tokens ) - index = bisect.bisect_left(self.capture_bs, max_batch_size) + bs = self._pad_to_bucket(int(max_batch_size), self.capture_bs) else: - index = bisect.bisect_left(self.capture_bs, raw_bs) - bs = self.capture_bs[index] + bs = self._pad_to_bucket(raw_bs, self.capture_bs) self.buffer_registry.fill_from( forward_batch, @@ -1124,7 +1004,6 @@ class CudaGraphRunner: and forward_batch.input_embeds is not None ): buffers.input_embeds[:raw_num_token].copy_(forward_batch.input_embeds) - # Padded tokens aren't read, so skip zeroing them. if self.enable_two_batch_overlap: self.tbo_plugin.replay_prepare( forward_mode=self.capture_forward_mode, @@ -1134,7 +1013,6 @@ class CudaGraphRunner: ) if forward_batch.forward_mode.is_idle() and forward_batch.spec_info is not None: forward_batch.spec_info.custom_mask = buffers.custom_mask - # Attention backend if self.enable_pdmux: stream_idx = get_current_stream_idx() attn_backend = self.model_runner.decode_attn_backend_group[stream_idx] @@ -1152,7 +1030,6 @@ class CudaGraphRunner: ) attn_backend.init_forward_metadata_out_graph(fb_view) - # Store fields self.raw_bs = raw_bs self.raw_num_token = raw_num_token self.bs = bs @@ -1160,47 +1037,30 @@ class CudaGraphRunner: if self.model_runner.hisparse_coordinator is not None: self.model_runner.hisparse_coordinator.num_real_reqs.fill_(raw_bs) + variant_label = self._resolve_lora_variant(forward_batch) + stream_idx = get_current_stream_idx() if self.enable_pdmux else None + self._replay_graph_key = self._make_graph_key( + self.bs, stream_idx, variant_label + ) + + # ----------------------------------------------------------------- + # replay + # ----------------------------------------------------------------- def replay( self, forward_batch: ForwardBatch, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> Union[LogitsProcessorOutput, PPProxyTensors]: - self.deepep_adapter.replay() - - if forward_batch.needs_forward_metadata_init(): - self.replay_prepare(forward_batch, pp_proxy_tensors) - else: - # Pre-planned (plan-stream replay_prepare already ran). - # In speculative decoding, these two fields are still needed. - self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids) - self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions) - if ( - self.model_runner.spec_algorithm.is_dflash() - and self.model_runner.is_draft_worker - and forward_batch.input_embeds is not None - ): - self.buffers.input_embeds[: self.raw_num_token].copy_( - forward_batch.input_embeds - ) - - # Replay - if self.enable_pdmux: - graph_key = f"{get_current_stream_idx()}_{self.bs}" - else: - graph_key = self.bs - ctx = ( + timer_ctx = ( self.model_runner.device_timer.wrap( - metadata={ - "category": forward_batch.forward_mode.name.lower(), - } + metadata={"category": forward_batch.forward_mode.name.lower()} ) if self.model_runner.device_timer else contextlib.nullcontext() ) - with ctx: - self.graphs[graph_key].replay() - - output = self.output_buffers[graph_key] + with timer_ctx, self.backend.replay_session(): + self.replay_prepare(forward_batch, pp_proxy_tensors) + output = self.backend.replay(self._replay_graph_key, forward_batch) if isinstance(output, LogitsProcessorOutput): if self.is_dllm: @@ -1232,6 +1092,9 @@ class CudaGraphRunner: assert isinstance(output, PPProxyTensors) return PPProxyTensors({k: v[: self.bs] for k, v in output.tensors.items()}) + # ----------------------------------------------------------------- + # spec info + # ----------------------------------------------------------------- def get_spec_info(self, num_tokens: int): spec_info = None if ( @@ -1270,8 +1133,6 @@ class CudaGraphRunner: resolve_dflash_verify_mask_policy, ) - # Avoid enabling custom-mask modes during graph capture for backends that - # can express DFLASH verify via their built-in causal path. _, build_custom_mask = resolve_dflash_verify_mask_policy( self.model_runner.attn_backend ) @@ -1306,33 +1167,3 @@ class CudaGraphRunner: spec_info.capture_hidden_mode = CaptureHiddenMode.NULL return spec_info - - -CUDA_GRAPH_CAPTURE_FAILED_MSG = ( - "Possible solutions:\n" - "1. set --mem-fraction-static to a smaller value (e.g., 0.8 or 0.7)\n" - "2. set --cuda-graph-max-bs to a smaller value (e.g., 16)\n" - "3. disable torch compile by not using --enable-torch-compile\n" - "4. disable CUDA graph by --disable-cuda-graph. (Not recommended. Huge performance loss)\n" - "Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n" -) - - -class DeepEPCudaGraphRunnerAdapter: - def __init__(self): - # Record DeepEP mode used during capture to ensure replay consistency - self._captured_deepep_mode = None - - def capture(self, is_extend_in_batch: bool): - 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): - 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) diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py new file mode 100644 index 000000000..e6d4d7276 --- /dev/null +++ b/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." + ) diff --git a/python/sglang/srt/model_executor/runner_backend/__init__.py b/python/sglang/srt/model_executor/runner_backend/__init__.py new file mode 100644 index 000000000..b50f02053 --- /dev/null +++ b/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, +) diff --git a/python/sglang/srt/model_executor/runner_backend/base_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/base_cuda_graph_backend.py new file mode 100644 index 000000000..0f0807576 --- /dev/null +++ b/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: ... diff --git a/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py new file mode 100644 index 000000000..51ad47aa2 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py new file mode 100644 index 000000000..c1b919731 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py new file mode 100644 index 000000000..d543bbb82 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/model_executor/runner_backend/utils.py b/python/sglang/srt/model_executor/runner_backend/utils.py new file mode 100644 index 000000000..d3a84266d --- /dev/null +++ b/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) diff --git a/python/sglang/srt/model_executor/runner_backend_utils/__init__.py b/python/sglang/srt/model_executor/runner_backend_utils/__init__.py new file mode 100644 index 000000000..505f436af --- /dev/null +++ b/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" +) diff --git a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/__init__.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/__init__.py new file mode 100644 index 000000000..009e49139 --- /dev/null +++ b/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, +) diff --git a/python/sglang/srt/model_executor/breakable_cuda_graph/breakable_cuda_graph.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py similarity index 92% rename from python/sglang/srt/model_executor/breakable_cuda_graph/breakable_cuda_graph.py rename to python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py index 7341301ef..114a5df6e 100644 --- a/python/sglang/srt/model_executor/breakable_cuda_graph/breakable_cuda_graph.py +++ b/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. """ diff --git a/python/sglang/srt/model_executor/breakable_cuda_graph/context.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/context.py similarity index 80% rename from python/sglang/srt/model_executor/breakable_cuda_graph/context.py rename to python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/context.py index 216d93ab8..a58713de1 100644 --- a/python/sglang/srt/model_executor/breakable_cuda_graph/context.py +++ b/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 diff --git a/python/sglang/srt/model_executor/breakable_cuda_graph/cuda_utils.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py similarity index 100% rename from python/sglang/srt/model_executor/breakable_cuda_graph/cuda_utils.py rename to python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py diff --git a/python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/__init__.py b/python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/__init__.py new file mode 100644 index 000000000..ed5beb9a8 --- /dev/null +++ b/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, +) diff --git a/python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py b/python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py new file mode 100644 index 000000000..59e699d8c --- /dev/null +++ b/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" +) diff --git a/python/sglang/srt/model_executor/runner_utils/__init__.py b/python/sglang/srt/model_executor/runner_utils/__init__.py new file mode 100644 index 000000000..278c61cae --- /dev/null +++ b/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, +) diff --git a/python/sglang/srt/model_executor/runner_utils/buffers.py b/python/sglang/srt/model_executor/runner_utils/buffers.py new file mode 100644 index 000000000..5bf76d3d5 --- /dev/null +++ b/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 + ) diff --git a/python/sglang/srt/model_executor/runner_utils/capture_mode.py b/python/sglang/srt/model_executor/runner_utils/capture_mode.py new file mode 100644 index 000000000..e1a2ff7bc --- /dev/null +++ b/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 diff --git a/python/sglang/srt/model_executor/runner_utils/deepep_adapter.py b/python/sglang/srt/model_executor/runner_utils/deepep_adapter.py new file mode 100644 index 000000000..f9d5fce35 --- /dev/null +++ b/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) diff --git a/python/sglang/srt/model_executor/runner_utils/pool.py b/python/sglang/srt/model_executor/runner_utils/pool.py new file mode 100644 index 000000000..cc3904f27 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index 6daafecf4..d0d066647 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/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, diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py index c0c6be3ca..c1bc85503 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/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 diff --git a/python/sglang/srt/models/cohere2_moe.py b/python/sglang/srt/models/cohere2_moe.py index 8326c35ef..e1d45e741 100644 --- a/python/sglang/srt/models/cohere2_moe.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py index 6dcf9bc79..6444aebae 100644 --- a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index d9fcbf807..14f164b26 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/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) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index c34681ced..b0c8bb870 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/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: diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 7494ce85a..deff2366b 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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: diff --git a/python/sglang/srt/models/exaone_moe.py b/python/sglang/srt/models/exaone_moe.py index ff0a02099..0eae269bc 100755 --- a/python/sglang/srt/models/exaone_moe.py +++ b/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 diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index 7ca23d6aa..ac6bffe8d 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/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 diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py index 55f1f7a5e..a36d4ee86 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/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, diff --git a/python/sglang/srt/models/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py index 2cb94e96a..f6010bcbb 100644 --- a/python/sglang/srt/models/gpt_oss.py +++ b/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) diff --git a/python/sglang/srt/models/grok.py b/python/sglang/srt/models/grok.py index 408811d71..a5eb64abc 100644 --- a/python/sglang/srt/models/grok.py +++ b/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 diff --git a/python/sglang/srt/models/hunyuan_v3.py b/python/sglang/srt/models/hunyuan_v3.py index d44e3ef21..f15a18cff 100644 --- a/python/sglang/srt/models/hunyuan_v3.py +++ b/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 diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index a96265baa..139f311cc 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/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, diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 7daf233d9..17a47641a 100644 --- a/python/sglang/srt/models/llada2.py +++ b/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, diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 9d4b90b45..a0a4f1e17 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/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: diff --git a/python/sglang/srt/models/mllama.py b/python/sglang/srt/models/mllama.py index 771ef54ac..9fca0b414 100644 --- a/python/sglang/srt/models/mllama.py +++ b/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) diff --git a/python/sglang/srt/models/moss_vl.py b/python/sglang/srt/models/moss_vl.py index 3a47b58c4..fcab2fe10 100644 --- a/python/sglang/srt/models/moss_vl.py +++ b/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 diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 3d50a20cf..e4ef5d579 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/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] diff --git a/python/sglang/srt/models/olmo2.py b/python/sglang/srt/models/olmo2.py index 36a7810fa..91dfd12cd 100644 --- a/python/sglang/srt/models/olmo2.py +++ b/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). """ diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 37552404b..a2147e750 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/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: diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index ce8c98d5b..14fb21dfc 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/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") diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 65f78e691..3d7cf345b 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/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: diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index 5b79a6b9f..f8fc8be00 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/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: diff --git a/python/sglang/srt/models/sarvam_moe.py b/python/sglang/srt/models/sarvam_moe.py index 83683933c..ab93eb3d9 100644 --- a/python/sglang/srt/models/sarvam_moe.py +++ b/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 ( diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index fa77d4afa..b7b08b5ce 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -29,9 +29,9 @@ from sglang.srt.environ import envs from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool -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_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.server_args import get_global_server_args from sglang.srt.utils import get_current_device_stream_fast, is_cuda, is_hip @@ -424,9 +424,10 @@ def _reshape_for_qk_norm(x: torch.Tensor, head_dim: int) -> torch.Tensor: inputs and fault on strided tensors (root cause of the #21734 revert in #23159). """ + if ( _is_cuda - and get_global_server_args().piecewise_cuda_graph_compiler == "inductor" + and get_global_server_args().cuda_graph_config.prefill.tc_compiler == "inductor" ): return x.view(*x.shape[:-1], -1, head_dim) return x.reshape(-1, head_dim) @@ -461,12 +462,13 @@ def apply_qk_norm( batch_size = q.size(0) q_eps = q_norm.variance_epsilon k_eps = k_norm.variance_epsilon + if ( _is_cuda # TODO(dark): have not tested on ROCm or other backends and allow_inplace # TODO(dark): this can be relaxed if needed and (q_eps == k_eps) # TODO(dark): this can also be relaxed and not envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get() - and get_global_server_args().piecewise_cuda_graph_compiler + and get_global_server_args().cuda_graph_config.prefill.tc_compiler != "inductor" # let inductor fuse QK norm and can_use_fused_inplace_qknorm(head_dim, q.dtype) ): diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 6c77ff64f..03f81a6b8 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -34,6 +34,7 @@ from sglang.jit_kernel.kv_canary.consts import RealKvHashMode from sglang.srt.arg_groups.argparse_actions import ( DeprecatedAction, DeprecatedAliasStoreAction, + DeprecatedStoreConstAction, DeprecatedStoreTrueAction, LoRAPathAction, ) @@ -48,6 +49,14 @@ from sglang.srt.environ import envs from sglang.srt.function_call.function_call_parser import FunctionCallParser from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.lora.lora_registry import LoRARef +from sglang.srt.model_executor.cuda_graph_config import ( + ALLOWED_BACKENDS_PER_PHASE, + Backend, + CudaGraphConfig, + Phase, + default_cuda_graph_config, + parse_cuda_graph_config_arg, +) from sglang.srt.parser.reasoning_parser import ReasoningParser from sglang.srt.platforms import current_platform from sglang.srt.utils.common import ( @@ -732,14 +741,39 @@ class ServerArgs: # Optimization/debug options disable_radix_cache: bool = False - cuda_graph_max_bs: Optional[int] = None - cuda_graph_bs: Optional[List[int]] = None - disable_cuda_graph: bool = False disable_cuda_graph_padding: bool = False - enable_breakable_cuda_graph: bool = False enable_profile_cuda_graph: bool = False enable_cudagraph_gc: bool = False debug_cuda_graph: bool = False + + # Accepts dict (CLI JSON / SDK) at construction time; normalized to + # CudaGraphConfig by _parse_cuda_graph_config. + cuda_graph_config: Optional[CudaGraphConfig] = None + + # Per-phase convenience CLI inputs that fold into cuda_graph_config. + cuda_graph_backend_decode: Optional[ + Literal["full", "breakable", "tc_piecewise", "disabled"] + ] = None + cuda_graph_backend_prefill: Optional[ + Literal["breakable", "tc_piecewise", "disabled"] + ] = None + cuda_graph_max_bs_decode: Optional[int] = None + cuda_graph_max_bs_prefill: Optional[int] = None + cuda_graph_bs_decode: Optional[List[int]] = None + cuda_graph_bs_prefill: Optional[List[int]] = None + cuda_graph_tc_compiler: Optional[Literal["eager", "inductor"]] = None + + # Legacy CLI inputs that fold into cuda_graph_config (with a CLI + # deprecation warning). Internal-only after parsing. + disable_cuda_graph: bool = False + disable_prefill_cuda_graph: bool = False + disable_decode_cuda_graph: bool = False + prefill_cuda_graph_backend: Optional[ + Literal["breakable", "tc_piecewise", "disabled"] + ] = None + decode_cuda_graph_backend: Optional[ + Literal["full", "breakable", "tc_piecewise", "disabled"] + ] = None enable_layerwise_nvtx_marker: bool = False enable_nccl_nvls: bool = False enable_symm_mem: bool = False @@ -762,13 +796,8 @@ class ServerArgs: enable_single_batch_overlap: bool = False tbo_token_distribution_threshold: float = 0.48 enable_torch_compile: bool = False - disable_piecewise_cuda_graph: bool = False - enforce_piecewise_cuda_graph: bool = False enable_torch_compile_debug_mode: bool = False torch_compile_max_bs: int = 32 - piecewise_cuda_graph_max_tokens: Optional[int] = None - piecewise_cuda_graph_tokens: Optional[List[int]] = None - piecewise_cuda_graph_compiler: str = "eager" torchao_config: str = "" enable_p2p_check: bool = False triton_attention_reduce_in_fp32: bool = False @@ -941,6 +970,8 @@ class ServerArgs: # Set missing default values. self._handle_missing_default_values() + self._handle_cuda_graph_config() + # Handle device-specific backends. self._handle_hpu_backends() self._handle_cpu_backends() @@ -951,9 +982,6 @@ class ServerArgs: # Allow OOT platform plugins to apply server args defaults. current_platform.apply_server_args_defaults(self) - # Handle piecewise CUDA graph. - self._handle_piecewise_cuda_graph() - # Get GPU memory capacity, which is a common dependency for several configuration steps. gpu_mem = get_device_memory_capacity(self.device) @@ -1237,11 +1265,11 @@ class ServerArgs: def _handle_modelscope_paths(self): """Resolve model / tokenizer / speculative-draft paths from the local - ModelScope cache when possible, falling back to ``snapshot_download`` + ModelScope cache when possible, falling back to snapshot_download for any path that is not already present on disk. - Note: ``speculative_token_map`` is intentionally NOT handled here - because its value uses ``repo_id/filename`` semantics rather than a + Note: speculative_token_map is intentionally NOT handled here + because its value uses repo_id/filename semantics rather than a plain repo ID. That resolution lives in :func:`sglang.srt.speculative.spec_utils.load_token_map`. """ @@ -1318,12 +1346,13 @@ class ServerArgs: set_default_server_args(self) - if self.piecewise_cuda_graph_compiler != "eager": + current = self.cuda_graph_config.prefill.tc_compiler + if current is not None and current != "eager": logger.warning( "At this moment Ascend platform only support prefill graph compilation with " - "piecewise_cuda_graph_compiler='eager', change piecewise_cuda_graph_compiler to 'eager'." + "cuda_graph_config[prefill].tc_compiler='eager'." ) - self.piecewise_cuda_graph_compiler = "eager" + self.cuda_graph_config.prefill.tc_compiler = "eager" def _handle_mps_backends(self): if self.device == "mps": @@ -1332,83 +1361,191 @@ class ServerArgs: def _handle_xpu_backends(self): if self.device == "xpu": - if not self.disable_piecewise_cuda_graph: + if self.cuda_graph_config.prefill.backend != Backend.DISABLED: logger.warning( - "XPU platform does not support piecewise CUDA graph, ignoring --disable-piecewise-cuda-graph" - " flag and disabling piecewise CUDA graph." + "XPU platform does not support piecewise CUDA graph, " + "disabling prefill cuda graph." ) - self.disable_piecewise_cuda_graph = True + self.cuda_graph_config.prefill.backend = Backend.DISABLED - def _handle_piecewise_cuda_graph(self): - # Skip auto-disable when enforce flag is set (for testing) - if self.enforce_piecewise_cuda_graph: - self.disable_piecewise_cuda_graph = False + # ------------------------------------------------------------------ + # CUDA graph configuration resolution + # ------------------------------------------------------------------ + # TODO: add unit tests in test/srt/test_server_args.py covering the + # precedence cascade + auto-disable matrix (follow-up PR). + def _handle_cuda_graph_config(self): + self._parse_cuda_graph_config() + self._apply_cuda_graph_compatibility() + self._validate_cuda_graph_config() + + def _parse_cuda_graph_config(self): + """Resolve cuda_graph_config from explicit JSON, per-phase + convenience flags, legacy global flags, and defaults. + Precedence (highest first): explicit JSON > convenience > legacy > defaults. + Also populates self._cuda_graph_config_locked — the set of + (phase, key) tuples that came from non-default sources; the + auto-disable cascade respects this lock (the old + --enforce-piecewise-cuda-graph semantics generalized). + """ + raw_input = self.cuda_graph_config + if isinstance(raw_input, CudaGraphConfig): + explicit_input = raw_input.to_dict() + else: + explicit_input = raw_input or {} + config = default_cuda_graph_config() + locked: set = set() + + def _set(phase: str, key: str, value: Any) -> None: + setattr(getattr(config, phase), key, value) + locked.add((phase, key)) + + # ---- Legacy global flags (lowest precedence above defaults) ---- + if self.disable_cuda_graph: + _set(Phase.DECODE, "backend", Backend.DISABLED) + _set(Phase.PREFILL, "backend", Backend.DISABLED) + + # ---- Legacy convenience flags ---- + if self.disable_prefill_cuda_graph: + _set(Phase.PREFILL, "backend", Backend.DISABLED) + if self.disable_decode_cuda_graph: + _set(Phase.DECODE, "backend", Backend.DISABLED) + if self.prefill_cuda_graph_backend is not None: + _set(Phase.PREFILL, "backend", self.prefill_cuda_graph_backend) + if self.decode_cuda_graph_backend is not None: + _set(Phase.DECODE, "backend", self.decode_cuda_graph_backend) + + # ---- Per-phase convenience flags ---- + if self.cuda_graph_backend_decode is not None: + _set(Phase.DECODE, "backend", self.cuda_graph_backend_decode) + if self.cuda_graph_backend_prefill is not None: + _set(Phase.PREFILL, "backend", self.cuda_graph_backend_prefill) + if self.cuda_graph_max_bs_decode is not None: + _set(Phase.DECODE, "max_bs", self.cuda_graph_max_bs_decode) + if self.cuda_graph_max_bs_prefill is not None: + _set(Phase.PREFILL, "max_bs", self.cuda_graph_max_bs_prefill) + if self.cuda_graph_bs_decode is not None: + _set(Phase.DECODE, "bs", self.cuda_graph_bs_decode) + if self.cuda_graph_bs_prefill is not None: + _set(Phase.PREFILL, "bs", self.cuda_graph_bs_prefill) + if self.cuda_graph_tc_compiler is not None: + # Written to both phases so the value is in place when TC_PIECEWISE + # decode is implemented; today decode ignores it. + _set(Phase.DECODE, "tc_compiler", self.cuda_graph_tc_compiler) + _set(Phase.PREFILL, "tc_compiler", self.cuda_graph_tc_compiler) + + # ---- Explicit JSON config (highest precedence) ---- + for phase, phase_config in explicit_input.items(): + if not isinstance(phase_config, dict): + continue + for key, value in phase_config.items(): + _set(phase, key, value) + + self.cuda_graph_config = config + self._cuda_graph_config_locked = locked + + def _apply_cuda_graph_compatibility(self): + """Auto-disable prefill cuda graph for incompatible configs. + Rules are split per backend — TcPiecewise and Breakable have + different constraints. Skipped when the user explicitly set the + prefill backend (this folds in the old + --enforce-piecewise-cuda-graph contract). + """ + if (Phase.PREFILL, "backend") in self._cuda_graph_config_locked: return + if self.cuda_graph_config.prefill.backend == Backend.TC_PIECEWISE: + self._disable_tc_piecewise_cudagraph_if_incompatible() + elif self.cuda_graph_config.prefill.backend == Backend.BREAKABLE: + self._disable_breakable_cudagraph_if_incompatible() - # Disable piecewise cuda graph with following conditions: - # 1. Disable Model Arch - if self.get_model_config().is_piecewise_cuda_graph_disabled_model: - self.disable_piecewise_cuda_graph = True - # 2. DP attention - if self.enable_dp_attention: - self.disable_piecewise_cuda_graph = True - # 3. Torch compile - if self.enable_torch_compile: - self.disable_piecewise_cuda_graph = True - # 4. Pipeline parallelism - if self.pp_size > 1: - self.disable_piecewise_cuda_graph = True - # 5. Non-CUDA hardware (AMD, NPU, CPU, MPS, XPU, etc.) - if is_hip() or is_npu() or is_cpu() or is_mps() or is_xpu(): - self.disable_piecewise_cuda_graph = True - # 5b. OOT platforms that don't support piecewise cuda graph - if current_platform.is_out_of_tree(): - if not current_platform.support_piecewise_cuda_graph(): - self.disable_piecewise_cuda_graph = True - # 6. MoE A2A backend - if self.moe_a2a_backend != "none": - self.disable_piecewise_cuda_graph = True - # 7. LoRA - if self.lora_paths or self.enable_lora: - self.disable_piecewise_cuda_graph = True - # 8. Multimodal / VLM models - if self.get_model_config().is_multimodal: - self.disable_piecewise_cuda_graph = True - # 9. GGUF quantized models (custom dequant ops unsupported by torch.compile) - if ( - self.load_format == "gguf" - or self.quantization == "gguf" - or check_gguf_file(self.model_path) - ): - self.disable_piecewise_cuda_graph = True - # 10. DLLM (diffusion LLM) models (context manager in forward breaks dynamo) - if self.dllm_algorithm is not None: - self.disable_piecewise_cuda_graph = True - # 11. CPU offload (breaks dynamo) - if self.cpu_offload_gb > 0 or self.enable_hierarchical_cache: - self.disable_piecewise_cuda_graph = True - # 12. Deterministic inference - if self.enable_deterministic_inference: - self.disable_piecewise_cuda_graph = True - # 13. PD disaggregation - if self.disaggregation_mode != "null": - self.disable_piecewise_cuda_graph = True - # 14. Symmetric memory (torch.cuda.use_mem_pool is untraceable by dynamo) - if self.enable_symm_mem: - self.disable_piecewise_cuda_graph = True - # 15. Expert distribution recorder - if self.enable_eplb or self.expert_distribution_recorder_mode is not None: - self.disable_piecewise_cuda_graph = True - # 16. Context parallel - if self.attn_cp_size > 1: - self.disable_piecewise_cuda_graph = True - # 18. CUDA Graph debug mode - if self.debug_cuda_graph: - self.disable_piecewise_cuda_graph = True - # 19. DSA prefill context parallelism (attn_cp_size is set later in - # _handle_model_specific_adjustments, so check the flag directly here) - if self.enable_dsa_prefill_context_parallel: - self.disable_piecewise_cuda_graph = True + def _disable_tc_piecewise_cudagraph_if_incompatible(self): + """TcPiecewise (torch.compile + piecewise) is incompatible with + these configurations. Most are torch.compile / dynamo limitations. + """ + + rules = [ + ( + "model-arch blacklist", + lambda: self.get_model_config().is_piecewise_cuda_graph_disabled_model, + ), + ("DP attention", lambda: self.enable_dp_attention), + ("full torch.compile mode", lambda: self.enable_torch_compile), + ("pipeline parallelism (pp_size > 1)", lambda: self.pp_size > 1), + ( + "non-CUDA hardware (HIP/NPU/CPU/MPS/XPU)", + lambda: is_hip() or is_npu() or is_cpu() or is_mps() or is_xpu(), + ), + ( + "OOT platform without piecewise support", + lambda: current_platform.is_out_of_tree() + and not current_platform.support_piecewise_cuda_graph(), + ), + ("MoE A2A backend", lambda: self.moe_a2a_backend != "none"), + ("LoRA", lambda: bool(self.lora_paths) or self.enable_lora), + ("multimodal model", lambda: self.get_model_config().is_multimodal), + ( + "GGUF quantization", + lambda: self.load_format == "gguf" + or self.quantization == "gguf" + or check_gguf_file(self.model_path), + ), + ("DLLM (diffusion LLM)", lambda: self.dllm_algorithm is not None), + ( + "CPU offload / hierarchical cache", + lambda: self.cpu_offload_gb > 0 or self.enable_hierarchical_cache, + ), + ( + "deterministic inference", + lambda: self.enable_deterministic_inference, + ), + ("PD disaggregation", lambda: self.disaggregation_mode != "null"), + ("symmetric memory", lambda: self.enable_symm_mem), + ( + "expert distribution recorder", + lambda: self.enable_eplb + or self.expert_distribution_recorder_mode is not None, + ), + ("context parallel (attn_cp_size > 1)", lambda: self.attn_cp_size > 1), + ("CUDA graph debug mode", lambda: self.debug_cuda_graph), + ( + "DSA prefill context parallelism", + lambda: self.enable_dsa_prefill_context_parallel, + ), + ] + for _name, predicate in rules: + if predicate(): + self.cuda_graph_config.prefill.backend = Backend.DISABLED + + def _disable_breakable_cudagraph_if_incompatible(self): + """Breakable (segmented capture, no torch.compile). Breakable enforces HIP + / memory-saver rejection in its own __init__; config-time + rules can be added here as they're discovered. + """ + rules = [ + # MLA prefill takes a different attn-forward path under BCG (no + # tc_piecewise gate), causing q.view shape mismatches. Disable + # until the MLA prefill path is BCG-aware. + ("MLA attention", lambda: self.use_mla_backend()), + ] + for name, predicate in rules: + if predicate(): + logger.warning( + "Breakable CUDA graph is incompatible with %s; " + "disabling prefill CUDA graph.", + name, + ) + self.cuda_graph_config.prefill.backend = Backend.DISABLED + return + + def _validate_cuda_graph_config(self): + if self.cuda_graph_config is None: + return + for phase in Phase.ALL: + backend = getattr(self.cuda_graph_config, phase).backend + if backend not in ALLOWED_BACKENDS_PER_PHASE[phase]: + raise ValueError( + f"--cuda-graph-config[{phase}].backend={backend!r} not allowed; " + f"allowed: {ALLOWED_BACKENDS_PER_PHASE[phase]}" + ) def _handle_multi_item_scoring(self): """Setup and validate multi-item scoring constraints. @@ -1421,10 +1558,10 @@ class ServerArgs: if not self.enable_mis: return - if not self.disable_cuda_graph: + if self.cuda_graph_config.decode.backend != Backend.DISABLED: logger.warning("CUDA graph is disabled because --enable-mis is set.") - self.disable_cuda_graph = True - self.disable_piecewise_cuda_graph = True + self.cuda_graph_config.decode.backend = Backend.DISABLED + self.cuda_graph_config.prefill.backend = Backend.DISABLED if not self.disable_radix_cache: logger.warning("Radix cache is disabled because --enable-mis is set.") @@ -1444,13 +1581,13 @@ class ServerArgs: def _handle_gpu_memory_settings(self, gpu_mem): """ Configure GPU memory-dependent settings including - chunked_prefill_size, cuda_graph_max_bs, and mem_fraction_static. + chunked_prefill_size, cuda_graph_config[decode].max_bs, and mem_fraction_static. Here are our heuristics: - - Set chunked_prefill_size and cuda_graph_max_bs based on the GPU memory capacity. + - Set chunked_prefill_size and cuda_graph_config[decode].max_bs based on the GPU memory capacity. This is because GPUs with more memory are generally more powerful, we need to use a larger - chunked_prefill_size and a larger cuda_graph_max_bs to fully utilize the GPU. - - Then set mem_fraction_static based on chunked_prefill_size and cuda_graph_max_bs. + chunked_prefill_size and a larger decode max_bs to fully utilize the GPU. + - Then set mem_fraction_static based on chunked_prefill_size and decode max_bs. GPU memory capacity = model weights + KV cache pool + activations + cuda graph buffers @@ -1459,134 +1596,140 @@ class ServerArgs: In order to compute mem_fraction_static, we need to estimate the size of activations and cuda graph buffers. The activation memory is proportional to the chunked_prefill_size. - The cuda graph memory is proportional to the cuda_graph_max_bs. - We use reserved_mem = chunked_prefill_size * 1.5 + cuda_graph_max_bs * 2 to estimate the size of activations and cuda graph buffers in GB. + The cuda graph memory is proportional to the decode max_bs. + We use reserved_mem = chunked_prefill_size * 1.5 + max_bs * 2 to estimate the size of activations and cuda graph buffers in GB, and set mem_fraction_static = (GPU memory capacity - reserved_mem) / GPU memory capacity. The coefficient 1.5 is a heuristic value, in the future, we can do better estimation by looking at the model types, hidden sizes or even do a dummy run. """ + decode_cuda_graph_config = self.cuda_graph_config.decode + prefill_cuda_graph_config = self.cuda_graph_config.prefill + if gpu_mem is not None: if gpu_mem < 20 * 1024: # T4, 4080 - # (chunked_prefill_size 2k, cuda_graph_max_bs 8) + # (chunked_prefill_size 2k, max_bs 8) if self.chunked_prefill_size is None: self.chunked_prefill_size = 2048 - if self.cuda_graph_max_bs is None: - self.cuda_graph_max_bs = 8 + if decode_cuda_graph_config.max_bs is None: + decode_cuda_graph_config.max_bs = 8 elif gpu_mem < 35 * 1024: # A10, 4090, 5090 - # (chunked_prefill_size 2k, cuda_graph_max_bs 24 if tp < 4 else 80) + # (chunked_prefill_size 2k, max_bs 24 if tp < 4 else 80) if self.chunked_prefill_size is None: self.chunked_prefill_size = 2048 - if self.cuda_graph_max_bs is None: - # Based on detailed statistics, when serving TP1/TP2 models on lower-end GPUs with HBM < 35GB, you can either disable cuda graph or set `cuda_graph_max_bs` to a very small value to reduce the memory overhead of creating cuda graphs, with almost no impact on performance. - # However, when serving models with TP4 or TP8, we need to enable cuda graph to maintain high performance. In this case, we can set `cuda_graph_max_bs` to 80 (half of the default value 160) to reduce the memory overhead of creating cuda graphs. Looking at the logs - # from TP4 serving of qwen2-72b, a value of 80 is sufficient and can reduce the memory overhead of creating cuda graphs on lower-end GPUs compared to the original 160, avoiding OOM issues. + if decode_cuda_graph_config.max_bs is None: if self.tp_size < 4: - self.cuda_graph_max_bs = 24 + decode_cuda_graph_config.max_bs = 24 else: - self.cuda_graph_max_bs = 80 + decode_cuda_graph_config.max_bs = 80 elif gpu_mem < 60 * 1024: # A100 (40GB), L40, - # (chunked_prefill_size 4k, cuda_graph_max_bs 32 if tp < 4 else 160) + # (chunked_prefill_size 4k, max_bs 32 if tp < 4 else 160) if self.chunked_prefill_size is None: self.chunked_prefill_size = 4096 - if self.cuda_graph_max_bs is None: + if decode_cuda_graph_config.max_bs is None: if self.tp_size < 4: - self.cuda_graph_max_bs = 32 + decode_cuda_graph_config.max_bs = 32 else: - self.cuda_graph_max_bs = 160 + decode_cuda_graph_config.max_bs = 160 elif gpu_mem < 90 * 1024: # H100, A100 - # (chunked_prefill_size 8k, cuda_graph_max_bs 256 if tp < 4 else 512) + # (chunked_prefill_size 8k, max_bs 256 if tp < 4 else 512) if self.chunked_prefill_size is None: self.chunked_prefill_size = 8192 - if self.cuda_graph_max_bs is None: + if decode_cuda_graph_config.max_bs is None: if self.tp_size < 4: - self.cuda_graph_max_bs = 256 + decode_cuda_graph_config.max_bs = 256 else: - self.cuda_graph_max_bs = 512 + decode_cuda_graph_config.max_bs = 512 elif gpu_mem < 160 * 1024: # H20, H200 - # (chunked_prefill_size 8k, cuda_graph_max_bs 256 if tp < 4 else 512) + # (chunked_prefill_size 8k, max_bs 256 if tp < 4 else 512) if self.chunked_prefill_size is None: self.chunked_prefill_size = 8192 - if self.cuda_graph_max_bs is None: + if decode_cuda_graph_config.max_bs is None: if self.tp_size < 4: - self.cuda_graph_max_bs = 256 + decode_cuda_graph_config.max_bs = 256 else: - self.cuda_graph_max_bs = 512 + decode_cuda_graph_config.max_bs = 512 else: # B200, MI300 - # (chunked_prefill_size 16k, cuda_graph_max_bs 512) + # (chunked_prefill_size 16k, max_bs 512) if self.chunked_prefill_size is None: self.chunked_prefill_size = 16384 - if self.cuda_graph_max_bs is None: - self.cuda_graph_max_bs = 512 + if decode_cuda_graph_config.max_bs is None: + decode_cuda_graph_config.max_bs = 512 else: # Fallback defaults when gpu_mem is None if self.chunked_prefill_size is None: self.chunked_prefill_size = 4096 - if self.cuda_graph_max_bs is None: - self.cuda_graph_max_bs = 160 + if decode_cuda_graph_config.max_bs is None: + decode_cuda_graph_config.max_bs = 160 # Set cuda graph batch sizes if self.device != "cpu": - if self.cuda_graph_bs is None: - self.cuda_graph_bs = self._generate_cuda_graph_batch_sizes() - else: - self.cuda_graph_max_bs = max(self.cuda_graph_bs) - else: - # Reuse cuda_graph_bs for cpu graph and use torch_compile_max_bs for cpu graph batch size limit, - # as cpu graph is based on torch.compile - if self.cuda_graph_bs is not None: - self.torch_compile_max_bs = max(self.cuda_graph_bs) - else: - # If cuda_graph_bs is not set, we will preferentially use torch_compile_max_bs - # to generate cuda_graph_bs - self.torch_compile_max_bs = ( - self.torch_compile_max_bs or self.cuda_graph_max_bs + if decode_cuda_graph_config.bs is None: + decode_cuda_graph_config.bs = ( + self._generate_decode_cuda_graph_batch_sizes( + decode_cuda_graph_config.max_bs + ) ) - self.cuda_graph_bs = self._generate_cpu_graph_batch_sizes() + else: + decode_cuda_graph_config.max_bs = max(decode_cuda_graph_config.bs) + else: + # Reuse decode_cuda_graph_config.bs for cpu graph and use torch_compile_max_bs for cpu graph batch size limit, + # as cpu graph is based on torch.compile + if decode_cuda_graph_config.bs is not None: + self.torch_compile_max_bs = max(decode_cuda_graph_config.bs) + else: + # If decode_cuda_graph_config.bs is not set, we will preferentially use torch_compile_max_bs + # to generate decode_cuda_graph_config.bs + self.torch_compile_max_bs = ( + self.torch_compile_max_bs or decode_cuda_graph_config.max_bs + ) + decode_cuda_graph_config.bs = self._generate_cpu_graph_batch_sizes() assert ( self.torch_compile_max_bs > 0 - ), "cuda_graph_bs should contain positive batch sizes" - self.cuda_graph_max_bs = self.torch_compile_max_bs + ), "cuda_graph_config[decode].bs should contain positive batch sizes" + decode_cuda_graph_config.max_bs = self.torch_compile_max_bs - if self.piecewise_cuda_graph_max_tokens is None: - # Refer to pr #15927, by default we set the piecewise cuda graph max tokens to the chunked prefill size by default. + if prefill_cuda_graph_config.max_bs is None: + # Refer to pr #15927, by default we set the prefill max_bs to the chunked prefill size. # For MLA backend, the introduction of piecewise cuda graph will influence the kernel dispatch difference compared to the original mode. - # To avoid the performance regression, we set the max tokens to 2048 by default. + # To avoid the performance regression, we set max_bs to 2048 by default. if not self.use_mla_backend(): - self.piecewise_cuda_graph_max_tokens = self.chunked_prefill_size + prefill_cuda_graph_config.max_bs = self.chunked_prefill_size else: - self.piecewise_cuda_graph_max_tokens = 2048 + prefill_cuda_graph_config.max_bs = 2048 - # If max_total_tokens is set, cap pcg tokens to not exceed max_total_tokens + # If max_total_tokens is set, cap prefill max_bs to not exceed max_total_tokens. if self.max_total_tokens is not None: - self.piecewise_cuda_graph_max_tokens = min( - self.piecewise_cuda_graph_max_tokens, self.max_total_tokens + prefill_cuda_graph_config.max_bs = min( + prefill_cuda_graph_config.max_bs, self.max_total_tokens ) - # For Llama2 series models, the max tokens is limited to 4096 + # For Llama2 series models, max_bs is limited to 4096. # TODO(yuwei): remove this after the issue is fixed if "llama-2" in self.model_path.lower(): - self.piecewise_cuda_graph_max_tokens = min( - self.piecewise_cuda_graph_max_tokens, 4096 + prefill_cuda_graph_config.max_bs = min( + prefill_cuda_graph_config.max_bs, 4096 ) - # Clamp to context_length if explicitly set — prevents PCG warmup - # from compiling graphs with more tokens than the model buffers - # can hold, which causes illegal memory access (#21112) + # Clamp to context_length if explicitly set — prevents prefill CG + # warmup from compiling graphs with more tokens than the model + # buffers can hold, which causes illegal memory access (#21112). if self.context_length is not None: - self.piecewise_cuda_graph_max_tokens = min( - self.piecewise_cuda_graph_max_tokens, self.context_length + prefill_cuda_graph_config.max_bs = min( + prefill_cuda_graph_config.max_bs, self.context_length ) - if self.piecewise_cuda_graph_tokens is None: - self.piecewise_cuda_graph_tokens = ( - self._generate_piecewise_cuda_graph_tokens() + if prefill_cuda_graph_config.bs is None: + prefill_cuda_graph_config.bs = ( + self._generate_prefill_cuda_graph_batch_sizes( + prefill_cuda_graph_config.max_bs + ) ) if self.mem_fraction_static is None: @@ -1598,25 +1741,25 @@ class ServerArgs: else: reserved_mem += max(self.max_prefill_tokens, 2048) * 1.5 # For cuda graphs - reserved_mem += self.cuda_graph_max_bs * 2 + reserved_mem += decode_cuda_graph_config.max_bs * 2 # Some adjustments for large parallel size reserved_mem += self.tp_size * self.pp_size / 8 * 1024 if self.enable_dp_attention: # DP attention needs more padding for some operations - reserved_mem += self.cuda_graph_max_bs * self.dp_size * 3 + reserved_mem += decode_cuda_graph_config.max_bs * self.dp_size * 3 # DP attention uses much more memory for large cuda graph max bs, # likely due to some inefficiencies in torch allocator or our implementation. # So we need to reserve more memory. - if self.cuda_graph_max_bs > 300: - reserved_mem += self.cuda_graph_max_bs * self.dp_size * 1.5 + if decode_cuda_graph_config.max_bs > 300: + reserved_mem += decode_cuda_graph_config.max_bs * self.dp_size * 1.5 # For piecewise cuda graphs - if not self.disable_piecewise_cuda_graph: + if prefill_cuda_graph_config.backend != Backend.DISABLED: if not self.use_mla_backend(): # Only calculate the memory overhead for Non-Torch Memory use since the Torch Memory can be reused with Cuda Graph Capture - reserved_mem += len(self.piecewise_cuda_graph_tokens) * 8 + reserved_mem += len(prefill_cuda_graph_config.bs) * 8 else: # For MLA backend the memory overhead is much higher than expected with fa3 reserved_mem += 1.5 * 1024 @@ -1652,21 +1795,21 @@ class ServerArgs: "Use environment variable SGLANG_SYMM_MEM_PREALLOC_GB_SIZE to change the prealloc size." ) - def _generate_cuda_graph_batch_sizes(self): + def _generate_decode_cuda_graph_batch_sizes(self, max_bs: int): """ - Generate the list of batch sizes for CUDA graph capture based on cuda_graph_max_bs. + Generate the list of batch sizes for CUDA graph capture based on max_bs. This integrates the logic from cuda_graph_runner.py. """ # Handle disable_cuda_graph_padding as the first condition for both spec and non-spec if self.disable_cuda_graph_padding: - capture_bs = list(range(1, self.cuda_graph_max_bs + 1)) + capture_bs = list(range(1, max_bs + 1)) elif self.speculative_algorithm is None: # Normal case: capture_bs = ( [1, 2, 4, 8, 12] + list(range(16, 257, 8)) + list(range(272, 512, 16)) - + list(range(512, self.cuda_graph_max_bs + 1, 32)) + + list(range(512, max_bs + 1, 32)) ) else: # Spec decoding case: less padding for smaller batch sizes @@ -1675,13 +1818,13 @@ class ServerArgs: + list(range(10, 33, 2)) + list(range(40, 65, 4)) + list(range(72, 257, 8)) - + list(range(272, self.cuda_graph_max_bs + 1, 16)) + + list(range(272, max_bs + 1, 16)) ) - capture_bs = [bs for bs in capture_bs if bs <= self.cuda_graph_max_bs] + capture_bs = [bs for bs in capture_bs if bs <= max_bs] - if self.cuda_graph_max_bs not in capture_bs: - capture_bs.append(self.cuda_graph_max_bs) + if max_bs not in capture_bs: + capture_bs.append(max_bs) return capture_bs @@ -1705,10 +1848,11 @@ class ServerArgs: return capture_bs - def _generate_piecewise_cuda_graph_tokens(self): + def _generate_prefill_cuda_graph_batch_sizes(self, max_bs: int): """ - Generate the list of batch sizes for piecewise CUDA graph capture - based on piecewise_cuda_graph_max_tokens. + Generate the list of batch sizes for prefill CUDA graph capture + based on max_bs. For tc_piecewise prefill, bs carries the + captured token count (one shape knob per phase). """ capture_sizes = ( list(range(4, 33, 4)) @@ -1716,12 +1860,10 @@ class ServerArgs: + list(range(288, 513, 32)) + list(range(576, 1024 + 1, 64)) + list(range(1280, 4096 + 1, 256)) - + list(range(4608, self.piecewise_cuda_graph_max_tokens + 1, 512)) + + list(range(4608, max_bs + 1, 512)) ) - capture_sizes = [ - s for s in capture_sizes if s <= self.piecewise_cuda_graph_max_tokens - ] + capture_sizes = [s for s in capture_sizes if s <= max_bs] return capture_sizes @@ -1905,7 +2047,7 @@ class ServerArgs: # DSACPLayerCommunicator does not all-reduce attention-TP # partial o_proj outputs before replicated dense FFNs. self.attn_cp_size = self.tp_size // self.dp_size - self.disable_piecewise_cuda_graph = True + self.cuda_graph_config.prefill.backend = Backend.DISABLED logger.warning( f"Enable DSA Context Parallel opt, " f"Setting dp_size == {self.dp_size} and " @@ -1914,7 +2056,7 @@ class ServerArgs: f"tp_size == {self.tp_size}, " f"kv_cache_dtype == {self.kv_cache_dtype}, " f"moe_a2a_backend {self.moe_a2a_backend}, " - f"disable_piecewise_cuda_graph=True" + f"cuda_graph_config[prefill].backend=disabled" ) else: # Pure TP and partial DP Attention mode is active for DSA, logging a warning @@ -1956,8 +2098,8 @@ class ServerArgs: ), "CP is only supported for prefill when PD disaggregation, please remove --enable-dsa-prefill-context-parallel." else: - # DeepSeek V3/R1/V3.1 and Kimi K2.5 - if not self.disable_piecewise_cuda_graph: + # DeepSeek V3/R1/V3.1 + if self.cuda_graph_config.prefill.backend != Backend.DISABLED: logger.info("Piecewise CUDA graph is enabled, use MLA for prefill.") if is_sm100_supported(): @@ -1990,7 +2132,7 @@ class ServerArgs: # DSACPLayerCommunicator does not all-reduce attention-TP # partial o_proj outputs before replicated dense FFNs. self.attn_cp_size = self.tp_size // self.dp_size - self.disable_piecewise_cuda_graph = True + self.cuda_graph_config.prefill.backend = Backend.DISABLED logger.warning( f"Enable Context Parallel opt for MLA, " f"Setting dp_size == {self.dp_size} and " @@ -1999,7 +2141,7 @@ class ServerArgs: f"ep_size == {self.ep_size}, " f"tp_size == {self.tp_size}, " f"moe_a2a_backend {self.moe_a2a_backend}, " - f"disable_piecewise_cuda_graph=True" + f"cuda_graph_config[prefill].backend=disabled" ) # Set moe backend for DeepSeek @@ -2853,13 +2995,15 @@ class ServerArgs: logger.warning( "Cuda graph is disabled because of using torch native attention backend" ) - self.disable_cuda_graph = True + self.cuda_graph_config.decode.backend = Backend.DISABLED + self.cuda_graph_config.prefill.backend = Backend.DISABLED if self.attention_backend == "flex_attention": logger.warning( "Cuda graph is disabled because of using torch Flex Attention backend" ) - self.disable_cuda_graph = True + self.cuda_graph_config.decode.backend = Backend.DISABLED + self.cuda_graph_config.prefill.backend = Backend.DISABLED assert ( self.speculative_algorithm is None ), "Speculative decoding is currently not supported with Flex Attention backend" @@ -3436,13 +3580,28 @@ class ServerArgs: self.ep_size == 1 ), "FP8/MXFP8 Cutlass MoE is only supported with ep_size == 1" + # TODO(yuwei): Fix piecewise cuda graph support for bypassed topk MoE backends. + # Exception: GptOssForCausalLM wraps the entire MoE block in its own + # custom op (moe_impl), so bypassed topk is handled inside the op body. + if ( + (Phase.PREFILL, "backend") not in self._cuda_graph_config_locked + and self.moe_runner_backend in ("flashinfer_trtllm", "flashinfer_mxfp4") + and self.get_model_config().hf_config.architectures[0] + != "GptOssForCausalLM" + ): + self.cuda_graph_config.prefill.backend = Backend.DISABLED + logger.info( + f"Piecewise cuda graph is disabled for MoE runner backend " + f"'{self.moe_runner_backend}' (bypassed topk is incompatible " + f"with torch.compile)." + ) + def cutedsl_moe_max_num_tokens(self) -> int: """Largest number of tokens a single forward routes through a CuteDSL MoE layer on one (DP) rank. Single source of truth for both the standard-allgather wrapper buffers and the FlashInfer A2A dispatcher budget. Max over the prefill (max_prefill_tokens), piecewise-prefill - capture (piecewise_cuda_graph_max_tokens), and decode/verify - (cuda_graph_max_bs * num_tokens_per_bs) bounds; num_tokens_per_bs is + capture, and decode/verify bounds; num_tokens_per_bs is speculative_num_draft_tokens under speculative decoding, else 1. """ if self.speculative_algorithm: @@ -3450,11 +3609,11 @@ class ServerArgs: else: num_tokens_per_bs = 1 prefill_tokens = self.max_prefill_tokens - if not self.disable_piecewise_cuda_graph: - prefill_tokens = max( - prefill_tokens, self.piecewise_cuda_graph_max_tokens or 0 - ) - decode_tokens = (self.cuda_graph_max_bs or 0) * num_tokens_per_bs + cg_config = self.cuda_graph_config + if cg_config is not None and cg_config.prefill.backend == Backend.TC_PIECEWISE: + prefill_tokens = max(prefill_tokens, cg_config.prefill.max_bs or 0) + decode_max_bs = (cg_config.decode.max_bs if cg_config is not None else 0) or 0 + decode_tokens = decode_max_bs * num_tokens_per_bs return max(prefill_tokens, decode_tokens) def _validate_cutedsl_a2a_token_budget(self): @@ -3522,7 +3681,8 @@ class ServerArgs: if self.moe_a2a_backend == "deepep": if self.deepep_mode == "normal": logger.warning("Cuda graph is disabled because deepep_mode=`normal`") - self.disable_cuda_graph = True + self.cuda_graph_config.decode.backend = Backend.DISABLED + self.cuda_graph_config.prefill.backend = Backend.DISABLED self.ep_size = self.tp_size logger.warning( f"DeepEP MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]." @@ -3914,14 +4074,14 @@ class ServerArgs: def _is_mistral_native_format(self) -> bool: """True iff the checkpoint requires load_format=mistral. - Looks for ``consolidated*.safetensors`` with no competing - ``model-*.safetensors``; when both weight formats ship in the + Looks for consolidated*.safetensors with no competing + model-*.safetensors; when both weight formats ship in the same checkpoint (e.g. Mistral-7B-Instruct-v0.3) the HF path is preferred to avoid loading Mistral-named weights into an HF-named architecture. - Name override: ``mistral-large-3`` / ``mistral-small-4`` / - ``leanstral`` always treat as Mistral-native when ``params.json`` + Name override: mistral-large-3 / mistral-small-4 / + leanstral always treat as Mistral-native when params.json is present -- those families need Mistral weight loading regardless of which weight files happen to be present. """ @@ -4305,11 +4465,15 @@ class ServerArgs: return # On AMD/HIP, disable cuda graph for DLLM and use triton backend if is_hip(): - if not self.disable_cuda_graph: + if ( + self.cuda_graph_config.decode.backend != Backend.DISABLED + or self.cuda_graph_config.prefill.backend != Backend.DISABLED + ): logger.warning( "Cuda graph is disabled for diffusion LLM inference on AMD GPUs" ) - self.disable_cuda_graph = True + self.cuda_graph_config.decode.backend = Backend.DISABLED + self.cuda_graph_config.prefill.backend = Backend.DISABLED if self.attention_backend not in ["triton", "aiter"]: logger.warning( "Attention backend is set to triton for diffusion LLM inference on AMD GPUs" @@ -4321,7 +4485,7 @@ class ServerArgs: "Attention backend is overridden to 'ascend' when running on NPU for diffusion LLM inference." ) self.attention_backend = "ascend" - elif not self.disable_cuda_graph: + elif self.cuda_graph_config.decode.backend != Backend.DISABLED: if self.attention_backend != "flashinfer": logger.warning( "Attention backend is set to flashinfer because of enabling cuda graph in diffusion LLM inference" @@ -4414,16 +4578,18 @@ class ServerArgs: logger.warning( "Cuda graph and server warmup are disabled because of using tensor dump mode" ) - self.disable_cuda_graph = True + self.cuda_graph_config.decode.backend = Backend.DISABLED + self.cuda_graph_config.prefill.backend = Backend.DISABLED self.skip_server_warmup = True if self.msprobe_dump_config is not None: logger.warning( "When msProbe is enabled, " - "cuda graph is disabled(disable_cuda_graph=True) because msProbe only supports dump in eager mode, " + "cuda graph is disabled because msProbe only supports dump in eager mode, " "warmup is disabled(skip_server_warmup=True) because there is no need to dump data for this stage." ) - self.disable_cuda_graph = True + self.cuda_graph_config.decode.backend = Backend.DISABLED + self.cuda_graph_config.prefill.backend = Backend.DISABLED self.skip_server_warmup = True # Validate limit_mm_per_prompt modalities @@ -6473,6 +6639,19 @@ class ServerArgs: action="store_true", help="Disable RadixAttention for prefix caching.", ) + # --- CUDA graph config: canonical JSON entry --------------------- + parser.add_argument( + "--cuda-graph-config", + type=parse_cuda_graph_config_arg, + default=ServerArgs.cuda_graph_config, + help="Per-phase CUDA graph settings as JSON, e.g. " + '\'{"decode":{"backend":"full","max_bs":256},"prefill":{"backend":"tc_piecewise","tc_compiler":"eager"}}\'. ' + "Allowed backends per phase: full, breakable, tc_piecewise, disabled " + "(full is decode-only). JSON wins over the per-phase --cuda-graph-* " + "convenience flags and over legacy flags.", + ) + + # --- KV canary debug flags (upstream PR #26818-26821) ------------ parser.add_argument( "--kv-canary", type=str, @@ -6503,33 +6682,62 @@ class ServerArgs: default=ServerArgs.kv_canary_sweep_interval, help="Every N forward steps, run a full-pool sweep.", ) + + # --- CUDA graph: per-phase convenience flags --------------------- parser.add_argument( - "--cuda-graph-max-bs", - type=int, - default=ServerArgs.cuda_graph_max_bs, - help="Set the maximum batch size for cuda graph. It will extend the cuda graph capture batch size to this value.", + "--cuda-graph-backend-decode", + type=str, + choices=Backend.ALL, + default=ServerArgs.cuda_graph_backend_decode, + help="Backend for the decode phase. Folds into cuda_graph_config[decode].backend.", ) parser.add_argument( - "--cuda-graph-bs", + "--cuda-graph-backend-prefill", + type=str, + choices=Backend.ALL, + default=ServerArgs.cuda_graph_backend_prefill, + help="Backend for the prefill phase. Folds into cuda_graph_config[prefill].backend.", + ) + parser.add_argument( + "--cuda-graph-max-bs-decode", + type=int, + default=ServerArgs.cuda_graph_max_bs_decode, + help="Maximum batch size captured for the decode cuda graph.", + ) + parser.add_argument( + "--cuda-graph-max-bs-prefill", + type=int, + default=ServerArgs.cuda_graph_max_bs_prefill, + help="Maximum batch size captured for the prefill cuda graph.", + ) + parser.add_argument( + "--cuda-graph-bs-decode", type=int, nargs="+", - help="Set the list of batch sizes for cuda graph.", + default=ServerArgs.cuda_graph_bs_decode, + help="Explicit list of batch sizes to capture for the decode cuda graph.", ) parser.add_argument( - "--disable-cuda-graph", - action="store_true", - help="Disable cuda graph.", + "--cuda-graph-bs-prefill", + type=int, + nargs="+", + default=ServerArgs.cuda_graph_bs_prefill, + help="Explicit list of batch sizes to capture for the prefill cuda graph.", ) + parser.add_argument( + "--cuda-graph-tc-compiler", + type=str, + choices=["eager", "inductor"], + default=ServerArgs.cuda_graph_tc_compiler, + help="Compiler used by the tc_piecewise backend (currently only the prefill phase consumes it).", + ) + + # --- CUDA graph: debug / profiling flags ------------------------- parser.add_argument( "--disable-cuda-graph-padding", action="store_true", help="Disable cuda graph when padding is needed. Still uses cuda graph when padding is not needed.", ) - parser.add_argument( - "--enable-breakable-cuda-graph", - action="store_true", - help="Use breakable CUDA graph for piecewise capture instead of torch.compile-based splitting.", - ) parser.add_argument( "--enable-profile-cuda-graph", action="store_true", @@ -6548,6 +6756,111 @@ class ServerArgs: "while still going through the CUDA graph capture / replay path. " "Useful for debugging CUDA graph capture / replay issues.", ) + + # --- CUDA graph related deprecated args. Remove them later. ----- + parser.add_argument( + "--cuda-graph-max-bs", + type=int, + action=DeprecatedAliasStoreAction, + new_flag="--cuda-graph-max-bs-decode", + dest="cuda_graph_max_bs_decode", + help="Deprecated alias for --cuda-graph-max-bs-decode.", + ) + parser.add_argument( + "--cuda-graph-bs", + type=int, + nargs="+", + action=DeprecatedAliasStoreAction, + new_flag="--cuda-graph-bs-decode", + dest="cuda_graph_bs_decode", + help="Deprecated alias for --cuda-graph-bs-decode.", + ) + parser.add_argument( + "--disable-cuda-graph", + action=DeprecatedStoreTrueAction, + new_flag="--cuda-graph-backend-{decode,prefill}=disabled", + help="Deprecated. Use --cuda-graph-backend-{decode,prefill}=disabled instead.", + ) + parser.add_argument( + "--enable-breakable-cuda-graph", + action=DeprecatedStoreConstAction, + dest="cuda_graph_backend_prefill", + const_value=Backend.BREAKABLE, + new_flag="--cuda-graph-backend-prefill=breakable", + help="Deprecated alias for --cuda-graph-backend-prefill=breakable.", + ) + parser.add_argument( + "--prefill-cuda-graph-backend", + type=str, + choices=Backend.ALL, + action=DeprecatedAliasStoreAction, + new_flag="--cuda-graph-backend-prefill", + help="Deprecated alias for --cuda-graph-backend-prefill.", + ) + parser.add_argument( + "--decode-cuda-graph-backend", + type=str, + choices=Backend.ALL, + action=DeprecatedAliasStoreAction, + new_flag="--cuda-graph-backend-decode", + help="Deprecated alias for --cuda-graph-backend-decode.", + ) + parser.add_argument( + "--disable-prefill-cuda-graph", + action=DeprecatedStoreTrueAction, + new_flag="--cuda-graph-backend-prefill=disabled", + help="Deprecated. Use --cuda-graph-backend-prefill=disabled instead.", + ) + parser.add_argument( + "--disable-decode-cuda-graph", + action=DeprecatedStoreTrueAction, + new_flag="--cuda-graph-backend-decode=disabled", + help="Deprecated. Use --cuda-graph-backend-decode=disabled instead.", + ) + parser.add_argument( + "--disable-piecewise-cuda-graph", + action=DeprecatedStoreConstAction, + dest="cuda_graph_backend_prefill", + const_value=Backend.DISABLED, + new_flag="--cuda-graph-backend-prefill=disabled", + help="Deprecated alias for --cuda-graph-backend-prefill=disabled.", + ) + parser.add_argument( + "--enforce-piecewise-cuda-graph", + action=DeprecatedStoreConstAction, + dest="cuda_graph_backend_prefill", + const_value=Backend.TC_PIECEWISE, + new_flag="--cuda-graph-backend-prefill=tc_piecewise", + help="Deprecated alias for --cuda-graph-backend-prefill=tc_piecewise. " + "Explicitly setting the prefill backend now skips the auto-disable " + "cascade automatically.", + ) + parser.add_argument( + "--piecewise-cuda-graph-tokens", + type=int, + nargs="+", + action=DeprecatedAliasStoreAction, + new_flag="--cuda-graph-bs-prefill", + dest="cuda_graph_bs_prefill", + help="Deprecated alias for --cuda-graph-bs-prefill.", + ) + parser.add_argument( + "--piecewise-cuda-graph-compiler", + type=str, + choices=["eager", "inductor"], + action=DeprecatedAliasStoreAction, + new_flag="--cuda-graph-tc-compiler", + dest="cuda_graph_tc_compiler", + help="Deprecated alias for --cuda-graph-tc-compiler.", + ) + parser.add_argument( + "--piecewise-cuda-graph-max-tokens", + type=int, + action=DeprecatedAliasStoreAction, + new_flag="--cuda-graph-max-bs-prefill", + dest="cuda_graph_max_bs_prefill", + help="Deprecated alias for --cuda-graph-max-bs-prefill.", + ) parser.add_argument( "--enable-layerwise-nvtx-marker", action="store_true", @@ -6656,46 +6969,12 @@ class ServerArgs: action="store_true", help="Enable debug mode for torch compile", ) - parser.add_argument( - "--disable-piecewise-cuda-graph", - action="store_true", - help="Disable piecewise cuda graph for extend/prefill.", - ) - parser.add_argument( - "--enable-piecewise-cuda-graph", - action=DeprecatedAction, - help="Deprecated: Piecewise cuda graph is enabled by default. Use --enforce-piecewise-cuda-graph to skip auto-disable conditions.", - ) - parser.add_argument( - "--enforce-piecewise-cuda-graph", - action="store_true", - help="Enforce piecewise cuda graph, skipping all auto-disable conditions. Used for testing.", - ) - parser.add_argument( - "--piecewise-cuda-graph-tokens", - type=int, - nargs="+", - help="Set the list of token lengths for piecewise cuda graph capture.", - ) - parser.add_argument( - "--piecewise-cuda-graph-compiler", - type=str, - default=ServerArgs.piecewise_cuda_graph_compiler, - help="Set the compiler for piecewise cuda graph. Choices are: eager, inductor.", - choices=["eager", "inductor"], - ) parser.add_argument( "--torch-compile-max-bs", type=int, default=ServerArgs.torch_compile_max_bs, help="Set the maximum batch size when using torch compile.", ) - parser.add_argument( - "--piecewise-cuda-graph-max-tokens", - type=int, - default=ServerArgs.piecewise_cuda_graph_max_tokens, - help="Set the maximum tokens when using piecewise cuda graph.", - ) parser.add_argument( "--torchao-config", type=str, @@ -7300,7 +7579,7 @@ class ServerArgs: return self.url(port=self.engine_info_bootstrap_port) def ssl_verify(self): - """Return the value for the requests library's ``verify=`` parameter. + """Return the value for the requests library's verify= parameter. When SSL is configured: - If a CA certificate file is provided, return its path so requests @@ -7874,8 +8153,8 @@ class ServerArgs: `/server_info` so KV-aware routers (e.g. the SGLang model gateway) can subscribe per-worker without operator-supplied port coordination. The router constructs the per-DP-rank SUB endpoint - as ``tcp://:`` for - every rank reported in ``dp_size``. + as tcp://: for + every rank reported in dp_size. Returned descriptor shape: @@ -7897,23 +8176,23 @@ class ServerArgs: # to open } - Returns ``None`` (i.e. "no publisher to describe") when any of: + Returns None (i.e. "no publisher to describe") when any of: - * ``--kv-events-config`` is unset / empty / malformed JSON, - * the configured publisher is ``"null"``, - * ``page_size`` is missing or non-positive (a placeholder - ``block_size`` would cause silent KV-cache misses by hashing + * --kv-events-config is unset / empty / malformed JSON, + * the configured publisher is "null", + * page_size is missing or non-positive (a placeholder + block_size would cause silent KV-cache misses by hashing prompts at the wrong granularity on the router side), - * the endpoint is not a routable TCP address (``inproc://`` / - ``ipc://``, missing port, non-integer port, or port outside - ``1..65535``). + * the endpoint is not a routable TCP address (inproc:// / + ipc://, missing port, non-integer port, or port outside + 1..65535). - Reuses ``KVEventsConfig.from_cli`` for JSON parsing; the inline - ``rfind(":")`` endpoint split mirrors - ``ZmqEventPublisher.offset_endpoint_port`` rather than adding a + Reuses KVEventsConfig.from_cli for JSON parsing; the inline + rfind(":") endpoint split mirrors + ZmqEventPublisher.offset_endpoint_port rather than adding a new module-level helper. """ - # Lazy import so loading ``server_args`` doesn't pull in + # Lazy import so loading server_args doesn't pull in # disaggregation / msgspec / zmq at module top level. from sglang.srt.disaggregation.kv_events import KVEventsConfig @@ -7925,7 +8204,7 @@ class ServerArgs: cfg = KVEventsConfig.from_cli(raw) except Exception: # Malformed JSON / schema mismatch. The publisher would - # have failed at server startup; ``/server_info`` must + # have failed at server startup; /server_info must # keep working, so just report "no publisher" to consumers. return None if cfg.publisher == "null" or not cfg.endpoint: diff --git a/python/sglang/srt/speculative/adaptive_runtime_state.py b/python/sglang/srt/speculative/adaptive_runtime_state.py index aeb0615dd..fa9211dd3 100644 --- a/python/sglang/srt/speculative/adaptive_runtime_state.py +++ b/python/sglang/srt/speculative/adaptive_runtime_state.py @@ -6,7 +6,7 @@ from sglang.srt.speculative.adaptive_spec_params import AdaptiveSpeculativeParam if TYPE_CHECKING: from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.model_executor.cpu_graph_runner import CPUGraphRunner - from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner + from sglang.srt.model_executor.runner import DecodeCudaGraphRunner from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( EAGLEDraftCudaGraphRunner, ) @@ -36,7 +36,7 @@ class SpecRuntimeState: # -- Verify stage: target model one-pass tree verification -- target_attn_backend: "AttentionBackend" - target_graph_runner: "CudaGraphRunner | CPUGraphRunner | None" + target_graph_runner: "DecodeCudaGraphRunner | CPUGraphRunner | None" # -- Extend stage: draft model KV cache catch-up after verify -- draft_extend_attn_backend: "AttentionBackend | None" @@ -61,14 +61,14 @@ class AdaptiveSpecWorker(Protocol): class AdaptiveController: """Facade that owns adaptive decision-making and runtime state switching. - Works with any worker that implements ``AdaptiveSpecWorker`` protocol: - - ``build_adaptive_runtime_state()`` → runtime state - - ``apply_runtime_state()`` → apply it to the worker + Works with any worker that implements AdaptiveSpecWorker protocol: + - build_adaptive_runtime_state(steps, draft_tokens) → runtime state + - apply_runtime_state(state) → apply it to the worker The worker only needs to: - 1. Call ``register()`` for the initial state, then ``init_states()`` + 1. Call register() for the initial state, then init_states() once during startup. - 2. Call ``on_verify_complete()`` after each decode verify. + 2. Call on_verify_complete(num_correct_drafts_per_req) after each decode verify. """ def __init__(self, worker: AdaptiveSpecWorker, config_path: str | None = None): @@ -86,7 +86,7 @@ class AdaptiveController: def register(self, state: SpecRuntimeState, steps: int | None = None) -> None: """Register a pre-built runtime state. - *steps* defaults to ``state.speculative_num_steps`` when not given. + *steps* defaults to state.speculative_num_steps when not given. """ key = steps if steps is not None else state.speculative_num_steps self._states[key] = state diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index 96c7286af..ee8842869 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -1,24 +1,17 @@ from __future__ import annotations -import bisect import contextlib from dataclasses import dataclass from typing import TYPE_CHECKING, Callable, Optional import torch +from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config from sglang.srt.environ import envs -from sglang.srt.layers.dp_attention import DpPaddingMode, set_dp_buffer_len -from sglang.srt.model_executor.cuda_graph_runner import ( - CUDA_GRAPH_CAPTURE_FAILED_MSG, - CudaGraphRunner, - DeepEPCudaGraphRunnerAdapter, - get_batch_sizes_to_capture, - get_global_graph_memory_pool, - model_capture_mode, - set_global_graph_memory_pool, +from sglang.srt.layers.dp_attention import ( + DpPaddingMode, + set_dp_buffer_len, set_is_extend_in_batch, - set_torch_compile_config, ) from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, @@ -27,6 +20,16 @@ from sglang.srt.model_executor.forward_batch_info import ( ) from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.input_buffers import ForwardInputBuffers +from sglang.srt.model_executor.runner import ( + DecodeCudaGraphRunner, + DeepEPCudaGraphRunnerAdapter, + get_batch_sizes_to_capture, + model_capture_mode, +) +from sglang.srt.model_executor.runner_backend import FullCudaGraphBackend +from sglang.srt.model_executor.runner_backend_utils import ( + CUDA_GRAPH_CAPTURE_FAILED_MSG, +) from sglang.srt.speculative.eagle_info import EagleDraftInput from sglang.srt.utils import ( require_attn_tp_gather, @@ -59,7 +62,22 @@ class EagleDraftInputBuffers(ForwardInputBuffers): global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] -class EAGLEDraftCudaGraphRunner: +class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): + """EAGLE draft cuda-graph runner. + + Subclasses DecodeCudaGraphRunner to inherit the outer capture + loop (capture()), bucket-padding helper (_pad_to_bucket), + and the backend-driven capture/replay scaffolding. EAGLE-specific + bits — buffer dataclass, dummy ForwardBatch construction in + capture_one_shape, replay output unwrap, and can_run — are + overridden. + + EAGLE does not call DecodeCudaGraphRunner.__init__ (that init + sets up many decode-only fields like SWA/encoder-decoder/MLA-aware + state). Instead it sets up its own state directly while making sure + the parent's capture() / backend contract is satisfied. + """ + def __init__( self, eagle_worker: EagleDraftWorker, @@ -74,16 +92,22 @@ class EAGLEDraftCudaGraphRunner: self.model_runner = model_runner = eagle_worker.draft_runner else: self.model_runner = model_runner = eagle_worker.model_runner - self.graphs = {} - self.output_buffers = {} + + # Fields the parent's capture() reads: + self.device = model_runner.device + self.device_module = torch.get_device_module(self.device) + self.tp_size = model_runner.tp_size + self.dp_size = model_runner.dp_size + self.pp_size = model_runner.server_args.pp_size self.enable_torch_compile = model_runner.server_args.enable_torch_compile self.disable_padding = model_runner.server_args.disable_cuda_graph_padding self.require_gathered_buffer = require_gathered_buffer(model_runner.server_args) self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args) self.require_mlp_sync = require_mlp_sync(model_runner.server_args) self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args) - self.tp_size = self.model_runner.tp_size - self.dp_size = self.model_runner.dp_size + self.enable_profile_cuda_graph = ( + model_runner.server_args.enable_profile_cuda_graph + ) self.speculative_num_steps = ( model_runner.server_args.speculative_num_steps if speculative_num_steps is None @@ -91,33 +115,41 @@ class EAGLEDraftCudaGraphRunner: ) self.topk = model_runner.server_args.speculative_eagle_topk self.draft_attn_backend = draft_attn_backend or model_runner.draft_attn_backend - self.enable_profile_cuda_graph = ( - model_runner.server_args.enable_profile_cuda_graph - ) + + # Patch_model in parent's capture() needs an attn_backend reference. + # EAGLE doesn't use it (capture_one_shape calls draft_forward instead), + # but the field must exist. + self.attn_backend = self.draft_attn_backend + + # Disable parent paths that don't apply to EAGLE. + self.compile_bs = [] # disables patch_model torch.compile wrapping self.enable_pdmux = False + self.record_nolora_graph = False + self.is_dllm = False + self.deepep_adapter = DeepEPCudaGraphRunnerAdapter() - # Batch sizes to capture - self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(model_runner) + # Capture-time globals required by parent's capture_one_shape signature. + self.capture_forward_mode = ForwardMode.DECODE + self.capture_hidden_mode = CaptureHiddenMode.LAST - # Attention backend + # Bucket sizes + self.capture_bs, _ = get_batch_sizes_to_capture(model_runner) self.num_tokens_per_bs = self.topk self.max_bs = max(self.capture_bs) self.max_num_token = self.max_bs * self.num_tokens_per_bs + # Attention backend init self.draft_attn_backend.init_cuda_graph_state(self.max_bs, self.max_num_token) self.seq_len_fill_value = self.draft_attn_backend.attn_backends[ 0 ].get_cuda_graph_seq_len_fill_value() - seq_lens_cpu = torch.full( - (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32 - ) self.extend_seq_lens_cpu = [self.seq_len_fill_value] * self.max_bs if self.enable_torch_compile: set_torch_compile_config() - # Graph inputs + # Static buffers with torch.device(model_runner.device): input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64) req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int64) @@ -171,6 +203,10 @@ class EAGLEDraftCudaGraphRunner: global_num_tokens_gpu = None global_num_tokens_for_logprob_gpu = None + seq_lens_cpu = torch.full( + (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32, device="cpu" + ) + self.buffers = EagleDraftInputBuffers( input_ids=input_ids, req_pool_indices=req_pool_indices, @@ -190,6 +226,12 @@ class EAGLEDraftCudaGraphRunner: ) self.buffers.share_buffers() + # Backend (Full CUDA graph capture) + self.backend = FullCudaGraphBackend( + self, + enable_memory_saver=model_runner.server_args.enable_memory_saver, + ) + # Capture try: with model_capture_mode(): @@ -199,9 +241,19 @@ class EAGLEDraftCudaGraphRunner: f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}" ) + # ----------------------------------------------------------------- + # Helpers + # ----------------------------------------------------------------- def _cache_loc_dtype(self): return torch.int64 + def _make_graph_key(self, bs, stream_idx=None, variant_label=None): + # EAGLE doesn't use stream_idx / lora variants; key is just bs. + return bs + + # ----------------------------------------------------------------- + # can_run + # ----------------------------------------------------------------- def can_run(self, forward_batch: ForwardBatch): if self.require_mlp_tp_gather: cuda_graph_bs = ( @@ -214,7 +266,7 @@ class EAGLEDraftCudaGraphRunner: cuda_graph_bs = forward_batch.batch_size is_bs_supported = ( - cuda_graph_bs in self.graphs + self.backend.can_run(forward_batch, cuda_graph_bs) if self.disable_padding else cuda_graph_bs <= self.max_bs ) @@ -224,45 +276,18 @@ class EAGLEDraftCudaGraphRunner: return is_bs_supported - def _create_graph(self): - return torch.cuda.CUDAGraph() - - def _capture_init(self, run_once_fn): - for _ in range(2): - torch.cuda.synchronize() - self.model_runner.tp_group.barrier() - run_once_fn() - hook = getattr( - self.model_runner.draft_attn_backend, - "on_after_cuda_graph_warmup", - None, - ) - if hook is not None: - hook() - - def _capture_graph(self, graph, pool, stream, run_once_fn): - with torch.cuda.graph(graph, pool=pool, stream=stream): - out = run_once_fn() - return out - - def _replay(self, forward_batch: ForwardBatch): - ctx = ( - self.model_runner.device_timer.wrap(metadata={"category": "eagle_draft"}) - if self.model_runner.device_timer - else contextlib.nullcontext() - ) - with ctx: - self.graphs[self.bs].replay() - - def capture(self): - CudaGraphRunner.capture(self) - - def capture_one_batch_size( - self, num_seqs: int, forward: Callable, stream_idx: int = 0 + # ----------------------------------------------------------------- + # Capture (per-shape) + # ----------------------------------------------------------------- + def capture_one_shape( + self, + size: int, + forward: Callable, + stream_idx: Optional[int] = None, + variant_label: Optional[str] = None, ): + num_seqs = size # EAGLE legacy name buffers = self.buffers - graph = self._create_graph() - stream = self.stream num_tokens = num_seqs * self.num_tokens_per_bs # Graph inputs @@ -323,7 +348,6 @@ class EAGLEDraftCudaGraphRunner: capture_hidden_mode=capture_mode, ) - # Forward batch forward_batch = ForwardBatch( forward_mode=ForwardMode.DECODE, batch_size=num_seqs, @@ -381,19 +405,23 @@ class EAGLEDraftCudaGraphRunner: # per-step forwards inside draft_forward must not re-plan. forward_batch.mark_forward_metadata_ready() self.deepep_adapter.capture(is_extend_in_batch=False) - self._capture_init(run_once) - out = self._capture_graph( - graph, get_global_graph_memory_pool(), stream, run_once + shape_key = self._make_graph_key(num_seqs) + self.backend.capture_one( + shape_key, + run_once, + dummies=None, + post_warmup_hook=getattr( + self.draft_attn_backend, "on_after_cuda_graph_warmup", None + ), ) - set_global_graph_memory_pool(graph.pool()) - return graph, out - def _postprocess_output_to_raw_bs(self, out, raw_bs): - # Keep the variables name for readability parent_list, top_scores_index, draft_tokens = (t[:raw_bs] for t in out) return parent_list, top_scores_index, draft_tokens + # ----------------------------------------------------------------- + # Replay + # ----------------------------------------------------------------- def replay(self, forward_batch: ForwardBatch): assert forward_batch.out_cache_loc is not None self.deepep_adapter.replay() @@ -402,7 +430,7 @@ class EAGLEDraftCudaGraphRunner: raw_bs = forward_batch.batch_size raw_num_token = raw_bs * self.num_tokens_per_bs - # Pad + # Pad to nearest captured shape if self.require_mlp_tp_gather: max_num_tokens = max(forward_batch.global_num_tokens_cpu) max_batch_size = ( @@ -411,11 +439,10 @@ class EAGLEDraftCudaGraphRunner: or self.model_runner.spec_algorithm.is_standalone() else max_num_tokens ) - index = bisect.bisect_left(self.capture_bs, max_batch_size) + bs = self._pad_to_bucket(int(max_batch_size), self.capture_bs) else: - index = bisect.bisect_left(self.capture_bs, raw_bs) + bs = self._pad_to_bucket(raw_bs, self.capture_bs) - bs = self.capture_bs[index] if bs != raw_bs: buffers.seq_lens.fill_(self.seq_len_fill_value) buffers.out_cache_loc.zero_() @@ -472,7 +499,6 @@ class EAGLEDraftCudaGraphRunner: buffers.global_num_tokens_gpu.fill_(bs * self.num_tokens_per_bs) buffers.global_num_tokens_for_logprob_gpu.fill_(bs * self.num_tokens_per_bs) - # Attention backend if bs != raw_bs: forward_batch.batch_size = bs forward_batch.seq_lens = buffers.seq_lens[:bs] @@ -498,11 +524,16 @@ class EAGLEDraftCudaGraphRunner: self.draft_attn_backend.init_forward_metadata_out_graph(forward_batch) self.raw_bs = raw_bs self.bs = bs - # TODO: The forward_batch.seq_len_sum might need to be updated to reflect the padding in the cuda graph - # Replay - self._replay(forward_batch) - out = self.output_buffers[bs] + # Replay via backend + shape_key = self._make_graph_key(bs) + timer_ctx = ( + self.model_runner.device_timer.wrap(metadata={"category": "eagle_draft"}) + if self.model_runner.device_timer + else contextlib.nullcontext() + ) + with timer_ctx: + out = self.backend.replay(shape_key, forward_batch) if bs != raw_bs: out = self._postprocess_output_to_raw_bs(out, raw_bs) diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 8143e60f0..6ed2a08ee 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -1,25 +1,18 @@ from __future__ import annotations -import bisect import contextlib from dataclasses import dataclass from typing import TYPE_CHECKING, Callable, Optional import torch -from sglang.srt.layers.dp_attention import DpPaddingMode, set_dp_buffer_len -from sglang.srt.model_executor.cuda_graph_runner import ( - CUDA_GRAPH_CAPTURE_FAILED_MSG, - CudaGraphRunner, - DeepEPCudaGraphRunnerAdapter, - LogitsProcessorOutput, - get_batch_sizes_to_capture, - get_global_graph_memory_pool, - model_capture_mode, - set_global_graph_memory_pool, +from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config +from sglang.srt.layers.dp_attention import ( + DpPaddingMode, + set_dp_buffer_len, set_is_extend_in_batch, - set_torch_compile_config, ) +from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, @@ -27,6 +20,16 @@ from sglang.srt.model_executor.forward_batch_info import ( ) from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.input_buffers import ForwardInputBuffers +from sglang.srt.model_executor.runner import ( + DecodeCudaGraphRunner, + DeepEPCudaGraphRunnerAdapter, + get_batch_sizes_to_capture, + model_capture_mode, +) +from sglang.srt.model_executor.runner_backend import FullCudaGraphBackend +from sglang.srt.model_executor.runner_backend_utils import ( + CUDA_GRAPH_CAPTURE_FAILED_MSG, +) from sglang.srt.speculative.eagle_info import EagleDraftExtendInput from sglang.srt.speculative.spec_utils import fast_topk from sglang.srt.utils import ( @@ -61,7 +64,14 @@ class EagleDraftExtendInputBuffers(ForwardInputBuffers): global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] -class EAGLEDraftExtendCudaGraphRunner: +class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): + """EAGLE draft-extend cuda-graph runner. + + Subclasses DecodeCudaGraphRunner to inherit the outer capture + loop + backend scaffolding. Overrides capture_one_shape, + replay, can_run for EAGLE-specific draft-extend semantics. + """ + def __init__( self, eagle_worker: EagleDraftWorker, @@ -72,23 +82,27 @@ class EAGLEDraftExtendCudaGraphRunner: # Parse args self.eagle_worker = eagle_worker if not hasattr(eagle_worker, "model_runner"): - # V2: EagleDraftWorker self.model_runner = model_runner = eagle_worker.draft_runner self.forward_mode = ForwardMode.DRAFT_EXTEND_V2 else: self.model_runner = model_runner = eagle_worker.model_runner self.forward_mode = ForwardMode.DRAFT_EXTEND - self.graphs = {} - self.output_buffers = {} + # Fields the parent's capture() reads: + self.device = model_runner.device + self.device_module = torch.get_device_module(self.device) + self.tp_size = model_runner.tp_size + self.dp_size = model_runner.dp_size + self.pp_size = model_runner.server_args.pp_size self.enable_torch_compile = model_runner.server_args.enable_torch_compile self.disable_padding = model_runner.server_args.disable_cuda_graph_padding self.require_gathered_buffer = require_gathered_buffer(model_runner.server_args) self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args) self.require_mlp_sync = require_mlp_sync(model_runner.server_args) self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args) - self.tp_size = self.model_runner.tp_size - self.dp_size = self.model_runner.dp_size + self.enable_profile_cuda_graph = ( + model_runner.server_args.enable_profile_cuda_graph + ) self.speculative_num_steps = ( model_runner.server_args.speculative_num_steps if speculative_num_steps is None @@ -98,16 +112,22 @@ class EAGLEDraftExtendCudaGraphRunner: self.draft_extend_attn_backend = ( draft_extend_attn_backend or eagle_worker.draft_extend_attn_backend ) - self.enable_profile_cuda_graph = ( - model_runner.server_args.enable_profile_cuda_graph - ) + self.attn_backend = self.draft_extend_attn_backend + + # Disable parent paths that don't apply. + self.compile_bs = [] self.enable_pdmux = False + self.record_nolora_graph = False + self.is_dllm = False + self.deepep_adapter = DeepEPCudaGraphRunnerAdapter() - self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(model_runner) + self.capture_forward_mode = self.forward_mode + self.capture_hidden_mode = CaptureHiddenMode.LAST + + self.capture_bs, _ = get_batch_sizes_to_capture(model_runner) self.padded_static_len = -1 - # Attention backend # Size cuda-graph buffers by num_draft_tokens (full tree width), not # num_steps + 1, or topk > 1 draft-extend overflows them. self.num_tokens_per_bs = model_runner.server_args.speculative_num_draft_tokens @@ -120,15 +140,11 @@ class EAGLEDraftExtendCudaGraphRunner: self.seq_len_fill_value = ( self.draft_extend_attn_backend.get_cuda_graph_seq_len_fill_value() ) - seq_lens_cpu = torch.full( - (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32 - ) self.extend_seq_lens_cpu = [self.num_tokens_per_bs] * self.max_bs if self.enable_torch_compile: set_torch_compile_config() - # Graph inputs with torch.device(model_runner.device): input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64) req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int64) @@ -204,6 +220,10 @@ class EAGLEDraftExtendCudaGraphRunner: dtype=torch.float, ) + seq_lens_cpu = torch.full( + (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32, device="cpu" + ) + self.buffers = EagleDraftExtendInputBuffers( input_ids=input_ids, req_pool_indices=req_pool_indices, @@ -222,7 +242,11 @@ class EAGLEDraftExtendCudaGraphRunner: ) self.buffers.share_buffers() - # Capture + self.backend = FullCudaGraphBackend( + self, + enable_memory_saver=model_runner.server_args.enable_memory_saver, + ) + try: with model_capture_mode(): self.capture() @@ -231,6 +255,12 @@ class EAGLEDraftExtendCudaGraphRunner: f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}" ) + def _cache_loc_dtype(self): + return torch.int64 + + def _make_graph_key(self, bs, stream_idx=None, variant_label=None): + return bs + def can_run(self, forward_batch: ForwardBatch): if self.require_mlp_tp_gather: cuda_graph_bs = ( @@ -243,7 +273,7 @@ class EAGLEDraftExtendCudaGraphRunner: cuda_graph_bs = forward_batch.seq_lens.numel() is_bs_supported = ( - cuda_graph_bs in self.graphs + self.backend.can_run(forward_batch, cuda_graph_bs) if self.disable_padding else cuda_graph_bs <= self.max_bs ) @@ -253,41 +283,15 @@ class EAGLEDraftExtendCudaGraphRunner: return is_bs_supported - def _create_graph(self): - return torch.cuda.CUDAGraph() - - def _cache_loc_dtype(self): - return torch.int64 - - def _capture_init(self, run_once_fn): - for _ in range(2): - torch.cuda.synchronize() - self.model_runner.tp_group.barrier() - run_once_fn() - - def _capture_graph(self, graph, pool, stream, run_once_fn): - with torch.cuda.graph(graph, pool=pool, stream=stream): - out = run_once_fn() - return out - - def _replay(self, forward_batch: ForwardBatch): - ctx = ( - self.model_runner.device_timer.wrap( - metadata={"category": "eagle_draft_extend"} - ) - if self.model_runner.device_timer - else contextlib.nullcontext() - ) - with ctx: - self.graphs[self.bs].replay() - - def capture(self): - CudaGraphRunner.capture(self) - - def capture_one_batch_size(self, bs: int, forward: Callable, stream_idx: int = 0): + def capture_one_shape( + self, + size: int, + forward: Callable, + stream_idx: Optional[int] = None, + variant_label: Optional[str] = None, + ): + bs = size buffers = self.buffers - graph = self._create_graph() - stream = self.stream num_tokens = bs * self.num_tokens_per_bs # Graph inputs @@ -349,7 +353,6 @@ class EAGLEDraftExtendCudaGraphRunner: num_accept_tokens=num_accept_tokens, ) - # Forward batch forward_batch = ForwardBatch( forward_mode=self.forward_mode, batch_size=bs, @@ -386,7 +389,6 @@ class EAGLEDraftExtendCudaGraphRunner: ) set_is_extend_in_batch(False) - # Backup two fields, which will be modified in-place in `draft_forward`. output_cache_loc_backup = forward_batch.out_cache_loc hidden_states_backup = forward_batch.spec_info.hidden_states @@ -418,29 +420,29 @@ class EAGLEDraftExtendCudaGraphRunner: forward_batch, in_capture=True ) self.deepep_adapter.capture(is_extend_in_batch=True) - canary_ctx = ( c.with_active_single_forward_manager(0) if (c := self.model_runner.canary_manager) is not None else contextlib.nullcontext() ) with canary_ctx: - self._capture_init(run_once) - - out = self._capture_graph( - graph, get_global_graph_memory_pool(), stream, run_once + shape_key = self._make_graph_key(bs) + self.backend.capture_one( + shape_key, + run_once, + dummies=None, + post_warmup_hook=getattr( + self.draft_extend_attn_backend, + "on_after_cuda_graph_warmup", + None, + ), ) - set_global_graph_memory_pool(graph.pool()) - return graph, out - def replay(self, forward_batch: ForwardBatch): assert forward_batch.out_cache_loc is not None self.deepep_adapter.replay() buffers = self.buffers - # batch_size and num_seqs can be different in case there are finished examples - # in the batch, which will not be counted as num_seqs raw_bs = forward_batch.batch_size num_tokens = forward_batch.input_ids.shape[0] if self.require_mlp_tp_gather: @@ -450,11 +452,10 @@ class EAGLEDraftExtendCudaGraphRunner: if self.model_runner.spec_algorithm.is_eagle() else max_num_tokens ) - index = bisect.bisect_left(self.capture_bs, max_batch_size) + bs = self._pad_to_bucket(int(max_batch_size), self.capture_bs) else: - index = bisect.bisect_left(self.capture_bs, raw_bs) + bs = self._pad_to_bucket(raw_bs, self.capture_bs) - bs = self.capture_bs[index] if bs * self.num_tokens_per_bs != num_tokens: buffers.seq_lens.fill_(self.seq_len_fill_value) buffers.out_cache_loc.zero_() @@ -466,7 +467,6 @@ class EAGLEDraftExtendCudaGraphRunner: buffers.num_accept_tokens.fill_(self.num_tokens_per_bs) buffers.extend_seq_lens.fill_(self.num_tokens_per_bs) - # Common inputs buffers.input_ids[:num_tokens].copy_(forward_batch.input_ids) buffers.seq_lens[:raw_bs].copy_(forward_batch.seq_lens) if forward_batch.extend_seq_lens is not None: @@ -496,7 +496,6 @@ class EAGLEDraftExtendCudaGraphRunner: # TODO(ch-wan): support num_token_non_padded if self.require_gathered_buffer: buffers.global_num_tokens_gpu.fill_(bs * self.num_tokens_per_bs) - # V1: pruned_states = bs; V2: pruned_states = num_tokens if self.forward_mode.is_draft_extend_v2(): buffers.global_num_tokens_for_logprob_gpu.fill_( bs * self.num_tokens_per_bs @@ -546,14 +545,20 @@ class EAGLEDraftExtendCudaGraphRunner: ) self.draft_extend_attn_backend.init_forward_metadata_out_graph(fb_view) - # Replay self.raw_bs = raw_bs self.bs = bs - self._replay(forward_batch) - out = self.output_buffers[bs] + shape_key = self._make_graph_key(bs) + timer_ctx = ( + self.model_runner.device_timer.wrap( + metadata={"category": "eagle_draft_extend"} + ) + if self.model_runner.device_timer + else contextlib.nullcontext() + ) + with timer_ctx: + out = self.backend.replay(shape_key, forward_batch) if self.forward_mode == ForwardMode.DRAFT_EXTEND_V2: - # DRAFT_EXTEND_V2: all tokens calculations whether accepted or not. unpadding_bs = num_tokens elif bs != raw_bs: forward_batch.spec_info.num_correct_drafts = buffers.num_correct_drafts[ diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index bbe7ed9c9..bac1a6ff9 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -430,11 +430,15 @@ class EagleVerifyInputV2Mixin: # Run attention backend plan and cuda graph preparation can_run_cuda_graph = bool( - target_worker.model_runner.graph_runner - and target_worker.model_runner.graph_runner.can_run(verify_forward_batch) + target_worker.model_runner.decode_cuda_graph_runner + and target_worker.model_runner.decode_cuda_graph_runner.can_run( + verify_forward_batch + ) ) if can_run_cuda_graph: - target_worker.model_runner.graph_runner.replay_prepare(verify_forward_batch) + target_worker.model_runner.decode_cuda_graph_runner.replay_prepare( + verify_forward_batch + ) verify_forward_batch.mark_forward_metadata_ready() # Non-cuda-graph: defer init to forward_extend, which runs after # `_forward_raw -> prepare_mlp_sync_batch` pads the batch. Initing diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index e7b9e4688..130ede959 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -34,9 +34,14 @@ from sglang.srt.managers.io_struct import ( from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker -from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + Phase, + check_cuda_graph_backend, +) from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardBatch from sglang.srt.model_executor.forward_context import ForwardContext, forward_context +from sglang.srt.model_executor.runner import DecodeCudaGraphRunner from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.adaptive_runtime_state import ( AdaptiveController, @@ -151,8 +156,8 @@ class EagleDraftWorker(BaseDraftWorker): # Do not capture cuda graph in `TpModelWorker` init, # will capture later with init_cuda_graphs() - backup_disable_cuda_graph = server_args.disable_cuda_graph - server_args.disable_cuda_graph = True + backup_decode_mode = server_args.cuda_graph_config.decode.backend + server_args.cuda_graph_config.decode.backend = Backend.DISABLED # Share the allocator with a target worker. # Draft and target worker own their own KV cache pools. @@ -199,7 +204,9 @@ class EagleDraftWorker(BaseDraftWorker): self.init_lm_head() # Init attention backend and cuda graphs - self.draft_runner.server_args.disable_cuda_graph = backup_disable_cuda_graph + self.draft_runner.server_args.cuda_graph_config.decode.backend = ( + backup_decode_mode + ) self.draft_tp_context = ( draft_tp_context if server_args.enable_dp_attention else empty_context ) @@ -209,10 +216,8 @@ class EagleDraftWorker(BaseDraftWorker): speculative_moe_a2a_backend_context(), ): self.init_attention_backend() - if server_args.enable_breakable_cuda_graph: - self.draft_runner.init_piecewise_cuda_graphs( - force_for_draft_worker=True - ) + if check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE): + self.draft_runner.init_prefill_cuda_graph(force_for_draft_worker=True) self.init_cuda_graphs() if (c := self.draft_runner.canary_manager) is not None: @@ -236,8 +241,13 @@ class EagleDraftWorker(BaseDraftWorker): ) num_steps = self.speculative_num_steps sa = self.server_args + decode_max_bs = ( + sa.cuda_graph_config.decode.max_bs + if sa.cuda_graph_config is not None + else None + ) max_bs = max( - sa.cuda_graph_max_bs or 0, + decode_max_bs or 0, sa.max_running_requests or 0, 1, ) @@ -323,7 +333,7 @@ class EagleDraftWorker(BaseDraftWorker): self.cuda_graph_runner = None self.cuda_graph_runner_for_draft_extend = None - if self.server_args.disable_cuda_graph: + if check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED): return if self.server_args.model_impl == "mindspore": @@ -885,7 +895,7 @@ class EAGLEWorkerV2(BaseSpecWorker): draft_attn_backend=self._draft_worker.draft_attn_backend, cuda_graph_runner=self._draft_worker.cuda_graph_runner, target_attn_backend=self._target_worker.model_runner.attn_backend, - target_graph_runner=self._target_worker.model_runner.graph_runner, + target_graph_runner=self._target_worker.model_runner.decode_cuda_graph_runner, draft_extend_attn_backend=self._draft_worker.draft_extend_attn_backend, cuda_graph_runner_for_draft_extend=self._draft_worker.cuda_graph_runner_for_draft_extend, ) @@ -894,7 +904,7 @@ class EAGLEWorkerV2(BaseSpecWorker): cuda_graph_bs=( None if self.server_args.disable_cuda_graph - else self.server_args.cuda_graph_bs + else self.server_args.cuda_graph_bs_decode ), ) @@ -1042,8 +1052,10 @@ class EAGLEWorkerV2(BaseSpecWorker): target_model_runner.init_new_workspace = backup_init target_graph_runner = None - if not self.server_args.disable_cuda_graph: - TargetGraphRunnerCls = NPUGraphRunner if _is_npu else CudaGraphRunner + if not check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED): + TargetGraphRunnerCls = ( + NPUGraphRunner if _is_npu else DecodeCudaGraphRunner + ) target_graph_runner = TargetGraphRunnerCls( target_model_runner, attn_backend=target_attn_backend, @@ -1102,7 +1114,9 @@ class EAGLEWorkerV2(BaseSpecWorker): # Target side self._target_worker.model_runner.attn_backend = state.target_attn_backend - self._target_worker.model_runner.graph_runner = state.target_graph_runner + self._target_worker.model_runner.decode_cuda_graph_runner = ( + state.target_graph_runner + ) # Sync server_args self.server_args.speculative_num_steps = state.speculative_num_steps @@ -1132,7 +1146,7 @@ class EAGLEWorkerV2(BaseSpecWorker): dw.cuda_graph_runner_for_draft_extend, sa.speculative_num_steps, sa.speculative_num_draft_tokens, - sa.cuda_graph_bs, + sa.cuda_graph_bs_decode, sa.disable_cuda_graph, ) @@ -1143,7 +1157,7 @@ class EAGLEWorkerV2(BaseSpecWorker): sa.speculative_num_steps = speculative_num_steps sa.speculative_num_draft_tokens = speculative_num_draft_tokens if cuda_graph_bs is not None: - sa.cuda_graph_bs = cuda_graph_bs + sa.cuda_graph_bs_decode = cuda_graph_bs # BS-aware adaptive spec may prune cuda_graph_bs to an empty list # for steps that no BS range uses (e.g. step=1). Disable graph # capture for those steps; restore in finally so subsequent steps @@ -1167,7 +1181,7 @@ class EAGLEWorkerV2(BaseSpecWorker): dw.cuda_graph_runner_for_draft_extend, sa.speculative_num_steps, sa.speculative_num_draft_tokens, - sa.cuda_graph_bs, + sa.cuda_graph_bs_decode, sa.disable_cuda_graph, ) = backup dw._rebuild_topk1_chain_buffers() @@ -1218,7 +1232,7 @@ class EAGLEWorkerV2(BaseSpecWorker): self.target_worker.model_runner.attn_backend.update_verify_buffers_to_fill_after_draft( verify_input, ( - self.target_worker.model_runner.graph_runner.bs + self.target_worker.model_runner.decode_cuda_graph_runner.bs if can_run_cuda_graph else None ), @@ -1421,7 +1435,7 @@ class EAGLEWorkerV2(BaseSpecWorker): batch: The batch to run. accept_index: The index of the accepted tokens (incl. bonus). num_correct_drafts: Per-req count of correct drafts (excludes bonus); - seq_lens is advanced by ``num_correct_drafts + 1`` to cover the bonus slot. + seq_lens is advanced by num_correct_drafts + 1 to cover the bonus slot. """ bs = len(batch.seq_lens) # accept_index element count, NOT bs * num_draft_tokens: for topk > 1 the diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index dffdcea45..91a9c3d83 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py @@ -5,18 +5,15 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Callable, Optional import torch +import tqdm -from sglang.srt.layers.dp_attention import DpPaddingMode, set_dp_buffer_len -from sglang.srt.model_executor.cuda_graph_runner import ( - CUDA_GRAPH_CAPTURE_FAILED_MSG, - CudaGraphRunner, - DeepEPCudaGraphRunnerAdapter, - get_batch_sizes_to_capture, - get_global_graph_memory_pool, - model_capture_mode, - set_global_graph_memory_pool, +from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config +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, - set_torch_compile_config, ) from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, @@ -25,6 +22,17 @@ from sglang.srt.model_executor.forward_batch_info import ( ) from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.input_buffers import ForwardInputBuffers +from sglang.srt.model_executor.runner import ( + DeepEPCudaGraphRunnerAdapter, + freeze_gc, + get_batch_sizes_to_capture, + get_global_graph_memory_pool, + model_capture_mode, + set_global_graph_memory_pool, +) +from sglang.srt.model_executor.runner_backend_utils import ( + CUDA_GRAPH_CAPTURE_FAILED_MSG, +) from sglang.srt.speculative.frozen_kv_mtp_info import FrozenKVMTPDraftInput from sglang.srt.utils import ( require_attn_tp_gather, @@ -192,7 +200,18 @@ class FrozenKVMTPCudaGraphRunner: self.graphs[self.bs].replay() def capture(self): - CudaGraphRunner.capture(self) + with freeze_gc(self.model_runner.server_args.enable_cudagraph_gc): + with graph_capture() as graph_capture_context: + self.stream = graph_capture_context.stream + capture_range = ( + tqdm.tqdm(list(reversed(self.capture_bs))) + if get_tensor_model_parallel_rank() == 0 + else reversed(self.capture_bs) + ) + for bs in capture_range: + graph, output_buffers = self.capture_one_batch_size(bs, None) + self.graphs[bs] = graph + self.output_buffers[bs] = output_buffers def capture_one_batch_size( self, num_seqs: int, forward: Callable, stream_idx: int = 0 @@ -288,7 +307,7 @@ class FrozenKVMTPCudaGraphRunner: # Swap the draft backend's token_to_kv_pool to the frozen target pool # for the capture; the single backend-attr swap is seen by both - # ``get_token_to_kv_pool()`` (via ``get_attn_backend()``) and the + # get_token_to_kv_pool() (via get_attn_backend()) and the # backend's own reads. target_pool = self.frozen_kv_mtp_worker.kv_context.target_token_to_kv_pool saved_backend_pool = self.draft_attn_backend.token_to_kv_pool diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index 8d9cbc19d..fb39f48d6 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -32,6 +32,10 @@ from sglang.srt.layers.moe.utils import ( ) from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + cuda_graph_fully_disabled, +) from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, @@ -114,8 +118,8 @@ class FrozenKVMTPDraftWorker(BaseDraftWorker, TpModelWorker): ) # Defer cuda graph capture; we do it ourselves below. - backup_disable_cuda_graph = server_args.disable_cuda_graph - server_args.disable_cuda_graph = True + backup_decode_mode = server_args.cuda_graph_config.decode.backend + server_args.cuda_graph_config.decode.backend = Backend.DISABLED # Draft attention uses target req_to_token + KV allocator (read-only). self.req_to_token_pool, self.token_to_kv_pool_allocator = ( @@ -166,8 +170,8 @@ class FrozenKVMTPDraftWorker(BaseDraftWorker, TpModelWorker): if hasattr(self.draft_model_runner.model, "bind_frozen_kv_context"): self._bind_kv_context() - self.draft_model_runner.server_args.disable_cuda_graph = ( - backup_disable_cuda_graph + self.draft_model_runner.server_args.cuda_graph_config.decode.backend = ( + backup_decode_mode ) self.draft_tp_context = ( @@ -324,7 +328,7 @@ class FrozenKVMTPDraftWorker(BaseDraftWorker, TpModelWorker): self.draft_attn_backend.init_forward_metadata_out_graph(fb_view) def init_cuda_graphs(self) -> None: - if self.server_args.disable_cuda_graph or self.speculative_num_steps <= 1: + if cuda_graph_fully_disabled() or self.speculative_num_steps <= 1: return if self.target_worker.device != "cuda": logger.info( diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index ec465e3b3..f138cac95 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py @@ -14,7 +14,6 @@ from __future__ import annotations -import bisect import logging import time from dataclasses import dataclass @@ -22,18 +21,17 @@ from typing import TYPE_CHECKING, Callable, List, Optional import torch -from sglang.srt.layers.dp_attention import DpPaddingMode, set_dp_buffer_len -from sglang.srt.model_executor.cuda_graph_runner import ( - CUDA_GRAPH_CAPTURE_FAILED_MSG, - CudaGraphRunner, - DeepEPCudaGraphRunnerAdapter, - LogitsProcessorOutput, - get_batch_sizes_to_capture, - get_global_graph_memory_pool, - model_capture_mode, - set_global_graph_memory_pool, +from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config +from sglang.srt.layers.dp_attention import ( + DpPaddingMode, + set_dp_buffer_len, set_is_extend_in_batch, - set_torch_compile_config, +) +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + Phase, + check_cuda_graph_backend, ) from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, @@ -46,6 +44,16 @@ from sglang.srt.model_executor.forward_context import ( get_req_to_token_pool, ) from sglang.srt.model_executor.input_buffers import ForwardInputBuffers +from sglang.srt.model_executor.runner import ( + DecodeCudaGraphRunner, + DeepEPCudaGraphRunnerAdapter, + get_batch_sizes_to_capture, + model_capture_mode, +) +from sglang.srt.model_executor.runner_backend import FullCudaGraphBackend +from sglang.srt.model_executor.runner_backend_utils import ( + CUDA_GRAPH_CAPTURE_FAILED_MSG, +) from sglang.srt.speculative.eagle_info import EagleDraftExtendInput from sglang.srt.speculative.multi_layer_eagle_utils import assign_new_state_triton from sglang.srt.speculative.spec_utils import fast_topk @@ -88,7 +96,17 @@ class MultiLayerEagleDraftExtendInputBuffers(ForwardInputBuffers): global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] -class MultiLayerEagleDraftExtendCudaGraphRunner: +class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): + """Per-step multi-layer EAGLE draft-extend runner. + + Subclasses DecodeCudaGraphRunner. Shares buffers across steps + via the composite MultiLayerEagleMultiStepDraftExtendCudaGraphRunner, + so initialization is split: __init__ does basic field setup, + init_buffers_and_capture (called by the composite once shared + buffers exist) finishes by allocating per-step buffers and running + capture. + """ + def __init__(self, eagle_worker: MultiLayerEagleDraftWorker, step: int): # Parse args self.step = step @@ -96,16 +114,18 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: self.model_runner = model_runner = eagle_worker.mtp_model_runner(self.step) self.forward_mode = ForwardMode.DRAFT_EXTEND_V2 - self.graphs = {} - self.output_buffers = {} + # Fields the parent's capture() reads: + self.device = model_runner.device + self.device_module = torch.get_device_module(self.device) + self.tp_size = model_runner.tp_size + self.dp_size = model_runner.server_args.dp_size + self.pp_size = model_runner.server_args.pp_size self.enable_torch_compile = model_runner.server_args.enable_torch_compile self.disable_padding = model_runner.server_args.disable_cuda_graph_padding self.require_gathered_buffer = require_gathered_buffer(model_runner.server_args) self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args) self.require_mlp_sync = require_mlp_sync(model_runner.server_args) self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args) - self.tp_size = self.model_runner.tp_size - self.dp_size = model_runner.server_args.dp_size self.enable_pdmux = model_runner.server_args.enable_pdmux self.speculative_num_steps = model_runner.server_args.speculative_num_steps self.speculative_num_draft_tokens = ( @@ -115,10 +135,21 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: self.enable_profile_cuda_graph = ( model_runner.server_args.enable_profile_cuda_graph ) - self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(model_runner) - self.padded_static_len = -1 + self.attn_backend = self.eagle_worker.draft_extend_attn_backend_list[self.step] + + # Disable parent paths that don't apply. + self.compile_bs = [] + self.record_nolora_graph = False + self.is_dllm = False + self.deepep_adapter = DeepEPCudaGraphRunnerAdapter() + self.capture_forward_mode = self.forward_mode + self.capture_hidden_mode = CaptureHiddenMode.FULL + + self.capture_bs, _ = get_batch_sizes_to_capture(model_runner) + self.padded_static_len = -1 + # For Attention Backend self.num_tokens_per_bs = self.speculative_num_steps + 1 + step self.max_bs = max(self.capture_bs) @@ -144,10 +175,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: if self.enable_torch_compile: set_torch_compile_config() - # Graph inputs with torch.device(self.model_runner.device): - # sliced buffers - # slice according to max_num_token input_ids = cuda_graph_buffers["input_ids"][ offset : offset + self.max_num_token ] @@ -158,7 +186,6 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: offset : offset + self.max_num_token ] - # shared states seq_lens = cuda_graph_buffers["seq_lens"] req_pool_indices = cuda_graph_buffers["req_pool_indices"] num_correct_drafts = cuda_graph_buffers["num_correct_drafts"] @@ -245,7 +272,11 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu, ) - # Capture + self.backend = FullCudaGraphBackend( + self, + enable_memory_saver=self.model_runner.server_args.enable_memory_saver, + ) + try: with model_capture_mode(): self.capture() @@ -254,6 +285,9 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}" ) + def _make_graph_key(self, bs, stream_idx=None, variant_label=None): + return bs + def can_run(self, forward_batch: ForwardBatch): if self.require_mlp_tp_gather: cuda_graph_bs = ( @@ -265,7 +299,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: cuda_graph_bs = forward_batch.seq_lens.numel() is_bs_supported = ( - cuda_graph_bs in self.graphs + self.backend.can_run(forward_batch, cuda_graph_bs) if self.disable_padding else cuda_graph_bs <= self.max_bs ) @@ -275,31 +309,10 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: return is_bs_supported - def _create_graph(self): - return torch.cuda.CUDAGraph() - - def _capture_init(self, run_once_fn): - for _ in range(2): - torch.cuda.synchronize() - self.model_runner.tp_group.barrier() - run_once_fn() - - def _capture_graph(self, graph, pool, stream, run_once_fn): - with torch.cuda.graph(graph, pool=pool, stream=stream): - out = run_once_fn() - return out - - def _replay(self, forward_batch: ForwardBatch): - self.graphs[self.bs].replay() - - def capture(self): - CudaGraphRunner.capture(self) - def get_forward_batch(self, bs: int) -> ForwardBatch: buffers = self.buffers num_tokens = bs * self.num_tokens_per_bs - # Graph inputs input_ids = buffers.input_ids[:num_tokens] req_pool_indices = buffers.req_pool_indices[:bs] seq_lens = buffers.seq_lens[:bs] @@ -383,7 +396,6 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: extend_seq_lens=extend_seq_lens, extend_seq_lens_cpu=extend_seq_lens_cpu, padded_static_len=self.padded_static_len, - # added args extend_start_loc=extend_start_loc, extend_num_tokens=self.num_tokens_per_bs * bs, num_token_non_padded_cpu=self.num_tokens_per_bs * bs, @@ -391,10 +403,15 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: ) return forward_batch - def capture_one_batch_size(self, bs: int, forward: Callable, stream_idx: int = 0): + def capture_one_shape( + self, + size: int, + forward: Callable, + stream_idx: Optional[int] = None, + variant_label: Optional[str] = None, + ): + bs = size buffers = self.buffers - graph = self._create_graph() - stream = self.stream num_tokens = bs * self.num_tokens_per_bs forward_batch = self.get_forward_batch(bs) @@ -411,7 +428,6 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: ) set_is_extend_in_batch(False) - # Backup two fields, which will be modified in-place in `draft_forward`. output_cache_loc_backup = forward_batch.out_cache_loc hidden_states_backup = forward_batch.spec_info.hidden_states @@ -421,10 +437,6 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: forward_batch, ) - # Chain-style MTP: overwrite buffers.hidden_states with the draft model's - # output (hidden_states_before_norm) so that assign_new_state_triton - # propagates each MTP layer's own output to the next MTP layer, - # rather than always feeding the target model's hidden states. if ( self.eagle_worker.chain_mtp_hidden_states and ret.hidden_states is not None @@ -479,19 +491,20 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: with forward_context(ForwardContext(attn_backend=attn_backend)): attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True) self.deepep_adapter.capture(is_extend_in_batch=True) - self._capture_init(run_once) - out = self._capture_graph( - graph, get_global_graph_memory_pool(), stream, run_once + shape_key = self._make_graph_key(bs) + self.backend.capture_one( + shape_key, + run_once, + dummies=None, + post_warmup_hook=getattr( + self.attn_backend, "on_after_cuda_graph_warmup", None + ), ) - set_global_graph_memory_pool(graph.pool()) - return graph, out - def init_replay_state( self, forward_batch: ForwardBatch, bs: int, raw_bs: int, num_tokens: int ): buffers = self.buffers - # Common inputs buffers.input_ids[:num_tokens].copy_(forward_batch.input_ids) buffers.seq_lens[:raw_bs].copy_(forward_batch.seq_lens) if forward_batch.extend_seq_lens is not None: @@ -528,18 +541,13 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: self.deepep_adapter.replay() buffers = self.buffers - # batch_size and num_seqs can be different in case there are finished examples - # in the batch, which will not be counted as num_seqs raw_bs = forward_batch.batch_size num_tokens = raw_bs * self.num_tokens_per_bs - # num_tokens = forward_batch.input_ids.shape[0] if self.require_mlp_tp_gather: max_batch_size = max(forward_batch.original_global_num_tokens_cpu) - index = bisect.bisect_left(self.capture_bs, max_batch_size) + bs = self._pad_to_bucket(int(max_batch_size), self.capture_bs) else: - index = bisect.bisect_left(self.capture_bs, raw_bs) - - bs = self.capture_bs[index] + bs = self._pad_to_bucket(raw_bs, self.capture_bs) if init_state: self.init_replay_state(forward_batch, bs, raw_bs, num_tokens) @@ -575,14 +583,12 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: self.step ].init_forward_metadata_out_graph(fb_view) - # Replay self.raw_bs = raw_bs self.bs = bs - self._replay(forward_batch) - out = self.output_buffers[bs] + shape_key = self._make_graph_key(bs) + out = self.backend.replay(shape_key, forward_batch) if self.forward_mode == ForwardMode.DRAFT_EXTEND_V2: - # DRAFT_EXTEND_V2: all tokens calculations whether accepted or not. unpadding_bs = num_tokens elif bs != raw_bs: forward_batch.spec_info.num_correct_drafts = buffers.num_correct_drafts[ @@ -607,6 +613,12 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner: + """Composite orchestrator that owns speculative_num_steps per-step + runners with shared input buffers. Not itself a + DecodeCudaGraphRunner — it only routes work to the per-step + runners. + """ + def __init__(self, eagle_worker: MultiLayerEagleDraftWorker): self.eagle_worker = eagle_worker self.device = eagle_worker.device @@ -625,14 +637,14 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner: self._init_and_capture() def _init_and_capture(self): - if self.eagle_worker.server_args.disable_cuda_graph: + if check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED): self.runners = [None] * self.speculative_num_steps return self.runners: List[Optional[MultiLayerEagleDraftExtendCudaGraphRunner]] = [] buffer_len_list: List[int] = [] - # 1. Capture loop + # 1. Construct per-step runners (cheap setup only). for step in range(self.speculative_num_steps): if self.draft_extend_attn_backend_list[step]: runner = MultiLayerEagleDraftExtendCudaGraphRunner( @@ -647,7 +659,7 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner: else: self.runners.append(None) - # 2. Allocate buffers + # 2. Allocate shared buffers. self.cuda_graph_buffers["seq_lens_cpu"] = torch.full( (self.max_bs,), self.seq_len_fill_value, @@ -655,7 +667,6 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner: ) with torch.device(self.device): - # Sliced buffers self.cuda_graph_buffers["input_ids"] = torch.zeros( (self.offsets[-1],), dtype=torch.int64 ) @@ -666,7 +677,6 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner: (self.offsets[-1],), dtype=torch.int64 ) - # Shared states self.cuda_graph_buffers["seq_lens"] = torch.full( (self.max_bs,), self.seq_len_fill_value, @@ -682,6 +692,8 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner: (self.max_bs,), 1, dtype=torch.int32 ) + # 3. Per-step capture, in reverse order so that next_cuda_graph_runner + # is already initialized when this step references it. for step in range(self.speculative_num_steps - 1, -1, -1): if self.runners[step] is not None: tic = time.perf_counter() diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 56bf32705..b794242d7 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -31,6 +31,11 @@ from sglang.srt.managers.io_struct import ( from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + Phase, + check_cuda_graph_backend, +) from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, @@ -128,8 +133,8 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): # Do not capture cuda graph in `TpModelWorker` init, # will capture later with init_cuda_graphs() - backup_disable_cuda_graph = server_args.disable_cuda_graph - server_args.disable_cuda_graph = True + backup_decode_mode = server_args.cuda_graph_config.decode.backend + server_args.cuda_graph_config.decode.backend = Backend.DISABLED # Share the allocator with a target worker. # Draft and target worker own their own KV cache pools. @@ -182,8 +187,8 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): # Init attention backend and cuda graphs for i in range(self.speculative_num_steps): - self.draft_runner_list[i].server_args.disable_cuda_graph = ( - backup_disable_cuda_graph + self.draft_runner_list[i].server_args.cuda_graph_config.decode.backend = ( + backup_decode_mode ) self.draft_tp_context = ( draft_tp_context if server_args.enable_dp_attention else empty_context @@ -230,7 +235,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): self.cuda_graph_runner = None self.cuda_graph_runner_for_draft_extend = None - if self.server_args.disable_cuda_graph: + if check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED): return if not _is_npu: @@ -534,7 +539,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): self.reset_cuda_graph_buffers(forward_batch, batch_result) else: logger.warning_once( - f"can't use cuda graph for draft extend! may have correctness issue!" + "can't use cuda graph for draft extend! may have correctness issue!" ) select_index = ( torch.arange(len(batch.seq_lens), device=self.device) @@ -796,7 +801,7 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): self.target_worker.model_runner.attn_backend.update_verify_buffers_to_fill_after_draft( verify_input, ( - self.target_worker.model_runner.graph_runner.bs + self.target_worker.model_runner.decode_cuda_graph_runner.bs if can_run_cuda_graph else None ), diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index 84a59515d..60a826a12 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/python/sglang/srt/speculative/standalone_worker_v2.py @@ -7,6 +7,7 @@ import torch from sglang.srt.environ import envs from sglang.srt.layers.moe.utils import speculative_moe_backend_context from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.adaptive_runtime_state import ( AdaptiveController, @@ -84,8 +85,8 @@ class StandaloneDraftWorker(EagleDraftWorker): # Do not capture cuda graph in `TpModelWorker` init, # will capture later with init_cuda_graphs() - backup_disable_cuda_graph = server_args.disable_cuda_graph - server_args.disable_cuda_graph = True + backup_decode_mode = server_args.cuda_graph_config.decode.backend + server_args.cuda_graph_config.decode.backend = Backend.DISABLED # Share the allocator with a target worker. # Draft and target worker own their own KV cache pools. @@ -117,7 +118,9 @@ class StandaloneDraftWorker(EagleDraftWorker): self.init_lm_head() # Init attention backend and cuda graphs - self.draft_runner.server_args.disable_cuda_graph = backup_disable_cuda_graph + self.draft_runner.server_args.cuda_graph_config.decode.backend = ( + backup_decode_mode + ) self.draft_tp_context = ( draft_tp_context if server_args.enable_dp_attention else empty_context ) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 4556d06b1..d2650cbb8 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -3074,13 +3074,14 @@ def dispose_tensor(x: torch.Tensor): interfering with torch.compile's memory tracking and graph recording. """ - # Skip disposal during piecewise CUDA graph to avoid torch.compile issues - # we do local import to avoid circular import - from sglang.srt.compilation.piecewise_context_manager import ( - is_in_piecewise_cuda_graph, + # Skip disposal during piecewise CUDA graph capture/replay: freeing the + # backing storage would invalidate addresses recorded in the graph. + # Local import avoids a circular dependency. + from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( + is_in_tc_piecewise_cuda_graph, ) - if is_in_piecewise_cuda_graph(): + if is_in_tc_piecewise_cuda_graph(): return x.set_(torch.empty((0,), device=x.device, dtype=x.dtype)) diff --git a/python/sglang/test/doc_patch.py b/python/sglang/test/doc_patch.py index 503ce86e7..6da9d635f 100644 --- a/python/sglang/test/doc_patch.py +++ b/python/sglang/test/doc_patch.py @@ -30,7 +30,7 @@ def patched_post_init(self): # Disable CUDA graphs to avoid memory spikes during capture. # Notebooks only run a few sample requests, so perf is not critical. self.disable_cuda_graph = True - self.cuda_graph_max_bs = 4 + self.cuda_graph_max_bs_decode = 4 server_args_mod.ServerArgs.__post_init__ = patched_post_init diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py index 84863d063..0e7ca5a32 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py @@ -11,6 +11,11 @@ from sglang.srt.layers import dp_attention as _dp_attention from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + CudaGraphConfig, + PhaseConfig, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.model_runner import ModelRunner @@ -328,8 +333,18 @@ class MockModelRunner(ModelRunner): self.server_args = make_mock_server_args( attention_backend=case.backend, chunked_prefill_size=-1, - disable_cuda_graph=disable_cuda_graph, - disable_piecewise_cuda_graph=disable_piecewise_cuda_graph, + cuda_graph_config=CudaGraphConfig( + decode=PhaseConfig( + backend=Backend.DISABLED if disable_cuda_graph else Backend.FULL, + ), + prefill=PhaseConfig( + backend=( + Backend.DISABLED + if (disable_cuda_graph or disable_piecewise_cuda_graph) + else Backend.TC_PIECEWISE + ), + ), + ), disable_radix_cache=False, dllm_algorithm=None, dllm_algorithm_config=None, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py index 68c7ad9d6..4150adcc8 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py @@ -9,6 +9,11 @@ from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS from sglang.srt.layers.attention.dsa import utils as _dsa_utils from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, ReqToTokenPool +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + CudaGraphConfig, + PhaseConfig, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.model_runner import ModelRunner @@ -305,8 +310,18 @@ class DSAMockModelRunner(ModelRunner): self.server_args = make_mock_server_args( attention_backend=case.backend, chunked_prefill_size=-1, - disable_cuda_graph=disable_cuda_graph, - disable_piecewise_cuda_graph=disable_piecewise_cuda_graph, + cuda_graph_config=CudaGraphConfig( + decode=PhaseConfig( + backend=Backend.DISABLED if disable_cuda_graph else Backend.FULL, + ), + prefill=PhaseConfig( + backend=( + Backend.DISABLED + if (disable_cuda_graph or disable_piecewise_cuda_graph) + else Backend.TC_PIECEWISE + ), + ), + ), disable_radix_cache=False, dllm_algorithm=None, dllm_algorithm_config=None, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py index b47c0e0f5..c0aa1a39c 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py @@ -27,6 +27,11 @@ from sglang.srt.layers.attention.dsv4.quant_k_cache import ( from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + CudaGraphConfig, + PhaseConfig, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.server_args import set_global_server_args_for_scheduler @@ -329,8 +334,18 @@ class MockDSV4ModelRunner: self.server_args = make_mock_server_args( attention_backend=case.backend, chunked_prefill_size=-1, - disable_cuda_graph=disable_cuda_graph, - disable_piecewise_cuda_graph=disable_piecewise_cuda_graph, + cuda_graph_config=CudaGraphConfig( + decode=PhaseConfig( + backend=Backend.DISABLED if disable_cuda_graph else Backend.FULL, + ), + prefill=PhaseConfig( + backend=( + Backend.DISABLED + if (disable_cuda_graph or disable_piecewise_cuda_graph) + else Backend.TC_PIECEWISE + ), + ), + ), disable_radix_cache=False, disaggregation_mode=None, dp_size=1, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py index ea4dfdd7d..e5b9ed125 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py @@ -10,6 +10,11 @@ from sglang.srt.layers.attention import ( from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + CudaGraphConfig, + PhaseConfig, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.model_runner import ModelRunner @@ -319,8 +324,18 @@ class DualChunkMockModelRunner(ModelRunner): self.server_args = make_mock_server_args( attention_backend=case.backend, chunked_prefill_size=-1, - disable_cuda_graph=disable_cuda_graph, - disable_piecewise_cuda_graph=disable_piecewise_cuda_graph, + cuda_graph_config=CudaGraphConfig( + decode=PhaseConfig( + backend=Backend.DISABLED if disable_cuda_graph else Backend.FULL, + ), + prefill=PhaseConfig( + backend=( + Backend.DISABLED + if (disable_cuda_graph or disable_piecewise_cuda_graph) + else Backend.TC_PIECEWISE + ), + ), + ), disable_radix_cache=False, dp_size=1, enable_dp_attention=False, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py index c404d303b..7f29b9e30 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py @@ -22,6 +22,11 @@ from sglang.srt.mem_cache.memory_pool import ( HybridReqToTokenPool, MHATokenToKVPool, ) +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + CudaGraphConfig, + PhaseConfig, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.model_runner import ModelRunner @@ -217,8 +222,18 @@ class MockGDNModelRunner(ModelRunner): self.server_args = make_mock_server_args( attention_backend=case.backend, chunked_prefill_size=-1, - disable_cuda_graph=disable_cuda_graph, - disable_piecewise_cuda_graph=disable_piecewise_cuda_graph, + cuda_graph_config=CudaGraphConfig( + decode=PhaseConfig( + backend=Backend.DISABLED if disable_cuda_graph else Backend.FULL, + ), + prefill=PhaseConfig( + backend=( + Backend.DISABLED + if (disable_cuda_graph or disable_piecewise_cuda_graph) + else Backend.TC_PIECEWISE + ), + ), + ), dllm_algorithm=None, dllm_algorithm_config=None, enable_deterministic_inference=False, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py index 5491cf153..dcb02d0dc 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py @@ -22,6 +22,11 @@ from sglang.srt.mem_cache.memory_pool import ( HybridReqToTokenPool, MHATokenToKVPool, ) +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + CudaGraphConfig, + PhaseConfig, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.model_runner import ModelRunner @@ -223,8 +228,18 @@ class MockKDAModelRunner(ModelRunner): self.server_args = make_mock_server_args( attention_backend=case.backend, chunked_prefill_size=-1, - disable_cuda_graph=disable_cuda_graph, - disable_piecewise_cuda_graph=disable_piecewise_cuda_graph, + cuda_graph_config=CudaGraphConfig( + decode=PhaseConfig( + backend=Backend.DISABLED if disable_cuda_graph else Backend.FULL, + ), + prefill=PhaseConfig( + backend=( + Backend.DISABLED + if (disable_cuda_graph or disable_piecewise_cuda_graph) + else Backend.TC_PIECEWISE + ), + ), + ), dllm_algorithm=None, dllm_algorithm_config=None, enable_deterministic_inference=False, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py index b3cace2b9..9e7b87bb0 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py @@ -21,6 +21,11 @@ from sglang.srt.mem_cache.memory_pool import ( HybridReqToTokenPool, MHATokenToKVPool, ) +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + CudaGraphConfig, + PhaseConfig, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.model_runner import ModelRunner @@ -232,8 +237,18 @@ class MockLightningModelRunner(ModelRunner): self.server_args = make_mock_server_args( attention_backend=case.backend, chunked_prefill_size=-1, - disable_cuda_graph=disable_cuda_graph, - disable_piecewise_cuda_graph=disable_piecewise_cuda_graph, + cuda_graph_config=CudaGraphConfig( + decode=PhaseConfig( + backend=Backend.DISABLED if disable_cuda_graph else Backend.FULL, + ), + prefill=PhaseConfig( + backend=( + Backend.DISABLED + if (disable_cuda_graph or disable_piecewise_cuda_graph) + else Backend.TC_PIECEWISE + ), + ), + ), dllm_algorithm=None, dllm_algorithm_config=None, enable_deterministic_inference=False, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py index 6ddbd2661..ae6caebe3 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py @@ -44,6 +44,11 @@ from sglang.srt.mem_cache.memory_pool import ( # noqa: E402 HybridReqToTokenPool, MHATokenToKVPool, ) +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + CudaGraphConfig, + PhaseConfig, +) from sglang.srt.model_executor.forward_batch_info import ( # noqa: E402 ForwardBatch, ForwardMode, @@ -335,8 +340,18 @@ class MockMamba2ModelRunner(ModelRunner): self.server_args = make_mock_server_args( attention_backend=case.backend, chunked_prefill_size=-1, - disable_cuda_graph=disable_cuda_graph, - disable_piecewise_cuda_graph=disable_piecewise_cuda_graph, + cuda_graph_config=CudaGraphConfig( + decode=PhaseConfig( + backend=Backend.DISABLED if disable_cuda_graph else Backend.FULL, + ), + prefill=PhaseConfig( + backend=( + Backend.DISABLED + if (disable_cuda_graph or disable_piecewise_cuda_graph) + else Backend.TC_PIECEWISE + ), + ), + ), dllm_algorithm=None, dllm_algorithm_config=None, enable_deterministic_inference=False, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py index 908574e28..68fb623c7 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py @@ -11,6 +11,11 @@ from sglang.srt.layers import dp_attention as _dp_attention from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool, ReqToTokenPool +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + CudaGraphConfig, + PhaseConfig, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ( ForwardContext, @@ -237,9 +242,19 @@ class MockMLAModelRunner(ModelRunner): self.server_args = make_mock_server_args( attention_backend=case.backend, chunked_prefill_size=-1, - disable_cuda_graph=disable_cuda_graph, + cuda_graph_config=CudaGraphConfig( + decode=PhaseConfig( + backend=Backend.DISABLED if disable_cuda_graph else Backend.FULL, + ), + prefill=PhaseConfig( + backend=( + Backend.DISABLED + if (disable_cuda_graph or disable_piecewise_cuda_graph) + else Backend.TC_PIECEWISE + ), + ), + ), disable_chunked_prefix_cache=True, - disable_piecewise_cuda_graph=disable_piecewise_cuda_graph, disable_radix_cache=False, disaggregation_mode=None, dllm_algorithm=None, diff --git a/python/sglang/test/kits/attention_unittest/mock_server_args.py b/python/sglang/test/kits/attention_unittest/mock_server_args.py index a129a018c..a347a4d15 100644 --- a/python/sglang/test/kits/attention_unittest/mock_server_args.py +++ b/python/sglang/test/kits/attention_unittest/mock_server_args.py @@ -21,6 +21,7 @@ validation it performs is irrelevant for module-level attention tests. import dataclasses +from sglang.srt.model_executor.cuda_graph_config import default_cuda_graph_config from sglang.srt.server_args import ServerArgs @@ -54,4 +55,8 @@ def make_mock_server_args(**overrides) -> ServerArgs: setattr(sa, f"_{k}", v) else: setattr(sa, k, v) + if sa.cuda_graph_config is None: + sa.cuda_graph_config = default_cuda_graph_config() + if not hasattr(sa, "_cuda_graph_config_locked"): + sa._cuda_graph_config_locked = set() return sa diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py index 1efefeff5..88ee9e76c 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py @@ -1154,19 +1154,19 @@ def _capture_eagle_draft_extend_graph_runner( ) -> EAGLEDraftExtendCudaGraphRunner: with ( patch( - "sglang.srt.model_executor.cuda_graph_runner.graph_capture", + "sglang.srt.model_executor.runner.decode_cuda_graph_runner.graph_capture", _single_rank_graph_capture, ), patch( - "sglang.srt.model_executor.cuda_graph_runner.get_tensor_model_parallel_rank", + "sglang.srt.model_executor.runner.decode_cuda_graph_runner.get_tensor_model_parallel_rank", lambda: 0, ), patch( - "sglang.srt.model_executor.cuda_graph_runner.get_available_gpu_memory", + "sglang.srt.model_executor.runner.decode_cuda_graph_runner.get_available_gpu_memory", lambda *args, **kwargs: 0.0, ), patch( - "sglang.srt.model_executor.cuda_graph_runner.get_attention_cp_size", + "sglang.srt.model_executor.runner.base_cuda_graph_runner.get_attention_cp_size", lambda: 1, ), ): diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py index 951b3b141..52bd0ec54 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py @@ -8,13 +8,14 @@ import torch from torch import nn from sglang.srt.layers.logits_processor import LogitsProcessorOutput -from sglang.srt.model_executor.cuda_graph_runner import set_global_graph_memory_pool +from sglang.srt.model_executor.cuda_graph_config import CudaGraphConfig, PhaseConfig from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, ForwardMode, ) from sglang.srt.model_executor.input_buffers import _forward_input_buffer_pool +from sglang.srt.model_executor.runner import set_global_graph_memory_pool from sglang.srt.server_args import set_global_server_args_for_scheduler from sglang.srt.speculative.draft_utils import DraftBackendFactory from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( @@ -296,7 +297,12 @@ def _configure_runner_for_eagle_draft( server_args = runner.server_args updates = { "attention_backend": case.backend, - "cuda_graph_bs": [settings.capture_batch_size], + "cuda_graph_config": CudaGraphConfig( + decode=PhaseConfig( + bs=[settings.capture_batch_size], + max_bs=settings.capture_batch_size, + ), + ), "debug_cuda_graph": False, "decode_attention_backend": case.backend, "disable_cuda_graph_padding": False, @@ -424,19 +430,19 @@ def _capture_eagle_draft_graph_runner( ) -> EAGLEDraftCudaGraphRunner: with ( patch( - "sglang.srt.model_executor.cuda_graph_runner.graph_capture", + "sglang.srt.model_executor.runner.decode_cuda_graph_runner.graph_capture", _single_rank_graph_capture, ), patch( - "sglang.srt.model_executor.cuda_graph_runner.get_tensor_model_parallel_rank", + "sglang.srt.model_executor.runner.decode_cuda_graph_runner.get_tensor_model_parallel_rank", lambda: 0, ), patch( - "sglang.srt.model_executor.cuda_graph_runner.get_available_gpu_memory", + "sglang.srt.model_executor.runner.decode_cuda_graph_runner.get_available_gpu_memory", lambda *args, **kwargs: 0.0, ), patch( - "sglang.srt.model_executor.cuda_graph_runner.get_attention_cp_size", + "sglang.srt.model_executor.runner.base_cuda_graph_runner.get_attention_cp_size", lambda: 1, ), ): @@ -453,19 +459,15 @@ def _capture_frozen_kv_mtp_graph_runner( ) -> FrozenKVMTPCudaGraphRunner: with ( patch( - "sglang.srt.model_executor.cuda_graph_runner.graph_capture", + "sglang.srt.speculative.frozen_kv_mtp_cuda_graph_runner.graph_capture", _single_rank_graph_capture, ), patch( - "sglang.srt.model_executor.cuda_graph_runner.get_tensor_model_parallel_rank", + "sglang.srt.speculative.frozen_kv_mtp_cuda_graph_runner.get_tensor_model_parallel_rank", lambda: 0, ), patch( - "sglang.srt.model_executor.cuda_graph_runner.get_available_gpu_memory", - lambda *args, **kwargs: 0.0, - ), - patch( - "sglang.srt.model_executor.cuda_graph_runner.get_attention_cp_size", + "sglang.srt.model_executor.runner.base_cuda_graph_runner.get_attention_cp_size", lambda: 1, ), ): diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/split_op_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/split_op_runner.py index a648727ae..8a1ef5452 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/split_op_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/split_op_runner.py @@ -3,16 +3,16 @@ from typing import Any, Callable import torch -from sglang.srt.compilation.piecewise_context_manager import ( - enable_piecewise_cuda_graph, -) -from sglang.srt.compilation.piecewise_context_manager import ( - set_forward_context as piecewise_forward_context, -) -from sglang.srt.model_executor.breakable_cuda_graph.context import ( +from sglang.srt.model_executor.forward_context import ForwardContext, forward_context +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( enable_breakable_cuda_graph, ) -from sglang.srt.model_executor.forward_context import ForwardContext, forward_context +from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph.context_manager import ( + enable_tc_piecewise_cuda_graph as enable_piecewise_cuda_graph, +) +from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph.context_manager import ( + set_tc_piecewise_forward_context as piecewise_forward_context, +) from ..attention_methods.dense_attention import DEFAULT_DEVICE as DENSE_DEFAULT_DEVICE from ..attention_methods.dense_attention import DEFAULT_DTYPE as DENSE_DEFAULT_DTYPE diff --git a/python/sglang/test/runners.py b/python/sglang/test/runners.py index e4a9fe044..e511db587 100644 --- a/python/sglang/test/runners.py +++ b/python/sglang/test/runners.py @@ -578,7 +578,7 @@ class SRTRunner: disable_overlap_schedule: bool = False, disable_custom_all_reduce: bool = False, torchao_config: Optional[str] = None, - cuda_graph_max_bs: int = 4, + cuda_graph_max_bs_decode: int = 4, sleep_on_idle=False, max_lora_rank: Optional[int] = None, lora_target_modules: Optional[List[str]] = None, @@ -635,7 +635,7 @@ class SRTRunner: dp_size=dp_size, tokenizer_path=tokenizer_path, disable_overlap_schedule=disable_overlap_schedule, - cuda_graph_max_bs=cuda_graph_max_bs, + cuda_graph_max_bs_decode=cuda_graph_max_bs_decode, disable_custom_all_reduce=disable_custom_all_reduce, sleep_on_idle=sleep_on_idle, max_lora_rank=max_lora_rank, diff --git a/python/sglang/test/scripted_runtime/http_server.py b/python/sglang/test/scripted_runtime/http_server.py index 7f9696740..d574c3c7e 100644 --- a/python/sglang/test/scripted_runtime/http_server.py +++ b/python/sglang/test/scripted_runtime/http_server.py @@ -230,7 +230,7 @@ def _spawn_server_process( kv_canary="raise", kv_canary_real_data="partial", kv_canary_sweep_interval=100, - disable_piecewise_cuda_graph=True, + disable_prefill_cuda_graph=True, ) launch_kwargs.update(engine_kwargs) http_port = launch_kwargs["port"] diff --git a/test/registered/ascend/basic_function/optimization_debug/test_npu_piecewise_graph_prefill.py b/test/registered/ascend/basic_function/optimization_debug/test_npu_piecewise_graph_prefill.py index 110e6d2b8..52be19f03 100644 --- a/test/registered/ascend/basic_function/optimization_debug/test_npu_piecewise_graph_prefill.py +++ b/test/registered/ascend/basic_function/optimization_debug/test_npu_piecewise_graph_prefill.py @@ -29,8 +29,8 @@ class TestPiecewiseGraphPrefillCorrectness(GSM8KAscendMixin, CustomTestCase): "ascend", "--cuda-graph-bs", 128, - "--enforce-piecewise-cuda-graph", - "--piecewise-cuda-graph-tokens", + "--cuda-graph-backend-prefill=tc_piecewise", + "--cuda-graph-bs-prefill", *TOKENS_TO_CAPTURE, ] accuracy = 0.84 @@ -45,8 +45,8 @@ class TestPiecewiseGraphPrefillBenchmark(CustomTestCase): 0.8, "--attention-backend", "ascend", - "--enforce-piecewise-cuda-graph", - "--piecewise-cuda-graph-tokens", + "--cuda-graph-backend-prefill=tc_piecewise", + "--cuda-graph-bs-prefill", ] + TOKENS_TO_CAPTURE latency = 0.045 diff --git a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py index f2877d153..8eaded05f 100644 --- a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py +++ b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py @@ -394,7 +394,15 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase): ) from sglang.srt.server_args import ServerArgs - self.assertFalse(ServerArgs(model_path="dummy").enable_breakable_cuda_graph) + # cg-refactor folded the legacy enable_breakable_cuda_graph flag + # into cuda_graph_config. Verify the per-phase backend selectors + # default to None (i.e. nothing opted into BREAKABLE without an + # explicit CLI flag). + sa = ServerArgs(model_path="dummy") + self.assertNotEqual(sa.cuda_graph_backend_decode, "breakable") + self.assertNotEqual(sa.cuda_graph_backend_prefill, "breakable") + self.assertNotEqual(sa.decode_cuda_graph_backend, "breakable") + self.assertNotEqual(sa.prefill_cuda_graph_backend, "breakable") self.assertFalse( AttentionBackend.use_captured_forward_metadata_for_breakable_cuda_graph ) diff --git a/test/registered/core/test_no_extra_forked_cuda_context.py b/test/registered/core/test_no_extra_forked_cuda_context.py index f64633cc0..064fcb0f0 100644 --- a/test/registered/core/test_no_extra_forked_cuda_context.py +++ b/test/registered/core/test_no_extra_forked_cuda_context.py @@ -34,8 +34,10 @@ class TestTPServerGPUProcesses(CustomTestCase): str(cls.tp_size), "--mem-fraction-static", "0.70", - "--disable-cuda-graph", - "--disable-piecewise-cuda-graph", + "--cuda-graph-backend-decode", + "disabled", + "--cuda-graph-backend-prefill", + "disabled", ], ) diff --git a/test/registered/breakable_cuda_graph/test_breakable_cuda_graph.py b/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py similarity index 94% rename from test/registered/breakable_cuda_graph/test_breakable_cuda_graph.py rename to test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py index 9b23bee62..12f7f8b7e 100644 --- a/test/registered/breakable_cuda_graph/test_breakable_cuda_graph.py +++ b/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py @@ -1,10 +1,10 @@ """Tests for the breakable CUDA graph (BCG) runner. Two test classes: -- ``TestBreakableCUDAGraphBasic`` / ``TestCopyOutput`` / ``TestBreakGraphHelper``: +- TestBreakableCUDAGraphBasic / TestCopyOutput / TestBreakGraphHelper: unit tests for the core capture / replay mechanism (simple tensor ops). -- ``TestBreakableCudaGraph``: integration test — spin up Qwen3-8B with - ``--enable-breakable-cuda-graph`` and check mgsm_en accuracy. +- TestBreakableCudaGraph: integration test — spin up Qwen3-8B with + --enable-breakable-cuda-graph and check mgsm_en accuracy. """ import unittest @@ -53,7 +53,7 @@ class TestBreakableCUDAGraphBasic(CustomTestCase): except ImportError: raise unittest.SkipTest("cuda-python not installed") - from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import ( + from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import ( BreakableCUDAGraph, BreakableCUDAGraphCapture, eager_on_graph, @@ -199,7 +199,7 @@ class TestCopyOutput(CustomTestCase): except ImportError: raise unittest.SkipTest("cuda-python not installed") - from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import ( + from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import ( _copy_output, ) @@ -261,7 +261,7 @@ class TestBreakGraphHelper(CustomTestCase): except ImportError: raise unittest.SkipTest("cuda-python not installed") - from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import ( + from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import ( BreakableCUDAGraph, BreakableCUDAGraphCapture, break_graph, @@ -303,7 +303,7 @@ class TestBreakableCudaGraph(CustomTestCase): cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, other_args=[ - "--enable-breakable-cuda-graph", + "--cuda-graph-backend-prefill=breakable", ], ) diff --git a/test/registered/piecewise_cuda_graph/test_pcg_glm5_fp4.py b/test/registered/cuda_graph/piecewise/test_pcg_glm5_fp4.py similarity index 97% rename from test/registered/piecewise_cuda_graph/test_pcg_glm5_fp4.py rename to test/registered/cuda_graph/piecewise/test_pcg_glm5_fp4.py index 6cce520de..5fb43c4ce 100644 --- a/test/registered/piecewise_cuda_graph/test_pcg_glm5_fp4.py +++ b/test/registered/cuda_graph/piecewise/test_pcg_glm5_fp4.py @@ -43,7 +43,7 @@ class TestPCGGlm5Fp4(CustomTestCase): "--quantization", "modelopt_fp4", "--disable-flashinfer-autotune", - "--enforce-piecewise-cuda-graph", + "--cuda-graph-backend-prefill=tc_piecewise", "--model-loader-extra-config", '{"enable_multithread_load": true, "num_threads": 64}', ], diff --git a/test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding.py b/test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding.py similarity index 96% rename from test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding.py rename to test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding.py index 91bab339e..9381defb5 100644 --- a/test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding.py +++ b/test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding.py @@ -22,7 +22,7 @@ class TestPCGWithEAGLE3(PCGSpecBase, unittest.TestCase): "--tp", "2", "--trust-remote-code", - "--enforce-piecewise-cuda-graph", + "--cuda-graph-backend-prefill=tc_piecewise", "--mem-fraction-static", "0.6", "--speculative-algorithm", diff --git a/test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding_dflash.py b/test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding_dflash.py similarity index 93% rename from test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding_dflash.py rename to test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding_dflash.py index 3095aacca..3c90e6630 100644 --- a/test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding_dflash.py +++ b/test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding_dflash.py @@ -26,7 +26,8 @@ class TestPCGWithDFlash(PCGSpecBase, CustomTestCase): "--trust-remote-code", "--attention-backend", "flashinfer", - "--enforce-piecewise-cuda-graph", + "--cuda-graph-backend-prefill", + "tc_piecewise", "--speculative-algorithm", "DFLASH", "--speculative-draft-model-path", @@ -35,7 +36,7 @@ class TestPCGWithDFlash(PCGSpecBase, CustomTestCase): "1", "--max-running-requests", "64", - "--cuda-graph-bs", + "--cuda-graph-bs-decode", *[str(i) for i in range(1, 65)], ] server_env = {"SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1"} diff --git a/test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding_extra.py b/test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding_extra.py similarity index 94% rename from test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding_extra.py rename to test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding_extra.py index 70fadf6e9..6a7d126bc 100644 --- a/test/registered/piecewise_cuda_graph/test_pcg_with_speculative_decoding_extra.py +++ b/test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding_extra.py @@ -24,7 +24,6 @@ class TestPCGWithMTP(PCGSpecBase, unittest.TestCase): "fp8", "--mamba-scheduler-strategy", "extra_buffer", - "--enable-piecewise-cuda-graph", "--speculative-algorithm", "NEXTN", "--reasoning-parser", @@ -42,7 +41,7 @@ class TestPCGWithSTANDALONE(PCGSpecBase, unittest.TestCase): model = "meta-llama/Llama-3.1-8B-Instruct" server_args = [ "--trust-remote-code", - "--enforce-piecewise-cuda-graph", + "--cuda-graph-backend-prefill=tc_piecewise", "--mem-fraction-static", "0.5", "--speculative-algorithm", @@ -65,7 +64,7 @@ class TestPCGWithNGRAM(PCGSpecBase, unittest.TestCase): model = "Qwen/Qwen2.5-Coder-7B-Instruct" server_args = [ "--trust-remote-code", - "--enforce-piecewise-cuda-graph", + "--cuda-graph-backend-prefill=tc_piecewise", "--speculative-algorithm", "NGRAM", "--speculative-num-draft-tokens", diff --git a/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_support_1_gpu.py b/test/registered/cuda_graph/piecewise/test_piecewise_cuda_graph_support_1_gpu.py similarity index 95% rename from test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_support_1_gpu.py rename to test/registered/cuda_graph/piecewise/test_piecewise_cuda_graph_support_1_gpu.py index e5707e14a..cfd22c857 100644 --- a/test/registered/piecewise_cuda_graph/test_piecewise_cuda_graph_support_1_gpu.py +++ b/test/registered/cuda_graph/piecewise/test_piecewise_cuda_graph_support_1_gpu.py @@ -33,7 +33,7 @@ class TestPiecewiseCudaGraphQwen25VL(CustomTestCase): cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, other_args=[ - "--enforce-piecewise-cuda-graph", + "--cuda-graph-backend-prefill=tc_piecewise", "--disable-radix-cache", ], ) @@ -69,7 +69,7 @@ class TestPiecewiseCudaGraphQwen25VLEmbedding(CustomTestCase): model_path=model_path, enable_multimodal=True, is_embedding=True, - enforce_piecewise_cuda_graph=True, + cuda_graph_backend_prefill="tc_piecewise", ) out = engine.encode([text], image_data=[DEFAULT_IMAGE_URL])[0]["embedding"] engine.shutdown() @@ -79,7 +79,7 @@ class TestPiecewiseCudaGraphQwen25VLEmbedding(CustomTestCase): model_path=model_path, enable_multimodal=True, is_embedding=True, - disable_piecewise_cuda_graph=True, + cuda_graph_backend_prefill="disabled", ) out_without_pcg = engine.encode([text], image_data=[DEFAULT_IMAGE_URL])[0][ "embedding" diff --git a/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py b/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py index 383da8aa2..a675b4bf5 100644 --- a/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py +++ b/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py @@ -325,7 +325,7 @@ def _run_server_and_generate( "--mem-fraction-static", "0.5", "--disable-cuda-graph", - "--disable-piecewise-cuda-graph", + "--cuda-graph-backend-prefill=disabled", "--disable-radix-cache", ] if extra_server_args: diff --git a/test/registered/kv_canary/test_self_unit_capacities.py b/test/registered/kv_canary/test_self_unit_capacities.py index 510f6120b..1c3495929 100644 --- a/test/registered/kv_canary/test_self_unit_capacities.py +++ b/test/registered/kv_canary/test_self_unit_capacities.py @@ -4,17 +4,24 @@ import unittest from types import SimpleNamespace from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + CudaGraphConfig, + PhaseConfig, +) +from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase -register_cuda_ci(est_time=45, stage="extra-a", runner_config="1-gpu-small") +register_cpu_ci(est_time=5, suite="base-a-test-cpu") class TestComputeLaunchCapacities(CustomTestCase): @staticmethod def _make_server_args(*, max_bs: int) -> SimpleNamespace: return SimpleNamespace( - cuda_graph_max_bs=max_bs, + cuda_graph_config=CudaGraphConfig( + decode=PhaseConfig(backend=Backend.FULL, max_bs=max_bs) + ), speculative_num_draft_tokens=0, chunked_prefill_size=None, max_prefill_tokens=128, diff --git a/test/registered/lora/test_lora_update.py b/test/registered/lora/test_lora_update.py index 977817961..835039978 100644 --- a/test/registered/lora/test_lora_update.py +++ b/test/registered/lora/test_lora_update.py @@ -879,7 +879,7 @@ class LoRAUpdateTestSessionBase: lora_target_modules: Optional[List[str]] = None, lora_backend: str = "csgmv", disable_cuda_graph: bool = False, - cuda_graph_max_bs: int = 4, + cuda_graph_max_bs_decode: int = 4, ): self.testcase = testcase self.model_path = model_path @@ -890,7 +890,7 @@ class LoRAUpdateTestSessionBase: self.max_loaded_loras = max_loaded_loras self.lora_backend = lora_backend self.disable_cuda_graph = disable_cuda_graph - self.cuda_graph_max_bs = cuda_graph_max_bs + self.cuda_graph_max_bs_decode = cuda_graph_max_bs_decode self.enable_lora = enable_lora self.expected_adapters = set() @@ -964,7 +964,7 @@ class LoRAUpdateEngineTestSession(LoRAUpdateTestSessionBase): max_loras_per_batch=self.max_loras_per_batch, max_loaded_loras=self.max_loaded_loras, disable_cuda_graph=self.disable_cuda_graph, - cuda_graph_max_bs=self.cuda_graph_max_bs, + cuda_graph_max_bs_decode=self.cuda_graph_max_bs_decode, enable_lora=self.enable_lora, disable_radix_cache=True, ) @@ -1101,8 +1101,8 @@ class LoRAUpdateServerTestSession(LoRAUpdateTestSessionBase): def __enter__(self): other_args = [ - "--cuda-graph-max-bs", - str(self.cuda_graph_max_bs), + "--cuda-graph-max-bs-decode", + str(self.cuda_graph_max_bs_decode), "--max-loras-per-batch", str(self.max_loras_per_batch), "--lora-backend", diff --git a/test/registered/model_loading/test_external_models.py b/test/registered/model_loading/test_external_models.py index e6f5a7ec7..efa1a4968 100644 --- a/test/registered/model_loading/test_external_models.py +++ b/test/registered/model_loading/test_external_models.py @@ -23,7 +23,7 @@ class TestExternalModels(CustomTestCase): engine = sgl.Engine( model_path=model_path, - cuda_graph_max_bs=1, + cuda_graph_max_bs_decode=1, max_total_tokens=64, enable_multimodal=True, ) diff --git a/test/registered/model_loading/test_load_weights_from_remote_instance.py b/test/registered/model_loading/test_load_weights_from_remote_instance.py index 42db54686..478653fc8 100644 --- a/test/registered/model_loading/test_load_weights_from_remote_instance.py +++ b/test/registered/model_loading/test_load_weights_from_remote_instance.py @@ -188,7 +188,7 @@ def init_process_dst( model_path=model_name, base_gpu_id=base_gpu_id, tp_size=tp_size, - cuda_graph_max_bs=2, + cuda_graph_max_bs_decode=2, tokenizer_path=model_name, remote_instance_weight_loader_seed_instance_ip=seed_instance_ip, remote_instance_weight_loader_seed_instance_service_port=seed_instance_service_port, diff --git a/test/registered/model_loading/test_load_weights_from_remote_instance_npu.py b/test/registered/model_loading/test_load_weights_from_remote_instance_npu.py index 6253610ea..6c9a16d57 100644 --- a/test/registered/model_loading/test_load_weights_from_remote_instance_npu.py +++ b/test/registered/model_loading/test_load_weights_from_remote_instance_npu.py @@ -197,7 +197,7 @@ def init_process_dst( model_path=model_name, base_gpu_id=base_gpu_id, tp_size=tp_size, - cuda_graph_max_bs=2, + cuda_graph_max_bs_decode=2, tokenizer_path=model_name, remote_instance_weight_loader_seed_instance_ip=seed_instance_ip, remote_instance_weight_loader_seed_instance_service_port=seed_instance_service_port, diff --git a/test/registered/model_loading/test_prefetch_checkpoints_multi_gpu.py b/test/registered/model_loading/test_prefetch_checkpoints_multi_gpu.py index a32357dbe..bad024095 100644 --- a/test/registered/model_loading/test_prefetch_checkpoints_multi_gpu.py +++ b/test/registered/model_loading/test_prefetch_checkpoints_multi_gpu.py @@ -26,7 +26,7 @@ class TestPrefetchCheckpointsMultiGPU(CustomTestCase): enable_dp_attention=True, disable_radix_cache=True, weight_loader_prefetch_checkpoints=True, - cuda_graph_max_bs=1, + cuda_graph_max_bs_decode=1, max_total_tokens=256, ) diff --git a/test/registered/model_loading/test_runai_model_loader.py b/test/registered/model_loading/test_runai_model_loader.py index cc1a65658..812646792 100644 --- a/test/registered/model_loading/test_runai_model_loader.py +++ b/test/registered/model_loading/test_runai_model_loader.py @@ -28,7 +28,7 @@ class TestRunaiModelLoader(CustomTestCase): cls.engine = sgl.Engine( model_path=TEST_GCS_MODEL, load_format="runai_streamer", - cuda_graph_max_bs=1, + cuda_graph_max_bs_decode=1, max_total_tokens=64, ) diff --git a/test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py b/test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py index de4ada42c..59380ae71 100644 --- a/test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py +++ b/test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py @@ -179,8 +179,8 @@ class TestDSV4FlashFP4BreakableCudaGraphB200( "4", "--enable-dp-attention", "--enable-mixed-chunk", - "--enable-breakable-cuda-graph", - "--enforce-piecewise-cuda-graph", + "--cuda-graph-backend-prefill", + "breakable", "--moe-a2a-backend", "deepep", "--deepep-config", diff --git a/test/registered/quant/test_fp8_gemm_sm120.py b/test/registered/quant/test_fp8_gemm_sm120.py index 143df6a4a..e4400955c 100644 --- a/test/registered/quant/test_fp8_gemm_sm120.py +++ b/test/registered/quant/test_fp8_gemm_sm120.py @@ -33,7 +33,7 @@ class FP8GemmSM120Base: "--trust-remote-code", "--fp8-gemm-backend", cls.backend, - "--disable-piecewise-cuda-graph", + "--cuda-graph-backend-prefill=disabled", ] if cls.quantization: other_args += ["--quantization", cls.quantization] diff --git a/test/registered/quant/test_gguf.py b/test/registered/quant/test_gguf.py index 083d814be..55115ab8a 100644 --- a/test/registered/quant/test_gguf.py +++ b/test/registered/quant/test_gguf.py @@ -19,7 +19,9 @@ class TestGGUF(CustomTestCase): filename="qwen2-1_5b-instruct-q4_k_m.gguf", ) - engine = sgl.Engine(model_path=model_path, random_seed=42, cuda_graph_max_bs=2) + engine = sgl.Engine( + model_path=model_path, random_seed=42, cuda_graph_max_bs_decode=2 + ) outputs = engine.generate(prompt, sampling_params)["text"] engine.shutdown() diff --git a/test/registered/quant/test_nvfp4_gemm_sm120.py b/test/registered/quant/test_nvfp4_gemm_sm120.py index 61dd5ea88..004266df6 100644 --- a/test/registered/quant/test_nvfp4_gemm_sm120.py +++ b/test/registered/quant/test_nvfp4_gemm_sm120.py @@ -32,7 +32,7 @@ class FP4GemmSM120Base: "modelopt_fp4", "--fp4-gemm-backend", cls.backend, - "--disable-piecewise-cuda-graph", + "--cuda-graph-backend-prefill=disabled", ] cls.process = popen_launch_server( cls.model, diff --git a/test/registered/radix_cache/test_swa_radix_cache_kl.py b/test/registered/radix_cache/test_swa_radix_cache_kl.py index 97afbd45b..810a3a1c6 100644 --- a/test/registered/radix_cache/test_swa_radix_cache_kl.py +++ b/test/registered/radix_cache/test_swa_radix_cache_kl.py @@ -18,7 +18,7 @@ class TestSWARadixCacheKL(KLDivergenceMixin, DefaultServerBase): "1", "--mem-fraction-static", "0.70", - "--disable-piecewise-cuda-graph", + "--cuda-graph-backend-prefill=disabled", ] diff --git a/test/registered/rl/test_update_weights_from_distributed.py b/test/registered/rl/test_update_weights_from_distributed.py index fe7e593a3..0aa9f82a6 100644 --- a/test/registered/rl/test_update_weights_from_distributed.py +++ b/test/registered/rl/test_update_weights_from_distributed.py @@ -315,7 +315,7 @@ def init_process_sgl( model_path=model_name, base_gpu_id=base_gpu_id, tp_size=tp_size, - cuda_graph_max_bs=2, + cuda_graph_max_bs_decode=2, ) else: if rank == 1: diff --git a/test/registered/sessions/test_session_control.py b/test/registered/sessions/test_session_control.py index 6aa4c91ae..38119dd70 100644 --- a/test/registered/sessions/test_session_control.py +++ b/test/registered/sessions/test_session_control.py @@ -45,7 +45,7 @@ class TestSessionControl(CustomTestCase): "--attention-backend", "triton", "--disable-cuda-graph", - "--disable-piecewise-cuda-graph", + "--cuda-graph-backend-prefill=disabled", ], ) diff --git a/test/registered/sessions/test_session_latency.py b/test/registered/sessions/test_session_latency.py index b47b2598a..7d85632c7 100644 --- a/test/registered/sessions/test_session_latency.py +++ b/test/registered/sessions/test_session_latency.py @@ -271,7 +271,7 @@ class TestSessionLatency(CustomTestCase): "--enable-streaming-session", "--mem-fraction-static", "0.70", - "--disable-piecewise-cuda-graph", + "--cuda-graph-backend-prefill=disabled", "--page-size", "4", ], diff --git a/test/registered/sessions/test_streaming_session_swa.py b/test/registered/sessions/test_streaming_session_swa.py index 889a12cf2..f6a9b66a1 100644 --- a/test/registered/sessions/test_streaming_session_swa.py +++ b/test/registered/sessions/test_streaming_session_swa.py @@ -21,7 +21,7 @@ SWA_MODEL = "openai/gpt-oss-20b" SWA_COMMON_ARGS = [ "--mem-fraction-static", "0.70", - "--disable-piecewise-cuda-graph", + "--cuda-graph-backend-prefill=disabled", ] diff --git a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py index 6f3546226..df5a0b1de 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py +++ b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py @@ -396,7 +396,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase): ) scheduler.server_args = SimpleNamespace( enable_two_batch_overlap=False, - disable_piecewise_cuda_graph=True, + cuda_graph_config=None, ) scheduler.spec_algorithm = SpeculativeAlgorithm.NONE scheduler.req_to_token_pool = ReqToTokenPool( diff --git a/test/registered/unit/managers/test_customized_info_streaming.py b/test/registered/unit/managers/test_customized_info_streaming.py index ebaccf2dc..52c998069 100644 --- a/test/registered/unit/managers/test_customized_info_streaming.py +++ b/test/registered/unit/managers/test_customized_info_streaming.py @@ -93,7 +93,6 @@ class TestCustomizedInfoStreaming(CustomTestCase): incremental_streaming_output=True, skip_tokenizer_init=True, disable_cuda_graph=True, - disable_piecewise_cuda_graph=True, disable_radix_cache=True, random_seed=0, log_level="error", diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 6464d65e9..4c8cc5409 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -3,10 +3,16 @@ import json import os import tempfile import unittest +from types import SimpleNamespace from unittest.mock import MagicMock, patch import sglang.srt.server_args as server_args_module from sglang.srt.arg_groups.speculative_hook import handle_speculative_decoding +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + CudaGraphConfig, + PhaseConfig, +) from sglang.srt.server_args import PortArgs, ServerArgs, prepare_server_args from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import ( @@ -621,9 +627,61 @@ class TestPrefillOnlyDisableKvCache(unittest.TestCase): ServerArgs(**self._base_kwargs(kv_cache_dtype="fp4_e2m1")) -class TestCutedslMoeMaxNumTokens(unittest.TestCase): +class TestCudaGraphConfigDataclassAccess(CustomTestCase): + def test_overlap_force_cpu_seq_lens_with_tc_piecewise_prefill(self): + from sglang.srt.managers.overlap_utils import decide_needs_cpu_seq_lens + + server_args = SimpleNamespace( + enable_two_batch_overlap=False, + cuda_graph_config=CudaGraphConfig( + prefill=PhaseConfig(backend=Backend.TC_PIECEWISE) + ), + ) + attn_backend = SimpleNamespace(needs_cpu_seq_lens=False) + + self.assertTrue(decide_needs_cpu_seq_lens(server_args, [attn_backend])) + + @patch( + "sglang.srt.model_executor.runner_backend." + "tc_piecewise_cuda_graph_backend.get_moe_a2a_backend" + ) + def test_tc_piecewise_build_config_reads_phase_config_dataclass( + self, mock_get_moe_a2a_backend + ): + from sglang.srt.model_executor.runner_backend.tc_piecewise_cuda_graph_backend import ( + TcPiecewiseCudaGraphBackend, + ) + + mock_backend = mock_get_moe_a2a_backend.return_value + mock_backend.is_deepep.return_value = False + mock_backend.is_mooncake.return_value = False + server_args = SimpleNamespace( + cuda_graph_config=CudaGraphConfig( + prefill=PhaseConfig( + backend=Backend.TC_PIECEWISE, + bs=[32, 64], + tc_compiler="eager", + ) + ), + enable_torch_compile_debug_mode=False, + ) + + config = TcPiecewiseCudaGraphBackend.build_compilation_config(server_args) + + self.assertEqual(config.get_capture_sizes(), [32, 64]) + self.assertEqual(config.compiler, "eager") + + +class TestCutedslMoeMaxNumTokens(CustomTestCase): """The shared CuteDSL MoE per-forward token bound. Fields are set directly - to exercise the math independently of __post_init__ resolution.""" + to exercise the math independently of __post_init__ resolution. + + cg-refactor: the legacy disable_piecewise_cuda_graph / + piecewise_cuda_graph_max_tokens / cuda_graph_max_bs fields were + consolidated into cuda_graph_config; the helper accepts the legacy + kwarg names for test readability and translates them to the per-phase + dataclasses. + """ def _args(self, **overrides): server_args = ServerArgs(model_path="dummy") @@ -636,8 +694,21 @@ class TestCutedslMoeMaxNumTokens(unittest.TestCase): cuda_graph_max_bs=512, ) fields.update(overrides) + disable_piecewise = fields.pop("disable_piecewise_cuda_graph") + piecewise_max = fields.pop("piecewise_cuda_graph_max_tokens") + cg_max_bs = fields.pop("cuda_graph_max_bs") for key, value in fields.items(): setattr(server_args, key, value) + server_args.cuda_graph_config = CudaGraphConfig( + decode=PhaseConfig(backend=Backend.FULL, max_bs=cg_max_bs), + prefill=PhaseConfig( + backend=( + Backend.DISABLED if disable_piecewise else Backend.TC_PIECEWISE + ), + max_bs=piecewise_max, + tc_compiler="eager", + ), + ) return server_args def test_prefill_dominates_in_default_config(self): diff --git a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py index fe24a63f2..a42cd162c 100644 --- a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py +++ b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py @@ -50,7 +50,10 @@ def _make_worker(num_steps: int, num_draft_tokens: int): worker.device = DEVICE worker.speculative_num_steps = num_steps worker.speculative_num_draft_tokens = num_draft_tokens - worker.server_args = SimpleNamespace(cuda_graph_max_bs=8, max_running_requests=8) + worker.server_args = SimpleNamespace( + cuda_graph_config=SimpleNamespace(decode=SimpleNamespace(max_bs=8)), + max_running_requests=8, + ) return worker