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

Enable Breakable Cuda Graph as Default (#29458)

Browse Source
This commit is contained in:
Yuwei An
2026-07-01 22:04:44 -07:00
committed by GitHub
parent d15b79e96e
commit 0543246184
10 changed files with 137 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/cuda_graph_config.py
+10 -1
View File
@@ -80,6 +80,15 @@ class PhaseConfig:
tc_compiler: str = "eager"
def default_prefill_backend() -> str:
"""BCG (breakable) is the prefill default on CUDA only; other platforms
(HIP/NPU/...) keep tc_piecewise until BCG is validated there. Lazy import
keeps this module's stdlib-only import invariant (see module docstring)."""
from sglang.srt.utils import is_cuda
return Backend.BREAKABLE if is_cuda() else Backend.TC_PIECEWISE
@dataclass
class CudaGraphConfig:
"""Top-level CUDA graph config: one PhaseConfig per phase."""
@@ -88,7 +97,7 @@ class CudaGraphConfig:
default_factory=lambda: PhaseConfig(backend=Backend.FULL)
)
prefill: PhaseConfig = field(
default_factory=lambda: PhaseConfig(backend=Backend.TC_PIECEWISE)
default_factory=lambda: PhaseConfig(backend=default_prefill_backend())
)
def __getitem__(self, phase: str) -> PhaseConfig:
python/sglang/srt/model_executor/runner/__init__.py
+1 -1
View File
@@ -38,7 +38,7 @@ from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import ( # noqa
)
from sglang.srt.model_executor.runner.shape_key import ShapeKey # noqa: F401
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( # noqa: F401
TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG,
TCPCG_FAILURE_HINT,
)
from sglang.srt.model_executor.runner_utils import ( # noqa: F401
DecodeInputBuffers,
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+51 -15
View File
@@ -70,7 +70,11 @@ from sglang.srt.model_executor.runner_backend.utils import (
from sglang.srt.model_executor.runner_backend_utils import (
PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG,
)
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
BCG_FAILURE_HINT,
)
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
TCPCG_FAILURE_HINT,
set_tc_piecewise_forward_context,
)
from sglang.srt.model_executor.runner_utils.buffers import (
@@ -93,6 +97,27 @@ logger = logging.getLogger(__name__)
_is_hip = is_hip()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
def prefill_failure_msg(backend_name: str) -> str:
"""Render PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG with a backend-specific
numbered suggestion list. The runner is only constructed for BREAKABLE
or TC_PIECEWISE; other values fall back to a generic OOM-style list."""
if backend_name == Backend.BREAKABLE:
hint = BCG_FAILURE_HINT
elif backend_name == Backend.TC_PIECEWISE:
hint = TCPCG_FAILURE_HINT
else:
hint = (
"1. disable the prefill CUDA graph by --cuda-graph-backend-prefill=disabled\n"
"2. if it is an OOM problem, set --mem-fraction-static to a smaller value "
"(e.g., 0.8 or 0.7) or set --cuda-graph-max-bs-prefill to a smaller value "
"(e.g., 2048)\n"
)
return PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG.format(
backend=backend_name, suggestions=hint
)
# 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.
@@ -156,6 +181,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
_prefill_backend_name = (
_cg_cfg.prefill.backend if _cg_cfg is not None else Backend.TC_PIECEWISE
)
self.prefill_backend_name = _prefill_backend_name
if (
_prefill_backend_name == Backend.BREAKABLE
and model_runner.spec_algorithm.is_eagle()
@@ -224,10 +250,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
try:
self.backend = resolve_prefill_backend(self)
except RuntimeError as e:
if _prefill_backend_name == Backend.TC_PIECEWISE:
if _prefill_backend_name in (Backend.TC_PIECEWISE, Backend.BREAKABLE):
raise Exception(
f"Capture prefill CUDA graph failed: {e}\n"
f"{PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG}"
f"{prefill_failure_msg(_prefill_backend_name)}"
)
raise
if isinstance(self.backend, BreakableCudaGraphBackend):
@@ -288,16 +314,18 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
language_model.model, "layers"
):
self.layer_model = language_model.model
params = list(inspect.signature(self.layer_model.forward).parameters)
self._input_embeds_arg_idx = (
params.index("input_embeds") if "input_embeds" in params else None
)
elif hasattr(language_model, "layers"):
self.layer_model = language_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."
)
params = list(inspect.signature(self.layer_model.forward).parameters)
self._input_embeds_arg_idx = (
params.index("input_embeds") if "input_embeds" in params else None
)
# --- aiter chip info pre-warming (AMD) -------------------------
if _use_aiter:
@@ -306,13 +334,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# --- capture --------------------------------------------------
self.device_module.synchronize()
self.model_runner.tp_group.barrier()
try:
self.capture()
except RuntimeError as e:
raise Exception(
f"Capture prefill CUDA graph failed: {e}\n"
f"{PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG}"
)
self.capture()
self.raw_num_tokens = 0
@@ -591,8 +613,17 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
forward_mode=ForwardMode.EXTEND,
batch_size=bs,
input_ids=_slot("input_ids"),
# BCG's graph is text-only, so it forces input_embeds=None;
# tc_piecewise keeps the slot so multimodal prefill keeps its
# image embeds (else NaN logits).
input_embeds=(
_slot("input_embeds") if registry.has_slot("input_embeds") else None
None
if self.prefill_backend_name == Backend.BREAKABLE
else (
_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"],
@@ -755,8 +786,13 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
)
input_ids = _slot("input_ids")
# BCG's graph is text-only, so it forces input_embeds=None; tc_piecewise
# keeps the slot so multimodal prefill keeps its image embeds (else NaN
# logits).
input_embeds = (
_slot("input_embeds") if registry.has_slot("input_embeds") else None
None
if self.prefill_backend_name == Backend.BREAKABLE
else (_slot("input_embeds") if registry.has_slot("input_embeds") else None)
)
positions = _slot("positions")
out_cache_loc = _slot("out_cache_loc")
python/sglang/srt/model_executor/runner_backend_utils/__init__.py
+3 -5
View File
@@ -21,10 +21,8 @@ CUDA_GRAPH_CAPTURE_FAILED_MSG = (
)
PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG = (
"Possible solutions:\n"
"1. set --mem-fraction-static to a smaller value (e.g., 0.8 or 0.7)\n"
"2. set --cuda-graph-max-bs-prefill to a smaller value (e.g., 2048)\n"
"3. disable prefill CUDA graph by --cuda-graph-backend-prefill=disabled. "
"(Not recommended. Performance loss)\n"
"Fail when using backend: {backend} for prefill runner.\n"
"Possible suggestions:\n"
"{suggestions}"
"Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
)
python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/context.py
+26
View File
@@ -15,8 +15,16 @@
from __future__ import annotations
import logging
from contextlib import contextmanager
from sglang.srt.model_executor.cuda_graph_config import Backend
from sglang.srt.model_executor.runner_backend_utils import (
PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG,
)
logger = logging.getLogger(__name__)
_in_breakable_cuda_graph = False
@@ -26,9 +34,27 @@ def is_in_breakable_cuda_graph() -> bool:
@contextmanager
def enable_breakable_cuda_graph():
"""Mark the enclosed scope as inside a BCG capture/replay. Any exception
raised inside is logged with the BCG-specific failure hint, then re-raised
for the caller to handle."""
global _in_breakable_cuda_graph
_in_breakable_cuda_graph = True
try:
yield
except Exception as exc:
msg = PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG.format(
backend=Backend.BREAKABLE, suggestions=BCG_FAILURE_HINT
)
logger.error(f"{type(exc).__name__}: {exc}\n{msg}")
raise
finally:
_in_breakable_cuda_graph = False
BCG_FAILURE_HINT = (
"1. change to tc_piecewise by --cuda-graph-backend-prefill=tc_piecewise\n"
"2. disable the prefill CUDA graph by --cuda-graph-backend-prefill=disabled\n"
"3. if it is an OOM problem, set --mem-fraction-static to a smaller value "
"(e.g., 0.8 or 0.7) or set --cuda-graph-max-bs-prefill to a smaller value "
"(e.g., 2048)\n"
)
python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/__init__.py
+3 -2
View File
@@ -4,7 +4,8 @@ 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.
- TCPCG_FAILURE_HINT — backend-switch suggestion plugged into
PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG by the prefill runner.
The torch.compile-warmup flag (is_in_torch_compile_warmup) lives in
sglang.srt.compilation.compile_phase — it is torch.compile-internal,
@@ -12,7 +13,7 @@ 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,
TCPCG_FAILURE_HINT,
TcPiecewiseForwardContext,
enable_tc_piecewise_cuda_graph,
get_tc_piecewise_forward_context,
python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py
+19 -4
View File
@@ -21,10 +21,12 @@ This module deliberately does **not** own torch.compile-specific state
from __future__ import annotations
import logging
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, List, Optional
from sglang.srt.model_executor.cuda_graph_config import Backend
from sglang.srt.model_executor.runner_backend_utils import (
PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG,
)
@@ -33,6 +35,8 @@ if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
logger = logging.getLogger(__name__)
_in_tc_piecewise_cuda_graph = False
@@ -43,13 +47,20 @@ def is_in_tc_piecewise_cuda_graph() -> bool:
@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.
"""Mark the enclosed scope as inside a tc_piecewise CUDA graph
capture/replay. Any exception raised inside is logged with the
PCG-specific failure hint, then re-raised for the caller to handle.
"""
global _in_tc_piecewise_cuda_graph
_in_tc_piecewise_cuda_graph = True
try:
yield
except Exception as exc:
msg = PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG.format(
backend=Backend.TC_PIECEWISE, suggestions=TCPCG_FAILURE_HINT
)
logger.error(f"{type(exc).__name__}: {exc}\n{msg}")
raise
finally:
_in_tc_piecewise_cuda_graph = False
@@ -101,6 +112,10 @@ def set_tc_piecewise_forward_context(
_tc_piecewise_forward_context = None
TC_PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG = (
"Piecewise CUDA graph failed.\n" + PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG
TCPCG_FAILURE_HINT = (
"1. change to breakable by --cuda-graph-backend-prefill=breakable\n"
"2. disable the prefill CUDA graph by --cuda-graph-backend-prefill=disabled\n"
"3. if it is an OOM problem, set --mem-fraction-static to a smaller value "
"(e.g., 0.8 or 0.7) or set --cuda-graph-max-bs-prefill to a smaller value "
"(e.g., 2048)\n"
)
python/sglang/srt/server_args.py
+19 -3
View File
@@ -3215,11 +3215,27 @@ class ServerArgs:
memory-saver rejection in its own __init__; config-time rules can be
added here as they're discovered.
"""
from sglang.srt.configs.model_config import is_deepseek_v4
rules = [
# MLA prefill takes a different attn-forward path under BCG (no
# tc_piecewise gate), causing q.view shape mismatches. Disable
# until the MLA prefill path is BCG-aware.
# MLA prefill takes a different attn-forward path under BCG.
("MLA attention", lambda: self.use_mla_backend()),
# DSV4 is BCG-compatible but introduces heavy memory pressure: the
# c4 indexer scratch is pinned in the capture pool and OOMs. Disable.
(
"DeepSeek-V4 (heavy capture-pool memory pressure)",
lambda: is_deepseek_v4(self.get_model_config().hf_config),
),
# CP all_gather replay size mismatch under BCG.
("context parallel (attn_cp_size > 1)", lambda: self.attn_cp_size > 1),
# BCG capture + LoRA adapter weights exceed host RAM headroom.
("LoRA", lambda: bool(self.lora_paths) or bool(self.enable_lora)),
# BCG bucket sizes exceed FlashInfer MoE A2A's dispatch cap.
("MoE A2A backend", lambda: self.moe_a2a_backend != "none"),
# DP-attn × BCG capture/replay not yet validated.
("DP attention", lambda: self.enable_dp_attention),
# Multimodal prefill replay faults under BCG.
("multimodal model", lambda: self.get_model_config().is_multimodal),
]
for name, predicate in rules:
if predicate():
test/registered/quant/test_fp8_blockwise_gemm.py
+3
View File
@@ -74,6 +74,9 @@ class MXFP8GemmBase:
"--trust-remote-code",
"--fp8-gemm-backend",
cls.backend,
# TODO: pin tc_piecewise — default `breakable` prefill runs MXFP8 RMSNorm in bf16, hurting accuracy; unrelated to BCG-default change.
"--cuda-graph-backend-prefill",
"tc_piecewise",
]
cls.process = popen_launch_server(
cls.model,
test/registered/quant/test_nvfp4_gemm.py
+2 -1
View File
@@ -58,7 +58,8 @@ class FP4GemmBase:
metrics = run_eval(args)
print(metrics)
self.assertGreater(metrics["score"], 0.64)
# TODO: restore 0.64 once the BCG-prefill RMSNorm fp32 fix lands.
self.assertGreater(metrics["score"], 0.63)
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")