[diffusion] fix: fix fsdp (#18187)

python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py

@@ -67,6 +67,9 @@ def _build_server_extra_args(case: DiffusionTestCase) -> str:
         a += f" --lora-path {server_args.lora_path}"
     if server_args.warmup:
         a += " --warmup"
+
+    for extra_arg in server_args.extras:
+        a += f" {extra_arg}"
     return a
 
 

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

@@ -1997,6 +1997,31 @@
             "expected_e2e_ms": 24895.28,
             "expected_avg_denoise_ms": 596.59,
             "expected_median_denoise_ms": 599.66
+        },
+        "fsdp-inference": {
+            "stages_ms": {
+                "InputValidationStage": 0.04,
+                "TextEncodingStage": 128.3,
+                "ConditioningStage": 0.01,
+                "TimestepPreparationStage": 1.44,
+                "LatentPreparationStage": 0.1,
+                "DenoisingStage": 1569.61,
+                "DecodingStage": 41.43
+            },
+            "denoise_step_ms": {
+                "0": 165.33,
+                "1": 158.34,
+                "2": 167.65,
+                "3": 179.11,
+                "4": 183.98,
+                "5": 175.08,
+                "6": 178.34,
+                "7": 178.53,
+                "8": 178.08
+            },
+            "expected_e2e_ms": 1742.7,
+            "expected_avg_denoise_ms": 173.83,
+            "expected_median_denoise_ms": 178.08
         }
     }
 }

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

@@ -90,6 +90,9 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
     if server_args.warmup:
         extra_args += f" --warmup"
 
+    for arg in server_args.extras:
+        extra_args += f" {arg}"
+
     # Build custom environment variables
     env_vars = {}
     if server_args.enable_cache_dit:

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

@@ -21,7 +21,7 @@ from __future__ import annotations
 import json
 import os
 import statistics
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Sequence
 
@@ -151,7 +151,7 @@ class BaselineConfig:
         return self
 
 
-@dataclass(frozen=True)
+@dataclass
 class DiffusionServerArgs:
     """Configuration for a single model/scenario test case."""
 
@@ -183,6 +183,14 @@ class DiffusionServerArgs:
     enable_cache_dit: bool = False
     text_encoder_cpu_offload: bool = False
 
+    extras: list[str] = field(default_factory=lambda: [])
+
+    def __post_init__(self):
+        if self.modality == "image":
+            self.custom_validator = "image"
+        elif self.modality == "video":
+            self.custom_validator = "video"
+
 
 @dataclass(frozen=True)
 class DiffusionSamplingParams:
@@ -331,6 +339,8 @@ TURBOWAN_I2V_sampling_params = DiffusionSamplingParams(
     fps=4,
 )
 
+DEFAULT_SMALL_MODEL = "Tongyi-MAI/Z-Image-Turbo"
+
 # All test cases with clean default values
 # To test different models, simply add more DiffusionCase entries
 ONE_GPU_CASES_A: list[DiffusionTestCase] = [
@@ -644,6 +654,17 @@ TWO_GPU_CASES_A = [
             prompt=T2V_PROMPT,
         ),
     ),
+    DiffusionTestCase(
+        "fsdp-inference",
+        DiffusionServerArgs(
+            model_path=DEFAULT_SMALL_MODEL,
+            modality="image",
+            num_gpus=2,
+            warmup=True,
+            extras=["--use-fsdp-inference"],
+        ),
+        T2I_sampling_params,
+    ),
 ]
 
 # Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536.