diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index f7f3e4b05..402e62518 100644 --- a/.buildkite/pipeline.yml +++ b/.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 diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index 7334df1d4..6e2c18d64 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/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. diff --git a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml index 4c6a67e8b..8c170dffd 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml +++ b/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 diff --git a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml index 67ff7309f..ef605a3a3 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml +++ b/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 diff --git a/tests/e2e/offline_inference/test_bagel_text2img.py b/tests/e2e/offline_inference/test_bagel_text2img.py index 777af52a7..9aef47a74 100644 --- a/tests/e2e/offline_inference/test_bagel_text2img.py +++ b/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) 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, ) diff --git a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py index 9d216363d..2f6d54aef 100644 --- a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py +++ b/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 diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index ad35f5f87..144c73f3a 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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 diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index fc91e02b4..cfee79157 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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, diff --git a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py index dcf9d54ac..dde690068 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py +++ b/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 diff --git a/vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py index 75fe5dbf4..1e759a8d2 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py +++ b/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)]