add

Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Signed-off-by: Roger Wang <hey@rogerw.io>

.buildkite/test-ready.yml

@@ -261,117 +261,39 @@ 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
-          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
+  # - 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

README.md

@@ -17,7 +17,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 *Latest News* 🔥
 
-- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
+- [2026/03] Joined vLLM HK First Meetup! Video on [youtube](https://youtu.be/sgwNfsNnR9I)
 - [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
 - [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
 - [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.

docs/community/governance.md

@@ -37,10 +37,12 @@ Committers have write access and merge rights. They typically have deep expertis
 - [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
 - [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
 - [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
+- [@lishunyang12](https://github.com/lishunyang12): 
 - [@princepride](https://github.com/princepride): Diffusion and Omni Support
 - [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
 - [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
-- [@wtomin](https://github.com/wtomin):
+- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism
+- [@ZeldaHuang](https://github.com/ZeldaHuang): 
 - [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
 
 ## Meetings

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 | 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        | ❌   |
+| 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        |
 
 !!! note "TP Limitations for Diffusion Models"
     We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.

tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml

@@ -10,7 +10,7 @@ stage_args:
       max_batch_size: 1
     engine_args:
       model_stage: thinker
-      model_arch: OmniBagelForConditionalGeneration
+      model_arch: BagelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.45

tests/e2e/offline_inference/test_bagel_img2img.py

@@ -1,184 +0,0 @@
-# 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()

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -14,6 +14,10 @@ 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
@@ -33,23 +37,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": (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)},
+    {"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)},
 ]
 
 # Maximum allowed difference per color channel
 PIXEL_TOLERANCE = 5
 
 # Default test prompt
-DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
+DEFAULT_PROMPT = "<|im_start|>A futuristic city skyline at twilight, cyberpunk style<|im_end|>"
 
 
 def _find_free_port() -> int:

tests/e2e/online_serving/test_bagel_online.py

@@ -1,247 +0,0 @@
-# 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}"

vllm_omni/diffusion/layers/rope.py

@@ -157,29 +157,3 @@ 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

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, apply_rope_to_qk
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
 logger = init_logger(__name__)
 
@@ -224,7 +224,12 @@ class FluxAttention(torch.nn.Module):
             key = torch.cat([encoder_key, key], dim=1)
             value = torch.cat([encoder_value, value], dim=1)
 
-        query, key = apply_rope_to_qk(self.rope, query, key, image_rotary_emb)  # [S, D/2]
+        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)
 
         hidden_states = self.attn(
             query,
@@ -470,12 +475,6 @@ 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"],

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
         )

vllm_omni/diffusion/models/flux2/flux2_transformer.py

@@ -553,12 +553,6 @@ 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,

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
         )

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

@@ -28,7 +28,6 @@ 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,
@@ -40,15 +39,8 @@ 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.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
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
-logger = init_logger(__name__)
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
@@ -106,7 +98,6 @@ class Flux2FeedForward(nn.Module):
 class Flux2Attention(nn.Module):
     def __init__(
         self,
-        parallel_config: DiffusionParallelConfig,
         query_dim: int,
         heads: int = 8,
         dim_head: int = 64,
@@ -121,7 +112,6 @@ 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
@@ -218,87 +208,33 @@ class Flux2Attention(nn.Module):
             encoder_query = self.norm_added_q(encoder_query)
             encoder_key = self.norm_added_k(encoder_key)
 
-            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
+            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 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)
+        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 = 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
+        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.attn(query, key, value, attn_metadata)
+        hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
 
-                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)
+        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)
 
         hidden_states = self.to_out[0](hidden_states)
         hidden_states = self.to_out[1](hidden_states)
@@ -315,7 +251,6 @@ class Flux2ParallelSelfAttention(nn.Module):
 
     def __init__(
         self,
-        parallel_config: DiffusionParallelConfig,
         query_dim: int,
         heads: int = 8,
         dim_head: int = 64,
@@ -330,7 +265,6 @@ 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
@@ -391,57 +325,20 @@ class Flux2ParallelSelfAttention(nn.Module):
         query = self.norm_q(query)
         key = self.norm_k(key)
 
-        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:
+        if image_rotary_emb is not None:
             cos, sin = image_rotary_emb
             cos = cos.to(query.dtype)
             sin = sin.to(query.dtype)
-            txt_cos, img_cos = cos[:text_seq_len], cos[text_seq_len:]
-            txt_sin, img_sin = sin[:text_seq_len], sin[text_seq_len:]
+            query = self.rope(query, cos, sin)
+            key = self.rope(key, cos, sin)
 
-            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_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 = attn_output.flatten(2, 3).to(query.dtype)
 
         mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states)
@@ -453,7 +350,6 @@ class Flux2ParallelSelfAttention(nn.Module):
 class Flux2SingleTransformerBlock(nn.Module):
     def __init__(
         self,
-        parallel_config: DiffusionParallelConfig,
         dim: int,
         num_attention_heads: int,
         attention_head_dim: int,
@@ -465,7 +361,6 @@ 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,
@@ -488,13 +383,6 @@ 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)
@@ -524,7 +412,6 @@ class Flux2SingleTransformerBlock(nn.Module):
 class Flux2TransformerBlock(nn.Module):
     def __init__(
         self,
-        parallel_config: DiffusionParallelConfig,
         dim: int,
         num_attention_heads: int,
         attention_head_dim: int,
@@ -538,7 +425,6 @@ 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,
@@ -639,47 +525,6 @@ 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,
@@ -735,29 +580,10 @@ 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"]
 
-    @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),
-        },
-        "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,
@@ -774,7 +600,6 @@ 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__()
@@ -797,13 +622,6 @@ 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,
@@ -819,12 +637,9 @@ 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,
@@ -840,7 +655,6 @@ 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,
@@ -881,10 +695,6 @@ 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
@@ -895,44 +705,21 @@ 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]
 
-        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)
-
+        image_rotary_emb = self.pos_embed(img_ids)
+        text_rotary_emb = self.pos_embed(txt_ids)
         concat_rotary_emb = (
-            torch.cat([txt_freqs_cos, img_freqs_cos], dim=0),
-            torch.cat([txt_freqs_sin, img_freqs_sin], dim=0),
+            torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
+            torch.cat([text_rotary_emb[1], image_rotary_emb[1]], 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,
@@ -952,7 +739,6 @@ 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:, ...]

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

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, apply_rope_to_qk
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
 ADALN_EMBED_DIM = 256
 SEQ_MULTI_OF = 32
@@ -329,7 +329,10 @@ class ZImageAttention(nn.Module):
         query = self.norm_q(query)
         key = self.norm_k(key)
 
-        query, key = apply_rope_to_qk(self.rope, query, key, (cos, sin))
+        cos = cos.to(query.dtype)
+        sin = sin.to(query.dtype)
+        query = self.rope(query, cos, sin)
+        key = self.rope(key, cos, sin)
         # Cast to correct dtype
         dtype = query.dtype
         query, key = query.to(dtype), key.to(dtype)