[CI] Refactor test_sequence_parallel.py and add a warmup run for more accurate performance stat (#1165)

Signed-off-by: mxuax <mxuax@connect.ust.hk>
Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -2,11 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-System test for Ulysses sequence parallel backend.
+System test for Sequence Parallel (SP) backends: Ulysses and Ring attention.
 
-This test verifies that Ulysses-SP (DeepSpeed Ulysses Sequence Parallel) works
-correctly with diffusion models. It uses minimal settings to keep test time
-short for CI.
+Tests verify that SP inference produces correct outputs compared to baseline.
 """
 
 import gc
@@ -14,6 +12,7 @@ import os
 import sys
 import time
 from pathlib import Path
+from typing import NamedTuple
 
 import numpy as np
 import pytest
@@ -32,24 +31,22 @@ from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.platforms import current_omni_platform
 
-models = ["riverclouds/qwen_image_random"]
-
+# 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
 
 
-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)
+class InferenceResult(NamedTuple):
+    """Result of an inference run."""
 
-
-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 = _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)
-    return abs_diff.mean().item(), abs_diff.max().item()
+    images: list[Image.Image]
+    elapsed_ms: float
 
 
 def _cleanup_distributed():
@@ -61,7 +58,6 @@ def _cleanup_distributed():
         os.environ.pop(key, None)
 
     gc.collect()
-
     if current_omni_platform.is_available():
         current_omni_platform.empty_cache()
         current_omni_platform.synchronize()
@@ -69,199 +65,216 @@ def _cleanup_distributed():
     time.sleep(5)
 
 
-def _run_baseline(model_name: str, dtype: torch.dtype, attn_backend: str, height: int, width: int, seed: int):
-    """Run baseline inference (no SP)."""
-    baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1)
-    baseline = Omni(
+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)
+    omni = Omni(
         model=model_name,
-        parallel_config=baseline_parallel_config,
+        parallel_config=parallel_config,
         dtype=dtype,
         attention_backend=attn_backend,
     )
 
     try:
-        outputs = baseline.generate(
+        # 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=4,
+                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,
             ),
         )
-        return outputs[0].request_output[0].images
+        elapsed_ms = (time.time() - start) * 1000
+
+        return InferenceResult(
+            images=outputs[0].request_output[0].images,
+            elapsed_ms=elapsed_ms,
+        )
     finally:
-        baseline.close()
+        omni.close()
         _cleanup_distributed()
 
 
-def _run_sp(
-    model_name: str,
-    dtype: torch.dtype,
-    attn_backend: str,
-    height: int,
-    width: int,
-    seed: int,
-    ulysses_degree: int,
-    ring_degree: int,
-):
-    """Run SP inference."""
-    sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
-    sp = Omni(
-        model=model_name,
-        parallel_config=sp_parallel_config,
-        dtype=dtype,
-        attention_backend=attn_backend,
-    )
+# =============================================================================
+# Correctness & Performance Tests
+# =============================================================================
 
-    try:
-        outputs = sp.generate(
-            PROMPT,
-            OmniDiffusionSamplingParams(
+# 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 = [
+    (2, 1, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True),  # Ulysses-2 - performance test
+    (1, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True),  # Ring-2 - performance test
+    (2, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, False, False),  # Hybrid - correctness only
+    (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.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:
+        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,
-                num_inference_steps=4,
-                guidance_scale=0.0,
-                generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
-                num_outputs_per_prompt=1,
-            ),
+                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,
         )
-        return outputs[0].request_output[0].images
-    finally:
-        sp.close()
-        _cleanup_distributed()
+        assert len(sp_result.images) == 1
 
+        # Compare outputs (correctness)
+        mean_diff, max_diff = _diff_metrics(baseline.images[0], sp_result.images[0])
 
-@pytest.mark.parametrize("model_name", models)
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("attn_backend", ["sdpa"])
-def test_baseline_only(model_name: str, dtype: torch.dtype, attn_backend: str):
-    """Test baseline inference only (no SP)."""
-    height = 256
-    width = 256
-    seed = 42
+        # 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)
 
-    images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
+        # 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)")
 
-    assert images is not None
-    assert len(images) == 1
-    assert images[0].width == width
-    assert images[0].height == height
+        print(f"[{sp_mode}] diff: mean={mean_diff:.6e}, max={max_diff:.6e}")
 
-
-@pytest.mark.parametrize("model_name", models)
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("attn_backend", ["sdpa"])
-def test_sp_ulysses2_only(model_name: str, dtype: torch.dtype, attn_backend: str):
-    """Test SP inference only (ulysses=2)."""
-    if current_omni_platform.get_device_count() < 2:
-        pytest.skip(f"Test requires 2 GPUs but only {current_omni_platform.get_device_count()} available")
-
-    height = 256
-    width = 256
-    seed = 42
-
-    images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree=2, ring_degree=1)
-
-    assert images is not None
-    assert len(images) == 1
-    assert images[0].width == width
-    assert images[0].height == height
-
-
-@pytest.mark.parametrize("model_name", models)
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("attn_backend", ["sdpa"])
-def test_sp_ring2_only(model_name: str, dtype: torch.dtype, attn_backend: str):
-    """Test SP inference only (ring=2)."""
-    if current_omni_platform.get_device_count() < 2:
-        pytest.skip(f"Test requires 2 GPUs but only {current_omni_platform.get_device_count()} available")
-
-    height = 256
-    width = 256
-    seed = 42
-
-    images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree=1, ring_degree=2)
-
-    assert images is not None
-    assert len(images) == 1
-    assert images[0].width == width
-    assert images[0].height == height
-
-
-@pytest.mark.parametrize("model_name", models)
-@pytest.mark.parametrize("ulysses_degree", [1, 2])
-@pytest.mark.parametrize("ring_degree", [1, 2])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("attn_backend", ["sdpa"])
-def test_sequence_parallel(
-    model_name: str,
-    ulysses_degree: int,
-    ring_degree: int,
-    dtype: torch.dtype,
-    attn_backend: str,
-):
-    """Compare baseline (ulysses_degree=1) vs SP (ulysses_degree>1) outputs."""
-    if ulysses_degree <= 1 and ring_degree <= 1:
-        pytest.skip(
-            "This test compares ulysses_degree * ring_degree = 1 vs ulysses_degree * ring_degree > 1; "
-            "provide ulysses_degree or ring_degree>1."
+        # 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}"
         )
 
-    sp_size = ulysses_degree * ring_degree
-    if current_omni_platform.get_device_count() < sp_size:
-        pytest.skip(f"Test requires {sp_size} GPUs but only {current_omni_platform.get_device_count()} available")
-
-    height = 256
-    width = 256
-    seed = 42
-
-    baseline_images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
-    assert baseline_images is not None and len(baseline_images) == 1
-
-    sp_images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree, ring_degree)
-    assert sp_images is not None and len(sp_images) == 1
-
-    mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0])
-
-    mean_threshold = 2e-2
-    max_threshold = 2e-1
-
-    assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
-        f"Image diff exceeded threshold: mean={mean_abs_diff:.6e}, max={max_abs_diff:.6e}"
-    )
-
-
-@pytest.mark.parametrize("model_name", models)
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("attn_backend", ["sdpa"])
-def test_sequence_parallel_ulysses4(model_name: str, dtype: torch.dtype, attn_backend: str):
-    """Test SP with ulysses_degree=4."""
-    ulysses_degree = 4
-    ring_degree = 1
-
-    if current_omni_platform.get_device_count() < ulysses_degree * ring_degree:
-        pytest.skip(
-            f"Test requires {ulysses_degree * ring_degree} GPUs but only {current_omni_platform.get_device_count()} available"
+    # 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}"
         )
-
-    height = 272
-    width = 272
-    seed = 42
-
-    baseline_images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
-    assert baseline_images is not None and len(baseline_images) == 1
-
-    sp_images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree, ring_degree)
-    assert sp_images is not None and len(sp_images) == 1
-
-    mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0])
-
-    mean_threshold = 2e-2
-    max_threshold = 2e-1
-
-    assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
-        f"Image diff exceeded threshold: mean={mean_abs_diff:.6e}, max={max_abs_diff:.6e}"
-    )
+    print("=" * 70)

vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py

@@ -7,11 +7,14 @@
 
 
 import torch
+from vllm.logger import init_logger
 
 from vllm_omni.diffusion.attention.backends.ring.ring_kernels import pytorch_attn_forward
 from vllm_omni.diffusion.attention.backends.ring.ring_utils import update_out_and_lse
 from vllm_omni.diffusion.distributed.comm import RingComm
 
+logger = init_logger(__name__)
+
 
 def ring_pytorch_attn_func(
     q,

vllm_omni/diffusion/attention/parallel/factory.py

@@ -3,12 +3,16 @@
 
 from __future__ import annotations
 
+from vllm.logger import init_logger
+
 from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention, ParallelAttentionStrategy
 from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
 from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
 from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size, get_sp_group
 from vllm_omni.diffusion.forward_context import get_forward_context
 
+logger = init_logger(__name__)
+
 
 def build_parallel_attention_strategy(
     *,
@@ -26,7 +30,8 @@ def build_parallel_attention_strategy(
     try:
         cfg = get_forward_context().omni_diffusion_config
         p = cfg.parallel_config
-    except Exception:
+    except Exception as e:
+        logger.debug(f"No forward context available for parallel attention strategy: {e}")
         return NoParallelAttention()
 
     ulysses_degree = getattr(p, "ulysses_degree", 1)
@@ -37,11 +42,18 @@ def build_parallel_attention_strategy(
         # Ensure SP group is initialized and world size > 1
         if get_sequence_parallel_world_size() <= 1:
             return NoParallelAttention()
-    except Exception:
+    except Exception as e:
+        # Log warning if SP is configured but group is not available
+        if ulysses_degree > 1 or ring_degree > 1:
+            logger.warning(
+                f"SP configured (ulysses={ulysses_degree}, ring={ring_degree}) but SP group not available: {e}. "
+                f"Falling back to NoParallelAttention. This may cause incorrect results."
+            )
         return NoParallelAttention()
 
     # Ulysses (or Hybrid Ulysses+Ring)
     if ulysses_degree > 1:
+        logger.debug(f"Using UlyssesParallelAttention (ulysses_degree={ulysses_degree})")
         return UlyssesParallelAttention(
             sp_group=sp_group,
             scatter_idx=scatter_idx,
@@ -51,6 +63,7 @@ def build_parallel_attention_strategy(
 
     # Pure Ring Attention
     if ring_degree > 1:
+        logger.debug(f"Using RingParallelAttention (ring_degree={ring_degree})")
         return RingParallelAttention(
             sp_group=sp_group,
         )