Merge branch 'main' into dev/vllm-align

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

.buildkite/pipeline.yml

@@ -233,43 +233,43 @@ steps:
         volumes:
         - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Omni Model Test with H100"
-    timeout_in_minutes: 30
-    depends_on: image-build
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
-      - pytest -s -v tests/e2e/online_serving/test_async_omni.py
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
+  # - label: "Omni Model Test with H100"
+  #   timeout_in_minutes: 30
+  #   depends_on: image-build
+  #   commands:
+  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #     - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+  #     - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+  #     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
+  #     - pytest -s -v tests/e2e/online_serving/test_async_omni.py
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #               resources:
+  #                 limits:
+  #                   nvidia.com/gpu: 2
+  #               volumeMounts:
+  #                 - name: devshm
+  #                   mountPath: /dev/shm
+  #                 - name: hf-cache
+  #                   mountPath: /root/.cache/huggingface
+  #               env:
+  #                 - name: HF_HOME
+  #                   value: /root/.cache/huggingface
+  #           nodeSelector:
+  #             node.kubernetes.io/instance-type: gpu-h100-sxm
+  #           volumes:
+  #             - name: devshm
+  #               emptyDir:
+  #                 medium: Memory
+  #             - name: hf-cache
+  #               hostPath:
+  #                 path: /mnt/hf-cache
+  #                 type: DirectoryOrCreate
 
   - label: "Diffusion Image Edit Test with H100 (1 GPU)"
     timeout_in_minutes: 20
@@ -306,37 +306,37 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: "Bagel Text2Img Model Test with H100"
-    timeout_in_minutes: 30
-    depends_on: image-build
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
+  # - label: "Bagel Text2Img Model Test with H100"
+  #   timeout_in_minutes: 30
+  #   depends_on: image-build
+  #   commands:
+  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #               resources:
+  #                 limits:
+  #                   nvidia.com/gpu: 1
+  #               volumeMounts:
+  #                 - name: devshm
+  #                   mountPath: /dev/shm
+  #                 - name: hf-cache
+  #                   mountPath: /root/.cache/huggingface
+  #               env:
+  #                 - name: HF_HOME
+  #                   value: /root/.cache/huggingface
+  #           nodeSelector:
+  #             node.kubernetes.io/instance-type: gpu-h100-sxm
+  #           volumes:
+  #             - name: devshm
+  #               emptyDir:
+  #                 medium: Memory
+  #             - name: hf-cache
+  #               hostPath:
+  #                 path: /mnt/hf-cache
+  #                 type: DirectoryOrCreate

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -18,20 +18,19 @@ The following table shows which models are currently supported by parallelism me
 
 ### ImageGen
 
-| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
-|-------|------------------|------------|---------|--------------|--------------------------|
-| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ❌ |
-| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ❌ |
-| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ |
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) |
-| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
-| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ |
-| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ❌ | ✅ |
-
+| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
+|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|
+| **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |
+| **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |
+| **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |
+| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |
+| **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |
+| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |
+| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |
+| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |
+| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ❌        |
+| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ❌      |    ❌    |      ❌       |        ✅        |
+| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |
 
 !!! note "TP Limitations for Diffusion Models"
     We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.

tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml

@@ -12,7 +12,7 @@ stage_args:
       model_arch: BagelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.60
+      gpu_memory_utilization: 0.35
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: text
@@ -44,7 +44,7 @@ stage_args:
       max_batch_size: 1
     engine_args:
       model_stage: dit
-      gpu_memory_utilization: 0.4
+      gpu_memory_utilization: 0.55
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: image

tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml

@@ -12,7 +12,7 @@ stage_args:
       model_arch: BagelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.65
+      gpu_memory_utilization: 0.35
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: text
@@ -43,7 +43,7 @@ stage_args:
       max_batch_size: 1
     engine_args:
       model_stage: dit
-      gpu_memory_utilization: 0.4
+      gpu_memory_utilization: 0.55
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: image

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -161,7 +161,7 @@ def _generate_bagel_image(omni: Omni, prompt: str = DEFAULT_PROMPT) -> Image.Ima
 def test_bagel_text2img_shared_memory_connector():
     """Test Bagel text2img with shared memory connector."""
     config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
-    omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path)
+    omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, stage_init_timeout=300)
 
     try:
         generated_image = _generate_bagel_image(omni)
