[diffusion] feat: support warmup with resolutions (#16330)

python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

@@ -268,7 +268,7 @@ class DiffGenerator:
         log_batch_completion(logger, len(results), total_gen_time)
 
         if results:
-            if self.server_args.enable_warmup:
+            if self.server_args.warmup:
                 total_duration_ms = results[0]["timings"]["total_duration_ms"]
                 logger.info(
                     f"Warmed-up request processed in {GREEN}%.2f{RESET} seconds (with warmup excluded)",

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

@@ -12,6 +12,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
     MergeLoraWeightsReq,
     SetLoraReq,
     UnmergeLoraWeightsReq,
+    _parse_size,
 )
 from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker
 from sglang.multimodal_gen.runtime.pipelines_core import Req
@@ -83,8 +84,11 @@ class Scheduler:
         # FIFO, new reqs are appended
         self.waiting_queue: deque[tuple[bytes, Req]] = deque()
 
+        # whether we've send the necessary warmup reqs
         self.warmed_up = False
 
+        self.prepare_server_warmup_reqs()
+
     def _handle_set_lora(self, reqs: List[Any]) -> OutputBatch:
         # TODO: return set status
         # TODO: return with SetLoRAResponse or something more appropriate
@@ -102,6 +106,9 @@ class Scheduler:
         return self.worker.unmerge_lora_weights(req.target)
 
     def _handle_generation(self, reqs: List[Req]):
+        has_warmup = any(req.is_warmup for req in reqs)
+        if has_warmup:
+            logger.info("Processing warmup req...")
         return self.worker.execute_forward(reqs)
 
     def return_result(
@@ -126,6 +133,49 @@ class Scheduler:
 
         return [item]
 
+    def prepare_server_warmup_reqs(self):
+        if (
+            self.server_args.warmup
+            and not self.warmed_up
+            and self.server_args.warmup_resolutions is not None
+        ):
+            # insert warmup reqs constructed with each warmup-resolution
+            for resolution in self.server_args.warmup_resolutions:
+                width, height = _parse_size(resolution)
+                req = Req(
+                    data_type=self.server_args.pipeline_config.task_type.data_type(),
+                    width=width,
+                    height=height,
+                    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
+
+    def process_received_reqs_with_req_based_warmup(
+        self, recv_reqs: List[tuple[bytes, Any]]
+    ) -> List[tuple[bytes, Any]]:
+        if (
+            self.warmed_up
+            or not self.server_args.warmup
+            or not recv_reqs
+            or self.server_args.warmup_resolutions is not None
+        ):
+            return recv_reqs
+
+        # handle server req-based warmup by inserting an identical req to the beginning of the waiting queue
+        # only the very first req through server's lifetime will be warmup
+        identity, req = recv_reqs[0]
+        if isinstance(req, Req):
+            warmup_req = deepcopy(req)
+            warmup_req.is_warmup = True
+            warmup_req.num_inference_steps = 1
+            recv_reqs.insert(0, (identity, warmup_req))
+            logger.info("Server warming up....")
+            self.warmed_up = True
+        return recv_reqs
+
     def recv_reqs(self) -> List[tuple[bytes, Any]]:
         """
         For non-main schedulers, reqs are broadcasted from main using broadcast_pyobj
@@ -178,22 +228,6 @@ class Scheduler:
 
         assert recv_reqs is not None
 
-        # handle server warmup by inserting an identical req to the beginning of the waiting queue
-        # only the very first req through server's lifetime will be warmup
-        if (
-            not self.warmed_up
-            and len(recv_reqs) == 1
-            and self.server_args.enable_warmup
-        ):
-            identity, req = recv_reqs[0]
-            if isinstance(req, Req):
-                warmup_req = deepcopy(req)
-                warmup_req.is_warmup = True
-                warmup_req.num_inference_steps = 1
-                recv_reqs.insert(0, (identity, warmup_req))
-                self.warmed_up = True
-                logger.info("Server warming up....")
-
         return recv_reqs
 
     def event_loop(self) -> None:
@@ -210,7 +244,7 @@ class Scheduler:
             # 1: receive requests
             try:
                 new_reqs = self.recv_reqs()
-                # after processing input reqs
+                new_reqs = self.process_received_reqs_with_req_based_warmup(new_reqs)
                 self.waiting_queue.extend(new_reqs)
             except Exception as e:
                 logger.error(
@@ -250,16 +284,20 @@ class Scheduler:
 
             # 3. return results
             try:
-                # TODO: Support sending back to multiple identities if batched
+                # log warmup info
                 is_warmup = (
                     processed_req.is_warmup if isinstance(processed_req, Req) else False
                 )
                 if is_warmup:
+                    if output_batch.error is None:
                         logger.info(
-                        f"Server warmup done in {GREEN}%.2f{RESET} seconds",
+                            f"Warmup req processed in {GREEN}%.2f{RESET} seconds",
                             output_batch.timings.total_duration_s,
                         )
+                    else:
+                        logger.info(f"Warmup req processing failed")
 
+                # TODO: Support sending back to multiple identities if batched
                 self.return_result(output_batch, identities[0], is_warmup=is_warmup)
             except zmq.ZMQError as e:
                 # Reply failed; log and keep loop alive to accept future requests

python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py

@@ -330,6 +330,7 @@ class ComposedPipelineBase(ABC):
         batch.log(server_args=server_args)
 
         # Execute each stage
+        if not batch.is_warmup:
             logger.info(
                 "Running pipeline stages: %s",
                 list(self._stage_name_mapping.keys()),

python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py

@@ -12,7 +12,6 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
 from sglang.multimodal_gen.runtime.pipelines_core import Req
 from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import (
     PipelineExecutor,
-    Timer,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
@@ -66,7 +65,6 @@ class ParallelExecutor(PipelineExecutor):
 
         # TODO: decide when to gather on main when CFG_PARALLEL -> MAIN_RANK_ONLY
         for stage in stages:
-            with Timer(stage.__class__.__name__):
             paradigm = stage.parallelism_type
 
             if paradigm == StageParallelismType.MAIN_RANK_ONLY:

python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py

@@ -30,7 +30,9 @@ class Timer(StageProfiler):
     """
 
     def __init__(self, name="Stage"):
-        super().__init__(stage_name=name, timings=None, simple_log=True, logger=logger)
+        super().__init__(
+            stage_name=name, timings=None, log_stage_start_end=True, logger=logger
+        )
 
 
 class PipelineExecutor(ABC):

python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py

@@ -9,7 +9,6 @@ from typing import List
 from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import (
     PipelineExecutor,
     SGLDiffusionProfiler,
-    Timer,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages import PipelineStage
@@ -31,9 +30,7 @@ class SyncExecutor(PipelineExecutor):
         Execute all pipeline stages sequentially.
         """
         for stage in stages:
-            with Timer(stage.__class__.__name__):
             batch = stage(batch, server_args)
-
             profiler = SGLDiffusionProfiler.get_instance()
             if profiler:
                 profiler.step_stage()

python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py

@@ -88,7 +88,7 @@ class Req:
 
     # Batch info
     num_outputs_per_prompt: int = 1
-    seed: int | None = None
+    seed: int | None = 42
     seeds: list[int] | None = None
     generator_device: str = (
         "cuda"  # Device for random generator: "cuda", "musa" or "cpu"
@@ -228,6 +228,9 @@ class Req:
 
         self.timings = RequestTimings(request_id=self.request_id)
 
+        if self.is_warmup:
+            self.num_inference_steps = 1
+
     def adjust_size(self, server_args: ServerArgs):
         pass
 

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

@@ -206,6 +206,7 @@ class PipelineStage(ABC):
             logger=logger,
             timings=batch.timings,
             perf_dump_path_provided=batch.perf_dump_path is not None,
+            log_stage_start_end=not batch.is_warmup,
         ):
             result = self.forward(batch, server_args)
 

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

@@ -10,6 +10,8 @@ This module contains implementations of timestep preparation stages for diffusio
 import inspect
 from typing import Any, Callable, Tuple
 
+import torch
+
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
@@ -144,6 +146,17 @@ class TimestepPreparationStage(PipelineStage):
 
     def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
         """Verify timestep preparation stage outputs."""
+        if (
+            batch.is_warmup
+            and isinstance(batch.timesteps, torch.Tensor)
+            and torch.isnan(batch.timesteps).any()
+        ):
+            # when num-inference-steps == 1, the last sigma being 1, the 1 / last_sigma could be nan
+            # this a workaround for warmup req only
+            batch.timesteps = torch.ones(
+                (1,), dtype=torch.float32, device=get_local_torch_device()
+            )
+
         result = VerificationResult()
         result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.with_dims(1)])
         return result

python/sglang/multimodal_gen/runtime/server_args.py

@@ -206,7 +206,10 @@ class ServerArgs:
 
     # Compilation
     enable_torch_compile: bool = False
-    enable_warmup: bool = False
+
+    # warmup
+    warmup: bool = False
+    warmup_resolutions: list[str] = None
 
     disable_autocast: bool | None = None
 
@@ -291,6 +294,15 @@ class ServerArgs:
         if self.attention_backend in ["fa3", "fa4"]:
             self.attention_backend = "fa"
 
+        # handle warmup
+        if self.warmup_resolutions is not None:
+            self.warmup = True
+
+        if self.warmup:
+            logger.info(
+                "Warmup enabled, the launch time is expected to be longer than usual"
+            )
+
         # network initialization: port and host
         self.port = self.settle_port(self.port)
         # Add randomization to avoid race condition when multiple servers start simultaneously
@@ -457,14 +469,24 @@ class ServerArgs:
             help="Use torch.compile to speed up DiT inference."
             + "However, will likely cause precision drifts. See (https://github.com/pytorch/pytorch/issues/145213)",
         )
+
+        # warmup
         parser.add_argument(
-            "--enable-warmup",
+            "--warmup",
             action=StoreBoolean,
-            default=ServerArgs.enable_warmup,
+            default=ServerArgs.warmup,
-            help="Perform a 1-step end-to-end warmup request before the actual request. "
+            help="Perform some warmup after server starts (if `--warmup-resolutions` is specified) or before processing the first request (if `--warmup-resolutions` is not specified)."
             "Recommended to enable when benchmarking to ensure fair comparison and best performance."
-            "When enabled, look for the line ending with `with warmup excluded` for actual processing time.",
+            "When enabled with `--warmup-resolutions` unspecified, look for the line ending with `(with warmup excluded)` for actual processing time.",
         )
+        parser.add_argument(
+            "--warmup-resolutions",
+            type=str,
+            nargs="+",
+            default=ServerArgs.warmup_resolutions,
+            help="Specify resolutions for server to warmup. e.g., `--warmup-resolutions 256x256, 720x720`",
+        )
+
         parser.add_argument(
             "--dit-cpu-offload",
             action=StoreBoolean,

python/sglang/multimodal_gen/runtime/utils/perf_logger.py

@@ -135,21 +135,21 @@ class StageProfiler:
         stage_name: str,
         logger: _SGLDiffusionLogger,
         timings: Optional["RequestTimings"],
-        simple_log: bool = False,
+        log_stage_start_end: bool = False,
         perf_dump_path_provided: bool = False,
     ):
         self.stage_name = stage_name
         self.timings = timings
         self.logger = logger
-        self.simple_log = simple_log
         self.start_time = 0.0
-        self.enabled = perf_dump_path_provided or envs.SGLANG_DIFFUSION_STAGE_LOGGING
+        self.log_timing = perf_dump_path_provided or envs.SGLANG_DIFFUSION_STAGE_LOGGING
+        self.log_stage_start_end = log_stage_start_end
 
     def __enter__(self):
-        if self.simple_log:
+        if self.log_stage_start_end:
             self.logger.info(f"[{self.stage_name}] started...")
 
-        if (self.enabled and self.timings) or self.simple_log:
+        if (self.log_timing and self.timings) or self.log_stage_start_end:
             if (
                 os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
                 and self.stage_name.startswith("denoising_step_")
@@ -161,7 +161,7 @@ class StageProfiler:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        if not ((self.enabled and self.timings) or self.simple_log):
+        if not ((self.log_timing and self.timings) or self.log_stage_start_end):
             return False
 
         if (
@@ -182,12 +182,12 @@ class StageProfiler:
             )
             return False
 
-        if self.simple_log:
+        if self.log_stage_start_end:
             self.logger.info(
                 f"[{self.stage_name}] finished in {execution_time_s:.4f} seconds",
             )
 
-        if self.enabled and self.timings:
+        if self.log_timing and self.timings:
             if "denoising_step_" in self.stage_name:
                 index = int(self.stage_name[len("denoising_step_") :])
                 self.timings.record_steps(index, execution_time_s)

python/sglang/srt/model_executor/cuda_graph_runner.py

@@ -534,7 +534,7 @@ class CudaGraphRunner:
         )
         graph_fn = (
             partial(memory_saver_adapter.cuda_graph, tag=GPU_MEMORY_TYPE_CUDA_GRAPH)
-            if memory_saver_adapter.enabled
+            if memory_saver_adapter.log_timing
             else self.device_module.graph
         )
         with graph_fn(cuda_graph=graph, pool=pool, stream=stream):