[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/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(