@@ -288,7 +288,7 @@ def test_bagel_text2img_mooncake_connector():
             http_port=MOONCAKE_HTTP_PORT,
         )
 
-        omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file)
+        omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file, stage_init_timeout=300)
 
         generated_image = _generate_bagel_image(omni)
         _validate_pixels(generated_image)

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,
         )

vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py

@@ -9,10 +9,11 @@ import torch.nn as nn
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
+from vllm.distributed import get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
+from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
@@ -35,9 +36,9 @@ class FeedForward(nn.Module):
         inner_dim = int(dim * mult)
         dim_out = dim_out if dim_out is not None else dim
 
-        self.w_in = ReplicatedLinear(dim, inner_dim, bias=bias, return_bias=False)
+        self.w_in = ColumnParallelLinear(dim, inner_dim, bias=bias, return_bias=False)
         self.act = get_act_fn("gelu_pytorch_tanh")
-        self.w_out = ReplicatedLinear(inner_dim, dim_out, bias=bias, return_bias=False)
+        self.w_out = RowParallelLinear(inner_dim, dim_out, bias=bias, return_bias=False)
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = self.w_in(hidden_states)
@@ -83,12 +84,11 @@ class LongCatImageAttention(nn.Module):
             hidden_size=query_dim,
             head_size=self.head_dim,
             total_num_heads=self.heads,
-            disable_tp=True,
             bias=bias,
         )
 
         if not self.pre_only:
-            self.to_out = torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)
+            self.to_out = RowParallelLinear(self.inner_dim, self.out_dim, bias=out_bias)
 
         if self.added_kv_proj_dim is not None:
             self.norm_added_q = RMSNorm(dim_head, eps=eps)
@@ -98,11 +98,10 @@ class LongCatImageAttention(nn.Module):
                 hidden_size=self.added_kv_proj_dim,
                 head_size=self.head_dim,
                 total_num_heads=self.heads,
-                disable_tp=True,
                 bias=added_proj_bias,
             )
 
