Revert "[diffusion] feat: support warmup with resolutions" (#16433)

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.warmup:
+            if self.server_args.enable_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,7 +12,6 @@ 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
@@ -84,11 +83,8 @@ 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
@@ -106,9 +102,6 @@ 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(
@@ -133,49 +126,6 @@ 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
@@ -228,6 +178,22 @@ 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:
@@ -244,7 +210,7 @@ class Scheduler:
             # 1: receive requests
             try:
                 new_reqs = self.recv_reqs()
-                new_reqs = self.process_received_reqs_with_req_based_warmup(new_reqs)
+                # after processing input reqs
                 self.waiting_queue.extend(new_reqs)
             except Exception as e:
                 logger.error(
@@ -284,20 +250,16 @@ class Scheduler:
 
             # 3. return results
             try:
-                # log warmup info
+                # TODO: Support sending back to multiple identities if batched
                 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"Warmup req processed in {GREEN}%.2f{RESET} seconds",
+                        f"Server warmup done 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,7 +330,6 @@ 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,6 +12,7 @@ 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 (
@@ -65,6 +66,7 @@ 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,9 +30,7 @@ class Timer(StageProfiler):
     """
 
     def __init__(self, name="Stage"):
-        super().__init__(
+        super().__init__(stage_name=name, timings=None, simple_log=True, logger=logger)
-            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,6 +9,7 @@ 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
@@ -30,7 +31,9 @@ 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 = 42
+    seed: int | None = None
     seeds: list[int] | None = None
     generator_device: str = (
         "cuda"  # Device for random generator: "cuda", "musa" or "cpu"
@@ -228,9 +228,6 @@ 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,7 +206,6 @@ 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,8 +10,6 @@ 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 (
@@ -146,17 +144,6 @@ 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,10 +206,7 @@ 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
 
@@ -294,15 +291,6 @@ 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
@@ -469,24 +457,14 @@ 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(
-            "--warmup",
+            "--enable-warmup",
             action=StoreBoolean,
-            default=ServerArgs.warmup,
+            default=ServerArgs.enable_warmup,
-            help="Perform some warmup after server starts (if `--warmup-resolutions` is specified) or before processing the first request (if `--warmup-resolutions` is not specified)."
+            help="Perform a 1-step end-to-end warmup request before the actual request. "
             "Recommended to enable when benchmarking to ensure fair comparison and best performance."
-            "When enabled with `--warmup-resolutions` unspecified, look for the line ending with `(with warmup excluded)` for actual processing time.",
+            "When enabled, 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"],
-        log_stage_start_end: bool = False,
+        simple_log: 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.log_timing = perf_dump_path_provided or envs.SGLANG_DIFFUSION_STAGE_LOGGING
+        self.enabled = perf_dump_path_provided or envs.SGLANG_DIFFUSION_STAGE_LOGGING
-        self.log_stage_start_end = log_stage_start_end
 
     def __enter__(self):
-        if self.log_stage_start_end:
+        if self.simple_log:
             self.logger.info(f"[{self.stage_name}] started...")
 
-        if (self.log_timing and self.timings) or self.log_stage_start_end:
+        if (self.enabled and self.timings) or self.simple_log:
             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.log_timing and self.timings) or self.log_stage_start_end):
+        if not ((self.enabled and self.timings) or self.simple_log):
             return False
 
         if (
@@ -182,12 +182,12 @@ class StageProfiler:
             )
             return False
 
-        if self.log_stage_start_end:
+        if self.simple_log:
             self.logger.info(
                 f"[{self.stage_name}] finished in {execution_time_s:.4f} seconds",
             )
 
-        if self.log_timing and self.timings:
+        if self.enabled 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)