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

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
@@ -8,6 +8,7 @@ import torch
from sglang.srt.arg_groups.overrides import declare_resolution from sglang.srt.arg_groups.overrides import declare_resolution
from sglang.srt.environ import envs 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 from sglang.srt.utils import get_npu_memory_capacity, is_npu
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -71,7 +72,6 @@ def set_default_server_args(args: "ServerArgs"):
) )
# NPU memory settings # NPU memory settings
decode = cfg.cuda_graph_config.decode
npu_mem = get_npu_memory_capacity() npu_mem = get_npu_memory_capacity()
if npu_mem <= 32 * 1024: if npu_mem <= 32 * 1024:
# Ascend 910B4,910B4_1 # Ascend 910B4,910B4_1
@@ -82,11 +82,27 @@ def set_default_server_args(args: "ServerArgs"):
"set_default_server_args", "set_default_server_args",
chunked_prefill_size=4 * 1024, 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: 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: 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: elif npu_mem <= 64 * 1024:
# Ascend 910B1,910B2,910B2C,910B3,910_9391,910_9392,910_9381,910_9382,910_9372,910_9362 # 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) # (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", "set_default_server_args",
chunked_prefill_size=8 * 1024, 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: 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: 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 # NPU does not support CustomAllReduce
declare_resolution( declare_resolution(
@@ -23,7 +23,7 @@ inside the function body to preserve that invariant.
import argparse import argparse
import dataclasses import dataclasses
import json import json
from dataclasses import dataclass, field from dataclasses import dataclass, field, replace
from typing import Any, Dict, List, Optional 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 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 @dataclass
class CudaGraphConfig: class CudaGraphConfig:
"""Top-level CUDA graph config: one PhaseConfig per phase.""" """Top-level CUDA graph config: one PhaseConfig per phase."""
+237 -43
View File
@@ -85,6 +85,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
Phase, Phase,
default_cuda_graph_config, default_cuda_graph_config,
parse_cuda_graph_config_arg, parse_cuda_graph_config_arg,
with_phase,
) )
from sglang.srt.parser.reasoning_parser import ReasoningParser from sglang.srt.parser.reasoning_parser import ReasoningParser
from sglang.srt.platforms import current_platform from sglang.srt.platforms import current_platform
@@ -4055,8 +4056,18 @@ class ServerArgs:
) )
# cuda_graph_config was already parsed from the legacy boolean, so # cuda_graph_config was already parsed from the legacy boolean, so
# flipping the boolean alone would not stop graph capture. # flipping the boolean alone would not stop graph capture.
cfg.cuda_graph_config.decode.backend = Backend.DISABLED self._declare(
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED "_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( logger.warning(
"HRM-Text (prefix_lm) detected: forcing --attention-backend " "HRM-Text (prefix_lm) detected: forcing --attention-backend "
"triton, --chunked-prefill-size -1, --disable-radix-cache, and " "triton, --chunked-prefill-size -1, --disable-radix-cache, and "
@@ -4139,9 +4150,19 @@ class ServerArgs:
prefill_only_disable_kv_cache=True, prefill_only_disable_kv_cache=True,
) )
self._validate_prefill_only_disable_kv_cache_args() 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: 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 # CUDA-graph sizing has already run by this point and derives
# its generic maximum from the 8K chunked-prefill default. # its generic maximum from the 8K chunked-prefill default.
# On the Hopper/Blackwell FA raw-K/V path, raise the unlocked # On the Hopper/Blackwell FA raw-K/V path, raise the unlocked
@@ -4157,21 +4178,32 @@ class ServerArgs:
self, "_cuda_graph_config_locked", set() self, "_cuda_graph_config_locked", set()
) )
if (Phase.PREFILL, "max_bs") not in cuda_graph_config_locked: if (Phase.PREFILL, "max_bs") not in cuda_graph_config_locked:
prefill_config.max_bs = max( sizing = {
"max_bs": max(
prefill_config.max_bs or 0, prefill_config.max_bs or 0,
model_config.context_len, model_config.context_len,
16384, 16384,
) )
}
if (Phase.PREFILL, "bs") not in cuda_graph_config_locked: if (Phase.PREFILL, "bs") not in cuda_graph_config_locked:
prefill_config.bs = ( sizing["bs"] = self._generate_prefill_cuda_graph_batch_sizes(
self._generate_prefill_cuda_graph_batch_sizes( sizing["max_bs"]
prefill_config.max_bs
) )
self._declare(
"_handle_model_capability_adjustments",
cuda_graph_config=with_phase(
cfg.cuda_graph_config, Phase.PREFILL, **sizing
),
) )
elif not is_cuda(): elif not is_cuda():
# BCG is CUDA-only. Other graph backends do not support this # BCG is CUDA-only. Other graph backends do not support this
# encoder-style prefill, so retain the eager Triton path. # 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( logger.info(
"EmbeddingGemma detected: disabling radix cache and chunked " "EmbeddingGemma detected: disabling radix cache and chunked "
"prefill; using breakable CUDA graph for CUDA prefill." "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 " "At this moment Ascend platform only support prefill graph compilation with "
"cuda_graph_config[prefill].tc_compiler='eager'." "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): def _handle_mps_backends(self):
cfg = resolving_view(self) cfg = resolving_view(self)
@@ -4734,7 +4771,12 @@ class ServerArgs:
# --cuda-graph-backend-decode (or --cuda-graph-config), keep it # --cuda-graph-backend-decode (or --cuda-graph-config), keep it
# disabled so the default startup doesn't require graph capture. # disabled so the default startup doesn't require graph capture.
if (Phase.DECODE, "backend") not in self._cuda_graph_config_locked: 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 ( elif cfg.cuda_graph_config.decode.backend not in (
Backend.DISABLED, Backend.DISABLED,
Backend.FULL, Backend.FULL,
@@ -4744,7 +4786,12 @@ class ServerArgs:
"disabling unsupported decode backend '%s'.", "disabling unsupported decode backend '%s'.",
cfg.cuda_graph_config.decode.backend, 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 # CUDA graph configuration resolution
@@ -4829,8 +4876,15 @@ class ServerArgs:
sorted(bs), sorted(bs),
aligned, aligned,
) )
cfg.cuda_graph_config.prefill.bs = aligned self._declare(
cfg.cuda_graph_config.prefill.max_bs = aligned[-1] "_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): def _parse_cuda_graph_config(self):
"""Resolve cuda_graph_config from explicit JSON, per-phase """Resolve cuda_graph_config from explicit JSON, per-phase
@@ -4926,7 +4980,12 @@ class ServerArgs:
"Using tc_piecewise CUDA graph for validated multimodal " "Using tc_piecewise CUDA graph for validated multimodal "
"decoder prefill." "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: if cfg.cuda_graph_config.prefill.backend == Backend.TC_PIECEWISE:
self._disable_tc_piecewise_cudagraph_if_incompatible() self._disable_tc_piecewise_cudagraph_if_incompatible()
@@ -4939,10 +4998,20 @@ class ServerArgs:
cfg = resolving_view(self) cfg = resolving_view(self)
if cfg.disaggregation_mode == "prefill": if cfg.disaggregation_mode == "prefill":
if (Phase.DECODE, "backend") not in self._cuda_graph_config_locked: 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": elif cfg.disaggregation_mode == "decode":
if (Phase.PREFILL, "backend") not in self._cuda_graph_config_locked: 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): def _disable_tc_piecewise_cudagraph_if_incompatible(self):
"""TcPiecewise (torch.compile + piecewise) is incompatible with """TcPiecewise (torch.compile + piecewise) is incompatible with
@@ -5018,7 +5087,15 @@ class ServerArgs:
] ]
for _name, predicate in rules: for _name, predicate in rules:
if predicate(): 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): def _disable_breakable_cudagraph_if_incompatible(self):
"""Breakable (segmented capture, no torch.compile). Breakable enforces """Breakable (segmented capture, no torch.compile). Breakable enforces
@@ -5071,7 +5148,12 @@ class ServerArgs:
"disabling prefill CUDA graph.", "disabling prefill CUDA graph.",
name, 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 return
def _disable_full_prefill_cudagraph_if_incompatible(self): def _disable_full_prefill_cudagraph_if_incompatible(self):
@@ -5085,7 +5167,12 @@ class ServerArgs:
"disabling prefill CUDA graph.", "disabling prefill CUDA graph.",
name, 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 return
def _disable_prefill_cuda_graph_for_deepseek_trtllm_mla(self): 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.", "backend explicitly (e.g. --cuda-graph-backend-prefill tc_piecewise) to override.",
cfg.cuda_graph_config.prefill.backend, 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): def _validate_cuda_graph_config(self):
cfg = resolving_view(self) cfg = resolving_view(self)
@@ -5143,8 +5235,18 @@ class ServerArgs:
if cfg.cuda_graph_config.decode.backend != Backend.DISABLED: if cfg.cuda_graph_config.decode.backend != Backend.DISABLED:
logger.warning("CUDA graph is disabled because --enable-mis is set.") logger.warning("CUDA graph is disabled because --enable-mis is set.")
cfg.cuda_graph_config.decode.backend = Backend.DISABLED self._declare(
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED "_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: if not cfg.disable_radix_cache:
logger.warning("Radix cache is disabled because --enable-mis is set.") 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. 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) cfg = resolving_view(self)
decode_cuda_graph_config = cfg.cuda_graph_config.decode # A copy, so an earlier declaration keeps the value it recorded.
prefill_cuda_graph_config = cfg.cuda_graph_config.prefill 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 is not None:
if gpu_mem < 20 * 1024: 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 cfg.mem_fraction_static is None:
if self.post_capture_kv_sizing_planned(): if self.post_capture_kv_sizing_planned():
# Post-capture sizing measures free memory after graph capture, so # 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 # The DSA CP field declarations moved to the override
# registry (arg_groups/overrides.py: # registry (arg_groups/overrides.py:
# _deepseek_family_overrides). # _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: else:
# Pure TP and partial DP Attention mode is active for DSA, logging a warning # Pure TP and partial DP Attention mode is active for DSA, logging a warning
if cfg.dp_size < cfg.tp_size: if cfg.dp_size < cfg.tp_size:
@@ -5862,7 +5978,14 @@ class ServerArgs:
# the override registry (arg_groups/overrides.py: # the override registry (arg_groups/overrides.py:
# _deepseek_family_overrides). # _deepseek_family_overrides).
if cfg.enable_prefill_cp and self.use_mla_backend(): 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 # Set moe backend for DeepSeek: the sm100 quant/moe resolution
# moved to the resolution pipeline (arg_groups/overrides.py: # moved to the resolution pipeline (arg_groups/overrides.py:
@@ -6358,15 +6481,35 @@ class ServerArgs:
logger.warning( logger.warning(
"Cuda graph is disabled because of using torch native attention backend" "Cuda graph is disabled because of using torch native attention backend"
) )
cfg.cuda_graph_config.decode.backend = Backend.DISABLED self._declare(
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED "_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": if attention_backend == "flex_attention":
logger.warning( logger.warning(
"Cuda graph is disabled because of using torch Flex Attention backend" "Cuda graph is disabled because of using torch Flex Attention backend"
) )
cfg.cuda_graph_config.decode.backend = Backend.DISABLED self._declare(
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED "_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 ( assert (
cfg.speculative_algorithm is None cfg.speculative_algorithm is None
), "Speculative decoding is currently not supported with Flex Attention backend" ), "Speculative decoding is currently not supported with Flex Attention backend"
@@ -7212,10 +7355,16 @@ class ServerArgs:
and prefill_cfg.max_bs > cfg.chunked_prefill_size and prefill_cfg.max_bs > cfg.chunked_prefill_size
and (Phase.PREFILL, "max_bs") not in self._cuda_graph_config_locked 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: if (Phase.PREFILL, "bs") not in self._cuda_graph_config_locked:
prefill_cfg.bs = self._generate_prefill_cuda_graph_batch_sizes( clamped["bs"] = self._generate_prefill_cuda_graph_batch_sizes(
prefill_cfg.max_bs 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 # Resolve the phase-aware TP LM-head default before validating the
@@ -7533,8 +7682,18 @@ class ServerArgs:
) )
if cfg.deepep_mode == "normal": if cfg.deepep_mode == "normal":
logger.warning("Cuda graph is disabled because deepep_mode=`normal`") logger.warning("Cuda graph is disabled because deepep_mode=`normal`")
cfg.cuda_graph_config.decode.backend = Backend.DISABLED self._declare(
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED "_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": if a2a_backend == "deepep_v2":
self._validate_deepep_v2_model_architecture() self._validate_deepep_v2_model_architecture()
@@ -7574,7 +7733,12 @@ class ServerArgs:
"--moe-a2a-backend deepep_v2." "--moe-a2a-backend deepep_v2."
) )
# Prefill reads host counts and is not graph-capturable. # 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( 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}]." 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( logger.warning(
"Cuda graph is disabled for diffusion LLM inference on AMD GPUs" "Cuda graph is disabled for diffusion LLM inference on AMD GPUs"
) )
cfg.cuda_graph_config.decode.backend = Backend.DISABLED self._declare(
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED "_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 ( from sglang.srt.arg_groups.overrides import (
_dllm_attention_backend, _dllm_attention_backend,
@@ -9436,8 +9610,18 @@ class ServerArgs:
logger.warning( logger.warning(
"Cuda graph and server warmup are disabled because of using tensor dump mode" "Cuda graph and server warmup are disabled because of using tensor dump mode"
) )
cfg.cuda_graph_config.decode.backend = Backend.DISABLED self._declare(
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED "_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) self._declare("_handle_other_validations", skip_server_warmup=True)
if cfg.msprobe_dump_config is not None: 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, " "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." "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 self._declare(
cfg.cuda_graph_config.prefill.backend = Backend.DISABLED "_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) self._declare("_handle_other_validations", skip_server_warmup=True)
# Validate limit_mm_per_prompt modalities # Validate limit_mm_per_prompt modalities
@@ -920,5 +920,103 @@ class TestResolutionDeclarations(CustomTestCase):
) )
class TestDeclaredValuesAreNotEditedLater(CustomTestCase):
"""A declaration records a value, not a handle on one.
The stash keeps whatever object the declaring handler passed, so a handler
that declares a mutable and then edits it in place rewrites an entry that
already went into the log. The projection still answers with the end state,
which is why nothing else notices: what is lost is *which* handler decided
what, and `validate_declarations` never sees the later change at all.
"""
def setUp(self):
super().setUp()
environment = dict(os.environ)
def restore():
os.environ.clear()
os.environ.update(environment)
self.addCleanup(restore)
def _resolve_recording_each_entry(self, **supplied):
"""Resolve, deep-copying every stash entry the moment it is appended."""
from sglang.srt.arg_groups import overrides
recorded = []
def watch(name):
original = getattr(overrides, name)
def wrapper(server_args, *args, **kwargs):
result = original(server_args, *args, **kwargs)
stash = getattr(server_args, "_resolved_overrides", None) or []
while len(recorded) < len(stash):
index = len(recorded)
recorded.append((index, copy.deepcopy(stash[index])))
return result
return original, wrapper
# Every path that appends to the stash.
patched = {}
for name in (
"declare_resolution",
"declare_late_resolution",
"declare_direct_writes",
"run_post_process_pass",
):
original, wrapper = watch(name)
patched[name] = original
setattr(overrides, name, wrapper)
try:
path = tempfile.mkdtemp(prefix="declared_values_")
self.addCleanup(shutil.rmtree, path, ignore_errors=True)
with open(os.path.join(path, "config.json"), "w") as handle:
json.dump(_MINI_CONFIG, handle)
server_args = ServerArgs(
model_path=path, device="cuda", random_seed=42, **supplied
)
server_args.resolve_once()
finally:
for name, original in patched.items():
setattr(overrides, name, original)
return server_args, recorded
def test_no_entry_changes_after_it_is_recorded(self):
# One shape per family of handlers that decides a graph setting.
for label, supplied in (
("plain", {}),
("cuda_graph_knobs", {"cuda_graph_max_bs_decode": 16}),
("chunked_prefill", {"chunked_prefill_size": 1024}),
("explicit_json", {"cuda_graph_config": {"decode": {"max_bs": 12}}}),
("disaggregation", {"disaggregation_mode": "prefill"}),
("deterministic", {"enable_deterministic_inference": True}),
("speculative", {"speculative_algorithm": "EAGLE"}),
("dp_attention", {"tp_size": 2, "dp_size": 2, "enable_dp_attention": True}),
):
with self.subTest(shape=label):
server_args, recorded = self._resolve_recording_each_entry(**supplied)
stash = server_args._resolved_overrides
self.assertGreater(
len(recorded),
0,
"nothing was recorded, so this case is not watching the "
"declaration paths it thinks it is",
)
drifted = [
(index, was, stash[index])
for index, was in recorded
if stash[index] != was
]
self.assertEqual(
[],
drifted,
"these entries changed after they were declared, so the log "
f"credits the wrong handler for the end state: {drifted}",
)
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()
@@ -2245,8 +2245,9 @@ class TestDeepEPv2Args(CustomTestCase):
for mode in ("direct", "hybrid"): for mode in ("direct", "hybrid"):
args = self._args(moe_runner_backend="deep_gemm", deepep_v2_mode=mode) args = self._args(moe_runner_backend="deep_gemm", deepep_v2_mode=mode)
args._handle_a2a_moe() args._handle_a2a_moe()
self.assertEqual(args.cuda_graph_config.decode.backend, Backend.FULL) declared = resolution_result(args, "cuda_graph_config")
self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.DISABLED) self.assertEqual(declared.decode.backend, Backend.FULL)
self.assertEqual(declared.prefill.backend, Backend.DISABLED)
def test_two_batch_overlap_rejected(self): def test_two_batch_overlap_rejected(self):
args = self._args(moe_runner_backend="deep_gemm", enable_two_batch_overlap=True) args = self._args(moe_runner_backend="deep_gemm", enable_two_batch_overlap=True)