[diffusion] fix: fix diffusion serve warmup defaults (#26247)

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -438,6 +438,10 @@ class PipelineConfig:
     def maybe_prepare_latent_ids(self, latents):
         return None
 
+    # called before vae encode
+    def preprocess_vae_encode(self, image, vae):
+        return image
+
     # called after vae encode
     def postprocess_vae_encode(self, image_latents, vae):
         return image_latents

python/sglang/multimodal_gen/configs/pipeline_configs/flux_finetuned.py

@@ -36,6 +36,19 @@ class Flux2FinetunedPipelineConfig(Flux2PipelineConfig):
     - 5D latents support for both single-frame and multi-frame generation
     """
 
+    def preprocess_vae_encode(self, image: torch.Tensor, vae) -> torch.Tensor:
+        if image.ndim == 5 and image.shape[2] == 1 and not self._check_vae_has_bn(vae):
+            return image.squeeze(2)
+        return image
+
+    def postprocess_vae_encode(self, image_latents: torch.Tensor, vae) -> torch.Tensor:
+        if (
+            not self._check_vae_has_bn(vae)
+            and image_latents.shape[1] == self.dit_config.arch_config.in_channels
+        ):
+            return image_latents
+        return super().postprocess_vae_encode(image_latents, vae)
+
     def preprocess_decoding(
         self, latents: torch.Tensor, server_args=None, vae=None
     ) -> torch.Tensor:

python/sglang/multimodal_gen/runtime/entrypoints/cli/serve.py

@@ -12,11 +12,8 @@ from sglang.multimodal_gen.runtime.launch_server import (
     dispatch_launch,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import FlexibleArgumentParser
 
-logger = init_logger(__name__)
-
 
 def add_multimodal_gen_serve_args(parser: argparse.ArgumentParser):
     """Add the arguments for the serve command."""
@@ -32,10 +29,9 @@ def add_multimodal_gen_serve_args(parser: argparse.ArgumentParser):
 
 def execute_serve_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None):
     """The entry point for the serve command."""
-    server_args = ServerArgs.from_cli_args(args, unknown_args)
-    if not server_args.is_arg_explicitly_set("warmup"):
-        server_args.warmup = True
-        logger.info("Warmup is enabled by default for sglang serve.")
+    server_args = ServerArgs.from_cli_args(
+        args, unknown_args, default_args={"warmup": True, "server_warmup": True}
+    )
 
     dispatch_launch(server_args)
 

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

@@ -3,10 +3,12 @@
 import asyncio
 import base64
 import os
+import signal
 import uuid
-from contextlib import asynccontextmanager
+from contextlib import asynccontextmanager, suppress
 from typing import TYPE_CHECKING
 
+import httpx
 import torch
 from fastapi import APIRouter, FastAPI, Request
 
@@ -26,6 +28,11 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import (
 )
 from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
 from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args
+from sglang.multimodal_gen.runtime.server_warmup import (
+    build_warmup_reqs,
+    prepare_warmup_image_path,
+    should_include_warmup_image,
+)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.srt.utils.json_response import orjson_response
 from sglang.version import __version__
@@ -36,6 +43,67 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 VERTEX_ROUTE = os.environ.get("AIP_PREDICT_ROUTE", "/vertex_generate")
+SERVER_WARMUP_BYPASS_PATHS = (
+    "/health",
+    "/health_generate",
+    "/model_info",
+    "/server_info",
+)
+
+
+async def _wait_until_http_ready(server_args: ServerArgs) -> None:
+    """for server warmup"""
+    health_url = f"{server_args.url()}/health"
+    async with httpx.AsyncClient() as client:
+        for _ in range(120):
+            try:
+                response = await client.get(health_url, timeout=5.0)
+                if response.status_code == 200:
+                    return
+            except httpx.HTTPError:
+                pass
+            await asyncio.sleep(1.0)
+    raise RuntimeError(f"HTTP server did not become ready at {health_url}")
+
+
+async def _run_server_warmup_after_http_ready(
+    server_args: ServerArgs, warmup_done: asyncio.Event
+) -> None:
+    try:
+        if (
+            not server_args.warmup
+            or not server_args.server_warmup
+            or server_args.warmup_resolutions is not None
+        ):
+            warmup_done.set()
+            return
+
+        await _wait_until_http_ready(server_args)
+
+        warmup_input_path = None
+        if should_include_warmup_image(server_args, server_based_warmup=True):
+            warmup_input_path = await prepare_warmup_image_path(server_args)
+
+        warmup_reqs = build_warmup_reqs(
+            server_args,
+            warmup_resolutions=None,
+            warmup_input_path=warmup_input_path,
+            return_warmup_result=True,
+            server_based_warmup=True,
+            use_model_sampling_defaults=True,
+        )
+        for req in warmup_reqs:
+            response = await async_scheduler_client.forward(req)
+            if response.error is not None:
+                raise RuntimeError(response.error)
+
+        logger.info("The server is fired up and ready to roll!")
+        warmup_done.set()
+    except asyncio.CancelledError:
+        raise
+    except Exception as e:
+        logger.error("Server warmup failed; aborting startup: %s", e, exc_info=True)
+        os.kill(os.getpid(), signal.SIGTERM)
 
 
 @asynccontextmanager
@@ -48,16 +116,31 @@ async def lifespan(app: FastAPI):
     # 1. Initialize the singleton client that connects to the backend Scheduler
     server_args = app.state.server_args
     async_scheduler_client.initialize(server_args)
+    warmup_done = asyncio.Event()
+    app.state.server_warmup_done = warmup_done
 
     # 2. Start the ZMQ Broker in the background to handle offline requests
     broker_task = asyncio.create_task(run_zeromq_broker(server_args))
+    warmup_task = None
+    if server_args.server_warmup:
+        warmup_task = asyncio.create_task(
+            _run_server_warmup_after_http_ready(server_args, warmup_done)
+        )
+    else:
+        warmup_done.set()
 
-    yield
+    try:
+        yield
+    finally:
+        if warmup_task is not None and not warmup_task.done():
+            warmup_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await warmup_task
 
-    # On shutdown
-    logger.info("FastAPI app is shutting down...")
-    broker_task.cancel()
-    async_scheduler_client.close()
+        # On shutdown
+        logger.info("FastAPI app is shutting down...")
+        broker_task.cancel()
+        async_scheduler_client.close()
 
 
 # Health router
@@ -299,6 +382,17 @@ def create_app(server_args: ServerArgs):
     """
     app = FastAPI(lifespan=lifespan)
 
+    @app.middleware("http")
+    async def wait_for_server_warmup(request: Request, call_next):
+        warmup_done = getattr(request.app.state, "server_warmup_done", None)
+        if (
+            warmup_done is not None
+            and not warmup_done.is_set()
+            and request.url.path not in SERVER_WARMUP_BYPASS_PATHS
+        ):
+            await warmup_done.wait()
+        return await call_next(request)
+
     app.include_router(health_router)
     app.include_router(vertex_router)
 

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

@@ -423,8 +423,6 @@ def prepare_request(
     if diffusers_kwargs and "max_sequence_length" in diffusers_kwargs:
         req.max_sequence_length = diffusers_kwargs["max_sequence_length"]
 
-    req.adjust_size(server_args)
-
     if not isinstance(req.prompt, str):
         raise TypeError(f"`prompt` must be a string, but got {type(req.prompt)}")
 

python/sglang/multimodal_gen/runtime/launch_server.py

@@ -318,6 +318,7 @@ def launch_pool_disagg_server(
                 "pool_result_endpoint": result_ep,
                 "num_gpus": num_role_gpus,
                 "warmup": role_type == RoleType.ENCODER,
+                "server_warmup": False,
                 "scheduler_port": find_port(port_cursor),
                 "master_port": find_port(port_cursor + 100),
                 # Per-role parallelism (None = auto-derive from num_gpus)
@@ -590,6 +591,7 @@ def launch_disagg_role(server_args: ServerArgs):
         "pool_work_endpoint": work_endpoint,
         "pool_result_endpoint": result_endpoint,
         "warmup": role_type == RoleType.ENCODER,
+        "server_warmup": False,
         "scheduler_port": internal_scheduler_port,
         # Per-role parallelism (None = auto-derive from num_gpus)
         "tp_size": role_par["tp_size"],

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

@@ -1,11 +1,8 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
 
 # SPDX-License-Identifier: Apache-2.0
-import asyncio
 import dataclasses
-import os
 import pickle
-import tempfile
 import time
 from collections import deque
 from contextlib import contextmanager
@@ -20,10 +17,6 @@ from sglang.multimodal_gen.runtime.disaggregation.scheduler_mixin import (
     SchedulerDisaggMixin,
 )
 from sglang.multimodal_gen.runtime.distributed import get_world_group
-from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
-    _parse_size,
-    save_image_to_path,
-)
 from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import (
     GetWeightsChecksumReqInput,
     UpdateWeightFromDiskReqInput,
@@ -55,6 +48,15 @@ from sglang.multimodal_gen.runtime.server_args import (
     ServerArgs,
     set_global_server_args,
 )
+from sglang.multimodal_gen.runtime.server_warmup import (
+    build_warmup_reqs,
+    get_first_generation_req,
+    is_server_based_warmup,
+    is_warmup_req,
+    prepare_warmup_image_path_sync,
+    should_include_warmup_image,
+    should_return_warmup_result,
+)
 from sglang.multimodal_gen.runtime.utils.common import get_zmq_socket
 from sglang.multimodal_gen.runtime.utils.distributed import broadcast_pyobj
 from sglang.multimodal_gen.runtime.utils.logging_utils import GREEN, RESET, init_logger
@@ -62,15 +64,6 @@ from sglang.multimodal_gen.runtime.utils.trace_wrapper import DiffStage, trace_s
 
 logger = init_logger(__name__)
 
-MINIMUM_PICTURE_BASE64_FOR_WARMUP = "data:image/jpg;base64,iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAAACXBIWXMAAA7EAAAOxAGVKw4bAAAAbUlEQVRYhe3VsQ2AMAxE0Y/lIgNQULD/OqyCMgCihCKSG4yRuKuiNH6JLsoEbMACOGBcua9HOR7Y6w6swBwMy0qLTpkeI77qdEBpBFAHBBDAGH8WrwJKI4AAegUCfAKgEgpQDvh3CR3oQCuav58qlAw73kKCSgAAAABJRU5ErkJggg=="
-
-# Placeholder negative_prompt used in synthesized warmup Reqs when
-# --enable-cfg-parallel is on. A non-empty, real word (vs "" or " ") so
-# every tokenizer backend emits a predictable, non-degenerate token
-# sequence — rank 1's uncond branch then produces a valid tensor for
-# _combine_cfg_parallel's all-reduce.
-DEFAULT_PLACEHOLDER_PROMPT = "warmup"
-
 _MAX_RECV_REQS_PER_POLL = 1024
 _BATCH_METRICS_LOG_INTERVAL = 5
 
@@ -237,22 +230,6 @@ class Scheduler(SchedulerDisaggMixin):
             return reqs[0]
         return reqs
 
-    @staticmethod
-    def _first_generation_req(req_or_group: Any) -> Req | None:
-        """Extract the first req"""
-        if isinstance(req_or_group, Req):
-            return req_or_group
-        if isinstance(req_or_group, list) and req_or_group:
-            first_req = req_or_group[0]
-            if isinstance(first_req, Req):
-                return first_req
-        return None
-
-    @classmethod
-    def _is_warmup_item(cls, req_or_group: Any) -> bool:
-        req = cls._first_generation_req(req_or_group)
-        return req.is_warmup if req is not None else False
-
     def _dispatch_single_request(self, req_or_group: Any) -> OutputBatch:
         if isinstance(req_or_group, list):
             if not all(isinstance(req, Req) for req in req_or_group):
@@ -277,10 +254,17 @@ class Scheduler(SchedulerDisaggMixin):
             return [self._dispatch_single_request(req) for req in reqs]
         return self._dispatch_single_request(reqs[0])
 
-    def _log_warmup_result(self, output_batch: OutputBatch, is_warmup: bool) -> None:
+    def _log_warmup_result(
+        self,
+        output_batch: OutputBatch,
+        req_or_group: Any,
+        is_warmup: bool,
+    ) -> None:
         if not is_warmup:
             return
 
+        server_based_warmup = is_server_based_warmup(req_or_group)
+
         if output_batch.error is None:
             total_duration_s = (
                 output_batch.metrics.total_duration_s
@@ -297,8 +281,13 @@ class Scheduler(SchedulerDisaggMixin):
                     f"Warmup req processed in {GREEN}%.2f{RESET} seconds",
                     total_duration_s,
                 )
-            if not self._logged_server_ready_after_warmup and (
-                self._warmup_total <= 0 or self._warmup_processed >= self._warmup_total
+            if (
+                not server_based_warmup
+                and not self._logged_server_ready_after_warmup
+                and (
+                    self._warmup_total <= 0
+                    or self._warmup_processed >= self._warmup_total
+                )
             ):
                 logger.info("The server is fired up and ready to roll!")
                 self._logged_server_ready_after_warmup = True
@@ -617,12 +606,12 @@ class Scheduler(SchedulerDisaggMixin):
         self,
         output_batch: OutputBatch,
         identity: bytes | None = None,
-        is_warmup: bool = False,
+        should_not_return: bool = False,
     ):
         """
         replies to client, only on rank 0
         """
-        if not is_warmup and self.receiver is not None and identity is not None:
+        if not should_not_return and self.receiver is not None and identity is not None:
             # if the server is local, use temp file to spill the frame array instead of
             # leaving it in OutputBatch to be pickled later
             if is_local_endpoint(self.server_args.scheduler_endpoint):
@@ -917,46 +906,29 @@ class Scheduler(SchedulerDisaggMixin):
 
     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
+            not self.server_args.warmup
+            or self.warmed_up
+            or self.server_args.warmup_resolutions is None
         ):
-            # insert warmup reqs constructed with each warmup-resolution
-            self._warmup_total = len(self.server_args.warmup_resolutions)
-            self._warmup_processed = 0
-            task_type = self.server_args.pipeline_config.task_type
+            return
 
-            requires_warmup_image = task_type.accepts_image_input()
-            warmup_input_path = None
-            if requires_warmup_image:
-                warmup_input_path = self._prepare_shared_warmup_image_path()
+        self._warmup_total = len(self.server_args.warmup_resolutions)
+        self._warmup_processed = 0
 
-            for resolution in self.server_args.warmup_resolutions:
-                width, height = _parse_size(resolution)
+        warmup_input_path = None
+        if should_include_warmup_image(self.server_args, server_based_warmup=False):
+            warmup_input_path = self._prepare_shared_warmup_image_path()
 
-                # CFG-parallel splits cond/uncond across ranks, so rank 1
-                # needs a real uncond pass. Force do_classifier_free_guidance
-                # + non-empty negative_prompt when cfg-parallel is on, so the
-                # synthesized warmup Req exercises both ranks' denoising paths.
-                # When cfg-parallel is off, the Req construction is
-                # byte-identical to the pre-fix behavior.
-                req_kwargs = dict(
-                    data_type=task_type.data_type(),
-                    width=width,
-                    height=height,
-                    prompt="",
-                )
-                if requires_warmup_image:
-                    req_kwargs["negative_prompt"] = ""
-                    req_kwargs["image_path"] = [warmup_input_path]
-                if self.server_args.enable_cfg_parallel:
-                    req_kwargs["negative_prompt"] = DEFAULT_PLACEHOLDER_PROMPT
-                    req_kwargs["do_classifier_free_guidance"] = True
-                req = Req(**req_kwargs)
-                req.set_as_warmup(self.server_args.warmup_steps)
-                self.waiting_queue.append((None, req, time.monotonic()))
-            # if server is warmed-up, set this flag to avoid req-based warmup
-            self.warmed_up = True
+        warmup_reqs = build_warmup_reqs(
+            self.server_args,
+            warmup_resolutions=self.server_args.warmup_resolutions,
+            warmup_input_path=warmup_input_path,
+        )
+        for req in warmup_reqs:
+            self.waiting_queue.append((None, req, time.monotonic()))
+
+        # if server is warmed-up, set this flag to avoid req-based warmup
+        self.warmed_up = True
 
     def _prepare_shared_warmup_image_path(self) -> str:
         world_group = get_world_group()
@@ -965,18 +937,7 @@ class Scheduler(SchedulerDisaggMixin):
         warmup_sync: dict[str, str | None]
         if world_group.rank == src_rank:
             try:
-                if self.server_args.input_save_path is not None:
-                    uploads_dir = self.server_args.input_save_path
-                    os.makedirs(uploads_dir, exist_ok=True)
-                else:
-                    uploads_dir = tempfile.mkdtemp(prefix="sglang_input_")
-                warmup_image_base = os.path.join(uploads_dir, "warmup_image")
-                input_path = asyncio.run(
-                    save_image_to_path(
-                        MINIMUM_PICTURE_BASE64_FOR_WARMUP,
-                        warmup_image_base,
-                    )
-                )
+                input_path = prepare_warmup_image_path_sync(self.server_args)
                 warmup_sync = {"input_path": input_path, "error": None}
             except Exception as e:
                 warmup_sync = {"input_path": None, "error": str(e)}
@@ -1013,13 +974,14 @@ class Scheduler(SchedulerDisaggMixin):
             or not self.server_args.warmup
             or not recv_reqs
             or self.server_args.warmup_resolutions is not None
+            or self.server_args.server_warmup
         ):
             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 warmed up
         identity, req_or_group = recv_reqs[0]
-        req = self._first_generation_req(req_or_group)
+        req = get_first_generation_req(req_or_group)
         if req is not None:
             warmup_req = req.copy_as_warmup(self.server_args.warmup_steps)
             recv_reqs.insert(0, (identity, warmup_req))
@@ -1200,10 +1162,19 @@ class Scheduler(SchedulerDisaggMixin):
                 for (identity, processed_req), output_batch in zip(
                     items, output_batches, strict=True
                 ):
-                    is_warmup = self._is_warmup_item(processed_req)
-                    self._log_warmup_result(output_batch, is_warmup)
+                    is_warmup = is_warmup_req(processed_req)
+                    self._log_warmup_result(output_batch, processed_req, is_warmup)
 
-                    self.return_result(output_batch, identity, is_warmup=is_warmup)
+                    if is_warmup and should_return_warmup_result(processed_req):
+                        # only keep the necessary lightweight payloads
+                        output_batch.drop_payload_for_warmup()
+                        self.return_result(
+                            output_batch, identity, should_not_return=False
+                        )
+                    else:
+                        self.return_result(
+                            output_batch, identity, should_not_return=is_warmup
+                        )
             except zmq.ZMQError as e:
                 # Reply failed; log and keep loop alive to accept future requests
                 logger.error(f"ZMQ error sending reply: {e}")

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

@@ -328,9 +328,6 @@ class Req:
 
         self.metrics = RequestMetrics(request_id=self.request_id)
 
-    def adjust_size(self, server_args: ServerArgs):
-        pass
-
     def __str__(self):
         return pprint.pformat(asdict(self), indent=2, width=120)
 
@@ -400,3 +397,13 @@ class OutputBatch:
     # For ComfyUI integration: noise prediction from denoising stage
     noise_pred: torch.Tensor | None = None
     peak_memory_mb: float = 0.0
+
+    def drop_payload_for_warmup(self) -> None:
+        self.output = None
+        self.audio = None
+        self.trajectory_timesteps = None
+        self.trajectory_latents = None
+        self.rollout_trajectory_data = None
+        self.trajectory_decoded = None
+        self.output_file_paths = None
+        self.noise_pred = None

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

@@ -904,6 +904,9 @@ class ImageVAEEncodingStage(PipelineStage):
                     #     self.vae.enable_parallel()
                     if not vae_autocast_enabled:
                         video_condition = video_condition.to(vae_dtype)
+                    video_condition = server_args.pipeline_config.preprocess_vae_encode(
+                        video_condition, self.vae
+                    )
                     latent_dist: DiagonalGaussianDistribution = self.vae.encode(
                         video_condition
                     )
@@ -939,15 +942,9 @@ class ImageVAEEncodingStage(PipelineStage):
                         )
                     )
 
