From 9494d691bf1f1d0bf41176c3ccbf00b0f7731403 Mon Sep 17 00:00:00 2001 From: XU Mingshi <91017482+mxuax@users.noreply.github.com> Date: Tue, 3 Feb 2026 12:23:15 +0800 Subject: [PATCH] [CI] Refactor test_sequence_parallel.py and add a warmup run for more accurate performance stat (#1165) Signed-off-by: mxuax Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- .../test_sequence_parallel.py | 381 +++++++++--------- .../attention/backends/ring_pytorch_attn.py | 3 + .../diffusion/attention/parallel/factory.py | 17 +- 3 files changed, 215 insertions(+), 186 deletions(-) diff --git a/tests/e2e/offline_inference/test_sequence_parallel.py b/tests/e2e/offline_inference/test_sequence_parallel.py index 262c80243..3e7bb5617 100644 --- a/tests/e2e/offline_inference/test_sequence_parallel.py +++ b/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) diff --git a/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py b/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py index 9ed2c2076..482cbc9f8 100644 --- a/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py +++ b/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, diff --git a/vllm_omni/diffusion/attention/parallel/factory.py b/vllm_omni/diffusion/attention/parallel/factory.py index 2d5ef2317..ce95d409a 100644 --- a/vllm_omni/diffusion/attention/parallel/factory.py +++ b/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, )