From 030496eb06472f76fcb11de53d93f10cefb4604f Mon Sep 17 00:00:00 2001
From: HuangJi <32611516+IPostYellow@users.noreply.github.com>
Date: Wed, 14 Jan 2026 14:44:10 +0800
Subject: [PATCH] [diffusion] fix: fix --warmup-resolutions' conflict with
 CacheDiT (#16962)

---
 .../sglang/multimodal_gen/runtime/managers/scheduler.py   | 4 +---
 .../runtime/pipelines_core/stages/denoising.py            | 8 ++++----
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py
index 962840768..1944ad789 100644
--- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py
+++ b/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
 
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 a22fa7f32..dd0a721b4 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
+++ b/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(