[diffusion] fix: warn when bcg disables cache-dit (#34242)

Co-authored-by: chilltongx <284668524+chilltongx@users.noreply.github.com>

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -671,6 +671,11 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         if self.server_args.enable_breakable_cuda_graph:
             # Cache-DiT wraps transformer.forward with step-skipping control
             # flow that must not be baked into a captured CUDA graph.
+            if self._cache_dit_requested():
+                logger.warning_once(
+                    "Cache-DiT was requested but is disabled because breakable "
+                    "CUDA graphs are enabled."
+                )
             return
         # NOTE: When a new request arrives, we need to refresh the cache-dit context.
         if self._cache_dit_enabled:

python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py

@@ -16,6 +16,9 @@ from sglang.multimodal_gen.runtime.layers.attention import (
 from sglang.multimodal_gen.runtime.models.dits.qwen_image import (
     _attn_mask_meta_local_pad,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages import (
+    denoising as denoising_module,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import (
     DenoisingStage,
 )
@@ -23,6 +26,7 @@ from sglang.multimodal_gen.runtime.server_args import (
     BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS,
     BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS,
 )
+from sglang.multimodal_gen.runtime.utils.logging_utils import _print_warning_once
 
 
 class QwenImageTransformer2DModel(torch.nn.Module):
@@ -428,6 +432,38 @@ class TestDiffusionBCGPadding(unittest.TestCase):
         self.stage._maybe_enable_cache_dit(1, SimpleNamespace(is_warmup=True))
         self.assertEqual(self.stage._bcg_runners, {})
 
+    def test_bcg_warns_when_cache_dit_is_requested(self):
+        self.stage.server_args = SimpleNamespace(enable_breakable_cuda_graph=True)
+        _print_warning_once.cache_clear()
+        self.addCleanup(_print_warning_once.cache_clear)
+
+        with (
+            patch.object(self.stage, "_cache_dit_requested", return_value=True),
+            patch.object(denoising_module.logger, "warning") as warning,
+        ):
+            self.stage._maybe_enable_cache_dit(1, SimpleNamespace(is_warmup=True))
+            self.stage._maybe_enable_cache_dit(1, SimpleNamespace(is_warmup=False))
+
+        warning.assert_called_once_with(
+            "Cache-DiT was requested but is disabled because breakable CUDA "
+            "graphs are enabled.",
+            stacklevel=2,
+        )
+
+    def test_bcg_does_not_warn_when_cache_dit_is_not_requested(self):
+        self.stage.server_args = SimpleNamespace(enable_breakable_cuda_graph=True)
+
+        with (
+            patch.object(self.stage, "_cache_dit_requested", return_value=False),
+            patch(
+                "sglang.multimodal_gen.runtime.pipelines_core.stages.denoising."
+                "logger.warning_once"
+            ) as warning_once,
+        ):
+            self.stage._maybe_enable_cache_dit(1, SimpleNamespace(is_warmup=True))
+
+        warning_once.assert_not_called()
+
     def test_bcg_runner_cache_is_per_model_module(self):
         self.stage.server_args = SimpleNamespace(enable_breakable_cuda_graph=True)
         self.stage._bcg_runners = {}

python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py

@@ -29,6 +29,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.m
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.resolved_plan import (
     minimax_h3_resolve_plan,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.stages.denoising import (
+    MiniMaxH3DenoisingStage,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import (
     partition_for_task,
 )
@@ -256,6 +259,29 @@ def _quality_server_args():
     )
 
 
+def test_high_quality_request_warns_when_bcg_suppresses_cache_dit():
+    stage = MiniMaxH3DenoisingStage.__new__(MiniMaxH3DenoisingStage)
+    stage.server_args = SimpleNamespace(enable_breakable_cuda_graph=True)
+    stage._cache_dit_enabled = False
+    batch = SimpleNamespace(
+        sampling_params=SimpleNamespace(
+            quality="high",
+            _explicit_fields={"quality"},
+        )
+    )
+
+    with patch(
+        "sglang.multimodal_gen.runtime.pipelines_core.stages.denoising."
+        "logger.warning_once"
+    ) as warning_once:
+        stage._maybe_enable_cache_dit(50, batch)
+
+    warning_once.assert_called_once_with(
+        "Cache-DiT was requested but is disabled because breakable CUDA graphs "
+        "are enabled."
+    )
+
+
 def test_quality_admission_fails_closed_outside_validated_request():
     metadata = MiniMaxH3ReleaseMetadata.from_model_index(
         {