-                    # apply shift & scale if needed
-                    if isinstance(shift_factor, torch.Tensor):
-                        shift_factor = shift_factor.to(latent_condition.device)
-
-                    if isinstance(scaling_factor, torch.Tensor):
-                        scaling_factor = scaling_factor.to(latent_condition.device)
-
-                    latent_condition -= shift_factor
-                    latent_condition = latent_condition * scaling_factor
+                    latent_condition = self.scale_and_shift_encode_latents(
+                        latent_condition, scaling_factor, shift_factor
+                    )
                 else:
                     latent_condition = normalized_latent_condition
 
@@ -965,6 +962,19 @@ class ImageVAEEncodingStage(PipelineStage):
 
         return batch
 
+    @staticmethod
+    def scale_and_shift_encode_latents(
+        latents: torch.Tensor, scaling_factor, shift_factor
+    ) -> torch.Tensor:
+        if shift_factor is not None:
+            if isinstance(shift_factor, torch.Tensor):
+                shift_factor = shift_factor.to(latents.device)
+            latents -= shift_factor
+
+        if isinstance(scaling_factor, torch.Tensor):
+            scaling_factor = scaling_factor.to(latents.device)
+        return latents * scaling_factor
+
     def build_dedup_fingerprint(
         self, batch: Req, server_args: ServerArgs
     ) -> ImageVAEEncodingFingerprint | int:
