[diffusion] CI: refactor CI (#28762)

.github/workflows/nightly-test-musa.yml

@@ -121,8 +121,8 @@ jobs:
         if: steps.gate.outputs.run_job == 'true'
         timeout-minutes: 30
         run: |
-          pytest python/sglang/multimodal_gen/test/layers/test_musa_rmsnorm.py
+          pytest python/sglang/multimodal_gen/test/unit/musa/layers/test_musa_rmsnorm.py
-          pytest python/sglang/multimodal_gen/test/layers/test_musa_silu_and_mul.py
+          pytest python/sglang/multimodal_gen/test/unit/musa/layers/test_musa_silu_and_mul.py
 
   # ==================== LLM server: 1-GPU ====================
   nightly-test-llm-server-1-gpu-musa:
@@ -219,7 +219,7 @@ jobs:
           RUNAI_STREAMER_MEMORY_LIMIT: 0
         run: |
           cd python
-          python3 sglang/multimodal_gen/test/run_suite_musa.py \
+          python3 sglang/multimodal_gen/test/server/musa/run_suite.py \
             --suite 1-gpu-musa-nightly \
             --partition-id ${{ matrix.part }} \
             --total-partitions 2 \
@@ -268,7 +268,7 @@ jobs:
           RUNAI_STREAMER_MEMORY_LIMIT: 0
         run: |
           cd python
-          python3 sglang/multimodal_gen/test/run_suite_musa.py \
+          python3 sglang/multimodal_gen/test/server/musa/run_suite.py \
             --suite 2-gpu-musa \
             --continue-on-error
 

.github/workflows/pr-test-amd-rocm720.yml

@@ -746,7 +746,7 @@ jobs:
       fail-fast: false
       matrix:
         runner: [linux-mi325-2gpu-sglang]
-        part: [0, 1, 2]  # 3 partitions: 2 parametrized + 1 standalone (test_disagg_server.py)
+        part: [0, 1, 2]  # 3 partitions: 2 parametrized + 1 standalone (single_test_file/test_disagg_server.py)
     runs-on: ${{matrix.runner}}
     steps:
       - name: Checkout code

.github/workflows/pr-test-amd.yml

@@ -767,7 +767,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        part: [0, 1, 2]  # 3 partitions: 2 parametrized + 1 standalone (test_disagg_server.py)
+        part: [0, 1, 2]  # 3 partitions: 2 parametrized + 1 standalone (single_test_file/test_disagg_server.py)
     runs-on: ${{ format('linux-{0}-2gpu-sglang', inputs.runner_arch || 'mi325') }}
     steps:
       - name: Checkout code

.github/workflows/pr-test-musa.yml

@@ -77,9 +77,9 @@ jobs:
             multimodal_gen:
               - "python/pyproject_other.toml"
               - "python/sglang/multimodal_gen/runtime/platforms/musa.py"
-              - "python/sglang/multimodal_gen/test/layers/test_musa_rmsnorm.py"
+              - "python/sglang/multimodal_gen/test/unit/musa/layers/test_musa_rmsnorm.py"
-              - "python/sglang/multimodal_gen/test/layers/test_musa_silu_and_mul.py"
+              - "python/sglang/multimodal_gen/test/unit/musa/layers/test_musa_silu_and_mul.py"
-              - "python/sglang/multimodal_gen/test/run_suite_musa.py"
+              - "python/sglang/multimodal_gen/test/server/musa/run_suite.py"
               - "python/sglang/multimodal_gen/test/server/musa/**"
             sgl_kernel:
               - ".github/workflows/pr-test-musa.yml"
@@ -127,7 +127,7 @@ jobs:
           RUNAI_STREAMER_MEMORY_LIMIT: 0
         run: |
           cd python
-          python3 sglang/multimodal_gen/test/run_suite_musa.py \
+          python3 sglang/multimodal_gen/test/server/musa/run_suite.py \
             --suite 1-gpu-musa \
             --partition-id ${{ matrix.part }} \
             --total-partitions 2
@@ -162,7 +162,7 @@ jobs:
           RUNAI_STREAMER_MEMORY_LIMIT: 0
         run: |
           cd python
-          python3 sglang/multimodal_gen/test/run_suite_musa.py \
+          python3 sglang/multimodal_gen/test/server/musa/run_suite.py \
             --suite 2-gpu-musa
 
   multimodal-gen-layer-unit-test-musa:
@@ -190,8 +190,8 @@ jobs:
       - name: Run multimodal gen layer unit test
         timeout-minutes: 30
         run: |
-          pytest python/sglang/multimodal_gen/test/layers/test_musa_rmsnorm.py
+          pytest python/sglang/multimodal_gen/test/unit/musa/layers/test_musa_rmsnorm.py
-          pytest python/sglang/multimodal_gen/test/layers/test_musa_silu_and_mul.py
+          pytest python/sglang/multimodal_gen/test/unit/musa/layers/test_musa_silu_and_mul.py
 
 # =============================================== sgl-kernel ====================================================
   sgl-kernel-unit-test-musa:

python/sglang/multimodal_gen/runtime/post_training/weights_updater.py

@@ -7,17 +7,16 @@ LLM engine's ModelRunner.update_weights_from_disk.
 
 Detailed usage of higher level API can be found in
 
-/python/sglang/multimodal_gen/test/server/test_update_weights_from_disk.py
+/python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_disk.py
 
 Key design decisions:
 
 - All-or-nothing with rollback: modules are updated sequentially.  If
   any module fails (shape mismatch, corrupted file, etc.), every module
   that was already updated is rolled back by reloading its weights from
-  pipeline.model_path (the last successfully-loaded checkpoint).  On
+  that module's last successfully-loaded weights directory.  On a full
-  success, pipeline.model_path is updated to the new model_path so
+  successful update, pipeline.model_path is updated to the new model_path;
-  that future rollbacks target the latest good checkpoint, not the
+  target_modules updates keep per-module rollback state for hybrid models.
-  originally-launched model.
 
 - Rollback failures propagate: if rollback itself fails, the exception is
   not caught so the caller knows the model is in an inconsistent state.
@@ -228,12 +227,16 @@ class WeightsUpdater:
 
     Args:
         pipeline: A ComposedPipelineBase (or DiffusersPipeline) instance
-            whose modules will be updated.  The pipeline's model_path
+            whose modules will be updated.
-            attribute is used for rollback on failure.
     """
 
     def __init__(self, pipeline):
         self.pipeline = pipeline
+        try:
+            self._module_weight_dirs = pipeline._weights_updater_module_weight_dirs
+        except AttributeError:
+            self._module_weight_dirs = {}
+            pipeline._weights_updater_module_weight_dirs = self._module_weight_dirs
 
     def update_weights_from_disk(
         self,
@@ -282,6 +285,12 @@ class WeightsUpdater:
 
         success, message = self._apply_weights(modules_to_update, weights_map)
 
+        if success:
+            for module_name, _ in modules_to_update:
+                self._module_weight_dirs[module_name] = weights_map[module_name]
+            if target_modules is None:
+                self.pipeline.model_path = local_model_path
+
         gc.collect()
         torch.cuda.empty_cache()
 
@@ -355,12 +364,17 @@ class WeightsUpdater:
         """
         if not updated_modules:
             return
-        original_path = maybe_download_model(self.pipeline.model_path)
+        original_path: str | None = None
         for name in updated_modules:
             module = self.pipeline.get_module(name)
             if module is None:
                 continue
-            weights_dir = Path(original_path) / name
+            weights_dir = self._module_weight_dirs.get(name)
+            if weights_dir is None:
+                if original_path is None:
+                    original_path = maybe_download_model(self.pipeline.model_path)
+                weights_dir = str(Path(original_path) / name)
+            weights_dir = Path(weights_dir)
             if not weights_dir.exists():
                 continue
             weights_iter = _get_weights_iter(str(weights_dir))

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

@@ -1,23 +0,0 @@
-# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
-
-import unittest
-
-from sglang.multimodal_gen.configs.sample.sampling_params import DataType
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-from sglang.multimodal_gen.test.cli.test_generate_common import CLIBase
-from sglang.multimodal_gen.test.test_utils import DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST
-
-logger = init_logger(__name__)
-
-
-class TestFlux_T2V(CLIBase):
-    model_path = DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST
-    extra_args = []
-    data_type: DataType = DataType.IMAGE
-
-
-del CLIBase
-
-
-if __name__ == "__main__":
-    unittest.main()

python/sglang/multimodal_gen/test/run_suite.py

@@ -13,7 +13,6 @@ import random
 import subprocess
 import sys
 import time
-import xml.etree.ElementTree as ET
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -25,6 +24,10 @@ from sglang.multimodal_gen.test.partitioning import (
     PartitionItem,
     partition_items_by_lpt,
 )
+from sglang.multimodal_gen.test.runner.pytest_runner import (
+    partition_items_by_index,
+    run_pytest,
+)
 from sglang.multimodal_gen.test.server.testcase_configs import (
     BASELINE_CONFIG,
     DiffusionTestCase,
@@ -316,93 +319,6 @@ def parse_args():
     return parser.parse_args()
 
 
-def collect_test_items(files: list[str], filter_expr: str | None = None) -> list[str]:
-    """Collect test node IDs from the given files using pytest --collect-only."""
-    cmd = [sys.executable, "-m", "pytest", "--collect-only", "-q"]
-    if filter_expr:
-        cmd.extend(["-k", filter_expr])
-    cmd.extend(files)
-
-    filter_note = f" with filter: {filter_expr}" if filter_expr else ""
-    print(f"Collecting tests from {len(files)} file(s){filter_note}")
-    result = subprocess.run(cmd, capture_output=True, text=True)
-
-    if result.returncode not in (0, 5):
-        error_msg = (
-            f"pytest --collect-only failed with exit code {result.returncode}\n"
-            f"Command: {' '.join(cmd)}\n"
-        )
-        if result.stderr:
-            error_msg += f"stderr:\n{result.stderr}\n"
-        if result.stdout:
-            error_msg += f"stdout:\n{result.stdout}\n"
-        logger.error(error_msg)
-        raise RuntimeError(error_msg)
-
-    if result.returncode == 5:
-        print(
-            "No tests were collected (exit code 5). This may be expected with filters."
-        )
-
-    test_items = []
-    for line in result.stdout.strip().split("\n"):
-        line = line.strip()
-        if line and "::" in line and not line.startswith(("=", "-", " ")):
-            test_id = line.split()[0] if " " in line else line
-            if "::" in test_id:
-                test_items.append(test_id)
-
-    print(f"Collected {len(test_items)} test items")
-    return test_items
-
-
-def parse_junit_xml_for_executed_cases(xml_path: str) -> list[str]:
-    if not Path(xml_path).exists():
-        return []
-
-    executed_cases = []
-    tree = ET.parse(xml_path)
-    root = tree.getroot()
-
-    for testcase in root.iter("testcase"):
-        if testcase.find("skipped") is not None:
-            continue
-
-        name = testcase.get("name", "")
-        if "[" in name and "]" in name:
-            case_id = name[name.index("[") + 1 : name.index("]")]
-            executed_cases.append(case_id)
-
-    return executed_cases
-
-
-def parse_junit_xml_for_case_results(xml_path: str) -> dict[str, str]:
-    if not Path(xml_path).exists():
-        return {}
-
-    case_results = {}
-    tree = ET.parse(xml_path)
-    root = tree.getroot()
-
-    for testcase in root.iter("testcase"):
-        if testcase.find("skipped") is not None:
-            continue
-
-        name = testcase.get("name", "")
-        if "[" not in name or "]" not in name:
-            continue
-
-        case_id = name[name.index("[") + 1 : name.index("]")]
-        if testcase.find("failure") is not None:
-            case_results[case_id] = "fail"
-        elif testcase.find("error") is not None:
-            case_results[case_id] = "error"
-        else:
-            case_results[case_id] = "pass"
-
-    return case_results
-
-
 def write_execution_report(
     suite: str,
     partition_id: int,
@@ -435,274 +351,6 @@ def write_execution_report(
     return str(report_path)
 
 
-def _run_pytest_attempt(cmd: list[str]) -> tuple[int, str]:
-    process = subprocess.Popen(
-        cmd,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        bufsize=0,
-    )
-
-    output_bytes = bytearray()
-    while True:
-        chunk = process.stdout.read(4096)
-        if not chunk:
-            break
-        sys.stdout.buffer.write(chunk)
-        sys.stdout.buffer.flush()
-        output_bytes.extend(chunk)
-
-    process.wait()
-    return process.returncode, output_bytes.decode("utf-8", errors="replace")
-
-
-def _extract_collection_line(full_output: str) -> str | None:
-    for line in full_output.splitlines():
-        stripped = line.strip()
-        if stripped.startswith("collected "):
-            return stripped
-    return None
-
-
-def _extract_short_test_summary(full_output: str) -> list[str]:
-    summary_lines = []
-    in_summary = False
-    for line in full_output.splitlines():
-        stripped = line.strip()
-        if "short test summary info" in stripped:
-            in_summary = True
-            continue
-        if not in_summary:
-            continue
-        if stripped.startswith("="):
-            break
-        if not stripped or stripped.startswith("!"):
-            continue
-        summary_lines.append(stripped)
-    return summary_lines
-
-
-def _extract_failure_tail(full_output: str, max_lines: int = 20) -> list[str]:
-    summary_lines = _extract_short_test_summary(full_output)
-    if summary_lines:
-        return summary_lines
-
-    lines = [line.rstrip() for line in full_output.splitlines() if line.strip()]
-    return lines[-max_lines:]
-
-
-def _summary_has_retryable_failure(summary_lines: list[str]) -> bool:
-    for line in summary_lines:
-        lowered = line.lower()
-        if (
-            "[performance]" in line
-            or "SafetensorError" in line
-            or "FileNotFoundError" in line
-            or "TimeoutError" in line
-            or "out of memory" in lowered
-            or "oom killer" in lowered
-        ):
-            return True
-    return False
-
-
-def _is_consistency_failure(full_output: str) -> bool:
-    summary_lines = _extract_short_test_summary(full_output)
-    for line in summary_lines:
-        if "Consistency check failed for" in line or "GT not found for" in line:
-            return True
-
-    return (
-        "Consistency check failed for " in full_output
-        or "GT not found for " in full_output
-        or "--- MISSING GROUND TRUTH DETECTED ---" in full_output
-    )
-
-
-def _is_retryable_failure(full_output: str) -> bool:
-    if _is_consistency_failure(full_output):
-        return False
-
-    summary_lines = _extract_short_test_summary(full_output)
-    is_perf_assertion = (
-        "multimodal_gen/test/server/test_server_utils.py" in full_output
-        and "AssertionError" in full_output
-    )
-    is_aggregated_retryable_failure = _summary_has_retryable_failure(summary_lines)
-
-    is_flaky_ci_assertion = (
-        "SafetensorError" in full_output
-        or "FileNotFoundError" in full_output
-        or "TimeoutError" in full_output
-    )
-
-    is_oom_error = (
-        "out of memory" in full_output.lower() or "oom killer" in full_output.lower()
-    )
-
-    return (
-        is_perf_assertion
-        or is_aggregated_retryable_failure
-        or is_flaky_ci_assertion
-        or is_oom_error
-    )
-
-
-def _print_attempt_tail_summary(
-    attempt_reports: list[dict], assigned_count: int
-) -> None:
-    if len(attempt_reports) == 1 and attempt_reports[0]["returncode"] in (0, 5):
-        return
-
-    rows = []
-    for report in attempt_reports:
-        if report["returncode"] in (0, 5):
-            result = "success"
-        elif report["retryable"]:
-            result = "retryable failure"
-        else:
-            result = "failure"
-        rows.append(
-            [
-                report["attempt"],
-                report["mode"],
-                result,
-                report["collection_line"] or "-",
-            ]
-        )
-
-    print("\n" + "=" * 32 + " Pytest Tail Summary " + "=" * 32, flush=True)
-    print(f"Assigned {assigned_count} test item(s)", flush=True)
-    print(
-        tabulate.tabulate(
-            rows,
-            headers=["Attempt", "Mode", "Result", "Collection"],
-            tablefmt="psql",
-        ),
-        flush=True,
-    )
-
-    for report in attempt_reports:
-        if not report["failure_tail"]:
-            continue
-        print(f"\nAttempt {report['attempt']} failure summary:", flush=True)
-        for line in report["failure_tail"]:
-            print(f"  {line}", flush=True)
-
-    print("=" * 84, flush=True)
-
-
-def run_pytest(
-    files: list[str],
-    filter_expr: str | None = None,
-    junit_xml_path: str | None = None,
-) -> tuple[int, list[str], dict[str, str]]:
-    if not files:
-        print("No files to run.")
-        return (0, [], {})
-
-    all_executed_cases: set[str] = set()
-    all_case_results: dict[str, str] = {}
-
-    base_cmd = [
-        sys.executable,
-        "-m",
-        "pytest",
-        "-s",
-        "-v",
-        "--tb=short",
-        "--no-header",
-    ]
-    if junit_xml_path:
-        base_cmd.extend(["--junit-xml", junit_xml_path])
-    if filter_expr:
-        base_cmd.extend(["-k", filter_expr])
-
-    max_retries = 6
-    attempt_reports = []
-    for i in range(max_retries + 1):
-        is_retry = i > 0
-        cmd = list(base_cmd)
-        if is_retry:
-            cmd.append("--last-failed")
-        cmd.extend(files)
-
-        mode = "retry failed items" if is_retry else "initial pass"
-        print(
-            f"Starting pytest attempt {i + 1}/{max_retries + 1}: {mode} "
-            f"for {len(files)} assigned item(s)"
-        )
-
-        returncode, full_output = _run_pytest_attempt(cmd)
-        retryable = returncode not in (0, 5) and _is_retryable_failure(full_output)
-        attempt_reports.append(
-            {
-                "attempt": i + 1,
-                "mode": mode,
-                "returncode": returncode,
-                "retryable": retryable,
-                "collection_line": _extract_collection_line(full_output),
-                "failure_tail": (
-                    _extract_failure_tail(full_output)
-                    if returncode not in (0, 5)
-                    else []
-                ),
-            }
-        )
-
-        if junit_xml_path:
-            all_executed_cases.update(
-                parse_junit_xml_for_executed_cases(junit_xml_path)
-            )
-            all_case_results.update(parse_junit_xml_for_case_results(junit_xml_path))
-
-        if returncode == 0:
-            if is_retry:
-                print(f"Recovered retryable failures on attempt {i + 1}.")
-            _print_attempt_tail_summary(attempt_reports, len(files))
-            return (0, list(all_executed_cases), all_case_results)
-        if returncode == 5:
-            print(
-                "No tests collected (exit code 5). This is expected when filters "
-                "deselect all tests in a partition. Treating as success."
-            )
-            _print_attempt_tail_summary(attempt_reports, len(files))
-            return (0, list(all_executed_cases), all_case_results)
-
-        if not retryable:
-            _print_attempt_tail_summary(attempt_reports, len(files))
-            return (returncode, list(all_executed_cases), all_case_results)
-
-        if i == max_retries:
-            print(f"Max retry exceeded ({max_retries})")
-            _print_attempt_tail_summary(attempt_reports, len(files))
-            return (returncode, list(all_executed_cases), all_case_results)
-
-        print(
-            f"Retryable failure detected on attempt {i + 1}. "
-            "Retrying only previously failed items."
-        )
-
-    _print_attempt_tail_summary(attempt_reports, len(files))
-    return (
-        attempt_reports[-1]["returncode"],
-        list(all_executed_cases),
-        all_case_results,
-    )
-
-
-def partition_items_by_index(
-    items: list[str], partition_id: int, total_partitions: int
-) -> list[str]:
-    return [
-        item for i, item in enumerate(items) if i % total_partitions == partition_id
-    ]
-
-
-def partition_test_files(files, partition_id, total_partitions):
-    return partition_items_by_index(files, partition_id, total_partitions)
-
-
 def run_component_accuracy_files(files, filter_expr=None, continue_on_error=False):
     exit_code = 0
     for file_path in files:
@@ -1107,7 +755,7 @@ def main():
             print(f"No valid test files found for suite '{args.suite}'.")
             sys.exit(1 if args.suite in STRICT_SUITES else 0)
 
-        my_files = partition_test_files(
+        my_files = partition_items_by_index(
             suite_files_abs, args.partition_id, args.total_partitions
         )
         partition_info = (

python/sglang/multimodal_gen/test/run_suite_musa.py

@@ -1,270 +0,0 @@
-"""
-Test runner for multimodal_gen MUSA suites that manages partitioned execution.
-
-Usage:
-    python3 run_suite_musa.py --suite <suite_name> --partition-id <id> --total-partitions <num>
-
-Example:
-    python3 run_suite_musa.py --suite 1-gpu-musa --partition-id 0 --total-partitions 2
-"""
-
-import argparse
-import os
-import subprocess
-import sys
-from pathlib import Path
-
-import tabulate
-
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-logger = init_logger(__name__)
-
-SUITES = {
-    "1-gpu-musa": [
-        "musa/test_server_1_gpu_musa.py",
-    ],
-    "1-gpu-musa-nightly": [
-        "musa/test_server_1_gpu_musa_nightly.py",
-    ],
-    "2-gpu-musa": [
-        "musa/test_server_2_gpu_musa.py",
-    ],
-}
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Run multimodal_gen MUSA test suite")
-    parser.add_argument(
-        "--suite",
-        type=str,
-        required=True,
-        choices=list(SUITES.keys()),
-        help="The test suite to run (valid names are defined in SUITES)",
-    )
-    parser.add_argument(
-        "--partition-id",
-        type=int,
-        default=0,
-        help="Index of the current partition (for parallel execution)",
-    )
-    parser.add_argument(
-        "--total-partitions",
-        type=int,
-        default=1,
-        help="Total number of partitions",
-    )
-    parser.add_argument(
-        "--base-dir",
-        type=str,
-        default="server",
-        help="Base directory for tests relative to this script's parent",
-    )
-    parser.add_argument(
-        "-k",
-        "--filter",
-        type=str,
-        default=None,
-        help="Pytest filter expression (passed to pytest -k)",
-    )
-    parser.add_argument(
-        "--continue-on-error",
-        action="store_true",
-        default=False,
-        help="Continue running remaining tests even if one fails.",
-    )
-    return parser.parse_args()
-
-
-def collect_test_items(files, filter_expr=None):
-    """Collect test item node IDs from the given files using pytest --collect-only."""
-    cmd = [sys.executable, "-m", "pytest", "--collect-only", "-q"]
-    if filter_expr:
-        cmd.extend(["-k", filter_expr])
-    cmd.extend(files)
-
-    print(f"Collecting tests with command: {' '.join(cmd)}")
-    result = subprocess.run(cmd, capture_output=True, text=True)
-
-    if result.returncode not in (0, 5):
-        error_msg = (
-            f"pytest --collect-only failed with exit code {result.returncode}\n"
-            f"Command: {' '.join(cmd)}\n"
-        )
-        if result.stderr:
-            error_msg += f"stderr:\n{result.stderr}\n"
-        if result.stdout:
-            error_msg += f"stdout:\n{result.stdout}\n"
-        logger.error(error_msg)
-        raise RuntimeError(error_msg)
-
-    if result.returncode == 5:
-        print(
-            "No tests were collected (exit code 5). This may be expected with filters."
-        )
-
-    test_items = []
-    for line in result.stdout.strip().split("\n"):
-        line = line.strip()
-        if line and "::" in line and not line.startswith(("=", "-", " ")):
-            test_id = line.split()[0] if " " in line else line
-            if "::" in test_id:
-                test_items.append(test_id)
-
-    print(f"Collected {len(test_items)} test items")
-    return test_items
-
-
-def run_pytest(files, filter_expr=None, exitfirst=False):
-    if not files:
-        print("No files to run.")
-        return 0
-
-    base_cmd = [sys.executable, "-m", "pytest", "-s", "-v"]
-    if exitfirst:
-        base_cmd.append("-x")
-
-    if filter_expr:
-        base_cmd.extend(["-k", filter_expr])
-
-    max_retries = 6
-    for i in range(max_retries + 1):
-        cmd = list(base_cmd)
-        if i > 0:
-            cmd.append("--last-failed")
-        cmd.extend(files)
-
-        if i > 0:
-            print(
-                f"Performance assertion failed. Retrying ({i}/{max_retries}) with --last-failed..."
-            )
-
-        print(f"Running command: {' '.join(cmd)}")
-
-        process = subprocess.Popen(
-            cmd,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,
-            bufsize=0,
-        )
-
-        output_bytes = bytearray()
-        while True:
-            chunk = process.stdout.read(4096)
-            if not chunk:
-                break
-            sys.stdout.buffer.write(chunk)
-            sys.stdout.buffer.flush()
-            output_bytes.extend(chunk)
-
-        process.wait()
-        returncode = process.returncode
-
-        if returncode == 0:
-            return 0
-
-        if returncode == 5:
-            print(
-                "No tests collected (exit code 5). This is expected when filters "
-                "deselect all tests in a partition. Treating as success."
-            )
-            return 0
-
-        full_output = output_bytes.decode("utf-8", errors="replace")
-        is_perf_assertion = (
-            "multimodal_gen/test/server/test_server_utils.py" in full_output
-            and "AssertionError" in full_output
-        )
-        is_flaky_ci_assertion = (
-            "SafetensorError" in full_output
-            or "FileNotFoundError" in full_output
-            or "TimeoutError" in full_output
-        )
-        is_oom_error = (
-            "out of memory" in full_output.lower()
-            or "oom killer" in full_output.lower()
-        )
-
-        if not (is_perf_assertion or is_flaky_ci_assertion or is_oom_error):
-            return returncode
-
-    print("Max retry exceeded")
-    return returncode
-
-
-def main():
-    args = parse_args()
-
-    current_file_path = Path(__file__).resolve()
-    test_root_dir = current_file_path.parent
-    target_dir = test_root_dir / args.base_dir
-
-    if not target_dir.exists():
-        print(f"Error: Target directory {target_dir} does not exist.")
-        sys.exit(1)
-
-    suite_files_rel = SUITES[args.suite]
-    suite_files_abs = []
-    for rel_path in suite_files_rel:
-        abs_path = target_dir / rel_path
-        if not abs_path.exists():
-            print(f"Warning: Test file {rel_path} not found in {target_dir}. Skipping.")
-            continue
-        suite_files_abs.append(str(abs_path))
-
-    if not suite_files_abs:
-        print(f"No valid test files found for suite '{args.suite}'.")
-        sys.exit(0)
-
-    all_test_items = collect_test_items(suite_files_abs, filter_expr=args.filter)
-    if not all_test_items:
-        print(f"No test items found for suite '{args.suite}'.")
-        sys.exit(0)
-
-    my_items = [
-        item
-        for i, item in enumerate(all_test_items)
-        if i % args.total_partitions == args.partition_id
-    ]
-
-    partition_info = (
-        f"{args.partition_id + 1}/{args.total_partitions} "
-        f"(0-based id={args.partition_id})"
-    )
-    rows = [[args.suite, partition_info]]
-    msg = (
-        tabulate.tabulate(rows, headers=["Suite", "Partition"], tablefmt="psql") + "\n"
-    )
-    msg += f"Enabled {len(my_items)} test(s):\n"
-    for item in my_items:
-        msg += f"  - {item}\n"
-    print(msg, flush=True)
-    print(
-        f"Suite: {args.suite} | Partition: {args.partition_id}/{args.total_partitions}"
-    )
-    print(f"Selected {len(suite_files_abs)} files:")
-    for file_path in suite_files_abs:
-        print(f"  - {os.path.basename(file_path)}")
-
-    if not my_items:
-        print("No items assigned to this partition. Exiting success.")
-        sys.exit(0)
-
-    print(f"Running {len(my_items)} items in this shard: {', '.join(my_items)}")
-    exit_code = run_pytest(my_items, exitfirst=not args.continue_on_error)
-
-    msg = (
-        "\n"
-        + tabulate.tabulate(rows, headers=["Suite", "Partition"], tablefmt="psql")
-        + "\n"
-    )
-    msg += f"Executed {len(my_items)} test(s):\n"
-    for item in my_items:
-        msg += f"  - {item}\n"
-    print(msg, flush=True)
-
-    sys.exit(exit_code)
-
-
-if __name__ == "__main__":
-    main()

python/sglang/multimodal_gen/test/runner/__init__.py

@@ -0,0 +1 @@
+"""Shared helpers for multimodal_gen test runners."""

python/sglang/multimodal_gen/test/runner/pytest_runner.py

@@ -0,0 +1,364 @@
+from __future__ import annotations
+
+import subprocess
+import sys
+import xml.etree.ElementTree as ET
+from pathlib import Path
+from typing import Sequence
+
+import tabulate
+
+
+def collect_test_items(
+    files: Sequence[str], filter_expr: str | None = None
+) -> list[str]:
+    """Collect pytest node IDs from the given files or node selectors."""
+    cmd = [sys.executable, "-m", "pytest", "--collect-only", "-q"]
+    if filter_expr:
+        cmd.extend(["-k", filter_expr])
+    cmd.extend(files)
+
+    filter_note = f" with filter: {filter_expr}" if filter_expr else ""
+    print(f"Collecting tests from {len(files)} item(s){filter_note}")
+    result = subprocess.run(cmd, capture_output=True, text=True)
+
+    if result.returncode not in (0, 5):
+        error_msg = (
+            f"pytest --collect-only failed with exit code {result.returncode}\n"
+            f"Command: {' '.join(cmd)}\n"
+        )
+        if result.stderr:
+            error_msg += f"stderr:\n{result.stderr}\n"
+        if result.stdout:
+            error_msg += f"stdout:\n{result.stdout}\n"
+        raise RuntimeError(error_msg)
+
+    if result.returncode == 5:
+        print(
+            "No tests were collected (exit code 5). This may be expected with filters."
+        )
+
+    test_items = []
+    for line in result.stdout.strip().split("\n"):
+        line = line.strip()
+        if line and "::" in line and not line.startswith(("=", "-", " ")):
+            test_id = line.split()[0] if " " in line else line
+            if "::" in test_id:
+                test_items.append(test_id)
+
+    print(f"Collected {len(test_items)} test items")
+    return test_items
+
+
+def parse_junit_xml_for_executed_cases(xml_path: str) -> list[str]:
+    if not Path(xml_path).exists():
+        return []
+
+    executed_cases = []
+    tree = ET.parse(xml_path)
+    root = tree.getroot()
+
+    for testcase in root.iter("testcase"):
+        if testcase.find("skipped") is not None:
+            continue
+
+        name = testcase.get("name", "")
+        if "[" in name and "]" in name:
+            case_id = name[name.index("[") + 1 : name.index("]")]
+            executed_cases.append(case_id)
+
+    return executed_cases
+
+
+def parse_junit_xml_for_case_results(xml_path: str) -> dict[str, str]:
+    if not Path(xml_path).exists():
+        return {}
+
+    case_results = {}
+    tree = ET.parse(xml_path)
+    root = tree.getroot()
+
+    for testcase in root.iter("testcase"):
+        if testcase.find("skipped") is not None:
+            continue
+
+        name = testcase.get("name", "")
+        if "[" not in name or "]" not in name:
+            continue
+
+        case_id = name[name.index("[") + 1 : name.index("]")]
+        if testcase.find("failure") is not None:
+            case_results[case_id] = "fail"
+        elif testcase.find("error") is not None:
+            case_results[case_id] = "error"
+        else:
+            case_results[case_id] = "pass"
+
+    return case_results
+
+
+def _run_pytest_attempt(cmd: list[str]) -> tuple[int, str]:
+    process = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        bufsize=0,
+    )
+
+    output_bytes = bytearray()
+    while True:
+        chunk = process.stdout.read(4096)
+        if not chunk:
+            break
+        sys.stdout.buffer.write(chunk)
+        sys.stdout.buffer.flush()
+        output_bytes.extend(chunk)
+
+    process.wait()
+    return process.returncode, output_bytes.decode("utf-8", errors="replace")
+
+
+def _extract_collection_line(full_output: str) -> str | None:
+    for line in full_output.splitlines():
+        stripped = line.strip()
+        if stripped.startswith("collected "):
+            return stripped
+    return None
+
+
+def _extract_short_test_summary(full_output: str) -> list[str]:
+    summary_lines = []
+    in_summary = False
+    for line in full_output.splitlines():
+        stripped = line.strip()
+        if "short test summary info" in stripped:
+            in_summary = True
+            continue
+        if not in_summary:
+            continue
+        if stripped.startswith("="):
+            break
+        if not stripped or stripped.startswith("!"):
+            continue
+        summary_lines.append(stripped)
+    return summary_lines
+
+
+def _extract_failure_tail(full_output: str, max_lines: int = 20) -> list[str]:
+    summary_lines = _extract_short_test_summary(full_output)
+    if summary_lines:
+        return summary_lines
+
+    lines = [line.rstrip() for line in full_output.splitlines() if line.strip()]
+    return lines[-max_lines:]
+
+
+def _summary_has_retryable_failure(summary_lines: list[str]) -> bool:
+    for line in summary_lines:
+        lowered = line.lower()
+        if (
+            "[performance]" in line
+            or "SafetensorError" in line
+            or "FileNotFoundError" in line
+            or "TimeoutError" in line
+            or "out of memory" in lowered
+            or "oom killer" in lowered
+        ):
+            return True
+    return False
+
+
+def _is_consistency_failure(full_output: str) -> bool:
+    summary_lines = _extract_short_test_summary(full_output)
+    for line in summary_lines:
+        if "Consistency check failed for" in line or "GT not found for" in line:
+            return True
+
+    return (
+        "Consistency check failed for " in full_output
+        or "GT not found for " in full_output
+        or "--- MISSING GROUND TRUTH DETECTED ---" in full_output
+    )
+
+
+def _is_retryable_failure(full_output: str) -> bool:
+    if _is_consistency_failure(full_output):
+        return False
+
+    summary_lines = _extract_short_test_summary(full_output)
+    is_perf_assertion = (
+        "multimodal_gen/test/server/test_server_utils.py" in full_output
+        and "AssertionError" in full_output
+    )
+    is_aggregated_retryable_failure = _summary_has_retryable_failure(summary_lines)
+
+    is_flaky_ci_assertion = (
+        "SafetensorError" in full_output
+        or "FileNotFoundError" in full_output
+        or "TimeoutError" in full_output
+    )
+
+    is_oom_error = (
+        "out of memory" in full_output.lower() or "oom killer" in full_output.lower()
+    )
+
+    return (
+        is_perf_assertion
+        or is_aggregated_retryable_failure
+        or is_flaky_ci_assertion
+        or is_oom_error
+    )
+
+
+def _print_attempt_tail_summary(
+    attempt_reports: list[dict], assigned_count: int
+) -> None:
+    if len(attempt_reports) == 1 and attempt_reports[0]["returncode"] in (0, 5):
+        return
+
+    rows = []
+    for report in attempt_reports:
+        if report["returncode"] in (0, 5):
+            result = "success"
+        elif report["retryable"]:
+            result = "retryable failure"
+        else:
+            result = "failure"
+        rows.append(
+            [
+                report["attempt"],
+                report["mode"],
+                result,
+                report["collection_line"] or "-",
+            ]
+        )
+
+    print("\n" + "=" * 32 + " Pytest Tail Summary " + "=" * 32, flush=True)
+    print(f"Assigned {assigned_count} test item(s)", flush=True)
+    print(
+        tabulate.tabulate(
+            rows,
+            headers=["Attempt", "Mode", "Result", "Collection"],
+            tablefmt="psql",
+        ),
+        flush=True,
+    )
+
+    for report in attempt_reports:
+        if not report["failure_tail"]:
+            continue
+        print(f"\nAttempt {report['attempt']} failure summary:", flush=True)
+        for line in report["failure_tail"]:
+            print(f"  {line}", flush=True)
+
+    print("=" * 84, flush=True)
+
+
+def run_pytest(
+    files: Sequence[str],
+    filter_expr: str | None = None,
+    junit_xml_path: str | None = None,
+    exitfirst: bool = False,
+) -> tuple[int, list[str], dict[str, str]]:
+    if not files:
+        print("No files to run.")
+        return (0, [], {})
+
+    all_executed_cases: set[str] = set()
+    all_case_results: dict[str, str] = {}
+
+    base_cmd = [
+        sys.executable,
+        "-m",
+        "pytest",
+        "-s",
+        "-v",
+        "--tb=short",
+        "--no-header",
+    ]
+    if exitfirst:
+        base_cmd.append("-x")
+    if junit_xml_path:
+        base_cmd.extend(["--junit-xml", junit_xml_path])
+    if filter_expr:
+        base_cmd.extend(["-k", filter_expr])
+
+    max_retries = 6
+    attempt_reports = []
+    for i in range(max_retries + 1):
+        is_retry = i > 0
+        cmd = list(base_cmd)
+        if is_retry:
+            cmd.append("--last-failed")
+        cmd.extend(files)
+
+        mode = "retry failed items" if is_retry else "initial pass"
+        print(
+            f"Starting pytest attempt {i + 1}/{max_retries + 1}: {mode} "
+            f"for {len(files)} assigned item(s)"
+        )
+
+        returncode, full_output = _run_pytest_attempt(cmd)
+        retryable = returncode not in (0, 5) and _is_retryable_failure(full_output)
+        attempt_reports.append(
+            {
+                "attempt": i + 1,
+                "mode": mode,
+                "returncode": returncode,
+                "retryable": retryable,
+                "collection_line": _extract_collection_line(full_output),
+                "failure_tail": (
+                    _extract_failure_tail(full_output)
+                    if returncode not in (0, 5)
+                    else []
+                ),
+            }
+        )
+
+        if junit_xml_path:
+            all_executed_cases.update(
+                parse_junit_xml_for_executed_cases(junit_xml_path)
+            )
+            all_case_results.update(parse_junit_xml_for_case_results(junit_xml_path))
+
+        if returncode == 0:
+            if is_retry:
+                print(f"Recovered retryable failures on attempt {i + 1}.")
+            _print_attempt_tail_summary(attempt_reports, len(files))
+            return (0, list(all_executed_cases), all_case_results)
+        if returncode == 5:
+            print(
+                "No tests collected (exit code 5). This is expected when filters "
+                "deselect all tests in a partition. Treating as success."
+            )
+            _print_attempt_tail_summary(attempt_reports, len(files))
+            return (0, list(all_executed_cases), all_case_results)
+
+        if not retryable:
+            _print_attempt_tail_summary(attempt_reports, len(files))
+            return (returncode, list(all_executed_cases), all_case_results)
+
+        if i == max_retries:
+            print(f"Max retry exceeded ({max_retries})")
+            _print_attempt_tail_summary(attempt_reports, len(files))
+            return (returncode, list(all_executed_cases), all_case_results)
+
+        print(
+            f"Retryable failure detected on attempt {i + 1}. "
+            "Retrying only previously failed items."
+        )
+
+    _print_attempt_tail_summary(attempt_reports, len(files))
+    return (
+        attempt_reports[-1]["returncode"],
+        list(all_executed_cases),
+        all_case_results,
+    )
+
+
+def partition_items_by_index(
+    items: Sequence[str], partition_id: int, total_partitions: int
+) -> list[str]:
+    return [
+        item for i, item in enumerate(items) if i % total_partitions == partition_id
+    ]

python/sglang/multimodal_gen/test/server/ascend/test_server_1_npu.py

@@ -7,23 +7,17 @@ If the actual run is significantly better than the baseline, the improved cases
 
 from __future__ import annotations
 
-import pytest
-
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.test.server.ascend.testcase_configs_npu import ONE_NPU_CASES
+from sglang.multimodal_gen.test.server.common.case_fixtures import (
+    diffusion_case_fixture,
+)
 from sglang.multimodal_gen.test.server.test_server_common import (  # noqa: F401
     DiffusionServerBase,
     diffusion_server,
 )
-from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
-
-logger = init_logger(__name__)
 
 
 class TestDiffusionServerOneNpu(DiffusionServerBase):
     """Performance tests for 1-NPU diffusion cases."""
 
-    @pytest.fixture(params=ONE_NPU_CASES, ids=lambda c: c.id)
+    case = diffusion_case_fixture(ONE_NPU_CASES)
-    def case(self, request) -> DiffusionTestCase:
-        """Provide a DiffusionTestCase for each 1-NPU test."""
-        return request.param

python/sglang/multimodal_gen/test/server/ascend/test_server_2_npu.py

@@ -7,23 +7,17 @@ If the actual run is significantly better than the baseline, the improved cases
 
 from __future__ import annotations
 
-import pytest
-
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.test.server.ascend.testcase_configs_npu import TWO_NPU_CASES
+from sglang.multimodal_gen.test.server.common.case_fixtures import (
+    diffusion_case_fixture,
+)
 from sglang.multimodal_gen.test.server.test_server_common import (  # noqa: F401
     DiffusionServerBase,
     diffusion_server,
 )
-from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
-
-logger = init_logger(__name__)
 
 
 class TestDiffusionServerTwoNpu(DiffusionServerBase):
     """Performance tests for 2-NPU diffusion cases."""
 
-    @pytest.fixture(params=TWO_NPU_CASES, ids=lambda c: c.id)
+    case = diffusion_case_fixture(TWO_NPU_CASES)
-    def case(self, request) -> DiffusionTestCase:
-        """Provide a DiffusionTestCase for each 2-NPU test."""
-        return request.param

python/sglang/multimodal_gen/test/server/common/__init__.py

@@ -0,0 +1 @@
+"""Common server test helpers."""

python/sglang/multimodal_gen/test/server/common/case_fixtures.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
+import pytest
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
+
+
+def diffusion_case_fixture(cases: Sequence[DiffusionTestCase]):
+    @pytest.fixture(params=cases, ids=lambda case: case.id)
+    def case(self, request) -> DiffusionTestCase:
+        return request.param
+
+    return case

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

@@ -232,6 +232,7 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [
         "wan2_1_t2v_1.3b",
         DiffusionServerArgs(
             model_path=DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST,
+            modality="video",
         ),
     ),
     DiffusionTestCase(
@@ -466,12 +467,8 @@ if not current_platform.is_hip():
             DiffusionServerArgs(
                 model_path="IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers",
             ),
-            # Pin CI's shared T2V_PROMPT ("A curious raccoon") instead of relying on
-            # prompt=None / unconditional generation — the latter drifts as the
-            # pipeline evolves, which is why this (previously planner-invisible) case
-            # diverged from its stale sglang_generated GT.
             T2V_sampling_params,
-        )
+        ),
     )
 # Skip all ModelOpt tests on AMD: FP8 requires torch._scaled_mm (HIPBLAS_STATUS_NOT_SUPPORTED
 # on ROCm), NVFP4 requires flashinfer or sgl_kernel FP4 kernels (CUDA-only).
@@ -864,14 +861,14 @@ def _discover_unit_tests() -> list[str]:
 FILE_SUITES = {
     "unit": _discover_unit_tests(),
     "component-accuracy": [
-        "test_component_accuracy_1_gpu.py",
+        "../single_test_file/component_accuracy/test_component_accuracy_1_gpu.py",
-        "test_component_accuracy_2_gpu.py",
+        "../single_test_file/component_accuracy/test_component_accuracy_2_gpu.py",
     ],
     "component-accuracy-1-gpu": [
-        "test_component_accuracy_1_gpu.py",
+        "../single_test_file/component_accuracy/test_component_accuracy_1_gpu.py",
     ],
     "component-accuracy-2-gpu": [
-        "test_component_accuracy_2_gpu.py",
+        "../single_test_file/component_accuracy/test_component_accuracy_2_gpu.py",
     ],
     "1-gpu-b200": [
         "test_server_b200.py",
@@ -889,13 +886,11 @@ PARAMETRIZED_CASE_GROUPS = {
 
 STANDALONE_FILES = {
     "1-gpu": [
-        "../cli/test_generate_t2i_perf.py",
+        "../single_test_file/test_generate_zimage_turbo_cli.py",
-        # Temporarily disabled: 24 timeout failures since 2026-04-09 across
+        "../single_test_file/test_update_weights_from_disk.py",
-        # multimodal-gen-test-1-gpu. Re-enable after the flakiness is fixed.
-        # "test_update_weights_from_disk.py",
     ],
     "2-gpu": [
-        "test_disagg_server.py",
+        "../single_test_file/test_disagg_server.py",
     ],
 }
 
@@ -904,14 +899,12 @@ STANDALONE_FILES = {
 # measured value that must be copied into STANDALONE_FILE_EST_TIMES.
 STANDALONE_FILE_EST_TIMES = {
     "1-gpu": {
-        "../cli/test_generate_t2i_perf.py": 240.0,
+        "../single_test_file/test_update_weights_from_disk.py": 1200.0,
-        # See STANDALONE_FILES note above — temporarily disabled.
-        # "test_update_weights_from_disk.py": 480.0,
     },
     "2-gpu": {
         # Two disagg clusters × (~3 min startup + ~1 min generate) ≈ 8 min.
         # Raise if CI reports a higher measured time.
-        "test_disagg_server.py": 600.0,
+        "../single_test_file/test_disagg_server.py": 600.0,
     },
 }
 
@@ -940,9 +933,8 @@ DEFAULT_EST_TIME_SECONDS = 300.0
 STARTUP_OVERHEAD_SECONDS = 120.0
 DEFAULT_STANDALONE_EST_TIME_SECONDS = 300.0
 
-_UPDATE_WEIGHTS_FROM_DISK_TEST_FILE = "test_update_weights_from_disk.py"
+_UPDATE_WEIGHTS_FROM_DISK_TEST_FILE = (
-_UPDATE_WEIGHTS_MODEL_PAIR_ENV = "SGLANG_MMGEN_UPDATE_WEIGHTS_PAIR"
+    "../single_test_file/test_update_weights_from_disk.py"
-_UPDATE_WEIGHTS_MODEL_PAIR_IDS = (
-    "FLUX.2-klein-base-4B",
-    "Qwen-Image",
 )
+_UPDATE_WEIGHTS_MODEL_PAIR_ENV = "SGLANG_MMGEN_UPDATE_WEIGHTS_PAIR"
+_UPDATE_WEIGHTS_MODEL_PAIR_IDS = ("FLUX.2-klein-base-4B",)

python/sglang/multimodal_gen/test/server/musa/run_suite.py

@@ -0,0 +1,172 @@
+"""
+Test runner for multimodal_gen MUSA suites that manages partitioned execution.
+
+Usage:
+    python3 -m sglang.multimodal_gen.test.server.musa.run_suite --suite <suite_name>
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+import tabulate
+
+TEST_ROOT = Path(__file__).resolve().parents[2]
+if str(TEST_ROOT) not in sys.path:
+    sys.path.insert(0, str(TEST_ROOT))
+
+from runner.pytest_runner import (  # noqa: E402
+    collect_test_items,
+    partition_items_by_index,
+    run_pytest,
+)
+
+SUITES = {
+    "1-gpu-musa": [
+        "test_server_1_gpu_musa.py",
+    ],
+    "1-gpu-musa-nightly": [
+        "test_server_1_gpu_musa_nightly.py",
+    ],
+    "2-gpu-musa": [
+        "test_server_2_gpu_musa.py",
+    ],
+}
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Run multimodal_gen MUSA test suite")
+    parser.add_argument(
+        "--suite",
+        type=str,
+        required=True,
+        choices=list(SUITES.keys()),
+        help="The test suite to run (valid names are defined in SUITES)",
+    )
+    parser.add_argument(
+        "--partition-id",
+        type=int,
+        default=0,
+        help="Index of the current partition (for parallel execution)",
+    )
+    parser.add_argument(
+        "--total-partitions",
+        type=int,
+        default=1,
+        help="Total number of partitions",
+    )
+    parser.add_argument(
+        "--base-dir",
+        type=str,
+        default=None,
+        help=(
+            "Base directory for tests relative to multimodal_gen/test. "
+            "Defaults to server/musa."
+        ),
+    )
+    parser.add_argument(
+        "-k",
+        "--filter",
+        type=str,
+        default=None,
+        help="Pytest filter expression (passed to pytest -k)",
+    )
+    parser.add_argument(
+        "--continue-on-error",
+        action="store_true",
+        default=False,
+        help="Continue running remaining tests even if one fails.",
+    )
+    return parser.parse_args()
+
+
+def _resolve_suite_files(
+    suite: str, test_root_dir: Path, musa_dir: Path, base_dir: str | None
+) -> tuple[Path, list[str]]:
+    target_dir = test_root_dir / base_dir if base_dir else musa_dir
+    suite_files_rel = SUITES[suite]
+    if target_dir == musa_dir:
+        return target_dir, suite_files_rel
+
+    musa_rel = musa_dir.relative_to(target_dir)
+    return target_dir, [str(musa_rel / rel_path) for rel_path in suite_files_rel]
+
+
+def main():
+    args = parse_args()
+
+    musa_dir = Path(__file__).resolve().parent
+    test_root_dir = musa_dir.parent.parent
+    target_dir, suite_files_rel = _resolve_suite_files(
+        args.suite, test_root_dir, musa_dir, args.base_dir
+    )
+
+    if not target_dir.exists():
+        print(f"Error: Target directory {target_dir} does not exist.")
+        sys.exit(1)
+
+    suite_files_abs = []
+    for rel_path in suite_files_rel:
+        abs_path = target_dir / rel_path
+        if not abs_path.exists():
+            print(f"Warning: Test file {rel_path} not found in {target_dir}. Skipping.")
+            continue
+        suite_files_abs.append(str(abs_path))
+
+    if not suite_files_abs:
+        print(f"No valid test files found for suite '{args.suite}'.")
+        sys.exit(0)
+
+    all_test_items = collect_test_items(suite_files_abs, filter_expr=args.filter)
+    if not all_test_items:
+        print(f"No test items found for suite '{args.suite}'.")
+        sys.exit(0)
+
+    my_items = partition_items_by_index(
+        all_test_items, args.partition_id, args.total_partitions
+    )
+
+    partition_info = (
+        f"{args.partition_id + 1}/{args.total_partitions} "
+        f"(0-based id={args.partition_id})"
+    )
+    rows = [[args.suite, partition_info]]
+    msg = (
+        tabulate.tabulate(rows, headers=["Suite", "Partition"], tablefmt="psql") + "\n"
+    )
+    msg += f"Enabled {len(my_items)} test(s):\n"
+    for item in my_items:
+        msg += f"  - {item}\n"
+    print(msg, flush=True)
+    print(
+        f"Suite: {args.suite} | Partition: {args.partition_id}/{args.total_partitions}"
+    )
+    print(f"Selected {len(suite_files_abs)} files:")
+    for file_path in suite_files_abs:
+        print(f"  - {os.path.basename(file_path)}")
+
+    if not my_items:
+        print("No items assigned to this partition. Exiting success.")
+        sys.exit(0)
+
+    print(f"Running {len(my_items)} items in this shard: {', '.join(my_items)}")
+    exit_code, _, _ = run_pytest(my_items, exitfirst=not args.continue_on_error)
+
+    msg = (
+        "\n"
+        + tabulate.tabulate(rows, headers=["Suite", "Partition"], tablefmt="psql")
+        + "\n"
+    )
+    msg += f"Executed {len(my_items)} test(s):\n"
+    for item in my_items:
+        msg += f"  - {item}\n"
+    print(msg, flush=True)
+
+    sys.exit(exit_code)
+
+
+if __name__ == "__main__":
+    main()

python/sglang/multimodal_gen/test/server/musa/test_server_1_gpu_musa.py

@@ -4,9 +4,9 @@ MUSA-specific 1-GPU diffusion performance tests.
 
 from __future__ import annotations
 
-import pytest
+from sglang.multimodal_gen.test.server.common.case_fixtures import (
-
+    diffusion_case_fixture,
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+)
 from sglang.multimodal_gen.test.server.musa.testcase_configs_musa import (
     ONE_GPU_MUSA_CASES,
 )
@@ -14,15 +14,9 @@ from sglang.multimodal_gen.test.server.test_server_common import (  # noqa: F401
     DiffusionServerBase,
     diffusion_server,
 )
-from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
-
-logger = init_logger(__name__)
 
 
 class TestDiffusionServerOneGpuMusa(DiffusionServerBase):
     """Performance tests for 1-GPU diffusion cases on MUSA."""
 
-    @pytest.fixture(params=ONE_GPU_MUSA_CASES, ids=lambda c: c.id)
+    case = diffusion_case_fixture(ONE_GPU_MUSA_CASES)
-    def case(self, request) -> DiffusionTestCase:
-        """Provide a DiffusionTestCase for each 1-GPU MUSA test."""
-        return request.param

python/sglang/multimodal_gen/test/server/musa/test_server_1_gpu_musa_nightly.py

@@ -4,9 +4,9 @@ MUSA-specific 1-GPU diffusion performance tests for nightly suite.
 
 from __future__ import annotations
 
-import pytest
+from sglang.multimodal_gen.test.server.common.case_fixtures import (
-
+    diffusion_case_fixture,
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+)
 from sglang.multimodal_gen.test.server.musa.testcase_configs_musa import (
     ONE_GPU_NIGHTLY_MUSA_CASES,
 )
@@ -14,15 +14,9 @@ from sglang.multimodal_gen.test.server.test_server_common import (  # noqa: F401
     DiffusionServerBase,
     diffusion_server,
 )
-from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
-
-logger = init_logger(__name__)
 
 
 class TestDiffusionServerOneGpuMusaNightly(DiffusionServerBase):
     """Performance tests for 1-GPU diffusion cases on MUSA (nightly-only)."""
 
-    @pytest.fixture(params=ONE_GPU_NIGHTLY_MUSA_CASES, ids=lambda c: c.id)
+    case = diffusion_case_fixture(ONE_GPU_NIGHTLY_MUSA_CASES)
-    def case(self, request) -> DiffusionTestCase:
-        """Provide a DiffusionTestCase for each 1-GPU MUSA nightly test."""
-        return request.param

python/sglang/multimodal_gen/test/server/musa/test_server_2_gpu_musa.py

@@ -4,9 +4,9 @@ MUSA-specific 2-GPU diffusion performance tests.
 
 from __future__ import annotations
 
-import pytest
+from sglang.multimodal_gen.test.server.common.case_fixtures import (
-
+    diffusion_case_fixture,
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+)
 from sglang.multimodal_gen.test.server.musa.testcase_configs_musa import (
     TWO_GPU_MUSA_CASES,
 )
@@ -14,15 +14,9 @@ from sglang.multimodal_gen.test.server.test_server_common import (  # noqa: F401
     DiffusionServerBase,
     diffusion_server,
 )
-from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
-
-logger = init_logger(__name__)
 
 
 class TestDiffusionServerTwoGpuMusa(DiffusionServerBase):
     """Performance tests for 2-GPU diffusion cases on MUSA."""
 
-    @pytest.fixture(params=TWO_GPU_MUSA_CASES, ids=lambda c: c.id)
+    case = diffusion_case_fixture(TWO_GPU_MUSA_CASES)
-    def case(self, request) -> DiffusionTestCase:
-        """Provide a DiffusionTestCase for each 2-GPU MUSA test."""
-        return request.param

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

@@ -7,23 +7,17 @@ If the actual run is significantly better than the baseline, the improved cases
 
 from __future__ import annotations
 
-import pytest
+from sglang.multimodal_gen.test.server.common.case_fixtures import (
-
+    diffusion_case_fixture,
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+)
 from sglang.multimodal_gen.test.server.gpu_cases import ONE_GPU_CASES
 from sglang.multimodal_gen.test.server.test_server_common import (  # noqa: F401
     DiffusionServerBase,
     diffusion_server,
 )
-from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
-
-logger = init_logger(__name__)
 
 
 class TestDiffusionServerOneGpu(DiffusionServerBase):
     """Performance tests for 1-GPU diffusion cases."""
 
-    @pytest.fixture(params=ONE_GPU_CASES, ids=lambda c: c.id)
+    case = diffusion_case_fixture(ONE_GPU_CASES)
-    def case(self, request) -> DiffusionTestCase:
-        """Provide a DiffusionTestCase for each 1-GPU test."""
-        return request.param

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

@@ -4,20 +4,17 @@
 
 from __future__ import annotations
 
-import pytest
+from sglang.multimodal_gen.test.server.common.case_fixtures import (
-
+    diffusion_case_fixture,
+)
 from sglang.multimodal_gen.test.server.gpu_cases import TWO_GPU_CASES
 from sglang.multimodal_gen.test.server.test_server_common import (  # noqa: F401
     DiffusionServerBase,
     diffusion_server,
 )
-from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
 
 
 class TestDiffusionServerTwoGpu(DiffusionServerBase):
     """Performance tests for 2-GPU diffusion cases."""
 
-    @pytest.fixture(params=TWO_GPU_CASES, ids=lambda c: c.id)
+    case = diffusion_case_fixture(TWO_GPU_CASES)
-    def case(self, request) -> DiffusionTestCase:
-        """Provide a DiffusionTestCase for each 2-GPU test."""
-        return request.param

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

@@ -4,23 +4,17 @@ Config-driven diffusion performance test with pytest parametrization.
 
 from __future__ import annotations
 
-import pytest
+from sglang.multimodal_gen.test.server.common.case_fixtures import (
-
+    diffusion_case_fixture,
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+)
 from sglang.multimodal_gen.test.server.gpu_cases import ONE_GPU_B200_CASES
 from sglang.multimodal_gen.test.server.test_server_common import (  # noqa: F401
     DiffusionServerBase,
     diffusion_server,
 )
-from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
-
-logger = init_logger(__name__)
 
 
 class TestDiffusionServerOneGpuB200(DiffusionServerBase):
     """B200-targeted CI tests for 1-GPU Blackwell-only diffusion cases."""
 
-    @pytest.fixture(params=ONE_GPU_B200_CASES, ids=lambda c: c.id)
+    case = diffusion_case_fixture(ONE_GPU_B200_CASES)
-    def case(self, request) -> DiffusionTestCase:
-        """Provide a DiffusionTestCase for each 1-GPU B200 test."""
-        return request.param

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

@@ -29,6 +29,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import (
     init_logger,
 )
 from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord
+from sglang.multimodal_gen.test.server.common.slack import upload_file_to_slack
 from sglang.multimodal_gen.test.server.realtime_consistency import (
     build_realtime_init_payload,
     collect_realtime_output,
@@ -44,7 +45,6 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
     ScenarioConfig,
     ToleranceConfig,
 )
-from sglang.multimodal_gen.test.slack_utils import upload_file_to_slack
 from sglang.multimodal_gen.test.test_utils import (
     get_expected_image_format,
     get_video_frame_count,

python/sglang/multimodal_gen/test/single_test_file/__init__.py

@@ -0,0 +1 @@
+"""Standalone multimodal_gen test files run as CI single-file suites."""

python/sglang/multimodal_gen/test/single_test_file/cli_generate_common.py

@@ -4,62 +4,68 @@
 Common generate cli test, one test for image and video each
 """
 
-import dataclasses
 import os
 import shlex
+import tempfile
 import unittest
+from typing import Any
 
 from PIL import Image
 
-from sglang.multimodal_gen.configs.sample.sampling_params import DataType
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-from sglang.multimodal_gen.test.test_utils import check_image_size, run_command
 
-logger = init_logger(__name__)
+def run_command(command: list[str]) -> bool:
+    from sglang.multimodal_gen.test.test_utils import run_command as _run_command
+
+    return _run_command(command)
 
 
-@dataclasses.dataclass
+def check_image_size(ut, image, width, height):
-class TestResult:
+    ut.assertEqual(image.size, (width, height))
-    name: str
-    key: str
-    succeed: bool
 
 
 class CLIBase(unittest.TestCase):
     model_path: str = None
-    extra_args = []
+    extra_args: tuple[str, ...] = ()
-    data_type: DataType = None
+    data_type: Any = None
-    log_level: str = "info"
-    # tested on h100
 
+    log_level: str = "info"
     width: int = 720
     height: int = 720
-    output_path: str = "test_outputs"
+    output_path: str | None = None
 
     def setUp(self):
         super().setUp()
-        if not os.path.exists(self.output_path):
+        self._temp_output_dir = None
+        if self.output_path is None:
+            self._temp_output_dir = tempfile.TemporaryDirectory(
+                prefix="sglang_cli_test_"
+            )
+            self.output_path = self._temp_output_dir.name
+        else:
             os.makedirs(self.output_path, exist_ok=True)
-        if os.path.exists(self.output_path):
+            self._clear_output_files()
-            for f in os.listdir(self.output_path):
-                path = os.path.join(self.output_path, f)
-                if os.path.isfile(path):
-                    os.remove(path)
 
     def tearDown(self):
-        super().tearDown()
+        try:
-        if os.path.exists(self.output_path):
+            if self._temp_output_dir is not None:
-            for f in os.listdir(self.output_path):
+                self._temp_output_dir.cleanup()
-                path = os.path.join(self.output_path, f)
+            elif self.output_path and os.path.exists(self.output_path):
-                if os.path.isfile(path):
+                self._clear_output_files()
-                    os.remove(path)
+        finally:
+            super().tearDown()
+
+    def _clear_output_files(self):
+        for filename in os.listdir(self.output_path):
+            path = os.path.join(self.output_path, filename)
+            if os.path.isfile(path):
+                os.remove(path)
 
     def get_base_command(self):
         return [
             "sglang",
             "generate",
             "--prompt",
-            "A curious raccoon",
+            "A red cube on a white table",
             "--save-output",
             f"--log-level={self.log_level}",
             f"--width={self.width}",
@@ -67,13 +73,13 @@ class CLIBase(unittest.TestCase):
             f"--output-path={self.output_path}",
         ]
 
-    def _run_command(self, name: str, model_path: str, args=[]):
+    def _run_command(self, name: str, model_path: str, args: str | None = None):
         command = (
             self.get_base_command()
             + [f"--model-path={model_path}"]
             + shlex.split(args or "")
             + ["--output-file-name", f"{name}"]
-            + self.extra_args
+            + list(self.extra_args)
         )
         succeed = run_command(command)
         status = "Success" if succeed else "Failed"
@@ -96,7 +102,7 @@ class CLIBase(unittest.TestCase):
             self.output_path, f"{name}.{self.data_type.get_default_extension()}"
         )
         self.assertTrue(os.path.exists(path), f"Output file not exist for {path}")
-        if self.data_type == DataType.IMAGE:
+        if self.data_type.get_default_extension() in ("png", "jpg", "jpeg", "webp"):
             with Image.open(path) as image:
                 check_image_size(self, image, self.width, self.height)
 

python/sglang/multimodal_gen/test/single_test_file/component_accuracy/__init__.py

@@ -0,0 +1 @@
+"""Component accuracy test helpers and suites."""

python/sglang/multimodal_gen/test/single_test_file/component_accuracy/engine.py

@@ -48,14 +48,15 @@ from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
     get_diffusers_component_config,
 )
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-from sglang.multimodal_gen.test.server.accuracy_config import (
+from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.config import (
     DEFAULT_TIMESTEP,
     ComponentType,
 )
-from sglang.multimodal_gen.test.server.accuracy_hooks import (
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.hooks import (
     resolve_component_native_profile,
 )
-from sglang.multimodal_gen.test.server.accuracy_utils import (
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import (
     build_accuracy_server_args,
     build_parameter_shard_contexts,
     build_state_lookup,
@@ -70,7 +71,6 @@ from sglang.multimodal_gen.test.server.accuracy_utils import (
     resolve_text_encoder_forward_module,
     select_component_source,
 )
-from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
 
 logger = init_logger(__name__)
 

python/sglang/multimodal_gen/test/single_test_file/component_accuracy/hooks.py

@@ -7,13 +7,13 @@ from typing import Any, Callable, Dict, Optional
 import torch
 import torch.nn as nn
 
-from sglang.multimodal_gen.test.server.accuracy_config import (
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.config import (
     DEFAULT_TIMESTEP,
     I2V_IMAGE_DIM,
     TIMESTEP_NORMALIZATION_FACTOR,
     ComponentType,
 )
-from sglang.multimodal_gen.test.server.accuracy_utils import (
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import (
     extract_output_tensor,
     seed_and_broadcast,
 )

python/sglang/multimodal_gen/test/single_test_file/component_accuracy/test_component_accuracy_1_gpu.py

@@ -1,19 +1,21 @@
 import pytest
 
-from sglang.multimodal_gen.test.server.accuracy_config import (
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.config import (
     ComponentType,
     get_skip_reason,
     should_skip_component,
 )
-from sglang.multimodal_gen.test.server.accuracy_testcase_configs import (
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.engine import (
+    AccuracyEngine,
+)
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.testcase_configs import (
     ACCURACY_ONE_GPU_CASES,
     get_component_duplicate_skip_reason,
 )
-from sglang.multimodal_gen.test.server.accuracy_utils import (
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import (
     run_native_component_accuracy_case,
     run_text_encoder_accuracy_case,
 )
-from sglang.multimodal_gen.test.server.component_accuracy import AccuracyEngine
 
 VAE_CHANNELS_LAST_3D_PARITY_CASE_IDS = {
     "wan2_1_t2v_1.3b",

python/sglang/multimodal_gen/test/single_test_file/component_accuracy/test_component_accuracy_2_gpu.py

@@ -1,19 +1,21 @@
 import pytest
 
-from sglang.multimodal_gen.test.server.accuracy_config import (
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.config import (
     ComponentType,
     get_skip_reason,
     should_skip_component,
 )
-from sglang.multimodal_gen.test.server.accuracy_testcase_configs import (
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.engine import (
+    AccuracyEngine,
+)
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.testcase_configs import (
     ACCURACY_TWO_GPU_CASES,
     get_component_duplicate_skip_reason,
 )
-from sglang.multimodal_gen.test.server.accuracy_utils import (
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import (
     run_native_component_accuracy_case,
     run_text_encoder_accuracy_case,
 )
-from sglang.multimodal_gen.test.server.component_accuracy import AccuracyEngine
 
 VAE_CHANNELS_LAST_3D_PARITY_CASE_IDS = {
     "wan2_2_i2v_a14b_2gpu",

python/sglang/multimodal_gen/test/single_test_file/component_accuracy/testcase_configs.py

@@ -1,18 +1,20 @@
 from __future__ import annotations
 
-from sglang.multimodal_gen.test.server.accuracy_config import (
-    ComponentType,
-    should_skip_component,
-)
-from sglang.multimodal_gen.test.server.accuracy_utils import (
-    extract_component_path_overrides,
-)
-from sglang.multimodal_gen.test.server.component_accuracy import COMPONENT_SPECS
 from sglang.multimodal_gen.test.server.gpu_cases import (
     ONE_GPU_CASES,
     TWO_GPU_CASES,
 )
 from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.config import (
+    ComponentType,
+    should_skip_component,
+)
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.engine import (
+    COMPONENT_SPECS,
+)
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import (
+    extract_component_path_overrides,
+)
 
 
 def _component_accuracy_key(case: DiffusionTestCase, component: ComponentType) -> tuple:

python/sglang/multimodal_gen/test/single_test_file/component_accuracy/utils.py

@@ -28,7 +28,7 @@ from sglang.multimodal_gen.runtime.utils.model_overlay import (
     load_overlay_manifest_if_present,
     resolve_model_overlay_target,
 )
-from sglang.multimodal_gen.test.server.accuracy_config import (
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.config import (
     DEFAULT_TEXT_ENCODER_VOCAB_SIZE,
     I2V_TEXT_ENCODER_DIM,
     TEXT_ENCODER_INPUT_SEED,
@@ -709,7 +709,7 @@ def _run_staged_native_component_accuracy_case(
     library: str,
     num_gpus: int,
 ) -> None:
-    from sglang.multimodal_gen.test.server.accuracy_hooks import (
+    from sglang.multimodal_gen.test.single_test_file.component_accuracy.hooks import (
         resolve_component_native_profile,
     )
 

python/sglang/multimodal_gen/test/single_test_file/test_disagg_server.py

@@ -16,7 +16,7 @@ Two configurations are covered:
 
 Run directly:
 
-    pytest -v python/sglang/multimodal_gen/test/server/test_disagg_server.py
+    pytest -v python/sglang/multimodal_gen/test/single_test_file/test_disagg_server.py
     pytest -v ... -k ZImage1Rank              # one class
     pytest -v ... -k test_generates_image     # one test
 """

python/sglang/multimodal_gen/test/single_test_file/test_generate_i2i.py

@@ -4,7 +4,10 @@ import unittest
 from PIL import Image
 
 from sglang.multimodal_gen.configs.sample.sampling_params import DataType
-from sglang.multimodal_gen.test.cli.test_generate_common import CLIBase, run_command
+from sglang.multimodal_gen.test.single_test_file.cli_generate_common import (
+    CLIBase,
+    run_command,
+)
 from sglang.multimodal_gen.test.test_utils import (
     DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST,
     check_image_size,

python/sglang/multimodal_gen/test/single_test_file/test_generate_zimage_turbo_cli.py

@@ -0,0 +1,45 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import os
+import shlex
+import unittest
+
+from PIL import Image
+
+from sglang.multimodal_gen.configs.sample.sampling_params import DataType
+from sglang.multimodal_gen.test.single_test_file.cli_generate_common import (
+    CLIBase,
+    check_image_size,
+)
+from sglang.multimodal_gen.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+
+
+class TestZImageTurboCLI(CLIBase):
+    model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+    extra_args = ("--num-inference-steps=4",)
+    data_type: DataType = DataType.IMAGE
+    width = 512
+    height = 512
+
+    def test_output_file_path_alias(self):
+        output_file_path = os.path.join(self.output_path, "zimage_turbo_alias.png")
+        _, status = self._run_command(
+            "ignored_output_file_name",
+            self.model_path,
+            args=f"--output-file-path {shlex.quote(output_file_path)}",
+        )
+
+        self.assertEqual(status, "Success", "output-file-path command failed")
+        self.assertTrue(
+            os.path.exists(output_file_path),
+            f"Output file not exist for {output_file_path}",
+        )
+        with Image.open(output_file_path) as image:
+            check_image_size(self, image, self.width, self.height)
+
+
+del CLIBase
+
+
+if __name__ == "__main__":
+    unittest.main()

python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_disk.py

@@ -26,15 +26,15 @@ encoder are still the same.
 To strictly verify the correctness of the refit API, we compare the checksum in
 SHA-256 on the disk and the server.
 
-NOTE and TODO: In the refit a specific module test, we randomly select one module
+NOTE and TODO: In the refit a specific module test, we update the transformer
-from the transformer and vae to refit the server and keep other modules the same.
+module and keep other modules the same. As described above, the vae's weights
-As described above, the vae's weights are perturbed. If we select the vae to be the
+are perturbed. If we add vae as a target module in the future, ideally speaking,
-target module, ideally speaking, we should assert that the refitted vae's checksum
+we should assert that the refitted vae's checksum is the same as directly
-is the same as directly computed from the perturbed vae weights in the disk. However,
+computed from the perturbed vae weights in the disk. However, since there is
-since the there is complex weight-name remapping and QKV merge during model loading,
+complex weight-name remapping and QKV merge during model loading, it is not easy
-it is not easy to compare the server-disk checksum for vae and text encoder directly.
+to compare the server-disk checksum for vae and text encoder directly. Therefore,
-Therefore, if the target module is vae, we only verify that the refitted vae's checksum
+if the target module is vae, we only verify that the refitted vae's checksum is
-is different from the base model's vae's checksum.
+different from the base model's vae's checksum.
 
 It should be good issue to solve for the community to adds comparison the server-disk
 checksum for vae and text encoder in this test.
@@ -43,12 +43,12 @@ checksum for vae and text encoder in this test.
 
 Test organization:
 
-7 test cases in 2 classes;
+5 test cases in 2 classes;
 two model pairs are tested locally, one in CI.
 
 =============================================================================
 
-Class 1: TestUpdateWeightsFromDisk                  (6 tests) — API contract, checksum & rollback
+Class 1: TestUpdateWeightsFromDisk                  (4 tests) — API contract, checksum & rollback
 Class 2: TestUpdateWeightsFromDiskWithOffload       (1 test) — Offload-aware update + checksum
 
 -----------------------------------------------------------------------------
@@ -72,32 +72,16 @@ base model first so behavior is order-independent and updates are real
 
   • test_update_weights_specific_modules
 
-    base -> perturbed with flush_cache=False.  Randomly selects one module
+    base -> perturbed with flush_cache=False. Updates only transformer as
-    from _DIFFERING_MODULES (transformer and vae) as target_modules, updates
+    target_modules. Verifies that:
-    only that module. Verifies that:
     (1) targeted module's in-memory checksum changed;
     (2) non-targeted modules' in-memory checksums are unchanged.
 
-  • test_update_weights_nonexistent_model
+  • test_update_weights_rejects_invalid_requests
 
-    model_path set to a non-existent path; must fail (400, success=False).
+    Invalid requests must fail (400, success=False) without mutating live
-
+    weights. Covers nonexistent model, missing model_path, and nonexistent
-    Ensure server is healthy after failed update and server's transformer
+    target module.
-    checksums equal base model's transformer disk checksum.
-
-  • test_update_weights_missing_model_path
-
-    Request body empty (no model_path); must fail (400, success=False).
-
-    Ensure server is healthy after failed update and server's transformer
-    checksums equal base model's transformer disk checksum.
-
-  • test_update_weights_nonexistent_module
-
-    target_modules=["nonexistent_module"]; must fail (400, success=False).
-
-    Verify server is healthy after failed update and server's checksums
-    equal base model's transformer disk checksum.
 
   • test_corrupted_weights_rollback
 
@@ -137,7 +121,6 @@ from __future__ import annotations
 
 import functools
 import os
-import random
 import shutil
 import sys
 import tempfile
@@ -194,6 +177,9 @@ _ALL_MODEL_PAIRS: list[tuple[str, str]] = [
 
 
 _CI_MODEL_PAIR_ENV = "SGLANG_MMGEN_UPDATE_WEIGHTS_PAIR"
+_CI_MODEL_PAIR_IDS = ("FLUX.2-klein-base-4B",)
+_UPDATE_TIMEOUT_SECONDS = 900
+_CHECKSUM_TIMEOUT_SECONDS = 600
 
 
 def _resolve_active_model_pairs() -> list[tuple[str, str]]:
@@ -203,7 +189,7 @@ def _resolve_active_model_pairs() -> list[tuple[str, str]]:
     pair_by_id = {pair[0].split("/")[-1]: pair for pair in _ALL_MODEL_PAIRS}
     selected_pair_id = os.environ.get(_CI_MODEL_PAIR_ENV)
     if selected_pair_id is None:
-        return [random.choice(_ALL_MODEL_PAIRS)]
+        return [pair_by_id[_CI_MODEL_PAIR_IDS[0]]]
 
     selected_pair = pair_by_id.get(selected_pair_id)
     if selected_pair is None:
@@ -296,7 +282,6 @@ def _truncate_safetensor(src_file: str, dst_file: str) -> None:
 
 
 def _perturb_safetensor(src_file: str, dst_file: str) -> None:
-
     tensors = load_file(src_file)
     perturbed = {
         k: (t + 0.01 if t.is_floating_point() else t) for k, t in tensors.items()
@@ -312,7 +297,7 @@ class _UpdateWeightsApiMixin:
         model_path: str,
         flush_cache: bool = True,
         target_modules: list[str] | None = None,
-        timeout: int = 300,
+        timeout: int = _UPDATE_TIMEOUT_SECONDS,
     ) -> tuple[dict, int]:
         payload = {"model_path": model_path, "flush_cache": flush_cache}
         if target_modules is not None:
@@ -328,7 +313,7 @@ class _UpdateWeightsApiMixin:
         self,
         base_url: str,
         module_names: list[str] | None = None,
-        timeout: int = 300,
+        timeout: int = _CHECKSUM_TIMEOUT_SECONDS,
     ) -> dict:
         payload = {}
         if module_names is not None:
@@ -361,7 +346,6 @@ class _UpdateWeightsApiMixin:
 
 
 class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin):
-
     @pytest.fixture(
         scope="class",
         params=_ACTIVE_MODEL_PAIRS,
@@ -460,7 +444,7 @@ class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin):
             base_url, module_names=_DIFFERING_MODULES
         )
 
-        target_modules = [random.choice(_DIFFERING_MODULES)]
+        target_modules = [_TRANSFORMER_MODULE]
         result, status_code = self._update_weights(
             base_url,
             perturbed_model_dir,
@@ -492,12 +476,13 @@ class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin):
                 f"  after:  {cs}"
             )
 
-    def test_update_weights_nonexistent_model(self, diffusion_server_no_offload):
+    def test_update_weights_rejects_invalid_requests(self, diffusion_server_no_offload):
-        """Nonexistent model path must fail (400). Server healthy, checksums == base disk."""
+        ctx, _, perturbed_model_dir, _ = diffusion_server_no_offload
-        ctx, default_model, _, _ = diffusion_server_no_offload
         base_url = f"http://localhost:{ctx.port}"
 
-        self._update_weights(base_url, default_model)
+        before_checksums = self._get_weights_checksum(
+            base_url, module_names=_DIFFERING_MODULES
+        )
 
         result, status_code = self._update_weights(
             base_url,
@@ -508,14 +493,6 @@ class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin):
 
         assert status_code == 400, f"Expected 400, got {status_code}"
         assert not result.get("success", True), "Should fail for nonexistent model"
-        self._assert_server_matches_model(base_url, default_model)
-
-    def test_update_weights_missing_model_path(self, diffusion_server_no_offload):
-        """Request without model_path must fail (400). Server healthy, checksums == base disk."""
-        ctx, default_model, _, _ = diffusion_server_no_offload
-        base_url = f"http://localhost:{ctx.port}"
-
-        self._update_weights(base_url, default_model)
 
         response = requests.post(
             f"{base_url}/update_weights_from_disk",
@@ -526,14 +503,6 @@ class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin):
         assert response.status_code == 400, f"Expected 400, got {response.status_code}"
         result = response.json()
         assert not result.get("success", True), "Should fail when model_path is missing"
-        self._assert_server_matches_model(base_url, default_model)
-
-    def test_update_weights_nonexistent_module(self, diffusion_server_no_offload):
-        """Nonexistent module must fail (400). Server healthy, checksums == base disk."""
-        ctx, default_model, perturbed_model_dir, _ = diffusion_server_no_offload
-        base_url = f"http://localhost:{ctx.port}"
-
-        self._update_weights(base_url, default_model)
 
         result, status_code = self._update_weights(
             base_url,
@@ -546,7 +515,11 @@ class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin):
         assert status_code == 400, f"Expected 400, got {status_code}"
         assert not result.get("success", True), "Should fail for nonexistent module"
         assert "not found in pipeline" in result.get("message", "")
-        self._assert_server_matches_model(base_url, default_model)
+
+        after_checksums = self._get_weights_checksum(
+            base_url, module_names=_DIFFERING_MODULES
+        )
+        assert after_checksums == before_checksums
 
     def test_corrupted_weights_rollback(self, diffusion_server_no_offload):
         ctx, default_model, perturbed_model_dir, corrupted_vae_model_dir = (

python/sglang/multimodal_gen/test/test_files/launch_flux.json

@@ -1,11 +0,0 @@
-{
-    "model_path": "black-forest-labs/FLUX.1-dev",
-    "prompt": "A beautiful woman in a red dress walking down a street",
-    "text_encoder_cpu_offload": true,
-    "pin_cpu_memory": true,
-    "save_output": true,
-    "width": 720,
-    "height": 720,
-    "output_path": "outputs",
-    "output_file_name": "FLUX.1-dev, single gpu"
-}

python/sglang/multimodal_gen/test/test_files/launch_wan.json

@@ -1,11 +0,0 @@
-{
-    "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-    "prompt": "A beautiful woman in a red dress walking down a street",
-    "text_encoder_cpu_offload": true,
-    "pin_cpu_memory": true,
-    "save_output": true,
-    "width": 720,
-    "height": 720,
-    "output_path": "outputs",
-    "output_file_name": "Wan2.1-T2V-1.3B-Diffusers, single gpu"
-}

python/sglang/multimodal_gen/test/test_utils.py

@@ -999,7 +999,7 @@ def _is_ascend_consistency_case(case_id: str) -> bool:
     return "npu" in case_id
 
 
-def _remote_file_exists(url: str) -> bool:
+def _remote_file_exists(url: str) -> bool | None:
     for _ in range(3):
         for method in ("head", "get"):
             try:
@@ -1016,6 +1016,8 @@ def _remote_file_exists(url: str) -> bool:
                 try:
                     if resp.status_code in (200, 206):
                         return True
+                    if resp.status_code == 404:
+                        return False
                     if (
                         resp.status_code not in (403, 405, 429)
                         and resp.status_code < 500
@@ -1025,7 +1027,7 @@ def _remote_file_exists(url: str) -> bool:
                     resp.close()
             except requests.RequestException:
                 pass
-    return False
+    return None
 
 
 def _load_remote_gt_image(url: str) -> np.ndarray:
@@ -1068,12 +1070,19 @@ def _find_remote_consistency_gt_files(
             base_url, case_id, num_gpus, is_video, output_format
         )
         if is_video:
-            if all(_remote_file_exists(url) for _, url in candidates):
+            exists = [_remote_file_exists(url) for _, url in candidates]
+            if all(status is not False for status in exists):
                 return candidates
         else:
+            uncertain_candidate = None
             for filename, url in candidates:
-                if _remote_file_exists(url):
+                exists = _remote_file_exists(url)
+                if exists is True:
                     return [(filename, url)]
+                if exists is None and uncertain_candidate is None:
+                    uncertain_candidate = (filename, url)
+            if uncertain_candidate is not None:
+                return [uncertain_candidate]
     return []
 
 

python/sglang/multimodal_gen/test/unit/test_cli_generate_common.py

@@ -0,0 +1,69 @@
+import os
+import unittest
+from unittest.mock import patch
+
+from PIL import Image
+
+from sglang.multimodal_gen.test.single_test_file import cli_generate_common
+
+
+class TestCLIBaseHelpers(unittest.TestCase):
+    def _make_case(self):
+        class _ImageDataType:
+            @staticmethod
+            def get_default_extension():
+                return "png"
+
+        class _ImageCLI(cli_generate_common.CLIBase):
+            model_path = "dummy/model"
+            extra_args = ("--dummy-extra",)
+            data_type = _ImageDataType()
+            width = 32
+            height = 16
+
+        case = _ImageCLI(methodName="test_single_gpu")
+        case.setUp()
+        self.addCleanup(case.tearDown)
+        return case
+
+    def test_run_command_builds_generate_command(self):
+        case = self._make_case()
+        captured = []
+
+        def fake_run_command(command):
+            captured.append(command)
+            return True
+
+        with patch(
+            "sglang.multimodal_gen.test.single_test_file.cli_generate_common.run_command",
+            side_effect=fake_run_command,
+        ):
+            name, status = case._run_command(
+                "sample",
+                model_path="dummy/model",
+                args='--negative-prompt "low quality"',
+            )
+
+        self.assertEqual(name, "sample")
+        self.assertEqual(status, "Success")
+        self.assertEqual(len(captured), 1)
+        self.assertEqual(captured[0][0:2], ["sglang", "generate"])
+        self.assertIn("--model-path=dummy/model", captured[0])
+        self.assertIn("--negative-prompt", captured[0])
+        self.assertIn("low quality", captured[0])
+        self.assertEqual(
+            captured[0][-3:], ["--output-file-name", "sample", "--dummy-extra"]
+        )
+
+    def test_verify_accepts_expected_image_output(self):
+        case = self._make_case()
+        output_path = os.path.join(case.output_path, "sample.png")
+        Image.new("RGB", (case.width, case.height), color=(255, 0, 0)).save(output_path)
+
+        case.verify("Success", "sample")
+
+    def test_verify_fails_when_output_missing(self):
+        case = self._make_case()
+
+        with self.assertRaises(AssertionError):
+            case.verify("Success", "missing")

python/sglang/multimodal_gen/test/unit/test_consistency_metrics.py

@@ -26,6 +26,34 @@ def test_consistency_gt_urls_are_pinned_to_ci_data_revision():
     assert revision_path in test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE
 
 
+def test_remote_file_exists_returns_false_for_definitive_404(monkeypatch):
+    class Response:
+        status_code = 404
+
+        def close(self):
+            pass
+
+    monkeypatch.setattr(test_utils.requests, "head", lambda *args, **kwargs: Response())
+
+    assert test_utils._remote_file_exists("https://example.com/missing.png") is False
+
+
+def test_remote_video_gt_candidates_survive_inconclusive_probe(monkeypatch):
+    monkeypatch.setattr(test_utils, "_remote_file_exists", lambda url: None)
+
+    files = test_utils._find_remote_consistency_gt_files(
+        "unit_video",
+        1,
+        is_video=True,
+    )
+
+    assert [filename for filename, _ in files] == [
+        "unit_video_1gpu_frame_0.png",
+        "unit_video_1gpu_frame_mid.png",
+        "unit_video_1gpu_frame_last.png",
+    ]
+
+
 def test_pixel_metrics_identical_image():
     image = _solid_image(128)