[Test] Unify L2/L3 test layout, Buildkite steps, and test helpers (#2556)

Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com>
Co-authored-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>

tests/diffusion/distributed/test_expert_parallel.py

@@ -0,0 +1,250 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+System test for Expert Parallel (EP).
+
+Tests verify that EP inference produces correct outputs compared to baseline.
+"""
+
+import gc
+import os
+import time
+from typing import NamedTuple
+
+import numpy as np
+import pytest
+import torch
+import torch.distributed as dist
+from PIL import Image
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+# Test configuration
+MODELS = ["tencent/HunyuanImage-3.0"]
+PROMPT = "A brown and white dog is running on the grass"
+DEFAULT_HEIGHT = 1024
+DEFAULT_WIDTH = 1024
+DEFAULT_SEED = 1234
+DEFAULT_STEPS = 50
+DEFAULT_GUIDANCE_SCALE = 0.0
+COS_SIM_MEAN_THRESHOLD = 1e-2
+COS_SIM_MAX_THRESHOLD = 1e-2
+MSE_THRESHOLD = 1e-2
+DIFF_MEAN_THRESHOLD = 5e-2
+DIFF_MAX_THRESHOLD = 1
+TOTAL_CARD_NUM = 8
+
+
+class InferenceResult(NamedTuple):
+    """Result of an inference run."""
+
+    images: list[Image.Image]
+    elapsed_ms: float
+
+
+def _cleanup_distributed():
+    """Clean up distributed environment and GPU resources."""
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+    for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
+        os.environ.pop(key, None)
+
+    gc.collect()
+    if current_omni_platform.is_available():
+        current_omni_platform.empty_cache()
+        current_omni_platform.synchronize()
+
+    time.sleep(5)
+
+
+def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float, float, float, float]:
+    """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1]."""
+    ta = torch.from_numpy(np.asarray(a.convert("RGB"), dtype=np.float32) / 255.0)
+    tb = torch.from_numpy(np.asarray(b.convert("RGB"), dtype=np.float32) / 255.0)
+    assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
+    abs_diff = torch.abs(ta - tb)
+
+    vec_a = ta.reshape(-1)
+    vec_b = tb.reshape(-1)
+    cos_sim = torch.nn.functional.cosine_similarity(vec_a, vec_b, dim=0, eps=1e-8)
+    mse = torch.nn.functional.mse_loss(ta, tb)
+    return abs_diff.mean().item(), abs_diff.max().item(), cos_sim.mean().item(), cos_sim.max().item(), mse.item()
+
+
+def _run_inference(
+    model_name: str,
+    tensor_parallel_size: int,
+    enable_expert_parallel: bool,
+    guidance_scale: float = DEFAULT_GUIDANCE_SCALE,
+    height: int = DEFAULT_HEIGHT,
+    width: int = DEFAULT_WIDTH,
+    seed: int = DEFAULT_SEED,
+    warmup: bool = True,
+) -> InferenceResult:
+    """Run inference with specified configuration.
+
+    Args:
+        warmup: If True, run one warmup iteration before the timed run.
+    """
+    parallel_config = DiffusionParallelConfig(
+        tensor_parallel_size=tensor_parallel_size,
+        enable_expert_parallel=enable_expert_parallel,
+    )
+    try:
+        with OmniRunner(model_name, parallel_config=parallel_config) as runner:
+            omni = runner.omni
+            # Warmup run (not timed)
+            if warmup:
+                _ = omni.generate(
+                    PROMPT,
+                    OmniDiffusionSamplingParams(
+                        height=height,
+                        width=width,
+                        num_inference_steps=DEFAULT_STEPS,
+                        guidance_scale=guidance_scale,
+                        generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
+                        num_outputs_per_prompt=1,
+                    ),
+                )
+
+            # Timed run
+            start = time.time()
+            outputs = omni.generate(
+                PROMPT,
+                OmniDiffusionSamplingParams(
+                    height=height,
+                    width=width,
+                    num_inference_steps=DEFAULT_STEPS,
+                    guidance_scale=guidance_scale,
+                    generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
+                    num_outputs_per_prompt=1,
+                ),
+            )
+            elapsed_ms = (time.time() - start) * 1000
+
+            return InferenceResult(
+                images=outputs[0].images,
+                elapsed_ms=elapsed_ms,
+            )
+    finally:
+        _cleanup_distributed()
+
+
+# EP configurations: (tensor_parallel_size, enable_ep, height, width, is_perf_test)
+EP_TEST_CONFIG = [
+    (TOTAL_CARD_NUM, True, DEFAULT_HEIGHT, DEFAULT_WIDTH, True),
+]
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": TOTAL_CARD_NUM})
+@pytest.mark.parametrize("model_name", MODELS)
+def test_ep(model_name):
+    device_count = current_omni_platform.get_device_count()
+    if device_count < TOTAL_CARD_NUM:
+        pytest.skip(f"Requires at least {TOTAL_CARD_NUM} GPUs, got {device_count}")
+
+    baseline_cache: dict[str, InferenceResult] = {}
+    results: list[dict] = []
+    # baseline_key = PROMPT + "_" + str(DEFAULT_HEIGHT) + "*" + str(DEFAULT_WIDTH) + "_" + str(DEFAULT_SEED) + "_" + str(DEFAULT_STEPS)
+    # baseline_cache[baseline_key] = _run_inference(
+    #     model_name=model_name,
+    #     tensor_parallel_size=TOTAL_CARD_NUM,
+    #     enable_expert_parallel=False,
+    #     warmup=True,
+    # )
+
+    print("\n" + "=" * 90)
+    print(f"Sequence Parallel Test - Model: {model_name}")
+    print(f"Available GPUs: {device_count}")
+    print("=" * 90)
+
+    for tensor_parallel_size, enable_ep, height, width, is_perf_test in EP_TEST_CONFIG:
+        cache_key = PROMPT + "_" + str(height) + "*" + str(width) + "_" + str(DEFAULT_SEED) + "_" + str(DEFAULT_STEPS)
+        if cache_key not in baseline_cache:
+            print(
+                f"\n--- Running baseline{{ prompt: {PROMPT}, resolution: {height}x{width}, seed: {DEFAULT_SEED}, steps: {DEFAULT_STEPS} }}---"
+            )
+            baseline = _run_inference(
+                model_name=model_name,
+                tensor_parallel_size=TOTAL_CARD_NUM,
+                enable_expert_parallel=False,
+                warmup=True,
+            )
+            assert len(baseline.images) == 1
+            baseline_cache[cache_key] = baseline
+            print(f"[baseline] {height}x{width}: {baseline.elapsed_ms:.0f}ms")
+        else:
+            baseline = baseline_cache[cache_key]
+
+        print(f"\n--- Running test: {{ enable_ep: {enable_ep} }} ---")
+        infer_result = _run_inference(
+            model_name=model_name,
+            tensor_parallel_size=tensor_parallel_size,
+            enable_expert_parallel=enable_ep,
+            warmup=is_perf_test,
+        )
+        assert len(infer_result.images) == 1
+        mean_diff, max_diff, cos_sim_mean, cos_sim_max, mse = _diff_metrics(baseline.images[0], infer_result.images[0])
+        print(f"{cos_sim_mean}, {cos_sim_max}, {mse}")
+
+        # Build result entry
+        result = {
+            "prompt": PROMPT,
+            "enable_ep": enable_ep,
+            "height": height,
+            "width": width,
+            "baseline_ms": baseline.elapsed_ms,
+            "ep_ms": infer_result.elapsed_ms,
+            "mean_diff": mean_diff,
+            "max_diff": max_diff,
+            "cos_sim_mean": cos_sim_mean,
+            "cos_sim_max": cos_sim_max,
+            "mse": mse,
+            "is_perf_test": is_perf_test,
+        }
+        results.append(result)
+
+        # Output based on test type
+        if is_perf_test:
+            speedup = baseline.elapsed_ms / infer_result.elapsed_ms if infer_result.elapsed_ms > 0 else 0
+            result["speedup"] = speedup
+            print(
+                f"[enable_ep: {enable_ep}] {tensor_parallel_size} GPUs | "
+                f"baseline: {baseline.elapsed_ms:.0f}ms, ep: {infer_result.elapsed_ms:.0f}ms, "
+                f"speedup: {speedup:.2f}x"
+            )
+        else:
+            print(
+                f"[enable_ep: {enable_ep}] {tensor_parallel_size} GPUs | ep: {infer_result.elapsed_ms:.0f}ms (correctness only)"
+            )
+
+        print(
+            f"[enable_ep: {enable_ep}] diff: [mean={mean_diff:.6e}, max={max_diff:.6e}], cos_sim: [mean={cos_sim_mean:.6e}, max={cos_sim_max:.6e}], mse: {mse:.6e}"
+        )
+        assert mean_diff <= DIFF_MEAN_THRESHOLD and max_diff <= DIFF_MAX_THRESHOLD, (
+            f"[enable_ep: {enable_ep}] output differs from baseline: mean={mean_diff:.6e}, max={max_diff:.6e}"
+        )
+
+    print("\n" + "=" * 90)
+    print("SUMMARY")
+    print("=" * 90)
+    print(f"{'Mode':<15} {'GPUs':<6} {'Size':<10} {'Baseline':<12} {'EP':<12} {'Speedup':<10} {'Status'}")
+    print("-" * 90)
+    for r in results:
+        speedup_str = f"{r['speedup']:.2f}x" if r.get("speedup") else "N/A"
+        baseline_str = f"{r['baseline_ms']:.0f}ms" if r["is_perf_test"] else "N/A"
+        status = "PASS" if r["mean_diff"] <= DIFF_MEAN_THRESHOLD else "FAIL"
+        print(
+            f"{r['prompt']:<15.10} {r['enable_ep']:<6} {r['height']}x{r['width']:<5} "
+            f"{baseline_str:<12} {r['ep_ms']:.0f}ms{'':<7} {speedup_str:<10} {status}"
+        )
+    print("=" * 90)

tests/diffusion/distributed/test_sequence_parallel.py

@@ -0,0 +1,414 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+System test for Sequence Parallel (SP) backends: Ulysses and Ring attention.
+
+Tests verify that SP inference produces correct outputs compared to baseline.
+"""
+
+import gc
+import os
+import sys
+import time
+from pathlib import Path
+from typing import NamedTuple
+
+import numpy as np
+import pytest
+import torch
+import torch.distributed as dist
+from PIL import Image
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+# ruff: noqa: E402
+REPO_ROOT = Path(__file__).resolve().parents[3]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+
+# Test configuration
+MODELS = ["riverclouds/qwen_image_random"]
+PROMPT = "a photo of a cat sitting on a laptop keyboard"
+DEFAULT_HEIGHT = 256
+DEFAULT_WIDTH = 256
+DEFAULT_SEED = 42
+DEFAULT_STEPS = 4
+DIFF_MEAN_THRESHOLD = 2e-2
+DIFF_MAX_THRESHOLD = 2e-1
+
+
+class InferenceResult(NamedTuple):
+    """Result of an inference run."""
+
+    images: list[Image.Image]
+    elapsed_ms: float
+
+
+def _cleanup_distributed():
+    """Clean up distributed environment and GPU resources."""
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+    for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
+        os.environ.pop(key, None)
+
+    gc.collect()
+    if current_omni_platform.is_available():
+        current_omni_platform.empty_cache()
+        current_omni_platform.synchronize()
+
+    time.sleep(5)
+
+
+def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
+    """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1]."""
+    ta = torch.from_numpy(np.asarray(a.convert("RGB"), dtype=np.float32) / 255.0)
+    tb = torch.from_numpy(np.asarray(b.convert("RGB"), dtype=np.float32) / 255.0)
+    assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
+    abs_diff = torch.abs(ta - tb)
+    return abs_diff.mean().item(), abs_diff.max().item()
+
+
+def _run_inference(
+    model_name: str,
+    dtype: torch.dtype,
+    attn_backend: str,
+    ulysses_degree: int = 1,
+    ring_degree: int = 1,
+    height: int = DEFAULT_HEIGHT,
+    width: int = DEFAULT_WIDTH,
+    seed: int = DEFAULT_SEED,
+    warmup: bool = True,
+) -> InferenceResult:
+    """Run inference with specified configuration.
+
+    Args:
+        warmup: If True, run one warmup iteration before the timed run.
+    """
+    parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
+    try:
+        with OmniRunner(
+            model_name,
+            parallel_config=parallel_config,
+            dtype=dtype,
+            attention_backend=attn_backend,
+        ) as runner:
+            omni = runner.omni
+            # Warmup run (not timed)
+            if warmup:
+                _ = omni.generate(
+                    PROMPT,
+                    OmniDiffusionSamplingParams(
+                        height=height,
+                        width=width,
+                        num_inference_steps=DEFAULT_STEPS,
+                        guidance_scale=0.0,
+                        generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed + 1000),
+                        num_outputs_per_prompt=1,
+                    ),
+                )
+
+            # Timed run
+            start = time.time()
+            outputs = omni.generate(
+                PROMPT,
+                OmniDiffusionSamplingParams(
+                    height=height,
+                    width=width,
+                    num_inference_steps=DEFAULT_STEPS,
+                    guidance_scale=0.0,
+                    generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
+                    num_outputs_per_prompt=1,
+                ),
+            )
+            elapsed_ms = (time.time() - start) * 1000
+
+            return InferenceResult(
+                images=outputs[0].request_output.images,
+                elapsed_ms=elapsed_ms,
+            )
+    finally:
+        _cleanup_distributed()
+
+
+# =============================================================================
+# Correctness & Performance Tests
+# =============================================================================
+
+# SP configurations: (ulysses_degree, ring_degree, height, width, warmup, is_perf_test)
+# - warmup: whether to run warmup for this SP config
+# - is_perf_test: whether this is a performance test (show speedup metrics)
+SP_CONFIGS_L2 = [
+    # Hybrid - correctness only
+    (2, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, False, False),
+]
+
+SP_CONFIGS_L3 = [
+    # Ulysses-2 - performance test
+    (2, 1, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True),
+    (1, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True),  # Ring-2 - performance test
+    # Hybrid - correctness only
+    (2, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, False, False),
+    (4, 1, 272, 272, False, False),  # Ulysses-4 - shape and correctness
+]
+
+
+def _get_sp_mode(ulysses_degree: int, ring_degree: int) -> str:
+    """Get SP mode name for logging."""
+    if ulysses_degree > 1 and ring_degree == 1:
+        return f"ulysses-{ulysses_degree}"
+    elif ring_degree > 1 and ulysses_degree == 1:
+        return f"ring-{ring_degree}"
+    else:
+        return f"hybrid-{ulysses_degree}x{ring_degree}"
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 2, "rocm": 2})
+@pytest.mark.parametrize("model_name", MODELS)
+def test_sp_correctness(model_name: str):
+    """Test that SP inference produces correct outputs and measure performance.
+
+    Runs baseline once per unique (height, width), then tests all SP configs.
+
+    Note: Run with `pytest -v -s` to see detailed output.
+    """
+    device_count = current_omni_platform.get_device_count()
+
+    # Cache baseline results by (height, width)
+    # Key: (height, width), Value: (result, warmup_used)
+    baseline_cache: dict[tuple[int, int], InferenceResult] = {}
+
+    # Collect results for summary
+    results: list[dict] = []
+
+    print("\n" + "=" * 70)
+    print(f"Sequence Parallel Test - Model: {model_name}")
+    print(f"Available GPUs: {device_count}")
+    print("=" * 70)
+
+    for ulysses_degree, ring_degree, height, width, sp_warmup, is_perf_test in SP_CONFIGS_L2:
+        sp_size = ulysses_degree * ring_degree
+        sp_mode = _get_sp_mode(ulysses_degree, ring_degree)
+
+        if device_count < sp_size:
+            print(f"\n[{sp_mode}] SKIPPED (requires {sp_size} GPUs)")
+            continue
+
+        # Determine baseline warmup: only for default size (performance tests)
+        cache_key = (height, width)
+        baseline_warmup = height == DEFAULT_HEIGHT and width == DEFAULT_WIDTH
+
+        # Get or compute baseline for this (height, width)
+        if cache_key not in baseline_cache:
+            print(f"\n--- Running baseline {height}x{width} (warmup={baseline_warmup}) ---")
+            baseline = _run_inference(
+                model_name,
+                torch.bfloat16,
+                "sdpa",
+                height=height,
+                width=width,
+                warmup=baseline_warmup,
+            )
+            assert len(baseline.images) == 1
+            baseline_cache[cache_key] = baseline
+            print(f"[baseline] {height}x{width}: {baseline.elapsed_ms:.0f}ms")
+        else:
+            baseline = baseline_cache[cache_key]
+
+        # Run SP
+        print(f"\n--- Running {sp_mode} (warmup={sp_warmup}) ---")
+        sp_result = _run_inference(
+            model_name,
+            torch.bfloat16,
+            "sdpa",
+            ulysses_degree=ulysses_degree,
+            ring_degree=ring_degree,
+            height=height,
+            width=width,
+            warmup=sp_warmup,
+        )
+        assert len(sp_result.images) == 1
+
+        # Compare outputs (correctness)
+        mean_diff, max_diff = _diff_metrics(baseline.images[0], sp_result.images[0])
+
+        # Build result entry
+        result = {
+            "mode": sp_mode,
+            "sp_size": sp_size,
+            "height": height,
+            "width": width,
+            "baseline_ms": baseline.elapsed_ms,
+            "sp_ms": sp_result.elapsed_ms,
+            "mean_diff": mean_diff,
+            "max_diff": max_diff,
+            "is_perf_test": is_perf_test,
+        }
+        results.append(result)
+
+        # Output based on test type
+        if is_perf_test:
+            speedup = baseline.elapsed_ms / sp_result.elapsed_ms if sp_result.elapsed_ms > 0 else 0
+            result["speedup"] = speedup
+            print(
+                f"[{sp_mode}] {sp_size} GPUs | "
+                f"baseline: {baseline.elapsed_ms:.0f}ms, sp: {sp_result.elapsed_ms:.0f}ms, "
+                f"speedup: {speedup:.2f}x"
+            )
+        else:
+            print(f"[{sp_mode}] {sp_size} GPUs | sp: {sp_result.elapsed_ms:.0f}ms (correctness only)")
+
+        print(f"[{sp_mode}] diff: mean={mean_diff:.6e}, max={max_diff:.6e}")
+
+        # Assert correctness
+        assert mean_diff <= DIFF_MEAN_THRESHOLD and max_diff <= DIFF_MAX_THRESHOLD, (
+            f"[{sp_mode}] SP output differs from baseline: mean={mean_diff:.6e}, max={max_diff:.6e}"
+        )
+
+    # Summary
+    print("\n" + "=" * 70)
+    print("SUMMARY")
+    print("=" * 70)
+    print(f"{'Mode':<15} {'GPUs':<6} {'Size':<10} {'Baseline':<12} {'SP':<12} {'Speedup':<10} {'Status'}")
+    print("-" * 70)
+    for r in results:
+        speedup_str = f"{r['speedup']:.2f}x" if r.get("speedup") else "N/A"
+        baseline_str = f"{r['baseline_ms']:.0f}ms" if r["is_perf_test"] else "N/A"
+        status = "PASS" if r["mean_diff"] <= DIFF_MEAN_THRESHOLD else "FAIL"
+        print(
+            f"{r['mode']:<15} {r['sp_size']:<6} {r['height']}x{r['width']:<5} "
+            f"{baseline_str:<12} {r['sp_ms']:.0f}ms{'':<7} {speedup_str:<10} {status}"
+        )
+    print("=" * 70)
+
+
+# TODO: After PR#1272 is merged, add markers
+# @pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
+@pytest.mark.parametrize("model_name", MODELS)
+def test_sp_correctness_advanced(model_name: str):
+    """Test that SP inference produces correct outputs and measure performance.
+
+    Runs baseline once per unique (height, width), then tests all SP configs.
+
+    Note: Run with `pytest -v -s` to see detailed output.
+    """
+    device_count = current_omni_platform.get_device_count()
+
+    # Cache baseline results by (height, width)
+    # Key: (height, width), Value: (result, warmup_used)
+    baseline_cache: dict[tuple[int, int], InferenceResult] = {}
+
+    # Collect results for summary
+    results: list[dict] = []
+
+    print("\n" + "=" * 70)
+    print(f"Sequence Parallel Test - Model: {model_name}")
+    print(f"Available GPUs: {device_count}")
+    print("=" * 70)
+
+    for ulysses_degree, ring_degree, height, width, sp_warmup, is_perf_test in SP_CONFIGS_L3:
+        sp_size = ulysses_degree * ring_degree
+        sp_mode = _get_sp_mode(ulysses_degree, ring_degree)
+
+        if device_count < sp_size:
+            print(f"\n[{sp_mode}] SKIPPED (requires {sp_size} GPUs)")
+            continue
+
+        # Determine baseline warmup: only for default size (performance tests)
+        cache_key = (height, width)
+        baseline_warmup = height == DEFAULT_HEIGHT and width == DEFAULT_WIDTH
+
+        # Get or compute baseline for this (height, width)
+        if cache_key not in baseline_cache:
+            print(f"\n--- Running baseline {height}x{width} (warmup={baseline_warmup}) ---")
+            baseline = _run_inference(
+                model_name,
+                torch.bfloat16,
+                "sdpa",
+                height=height,
+                width=width,
+                warmup=baseline_warmup,
+            )
+            assert len(baseline.images) == 1
+            baseline_cache[cache_key] = baseline
+            print(f"[baseline] {height}x{width}: {baseline.elapsed_ms:.0f}ms")
+        else:
+            baseline = baseline_cache[cache_key]
+
+        # Run SP
+        print(f"\n--- Running {sp_mode} (warmup={sp_warmup}) ---")
+        sp_result = _run_inference(
+            model_name,
+            torch.bfloat16,
+            "sdpa",
+            ulysses_degree=ulysses_degree,
+            ring_degree=ring_degree,
+            height=height,
+            width=width,
+            warmup=sp_warmup,
+        )
+        assert len(sp_result.images) == 1
+
+        # Compare outputs (correctness)
+        mean_diff, max_diff = _diff_metrics(baseline.images[0], sp_result.images[0])
+
+        # Build result entry
+        result = {
+            "mode": sp_mode,
+            "sp_size": sp_size,
+            "height": height,
+            "width": width,
+            "baseline_ms": baseline.elapsed_ms,
+            "sp_ms": sp_result.elapsed_ms,
+            "mean_diff": mean_diff,
+            "max_diff": max_diff,
+            "is_perf_test": is_perf_test,
+        }
+        results.append(result)
+
+        # Output based on test type
+        if is_perf_test:
+            speedup = baseline.elapsed_ms / sp_result.elapsed_ms if sp_result.elapsed_ms > 0 else 0
+            result["speedup"] = speedup
+            print(
+                f"[{sp_mode}] {sp_size} GPUs | "
+                f"baseline: {baseline.elapsed_ms:.0f}ms, sp: {sp_result.elapsed_ms:.0f}ms, "
+                f"speedup: {speedup:.2f}x"
+            )
+        else:
+            print(f"[{sp_mode}] {sp_size} GPUs | sp: {sp_result.elapsed_ms:.0f}ms (correctness only)")
+
+        print(f"[{sp_mode}] diff: mean={mean_diff:.6e}, max={max_diff:.6e}")
+
+        # Assert correctness
+        assert mean_diff <= DIFF_MEAN_THRESHOLD and max_diff <= DIFF_MAX_THRESHOLD, (
+            f"[{sp_mode}] SP output differs from baseline: mean={mean_diff:.6e}, max={max_diff:.6e}"
+        )
+
+    # Summary
+    print("\n" + "=" * 70)
+    print("SUMMARY")
+    print("=" * 70)
+    print(f"{'Mode':<15} {'GPUs':<6} {'Size':<10} {'Baseline':<12} {'SP':<12} {'Speedup':<10} {'Status'}")
+    print("-" * 70)
+    for r in results:
+        speedup_str = f"{r['speedup']:.2f}x" if r.get("speedup") else "N/A"
+        baseline_str = f"{r['baseline_ms']:.0f}ms" if r["is_perf_test"] else "N/A"
+        status = "PASS" if r["mean_diff"] <= DIFF_MEAN_THRESHOLD else "FAIL"
+        print(
+            f"{r['mode']:<15} {r['sp_size']:<6} {r['height']}x{r['width']:<5} "
+            f"{baseline_str:<12} {r['sp_ms']:.0f}ms{'':<7} {speedup_str:<10} {status}"
+        )
+    print("=" * 70)