@@ -992,8 +1002,20 @@ class ImageVAEEncodingStage(PipelineStage):
         sample_mode: str = "sample",
     ):
         if sample_mode == "sample":
+            if hasattr(encoder_output, "latent_dist"):
+                return encoder_output.latent_dist.sample(generator)
+            if hasattr(encoder_output, "latent"):
+                return encoder_output.latent
+            if hasattr(encoder_output, "latents"):
+                return encoder_output.latents
             return encoder_output.sample(generator)
         elif sample_mode == "argmax":
+            if hasattr(encoder_output, "latent_dist"):
+                return encoder_output.latent_dist.mode()
+            if hasattr(encoder_output, "latent"):
+                return encoder_output.latent
+            if hasattr(encoder_output, "latents"):
+                return encoder_output.latents
             return encoder_output.mode()
         else:
             raise AttributeError("Could not access latents of provided encoder_output")

python/sglang/multimodal_gen/runtime/server_args.py

@@ -57,10 +57,6 @@ from sglang.multimodal_gen.runtime.utils.common import (
     is_valid_ipv6_address,
 )
 from sglang.multimodal_gen.runtime.utils.logging_utils import (
-    CYAN,
-    GREEN,
-    RED,
-    RESET,
     _sanitize_for_logging,
     configure_logger,
     init_logger,
@@ -224,6 +220,7 @@ class ServerArgs(DisaggArgsMixin):
 
     # warmup
     warmup: bool = False
+    server_warmup: bool = False
     warmup_resolutions: list[str] = None
     warmup_steps: int = 1
 
@@ -427,9 +424,6 @@ class ServerArgs(DisaggArgsMixin):
             if self.image_encoder_cpu_offload is None:
                 self.image_encoder_cpu_offload = True
         elif self.pipeline_config.task_type.is_image_gen():
-            logger.info(
-                "Disabling some offloading (except dit, text_encoder) for image generation model"
-            )
             if self.dit_cpu_offload is None:
                 self.dit_cpu_offload = True
             if self.text_encoder_cpu_offload is None:
@@ -671,11 +665,13 @@ class ServerArgs(DisaggArgsMixin):
     def _adjust_warmup(self):
         if self.warmup_resolutions is not None:
             self.warmup = True
+            self.server_warmup = False
 
-        if self.warmup:
-            logger.info(
-                "Warmup enabled, the launch time is expected to be longer than usual"
-            )
+        if self.disagg_role != RoleType.MONOLITHIC:
+            self.server_warmup = False
+
+        if not self.warmup:
+            self.server_warmup = False
 
     @staticmethod
     def _require_port(port: int, name: str) -> None:
@@ -928,11 +924,18 @@ class ServerArgs(DisaggArgsMixin):
             self.vae_cpu_offload = False
             disabled_flag_names.append("vae_cpu_offload")
 
-        if disabled_flag_names:
+        explicit_disabled_flag_names = [
+            flag_name
+            for flag_name in disabled_flag_names
+            if self.is_arg_explicitly_set(flag_name)
+        ]
+        if explicit_disabled_flag_names:
             logger.info(
-                "Disabling %s because the selected layerwise offload components "
-                "manage the same weights.",
-                ", ".join(disabled_flag_names),
+                "Ignoring explicit CPU-offload flags because layerwise offload "
+                "manages the same component weights: %s",
+                ", ".join(
+                    f"{flag_name}=False" for flag_name in explicit_disabled_flag_names
+                ),
             )
 
     def _adjust_autocast(self):
@@ -1209,9 +1212,15 @@ class ServerArgs(DisaggArgsMixin):
             "--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)."
-            "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.",
+            help=(
+                "Perform warmup before normal traffic. `sglang serve` runs a "
+                "lightweight server warmup after HTTP is ready; other entrypoints "
+                "use request-based warmup unless `--warmup-resolutions` is "
+                "specified. 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."
+            ),
         )
         parser.add_argument(
             "--warmup-resolutions",
@@ -1226,6 +1235,12 @@ class ServerArgs(DisaggArgsMixin):
             default=ServerArgs.warmup_steps,
             help="The number of warmup steps to perform for each resolution.",
         )
