From 447048dba26cd953dc156c4f7f416da144ce9b63 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Mon, 24 Aug 2026 08:50:26 +0800 Subject: [PATCH] [diffusion] Reject unsafe quality=high BCG replay (#36008) --- docs/docs/sglang-diffusion/api/cli.mdx | 2 +- docs/docs/sglang-diffusion/fused_kernels.mdx | 9 +++++ .../pipelines_core/stages/denoising.py | 13 +++++++ .../test/unit/test_diffusion_bcg_padding.py | 38 +++++++++++++++++++ 4 files changed, 61 insertions(+), 1 deletion(-) diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 9697275b9..90e46ad5a 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -97,7 +97,7 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis - `--warmup-mode {off|request|server}`: control startup warmup for `sglang serve`; `off` skips warmup, `request` primes the request path, and `server` runs a full synthetic server warmup before serving traffic - `--enable-torch-compile {true|false}`: compile native diffusion hot paths. When no warmup mode is configured, this also enables server warmup so first real requests do not pay compile latency. - `--offload-during-compile {true|false}`: when compile warmup is active, temporarily layerwise-offload DiT weights and move resident non-DiT components off-device so `max-autotune` fits on tighter-memory GPUs; the configured serving residency is restored before real traffic. Skipped under existing layerwise offload, Cache-DiT, or FSDP. -- `--enable-breakable-cuda-graph {true|false}`: capture supported DiT forwards as breakable CUDA graph segments to reduce launch overhead. Requires `--warmup-resolutions` for every served resolution because each resolution is captured separately. +- `--enable-breakable-cuda-graph {true|false}`: capture supported DiT forwards as breakable CUDA graph segments to reduce launch overhead. Requires `--warmup-resolutions` for every served resolution because each resolution is captured separately. A `quality=high` request is rejected when it would mount request-scoped DiT fusions that were not present during lossless graph capture; VAE-only high-quality paths remain compatible. - `--bcg-text-buckets {N...}`: prompt-length padding buckets for breakable CUDA graph capture/replay reuse. - `--attention-backend {BACKEND}`: attention backend for native SGLang and diffusers pipelines - `--component-attention-backends {MAP}`: per-component attention backend overrides, for example `text_encoder=torch_sdpa,transformer=fa` diff --git a/docs/docs/sglang-diffusion/fused_kernels.mdx b/docs/docs/sglang-diffusion/fused_kernels.mdx index e2ef68bf3..af3201324 100644 --- a/docs/docs/sglang-diffusion/fused_kernels.mdx +++ b/docs/docs/sglang-diffusion/fused_kernels.mdx @@ -48,6 +48,15 @@ The `quality` field in a **video response** body is unrelated. It is Sora-compat `quality` participates in the dynamic-batch signature, so mixed-quality traffic is batched separately and the transition happens safely at a batch boundary. Mounting is all-or-nothing: if any marked site on a transformer fails its static guards, no site on that transformer is fused. + +Do not combine request-gated DiT fusions with `--enable-breakable-cuda-graph`. +BCG warmup captures the lossless module branches before a high-quality request +mounts its DiT fusions, so replay would bypass the requested kernels. SGLang +rejects this combination for models with eligible DiT quality sites. Models +whose high-quality path changes only VAE decode remain allowed because BCG +captures the DiT only. + + These fusion families mount under `quality="high"`: | Fusion | What it folds | diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 8bfb544ed..6b787333f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -659,6 +659,19 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): mounted_fusions.add(description) else: unmount(transformer) + + if want and mounted_fusions and self.server_args.enable_breakable_cuda_graph: + for transformer in filter(None, [self.transformer, self.transformer_2]): + for _, _, unmount in _QUALITY_FUSION_HANDLERS: + unmount(transformer) + descriptions = ", ".join(sorted(mounted_fusions)) + raise ValueError( + "quality='high' cannot be used with breakable CUDA graphs for " + f"this model because its request-scoped DiT fusions " + f"({descriptions}) do not match the lossless warmup graphs. " + "Disable breakable CUDA graphs or use quality='lossless'." + ) + self._quality_fusions_mounted = want for description in sorted(mounted_fusions): logger.info("Mounted %s for quality=high", description) diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py index 8dfb7e69b..79e00feb3 100644 --- a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py +++ b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py @@ -58,6 +58,44 @@ class SanaVideoTransformer3DModel(torch.nn.Module): pass +class TestQualityFusionBCGCompatibility(unittest.TestCase): + def setUp(self): + self.stage = DenoisingStage.__new__(DenoisingStage) + self.stage.server_args = SimpleNamespace(enable_breakable_cuda_graph=True) + self.stage.transformer = OtherTransformer2DModel() + self.stage.transformer_2 = None + self.stage._quality_fusions_mounted = False + + @staticmethod + def _batch(quality: str): + return SimpleNamespace(sampling_params=SimpleNamespace(quality=quality)) + + def test_rejects_high_when_dit_fusion_would_replace_captured_graph(self): + unmounted = [] + handlers = ( + ( + "test fusion", + lambda _: True, + lambda transformer: unmounted.append(transformer), + ), + ) + + with patch.object(denoising_module, "_QUALITY_FUSION_HANDLERS", handlers): + with self.assertRaisesRegex(ValueError, "lossless warmup graphs"): + self.stage._maybe_toggle_quality_fusions(self._batch("high")) + + self.assertEqual(unmounted, [self.stage.transformer]) + self.assertFalse(self.stage._quality_fusions_mounted) + + def test_allows_high_when_model_has_no_dit_quality_fusions(self): + handlers = (("test fusion", lambda _: False, lambda _: None),) + + with patch.object(denoising_module, "_QUALITY_FUSION_HANDLERS", handlers): + self.stage._maybe_toggle_quality_fusions(self._batch("high")) + + self.assertTrue(self.stage._quality_fusions_mounted) + + def _fake_cache_dit_batch(*, is_warmup: bool) -> SimpleNamespace: return SimpleNamespace( is_warmup=is_warmup,