[diffusion] fix: fix serving with dit-layerwise-offload enabled (#16066)

Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>

python/sglang/multimodal_gen/benchmarks/bench_serving.py

@@ -750,7 +750,7 @@ async def benchmark(args):
             )
         )
 
-    print("\n" + "=" * 60)
+    print("=" * 60)
 
     if args.output_file:
         with open(args.output_file, "w") as f:

python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py

@@ -121,8 +121,6 @@ async def generations(
         server_args=get_global_server_args(),
         sampling_params=sampling,
     )
-
-    # Run synchronously for images and save to disk
     save_file_path, result = await process_generation_batch(
         async_scheduler_client, batch
     )

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

@@ -12,7 +12,7 @@ import time
 import weakref
 from collections.abc import Iterable
 from functools import lru_cache
-from typing import Any
+from typing import Any, Optional
 
 import torch
 from einops import rearrange
@@ -61,6 +61,9 @@ from sglang.multimodal_gen.runtime.platforms import (
     current_platform,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.layerwise_offload import (
+    LayerwiseOffloadManager,
+)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
 from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler
@@ -721,6 +724,14 @@ class DenoisingStage(PipelineStage):
                 torch.mps.current_allocated_memory(),
             )
 
+        # reset offload manager with prefetching first layer for next forward
+        offload_mgr: Optional[LayerwiseOffloadManager] = None
+        for transformer in filter(None, [self.transformer, self.transformer_2]):
+            if (
+                offload_mgr := getattr(transformer, "_layerwise_offload_manager", None)
+            ) is not None:
+                offload_mgr.prepare_for_next_denoise(non_blocking=True)
+
     def _preprocess_sp_latents(self, batch: Req, server_args: ServerArgs):
         """Shard latents for Sequence Parallelism if applicable."""
         if get_sp_world_size() <= 1:

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

@@ -98,8 +98,11 @@ class LayerwiseOffloadManager:
                 continue
             self._offload_tensor(name, buf, layer_idx)
 
-        self.prefetch_layer(0, non_blocking=False)
+        self.prepare_for_next_denoise(non_blocking=False)
-        if self.copy_stream is not None:
+
+    def prepare_for_next_denoise(self, non_blocking=True):
+        self.prefetch_layer(0, non_blocking=non_blocking)
+        if not non_blocking and self.copy_stream is not None:
             torch.cuda.current_stream().wait_stream(self.copy_stream)
 
     @torch.compiler.disable

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

@@ -304,6 +304,74 @@
             "expected_avg_denoise_ms": 520.09,
             "expected_median_denoise_ms": 528.0
         },
