[Diffusion]: Diffusion Ring Attention support (#273)

Signed-off-by: mxuax <mxuax@connect.ust.hk>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -11,11 +11,13 @@ short for CI.
 
 import os
 import sys
+import time
 from pathlib import Path
 
 import numpy as np
 import pytest
 import torch
+import torch.distributed as dist
 from PIL import Image
 
 # ruff: noqa: E402
@@ -51,18 +53,61 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
     return abs_diff.mean().item(), abs_diff.max().item()
 
 
+def _get_images(output):
+    """Extract images from output, handling both dict and SimpleNamespace types.
+
+    The output structure varies depending on serialization path:
+    - Direct memory: SimpleNamespace with .images attribute
+    - SHM serialization: dict with "images" key (dataclass converted via asdict)
+    - Wrapped output: SimpleNamespace(output=...) which needs unwrapping
+    """
+    # Check if output has direct images attribute (diffusion mode)
+    if hasattr(output, "images") and output.images:
+        return output.images
+
+    # Check request_output for pipeline mode
+    if output.request_output is None:
+        return None
+
+    if isinstance(output.request_output, list) and len(output.request_output) == 0:
+        return None
+
+    item = output.request_output[0]
+
+    # Handle wrapped SimpleNamespace (e.g. from omni_stage.py)
+    # Some items are wrapped as SimpleNamespace(request_id=..., output=...)
+    while hasattr(item, "output") and not hasattr(item, "images"):
+        item = item.output
+        if item is None:
+            return None
+
+    # Handle both dict (from SHM serialization) and object (direct) types
+    if isinstance(item, dict):
+        return item.get("images")
+    return getattr(item, "images", None)
+
+
 @pytest.mark.parametrize("model_name", models)
-@pytest.mark.parametrize("ulysses_degree", [2])
-@pytest.mark.parametrize("ring_degree", [1])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-def test_sequence_parallel(model_name: str, ulysses_degree: int, ring_degree: int, dtype: torch.dtype):
+@pytest.mark.parametrize("ulysses_degree", [1, 2])
+@pytest.mark.parametrize("ring_degree", [1, 2])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])  # Only test bfloat16 to reduce CI time
+@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:
-        pytest.skip("This test compares ulysses_degree=1 vs ulysses_degree>1; provide ulysses_degree>1.")
+    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."
+        )
 
     # Skip if not enough GPUs available for SP run
-    if device_count() < ulysses_degree:
-        pytest.skip(f"Test requires {ulysses_degree} GPUs but only {device_count()} available")
+    if device_count() < ulysses_degree * ring_degree:
+        pytest.skip(f"Test requires {ulysses_degree * ring_degree} GPUs but only {device_count()} available")
 
     # Use minimal settings for fast testing
     height = 256
@@ -76,6 +121,7 @@ def test_sequence_parallel(model_name: str, ulysses_degree: int, ring_degree: in
         model=model_name,
         parallel_config=baseline_parallel_config,
         dtype=dtype,
+        attention_backend=attn_backend,
     )
     try:
         outputs = baseline.generate(
@@ -90,6 +136,11 @@ def test_sequence_parallel(model_name: str, ulysses_degree: int, ring_degree: in
         baseline_images = list(outputs)[0].request_output[0].images
     finally:
         baseline.close()
+        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)
+        time.sleep(2)  # Wait for resources to release
 
     assert baseline_images is not None
     assert len(baseline_images) == 1
@@ -102,6 +153,7 @@ def test_sequence_parallel(model_name: str, ulysses_degree: int, ring_degree: in
         model=model_name,
         parallel_config=sp_parallel_config,
         dtype=dtype,
+        attention_backend=attn_backend,
     )
     try:
         outputs = sp.generate(
@@ -116,6 +168,11 @@ def test_sequence_parallel(model_name: str, ulysses_degree: int, ring_degree: in
         sp_images = list(outputs)[0].request_output[0].images
     finally:
         sp.close()
+        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)
+        time.sleep(2)
 
     assert sp_images is not None
     assert len(sp_images) == 1
@@ -134,7 +191,7 @@ def test_sequence_parallel(model_name: str, ulysses_degree: int, ring_degree: in
         max_threshold = 1e-1
 
     print(
-        "Image diff stats (baseline ulysses_degree=1 vs SP): "
+        "Image diff stats (baseline ulysses_degree*ring_degree=1 vs SP): "
         f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; "
         f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
         f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"