+        parser.add_argument(
+            "--server-warmup",
+            action=StoreBoolean,
+            default=ServerArgs.server_warmup,
+            help="Send a warmup request after server ready",
+        )
 
         # layerwise offload
         parser.add_argument(
@@ -1651,7 +1666,10 @@ class ServerArgs(DisaggArgsMixin):
 
     @classmethod
     def from_cli_args(
-        cls, args: argparse.Namespace, unknown_args: list[str] | None = None
+        cls,
+        args: argparse.Namespace,
+        unknown_args: list[str] | None = None,
+        default_args: dict[str, Any] | None = None,
     ) -> "ServerArgs":
         if unknown_args is None:
             unknown_args = []
@@ -1665,24 +1683,33 @@ class ServerArgs(DisaggArgsMixin):
             raise SystemExit(f"error: unrecognized arguments: {' '.join(remaining)}")
 
         provided_args = cls.get_provided_args(args, unknown_args)
+        explicit_arg_names = set(provided_args)
 
         # Handle config file
         config_file = provided_args.get("config")
         if config_file:
             config_args = cls.load_config_file(config_file)
+            explicit_arg_names.update(config_args)
             provided_args = {**config_args, **provided_args}
 
+        if default_args:
+            for key, value in default_args.items():
+                provided_args.setdefault(key, value)
+
         if dynamic_paths:
             existing = dict(provided_args.get("component_paths") or {})
             existing.update(dynamic_paths)
             provided_args["component_paths"] = existing
+            explicit_arg_names.add("component_paths")
         if dynamic_attention_backends:
             existing = cls._parse_component_attention_backend_map(
                 provided_args.get("component_attention_backends")
             )
             existing.update(dynamic_attention_backends)
             provided_args["component_attention_backends"] = existing
+            explicit_arg_names.add("component_attention_backends")
 
+        provided_args["_explicit_arg_names"] = explicit_arg_names
         return cls.from_dict(provided_args)
 
     @classmethod
@@ -1690,14 +1717,19 @@ class ServerArgs(DisaggArgsMixin):
         """Create a ServerArgs object from a dictionary."""
         attrs = [attr.name for attr in dataclasses.fields(cls) if attr.init]
         server_args_kwargs: dict[str, Any] = {}
+        explicit_arg_names = kwargs.get("_explicit_arg_names")
+        if explicit_arg_names is None:
+            explicit_arg_names = set(kwargs)
 
         component_paths = dict(kwargs.get("component_paths") or {})
         if component_paths:
             server_args_kwargs["component_paths"] = component_paths
-        server_args_kwargs["_explicit_arg_names"] = set(kwargs)
+        server_args_kwargs["_explicit_arg_names"] = set(explicit_arg_names)
 
         for attr in attrs:
-            if attr == "pipeline_config":
+            if attr == "_explicit_arg_names":
+                continue
+            elif attr == "pipeline_config":
                 pipeline_config = PipelineConfig.from_kwargs(kwargs)
                 logger.debug(f"Using PipelineConfig: {type(pipeline_config)}")
                 server_args_kwargs["pipeline_config"] = pipeline_config
@@ -1824,17 +1856,17 @@ class ServerArgs(DisaggArgsMixin):
                     "or disable SGLANG_CACHE_DIT_ENABLED."
                 )
 
-            logger.warning(
-                "layerwise offload components are selected: %slower GPU memory usage%s, but %smay reduce throughput or increase latency%s. "
-                "%sIf you are using multi-GPU deployment and already have enough memory headroom, prefer keeping layerwise offload disabled.%s "
-                "Please tune this based on your memory headroom and performance target.",
-                GREEN,
-                RESET,
-                RED,
-                RESET,
-                CYAN,
-                RESET,
-            )
+            if (
+                self.performance_mode == "memory"
+                or self.is_arg_explicitly_set("layerwise_offload_components")
+                or self.dit_layerwise_offload
+            ):
+                logger.info_once(
+                    "Using layerwise offload components: "
+                    f"{', '.join(self.layerwise_offload_components)}. "
+                    "This reduces peak GPU memory and can increase latency; use "
+                    "--performance-mode speed for GPU-resident defaults when memory allows."
+                )
 
     def _validate_parallelism(self):
         if self.sp_degree > self.num_gpus or self.num_gpus % self.sp_degree != 0:

python/sglang/multimodal_gen/runtime/server_args_auto_tune.py

@@ -220,9 +220,9 @@ class ServerArgsAutoTuner:
             return
 
         logger.info(
-            "Automatically enable default layerwise offload for %s: %s",
+            "Auto memory policy for %s selected layerwise offload components: %s",
             args.pipeline_config.__class__.__name__,
-            layerwise_components,
+            ", ".join(layerwise_components),
         )
         args.layerwise_offload_components = layerwise_components
 

