[diffusion] BCG: auto-capture the default warmup resolution instead of hard-requiring --warmup-resolutions (H200 SANA denoise 0.73->0.457 s with a single flag) (#34174)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-08-10 08:37:32 +08:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 7c90840bad
commit 56ef810cad
2 changed files with 54 additions and 5 deletions
@@ -549,11 +549,16 @@ class ServerArgs(DisaggServerArgsMixin):
# latent shape, so the user must declare the resolutions up front. We
# capture every one of them at warmup; serving then never re-captures.
if not self.warmup_resolutions:
raise ValueError(
"--enable-breakable-cuda-graph requires --warmup-resolutions: "
"diffusion CUDA graphs only replay for a fixed resolution, so "
"every served resolution must be declared and captured at "
"warmup, e.g. --warmup-resolutions 1024x1024 1328x1328."
# No explicit resolutions: capture the model's default warmup
# resolution (derived by build_warmup_reqs) so
# --enable-breakable-cuda-graph works standalone. BCG graphs are
# resolution-specific; a request at any other resolution simply
# falls back to eager (the runner never re-captures at serving
# time). Pass --warmup-resolutions to capture additional shapes.
logger.info(
"[Diffusion BCG] --warmup-resolutions unset; capturing the "
"model default warmup resolution. Requests at other "
"resolutions run eager."
)
if self.bcg_text_buckets is not None and not any(
int(b) > 0 for b in self.bcg_text_buckets
@@ -572,6 +577,8 @@ class ServerArgs(DisaggServerArgsMixin):
pipeline_config_name in BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS
and self._is_breakable_cuda_graph_supported_model()
):
if not self.warmup_resolutions:
self._default_bcg_warmup_resolution()
return
logger.warning(
@@ -588,6 +595,39 @@ class ServerArgs(DisaggServerArgsMixin):
refs.update(_normalized_bcg_model_refs(self.model_path))
return bool(refs & BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS)
def _default_bcg_warmup_resolution(self) -> None:
"""Seed --warmup-resolutions with the model default for BCG.
BCG graphs are resolution-specific and captured at warmup. When the
user does not pre-declare resolutions we capture the model's default
warmup resolution so --enable-breakable-cuda-graph works standalone;
requests at any other resolution fall back to eager.
"""
from sglang.multimodal_gen.runtime.warmup_request_builder import (
_resolve_default_warmup_resolution,
get_model_sampling_defaults,
)
try:
sampling_defaults = get_model_sampling_defaults(self)
width, height = _resolve_default_warmup_resolution(
self, sampling_defaults, server_based_warmup=True
)
except Exception as exc: # pragma: no cover - defensive
logger.warning(
"[Diffusion BCG] could not derive a default warmup resolution "
"(%s); no graph will be captured and serving runs eager.",
exc,
)
return
self.warmup_resolutions = [f"{width}x{height}"]
logger.info(
"[Diffusion BCG] --warmup-resolutions unset; capturing the model "
"default %dx%d. Requests at other resolutions run eager.",
width,
height,
)
def _adjust_save_paths(self):
"""Normalize empty-string save paths to None (disabled)."""
if self.output_path is not None and self.output_path.strip() == "":
@@ -588,6 +588,15 @@ class TestWarmupModeNormalization(unittest.TestCase):
sa = self._resolve(enable_breakable_cuda_graph=True)
self.assertEqual(sa.warmup_mode, "server")
def test_breakable_cuda_graph_allows_unset_resolutions(self):
# BCG no longer hard-requires --warmup-resolutions; the model
# default warmup resolution is captured at warmup instead.
sa = ServerArgs.__new__(ServerArgs)
sa.enable_breakable_cuda_graph = True
sa.warmup_resolutions = None
sa.bcg_text_buckets = None
sa._validate_breakable_cuda_graph() # must not raise
def test_disagg_role_disables_server_warmup(self):
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType