minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
This commit is contained in:
Yuwei An
2026-06-09 21:36:57 -07:00
committed by GitHub
co-authored by BBuf Cheng Wan Lianmin Zheng
parent 56f06278c6
commit 2495c02c2c
160 changed files with 5212 additions and 3083 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs_new/docs/advanced_features/server_arguments.mdx
+117 -33
View File
@@ -1990,22 +1990,52 @@ Please consult the documentation below and [server_args.py](https://github.com/s
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--cuda-graph-max-bs</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Set the maximum batch size for cuda graph. It will extend the cuda graph capture batch size to this value.</td>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--cuda-graph-config</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Canonical per-phase CUDA graph settings as JSON, e.g. <code>{`{"decode":{"backend":"full","max_bs":256},"prefill":{"backend":"tc_piecewise","tc_compiler":"eager"}}`}</code>. JSON wins over the per-phase <code>--cuda-graph-*</code> convenience flags and over the legacy flags. Allowed backends: <code>full</code>, <code>breakable</code>, <code>tc_piecewise</code>, <code>disabled</code> (<code>full</code> is decode-only).</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: JSON (dict-of-dicts)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--cuda-graph-backend-decode</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Backend for the decode phase. Folds into <code>cuda_graph_config[decode].backend</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>full</code>, <code>breakable</code>, <code>tc_piecewise</code>, <code>disabled</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--cuda-graph-backend-prefill</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Backend for the prefill phase. Folds into <code>cuda_graph_config[prefill].backend</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>breakable</code>, <code>tc_piecewise</code>, <code>disabled</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--cuda-graph-max-bs-decode</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Maximum batch size captured for the decode CUDA graph.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: int</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--cuda-graph-bs</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Set the list of batch sizes for cuda graph.</td>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--cuda-graph-max-bs-prefill</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Maximum batch size captured for the prefill CUDA graph.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: int</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--cuda-graph-bs-decode</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Explicit list of batch sizes to capture for the decode CUDA graph.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>List[int]</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--disable-cuda-graph</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Disable cuda graph.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--cuda-graph-bs-prefill</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Explicit list of batch sizes to capture for the prefill CUDA graph.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>List[int]</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--cuda-graph-tc-compiler</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Compiler used by the <code>tc_piecewise</code> backend (only the prefill phase consumes it today).</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>eager</code>, <code>inductor</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--disable-cuda-graph-padding</code></td>
@@ -2019,6 +2049,12 @@ Please consult the documentation below and [server_args.py](https://github.com/s
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--debug-cuda-graph</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>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.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--enable-cudagraph-gc</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Enable garbage collection during CUDA graph capture. If disabled (default), GC is frozen during capture to speed up the process.</td>
@@ -2139,39 +2175,87 @@ Please consult the documentation below and [server_args.py](https://github.com/s
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--disable-piecewise-cuda-graph</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Disable piecewise cuda graph for extend/prefill. PCG is enabled by default.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to disable)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--enforce-piecewise-cuda-graph</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Enforce piecewise cuda graph, skipping all auto-disable conditions. For testing only.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--piecewise-cuda-graph-tokens</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Set the list of tokens when using piecewise cuda graph.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: JSON list</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--piecewise-cuda-graph-compiler</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Set the compiler for piecewise cuda graph. Choices are: eager, inductor.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>eager</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>eager</code>, <code>inductor</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--torch-compile-max-bs</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Set the maximum batch size when using torch compile.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>32</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: int</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--cuda-graph-max-bs</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><strong>Deprecated alias</strong> for <code>--cuda-graph-max-bs-decode</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: int</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--cuda-graph-bs</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><strong>Deprecated alias</strong> for <code>--cuda-graph-bs-decode</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>List[int]</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--disable-cuda-graph</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><strong>Deprecated.</strong> Use <code>--cuda-graph-backend-decode=disabled</code> and/or <code>--cuda-graph-backend-prefill=disabled</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--enable-breakable-cuda-graph</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><strong>Deprecated alias</strong> for <code>--cuda-graph-backend-prefill=breakable</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--prefill-cuda-graph-backend</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><strong>Deprecated alias</strong> for <code>--cuda-graph-backend-prefill</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>breakable</code>, <code>tc_piecewise</code>, <code>disabled</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--decode-cuda-graph-backend</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><strong>Deprecated alias</strong> for <code>--cuda-graph-backend-decode</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>full</code>, <code>breakable</code>, <code>tc_piecewise</code>, <code>disabled</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--disable-prefill-cuda-graph</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><strong>Deprecated.</strong> Use <code>--cuda-graph-backend-prefill=disabled</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--disable-decode-cuda-graph</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><strong>Deprecated.</strong> Use <code>--cuda-graph-backend-decode=disabled</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--disable-piecewise-cuda-graph</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><strong>Deprecated alias</strong> for <code>--cuda-graph-backend-prefill=disabled</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--enforce-piecewise-cuda-graph</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><strong>Deprecated alias</strong> for <code>--cuda-graph-backend-prefill=tc_piecewise</code>. Explicitly setting the prefill backend now skips the auto-disable cascade automatically.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--piecewise-cuda-graph-tokens</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><strong>Deprecated alias</strong> for <code>--cuda-graph-bs-prefill</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>List[int]</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--piecewise-cuda-graph-compiler</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><strong>Deprecated alias</strong> for <code>--cuda-graph-tc-compiler</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>eager</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>eager</code>, <code>inductor</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--piecewise-cuda-graph-max-tokens</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Set the maximum tokens when using piecewise cuda graph.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><strong>Deprecated alias</strong> for <code>--cuda-graph-max-bs-prefill</code>.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>4096</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: int</td>
</tr>
python/sglang/auto_benchmark_lib.py
+2 -2
View File
@@ -58,7 +58,7 @@ PROGRESS_FLAG_KEYS = (
"num_continuous_decode_steps",
"stream_interval",
"page_size",
"cuda_graph_max_bs",
"cuda_graph_max_bs_decode",
"speculative_num_steps",
"speculative_eagle_topk",
"speculative_num_draft_tokens",
@@ -84,7 +84,7 @@ PROGRESS_FLAG_ALIASES = {
"num_continuous_decode_steps": "decode_steps",
"stream_interval": "stream",
"page_size": "page",
"cuda_graph_max_bs": "cg_bs",
"cuda_graph_max_bs_decode": "cg_bs",
"speculative_num_steps": "spec_steps",
"speculative_eagle_topk": "eagle_topk",
"speculative_num_draft_tokens": "draft_tok",
python/sglang/bench_one_batch.py
+5 -1
View File
@@ -77,6 +77,7 @@ from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
from sglang.srt.managers.scheduler_components.dp_attn import prepare_mlp_sync_batch_raw
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
from sglang.srt.model_executor.cuda_graph_config import Phase
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.sampling.sampling_params import SamplingParams
@@ -944,7 +945,10 @@ def latency_test(
def main(server_args, bench_args):
server_args.cuda_graph_max_bs = max(bench_args.batch_size)
# Post-init write to the legacy cuda_graph_max_bs_decode field would
# not propagate to cuda_graph_config; update the decode phase directly.
if server_args.cuda_graph_config is not None:
server_args.cuda_graph_config[Phase.DECODE].max_bs = max(bench_args.batch_size)
_set_envs_and_config(server_args)
python/sglang/compile_deep_gemm.py
+6 -2
View File
@@ -23,6 +23,7 @@ from sglang.srt.entrypoints.warmup import warmup
from sglang.srt.environ import envs
from sglang.srt.managers.io_struct import GenerateReqInput
from sglang.srt.managers.tokenizer_manager import TokenizerManager
from sglang.srt.model_executor.cuda_graph_config import Backend, Phase
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import kill_process_tree
@@ -169,8 +170,11 @@ def launch_server_process_and_send_one_request(
def refine_server_args(server_args: ServerArgs, compile_args: CompileArgs):
# Disable cuda graph and torch compile to save time
server_args.disable_cuda_graph = True
# Disable cuda graph and torch compile to save time. Writes after
# ServerArgs.__post_init__ don't propagate to cuda_graph_config via the
# legacy disable_cuda_graph field, so flip both phases directly.
server_args.cuda_graph_config[Phase.DECODE].backend = Backend.DISABLED
server_args.cuda_graph_config[Phase.PREFILL].backend = Backend.DISABLED
server_args.enable_torch_compile = False
print(f"Disable CUDA Graph and Torch Compile to save time...")
python/sglang/srt/arg_groups/argparse_actions.py
+28
View File
@@ -67,6 +67,34 @@ class DeprecatedStoreTrueAction(argparse.Action):
setattr(namespace, self.dest, True)
class DeprecatedStoreConstAction(argparse.Action):
"""Deprecated boolean flag that stores a fixed string/value into ``dest``
and prints a warning. Used to translate a legacy boolean flag into a
setting on the new per-phase config dict (e.g.
``--disable-piecewise-cuda-graph`` -> ``cuda_graph_backend_prefill="disabled"``)."""
def __init__(
self,
option_strings,
dest,
new_flag=None,
const_value=None,
nargs=0,
default=None,
**kwargs,
):
self.new_flag = new_flag
self.const_value = const_value
super().__init__(option_strings, dest, nargs=nargs, default=default, **kwargs)
def __call__(self, parser, namespace, values, option_string=None):
replacement = f" Use '{self.new_flag}' instead." if self.new_flag else ""
print_deprecated_warning(
f"'{option_string}' is deprecated and will be removed in a future release.{replacement}"
)
setattr(namespace, self.dest, self.const_value)
class DeprecatedAliasStoreAction(argparse.Action):
"""Deprecated alias that stores its value and prints a warning."""
python/sglang/srt/compilation/compile.py
+4 -2
View File
@@ -9,7 +9,9 @@ from typing import Any, Callable, Optional, Union
import torch
from sglang.srt.compilation.compilation_config import CompilationConfig
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
is_in_tc_piecewise_cuda_graph,
)
logger = logging.getLogger(__name__)
@@ -186,7 +188,7 @@ def install_torch_compiled(
state["compiled_callable"] = compiled_callable
def trampoline(self, *args, **kwargs):
use_compiled = is_in_piecewise_cuda_graph()
use_compiled = is_in_tc_piecewise_cuda_graph()
if use_compiled:
if not state["compiled"]:
_ensure_compiled(self, *args, **kwargs)
python/sglang/srt/compilation/compile_phase.py
+57
View File
@@ -0,0 +1,57 @@
"""torch.compile-internal phase markers used by the tc_piecewise backend.
Two pieces of state, both private to the torch.compile path (the
``cuda_piecewise_backend`` FX backend and the runner that drives it):
* ``_in_torch_compile_warmup`` — true during the warmup-compile loop
where we run the compiled callable to trigger inductor compilation
but explicitly do **not** capture into a CUDA graph yet.
``cuda_piecewise_backend`` reads this to short-circuit the capture
branch.
* ``_pcg_capture_stream`` — the CUDA stream on which the runner is
performing capture, surfaced so the FX backend can use the same
stream for its own ``torch.cuda.graph(...)`` calls.
"""
from __future__ import annotations
from contextlib import contextmanager
import torch
_in_torch_compile_warmup = False
_pcg_capture_stream: "torch.cuda.Stream | None" = None
def is_in_torch_compile_warmup() -> bool:
"""True while inside the tc_piecewise warmup-compile pass. Strict subset of
``torch.compiler.is_compiling()``.
"""
return _in_torch_compile_warmup
@contextmanager
def enable_torch_compile_warmup():
"""Mark the enclosed scope as the tc_piecewise warmup-compile pass. The FX
piecewise backend uses this to skip CUDA graph capture during warmup.
"""
global _in_torch_compile_warmup
_in_torch_compile_warmup = True
try:
yield
finally:
_in_torch_compile_warmup = False
def get_pcg_capture_stream() -> "torch.cuda.Stream | None":
return _pcg_capture_stream
@contextmanager
def set_pcg_capture_stream(stream: torch.cuda.Stream):
global _pcg_capture_stream
_pcg_capture_stream = stream
try:
yield
finally:
_pcg_capture_stream = None
python/sglang/srt/compilation/cuda_piecewise_backend.py
+3 -3
View File
@@ -13,9 +13,9 @@ import torch.fx as fx
from sglang.srt.compilation.compilation_config import CompilationConfig
from sglang.srt.compilation.compilation_counter import compilation_counter
from sglang.srt.compilation.piecewise_context_manager import (
from sglang.srt.compilation.compile_phase import (
get_pcg_capture_stream,
is_in_pcg_torch_compile,
is_in_torch_compile_warmup,
)
from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors
from sglang.srt.utils import is_hip
@@ -145,7 +145,7 @@ class CUDAPiecewiseBackend:
if self.is_last_graph and not self.to_be_compiled_sizes:
self.check_for_ending_compilation()
if is_in_pcg_torch_compile():
if is_in_torch_compile_warmup():
return entry.runnable(*args)
if entry.cudagraph is None:
python/sglang/srt/compilation/piecewise_context_manager.py
-138
View File
@@ -1,138 +0,0 @@
from __future__ import annotations
import logging
from contextlib import contextmanager
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, List, Optional
import torch
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
_in_piecewise_cuda_graph = False
_in_pcg_torch_compile = False
_pcg_capture_stream = None
def is_in_piecewise_cuda_graph():
return _in_piecewise_cuda_graph
def is_in_pcg_torch_compile():
return _in_pcg_torch_compile
def get_pcg_capture_stream():
return _pcg_capture_stream
@contextmanager
def enable_piecewise_cuda_graph_compile():
global _in_pcg_torch_compile
_in_pcg_torch_compile = True
yield
_in_pcg_torch_compile = False
@contextmanager
def enable_piecewise_cuda_graph():
global _in_piecewise_cuda_graph
_in_piecewise_cuda_graph = True
try:
yield
except Exception as e:
logger.error(
"Piecewise CUDA Graph failed with error: %s\n%s",
e,
PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG,
)
raise
finally:
_in_piecewise_cuda_graph = False
@contextmanager
def set_pcg_capture_stream(stream: torch.cuda.Stream):
global _pcg_capture_stream
_pcg_capture_stream = stream
yield
_pcg_capture_stream = None
@dataclass
class ForwardContext:
def __init__(self):
self.forward_batch = None
self.attention_layers = None
self.quant_config = None
self.moe_layers = None
self.moe_fusions = None
self.dsa_indexers = None
self.num_tokens: Optional[int] = None
self.raw_num_tokens: Optional[int] = None
def set_forward_batch(self, forward_batch: ForwardBatch):
self.forward_batch = forward_batch
def set_attention_layers(self, layers: List[Any]):
self.attention_layers = layers
def set_quant_config(self, quant_config: Any):
self.quant_config = quant_config
def set_moe_layers(self, layers: List[Any]):
self.moe_layers = layers
def set_moe_fusions(self, fusions: List[Any]):
self.moe_fusions = fusions
def set_dsa_indexers(self, indexers: List[Any]):
self.dsa_indexers = indexers
_forward_context: Optional[ForwardContext] = None
def get_forward_context() -> Optional[ForwardContext]:
if _forward_context is None:
return None
return _forward_context
@contextmanager
def set_forward_context(
forward_batch: ForwardBatch,
attention_layers: List[Any],
quant_config: Any,
moe_layers: List[Any],
moe_fusions: List[Any],
dsa_indexers: Optional[List[Any]] = None,
num_tokens: Optional[int] = None,
raw_num_tokens: Optional[int] = None,
):
global _forward_context
_forward_context = ForwardContext()
_forward_context.set_forward_batch(forward_batch)
_forward_context.set_attention_layers(attention_layers)
_forward_context.set_quant_config(quant_config)
_forward_context.set_moe_layers(moe_layers)
_forward_context.set_moe_fusions(moe_fusions)
if dsa_indexers is not None:
_forward_context.set_dsa_indexers(dsa_indexers)
_forward_context.num_tokens = num_tokens
_forward_context.raw_num_tokens = raw_num_tokens
try:
yield
finally:
_forward_context = None
PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG = (
"Piecewise CUDA Graph is enabled by default as an experimental feature.\n"
"To work around this error, add --disable-piecewise-cuda-graph to your launch command.\n"
"Please report this issue at https://github.com/sgl-project/sglang/issues/new/choose"
)
python/sglang/srt/compilation/torch_compile_decoration.py
+83
View File
@@ -0,0 +1,83 @@
"""torch.compile decoration helpers used by the decode-Full path under
``--enable-torch-compile``.
``patch_model`` wraps the model forward with ``torch.compile`` for batch
sizes that fall in the compile bucket list and returns the raw forward
otherwise. ``set_torch_compile_config`` flips the inductor/dynamo config
flags expected by that path.
Note: the prefill-tc_piecewise path (``TcPiecewiseCudaGraphBackend``) does NOT
use ``patch_model`` — it goes through ``compilation/compile.py``'s
``install_torch_compiled``. ``_to_torch`` here is duplicated by
tc_piecewise's local ``_toggle_multi_platform_ops``; the duplication is kept
because the two paths have different lifecycle requirements.
"""
from __future__ import annotations
import os
from contextlib import contextmanager
import torch
from sglang.srt.distributed.parallel_state import GroupCoordinator
from sglang.srt.layers.utils import MultiPlatformOp
from sglang.srt.utils import get_bool_env_var, is_hip
from sglang.srt.utils.patch_torch import monkey_patch_torch_compile
_is_hip = is_hip()
def _to_torch(model: torch.nn.Module, reverse: bool, num_tokens: int) -> None:
for sub in model._modules.values():
if isinstance(sub, MultiPlatformOp):
if reverse:
sub.leave_torch_compile()
else:
sub.enter_torch_compile(num_tokens=num_tokens)
if isinstance(sub, torch.nn.Module):
_to_torch(sub, reverse, num_tokens)
@contextmanager
def patch_model(
model: torch.nn.Module,
enable_compile: bool,
num_tokens: int,
tp_group: GroupCoordinator,
):
"""Patch the model to make it compatible with torch.compile."""
backup_ca_comm = None
try:
if enable_compile:
_to_torch(model, reverse=False, num_tokens=num_tokens)
backup_ca_comm = tp_group.ca_comm
yield torch.compile(
torch.no_grad()(model.forward),
mode=os.environ.get(
"SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs"
),
dynamic=_is_hip and get_bool_env_var("SGLANG_TORCH_DYNAMIC_SHAPE"),
)
else:
yield model.forward
finally:
if enable_compile:
_to_torch(model, reverse=True, num_tokens=num_tokens)
tp_group.ca_comm = backup_ca_comm
def set_torch_compile_config() -> None:
import torch._dynamo.config
import torch._inductor.config
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.triton.unique_kernel_names = True
torch._inductor.config.fx_graph_cache = True
torch._dynamo.config.accumulated_cache_size_limit = 1024
if hasattr(torch._dynamo.config, "cache_size_limit"):
torch._dynamo.config.cache_size_limit = 1024
monkey_patch_torch_compile()
python/sglang/srt/debug_utils/pr_fix_toggle.py
+1 -1
View File
@@ -24,7 +24,7 @@ patches:
replacement: |
hidden_states = logits_output.hidden_states
- target: sglang.srt.speculative.eagle_draft_cuda_graph_runner.EAGLEDraftCudaGraphRunner.capture_one_batch_size
- target: sglang.srt.speculative.eagle_draft_cuda_graph_runner.EAGLEDraftCudaGraphRunner.capture_one_shape
edits:
- match: |
forward_batch.spec_info.hidden_states = hidden_states_backup
python/sglang/srt/distributed/device_communicators/custom_all_reduce.py
+6 -4
View File
@@ -13,13 +13,15 @@ import torch.distributed as dist
from torch.distributed import ProcessGroup
import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
can_use_custom_all_reduce_with_nvlink,
is_weak_contiguous,
)
from sglang.srt.environ import envs
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.utils import (
get_bool_env_var,
is_cuda,
@@ -316,7 +318,7 @@ class CustomAllreduce:
# Could be warmup OR piecewise cuda graph split op execution.
# In piecewise cuda graph, split ops run eagerly outside the graph
# but _IS_CAPTURING is still True. We need to do real all-reduce.
if is_in_piecewise_cuda_graph():
if is_in_tc_piecewise_cuda_graph():
# Split op execution - do real all-reduce
return self._all_reduce_impl(input, registered=False)
else:
@@ -350,8 +352,8 @@ def dispatch_custom_allreduce(
On CUDA, the JIT-compiled v2 implementation is used by default.
Set SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2=0 to fall back to the legacy CustomAllreduce.
Note: ``ServerArgs._handle_environment_variables`` forces this env to "0" when
``nnodes > 1`` since custom AR is intra-node only.
Note: ServerArgs._handle_environment_variables forces this env to "0" when
nnodes > 1 since custom AR is intra-node only.
"""
if _is_cuda and envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get():
from .custom_all_reduce_v2 import (
python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py
+4 -2
View File
@@ -8,11 +8,13 @@ import torch.distributed as dist
from torch.distributed import ProcessGroup
from sglang.jit_kernel.all_reduce import AllReduceAlgo, get_custom_all_reduce_cls
from sglang.srt.distributed import is_in_piecewise_cuda_graph
from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
can_use_custom_all_reduce_with_nvlink,
is_weak_contiguous,
)
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.utils import is_sm100_supported, log_info_on_rank0
logger = logging.getLogger(__name__)
@@ -121,7 +123,7 @@ class CustomAllReduceV2:
return inp_size <= self.max_size
def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor:
if is_in_piecewise_cuda_graph(): # disable inplace optimization
if is_in_tc_piecewise_cuda_graph(): # disable inplace optimization
try:
self.obj.set_cuda_graph_capture(False)
return self._all_reduce(input)
python/sglang/srt/distributed/device_communicators/pymscclpp.py
+7 -5
View File
@@ -7,10 +7,12 @@ import torch
import torch.distributed as dist
from torch.distributed import ProcessGroup, ReduceOp
from sglang.srt.compilation.piecewise_context_manager import (
from sglang.srt.compilation.compile_phase import (
get_pcg_capture_stream,
is_in_pcg_torch_compile,
is_in_piecewise_cuda_graph,
is_in_torch_compile_warmup,
)
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.server_args import get_global_server_args
@@ -343,8 +345,8 @@ class PyMscclppCommunicator:
# (compile, capture, or replay) as it changes the allreduce dispatch
# path and triggers recompilation.
if (
is_in_piecewise_cuda_graph()
or is_in_pcg_torch_compile()
is_in_tc_piecewise_cuda_graph()
or is_in_torch_compile_warmup()
or get_pcg_capture_stream() is not None
):
return False
python/sglang/srt/distributed/parallel_state.py
+6 -4
View File
@@ -43,9 +43,11 @@ import torch.distributed
from torch.distributed import Backend, ProcessGroup
from sglang.srt.compilation.compilation_config import register_split_op
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.distributed.utils import set_global_tcp_store
from sglang.srt.environ import envs
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.utils import (
get_current_device_stream_fast,
get_int_env_var,
@@ -641,7 +643,7 @@ class GroupCoordinator:
and self.torch_symm_mem_comm.should_torch_symm_mem_allreduce(input_)
):
outplace_all_reduce_method = "torch_symm_mem"
elif is_in_piecewise_cuda_graph() and self.pynccl_comm is not None:
elif is_in_tc_piecewise_cuda_graph() and self.pynccl_comm is not None:
# For piecewise cuda graph, we use pynccl outplace allreduce
outplace_all_reduce_method = "pynccl"
if outplace_all_reduce_method is not None:
@@ -708,7 +710,7 @@ class GroupCoordinator:
if (
getattr(ca_comm, "_IS_CAPTURING", False)
and not torch.cuda.is_current_stream_capturing()
and is_in_piecewise_cuda_graph()
and is_in_tc_piecewise_cuda_graph()
):
if not hasattr(ca_comm, "fused_ar_rms"):
return None
@@ -863,7 +865,7 @@ class GroupCoordinator:
if getattr(ca_comm, "_IS_CAPTURING", False):
if torch.cuda.is_current_stream_capturing():
ca_comm.all_gather_reg(input, out=output, dim=0)
elif is_in_piecewise_cuda_graph():
elif is_in_tc_piecewise_cuda_graph():
ca_comm.all_gather_unreg(input, out=output, dim=0)
else:
# True CUDA graph warmup: avoid a different host collective.
python/sglang/srt/hardware_backend/mlx/model_runner_stub.py
+1 -1
View File
@@ -186,7 +186,7 @@ class MlxModelRunnerStub(ModelRunner):
)
# No CUDA graphs, no attention backend
self.graph_runner = None
self.decode_cuda_graph_runner = None
self.graph_mem_usage = 0
self.attn_backend = None
python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py
+2 -1
View File
@@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, List, Optional
import torch
from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
MultiLayerEagleDraftExtendCudaGraphRunner,
@@ -82,7 +83,7 @@ class MultiLayerEagleMultiStepDraftExtendNpuGraphRunner(
super().__init__(eagle_worker)
def _init_and_capture(self):
if self.eagle_worker.server_args.disable_cuda_graph:
if cuda_graph_fully_disabled():
self.runners = [None] * self.speculative_num_steps
return
python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py
+166
View File
@@ -0,0 +1,166 @@
"""NPUCudaGraphBackend — Ascend NPU full-graph capture (torch.npu.NPUGraph).
Mirrors FullCudaGraphBackend with two differences:
- Captures via torch.npu.graph(...) into torch.npu.NPUGraph.
- replay_with_input_update(shape_key, seq_lens, attr_name) rebinds
the recorded graph's input bindings for variable seq_lens at replay
time (NPU's NPUGraph.update(...) API).
torch.npu is imported lazily inside methods so the module loads on
non-NPU hosts.
"""
from __future__ import annotations
import threading
from contextlib import AbstractContextManager, contextmanager
from functools import partial
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
import numpy as np
import torch
from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
set_graph_pool_id,
)
from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
BaseCudaGraphBackend,
)
from sglang.srt.utils import empty_context, get_bool_env_var
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
BaseCudaGraphRunner,
)
class NPUCudaGraphBackend(BaseCudaGraphBackend):
"""One torch.npu.NPUGraph per shape; attention metadata captured
inside the graph. replay_with_input_update substitutes fresh
seq_lens without re-recording."""
def __init__(
self,
cuda_graph_runner: BaseCudaGraphRunner,
*,
enable_memory_saver: bool = False,
) -> None:
self._graphs: Dict[Any, Any] = {}
self._outputs: Dict[Any, Any] = {}
self._pool = None
self._device_module = cuda_graph_runner.device_module
self._tp_group = cuda_graph_runner.model_runner.tp_group
self._capture_stream = None
self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
enable=enable_memory_saver
and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
)
self._enable_torch_compile = getattr(
cuda_graph_runner, "enable_torch_compile", False
)
@contextmanager
def capture_session(self, stream):
if self._pool is None:
self._pool = self._device_module.graph_pool_handle()
set_graph_pool_id(self._pool)
self._capture_stream = stream
try:
yield
finally:
self._capture_stream = None
def capture_one(
self,
shape_key: Any,
forward_fn: Callable[[], Any],
dummies: Optional[Any] = None,
post_warmup_hook: Optional[Callable[[], None]] = None,
) -> None:
import torch_npu # noqa: F401 (verifies NPU availability)
# Two warmups so kernels are loaded and one-time setup is paid before capture.
# post_warmup_hook lets the attention backend reset state that warmup mutated.
for _ in range(2):
self._device_module.synchronize()
self._tp_group.barrier()
forward_fn()
if post_warmup_hook is not None:
post_warmup_hook()
graph = torch.npu.NPUGraph()
if self._enable_torch_compile:
skip_guard_context = torch.compiler.set_stance(skip_guard_eval_unsafe=True)
else:
skip_guard_context = empty_context()
graph_ctx: Callable[..., AbstractContextManager]
if (
self._memory_saver_adapter is not None
and self._memory_saver_adapter.enabled
):
graph_ctx = partial(
self._memory_saver_adapter.cuda_graph,
tag=GPU_MEMORY_TYPE_CUDA_GRAPH,
)
else:
graph_ctx = torch.npu.graph
with skip_guard_context, graph_ctx(
graph,
pool=self._pool,
stream=self._capture_stream,
auto_dispatch_capture=True,
):
out = forward_fn()
self._graphs[shape_key] = graph
self._outputs[shape_key] = out
def can_run(self, forward_batch: ForwardBatch, shape_key: Any) -> bool:
return shape_key in self._graphs
@contextmanager
def replay_session(self):
yield
def replay(
self,
shape_key: Any,
static_forward_batch: ForwardBatch,
**kwargs,
) -> Any:
self._graphs[shape_key].replay()
return self._outputs[shape_key]
def replay_with_input_update(
self,
shape_key: Any,
seq_lens: list,
attr_name: str,
attr_type: Any,
) -> Any:
"""Rebind seq_lens on the recorded NPU graph in a background
thread, then replay. Used when the model is not deepseek-nsa."""
if isinstance(attr_type, torch.Tensor):
seq_lens = torch.from_numpy(np.array(seq_lens).astype(np.int32))
graph = self._graphs[shape_key]
def _update():
graph.update(cpu_update_input=[{attr_name: seq_lens}])
thread = threading.Thread(target=_update)
thread.start()
graph.replay()
thread.join()
return self._outputs[shape_key]
def cleanup(self) -> None:
self._graphs.clear()
self._outputs.clear()
self._pool = None
python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py
+29 -16
View File
@@ -11,13 +11,23 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Run the model with npu graph and torch.compile."""
"""Run the model with NPU graph and torch.compile.
NPUGraphRunner is a thin subclass of DecodeCudaGraphRunner: the
factory returns NPUCudaGraphBackend for NPU devices, so all
capture/replay mechanics live in the backend. This class adds:
- NPU-specific patch_model monkey-patch for the decode-Full +
torch.compile path.
- Profile context override (NPU profiler emits to disk, not in-mem).
- Replay override that issues an async NPUGraph.update for
seq_lens before replay (skipped for deepseek-nsa).
- Smaller cache_loc dtype (int32 instead of int64).
"""
from __future__ import annotations
import logging
import os
import threading
from contextlib import contextmanager
from pathlib import Path
from typing import TYPE_CHECKING, Dict, Optional, Union
@@ -25,11 +35,10 @@ from typing import TYPE_CHECKING, Dict, Optional, Union
import numpy as np
import torch
import sglang
from sglang.srt.configs.model_config import AttentionArch, is_deepseek_dsa
from sglang.srt.distributed.parallel_state import GroupCoordinator
from sglang.srt.environ import envs
from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
from sglang.srt.model_executor.runner import DecodeCudaGraphRunner
from sglang.srt.utils import (
empty_context,
get_bool_env_var,
@@ -71,8 +80,8 @@ def patch_model_npu(
yield model.forward
class NPUGraphRunner(CudaGraphRunner):
"""A NPUGraphRunner runs the forward pass of a model with npu graph and torch.compile."""
class NPUGraphRunner(DecodeCudaGraphRunner):
"""A NPUGraphRunner runs the forward pass of a model with NPU graph and torch.compile."""
def __init__(
self,
@@ -82,7 +91,11 @@ class NPUGraphRunner(CudaGraphRunner):
speculative_num_steps: Optional[int] = None,
speculative_num_draft_tokens: Optional[int] = None,
):
sglang.srt.model_executor.cuda_graph_runner.patch_model = patch_model_npu
# NPU patch_model override: monkey-patch torch_compile_decoration's
# patch_model with the NPU-specific version.
from sglang.srt.compilation import torch_compile_decoration
torch_compile_decoration.patch_model = patch_model_npu
super().__init__(
model_runner,
attn_backend=attn_backend,
@@ -215,9 +228,8 @@ class NPUGraphRunner(CudaGraphRunner):
forward_batch.mrope_positions
)
self.update_attr_name = self._get_update_attr_name()
self.update_attr_type = self._get_update_attr_type()
# Replay
graph_key = self._make_graph_key(self.bs)
if not is_deepseek_dsa(self.model_runner.model_config.hf_config):
if forward_batch.forward_mode.is_target_verify():
seq_lens_cpu = forward_batch.seq_lens.cpu() + self.num_tokens_per_bs
@@ -226,14 +238,15 @@ class NPUGraphRunner(CudaGraphRunner):
seq_lens = forward_batch.seq_lens.cpu().tolist() + [0] * (
self.bs - self.raw_bs
)
thread = threading.Thread(target=self._update_inputs, args=(seq_lens,))
thread.start()
self.graphs[self.bs].replay()
thread.join()
output = self.backend.replay_with_input_update(
graph_key,
seq_lens=seq_lens,
attr_name=self.attr_name[AttentionArch.MLA],
attr_type=self.attr_type[AttentionArch.MLA],
)
else:
self.graphs[self.bs].replay()
output = self.backend.replay(graph_key, forward_batch)
output = self.output_buffers[self.bs]
if isinstance(output, LogitsProcessorOutput):
if self.is_dllm:
next_token_logits = None
python/sglang/srt/hardware_backend/npu/utils.py
+11 -10
View File
@@ -56,27 +56,28 @@ def set_default_server_args(args: "ServerArgs"):
args.page_size = 128
# NPU memory settings
decode = args.cuda_graph_config.decode
npu_mem = get_npu_memory_capacity()
if npu_mem <= 32 * 1024:
# Ascend 910B4,910B4_1
# (chunked_prefill_size 4k, cuda_graph_max_bs 16 if tp < 4 else 64)
# (chunked_prefill_size 4k, max_bs 16 if tp < 4 else 64)
if args.chunked_prefill_size is None:
args.chunked_prefill_size = 4 * 1024
if args.cuda_graph_max_bs is None:
if decode.max_bs is None:
if args.tp_size < 4:
args.cuda_graph_max_bs = 16
decode.max_bs = 16
else:
args.cuda_graph_max_bs = 64
decode.max_bs = 64
elif npu_mem <= 64 * 1024:
# Ascend 910B1,910B2,910B2C,910B3,910_9391,910_9392,910_9381,910_9382,910_9372,910_9362
# (chunked_prefill_size 8k, cuda_graph_max_bs 64 if tp < 4 else 256)
# (chunked_prefill_size 8k, max_bs 64 if tp < 4 else 256)
if args.chunked_prefill_size is None:
args.chunked_prefill_size = 8 * 1024
if args.cuda_graph_max_bs is None:
if decode.max_bs is None:
if args.tp_size < 4:
args.cuda_graph_max_bs = 64
decode.max_bs = 64
else:
args.cuda_graph_max_bs = 256
decode.max_bs = 256
# NPU does not support CustomAllReduce
args.disable_custom_all_reduce = True
@@ -216,7 +217,7 @@ def init_zbal(world_size, gpu_id, world_rank, do_check=True):
gva_is_inited = True
if do_check and not ret:
logger.error(f"[ZBAL] zbal init failed!")
logger.error("[ZBAL] zbal init failed!")
sys.exit(-1)
return ret
@@ -271,7 +272,7 @@ def lazy_init_zbal_gva_mem(
gva_is_inited = True
if do_check and not res:
logger.error(f"[ZBAL] zbal lazy init failed!")
logger.error("[ZBAL] zbal lazy init failed!")
sys.exit(-1)
return res
python/sglang/srt/kv_canary/api.py
+9 -3
View File
@@ -12,6 +12,11 @@ from sglang.srt.kv_canary.pool_patcher.api import attach_canary_buffers
from sglang.srt.kv_canary.pool_patcher.utils import wrap_method
from sglang.srt.kv_canary.runner.canary_manager import CanaryManager
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
if TYPE_CHECKING:
@@ -32,10 +37,11 @@ def install_canary(
if config.mode is CanaryMode.NONE:
return None
assert server_args.disable_piecewise_cuda_graph, (
assert not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE), (
"kv-canary: piecewise cuda graph is not supported by the current "
"SingleForwardManager design; pass --disable-piecewise-cuda-graph "
"when canary is enabled"
"SingleForwardManager design; set "
"--cuda-graph-backend-prefill=disabled (or =breakable) when canary "
"is enabled"
)
perturb_config = PerturbConfig.from_env()
python/sglang/srt/kv_canary/capacities.py
+4 -1
View File
@@ -63,7 +63,10 @@ class CanaryLaunchCapacities:
f"kv-canary: pool_slot_count must be positive, got {pool_slot_count}"
)
cuda_graph_max_bs = server_args.cuda_graph_max_bs or 0
cuda_graph_config = server_args.cuda_graph_config
cuda_graph_max_bs = (
cuda_graph_config.decode.max_bs if cuda_graph_config is not None else 0
) or 0
if cuda_graph_max_bs < 0:
raise ValueError(
f"kv-canary: cuda_graph_max_bs must be non-negative, got {cuda_graph_max_bs}"
python/sglang/srt/layers/activation.py
+8
View File
@@ -30,6 +30,11 @@ from sglang.srt.distributed import (
from sglang.srt.environ import envs
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils import MultiPlatformOp
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import (
cpu_has_amx_support,
@@ -128,6 +133,9 @@ class SiluAndMul(MultiPlatformOp):
return out
def forward_musa(self, x: torch.Tensor) -> torch.Tensor:
if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE):
return self.forward_native(x)
if not hasattr(self, "_musa_swish_glu"):
# XXX (MUSA): nn.SwishGLU seems to have better performance than silu_and_mul on MUSA, we can switch to it for now. We can consider implementing a silu_and_mul kernel for MUSA in the future if needed.
self._musa_swish_glu = nn.SwishGLU()
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+17 -17
View File
@@ -12,10 +12,6 @@ from sglang.jit_kernel.fused_store_index_cache import (
can_use_dsa_fused_store,
fused_store_index_k_cache,
)
from sglang.srt.compilation.piecewise_context_manager import (
get_forward_context,
is_in_piecewise_cuda_graph,
)
from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa.utils import (
aiter_can_use_preshuffle_paged_mqa,
@@ -26,6 +22,10 @@ from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor
from sglang.srt.layers.layernorm import LayerNorm
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz
from sglang.srt.layers.utils import MultiPlatformOp
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
get_tc_piecewise_forward_context,
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.state_capturer.indexer_topk import (
maybe_capture_indexer_topk,
)
@@ -83,13 +83,13 @@ from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_output
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_context import (
get_attn_backend,
get_req_to_token_pool,
get_token_to_kv_pool,
)
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.server_args import get_global_server_args
_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get()
@@ -118,8 +118,8 @@ if _is_cuda:
), "Internal error: piecewise CUDA graph is only supported on CUDA"
from sglang.srt.layers.attention.dsa.triton_kernel import act_quant
forward_batch = get_forward_context().forward_batch
indexer = get_forward_context().dsa_indexers[layer_id]
forward_batch = get_tc_piecewise_forward_context().forward_batch
indexer = get_tc_piecewise_forward_context().dsa_indexers[layer_id]
metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch)
# slice off padding from piecewise CUDA graph
@@ -198,7 +198,7 @@ def _broadcast_indexer_topk_from_rank0(
if topk_indices is None or not envs.SGLANG_DSA_TOPK_BROADCAST.get():
return topk_indices
if is_in_piecewise_cuda_graph():
if is_in_tc_piecewise_cuda_graph():
broadcast_indexer_topk_from_rank0_(topk_indices)
else:
_broadcast_indexer_topk_from_rank0_impl(topk_indices)
@@ -999,7 +999,7 @@ class Indexer(MultiPlatformOp):
cp_index: List[Tuple[int, int, int]] = None,
) -> torch.Tensor:
assert (
not is_in_piecewise_cuda_graph()
not is_in_tc_piecewise_cuda_graph()
), "DSA context parallel (_get_topk_ragged_with_cp) not supported under piecewise CUDA graph"
if TYPE_CHECKING:
assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool)
@@ -1148,7 +1148,7 @@ class Indexer(MultiPlatformOp):
layer_id: int,
) -> Optional[torch.Tensor]:
assert (
not is_in_piecewise_cuda_graph()
not is_in_tc_piecewise_cuda_graph()
), "DSA forward_indexer (non-CUDA loop path) not supported under piecewise CUDA graph"
if not _is_npu:
from sglang.srt.layers.attention.dsa.tilelang_kernel import fp8_index
@@ -1339,10 +1339,10 @@ class Indexer(MultiPlatformOp):
# a tuple like (x_fp8, x_scale[, y]). Use `x_meta` for shape/device queries.
x_meta = x[0] if isinstance(x, tuple) else x
# In piecewise CUDA graph mode, metadata is fetched inside custom ops via get_forward_context() to
# In piecewise CUDA graph mode, metadata is fetched inside custom ops via get_tc_piecewise_forward_context() to
# prevent Dynamo from guarding on forward_metadata identity (which changes each
# replay when init_forward_metadata creates a new ForwardMetadata object).
if not is_in_piecewise_cuda_graph():
if not is_in_tc_piecewise_cuda_graph():
metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch)
if metadata is None:
return None
@@ -1360,7 +1360,7 @@ class Indexer(MultiPlatformOp):
# We can only skip the logits computation if cuda graph is not involved
skip_logits_computation = False
if (
not is_in_piecewise_cuda_graph()
not is_in_tc_piecewise_cuda_graph()
and forward_batch.forward_mode.is_extend_without_speculative()
):
if forward_batch.seq_lens_cpu is not None:
@@ -1417,7 +1417,7 @@ class Indexer(MultiPlatformOp):
act_quant=act_quant,
)
current_stream.wait_stream(self.alt_stream)
elif not is_in_piecewise_cuda_graph():
elif not is_in_tc_piecewise_cuda_graph():
q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
self._store_index_k_cache(
forward_batch=forward_batch,
@@ -1471,7 +1471,7 @@ class Indexer(MultiPlatformOp):
else:
x_for_gate = x
if is_in_piecewise_cuda_graph():
if is_in_tc_piecewise_cuda_graph():
weights = logits_head_gate_pcg(
x_for_gate,
self.weights_proj.weight,
@@ -1485,7 +1485,7 @@ class Indexer(MultiPlatformOp):
if _is_cuda or _is_hip:
# In piecewise CUDA graph, any access to seq_lens_cpu creates a Dynamo shape guard.
# Piecewise CUDA graph never has empty batches.
if not is_in_piecewise_cuda_graph():
if not is_in_tc_piecewise_cuda_graph():
assert forward_batch.seq_lens_cpu is not None
if len(forward_batch.seq_lens_cpu) == 0:
# this seems b/c max-pad, no worries?
@@ -1556,7 +1556,7 @@ class Indexer(MultiPlatformOp):
topk_result = torch.cat([topk_result_prev, topk_result_next], dim=0)
topk_result = _broadcast_indexer_topk_from_rank0(topk_result)
return maybe_capture_indexer_topk(layer_id, topk_result)
elif is_in_piecewise_cuda_graph():
elif is_in_tc_piecewise_cuda_graph():
assert (
not enable_dual_stream
), "Internal error: piecewise CUDA graph should not be enabled with dual stream"
python/sglang/srt/layers/attention/dsa_backend.py
+3 -3
View File
@@ -2252,13 +2252,13 @@ class DeepseekSparseAttnBackend(
"""
Decide all attention prefill dispatch strategies for this batch.
"""
from sglang.srt.compilation.piecewise_context_manager import (
is_in_piecewise_cuda_graph,
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.utils import get_device_sm, is_blackwell
# Decide MHA vs MLA
if is_in_piecewise_cuda_graph():
if is_in_tc_piecewise_cuda_graph():
# Can't branch on seq_lens_cpu in PCG, force mha off to guarantee correctness.
self.use_mha = False
elif (
python/sglang/srt/layers/attention/fla/layernorm_gated.py
+7 -7
View File
@@ -15,7 +15,11 @@ import triton.language as tl
from einops import rearrange
from sglang.jit_kernel.utils import is_arch_support_pdl
from sglang.srt.server_args import get_global_server_args
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
)
from sglang.srt.utils import (
cdiv,
cpu_has_amx_support,
@@ -190,12 +194,8 @@ def _get_sm_count(device: torch.device) -> int:
def calc_rows_per_block(M: int, device: torch.device) -> int:
# When piecewise cuda graph is enabled, use a constant value to avoid
# torch.compile creating guards on the dynamic batch dimension.
try:
if not get_global_server_args().disable_piecewise_cuda_graph:
return MAX_ROWS_PER_BLOCK
except ValueError:
# Global server args not initialized (e.g., in unit tests)
pass
if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE):
return MAX_ROWS_PER_BLOCK
sm_count = _get_sm_count(device)
rows_per_block = next_power_of_2(cdiv(M, 2 * sm_count))
rows_per_block = min(rows_per_block, MAX_ROWS_PER_BLOCK)
python/sglang/srt/layers/attention/flashattention_backend.py
+3 -5
View File
@@ -34,6 +34,7 @@ from sglang.jit_kernel.flash_attention import (
flash_attn_varlen_func,
flash_attn_with_kvcache,
)
from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
@dataclass
@@ -212,10 +213,7 @@ class FlashAttentionBackend(AttentionBackend):
self.num_splits = (
1
if model_runner.server_args.enable_deterministic_inference
or (
self.fa_impl_ver == 4
and not model_runner.server_args.disable_cuda_graph
)
or (self.fa_impl_ver == 4 and not cuda_graph_fully_disabled())
else 0
)
@@ -2028,7 +2026,7 @@ class FlashAttentionBackend(AttentionBackend):
"""Shared capture+replay body for the cuda-graph init path.
Public entry: :py:meth:`init_forward_metadata_out_graph`. This helper
formerly lived as the legacy ``init_forward_metadata_replay_cuda_graph``;
formerly lived as the legacy init_forward_metadata_replay_cuda_graph;
the capture path used to wrap it. Both legacy method overrides
are gone.
"""
python/sglang/srt/layers/attention/flashinfer_backend.py
+12 -3
View File
@@ -17,7 +17,6 @@ from typing import TYPE_CHECKING, Callable, List, Optional, Union
import torch
from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.dllm.config import DllmConfig
from sglang.srt.environ import envs
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
@@ -28,7 +27,15 @@ from sglang.srt.layers.attention.utils import (
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.layers.radix_attention import AttentionType
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.speculative.spec_info import SpecInput
from sglang.srt.speculative.spec_utils import (
draft_kv_indices_buffer_width,
@@ -281,7 +288,9 @@ class FlashInferAttnBackend(AttentionBackend):
fmha_backend = "auto"
if is_sm100_supported():
if not model_runner.server_args.disable_piecewise_cuda_graph:
# Disable CUTLASS backend when piecewise cuda graph is enabled
# due to TMA descriptor initialization issues on B200
if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE):
logger.info(
"CUTLASS backend is disabled when piecewise cuda graph is enabled "
"due to TMA descriptor initialization issues on SM100 GPUs. "
@@ -592,7 +601,7 @@ class FlashInferAttnBackend(AttentionBackend):
else:
use_ragged = (
not self.enable_deterministic
and not is_in_piecewise_cuda_graph()
and not is_in_tc_piecewise_cuda_graph()
and not self.use_paged
)
extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu)
python/sglang/srt/layers/attention/flashinfer_mla_backend.py
+4 -2
View File
@@ -15,7 +15,6 @@ from typing import TYPE_CHECKING, Callable, Optional, Union
import torch
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.environ import envs
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.flashinfer_backend import (
@@ -24,6 +23,9 @@ from sglang.srt.layers.attention.flashinfer_backend import (
from sglang.srt.layers.attention.utils import assert_buffer_fits
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.speculative.spec_info import SpecInput
from sglang.srt.speculative.spec_utils import (
@@ -408,7 +410,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
not get_global_server_args().flashinfer_mla_disable_ragged
and extend_no_prefix
# Piecewise cuda graph should use paged prefill to be compatible with prefix cache
and not is_in_piecewise_cuda_graph()
and not is_in_tc_piecewise_cuda_graph()
)
self.indices_updater_prefill.update(
python/sglang/srt/layers/attention/tokenspeed_mla_backend.py
+1 -1
View File
@@ -222,7 +222,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
k_pe: torch.Tensor,
positions: torch.Tensor,
layer: "DeepseekV2AttentionMLA",
forward_batch: "ForwardBatch",
forward_batch: ForwardBatch,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Build FP8 (Q, K, V) for the FMHA kernel and write FP8 KV cache."""
kv = layer.kv_b_proj(kv_a)[0]
python/sglang/srt/layers/attention/triton_backend.py
+10 -9
View File
@@ -15,6 +15,7 @@ from sglang.srt.layers.attention.triton_ops.metadata import get_num_kv_splits_tr
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.layers.radix_attention import AttentionType
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.speculative.spec_utils import (
draft_kv_indices_buffer_width,
@@ -184,8 +185,8 @@ class TritonAttnBackend(AttentionBackend):
self.use_pdl = False
self.allow_bidirectional_attention_in_extend = (
model_runner.server_args.disable_cuda_graph
and (model_runner.server_args.chunked_prefill_size == -1)
cuda_graph_fully_disabled()
and model_runner.server_args.chunked_prefill_size == -1
)
# Decide whether enable deterministic inference with batch-invariant operations
@@ -332,8 +333,8 @@ class TritonAttnBackend(AttentionBackend):
):
"""Fill KV (and SWA) cuda-graph buffers for decode/idle mode.
Returns ``(kv_indptr, window_kv_indptr, window_kv_lens)`` where
``window_kv_lens`` is ``None`` when sliding-window is disabled.
Returns (kv_indptr, window_kv_indptr, window_kv_lens) where
window_kv_lens is None when sliding-window is disabled.
"""
seq_lens = seq_lens[:bs]
req_pool_indices = req_pool_indices[:bs]
@@ -431,7 +432,7 @@ class TritonAttnBackend(AttentionBackend):
):
"""Fill QO + KV cuda-graph buffers for draft_extend mode.
Returns ``(qo_indptr, kv_indptr, num_tokens_per_bs)``.
Returns (qo_indptr, kv_indptr, num_tokens_per_bs).
"""
seq_lens = seq_lens[:bs]
# V2 draft-extend fills num_draft_tokens per req (the cuda-graph runner's
@@ -848,7 +849,7 @@ class TritonAttnBackend(AttentionBackend):
Called by capture after the buffer-update helpers have already run
(either via replay or directly). All fields reference the same
``self.cuda_graph_*`` tensors that the captured graph kernels will
self.cuda_graph_* tensors that the captured graph kernels will
read — the Python object is rebuilt each capture, but the underlying
GPU memory addresses are stable.
"""
@@ -1514,9 +1515,9 @@ def update_sliding_window_buffer(
):
"""Fill window KV buffers for sliding-window attention.
Pass ``window_kv_indices`` to write into a pre-allocated buffer (CUDA-graph
path); omit it (or pass ``None``) to allocate a fresh tensor (eager path,
requires ``device``).
Pass window_kv_indices to write into a pre-allocated buffer (CUDA-graph
path); omit it (or pass None) to allocate a fresh tensor (eager path,
requires device).
"""
window_kv_lens = torch.minimum(
seq_lens,
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+8 -6
View File
@@ -14,7 +14,6 @@ import triton
import triton.language as tl
from sglang.jit_kernel.fixup_zero_kv import fixup_zero_kv_rows
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.environ import envs
from sglang.srt.layers.attention.flashinfer_mla_backend import (
FlashInferMLAAttnBackend,
@@ -30,6 +29,9 @@ from sglang.srt.layers.attention.utils import (
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import is_flashinfer_available, is_float4_e2m1fn_x2
@@ -562,11 +564,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
"""Get the fill value for sequence lengths in CUDA graph."""
return 1
def init_mha_chunk_metadata(self, forward_batch: "ForwardBatch") -> None:
def init_mha_chunk_metadata(self, forward_batch: ForwardBatch) -> None:
has_prefix = any(forward_batch.extend_prefix_lens_cpu)
fallback_to_flashinfer_impl = (
self.disable_chunked_prefix_cache and has_prefix
) or is_in_piecewise_cuda_graph()
) or is_in_tc_piecewise_cuda_graph()
if fallback_to_flashinfer_impl:
super().init_mha_chunk_metadata(
forward_batch, disable_flashinfer_ragged=True
@@ -627,7 +629,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
has_prefix = any(forward_batch.extend_prefix_lens_cpu)
fallback_to_flashinfer_impl = (
self.disable_chunked_prefix_cache and has_prefix
) or is_in_piecewise_cuda_graph()
) or is_in_tc_piecewise_cuda_graph()
if fallback_to_flashinfer_impl:
super().init_forward_metadata(forward_batch)
@@ -789,7 +791,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
return output[:total_tokens, :, :]
def _compute_decode_bmm1_scale(self, layer: RadixAttention) -> float:
"""BMM1 scale ``q_scale * k_scale * softmax_scale``. k_scale only
"""BMM1 scale q_scale * k_scale * softmax_scale. k_scale only
applies when the KV cache stores FP8."""
q_scale = 1.0
if self.data_type == torch.float8_e4m3fn:
@@ -867,7 +869,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
):
"""Hook for subclasses to swap the ragged prefill kernel. Q/K/V arrive
in model-native dtype; subclasses do any kernel-specific quantization.
Returns the output tensor or ``(output, lse)`` if ``return_lse``."""
Returns the output tensor or (output, lse) if return_lse."""
q_scale = k_scale = v_scale = 1.0
if self.data_type == torch.float8_e4m3fn:
q, k, v, k_scale, v_scale = _quantize_fp8_qkv(q, k, v, layer)
python/sglang/srt/layers/communicator.py
+6 -1
View File
@@ -70,6 +70,11 @@ from sglang.srt.layers.utils.cp_utils import (
is_mla_prefill_cp_enabled,
mla_use_prefill_cp,
)
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.server_args import get_global_server_args
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -274,7 +279,7 @@ class AttnTpContext:
and not is_dp_attention_enabled()
and get_moe_a2a_backend().is_none()
and not enable_moe_dense_fully_dp()
and get_global_server_args().disable_piecewise_cuda_graph
and not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
and get_global_server_args().speculative_algorithm != "EAGLE3"
)
if get_global_server_args().enable_attn_tp_input_scattered:
python/sglang/srt/layers/layernorm.py
+8
View File
@@ -26,6 +26,11 @@ from sglang.srt.batch_invariant_ops import (
)
from sglang.srt.environ import envs
from sglang.srt.layers.utils import MultiPlatformOp
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import (
cpu_has_amx_support,
@@ -431,6 +436,9 @@ class RMSNorm(MultiPlatformOp):
residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE):
return self.forward_native(x, residual, post_residual_addition)
if not x.is_contiguous():
x = x.contiguous()
python/sglang/srt/layers/moe/ep_moe/layer.py
+4 -2
View File
@@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, Any, Dict, Optional
import torch
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.environ import envs
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.moe import (
@@ -26,6 +25,9 @@ from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.quantization.fp8 import Fp8Config
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config, W4AFp8MoEMethod
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.utils import get_bool_env_var, is_hip, is_npu
if TYPE_CHECKING:
@@ -150,7 +152,7 @@ class DeepEPMoE(FusedMoE):
hidden_states: torch.Tensor,
topk_output: TopKOutput,
):
if is_in_piecewise_cuda_graph():
if is_in_tc_piecewise_cuda_graph():
assert TopKOutputChecker.format_is_standard(
topk_output
), "Only standard topk output is supported for piecewise cuda graph"
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+7 -7
View File
@@ -12,10 +12,6 @@ from torch.nn.parameter import UninitializedParameter
from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs
from sglang.srt.batch_overlap.two_batch_overlap import MaybeTboDeepEPDispatcher
from sglang.srt.compilation.piecewise_context_manager import (
get_forward_context,
is_in_piecewise_cuda_graph,
)
from sglang.srt.distributed import (
get_moe_expert_parallel_rank,
get_moe_expert_parallel_world_size,
@@ -63,6 +59,10 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import (
from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod
from sglang.srt.layers.quantization.modelopt_quant import ModelOptNvFp4FusedMoEMethod
from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
get_tc_piecewise_forward_context,
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.model_loader.weight_utils import narrow_padded_param_and_loaded_weight
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import (
@@ -1071,7 +1071,7 @@ class FusedMoE(torch.nn.Module):
from sglang.srt.hardware_backend.npu.moe.fuseep import forward_fuseep
return forward_fuseep(self, hidden_states, topk_output)
if is_in_piecewise_cuda_graph():
if is_in_tc_piecewise_cuda_graph():
if TopKOutputChecker.format_is_standard(topk_output):
return moe_forward_piecewise_cuda_graph_impl(
hidden_states,
@@ -1308,7 +1308,7 @@ def moe_forward_piecewise_cuda_graph_impl(
topk_output = StandardTopKOutput(
topk_weights=topk_weights, topk_ids=topk_ids, router_logits=router_logits
)
forward_context = get_forward_context()
forward_context = get_tc_piecewise_forward_context()
moe_layer = forward_context.moe_layers[layer_id]
return moe_layer.forward_impl(hidden_states, topk_output)
@@ -1335,6 +1335,6 @@ def fused_moe_bypassed_piecewise_cuda_graph_impl(
renormalize=renormalize,
),
)
forward_context = get_forward_context()
forward_context = get_tc_piecewise_forward_context()
moe_layer = forward_context.moe_layers[layer_id]
return moe_layer.forward_impl(hidden_states, topk_output)
python/sglang/srt/layers/moe/mega_moe.py
+3 -3
View File
@@ -26,7 +26,7 @@ from sglang.srt.environ import envs
from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
from sglang.srt.layers.dp_attention import get_dp_global_num_tokens
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.runner import get_is_capture_mode
if TYPE_CHECKING:
from deep_gemm import SymmBuffer
@@ -114,7 +114,7 @@ def should_use_mega_moe(moe: "DeepseekV2MoE", hidden_states: torch.Tensor) -> bo
def forward_mega_moe(
moe: "DeepseekV2MoE",
hidden_states: torch.Tensor,
forward_batch: Optional["ForwardBatch"] = None,
forward_batch: Optional[ForwardBatch] = None,
input_ids_global: Optional[torch.Tensor] = None,
) -> torch.Tensor:
num_tokens = hidden_states.shape[0]
@@ -151,7 +151,7 @@ def forward_mega_moe(
def _run_mega_routed(
moe: "DeepseekV2MoE",
hidden_states: torch.Tensor,
forward_batch: Optional["ForwardBatch"],
forward_batch: Optional[ForwardBatch],
input_ids_global: Optional[torch.Tensor],
num_tokens: int,
) -> torch.Tensor:
python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py
+16 -13
View File
@@ -11,6 +11,7 @@ from sglang.srt.layers.moe.moe_runner.base import (
MoeRunnerConfig,
register_fused_func,
)
from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
from sglang.srt.utils.common import log_info_on_rank0, print_warning_once
if TYPE_CHECKING:
@@ -251,7 +252,9 @@ def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None:
)
server_args = get_global_server_args()
use_cuda_graph = not server_args.disable_cuda_graph
# CuteDSL wrapper preallocates CG buffers used by any captured graph
# that routes through this MoE — decode and prefill alike.
use_cuda_graph = not cuda_graph_fully_disabled()
# Size the wrapper's CUDA-graph buffers for the largest number of tokens a
# single forward can route through this layer.
@@ -301,17 +304,17 @@ class CuteDslFp4MoeQuantInfo(MoeQuantInfo):
Shared by the two CuteDSL runner entries:
* "v2" standard path (a2a=``none``/``flashinfer``): consumed by the
``@register_fused_func("none", "flashinfer_cutedsl")`` entry, which
drives ``CuteDslMoEWrapper.run``. Weights are ``[Up, Gate]``
interleaved with MMA-layout blockscales. ``wrapper`` is set;
``w*_scale`` are scalarized.
* "v2" standard path (a2a=none/flashinfer): consumed by the
@register_fused_func("none", "flashinfer_cutedsl") entry, which
drives CuteDslMoEWrapper.run. Weights are [Up, Gate]
interleaved with MMA-layout blockscales. wrapper is set;
w*_scale are scalarized.
* "v1" DeepEP low-latency path (a2a=``deepep``): consumed by the
``@register_fused_func("deepep", "flashinfer_cutedsl")`` entry,
which drives ``flashinfer_cutedsl_moe_masked``. Weights are
``[Gate, Up]`` non-interleaved with swizzled blockscales.
``wrapper`` is ``None``; ``w*_scale`` are per-expert.
* "v1" DeepEP low-latency path (a2a=deepep): consumed by the
@register_fused_func("deepep", "flashinfer_cutedsl") entry,
which drives flashinfer_cutedsl_moe_masked. Weights are
[Gate, Up] non-interleaved with swizzled blockscales.
wrapper is None; w*_scale are per-expert.
"""
# FP4 packed weights (uint8)
@@ -332,10 +335,10 @@ class CuteDslFp4MoeQuantInfo(MoeQuantInfo):
a1_scale: torch.Tensor
a2_scale: torch.Tensor
# v2 only: lazily-created CuteDslMoEWrapper (``None`` on the v1 path).
# v2 only: lazily-created CuteDslMoEWrapper (None on the v1 path).
wrapper: Optional[Any] = None
# v1 only: ``True`` when DeepEP pre-quantizes activations to NVFP4.
# v1 only: True when DeepEP pre-quantizes activations to NVFP4.
use_nvfp4_dispatch: bool = False
# v1 only: SBO down-GEMM overlap args.
python/sglang/srt/layers/quantization/fp8_utils.py
+5 -3
View File
@@ -1502,12 +1502,14 @@ def apply_fp8_linear(
# eliminating a separate kernel launch per linear layer.
# weight_scale shape does not matter here -- it is only used in the
# GEMM epilogue, not in the activation quant fusion. Only activates when
# piecewise_cuda_graph_compiler=inductor; eager PCG and decode both
# use the faster custom kernel.
# cuda_graph_config[prefill].tc_compiler=inductor; eager PCG and
# decode both use the faster custom kernel.
if (
input_scale is not None
and input_scale.numel() == 1
and get_global_server_args().piecewise_cuda_graph_compiler == "inductor"
and get_global_server_args().cuda_graph_config.prefill.tc_compiler
== "inductor"
):
qinput = (
(input_2d * input_scale.reciprocal())
python/sglang/srt/layers/quantization/marlin_utils.py
+6 -4
View File
@@ -33,7 +33,9 @@ if TYPE_CHECKING:
from sglang.srt.layers.linear import LinearBase
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.compilation.piecewise_context_manager import get_forward_context
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
get_tc_piecewise_forward_context,
)
try:
from vllm import _custom_ops as ops
@@ -499,7 +501,7 @@ def apply_gptq_marlin_linear(
dtype=input.dtype,
)
forward_context = get_forward_context()
forward_context = get_tc_piecewise_forward_context()
if forward_context is None:
output = gptq_marlin_gemm(
reshaped_x,
@@ -569,7 +571,7 @@ def apply_awq_marlin_linear(
dtype=input.dtype,
)
forward_context = get_forward_context()
forward_context = get_tc_piecewise_forward_context()
if forward_context is None:
output = gptq_marlin_gemm(
reshaped_x,
@@ -906,7 +908,7 @@ def unified_apply_gptq_marlin_gemm(
use_fp32_reduce: bool,
is_zp_float: bool,
) -> torch.Tensor:
quant_config = get_forward_context().quant_config
quant_config = get_tc_piecewise_forward_context().quant_config
quant_type = quant_config.quant_type
return gptq_marlin_gemm(
input,
python/sglang/srt/layers/radix_attention.py
+14 -9
View File
@@ -22,14 +22,14 @@ import torch
from torch import nn
from sglang.srt.compilation.compilation_config import register_split_op
from sglang.srt.compilation.piecewise_context_manager import get_forward_context
from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import (
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
eager_on_graph,
)
from sglang.srt.model_executor.breakable_cuda_graph.context import (
is_in_breakable_cuda_graph,
)
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
get_tc_piecewise_forward_context,
)
from sglang.srt.utils import is_hip
from sglang.srt.utils.custom_op import register_custom_op
@@ -124,13 +124,16 @@ class RadixAttention(nn.Module):
else:
k = k.view(-1, self.tp_k_head_num, self.v_head_dim)
if forward_batch.forward_mode.is_extend() and get_forward_context() is not None:
if (
forward_batch.forward_mode.is_extend()
and get_tc_piecewise_forward_context() is not None
):
if self.qk_head_dim != self.v_head_dim:
output = q.new_empty((q.shape[0], self.tp_q_head_num * self.v_head_dim))
else:
output = torch.empty_like(q)
if is_in_breakable_cuda_graph():
bcg_unified_attention_with_output(
breakable_unified_attention_with_output(
q, k, v, output, save_kv_cache, self.layer_id, **kwargs
)
else:
@@ -170,7 +173,7 @@ def unified_attention_with_output(
llama_4_scaling: Optional[torch.Tensor] = None,
topk_indices: Optional[torch.Tensor] = None,
) -> None:
context = get_forward_context()
context = get_tc_piecewise_forward_context()
forward_batch = context.forward_batch
attention_layers = context.attention_layers
attention_layer = attention_layers[layer_id]
@@ -250,4 +253,6 @@ def unified_attention_with_output(
return
bcg_unified_attention_with_output = eager_on_graph(True)(unified_attention_with_output)
breakable_unified_attention_with_output = eager_on_graph(True)(
unified_attention_with_output
)
python/sglang/srt/layers/radix_linear_attention.py
+10 -7
View File
@@ -21,14 +21,14 @@ import torch
from torch import nn
from sglang.srt.compilation.compilation_config import register_split_op
from sglang.srt.compilation.piecewise_context_manager import get_forward_context
from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import (
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
eager_on_graph,
)
from sglang.srt.model_executor.breakable_cuda_graph.context import (
is_in_breakable_cuda_graph,
)
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
get_tc_piecewise_forward_context,
)
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
@@ -82,7 +82,10 @@ class RadixLinearAttention(nn.Module):
a: torch.Tensor,
b: torch.Tensor,
) -> torch.Tensor:
if forward_batch.forward_mode.is_extend() and get_forward_context() is not None:
if (
forward_batch.forward_mode.is_extend()
and get_tc_piecewise_forward_context() is not None
):
# Output shape from linear attention: (1, seq_len, num_v_heads, head_v_dim)
seq_len = mixed_qkv.shape[0]
output = torch.empty(
@@ -129,7 +132,7 @@ def unified_linear_attention_with_output(
"""
Custom op wrapper for linear attention computation only.
"""
context = get_forward_context()
context = get_tc_piecewise_forward_context()
forward_batch = context.forward_batch
attention_layers = context.attention_layers
attention_layer = attention_layers[layer_id]
python/sglang/srt/lora/lora_moe_runners.py
+1 -1
View File
@@ -29,7 +29,7 @@ from typing import Callable
import torch
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.utils import is_cuda, is_hip, is_xpu, next_power_of_2
_is_cuda = is_cuda()
python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py
+2 -2
View File
@@ -66,7 +66,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora(
from sglang.srt.lora.trtllm_lora_temp.triton_ops import (
merged_experts_fused_moe_lora_add,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode
assert runner_config.activation == "silu" and runner_config.is_gated, (
"experimental_sgl_trtllm LoRA currently supports the gated SwiGLU FP8 "
@@ -327,7 +327,7 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora(
from sglang.srt.lora.trtllm_lora_temp.triton_ops import (
merged_experts_fused_moe_lora_add,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode
assert (
runner_config.activation == "silu" and runner_config.is_gated
python/sglang/srt/managers/overlap_utils.py
+6 -1
View File
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Sequence, Union
import torch
from sglang.srt.environ import envs
from sglang.srt.model_executor.cuda_graph_config import Backend
from sglang.srt.speculative.spec_utils import spec_need_hidden_states
from sglang.srt.speculative.triton_ops.gather_spec_extras import gather_spec_extras
from sglang.srt.utils import is_cuda, is_hip, is_npu
@@ -30,7 +31,11 @@ def decide_needs_cpu_seq_lens(
if server_args.enable_two_batch_overlap:
# FIXME: support TBO without seq lens cpu value
return True
if not server_args.disable_piecewise_cuda_graph:
cuda_graph_config = server_args.cuda_graph_config
if (
cuda_graph_config is not None
and cuda_graph_config.prefill.backend == Backend.TC_PIECEWISE
):
# FIXME: support PCG without seq lens cpu value
return True
# Skip unset slots (e.g. draft_extend_attn_backend on some spec configs);
python/sglang/srt/managers/scheduler_components/dp_attn.py
+2 -1
View File
@@ -14,6 +14,7 @@ from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.observability.metrics_collector import DPCooperationInfo
from sglang.srt.server_args import ServerArgs
@@ -276,7 +277,7 @@ class SchedulerDPAttnAdapter:
attn_cp_size=self.ps.attn_cp_size,
tp_group=self.tp_group,
get_idle_batch=self.get_idle_batch,
disable_cuda_graph=self.server_args.disable_cuda_graph,
disable_cuda_graph=cuda_graph_fully_disabled(),
require_mlp_tp_gather=require_mlp_tp_gather(self.server_args),
disable_overlap_schedule=self.server_args.disable_overlap_schedule,
offload_tags=self.offload_tags,
python/sglang/srt/mem_cache/memory_pool.py
+2 -2
View File
@@ -129,7 +129,7 @@ def _set_kv_buffer_impl(
row_dim,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.runner import get_is_capture_mode
if get_is_capture_mode() and alt_stream is not None:
current_stream = device_module.current_stream()
@@ -1531,7 +1531,7 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool):
layer_id_override: Optional[int] = None,
):
maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MHA-FP4)")
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.runner import get_is_capture_mode
if layer_id_override is not None:
layer_id = layer_id_override
python/sglang/srt/model_executor/breakable_cuda_graph/__init__.py
View File
python/sglang/srt/model_executor/breakable_cuda_graph_runner.py
-541
View File
@@ -1,541 +0,0 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Breakable CUDA graph (BCG) runner.
Captures the model forward as a sequence of ``torch.cuda.CUDAGraph`` segments
split at attention layers. Functionally parallel to the torch.compile-based
PCG runner but does not depend on torch.compile or FX graph splitting — graph
breaks are inserted eagerly via :func:`eager_on_graph` decorated callables
(radix attention for dense models, mamba for hybrid models).
"""
from __future__ import annotations
import bisect
import inspect
import logging
from typing import TYPE_CHECKING, Union
import torch
import tqdm
from sglang.srt.compilation.piecewise_context_manager import set_forward_context
from sglang.srt.distributed import get_tensor_model_parallel_rank
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
set_graph_pool_id,
)
from sglang.srt.distributed.parallel_state import graph_capture
from sglang.srt.layers.dp_attention import set_dp_buffer_len, set_is_extend_in_batch
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.pooler import EmbeddingPoolerOutput
from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import (
BreakableCUDAGraph,
BreakableCUDAGraphCapture,
)
from sglang.srt.model_executor.breakable_cuda_graph.context import (
enable_breakable_cuda_graph,
)
from sglang.srt.model_executor.cuda_graph_runner import (
get_global_graph_memory_pool,
set_global_graph_memory_pool,
)
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
PPProxyTensors,
)
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
PiecewiseCudaGraphRunner,
freeze_gc,
)
from sglang.srt.utils import get_available_gpu_memory, log_info_on_rank0
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.model_runner import ModelRunner
class BreakableCudaGraphRunner:
"""Breakable CUDA graph runner.
Captures the model forward as a series of ``torch.cuda.CUDAGraph`` segments
with graph breaks at attention layers. Simpler than the torch.compile-based
PCG runner: no FX tracing, no compiled-kernel fusion — just segment-level
graph capture of the eager kernel stream.
"""
def __init__(self, model_runner: ModelRunner):
self.model_runner = model_runner
self.device = model_runner.device
self.device_module = torch.get_device_module(self.device)
self.graphs = {}
self.output_buffers = {}
self.quant_config = getattr(model_runner.model, "quant_config", None)
self.is_multimodal = model_runner.is_multimodal
# Read by the shared replay_prepare (bound from PiecewiseCudaGraphRunner).
self.capture_return_pooled_hidden_states = not model_runner.is_generation
# Capture sizes
capture_tokens = model_runner.server_args.piecewise_cuda_graph_tokens
assert capture_tokens is not None
self.capture_num_tokens = sorted(capture_tokens)
self.max_num_tokens = (
max(self.capture_num_tokens) if self.capture_num_tokens else 8192
)
self.max_bs = model_runner.req_to_token_pool.size
self.capture_hidden_mode = CaptureHiddenMode.NULL
if model_runner.server_args.enable_return_hidden_states:
self.capture_hidden_mode = CaptureHiddenMode.FULL
if (
model_runner.spec_algorithm is not None
and model_runner.spec_algorithm.is_eagle()
):
if model_runner.is_draft_worker:
self.capture_hidden_mode = CaptureHiddenMode.LAST
else:
self.capture_hidden_mode = CaptureHiddenMode.FULL
log_info_on_rank0(
logger,
f"[BCG] Capture num tokens: {self.capture_num_tokens}",
)
self._init_buffers(model_runner)
self.attention_layers = model_runner.attention_layers
self.moe_layers = model_runner.moe_layers
self.moe_fusions = model_runner.moe_fusions
self.use_captured_attn_metadata = (
model_runner.attn_backend.use_captured_forward_metadata_for_breakable_cuda_graph
)
self.attn_metadata_buffers = {} if self.use_captured_attn_metadata else None
# Resolve the inner transformer-stack module (the same boundary PCG draws
# via patch_model). At replay we monkey-patch this module's forward with
# a closure that replays the captured CUDAGraph and returns the captured
# hidden_states; the outer model.forward then runs logits_processor /
# pooler eagerly with the live (multi-req) forward_batch.
language_model = getattr(
model_runner.model, "language_model", model_runner.model
)
if hasattr(language_model, "model") and hasattr(language_model.model, "layers"):
self.layer_model = language_model.model
else:
# If we can't find the inner layer_model, disable BCG.
self.layer_model = None
logger.warning(
"[BCG] Could not resolve inner layer_model on %s. BCG is "
"disabled for this model; prefill will fall back to eager.",
type(language_model).__name__,
)
return
self.use_input_embeds = self.is_multimodal
if self.use_input_embeds:
sig = inspect.signature(self.layer_model.forward)
params = list(sig.parameters)
if "input_embeds" not in params:
raise ValueError(
f"layer_model.forward must accept 'input_embeds' for "
f"multimodal BCG, got params: {params}"
)
self._input_embeds_arg_idx = params.index("input_embeds")
# Memory pool
if get_global_graph_memory_pool() is None:
set_global_graph_memory_pool(self.device_module.graph_pool_handle())
set_graph_pool_id(get_global_graph_memory_pool())
# Warmup then capture
self._warmup()
self.device_module.synchronize()
self.model_runner.tp_group.barrier()
self._capture_all()
self.raw_num_tokens = 0
def _has_inactive_dp_rank(self, forward_batch: "ForwardBatch") -> bool:
global_num_tokens = forward_batch.global_num_tokens_cpu
if global_num_tokens is None:
return False
# DSV4 DP attention / DeepEP collectives need every DP rank to enter
# the same replay path. Sparse-DP batches fall back to eager to avoid
# hanging ranks that have zero local tokens.
return len(global_num_tokens) > 1 and any(
int(num_tokens) == 0 for num_tokens in global_num_tokens
)
def _init_buffers(self, model_runner):
"""Initialize input buffers."""
from sglang.srt.model_executor.cuda_graph_buffer_registry import (
build_prefill_registry,
)
from sglang.srt.utils import is_npu
cache_loc_dtype = torch.int64 if not is_npu() else torch.int32
if model_runner.is_draft_worker:
from sglang.srt.speculative.eagle_utils import get_draft_hidden_dim
hidden_dim = get_draft_hidden_dim(model_runner)
self.static_draft_hidden_states = torch.zeros(
(self.max_num_tokens, hidden_dim),
dtype=model_runner.dtype,
device=self.device,
)
# Registry owns (allocates + pools) the token-axis input buffers.
self.buffer_registry = build_prefill_registry(
device=self.device,
max_bs=1,
max_num_token=self.max_num_tokens,
cache_loc_dtype=cache_loc_dtype,
is_multimodal=self.is_multimodal,
hidden_size=model_runner.model_config.hidden_size,
embed_dtype=model_runner.dtype,
enable_mamba_track=False,
share_pool=not is_npu(),
source=None,
)
@torch.no_grad()
def _run_forward(self, forward_batch, num_tokens):
"""Run layer-stack forward with proper context.
Captures only the inner transformer stack (layer_model). The outer
model.forward's tail (logits_processor / pooler) is intentionally
excluded — it has bs-shaped kernels that would bake batch_size=1
into the captured graph.
``@torch.no_grad`` mirrors the decorator on the outer ``*ForCausalLM.forward``
(e.g. qwen3.py:507). Calling ``layer_model.forward`` directly skips that
decorator, so we apply it here — without it some MoE @torch.compile
kernels (``torch.sum(out=...)``) fail dynamo with "out= doesn't support
autograd", and mamba state ops can spuriously track gradients.
"""
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
set_dp_buffer_len(None, num_tokens, forward_batch.dp_padding_mode.is_max_len())
set_is_extend_in_batch(False)
with set_forward_context(
forward_batch,
self.attention_layers,
self.quant_config,
self.moe_layers,
self.moe_fusions,
):
output = self.layer_model.forward(
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
input_embeds=forward_batch.input_embeds,
)
return output
def _build_capture_forward_batch(self, num_tokens):
"""Build a bs=1 placeholder ForwardBatch for capture.
bs=1 here is only a placeholder for attention/mamba breaks' metadata
shapes; replay supplies live multi-req metadata via replay_prepare.
Captured kernels run only on the token-major layer stack and are
bs-invariant.
"""
from sglang.srt.layers.dp_attention import DpPaddingMode
from sglang.srt.model_executor.forward_batch_info import (
ForwardBatch,
ForwardMode,
)
spec_info = None
if self.model_runner.is_draft_worker:
from sglang.srt.speculative.eagle_info import EagleDraftInput
spec_info = EagleDraftInput(
hidden_states=self.static_draft_hidden_states[:num_tokens],
)
registry = self.buffer_registry
bs = 1
def _slot(name):
return registry.get_slot(name).slice_for(bs, num_tokens)
with torch.device(self.device):
seq_lens = torch.full((bs,), num_tokens, dtype=torch.int64)
extend_seq_lens = torch.full((bs,), num_tokens, dtype=torch.int64)
extend_prefix_lens = torch.zeros((bs,), dtype=torch.int64)
extend_start_loc = torch.zeros((bs,), dtype=torch.int64)
req_pool_indices = torch.arange(bs, dtype=torch.int64)
orig_seq_lens = torch.full((bs,), num_tokens, dtype=torch.int64)
return ForwardBatch(
forward_mode=ForwardMode.EXTEND,
batch_size=bs,
input_ids=_slot("input_ids"),
input_embeds=(
_slot("input_embeds") if registry.has_slot("input_embeds") else None
),
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
next_token_logits_buffer=None,
orig_seq_lens=orig_seq_lens,
seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
out_cache_loc=_slot("out_cache_loc"),
seq_lens_sum=num_tokens,
mamba_track_indices=None,
mamba_track_mask=None,
mamba_track_seqlens=None,
encoder_lens=None,
return_logprob=False,
extend_num_tokens=num_tokens,
extend_seq_lens=extend_seq_lens,
extend_prefix_lens=extend_prefix_lens,
extend_start_loc=extend_start_loc,
extend_prefix_lens_cpu=torch.tensor([0], device="cpu"),
extend_seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
extend_logprob_start_lens_cpu=torch.tensor([num_tokens], device="cpu"),
positions=_slot("positions"),
global_num_tokens_gpu=None,
global_num_tokens_for_logprob_gpu=None,
dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
global_dp_buffer_len=None,
mrope_positions=(
_slot("mrope_positions")
if registry.has_slot("mrope_positions")
else None
),
spec_algorithm=None,
spec_info=spec_info,
capture_hidden_mode=self.capture_hidden_mode,
num_token_non_padded=None,
global_forward_mode=ForwardMode.EXTEND,
lora_ids=None,
)
def _warmup(self):
"""Warmup the model with a forward pass."""
num_tokens = self.capture_num_tokens[0]
forward_batch = self._build_capture_forward_batch(num_tokens)
with (
forward_context(
ForwardContext(attn_backend=self.model_runner.attn_backend)
),
set_forward_context(
forward_batch,
self.attention_layers,
self.quant_config,
self.moe_layers,
self.moe_fusions,
),
):
self._init_forward_metadata_for_capture(forward_batch, num_tokens)
self._run_forward(forward_batch, num_tokens)
def _init_forward_metadata_for_capture(self, forward_batch, num_tokens):
attn_backend = self.model_runner.attn_backend
if not self.use_captured_attn_metadata:
attn_backend.init_forward_metadata(forward_batch)
return
metadata = attn_backend.init_forward_metadata_for_breakable_cuda_graph_capture(
forward_batch
)
assert self.attn_metadata_buffers is not None
self.attn_metadata_buffers[num_tokens] = metadata
def _prepare_forward_metadata_for_replay(
self, forward_batch, static_forward_batch, num_tokens
):
attn_backend = self.model_runner.attn_backend
if not self.use_captured_attn_metadata:
attn_backend.init_forward_metadata(forward_batch)
return
assert self.attn_metadata_buffers is not None
metadata = self.attn_metadata_buffers[num_tokens]
attn_backend.prepare_forward_metadata_for_breakable_cuda_graph_replay(
metadata,
forward_batch,
static_forward_batch=static_forward_batch,
)
def _capture_all(self):
"""Capture breakable CUDA graphs for all token sizes."""
with (
freeze_gc(self.model_runner.server_args.enable_cudagraph_gc),
graph_capture() as graph_capture_context,
enable_breakable_cuda_graph(),
):
stream = graph_capture_context.stream
pool = get_global_graph_memory_pool()
capture_range = (
tqdm.tqdm(list(reversed(self.capture_num_tokens)))
if get_tensor_model_parallel_rank() == 0
else reversed(self.capture_num_tokens)
)
for num_tokens in capture_range:
if get_tensor_model_parallel_rank() == 0:
avail_mem = get_available_gpu_memory(
self.model_runner.device,
self.model_runner.gpu_id,
empty_cache=False,
)
capture_range.set_description(
f"[BCG] Capturing ({num_tokens=} {avail_mem=:.2f} GB)"
)
graph, output = self._capture_one(num_tokens, pool, stream)
self.graphs[num_tokens] = graph
self.output_buffers[num_tokens] = output
def can_run(self, forward_batch: "ForwardBatch"):
if self.layer_model is None:
return False
if forward_batch.forward_mode.is_target_verify():
return False
if forward_batch.capture_hidden_mode != self.capture_hidden_mode:
return False
if forward_batch.input_embeds is not None:
return False
if forward_batch.replace_embeds is not None:
return False
if self._has_inactive_dp_rank(forward_batch):
return False
if (
forward_batch.global_num_tokens_cpu is not None
and not forward_batch.can_run_dp_breakable_cuda_graph
):
return False
num_tokens = len(forward_batch.input_ids)
if forward_batch.return_logprob:
for start_len, seq_len in zip(
forward_batch.extend_logprob_start_lens_cpu,
forward_batch.extend_seq_lens_cpu,
):
if start_len is not None and start_len < seq_len:
return False
return num_tokens <= self.max_num_tokens
def _capture_one(self, num_tokens, pool, stream):
"""Capture a breakable CUDA graph for one token size."""
forward_batch = self._build_capture_forward_batch(num_tokens)
self._init_forward_metadata_for_capture(forward_batch, num_tokens)
def run_once():
return self._run_forward(forward_batch, num_tokens)
with forward_context(
ForwardContext(attn_backend=self.model_runner.attn_backend)
):
for _ in range(2):
self.device_module.synchronize()
self.model_runner.tp_group.barrier()
run_once()
graph = BreakableCUDAGraph()
with BreakableCUDAGraphCapture(cuda_graph=graph, pool=pool, stream=stream):
output = run_once()
return graph, output
def replay_prepare(self, forward_batch, **kwargs):
# TODO: fix PiecewiseCudaGraphRunner to support draft workers as well.
static_forward_batch = PiecewiseCudaGraphRunner.replay_prepare(
self, forward_batch, **kwargs
)
if self.model_runner.is_draft_worker and forward_batch.spec_info is not None:
num_tokens = len(forward_batch.input_ids)
self.static_draft_hidden_states[:num_tokens].copy_(
forward_batch.spec_info.hidden_states
)
return static_forward_batch
def replay(
self,
forward_batch: ForwardBatch,
**kwargs,
) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
num_tokens = len(forward_batch.input_ids)
index = bisect.bisect_left(self.capture_num_tokens, num_tokens)
static_num_tokens = self.capture_num_tokens[index]
captured_graph = self.graphs[static_num_tokens]
captured_hidden = self.output_buffers[static_num_tokens]
# Closure replaces layer_model.forward for the duration of the outer
# model.forward call. Replays the captured CUDAGraph and hands the
# outer forward the captured hidden_states; logits_processor / pooler
# then runs eagerly on top with the live multi-req forward_batch.
def replay_layer_forward(*args, **layer_kwargs):
ie = layer_kwargs.get("input_embeds") or (
args[self._input_embeds_arg_idx]
if self.use_input_embeds and len(args) > self._input_embeds_arg_idx
else None
)
if self.use_input_embeds:
if ie is None:
raise ValueError("BCG replay expects input_embeds but got None")
self.buffer_registry.get_slot("input_embeds").slice_for(
1, static_num_tokens
).copy_(ie[:static_num_tokens])
else:
if ie is not None:
raise ValueError(
"BCG replay got unexpected input_embeds on non-multimodal model"
)
captured_graph.replay()
return captured_hidden
with enable_breakable_cuda_graph():
static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
original_layer_forward = self.layer_model.forward
self.layer_model.forward = replay_layer_forward
try:
self._prepare_forward_metadata_for_replay(
forward_batch, static_forward_batch, static_num_tokens
)
with set_forward_context(
static_forward_batch,
self.attention_layers,
self.quant_config,
self.moe_layers,
self.moe_fusions,
):
output = self.model_runner.model.forward(
static_forward_batch.input_ids,
static_forward_batch.positions,
static_forward_batch,
**kwargs,
)
finally:
self.layer_model.forward = original_layer_forward
if isinstance(output, LogitsProcessorOutput):
return LogitsProcessorOutput(
next_token_logits=output.next_token_logits[: self.raw_num_tokens],
hidden_states=(
output.hidden_states[: self.raw_num_tokens]
if output.hidden_states is not None
else None
),
)
elif isinstance(output, EmbeddingPoolerOutput):
return output
else:
assert isinstance(output, PPProxyTensors)
raise NotImplementedError(
"PPProxyTensors is not supported in BreakableCudaGraphRunner."
)
python/sglang/srt/model_executor/cpu_graph_runner.py
+4 -4
View File
@@ -29,7 +29,6 @@ import tqdm
from sglang.srt.distributed import get_tensor_model_parallel_rank
from sglang.srt.distributed.parallel_state import GroupCoordinator
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.model_executor.cuda_graph_runner import model_capture_mode
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
@@ -38,6 +37,7 @@ from sglang.srt.model_executor.forward_batch_info import (
enable_num_token_non_padded,
)
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.runner_utils.capture_mode import model_capture_mode
from sglang.srt.utils import (
empty_context,
log_info_on_rank0,
@@ -126,10 +126,10 @@ def set_torch_compile_config():
def get_batch_sizes_to_capture(model_runner: ModelRunner):
# torch compile speeds up decoding by reducing python overhead on CPU
server_args = model_runner.server_args
# Note that we reuse server_args.cuda_graph_bs here.
# Reuse cuda_graph_config[decode].bs here.
# Users can customize the batch sizes supported by cpu_graph, such as:
# --cuda-graph-bs 1 2 4 8 16
capture_bs = server_args.cuda_graph_bs
# --cuda-graph-bs-decode 1 2 4 8 16
capture_bs = server_args.cuda_graph_config.decode.bs
assert (
max(capture_bs) <= server_args.torch_compile_max_bs
), f"{capture_bs=}, {server_args.torch_compile_max_bs=}"
python/sglang/srt/model_executor/cuda_graph_config.py
+211
View File
@@ -0,0 +1,211 @@
"""Phase / backend identifiers, the canonical default for
cuda_graph_config, and the --cuda-graph-config JSON CLI parser.
Module-level imports are pure stdlib — no torch / sglang.srt deps — so
ServerArgs can import everything here without pulling in backend
classes. check_cuda_graph_backend lazy-imports get_global_server_args
inside the function body to preserve that invariant.
"""
import argparse
import dataclasses
import json
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
class Phase:
"""The two phases of model forward."""
DECODE = "decode"
PREFILL = "prefill"
ALL = (DECODE, PREFILL)
class Backend:
"""CUDA graph capture backends a phase can use."""
FULL = "full"
BREAKABLE = "breakable"
TC_PIECEWISE = "tc_piecewise"
DISABLED = "disabled"
ALL = (FULL, BREAKABLE, TC_PIECEWISE, DISABLED)
ALLOWED_BACKENDS_PER_PHASE = {
Phase.DECODE: (
Backend.FULL,
Backend.BREAKABLE,
Backend.TC_PIECEWISE,
Backend.DISABLED,
),
# full is rejected for prefill — full CUDA graph capture only
# fits fixed-shape and prefill is variable-shape. Use breakable
# or tc_piecewise for prefill.
Phase.PREFILL: (Backend.BREAKABLE, Backend.TC_PIECEWISE, Backend.DISABLED),
}
# Per-phase settings schema. Keys other than backend are runner-level
# (read by any backend in that phase); tc_compiler is the lone
# backend-specific knob (only meaningful when backend == tc_piecewise).
# For prefill, bs carries the captured shape size (token count for
# tc_piecewise, request count for breakable) — one shape knob per phase.
ALLOWED_KEYS_PER_PHASE = {
Phase.DECODE: ("backend", "max_bs", "bs", "tc_compiler"),
Phase.PREFILL: ("backend", "max_bs", "bs", "tc_compiler"),
}
@dataclass
class PhaseConfig:
"""Per-phase CUDA graph settings."""
backend: str = Backend.DISABLED
max_bs: Optional[int] = None
bs: Optional[List[int]] = None
# Only meaningful when backend == tc_piecewise; ignored otherwise.
tc_compiler: str = "eager"
@dataclass
class CudaGraphConfig:
"""Top-level CUDA graph config: one PhaseConfig per phase."""
decode: PhaseConfig = field(
default_factory=lambda: PhaseConfig(backend=Backend.FULL)
)
prefill: PhaseConfig = field(
default_factory=lambda: PhaseConfig(backend=Backend.TC_PIECEWISE)
)
def __getitem__(self, phase: str) -> PhaseConfig:
"""Phase-string lookup; kept for migration ergonomics."""
if phase not in Phase.ALL:
raise KeyError(phase)
return getattr(self, phase)
def to_dict(self) -> Dict[str, Dict[str, Any]]:
# Diff-only, not asdict: the parser locks every (phase, key) it sees,
# so emitting defaults would lock fields the caller never set.
baseline = default_cuda_graph_config()
return {
Phase.DECODE: _diff_phase(self.decode, baseline.decode),
Phase.PREFILL: _diff_phase(self.prefill, baseline.prefill),
}
@classmethod
def from_dict(cls, raw: Optional[Dict[str, Dict[str, Any]]]) -> "CudaGraphConfig":
"""Build from a (partial) dict of overrides, defaults fill the rest.
Unknown phases / keys are silently dropped — the JSON-input
validator (parse_cuda_graph_config_arg) rejects them upstream."""
cfg = cls()
if not raw:
return cfg
for phase, phase_settings in raw.items():
if phase not in Phase.ALL or not isinstance(phase_settings, dict):
continue
phase_cfg = getattr(cfg, phase)
allowed = ALLOWED_KEYS_PER_PHASE[phase]
for key, value in phase_settings.items():
if key in allowed:
setattr(phase_cfg, key, value)
return cfg
def default_cuda_graph_config() -> CudaGraphConfig:
"""Fresh CudaGraphConfig populated with canonical defaults."""
return CudaGraphConfig()
def _diff_phase(actual: PhaseConfig, baseline: PhaseConfig) -> Dict[str, Any]:
"""Return only fields whose value differs from the per-phase default."""
return {
f.name: getattr(actual, f.name)
for f in dataclasses.fields(actual)
if getattr(actual, f.name) != getattr(baseline, f.name)
}
def check_cuda_graph_backend(phase: str, backend: str) -> bool:
"""True if cuda_graph_config[phase].backend == backend on the
global server args. Returns False if the global server args have not
been initialized yet (e.g. unit tests, early startup)."""
from sglang.srt.server_args import get_global_server_args
try:
server_args = get_global_server_args()
except ValueError:
return False
cfg = server_args.cuda_graph_config
if cfg is None or phase not in Phase.ALL:
return False
return getattr(cfg, phase).backend == backend
def cuda_graph_fully_disabled() -> bool:
"""True iff cuda_graph_config has Backend.DISABLED on every phase.
Use at sites that ask the legacy server_args.disable_cuda_graph
question ("no CG anywhere globally") — e.g., preallocating buffers
that any captured graph would otherwise reuse, or one-shot init
that's a no-op when CG is completely off.
"""
return check_cuda_graph_backend(
Phase.DECODE, Backend.DISABLED
) and check_cuda_graph_backend(Phase.PREFILL, Backend.DISABLED)
def parse_cuda_graph_config_arg(raw: str) -> Dict[str, Dict[str, Any]]:
"""argparse type for --cuda-graph-config: parse JSON dict of
phase → settings dict. Each phase's settings dict is itself validated
against ALLOWED_KEYS_PER_PHASE. Returns a plain dict — the
precedence pipeline in ServerArgs converts to CudaGraphConfig
after merging."""
try:
parsed = json.loads(raw)
except json.JSONDecodeError as e:
raise argparse.ArgumentTypeError(f"--cuda-graph-config must be JSON: {e}")
if not isinstance(parsed, dict):
raise argparse.ArgumentTypeError(
f"--cuda-graph-config must be a JSON object, got {type(parsed).__name__}"
)
result: Dict[str, Dict[str, Any]] = {}
for phase, phase_settings in parsed.items():
phase = str(phase)
if phase not in Phase.ALL:
raise argparse.ArgumentTypeError(
f"--cuda-graph-config: unknown phase '{phase}', expected one of {Phase.ALL}"
)
if not isinstance(phase_settings, dict):
raise argparse.ArgumentTypeError(
f"--cuda-graph-config['{phase}'] must be a JSON object, got "
f"{type(phase_settings).__name__}"
)
allowed = ALLOWED_KEYS_PER_PHASE[phase]
result[phase] = {}
for key, value in phase_settings.items():
if key not in allowed:
raise argparse.ArgumentTypeError(
f"--cuda-graph-config['{phase}']: unknown key '{key}', expected one of {allowed}"
)
result[phase][key] = value
return result
def explicit_keys_in(
settings: Optional[Dict[str, Dict[str, Any]]],
) -> set:
"""Return the set of (phase, key) tuples present in settings
(the raw dict form, as it arrives from CLI/SDK). Used by ServerArgs
to track keys the user explicitly set so the auto-disable cascade can
skip them."""
out: set = set()
if not settings:
return out
for phase, phase_settings in settings.items():
if not isinstance(phase_settings, dict):
continue
for key in phase_settings.keys():
out.add((phase, key))
return out
python/sglang/srt/model_executor/forward_context.py
+13 -13
View File
@@ -1,23 +1,23 @@
"""Per-forward-call control context.
Owns ``ForwardContext`` — a frozen dataclass holding control configs the model
layer reads at depth via ``get_forward_context()``. The only mandatory field
today is ``attn_backend``; pool refs are derived from ``attn_backend.*``
(every backend caches them at ``__init__``), so a published ``ForwardContext``
Owns ForwardContext — a frozen dataclass holding control configs the model
layer reads at depth via get_forward_context(). The only mandatory field
today is attn_backend; pool refs are derived from attn_backend.*
(every backend caches them at __init__), so a published ForwardContext
is enough to resolve the active pools without a separate global.
``ModelRunner._forward_raw`` publishes a fresh ``ForwardContext`` for the
ModelRunner._forward_raw publishes a fresh ForwardContext for the
duration of each forward; callers that need a per-call override (PDmux
per-stream backend, frozen-KV MTP draft loop, TBO per-child dispatch) use
``dataclasses.replace`` and wrap the override scope with ``forward_context()``.
dataclasses.replace and wrap the override scope with forward_context().
Distinct from ``sglang.srt.compilation.piecewise_context_manager.ForwardContext``,
Distinct from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph.TcPiecewiseForwardContext,
which collects compilation-time refs for the piecewise CUDA graph backend.
Concurrency: ``_current`` is a plain module-level global, not thread-local.
This matches the ``global_server_args`` precedent and is safe because each
Concurrency: _current is a plain module-level global, not thread-local.
This matches the global_server_args precedent and is safe because each
forward runs synchronously on a single Python thread per worker process. If
worker threads ever share a process, migrate to ``contextvars.ContextVar``.
worker threads ever share a process, migrate to contextvars.ContextVar.
"""
from __future__ import annotations
@@ -33,9 +33,9 @@ if TYPE_CHECKING:
@dataclass(frozen=True, slots=True)
class ForwardContext:
"""Per-forward-call control configs. Read via ``get_forward_context()``;
"""Per-forward-call control configs. Read via get_forward_context();
extend by adding fields here. Frozen so accidental mutation raises at
write time — use ``dataclasses.replace`` for per-call overrides."""
write time — use dataclasses.replace for per-call overrides."""
attn_backend: AttentionBackend
@@ -45,7 +45,7 @@ _current: Optional[ForwardContext] = None
def set_forward_context(ctx: Optional[ForwardContext]) -> Optional[ForwardContext]:
"""Set the active context; return the previous one for explicit
save/restore. Prefer the ``forward_context()`` context manager."""
save/restore. Prefer the forward_context() context manager."""
global _current
prev, _current = _current, ctx
return prev
python/sglang/srt/model_executor/model_runner.py
+66 -94
View File
@@ -35,10 +35,7 @@ import torch.distributed as dist
from torch import nn
from sglang.jit_kernel.ngram_embedding import update_token_table_decode
from sglang.srt.compilation.piecewise_context_manager import (
enable_piecewise_cuda_graph,
set_forward_context,
)
from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config
from sglang.srt.configs import (
BailingHybridConfig,
FalconH1Config,
@@ -140,19 +137,17 @@ from sglang.srt.lora.lora_registry import LoRARef
from sglang.srt.managers.schedule_batch import sanity_check_mm_pad_shift_value
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.model_executor.breakable_cuda_graph_runner import (
BreakableCudaGraphRunner,
)
from sglang.srt.model_executor.cpu_graph_runner import CPUGraphRunner
from sglang.srt.model_executor.cuda_graph_buffer_registry import (
CudaGraphBufferRegistry,
build_decode_registry,
build_prefill_registry,
)
from sglang.srt.model_executor.cuda_graph_runner import (
CudaGraphRunner,
_allocate_decode_buffers,
set_torch_compile_config,
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
cuda_graph_fully_disabled,
)
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
@@ -169,10 +164,13 @@ from sglang.srt.model_executor.hook_manager import register_forward_hooks
from sglang.srt.model_executor.model_runner_kv_cache_mixin import (
ModelRunnerKVCacheMixin,
)
from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
PiecewiseCudaGraphRunner,
)
from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig
from sglang.srt.model_executor.runner import (
PrefillCudaGraphRunner,
)
from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
_allocate_decode_buffers,
)
from sglang.srt.model_loader.loader import DefaultModelLoader, get_model_loader
from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
RemoteInstanceWeightLoaderBackend,
@@ -763,12 +761,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
# Init lora
if server_args.enable_lora:
self.init_lora_manager()
if not server_args.disable_cuda_graph:
if not cuda_graph_fully_disabled():
# Phase 1 of LoRA CUDA graph init: pre-allocate large MoE
# intermediate buffers before init_memory_pool() so memory
# profiling accounts for them. Phase 2 (dense LoRA batch
# metadata) is handled in CudaGraphRunner.__init__() via
# lora_manager.init_cuda_graph_batch_info().
# profiling accounts for them. The buffers are reused by
# any captured graph (decode today; widen here so any
# future prefill capture path also picks them up).
self._init_lora_cuda_graph_moe_buffers()
# Enable batch invariant mode
@@ -801,7 +799,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.init_indexer_capturer()
# TODO: Refactor device-specific init branches into platform interface (separate PR).
# Must be called BEFORE init_device_graphs() so CUDA graph capture
# Must be called BEFORE init_decode_cuda_graph() so CUDA graph capture
# runs with aux hidden state capture enabled.
self.init_aux_hidden_state_capture()
@@ -831,13 +829,13 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.init_attention_backend()
self.kernel_warmup()
self._pre_initialize_flashinfer_allreduce_workspace()
self.init_device_graphs()
self.init_decode_cuda_graph()
elif self.device == "cpu":
self.init_attention_backend()
self.init_device_graphs()
self.init_decode_cuda_graph()
elif self.device == "npu":
self.init_attention_backend()
# lazy init for zbal with mix mode(before graph capture when enable_cuda_graph)
# lazy init for zbal with mix mode (before graph capture when enable_cuda_graph)
if envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() > 0 and not self.is_draft_worker:
from sglang.srt.hardware_backend.npu.utils import lazy_init_zbal_gva_mem
@@ -848,16 +846,16 @@ class ModelRunner(ModelRunnerKVCacheMixin):
get_world_group().world_size,
get_world_group().cpu_group,
)
self.init_device_graphs()
self.init_decode_cuda_graph()
elif current_platform.is_out_of_tree():
self.init_attention_backend()
if current_platform.support_cuda_graph():
self.init_device_graphs()
self.init_decode_cuda_graph()
else:
self.graph_runner = None
self.decode_cuda_graph_runner = None
self.graph_mem_usage = 0
else:
self.graph_runner = None
self.decode_cuda_graph_runner = None
self.graph_mem_usage = 0
self.init_attention_backend()
@@ -865,7 +863,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
register_forward_hooks(self.model, server_args.forward_hooks)
# Initialize piecewise CUDA graph
self.init_piecewise_cuda_graphs()
self.init_prefill_cuda_graph()
self.prealloc_symmetric_memory_pool()
@@ -1720,7 +1718,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
and current_platform.support_cuda_graph()
)
):
self.init_device_graphs()
self.init_decode_cuda_graph()
logger.info("Update weights end.")
return True, "Succeeded to update model weights."
@@ -2082,7 +2080,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
"""
from sglang.srt.lora.layers import FusedMoEWithLoRA
max_bs = self.server_args.cuda_graph_max_bs
max_bs = self.server_args.cuda_graph_config.decode.max_bs
max_loras = self.server_args.max_loras_per_batch
for module in self.model.modules():
if isinstance(module, FusedMoEWithLoRA):
@@ -2546,8 +2544,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
):
"""Run a dummy forward pass for warmup/profiling.
``forward_mode_override`` forces EXTEND/DECODE regardless of
``is_generation`` (used by the PP-parallel DeepGEMM warmup).
forward_mode_override forces EXTEND/DECODE regardless of
is_generation (used by the PP-parallel DeepGEMM warmup).
"""
if forward_mode_override is not None:
capture_forward_mode = forward_mode_override
@@ -2845,7 +2843,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
def maybe_update_ngram_token_table(
self,
next_token_ids: torch.Tensor,
forward_batch: "ForwardBatch",
forward_batch: ForwardBatch,
):
"""Update the ngram embedding token table after sampling."""
ngram_embedding_info = forward_batch.ngram_embedding_info
@@ -2862,9 +2860,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
column_starts=ngram_embedding_info.out_column_starts,
)
def init_device_graphs(self):
def init_decode_cuda_graph(self):
"""Capture device graphs."""
self.graph_runner = None
self.decode_cuda_graph_runner = None
self.graph_mem_usage = 0
if not self.is_generation:
@@ -2874,7 +2872,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
if self.server_args.model_impl.lower() == ModelImpl.MINDSPORE:
return
if self.device != "cpu" and self.server_args.disable_cuda_graph:
if self.device != "cpu" and check_cuda_graph_backend(
Phase.DECODE, Backend.DISABLED
):
return
if self.device == "cpu" and not self.server_args.enable_torch_compile:
@@ -2896,16 +2896,20 @@ class ModelRunner(ModelRunnerKVCacheMixin):
)
if current_platform.is_out_of_tree():
GraphRunnerCls = current_platform.get_graph_runner_cls()
self.graph_runner = GraphRunnerCls(self)
self.decode_cuda_graph_runner = GraphRunnerCls(self)
else:
from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
DecodeCudaGraphRunner,
)
graph_runners = defaultdict(
lambda: CudaGraphRunner,
lambda: DecodeCudaGraphRunner,
{
"cpu": CPUGraphRunner,
"npu": NPUGraphRunner,
},
)
self.graph_runner = graph_runners[self.device](self)
self.decode_cuda_graph_runner = graph_runners[self.device](self)
after_mem = get_available_gpu_memory(self.device, self.gpu_id)
self.graph_mem_usage = before_mem - after_mem
@@ -2914,13 +2918,15 @@ class ModelRunner(ModelRunnerKVCacheMixin):
f"mem usage={self.graph_mem_usage:.2f} GB. avail mem={after_mem:.2f} GB."
)
def init_piecewise_cuda_graphs(self, force_for_draft_worker: bool = False):
def init_prefill_cuda_graph(self, force_for_draft_worker: bool = False):
"""Initialize piecewise CUDA graph runner."""
self.piecewise_cuda_graph_runner = None
self.prefill_cuda_graph_runner = None
if self.server_args.disable_piecewise_cuda_graph:
if check_cuda_graph_backend(Phase.PREFILL, Backend.DISABLED):
logger.info(
"Disable piecewise CUDA graph because --disable-piecewise-cuda-graph is set"
"Disable prefill CUDA graph because cuda_graph_config "
"resolved prefill.backend='disabled' (e.g. via "
"--cuda-graph-backend-prefill=disabled or auto-disable rules)."
)
return
@@ -2938,7 +2944,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
return
# Disable piecewise CUDA graph for non capture size
if not self.server_args.piecewise_cuda_graph_tokens:
if not self.server_args.cuda_graph_config.prefill.bs:
logger.warning(
"Disable piecewise CUDA graph because the capture size is not set"
)
@@ -2971,8 +2977,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
elif hasattr(layer.self_attn, "attn_mqa"):
# For DeepSeek model
attn_layer = layer.self_attn.attn_mqa
if _is_hip and hasattr(layer.self_attn, "attn_mha"):
attn_layer._pcg_mha_companion = layer.self_attn.attn_mha
# For hybrid model
elif hasattr(layer, "attn"):
attn_layer = layer.attn
@@ -3037,11 +3041,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
f"Capture piecewise CUDA graph begin. avail mem={before_mem:.2f} GB"
)
if self.server_args.enable_breakable_cuda_graph:
# Experimental feature
self.piecewise_cuda_graph_runner = BreakableCudaGraphRunner(self)
else:
self.piecewise_cuda_graph_runner = PiecewiseCudaGraphRunner(self)
self.prefill_cuda_graph_runner = PrefillCudaGraphRunner(self)
after_mem = get_available_gpu_memory(self.device, self.gpu_id)
mem_usage = before_mem - after_mem
@@ -3278,12 +3278,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
# Check piecewies cuda graph
can_run_graph = (
self.piecewise_cuda_graph_runner is not None
and self.piecewise_cuda_graph_runner.can_run(forward_batch)
self.prefill_cuda_graph_runner is not None
and self.prefill_cuda_graph_runner.can_run(forward_batch)
)
if can_run_graph:
# TODO: device_timer.wrap is too broad here — it also includes
# replay_prepare time. Move timing into the piecewise cuda graph
# replay_prepare time. Move timing into the prefill cuda graph
# runner to capture only the model.forward part.
ctx = (
self.device_timer.wrap(metadata={"category": "extend"})
@@ -3291,7 +3291,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
else contextlib.nullcontext()
)
with ctx:
ret = self.piecewise_cuda_graph_runner.replay(forward_batch, **kwargs)
ret = self.prefill_cuda_graph_runner.replay(forward_batch, **kwargs)
return (ret, can_run_graph)
if not self.server_args.enable_pdmux and self.device == "cuda":
@@ -3311,37 +3311,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
else contextlib.nullcontext()
)
with ctx:
if _is_hip and self.piecewise_cuda_graph_runner is not None:
# AMD/HIP: when PCG is enabled but the batch exceeds max captured
# size, run eagerly under enable_piecewise_cuda_graph() and
# set_forward_context() so that (a) Dynamo guards on
# _in_piecewise_cuda_graph stay consistent with the PCG-traced
# graph (preventing runtime recompilation) and (b) PCG-specific
# code paths (MoE, attention) can access their layer objects.
with (
enable_piecewise_cuda_graph(),
set_forward_context(
forward_batch,
self.attention_layers,
getattr(self.model, "quant_config", None),
self.moe_layers,
self.moe_fusions,
dsa_indexers=self.dsa_indexers,
),
):
ret = self.model.forward(
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
**kwargs,
)
else:
ret = self.model.forward(
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
**kwargs,
)
ret = self.model.forward(
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
**kwargs,
)
return (ret, can_run_graph)
def forward_idle(
@@ -3472,7 +3447,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
output.routed_experts_output = experts_capturer.on_forward_end(
forward_batch=forward_batch,
can_run_graph=output.can_run_graph,
cuda_graph_batch=getattr(self.graph_runner, "bs", None),
cuda_graph_batch=getattr(self.decode_cuda_graph_runner, "bs", None),
no_copy_to_cpu=no_copy_to_cpu,
)
@@ -3480,7 +3455,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
output.indexer_topk_output = indexer_capturer.on_forward_end(
forward_batch=forward_batch,
can_run_graph=output.can_run_graph,
cuda_graph_batch=getattr(self.graph_runner, "bs", None),
cuda_graph_batch=getattr(self.decode_cuda_graph_runner, "bs", None),
no_copy_to_cpu=no_copy_to_cpu,
)
@@ -3506,9 +3481,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
reinit_attn_backend: bool = False,
split_forward_count: int = 1,
) -> ModelRunnerOutput:
# Honor an outer-published context (spec workers wrap each per-step
# draft forward with the i-th child backend); otherwise publish this
# runner's own attn_backend for the forward.
if has_forward_context():
ctx_mgr = contextlib.nullcontext()
else:
@@ -3521,21 +3493,21 @@ class ModelRunner(ModelRunnerKVCacheMixin):
)
can_run_graph = bool(
mode_check()
and self.graph_runner
and self.graph_runner.can_run(forward_batch)
and self.decode_cuda_graph_runner
and self.decode_cuda_graph_runner.can_run(forward_batch)
)
# Hisparse coordinator — backends now read it from self.model_runner.
if (
forward_batch.forward_mode.is_decode()
and self.hisparse_coordinator is not None
):
forward_batch.hisparse_coordinator = self.hisparse_coordinator
self.hisparse_coordinator.wait_for_pending_backup()
self.hisparse_coordinator.num_real_reqs.fill_(forward_batch.batch_size)
# Replay cuda graph if applicable
if can_run_graph:
ret = self.graph_runner.replay(
ret = self.decode_cuda_graph_runner.replay(
forward_batch,
pp_proxy_tensors=pp_proxy_tensors,
)
python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py
-860
View File
@@ -1,860 +0,0 @@
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Run the model with cuda graph and torch.compile."""
from __future__ import annotations
import bisect
import gc
import logging
import warnings
from contextlib import contextmanager
from typing import TYPE_CHECKING, Union
import torch
import tqdm
from sglang.srt.batch_overlap.two_batch_overlap import TboCudaGraphRunnerPlugin
from sglang.srt.compilation.compilation_config import CompilationConfig
from sglang.srt.compilation.compile import install_torch_compiled
from sglang.srt.compilation.piecewise_context_manager import (
enable_piecewise_cuda_graph,
enable_piecewise_cuda_graph_compile,
set_forward_context,
set_pcg_capture_stream,
)
from sglang.srt.distributed import get_tensor_model_parallel_rank
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
set_graph_pool_id,
)
from sglang.srt.distributed.parallel_state import graph_capture
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
get_attention_cp_size,
get_attention_tp_rank,
get_attention_tp_size,
set_dp_buffer_len,
set_is_extend_in_batch,
)
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
from sglang.srt.layers.pooler import EmbeddingPoolerOutput
from sglang.srt.layers.utils import MultiPlatformOp
from sglang.srt.model_executor.cuda_graph_buffer_registry import build_prefill_registry
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
ForwardMode,
PPProxyTensors,
)
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.utils import (
get_available_gpu_memory,
get_bool_env_var,
is_hip,
is_musa,
is_npu,
log_info_on_rank0,
require_gathered_buffer,
)
_is_hip = is_hip()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
# Suppress Dynamo warning about tracing through lru_cache-wrapped functions (e.g., is_arch_support_pdl).
warnings.filterwarnings("ignore", message=".*lru_cache.*", module="torch._dynamo")
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
_is_musa = is_musa()
@contextmanager
def freeze_gc(enable_cudagraph_gc: bool):
"""
Optimize garbage collection during CUDA graph capture.
Clean up, then freeze all remaining objects from being included
in future collections if GC is disabled during capture.
"""
gc.collect()
should_freeze = not enable_cudagraph_gc
if should_freeze:
gc.freeze()
try:
yield
finally:
if should_freeze:
gc.unfreeze()
def _to_torch(model: torch.nn.Module, reverse: bool, num_tokens: int):
for sub in model._modules.values():
if isinstance(sub, MultiPlatformOp):
if reverse:
sub.leave_torch_compile()
else:
sub.enter_torch_compile(num_tokens=num_tokens)
if isinstance(sub, torch.nn.Module):
_to_torch(sub, reverse, num_tokens)
@contextmanager
def patch_model(model: torch.nn.Module, compiler: str):
try:
if compiler != "eager":
_to_torch(model, reverse=False, num_tokens=16)
yield model
finally:
_to_torch(model, reverse=True, num_tokens=16)
# Reuse this memory pool across all cuda graph runners.
global_graph_memory_pool = None
def get_global_graph_memory_pool():
return global_graph_memory_pool
def set_global_graph_memory_pool(val):
global global_graph_memory_pool
global_graph_memory_pool = val
def set_torch_compile_config():
import torch._dynamo.config
# Resolve torch._dynamo.exc.FailOnRecompileLimitHit
torch._dynamo.config.accumulated_cache_size_limit = 1024
if hasattr(torch._dynamo.config, "cache_size_limit"):
torch._dynamo.config.cache_size_limit = 1024
if _is_musa:
from sglang.srt.hardware_backend.musa.utils.patch_torch import (
patch_fx_custom_device,
)
patch_fx_custom_device()
class PiecewiseCudaGraphRunner:
"""A PiecewiseCudaGraphRunner runs the forward pass of a model with cuda graph and torch.compile."""
def is_mamba_track_enabled(self):
return (
self.model_runner.server_args.enable_mamba_extra_buffer()
and not self.model_runner.server_args.disable_radix_cache
and self.model_runner.spec_algorithm.is_none()
)
def __init__(self, model_runner: ModelRunner):
# Parse args
self.model_runner = model_runner
self.device = model_runner.device
self.device_module = torch.get_device_module(self.device)
self.graphs = {}
self.output_buffers = {}
self.tp_size = model_runner.server_args.tp_size
self.dp_size = model_runner.server_args.dp_size
self.pp_size = model_runner.server_args.pp_size
self.attn_tp_size = get_attention_tp_size()
self.attn_tp_rank = get_attention_tp_rank()
set_torch_compile_config()
assert (
self.model_runner.server_args.piecewise_cuda_graph_tokens is not None
), "piecewise_cuda_graph_tokens is not set"
assert self.model_runner.server_args.piecewise_cuda_graph_compiler in [
"eager",
"inductor",
], "By now, only eager and inductor are supported for piecewise cuda graph compiler."
self.compile_config = CompilationConfig(
self.model_runner.server_args.piecewise_cuda_graph_tokens,
self.model_runner.server_args.piecewise_cuda_graph_compiler,
self.model_runner.server_args.enable_torch_compile_debug_mode,
)
if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
self.compile_config.add_split_op(
"sglang.moe_forward_piecewise_cuda_graph_impl"
)
self.quant_config = getattr(self.model_runner.model, "quant_config", None)
# Batch sizes to capture
self.capture_num_tokens = self.compile_config.get_capture_sizes()
# When the layer communicator scatters/gathers across the attention TP
# group (e.g. with --moe-dense-tp-size 1), the model's reduce_scatter
# requires the token count to be divisible by attn_tp_size * attn_cp_size.
# Drop captures that would violate this (mirrors the filter used by
# the regular CUDA graph runner in get_batch_sizes_to_capture).
if require_gathered_buffer(self.model_runner.server_args):
mul_base = self.attn_tp_size
attn_cp_size = get_attention_cp_size()
if mul_base % attn_cp_size != 0:
mul_base *= attn_cp_size
filtered = [n for n in self.capture_num_tokens if n % mul_base == 0]
assert (
len(filtered) > 0
), f"No piecewise CUDA graph capture sizes are multiples of {mul_base}"
self.capture_num_tokens = filtered
log_info_on_rank0(
logger, f"Capture cuda graph num tokens {self.capture_num_tokens}"
)
self.capture_forward_mode = ForwardMode.EXTEND
self.capture_hidden_mode = CaptureHiddenMode.NULL
# If returning hidden states is enabled, or if speculative prefill needs
# aux hidden states (DFLASH), capture the FULL variant up front.
if (
model_runner.server_args.enable_return_hidden_states
or model_runner.spec_algorithm.is_dflash()
):
self.capture_hidden_mode = CaptureHiddenMode.FULL
self.max_num_tokens = (
max(self.capture_num_tokens) if self.capture_num_tokens else 8192
)
self.max_bs = model_runner.req_to_token_pool.size
self.is_multimodal = model_runner.is_multimodal
self.mamba_track_enabled = self.is_mamba_track_enabled()
# Classification/reward forwards branch on return_pooled_hidden_states; piecewise
# CUDA graph capture must use the same flag value as replay for those models.
self.capture_return_pooled_hidden_states = not model_runner.is_generation
with torch.device(self.device):
self.tbo_plugin = TboCudaGraphRunnerPlugin()
# Registry owns (allocates + pools) the token-axis input buffers.
self.buffer_registry = build_prefill_registry(
device=self.device,
max_bs=self.max_bs,
max_num_token=self.max_num_tokens,
cache_loc_dtype=self._cache_loc_dtype(),
is_multimodal=self.is_multimodal,
hidden_size=self.model_runner.model_config.hidden_size,
embed_dtype=self.model_runner.dtype,
enable_mamba_track=self.mamba_track_enabled,
share_pool=not is_npu(),
source=None,
)
self.attention_layers = self.model_runner.attention_layers
self.moe_layers = self.model_runner.moe_layers
self.moe_fusions = self.model_runner.moe_fusions
self.dsa_indexers = getattr(self.model_runner, "dsa_indexers", None)
if get_global_graph_memory_pool() is None:
set_global_graph_memory_pool(self.device_module.graph_pool_handle())
# Set graph pool id globally to be able to use symmetric memory
set_graph_pool_id(get_global_graph_memory_pool())
with enable_piecewise_cuda_graph():
language_model = getattr(
self.model_runner.model, "language_model", self.model_runner.model
)
layer_model = (
language_model.model
if hasattr(language_model, "model")
and hasattr(language_model.model, "layers")
else language_model
)
with patch_model(
layer_model, self.compile_config.compiler
) as patched_model:
# Dummy warmup for jit kernel
self.warmup_compile(num_tokens=self.capture_num_tokens[0])
install_torch_compiled(
patched_model,
fullgraph=True,
dynamic_arg_dims=None,
compile_config=self.compile_config,
graph_pool=get_global_graph_memory_pool(),
)
if _is_hip:
# AMD: single Dynamo trace is sufficient; the capture
# phase does per-shape JIT kernel warmup before each
# CUDA graph recording. The N-iteration loop is
# redundant and extremely slow on ROCm (~30 min).
with enable_piecewise_cuda_graph_compile():
self.warmup_compile(num_tokens=self.capture_num_tokens[-1])
else:
with enable_piecewise_cuda_graph_compile():
compile_range = (
tqdm.tqdm(list(reversed(self.capture_num_tokens)))
if get_tensor_model_parallel_rank() == 0
else reversed(self.capture_num_tokens)
)
for _, num_tokens in enumerate(compile_range):
if get_tensor_model_parallel_rank() == 0:
compile_range.set_description(
f"Compiling num tokens ({num_tokens=})"
)
self.warmup_compile(num_tokens=num_tokens)
set_global_graph_memory_pool(self.device_module.graph_pool_handle())
set_graph_pool_id(get_global_graph_memory_pool())
if _use_aiter:
self._pre_warm_aiter_chip_info()
self.device_module.synchronize()
self.model_runner.tp_group.barrier()
# Capture
self.capture()
self.raw_num_tokens = 0
_aiter_chip_info_cached = False
@classmethod
def _pre_warm_aiter_chip_info(cls):
"""Pre-populate aiter chip info env vars before CUDA graph capture.
aiter's get_cu_num_custom_op and get_gfx_custom_op call
subprocess.run(rocminfo) to query GPU info. During CUDA graph capture
the GPU context is locked, so rocminfo hangs indefinitely. Pre-calling
them here caches the results as environment variables so the subprocess
is never invoked during capture. Only runs once per process.
"""
if cls._aiter_chip_info_cached:
return
cls._aiter_chip_info_cached = True
import os
try:
from aiter.jit.utils.chip_info import get_cu_num, get_gfx
if not os.environ.get("CU_NUM"):
cu_num = get_cu_num()
os.environ["CU_NUM"] = str(cu_num)
logger.info(f"Pre-warmed aiter CU_NUM={cu_num}")
if not os.environ.get("GPU_ARCHS"):
gfx = get_gfx()
os.environ["GPU_ARCHS"] = gfx
logger.info(f"Pre-warmed aiter GPU_ARCHS={gfx}")
except ImportError:
pass
except Exception as e:
logger.warning(f"Failed to pre-warm aiter chip info: {e}")
def warmup_compile(self, num_tokens: int):
"""Warmup the model with a simple forward pass before CUDA graph capture."""
registry = self.buffer_registry
bs = 1
def _slot(name):
return registry.get_slot(name).slice_for(bs, num_tokens)
input_ids = _slot("input_ids")
positions = _slot("positions")
out_cache_loc = _slot("out_cache_loc")
input_embeds = (
_slot("input_embeds") if registry.has_slot("input_embeds") else None
)
mrope_positions = (
_slot("mrope_positions") if registry.has_slot("mrope_positions") else None
)
mamba_track_indices = (
_slot("mamba_track_indices")
if registry.has_slot("mamba_track_indices")
else None
)
mamba_track_mask = (
_slot("mamba_track_mask") if registry.has_slot("mamba_track_mask") else None
)
mamba_track_seqlens = (
_slot("mamba_track_seqlens")
if registry.has_slot("mamba_track_seqlens")
else None
)
with torch.device(self.device):
forward_batch = ForwardBatch(
forward_mode=ForwardMode.EXTEND,
batch_size=1,
input_ids=input_ids,
input_embeds=input_embeds,
req_pool_indices=torch.arange(1, device=self.device),
seq_lens=torch.tensor([num_tokens], device=self.device),
next_token_logits_buffer=None,
orig_seq_lens=torch.tensor([num_tokens], device=self.device),
seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
out_cache_loc=out_cache_loc,
seq_lens_sum=num_tokens,
mamba_track_indices=mamba_track_indices,
mamba_track_mask=mamba_track_mask,
mamba_track_seqlens=mamba_track_seqlens,
encoder_lens=None,
return_logprob=False,
extend_num_tokens=num_tokens,
extend_seq_lens=torch.tensor([num_tokens], device=self.device),
extend_prefix_lens=torch.tensor([0], device=self.device),
extend_start_loc=torch.tensor([0], device=self.device),
extend_prefix_lens_cpu=torch.tensor([0], device="cpu"),
extend_seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
extend_logprob_start_lens_cpu=torch.tensor([num_tokens], device="cpu"),
positions=positions,
global_num_tokens_gpu=None,
global_num_tokens_for_logprob_gpu=None,
dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
global_dp_buffer_len=None,
mrope_positions=mrope_positions,
spec_algorithm=None,
spec_info=None,
capture_hidden_mode=self.capture_hidden_mode,
num_token_non_padded=None,
num_token_non_padded_cpu=num_tokens,
global_forward_mode=ForwardMode.EXTEND,
lora_ids=None,
return_pooled_hidden_states=self.capture_return_pooled_hidden_states,
)
self.model_runner.attn_backend.init_forward_metadata(forward_batch)
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
set_dp_buffer_len(None, num_tokens, forward_batch.dp_padding_mode.is_max_len())
set_is_extend_in_batch(False)
with forward_context(
ForwardContext(attn_backend=self.model_runner.attn_backend)
):
with set_forward_context(
forward_batch,
self.attention_layers,
self.quant_config,
self.moe_layers,
self.moe_fusions,
dsa_indexers=self.dsa_indexers,
):
_ = self.model_runner.model.forward(
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
)
def _cache_loc_dtype(self):
return torch.int64 if not is_npu() else torch.int32
def can_run(self, forward_batch: ForwardBatch):
# Disable piecewise cuda graph for input embeddings
# TODO(yuwei): fix it
if forward_batch.input_embeds is not None:
return False
# PCG graphs are captured with ForwardMode.EXTEND and spec_info=None.
# TARGET_VERIFY has different spec_info and capture_hidden_mode,
# so it must not use PCG-captured graphs.
if forward_batch.forward_mode.is_target_verify():
return False
# PCG graphs are captured with the runner's capture_hidden_mode.
# If the batch needs a different mode (e.g. FULL for speculative
# decoding), PCG replay would return wrong/missing hidden_states.
if forward_batch.capture_hidden_mode != self.capture_hidden_mode:
return False
# Disable for token embedding overrides (dynamic per-request)
if forward_batch.replace_embeds is not None:
return False
num_tokens = len(forward_batch.input_ids)
if forward_batch.return_logprob:
for start_len, seq_len in zip(
forward_batch.extend_logprob_start_lens_cpu,
forward_batch.extend_seq_lens_cpu,
):
if start_len is not None and start_len < seq_len:
return False
if num_tokens <= self.max_num_tokens:
return True
return False
def capture(self) -> None:
# Trigger CUDA graph capture for specific shapes.
# Capture the large shapes first so that the smaller shapes
# can reuse the memory pool allocated for the large shapes.
with (
freeze_gc(self.model_runner.server_args.enable_cudagraph_gc),
graph_capture() as graph_capture_context,
):
stream = graph_capture_context.stream
with set_pcg_capture_stream(stream):
avail_mem = get_available_gpu_memory(
self.model_runner.device,
self.model_runner.gpu_id,
empty_cache=False,
)
# Reverse the order to enable better memory sharing across cuda graphs.
capture_range = (
tqdm.tqdm(list(reversed(self.capture_num_tokens)))
if get_tensor_model_parallel_rank() == 0
else reversed(self.capture_num_tokens)
)
for i, num_tokens in enumerate(capture_range):
if get_tensor_model_parallel_rank() == 0:
avail_mem = get_available_gpu_memory(
self.model_runner.device,
self.model_runner.gpu_id,
empty_cache=False,
)
capture_range.set_description(
f"Capturing num tokens ({num_tokens=} {avail_mem=:.2f} GB)"
)
self.capture_one_batch_size(num_tokens)
def capture_one_batch_size(self, num_tokens: int):
registry = self.buffer_registry
bs = 1
# Graph inputs — views into the registry's (adopted) graph-resident
# slots; capture burns these addresses into the graph.
def _slot(name):
return registry.get_slot(name).slice_for(bs, num_tokens)
input_ids = _slot("input_ids")
positions = _slot("positions")
out_cache_loc = _slot("out_cache_loc")
input_embeds = (
_slot("input_embeds") if registry.has_slot("input_embeds") else None
)
mrope_positions = (
_slot("mrope_positions") if registry.has_slot("mrope_positions") else None
)
mamba_track_indices = (
_slot("mamba_track_indices")
if registry.has_slot("mamba_track_indices")
else None
)
mamba_track_mask = (
_slot("mamba_track_mask") if registry.has_slot("mamba_track_mask") else None
)
mamba_track_seqlens = (
_slot("mamba_track_seqlens")
if registry.has_slot("mamba_track_seqlens")
else None
)
global_dp_buffer_len = None
global_num_tokens_cpu = None
if self.model_runner.server_args.enable_lora:
# It is safe to capture CUDA graph using empty LoRA id, as the LoRA kernels will always be launched whenever
# `--enable-lora` is set to True (and return immediately if the LoRA id is empty for perf optimization).
lora_ids = [None] * bs
else:
lora_ids = None
with torch.device(self.device):
forward_batch = ForwardBatch(
forward_mode=ForwardMode.EXTEND,
batch_size=bs,
input_ids=input_ids,
input_embeds=input_embeds,
req_pool_indices=torch.arange(bs, device=self.device),
seq_lens=torch.tensor([num_tokens], device=self.device),
next_token_logits_buffer=None,
orig_seq_lens=torch.tensor([num_tokens], device=self.device),
seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
out_cache_loc=out_cache_loc,
seq_lens_sum=num_tokens,
mamba_track_indices=mamba_track_indices,
mamba_track_mask=mamba_track_mask,
mamba_track_seqlens=mamba_track_seqlens,
encoder_lens=None,
return_logprob=False,
extend_num_tokens=num_tokens,
extend_seq_lens=torch.tensor([num_tokens], device=self.device),
extend_prefix_lens=torch.tensor([0], device=self.device),
extend_start_loc=torch.tensor([0], device=self.device),
extend_prefix_lens_cpu=torch.tensor([0], device="cpu"),
extend_seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
extend_logprob_start_lens_cpu=torch.tensor([num_tokens], device="cpu"),
positions=positions,
global_num_tokens_gpu=None,
global_num_tokens_for_logprob_gpu=None,
dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
global_dp_buffer_len=None,
mrope_positions=mrope_positions,
spec_algorithm=None,
spec_info=None,
capture_hidden_mode=self.capture_hidden_mode,
num_token_non_padded=None,
num_token_non_padded_cpu=num_tokens,
global_forward_mode=ForwardMode.EXTEND,
lora_ids=None,
return_pooled_hidden_states=self.capture_return_pooled_hidden_states,
)
# Setup hooks below read get_attn_backend() and must run inside the
# same ForwardContext as the warmup/capture forward.
with forward_context(
ForwardContext(attn_backend=self.model_runner.attn_backend)
):
self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
if lora_ids is not None:
self.model_runner.lora_manager.prepare_lora_batch(forward_batch)
self.model_runner.attn_backend.init_forward_metadata(forward_batch)
# Run and capture
def run_once():
# Clean intermediate result cache for DP attention
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = (
None
)
set_dp_buffer_len(
global_dp_buffer_len,
num_tokens,
forward_batch.dp_padding_mode.is_max_len(),
global_num_tokens_cpu,
)
# FIXME: the implementation is hacky. `is_extend_in_batch`` is for determining the deepep mode.
# It is True in this context but we need to set it to use low latency deepep mode.
set_is_extend_in_batch(False)
kwargs = {}
with set_forward_context(
forward_batch,
self.attention_layers,
self.quant_config,
self.moe_layers,
self.moe_fusions,
dsa_indexers=self.dsa_indexers,
):
self.model_runner.model.forward(
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
**kwargs,
)
return
# run twice for warmup at the first time and cuda graph capture at the second time
# detail lies in sglang/python/sglang/srt/compilation/cuda_piecewise_backend.py
for _ in range(2):
self.device_module.synchronize()
self.model_runner.tp_group.barrier()
run_once()
return
def replay_prepare(
self,
forward_batch: ForwardBatch,
**kwargs,
):
num_tokens = len(forward_batch.input_ids)
index = bisect.bisect_left(self.capture_num_tokens, num_tokens)
static_num_tokens = self.capture_num_tokens[index]
self.raw_num_tokens = num_tokens
bs = forward_batch.batch_size
registry = self.buffer_registry
# Reset the padded token tail (ZERO) + copy the [:num_tokens] head for
# every graph-resident slot in one grouped pass. input_embeds is
# reset-only (the model writes embeds into it inside the graph).
registry.fill_from(
forward_batch,
raw_bs=bs,
padded_bs=bs,
raw_num_tokens=num_tokens,
padded_num_tokens=static_num_tokens,
)
def _slot(name):
return registry.get_slot(name).slice_for(bs, static_num_tokens)
input_ids = _slot("input_ids")
positions = _slot("positions")
out_cache_loc = _slot("out_cache_loc")
mamba_track_indices = (
_slot("mamba_track_indices")
if registry.has_slot("mamba_track_indices")
else None
)
mamba_track_mask = (
_slot("mamba_track_mask") if registry.has_slot("mamba_track_mask") else None
)
mamba_track_seqlens = (
_slot("mamba_track_seqlens")
if registry.has_slot("mamba_track_seqlens")
else None
)
input_embeds = (
_slot("input_embeds") if registry.has_slot("input_embeds") else None
)
mrope_positions = (
_slot("mrope_positions")
if (
registry.has_slot("mrope_positions")
and forward_batch.mrope_positions is not None
)
else None
)
next_token_logits_buffer = None
# Normalize MIXED→EXTEND so dynamo's guard (captured with EXTEND=1) doesn't fail on MIXED=3.
pcg_forward_mode = (
ForwardMode.EXTEND
if forward_batch.forward_mode == ForwardMode.MIXED
else forward_batch.forward_mode
)
pcg_global_forward_mode = (
ForwardMode.EXTEND
if forward_batch.global_forward_mode == ForwardMode.MIXED
else forward_batch.global_forward_mode
)
static_forward_batch = ForwardBatch(
forward_mode=pcg_forward_mode,
batch_size=bs,
input_ids=input_ids,
input_embeds=input_embeds,
req_pool_indices=forward_batch.req_pool_indices,
seq_lens=forward_batch.seq_lens,
next_token_logits_buffer=next_token_logits_buffer,
orig_seq_lens=forward_batch.orig_seq_lens,
seq_lens_cpu=forward_batch.seq_lens_cpu,
out_cache_loc=out_cache_loc,
seq_lens_sum=forward_batch.seq_lens_sum,
mamba_track_indices=mamba_track_indices,
mamba_track_mask=mamba_track_mask,
mamba_track_seqlens=mamba_track_seqlens,
encoder_lens=forward_batch.encoder_lens,
return_logprob=False,
extend_seq_lens=forward_batch.extend_seq_lens,
extend_prefix_lens=forward_batch.extend_prefix_lens,
extend_start_loc=forward_batch.extend_start_loc,
extend_prefix_lens_cpu=forward_batch.extend_prefix_lens_cpu,
extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
extend_logprob_start_lens_cpu=forward_batch.extend_logprob_start_lens_cpu,
extend_num_tokens=forward_batch.extend_num_tokens,
extend_input_logprob_token_ids_gpu=forward_batch.extend_input_logprob_token_ids_gpu,
positions=positions,
global_num_tokens_gpu=forward_batch.global_num_tokens_gpu,
global_num_tokens_for_logprob_gpu=forward_batch.global_num_tokens_for_logprob_gpu,
dp_padding_mode=forward_batch.dp_padding_mode,
global_dp_buffer_len=forward_batch.global_dp_buffer_len,
mrope_positions=mrope_positions,
spec_algorithm=forward_batch.spec_algorithm,
spec_info=forward_batch.spec_info,
capture_hidden_mode=forward_batch.capture_hidden_mode,
num_token_non_padded=forward_batch.num_token_non_padded,
num_token_non_padded_cpu=forward_batch.num_token_non_padded_cpu,
global_forward_mode=pcg_global_forward_mode,
lora_ids=forward_batch.lora_ids,
sampling_info=forward_batch.sampling_info,
mm_inputs=forward_batch.mm_inputs,
temperature=forward_batch.temperature,
top_p=forward_batch.top_p,
dimensions=forward_batch.dimensions,
return_pooled_hidden_states=(
self.capture_return_pooled_hidden_states
or forward_batch.return_pooled_hidden_states
),
)
return static_forward_batch
def replay(
self,
forward_batch: ForwardBatch,
**kwargs,
) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
with enable_piecewise_cuda_graph():
static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
static_num_tokens = len(static_forward_batch.input_ids)
raw_num_tokens = self.raw_num_tokens
# Replay
with set_forward_context(
static_forward_batch,
self.attention_layers,
self.quant_config,
self.moe_layers,
self.moe_fusions,
dsa_indexers=self.dsa_indexers,
num_tokens=static_num_tokens,
raw_num_tokens=raw_num_tokens,
):
self.model_runner.attn_backend.init_forward_metadata(forward_batch)
output = self.model_runner.model.forward(
static_forward_batch.input_ids,
static_forward_batch.positions,
static_forward_batch,
**kwargs,
)
if isinstance(output, LogitsProcessorOutput):
# Preserve mm_input_embeds when speculative decoding is
# enabled. The speculative draft's prefill path
# (eagle_worker_v2._draft_extend_for_prefill) reads
# mm_input_embeds off this LogitsProcessorOutput to reuse
# the target's encoder embeddings instead of re-embedding
# multimodal placeholder token ids.
mm_input_embeds = None
if (
self.model_runner.spec_algorithm.is_speculative()
and output.mm_input_embeds is not None
):
mm_input_embeds = output.mm_input_embeds[: self.raw_num_tokens]
return LogitsProcessorOutput(
next_token_logits=output.next_token_logits[
: self.raw_num_tokens
],
hidden_states=(
output.hidden_states[: self.raw_num_tokens]
if output.hidden_states is not None
else None
),
mm_input_embeds=mm_input_embeds,
)
elif isinstance(output, EmbeddingPoolerOutput):
return output
else:
assert isinstance(output, PPProxyTensors)
# TODO(Yuwei): support PP Support
raise NotImplementedError(
"PPProxyTensors is not supported in PiecewiseCudaGraphRunner yet."
)
def get_spec_info(self, num_tokens: int):
spec_info = None
if (
self.model_runner.spec_algorithm.is_eagle()
or self.model_runner.spec_algorithm.is_standalone()
):
from sglang.srt.speculative.eagle_utils import EagleVerifyInput
if self.model_runner.is_draft_worker:
raise RuntimeError("This should not happen.")
else:
spec_info = EagleVerifyInput(
draft_token=None,
custom_mask=self.custom_mask,
positions=None,
retrieve_index=None,
retrieve_next_token=None,
retrieve_next_sibling=None,
retrieve_cum_len=None,
spec_steps=self.model_runner.server_args.speculative_num_steps,
topk=self.model_runner.server_args.speculative_eagle_topk,
draft_token_num=self.model_runner.server_args.speculative_num_draft_tokens,
capture_hidden_mode=CaptureHiddenMode.FULL,
seq_lens_sum=None,
seq_lens_cpu=None,
)
return spec_info
python/sglang/srt/model_executor/runner/__init__.py
+49
View File
@@ -0,0 +1,49 @@
"""Phase-aware CUDA graph runners.
One concrete runner per phase. Each runner owns its phase-specific
shape semantics (decode → batch size; prefill → token count) and
delegates capture/replay mechanics to a pluggable
BaseCudaGraphBackend chosen via cuda_graph_config.
Public API:
- BaseCudaGraphRunner — abstract base; shared init + bucket
padding + capture-loop scaffolding.
- DecodeCudaGraphRunner — concrete decode-phase runner.
- PrefillCudaGraphRunner — concrete prefill-phase runner.
- Buffer dataclasses, capture-mode flags, the global memory pool,
and the DeepEP adapter live in
sglang.srt.model_executor.runner_utils; they are
re-exported here for the EAGLE / multi-step draft cuda graph
runners that were authored against the legacy public surface.
"""
from sglang.srt.model_executor.runner.base_cuda_graph_runner import ( # noqa: F401
BaseCudaGraphRunner,
freeze_gc,
get_batch_sizes_to_capture,
)
from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
DecodeCudaGraphRunner,
)
from sglang.srt.model_executor.runner.decode_cuda_graph_runner import ( # noqa: F401
_make_graph_key as _default_make_graph_key,
)
from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import ( # noqa: F401
PrefillCudaGraphRunner,
)
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( # noqa: F401
TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG,
)
from sglang.srt.model_executor.runner_utils import ( # noqa: F401
DecodeInputBuffers,
DeepEPCudaGraphRunnerAdapter,
PrefillInputBuffers,
_grouped_foreach_copy_,
_set_capture_lora_variant,
compile_in_capture_mode,
get_capture_lora_variant,
get_global_graph_memory_pool,
get_is_capture_mode,
model_capture_mode,
set_global_graph_memory_pool,
)
python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py
+177
View File
@@ -0,0 +1,177 @@
"""Shared scaffolding for the prefill and decode CUDA graph runners."""
from __future__ import annotations
import bisect
import gc
import logging
from abc import ABC, abstractmethod
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, List, Sequence, Tuple
import torch
from sglang.srt.batch_overlap.two_batch_overlap import TboCudaGraphRunnerPlugin
from sglang.srt.layers.dp_attention import (
get_attention_cp_size,
get_attention_tp_rank,
get_attention_tp_size,
)
from sglang.srt.utils import require_gathered_buffer
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
BaseCudaGraphBackend,
)
logger = logging.getLogger(__name__)
@contextmanager
def freeze_gc(enable_cudagraph_gc: bool):
"""Optimize garbage collection during CUDA graph capture.
Clean up first, then freeze remaining objects from being included in
future collections if GC is disabled during capture.
"""
gc.collect()
should_freeze = not enable_cudagraph_gc
if should_freeze:
gc.freeze()
try:
yield
finally:
if should_freeze:
gc.unfreeze()
gc.collect()
def get_batch_sizes_to_capture(
model_runner: ModelRunner, num_tokens_per_bs: int = 1
) -> Tuple[List[int], List[int]]:
"""Build the (capture_bs, compile_bs) lists for the decode runner.
Filters cuda_graph_config[decode].bs by attention-tp/cp alignment
constraints and clamps to req_to_token_pool.size.
"""
server_args = model_runner.server_args
capture_bs = server_args.cuda_graph_config.decode.bs
num_max_requests = model_runner.req_to_token_pool.size
mul_base = 1
if server_args.enable_two_batch_overlap:
mul_base *= 2
num_tokens_per_bs = 1
if require_gathered_buffer(server_args):
mul_base *= get_attention_tp_size()
if mul_base % get_attention_cp_size() != 0:
mul_base *= get_attention_cp_size()
num_max_requests = (num_max_requests + mul_base - 1) // mul_base * mul_base
if max(capture_bs) > num_max_requests:
capture_bs += [num_max_requests]
capture_bs = [bs for bs in capture_bs if bs * num_tokens_per_bs % mul_base == 0]
capture_bs = [bs for bs in capture_bs if bs <= num_max_requests]
capture_bs = list(sorted(set(capture_bs)))
assert len(capture_bs) > 0 and capture_bs[0] > 0, f"{capture_bs=}"
compile_bs = (
[bs for bs in capture_bs if bs <= server_args.torch_compile_max_bs]
if server_args.enable_torch_compile
else []
)
return capture_bs, compile_bs
class BaseCudaGraphRunner(ABC):
"""Abstract base for phase-specific cuda-graph runners.
A subclass (DecodeCudaGraphRunner / PrefillCudaGraphRunner) owns one
phase and plugs in a BaseCudaGraphBackend that handles the
capture / replay mechanics. The runner orchestrates bucket
selection, static buffer population, attention metadata init,
replay dispatch, and output slicing.
Methods:
- can_run(forward_batch) — should forward_batch go through cuda
graph replay (vs eager fallback)?
- capture_prepare(size, ...) — build the dummy ForwardBatch and
per-capture local state needed by capture_one_shape.
- capture() — outer capture loop; iterates over shapes and calls
capture_one_shape for each.
- capture_one_shape(size, ...) — drive one model forward at this
shape into the backend's captured artifact.
- replay_prepare(forward_batch, ...) — pad to the nearest captured
bucket, populate static input buffers, init attention metadata.
- replay(forward_batch, ...) — dispatch one batch through cuda
graph replay.
Notes:
- buffers and backend are populated by the subclass before
capture(); the base only declares them.
"""
# Subclasses populate before calling capture().
buffers: ForwardInputBuffers
backend: BaseCudaGraphBackend
def __init__(self, model_runner: ModelRunner) -> None:
self.model_runner = model_runner
self.device = model_runner.device
self.device_module = torch.get_device_module(self.device)
self.tp_size = model_runner.server_args.tp_size
self.dp_size = model_runner.server_args.dp_size
self.pp_size = model_runner.server_args.pp_size
self.attn_tp_size = get_attention_tp_size()
self.attn_tp_rank = get_attention_tp_rank()
self.tbo_plugin = TboCudaGraphRunnerPlugin()
@staticmethod
def _pad_to_bucket(raw_size: int, buckets: Sequence[int]) -> int:
"""Return the smallest buckets[i] >= raw_size.
Caller's can_run must reject raw_size > max(buckets) before
reaching replay_prepare; this assertion makes the contract
explicit (bisect_left returns len(buckets) when the value
exceeds all buckets, which would otherwise IndexError below
with no diagnostic).
"""
assert raw_size <= buckets[-1], (
f"size {raw_size} exceeds max captured bucket {buckets[-1]}; "
f"can_run should have rejected this batch"
)
index = bisect.bisect_left(buckets, raw_size)
return buckets[index]
@abstractmethod
def can_run(self, forward_batch: ForwardBatch) -> bool: ...
@abstractmethod
def capture_prepare(self, size: int, *args, **kwargs) -> Any: ...
@abstractmethod
def capture(self) -> None: ...
@abstractmethod
def capture_one_shape(self, size: int, *args, **kwargs) -> Any: ...
@abstractmethod
def replay_prepare(
self,
forward_batch: ForwardBatch,
**kwargs,
) -> Any: ...
@abstractmethod
def replay(
self,
forward_batch: ForwardBatch,
**kwargs,
) -> Any: ...
python/sglang/srt/model_executor/cuda_graph_runner.py → python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+292 -461
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+846
View File
@@ -0,0 +1,846 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""PrefillCudaGraphRunner — runs the EXTEND phase under a pluggable backend.
Backend selection comes from cuda_graph_config.prefill:
- "tc_piecewise" — default, TcPiecewiseCudaGraphBackend: torch.compile
wraps the model; per-shape graphs live in
torch.compile's internal cache. Multi-batch supported.
- "breakable" — BreakableCudaGraphBackend: segmented capture (no
torch.compile). Captures with bs=1; rejects multi-req
prefill in can_run.
- "full" — rejected at config validation; not supported for prefill.
- "disabled" — handled at the model_runner level — runner not
constructed.
"""
from __future__ import annotations
import logging
import warnings
from typing import TYPE_CHECKING, Dict, Optional, Union
import torch
import tqdm
from sglang.srt.distributed import get_tensor_model_parallel_rank
from sglang.srt.distributed.parallel_state import graph_capture
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
set_dp_buffer_len,
set_is_extend_in_batch,
)
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.pooler import EmbeddingPoolerOutput
from sglang.srt.model_executor.cuda_graph_buffer_registry import (
CudaGraphBufferRegistry,
build_prefill_registry,
)
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
ForwardMode,
PPProxyTensors,
)
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
BaseCudaGraphRunner,
freeze_gc,
)
from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import (
BreakableCudaGraphBackend,
)
from sglang.srt.model_executor.runner_backend.utils import (
resolve_prefill_backend,
)
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
set_tc_piecewise_forward_context,
)
from sglang.srt.model_executor.runner_utils.buffers import (
PrefillInputBuffers,
)
from sglang.srt.utils import (
get_available_gpu_memory,
is_npu,
log_info_on_rank0,
require_attn_tp_gather,
require_mlp_tp_gather,
)
# Suppress Dynamo warning about tracing through lru_cache-wrapped functions.
warnings.filterwarnings("ignore", message=".*lru_cache.*", module="torch._dynamo")
logger = logging.getLogger(__name__)
# Names of the static prefill input tensors a Breakable-backed prefill
# runner owns. Each is a 1-D int64 tensor of length max_bs; captured
# Breakable segments read from these stable addresses.
_PREFILL_STATIC_FIELDS = (
"seq_lens",
"extend_seq_lens",
"extend_prefix_lens",
"extend_start_loc",
"req_pool_indices",
"orig_seq_lens",
)
if TYPE_CHECKING:
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.model_executor.model_runner import ModelRunner
class PrefillCudaGraphRunner(BaseCudaGraphRunner):
"""Prefill-phase CUDA graph runner.
Owns: PrefillInputBuffers, capture-num-tokens list, attention layers
snapshot, and the pluggable self.backend. The backend handles capture
+ replay mechanics; this runner handles dummy ForwardBatch construction,
buffer population, attention metadata init, and output slicing.
"""
def __init__(self, model_runner: ModelRunner):
super().__init__(model_runner)
# --- core state ------------------------------------------------
self.quant_config = getattr(self.model_runner.model, "quant_config", None)
self.is_multimodal = model_runner.is_multimodal
# Classification/reward forwards branch on return_pooled_hidden_states;
# capture must use the same flag value as replay for those models.
self.capture_return_pooled_hidden_states = not model_runner.is_generation
# --- bucket sizes ---------------------------------------------
# bs in prefill carries the captured shape (token count for
# tc_piecewise) — one shape knob per phase.
capture_tokens = model_runner.server_args.cuda_graph_config.prefill.bs
assert capture_tokens is not None, "cuda_graph_config[prefill].bs is not set"
self.capture_num_tokens = sorted(capture_tokens)
self.max_num_tokens = (
max(self.capture_num_tokens) if self.capture_num_tokens else 8192
)
self.max_bs = model_runner.req_to_token_pool.size
log_info_on_rank0(
logger, f"Capture cuda graph num tokens {self.capture_num_tokens}"
)
self.capture_forward_mode = ForwardMode.EXTEND
self.capture_hidden_mode = CaptureHiddenMode.NULL
# If returning hidden states is enabled, or if speculative prefill
# needs aux hidden states (DFLASH), capture the FULL variant up front.
# Ported from main #27468.
if (
model_runner.server_args.enable_return_hidden_states
or model_runner.spec_algorithm.is_dflash()
):
self.capture_hidden_mode = CaptureHiddenMode.FULL
self.mamba_track_enabled = self._is_mamba_track_enabled()
# --- buffers ---------------------------------------------------
self.buffers: PrefillInputBuffers = PrefillInputBuffers.create(
device=self.device,
max_bs=self.max_bs,
max_num_tokens=self.max_num_tokens,
cache_loc_dtype=self._cache_loc_dtype(),
is_hybrid_swa=model_runner.is_hybrid_swa,
is_multimodal=self.is_multimodal,
hidden_size=self.model_runner.model_config.hidden_size,
dtype=self.model_runner.dtype,
enable_mamba_track=self.mamba_track_enabled,
)
self.buffers.share_buffers()
# Token-axis FB-shared slot registry adopting PrefillInputBuffers
# storage; same physical tensors, stable data_ptr for capture vs
# replay. Replaces populate_from_forward_batch on capture/replay paths.
self.buffer_registry: CudaGraphBufferRegistry = build_prefill_registry(
device=self.device,
max_bs=self.max_bs,
max_num_token=self.max_num_tokens,
cache_loc_dtype=self._cache_loc_dtype(),
is_multimodal=self.is_multimodal,
hidden_size=self.model_runner.model_config.hidden_size,
embed_dtype=self.model_runner.dtype,
enable_mamba_track=self.mamba_track_enabled,
source=self.buffers,
)
self.attention_layers = self.model_runner.attention_layers
self.moe_layers = self.model_runner.moe_layers
self.moe_fusions = self.model_runner.moe_fusions
self.dsa_indexers = getattr(self.model_runner, "dsa_indexers", None)
self.dp_size = model_runner.server_args.dp_size
self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args)
self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args)
# --- backend ---------------------------------------------------
# When the backend is Breakable, captured segments need stable
# tensor addresses, so we own a set of static int64 buffers here
# and rebind them into capture-time dummy inputs / replay-time
# serving inputs below. Other backends don't need this.
# Initialize the slot to None BEFORE constructing the backend:
# TcPiecewise runs its compile pass during __init__ which calls
# _run_dummy_forward -> capture_prepare, and capture_prepare reads
# self._prefill_static_buffers. self.layer_model has the same
# ordering requirement: _run_forward checks `self.layer_model is
# not None` to decide whether to call the inner stack or outer
# model.forward, and that check fires inside TcPiecewise's
# _run_compile_pass before backend resolution returns.
self._prefill_static_buffers: Optional[Dict[str, torch.Tensor]] = None
self.layer_model = None
self.backend = resolve_prefill_backend(self)
if isinstance(self.backend, BreakableCudaGraphBackend):
with torch.device(self.device):
self._prefill_static_buffers = {
name: torch.zeros((self.max_bs,), dtype=torch.int64)
for name in _PREFILL_STATIC_FIELDS
}
# Some attention backends (e.g. DSV4) opt into a captured-metadata
# contract under BCG: capture-time builds a per-bucket metadata
# object the backend then refreshes in place at replay. We honor
# the contract only when the backend is Breakable; FullCG and
# TC_PIECEWISE use the eager init_forward_metadata path.
if isinstance(self.backend, BreakableCudaGraphBackend):
self.use_captured_attn_metadata = (
model_runner.attn_backend.use_captured_forward_metadata_for_breakable_cuda_graph
)
else:
self.use_captured_attn_metadata = False
self.attn_metadata_buffers: Optional[Dict[int, object]] = (
{} if self.use_captured_attn_metadata else None
)
# --- BCG: resolve inner layer_model for capture/replay --------
# BCG captures only the inner transformer stack (layer_model.forward)
# — not the outer model.forward. The outer's tail (logits_processor /
# pooler) has bs-shaped kernels that would bake bs=1 into the captured
# graph and break multi-req replay. At replay, we monkey-patch
# layer_model.forward to replay the captured graph and return the
# captured hidden states; the outer model.forward then runs
# logits_processor eagerly on top with the live multi-req metadata.
# Mirrors main's BreakableCudaGraphRunner. (Slot pre-init lives
# above next to _prefill_static_buffers — TcPiecewise's compile
# pass runs during backend construction and reads self.layer_model.)
if isinstance(self.backend, BreakableCudaGraphBackend):
language_model = getattr(
self.model_runner.model, "language_model", self.model_runner.model
)
if hasattr(language_model, "model") and hasattr(
language_model.model, "layers"
):
self.layer_model = language_model.model
else:
raise RuntimeError(
f"BCG could not resolve inner layer_model on "
f"{type(language_model).__name__}; BCG is unsupported for "
f"this model architecture."
)
# --- capture --------------------------------------------------
self.device_module.synchronize()
self.model_runner.tp_group.barrier()
self.capture()
self.raw_num_tokens = 0
# -----------------------------------------------------------------
# Helpers
# -----------------------------------------------------------------
def _is_mamba_track_enabled(self) -> bool:
return (
self.model_runner.server_args.enable_mamba_extra_buffer()
and not self.model_runner.server_args.disable_radix_cache
and self.model_runner.spec_algorithm.is_none()
)
def _cache_loc_dtype(self):
return torch.int64 if not is_npu() else torch.int32
@torch.no_grad()
def _run_forward(self, forward_batch: ForwardBatch, num_tokens: int):
"""Run forward inside the prefill set_tc_piecewise_forward_context.
BCG path: captures only the inner layer_model.forward (transformer
stack), excluding the outer model.forward tail (logits_processor /
pooler). The captured output is bs=1 hidden states; replay then runs
the outer tail eagerly with live multi-req metadata.
TC_PIECEWISE path: captures the outer model.forward; torch.compile
FX-traces produce bs-invariant kernels.
``@torch.no_grad`` mirrors the decorator on the outer
``*ForCausalLM.forward``. For BCG, calling ``layer_model.forward``
directly skips that decorator, so we apply it here — without it
some MoE ``@torch.compile`` kernels (``torch.sum(out=...)``) fail
dynamo with "out= doesn't support autograd", and mamba state ops
can spuriously track gradients.
"""
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
set_dp_buffer_len(
forward_batch.global_dp_buffer_len,
num_tokens,
forward_batch.dp_padding_mode.is_max_len(),
forward_batch.global_num_tokens_cpu,
)
set_is_extend_in_batch(False)
with forward_context(
ForwardContext(attn_backend=self.model_runner.attn_backend)
), set_tc_piecewise_forward_context(
forward_batch,
self.attention_layers,
self.quant_config,
self.moe_layers,
self.moe_fusions,
dsa_indexers=self.dsa_indexers,
):
if self.layer_model is not None:
return self.layer_model.forward(
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
forward_batch.input_embeds,
)
return self.model_runner.model.forward(
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
)
def _run_dummy_forward(self, num_tokens: int) -> None:
"""Build a dummy ForwardBatch at this shape, init attn metadata,
run forward once. Used by TcPiecewiseCudaGraphBackend.prepare
for both the JIT-activate forward (single shape, before
torch.compile install) and the compile-loop pass (every shape,
inside enable_torch_compile_warmup).
"""
fb, attn_backend = self.capture_prepare(num_tokens)
attn_backend.init_forward_metadata(fb)
self._run_forward(fb, num_tokens)
def _has_inactive_dp_rank(self, forward_batch: ForwardBatch) -> bool:
# DSV4 DP attention / DeepEP collectives need every DP rank to enter
# the same replay path. Sparse-DP batches (one or more ranks with
# zero local tokens) fall back to eager to avoid hanging ranks.
global_num_tokens = forward_batch.global_num_tokens_cpu
if global_num_tokens is None:
return False
return len(global_num_tokens) > 1 and any(
int(num_tokens) == 0 for num_tokens in global_num_tokens
)
def _init_forward_metadata_for_capture(
self, forward_batch: ForwardBatch, num_tokens: int
) -> None:
"""Capture-time metadata init for the BCG-with-captured-metadata
contract. For opt-in backends (DSV4), call the BCG-specific entry
and stash the returned per-bucket metadata object; otherwise fall
back to the generic eager init that BCG/TC_PIECEWISE use today."""
attn_backend = self.model_runner.attn_backend
if not self.use_captured_attn_metadata:
attn_backend.init_forward_metadata(forward_batch)
return
metadata = attn_backend.init_forward_metadata_for_breakable_cuda_graph_capture(
forward_batch
)
assert self.attn_metadata_buffers is not None
self.attn_metadata_buffers[num_tokens] = metadata
def _prepare_forward_metadata_for_replay(
self,
forward_batch: ForwardBatch,
static_forward_batch: ForwardBatch,
num_tokens: int,
) -> None:
"""Replay-time metadata refresh for the BCG-with-captured-metadata
contract. For opt-in backends, refresh the stashed per-bucket
metadata in place against the current batch; otherwise fall back
to the generic eager init."""
attn_backend = self.model_runner.attn_backend
if not self.use_captured_attn_metadata:
attn_backend.init_forward_metadata(forward_batch)
return
assert self.attn_metadata_buffers is not None
metadata = self.attn_metadata_buffers[num_tokens]
attn_backend.prepare_forward_metadata_for_breakable_cuda_graph_replay(
metadata,
forward_batch,
static_forward_batch=static_forward_batch,
)
# -----------------------------------------------------------------
# can_run
# -----------------------------------------------------------------
def can_run(self, forward_batch: ForwardBatch) -> bool:
if forward_batch.input_embeds is not None:
return False
if forward_batch.replace_embeds is not None:
return False
# tc_piecewise captures with ForwardMode.EXTEND and spec_info=None.
if forward_batch.forward_mode.is_target_verify():
return False
if forward_batch.capture_hidden_mode != self.capture_hidden_mode:
return False
# BCG-with-captured-metadata under DP attention: every rank must
# have local tokens, and the batch must declare itself replayable.
# These gates are no-ops for non-DP / non-opt-in paths because
# global_num_tokens_cpu stays None.
if self._has_inactive_dp_rank(forward_batch):
return False
if (
forward_batch.global_num_tokens_cpu is not None
and not forward_batch.can_run_dp_breakable_cuda_graph
):
return False
num_tokens = len(forward_batch.input_ids)
if forward_batch.return_logprob:
for start_len, seq_len in zip(
forward_batch.extend_logprob_start_lens_cpu,
forward_batch.extend_seq_lens_cpu,
):
if start_len is not None and start_len < seq_len:
return False
if num_tokens > self.max_num_tokens:
return False
# No backend-level shape check here: replay_prepare bucket-pads
# num_tokens up to the nearest captured shape, so eligibility is
# bounded by num_tokens <= self.max_num_tokens (already
# checked above), not by exact shape membership.
#
# Multi-req replay is supported by BCG via the layer_model.forward
# monkey-patch in replay(): the captured bs=1 graph runs the
# transformer stack, then the outer model.forward runs
# logits_processor eagerly on top with live multi-req metadata.
return True
# -----------------------------------------------------------------
# capture_prepare
# -----------------------------------------------------------------
def capture_prepare(
self, num_tokens: int
) -> tuple[ForwardBatch, "AttentionBackend"]:
"""Build a dummy prefill ForwardBatch for capture/warmup at this shape.
Default tensor inputs are fresh literals; under a Breakable
backend, we swap in slices of our static buffers so captured
segments read from stable addresses.
Returns ``(forward_batch, attn_backend)`` to mirror decode's
capture_prepare signature.
"""
buffers = self.buffers
bs = 1
with torch.device(self.device):
shape_inputs = {
"req_pool_indices": torch.arange(bs, device=self.device),
"seq_lens": torch.tensor([num_tokens], device=self.device),
"orig_seq_lens": torch.tensor([num_tokens], device=self.device),
"extend_seq_lens": torch.tensor([num_tokens], device=self.device),
"extend_prefix_lens": torch.tensor([0], device=self.device),
"extend_start_loc": torch.tensor([0], device=self.device),
}
if self._prefill_static_buffers is not None:
s = self._prefill_static_buffers
s["seq_lens"][:bs].fill_(num_tokens)
s["extend_seq_lens"][:bs].fill_(num_tokens)
s["extend_prefix_lens"][:bs].zero_()
s["extend_start_loc"][:bs].zero_()
s["req_pool_indices"][:bs].copy_(
torch.arange(bs, device=s["req_pool_indices"].device)
)
s["orig_seq_lens"][:bs].fill_(num_tokens)
for name in _PREFILL_STATIC_FIELDS:
shape_inputs[name] = s[name][:bs]
registry = self.buffer_registry
def _slot(name):
return registry.get_slot(name).slice_for(bs, num_tokens)
if self.require_mlp_tp_gather:
global_num_tokens_cpu = [num_tokens] * self.dp_size
elif self.require_attn_tp_gather:
global_num_tokens_cpu = [num_tokens]
else:
global_num_tokens_cpu = None
if global_num_tokens_cpu is not None:
global_dp_buffer_len = sum(global_num_tokens_cpu)
num_tokens_tensor = torch.tensor(
global_num_tokens_cpu, dtype=torch.int32, device=self.device
)
global_num_tokens_gpu = num_tokens_tensor
global_num_tokens_for_logprob_gpu = num_tokens_tensor
else:
global_dp_buffer_len = None
global_num_tokens_gpu = None
global_num_tokens_for_logprob_gpu = None
with torch.device(self.device):
forward_batch = ForwardBatch(
forward_mode=ForwardMode.EXTEND,
batch_size=bs,
input_ids=_slot("input_ids"),
input_embeds=(
_slot("input_embeds") if registry.has_slot("input_embeds") else None
),
req_pool_indices=shape_inputs["req_pool_indices"],
seq_lens=shape_inputs["seq_lens"],
next_token_logits_buffer=None,
orig_seq_lens=shape_inputs["orig_seq_lens"],
seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
out_cache_loc=_slot("out_cache_loc"),
seq_lens_sum=num_tokens,
mamba_track_indices=(
_slot("mamba_track_indices")
if registry.has_slot("mamba_track_indices")
else None
),
mamba_track_mask=(
_slot("mamba_track_mask")
if registry.has_slot("mamba_track_mask")
else None
),
mamba_track_seqlens=(
_slot("mamba_track_seqlens")
if registry.has_slot("mamba_track_seqlens")
else None
),
encoder_lens=None,
return_logprob=False,
extend_num_tokens=num_tokens,
extend_seq_lens=shape_inputs["extend_seq_lens"],
extend_prefix_lens=shape_inputs["extend_prefix_lens"],
extend_start_loc=shape_inputs["extend_start_loc"],
extend_prefix_lens_cpu=torch.tensor([0], device="cpu"),
extend_seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
extend_logprob_start_lens_cpu=torch.tensor([num_tokens], device="cpu"),
positions=_slot("positions"),
global_num_tokens_gpu=global_num_tokens_gpu,
global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
global_num_tokens_cpu=global_num_tokens_cpu,
dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
global_dp_buffer_len=global_dp_buffer_len,
mrope_positions=(
_slot("mrope_positions")
if registry.has_slot("mrope_positions")
else None
),
spec_algorithm=None,
spec_info=None,
# Use self.capture_hidden_mode so dflash spec (which needs
# FULL aux hidden states) captures with the right mode.
# Ported from main #27468.
capture_hidden_mode=self.capture_hidden_mode,
num_token_non_padded=None,
num_token_non_padded_cpu=num_tokens,
global_forward_mode=ForwardMode.EXTEND,
lora_ids=None,
return_pooled_hidden_states=self.capture_return_pooled_hidden_states,
)
self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
return forward_batch, self.model_runner.attn_backend
# -----------------------------------------------------------------
# capture
# -----------------------------------------------------------------
def capture(self) -> None:
with freeze_gc(self.model_runner.server_args.enable_cudagraph_gc):
with graph_capture() as graph_capture_context:
self.stream = graph_capture_context.stream
with self.backend.capture_session(self.stream):
self._capture_one_stream()
def _capture_one_stream(self) -> None:
avail_mem = get_available_gpu_memory(
self.model_runner.device,
self.model_runner.gpu_id,
empty_cache=False,
)
capture_range = (
tqdm.tqdm(list(reversed(self.capture_num_tokens)))
if get_tensor_model_parallel_rank() == 0
else reversed(self.capture_num_tokens)
)
for num_tokens in capture_range:
if get_tensor_model_parallel_rank() == 0:
avail_mem = get_available_gpu_memory(
self.model_runner.device,
self.model_runner.gpu_id,
empty_cache=False,
)
capture_range.set_description(
f"Capturing num tokens ({num_tokens=} {avail_mem=:.2f} GB)"
)
self.capture_one_shape(num_tokens)
# -----------------------------------------------------------------
# capture_one_shape
# -----------------------------------------------------------------
def capture_one_shape(self, size: int) -> None:
"""Per-shape capture: build dummy ForwardBatch + run_once,
delegate to backend. size is the prefill token count.
"""
num_tokens = size
forward_batch, attn_backend = self.capture_prepare(num_tokens)
self._init_forward_metadata_for_capture(forward_batch, num_tokens)
def run_once():
return self._run_forward(forward_batch, num_tokens)
# Main's monolithic BCG runner never invokes
# on_after_cuda_graph_warmup between warmup iterations — the BCG
# contract is to keep warmup state untouched and let
# init_forward_metadata_in_graph (recorded inside the captured
# forward) do any raw->full upgrade. cg-refactor's runner_backend
# abstraction exposes a post_warmup_hook for backends that need
# workspace cleanup between iterations; suppress it for BCG so
# DSV4's hook (which restores forward_metadata to a stale
# _current_capture_raw left over from decode CG capture) doesn't
# corrupt warmup iter 2's metadata read.
if isinstance(self.backend, BreakableCudaGraphBackend):
post_warmup_hook = None
else:
post_warmup_hook = getattr(attn_backend, "on_after_cuda_graph_warmup", None)
self.backend.capture_one(
num_tokens,
run_once,
dummies=None,
post_warmup_hook=post_warmup_hook,
)
# -----------------------------------------------------------------
# replay_prepare
# -----------------------------------------------------------------
def replay_prepare(self, forward_batch: ForwardBatch, **kwargs) -> ForwardBatch:
"""Pad, populate static buffers, and build the static_forward_batch
the model code reads during replay.
"""
buffers = self.buffers
num_tokens = len(forward_batch.input_ids)
static_num_tokens = self._pad_to_bucket(num_tokens, self.capture_num_tokens)
self.raw_num_tokens = num_tokens
bs = forward_batch.batch_size
self.buffer_registry.fill_from(
forward_batch,
raw_bs=bs,
padded_bs=bs,
raw_num_tokens=num_tokens,
padded_num_tokens=static_num_tokens,
)
registry = self.buffer_registry
def _slot(name):
return registry.get_slot(name).slice_for(bs, static_num_tokens)
mamba_track_indices = (
_slot("mamba_track_indices")
if registry.has_slot("mamba_track_indices")
else None
)
mamba_track_mask = (
_slot("mamba_track_mask") if registry.has_slot("mamba_track_mask") else None
)
mamba_track_seqlens = (
_slot("mamba_track_seqlens")
if registry.has_slot("mamba_track_seqlens")
else None
)
input_ids = _slot("input_ids")
input_embeds = (
_slot("input_embeds") if registry.has_slot("input_embeds") else None
)
positions = _slot("positions")
out_cache_loc = _slot("out_cache_loc")
mrope_positions = (
_slot("mrope_positions")
if registry.has_slot("mrope_positions")
and forward_batch.mrope_positions is not None
else None
)
# Normalize MIXED→EXTEND so dynamo's guard (captured with EXTEND=1)
# doesn't fail on MIXED=3.
pcg_forward_mode = (
ForwardMode.EXTEND
if forward_batch.forward_mode == ForwardMode.MIXED
else forward_batch.forward_mode
)
pcg_global_forward_mode = (
ForwardMode.EXTEND
if forward_batch.global_forward_mode == ForwardMode.MIXED
else forward_batch.global_forward_mode
)
static_forward_batch = ForwardBatch(
forward_mode=pcg_forward_mode,
batch_size=bs,
input_ids=input_ids,
input_embeds=input_embeds,
req_pool_indices=forward_batch.req_pool_indices,
seq_lens=forward_batch.seq_lens,
next_token_logits_buffer=None,
orig_seq_lens=forward_batch.orig_seq_lens,
seq_lens_cpu=forward_batch.seq_lens_cpu,
out_cache_loc=out_cache_loc,
seq_lens_sum=forward_batch.seq_lens_sum,
mamba_track_indices=mamba_track_indices,
mamba_track_mask=mamba_track_mask,
mamba_track_seqlens=mamba_track_seqlens,
encoder_lens=forward_batch.encoder_lens,
return_logprob=False,
extend_seq_lens=forward_batch.extend_seq_lens,
extend_prefix_lens=forward_batch.extend_prefix_lens,
extend_start_loc=forward_batch.extend_start_loc,
extend_prefix_lens_cpu=forward_batch.extend_prefix_lens_cpu,
extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
extend_logprob_start_lens_cpu=forward_batch.extend_logprob_start_lens_cpu,
extend_num_tokens=forward_batch.extend_num_tokens,
extend_input_logprob_token_ids_gpu=forward_batch.extend_input_logprob_token_ids_gpu,
positions=positions,
global_num_tokens_gpu=forward_batch.global_num_tokens_gpu,
global_num_tokens_for_logprob_gpu=forward_batch.global_num_tokens_for_logprob_gpu,
dp_padding_mode=forward_batch.dp_padding_mode,
global_dp_buffer_len=forward_batch.global_dp_buffer_len,
mrope_positions=mrope_positions,
spec_algorithm=forward_batch.spec_algorithm,
spec_info=forward_batch.spec_info,
capture_hidden_mode=forward_batch.capture_hidden_mode,
num_token_non_padded=forward_batch.num_token_non_padded,
num_token_non_padded_cpu=forward_batch.num_token_non_padded_cpu,
global_forward_mode=pcg_global_forward_mode,
lora_ids=forward_batch.lora_ids,
sampling_info=forward_batch.sampling_info,
mm_inputs=forward_batch.mm_inputs,
temperature=forward_batch.temperature,
top_p=forward_batch.top_p,
dimensions=forward_batch.dimensions,
return_pooled_hidden_states=(
self.capture_return_pooled_hidden_states
or forward_batch.return_pooled_hidden_states
),
)
# Under Breakable, copy serving-time values into the static
# buffers so the addresses captured segments hold stay live with
# current data.
if self._prefill_static_buffers is not None:
bs = forward_batch.batch_size
s = self._prefill_static_buffers
s["seq_lens"][:bs].copy_(forward_batch.seq_lens)
s["extend_seq_lens"][:bs].copy_(forward_batch.extend_seq_lens)
s["extend_prefix_lens"][:bs].copy_(forward_batch.extend_prefix_lens)
s["extend_start_loc"][:bs].copy_(forward_batch.extend_start_loc)
s["req_pool_indices"][:bs].copy_(forward_batch.req_pool_indices)
if forward_batch.orig_seq_lens is not None:
s["orig_seq_lens"][:bs].copy_(forward_batch.orig_seq_lens)
self._prepare_forward_metadata_for_replay(
forward_batch, static_forward_batch, static_num_tokens
)
self._static_num_tokens = static_num_tokens
return static_forward_batch
# -----------------------------------------------------------------
# replay
# -----------------------------------------------------------------
def replay(
self, forward_batch: ForwardBatch, **kwargs
) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
with self.backend.replay_session():
static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
if self.layer_model is not None:
# BCG path. The captured graph is a bs=1 replay of
# layer_model.forward. Monkey-patch layer_model.forward to
# call backend.replay (which fires the captured graph and
# returns the captured hidden_states), then drive the outer
# model.forward eagerly with the live multi-req
# static_forward_batch. The outer's logits_processor /
# pooler then runs on top with live multi-req metadata.
shape_key = self._static_num_tokens
def replay_layer_forward(*args, **layer_kwargs):
return self.backend.replay(
shape_key, static_forward_batch, **kwargs
)
original_layer_forward = self.layer_model.forward
self.layer_model.forward = replay_layer_forward
try:
with forward_context(
ForwardContext(attn_backend=self.model_runner.attn_backend)
), set_tc_piecewise_forward_context(
static_forward_batch,
self.attention_layers,
self.quant_config,
self.moe_layers,
self.moe_fusions,
dsa_indexers=self.dsa_indexers,
):
output = self.model_runner.model.forward(
static_forward_batch.input_ids,
static_forward_batch.positions,
static_forward_batch,
**kwargs,
)
finally:
self.layer_model.forward = original_layer_forward
else:
# TC_PIECEWISE path. backend.replay calls the compiled
# outer model.forward directly (torch.compile handles
# multi-req via bs-invariant FX-traced kernels).
with forward_context(
ForwardContext(attn_backend=self.model_runner.attn_backend)
), set_tc_piecewise_forward_context(
static_forward_batch,
self.attention_layers,
self.quant_config,
self.moe_layers,
self.moe_fusions,
dsa_indexers=self.dsa_indexers,
):
output = self.backend.replay(
self._static_num_tokens, static_forward_batch, **kwargs
)
if isinstance(output, LogitsProcessorOutput):
# Preserve mm_input_embeds for speculative decoding.
mm_input_embeds = None
if (
self.model_runner.spec_algorithm.is_speculative()
and output.mm_input_embeds is not None
):
mm_input_embeds = output.mm_input_embeds[: self.raw_num_tokens]
return LogitsProcessorOutput(
next_token_logits=output.next_token_logits[: self.raw_num_tokens],
hidden_states=(
output.hidden_states[: self.raw_num_tokens]
if output.hidden_states is not None
else None
),
mm_input_embeds=mm_input_embeds,
)
elif isinstance(output, EmbeddingPoolerOutput):
return output
else:
assert isinstance(output, PPProxyTensors)
raise NotImplementedError(
"PPProxyTensors is not supported in PrefillCudaGraphRunner yet."
)
python/sglang/srt/model_executor/runner_backend/__init__.py
+31
View File
@@ -0,0 +1,31 @@
"""Capture-mechanism backends for CUDA graphs.
A backend owns *how* a captured artifact is produced and replayed for
one shape; it is phase-agnostic. Runners (cuda_graph_runner/) own
*what* data flows in and out.
Public API:
- BaseCudaGraphBackend — abstract interface.
- FullCudaGraphBackend — single torch.cuda.CUDAGraph per shape.
- BreakableCudaGraphBackend — segmented capture with eager break
markers; no torch.compile.
- TcPiecewiseCudaGraphBackend — torch.compile-driven piecewise
capture; FX-splits the model at attention layers.
"""
from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import ( # noqa: F401
BaseCudaGraphBackend,
)
from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import ( # noqa: F401
BreakableCudaGraphBackend,
)
from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import ( # noqa: F401
FullCudaGraphBackend,
)
from sglang.srt.model_executor.runner_backend.tc_piecewise_cuda_graph_backend import ( # noqa: F401
TcPiecewiseCudaGraphBackend,
)
from sglang.srt.model_executor.runner_backend.utils import ( # noqa: F401
resolve_decode_backend,
resolve_prefill_backend,
)
python/sglang/srt/model_executor/runner_backend/base_cuda_graph_backend.py
+67
View File
@@ -0,0 +1,67 @@
"""Backend interface for CUDA graph capture/replay."""
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Callable, Iterator, Optional
import torch
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
class BaseCudaGraphBackend(ABC):
"""Pure ABC: no state, no defaults. Each implementation owns its
per-backend state and binds the handles it needs from the
cuda_graph_runner passed to its __init__.
Methods:
- capture_session(stream) — context wrapping the runner's outer
capture loop; backends bind stream / pool and open per-backend
capture flags here.
- capture_one(shape_key, forward_fn, dummies, post_warmup_hook)
— record the replayable artifact for shape_key; one call per
shape inside capture_session.
- can_run(forward_batch, shape_key) — can this backend replay
for the given batch at the given shape.
- replay_session() — context wrapping replay-time model code;
backends open the "we are replaying" flag here when they have
one.
- replay(shape_key, static_forward_batch, **kwargs) — invoke
the captured artifact.
- cleanup() — release pool and drop captured artifacts.
Notes:
- The outer capture loop is runner-specific; it lives on the
runner, not here.
"""
@abstractmethod
def capture_session(self, stream: torch.cuda.Stream) -> Iterator[None]: ...
@abstractmethod
def capture_one(
self,
shape_key: Any,
forward_fn,
dummies: Optional[Any] = None,
post_warmup_hook: Optional[Callable[[], None]] = None,
) -> None: ...
@abstractmethod
def can_run(self, forward_batch: ForwardBatch, shape_key: Any) -> bool: ...
@abstractmethod
def replay_session(self) -> Iterator[None]: ...
@abstractmethod
def replay(
self,
shape_key: Any,
static_forward_batch: ForwardBatch,
**kwargs,
) -> Any: ...
@abstractmethod
def cleanup(self) -> None: ...
python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py
+127
View File
@@ -0,0 +1,127 @@
"""BreakableCudaGraphBackend — segment-captured graphs with eager break
markers (eager_on_graph decorators on attention / mamba layers).
No torch.compile.
"""
from __future__ import annotations
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
import torch
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
set_graph_pool_id,
)
from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
BaseCudaGraphBackend,
)
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
BreakableCUDAGraph,
BreakableCUDAGraphCapture,
eager_on_graph,
enable_breakable_cuda_graph,
)
from sglang.srt.utils import get_bool_env_var, is_hip
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
BaseCudaGraphRunner,
)
class BreakableCudaGraphBackend(BaseCudaGraphBackend):
"""Segmented capture: graphs break at attention / mamba boundaries;
attention metadata is recomputed at replay outside captured segments.
"""
def __init__(
self,
cuda_graph_runner: BaseCudaGraphRunner,
*,
enable_memory_saver: bool = False,
debug_eager: bool = False,
) -> None:
if is_hip():
raise RuntimeError("Breakable CUDA graph is not supported on ROCm/HIP")
self._graphs: Dict[Any, BreakableCUDAGraph] = {}
self._outputs: Dict[Any, Any] = {}
self._pool = None
self._device_module = cuda_graph_runner.device_module
self._tp_group = cuda_graph_runner.model_runner.tp_group
self._capture_stream: Optional[torch.cuda.Stream] = None
self._debug_eager = debug_eager
self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
enable=enable_memory_saver
and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
)
if (
self._memory_saver_adapter is not None
and self._memory_saver_adapter.enabled
):
raise NotImplementedError(
"Breakable CUDA graph is not compatible with memory saver mode"
)
@contextmanager
def capture_session(self, stream: torch.cuda.Stream):
if self._pool is None:
self._pool = self._device_module.graph_pool_handle()
set_graph_pool_id(self._pool)
self._capture_stream = stream
try:
with self.replay_session():
yield
finally:
self._capture_stream = None
def capture_one(
self,
shape_key: Any,
forward_fn: Callable[[], Any],
dummies: Optional[Any] = None,
post_warmup_hook: Optional[Callable[[], None]] = None,
) -> None:
for _ in range(2):
self._device_module.synchronize()
self._tp_group.barrier()
forward_fn()
if post_warmup_hook is not None:
post_warmup_hook()
graph = BreakableCUDAGraph()
captured_fn = (
eager_on_graph(True)(forward_fn) if self._debug_eager else forward_fn
)
with BreakableCUDAGraphCapture(
cuda_graph=graph,
pool=self._pool,
stream=self._capture_stream,
):
out = captured_fn()
self._graphs[shape_key] = graph
self._outputs[shape_key] = out
def can_run(self, forward_batch: ForwardBatch, shape_key: Any) -> bool:
return shape_key in self._graphs
@contextmanager
def replay_session(self):
with enable_breakable_cuda_graph():
yield
def replay(
self,
shape_key: Any,
static_forward_batch: ForwardBatch,
**kwargs,
) -> Any:
self._graphs[shape_key].replay()
return self._outputs[shape_key]
def cleanup(self) -> None:
self._graphs.clear()
self._outputs.clear()
self._pool = None
python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py
+118
View File
@@ -0,0 +1,118 @@
"""FullCudaGraphBackend — captures the entire model forward as one
torch.cuda.CUDAGraph per shape.
"""
from __future__ import annotations
from contextlib import AbstractContextManager, contextmanager
from functools import partial
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
import torch
from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
set_graph_pool_id,
)
from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
BaseCudaGraphBackend,
)
from sglang.srt.utils import get_bool_env_var
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
BaseCudaGraphRunner,
)
class FullCudaGraphBackend(BaseCudaGraphBackend):
"""One torch.cuda.CUDAGraph per shape; attention metadata is
captured inside the graph. Memory-saver-aware.
"""
def __init__(
self,
cuda_graph_runner: BaseCudaGraphRunner,
*,
enable_memory_saver: bool = False,
) -> None:
self._graphs: Dict[Any, torch.cuda.CUDAGraph] = {}
self._outputs: Dict[Any, Any] = {}
self._pool = None
self._device_module = cuda_graph_runner.device_module
self._tp_group = cuda_graph_runner.model_runner.tp_group
self._capture_stream: Optional[torch.cuda.Stream] = None
self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
enable=enable_memory_saver
and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
)
@contextmanager
def capture_session(self, stream: torch.cuda.Stream):
if self._pool is None:
self._pool = self._device_module.graph_pool_handle()
set_graph_pool_id(self._pool)
self._capture_stream = stream
try:
yield
finally:
self._capture_stream = None
def capture_one(
self,
shape_key: Any,
forward_fn: Callable[[], Any],
dummies: Optional[Any] = None,
post_warmup_hook: Optional[Callable[[], None]] = None,
) -> None:
# Two warmups so kernels are loaded and one-time setup is paid before capture.
# post_warmup_hook lets the attention backend reset state that warmup mutated.
for _ in range(2):
self._device_module.synchronize()
self._tp_group.barrier()
forward_fn()
if post_warmup_hook is not None:
post_warmup_hook()
graph = torch.cuda.CUDAGraph()
graph_ctx: Callable[..., AbstractContextManager]
if (
self._memory_saver_adapter is not None
and self._memory_saver_adapter.enabled
):
graph_ctx = partial(
self._memory_saver_adapter.cuda_graph,
tag=GPU_MEMORY_TYPE_CUDA_GRAPH,
)
else:
graph_ctx = self._device_module.graph
with graph_ctx(cuda_graph=graph, pool=self._pool, stream=self._capture_stream):
out = forward_fn()
self._graphs[shape_key] = graph
self._outputs[shape_key] = out
def can_run(self, forward_batch: ForwardBatch, shape_key: Any) -> bool:
return shape_key in self._graphs
@contextmanager
def replay_session(self):
yield
def replay(
self,
shape_key: Any,
static_forward_batch: ForwardBatch,
**kwargs,
) -> Any:
self._graphs[shape_key].replay()
return self._outputs[shape_key]
def cleanup(self) -> None:
self._graphs.clear()
self._outputs.clear()
self._pool = None
python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py
+225
View File
@@ -0,0 +1,225 @@
"""TcPiecewiseCudaGraphBackend — torch.compile-driven piecewise CUDA graph.
FX-splits the model forward at attention layers; per-shape compiled
callables internally capture sub-graphs via
compilation/cuda_piecewise_backend. torch.compile owns the per-shape
cache so this backend has no _graphs table — only a single
_compiled_fn reused for every shape.
"""
from __future__ import annotations
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, Callable, Optional
import torch
import tqdm
from sglang.srt.compilation.compilation_config import CompilationConfig
from sglang.srt.compilation.compile import install_torch_compiled
from sglang.srt.compilation.compile_phase import (
enable_torch_compile_warmup,
set_pcg_capture_stream,
)
from sglang.srt.distributed import get_tensor_model_parallel_rank
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
set_graph_pool_id,
)
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
from sglang.srt.layers.utils import MultiPlatformOp
from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
BaseCudaGraphBackend,
)
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
enable_tc_piecewise_cuda_graph,
)
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
BaseCudaGraphRunner,
)
from sglang.srt.server_args import ServerArgs
_VALID_COMPILERS = ("eager", "inductor")
def _toggle_multi_platform_ops(
model: torch.nn.Module, *, reverse: bool, num_tokens: int
) -> None:
"""Recursively flip MultiPlatformOp submodules into / out of
torch.compile mode."""
for sub in model._modules.values():
if isinstance(sub, MultiPlatformOp):
if reverse:
sub.leave_torch_compile()
else:
sub.enter_torch_compile(num_tokens=num_tokens)
if isinstance(sub, torch.nn.Module):
_toggle_multi_platform_ops(sub, reverse=reverse, num_tokens=num_tokens)
class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
"""torch.compile-driven piecewise capture; attention metadata
recomputed at replay outside the compiled callable's sub-graphs.
"""
def __init__(self, cuda_graph_runner: BaseCudaGraphRunner) -> None:
model_runner = cuda_graph_runner.model_runner
self._pool = None
self._device_module = cuda_graph_runner.device_module
self._tp_group = model_runner.tp_group
self._capture_stream: Optional[torch.cuda.Stream] = None
self._compile_config: CompilationConfig = self.build_compilation_config(
model_runner.server_args
)
self._language_model: torch.nn.Module = getattr(
model_runner.model, "language_model", model_runner.model
)
self._run_compile_pass(cuda_graph_runner)
# model_runner.model.forward is the wrapper that builds LogitsProcessorOutput.
# The compiled trampoline is dispatched internally by it.
self._compiled_fn: Callable = model_runner.model.forward
@staticmethod
def build_compilation_config(server_args: ServerArgs) -> CompilationConfig:
"""Construct a CompilationConfig from ServerArgs and
register the MoE A2A split-op when DeepEP / Mooncake is in use."""
prefill = server_args.cuda_graph_config.prefill
num_tokens = prefill.bs
compiler = prefill.tc_compiler
assert num_tokens is not None, "cuda_graph_config[prefill].bs is not set"
assert compiler in _VALID_COMPILERS, (
f"By now, only {_VALID_COMPILERS} are supported for the "
"tc_piecewise prefill compiler."
)
config = CompilationConfig(
num_tokens,
compiler,
server_args.enable_torch_compile_debug_mode,
)
if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
config.add_split_op("sglang.moe_forward_piecewise_cuda_graph_impl")
return config
@staticmethod
def install_compile(
language_model: Any,
*,
compile_config: CompilationConfig,
graph_pool: Any,
fullgraph: bool = True,
dynamic_arg_dims: Optional[Any] = None,
) -> None:
"""Wrap language_model.model.forward with torch.compile."""
install_torch_compiled(
language_model,
fullgraph=fullgraph,
dynamic_arg_dims=dynamic_arg_dims,
compile_config=compile_config,
graph_pool=graph_pool,
)
def _run_compile_pass(self, cuda_graph_runner: BaseCudaGraphRunner) -> None:
"""JIT-activate kernels at the smallest shape, install
torch.compile, then run one forward per shape inside
enable_torch_compile_warmup to drive FX / inductor through
every shape without capturing cuda graphs yet."""
language_model = self._language_model
compiler = self._compile_config.compiler
with enable_tc_piecewise_cuda_graph():
try:
if compiler != "eager":
_toggle_multi_platform_ops(
language_model.model, reverse=False, num_tokens=16
)
cuda_graph_runner._run_dummy_forward(
num_tokens=cuda_graph_runner.capture_num_tokens[0]
)
if self._pool is None:
self._pool = self._device_module.graph_pool_handle()
set_graph_pool_id(self._pool)
self.install_compile(
language_model.model,
compile_config=self._compile_config,
graph_pool=self._pool,
)
with enable_torch_compile_warmup():
compile_range = (
tqdm.tqdm(list(reversed(cuda_graph_runner.capture_num_tokens)))
if get_tensor_model_parallel_rank() == 0
else reversed(cuda_graph_runner.capture_num_tokens)
)
for num_tokens in compile_range:
if get_tensor_model_parallel_rank() == 0:
compile_range.set_description(
f"Compiling num tokens ({num_tokens=})"
)
cuda_graph_runner._run_dummy_forward(num_tokens=num_tokens)
finally:
_toggle_multi_platform_ops(
language_model.model, reverse=True, num_tokens=16
)
@contextmanager
def capture_session(self, stream: torch.cuda.Stream):
self._capture_stream = stream
try:
with self.replay_session():
with set_pcg_capture_stream(stream):
yield
finally:
self._capture_stream = None
def capture_one(
self,
shape_key: Any,
forward_fn: Callable[[], Any],
dummies: Optional[Any] = None,
post_warmup_hook: Optional[Callable[[], None]] = None,
) -> None:
# Call 1 warms FX state; call 2 captures the cuda graph inside capture_session.
# See cuda_piecewise_backend.py for the FX backend that drives the capture.
for _ in range(2):
self._device_module.synchronize()
self._tp_group.barrier()
forward_fn()
if post_warmup_hook is not None:
post_warmup_hook()
def can_run(self, forward_batch: ForwardBatch, shape_key: Any) -> bool:
# torch.compile manages its per-shape cache internally.
# _run_compile_pass warms every shape in capture_num_tokens at __init__.
return True
@contextmanager
def replay_session(self):
with enable_tc_piecewise_cuda_graph():
yield
def replay(
self,
shape_key: Any,
static_forward_batch: ForwardBatch,
**kwargs,
) -> Any:
return self._compiled_fn(
static_forward_batch.input_ids,
static_forward_batch.positions,
static_forward_batch,
**kwargs,
)
def cleanup(self) -> None:
self._compiled_fn = None
self._compile_config = None
self._language_model = None
self._pool = None
python/sglang/srt/model_executor/runner_backend/utils.py
+96
View File
@@ -0,0 +1,96 @@
"""runner_backend utilities — phase → BaseCudaGraphBackend resolution.
Centralizes per-phase backend resolution so platform overrides (NPU,
out-of-tree) and future backend additions can plug in without
modifying the runner files. Phase / backend identifiers used here
live in :mod:`.cuda_graph_config`.
"""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
from sglang.srt.model_executor.cuda_graph_config import Backend
from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
BaseCudaGraphBackend,
)
from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import (
BreakableCudaGraphBackend,
)
from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import (
FullCudaGraphBackend,
)
from sglang.srt.model_executor.runner_backend.tc_piecewise_cuda_graph_backend import (
TcPiecewiseCudaGraphBackend,
)
if TYPE_CHECKING:
from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
BaseCudaGraphRunner,
)
logger = logging.getLogger(__name__)
# Track first occurrence of each fallback warning to avoid log spam.
_TC_PIECEWISE_DECODE_FALLBACK_LOGGED = False
def resolve_decode_backend(
cuda_graph_runner: BaseCudaGraphRunner,
) -> BaseCudaGraphBackend:
"""Pick a backend instance from cuda_graph_config['decode']['backend'].
NPU device returns NPUCudaGraphBackend regardless of mode (only
the Full-style backend is wired for NPU today).
"""
model_runner = cuda_graph_runner.model_runner
cfg = model_runner.server_args.cuda_graph_config
backend_name = cfg.decode.backend if cfg is not None else Backend.FULL
enable_memory_saver = model_runner.server_args.enable_memory_saver
if model_runner.device == "npu":
from sglang.srt.hardware_backend.npu.graph_runner.npu_cudagraph_backend import (
NPUCudaGraphBackend,
)
return NPUCudaGraphBackend(
cuda_graph_runner, enable_memory_saver=enable_memory_saver
)
if backend_name == Backend.BREAKABLE:
return BreakableCudaGraphBackend(
cuda_graph_runner,
enable_memory_saver=enable_memory_saver,
debug_eager=model_runner.server_args.debug_cuda_graph,
)
if backend_name == Backend.TC_PIECEWISE:
global _TC_PIECEWISE_DECODE_FALLBACK_LOGGED
if not _TC_PIECEWISE_DECODE_FALLBACK_LOGGED:
logger.warning(
"cuda_graph_config decode='tc_piecewise' is not yet implemented; "
"falling back to 'full'."
)
_TC_PIECEWISE_DECODE_FALLBACK_LOGGED = True
return FullCudaGraphBackend(
cuda_graph_runner, enable_memory_saver=enable_memory_saver
)
def resolve_prefill_backend(
cuda_graph_runner: BaseCudaGraphRunner,
) -> BaseCudaGraphBackend:
"""Pick a backend instance from cuda_graph_config['prefill']['backend']."""
model_runner = cuda_graph_runner.model_runner
cfg = model_runner.server_args.cuda_graph_config
backend_name = cfg.prefill.backend if cfg is not None else Backend.TC_PIECEWISE
if backend_name == Backend.BREAKABLE:
return BreakableCudaGraphBackend(
cuda_graph_runner,
enable_memory_saver=model_runner.server_args.enable_memory_saver,
debug_eager=model_runner.server_args.debug_cuda_graph,
)
# Default: tc_piecewise. (prefill, full) is rejected at config validation.
return TcPiecewiseCudaGraphBackend(cuda_graph_runner)
python/sglang/srt/model_executor/runner_backend_utils/__init__.py
+22
View File
@@ -0,0 +1,22 @@
"""Low-level primitives used by the CUDA graph backends.
Subpackages:
- breakable_cuda_graph: BreakableCUDAGraph + capture context,
eager_on_graph decorator, is_in_breakable_cuda_graph flag.
- piecewise_cuda_graph: shared piecewise context manager
(set_tc_piecewise_forward_context, is_in_tc_piecewise_cuda_graph).
Backends in cuda_graph_backend/ import from here. Runners do not.
"""
# Generic failure-message hint for non-piecewise CUDA graph capture
# paths (Full backend used by decode + EAGLE draft runners). The
# piecewise-specific variant lives in
# piecewise_cuda_graph.context_manager and points users at
# --disable-piecewise-cuda-graph, which doesn't apply here.
CUDA_GRAPH_CAPTURE_FAILED_MSG = (
"CUDA graph capture failed.\n"
"To work around this error, add --disable-cuda-graph to your launch command\n"
"(or use --disable-decode-cuda-graph to disable only the decode phase).\n"
"Please report this issue at https://github.com/sgl-project/sglang/issues/new/choose"
)
python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/__init__.py
+21
View File
@@ -0,0 +1,21 @@
"""Breakable primitives — segmented CUDA graph capture with eager break points.
Public API (also reachable via the deeper module paths):
- BreakableCUDAGraph, BreakableCUDAGraphCapture — capture/replay
- eager_on_graph — decorator that marks a callable as a graph break
- enable_breakable_cuda_graph — context that flips the Breakable runtime flag
- is_in_breakable_cuda_graph — runtime flag getter
The legacy model_executor/breakable_cuda_graph/ package is a
backwards-compat shim that re-exports from here.
"""
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import ( # noqa: F401
BreakableCUDAGraph,
BreakableCUDAGraphCapture,
eager_on_graph,
)
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( # noqa: F401
enable_breakable_cuda_graph,
is_in_breakable_cuda_graph,
)
python/sglang/srt/model_executor/breakable_cuda_graph/breakable_cuda_graph.py → python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py
+14 -12
View File
@@ -12,12 +12,12 @@
# limitations under the License.
# ==============================================================================
"""Breakable CUDA Graph: capture a region as a sequence of
``torch.cuda.CUDAGraph`` segments separated by eager break points.
torch.cuda.CUDAGraph segments separated by eager break points.
Each segment is a real ``torch.cuda.CUDAGraph``. Its destructor calls
``releasePool`` on the shared mempool, so the pool's ``use_count`` tracks how
Each segment is a real torch.cuda.CUDAGraph. Its destructor calls
releasePool on the shared mempool, so the pool's use_count tracks how
many segments are alive; the pool stays pinned as long as any segment graph
is alive. This lets ``weak_ref_tensor`` views of intermediate pool-allocated
is alive. This lets weak_ref_tensor views of intermediate pool-allocated
tensors remain valid across replays — we don't need Python-managed bridge
buffers to keep break-point tensors at stable addresses.
"""
@@ -34,7 +34,9 @@ try:
except ImportError:
rt = None
from sglang.srt.model_executor.breakable_cuda_graph.cuda_utils import checkCudaErrors
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.cuda_utils import (
checkCudaErrors,
)
logger = logging.getLogger(__name__)
@@ -153,9 +155,9 @@ def _weak_ref_if_tensor(x):
mempool reclaim per-layer intermediates between segments — storage stays
alive for each segment CUDAGraph's lifetime via its pool use_count.
``weak_ref_tensors`` is imported lazily: the module hard-raises on
weak_ref_tensors is imported lazily: the module hard-raises on
non-CUDA/NPU platforms, and we only reach this code during an active
BCG capture (which can't happen on CPU-only runners anyway)."""
Breakable capture (which can't happen on CPU-only runners anyway)."""
if torch.is_tensor(x):
from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors
@@ -238,7 +240,7 @@ def eager_on_graph(enable: bool):
class BreakableCUDAGraph:
"""Container holding one ``torch.cuda.CUDAGraph`` per segment plus an
"""Container holding one torch.cuda.CUDAGraph per segment plus an
eager break function between consecutive segments."""
def __init__(self) -> None:
@@ -259,12 +261,12 @@ class BreakableCUDAGraph:
class BreakableCUDAGraphCapture:
"""Context manager that captures the enclosed code as one or more
``torch.cuda.CUDAGraph`` segments separated by eager break points.
torch.cuda.CUDAGraph segments separated by eager break points.
Each segment shares the supplied ``pool`` (``MempoolId_t`` tuple) so
Each segment shares the supplied pool (MempoolId_t tuple) so
pool-allocated intermediates can be reused across segments. While any
segment is alive, its ``beginAllocateToPool`` call keeps the mempool's
``use_count`` > 0, which makes ``weak_ref_tensor`` of segment-allocated
segment is alive, its beginAllocateToPool call keeps the mempool's
use_count > 0, which makes weak_ref_tensor of segment-allocated
tensors safe across subsequent replays.
"""
python/sglang/srt/model_executor/breakable_cuda_graph/context.py → python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/context.py
+1 -6
View File
@@ -11,12 +11,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Runtime state for the breakable CUDA graph (BCG) runner.
Kept intentionally separate from ``compilation/piecewise_context_manager.py``:
BCG no longer inherits from the torch.compile-based PCG path, so its
capture/replay lifecycle is managed on its own.
"""
"""Runtime state for the breakable CUDA graph runner."""
from __future__ import annotations
python/sglang/srt/model_executor/breakable_cuda_graph/cuda_utils.py → python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py
View File
python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/__init__.py
+21
View File
@@ -0,0 +1,21 @@
"""Piecewise CUDA graph utilities — shared between Breakable and tc_piecewise backends.
Public API:
- is_in_tc_piecewise_cuda_graph() — true while inside any piecewise capture.
- enable_tc_piecewise_cuda_graph() — context manager that toggles the flag.
- TcPiecewiseForwardContext + set_tc_piecewise_forward_context + get_tc_piecewise_forward_context.
- TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG.
The torch.compile-warmup flag (is_in_torch_compile_warmup) lives in
sglang.srt.compilation.compile_phase — it is torch.compile-internal,
not piecewise-shared.
"""
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph.context_manager import ( # noqa: F401
TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG,
TcPiecewiseForwardContext,
enable_tc_piecewise_cuda_graph,
get_tc_piecewise_forward_context,
is_in_tc_piecewise_cuda_graph,
set_tc_piecewise_forward_context,
)
python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py
+112
View File
@@ -0,0 +1,112 @@
"""CUDA graph capture context manager + forward-context propagation.
Owns two pieces of cross-cutting state used by *every* piecewise-style
backend (currently breakable + tc_piecewise):
* _in_tc_piecewise_cuda_graph — a process-global flag set true while we
are inside the capture or replay window of a piecewise CUDA graph.
Read by model code that needs to take the static-buffer / fixed-shape
branch. See refactor/plan.md §6.5 for the full semantics.
* TcPiecewiseForwardContext — a dataclass propagated across attention/MoE
layers during capture and replay so that submodules can reach the
current ForwardBatch and per-layer metadata without threading
arguments through every call site. Named TcPiecewise… (matches
Backend.TC_PIECEWISE + enable_tc_piecewise_cuda_graph) to
disambiguate from the per-forward-call
sglang.srt.model_executor.forward_context.ForwardContext.
This module deliberately does **not** own torch.compile-specific state
(warmup flag, capture stream); those live in compilation/compile_phase.py.
"""
from __future__ import annotations
import logging
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, List, Optional
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
_in_tc_piecewise_cuda_graph = False
def is_in_tc_piecewise_cuda_graph() -> bool:
"""True while inside tc_piecewise CUDA graph capture/replay."""
return _in_tc_piecewise_cuda_graph
@contextmanager
def enable_tc_piecewise_cuda_graph():
"""Mark the enclosed scope as "we are inside a piecewise CUDA graph
capture/replay". Sets _in_tc_piecewise_cuda_graph true for the duration.
Errors during capture surface a hint that lets users disable the
feature while filing a bug.
"""
global _in_tc_piecewise_cuda_graph
_in_tc_piecewise_cuda_graph = True
try:
yield
except Exception as e:
logger.error(
"Piecewise CUDA Graph failed with error: %s\n%s",
e,
TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG,
)
raise
finally:
_in_tc_piecewise_cuda_graph = False
@dataclass
class TcPiecewiseForwardContext:
forward_batch: Optional[ForwardBatch] = None
attention_layers: Optional[List[Any]] = field(default=None)
quant_config: Any = None
moe_layers: Optional[List[Any]] = field(default=None)
moe_fusions: Optional[List[Any]] = field(default=None)
dsa_indexers: Optional[List[Any]] = field(default=None)
_tc_piecewise_forward_context: Optional[TcPiecewiseForwardContext] = None
def get_tc_piecewise_forward_context() -> Optional[TcPiecewiseForwardContext]:
return _tc_piecewise_forward_context
@contextmanager
def set_tc_piecewise_forward_context(
forward_batch: ForwardBatch,
attention_layers: List[Any],
quant_config: Any,
moe_layers: List[Any],
moe_fusions: List[Any],
dsa_indexers: Optional[List[Any]] = None,
):
global _tc_piecewise_forward_context
_tc_piecewise_forward_context = TcPiecewiseForwardContext(
forward_batch=forward_batch,
attention_layers=attention_layers,
quant_config=quant_config,
moe_layers=moe_layers,
moe_fusions=moe_fusions,
dsa_indexers=dsa_indexers,
)
try:
yield
finally:
_tc_piecewise_forward_context = None
TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG = (
"Piecewise CUDA Graph is enabled by default as an experimental feature.\n"
"To work around this error, add --disable-piecewise-cuda-graph to your launch command.\n"
"Please report this issue at https://github.com/sgl-project/sglang/issues/new/choose"
)
python/sglang/srt/model_executor/runner_utils/__init__.py
+28
View File
@@ -0,0 +1,28 @@
"""Low-level utilities used by the CUDA graph runners.
Mirror of cuda_graph_backend_utils/ for runner-side state — buffer
dataclasses, process-global capture flags, the speculative-shared
graph memory pool, and the DeepEP capture/replay adapter. Runners in
cuda_graph_runner/ import from here; nothing here should import
back into cuda_graph_runner/.
"""
from sglang.srt.model_executor.runner_utils.buffers import ( # noqa: F401
DecodeInputBuffers,
PrefillInputBuffers,
_grouped_foreach_copy_,
)
from sglang.srt.model_executor.runner_utils.capture_mode import ( # noqa: F401
_set_capture_lora_variant,
compile_in_capture_mode,
get_capture_lora_variant,
get_is_capture_mode,
model_capture_mode,
)
from sglang.srt.model_executor.runner_utils.deepep_adapter import ( # noqa: F401
DeepEPCudaGraphRunnerAdapter,
)
from sglang.srt.model_executor.runner_utils.pool import ( # noqa: F401
get_global_graph_memory_pool,
set_global_graph_memory_pool,
)
python/sglang/srt/model_executor/runner_utils/buffers.py
+442
View File
@@ -0,0 +1,442 @@
"""Static-buffer dataclasses used by the CUDA graph runners.
DecodeInputBuffers backs the decode-phase capture/replay path.
PrefillInputBuffers backs the prefill-phase capture/replay path.
Both subclass ForwardInputBuffers so that buffer-pool sharing works
the same way as for non-cuda-graph forward paths.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
import torch
from sglang.srt.environ import envs
from sglang.srt.model_executor.forward_batch_info import (
ForwardBatch,
NgramEmbeddingInfo,
PPProxyTensors,
compute_local_num_token_non_padded,
)
from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
_has_foreach_copy = hasattr(torch, "_foreach_copy_")
def _grouped_foreach_copy_(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:
"""Call torch._foreach_copy_ grouped by (dst_dtype, src_dtype) pairs."""
def foreach_copy(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:
if _has_foreach_copy:
torch._foreach_copy_(dsts, srcs)
else:
for dst, src in zip(dsts, srcs):
dst.copy_(src)
groups: Dict[Tuple[torch.dtype, torch.dtype], Tuple[List, List]] = {}
for dst, src in zip(dsts, srcs):
key = (dst.dtype, src.dtype)
if key not in groups:
groups[key] = ([], [])
groups[key][0].append(dst)
groups[key][1].append(src)
for group_dsts, group_srcs in groups.values():
foreach_copy(group_dsts, group_srcs)
@dataclass
class DecodeInputBuffers(ForwardInputBuffers):
input_ids: torch.Tensor
input_embeds: torch.Tensor
req_pool_indices: torch.Tensor
seq_lens: torch.Tensor
seq_lens_cpu: torch.Tensor
out_cache_loc: torch.Tensor
out_cache_loc_swa: Optional[torch.Tensor]
positions: torch.Tensor
mrope_positions: torch.Tensor
num_token_non_padded: torch.Tensor
custom_mask: torch.Tensor
next_token_logits_buffer: torch.Tensor
mamba_track_indices: Optional[torch.Tensor]
mamba_track_mask: Optional[torch.Tensor]
global_num_tokens_gpu: torch.Tensor
global_num_tokens_for_logprob_gpu: torch.Tensor
encoder_lens: Optional[torch.Tensor]
pp_proxy_tensors: Optional[Dict[str, torch.Tensor]]
ngram_embedding_info: Optional["NgramEmbeddingInfo"]
rids_int: Optional[torch.Tensor]
bootstrap_room_ids_int: Optional[torch.Tensor]
@classmethod
def create(
cls,
*,
device: torch.device,
max_bs: int,
max_num_token: int,
hidden_size: int,
vocab_size: int,
dtype: torch.dtype,
dp_size: int,
pp_size: int,
is_encoder_decoder: bool,
require_mlp_tp_gather: bool,
seq_len_fill_value: int,
encoder_len_fill_value: int,
num_tokens_per_bs: int,
cache_loc_dtype: torch.dtype,
enable_mamba_track: bool,
ne_token_table: Optional[torch.Tensor] = None,
is_hybrid_swa: bool = False,
hc_hidden_size: Optional[int] = None,
) -> "DecodeInputBuffers":
with torch.device(device):
input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
input_embeds = torch.zeros((max_num_token, hidden_size), dtype=dtype)
req_pool_indices = torch.zeros((max_bs,), dtype=torch.int64)
seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int32)
out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype)
out_cache_loc_swa = (
torch.zeros((max_num_token,), dtype=torch.int64)
if is_hybrid_swa
else None
)
positions = torch.zeros((max_num_token,), dtype=torch.int64)
mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64)
num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
custom_mask = torch.ones(
(max_bs * seq_len_fill_value + max_num_token) * num_tokens_per_bs,
dtype=torch.bool,
)
next_token_logits_buffer = torch.zeros(
(max_num_token, vocab_size),
dtype=torch.float,
)
mamba_track_indices = (
torch.zeros((max_bs,), dtype=torch.int64)
if enable_mamba_track
else None
)
mamba_track_mask = (
torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None
)
if pp_size > 1:
is_mhc = hc_hidden_size is not None
hs = hc_hidden_size if is_mhc else hidden_size
pp_proxy_tensors = {
"hidden_states": torch.zeros((max_bs, hs), dtype=dtype),
}
if not is_mhc:
pp_proxy_tensors["residual"] = torch.zeros(
(max_bs, hidden_size), dtype=dtype
)
else:
pp_proxy_tensors = None
if is_encoder_decoder:
encoder_lens = torch.full(
(max_bs,), encoder_len_fill_value, dtype=torch.int32
)
else:
encoder_lens = None
if require_mlp_tp_gather:
global_num_tokens_gpu = torch.zeros((dp_size,), dtype=torch.int32)
global_num_tokens_for_logprob_gpu = torch.zeros(
(dp_size,), dtype=torch.int32
)
else:
global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
global_num_tokens_for_logprob_gpu = torch.zeros((1,), dtype=torch.int32)
ngram_embedding_info = (
NgramEmbeddingInfo(
token_table=ne_token_table,
column_starts=torch.zeros([max_bs], dtype=torch.int32),
req_lens=torch.ones([max_bs], dtype=torch.int32),
out_column_starts=torch.zeros([max_bs], dtype=torch.int32),
out_req_lens=torch.ones([max_bs], dtype=torch.int32),
)
if ne_token_table is not None
else None
)
if envs.SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE.get():
rids_int = torch.zeros((max_bs,), dtype=torch.int64)
bootstrap_room_ids_int = torch.full((max_bs,), -1, dtype=torch.int64)
else:
rids_int = None
bootstrap_room_ids_int = None
seq_lens_cpu = torch.full(
(max_bs,),
seq_len_fill_value,
dtype=torch.int32,
device="cpu",
)
return cls(
input_ids=input_ids,
input_embeds=input_embeds,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
seq_lens_cpu=seq_lens_cpu,
out_cache_loc=out_cache_loc,
out_cache_loc_swa=out_cache_loc_swa,
positions=positions,
mrope_positions=mrope_positions,
num_token_non_padded=num_token_non_padded,
custom_mask=custom_mask,
next_token_logits_buffer=next_token_logits_buffer,
mamba_track_indices=mamba_track_indices,
mamba_track_mask=mamba_track_mask,
encoder_lens=encoder_lens,
global_num_tokens_gpu=global_num_tokens_gpu,
global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
pp_proxy_tensors=pp_proxy_tensors,
ngram_embedding_info=ngram_embedding_info,
rids_int=rids_int,
bootstrap_room_ids_int=bootstrap_room_ids_int,
)
def populate_from_forward_batch(
self,
*,
forward_batch: ForwardBatch,
raw_bs: int,
raw_num_token: int,
bs: int,
seq_len_fill_value: int,
require_gathered_buffer: bool,
num_tokens_per_bs: int,
dsa_enable_prefill_cp: bool,
enable_num_token_non_padded_flag: bool,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
):
if bs != raw_bs:
self.seq_lens.fill_(seq_len_fill_value)
self.out_cache_loc.zero_()
if self.mamba_track_indices is not None:
self.mamba_track_indices.zero_()
if self.mamba_track_mask is not None:
self.mamba_track_mask.fill_(False)
# Build batched copy lists for all GPU tensors.
dsts = [
self.input_ids[:raw_num_token],
self.req_pool_indices[:raw_bs],
self.seq_lens[:raw_bs],
self.out_cache_loc[:raw_num_token],
self.positions[:raw_num_token],
]
srcs = [
forward_batch.input_ids,
forward_batch.req_pool_indices,
forward_batch.seq_lens,
forward_batch.out_cache_loc,
forward_batch.positions,
]
if self.ngram_embedding_info is not None:
ngram_embedding_info = forward_batch.ngram_embedding_info
self.ngram_embedding_info.column_starts[:raw_bs].copy_(
ngram_embedding_info.column_starts
)
self.ngram_embedding_info.req_lens[:raw_bs].copy_(
ngram_embedding_info.req_lens
)
if (
self.mamba_track_indices is not None
and forward_batch.mamba_track_indices is not None
):
dsts.append(self.mamba_track_indices[:raw_bs])
srcs.append(forward_batch.mamba_track_indices)
if (
self.mamba_track_mask is not None
and forward_batch.mamba_track_mask is not None
):
dsts.append(self.mamba_track_mask[:raw_bs])
srcs.append(forward_batch.mamba_track_mask)
if self.encoder_lens is not None and forward_batch.encoder_lens is not None:
dsts.append(self.encoder_lens[:raw_bs])
srcs.append(forward_batch.encoder_lens)
if forward_batch.mrope_positions is not None:
dsts.append(self.mrope_positions[:, :raw_num_token])
srcs.append(forward_batch.mrope_positions)
if self.rids_int is not None and forward_batch.rids_int is not None:
dsts.append(self.rids_int[:raw_bs])
srcs.append(forward_batch.rids_int)
if (
self.bootstrap_room_ids_int is not None
and forward_batch.bootstrap_room_ids_int is not None
):
dsts.append(self.bootstrap_room_ids_int[:raw_bs])
srcs.append(forward_batch.bootstrap_room_ids_int)
if require_gathered_buffer:
self.global_num_tokens_gpu.fill_(bs * num_tokens_per_bs)
self.global_num_tokens_for_logprob_gpu.fill_(bs * num_tokens_per_bs)
if enable_num_token_non_padded_flag:
if require_gathered_buffer and not dsa_enable_prefill_cp:
num_tokens_per_dp = bs * num_tokens_per_bs
local = compute_local_num_token_non_padded(
global_num_token_non_padded=forward_batch.num_token_non_padded,
num_tokens_per_dp=num_tokens_per_dp,
)
dsts.append(self.num_token_non_padded)
srcs.append(local)
else:
dsts.append(self.num_token_non_padded)
srcs.append(forward_batch.num_token_non_padded)
# Pipeline-parallel proxy tensors.
if pp_proxy_tensors is not None and self.pp_proxy_tensors is not None:
for key, buf in self.pp_proxy_tensors.items():
src = pp_proxy_tensors.tensors[key]
dim = src.shape[0]
dsts.append(buf[:dim])
srcs.append(src)
# SWA cache location (int32, separate from the int64 batch above).
if (
self.out_cache_loc_swa is not None
and forward_batch.out_cache_loc_swa is not None
):
dsts.append(self.out_cache_loc_swa[:raw_num_token])
srcs.append(forward_batch.out_cache_loc_swa[:raw_num_token])
# Batch all GPU copies, grouped by dtype pair.
_grouped_foreach_copy_(dsts, srcs)
if forward_batch.seq_lens_cpu is not None:
if bs != raw_bs:
self.seq_lens_cpu.fill_(seq_len_fill_value)
self.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
@dataclass
class PrefillInputBuffers(ForwardInputBuffers):
input_ids: torch.Tensor
out_cache_loc: torch.Tensor
out_cache_loc_swa: Optional[torch.Tensor]
mamba_track_indices: Optional[torch.Tensor]
mamba_track_mask: Optional[torch.Tensor]
mamba_track_seqlens: Optional[torch.Tensor]
positions: torch.Tensor
input_embeds: Optional[torch.Tensor]
mrope_positions: Optional[torch.Tensor]
@classmethod
def create(
cls,
*,
device: torch.device,
max_bs: int,
max_num_tokens: int,
cache_loc_dtype: torch.dtype,
is_hybrid_swa: bool,
is_multimodal: bool,
hidden_size: int,
dtype: torch.dtype,
enable_mamba_track: bool,
) -> "PrefillInputBuffers":
with torch.device(device):
input_ids = torch.zeros((max_num_tokens,), dtype=torch.int64)
out_cache_loc = torch.zeros((max_num_tokens,), dtype=cache_loc_dtype)
out_cache_loc_swa = (
torch.zeros((max_num_tokens,), dtype=torch.int64)
if is_hybrid_swa
else None
)
mamba_track_indices = (
torch.zeros((max_bs,), dtype=torch.int64)
if enable_mamba_track
else None
)
mamba_track_mask = (
torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None
)
mamba_track_seqlens = (
torch.zeros((max_bs,), dtype=torch.int32)
if enable_mamba_track
else None
)
positions = torch.zeros((max_num_tokens,), dtype=torch.int64)
if is_multimodal:
input_embeds = torch.zeros((max_num_tokens, hidden_size), dtype=dtype)
mrope_positions = torch.zeros((3, max_num_tokens), dtype=torch.int64)
else:
input_embeds = None
mrope_positions = None
return cls(
input_ids=input_ids,
out_cache_loc=out_cache_loc,
out_cache_loc_swa=out_cache_loc_swa,
mamba_track_indices=mamba_track_indices,
mamba_track_mask=mamba_track_mask,
mamba_track_seqlens=mamba_track_seqlens,
positions=positions,
input_embeds=input_embeds,
mrope_positions=mrope_positions,
)
def populate_from_forward_batch(
self,
*,
forward_batch: ForwardBatch,
raw_num_tokens: int,
static_num_tokens: int,
is_multimodal: bool,
) -> None:
"""Copy serving-batch values into static buffers and zero out
the padding region between raw_num_tokens and
static_num_tokens.
"""
if static_num_tokens != raw_num_tokens:
self.out_cache_loc.zero_()
self.input_ids[raw_num_tokens:static_num_tokens].zero_()
self.positions[raw_num_tokens:static_num_tokens].zero_()
if is_multimodal:
self.input_embeds[raw_num_tokens:static_num_tokens].zero_()
if forward_batch.mrope_positions is not None:
self.mrope_positions[:, raw_num_tokens:static_num_tokens].zero_()
bs = forward_batch.batch_size
self.input_ids[:raw_num_tokens].copy_(forward_batch.input_ids)
self.positions[:raw_num_tokens].copy_(forward_batch.positions)
self.out_cache_loc[:raw_num_tokens].copy_(forward_batch.out_cache_loc)
if (
self.mamba_track_indices is not None
and forward_batch.mamba_track_indices is not None
):
self.mamba_track_indices[:bs].copy_(forward_batch.mamba_track_indices)
if (
self.mamba_track_mask is not None
and forward_batch.mamba_track_mask is not None
):
self.mamba_track_mask[:bs].copy_(forward_batch.mamba_track_mask)
if (
self.mamba_track_seqlens is not None
and forward_batch.mamba_track_seqlens is not None
):
self.mamba_track_seqlens[:bs].copy_(forward_batch.mamba_track_seqlens)
if forward_batch.mrope_positions is not None:
self.mrope_positions[:, :raw_num_tokens].copy_(
forward_batch.mrope_positions
)
python/sglang/srt/model_executor/runner_utils/capture_mode.py
+56
View File
@@ -0,0 +1,56 @@
"""Process-global capture-mode flags shared by the decode runner and the
speculative-draft runners. Read by model code that needs to take a
capture-time branch (e.g. lora dual-graph capture decides per-batch
which variant to use).
"""
from __future__ import annotations
from contextlib import contextmanager
from typing import Optional
import torch
# Detect whether the current forward pass is in capture mode.
is_capture_mode = False
# When capturing dual MoE backends, tracks which variant is being captured.
# None = not dual, "lora" = capturing lora variant, "nolora" = capturing nolora variant.
_capture_lora_variant: Optional[str] = None
def get_is_capture_mode() -> bool:
return is_capture_mode
def compile_in_capture_mode(func):
"""Decorator: wrap func with torch.compile only when defined
inside model capture mode; passthrough otherwise.
Used by model code (e.g. DeepSeek-V4) to opt nested helpers into
torch.compile during cuda-graph capture without paying the
compilation cost in the eager forward path.
"""
if is_capture_mode:
return torch.compile(func)
return func
def get_capture_lora_variant() -> Optional[str]:
"""Return the lora variant being captured, or None if not in dual capture."""
return _capture_lora_variant
def _set_capture_lora_variant(variant: Optional[str]) -> None:
global _capture_lora_variant
_capture_lora_variant = variant
@contextmanager
def model_capture_mode():
global is_capture_mode
is_capture_mode = True
try:
yield
finally:
is_capture_mode = False
python/sglang/srt/model_executor/runner_utils/deepep_adapter.py
+29
View File
@@ -0,0 +1,29 @@
"""DeepEP capture/replay adapter — records the dispatch mode used during
capture and re-applies it during replay so DeepEP all-to-all has
consistent expert routing across the captured graph.
"""
from __future__ import annotations
from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPBuffer
from sglang.srt.layers.moe.utils import get_deepep_mode, get_moe_a2a_backend
class DeepEPCudaGraphRunnerAdapter:
def __init__(self) -> None:
# Record DeepEP mode used during capture to ensure replay consistency.
self._captured_deepep_mode = None
def capture(self, is_extend_in_batch: bool) -> None:
if not get_moe_a2a_backend().is_deepep():
return
self._captured_deepep_mode = get_deepep_mode().resolve(
is_extend_in_batch=is_extend_in_batch
)
DeepEPBuffer.set_dispatch_mode(self._captured_deepep_mode)
def replay(self) -> None:
if not get_moe_a2a_backend().is_deepep():
return
assert self._captured_deepep_mode is not None
DeepEPBuffer.set_dispatch_mode(self._captured_deepep_mode)
python/sglang/srt/model_executor/runner_utils/pool.py
+21
View File
@@ -0,0 +1,21 @@
"""Shared graph memory pool used by the speculative-draft cuda graph
runners. The new DecodeCudaGraphRunner and PrefillCudaGraphRunner
backends each own their pool internally; this global is retained for the
EAGLE / multi-step draft runners that haven't been folded into the new
backend interface.
"""
from __future__ import annotations
from typing import Any, Optional
_global_graph_memory_pool: Optional[Any] = None
def get_global_graph_memory_pool() -> Optional[Any]:
return _global_graph_memory_pool
def set_global_graph_memory_pool(val: Any) -> None:
global _global_graph_memory_pool
_global_graph_memory_pool = val
python/sglang/srt/models/bailing_moe.py
+1 -1
View File
@@ -73,8 +73,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.utils import (
apply_qk_norm,
python/sglang/srt/models/bailing_moe_linear.py
+1 -1
View File
@@ -57,8 +57,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA, DeepseekV2MLP, _is_hip
from sglang.srt.models.utils import WeightsMapper
python/sglang/srt/models/cohere2_moe.py
+1 -1
View File
@@ -29,8 +29,8 @@ from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.rotary_embedding import get_rope
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.utils import add_prefix, get_compiler_backend, is_cuda, make_layers
python/sglang/srt/models/deepseek_common/attention_backend_handler.py
+6 -4
View File
@@ -1,7 +1,9 @@
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.models.deepseek_common.attention_forward_methods.forward_methods import (
AttnForwardMethod,
)
@@ -72,7 +74,7 @@ def _support_mha_one_shot(attn, forward_batch, backend_name):
def _handle_attention_backend(attn, forward_batch, backend_name):
if is_in_piecewise_cuda_graph():
if is_in_tc_piecewise_cuda_graph():
return AttnForwardMethod.MLA
# MLA prefill CP forces absorbed MLA regardless of prefix length: the
@@ -130,7 +132,7 @@ def handle_attention_fa4(attn, forward_batch):
def handle_attention_trtllm_mla(attn, forward_batch):
if is_in_piecewise_cuda_graph():
if is_in_tc_piecewise_cuda_graph():
return AttnForwardMethod.MLA
sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch)
@@ -170,7 +172,7 @@ def handle_attention_dsa(attn, forward_batch):
def handle_attention_triton(attn, forward_batch):
if is_in_piecewise_cuda_graph():
if is_in_tc_piecewise_cuda_graph():
return AttnForwardMethod.MLA
# when deterministic inference is enabled, use MLA
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
+5 -3
View File
@@ -4,7 +4,6 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
from sglang.srt.environ import envs
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
@@ -29,6 +28,9 @@ from sglang.srt.model_executor.forward_context import (
get_attn_backend,
get_token_to_kv_pool,
)
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.models.deepseek_common.utils import (
FORWARD_ABSORB_CORE_ATTENTION_BACKENDS,
_is_cpu,
@@ -147,7 +149,7 @@ class DeepseekMLAForwardMixin:
llama_4_scaling: Optional[torch.Tensor] = None,
prev_topk_indices: Optional[torch.Tensor] = None,
):
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.runner import get_is_capture_mode
q_lora = None
topk_indices = None
@@ -709,7 +711,7 @@ class DeepseekMLAForwardMixin:
)
attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
else:
if is_in_piecewise_cuda_graph():
if is_in_tc_piecewise_cuda_graph():
# torch dynamo requires out= op was called where output tensor was non-contiguous
attn_bmm_output = (
torch.bmm(attn_output.transpose(0, 1), self.w_vc)
python/sglang/srt/models/deepseek_v2.py
+8 -3
View File
@@ -135,8 +135,13 @@ from sglang.srt.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
get_embedding_tp_kwargs,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.models.deepseek_common.attention_backend_handler import (
AttentionBackendRegistry,
)
@@ -724,7 +729,7 @@ class DeepseekV2MoE(nn.Module):
ModelOptFp4LinearMethod,
)
and fc1_n % 128 == 0
and get_global_server_args().disable_piecewise_cuda_graph
and not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
):
self.shared_experts.gate_up_proj._interleave_for_swiglu_fusion = True
self.shared_experts._enable_nvfp4_gemm_swiglu_fusion = True
@@ -2439,7 +2444,7 @@ class DeepseekV2Model(nn.Module):
# NOTE: torch dynamo does not support graph break in context manager
ctx = (
nullcontext()
if not get_global_server_args().disable_piecewise_cuda_graph
if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
else get_global_expert_distribution_recorder().with_current_layer(i)
)
with ctx:
python/sglang/srt/models/deepseek_v4.py
+19 -12
View File
@@ -28,7 +28,6 @@ from sglang.jit_kernel.dsv4 import (
fused_rope_inplace,
)
from sglang.srt.compilation.compilation_config import register_split_op
from sglang.srt.compilation.piecewise_context_manager import get_forward_context
from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
from sglang.srt.distributed import (
get_pp_group,
@@ -84,21 +83,29 @@ from sglang.srt.layers.utils.cp_utils import (
)
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
from sglang.srt.mem_cache.memory_pool import RadixAttention
from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import (
eager_on_graph,
)
from sglang.srt.model_executor.breakable_cuda_graph.context import (
is_in_breakable_cuda_graph,
)
from sglang.srt.model_executor.cuda_graph_runner import (
compile_in_capture_mode,
get_is_capture_mode,
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
)
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
from sglang.srt.model_executor.forward_context import (
get_attn_backend,
get_token_to_kv_pool,
)
from sglang.srt.model_executor.runner import (
compile_in_capture_mode,
get_is_capture_mode,
)
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
eager_on_graph,
)
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
is_in_breakable_cuda_graph,
)
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
get_tc_piecewise_forward_context,
)
from sglang.srt.model_loader.utils import maybe_executor_submit, should_async_load
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.dbrx import ReplicatedLinear
@@ -211,7 +218,7 @@ def deepseek_v4_attention_with_output(
attn_sink: torch.Tensor,
save_kv_cache: bool,
) -> None:
context = get_forward_context()
context = get_tc_piecewise_forward_context()
forward_batch = context.forward_batch
attention_layers = context.attention_layers
attention_layer = attention_layers[layer_id]
@@ -1678,7 +1685,7 @@ class DeepseekV4Model(nn.Module):
last_layer = layer
ctx = (
nullcontext()
if not get_global_server_args().disable_piecewise_cuda_graph
if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
else get_global_expert_distribution_recorder().with_current_layer(i)
)
with ctx:
python/sglang/srt/models/exaone_moe.py
+1 -1
View File
@@ -63,8 +63,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_layers
python/sglang/srt/models/glm4_moe.py
+1 -1
View File
@@ -82,8 +82,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
from sglang.srt.models.utils import apply_qk_norm
python/sglang/srt/models/glm4_moe_lite.py
+1 -1
View File
@@ -68,8 +68,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
DeepseekV2WeightLoaderMixin,
python/sglang/srt/models/gpt_oss.py
+6 -6
View File
@@ -26,10 +26,6 @@ import torch
from torch import nn
from transformers import PretrainedConfig
from sglang.srt.compilation.piecewise_context_manager import (
get_forward_context,
is_in_piecewise_cuda_graph,
)
from sglang.srt.distributed import (
get_moe_expert_parallel_rank,
get_moe_expert_parallel_world_size,
@@ -70,6 +66,10 @@ from sglang.srt.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
get_tc_piecewise_forward_context,
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.utils import (
create_fused_set_kv_buffer_arg,
@@ -299,7 +299,7 @@ class GptOssSparseMoeBlock(nn.Module):
else:
router_input = hidden_states
if is_in_piecewise_cuda_graph():
if is_in_tc_piecewise_cuda_graph():
final_hidden_states = moe_impl(self.layer_id, hidden_states)
else:
router_logits, _ = self.router(router_input)
@@ -326,7 +326,7 @@ class GptOssSparseMoeBlock(nn.Module):
@register_custom_op(out_shape="hidden_states")
def moe_impl(layer_id: int, hidden_states: torch.Tensor) -> torch.Tensor:
forward_context = get_forward_context()
forward_context = get_tc_piecewise_forward_context()
moe_fusion = forward_context.moe_fusions[layer_id]
router_logits, _ = moe_fusion.router(hidden_states)
topk_output = moe_fusion.topk(hidden_states, router_logits)
python/sglang/srt/models/grok.py
+1 -1
View File
@@ -56,8 +56,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.loader import DefaultModelLoader
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.utils import add_prefix, is_npu
python/sglang/srt/models/hunyuan_v3.py
+1 -1
View File
@@ -45,7 +45,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from sglang.srt.managers.schedule_batch import ForwardBatch
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.utils import is_cuda
from sglang.srt.utils.hf_transformers_utils import get_rope_config
python/sglang/srt/models/kimi_linear.py
+1 -1
View File
@@ -39,8 +39,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import (
default_weight_loader,
maybe_remap_kv_scale_name,
python/sglang/srt/models/llada2.py
+1 -1
View File
@@ -72,8 +72,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.utils import (
apply_qk_norm,
python/sglang/srt/models/minimax_m2.py
+8 -3
View File
@@ -74,6 +74,11 @@ from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_loader.weight_utils import (
default_weight_loader,
@@ -633,7 +638,7 @@ class MiniMaxM2MoE(nn.Module):
if router_logits is not None:
ctx = (
nullcontext()
if not get_global_server_args().disable_piecewise_cuda_graph
if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
else get_global_expert_distribution_recorder().with_current_layer(
self.layer_id
)
@@ -671,7 +676,7 @@ class MiniMaxM2MoE(nn.Module):
if self.ep_size > 1:
ctx = (
nullcontext()
if not get_global_server_args().disable_piecewise_cuda_graph
if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
else get_global_expert_distribution_recorder().with_current_layer(
self.layer_id
)
@@ -1171,7 +1176,7 @@ class MiniMaxM2Model(nn.Module):
for i in range(self.start_layer, self.end_layer):
ctx = (
nullcontext()
if not get_global_server_args().disable_piecewise_cuda_graph
if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
else get_global_expert_distribution_recorder().with_current_layer(i)
)
with ctx:
python/sglang/srt/models/mllama.py
+1 -1
View File
@@ -957,7 +957,7 @@ class MllamaForConditionalGeneration(nn.Module):
positions: torch.Tensor,
forward_batch: ForwardBatch,
) -> Union[Tuple, CausalLMOutputWithPast]:
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.runner import get_is_capture_mode
batched_images, batched_ar_ids, batched_ar_mask, encoder_lens_need = (
self._batch_image_inputs(forward_batch)
python/sglang/srt/models/moss_vl.py
+1 -1
View File
@@ -44,8 +44,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from sglang.srt.managers.schedule_batch import MultimodalInputs
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import add_prefix
python/sglang/srt/models/nemotron_h.py
+14 -15
View File
@@ -24,10 +24,6 @@ import torch
from torch import nn
from sglang.srt.compilation.compilation_config import register_split_op
from sglang.srt.compilation.piecewise_context_manager import (
get_forward_context,
is_in_piecewise_cuda_graph,
)
from sglang.srt.configs import NemotronHConfig
from sglang.srt.configs.nemotron_h import ATTENTION, MAMBA, MLP, MOE
from sglang.srt.distributed import (
@@ -62,14 +58,16 @@ from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import (
eager_on_graph,
)
from sglang.srt.model_executor.breakable_cuda_graph.context import (
is_in_breakable_cuda_graph,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
eager_on_graph,
is_in_breakable_cuda_graph,
)
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
get_tc_piecewise_forward_context,
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.model_loader.weight_utils import (
default_weight_loader,
maybe_remap_kv_scale_name,
@@ -232,9 +230,10 @@ class NemotronHMoE(nn.Module):
self,
hidden_states: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor | None]:
# torch.compile cannot trace CUDA streams, so use the non-overlapping
# path when inside piecewise CUDA graph compilation.
if _is_cuda and not is_in_piecewise_cuda_graph():
# torch.compile cannot trace CUDA streams. Take the
# non-overlapping path only during dynamo tracing; replay can
# use the overlapping fast path since dynamo is no longer active.
if _is_cuda and not torch.compiler.is_compiling():
return self._forward_core_shared_routed_overlap(hidden_states)
else:
return self._forward_core_normal(hidden_states)
@@ -447,7 +446,7 @@ class NemotronHMambaDecoderLayer(nn.Module):
breakable_nemotron_mamba2_with_output(hidden_states, output, self.layer_id)
return output, residual
if is_in_piecewise_cuda_graph():
if is_in_tc_piecewise_cuda_graph():
output = torch.empty_like(hidden_states)
nemotron_mamba2_with_output(hidden_states, output, self.layer_id)
return output, residual
@@ -1022,7 +1021,7 @@ def nemotron_mamba2_with_output(
layer_id: int,
) -> None:
"""Split op for Mamba2 forward in piecewise CUDA graph mode."""
context = get_forward_context()
context = get_tc_piecewise_forward_context()
forward_batch = context.forward_batch
attention_layers = context.attention_layers
mamba_layer = attention_layers[layer_id]
python/sglang/srt/models/olmo2.py
+4 -4
View File
@@ -46,8 +46,8 @@ from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.utils import add_prefix, is_cuda, make_layers
@@ -63,7 +63,7 @@ def get_attention_sliding_window_size(config):
class Olmo2Attention(nn.Module):
"""
This is the attention block where the output is computed as
``Attention(LN(x))`` in ``MLP(LN(x + Attention(LN(x))))``
Attention(LN(x)) in MLP(LN(x + Attention(LN(x))))
(plus another skip connection).
"""
@@ -216,7 +216,7 @@ class Olmo2Attention(nn.Module):
class Olmo2MLP(nn.Module):
"""
This is the MLP block where the output is computed as
``MLP(x)`` in ``LN(MLP(x + LN(Attention(x))))``
MLP(x) in LN(MLP(x + LN(Attention(x))))
(plus another skip connection).
"""
@@ -265,7 +265,7 @@ class Olmo2MLP(nn.Module):
class Olmo2DecoderLayer(nn.Module):
"""
This is a typical transformer block where the output is
computed as ``MLP(LN(x + Attention(LN(x))))``
computed as MLP(LN(x + Attention(LN(x))))
(plus another skip connection).
"""
python/sglang/srt/models/qwen2_moe.py
+8 -3
View File
@@ -87,8 +87,13 @@ from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import (
@@ -129,7 +134,7 @@ def can_fuse_shared_expert(
) -> bool:
"""Whether the shared expert may be fused as an extra MoE expert (Qwen3.5 + Aiter).
Caller must still gate on ``support_shared_expert_fusion`` and ``_use_aiter``.
Caller must still gate on support_shared_expert_fusion and _use_aiter.
"""
if (
get_global_server_args().disable_shared_experts_fusion is True
@@ -868,7 +873,7 @@ class Qwen2MoeModel(nn.Module):
for i in range(self.start_layer, self.end_layer):
ctx = (
nullcontext()
if not get_global_server_args().disable_piecewise_cuda_graph
if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
else get_global_expert_distribution_recorder().with_current_layer(i)
)
with ctx:
python/sglang/srt/models/qwen3.py
+6 -1
View File
@@ -22,6 +22,11 @@ from sglang.srt.layers.rotary_embedding import get_rope
from sglang.srt.layers.rotary_embedding.mrope import MRotaryEmbedding
from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
from sglang.srt.model_loader.weight_utils import (
@@ -416,7 +421,7 @@ class Qwen3DecoderLayer(nn.Module):
cache=(
[self.mlp.gate_up_proj.weight, self.mlp.down_proj.weight]
if _is_npu
and not get_global_server_args().disable_piecewise_cuda_graph
and check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
and (
hasattr(self.mlp.gate_up_proj, "weight")
and hasattr(self.mlp.down_proj, "weight")
python/sglang/srt/models/qwen3_5.py
+7 -2
View File
@@ -69,8 +69,13 @@ from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
from sglang.srt.layers.rotary_embedding import get_rope
from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import (
default_weight_loader,
sharded_weight_loader,
@@ -448,7 +453,7 @@ class Qwen3_5GatedDeltaNet(nn.Module):
if (
_is_cpu
or _is_npu
or not get_global_server_args().disable_piecewise_cuda_graph
or check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
):
DUAL_STREAM_TOKEN_THRESHOLD = 0
else:
python/sglang/srt/models/qwen3_next.py
+8 -3
View File
@@ -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:

Some files were not shown because too many files have changed in this diff Show More