python/sglang/multimodal_gen/runtime/server_warmup.py

@@ -0,0 +1,219 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+import asyncio
+import os
+import tempfile
+from copy import copy
+from typing import Any
+
+from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.registry import get_pipeline_config_classes
+from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
+    _parse_size,
+    save_image_to_path,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+MINIMUM_PICTURE_BASE64_FOR_WARMUP = "data:image/jpg;base64,iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAAACXBIWXMAAA7EAAAOxAGVKw4bAAAAbUlEQVRYhe3VsQ2AMAxE0Y/lIgNQULD/OqyCMgCihCKSG4yRuKuiNH6JLsoEbMACOGBcua9HOR7Y6w6swBwMy0qLTpkeI77qdEBpBFAHBBDAGH8WrwJKI4AAegUCfAKgEgpQDvh3CR3oQCuav58qlAw73kKCSgAAAABJRU5ErkJggg=="
+
+DEFAULT_PLACEHOLDER_PROMPT = "warmup"
+DEFAULT_LIGHTWEIGHT_IMAGE_RESOLUTION = (64, 64)
+
+
+def get_first_generation_req(req_or_group: Any) -> Req | None:
+    """Extract the first req"""
+    if isinstance(req_or_group, Req):
+        return req_or_group
+    if isinstance(req_or_group, list) and req_or_group:
+        first_req = req_or_group[0]
+        if isinstance(first_req, Req):
+            return first_req
+    return None
+
+
+def is_warmup_req(req_or_group: Any) -> bool:
+    """either server-based or req-based"""
+    req = get_first_generation_req(req_or_group)
+    return req.is_warmup if req is not None else False
+
+
+def is_server_based_warmup(req_or_group: Any) -> bool:
+    req = get_first_generation_req(req_or_group)
+    return (
+        req is not None and req.is_warmup and bool(req.extra.get("server_based_warmup"))
+    )
+
+
+def should_return_warmup_result(req_or_group: Any) -> bool:
+    # server-based warmup needs to return to the http server to finish the startup
+    req = get_first_generation_req(req_or_group)
+    return (
+        req is not None
+        and req.is_warmup
+        and bool(req.extra.get("return_warmup_result"))
+    )
+
+
+def get_model_sampling_defaults(server_args: ServerArgs) -> SamplingParams:
+    pipeline_class_name = server_args.pipeline_class_name
+    try:
+        if pipeline_class_name:
+            config_classes = get_pipeline_config_classes(pipeline_class_name)
+            if config_classes is not None:
+                _, sampling_params_cls = config_classes
+                return sampling_params_cls()
+
+        return SamplingParams.from_pretrained(
+            server_args.model_path,
+            backend=server_args.backend,
+            model_id=server_args.model_id,
+        )
+    except Exception:
+        logger.debug("Falling back to base SamplingParams for server warmup")
+        return SamplingParams()
+
+
+async def prepare_warmup_image_path(server_args: ServerArgs) -> str:
+    if server_args.input_save_path is not None:
+        uploads_dir = server_args.input_save_path
+        os.makedirs(uploads_dir, exist_ok=True)
+    else:
+        uploads_dir = tempfile.mkdtemp(prefix="sglang_input_")
+
+    warmup_image_base = os.path.join(uploads_dir, "warmup_image")
+    return await save_image_to_path(
+        MINIMUM_PICTURE_BASE64_FOR_WARMUP, warmup_image_base
+    )
+
+
+def prepare_warmup_image_path_sync(server_args: ServerArgs) -> str:
+    return asyncio.run(prepare_warmup_image_path(server_args))
+
+
+def _resolve_default_warmup_resolution(
+    server_args: ServerArgs,
+    sampling_defaults: SamplingParams,
+) -> tuple[int, int]:
+    supported_resolutions = sampling_defaults.supported_resolutions
+    if supported_resolutions:
+        return min(supported_resolutions, key=lambda size: size[0] * size[1])
+
+    width = sampling_defaults.width
+    height = sampling_defaults.height
+    if width is not None and height is not None:
+        return width, height
+
+    if server_args.pipeline_config.task_type.is_image_gen():
+        return DEFAULT_LIGHTWEIGHT_IMAGE_RESOLUTION
+
+    return (
+        width or DEFAULT_LIGHTWEIGHT_IMAGE_RESOLUTION[0],
+        height or DEFAULT_LIGHTWEIGHT_IMAGE_RESOLUTION[1],
+    )
+
+
+def _effective_cfg_scale(sampling_defaults: SamplingParams) -> float | None:
+    if sampling_defaults.true_cfg_scale is not None:
+        return sampling_defaults.true_cfg_scale
+    return sampling_defaults.guidance_scale
+
+
+def should_include_warmup_image(
+    server_args: ServerArgs, server_based_warmup: bool
+) -> bool:
+    task_type = server_args.pipeline_config.task_type
+    if not task_type.accepts_image_input():
+        return False
+    if task_type.requires_image_input():
+        return True
+    if server_based_warmup:
+        return task_type in (ModelTaskType.TI2I, ModelTaskType.TI2V)
+    return True
+
+
+def build_warmup_reqs(
+    server_args: ServerArgs,
+    *,
+    warmup_resolutions: list[str] | None,
+    warmup_input_path: str | None = None,
+    return_warmup_result: bool = False,
+    server_based_warmup: bool = False,
+    use_model_sampling_defaults: bool = False,
+) -> list[Req]:
+    task_type = server_args.pipeline_config.task_type
+    if warmup_resolutions is None or use_model_sampling_defaults:
+        sampling_defaults = get_model_sampling_defaults(server_args)
+    else:
+        sampling_defaults = SamplingParams()
+
+    if warmup_resolutions is None:
+        width, height = _resolve_default_warmup_resolution(
+            server_args, sampling_defaults
+        )
+        resolutions: list[tuple[int, int]] = [(width, height)]
+    else:
+        resolutions = [_parse_size(resolution) for resolution in warmup_resolutions]
+
+    negative_prompt: Any = (
+        sampling_defaults.negative_prompt if use_model_sampling_defaults else None
+    )
+    cfg_scale = (
+        _effective_cfg_scale(sampling_defaults) if use_model_sampling_defaults else None
+    )
+
+    warmup_reqs = []
+    include_warmup_image = should_include_warmup_image(server_args, server_based_warmup)
+    for width, height in resolutions:
+        req_kwargs = dict(
+            data_type=task_type.data_type(),
+            width=width,
+            height=height,
+            prompt=DEFAULT_PLACEHOLDER_PROMPT,
+        )
+        if use_model_sampling_defaults:
+            req_kwargs["sampling_params"] = copy(sampling_defaults)
+            req_kwargs.update(
+                negative_prompt=negative_prompt,
+                guidance_scale=sampling_defaults.guidance_scale,
+                guidance_scale_2=sampling_defaults.guidance_scale_2,
+                true_cfg_scale=sampling_defaults.true_cfg_scale,
+                num_inference_steps=sampling_defaults.num_inference_steps,
+            )
+        if include_warmup_image:
+            if warmup_input_path is None:
+                raise RuntimeError(
+                    "Warmup image path is required for image-input model"
+                )
+            req_kwargs["prompt"] = DEFAULT_PLACEHOLDER_PROMPT
+            if not use_model_sampling_defaults:
+                req_kwargs["negative_prompt"] = ""
+            req_kwargs["image_path"] = [warmup_input_path]
+        if (
+            server_args.enable_cfg_parallel
+            and req_kwargs.get("negative_prompt") is None
+        ):
+            req_kwargs["negative_prompt"] = DEFAULT_PLACEHOLDER_PROMPT
+            req_kwargs["do_classifier_free_guidance"] = True
+        elif (
+            use_model_sampling_defaults
+            and negative_prompt is not None
+            and cfg_scale is not None
+            and cfg_scale > 1.0
+        ):
+            req_kwargs["do_classifier_free_guidance"] = True
+
+        req = Req(**req_kwargs)
+        req.set_as_warmup(server_args.warmup_steps)
+        if return_warmup_result:
+            req.extra["return_warmup_result"] = True
+        if server_based_warmup:
+            req.extra["server_based_warmup"] = True
+        warmup_reqs.append(req)
+
+    return warmup_reqs

