chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Kangyan-Zhou
2026-01-04 18:44:05 -08:00
committed by GitHub
parent 1e7b326482
commit ca80c19b55
11 changed files with 66 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
View File

@@ -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)",

82
python/sglang/multimodal_gen/runtime/managers/scheduler.py
View File

@@ -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

11
python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
View File

@@ -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)

36
python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py
View File

@@ -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(

4
python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py
View File

@@ -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):

5
python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py
View File

@@ -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()

5
python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py
View File

@@ -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

1
python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py
View File

@@ -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)

13
python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py
View File

@@ -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

32
python/sglang/multimodal_gen/runtime/server_args.py
View File

@@ -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,

16
python/sglang/multimodal_gen/runtime/utils/perf_logger.py
View File

@@ -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)