tests/diffusion/distributed/test_vae_decode_parallelism.py

@@ -0,0 +1,152 @@
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+# ruff: noqa: E402
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+import time
+
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.platforms import current_omni_platform
+
+# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+all_model_configs = [
+    {
+        "model_name": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+        "out_width": 1280,
+        "out_height": 720,
+        "out_frames": 16,
+    },
+    {
+        "model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+        "out_width": 1280,
+        "out_height": 720,
+        "out_frames": 16,
+    },
+    {
+        "model_name": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+        "out_width": 1280,
+        "out_height": 720,
+        "out_frames": 16,
+    },
+    {
+        "model_name": "Qwen/Qwen-Image",
+        "out_width": 3072,
+        "out_height": 3072,
+        "out_frames": 1,
+    },
+    # {
+    #     "model_name": "stepfun-ai/NextStep-1.1",
+    #     "out_width": 3072,
+    #     "out_height": 3072,
+    #     "out_frames": 1,
+    # },
+]
+
+
+def is_nextstep_model(model_name: str) -> bool:
+    """Check if the model is a NextStep model by reading its config."""
+    from vllm.transformers_utils.config import get_hf_file_to_dict
+
+    try:
+        cfg = get_hf_file_to_dict("config.json", model_name)
+        if cfg and cfg.get("model_type") == "nextstep":
+            return True
+    except Exception:
+        pass
+    return False
+
+
+def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1):
+    parallel_config = DiffusionParallelConfig(
+        tensor_parallel_size=tp,
+        vae_patch_parallel_size=vae_patch_parallel_size,
+    )
+
+    omni_kwargs = {
+        "vae_use_tiling": using_tile,
+        "parallel_config": parallel_config,
+    }
+    use_nextstep = is_nextstep_model(model_configs["model_name"])
+    if use_nextstep:
+        # NextStep-1.1 requires explicit pipeline class
+        omni_kwargs["model_class_name"] = "NextStep11Pipeline"
+
+    with OmniRunner(model_configs["model_name"], **omni_kwargs) as runner:
+        m = runner.omni
+        image = Image.new("RGB", (out_width, out_height), (0, 0, 0))
+        start = time.perf_counter()
+        outputs = m.generate(
+            {
+                "prompt": "A cat sitting on a table",
+                "multi_modal_data": {"image": image},
+            },
+            sampling_params_list=OmniDiffusionSamplingParams(
+                height=model_configs["out_height"],
+                width=model_configs["out_width"],
+                num_frames=model_configs["out_frames"],
+                num_inference_steps=2,
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
+            ),
+        )
+        end = time.perf_counter()
+        first_output = outputs[0]
+        req_out = first_output.request_output
+        frames = req_out.images[0]
+        if isinstance(frames, torch.Tensor):
+            frames = frames.detach().cpu().numpy()
+        elif isinstance(frames, Image.Image):
+            frames = np.array(frames)
+        # frames shape: (batch, num_frames, height, width, channels)
+        cost = (end - start) * 1000
+        return frames, cost
+        cleanup_dist_env_and_memory()
+
+
+@pytest.mark.parametrize("model_configs", all_model_configs)
+def test_vae_parallel_model(model_configs: str):
+    parallel_result, parallel_time = model_run(
+        model_configs=model_configs,
+        tp=2,
+        out_width=1280,
+        out_height=704,
+        out_frames=5,
+        using_tile=True,
+        vae_patch_parallel_size=2,
+    )
+    non_parallel_result, non_parallel_time = model_run(
+        model_configs=model_configs,
+        tp=2,
+        out_width=1280,
+        out_height=704,
+        out_frames=5,
+        using_tile=True,
+        vae_patch_parallel_size=1,
+    )
+    result_diff = np.abs(non_parallel_result - parallel_result)
+
+    mean_threshold = 3e-2
+    max_threshold = 3e-2  # they should be totally same
+    print(
+        f"{model_configs['model_name']} TP = 2 (tile + parallel vs tile): "
+        f"mean_abs_diff={result_diff.mean():.6e}, max_abs_diff={result_diff.max():.6e}; "
+        f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
+        f"parallel generate take time: {parallel_time:.2f} ms, non-parallel take time: {non_parallel_time:.2f} ms"
+    )
+    # assert non_parallel_time > parallel_time
+    assert result_diff.mean() < mean_threshold
+    assert result_diff.max() < max_threshold

