From 4280a18a133d2ca84508f6f7cf1e701a072e6699 Mon Sep 17 00:00:00 2001 From: Yuhao Yang <47235274+yhyang201@users.noreply.github.com> Date: Wed, 31 Dec 2025 19:58:10 +0800 Subject: [PATCH] [diffusion] CI: add test for cache-dit (#16204) Co-authored-by: Mick --- .../test/server/perf_baselines.json | 69 ++++++++++++++++++- .../test/server/test_server_common.py | 6 ++ .../test/server/test_server_utils.py | 5 ++ .../test/server/testcase_configs.py | 28 +++++--- 4 files changed, 98 insertions(+), 10 deletions(-) diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 0e4fc3fd4..f680c5eb1 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/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, diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 7e6b46b0c..969998aa1 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/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() diff --git a/python/sglang/multimodal_gen/test/server/test_server_utils.py b/python/sglang/multimodal_gen/test/server/test_server_utils.py index 2be6cde17..b4c170b2e 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_utils.py +++ b/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)}") diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index be210996a..c3a3763b4 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/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", - # model_path="hunyuanvideo-community/HunyuanVideo", - # modality="video", - # prompt="A curious raccoon", - # output_size="720x480", - # warmup_text=0, - # warmup_edit=0, - # custom_validator="video", + # "hunyuan_video", + # DiffusionServerArgs( + # model_path="hunyuanvideo-community/HunyuanVideo", + # modality="video", + # ), + # DiffusionSamplingParams( + # prompt=T2V_PROMPT, + # ), # ), DiffusionTestCase( "flux_2_ti2i",