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 1d5977532..4c17f0e10 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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: 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 7d774c209..07ab70cb0 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 @@ -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 = {} diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py index 8bc4b08dd..170f12bbd 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py +++ b/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( {