+
+        "flux_2_image_t2i_layerwise_offload": {
+            "stages_ms": {
+                "InputValidationStage": 0.06,
+                "TextEncodingStage": 513.58,
+                "ImageVAEEncodingStage": 0.0,
+                "ConditioningStage": 0.03,
+                "LatentPreparationStage": 0.46,
+                "TimestepPreparationStage": 2.38,
+                "DenoisingStage": 52187.62,
+                "DecodingStage": 190.31
+            },
+            "denoise_step_ms": {
+                "0": 1033.45,
+                "1": 137.03,
+                "2": 1046.96,
+                "3": 1039.28,
+                "4": 1039.05,
+                "5": 1043.91,
+                "6": 1041.75,
+                "7": 1037.6,
+                "8": 1043.54,
+                "9": 1048.63,
+                "10": 1039.8,
+                "11": 1042.25,
+                "12": 1041.54,
+                "13": 1045.89,
+                "14": 1038.99,
+                "15": 1041.82,
+                "16": 1038.32,
+                "17": 1045.53,
+                "18": 1046.54,
+                "19": 1041.22,
+                "20": 1044.55,
+                "21": 1041.31,
+                "22": 1051.28,
+                "23": 1043.12,
+                "24": 1044.65,
+                "25": 1042.25,
+                "26": 1046.47,
+                "27": 1052.9,
+                "28": 1039.04,
+                "29": 1042.39,
+                "30": 1045.33,
+                "31": 1038.05,
+                "32": 1037.76,
+                "33": 1037.93,
+                "34": 1052.85,
+                "35": 1045.59,
+                "36": 1054.32,
+                "37": 1044.59,
+                "38": 1043.57,
+                "39": 1041.93,
+                "40": 1043.59,
+                "41": 1046.17,
+                "42": 1046.92,
+                "43": 1047.04,
+                "44": 1046.8,
+                "45": 1041.86,
+                "46": 1041.05,
+                "47": 1044.04,
+                "48": 1039.77,
+                "49": 1047.12
+            },
+            "expected_e2e_ms": 53290.15,
+            "expected_avg_denoise_ms": 1025.35,
+            "expected_median_denoise_ms": 1043.33
+        },
         "flux_2_ti2i": {
             "stages_ms": {
                 "InputValidationStage": 99.82,

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

@@ -73,6 +73,9 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
     if server_args.ulysses_degree is not None:
         extra_args += f" --ulysses-degree {server_args.ulysses_degree}"
 
+    if server_args.dit_layerwise_offload:
+        extra_args += f" --dit-layerwise-offload true"
+
     if server_args.ring_degree is not None:
         extra_args += f" --ring-degree {server_args.ring_degree}"
 

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

@@ -718,6 +718,7 @@ def get_generate_fn(
     """Return appropriate generation function for the case."""
     # Allow override via environment variable (useful for AMD where large resolutions cause slow VAE)
     output_size = os.environ.get("SGLANG_TEST_OUTPUT_SIZE", sampling_params.output_size)
+    n = sampling_params.num_outputs_per_prompt
 
     def _create_and_download_video(
         client,
@@ -837,7 +838,7 @@ def get_generate_fn(
         response = client.images.with_raw_response.generate(
             model=model_path,
             prompt=sampling_params.prompt,
-            n=1,
+            n=n,
             size=output_size,
             response_format="b64_json",
         )
@@ -906,7 +907,7 @@ def get_generate_fn(
                 model=model_path,
                 image=images,
                 prompt=sampling_params.prompt,
-                n=1,
+                n=n,
                 size=output_size,
                 response_format="b64_json",
             )
@@ -974,7 +975,7 @@ def get_generate_fn(
             model=model_path,
             prompt=sampling_params.prompt,
             image=[],  # Only for OpenAI verification
-            n=1,
+            n=n,
             size=sampling_params.output_size,
             response_format="b64_json",
             extra_body={"url": image_urls},

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

@@ -152,6 +152,8 @@ class DiffusionServerArgs:
     # LoRA
     lora_path: str | None = None  # LoRA adapter path (HF repo or local path)
 
+    dit_layerwise_offload: bool = False
+
 
 @dataclass(frozen=True)
 class DiffusionSamplingParams:
@@ -171,6 +173,8 @@ class DiffusionSamplingParams:
     # URL direct test flag - if True, don't pre-download URL images
     direct_url_test: bool = False
 
+    num_outputs_per_prompt: int = 1
+
 
 @dataclass(frozen=True)
 class DiffusionTestCase:
@@ -308,6 +312,19 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
         ),
         T2I_sampling_params,
     ),
+    # TODO: replace with a faster model to test the --dit-layerwise-offload
+    # TODO: currently, we don't support sending more than one request in test, and setting `num_outputs_per_prompt` to 2 doesn't guarantee the denoising be executed twice,
+    # so we do one warmup and send one request instead
+    DiffusionTestCase(
+        "flux_2_image_t2i_layerwise_offload",
+        DiffusionServerArgs(
+            model_path="black-forest-labs/FLUX.2-dev",
+            modality="image",
+            dit_layerwise_offload=True,
+            warmup_text=1,
+        ),
+        T2I_sampling_params,
+    ),
     DiffusionTestCase(
         "zimage_image_t2i",
         DiffusionServerArgs(