[diffusion] fix: fix --warmup-resolutions' conflict with CacheDiT (#16962)

python/sglang/multimodal_gen/runtime/managers/scheduler.py

@@ -202,7 +202,6 @@ class Scheduler:
                         prompt="",
                         is_warmup=True,
                     )
-
                 self.waiting_queue.append((None, req))
             # if server is warmed-up, set this flag to avoid req-based warmup
             self.warmed_up = True
@@ -226,8 +225,7 @@ class Scheduler:
             warmup_req.set_as_warmup()
             recv_reqs.insert(0, (identity, warmup_req))
             self._warmup_total = 1
-            self._warmup_processed = 1
-            logger.info("Processing warmup req... (1/1)")
+            self._warmup_processed = 0
             self.warmed_up = True
         return recv_reqs
 

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

@@ -137,7 +137,7 @@ class DenoisingStage(PipelineStage):
         setattr(module, "forward", compiled_forward)
         return module
 
-    def _maybe_enable_cache_dit(self, num_inference_steps: int) -> None:
+    def _maybe_enable_cache_dit(self, num_inference_steps: int, batch: Req) -> None:
         """Enable cache-dit on the transformers if configured (idempotent).
 
         This method should be called after the transformer is fully loaded
@@ -158,7 +158,7 @@ class DenoisingStage(PipelineStage):
                 )
             return
         # check if cache-dit is enabled in config
-        if not envs.SGLANG_CACHE_DIT_ENABLED:
+        if not envs.SGLANG_CACHE_DIT_ENABLED or batch.is_warmup:
             return
 
         from sglang.multimodal_gen.runtime.distributed import (
@@ -496,13 +496,13 @@ class DenoisingStage(PipelineStage):
                 server_args.model_paths["transformer"], server_args, "transformer"
             )
             # enable cache-dit before torch.compile (delayed mounting)
-            self._maybe_enable_cache_dit(cache_dit_num_inference_steps)
+            self._maybe_enable_cache_dit(cache_dit_num_inference_steps, batch)
             self.compile_module_with_torch_compile(self.transformer)
             if pipeline:
                 pipeline.add_module("transformer", self.transformer)
             server_args.model_loaded["transformer"] = True
         else:
-            self._maybe_enable_cache_dit(cache_dit_num_inference_steps)
+            self._maybe_enable_cache_dit(cache_dit_num_inference_steps, batch)
 
         # Prepare extra step kwargs for scheduler
         extra_step_kwargs = self.prepare_extra_func_kwargs(