Merge remote-tracking branch 'origin/main' into dev/vllm-align

2026-03-16 01:49:46 +00:00
commit f137a1e3ad
@@ -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 </dev/urandom | head -c 10
+    echo
+)"
+
+cd "${omni_source_dir}"
+if [ -z "$(docker images -q "${base_image_name}")" ]; then
+    docker build --target vllm-base -t "${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -f docker/Dockerfile.xpu .
+fi
+
+# Try building the docker image
+docker build --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -t "${image_name}" -f docker/Dockerfile.xpu .
+
+# Setup cleanup
+remove_docker_container() {
+    docker rm -f "${container_name}" || true
+    docker image rm -f "${image_name}" || true
+    docker system prune -f || true
+}
+trap remove_docker_container EXIT
+
+HF_CACHE="${HF_CACHE:-$(realpath ~)/.cache/huggingface}"
+mkdir -p "${HF_CACHE}"
+HF_MOUNT="/root/.cache/huggingface"
+
+time timeout -k 30 30m docker run \
+    --device /dev/dri:/dev/dri \
+    --net=host \
+    --ipc=host \
+    -v /dev/dri/by-path:/dev/dri/by-path \
+    -v "${HF_CACHE}:${HF_MOUNT}" \
+    --security-opt seccomp=unconfined \
+    --entrypoint="" \
+    -e VLLM_LOGGING_LEVEL \
+    -e VLLM_OMNI_LOGGING_LEVEL \
+    -e HF_TOKEN \
+    -e ZE_AFFINITY_MASK \
+    --name "${container_name}" \
+    "${image_name}" \
+    bash -c '
+    set -e
+    echo $ZE_AFFINITY_MASK
+    pip install tblib==3.1.0
+    cd /workspace/vllm-omni
+    pytest -v -s -m "core_model and xpu and B60"
+    pytest -v -s -m "advanced_model and xpu and B60"
+'
@@ -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
@@ -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
@@ -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):

@@ -76,7 +76,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"

-ARG VLLM_VERSION=v0.16.0
+ARG VLLM_VERSION=v0.17.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm

@@ -17,6 +17,7 @@ Defined in `pyproject.toml`:
 | `gpu`              | Tests that run on GPU *                                   |
 | `cuda`             | Tests that run on CUDA *                                  |
 | `rocm`             | Tests that run on AMD/ROCm *                              |
+| `xpu`              | Tests that run on Intel XPU *                             |
 | `npu`              | Tests that run on NPU/Ascend *                            |
 | `H100`             | Tests that require H100 GPU  *                            |
 | `L4`               | Tests that require L4 GPU *                               |
@@ -55,7 +56,7 @@ def test_video_to_audio()
 This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions:

 1. **Applies platform and resource markers**  
-   Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `A2`, `A3`).
+   Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`).
   ```
   @pytest.mark.cuda
   @pytest.mark.L4
@@ -80,14 +81,14 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
 - Single call for multiple platforms:
    ```python
    @hardware_test(
-        res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
-        num_cards={"cuda": 2, "rocm": 2, "npu": 2},
+        res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
+        num_cards={"cuda": 2, "rocm": 2, "xpu": 2, "npu": 2},
    )
    ```
    or
    ```python
    @hardware_test(
-        res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+        res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
        num_cards=2,
    )
    ```
@@ -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.
@@ -167,6 +167,7 @@ markers = [
    "gpu: Tests that run on GPU (auto-added)",
    "cuda: Tests that run on CUDA (auto-added)",
    "rocm: Tests that run on AMD/ROCm (auto-added)",
+    "xpu: Tests that run on XPU (auto-added)",
    "npu: Tests that run on NPU/Ascend (auto-added)",
    # specified computation resources marks (auto-added)
    "H100: Tests that require H100 GPU",
@@ -176,9 +177,11 @@ markers = [
    "A3: Tests that require A3 NPU",
    "distributed_cuda: Tests that require multi cards on CUDA platform",
    "distributed_rocm: Tests that require multi cards on ROCm platform",
+    "distributed_xpu: Tests that require multi cards on XPU platform",
    "distributed_npu: Tests that require multi cards on NPU platform",
    "skipif_cuda: Skip if the num of CUDA cards is less than the required",
    "skipif_rocm: Skip if the num of ROCm cards is less than the required",
+    "skipif_xpu: Skip if the num of XPU cards is less than the required",
    "skipif_npu: Skip if the num of NPU cards is less than the required",
    # more detailed markers
    "slow: Slow tests (may skip in quick CI)",
@@ -5,10 +5,14 @@ import pytest

 from tests.conftest import OmniServerParams
 from tests.utils import hardware_test
+from vllm_omni.platforms import current_omni_platform

 models = ["Qwen/Qwen2.5-Omni-7B"]
 stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_ci.yaml")]

+if current_omni_platform.is_xpu():
+    stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "xpu" / "qwen2_5_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
@@ -17,7 +21,7 @@ test_params = [

@pytest.mark.core_model
@pytest.mark.benchmark
-@hardware_test(res={"cuda": "L4"}, num_cards=3)
+@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=3)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_bench_serve_chat(omni_server):
    command = [
@@ -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)
@@ -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
@@ -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()
@@ -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:
@@ -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_cache_dit(model_name: str):
    """Test cache-dit backend with diffusion model."""
@@ -39,6 +39,9 @@ if current_omni_platform.is_npu():
 elif current_omni_platform.is_rocm():
    # ROCm stage config optimized for MI325 GPU
    stage_config = str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml")
+elif current_omni_platform.is_xpu():
+    # Intel XPU stage config optimized for B60 GPU
+    stage_config = str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen2_5_omni_ci.yaml")
 else:
    stage_config = get_cuda_graph_config()

@@ -56,7 +59,7 @@ def get_question(prompt_type="mix"):

@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_mix_to_audio(omni_runner, omni_runner_handler) -> 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:
    """
@@ -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()]

@@ -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)
@@ -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":
@@ -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."""
@@ -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"}, 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.")
+@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})"
+    )
@@ -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}"
@@ -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())
@@ -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()]

@@ -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
@@ -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
@@ -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
@@ -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 = [
@@ -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 = [
@@ -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
+}
@@ -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
@@ -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"
@@ -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)
@@ -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"

@@ -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:
@@ -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
@@ -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,
@@ -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:
@@ -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,
@@ -475,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"],
@@ -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
        )
@@ -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,
@@ -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
        )
@@ -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,29 @@ 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,
@@ -600,6 +774,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 +797,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 +819,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 +840,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 +881,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 +895,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 +952,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:, ...]
@@ -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",
@@ -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)
@@ -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()

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