From e17130090c9aead4df892b4dfe8644406e628d01 Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Fri, 13 Mar 2026 20:42:54 -0600 Subject: [PATCH 01/10] [BugFix] Fix Max Rank Handling in LoRA (#1397) Signed-off-by: Alex Brooks --- tests/diffusion/lora/test_lora_manager.py | 47 +++++++++++++++++++++++ vllm_omni/diffusion/lora/manager.py | 24 ++++++++++-- 2 files changed, 67 insertions(+), 4 deletions(-) diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index fe4e654e8..8d4a1487f 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -508,3 +508,50 @@ def test_scale_keys_are_rounded(): # Currently we round keys to 3 decimal places manager._update_adapter_scale(adapter_id, 0.0031) assert manager._adapter_scales[adapter_id] == 0.003 + + +def test_lora_manager_uses_valid_max_rank(monkeypatch): + """Ensure that the LoRA manager uses a valid max rank for vLLM.""" + manager = DiffusionLoRAManager( + pipeline=torch.nn.Module(), + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + + # Ensure that the rank is correctly adjusted to the smallest valid max rank + supported_max_rank = 64 + unsupported_max_rank = 63 + assert supported_max_rank in DiffusionLoRAManager._VALID_MAX_RANKS + assert unsupported_max_rank not in DiffusionLoRAManager._VALID_MAX_RANKS + + def _fake_load(_req: LoRARequest): + lora_model = type("LM", (), {"id": _req.lora_int_id})() + peft_helper = type("PH", (), {"r": unsupported_max_rank})() + return lora_model, peft_helper + + monkeypatch.setattr(manager, "_load_adapter", _fake_load) + req1 = _dummy_lora_request(1) + manager.add_adapter(req1) + assert manager._max_lora_rank == supported_max_rank + + +@pytest.mark.parametrize("rank", [-1, 0, DiffusionLoRAManager._VALID_MAX_RANKS[-1] + 1]) +def test_lora_manager_max_rank_validation(monkeypatch, rank): + """Check that invalid max ranks are handled correctly.""" + manager = DiffusionLoRAManager( + pipeline=torch.nn.Module(), + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + + lora_rank = rank + + def _fake_load(_req: LoRARequest): + lora_model = type("LM", (), {"id": _req.lora_int_id})() + peft_helper = type("PH", (), {"r": lora_rank})() + return lora_model, peft_helper + + monkeypatch.setattr(manager, "_load_adapter", _fake_load) + req1 = _dummy_lora_request(1) + with pytest.raises(ValueError): + manager.add_adapter(req1) diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 0dd5ba986..1466a3358 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -4,9 +4,11 @@ import math import time from collections import OrderedDict +from typing import get_args import torch import torch.nn as nn +from vllm.config.lora import MaxLoRARanks from vllm.logger import init_logger from vllm.lora.layers import BaseLayerWithLoRA from vllm.lora.lora_model import LoRAModel @@ -38,6 +40,9 @@ class DiffusionLoRAManager: Uses LRU cache management similar to LRUCacheLoRAModelManager. """ + # Valid max allowed ranks for LoRA in vLLM + _VALID_MAX_RANKS: list[int] = sorted(get_args(MaxLoRARanks)) + def __init__( self, pipeline: nn.Module, @@ -413,11 +418,10 @@ class DiffusionLoRAManager: if min_rank <= self._max_lora_rank: return - if min_rank <= 0: - raise ValueError(f"Invalid LoRA rank: {min_rank}") + valid_max_rank = self._get_smallest_valid_max_rank(min_rank) - logger.info("Increasing max LoRA rank: %d -> %d", self._max_lora_rank, min_rank) - self._max_lora_rank = min_rank + logger.info("Increasing max LoRA rank: %d -> %d", self._max_lora_rank, valid_max_rank) + self._max_lora_rank = valid_max_rank if not self._lora_modules: return @@ -441,6 +445,18 @@ class DiffusionLoRAManager: self._active_adapter_id = None self._activate_adapter(active_id, active_scale) + @classmethod + def _get_smallest_valid_max_rank(cls, min_rank: int) -> int: + """Given a LoRA rank, get the smallest max rank that can support it.""" + if min_rank <= 0: + raise ValueError(f"Invalid LoRA rank: {min_rank}") + + allowed_ranks = [rank for rank in cls._VALID_MAX_RANKS if rank >= min_rank] + if not allowed_ranks: + raise ValueError(f"LoRA rank of {min_rank} exceeds max allowed rank of {max(cls._VALID_MAX_RANKS)}") + + return min(allowed_ranks) + def _get_lora_weights( self, lora_model: LoRAModel, From c107f0b33671cc04087fdcdef180c1f62a4d3a47 Mon Sep 17 00:00:00 2001 From: Daniel Huang Date: Fri, 13 Mar 2026 20:24:05 -0700 Subject: [PATCH 02/10] Buildkite hardware ci xpu test (#1340) Signed-off-by: Daniel Huang --- .../scripts/hardware_ci/run-xpu-test.sh | 56 +++++++++ docker/Dockerfile.xpu | 2 +- docs/contributing/ci/tests_markers.md | 9 +- pyproject.toml | 3 + tests/benchmarks/test_serve_cli.py | 6 +- tests/e2e/offline_inference/test_cache_dit.py | 2 +- .../offline_inference/test_qwen2_5_omni.py | 7 +- .../e2e/offline_inference/test_qwen3_omni.py | 2 + .../test_stable_audio_model.py | 2 +- tests/e2e/offline_inference/test_t2i_model.py | 2 +- tests/e2e/offline_inference/test_teacache.py | 2 +- .../test_zimage_parallelism.py | 2 +- .../test_images_generations_lora.py | 2 +- tests/e2e/online_serving/test_qwen3_omni.py | 2 + .../test_qwen3_omni_expansion.py | 6 + .../stage_configs/xpu/qwen2_5_omni_ci.yaml | 110 ++++++++++++++++++ .../e2e/stage_configs/xpu/qwen3_omni_ci.yaml | 109 +++++++++++++++++ .../online_serving/test_qwen2_5_omni.py | 7 ++ .../online_serving/test_qwen3_omni.py | 6 + tests/utils.py | 38 +++++- 20 files changed, 357 insertions(+), 18 deletions(-) create mode 100755 .buildkite/scripts/hardware_ci/run-xpu-test.sh create mode 100644 tests/e2e/stage_configs/xpu/qwen2_5_omni_ci.yaml create mode 100644 tests/e2e/stage_configs/xpu/qwen3_omni_ci.yaml diff --git a/.buildkite/scripts/hardware_ci/run-xpu-test.sh b/.buildkite/scripts/hardware_ci/run-xpu-test.sh new file mode 100755 index 000000000..e1d464ef0 --- /dev/null +++ b/.buildkite/scripts/hardware_ci/run-xpu-test.sh @@ -0,0 +1,56 @@ +#!/bin/bash + +# This script build the XPU docker image and run the offline inference inside the container. +set -ex + +omni_source_dir=$(git rev-parse --show-toplevel) + +base_image_name="xpu/vllm-omni-ci-base:${VLLM_VERSION:?VLLM_VERSION must be set}" +image_name="xpu/vllm-omni-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set}" +container_name="xpu_${BUILDKITE_COMMIT}_$( + tr -dc A-Za-z0-9 None: """ @@ -87,7 +90,7 @@ def test_mix_to_audio(omni_runner, omni_runner_handler) -> None: @pytest.mark.core_model @pytest.mark.omni -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) +@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 3}) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) def test_text_to_text(omni_runner, omni_runner_handler) -> None: """ diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py index 8c57becef..01be0486f 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni.py +++ b/tests/e2e/offline_inference/test_qwen3_omni.py @@ -40,6 +40,8 @@ def get_cuda_graph_config(): if current_omni_platform.is_rocm(): # ROCm stage config optimized for MI325 GPU stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")] +elif current_omni_platform.is_xpu(): + stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")] else: stage_configs = [get_cuda_graph_config()] diff --git a/tests/e2e/offline_inference/test_stable_audio_model.py b/tests/e2e/offline_inference/test_stable_audio_model.py index 2acfe5097..c36f06493 100644 --- a/tests/e2e/offline_inference/test_stable_audio_model.py +++ b/tests/e2e/offline_inference/test_stable_audio_model.py @@ -23,7 +23,7 @@ models = ["linyueqian/stable_audio_random"] @pytest.mark.core_model @pytest.mark.diffusion -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": "L4", "xpu": "B60"}) @pytest.mark.parametrize("model_name", models) def test_stable_audio_model(model_name: str): m = Omni(model=model_name) diff --git a/tests/e2e/offline_inference/test_t2i_model.py b/tests/e2e/offline_inference/test_t2i_model.py index 33960f560..0cb0dafeb 100644 --- a/tests/e2e/offline_inference/test_t2i_model.py +++ b/tests/e2e/offline_inference/test_t2i_model.py @@ -36,7 +36,7 @@ elif current_omni_platform.is_rocm(): @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.diffusion -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 2}) +@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 1, "rocm": 2, "xpu": 2}) @pytest.mark.parametrize("model_name", models) def test_diffusion_model(model_name: str, run_level): if run_level == "core_model" and model_name != "riverclouds/qwen_image_random": diff --git a/tests/e2e/offline_inference/test_teacache.py b/tests/e2e/offline_inference/test_teacache.py index 96254f5b6..f2d110125 100644 --- a/tests/e2e/offline_inference/test_teacache.py +++ b/tests/e2e/offline_inference/test_teacache.py @@ -36,7 +36,7 @@ models = ["riverclouds/qwen_image_random"] @pytest.mark.core_model @pytest.mark.diffusion @pytest.mark.cache -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) +@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}) @pytest.mark.parametrize("model_name", models) def test_teacache(model_name: str): """Test TeaCache backend with diffusion model.""" diff --git a/tests/e2e/offline_inference/test_zimage_parallelism.py b/tests/e2e/offline_inference/test_zimage_parallelism.py index d5b5e39e5..b11176070 100644 --- a/tests/e2e/offline_inference/test_zimage_parallelism.py +++ b/tests/e2e/offline_inference/test_zimage_parallelism.py @@ -160,7 +160,7 @@ def _run_zimage_generate( # @pytest.mark.advanced_model # @pytest.mark.diffusion # @pytest.mark.parallel -# @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) +# @hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 2}) # def test_zimage_tensor_parallel_tp2(tmp_path: Path): # if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): # pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.") diff --git a/tests/e2e/online_serving/test_images_generations_lora.py b/tests/e2e/online_serving/test_images_generations_lora.py index ac5d81ce0..a1f1d1450 100644 --- a/tests/e2e/online_serving/test_images_generations_lora.py +++ b/tests/e2e/online_serving/test_images_generations_lora.py @@ -147,7 +147,7 @@ def _basic_payload() -> dict: @pytest.mark.advanced_model @pytest.mark.diffusion -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) +@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}) def test_images_generations_per_request_lora_switching(omni_server: OmniServer, tmp_path: Path) -> None: # Base generation. base_img = _post_images(omni_server, _basic_payload()) diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index 9227928f3..ef4d40198 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -48,6 +48,8 @@ def get_chunk_config(): if current_omni_platform.is_rocm(): # ROCm stage config optimized for MI325 GPU stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")] +elif current_omni_platform.is_xpu(): + stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")] else: stage_configs = [get_chunk_config()] diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index a28039888..6fb6a069e 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -6,6 +6,8 @@ E2E Online tests for Qwen3-Omni model. import os +from vllm_omni.platforms import current_omni_platform + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" from pathlib import Path @@ -50,6 +52,10 @@ def get_chunk_config(default_path): # CI stage config for 2*H100-80G GPUs default_path = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml") stage_configs = [default_path, get_chunk_config(default_path)] + +if current_omni_platform.is_xpu(): + stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")] + # Create parameter combinations for model and stage config test_params = [ OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs diff --git a/tests/e2e/stage_configs/xpu/qwen2_5_omni_ci.yaml b/tests/e2e/stage_configs/xpu/qwen2_5_omni_ci.yaml new file mode 100644 index 000000000..de7e9d901 --- /dev/null +++ b/tests/e2e/stage_configs/xpu/qwen2_5_omni_ci.yaml @@ -0,0 +1,110 @@ +# stage config for running qwen2.5-omni with architecture of OmniLLM. + +# The following config is verified with 2 * Intel Arc Pro B60 XPU. +stage_args: + - stage_id: 0 + stage_type: llm # Use llm stage type to launch OmniLLM + runtime: + process: true # Run this stage in a separate process + devices: "0" # Visible devices for this stage + max_batch_size: 1 + engine_args: + model_stage: thinker + model_arch: Qwen2_5OmniForConditionalGeneration + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + max_model_len: 16384 + max_num_batched_tokens: 16384 + max_num_seqs: 1 + gpu_memory_utilization: 0.9 # thinker weight is around 16.74GB for Qwen2.5-Omni-7B + skip_mm_profiling: true + enforce_eager: true + trust_remote_code: true + engine_output_type: latent + enable_prefix_caching: false + mm_processor_cache_gb: 0 + is_comprehension: true + final_output: true + final_output_type: text + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 128 + seed: 42 + detokenize: True + repetition_penalty: 1.1 + - stage_id: 1 + stage_type: llm # Use llm stage type to launch OmniLLM + runtime: + process: true + devices: "1" + max_batch_size: 1 + engine_args: + model_stage: talker + model_arch: Qwen2_5OmniForConditionalGeneration + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + max_model_len: 16384 + max_num_batched_tokens: 16384 + max_num_seqs: 1 + gpu_memory_utilization: 0.5 # talker weight is 6.03GB for Qwen2.5-Omni-7B + skip_mm_profiling: true + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + engine_output_type: latent + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker + default_sampling_params: + temperature: 0.9 + top_p: 0.8 + top_k: 40 + max_tokens: 4096 + seed: 42 + detokenize: True + repetition_penalty: 1.05 + stop_token_ids: [8294] + + - stage_id: 2 + stage_type: llm # Use llm stage type to launch OmniLLM + runtime: + process: true + devices: "2" + max_batch_size: 1 + engine_args: + model_stage: code2wav + model_arch: Qwen2_5OmniForConditionalGeneration + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + gpu_memory_utilization: 0.3 # code2wav weight is around 1.46GB for Qwen2.5-Omni-7B + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + engine_output_type: audio + engine_input_source: [1] + final_output: true + final_output_type: audio + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 2048 + seed: 42 + detokenize: True + repetition_penalty: 1.1 + +# Top-level runtime config (concise): default windows and stage edges +runtime: + enabled: true + defaults: + window_size: -1 # Simplified: trigger downstream only after full upstream completion + max_inflight: 1 # Simplified: process serially within each stage + + edges: + - from: 0 # thinker → talker: trigger only after receiving full input (-1) + to: 1 + window_size: -1 + - from: 1 # talker → code2wav: trigger only after receiving full input (-1) + to: 2 + window_size: -1 diff --git a/tests/e2e/stage_configs/xpu/qwen3_omni_ci.yaml b/tests/e2e/stage_configs/xpu/qwen3_omni_ci.yaml new file mode 100644 index 000000000..b362a8667 --- /dev/null +++ b/tests/e2e/stage_configs/xpu/qwen3_omni_ci.yaml @@ -0,0 +1,109 @@ +# Stage config for running Qwen3-Omni-MoE with 3-stage architecture +# Stage 0: Thinker (multimodal understanding + text generation) +# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes) +# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform) + +# The following config is verified with 8 * Intel Arc Pro B60 XPU. +stage_args: +- stage_id: 0 + stage_type: llm # Use llm stage type to launch OmniLLM + runtime: + devices: "0,1,2,3" + max_batch_size: 1 + engine_args: + model_stage: thinker + model_arch: Qwen3OmniMoeForConditionalGeneration + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.85 # thinker weight is around 61.08GB for Qwen3-Omni-30B-A3B-Instruct + skip_mm_profiling: true + enforce_eager: true + trust_remote_code: true + engine_output_type: latent # Output hidden states for talker + distributed_executor_backend: "mp" + max_num_batched_tokens: 4096 + max_model_len: 4096 + enable_prefix_caching: false + hf_config_name: thinker_config + tensor_parallel_size: 4 + max_cudagraph_capture_size: 0 + final_output: true + final_output_type: text + is_comprehension: true + default_sampling_params: + temperature: 0.4 + top_p: 0.9 + top_k: 1 + max_tokens: 100 + seed: 42 + ignore_eos: False + detokenize: True + repetition_penalty: 1.05 + +- stage_id: 1 + stage_type: llm # Use llm stage type to launch OmniLLM + runtime: + devices: "4" + max_batch_size: 1 + engine_args: + model_stage: talker + model_arch: Qwen3OmniMoeForConditionalGeneration + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.6 # talker weight is around 8.5GB for Qwen3-Omni-30B-A3B-Instruct + skip_mm_profiling: true + enforce_eager: true + trust_remote_code: true + engine_output_type: latent # Output codec codes for code2wav + enable_prefix_caching: false + max_num_batched_tokens: 4096 + max_model_len: 4096 + distributed_executor_backend: "mp" + hf_config_name: talker_config + max_cudagraph_capture_size: 0 + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker + # final_output: true + # final_output_type: text + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 4096 + seed: 42 + detokenize: False + repetition_penalty: 1.05 + stop_token_ids: [2150] + +- stage_id: 2 + stage_type: llm # Use llm stage type to launch OmniLLM + runtime: + devices: "5" + max_batch_size: 1 + engine_args: + model_stage: code2wav + model_arch: Qwen3OmniMoeForConditionalGeneration + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + engine_output_type: audio # Final output: audio waveform + gpu_memory_utilization: 0.3 # code2wav weight is around 0.4GB for Qwen3-Omni-30B-A3B-Instruct + skip_mm_profiling: true + distributed_executor_backend: "mp" + max_num_batched_tokens: 100000 + hf_config_name: thinker_config + async_scheduling: false + max_cudagraph_capture_size: 0 + engine_input_source: [1] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav + final_output: true + final_output_type: audio + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 2000 + seed: 42 + detokenize: True + repetition_penalty: 1.1 diff --git a/tests/examples/online_serving/test_qwen2_5_omni.py b/tests/examples/online_serving/test_qwen2_5_omni.py index 58b011f19..6d6ebd318 100644 --- a/tests/examples/online_serving/test_qwen2_5_omni.py +++ b/tests/examples/online_serving/test_qwen2_5_omni.py @@ -4,6 +4,8 @@ Example online tests for Qwen2.5-Omni-7B model. import os +from vllm_omni.platforms import current_omni_platform + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" import re @@ -20,6 +22,11 @@ models = ["Qwen/Qwen2.5-Omni-7B"] stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_ci.yaml")] +if current_omni_platform.is_xpu(): + stage_configs = [ + str(Path(__file__).parent.parent.parent / "e2e" / "stage_configs" / "xpu" / "qwen2_5_omni_ci.yaml") + ] + example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen2_5_omni") # Create parameter combinations for model and stage config test_params = [ diff --git a/tests/examples/online_serving/test_qwen3_omni.py b/tests/examples/online_serving/test_qwen3_omni.py index 568f1f9b4..abe5322dd 100644 --- a/tests/examples/online_serving/test_qwen3_omni.py +++ b/tests/examples/online_serving/test_qwen3_omni.py @@ -4,6 +4,8 @@ Example Online tests for Qwen3-Omni model. import os +from vllm_omni.platforms import current_omni_platform + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" import re @@ -20,6 +22,10 @@ models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_configs" / "qwen3_omni_ci.yaml")] +if current_omni_platform.is_xpu(): + stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")] + + example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen3_omni") # Create parameter combinations for model and stage config test_params = [ diff --git a/tests/utils.py b/tests/utils.py index e8a9d5c05..72fc6639a 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -402,13 +402,41 @@ def rocm_marks(*, res: str, num_cards: int): return marks + [test_distributed] +def xpu_marks(*, res: str, num_cards: int): + """ + Get a collection of pytest marks to apply for `@xpu_test`. + + Args: + res: Resource type, e.g., "B60". + num_cards: Number of GPU cards required. + + Returns: + List of pytest marks to apply. + """ + test_platform_detail = pytest.mark.xpu + + if res == "B60": + test_resource = pytest.mark.B60 + else: + raise ValueError(f"Invalid XPU resource type: {res}. Supported: B60") + + marks = [test_resource, test_platform_detail] + + if num_cards == 1: + return marks + else: + test_distributed = pytest.mark.distributed_rocm(num_cards=num_cards) + # TODO: add XPU support for `skipif_xpu` marker + return marks + [test_distributed] + + def gpu_marks(*, res: str, num_cards: int): """ Get a collection of pytest marks to apply for `@gpu_test`. Platform is automatically determined based on resource type. Args: - res: Resource type, e.g., "L4", "H100" for CUDA, or "MI325" for ROCm. + res: Resource type, e.g., "L4", "H100" for CUDA, or "MI325" for ROCm, or "B60" for XPU. num_cards: Number of GPU cards required. Returns: @@ -419,6 +447,8 @@ def gpu_marks(*, res: str, num_cards: int): return [test_platform] + cuda_marks(res=res, num_cards=num_cards) if res == "MI325": return [test_platform] + rocm_marks(res=res, num_cards=num_cards) + if res == "B60": + return [test_platform] + xpu_marks(res=res, num_cards=num_cards) raise ValueError(f"Invalid resource type: {res}. Supported: L4, H100, MI325") @@ -446,13 +476,13 @@ def npu_marks(*, res: str, num_cards: int): def hardware_marks(*, res: dict[str, str], num_cards: int | dict[str, int] = 1): """ Get a collection of pytest marks to apply for `@hardware_test`, - including CUDA, ROCm, and NPU, + including CUDA, ROCm, XPU, and NPU, based on the specified platforms and resources. """ # Validate platforms # Don't validate platform details in this decorator for platform, _ in res.items(): - if platform not in ("cuda", "rocm", "npu"): + if platform not in ("cuda", "rocm", "xpu", "npu"): raise ValueError(f"Unsupported platform: {platform}") # Normalize num_cards @@ -473,7 +503,7 @@ def hardware_marks(*, res: dict[str, str], num_cards: int | dict[str, int] = 1): all_marks: list[pytest.MarkDecorator] = [] for platform, resource in res.items(): cards = num_cards_dict[platform] - if platform == "cuda" or platform == "rocm": + if platform == "cuda" or platform == "rocm" or platform == "xpu": marks = gpu_marks(res=resource, num_cards=cards) elif platform == "npu": marks = npu_marks(res=resource, num_cards=cards) From dc72aede400e263031158cd76da7e1d8c78226f4 Mon Sep 17 00:00:00 2001 From: zhou zhuoxin Date: Sat, 14 Mar 2026 14:15:25 +0800 Subject: [PATCH 03/10] [CI] add multimodal processing correctness tests for Omni models (#1445) Signed-off-by: zhou zhuoxin Co-authored-by: princepride Co-authored-by: Ding Zuhao --- tests/model_executor/models/registry.py | 179 ++++++++++ .../models/test_omni_processing.py | 327 ++++++++++++++++++ 2 files changed, 506 insertions(+) create mode 100644 tests/model_executor/models/registry.py create mode 100644 tests/model_executor/models/test_omni_processing.py diff --git a/tests/model_executor/models/registry.py b/tests/model_executor/models/registry.py new file mode 100644 index 000000000..98f9aefb5 --- /dev/null +++ b/tests/model_executor/models/registry.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +""" +Test metadata for Omni models. + +Mirrors vllm's ``tests/models/registry.py`` (_HfExamplesInfo) but adapted for +the Omni multi-stage architecture. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +import pytest +from packaging.version import Version +from transformers import __version__ as TRANSFORMERS_VERSION + + +@dataclass(frozen=True) +class _OmniExamplesInfo: + """Metadata used by processing-correctness tests for Omni models.""" + + default: str + """The default HF model ID to use for testing this architecture.""" + + model_stage: str = "thinker" + """Stage type identifier, e.g. ``"thinker"`` or ``"talker"``.""" + + hf_config_name: str | None = None + """ + Name of the sub-config inside the top-level HF config + (e.g. ``"thinker_config"``). ``None`` means use the root config. + """ + + has_multimodal_processor: bool = False + """Whether the architecture registers a multimodal processor.""" + + min_transformers_version: str | None = None + """Minimum ``transformers`` version required to load this model.""" + + max_transformers_version: str | None = None + """Maximum ``transformers`` version that this model supports.""" + + is_available_online: bool = True + """Set to ``False`` if the model is no longer hosted on HF Hub.""" + + trust_remote_code: bool = False + """Whether ``trust_remote_code=True`` is needed to load the model.""" + + max_model_len: int | None = None + """ + An explicit ``max_model_len`` override. Useful when the default value + from the HF config is too large for CI. + """ + + # ----- helper methods (same API as _HfExamplesInfo) -------------------- + + def check_transformers_version( + self, + *, + on_fail: Literal["error", "skip", "return"], + ) -> str | None: + """Skip / error / return a message when the installed + ``transformers`` does not satisfy the version constraints.""" + if self.min_transformers_version is None and self.max_transformers_version is None: + return None + + current = TRANSFORMERS_VERSION + cur_base = Version(current).base_version + min_ver = self.min_transformers_version + max_ver = self.max_transformers_version + + msg = f"`transformers=={current}` installed, but `transformers" + if min_ver and Version(cur_base) < Version(min_ver): + msg += f">={min_ver}` is required to run this model." + elif max_ver and Version(cur_base) > Version(max_ver): + msg += f"<={max_ver}` is required to run this model." + else: + return None + + if on_fail == "error": + raise RuntimeError(msg) + elif on_fail == "skip": + pytest.skip(msg) + return msg + + def check_available_online( + self, + *, + on_fail: Literal["error", "skip"], + ) -> None: + """Skip / error when the model is not available online.""" + if not self.is_available_online: + msg = "Model is not available online" + if on_fail == "error": + raise RuntimeError(msg) + else: + pytest.skip(msg) + + +# --------------------------------------------------------------------------- +# Example model registry for tests +# --------------------------------------------------------------------------- + +_OMNI_EXAMPLE_MODELS: dict[str, _OmniExamplesInfo] = { + # ---- Qwen2.5-Omni ---- + "Qwen2_5OmniForConditionalGeneration": _OmniExamplesInfo( + default="Qwen/Qwen2.5-Omni-7B", + model_stage="thinker", + has_multimodal_processor=True, + ), + "Qwen2_5OmniThinkerModel": _OmniExamplesInfo( + default="Qwen/Qwen2.5-Omni-7B", + model_stage="thinker", + hf_config_name="thinker_config", + has_multimodal_processor=True, + ), + "Qwen2_5OmniTalkerModel": _OmniExamplesInfo( + default="Qwen/Qwen2.5-Omni-7B", + model_stage="talker", + hf_config_name="talker_config", + has_multimodal_processor=True, + ), + "Qwen2_5OmniToken2WavModel": _OmniExamplesInfo( + default="Qwen/Qwen2.5-Omni-7B", + model_stage="token2wav", + ), + "Qwen2_5OmniToken2WavDiTModel": _OmniExamplesInfo( + default="Qwen/Qwen2.5-Omni-7B", + model_stage="token2wav", + ), + # Internal sub-model key hardcoded by Qwen2_5OmniTalkerForConditionalGeneration + # to load its language model component (see qwen2_5_omni_talker.py). + "Qwen2ForCausalLM_old": _OmniExamplesInfo( + default="Qwen/Qwen2.5-Omni-7B", + model_stage="thinker", + ), + # ---- Qwen3-Omni MoE ---- + "Qwen3OmniMoeForConditionalGeneration": _OmniExamplesInfo( + default="Qwen/Qwen3-Omni-30B-A3B-Instruct", + model_stage="thinker", + has_multimodal_processor=True, + ), + "Qwen3OmniMoeThinkerForConditionalGeneration": _OmniExamplesInfo( + default="Qwen/Qwen3-Omni-30B-A3B-Instruct", + model_stage="thinker", + hf_config_name="thinker_config", + has_multimodal_processor=True, + ), + "Qwen3OmniMoeTalkerForConditionalGeneration": _OmniExamplesInfo( + default="Qwen/Qwen3-Omni-30B-A3B-Instruct", + model_stage="talker", + hf_config_name="talker_config", + ), + "Qwen3OmniMoeCode2Wav": _OmniExamplesInfo( + default="Qwen/Qwen3-Omni-30B-A3B-Instruct", + model_stage="code2wav", + ), + # ---- Qwen3-TTS ---- + "Qwen3TTSForConditionalGeneration": _OmniExamplesInfo( + default="Qwen/Qwen3-TTS", + model_stage="talker", + ), + "Qwen3TTSTalkerForConditionalGeneration": _OmniExamplesInfo( + default="Qwen/Qwen3-TTS", + model_stage="talker", + ), + "Qwen3TTSCode2Wav": _OmniExamplesInfo( + default="Qwen/Qwen3-TTS", + model_stage="code2wav", + ), +} + +# Only architectures that register a multimodal processor are relevant for +# processing-correctness tests. +_MULTIMODAL_OMNI_EXAMPLE_MODELS: dict[str, _OmniExamplesInfo] = { + arch: info for arch, info in _OMNI_EXAMPLE_MODELS.items() if info.has_multimodal_processor +} diff --git a/tests/model_executor/models/test_omni_processing.py b/tests/model_executor/models/test_omni_processing.py new file mode 100644 index 000000000..fe9c63b82 --- /dev/null +++ b/tests/model_executor/models/test_omni_processing.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +""" +Test multimodal processing correctness for Omni models. + +Tests that: +1. Cached vs non-cached processor outputs are identical. +2. Text prompt vs token prompt inputs produce identical results after processing. + +Adapted from vllm/tests/models/multimodal/processing/test_common.py +""" + +from functools import partial + +import numpy as np +import pytest +from PIL import Image +from vllm.config.multimodal import ( + AudioDummyOptions, + BaseDummyOptions, + ImageDummyOptions, + VideoDummyOptions, +) +from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalDataDict +from vllm.multimodal.cache import MultiModalProcessorOnlyCache +from vllm.multimodal.inputs import MultiModalInputs, batched_tensors_equal +from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext +from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config + +from tests.model_executor.models.registry import ( + _MULTIMODAL_OMNI_EXAMPLE_MODELS, + _OmniExamplesInfo, +) +from vllm_omni.config import OmniModelConfig + + +def random_image(rng: np.random.RandomState, min_wh: int, max_wh: int): + w, h = rng.randint(min_wh, max_wh, size=(2,)) + arr = rng.randint(0, 255, size=(w, h, 3), dtype=np.uint8) + return Image.fromarray(arr) + + +def random_video( + rng: np.random.RandomState, + min_frames: int, + max_frames: int, + min_wh: int, + max_wh: int, +): + num_frames = rng.randint(min_frames, max_frames) + w, h = rng.randint(min_wh, max_wh, size=(2,)) + return rng.randint(0, 255, size=(num_frames, w, h, 3), dtype=np.uint8) + + +def random_audio( + rng: np.random.RandomState, + min_len: int, + max_len: int, + sr: int, +): + audio_len = rng.randint(min_len, max_len) + return rng.rand(audio_len), sr + + +_IGNORE_MM_KEYS: dict[str, set[str]] = {} + + +def _get_model_archs_to_test() -> list[str]: + """Return architecture strings for parametrization.""" + return list(_MULTIMODAL_OMNI_EXAMPLE_MODELS.keys()) + + +def get_text_token_prompts( + processor: BaseMultiModalProcessor, + mm_data: MultiModalDataDict, +): + """Return ``(text_prompt, token_prompt)`` for the given multimodal data.""" + dummy_inputs = processor.dummy_inputs + tokenizer: TokenizerLike = processor.info.get_tokenizer() + model_config = processor.info.ctx.model_config + + parsed_data = processor.info.parse_mm_data(mm_data) + mm_counts = {k: len(vs) for k, vs in parsed_data.items()} + + inputs = dummy_inputs.get_dummy_processor_inputs( + model_config.max_model_len, + mm_counts, + mm_options={}, + ) + + text_prompt: str | None + token_prompt: list[int] + if isinstance(inputs.prompt, list): + text_prompt = None + token_prompt = inputs.prompt + elif isinstance(inputs.prompt, str): + text_prompt = inputs.prompt + token_prompt = tokenizer.encode( + text_prompt, + **processor.info.get_default_tok_params().get_encode_kwargs(), + ) + else: + raise TypeError(type(inputs.prompt)) + + return text_prompt, token_prompt + + +def _build_model_config(model_arch: str, info: _OmniExamplesInfo) -> OmniModelConfig: + """Create an ``OmniModelConfig`` suitable for processor testing.""" + kwargs: dict = dict( + model=info.default, + tokenizer=info.default, + tokenizer_mode="auto", + trust_remote_code=info.trust_remote_code, + model_arch=model_arch, + model_stage=info.model_stage, + max_model_len=info.max_model_len, + enforce_eager=True, + dtype="auto", + ) + if info.hf_config_name is not None: + kwargs["hf_config_name"] = info.hf_config_name + + model_config = OmniModelConfig(**kwargs) + # Ensure cache is effectively unlimited to prevent eviction during test + model_config.multimodal_config.mm_processor_cache_gb = 2048 + return model_config + + +def _test_processing_correctness( + model_arch: str, + hit_rate: float, + num_batches: int, + simplify_rate: float, +): + """Run the full processing-correctness test for one Omni model.""" + info = _MULTIMODAL_OMNI_EXAMPLE_MODELS[model_arch] + info.check_transformers_version(on_fail="skip") + info.check_available_online(on_fail="skip") + + model_config = _build_model_config(model_arch, info) + + model_cls = MULTIMODAL_REGISTRY._get_model_cls(model_config) + factories = model_cls._processor_factory + ctx = InputProcessingContext( + model_config, + tokenizer=cached_tokenizer_from_config(model_config), + ) + cache = MultiModalProcessorOnlyCache(model_config) + + processing_info = factories.info(ctx) + supported_mm_limits = processing_info.get_supported_mm_limits() + limit_mm_per_prompt_ints = { + modality: 3 if limit is None else limit for modality, limit in supported_mm_limits.items() + } + + def _to_dummy_options(modality: str, count: int) -> BaseDummyOptions: + if modality == "video": + return VideoDummyOptions(count=count) + if modality == "image": + return ImageDummyOptions(count=count) + if modality == "audio": + return AudioDummyOptions(count=count) + return BaseDummyOptions(count=count) + + model_config.get_multimodal_config().limit_per_prompt = { + modality: _to_dummy_options(modality, count) for modality, count in limit_mm_per_prompt_ints.items() + } + + baseline_processor = factories.build_processor(ctx, cache=None) + cached_processor = factories.build_processor(ctx, cache=cache) + + rng = np.random.RandomState(0) + + input_to_hit = { + "image": Image.new("RGB", size=(128, 128)), + "video": np.zeros((4, 128, 128, 3), dtype=np.uint8), + "audio": (np.zeros((512,)), 16000), + } + input_factory = { + "image": partial(random_image, rng, min_wh=128, max_wh=256), + "video": partial(random_video, rng, min_frames=2, max_frames=16, min_wh=128, max_wh=256), + "audio": partial(random_audio, rng, min_len=512, max_len=1024, sr=16000), + } + + for batch_idx in range(num_batches): + mm_data = { + k: [ + (input_to_hit[k] if rng.rand() < hit_rate else input_factory[k]()) + for _ in range(rng.randint(limit + 1)) + ] + for k, limit in limit_mm_per_prompt_ints.items() + } + + # Drop unnecessary keys and test single -> multi conversion + if rng.rand() < simplify_rate: + for k in list(mm_data.keys()): + if not mm_data[k]: + del mm_data[k] + elif len(mm_data[k]) == 1: + mm_data[k] = mm_data[k][0] + + _test_processing_correctness_one( + model_config, + mm_data, + baseline_processor, + cached_processor, + batch_idx, + ) + + +def _test_processing_correctness_one( + model_config: OmniModelConfig, + mm_data: MultiModalDataDict, + baseline_processor: BaseMultiModalProcessor, + cached_processor: BaseMultiModalProcessor, + batch_idx: int, +): + model_type = model_config.hf_config.model_type + + text_prompt, token_prompt = get_text_token_prompts(baseline_processor, mm_data) + mm_items = baseline_processor.info.parse_mm_data(mm_data) + ignore_mm_keys = _IGNORE_MM_KEYS.get(model_type, set[str]()) + + baseline_tokenized_result = baseline_processor( + token_prompt, + mm_items=mm_items, + hf_processor_mm_kwargs={}, + ) + + cached_tokenized_result = cached_processor( + token_prompt, + mm_items=mm_items, + hf_processor_mm_kwargs={}, + ) + + _assert_inputs_equal( + baseline_tokenized_result, + cached_tokenized_result, + ignore_mm_keys=ignore_mm_keys, + msg=f"Failed ({batch_idx=}, {token_prompt=}, {mm_data=})", + ) + + if text_prompt is not None: + baseline_text_result = baseline_processor( + text_prompt, + mm_items=mm_items, + hf_processor_mm_kwargs={}, + ) + cached_text_result = cached_processor( + text_prompt, + mm_items=mm_items, + hf_processor_mm_kwargs={}, + ) + + _assert_inputs_equal( + baseline_text_result, + cached_text_result, + ignore_mm_keys=ignore_mm_keys, + msg=f"Failed ({batch_idx=}, {text_prompt=}, {mm_data=})", + ) + + _assert_inputs_equal( + baseline_text_result, + baseline_tokenized_result, + ignore_mm_keys=ignore_mm_keys, + msg=f"Failed ({batch_idx=}, {text_prompt=}, {token_prompt=}, {mm_data=})", + ) + + _assert_inputs_equal( + cached_text_result, + cached_tokenized_result, + ignore_mm_keys=ignore_mm_keys, + msg=f"Failed ({batch_idx=}, {text_prompt=}, {token_prompt=}, {mm_data=})", + ) + + +@pytest.mark.core_model +@pytest.mark.omni +@pytest.mark.cpu +@pytest.mark.parametrize("model_arch", _get_model_archs_to_test()) +@pytest.mark.parametrize("hit_rate", [0.3, 0.5, 1.0]) +@pytest.mark.parametrize("num_batches", [32]) +@pytest.mark.parametrize("simplify_rate", [1.0]) +def test_omni_processing_correctness( + model_arch: str, + hit_rate: float, + num_batches: int, + simplify_rate: float, +): + """ + For each Omni model verify that: + - cached and non-cached processor outputs are identical. + - text-prompt and token-prompt inputs yield the same result. + """ + _test_processing_correctness( + model_arch, + hit_rate=hit_rate, + num_batches=num_batches, + simplify_rate=simplify_rate, + ) + + +def _assert_inputs_equal( + a: MultiModalInputs, + b: MultiModalInputs, + *, + ignore_mm_keys: set[str] | None = None, + msg: str = "", +): + if ignore_mm_keys is None: + ignore_mm_keys = set() + + ignore_prompt_keys = ("prompt", "mm_kwargs") + a_rest = {k: v for k, v in a.items() if k not in ignore_prompt_keys} + b_rest = {k: v for k, v in b.items() if k not in ignore_prompt_keys} + + assert a_rest == b_rest, msg + + a_data = a["mm_kwargs"].get_data() + b_data = b["mm_kwargs"].get_data() + + for key in ignore_mm_keys: + a_data.pop(key, None) + b_data.pop(key, None) + + assert batched_tensors_equal(a_data, b_data), msg From 5377619c0bd1953e785e1bccb14e8a0bb4058cdc Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Sat, 14 Mar 2026 17:07:13 +0800 Subject: [PATCH 04/10] fix: propagate parallel_config through create_default_diffusion (#1878) Signed-off-by: lishunyang Signed-off-by: Hongsheng Liu Co-authored-by: Gao Han Co-authored-by: Hongsheng Liu --- .../test_zimage_parallelism.py | 219 +++++++++--------- tests/test_diffusion_config_propagation.py | 108 +++++++++ vllm_omni/config/stage_config.py | 7 +- vllm_omni/diffusion/data.py | 9 +- 4 files changed, 225 insertions(+), 118 deletions(-) create mode 100644 tests/test_diffusion_config_propagation.py diff --git a/tests/e2e/offline_inference/test_zimage_parallelism.py b/tests/e2e/offline_inference/test_zimage_parallelism.py index b11176070..279710195 100644 --- a/tests/e2e/offline_inference/test_zimage_parallelism.py +++ b/tests/e2e/offline_inference/test_zimage_parallelism.py @@ -17,14 +17,12 @@ import time from pathlib import Path import numpy as np - -# import pytest +import pytest import torch from PIL import Image from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -# from tests.utils import DeviceMemoryMonitor, hardware_test -from tests.utils import DeviceMemoryMonitor +from tests.utils import DeviceMemoryMonitor, hardware_test from vllm_omni import Omni from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -37,7 +35,6 @@ if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) -# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" PROMPT = "a photo of a cat sitting on a laptop keyboard" @@ -157,124 +154,124 @@ def _run_zimage_generate( cleanup_dist_env_and_memory() -# @pytest.mark.advanced_model -# @pytest.mark.diffusion -# @pytest.mark.parallel -# @hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 2}) -# def test_zimage_tensor_parallel_tp2(tmp_path: Path): -# if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): -# pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.") -# if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: -# pytest.skip("Z-Image TP=2 requires >= 2 devices.") +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) +def test_zimage_tensor_parallel_tp2(tmp_path: Path): + if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): + pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.") + if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: + pytest.skip("Z-Image TP=2 requires >= 2 devices.") -# enforce_eager = False + enforce_eager = False -# height = 512 -# width = 512 -# num_inference_steps = 2 -# seed = 42 + height = 512 + width = 512 + num_inference_steps = 2 + seed = 42 -# tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate( -# tp_size=1, -# height=height, -# width=width, -# num_inference_steps=num_inference_steps, -# seed=seed, -# enforce_eager=enforce_eager, -# ) -# tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate( -# tp_size=2, -# height=height, -# width=width, -# num_inference_steps=num_inference_steps, -# seed=seed, -# enforce_eager=enforce_eager, -# ) + tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate( + tp_size=1, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=seed, + enforce_eager=enforce_eager, + ) + tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate( + tp_size=2, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=seed, + enforce_eager=enforce_eager, + ) -# tp1_path = tmp_path / "zimage_tp1.png" -# tp2_path = tmp_path / "zimage_tp2.png" -# tp1_img.save(tp1_path) -# tp2_img.save(tp2_path) + tp1_path = tmp_path / "zimage_tp1.png" + tp2_path = tmp_path / "zimage_tp2.png" + tp1_img.save(tp1_path) + tp2_img.save(tp2_path) -# assert tp1_img.width == width and tp1_img.height == height -# assert tp2_img.width == width and tp2_img.height == height + assert tp1_img.width == width and tp1_img.height == height + assert tp2_img.width == width and tp2_img.height == height -# mean_abs_diff, p99_abs_diff = _diff_metrics(tp1_img, tp2_img) -# mean_threshold = 3e-2 -# p99_threshold = 2.5e-1 -# print( -# "Z-Image TP image diff stats (TP=1 vs TP=2): " -# f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; " -# f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; " -# f"tp1_img={tp1_path}, tp2_img={tp2_path}" -# ) -# assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, ( -# f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} " -# f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" -# ) + mean_abs_diff, p99_abs_diff = _diff_metrics(tp1_img, tp2_img) + mean_threshold = 3e-2 + p99_threshold = 2.5e-1 + print( + "Z-Image TP image diff stats (TP=1 vs TP=2): " + f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; " + f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; " + f"tp1_img={tp1_path}, tp2_img={tp2_path}" + ) + assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, ( + f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} " + f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" + ) -# print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}") -# assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})" + print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}") + assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})" -# print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}") -# assert tp2_peak_mem < tp1_peak_mem, ( -# f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})" -# ) + print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}") + assert tp2_peak_mem < tp1_peak_mem, ( + f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})" + ) -# @pytest.mark.integration -# def test_zimage_vae_patch_parallel_tp2(tmp_path: Path): -# if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): -# pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA for now.") -# if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: -# pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 devices.") +@pytest.mark.integration +def test_zimage_vae_patch_parallel_tp2(tmp_path: Path): + if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): + pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA for now.") + if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: + pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 devices.") -# enforce_eager = False + enforce_eager = False -# # Use a larger image to ensure there are multiple VAE tiles. -# height = 1152 -# width = 1152 -# num_inference_steps = 2 -# seed = 42 + # Use a larger image to ensure there are multiple VAE tiles. + height = 1152 + width = 1152 + num_inference_steps = 2 + seed = 42 -# baseline_img, _baseline_time_s, _baseline_peak_mem = _run_zimage_generate( -# tp_size=2, -# height=height, -# width=width, -# num_inference_steps=num_inference_steps, -# seed=seed, -# enforce_eager=enforce_eager, -# vae_use_tiling=True, -# vae_patch_parallel_size=1, -# num_requests=2, -# ) -# pp2_img, _pp2_time_s, _pp2_peak_mem = _run_zimage_generate( -# tp_size=2, -# height=height, -# width=width, -# num_inference_steps=num_inference_steps, -# seed=seed, -# enforce_eager=enforce_eager, -# vae_use_tiling=True, -# vae_patch_parallel_size=2, -# num_requests=2, -# ) + baseline_img, _baseline_time_s, _baseline_peak_mem = _run_zimage_generate( + tp_size=2, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=seed, + enforce_eager=enforce_eager, + vae_use_tiling=True, + vae_patch_parallel_size=1, + num_requests=2, + ) + pp2_img, _pp2_time_s, _pp2_peak_mem = _run_zimage_generate( + tp_size=2, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=seed, + enforce_eager=enforce_eager, + vae_use_tiling=True, + vae_patch_parallel_size=2, + num_requests=2, + ) -# baseline_path = tmp_path / "zimage_tp2_vae_pp1.png" -# pp2_path = tmp_path / "zimage_tp2_vae_pp2.png" -# baseline_img.save(baseline_path) -# pp2_img.save(pp2_path) + baseline_path = tmp_path / "zimage_tp2_vae_pp1.png" + pp2_path = tmp_path / "zimage_tp2_vae_pp2.png" + baseline_img.save(baseline_path) + pp2_img.save(pp2_path) -# mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img) -# mean_threshold = 5e-3 -# p99_threshold = 1e-1 -# print( -# "Z-Image VAE patch parallel image diff stats (TP=2, pp=1 vs pp=2): " -# f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; " -# f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; " -# f"pp1_img={baseline_path}, pp2_img={pp2_path}" -# ) -# assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, ( -# f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} " -# f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" -# ) + mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img) + mean_threshold = 5e-3 + p99_threshold = 1e-1 + print( + "Z-Image VAE patch parallel image diff stats (TP=2, pp=1 vs pp=2): " + f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; " + f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; " + f"pp1_img={baseline_path}, pp2_img={pp2_path}" + ) + assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, ( + f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} " + f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" + ) diff --git a/tests/test_diffusion_config_propagation.py b/tests/test_diffusion_config_propagation.py new file mode 100644 index 000000000..58eb6097c --- /dev/null +++ b/tests/test_diffusion_config_propagation.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests that parallel_config survives the create_default_diffusion roundtrip. + +Regression tests for https://github.com/vllm-project/vllm-omni/issues/1862 +""" + +from collections.abc import Mapping + +import torch + +from vllm_omni.config.stage_config import StageConfigFactory +from vllm_omni.diffusion.data import ( + DiffusionParallelConfig, + OmniDiffusionConfig, +) + + +def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig: + """Simulate the real path: create_default_diffusion → OmniDiffusionConfig. + + Does NOT manually reconstruct parallel_config — relies on + OmniDiffusionConfig.__post_init__ to handle the dict, just like + the production code path does. + """ + stages = StageConfigFactory.create_default_diffusion(kwargs) + engine_args = dict(stages[0]["engine_args"]) + return OmniDiffusionConfig.from_kwargs(**engine_args) + + +class TestParallelConfigPropagation: + """Core regression tests: parallel_config must survive serialization.""" + + def test_tp2_roundtrip(self): + pc = DiffusionParallelConfig(tensor_parallel_size=2) + od = _roundtrip_diffusion_config(model="test-model", parallel_config=pc) + assert od.parallel_config.tensor_parallel_size == 2 + assert od.parallel_config.world_size == 2 + + def test_tp4_devices_and_config(self): + pc = DiffusionParallelConfig(tensor_parallel_size=4) + stages = StageConfigFactory.create_default_diffusion({"parallel_config": pc, "model": "x"}) + assert stages[0]["runtime"]["devices"] == "0,1,2,3" + + # Let __post_init__ reconstruct from dict (real code path) + ea = dict(stages[0]["engine_args"]) + od = OmniDiffusionConfig.from_kwargs(**ea) + assert od.parallel_config.tensor_parallel_size == 4 + assert od.parallel_config.world_size == 4 + + def test_sp_config_roundtrip(self): + pc = DiffusionParallelConfig( + tensor_parallel_size=2, + ulysses_degree=2, + ring_degree=1, + ) + od = _roundtrip_diffusion_config(model="x", parallel_config=pc) + assert od.parallel_config.ulysses_degree == 2 + assert od.parallel_config.ring_degree == 1 + + def test_cfg_parallel_roundtrip(self): + pc = DiffusionParallelConfig(cfg_parallel_size=2) + od = _roundtrip_diffusion_config(model="x", parallel_config=pc) + assert od.parallel_config.cfg_parallel_size == 2 + assert od.parallel_config.world_size == 2 + + def test_no_parallel_config_defaults_to_tp1(self): + od = _roundtrip_diffusion_config(model="x") + assert od.parallel_config.tensor_parallel_size == 1 + assert od.parallel_config.world_size == 1 + + def test_num_gpus_derived_from_world_size(self): + pc = DiffusionParallelConfig(tensor_parallel_size=2) + od = _roundtrip_diffusion_config(model="x", parallel_config=pc) + assert od.num_gpus == 2 + + +class TestCreateDefaultDiffusion: + """Verify engine_args structure from create_default_diffusion.""" + + def test_parallel_config_serialized_as_dict(self): + """The key fix: parallel_config must appear in engine_args as a dict.""" + pc = DiffusionParallelConfig(tensor_parallel_size=2) + stages = StageConfigFactory.create_default_diffusion({"model": "x", "parallel_config": pc}) + ea = stages[0]["engine_args"] + assert "parallel_config" in ea + assert isinstance(ea["parallel_config"], Mapping) + assert ea["parallel_config"]["tensor_parallel_size"] == 2 + + def test_dtype_serialized_as_string(self): + stages = StageConfigFactory.create_default_diffusion({"dtype": torch.float16, "model": "x"}) + assert stages[0]["engine_args"]["dtype"] == "torch.float16" + + def test_cache_backend_defaults_to_none(self): + stages = StageConfigFactory.create_default_diffusion({"model": "x"}) + assert stages[0]["engine_args"]["cache_backend"] == "none" + + def test_single_gpu_default_devices(self): + stages = StageConfigFactory.create_default_diffusion({"model": "x"}) + assert stages[0]["runtime"]["devices"] == "0" + + def test_extra_kwargs_forwarded(self): + stages = StageConfigFactory.create_default_diffusion( + {"model": "x", "enforce_eager": True, "lora_path": "/tmp/lora"} + ) + ea = stages[0]["engine_args"] + assert ea["enforce_eager"] is True + assert ea["lora_path"] == "/tmp/lora" diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index ad298aa9d..e457c3d99 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -11,7 +11,7 @@ Runtime parameters (gpu_memory_utilization, tp_size, etc.) come from CLI. from __future__ import annotations import re -from dataclasses import dataclass, field +from dataclasses import asdict, dataclass, field from enum import Enum from pathlib import Path from typing import Any @@ -279,13 +279,16 @@ class StageConfigFactory: for i in range(1, num_devices): devices += f",{i}" - # Collect engine args – skip non-serializable objects engine_args: dict[str, Any] = {} for key, value in kwargs.items(): if key in ("parallel_config",): continue engine_args[key] = value + # Serialize parallel_config as dict for OmegaConf compatibility + if "parallel_config" in kwargs: + engine_args["parallel_config"] = asdict(kwargs["parallel_config"]) + engine_args.setdefault("cache_backend", "none") engine_args["model_stage"] = "diffusion" diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 3e44f861b..b67736c81 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -514,12 +514,11 @@ class OmniDiffusionConfig: initial_master_port = (self.master_port or 30005) + random.randint(0, 100) self.master_port = self.settle_port(initial_master_port, 37) - # Convert parallel_config dict to DiffusionParallelConfig if needed - # This must be done before accessing parallel_config.world_size - if isinstance(self.parallel_config, dict): - self.parallel_config = DiffusionParallelConfig.from_dict(self.parallel_config) + # Convert parallel_config dict/DictConfig to DiffusionParallelConfig + # Use Mapping to handle both plain dicts and OmegaConf DictConfig + if isinstance(self.parallel_config, Mapping): + self.parallel_config = DiffusionParallelConfig.from_dict(dict(self.parallel_config)) elif not isinstance(self.parallel_config, DiffusionParallelConfig): - # If it's neither dict nor DiffusionParallelConfig, use default config self.parallel_config = DiffusionParallelConfig() if self.num_gpus is None: From 365dfd5f8e00b471b785cba70492c7990671c8ce Mon Sep 17 00:00:00 2001 From: Gao Han Date: Sat, 14 Mar 2026 22:56:43 +0800 Subject: [PATCH 05/10] [CI pipeline] Re-enable Diffusion Tensor Parallelism Test in pipeline (#1892) Signed-off-by: Gao Han --- .buildkite/test-merge.yml | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 7fa510716..5479f8ac1 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -119,23 +119,23 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" - # - label: "Diffusion Tensor Parallelism Test" - # timeout_in_minutes: 20 - # depends_on: upload-merge-pipeline - # commands: - # - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py - # agents: - # queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU - # plugins: - # - docker#v5.2.0: - # image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - # always-pull: true - # propagate-environment: true - # shm-size: "8gb" - # environment: - # - "HF_HOME=/fsx/hf_cache" - # volumes: - # - "/fsx/hf_cache:/fsx/hf_cache" + - label: "Diffusion Tensor Parallelism Test" + timeout_in_minutes: 20 + depends_on: upload-merge-pipeline + commands: + - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py + agents: + queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" - label: "Diffusion GPU Worker Test" timeout_in_minutes: 20 From 4de0c8d41bdf565a8a23d6ba4661b5470e0e6612 Mon Sep 17 00:00:00 2001 From: Wu JIAZHEN <83007646+asukaqaq-s@users.noreply.github.com> Date: Sat, 14 Mar 2026 23:31:18 +0800 Subject: [PATCH 06/10] [skip CI][Docs][Benchmark]: clarify vbench parameter behavior and add t2v example (#1497) Signed-off-by: asuka <1311722138@qq.com> --- benchmarks/diffusion/README.md | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/benchmarks/diffusion/README.md b/benchmarks/diffusion/README.md index 3eb775fe2..06bf01726 100644 --- a/benchmarks/diffusion/README.md +++ b/benchmarks/diffusion/README.md @@ -43,6 +43,27 @@ The benchmark supports three dataset modes via `--dataset`: ### VBench dataset +`vbench` only provides prompt data (and image paths for i2v/i2i); it does not carry +per-request generation fields. In this mode, all requests share CLI values: +`--width --height --num-frames --fps --num-inference-steps` +(pass `--width` and `--height` together). + +Example (`t2v`): + +```bash +python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ + --base-url http://localhost:8099 \ + --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --task t2v \ + --dataset vbench \ + --num-prompts 50 \ + --width 640 --height 480 \ + --num-frames 81 --fps 16 \ + --num-inference-steps 40 +``` + +Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths). + If you use i2v/i2i bench datasets and need auto-download support, you may need: ```bash @@ -90,7 +111,7 @@ Common optional flags: Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`. - For `vbench` / `random`: these CLI flags act as global defaults for all generated requests. -- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`). +- For `trace`: requests can carry their own fields (e.g., `width/height/num_frames/num_inference_steps`), with overrides/fallbacks as below. Precedence rules for `trace` (i.e., what actually gets sent): From af012babf8e1b5a272806566b423cddc9a681737 Mon Sep 17 00:00:00 2001 From: Lancer Date: Sun, 15 Mar 2026 02:55:30 +0800 Subject: [PATCH 07/10] [Bugfix] Fix cpu offload and quantization compatibility (#1473) Signed-off-by: Lancer --- vllm_omni/diffusion/model_loader/diffusers_loader.py | 6 ++++++ vllm_omni/diffusion/worker/diffusion_model_runner.py | 1 + vllm_omni/entrypoints/async_omni.py | 1 + 3 files changed, 8 insertions(+) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 78d9ffd03..daccef343 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -256,8 +256,14 @@ class DiffusersPipelineLoader: load_device: str, load_format: str = "default", custom_pipeline_name: str | None = None, + device: torch.device | None = None, ) -> nn.Module: """Load a model with the given configurations.""" + # CPU offload + FP8: load weights on device for FP8 quantization + if load_device == "cpu" and od_config.quantization and od_config.quantization.lower() != "none": + load_device = device.type + logger.info(f"Quantization enabled with CPU offload, using {load_device} for weight loading") + target_device = torch.device(load_device) with set_default_torch_dtype(od_config.dtype): if od_config.parallel_config.use_hsdp: diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 105629f34..accb173e1 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -127,6 +127,7 @@ class DiffusionModelRunner: load_device=load_device, load_format=load_format, custom_pipeline_name=custom_pipeline_name, + device=self.device, ) time_after_load = time.perf_counter() diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index b98679033..474834d16 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -225,6 +225,7 @@ class AsyncOmni(OmniBase): "enforce_eager": kwargs.get("enforce_eager", False), "diffusion_load_format": kwargs.get("diffusion_load_format", "default"), "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), + "quantization": kwargs.get("quantization", None), "worker_extension_cls": kwargs.get("worker_extension_cls", None), "enable_sleep_mode": kwargs.get("enable_sleep_mode", False), "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), From 2472c360aa8d1bf6e64036d6c23b76ab728f6468 Mon Sep 17 00:00:00 2001 From: Lancer Date: Sun, 15 Mar 2026 12:18:12 +0800 Subject: [PATCH 08/10] [Feat] support SP for FLUX.2-klein (#1250) Signed-off-by: Lancer --- .../diffusion/parallelism_acceleration.md | 2 +- vllm_omni/diffusion/layers/rope.py | 26 ++ .../diffusion/models/flux/flux_transformer.py | 9 +- .../flux2_klein/flux2_klein_transformer.py | 289 +++++++++++++++--- .../models/z_image/z_image_transformer.py | 7 +- 5 files changed, 280 insertions(+), 53 deletions(-) diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index 38cd5e67c..37fd5b1d1 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -35,7 +35,7 @@ The following table shows which models are currently supported by parallelism me | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | | **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | N/A | | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ | N/A | -| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | +| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | | **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | N/A | | **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | | **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 461c25652..65d37d0b0 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -157,3 +157,29 @@ class RotaryEmbedding(CustomOp): sin, interleaved=self.interleaved, ) + + +def apply_rope_to_qk( + rope: RotaryEmbedding, + query: torch.Tensor, + key: torch.Tensor, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply rotary positional embeddings to query and key tensors. + + Args: + rope: RotaryEmbedding instance for applying position embeddings + query: Query tensor [B, S, H, D] + key: Key tensor [B, S, H, D] + image_rotary_emb: Tuple of (cos, sin) tensors or None + + Returns: + Tuple of (query, key) with RoPE applied if rotary embeddings provided + """ + if image_rotary_emb is not None: + cos, sin = image_rotary_emb + cos = cos.to(query.dtype) + sin = sin.to(query.dtype) + query = rope(query, cos, sin) + key = rope(key, cos, sin) + return query, key diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py index db6f0d34e..2979fd4f6 100644 --- a/vllm_omni/diffusion/models/flux/flux_transformer.py +++ b/vllm_omni/diffusion/models/flux/flux_transformer.py @@ -30,7 +30,7 @@ if TYPE_CHECKING: from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.data import OmniDiffusionConfig -from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk logger = init_logger(__name__) @@ -224,12 +224,7 @@ class FluxAttention(torch.nn.Module): key = torch.cat([encoder_key, key], dim=1) value = torch.cat([encoder_value, value], dim=1) - if image_rotary_emb is not None: - cos, sin = image_rotary_emb # [S, D/2] - cos = cos.to(query.dtype) - sin = sin.to(query.dtype) - query = self.rope(query, cos, sin) - key = self.rope(key, cos, sin) + query, key = apply_rope_to_qk(self.rope, query, key, image_rotary_emb) # [S, D/2] hidden_states = self.attn( query, diff --git a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py index 2499b9c16..c10f06751 100644 --- a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py +++ b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py @@ -28,6 +28,7 @@ from diffusers.models.embeddings import ( ) from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import AdaLayerNormContinuous +from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -39,8 +40,15 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention -from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, + SequenceParallelOutput, +) +from vllm_omni.diffusion.forward_context import get_forward_context +from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk +logger = init_logger(__name__) if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -98,6 +106,7 @@ class Flux2FeedForward(nn.Module): class Flux2Attention(nn.Module): def __init__( self, + parallel_config: DiffusionParallelConfig, query_dim: int, heads: int = 8, dim_head: int = 64, @@ -112,6 +121,7 @@ class Flux2Attention(nn.Module): quant_config: "QuantizationConfig | None" = None, ): super().__init__() + self.parallel_config = parallel_config self.head_dim = dim_head self.inner_dim = out_dim if out_dim is not None else dim_head * heads self.query_dim = query_dim @@ -208,33 +218,87 @@ class Flux2Attention(nn.Module): encoder_query = self.norm_added_q(encoder_query) encoder_key = self.norm_added_k(encoder_key) - query = torch.cat([encoder_query, query], dim=1) - key = torch.cat([encoder_key, key], dim=1) - value = torch.cat([encoder_value, value], dim=1) + sp_size = self.parallel_config.sequence_parallel_size + forward_ctx = get_forward_context() + use_sp_joint_attention = sp_size is not None and sp_size > 1 and not forward_ctx.split_text_embed_in_sp - if image_rotary_emb is not None: - cos, sin = image_rotary_emb - cos = cos.to(query.dtype) - sin = sin.to(query.dtype) - query = self.rope(query, cos, sin) - key = self.rope(key, cos, sin) + if use_sp_joint_attention and image_rotary_emb is not None: + cos, sin = image_rotary_emb + cos = cos.to(query.dtype) + sin = sin.to(query.dtype) + txt_len = encoder_query.shape[1] + txt_cos, img_cos = cos[:txt_len], cos[txt_len:] + txt_sin, img_sin = sin[:txt_len], sin[txt_len:] + query = self.rope(query, img_cos, img_sin) + key = self.rope(key, img_cos, img_sin) + encoder_query = self.rope(encoder_query, txt_cos, txt_sin) + encoder_key = self.rope(encoder_key, txt_cos, txt_sin) - attn_metadata = None - if attention_mask is not None: - if attention_mask.dim() == 3: - attention_mask = attention_mask.unsqueeze(1) - attn_metadata = AttentionMetadata(attn_mask=attention_mask) + attn_metadata = AttentionMetadata( + joint_query=encoder_query, + joint_key=encoder_key, + joint_value=encoder_value, + joint_strategy="front", + ) + hidden_states_mask: torch.Tensor | None = kwargs.get("hidden_states_mask", None) + encoder_hidden_states_mask: torch.Tensor | None = kwargs.get("encoder_hidden_states_mask", None) + if hidden_states_mask is not None: + attn_metadata.attn_mask = hidden_states_mask + if encoder_hidden_states_mask is not None: + attn_metadata.joint_attn_mask = encoder_hidden_states_mask - hidden_states = self.attn(query, key, value, attn_metadata) - hidden_states = hidden_states.flatten(2, 3).to(query.dtype) + hidden_states = self.attn(query, key, value, attn_metadata) + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) - if encoder_hidden_states is not None: - context_len = encoder_hidden_states.shape[1] - encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( - [context_len, hidden_states.shape[1] - context_len], - dim=1, - ) - encoder_hidden_states = self.to_add_out(encoder_hidden_states) + txt_len = encoder_hidden_states.shape[1] + encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( + [txt_len, hidden_states.shape[1] - txt_len], + dim=1, + ) + encoder_hidden_states = self.to_add_out(encoder_hidden_states) + else: + query = torch.cat([encoder_query, query], dim=1) + key = torch.cat([encoder_key, key], dim=1) + value = torch.cat([encoder_value, value], dim=1) + + if image_rotary_emb is not None: + cos, sin = image_rotary_emb + cos = cos.to(query.dtype) + sin = sin.to(query.dtype) + query = self.rope(query, cos, sin) + key = self.rope(key, cos, sin) + + attn_metadata = None + if attention_mask is not None: + if attention_mask.dim() == 3: + attention_mask = attention_mask.unsqueeze(1) + attn_metadata = AttentionMetadata(attn_mask=attention_mask) + + hidden_states = self.attn(query, key, value, attn_metadata) + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) + + context_len = encoder_hidden_states.shape[1] + encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( + [context_len, hidden_states.shape[1] - context_len], + dim=1, + ) + encoder_hidden_states = self.to_add_out(encoder_hidden_states) + else: + if image_rotary_emb is not None: + cos, sin = image_rotary_emb + cos = cos.to(query.dtype) + sin = sin.to(query.dtype) + query = self.rope(query, cos, sin) + key = self.rope(key, cos, sin) + + attn_metadata = None + if attention_mask is not None: + if attention_mask.dim() == 3: + attention_mask = attention_mask.unsqueeze(1) + attn_metadata = AttentionMetadata(attn_mask=attention_mask) + + hidden_states = self.attn(query, key, value, attn_metadata) + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) hidden_states = self.to_out[0](hidden_states) hidden_states = self.to_out[1](hidden_states) @@ -251,6 +315,7 @@ class Flux2ParallelSelfAttention(nn.Module): def __init__( self, + parallel_config: DiffusionParallelConfig, query_dim: int, heads: int = 8, dim_head: int = 64, @@ -265,6 +330,7 @@ class Flux2ParallelSelfAttention(nn.Module): quant_config: "QuantizationConfig | None" = None, ): super().__init__() + self.parallel_config = parallel_config self.head_dim = dim_head self.inner_dim = out_dim if out_dim is not None else dim_head * heads self.query_dim = query_dim @@ -325,20 +391,57 @@ class Flux2ParallelSelfAttention(nn.Module): query = self.norm_q(query) key = self.norm_k(key) - if image_rotary_emb is not None: + sp_size = self.parallel_config.sequence_parallel_size + forward_ctx = get_forward_context() + text_seq_len = kwargs.get("text_seq_len", None) + use_sp_single_stream = ( + sp_size is not None and sp_size > 1 and not forward_ctx.split_text_embed_in_sp and text_seq_len is not None + ) + + if use_sp_single_stream and image_rotary_emb is not None: cos, sin = image_rotary_emb cos = cos.to(query.dtype) sin = sin.to(query.dtype) - query = self.rope(query, cos, sin) - key = self.rope(key, cos, sin) + txt_cos, img_cos = cos[:text_seq_len], cos[text_seq_len:] + txt_sin, img_sin = sin[:text_seq_len], sin[text_seq_len:] - attn_metadata = None - if attention_mask is not None: - if attention_mask.dim() == 3: - attention_mask = attention_mask.unsqueeze(1) - attn_metadata = AttentionMetadata(attn_mask=attention_mask) + img_query = query[:, text_seq_len:] + img_key = key[:, text_seq_len:] + img_value = value[:, text_seq_len:] + text_query = query[:, :text_seq_len] + text_key = key[:, :text_seq_len] + text_value = value[:, :text_seq_len] + + img_query = self.rope(img_query, img_cos, img_sin) + img_key = self.rope(img_key, img_cos, img_sin) + text_query = self.rope(text_query, txt_cos, txt_sin) + text_key = self.rope(text_key, txt_cos, txt_sin) + + attn_metadata = AttentionMetadata( + joint_query=text_query, + joint_key=text_key, + joint_value=text_value, + joint_strategy="front", + ) + hidden_states_mask: torch.Tensor | None = kwargs.get("hidden_states_mask", None) + encoder_hidden_states_mask: torch.Tensor | None = kwargs.get("encoder_hidden_states_mask", None) + if hidden_states_mask is not None: + attn_metadata.attn_mask = hidden_states_mask + if encoder_hidden_states_mask is not None: + attn_metadata.joint_attn_mask = encoder_hidden_states_mask + + attn_output = self.attn(img_query, img_key, img_value, attn_metadata) + else: + query, key = apply_rope_to_qk(self.rope, query, key, image_rotary_emb) + + attn_metadata = None + if attention_mask is not None: + if attention_mask.dim() == 3: + attention_mask = attention_mask.unsqueeze(1) + attn_metadata = AttentionMetadata(attn_mask=attention_mask) + + attn_output = self.attn(query, key, value, attn_metadata) - attn_output = self.attn(query, key, value, attn_metadata) attn_output = attn_output.flatten(2, 3).to(query.dtype) mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states) @@ -350,6 +453,7 @@ class Flux2ParallelSelfAttention(nn.Module): class Flux2SingleTransformerBlock(nn.Module): def __init__( self, + parallel_config: DiffusionParallelConfig, dim: int, num_attention_heads: int, attention_head_dim: int, @@ -361,6 +465,7 @@ class Flux2SingleTransformerBlock(nn.Module): super().__init__() self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) self.attn = Flux2ParallelSelfAttention( + parallel_config=parallel_config, query_dim=dim, dim_head=attention_head_dim, heads=num_attention_heads, @@ -383,6 +488,13 @@ class Flux2SingleTransformerBlock(nn.Module): split_hidden_states: bool = False, text_seq_len: int | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """ + Forward pass for Flux2SingleTransformerBlock with SP support. + + In SP mode: image hidden_states is chunked (B, img_len/SP, D), + text encoder_hidden_states is full (B, txt_len, D). + The block concatenates them for joint attention. + """ if encoder_hidden_states is not None: text_seq_len = encoder_hidden_states.shape[1] hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) @@ -412,6 +524,7 @@ class Flux2SingleTransformerBlock(nn.Module): class Flux2TransformerBlock(nn.Module): def __init__( self, + parallel_config: DiffusionParallelConfig, dim: int, num_attention_heads: int, attention_head_dim: int, @@ -425,6 +538,7 @@ class Flux2TransformerBlock(nn.Module): self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) self.attn = Flux2Attention( + parallel_config=parallel_config, query_dim=dim, added_kv_proj_dim=dim, dim_head=attention_head_dim, @@ -525,6 +639,47 @@ class Flux2PosEmbed(nn.Module): return freqs_cos, freqs_sin +class Flux2RopePrepare(nn.Module): + """Prepares RoPE embeddings for sequence parallel. + + This module encapsulates the RoPE computation for Flux.2-klein. + For dual-stream attention, text components (outputs 0, 1) are replicated + across SP ranks, while image components (outputs 2, 3) are sharded. + + NOTE: The hidden_states projection is handled separately in forward() + so that _sp_plan can shard it at the root level. + """ + + def __init__(self, pos_embed: Flux2PosEmbed): + super().__init__() + self.pos_embed = pos_embed + + def forward( + self, + img_ids: torch.Tensor, + txt_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Compute RoPE embeddings for text and image sequences. + + Args: + img_ids: Image position IDs (img_seq_len, n_axes) + txt_ids: Text position IDs (txt_seq_len, n_axes) + + Returns: + Tuple of cosine / sine components for text & image + in the order: (txt_cos, txt_sin, img_cos, img_sin) + + NOTE: careful about output orders if this is refactored in the + future; we need to match the _sp_plan indices, since text + components (0 & 1) need to be replicated across SP ranks, + while image components (2 & 3) must be sharded. + """ + img_freqs_cos, img_freqs_sin = self.pos_embed(img_ids) + txt_freqs_cos, txt_freqs_sin = self.pos_embed(txt_ids) + return txt_freqs_cos, txt_freqs_sin, img_freqs_cos, img_freqs_sin + + class Flux2TimestepGuidanceEmbeddings(nn.Module): def __init__( self, @@ -580,10 +735,24 @@ class Flux2Modulation(nn.Module): class Flux2Transformer2DModel(nn.Module): """ The Transformer model introduced in Flux 2. + + Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig. """ _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"] + _sp_plan = { + "": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True), + }, + "rope_prepare": { + 2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), + 3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), + }, + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + """SP plan: shard hidden_states at root level, shard img_freqs at rope_prepare, gather output at proj_out.""" + def __init__( self, patch_size: int = 1, @@ -600,6 +769,7 @@ class Flux2Transformer2DModel(nn.Module): rope_theta: int = 2000, eps: float = 1e-6, guidance_embeds: bool = True, + od_config: OmniDiffusionConfig = None, quant_config: "QuantizationConfig | None" = None, ): super().__init__() @@ -622,6 +792,13 @@ class Flux2Transformer2DModel(nn.Module): guidance_embeds=guidance_embeds, ) + if od_config is not None: + self.parallel_config = od_config.parallel_config + else: + from vllm_omni.diffusion.data import DiffusionParallelConfig + + self.parallel_config = DiffusionParallelConfig() + self.pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=list(axes_dims_rope)) self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings( in_channels=timestep_guidance_channels, @@ -637,9 +814,12 @@ class Flux2Transformer2DModel(nn.Module): self.x_embedder = nn.Linear(in_channels, self.inner_dim, bias=False) self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim, bias=False) + self.rope_prepare = Flux2RopePrepare(self.pos_embed) + self.transformer_blocks = nn.ModuleList( [ Flux2TransformerBlock( + parallel_config=self.parallel_config, dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, @@ -655,6 +835,7 @@ class Flux2Transformer2DModel(nn.Module): self.single_transformer_blocks = nn.ModuleList( [ Flux2SingleTransformerBlock( + parallel_config=self.parallel_config, dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, @@ -695,6 +876,10 @@ class Flux2Transformer2DModel(nn.Module): num_txt_tokens = encoder_hidden_states.shape[1] + sp_size = self.parallel_config.sequence_parallel_size + if sp_size is not None and sp_size > 1: + get_forward_context().split_text_embed_in_sp = False + timestep = timestep.to(hidden_states.dtype) * 1000 if guidance is not None: guidance = guidance.to(hidden_states.dtype) * 1000 @@ -705,21 +890,44 @@ class Flux2Transformer2DModel(nn.Module): double_stream_mod_txt = self.double_stream_modulation_txt(temb) single_stream_mod = self.single_stream_modulation(temb)[0] - hidden_states = self.x_embedder(hidden_states) - encoder_hidden_states = self.context_embedder(encoder_hidden_states) - if img_ids.ndim == 3: img_ids = img_ids[0] if txt_ids.ndim == 3: txt_ids = txt_ids[0] - image_rotary_emb = self.pos_embed(img_ids) - text_rotary_emb = self.pos_embed(txt_ids) + hidden_states = self.x_embedder(hidden_states) + encoder_hidden_states = self.context_embedder(encoder_hidden_states) + + txt_freqs_cos, txt_freqs_sin, img_freqs_cos, img_freqs_sin = self.rope_prepare(img_ids, txt_ids) + concat_rotary_emb = ( - torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0), - torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0), + torch.cat([txt_freqs_cos, img_freqs_cos], dim=0), + torch.cat([txt_freqs_sin, img_freqs_sin], dim=0), ) + # Create separate masks for image and text portions for Ulysses SP joint attention + hidden_states_mask = None + encoder_hidden_states_mask = None + ctx = get_forward_context() + if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: + batch_size = hidden_states.shape[0] + img_padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size + + hidden_states_mask = torch.ones( + batch_size, + img_padded_seq_len, + dtype=torch.bool, + device=hidden_states.device, + ) + hidden_states_mask[:, ctx.sp_original_seq_len :] = False + if hidden_states_mask.all(): + hidden_states_mask = None + + if hidden_states_mask is not None: + joint_attention_kwargs["hidden_states_mask"] = hidden_states_mask + if encoder_hidden_states_mask is not None: + joint_attention_kwargs["encoder_hidden_states_mask"] = encoder_hidden_states_mask + for block in self.transformer_blocks: encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, @@ -739,6 +947,7 @@ class Flux2Transformer2DModel(nn.Module): temb_mod_params=single_stream_mod, image_rotary_emb=concat_rotary_emb, joint_attention_kwargs=joint_attention_kwargs, + text_seq_len=num_txt_tokens, ) hidden_states = hidden_states[:, num_txt_tokens:, ...] diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index 3e61610b1..b68b6f70d 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -49,7 +49,7 @@ from vllm_omni.diffusion.forward_context import ( get_forward_context, is_forward_context_available, ) -from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk ADALN_EMBED_DIM = 256 SEQ_MULTI_OF = 32 @@ -329,10 +329,7 @@ class ZImageAttention(nn.Module): query = self.norm_q(query) key = self.norm_k(key) - cos = cos.to(query.dtype) - sin = sin.to(query.dtype) - query = self.rope(query, cos, sin) - key = self.rope(key, cos, sin) + query, key = apply_rope_to_qk(self.rope, query, key, (cos, sin)) # Cast to correct dtype dtype = query.dtype query, key = query.to(dtype), key.to(dtype) From 284575ab8231812685afbd1458568a3dd7aba204 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B1=AA=E5=BF=97=E9=B9=8F?= Date: Sun, 15 Mar 2026 18:27:07 +0800 Subject: [PATCH 09/10] [CI]: Add/Fix bagel e2e online/offline test (#1895) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: princepride Signed-off-by: 汪志鹏 --- .buildkite/test-ready.yml | 150 ++++++++--- .../stage_configs/bagel_sharedmemory_ci.yaml | 2 +- .../offline_inference/test_bagel_img2img.py | 184 +++++++++++++ .../offline_inference/test_bagel_text2img.py | 26 +- tests/e2e/online_serving/test_bagel_online.py | 247 ++++++++++++++++++ 5 files changed, 557 insertions(+), 52 deletions(-) create mode 100644 tests/e2e/offline_inference/test_bagel_img2img.py create mode 100644 tests/e2e/online_serving/test_bagel_online.py diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 66598f3ea..c8579979c 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -261,39 +261,117 @@ steps: # path: /mnt/hf-cache # type: DirectoryOrCreate - # - label: "Bagel Text2Img Model Test with H100" - # depends_on: upload-ready-pipeline - # commands: - # - | - # timeout 30m bash -c ' - # 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: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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" + depends_on: upload-ready-pipeline + commands: + - | + timeout 30m bash -c ' + export VLLM_WORKER_MULTIPROC_METHOD=spawn + export VLLM_TEST_CLEAN_GPU_MEMORY=1 + pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py + ' + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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 Img2Img Model Test with H100" + depends_on: upload-ready-pipeline + commands: + - | + timeout 30m bash -c ' + export VLLM_WORKER_MULTIPROC_METHOD=spawn + export VLLM_TEST_CLEAN_GPU_MEMORY=1 + pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py + ' + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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 Online Serving Test with H100" + depends_on: upload-ready-pipeline + commands: + - | + timeout 40m bash -c ' + export VLLM_WORKER_MULTIPROC_METHOD=spawn + export VLLM_TEST_CLEAN_GPU_MEMORY=1 + export VLLM_IMAGE_FETCH_TIMEOUT=60 + pytest -s -v tests/e2e/online_serving/test_bagel_online.py + ' + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml index 686600f72..721c50248 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml @@ -10,7 +10,7 @@ stage_args: max_batch_size: 1 engine_args: model_stage: thinker - model_arch: BagelForConditionalGeneration + model_arch: OmniBagelForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.45 diff --git a/tests/e2e/offline_inference/test_bagel_img2img.py b/tests/e2e/offline_inference/test_bagel_img2img.py new file mode 100644 index 000000000..eef0b7d6c --- /dev/null +++ b/tests/e2e/offline_inference/test_bagel_img2img.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +End-to-end test for Bagel img2img generation. + +This test validates that the Bagel model generates images from an input image +and text prompt that match expected reference pixel values within a ±5 tolerance. + +Equivalent to running: + python3 examples/offline_inference/bagel/end2end.py \ + --prompts "Change the grass color to red" \ + --modality img2img --step 15 \ + --image-path 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg +""" + +import socket +from pathlib import Path +from typing import Any + +import pytest +from PIL import Image +from vllm.assets.image import ImageAsset + +from tests.utils import hardware_test +from vllm_omni.entrypoints.omni import Omni + +# Reference pixel data extracted from the known-good output image +# Generated with seed=52, num_inference_steps=15, +# prompt='Change the grass color to red', +# input image: 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg +REFERENCE_PIXELS = [ + {"position": (100, 100), "rgb": (157, 172, 217)}, + {"position": (400, 50), "rgb": (105, 144, 218)}, + {"position": (700, 100), "rgb": (118, 159, 233)}, + {"position": (150, 400), "rgb": (195, 34, 60)}, + {"position": (512, 336), "rgb": (222, 214, 193)}, + {"position": (700, 400), "rgb": (197, 15, 43)}, + {"position": (100, 600), "rgb": (105, 13, 18)}, + {"position": (400, 600), "rgb": (169, 33, 44)}, + {"position": (700, 600), "rgb": (101, 86, 93)}, + {"position": (256, 256), "rgb": (181, 202, 222)}, +] + +PIXEL_TOLERANCE = 5 + +DEFAULT_PROMPT = "<|fim_middle|><|im_start|>Change the grass color to red<|im_end|>" + +EXPECTED_OUTPUT_SIZE = (1024, 672) + + +def _load_input_image() -> Image.Image: + """Load the test input image via vllm's ImageAsset.""" + return ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") + + +def _find_free_port() -> int: + """Find and return a free ephemeral port by binding to port 0.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + s.listen(1) + port = s.getsockname()[1] + return port + + +def _configure_sampling_params(omni: Omni, max_tokens: int = 1, num_inference_steps: int = 15) -> list: + """Configure sampling parameters for Bagel img2img generation. + + Args: + omni: The Omni instance to get default params from. + max_tokens: Maximum tokens for the first stage. + num_inference_steps: Number of inference steps for the diffusion stage. + + Returns: + Configured sampling params list. + """ + params_list = omni.default_sampling_params_list + params_list[0].max_tokens = max_tokens # type: ignore + if len(params_list) > 1: + params_list[1].num_inference_steps = num_inference_steps # type: ignore + params_list[1].extra_args = { # type: ignore + "cfg_text_scale": 4.0, + "cfg_img_scale": 1.5, + } + return params_list + + +def _extract_generated_image(omni_outputs: list) -> Image.Image | None: + """Extract the generated image from Omni outputs. + + Args: + omni_outputs: List of outputs from omni.generate(). + + Returns: + The first generated PIL Image, or None if no image found. + """ + for req_output in omni_outputs: + if images := getattr(req_output, "images", None): + return images[0] + if hasattr(req_output, "request_output") and req_output.request_output: + for stage_out in req_output.request_output: + if hasattr(stage_out, "images") and stage_out.images: + return stage_out.images[0] + return None + + +def _validate_pixels( + image: Image.Image, + reference_pixels: list[dict[str, Any]] = REFERENCE_PIXELS, + tolerance: int = PIXEL_TOLERANCE, +) -> None: + """Validate that image pixels match expected reference values. + + Args: + image: The PIL Image to validate. + reference_pixels: List of dicts with 'position' (x, y) and 'rgb' (R, G, B). + tolerance: Maximum allowed difference per color channel. + + Raises: + AssertionError: If any pixel differs beyond tolerance. + """ + for ref in reference_pixels: + x, y = ref["position"] + expected = ref["rgb"] + actual = image.getpixel((x, y))[:3] + assert all(abs(a - e) <= tolerance for a, e in zip(actual, expected)), ( + f"Pixel mismatch at ({x}, {y}): expected {expected}, got {actual}" + ) + + +def _generate_bagel_img2img( + omni: Omni, + input_image: Image.Image, + prompt: str = DEFAULT_PROMPT, +) -> Image.Image: + """Generate an image using Bagel model with img2img pipeline. + + Args: + omni: The Omni instance to use for generation. + input_image: The input PIL Image for img2img. + prompt: The text prompt for image editing. + + Returns: + The generated PIL Image. + + Raises: + AssertionError: If no image is generated or size is incorrect. + """ + params_list = _configure_sampling_params(omni) + + omni_outputs = list( + omni.generate( + prompts=[ + { + "prompt": prompt, + "multi_modal_data": {"img2img": input_image}, + "modalities": ["img2img"], + } + ], + sampling_params_list=params_list, + ) + ) + + generated_image = _extract_generated_image(omni_outputs) + assert generated_image is not None, "No images generated" + assert generated_image.size == EXPECTED_OUTPUT_SIZE, f"Expected {EXPECTED_OUTPUT_SIZE}, got {generated_image.size}" + + return generated_image + + +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}) +def test_bagel_img2img_shared_memory_connector(): + """Test Bagel img2img with shared memory connector.""" + input_image = _load_input_image() + 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, stage_init_timeout=300) + + try: + generated_image = _generate_bagel_img2img(omni, input_image) + _validate_pixels(generated_image) + finally: + omni.close() diff --git a/tests/e2e/offline_inference/test_bagel_text2img.py b/tests/e2e/offline_inference/test_bagel_text2img.py index 20f788895..360d49bb1 100644 --- a/tests/e2e/offline_inference/test_bagel_text2img.py +++ b/tests/e2e/offline_inference/test_bagel_text2img.py @@ -14,10 +14,6 @@ Equivalent to running: """ import os - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" - import signal import socket import subprocess @@ -37,23 +33,23 @@ from vllm_omni.entrypoints.omni import Omni # "Generated with seed=52, num_inference_steps=15, # prompt='A futuristic city skyline at twilight, cyberpunk style'" REFERENCE_PIXELS = [ - {"position": (100, 100), "rgb": (49, 96, 134)}, - {"position": (400, 50), "rgb": (63, 127, 167)}, - {"position": (700, 100), "rgb": (70, 101, 141)}, - {"position": (150, 400), "rgb": (115, 90, 150)}, - {"position": (512, 512), "rgb": (98, 86, 119)}, - {"position": (700, 400), "rgb": (29, 42, 91)}, - {"position": (100, 700), "rgb": (47, 50, 88)}, - {"position": (400, 700), "rgb": (36, 52, 91)}, - {"position": (700, 700), "rgb": (45, 58, 99)}, - {"position": (256, 256), "rgb": (62, 94, 135)}, + {"position": (100, 100), "rgb": (121, 118, 100)}, + {"position": (400, 50), "rgb": (163, 162, 143)}, + {"position": (700, 100), "rgb": (170, 156, 127)}, + {"position": (150, 400), "rgb": (129, 127, 112)}, + {"position": (512, 512), "rgb": (135, 61, 59)}, + {"position": (700, 400), "rgb": (205, 107, 43)}, + {"position": (100, 700), "rgb": (197, 177, 157)}, + {"position": (400, 700), "rgb": (139, 107, 86)}, + {"position": (700, 700), "rgb": (247, 205, 146)}, + {"position": (256, 256), "rgb": (171, 160, 153)}, ] # Maximum allowed difference per color channel PIXEL_TOLERANCE = 5 # Default test prompt -DEFAULT_PROMPT = "<|im_start|>A futuristic city skyline at twilight, cyberpunk style<|im_end|>" +DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>" def _find_free_port() -> int: diff --git a/tests/e2e/online_serving/test_bagel_online.py b/tests/e2e/online_serving/test_bagel_online.py new file mode 100644 index 000000000..4056cfdef --- /dev/null +++ b/tests/e2e/online_serving/test_bagel_online.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +End-to-end online serving test for Bagel text2img and img2img generation. + +This test validates that the Bagel model can serve image generation requests +via the OpenAI-compatible chat completions API. + +Equivalent to running: + vllm-omni serve "ByteDance-Seed/BAGEL-7B-MoT" --omni --port 8091 + + # text2img + python3 examples/online_serving/bagel/openai_chat_client.py \\ + --prompt "A cute cat" --modality text2img + + # img2img + python3 examples/online_serving/bagel/openai_chat_client.py \\ + --prompt "Let the woman wear a blue dress" --modality img2img \\ + --image-url women.jpg +""" + +import base64 +import os +import signal +import socket +import subprocess +import sys +import time +from io import BytesIO +from pathlib import Path +from typing import Any + +import pytest +import requests +from PIL import Image +from vllm.assets.image import ImageAsset + +from tests.utils import hardware_test + +MODEL = "ByteDance-Seed/BAGEL-7B-MoT" +STAGE_CONFIGS_PATH = str( + Path(__file__).parent.parent / "offline_inference" / "stage_configs" / "bagel_sharedmemory_ci.yaml" +) + +TEXT2IMG_PROMPT = "A cute cat" +IMG2IMG_PROMPT = "Change the grass color to red" + + +class BagelOmniServer: + """Context manager to start/stop a vLLM-Omni server for Bagel model tests.""" + + def __init__( + self, + model: str = MODEL, + stage_configs_path: str = STAGE_CONFIGS_PATH, + env_dict: dict[str, str] | None = None, + ) -> None: + self.model = model + self.stage_configs_path = stage_configs_path + self.env_dict = env_dict + self.proc: subprocess.Popen | None = None + self.host = "127.0.0.1" + self.port = _find_free_port() + + @property + def base_url(self) -> str: + return f"http://{self.host}:{self.port}" + + def _start_server(self) -> None: + env = os.environ.copy() + if self.env_dict is not None: + env.update(self.env_dict) + + cmd = [ + sys.executable, + "-m", + "vllm_omni.entrypoints.cli.main", + "serve", + self.model, + "--omni", + "--host", + self.host, + "--port", + str(self.port), + "--stage-configs-path", + self.stage_configs_path, + "--stage-init-timeout", + "300", + ] + + self.proc = subprocess.Popen( + cmd, + env=env, + start_new_session=True, + ) + + try: + if not _wait_for_port(self.host, self.port, timeout=600, proc=self.proc): + self.terminate() + raise RuntimeError(f"Server failed to start within 600 seconds on {self.host}:{self.port}") + except Exception: + self.terminate() + raise + + def __enter__(self): + self._start_server() + return self + + def terminate(self) -> None: + if self.proc: + try: + os.killpg(os.getpgid(self.proc.pid), signal.SIGTERM) + except ProcessLookupError: + pass + try: + self.proc.wait(timeout=30) + except subprocess.TimeoutExpired: + try: + os.killpg(os.getpgid(self.proc.pid), signal.SIGKILL) + except ProcessLookupError: + pass + self.proc.wait() + self.proc = None + + def __exit__(self, exc_type, exc_val, exc_tb): + self.terminate() + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + s.listen(1) + return s.getsockname()[1] + + +def _wait_for_port(host: str, port: int, timeout: int = 600, proc: subprocess.Popen | None = None) -> bool: + start = time.time() + while time.time() - start < timeout: + if proc is not None and proc.poll() is not None: + # Server process exited early + return False + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(1) + if sock.connect_ex((host, port)) == 0: + return True + except Exception: + pass + time.sleep(2) + return False + + +def _send_chat_request( + server_url: str, + prompt: str, + *, + modality: str = "text2img", + image: Image.Image | None = None, + timeout: int = 300, +) -> dict[str, Any]: + """Send a chat completion request matching the openai_chat_client.py format.""" + content: list[dict[str, Any]] = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}] + + if image is not None: + buffer = BytesIO() + image.save(buffer, format="JPEG") + b64_data = base64.b64encode(buffer.getvalue()).decode("utf-8") + content.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{b64_data}"}, + } + ) + + payload: dict[str, Any] = { + "messages": [{"role": "user", "content": content}], + } + + if modality in ("text2img", "img2img"): + payload["modalities"] = ["image"] + + resp = requests.post( + f"{server_url}/v1/chat/completions", + headers={"Content-Type": "application/json"}, + json=payload, + timeout=timeout, + ) + resp.raise_for_status() + return resp.json() + + +def _extract_image_from_response(data: dict[str, Any]) -> Image.Image | None: + """Extract the generated PIL Image from a chat completion response.""" + for choice in data.get("choices", []): + content = choice.get("message", {}).get("content") + if isinstance(content, list) and content: + first_item = content[0] + if isinstance(first_item, dict) and "image_url" in first_item: + url = first_item["image_url"].get("url", "") + if url.startswith("data:image"): + _, b64 = url.split(",", 1) + return Image.open(BytesIO(base64.b64decode(b64))) + return None + + +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}) +def test_bagel_text2img_online(): + """Test Bagel text2img via OpenAI-compatible chat completions API.""" + with BagelOmniServer() as server: + response_data = _send_chat_request( + server.base_url, + TEXT2IMG_PROMPT, + modality="text2img", + ) + + image = _extract_image_from_response(response_data) + assert image is not None, f"No image in response: {response_data}" + image.load() + + w, h = image.size + assert w > 0 and h > 0, f"Invalid image size: {image.size}" + + +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}) +def test_bagel_img2img_online(): + """Test Bagel img2img via OpenAI-compatible chat completions API.""" + input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") + + with BagelOmniServer() as server: + response_data = _send_chat_request( + server.base_url, + IMG2IMG_PROMPT, + modality="img2img", + image=input_image, + ) + + image = _extract_image_from_response(response_data) + assert image is not None, f"No image in response: {response_data}" + image.load() + + w, h = image.size + assert w > 0 and h > 0, f"Invalid image size: {image.size}" From 9aae75c3ee0729e4ae454778a4d13c375bc2da6a Mon Sep 17 00:00:00 2001 From: Lancer Date: Sun, 15 Mar 2026 22:48:12 +0800 Subject: [PATCH 10/10] [Feat] support HSDP for Flux family (#1900) Signed-off-by: Lancer --- .../diffusion/parallelism_acceleration.md | 32 +++++++++---------- .../diffusion/models/flux/flux_transformer.py | 6 ++++ .../diffusion/models/flux/pipeline_flux.py | 4 +-- .../models/flux2/flux2_transformer.py | 6 ++++ .../diffusion/models/flux2/pipeline_flux2.py | 2 +- .../flux2_klein/flux2_klein_transformer.py | 5 +++ .../flux2_klein/pipeline_flux2_klein.py | 2 +- 7 files changed, 37 insertions(+), 20 deletions(-) diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index 37fd5b1d1..20f09d2c0 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -24,22 +24,22 @@ The following table shows which models are currently supported by parallelism me ### ImageGen -| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel | -|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:| -| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | -| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | -| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | N/A | -| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | N/A | -| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | -| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | -| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | -| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | N/A | -| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ | N/A | -| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | -| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | N/A | -| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | -| **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | -| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ✅ | ❌ | ❌ | N/A | +| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel | HSDP | +|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|:----:| +| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ❌ | +| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ❌ | +| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | N/A | ❌ | +| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | N/A | ❌ | +| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ | +| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ | +| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ | +| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | N/A | ❌ | +| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ | N/A | ❌ | +| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ✅ | +| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | N/A | ✅ | +| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | ✅ | +| **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | +| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ✅ | ❌ | ❌ | N/A | ❌ | !!! 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/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py index 2979fd4f6..1a003b86c 100644 --- a/vllm_omni/diffusion/models/flux/flux_transformer.py +++ b/vllm_omni/diffusion/models/flux/flux_transformer.py @@ -470,6 +470,12 @@ class FluxTransformer2DModel(nn.Module): # -- typically a transformer layer # used for torch compile optimizations _repeated_blocks = ["FluxTransformerBlock"] + + @staticmethod + def _is_transformer_block(name: str, module) -> bool: + return ("transformer_blocks" in name or "single_transformer_blocks" in name) and name.split(".")[-1].isdigit() + + _hsdp_shard_conditions = [_is_transformer_block] packed_modules_mapping = { "to_qkv": ["to_q", "to_k", "to_v"], "add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"], diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux.py b/vllm_omni/diffusion/models/flux/pipeline_flux.py index d4793d34e..3955fee12 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux.py +++ b/vllm_omni/diffusion/models/flux/pipeline_flux.py @@ -160,10 +160,10 @@ class FluxPipeline(nn.Module, CFGParallelMixin): ) self.text_encoder = CLIPTextModel.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only - ) + ).to(self.device) self.text_encoder_2 = T5EncoderModel.from_pretrained( model, subfolder="text_encoder_2", local_files_only=local_files_only - ) + ).to(self.device) self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.device ) diff --git a/vllm_omni/diffusion/models/flux2/flux2_transformer.py b/vllm_omni/diffusion/models/flux2/flux2_transformer.py index 040f2779a..62f415393 100644 --- a/vllm_omni/diffusion/models/flux2/flux2_transformer.py +++ b/vllm_omni/diffusion/models/flux2/flux2_transformer.py @@ -553,6 +553,12 @@ class Flux2Transformer2DModel(nn.Module): _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"] + @staticmethod + def _is_transformer_block(name: str, module) -> bool: + return ("transformer_blocks" in name or "single_transformer_blocks" in name) and name.split(".")[-1].isdigit() + + _hsdp_shard_conditions = [_is_transformer_block] + def __init__( self, patch_size: int = 1, diff --git a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py index 7f07adf6f..1da0f0cda 100644 --- a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py +++ b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py @@ -366,7 +366,7 @@ class Flux2Pipeline(nn.Module, SupportImageInput): ) self.text_encoder = Mistral3ForConditionalGeneration.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only - ) + ).to(self._execution_device) self.tokenizer = PixtralProcessor.from_pretrained( model, subfolder="tokenizer", local_files_only=local_files_only ) diff --git a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py index c10f06751..fa7bbc4fe 100644 --- a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py +++ b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py @@ -741,6 +741,11 @@ class Flux2Transformer2DModel(nn.Module): _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"] + @staticmethod + def _is_transformer_block(name: str, module) -> bool: + return ("transformer_blocks" in name or "single_transformer_blocks" in name) and name.split(".")[-1].isdigit() + + _hsdp_shard_conditions = [_is_transformer_block] _sp_plan = { "": { "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True), diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py index d43748380..551be4f06 100644 --- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py +++ b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py @@ -218,7 +218,7 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput): model, subfolder="text_encoder", local_files_only=local_files_only, - ) + ).to(self._execution_device) self.tokenizer = Qwen2TokenizerFast.from_pretrained( model, subfolder="tokenizer",