[diffusion] feat: support lightweight e2e warmup for benchmarking (#16213)

python/sglang/multimodal_gen/test/cli/test_generate_common.py

@@ -24,19 +24,13 @@ logger = init_logger(__name__)
 class TestResult:
     name: str
     key: str
-    duration: Optional[float]
     succeed: bool
 
-    @property
-    def duration_str(self):
-        return f"{self.duration:.4f}" if self.duration else "NA"
-
 
 def run_command(command) -> Optional[float]:
     """Runs a command and returns the execution time and status."""
     print(f"Running command: {shlex.join(command)}")
 
-    duration = None
     with subprocess.Popen(
         command,
         stdout=subprocess.PIPE,
@@ -46,15 +40,11 @@ def run_command(command) -> Optional[float]:
     ) as process:
         for line in process.stdout:
             sys.stdout.write(line)
-            if "Pixel data generated" in line:
-                words = line.split(" ")
-                duration = float(words[-2])
-
-    if process.returncode == 0:
-        return duration
-    else:
+        process.wait()
+        if process.returncode == 0:
+            return True
         print(f"Command failed with exit code {process.returncode}")
-        return None
+    return False
 
 
 class CLIBase(unittest.TestCase):
@@ -80,13 +70,7 @@ class CLIBase(unittest.TestCase):
             f"--output-path={self.output_path}",
         ]
 
-    results = []
-
-    @classmethod
-    def setUpClass(cls):
-        cls.results = []
-
-    def _run_command(self, name: str, model_path: str, test_key: str = "", args=[]):
+    def _run_command(self, name: str, model_path: str, args=[]):
         command = (
             self.get_base_command()
             + [f"--model-path={model_path}"]
@@ -94,28 +78,21 @@ class CLIBase(unittest.TestCase):
             + ["--output-file-name", f"{name}"]
             + self.extra_args
         )
-        duration = run_command(command)
-        status = "Success" if duration else "Failed"
-        succeed = duration is not None
+        succeed = run_command(command)
+        status = "Success" if succeed else "Failed"
 
-        duration = float(duration) if succeed else None
-        self.results.append(TestResult(name, test_key, duration, succeed))
-
-        return name, duration, status
+        return name, status
 
     def _run_test(self, name: str, args, model_path: str, test_key: str):
-        name, duration, status = self._run_command(
-            name, args=args, model_path=model_path, test_key=test_key
-        )
-        self.verify(status, name, duration)
+        name, status = self._run_command(name, args=args, model_path=model_path)
+        self.verify(status, name)
 
-    def verify(self, status, name, duration):
+    def verify(self, status, name):
         print("-" * 80)
         print("\n" * 3)
 
         # test task status
         self.assertEqual(status, "Success", f"{name} command failed")
-        self.assertIsNotNone(duration, f"Could not parse duration for {name}")
 
         # test output file
         path = os.path.join(
@@ -125,7 +102,6 @@ class CLIBase(unittest.TestCase):
         if self.data_type == DataType.IMAGE:
             with Image.open(path) as image:
                 check_image_size(self, image, self.width, self.height)
-        logger.info(f"{name} passed in {duration:.4f}s")
 
     def model_name(self):
         return self.model_path.split("/")[-1]

python/sglang/multimodal_gen/test/run_suite.py

@@ -198,7 +198,9 @@ def run_pytest(files, filter_expr=None):
             and "AssertionError" in full_output
         )
 
-        is_flaky_ci_assertion = "SafetensorError" in full_output
+        is_flaky_ci_assertion = (
+            "SafetensorError" in full_output or "FileNotFoundError" in full_output
+        )
 
         is_oom_error = (
             "out of memory" in full_output.lower()

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

@@ -529,6 +529,31 @@
             "expected_avg_denoise_ms": 94.15,
             "expected_median_denoise_ms": 102.03
         },
+        "zimage_image_t2i_warmup": {
+            "stages_ms": {
+                "InputValidationStage": 0.02,
+                "TextEncodingStage": 100.65,
+                "ConditioningStage": 0.01,
+                "TimestepPreparationStage": 0.98,
+                "LatentPreparationStage": 0.06,
+                "DenoisingStage": 889.42,
+                "DecodingStage": 37.81
+            },
+            "denoise_step_ms": {
+                "0": 16.49,
+                "1": 94.63,
+                "2": 109.65,
+                "3": 110.05,
+                "4": 109.39,
+                "5": 110.58,
+                "6": 109.52,
+                "7": 110.54,
+                "8": 115.24
+            },
+            "expected_e2e_ms": 1029.96,
+            "expected_avg_denoise_ms": 98.46,
+            "expected_median_denoise_ms": 109.65
+        },
         "qwen_image_edit_ti2i": {
             "stages_ms": {
                 "InputValidationStage": 38.62,

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

@@ -83,6 +83,9 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
     if server_args.lora_path:
         extra_args += f" --lora-path {server_args.lora_path}"
 
+    if server_args.enable_warmup:
+        extra_args += f" --enable-warmup"
+
     # Build custom environment variables
     env_vars = {}
     if server_args.enable_cache_dit:

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

@@ -151,6 +151,8 @@ class DiffusionServerArgs:
     ring_degree: int | None = None
     # LoRA
     lora_path: str | None = None  # LoRA adapter path (HF repo or local path)
+    # misc
+    enable_warmup: bool = False
 
     dit_layerwise_offload: bool = False
     enable_cache_dit: bool = False
@@ -359,6 +361,13 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
         ),
         T2I_sampling_params,
     ),
+    DiffusionTestCase(
+        "zimage_image_t2i_warmup",
+        DiffusionServerArgs(
+            model_path="Tongyi-MAI/Z-Image-Turbo", modality="image", enable_warmup=True
+        ),
+        T2I_sampling_params,
+    ),
     # === Text and Image to Image (TI2I) ===
     DiffusionTestCase(
         "qwen_image_edit_ti2i",