python/sglang/multimodal_gen/test/server/perf_baselines.json

@@ -465,7 +465,7 @@
                 "ImageVAEEncodingStage": 69.59
             },
             "denoise_step_ms": {
-                "0": 129.61,
+                "0": 910.0,
                 "1": 236.54,
                 "2": 934.61,
                 "3": 933.62,

python/sglang/multimodal_gen/test/server/test_server_utils.py

@@ -50,6 +50,11 @@ logger = init_logger(__name__)
 
 globally_suppress_loggers()
 
+FIRST_DENOISE_STEP_TOLERANCE = 4.0
+FIRST_DENOISE_STEP_MIN_ABS_TOLERANCE_MS = 80.0
+DECODING_STAGE_MIN_ABS_TOLERANCE_MS = 450.0
+VIDEO_DENOISE_STEP_MIN_ABS_TOLERANCE_MS = 160.0
+
 # Tracks mesh output file paths from generate_mesh for later correctness validation.
 # Keyed by case_id, cleaned up after use.
 MESH_OUTPUT_PATHS: dict[str, str] = {}
@@ -603,14 +608,23 @@ class PerformanceValidator:
             expected = self.scenario.denoise_step_ms.get(idx)
             if expected is None:
                 continue
-            # FIXME: hardcode, looser for first step
-            tolerance = 0.4 if idx == 0 else self.tolerances.denoise_step
+            if idx == 0:
+                # server warmup is generic, so the first real step can still
+                # pay request-shape/path lazy init that is not a steady-state signal
+                self._assert_le(
+                    f"Denoise Step {idx}",
+                    actual,
+                    expected,
+                    FIRST_DENOISE_STEP_TOLERANCE,
+                    min_abs_tolerance_ms=FIRST_DENOISE_STEP_MIN_ABS_TOLERANCE_MS,
+                )
+                continue
 
             self._assert_le(
                 f"Denoise Step {idx}",
                 actual,
                 expected,
-                tolerance,
+                self.tolerances.denoise_step,
             )
 
     def _validate_stages(self, summary: PerformanceSummary) -> None:
@@ -629,7 +643,7 @@ class PerformanceValidator:
             )
             if stage.endswith("DecodingStage"):
                 tolerance = max(tolerance, 0.9)
