[diffusion] CI: add test for cache-dit (#16204)

Co-authored-by: Mick <mickjagger19@icloud.com>

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

@@ -304,7 +304,6 @@
             "expected_avg_denoise_ms": 520.09,
             "expected_median_denoise_ms": 528.0
         },
-
         "flux_2_image_t2i_layerwise_offload": {
             "stages_ms": {
                 "InputValidationStage": 0.06,
@@ -597,6 +596,72 @@
             "expected_avg_denoise_ms": 850.73,
             "expected_median_denoise_ms": 854.0
         },
+        "qwen_image_t2i_cache_dit_enabled": {
+            "stages_ms": {
+                "InputValidationStage": 0.05,
+                "TextEncodingStage": 675.95,
+                "ConditioningStage": 0.02,
+                "TimestepPreparationStage": 3.21,
+                "LatentPreparationStage": 0.2,
+                "DenoisingStage": 5248.83,
+                "DecodingStage": 52.24
+            },
+            "denoise_step_ms": {
+                "0": 239.28,
+                "1": 285.07,
+                "2": 286.65,
+                "3": 299.09,
+                "4": 191.11,
+                "5": 101.8,
+                "6": 51.77,
+                "7": 147.26,
+                "8": 101.68,
+                "9": 51.4,
+                "10": 150.58,
+                "11": 106.56,
+                "12": 53.7,
+                "13": 7.33,
+                "14": 148.41,
+                "15": 102.08,
+                "16": 52.04,
+                "17": 7.3,
+                "18": 145.91,
+                "19": 100.31,
+                "20": 51.68,
+                "21": 7.27,
+                "22": 145.92,
+                "23": 104.35,
+                "24": 54.33,
+                "25": 7.46,
+                "26": 150.96,
+                "27": 104.39,
+                "28": 53.21,
+                "29": 7.22,
+                "30": 148.3,
+                "31": 102.06,
+                "32": 51.88,
+                "33": 7.31,
+                "34": 145.8,
+                "35": 101.95,
+                "36": 52.15,
+                "37": 7.42,
+                "38": 148.87,
+                "39": 103.01,
+                "40": 52.21,
+                "41": 7.28,
+                "42": 147.39,
+                "43": 103.89,
+                "44": 51.44,
+                "45": 7.12,
+                "46": 144.78,
+                "47": 100.34,
+                "48": 195.54,
+                "49": 246.97
+            },
+            "expected_e2e_ms": 5982.78,
+            "expected_avg_denoise_ms": 104.84,
+            "expected_median_denoise_ms": 102.01
+        },
         "wan2_1_t2v_1.3b": {
             "stages_ms": {
                 "InputValidationStage": 0.07,
@@ -853,7 +918,7 @@
         },
         "fastwan2_2_ti2v_5b": {
             "stages_ms": {
-                "InputValidationStage": 88.86,
+                "InputValidationStage": 300.00,
                 "TextEncodingStage": 2327.87,
                 "ConditioningStage": 0.01,
                 "TimestepPreparationStage": 58.66,

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

@@ -83,12 +83,18 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
     if server_args.lora_path:
         extra_args += f" --lora-path {server_args.lora_path}"
 
+    # Build custom environment variables
+    env_vars = {}
+    if server_args.enable_cache_dit:
+        env_vars["SGLANG_CACHE_DIT_ENABLED"] = "true"
+
     # start server
     manager = ServerManager(
         model=server_args.model_path,
         port=port,
         wait_deadline=float(os.environ.get("SGLANG_TEST_WAIT_SECS", "1200")),
         extra_args=extra_args,
+        env_vars=env_vars,
     )
     ctx = manager.start()
 

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

@@ -250,11 +250,13 @@ class ServerManager:
         port: int,
         wait_deadline: float = 1200.0,
         extra_args: str = "",
+        env_vars: dict[str, str] | None = None,
     ):
         self.model = model
         self.port = port
         self.wait_deadline = wait_deadline
         self.extra_args = extra_args
+        self.env_vars = env_vars or {}
 
     def _wait_for_rocm_gpu_memory_clear(self, max_wait: float = 60.0) -> None:
         """ROCm-specific: Wait for GPU memory to be mostly free before starting.
@@ -347,6 +349,9 @@ class ServerManager:
         env["SGLANG_DIFFUSION_STAGE_LOGGING"] = "1"
         env["SGLANG_PERF_LOG_DIR"] = log_dir.as_posix()
 
+        # Apply custom environment variables
+        env.update(self.env_vars)
+
         # TODO: unify with run_command
         logger.info(f"Running command: {shlex.join(command)}")
 

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

@@ -153,6 +153,7 @@ class DiffusionServerArgs:
     lora_path: str | None = None  # LoRA adapter path (HF repo or local path)
 
     dit_layerwise_offload: bool = False
+    enable_cache_dit: bool = False
 
 
 @dataclass(frozen=True)
@@ -304,6 +305,17 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
         ),
         T2I_sampling_params,
     ),
+    DiffusionTestCase(
+        "qwen_image_t2i_cache_dit_enabled",
+        DiffusionServerArgs(
+            model_path="Qwen/Qwen-Image",
+            modality="image",
+            warmup_text=1,
+            warmup_edit=0,
+            enable_cache_dit=True,
+        ),
+        T2I_sampling_params,
+    ),
     DiffusionTestCase(
         "flux_image_t2i",
         DiffusionServerArgs(
@@ -414,14 +426,14 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
     ),
     # NOTE(mick): flaky
     # DiffusionTestCase(
-    #     id="hunyuan_video",
+    #     "hunyuan_video",
+    #     DiffusionServerArgs(
     #         model_path="hunyuanvideo-community/HunyuanVideo",
     #         modality="video",
-    #     prompt="A curious raccoon",
+    #     ),
-    #     output_size="720x480",
+    #     DiffusionSamplingParams(
-    #     warmup_text=0,
+    #         prompt=T2V_PROMPT,
-    #     warmup_edit=0,
+    #     ),
-    #     custom_validator="video",
     # ),
     DiffusionTestCase(
         "flux_2_ti2i",