diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index d91804181..8551e1ba5 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.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)", diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index 0a01b806e..52fc3ce37 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/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: - logger.info( - f"Server warmup done in {GREEN}%.2f{RESET} seconds", - output_batch.timings.total_duration_s, - ) + 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") + # 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 6fa8b8f11..02dd23631 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,10 +330,11 @@ class ComposedPipelineBase(ABC): batch.log(server_args=server_args) # Execute each stage - logger.info( - "Running pipeline stages: %s", - list(self._stage_name_mapping.keys()), - main_process_only=True, - ) + if not batch.is_warmup: + 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 72622e254..7f04525a7 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,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,28 +65,27 @@ 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 + paradigm = stage.parallelism_type - 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] + if paradigm == StageParallelismType.MAIN_RANK_ONLY: + if rank == 0: + # Only main rank executes, others just wait batch = stage(batch, server_args) + torch.distributed.barrier() - 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) - elif paradigm == StageParallelismType.REPLICATED: - batch = stage(batch, server_args) + torch.distributed.barrier() + + 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 f77e06087..a0f19834e 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,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): 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 b7eb3ca7b..9af35e9be 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,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) - + 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 56108b290..eca3a893f 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 = 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 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 5f8ef94e0..f40b6f579 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py +++ b/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) 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 13bc6a066..0a60cc9a8 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,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 diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index de823e3bd..f76ed7c47 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/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, - help="Perform a 1-step end-to-end warmup request before the actual request. " + 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)." "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, diff --git a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py index 8435571bb..7c60d9352 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"], - 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)