-            self.to_add_out = ReplicatedLinear(self.inner_dim, query_dim, bias=out_bias)
+            self.to_add_out = RowParallelLinear(self.inner_dim, query_dim, bias=out_bias)
 
         self.attn = Attention(
             num_heads=heads,
@@ -185,22 +184,26 @@ class LongCatImageAttention(nn.Module):
         """
         qkv, _ = self.to_qkv(hidden_states)
 
-        query, key, value = qkv.chunk(3, dim=-1)
+        q_size = self.to_qkv.num_heads * self.head_dim
+        kv_size = self.to_qkv.num_kv_heads * self.head_dim
+        query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
 
-        query = query.unflatten(-1, (self.heads, -1))
-        key = key.unflatten(-1, (self.heads, -1))
-        value = value.unflatten(-1, (self.heads, -1))
+        query = query.unflatten(-1, (self.to_qkv.num_heads, -1))
+        key = key.unflatten(-1, (self.to_qkv.num_kv_heads, -1))
+        value = value.unflatten(-1, (self.to_qkv.num_kv_heads, -1))
 
         query = self.norm_q(query)
         key = self.norm_k(key)
 
         if self.added_kv_proj_dim is not None:
             encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states)
-            encoder_query, encoder_key, encoder_value = encoder_qkv.chunk(3, dim=-1)
+            q_size = self.add_kv_proj.num_heads * self.head_dim
+            kv_size = self.add_kv_proj.num_kv_heads * self.head_dim
+            encoder_query, encoder_key, encoder_value = encoder_qkv.split([q_size, kv_size, kv_size], dim=-1)
 
-            encoder_query = encoder_query.unflatten(-1, (self.heads, -1))
-            encoder_key = encoder_key.unflatten(-1, (self.heads, -1))
-            encoder_value = encoder_value.unflatten(-1, (self.heads, -1))
+            encoder_query = encoder_query.unflatten(-1, (self.add_kv_proj.num_heads, -1))
+            encoder_key = encoder_key.unflatten(-1, (self.add_kv_proj.num_kv_heads, -1))
+            encoder_value = encoder_value.unflatten(-1, (self.add_kv_proj.num_kv_heads, -1))
 
             # Apply RMSNorm to text Q/K
             encoder_query = self.norm_added_q(encoder_query)
@@ -284,11 +287,14 @@ class LongCatImageAttention(nn.Module):
             encoder_hidden_states, hidden_states = hidden_states.split_with_sizes(
                 [encoder_hidden_states.shape[1], hidden_states.shape[1] - encoder_hidden_states.shape[1]], dim=1
             )
-            hidden_states = self.to_out(hidden_states)
+            hidden_states, _ = self.to_out(hidden_states)
             encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states)
 
             return hidden_states, encoder_hidden_states
         else:
+            # For single-stream blocks, there's no to_out (RowParallelLinear) to handle the reduction
+            if get_tensor_model_parallel_world_size() > 1:
+                hidden_states = tensor_model_parallel_all_gather(hidden_states, dim=-1)
             return hidden_states
 
 

vllm_omni/entrypoints/async_omni.py

@@ -681,9 +681,7 @@ class AsyncOmni(OmniBase):
         OMNI: Required by upstream OpenAIServingModels.__init__ which
         accesses engine_client.renderer.
         """
-        if self.input_processor is not None:
-            return self.input_processor.renderer
-        return None
+        return self.input_processor.renderer
 
     async def do_log_stats(self) -> None:
         pass

vllm_omni/entrypoints/openai/serving_chat.py

@@ -13,6 +13,7 @@ import torch
 from fastapi import Request
 from PIL import Image
 from pydantic import TypeAdapter
+from vllm.renderers import RendererLike
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
@@ -328,7 +329,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
     async def _preprocess_chat(
         self,
         request: ChatLikeRequest | ResponsesRequest,
-        renderer: Any,
+        renderer: RendererLike,
         messages: list[ChatCompletionMessageParam],
         chat_template: str | None,
         chat_template_content_format: ChatTemplateContentFormatOption,

vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py

@@ -504,12 +504,19 @@ class Qwen3TTSConfig(PretrainedConfig):
         self.tts_bos_token_id = tts_bos_token_id
         self.tts_eos_token_id = tts_eos_token_id
 
+        # TODO: remove these dummy values after
+        self.image_token_id = 0  # dummy image token id
+        self.video_token_id = 0  # dummy video token id
+        self.vision_start_token_id = 0  # dummy vision start token id
+        self.vision_config = PretrainedConfig()  # dummy vision config
+        self.vision_config.spatial_merge_size = 1
+
     def get_text_config(self, **kwargs):
         # vLLM expects text config to expose hidden_size/num_attention_heads.
         # For Qwen3 TTS, the talker config is the text model config.
         config = self.talker_config
-        if hasattr(config, "rope_parameters"):
-            delattr(config, "rope_parameters")
+        # if hasattr(config, "rope_parameters"):
+        #     delattr(config, "rope_parameters")
         return config
 
 

vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py

@@ -430,6 +430,22 @@ def mel_spectrogram(
     return mel_spec
 
 
+def _compute_default_rope_parameters(
+    config,
+    device,
+):
+    base = config.rope_theta
+    partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
+    head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+    dim = int(head_dim * partial_rotary_factor)
+
+    attention_factor = 1.0  # Unused in this type of RoPE
+
+    # Compute the inverse frequencies
+    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim))
+    return inv_freq, attention_factor
+
+
 class Qwen3TTSPreTrainedModel(PreTrainedModel):
     config_class = Qwen3TTSConfig
     base_model_prefix = "model"
@@ -502,7 +518,9 @@ class Qwen3TTSTalkerRotaryEmbedding(nn.Module):
         self.original_max_seq_len = config.max_position_embeddings
 
         self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        self.rope_init_fn: Callable = _compute_default_rope_parameters
+        if self.rope_type != "default":
+            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
 
         inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
         self.register_buffer("inv_freq", inv_freq, persistent=False)
@@ -538,7 +556,9 @@ class Qwen3TTSRotaryEmbedding(nn.Module):
         self.original_max_seq_len = config.max_position_embeddings
 
         self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        self.rope_init_fn: Callable = _compute_default_rope_parameters
+        if self.rope_type != "default":
+            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
 
         inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
         self.register_buffer("inv_freq", inv_freq, persistent=False)
@@ -1389,7 +1409,6 @@ class Qwen3TTSTalkerModel(Qwen3TTSTalkerTextPreTrainedModel):
 
     def __init__(self, config):
         super().__init__(config)
-        self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
         self.layers = nn.ModuleList(
             [Qwen3TTSTalkerDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]