-                min_abs_tolerance_ms = 250.0
+                min_abs_tolerance_ms = DECODING_STAGE_MIN_ABS_TOLERANCE_MS
             else:
                 min_abs_tolerance_ms = 120.0
             self._assert_le(
@@ -646,6 +660,32 @@ class VideoPerformanceValidator(PerformanceValidator):
 
     is_video_gen = True
 
+    def _validate_denoise_steps(self, summary: PerformanceSummary) -> None:
+        """Validate individual denoising steps."""
+        for idx, actual in summary.sampled_steps.items():
+            expected = self.scenario.denoise_step_ms.get(idx)
+            if expected is None:
+                continue
+            if idx == 0:
+                self._assert_le(
+                    f"Denoise Step {idx}",
+                    actual,
+                    expected,
+                    FIRST_DENOISE_STEP_TOLERANCE,
+                    min_abs_tolerance_ms=FIRST_DENOISE_STEP_MIN_ABS_TOLERANCE_MS,
+                )
+                continue
+
+            # video per-step samples can catch one-off scheduling/offload jitter;
+            # avg and median denoise checks remain the steady-state guard
+            self._assert_le(
+                f"Denoise Step {idx}",
+                actual,
+                expected,
+                self.tolerances.denoise_step,
+                min_abs_tolerance_ms=VIDEO_DENOISE_STEP_MIN_ABS_TOLERANCE_MS,
+            )
+
     def validate(
         self,
         perf_record: RequestPerfRecord,

python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py

@@ -1,27 +1,42 @@
 """Unit tests for the --enable-cfg-parallel warmup fix and guard.
 
-Covers two code paths introduced alongside this file:
+Covers three code paths introduced alongside this file:
 - Scheduler.prepare_server_warmup_reqs synthesizes warmup Reqs that
   actually enable classifier-free guidance when cfg-parallel is on.
 - InputValidationStage.forward rejects non-CFG requests when the server
   has cfg-parallel on.
+- Server-based warmup can opt into model-default negative prompts so warmup
+  populates the negative text embedding cache.
 
 All tests are CPU-only; no model loading, no distributed init.
 """
 
 import unittest
 from collections import deque
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
+import torch
+
 from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType
-from sglang.multimodal_gen.runtime.managers.scheduler import (
-    DEFAULT_PLACEHOLDER_PROMPT,
-    Scheduler,
+from sglang.multimodal_gen.configs.pipeline_configs.flux_finetuned import (
+    Flux2FinetunedPipelineConfig,
 )
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.runtime.managers.scheduler import Scheduler
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.image_encoding import (
+    ImageVAEEncodingStage,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import (
     InputValidationStage,
 )
+from sglang.multimodal_gen.runtime.server_warmup import (
+    DEFAULT_LIGHTWEIGHT_IMAGE_RESOLUTION,
+    DEFAULT_PLACEHOLDER_PROMPT,
+    build_warmup_reqs,
+    should_include_warmup_image,
+)
 
 # Patch path for get_global_server_args used by Stage.__init__
 _GLOBAL_ARGS_PATCH = (
@@ -48,6 +63,7 @@ def _make_bare_scheduler(enable_cfg_parallel: bool) -> Scheduler:
     # branch entirely, so we don't need to mock
     # _prepare_shared_warmup_image_path.
     task_type = MagicMock()
+    task_type.requires_image_input.return_value = False
     task_type.accepts_image_input.return_value = False
     task_type.data_type.return_value = ModelTaskType.T2I.data_type()
     server_args.pipeline_config.task_type = task_type
@@ -99,6 +115,272 @@ class TestWarmupReqCfgParallel(unittest.TestCase):
         self.assertIs(req.do_classifier_free_guidance, False)
         self.assertNotEqual(req.negative_prompt, DEFAULT_PLACEHOLDER_PROMPT)
 
+    def test_server_based_warmup_uses_model_default_negative_prompt(self):
+        server_args = MagicMock()
+        server_args.warmup_steps = 1
+        server_args.enable_cfg_parallel = False
+
+        task_type = MagicMock()
+        task_type.requires_image_input.return_value = False
+        task_type.accepts_image_input.return_value = False
+        task_type.is_image_gen.return_value = True
+        task_type.data_type.return_value = ModelTaskType.T2I.data_type()
+        server_args.pipeline_config.task_type = task_type
+
+        sampling_defaults = SamplingParams(
+            negative_prompt="model default negative",
+            guidance_scale=4.0,
+            num_inference_steps=20,
+        )
+        with patch(
+            "sglang.multimodal_gen.runtime.server_warmup.get_model_sampling_defaults",
+            return_value=sampling_defaults,
+        ):
+            reqs = build_warmup_reqs(
+                server_args,
+                warmup_resolutions=None,
+                use_model_sampling_defaults=True,
+                return_warmup_result=True,
+                server_based_warmup=True,
+            )
+
+        self.assertEqual(len(reqs), 1)
+        req = reqs[0]
+        self.assertTrue(req.is_warmup)
+        self.assertEqual(req.negative_prompt, "model default negative")
+        self.assertIs(req.do_classifier_free_guidance, True)
+        self.assertTrue(req.extra["return_warmup_result"])
+        self.assertTrue(req.extra["server_based_warmup"])
+
+    def test_server_based_warmup_uses_model_default_resolution(self):
+        server_args = MagicMock()
+        server_args.warmup_steps = 1
+        server_args.enable_cfg_parallel = False
+
+        task_type = MagicMock()
+        task_type.requires_image_input.return_value = False
+        task_type.accepts_image_input.return_value = False
+        task_type.is_image_gen.return_value = True
+        task_type.data_type.return_value = ModelTaskType.T2I.data_type()
+        server_args.pipeline_config.task_type = task_type
+
+        sampling_defaults = SamplingParams(width=640, height=640)
+        with patch(
+            "sglang.multimodal_gen.runtime.server_warmup.get_model_sampling_defaults",
+            return_value=sampling_defaults,
+        ):
+            reqs = build_warmup_reqs(
+                server_args,
+                warmup_resolutions=None,
+                use_model_sampling_defaults=True,
+                server_based_warmup=True,
+            )
+
+        req = reqs[0]
+        self.assertEqual(req.width, 640)
+        self.assertEqual(req.height, 640)
+
+    def test_server_based_warmup_keeps_lightweight_image_fallback(self):
+        server_args = MagicMock()
+        server_args.warmup_steps = 1
+        server_args.enable_cfg_parallel = False
+
+        task_type = MagicMock()
+        task_type.requires_image_input.return_value = False
+        task_type.accepts_image_input.return_value = False
+        task_type.is_image_gen.return_value = True
+        task_type.data_type.return_value = ModelTaskType.T2I.data_type()
+        server_args.pipeline_config.task_type = task_type
+
+        with patch(
+            "sglang.multimodal_gen.runtime.server_warmup.get_model_sampling_defaults",
+            return_value=SamplingParams(),
+        ):
+            reqs = build_warmup_reqs(
+                server_args,
+                warmup_resolutions=None,
+                use_model_sampling_defaults=True,
+                server_based_warmup=True,
+            )
+
+        req = reqs[0]
+        self.assertEqual(
+            (req.width, req.height),
+            DEFAULT_LIGHTWEIGHT_IMAGE_RESOLUTION,
+        )
+
+    def test_warmup_image_inclusion_policy_all_task_types(self):
+        server_based_expected = {
+            ModelTaskType.T2I: False,
+            ModelTaskType.T2V: False,
+            ModelTaskType.TI2I: True,
+            ModelTaskType.TI2V: True,
+            ModelTaskType.I2I: True,
+            ModelTaskType.I2V: True,
+            ModelTaskType.I2M: True,
+        }
+
+        for task_type in ModelTaskType:
+            server_args = MagicMock()
+            server_args.pipeline_config.task_type = task_type
+
+            self.assertEqual(
+                should_include_warmup_image(server_args, server_based_warmup=True),
+                server_based_expected[task_type],
+                task_type.name,
+            )
+            self.assertEqual(
+                should_include_warmup_image(server_args, server_based_warmup=False),
+                task_type.accepts_image_input(),
+                task_type.name,
+            )
+
+    def test_server_based_warmup_keeps_ti2i_image_input(self):
+        server_args = MagicMock()
+        server_args.warmup_steps = 1
+        server_args.enable_cfg_parallel = False
+        server_args.pipeline_config.task_type = ModelTaskType.TI2I
+
+        with patch(
+            "sglang.multimodal_gen.runtime.server_warmup.get_model_sampling_defaults",
+            return_value=SamplingParams(width=512, height=512),
+        ):
+            reqs = build_warmup_reqs(
+                server_args,
+                warmup_resolutions=None,
+                warmup_input_path="/tmp/warmup.png",
+                use_model_sampling_defaults=True,
+                server_based_warmup=True,
+            )
+
+        self.assertEqual(reqs[0].image_path, ["/tmp/warmup.png"])
+
+    def test_server_based_warmup_keeps_required_image_input(self):
+        server_args = MagicMock()
+        server_args.warmup_steps = 1
+        server_args.enable_cfg_parallel = False
+        server_args.pipeline_config.task_type = ModelTaskType.I2I
+
+        with patch(
+            "sglang.multimodal_gen.runtime.server_warmup.get_model_sampling_defaults",
+            return_value=SamplingParams(width=512, height=512),
+        ):
+            reqs = build_warmup_reqs(
+                server_args,
+                warmup_resolutions=None,
+                warmup_input_path="/tmp/warmup.png",
+                use_model_sampling_defaults=True,
+                server_based_warmup=True,
+            )
+
+        self.assertEqual(reqs[0].image_path, ["/tmp/warmup.png"])
+
+    def test_server_based_warmup_keeps_ti2v_image_input(self):
+        server_args = MagicMock()
+        server_args.warmup_steps = 1
+        server_args.enable_cfg_parallel = False
+        server_args.pipeline_config.task_type = ModelTaskType.TI2V
+
+        with patch(
+            "sglang.multimodal_gen.runtime.server_warmup.get_model_sampling_defaults",
+            return_value=SamplingParams(width=512, height=512),
+        ):
+            reqs = build_warmup_reqs(
+                server_args,
+                warmup_resolutions=None,
+                warmup_input_path="/tmp/warmup.png",
+                use_model_sampling_defaults=True,
+                server_based_warmup=True,
+            )
+
+        self.assertEqual(reqs[0].image_path, ["/tmp/warmup.png"])
+
+
+class TestFlux2FinetunedVaeEncodePreprocess(unittest.TestCase):
+    def test_single_frame_custom_vae_encode_input_is_4d(self):
+        config = Flux2FinetunedPipelineConfig()
+        vae = MagicMock()
+        vae.bn = None
+
+        image = torch.zeros(1, 3, 1, 32, 32)
+        output = config.preprocess_vae_encode(image, vae)
+
+        self.assertEqual(tuple(output.shape), (1, 3, 32, 32))
+
+    def test_standard_flux2_vae_encode_input_stays_5d(self):
+        config = Flux2FinetunedPipelineConfig()
+        vae = MagicMock()
+        vae.bn = object()
+
+        image = torch.zeros(1, 3, 1, 32, 32)
+        output = config.preprocess_vae_encode(image, vae)
+
+        self.assertIs(output, image)
+
+    def test_custom_vae_already_patchified_encode_latents_stay_128_channels(self):
+        config = Flux2FinetunedPipelineConfig()
+        config.dit_config.arch_config.in_channels = 128
+        vae = MagicMock()
+        vae.bn = None
+
+        image_latents = torch.zeros(1, config.dit_config.arch_config.in_channels, 8, 8)
+        output = config.postprocess_vae_encode(image_latents, vae)
+
+        self.assertIs(output, image_latents)
+
+    def test_standard_flux2_vae_encode_latents_are_patchified(self):
+        config = Flux2FinetunedPipelineConfig()
+        vae = MagicMock()
+        vae.bn = object()
+
+        image_latents = torch.zeros(1, 32, 8, 8)
+        output = config.postprocess_vae_encode(image_latents, vae)
+
+        self.assertEqual(
+            tuple(output.shape),
+            (1, image_latents.shape[1] * 4, 4, 4),
+        )
+
+
+class TestImageVaeEncodingLatentRetrieval(unittest.TestCase):
+    def test_encode_scale_and_shift_allows_missing_shift(self):
+        latents = torch.ones(1, 4, 2, 2)
+        scaling_factor = torch.full((1, 1, 1, 1), 2.0)
+
+        output = ImageVAEEncodingStage.scale_and_shift_encode_latents(
+            latents, scaling_factor, None
+        )
+
+        self.assertTrue(torch.equal(output, torch.full_like(latents, 2.0)))
+
+    def test_retrieve_latents_accepts_encoder_output_latent(self):
+        stage = object.__new__(ImageVAEEncodingStage)
+        latents = torch.zeros(1, 32, 8, 8)
+        encoder_output = SimpleNamespace(latent=latents)
+
+        self.assertIs(
+            stage.retrieve_latents(encoder_output, sample_mode="argmax"),
+            latents,
+        )
+        self.assertIs(
+            stage.retrieve_latents(encoder_output, sample_mode="sample"),
+            latents,
+        )
+
+    def test_retrieve_latents_accepts_encoder_output_latents(self):
+        stage = object.__new__(ImageVAEEncodingStage)
+        latents = torch.zeros(1, 32, 8, 8)
+        encoder_output = SimpleNamespace(latents=latents)
+
+        self.assertIs(
+            stage.retrieve_latents(encoder_output, sample_mode="argmax"),
+            latents,
+        )
+        self.assertIs(
+            stage.retrieve_latents(encoder_output, sample_mode="sample"),
+            latents,
+        )
+
 
 class TestInputValidationCfgParallelGuard(unittest.TestCase):
     """Commit 2: per-request cfg-parallel check.

python/sglang/multimodal_gen/test/unit/test_server_args.py

@@ -1,5 +1,7 @@
+import json
 import os
 import sys
+import tempfile
 import unittest
 from unittest.mock import patch
 
@@ -271,6 +273,173 @@ class TestServerArgsPathExpansion(unittest.TestCase):
             server_args.layerwise_offload_components, ["transformer", "text_encoder"]
         )
 
+    def test_serve_cli_preserves_config_and_dynamic_unknown_args(self):
+        from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
+            add_multimodal_gen_serve_args,
+        )
+
+        with tempfile.NamedTemporaryFile("w", suffix=".json") as config_file:
+            json.dump({"model_path": "/from/config", "num_gpus": 2}, config_file)
+            config_file.flush()
+            parser = FlexibleArgumentParser()
+            add_multimodal_gen_serve_args(parser)
+            argv = [
+                "--config",
+                config_file.name,
+                "--model-path",
+                "/from/cli",
+                "--vae-path",
+                "/custom/vae",
+                "--component-attention-backends.transformer",
+                "fa3",
+            ]
+
+            with patch.object(sys, "argv", ["sglang", "serve"] + argv):
+                args, unknown_args = parser.parse_known_args(argv)
+                with (
+                    patch.object(
+                        PipelineConfig,
+                        "from_kwargs",
+                        return_value=QwenImagePipelineConfig(),
+                    ),
+                    patch(
+                        "sglang.multimodal_gen.registry.get_model_info",
+                        return_value=None,
+                    ),
+                    patch(
+                        "sglang.multimodal_gen.runtime.server_args.current_platform.get_device_total_memory",
+                        return_value=80 * 1024**3,
+                    ),
+                    patch(
+                        "sglang.multimodal_gen.runtime.server_args.current_platform.get_available_gpu_memory",
+                        return_value=80,
+                    ),
+                ):
+                    server_args = ServerArgs.from_cli_args(args, unknown_args)
+
+        self.assertEqual("/from/cli", server_args.model_path)
+        self.assertEqual(2, server_args.num_gpus)
+        self.assertEqual("/custom/vae", server_args.component_paths["vae"])
+        self.assertEqual(
+            {"transformer": "fa"},
+            server_args.component_attention_backends,
+        )
+
+    def test_serve_cli_defaults_warmup_on(self):
+        from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
+            add_multimodal_gen_serve_args,
+            execute_serve_cmd,
+        )
+
+        parser = FlexibleArgumentParser()
+        add_multimodal_gen_serve_args(parser)
+        argv = [
+            "--model-path",
+            "/fake",
+        ]
+
+        with (
+            patch.object(sys, "argv", ["sglang", "serve"] + argv),
+            patch.object(
+                PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.entrypoints.cli.serve.dispatch_launch"
+            ) as dispatch_launch,
+        ):
+            args, unknown_args = parser.parse_known_args(argv)
+            execute_serve_cmd(args, unknown_args)
+
+        server_args = dispatch_launch.call_args.args[0]
+        self.assertTrue(server_args.warmup)
+        self.assertTrue(server_args.server_warmup)
+        self.assertFalse(server_args.is_arg_explicitly_set("warmup"))
+        self.assertFalse(server_args.is_arg_explicitly_set("server_warmup"))
+
+    def test_serve_cli_preserves_explicit_warmup_false(self):
+        from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
+            add_multimodal_gen_serve_args,
+            execute_serve_cmd,
+        )
+
+        parser = FlexibleArgumentParser()
+        add_multimodal_gen_serve_args(parser)
+        argv = [
+            "--model-path",
+            "/fake",
+            "--warmup",
+            "false",
+        ]
+
+        with (
+            patch.object(sys, "argv", ["sglang", "serve"] + argv),
+            patch.object(
+                PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.entrypoints.cli.serve.dispatch_launch"
+            ) as dispatch_launch,
+        ):
+            args, unknown_args = parser.parse_known_args(argv)
+            execute_serve_cmd(args, unknown_args)
+
+        server_args = dispatch_launch.call_args.args[0]
+        self.assertFalse(server_args.warmup)
+        self.assertFalse(server_args.server_warmup)
+        self.assertTrue(server_args.is_arg_explicitly_set("warmup"))
+
+    def test_serve_cli_preserves_config_warmup_false(self):
+        from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
+            add_multimodal_gen_serve_args,
+            execute_serve_cmd,
+        )
+
+        with tempfile.NamedTemporaryFile("w", suffix=".json") as config_file:
+            json.dump({"model_path": "/fake", "warmup": False}, config_file)
+            config_file.flush()
+
+            parser = FlexibleArgumentParser()
+            add_multimodal_gen_serve_args(parser)
+            argv = [
+                "--config",
+                config_file.name,
+            ]
+
+            with (
+                patch.object(sys, "argv", ["sglang", "serve"] + argv),
+                patch.object(
+                    PipelineConfig,
+                    "from_kwargs",
+                    return_value=QwenImagePipelineConfig(),
+                ),
+                patch(
+                    "sglang.multimodal_gen.runtime.entrypoints.cli.serve.dispatch_launch"
+                ) as dispatch_launch,
+            ):
+                args, unknown_args = parser.parse_known_args(argv)
+                execute_serve_cmd(args, unknown_args)
+
+        server_args = dispatch_launch.call_args.args[0]
+        self.assertFalse(server_args.warmup)
+        self.assertFalse(server_args.server_warmup)
+        self.assertTrue(server_args.is_arg_explicitly_set("warmup"))
+
+    def test_disagg_role_disables_server_warmup(self):
+        with patch.object(
+            PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
+        ):
+            server_args = ServerArgs.from_dict(
+                {
+                    "model_path": "/fake",
+                    "warmup": True,
+                    "server_warmup": True,
+                    "disagg_role": "server",
+                }
+            )
+
+        self.assertTrue(server_args.warmup)
+        self.assertFalse(server_args.server_warmup)
+
 
 class TestOffloadDefaults(unittest.TestCase):
     def _from_dict_with_pipeline_config(