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

config: every handler declares its cuda-graph decisions (#36725)

Browse Source 
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-08-27 12:55:34 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent bd4bb1781a
commit 7c3b5a6732
5 changed files with 403 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/utils.py
+39 -7
View File
@@ -8,6 +8,7 @@ import torch
from sglang.srt.arg_groups.overrides import declare_resolution
from sglang.srt.environ import envs
from sglang.srt.model_executor.cuda_graph_config import Phase, with_phase
from sglang.srt.utils import get_npu_memory_capacity, is_npu
if TYPE_CHECKING:
@@ -71,7 +72,6 @@ def set_default_server_args(args: "ServerArgs"):
)
# NPU memory settings
decode = cfg.cuda_graph_config.decode
npu_mem = get_npu_memory_capacity()
if npu_mem <= 32 * 1024:
# Ascend 910B4,910B4_1
@@ -82,11 +82,27 @@ def set_default_server_args(args: "ServerArgs"):
"set_default_server_args",
chunked_prefill_size=4 * 1024,
)
if decode.max_bs is None:
if cfg.cuda_graph_config.decode.max_bs is None:
if cfg.tp_size < 4:
decode.max_bs = 16
declare_resolution(
args,
"set_default_server_args",
cuda_graph_config=with_phase(
cfg.cuda_graph_config,
Phase.DECODE,
max_bs=16,
),
)
else:
decode.max_bs = 64
declare_resolution(
args,
"set_default_server_args",
cuda_graph_config=with_phase(
cfg.cuda_graph_config,
Phase.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, max_bs 64 if tp < 4 else 256)
@@ -96,11 +112,27 @@ def set_default_server_args(args: "ServerArgs"):
"set_default_server_args",
chunked_prefill_size=8 * 1024,
)
if decode.max_bs is None:
if cfg.cuda_graph_config.decode.max_bs is None:
if cfg.tp_size < 4:
decode.max_bs = 64
declare_resolution(
args,
"set_default_server_args",
cuda_graph_config=with_phase(
cfg.cuda_graph_config,
Phase.DECODE,
max_bs=64,
),
)
else:
decode.max_bs = 256
declare_resolution(
args,
"set_default_server_args",
cuda_graph_config=with_phase(
cfg.cuda_graph_config,
Phase.DECODE,
max_bs=256,
),
)
# NPU does not support CustomAllReduce
declare_resolution(
python/sglang/srt/model_executor/cuda_graph_config.py
+20 -1
View File
@@ -23,7 +23,7 @@ inside the function body to preserve that invariant.
import argparse
import dataclasses
import json
from dataclasses import dataclass, field
from dataclasses import dataclass, field, replace
from typing import Any, Dict, List, Optional
@@ -119,6 +119,25 @@ def default_prefill_backend() -> str:
return Backend.BREAKABLE if is_cuda() else Backend.TC_PIECEWISE
def with_phase(config: "CudaGraphConfig", phase: str, **changes) -> "CudaGraphConfig":
"""A copy of ``config`` with ``changes`` applied to one phase.
Resolution declares values, so a handler that decides a graph setting hands
the stash a new config instead of editing the one an earlier handler
declared.
"""
if phase not in Phase.ALL:
raise KeyError(phase)
# Not a deep copy: `dataclasses.replace` copies field references, so a
# list-valued `bs` is shared. Rebind `bs`, never mutate it in place.
return CudaGraphConfig(
**{
name: replace(getattr(config, name), **(changes if name == phase else {}))
for name in Phase.ALL
}
)
@dataclass
class CudaGraphConfig:
"""Top-level CUDA graph config: one PhaseConfig per phase."""
python/sglang/srt/server_args.py
+243 -49
View File
@@ -85,6 +85,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
Phase,
default_cuda_graph_config,
parse_cuda_graph_config_arg,
with_phase,
)
from sglang.srt.parser.reasoning_parser import ReasoningParser
from sglang.srt.platforms import current_platform
@@ -4055,8 +4056,18 @@ class ServerArgs:
)
# cuda_graph_config was already parsed from the legacy boolean, so
# flipping the boolean alone would not stop graph capture.
cfg.cuda_graph_config.decode.backend = Backend.DISABLED
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_handle_model_capability_adjustments",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED
),
)
self._declare(
"_handle_model_capability_adjustments",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
logger.warning(
"HRM-Text (prefix_lm) detected: forcing --attention-backend "
"triton, --chunked-prefill-size -1, --disable-radix-cache, and "
@@ -4139,9 +4150,19 @@ class ServerArgs:
prefill_only_disable_kv_cache=True,
)
self._validate_prefill_only_disable_kv_cache_args()
cfg.cuda_graph_config.decode.backend = Backend.DISABLED
self._declare(
"_handle_model_capability_adjustments",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED
),
)
if is_cuda() and cfg.cuda_graph_config.prefill.backend != Backend.DISABLED:
cfg.cuda_graph_config.prefill.backend = Backend.BREAKABLE
self._declare(
"_handle_model_capability_adjustments",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.BREAKABLE
),
)
# CUDA-graph sizing has already run by this point and derives
# its generic maximum from the 8K chunked-prefill default.
# On the Hopper/Blackwell FA raw-K/V path, raise the unlocked
@@ -4157,21 +4178,32 @@ class ServerArgs:
self, "_cuda_graph_config_locked", set()
)
if (Phase.PREFILL, "max_bs") not in cuda_graph_config_locked:
prefill_config.max_bs = max(
prefill_config.max_bs or 0,
model_config.context_len,
16384,
)
if (Phase.PREFILL, "bs") not in cuda_graph_config_locked:
prefill_config.bs = (
self._generate_prefill_cuda_graph_batch_sizes(
prefill_config.max_bs
)
sizing = {
"max_bs": max(
prefill_config.max_bs or 0,
model_config.context_len,
16384,
)
}
if (Phase.PREFILL, "bs") not in cuda_graph_config_locked:
sizing["bs"] = self._generate_prefill_cuda_graph_batch_sizes(
sizing["max_bs"]
)
self._declare(
"_handle_model_capability_adjustments",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, **sizing
),
)
elif not is_cuda():
# BCG is CUDA-only. Other graph backends do not support this
# encoder-style prefill, so retain the eager Triton path.
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_handle_model_capability_adjustments",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
logger.info(
"EmbeddingGemma detected: disabling radix cache and chunked "
"prefill; using breakable CUDA graph for CUDA prefill."
@@ -4716,7 +4748,12 @@ class ServerArgs:
"At this moment Ascend platform only support prefill graph compilation with "
"cuda_graph_config[prefill].tc_compiler='eager'."
)
cfg.cuda_graph_config.prefill.tc_compiler = "eager"
self._declare(
"_handle_npu_backends",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, tc_compiler="eager"
),
)
def _handle_mps_backends(self):
cfg = resolving_view(self)
@@ -4734,7 +4771,12 @@ class ServerArgs:
# --cuda-graph-backend-decode (or --cuda-graph-config), keep it
# disabled so the default startup doesn't require graph capture.
if (Phase.DECODE, "backend") not in self._cuda_graph_config_locked:
cfg.cuda_graph_config.decode.backend = Backend.DISABLED
self._declare(
"_handle_xpu_backends",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED
),
)
elif cfg.cuda_graph_config.decode.backend not in (
Backend.DISABLED,
Backend.FULL,
@@ -4744,7 +4786,12 @@ class ServerArgs:
"disabling unsupported decode backend '%s'.",
cfg.cuda_graph_config.decode.backend,
)
cfg.cuda_graph_config.decode.backend = Backend.DISABLED
self._declare(
"_handle_xpu_backends",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED
),
)
# ------------------------------------------------------------------
# CUDA graph configuration resolution
@@ -4829,8 +4876,15 @@ class ServerArgs:
sorted(bs),
aligned,
)
cfg.cuda_graph_config.prefill.bs = aligned
cfg.cuda_graph_config.prefill.max_bs = aligned[-1]
self._declare(
"_apply_deepep_adjustments",
cuda_graph_config=with_phase(
cfg.cuda_graph_config,
Phase.PREFILL,
bs=aligned,
max_bs=aligned[-1],
),
)
def _parse_cuda_graph_config(self):
"""Resolve cuda_graph_config from explicit JSON, per-phase
@@ -4926,7 +4980,12 @@ class ServerArgs:
"Using tc_piecewise CUDA graph for validated multimodal "
"decoder prefill."
)
cfg.cuda_graph_config.prefill.backend = Backend.TC_PIECEWISE
self._declare(
"_apply_cuda_graph_compatibility",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.TC_PIECEWISE
),
)
if cfg.cuda_graph_config.prefill.backend == Backend.TC_PIECEWISE:
self._disable_tc_piecewise_cudagraph_if_incompatible()
@@ -4939,10 +4998,20 @@ class ServerArgs:
cfg = resolving_view(self)
if cfg.disaggregation_mode == "prefill":
if (Phase.DECODE, "backend") not in self._cuda_graph_config_locked:
cfg.cuda_graph_config.decode.backend = Backend.DISABLED
self._declare(
"_apply_cuda_graph_disaggregation_roles",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED
),
)
elif cfg.disaggregation_mode == "decode":
if (Phase.PREFILL, "backend") not in self._cuda_graph_config_locked:
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_apply_cuda_graph_disaggregation_roles",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
def _disable_tc_piecewise_cudagraph_if_incompatible(self):
"""TcPiecewise (torch.compile + piecewise) is incompatible with
@@ -5018,7 +5087,15 @@ class ServerArgs:
]
for _name, predicate in rules:
if predicate():
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_disable_tc_piecewise_cudagraph_if_incompatible",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
# One decision, one declaration: every rule declares the same
# value, so a later match would only append a duplicate entry.
break
def _disable_breakable_cudagraph_if_incompatible(self):
"""Breakable (segmented capture, no torch.compile). Breakable enforces
@@ -5071,7 +5148,12 @@ class ServerArgs:
"disabling prefill CUDA graph.",
name,
)
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_disable_breakable_cudagraph_if_incompatible",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
return
def _disable_full_prefill_cudagraph_if_incompatible(self):
@@ -5085,7 +5167,12 @@ class ServerArgs:
"disabling prefill CUDA graph.",
name,
)
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_disable_full_prefill_cudagraph_if_incompatible",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
return
def _disable_prefill_cuda_graph_for_deepseek_trtllm_mla(self):
@@ -5115,7 +5202,12 @@ class ServerArgs:
"backend explicitly (e.g. --cuda-graph-backend-prefill tc_piecewise) to override.",
cfg.cuda_graph_config.prefill.backend,
)
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_disable_prefill_cuda_graph_for_deepseek_trtllm_mla",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
def _validate_cuda_graph_config(self):
cfg = resolving_view(self)
@@ -5143,8 +5235,18 @@ class ServerArgs:
if cfg.cuda_graph_config.decode.backend != Backend.DISABLED:
logger.warning("CUDA graph is disabled because --enable-mis is set.")
cfg.cuda_graph_config.decode.backend = Backend.DISABLED
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_handle_multi_item_scoring",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED
),
)
self._declare(
"_handle_multi_item_scoring",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
if not cfg.disable_radix_cache:
logger.warning("Radix cache is disabled because --enable-mis is set.")
@@ -5192,8 +5294,10 @@ class ServerArgs:
The coefficient 1.5 is a heuristic value, in the future, we can do better estimation by looking at the model types, hidden sizes or even do a dummy run.
"""
cfg = resolving_view(self)
decode_cuda_graph_config = cfg.cuda_graph_config.decode
prefill_cuda_graph_config = cfg.cuda_graph_config.prefill
# A copy, so an earlier declaration keeps the value it recorded.
cuda_graph_config = copy.deepcopy(cfg.cuda_graph_config)
decode_cuda_graph_config = cuda_graph_config.decode
prefill_cuda_graph_config = cuda_graph_config.prefill
if gpu_mem is not None:
if gpu_mem < 20 * 1024:
@@ -5340,6 +5444,11 @@ class ServerArgs:
)
)
if cuda_graph_config != cfg.cuda_graph_config:
self._declare(
"_handle_gpu_memory_settings", cuda_graph_config=cuda_graph_config
)
if cfg.mem_fraction_static is None:
if self.post_capture_kv_sizing_planned():
# Post-capture sizing measures free memory after graph capture, so
@@ -5779,7 +5888,14 @@ class ServerArgs:
# The DSA CP field declarations moved to the override
# registry (arg_groups/overrides.py:
# _deepseek_family_overrides).
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_handle_model_specific_adjustments",
cuda_graph_config=with_phase(
cfg.cuda_graph_config,
Phase.PREFILL,
backend=Backend.DISABLED,
),
)
else:
# Pure TP and partial DP Attention mode is active for DSA, logging a warning
if cfg.dp_size < cfg.tp_size:
@@ -5862,7 +5978,14 @@ class ServerArgs:
# the override registry (arg_groups/overrides.py:
# _deepseek_family_overrides).
if cfg.enable_prefill_cp and self.use_mla_backend():
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_handle_model_specific_adjustments",
cuda_graph_config=with_phase(
cfg.cuda_graph_config,
Phase.PREFILL,
backend=Backend.DISABLED,
),
)
# Set moe backend for DeepSeek: the sm100 quant/moe resolution
# moved to the resolution pipeline (arg_groups/overrides.py:
@@ -6358,15 +6481,35 @@ class ServerArgs:
logger.warning(
"Cuda graph is disabled because of using torch native attention backend"
)
cfg.cuda_graph_config.decode.backend = Backend.DISABLED
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_handle_attention_backend_compatibility",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED
),
)
self._declare(
"_handle_attention_backend_compatibility",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
if attention_backend == "flex_attention":
logger.warning(
"Cuda graph is disabled because of using torch Flex Attention backend"
)
cfg.cuda_graph_config.decode.backend = Backend.DISABLED
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_handle_attention_backend_compatibility",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED
),
)
self._declare(
"_handle_attention_backend_compatibility",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
assert (
cfg.speculative_algorithm is None
), "Speculative decoding is currently not supported with Flex Attention backend"
@@ -7212,11 +7355,17 @@ class ServerArgs:
and prefill_cfg.max_bs > cfg.chunked_prefill_size
and (Phase.PREFILL, "max_bs") not in self._cuda_graph_config_locked
):
prefill_cfg.max_bs = cfg.chunked_prefill_size
clamped = {"max_bs": cfg.chunked_prefill_size}
if (Phase.PREFILL, "bs") not in self._cuda_graph_config_locked:
prefill_cfg.bs = self._generate_prefill_cuda_graph_batch_sizes(
prefill_cfg.max_bs
clamped["bs"] = self._generate_prefill_cuda_graph_batch_sizes(
clamped["max_bs"]
)
self._declare(
"_handle_data_parallelism",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, **clamped
),
)
# Resolve the phase-aware TP LM-head default before validating the
# resulting DP/TP LM-head configuration.
@@ -7533,8 +7682,18 @@ class ServerArgs:
)
if cfg.deepep_mode == "normal":
logger.warning("Cuda graph is disabled because deepep_mode=`normal`")
cfg.cuda_graph_config.decode.backend = Backend.DISABLED
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_handle_a2a_moe",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED
),
)
self._declare(
"_handle_a2a_moe",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
if a2a_backend == "deepep_v2":
self._validate_deepep_v2_model_architecture()
@@ -7574,7 +7733,12 @@ class ServerArgs:
"--moe-a2a-backend deepep_v2."
)
# Prefill reads host counts and is not graph-capturable.
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_handle_a2a_moe",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
logger.warning(
f"DeepEP v2 MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{cfg.tp_size}]."
)
@@ -9299,8 +9463,18 @@ class ServerArgs:
logger.warning(
"Cuda graph is disabled for diffusion LLM inference on AMD GPUs"
)
cfg.cuda_graph_config.decode.backend = Backend.DISABLED
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_handle_dllm_inference",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED
),
)
self._declare(
"_handle_dllm_inference",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
from sglang.srt.arg_groups.overrides import (
_dllm_attention_backend,
@@ -9436,8 +9610,18 @@ class ServerArgs:
logger.warning(
"Cuda graph and server warmup are disabled because of using tensor dump mode"
)
cfg.cuda_graph_config.decode.backend = Backend.DISABLED
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_handle_other_validations",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED
),
)
self._declare(
"_handle_other_validations",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
self._declare("_handle_other_validations", skip_server_warmup=True)
if cfg.msprobe_dump_config is not None:
@@ -9446,8 +9630,18 @@ class ServerArgs:
"cuda graph is disabled because msProbe only supports dump in eager mode, "
"warmup is disabled(skip_server_warmup=True) because there is no need to dump data for this stage."
)
cfg.cuda_graph_config.decode.backend = Backend.DISABLED
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED
self._declare(
"_handle_other_validations",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED
),
)
self._declare(
"_handle_other_validations",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
),
)
self._declare("_handle_other_validations", skip_server_warmup=True)
# Validate limit_mm_per_prompt modalities