From ca80c19b55092945f70629734a3f4b14adde1e7f Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Sun, 4 Jan 2026 18:44:05 -0800 Subject: [PATCH] Revert "[diffusion] feat: support warmup with resolutions" (#16433) --- .../entrypoints/diffusion_generator.py | 2 +- .../runtime/managers/scheduler.py | 82 +++++-------------- .../pipelines_core/composed_pipeline_base.py | 11 ++- .../executors/parallel_executor.py | 36 ++++---- .../executors/pipeline_executor.py | 4 +- .../pipelines_core/executors/sync_executor.py | 5 +- .../runtime/pipelines_core/schedule_batch.py | 5 +- .../runtime/pipelines_core/stages/base.py | 1 - .../stages/timestep_preparation.py | 13 --- .../multimodal_gen/runtime/server_args.py | 32 ++------ .../runtime/utils/perf_logger.py | 16 ++-- 11 files changed, 66 insertions(+), 141 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 8551e1ba5..d91804181 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/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)", diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index 52fc3ce37..0a01b806e 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/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", - output_batch.timings.total_duration_s, - ) - else: - logger.info(f"Warmup req processing failed") + logger.info( + f"Server warmup done in {GREEN}%.2f{RESET} seconds", + output_batch.timings.total_duration_s, + ) - # 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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index 02dd23631..6fa8b8f11 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -330,11 +330,10 @@ 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()), - main_process_only=True, - ) + logger.info( + "Running pipeline stages: %s", + list(self._stage_name_mapping.keys()), + main_process_only=True, + ) return self.executor.execute_with_profiling(self.stages, batch, server_args) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py index 7f04525a7..72622e254 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py +++ b/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,27 +66,28 @@ class ParallelExecutor(PipelineExecutor): # TODO: decide when to gather on main when CFG_PARALLEL -> MAIN_RANK_ONLY for stage in stages: - paradigm = stage.parallelism_type + with Timer(stage.__class__.__name__): + paradigm = stage.parallelism_type - if paradigm == StageParallelismType.MAIN_RANK_ONLY: - if rank == 0: - # Only main rank executes, others just wait + if paradigm == StageParallelismType.MAIN_RANK_ONLY: + if rank == 0: + # Only main rank executes, others just wait + batch = stage(batch, server_args) + torch.distributed.barrier() + + elif paradigm == StageParallelismType.CFG_PARALLEL: + obj_list = [batch] if rank == 0 else [] + broadcasted_list = broadcast_pyobj( + obj_list, rank=rank, dist_group=cfg_group.cpu_group, src=0 + ) + if rank != 0: + batch = broadcasted_list[0] batch = stage(batch, server_args) - torch.distributed.barrier() - elif paradigm == StageParallelismType.CFG_PARALLEL: - obj_list = [batch] if rank == 0 else [] - broadcasted_list = broadcast_pyobj( - obj_list, rank=rank, dist_group=cfg_group.cpu_group, src=0 - ) - if rank != 0: - batch = broadcasted_list[0] - batch = stage(batch, server_args) + torch.distributed.barrier() - torch.distributed.barrier() - - elif paradigm == StageParallelismType.REPLICATED: - batch = stage(batch, server_args) + elif paradigm == StageParallelismType.REPLICATED: + batch = stage(batch, server_args) return batch def execute( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py index a0f19834e..f77e06087 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py +++ b/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__( - stage_name=name, timings=None, log_stage_start_end=True, logger=logger - ) + super().__init__(stage_name=name, timings=None, simple_log=True, logger=logger) class PipelineExecutor(ABC): diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py index 9af35e9be..b7eb3ca7b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py +++ b/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: - batch = stage(batch, server_args) + with Timer(stage.__class__.__name__): + batch = stage(batch, server_args) + profiler = SGLDiffusionProfiler.get_instance() if profiler: profiler.step_stage() diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index eca3a893f..56108b290 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py index f40b6f579..5f8ef94e0 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py index 0a60cc9a8..13bc6a066 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index f76ed7c47..de823e3bd 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -206,10 +206,7 @@ class ServerArgs: # Compilation enable_torch_compile: bool = False - - # warmup - warmup: bool = False - warmup_resolutions: list[str] = None + enable_warmup: bool = False 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, - help="Perform some warmup after server starts (if `--warmup-resolutions` is specified) or before processing the first request (if `--warmup-resolutions` is not specified)." + default=ServerArgs.enable_warmup, + 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, diff --git a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py index 7c60d9352..8435571bb 100644 --- a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py +++ b/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.log_stage_start_end = log_stage_start_end + self.enabled = perf_dump_path_provided or envs.SGLANG_DIFFUSION_STAGE_LOGGING 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)