[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:
co-authored by
Claude Opus 4.8
parent
7c90840bad
commit
56ef810cad
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user