[diffusion] fix: fix fsdp (#18187)

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.