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

Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Yuhao Yang
2025-12-31 19:58:10 +08:00
committed by GitHub
parent 386e541520
commit 4280a18a13
4 changed files with 98 additions and 10 deletions

View File

@@ -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,

View File

@@ -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()

View File

@@ -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)}")

View File

@@ -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",