[diffusion] CI: add test for cache-dit (#16204)
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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)}")
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user