tests/diffusion/distributed/test_zimage_parallelism.py

@@ -0,0 +1,275 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Z-Image end-to-end tests for diffusion parallelism.
+
+This file currently covers:
+- DiT tensor parallelism (TP=2) vs TP=1.
+- VAE patch parallelism (vae_patch_parallel_size=2) vs baseline on TP=2.
+
+Note: CUDA-only (>=2 GPUs). We use `enforce_eager=False` (default) to enable
+`torch.compile`.
+"""
+
+import os
+import time
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+PROMPT = "a photo of a cat sitting on a laptop keyboard"
+
+
+def _get_zimage_model() -> str:
+    # Allow overriding the model for local/offline environments.
+    # Can be either a HuggingFace repo id or a local path.
+    return os.environ.get("VLLM_TEST_ZIMAGE_MODEL", "Tongyi-MAI/Z-Image-Turbo")
+
+
+def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor:
+    """Convert PIL image to float32 RGB tensor in [0, 1] with shape [H, W, 3]."""
+    arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
+    return torch.from_numpy(arr)
+
+
+def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
+    """Return (mean_abs_diff, p99_abs_diff) over RGB pixels in [0, 1]."""
+    ta = _pil_to_float_rgb_tensor(a)
+    tb = _pil_to_float_rgb_tensor(b)
+    assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
+    abs_diff = torch.abs(ta - tb)
+    p99_abs_diff = torch.quantile(abs_diff.flatten(), 0.99).item()
+    return abs_diff.mean().item(), p99_abs_diff
+
+
+def _extract_single_image(outputs) -> Image.Image:
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    if not hasattr(first_output, "request_output") or not first_output.request_output:
+        raise ValueError("No request_output found in OmniRequestOutput")
+
+    req_out = first_output.request_output
+    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
+        raise ValueError("Invalid request_output structure or missing 'images' key")
+
+    images = req_out.images
+    if images is None or len(images) != 1:
+        raise ValueError(f"Expected 1 image, got {0 if images is None else len(images)}")
+    return images[0]
+
+
+def _run_zimage_generate(
+    *,
+    tp_size: int,
+    height: int,
+    width: int,
+    num_inference_steps: int,
+    seed: int,
+    enforce_eager: bool,
+    vae_use_tiling: bool = False,
+    vae_patch_parallel_size: int = 1,
+    num_requests: int = 4,
+) -> tuple[Image.Image, float, float]:
+    if num_requests < 2:
+        raise ValueError("num_requests must be >= 2 (1 warmup + >=1 timed)")
+
+    current_omni_platform.empty_cache()
+    device_index = current_omni_platform.current_device()
+    monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
+    monitor.start()
+    try:
+        # Each run needs a distinct DiffusionParallelConfig; use OmniRunner per call (not the
+        # parametrized omni_runner fixture, which is fixed per module).
+        with OmniRunner(
+            _get_zimage_model(),
+            parallel_config=DiffusionParallelConfig(
+                tensor_parallel_size=tp_size,
+                vae_patch_parallel_size=vae_patch_parallel_size,
+            ),
+            enforce_eager=enforce_eager,
+            vae_use_tiling=vae_use_tiling,
+        ) as runner:
+            # NOTE: Omni closes itself when a generate() call is exhausted.
+            # To avoid measuring teardown time (process shutdown, memory cleanup),
+            # we measure the latency to produce *subsequent* outputs within a single
+            # generator run.
+            #
+            # This also serves as a warmup: the first output may include extra
+            # compilation/caching overhead, while later outputs are closer to
+            # steady-state inference.
+            gen = runner.omni.generate(
+                [PROMPT] * num_requests,
+                OmniDiffusionSamplingParams(
+                    height=height,
+                    width=width,
+                    num_inference_steps=num_inference_steps,
+                    guidance_scale=0.0,
+                    seed=seed,
+                    num_outputs_per_prompt=1,
+                ),
+                py_generator=True,
+            )
+
+            warmup_output = next(gen)
+
+            t_prev = time.perf_counter()
+            per_request_times_s: list[float] = []
+            last_output = warmup_output
+            for _ in range(num_requests - 1):
+                last_output = next(gen)
+                t_now = time.perf_counter()
+                per_request_times_s.append(t_now - t_prev)
+                t_prev = t_now
+
+            # Ensure the generator is fully consumed so it can clean up.
+            for _ in gen:
+                pass
+
+            median_time_s = float(np.median(per_request_times_s))
+
+            peak_memory_mb = monitor.peak_used_mb
+
+            return _extract_single_image([last_output]), median_time_s, peak_memory_mb
+    finally:
+        monitor.stop()
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
+def test_zimage_tensor_parallel_tp2(tmp_path: Path):
+    if current_omni_platform.is_npu():
+        pytest.skip("Z-Image TP e2e test is only supported on CUDA and ROCm for now.")
+    if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
+        pytest.skip("Z-Image TP=2 requires >= 2 devices.")
+
+    enforce_eager = False
+
+    height = 512
+    width = 512
+    num_inference_steps = 2
+    seed = 42
+
+    tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate(
+        tp_size=1,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+        enforce_eager=enforce_eager,
+    )
+    tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate(
+        tp_size=2,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+        enforce_eager=enforce_eager,
+    )
+
+    tp1_path = tmp_path / "zimage_tp1.png"
+    tp2_path = tmp_path / "zimage_tp2.png"
+    tp1_img.save(tp1_path)
+    tp2_img.save(tp2_path)
+
+    assert tp1_img.width == width and tp1_img.height == height
+    assert tp2_img.width == width and tp2_img.height == height
+
+    mean_abs_diff, p99_abs_diff = _diff_metrics(tp1_img, tp2_img)
+    mean_threshold = 3e-2
+    p99_threshold = 2.5e-1
+    print(
+        "Z-Image TP image diff stats (TP=1 vs TP=2): "
+        f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
+        f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
+        f"tp1_img={tp1_path}, tp2_img={tp2_path}"
+    )
+    assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
+        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
+        f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
+    )
+
+    print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
+    # ROCm is not optimized TP2 can be slower than TP1
+    if not current_omni_platform.is_rocm():
+        assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})"
+
+    print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}")
+    assert tp2_peak_mem < tp1_peak_mem, (
+        f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})"
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
+def test_zimage_vae_patch_parallel_tp2(tmp_path: Path):
+    if current_omni_platform.is_npu():
+        pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA and ROCm for now.")
+    if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
+        pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 devices.")
+
+    enforce_eager = False
+
+    # Use a larger image to ensure there are multiple VAE tiles.
+    height = 1152
+    width = 1152
+    num_inference_steps = 2
+    seed = 42
+
+    baseline_img, _baseline_time_s, _baseline_peak_mem = _run_zimage_generate(
+        tp_size=2,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+        enforce_eager=enforce_eager,
+        vae_use_tiling=True,
+        vae_patch_parallel_size=1,
+        num_requests=2,
+    )
+    pp2_img, _pp2_time_s, _pp2_peak_mem = _run_zimage_generate(
+        tp_size=2,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+        enforce_eager=enforce_eager,
+        vae_use_tiling=True,
+        vae_patch_parallel_size=2,
+        num_requests=2,
+    )
+
+    baseline_path = tmp_path / "zimage_tp2_vae_pp1.png"
+    pp2_path = tmp_path / "zimage_tp2_vae_pp2.png"
+    baseline_img.save(baseline_path)
+    pp2_img.save(pp2_path)
+
+    mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img)
+    mean_threshold = 5e-3
+    p99_threshold = 1e-1
+    print(
+        "Z-Image VAE patch parallel image diff stats (TP=2, pp=1 vs pp=2): "
+        f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
+        f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
+        f"pp1_img={baseline_path}, pp2_img={pp2_path}"
+    )
+    assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
+        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
+        f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
+    )