Fix griffe warning: add type annotation for **kwargs parameter

Co-authored-by: hsliuustc0106 <222337142+hsliuustc0106@users.noreply.github.com>

.buildkite/pipeline.yml

@@ -42,6 +42,40 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+  - label: "Diffusion Model CPU offloading Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Audio Generation Model Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Diffusion Cache Backend Test"
     timeout_in_minutes: 15
     depends_on: image-build
@@ -60,10 +94,29 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Parallelism Test"
-    timeout_in_minutes: 15
+    timeout_in_minutes: 25
     depends_on: image-build
     commands:
       - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+      - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.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
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:

.buildkite/scripts/docker_login_ecr_public.sh

@@ -0,0 +1,73 @@
+#!/bin/bash
+# Helper function to safely login to ECR Public with per-job config isolation
+# Uses DOCKER_CONFIG environment variable to prevent race conditions
+#
+# This script prevents the "device or resource busy" error by giving each
+# Buildkite job its own isolated Docker config directory.
+#
+# Usage:
+#   source docker_login_ecr_public.sh && safe_docker_login_ecr_public
+
+set -euo pipefail
+
+# Configuration
+ECR_REGISTRY="public.ecr.aws"
+
+setup_isolated_docker_config() {
+    # Use BUILDKITE_JOB_ID for job-specific isolation
+    # Fallback to PID if running outside Buildkite
+    local job_id="${BUILDKITE_JOB_ID:-$$}"
+
+    # Set Docker config to job-specific directory
+    export DOCKER_CONFIG="/tmp/docker-config-${job_id}"
+
+    # Create directory if it doesn't exist
+    mkdir -p "$DOCKER_CONFIG"
+
+    echo "[docker-config] Using isolated Docker config: $DOCKER_CONFIG"
+}
+
+check_docker_auth() {
+    # Check if already authenticated to the given registry
+    # Returns 0 if authenticated, 1 if not
+    local registry="$1"
+
+    # Check if credentials exist in the isolated config
+    if [[ -f "$DOCKER_CONFIG/config.json" ]]; then
+        # Check if registry is present in config
+        if grep -q "$registry" "$DOCKER_CONFIG/config.json" 2>/dev/null; then
+            return 0
+        fi
+    fi
+
+    return 1
+}
+
+safe_docker_login_ecr_public() {
+    # Setup isolated config first
+    setup_isolated_docker_config
+
+    local registry="$ECR_REGISTRY"
+
+    # Check if already authenticated (within this job)
+    if check_docker_auth "$registry"; then
+        echo "[docker-login] Already authenticated to $registry in this job"
+        return 0
+    fi
+
+    # Perform login to isolated config directory
+    echo "[docker-login] Logging in to $ECR_REGISTRY (isolated config)..."
+    if aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$ECR_REGISTRY"; then
+        echo "[docker-login] Login successful (config: $DOCKER_CONFIG)"
+        return 0
+    else
+        local exit_code=$?
+        echo "[docker-login] ERROR: Login failed with exit code $exit_code" >&2
+        return $exit_code
+    fi
+}
+
+# Execute if run as script (not sourced)
+if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
+    safe_docker_login_ecr_public
+fi

.buildkite/scripts/hardware_ci/run-amd-test.sh

@@ -66,6 +66,14 @@ done
 echo "--- Pulling container"
 image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
 container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
+
+# Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
+sudo apt-get update && sudo apt-get install -y awscli
+# Use safe docker login helper to prevent race conditions
+source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
+safe_docker_login_ecr_public
+# Pull the container from ECR Public Gallery
+
 docker pull "${image_name}"
 
 remove_docker_container() {

.buildkite/scripts/simple_test.sh

@@ -52,3 +52,4 @@ VENV_PYTHON="${VENV_DIR}/bin/python"
 "${VENV_PYTHON}" -m pytest -v -s tests/entrypoints/
 "${VENV_PYTHON}" -m pytest -v -s tests/diffusion/cache/
 "${VENV_PYTHON}" -m pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
+"${VENV_PYTHON}" -m pytest -v -s tests/worker/

.buildkite/test-amd.yaml

@@ -1,10 +1,10 @@
 steps:
 
 - label: "Diffusion Model Test"
-  timeout_in_minutes: 15
+  timeout_in_minutes: 20
   agent_pool: mi325_2
   depends_on: amd-build
-  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
+  mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
     - export GPU_ARCHS=gfx942
@@ -20,7 +20,7 @@ steps:
   timeout_in_minutes: 15
   agent_pool: mi325_1
   depends_on: amd-build
-  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
+  mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
     - export GPU_ARCHS=gfx942
@@ -34,11 +34,33 @@ steps:
     - export VLLM_ROCM_USE_AITER_RMSNORM=0
     - pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
 
+- label: "Diffusion Parallelism Test"
+  timeout_in_minutes: 15
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export MIOPEN_DEBUG_CONV_DIRECT=0
+    - export MIOPEN_DEBUG_CONV_GEMM=0
+    - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+
+- label: "Diffusion GPU Worker Test"
+  timeout_in_minutes: 20
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export MIOPEN_DEBUG_CONV_DIRECT=0
+    - export MIOPEN_DEBUG_CONV_GEMM=0
+    - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
+
 - label: "Omni Model Test Qwen2-5-Omni"
   timeout_in_minutes: 15
   agent_pool: mi325_2
   depends_on: amd-build
-  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
+  mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
     - export GPU_ARCHS=gfx942
@@ -51,3 +73,34 @@ steps:
     - export VLLM_ROCM_USE_AITER_LINEAR=0
     - export VLLM_ROCM_USE_AITER_RMSNORM=0
     - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+
+- label: "Omni Model Test Qwen3-Omni"
+  timeout_in_minutes: 15
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export MIOPEN_DEBUG_CONV_DIRECT=0
+    - export MIOPEN_DEBUG_CONV_GEMM=0
+    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py
+
+- label: "Diffusion Image Edit Test"
+  timeout_in_minutes: 15
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export MIOPEN_DEBUG_CONV_DIRECT=0
+    - export MIOPEN_DEBUG_CONV_GEMM=0
+    - export VLLM_ROCM_USE_AITER=1
+    - export VLLM_ROCM_USE_AITER_MHA=1
+    - export VLLM_ROCM_USE_AITER_LINEAR=0
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - pytest -s -v tests/e2e/online_serving/test_i2i_multi_image_input.py

.buildkite/test-template-amd-omni.j2

@@ -13,7 +13,7 @@
         depends_on: ~
         soft_fail: false
         commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+          - "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
           - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --progress plain ."
           - "docker push {{ docker_image_amd }}"
         key: "amd-build"

.github/ISSUE_TEMPLATE/400-bug-report.yml

@@ -40,6 +40,31 @@ body:
       </details>
   validations:
     required: true
+
+- type: textarea
+  attributes:
+    label: Your code version
+    description: |
+      Please provide a description of which commit id or version the code is.
+
+    value: |
+      <details>
+      <summary>The commit id or version of vllm</summary>
+
+      ```text
+
+      ```
+      </details>
+      <details>
+      <summary>The commit id or version of vllm-omni</summary>
+
+      ```text
+
+      ```
+      </details>
+  validations:
+    required: true
+
 - type: textarea
   attributes:
     label: 🐛 Describe the bug

.github/ISSUE_TEMPLATE/750-RFC.yml

@@ -20,6 +20,8 @@ body:
     label: Proposed Change.
     description: >
       The proposed change of the RFC.
+    value: |
+      Please provide the detailed design document of the RFC using the [template](https://docs.google.com/document/d/12YxSsVeD1jvL-InClkeAEnZyWFDndz_65JmXvsamuV4/edit?tab=t.0#heading=h.4ef4szrsgspp).
   validations:
     required: true
 - type: textarea

README.md

@@ -9,13 +9,14 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> |
 </p>
 
 ---
 
 *Latest News* 🔥
 
+- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm), please check our latest [design](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
 
 ---
@@ -69,6 +70,10 @@ Please check out [Contributing to vLLM-Omni](https://vllm-omni.readthedocs.io/en
 ## Join the Community
 Feel free to ask questions, provide feedbacks and discuss with fellow users of vLLM-Omni in `#sig-omni` slack channel at [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai).
 
+## Star History
+
+[![Star History Chart](https://api.star-history.com/svg?repos=vllm-project/vllm-omni&type=date&legend=top-left)](https://www.star-history.com/#vllm-project/vllm-omni&type=date&legend=top-left)
+
 ## License
 
 Apache License 2.0, as found in the [LICENSE](./LICENSE) file.

benchmarks/diffusion/README.md

@@ -0,0 +1,117 @@
+
+# Diffusion Serving Benchmark (Image/Video)
+
+This folder contains an online-serving benchmark script for diffusion models.
+It sends requests to a vLLM OpenAI-compatible endpoint and reports throughput,
+latency percentiles, and optional SLO attainment.
+
+The main entrypoint is:
+
+- `benchmarks/diffusion/diffusion_benchmark_serving.py`
+
+## 1. Quick Start
+
+1. Start the server:
+
+```bash
+vllm serve Qwen/Qwen-Image --omni --port 8099
+```
+
+2. Run a minimal benchmark:
+
+```bash
+python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+	--base-url http://localhost:8099 \
+	--model Qwen/Qwen-Image \
+	--task t2i \
+	--dataset vbench \
+	--num-prompts 5
+```
+
+**Notes**
+
+- The benchmark talks to `http://<host>:<port>/v1/chat/completions`.
+- If you run the server on another host or port, pass `--base-url` accordingly.
+
+## 2. Supported Datasets
+
+The benchmark supports three dataset modes via `--dataset`:
+
+- `vbench`: Built-in prompt/data loader.
+- `trace`: Heterogeneous request traces (each request can have different resolution/frames/steps).
+- `random`: Synthetic prompts for quick smoke tests.
+
+### VBench dataset
+
+If you use i2v/i2i bench datasets and need auto-download support, you may need:
+
+```bash
+uv pip install gdown
+```
+
+### Trace dataset
+
+Use `--dataset trace` to replay a trace file. The trace can specify per-request fields such as:
+
+- `width`, `height`
+- `num_frames` (video)
+- `num_inference_steps`
+- `seed`, `fps`
+- optional `slo_ms` (per-request SLO target)
+
+By default (when `--dataset-path` is not provided), the script downloads a default trace from
+the HuggingFace dataset repo `asukaqaqzz/Dit_Trace`. The default filename can depend on `--task`
+(e.g., `t2v` uses a video trace).
+
+Current defaults:
+
+- `--task t2i` -> `sd3_trace.txt`
+- `--task t2v` -> `cogvideox_trace.txt`
+
+You can point to your own trace using `--dataset-path`.
+
+## 3. Benchmark Parameters
+
+### Basic flags
+
+- `--base-url`: Server address (the script calls `.../v1/chat/completions`).
+- `--model`: The OpenAI-compatible `model` field.
+- `--task`: Task type (e.g., `t2i`, `t2v`, `i2i`, `i2v`).
+- `--dataset`: Dataset mode (`vbench` / `trace` / `random`).
+- `--num-prompts`: Number of requests to send.
+
+Common optional flags:
+
+- `--output-file`: Write metrics to a JSON file.
+- `--disable-tqdm`: Disable the progress bar.
+
+### Resolution / frames / steps: CLI defaults vs dataset fields
+
+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`).
+
+Precedence rules for `trace` (i.e., what actually gets sent):
+
+- `width/height`: if either `--width` or `--height` is explicitly set, it overrides per-request values from the trace; otherwise per-request values are used when present.
+- `num_frames`: per-request `num_frames` takes precedence; otherwise fall back to `--num-frames`.
+- `num_inference_steps`: per-request `num_inference_steps` takes precedence; otherwise fall back to `--num-inference-steps`.
+
+### SLO, warmup, and max concurrency
+
+Enable SLO evaluation with `--slo`.
+
+- If a request in the trace already has `slo_ms`, that value is used.
+- Otherwise, the script runs warmup requests to infer a base unit time, estimates `expected_ms` by linearly scaling with area/frames/steps, and then sets `slo_ms = expected_ms * --slo-scale`.
+
+Warmup flags:
+
+- `--warmup-requests`: Number of warmup requests.
+- `--warmup-num-inference-steps`: Steps used during warmup.
+- For `--task t2v`: warmup requests are forced to use `num_frames=1` to make warmup faster and less noisy.
+
+Traffic / concurrency flags:
+
+- `--request-rate`: Target request rate (requests/second). If set to `inf`, the script sends all requests immediately.
+- `--max-concurrency`: Max number of in-flight requests (default: `1`). This can hard-cap the achieved QPS: if it is too small, requests will queue behind the semaphore, and both achieved throughput and observed SLO attainment can be skewed.

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -9,28 +9,29 @@ If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
 Usage:
     # Video
     t2v:
-    python3 benchmarks/diffusion_benchmark_serving.py \
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
         --dataset vbench --task t2v --num-prompts 10 \
         --height 480 --width 640 --fps 16 --num-frames 80
 
     i2v:
-    python3 benchmarks/diffusion_benchmark_serving.py \
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
         --dataset vbench --task i2v --num-prompts 10
 
 
     # Image
     t2i:
-    python3 benchmarks/diffusion_benchmark_serving.py \
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
         --dataset vbench --task t2i --num-prompts 10 \
         --height 1024 --width 1024
 
     i2i:
-    python3 benchmarks/diffusion_benchmark_serving.py \
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
         --dataset vbench --task i2i --num-prompts 10
 
 """
 
 import argparse
+import ast
 import asyncio
 import base64
 import glob
@@ -40,6 +41,7 @@ import os
 import time
 import uuid
 from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field, replace
 from typing import Any
 
@@ -60,6 +62,8 @@ class RequestFuncInput:
     num_inference_steps: int | None = None
     seed: int | None = None
     fps: int | None = None
+    timestamp: float | None = None
+    slo_ms: float | None = None
     extra_body: dict[str, Any] = field(default_factory=dict)
     image_paths: list[str] | None = None
     request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
@@ -73,6 +77,7 @@ class RequestFuncOutput:
     start_time: float = 0.0
     response_body: dict[str, Any] = field(default_factory=dict)
     peak_memory_mb: float = 0.0
+    slo_achieved: bool | None = None
 
 
 class BaseDataset(ABC):
@@ -310,6 +315,226 @@ class VBenchDataset(BaseDataset):
         return [self[i] for i in range(len(self))]
 
 
+class TraceDataset(BaseDataset):
+    """Trace-based dataset loader for heterogeneous diffusion requests."""
+
+    DEFAULT_REPO_ID = "asukaqaqzz/Dit_Trace"
+    DEFAULT_FILENAME = "sd3_trace.txt"
+    DEFAULT_FILENAME_BY_TASK: dict[str, str] = {
+        # Text-to-image traces (e.g., SD3)
+        "t2i": "sd3_trace.txt",
+        # Text-to-video traces (e.g., CogVideoX)
+        "t2v": "cogvideox_trace.txt",
+    }
+
+    def __init__(self, args, api_url: str, model: str):
+        super().__init__(args, api_url, model)
+        self.cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "vllm-omni", "trace")
+        self.default_filename = self.DEFAULT_FILENAME_BY_TASK.get(getattr(args, "task", ""), self.DEFAULT_FILENAME)
+        dataset_root = args.dataset_path
+        if not dataset_root:
+            dataset_root = self._download_default_trace()
+        self.items = self._load_items(dataset_root)
+
+    @staticmethod
+    def _coerce_int(x: Any) -> int | None:
+        if x is None:
+            return None
+        if isinstance(x, bool):
+            return None
+        if isinstance(x, int):
+            return x
+        try:
+            s = str(x).strip()
+            if not s:
+                return None
+            return int(float(s))
+        except Exception:
+            return None
+
+    @staticmethod
+    def _coerce_float(x: Any) -> float | None:
+        if x is None:
+            return None
+        if isinstance(x, float):
+            return x
+        if isinstance(x, int):
+            return float(x)
+        try:
+            s = str(x).strip()
+            if not s:
+                return None
+            return float(s)
+        except Exception:
+            return None
+
+    def _download_default_trace(self) -> str:
+        """Download default trace file from HuggingFace Hub if not provided."""
+
+        try:
+            from huggingface_hub import hf_hub_download
+        except ImportError as exc:
+            raise ImportError(
+                "huggingface_hub is required to download the default trace dataset. "
+                "Install via `pip install huggingface_hub`."
+            ) from exc
+
+        os.makedirs(self.cache_dir, exist_ok=True)
+        return hf_hub_download(
+            repo_id=self.DEFAULT_REPO_ID,
+            filename=self.default_filename,
+            repo_type="dataset",
+            local_dir=self.cache_dir,
+            local_dir_use_symlinks=False,
+        )
+
+    def _expand_paths(self, dataset_path: str | None) -> list[str]:
+        if not dataset_path:
+            return []
+
+        parts = [p.strip() for p in str(dataset_path).split(",") if p.strip()]
+        paths: list[str] = []
+        for p in parts:
+            if any(ch in p for ch in ["*", "?", "["]):
+                paths.extend(sorted(glob.glob(p)))
+            elif os.path.isdir(p):
+                paths.extend(sorted(glob.glob(os.path.join(p, "**", "*.txt"), recursive=True)))
+            else:
+                paths.append(p)
+
+        seen = set()
+        unique_paths = []
+        for p in paths:
+            if p not in seen:
+                seen.add(p)
+                unique_paths.append(p)
+        return unique_paths
+
+    def _parse_trace_file(self, path: str) -> list[dict[str, Any]]:
+        rows: list[dict[str, Any]] = []
+
+        def parse_request_repr_line(line: str) -> dict[str, Any] | None:
+            text = line.strip()
+            if not text:
+                return None
+            if not (text.startswith("Request(") and text.endswith(")")):
+                return None
+            inner = text[len("Request(") : -1]
+            try:
+                expr = ast.parse(f"f({inner})", mode="eval")
+                if not isinstance(expr.body, ast.Call):
+                    return None
+                call = expr.body
+                out: dict[str, Any] = {}
+                for kw in call.keywords:
+                    if kw.arg is None:
+                        continue
+                    out[kw.arg] = ast.literal_eval(kw.value)
+                return out
+            except Exception:
+                return None
+
+        # detect first non-empty line to pick parser
+        first_non_empty = None
+        with open(path, encoding="utf-8") as f:
+            for _ in range(50):
+                pos = f.tell()
+                line = f.readline()
+                if not line:
+                    break
+                if line.strip():
+                    first_non_empty = line.strip()
+                    f.seek(pos)
+                    break
+
+        if first_non_empty is None:
+            return rows
+
+        if first_non_empty.startswith("Request("):
+            with open(path, encoding="utf-8") as f:
+                for line in f:
+                    parsed = parse_request_repr_line(line)
+                    if isinstance(parsed, dict):
+                        rows.append(parsed)
+            return rows
+
+        # txt fallback: parse Request(...) lines only
+        with open(path, encoding="utf-8") as f:
+            for line in f:
+                parsed = parse_request_repr_line(line)
+                if isinstance(parsed, dict):
+                    rows.append(parsed)
+        return rows
+
+    def _load_items(self, dataset_root: str) -> list[dict[str, Any]]:
+        paths = self._expand_paths(dataset_root)
+        if not paths:
+            raise ValueError("No trace files found. Provide --dataset-path or rely on default HuggingFace download.")
+
+        items: list[dict[str, Any]] = []
+        for p in paths:
+            if not os.path.exists(p):
+                continue
+            for row in self._parse_trace_file(p):
+                if isinstance(row, dict):
+                    row = dict(row)
+                    row.setdefault("_source", p)
+                    items.append(row)
+
+        if not items:
+            raise ValueError("Trace dataset is empty after parsing provided paths.")
+
+        if self.args.num_prompts is not None:
+            items = items[: self.args.num_prompts]
+
+        return items
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __getitem__(self, idx: int) -> RequestFuncInput:
+        row = self.items[idx]
+        prompt = row.get("prompt") or row.get("text") or ""
+
+        row_height = self._coerce_int(row.get("height"))
+        row_width = self._coerce_int(row.get("width"))
+        num_frames = self._coerce_int(row.get("num_frames"))
+        num_steps = self._coerce_int(row.get("num_inference_steps"))
+        seed = self._coerce_int(row.get("seed"))
+        fps = self._coerce_int(row.get("fps"))
+        timestamp = self._coerce_float(row.get("timestamp"))
+        slo_ms = self._coerce_float(row.get("slo_ms"))
+        image_paths = row.get("image_paths")
+
+        override_w = self.args.width
+        override_h = self.args.height
+        if override_w is not None or override_h is not None:
+            width = override_w
+            height = override_h
+        else:
+            width = row_width
+            height = row_height
+
+        return RequestFuncInput(
+            prompt=str(prompt),
+            api_url=self.api_url,
+            model=self.model,
+            width=width,
+            height=height,
+            num_frames=num_frames if num_frames is not None else self.args.num_frames,
+            num_inference_steps=num_steps if num_steps is not None else self.args.num_inference_steps,
+            seed=seed if seed is not None else self.args.seed,
+            fps=fps if fps is not None else self.args.fps,
+            timestamp=timestamp,
+            slo_ms=slo_ms,
+            image_paths=image_paths,
+            request_id=str(row.get("request_id")) if row.get("request_id") is not None else str(uuid.uuid4()),
+        )
+
+    def get_requests(self) -> list[RequestFuncInput]:
+        return [self[i] for i in range(len(self))]
+
+
 class RandomDataset(BaseDataset):
     def __init__(self, args, api_url: str, model: str):
         self.args = args
@@ -337,6 +562,128 @@ class RandomDataset(BaseDataset):
         return [self[i] for i in range(len(self))]
 
 
+def _compute_expected_latency_ms_from_base(req: RequestFuncInput, args, base_time_ms: float | None) -> float | None:
+    """Compute expected execution time (ms) based on a base per-step-per-frame unit time.
+
+    Assumes linear scaling with pixel area, frame count, and num_inference_steps.
+    The base unit represents latency for a 16x16 resolution, single frame, single step.
+    """
+
+    if base_time_ms is None:
+        return None
+
+    width = req.width if req.width is not None else args.width
+    height = req.height if req.height is not None else args.height
+    if width is None or height is None:
+        return None
+
+    frames = req.num_frames if req.num_frames is not None else args.num_frames
+    steps = req.num_inference_steps if req.num_inference_steps is not None else args.num_inference_steps
+
+    frame_scale = frames if isinstance(frames, int) and frames > 0 else 1
+    step_scale = steps if isinstance(steps, int) and steps > 0 else 1
+
+    area_units = max((float(width) * float(height)) / float(16 * 16), 1.0)
+    return float(base_time_ms) * area_units * frame_scale * step_scale
+
+
+def _infer_slo_base_time_ms_from_warmups(
+    warmup_pairs: list[tuple[RequestFuncInput, RequestFuncOutput]],
+    args,
+) -> float | None:
+    """Infer base SLO unit time from warmup requests.
+
+    Returns the median base latency (ms) for a 16x16 resolution, single-frame,
+    single-step request. Only uses warmups that succeeded and have resolvable
+    width/height.
+    """
+
+    candidates_ms: list[float] = []
+    for req, out in warmup_pairs:
+        if not out.success or out.latency <= 0:
+            continue
+
+        width = req.width if req.width is not None else args.width
+        height = req.height if req.height is not None else args.height
+        if width is None or height is None:
+            continue
+
+        frames = req.num_frames if req.num_frames is not None else args.num_frames
+        steps = req.num_inference_steps if req.num_inference_steps is not None else args.num_inference_steps
+
+        frame_scale = int(frames) if isinstance(frames, int) and frames > 0 else 1
+        step_scale = int(steps) if isinstance(steps, int) and steps > 0 else 1
+
+        area_units = max((float(width) * float(height)) / float(16 * 16), 1.0)
+        denom = area_units * float(frame_scale) * float(step_scale)
+        if denom <= 0:
+            continue
+
+        candidates_ms.append((out.latency * 1000.0) / denom)
+
+    if not candidates_ms:
+        return None
+    return float(np.median(candidates_ms))
+
+
+def _populate_slo_ms_from_warmups(
+    requests_list: list[RequestFuncInput],
+    warmup_pairs: list[tuple[RequestFuncInput, RequestFuncOutput]],
+    args,
+) -> list[RequestFuncInput]:
+    """Populate missing RequestFuncInput.slo_ms using warmup outputs.
+
+    - If a request already has slo_ms (e.g., trace-provided), it is kept as-is.
+    - If any request has slo_ms is None and we can infer base time from warmups,
+      we estimate each missing request's expected execution time and set:
+        req.slo_ms = expected_latency_ms * args.slo_scale
+
+    Returns updated requests_list.
+    """
+
+    if not any(req.slo_ms is None for req in requests_list):
+        return requests_list
+
+    base_time_ms = _infer_slo_base_time_ms_from_warmups(warmup_pairs, args)
+    if base_time_ms is None:
+        return requests_list
+
+    slo_scale = float(getattr(args, "slo_scale", 3.0))
+    if slo_scale <= 0:
+        raise ValueError(f"slo_scale must be positive, got {slo_scale}.")
+
+    updated: list[RequestFuncInput] = []
+    for req in requests_list:
+        if req.slo_ms is not None:
+            updated.append(req)
+            continue
+        expected_ms = _compute_expected_latency_ms_from_base(req, args, base_time_ms)
+        updated.append(replace(req, slo_ms=(expected_ms * slo_scale) if expected_ms is not None else None))
+
+    return updated
+
+
+async def iter_requests(
+    requests_list: list[RequestFuncInput],
+    request_rate: float,
+) -> AsyncGenerator[RequestFuncInput, None]:
+    """Yield requests using a fixed interval if request_rate is set.
+
+    - If request_rate is inf, all requests are yielded immediately (no sleep).
+    - Otherwise, requests are emitted at a fixed cadence of 1 / request_rate seconds.
+    """
+
+    if request_rate != float("inf"):
+        if request_rate <= 0:
+            raise ValueError(f"request_rate must be positive or inf, got {request_rate}.")
+        interval_s = 1.0 / float(request_rate)
+
+    for i, req in enumerate(requests_list):
+        if request_rate != float("inf") and i > 0:
+            await asyncio.sleep(interval_s)
+        yield req
+
+
 def _guess_mime_type(path: str) -> str:
     mime, _ = mimetypes.guess_type(path)
     return mime or "application/octet-stream"
@@ -415,12 +762,21 @@ async def async_request_chat_completions(
 
     output.latency = time.perf_counter() - output.start_time
 
+    if output.success and input.slo_ms is not None:
+        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
+
     if pbar:
         pbar.update(1)
     return output
 
 
-def calculate_metrics(outputs: list[RequestFuncOutput], total_duration: float):
+def calculate_metrics(
+    outputs: list[RequestFuncOutput],
+    total_duration: float,
+    requests_list: list[RequestFuncInput],
+    args,
+    slo_enabled: bool,
+):
     success_outputs = [o for o in outputs if o.success]
     error_outputs = [o for o in outputs if not o.success]
 
@@ -442,6 +798,29 @@ def calculate_metrics(outputs: list[RequestFuncOutput], total_duration: float):
         "peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0,
     }
 
+    if slo_enabled:
+        slo_defined_total = 0
+        slo_met_success = 0
+
+        for req, out in zip(requests_list, outputs):
+            if req.slo_ms is None:
+                continue
+            slo_defined_total += 1
+            if out.slo_achieved is None:
+                continue
+            if out.slo_achieved:
+                slo_met_success += 1
+
+        slo_attain_all = (slo_met_success / slo_defined_total) if slo_defined_total > 0 else 0.0
+
+        metrics.update(
+            {
+                "slo_attainment_rate": slo_attain_all,
+                "slo_met_success": slo_met_success,
+                "slo_scale": getattr(args, "slo_scale", 3.0),
+            }
+        )
+
     return metrics
 
 
@@ -475,6 +854,8 @@ async def benchmark(args):
 
     if args.dataset == "vbench":
         dataset = VBenchDataset(args, api_url, args.model)
+    elif args.dataset == "trace":
+        dataset = TraceDataset(args, api_url, args.model)
     elif args.dataset == "random":
         dataset = RandomDataset(args, api_url, args.model)
     else:
@@ -501,6 +882,7 @@ async def benchmark(args):
     pbar = tqdm(total=len(requests_list), disable=args.disable_tqdm)
 
     async with aiohttp.ClientSession() as session:
+        warmup_pairs: list[tuple[RequestFuncInput, RequestFuncOutput]] = []
         if args.warmup_requests and requests_list:
             print(
                 f"Running {args.warmup_requests} warmup request(s) \
@@ -513,16 +895,20 @@ async def benchmark(args):
                         warm_req,
                         num_inference_steps=args.warmup_num_inference_steps,
                     )
-                await limited_request_func(warm_req, session, None)
+                warm_out = await limited_request_func(warm_req, session, None)
+                warmup_pairs.append((warm_req, warm_out))
+
+        if args.slo:
+            # Prefer trace-provided per-request slo_ms. Only populate when missing.
+            requests_list = _populate_slo_ms_from_warmups(
+                requests_list=requests_list,
+                warmup_pairs=warmup_pairs,
+                args=args,
+            )
 
         start_time = time.perf_counter()
         tasks = []
-        for req in requests_list:
-            if args.request_rate != float("inf"):
-                # Poisson process: inter-arrival times follow exponential distribution
-                interval = np.random.exponential(1.0 / args.request_rate)
-                await asyncio.sleep(interval)
-
+        async for req in iter_requests(requests_list=requests_list, request_rate=args.request_rate):
             task = asyncio.create_task(limited_request_func(req, session, pbar))
             tasks.append(task)
 
@@ -532,7 +918,7 @@ async def benchmark(args):
     pbar.close()
 
     # Calculate metrics
-    metrics = calculate_metrics(outputs, total_duration)
+    metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo)
 
     print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="="))
 
@@ -561,6 +947,12 @@ async def benchmark(args):
     print("{:<40} {:<15.4f}".format("Latency Median (s):", metrics["latency_median"]))
     print("{:<40} {:<15.4f}".format("Latency P99 (s):", metrics["latency_p99"]))
 
+    if args.slo:
+        print(f"{'-' * 50}")
+        print("{:<40} {:<15.2%}".format("SLO Attainment Rate (all):", metrics.get("slo_attainment_rate", 0.0)))
+        print("{:<40} {:<15}".format("SLO Met (success count):", str(metrics.get("slo_met_success", 0))))
+        print("{:<40} {:<15}".format("SLO Scale:", str(metrics.get("slo_scale", 3.0))))
+
     if metrics["peak_memory_mb_max"] > 0:
         print(f"{'-' * 50}")
         print("{:<40} {:<15.2f}".format("Peak Memory Max (MB):", metrics["peak_memory_mb_max"]))
@@ -590,7 +982,7 @@ if __name__ == "__main__":
         "--dataset",
         type=str,
         default="vbench",
-        choices=["vbench", "random"],
+        choices=["vbench", "trace", "random"],
         help="Dataset to use.",
     )
     parser.add_argument(
@@ -645,7 +1037,7 @@ if __name__ == "__main__":
     parser.add_argument(
         "--num-inference-steps",
         type=int,
-        default=None,
+        default=50,
         help="Number of inference steps (for diffusion models).",
     )
     parser.add_argument(
@@ -656,6 +1048,21 @@ if __name__ == "__main__":
     )
     parser.add_argument("--fps", type=int, default=None, help="FPS (for video).")
     parser.add_argument("--output-file", type=str, default=None, help="Output JSON file for metrics.")
+    parser.add_argument(
+        "--slo",
+        action="store_true",
+        help=(
+            "Enable SLO calculation and reporting. If trace provides per-request slo_ms, it is used. "
+            "Otherwise, warmup request(s) are used to infer expected execution time assuming linear "
+            "scaling by resolution, frames, and steps, then slo_ms = expected_time * --slo-scale."
+        ),
+    )
+    parser.add_argument(
+        "--slo-scale",
+        type=float,
+        default=3.0,
+        help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).",
+    )
     parser.add_argument("--disable-tqdm", action="store_true", help="Disable progress bar.")
 
     args = parser.parse_args()

collect_env.py

@@ -57,6 +57,7 @@ SystemEnv = namedtuple(
         "cpu_info",
         "rocm_version",  # vllm specific field
         "vllm_version",  # vllm specific field
+        "vllm_omni_version",  # vllm-omni specific field
         "vllm_build_flags",  # vllm specific field
         "gpu_topo",  # vllm specific field
         "env_vars",
@@ -289,6 +290,31 @@ def get_vllm_version():
     return __version__
 
 
+def get_vllm_omni_version(run_lambda):
+    try:
+        import vllm_omni
+        from vllm_omni import __version__, __version_tuple__
+
+        version_str = __version_tuple__[-1]
+        if isinstance(version_str, str) and version_str.startswith("g"):
+            if "." in version_str:
+                git_sha = version_str.split(".")[0][1:]
+                date = version_str.split(".")[-1][1:]
+                return f"{__version__} (git sha: {git_sha}, date: {date})"
+            else:
+                git_sha = version_str[1:]
+                return f"{__version__} (git sha: {git_sha})"
+
+        package_dir = os.path.dirname(os.path.abspath(vllm_omni.__file__))
+        git_sha = run_and_read_all(run_lambda, f"git -C {package_dir} rev-parse --short HEAD")
+        if git_sha:
+            return f"{__version__} (git sha: {git_sha})"
+
+        return __version__
+    except ImportError:
+        return "N/A (vllm_omni not installed)"
+
+
 def summarize_vllm_build_flags():
     # This could be a static method if the flags are constant, or dynamic if you need to check environment variables, etc.
     return "CUDA Archs: {}; ROCm: {}".format(
@@ -524,6 +550,7 @@ def get_env_info():
 
     rocm_version = get_rocm_version(run_lambda)
     vllm_version = get_vllm_version()
+    vllm_omni_version = get_vllm_omni_version(run_lambda)
     vllm_build_flags = summarize_vllm_build_flags()
     gpu_topo = get_gpu_topo(run_lambda)
 
@@ -555,6 +582,7 @@ def get_env_info():
         cpu_info=get_cpu_info(run_lambda),
         rocm_version=rocm_version,
         vllm_version=vllm_version,
+        vllm_omni_version=vllm_omni_version,
         vllm_build_flags=vllm_build_flags,
         gpu_topo=gpu_topo,
         env_vars=get_env_vars(),
@@ -621,6 +649,7 @@ env_info_fmt += """
 ==============================
 ROCM Version                 : {rocm_version}
 vLLM Version                 : {vllm_version}
+vLLM-Omni Version            : {vllm_omni_version}
 vLLM Build Flags:
   {vllm_build_flags}
 GPU Topology:

docker/Dockerfile.ci

@@ -6,6 +6,12 @@ WORKDIR ${APP_DIR}
 
 COPY . .
 
+# Install system dependencies
+RUN apt-get update && \
+    apt-get install -y ffmpeg && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
 # Install vllm-omni into the same uv-managed Python environment used by the base image.
 RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
 

docs/.nav.yml

@@ -3,7 +3,7 @@ nav:
 - User Guide:
   - Getting Started:
     - getting_started/quickstart.md
-    - getting_started/installation
+    - getting_started/installation/*
   - Serving:
     - OpenAI-Compatible API:
       - Image Generation: serving/image_generation_api.md
@@ -11,6 +11,8 @@ nav:
     - examples/README.md
     - Offline Inference:
       - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
+      - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
+      - LTX-2: user_guide/examples/offline_inference/ltx2.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
       - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
@@ -25,14 +27,16 @@ nav:
   - Configuration:
     - configuration/README.md
     - configuration/*
-  - Diffusion Acceleration:
-    - Overview: user_guide/diffusion_acceleration.md
-    - Acceleration Methods:
-      - TeaCache: user_guide/acceleration/teacache.md
-      - Cache-DiT: user_guide/acceleration/cache_dit_acceleration.md
-      - Parallelism Acceleration: user_guide/acceleration/parallelism_acceleration.md
   - Models:
     - models/supported_models.md
+  - Features:
+    - Sleep Mode: features/sleep_mode.md
+    - Diffusion Features:
+      - Overview: user_guide/diffusion_acceleration.md
+      - TeaCache: user_guide/diffusion/teacache.md
+      - Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
+      - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
+      - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
 - Developer Guide:
   - General:
     - contributing/README.md
@@ -43,13 +47,11 @@ nav:
     - contributing/model/adding_omni_model.md
     - contributing/model/adding_diffusion_model.md
   - CI: contributing/ci
-  - Tests: contributing/tests
   - Design Documents:
     - design/index.md
     - design/architecture_overview.md
     - Feature Design:
       - design/feature/disaggregated_inference.md
-      - design/feature/multi_request_streaming.md
       - design/feature/ray_based_execution.md
     - Module Design:
       - design/module/ar_module.md

docs/README.md

@@ -60,5 +60,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
 
 For more information, checkout the following:
 
-- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1Y7t2Zm3BIISPN-_X_sDpxZMNCWir2rAx/edit?usp=drive_link&ouid=110473603432222024453&rtpof=true&sd=true)
+- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
 - [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)

docs/api/README.md

@@ -82,6 +82,7 @@ Model execution components.
 - [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniMoeTalkerCodePredictor][]
 - [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeModel][]
 - [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeTalkerForConditionalGeneration][]
+- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeTalkerSharedExpertWrapper][]
 - [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMForCausalLM][]
 - [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMModel][]
 - [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeConditionalGenerationMixin][]
@@ -102,8 +103,9 @@ Configuration classes.
 
 Worker classes and model runners for distributed inference.
 
-- [vllm_omni.diffusion.worker.gpu_worker.GPUWorker][]
-- [vllm_omni.diffusion.worker.gpu_worker.WorkerProc][]
+- [vllm_omni.diffusion.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]
+- [vllm_omni.diffusion.worker.gpu_diffusion_worker.GPUDiffusionWorker][]
+- [vllm_omni.diffusion.worker.gpu_diffusion_worker.WorkerProc][]
 - [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorker][]
 - [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorkerProc][]
 - [vllm_omni.worker.gpu_ar_model_runner.ExecuteModelState][]

docs/configuration/README.md

@@ -2,7 +2,7 @@
 
 This section lists the most common options for running vLLM-Omni.
 
-For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.11.0/configuration/index.html)
+For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.12.0/configuration/index.html)
 
 Currently, the main options are maintained by stage configs for each model.
 
@@ -16,6 +16,6 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
 
 ## Optimization Features
 
-- **[TeaCache Configuration](../user_guide/acceleration/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
-- **[Cache-DiT Configuration](../user_guide/acceleration/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
-- **[Parallelism Configuration](../user_guide/acceleration/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models
+- **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
+- **[Cache-DiT Configuration](../user_guide/diffusion/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
+- **[Parallelism Configuration](../user_guide/diffusion/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models

docs/contributing/ci/tests_markers.md

@@ -0,0 +1,160 @@
+# Markers for Tests
+
+By adding markers before test functions, tests can later be executed uniformly by simply declaring the corresponding marker type.
+
+## Current Markers
+Defined in `pyproject.toml`:
+
+| Marker             | Description                                             |
+| ------------------ | ------------------------------------------------------- |
+| `core_model`       | Core model tests (run in each PR)                       |
+| `diffusion`        | Diffusion model tests                                   |
+| `omni`             | Omni model tests                                        |
+| `cache`            | Cache backend tests                                     |
+| `parallel`         | Parallelism/distributed tests                           |
+| `cpu`              | Tests that run on CPU                                   |
+| `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)                 |
+| `npu`              | Tests that run on NPU/Ascend (auto-added)               |
+| `H100`             | Tests that require H100 GPU                             |
+| `L4`               | Tests that require L4 GPU                               |
+| `MI325`            | Tests that require MI325 GPU (AMD/ROCm)                 |
+| `A2`               | Tests that require A2 NPU                               |
+| `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_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_npu`       | Skip if the num of NPU cards is less than the required  |
+| `slow`             | Slow tests (may skip in quick CI)                       |
+| `benchmark`        | Benchmark tests                                         |
+
+For those markers shown as auto-added, they will be added by the `@hardware_test` decorator.
+
+### Example usage for markers
+
+```python
+from tests.utils import hardware_test
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(
+   res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+   num_cards=2,
+)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_video_to_audio()
+    ...
+```
+### Decorator: `@hardware_test`
+
+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`).
+   ```
+   @pytest.mark.cuda
+   @pytest.mark.L4
+   ```
+2. **Handles multi-card (distributed) scenarios**  
+   For tests requiring multiple cards, it automatically adds distributed markers such as `distributed_cuda`, `distributed_rocm`, or `distributed_npu`.
+   ```
+   @pytest.mark.distributed_cuda(num_cards=num_cards)
+   ```
+3. **Supports flexible card requirements**  
+   Accepts `num_cards` as either a single integer for all platforms or as a dictionary with per-platform values. If not specified, defaults to 1 card per platform.
+
+4. **Integrates resource validation**  
+   On CUDA, adds a skip marker (`skipif_cuda`) if the system does not have the required number of devices.
+   Support for `skipif_rocm` and `skipif_npu` will be implemented later.
+
+
+5. **Runs each test in a new process**  
+   Automatically wraps the distributed test with a decorator (`@create_new_process_for_each_test`) to ensure isolation and compatibility with multi-process hardware backends.
+
+6. **Works with pytest filtering**  
+   Allows tests to be filtered and selected at runtime using standard pytest marker expressions (e.g., `-m "distributed_cuda and L4"`).
+
+#### Example usage for decorator
+- Single call for multiple platforms:
+    ```python
+    @hardware_test(
+        res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+        num_cards={"cuda": 2, "rocm": 2, "npu": 2},
+    )
+    ```
+    or
+    ```python
+    @hardware_test(
+        res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+        num_cards=2,
+    )
+    ```
+- `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), NPU (A2/A3)
+- `num_cards` can be int (all platforms) or dict (per platform); defaults to 1 when missing
+- `hardware_test` automatically applies `@create_new_process_for_each_test` for distributed tests.
+- Distributed markers (`distributed_cuda`, `distributed_rocm`, `distributed_npu`) are auto-added for multi-card cases
+- Filtering examples:
+    - CUDA only: `pytest -m "distributed_cuda and L4"`
+    - ROCm only: `pytest -m "distributed_rocm and MI325"`
+    - NPU only: `pytest -m "distributed_npu"`
+
+## Add Support for a New Platform
+
+If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps:
+
+1. **Extend the marker list in your pytest config** so that platform/resource markers are defined:
+   ```toml
+   # In pyproject.toml or pytest.ini
+   [tool.pytest.ini_options]
+   markers = [
+       # ... existing markers ...
+       "tpu: Tests that require TPU device",
+       "TPU_V3: Tests that require TPU v3 hardware",
+       "distributed_tpu: Tests that require multiple TPU devices",
+   ]
+   ```
+2. **Implement a marker construction function for your platform** in `vllm-omni/tests/utils.py`:
+   ```python
+   # In vllm-omni/tests/utils.py
+
+   def tpu_marks(*, res: str, num_cards: int):
+       test_platform = pytest.mark.tpu
+       if res == "TPU_V3":
+           test_resource = pytest.mark.TPU_V3
+       else:
+           raise ValueError(
+               f"Invalid TPU resource type: {res}. Supported: TPU_V3")
+
+       if num_cards == 1:
+           return [test_platform, test_resource]
+       else:
+           test_distributed = pytest.mark.distributed_tpu(num_cards=num_cards)
+           # Optionally: add skipif_tpu when implemented
+           return [test_platform, test_resource, test_distributed]
+   ```
+3. **Update `hardware_test` to recognize your new platform**:
+    In the relevant place (see the `hardware_test` implementation), add:
+    ```python
+    if platform == "tpu":
+        marks = tpu_marks(res=resource, num_cards=cards)
+    ```
+4. **(Recommended) Add a test using your new markers**:
+   ```python
+   @hardware_test(
+       res={"tpu": "TPU_V3"},
+       num_cards=2,
+   )
+   def test_my_tpu_feature():
+       ...
+   ```
+
+**Summary**:  
+- Add pytest markers for your new platform/resources  
+- Implement a marker function (`xxx_marks`)  
+- Plug into `hardware_test`  
+- You're done: tests decorated with `@hardware_test` using your platform now automatically get the correct markers, distribution, and isolation!
+
+See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).

docs/contributing/ci/tests_style.md

@@ -113,7 +113,7 @@ vllm_omni/                          tests/
 
 ### Naming Conventions
 
-- **Unit/System Tests**: Use `test_<module_name>.py` format. Example: `omni_llm.py` → `test_omni_llm.py`
+- **Unit Tests**: Use `test_<module_name>.py` format. Example: `omni_llm.py` → `test_omni_llm.py`
 
 - **E2E Tests**: Place in `tests/e2e/offline_inference/` or `tests/e2e/online_serving/` with descriptive names. Example: `tests/e2e/offline_inference/test_qwen3_omni.py`, `tests/e2e/offline_inference/test_diffusion_model.py`
 
@@ -139,7 +139,7 @@ vllm_omni/                          tests/
 4. **Documentation**: Add docstrings to all test functions
 5. **Environment variables**: Set uniformly in `conftest.py` or at the top of files
 6. **Type annotations**: Add type annotations to all test function parameters
-7. **Resources**, Using pytest tag to specify the computation resources the test required.
+7. **Pytest Markers**: Add necessary markers like `@pytest.mark.core_model` and use `@hardware_test` to declare hardware requirements (check detailed in [Markers for Tests](../ci/tests_markers.md)).
 
 ### Template
 #### E2E - Online serving
@@ -155,6 +155,7 @@ from pathlib import Path
 import pytest
 import openai
 
+from tests.utils import hardware_test
 
 # Optional: set process start method for workers
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -184,6 +185,12 @@ def base64_encoded_video() -> str:
 def dummy_messages_from_video_data(video_data_url: str, content_text: str) -> str:
     xxx
 
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(
+    res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+    num_cards={"cuda": 2, "rocm": 2, "npu": 4},
+)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_video_to_audio(
     client: openai.OpenAI,
@@ -226,6 +233,7 @@ from pathlib import Path
 import pytest
 from vllm.assets.video import VideoAsset
 
+from tests.utils import hardware_test
 from ..multi_stages.conftest import OmniRunner
 
 # Optional: set process start method for workers
@@ -239,7 +247,12 @@ test_params = [(model, stage_config) for model in models for stage_config in sta
 
 # function name: test_{input_modality}_to_{output_modality}
 # modality candidate: text, image, audio, video, mixed_modalities
-@pytest.mark.gpu_mem_high  # requires high-memory GPU node
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(
+    res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+    num_cards=2,
+)
 @pytest.mark.parametrize("test_config", test_params)
 def test_video_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
     """Offline inference: video input, audio output."""
@@ -263,4 +276,5 @@ def test_video_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
 
 1. The file is saved in an appropriate place and the file name is clear.
 2. The coding style follows the requirements outlined above.
-3. For e2e model test, please ensure the test is configured under the `./buildkite/` folder.
+3. **All test functions have appropriate pytest markers**
+4. For tests that need run in CI, please ensure the test is configured under the `./buildkite/` folder.

docs/contributing/model/adding_diffusion_model.md

@@ -140,7 +140,7 @@ Key point for writing the example:
 + Save or display the generated results so users can validate the integration.
 
 # Testing
-For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../tests/tests_style.md).
+For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
 
 
 ## Adding a Model Recipe

docs/contributing/model/adding_omni_model.md

@@ -572,7 +572,7 @@ def talker2code2wav(
 
 ## Testing
 
-For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../tests/tests_style.md).
+For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
 
 ## Adding a Model Recipe
 

docs/contributing/profiling.md

@@ -1,3 +1,99 @@
 # Profiling vLLM-Omni
 
-Profiling is only intended for vLLM-Omni developers and maintainers to understand the proportion of time spent in different parts of the codebase. **vLLM-Omni end-users should never turn on profiling** as it will significantly slow down the inference.
+Performance Profiling Guidelines Profiling capabilities in vLLM-Omni are reserved for development and maintenance tasks aimed at temporal analysis of the codebase. Production use is **strongly discouraged**; enabling the profiler incurs a substantial overhead that negatively impacts inference latency.
+
+**Mechanism**: vLLM-Omni implements cross-stage profiling via the PyTorch Profiler. To accommodate the architecture—where stages operate as distinct engine instances in separate processes—the profiling interface supports both holistic capturing (all stages) and targeted capturing (specific stages).
+
+**1. Enabling the Profiler**
+
+Before running your script, you must set the ```VLLM_TORCH_PROFILER_DIR``` environment
+variable.
+
+```Bash
+export VLLM_TORCH_PROFILER_DIR=/path/to/save/traces
+```
+
+**Highly Recommended: Limit Profiling to a Single Iteration**  
+For most use cases (especially when profiling audio stages), you should limit the profiler to just **one iteration** to keep trace files small and readable.
+
+
+```bash
+export VLLM_PROFILER_MAX_ITERS=1
+```
+
+**2. Offline Inference**
+
+For offlinie processing using ```OmniLLM```, you can wrap your ```generate``` calls with ```start_profile``` and ```stop_profile()```.
+
+Basic Usage(All Stages)
+```Python
+from vllm_omni import OmniLLM
+
+omni_llm = OmniLLM.from_engine_args(engine_args)
+
+# Start profiling all active stages
+omni_llm.start_profile()
+
+outputs = omni_llm.generate(prompts, sampling_params)
+
+# Stop profiling and save traces
+omni_llm.stop_profile()
+```
+
+**Selective Stage Profiling**
+The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
+```python
+# Only profile Stage 1
+omni_llm.start_profile(stages=[1])
+```
+
+```python
+# Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
+omni_llm.start_profile(stages=[0, 2])
+```
+
+**Examples**:
+
+1. **Qwen-omni 2.5**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
+
+2. **Qwen-omni 3.0**:   [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
+
+**3. Online Inference(Async)**
+
+For online serving using AsyncOmni, the methods are asynchronous. This allows you to toggle profiling dynamically without restarting the server.
+
+```python
+from vllm_omni import AsyncOmni
+
+# Inside an async function:
+async_omni = AsyncOmni.from_engine_args(engine_args)
+
+await async_omni.start_profile()
+
+async for output in async_omni.generate(prompt, sampling_params, request_id):
+    # Process outputs...
+    pass
+
+await async_omni.stop_profile()
+```
+
+**4. Analyzing Omni Traces**
+
+After ``stop_profile()`` completes (and the file write wait time has passed), the directory specified in ```VLLM_TORCH_PROFILER_DIR``` will contain the trace files.
+
+```
+Output/
+│── ...rank-0.pt.trace.json.gz   # GPU 0 trace (TP=2 Example)
+│── ...rank-1.pt.trace.json.gz   # GPU 1 trace (TP=2 Example)
+│       # Load these into Perfetto to visualize synchronization
+│
+│── profiler_out_.txt            # Summary tables (CPU/CUDA time %)
+```
+
+**Viewing Tools:**
+     - [Perfetto](https://ui.perfetto.dev/): (Recommended): Best for handling large audio trace files.
+     - ```chrome://tracing```: Good for smaller text-only traces.
+
+
+**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM.  
+For more advanced configuration options (memory profiling, custom activities, etc.), see the official vLLM profiler documentation:  [vLLM Profiling Guide](https://docs.vllm.ai/en/latest/dev/profiling.html)

docs/design/architecture_overview.md

@@ -64,13 +64,13 @@ According to analysis for current popular open-source models, most of them have
 
 ## Key Components
 
-| Component | Description |
-|-----------|-------------|
-| **OmniRouter** | provide an intelligent router for Omni-modality requests dispatch |
-| **EntryPoints** | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni) and provide the OmniStage abstraction for different AR/DiT stages |
-| **AR** | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management |
-| **Diffusion** | natively implemented and optimized using acceleration components |
-| **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages |
+| Component         | Description                                                                                                                              |
+| ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
+| **OmniRouter**    | provide an intelligent router for Omni-modality requests dispatch                                                                        |
+| **EntryPoints**   | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni) and provide the OmniStage abstraction for different AR/DiT stages |
+| **AR**            | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management                                 |
+| **Diffusion**     | natively implemented and optimized using acceleration components                                                                         |
+| **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages                    |
 
 Disaggregated stages are managed through configuration, such as in the Qwen3-Omni example, where stages like Thinker, Talker, and Code2wav are defined as separate OmniStage instances with specific resources and input/output type.
 
@@ -192,4 +192,4 @@ curl -sS -X POST http://localhost:8091/v1/chat/completions \
 }
 ```
 
-For more usages, please refer to [examples](../user_guide/examples/).
+For more usages, please refer to [examples](../examples/README.md).

docs/design/feature/multi_request_streaming.md

@@ -1,47 +0,0 @@
-## Multi-Request Streaming (MRS) on a Single Machine
-
-### 1. Background & Scope
-- All processing runs on a single physical machine with multi-process, per-stage workers. No proxy or network transport involved.
-- Current alignment with vllm-omni: `OmniLLM` supports multiple stages (`OmniStage`). GPU runners already expose streamable steps (prefill/decoding/diffusion), but the entry layer still collects lists and lacks intra-stage streaming and window scheduling.
-- Goal: implement multi-stage, multi-request streaming (MRS) locally. Each stage outputs segments; downstream stages stitch and trigger compute based on configured windows. Shared memory and zero-copy strategies reduce data movement overhead.
-
-### 2. Key Constraints
-- Multi-process per stage: each stage is an independent process with a while loop; device visibility can be configured (`CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`).
-- Simple IPC (copy-based): use `multiprocessing.Queue`/Pipe for inter-process communication with CPU copies/serialization; do not rely on CUDA IPC/SHM zero-copy in this version.
-- Cross-stage pipeline: different stages can process different requests concurrently (e.g., stage A handles request 1 while stage B handles request 0).
-
-### 3. Architecture Overview
-- Processes & IPC queues
-  - Each "sub-stage" is an OS process (worker). The loop: take from input_queue → compute → put to output_queue.
-  - Inter-stage connection via IPC: copy-based `multiprocessing.Queue` passing dict payloads; use shared memory for large objects.
-  - Each link is SPSC (single-producer/single-consumer): the upstream is the orchestrator and the downstream is a single stage process; queues are unbounded (maxsize=0) on the orchestrator side.
-- Device visibility
-  - Each stage sets `CUDA_VISIBLE_DEVICES` or calls `torch.cuda.set_device` to bind to GPU sets.
-  - A stage may use multiple GPUs internally (TP/PP/DP) but presents as a single stage unit.
-- Simplified IPC: copy-based queues/pipes for data transfer; zero-copy is future work.
-- Pipeline progression: when a stage finishes a request, it enqueues outputs to the downstream stage; if downstream is idle, it starts immediately.
-- Scheduling
-  - A downstream stage triggers only after the upstream completes the request.
-  - Windowed segmentation/stitched triggering is not implemented; intra-stage streaming is not provided.
-
-### 4. IPC Implementation (simplified: copy-based)
-- Use `multiprocessing.Queue`/Pipe for inter-process communication (control + data).
-- Data is serialized/copied via CPU; no CUDA IPC/SHM zero-copy in this version.
-- Backpressure: queues are unbounded; pressure manifests as compute-rate differences. Optional SHM reduces large-object transfer cost; RX/decoding overhead is recorded for observability.
-
-### 5. Scheduling & Cancellation (simplified)
-- Pipeline: when a stage finishes a request, it enqueues to the next stage; that stage immediately pulls the next request from its input queue, enabling cross-stage concurrency.
-- Cancellation/timeout: explicit cancellation/timeouts are not provided; graceful shutdown uses a `None` sentinel sent to each stage input queue.
-
-#### Short sequence example (req0/req1, stage A→B)
-1) t0: stage A handles req0
-2) t1: req0 completes on A → enters B; A immediately starts req1
-3) t2: B handles req0 while A handles req1 (parallel across stages)
-
-### 6. Integration Points (by file)
-- `vllm_omni/entrypoints/omni.py` (Orchestrator)
-  - Class `Omni` orchestrates multi-process stages; constructs `OmniStage` instances in parallel and spawns per-stage workers.
-  - Spawns stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), creates control/data channels, builds simple full-trigger flow.
-  - Stats/logging are disabled by default; per-stage and orchestrator stats are only written when explicitly enabled.
-  - Manages process lifecycle: start/wait for readiness, graceful shutdown; forwards results between stages using copy-based IPC and optional SHM.
-  - Stage readiness: each stage emits `{"type": "stage_ready"}` after initialization; the orchestrator waits for all stages or times out and logs diagnostic suggestions.

docs/design/feature/ray_based_execution.md

@@ -1,14 +1,17 @@
 # Distributed utils
 
 This directory (vllm_omni/distributed/ray_utils) contains utilities for distributed execution in vllm-omni, supporting both **Ray** and **Multiprocessing** backends.
-
-## 1. Ray Utils
+## 1. Installation
+```bash
+pip install "ray[default]"
+```
+## 2. Ray Utils
 
 The `ray_utils` module provides helper functions for managing Ray clusters and actors, which is essential for:
 *   **Multi-node deployment**: Running pipeline stages across different physical machines.
 *   **Resource management**: Efficient GPU/CPU allocation.
 
-### 1.1 Basic Usage
+### 2.1 Basic Usage
 
 To use the Ray backend, specify `worker_backend="ray"` when initializing the engine.
 
@@ -21,7 +24,7 @@ vllm serve Qwen/Qwen2.5-Omni-7B \
   --ray-address auto
 ```
 
-### 1.2 Cluster Setup
+### 2.2 Cluster Setup
 
 **Step 1: Start Head Node**
 Run this on your primary machine:
@@ -38,7 +41,7 @@ ray start --address=<HEAD_NODE_IP>:6399
 > **Tip**: For a complete cluster setup script, refer to the vLLM example:
 > [run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/run_cluster.sh)
 
-### 1.3 Distributed Connector Support
+### 2.3 Distributed Connector Support
 
 When running on Ray, the system automatically adapts its communication strategy:
 
@@ -46,16 +49,11 @@ When running on Ray, the system automatically adapts its communication strategy:
 *   **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
 *   **SHM threshold default differs**: when `worker_backend="ray"`, the SharedMemoryConnector default threshold is set to `sys.maxsize`, which forces payloads to go inline (no SHM). Override `shm_threshold_bytes` in the connector config if you want SHM for Ray runs.
 
-### 1.4 Internal Helpers
+### 2.4 Internal Helpers
 
 *   **`initialize_ray_cluster`**: Connects to an existing Ray cluster or starts a local one.
 
-## 2. Troubleshooting
+## 3. Troubleshooting
 
 *   **Connection Issues**: Ensure the Ray head node is accessible and ports (default 6399 in this example) are open.
 *   **Version Mismatch**: Ensure all nodes run the same version of Ray and Python.
-
-### Installation
-```bash
-pip install "ray[default]"
-```

docs/design/index.md

@@ -9,7 +9,6 @@ This section contains design documents and architecture specifications for vLLM-
 ## Feature Design Documents
 
 - [Disaggregated Inference](feature/disaggregated_inference.md)
-- [Multi-Request Streaming](feature/multi_request_streaming.md)
 - [Ray-based Execution](feature/ray_based_execution.md)
 
 ## Module Design Documents

docs/design/module/dit_module.md

@@ -1,3 +1,7 @@
+---
+toc_depth: 4
+---
+
 # Diffusion Module Architecture Design
 
 The vLLM-Omni diffusion module (`vllm_omni/diffusion`) is a high-performance inference engine for diffusion models, designed with a modular architecture that separates concerns across multiple components. It provides efficient execution for non-autoregressive generation tasks such as image and video generation.
@@ -5,7 +9,7 @@ The vLLM-Omni diffusion module (`vllm_omni/diffusion`) is a high-performance inf
 This document describes the architecture design of the diffusion module, including the diffusion engine, scheduler, worker, diffusion pipeline, and acceleration components.
 
 <p align="center">
-   <img src="https://github.com/user-attachments/assets/cde1ebea-4006-4d30-8283-15a5721a09d2" alt="vLLM-Omni Diffusion Module Components" width="60%">
+   <img src="https://github.com/user-attachments/assets/cde1ebea-4006-4d30-8283-15a5721a09d2" alt="vLLM-Omni Diffusion Module Components" width="80%">
 </p>
 <p align="center">
   <em> Main Components of the Diffusion Module </em>
@@ -20,10 +24,10 @@ This document describes the architecture design of the diffusion module, includi
 - [Worker](#3-worker)
 - [Diffusion Pipeline](#4-diffusion-pipeline)
 - [Acceleration Components](#5-acceleration-components)
-   - [Attention Backends](#51-attention-backends)
-   - [Parallel Attention](#52-parallel-attention)
-   - [Cache Backends](#53-cache-backends)
-   - [Parallel Strategies](#54-parallel-strategies)
+    - [Attention Backends](#51-attention-backends)
+    - [Parallel Attention](#52-parallel-attention)
+    - [Cache Backends](#53-cache-backends)
+    - [Parallel Strategies](#54-parallel-strategies)
 - [Data Flow](#6-data-flow)
 
 ---
@@ -65,8 +69,11 @@ class DiffusionEngine:
 ```
 
 **Key Features**:
+
 - **Pre/Post Processing**: Registers model-specific pre-processing and post-processing functions via registry pattern
+
 - **Worker Management**: Launches and manages multiple worker processes (one per GPU)
+
 - **Process Isolation**: Uses multiprocessing for true parallelism
 
 #### 1.2 Worker Launch Process
@@ -86,8 +93,11 @@ def _launch_workers(self, broadcast_handle):
 ```
 
 **Design Decisions**:
+
 - **Spawn Method**: Ensures clean state for each worker (no shared memory issues)
+
 - **Pipe Communication**: Uses `mp.Pipe` for initialization handshake
+
 - **Device Selection**: Each worker is assigned a specific GPU (`cuda:{rank}`)
 
 #### 1.3 Request Processing Flow
@@ -106,8 +116,11 @@ def step(self, requests: list[OmniDiffusionRequest]):
 ```
 
 **Flow**:
+
 1. **Pre-processing**: Applies model-specific transformations
+
 2. **Scheduling**: Delegates to scheduler for distribution
+
 3. **Post-processing**: Converts raw outputs to final format (e.g., PIL images)
 
 ---
@@ -139,8 +152,11 @@ class Scheduler:
 ```
 
 **Communication Pattern**:
+
 - **Broadcast Queue**: One-to-many communication (scheduler → all workers)
+
 - **Result Queue**: One-to-one communication (rank 0 → scheduler)
+
 - **Shared Memory**: Uses `MessageQueue` (ZMQ-based) for efficient IPC
 
 #### 2.2 Request Distribution
@@ -156,8 +172,11 @@ def add_req(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
 ```
 
 **Design Features**:
+
 - **Broadcast Model**: All workers receive the same request (for tensor parallelism)
+
 - **Single Response**: Only rank 0 sends results back (avoids duplicate outputs)
+
 - **Synchronous**: Blocks until result is received (can be made async)
 
 #### 2.3 Singleton Pattern
@@ -176,8 +195,11 @@ scheduler = Scheduler()
 ```
 
 **Benefits**:
+
 - **Single Point of Control**: Ensures consistent state
+
 - **Easy Access**: Global `scheduler` instance accessible everywhere
+
 - **Resource Management**: Centralized queue management
 
 ---
@@ -212,12 +234,19 @@ class WorkerProc:
 ```
 
 **Initialization Steps**:
+
 1. **IPC Setup**: Creates ZMQ context and message queues
+
 2. **Distributed Environment Setup**: Initializes PyTorch distributed communication
-     - For CUDA GPUs: Uses NCCL (fast GPU communication)
-     - For NPU: Uses HCCL (Huawei Collective Communications Library)
-     - For other devices: Uses appropriate backend (GLOO, MCCL, etc.)
+
+    - For CUDA GPUs: Uses NCCL (fast GPU communication)
+
+    - For NPU: Uses HCCL (Huawei Collective Communications Library)
+
+    - For other devices: Uses appropriate backend (GLOO, MCCL, etc.)
+
 3. **Model Loading**: Loads diffusion pipeline on assigned GPU
+
 4. **Cache Setup**: Enables cache backend if configured.
 
 #### 3.2 GPU Worker
@@ -253,8 +282,11 @@ class GPUWorker:
 ```
 
 **Key Features**:
+
 - **Tensor Parallelism**: Supports multi-GPU tensor parallelism via PyTorch distributed
+
 - **Model Loading**: Uses `DiffusersPipelineLoader` for efficient weight loading
+
 - **Cache Integration**: Enables cache backends (TeaCache, cache-dit, etc.) transparently
 
 #### 3.3 Worker Busy Loop
@@ -283,9 +315,13 @@ def worker_busy_loop(self):
 ```
 
 **Execution Flow**:
+
 1. **Receive**: Dequeues unified messages from shared memory queue
+
 2. **Route**: Handles different message types (generation, RPC, shutdown)
+
 3. **Execute**: Runs forward pass through pipeline for generation requests
+
 4. **Respond**: Sends results back (rank 0 for generation, specified rank for RPC)
 
 #### 3.4 Model Execution
@@ -315,10 +351,14 @@ The model execution leverages multiple parallelism strategies that are transpare
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_sp_group, get_dp_group, get_cfg_group, get_pp_group
 )
+```
 
 **Optimizations**:
+
 - **Cache Refresh**: Clears cache state before each generation for clean state
+
 - **Context Management**: Forward context ensures parallel groups are available during execution
+
 - **Single Request**: Currently processes one request at a time (batching TODO)
 
 ---
@@ -327,66 +367,9 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 
 **Location**: `vllm_omni/diffusion/models/*/pipeline_*.py`
 
-### Architecture
-
 The pipeline is the **model-specific implementation** that orchestrates the diffusion process. Different models (QwenImage, Wan2.2, Z-Image) have their own pipeline implementations.
 
-### Key Components (QwenImage Example)
-
-#### 4.1 Pipeline Structure
-
-```python
-class QwenImagePipeline(nn.Module):
-    def __init__(self, od_config: OmniDiffusionConfig):
-        # Core components
-        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(...)
-        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(...)
-        self.vae = AutoencoderKLQwenImage.from_pretrained(...)
-        self.transformer = QwenImageTransformer2DModel(od_config=od_config)
-        self.tokenizer = Qwen2Tokenizer.from_pretrained(...)
-
-        # Cache backend (set by worker)
-        self._cache_backend = None
-```
-
-**Component Responsibilities**:
-- **Text Encoder**: Encodes prompts into embeddings
-- **Transformer**: Main denoising network (DiT-style)
-- **VAE**: Encodes/decodes between pixel and latent space
-- **Scheduler**: Manages timestep scheduling
-
-#### 4.2 Forward Pass
-
-```python
-def forward(self, req: OmniDiffusionRequest, ...) -> DiffusionOutput:
-    # 1. Encode prompts
-    prompt_embeds, prompt_embeds_mask = self.encode_prompt(prompt)
-
-    # 2. Prepare latents
-    latents = self.prepare_latents(batch_size, num_channels, height, width, ...)
-
-    # 3. Prepare timesteps
-    timesteps, num_inference_steps = self.prepare_timesteps(...)
-
-    # 4. Diffusion loop
-    latents = self.diffuse(
-        prompt_embeds, prompt_embeds_mask,
-        negative_prompt_embeds, negative_prompt_embeds_mask,
-        latents, img_shapes, txt_seq_lens,
-        timesteps, do_true_cfg, guidance, true_cfg_scale,
-    )
-
-    # 5. Decode latents
-    image = self.vae.decode(latents)
-
-    return DiffusionOutput(output=image)
-```
-
-The main steps of the forward pass are referred from `diffusers`.
-
-#### 4.3 Diffusion Loop
-
-The multi-step diffusion loop is usually the most time-consuming part during the overall inference process.
+Most pipeline implementation are referred from `diffusers`. The multi-step diffusion loop is usually the most time-consuming part during the overall inference process, which is defined by the `diffuse` function in the pipeline class. An example is as follows:
 
 ```python
 def diffuse(self, ...):
@@ -416,11 +399,14 @@ def diffuse(self, ...):
 ```
 
 **Key Features**:
+
 - **CFG Support**: Handles classifier-free guidance with separate forward passes
+
 - **Cache Branching**: Uses `cache_branch` parameter for cache-aware execution
+
 - **True CFG**: Implements advanced CFG with norm preservation
 
-To learn how to add a new diffusion pipeline, please view [Adding Diffusion Model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_diffusion_model)
+To learn more about the diffusion pipeline and how to add a new diffusion pipeline, please view [Adding Diffusion Model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_diffusion_model)
 
 ---
 
@@ -448,9 +434,13 @@ class Attention(nn.Module):
 ```
 
 **Available Backends**:
+
 - **FlashAttention**: Optimized CUDA kernel (FA2/FA3) - memory efficient via tiling
+
 - **SDPA**: PyTorch's scaled dot-product attention - default, cross-platform
+
 - **SageAttention**: Sparse attention implementation from SageAttention library
+
 - **AscendAttention**: NPU-optimized attention for Ascend hardware
 
 These backends provide the **kernel implementations** for attention computation. For attention-level sequence parallelism strategies (Ring Attention, Ulysses), see [Parallel Attention](#52-parallel-attention).
@@ -470,16 +460,25 @@ def get_attn_backend(head_size: int) -> type[AttentionBackend]:
 ```
 
 **Selection Priority**:
+
 1. **Environment Variable**: `DIFFUSION_ATTENTION_BACKEND` for manual override
-   - Valid values: `FLASH_ATTN`, `TORCH_SDPA`, `SAGE_ATTN`, `ASCEND`
-   - Example: `export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN`
+
+    - Valid values: `FLASH_ATTN`, `TORCH_SDPA`, `SAGE_ATTN`, `ASCEND`
+
+    - Example: `export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN`
+
 2. **Automatic Fallback**: Falls back to SDPA if selected backend unavailable
+
 3. **Hardware Detection**: Can select based on device type (NPU, CUDA, etc.)
 
 **Backend Availability**:
+
 - **SDPA**: Always available (PyTorch built-in)
+
 - **FlashAttention**: Requires `flash-attn` package installed
+
 - **SageAttention**: Requires `sage-attention` package (from THU-ML GitHub)
+
 - **AscendAttention**: Only available on Ascend NPU hardware
 
 #### Attention Backend Registry
@@ -563,8 +562,11 @@ Parallel attention strategies implement **Sequence Parallelism (SP) at the atten
 **Key Distinction**: Unlike AttentionBackend (which provides kernel implementations), ParallelAttentionStrategy provides communication patterns for multi-GPU attention parallelism. These strategies implement the `ParallelAttentionStrategy` interface and use AttentionBackend implementations internally.
 
 Both Ring Attention and Ulysses are forms of Sequence Parallelism (SP) that:
+
 - Split the sequence dimension across GPUs
+
 - Contribute to `sequence_parallel_size` (via `ring_degree` and `ulysses_degree`)
+
 - Work at the attention layer level (not model/pipeline level)
 
 #### Ulysses Sequence Parallelism (USP)
@@ -574,7 +576,9 @@ Both Ring Attention and Ulysses are forms of Sequence Parallelism (SP) that:
 USP is a sequence-parallel attention strategy that splits attention computation across multiple GPUs by distributing both the sequence dimension and attention heads. It uses **all-to-all communication** to efficiently parallelize attention for very long sequences. Specifically, it uses **all-to-all** collective operations to redistribute Q/K/V tensors before attention computation and gather results afterward.
 
 Ulysses splits attention computation in two dimensions:
+
 1. **Sequence Dimension**: Splits the sequence length across GPUs
+
 2. **Head Dimension**: Splits attention heads across GPUs
 
 **Configuration**: `ulysses_degree` contributes to `sequence_parallel_size`
@@ -601,9 +605,13 @@ class RingParallelAttention:
 ```
 
 **Integration**:
+
 - Ring Attention is activated when `ring_degree > 1` in parallel config
+
 - It's selected by `build_parallel_attention_strategy()` in the attention layer
+
 - The `Attention` layer routes to `_run_ring_attention()` when Ring is enabled
+
 - Works alongside attention backends: Ring handles communication, backends handle computation
 
 **Configuration**: `ring_degree` contributes to `sequence_parallel_size`
@@ -611,8 +619,11 @@ class RingParallelAttention:
 #### Relationship with AttentionBackend
 
 Parallel attention strategies (Ring, Ulysses) work **on top of** AttentionBackend implementations:
+
 - They use AttentionBackend for the actual attention computation (FlashAttention, SDPA, etc.)
+
 - They handle the multi-GPU communication/parallelization layer
+
 - They implement `ParallelAttentionStrategy` interface (not `AttentionBackend`)
 
 For general parallelism strategies (Data Parallelism, Tensor Parallelism, Pipeline Parallelism), see [Parallel Strategies](#54-parallel-strategies).
@@ -651,8 +662,11 @@ class CacheBackend(ABC):
 ```
 
 **Design Pattern**:
+
 - **Abstract Base Class**: Defines contract for all cache backends
+
 - **Pipeline-based**: Works with pipeline instances (not just transformers)
+
 - **State Management**: Provides refresh mechanism for clean state between generations
 
 #### Available Backends
@@ -686,9 +700,13 @@ class TeaCacheBackend(CacheBackend):
 ```
 
 **TeaCache Features**:
+
 - **Timestep-aware**: Caches based on timestep embedding similarity
+
 - **Adaptive**: Dynamically decides when to reuse cached computations
+
 - **CFG-aware**: Handles positive/negative branches separately
+
 - **Custom Hook System**: Uses a custom forward interception mechanism (via `HookRegistry`) that wraps the module's `forward` method, allowing transparent integration without modifying model code
 
 **2. Cache-DiT Backend**
@@ -710,9 +728,13 @@ class CacheDiTBackend(CacheBackend):
 ```
 
 **Cache-DiT Features**:
+
 - **DBCache**: Dynamic block caching with configurable compute blocks
+
 - **SCM**: Step Computation Masking for additional speedup
+
 - **TaylorSeer**: Advanced calibration for cache accuracy
+
 - **Dual-transformer Support**: Handles models like Wan2.2 with two transformers
 
 #### Cache Backend Selector
@@ -748,9 +770,13 @@ def get_cache_backend(
 ```
 
 **Usage Flow**:
+
 1. **Selection**: `get_cache_backend()` returns appropriate backend instance
+
 2. **Enable**: `backend.enable(pipeline)` called during worker initialization
+
 3. **Refresh**: `backend.refresh(pipeline, num_inference_steps)` called before each generation
+
 4. **Check**: `backend.is_enabled()` verifies cache is active
 
 ### 5.4 Parallel Strategies
@@ -762,23 +788,35 @@ def get_cache_backend(
 The system supports multiple orthogonal parallelism strategies:
 
 **Sequence Parallelism (SP)**
+
 - **Purpose**: Split sequence dimension across GPUs
+
 - **Attention-level SP**: Ring Attention and Ulysses (USP) implement SP at the attention layer level
-  - See [Parallel Attention](#52-parallel-attention) for details
-  - Configuration: `ulysses_degree` × `ring_degree` = `sequence_parallel_size`
+
+    - See [Parallel Attention](#52-parallel-attention) for details
+
+    - Configuration: `ulysses_degree` × `ring_degree` = `sequence_parallel_size`
+
 - **Use Case**: Very long sequences (e.g., high-resolution images)
 
 **Data Parallelism (DP)**
+
 - **Purpose**: Replicate model across GPUs, split batch
+
 - **Use Case**: Batch processing, throughput optimization
 
 **Tensor Parallelism (TP)** (Experimental)
+
 - **Purpose**: Split model weights across GPUs
+
 - **Implementation**: Uses vLLM's tensor parallel groups
+
 - **Use Case**: Large models that don't fit on single GPU
 
 **CFG Parallelism**  (under development)
+
 - **Purpose**: Parallelize Classifier-Free Guidance (positive/negative prompts)
+
 - **Infrastructure**: CFG parallel groups are initialized and available via `get_cfg_group()`
 
 #### Parallel Group Management

docs/features/sleep_mode.md

@@ -0,0 +1,39 @@
+# Sleep Mode
+
+vLLM-Omni’s **Sleep Mode** allows you to temporarily release most GPU memory used by a model—such as model weights and key-value (KV) caches (for autoregressive models)—**without stopping the server or unloading the Docker container**.
+
+This feature is inherited from [vLLM’s Sleep Mode](https://blog.vllm.ai/2025/10/26/sleep-mode.html), which provides zero-reload model switching for multi-model serving.  
+
+It is especially useful in **RLHF**, **training**, or **cost-saving scenarios**, where GPU resources must be freed between inference workloads.
+
+---
+
+## Omni Model
+
+Omni model inherit the feature from vLLM' Sleep Mode
+
+This means:
+
+- Support both Level 1 and Level 2 sleep, allow to release and reset both model weights and KV Cache
+
+## Diffusion Model Extension
+
+We added Sleep Mode support for **diffusion models**, which previously lacked this functionality.  
+In diffusion pipelines, this currently only offloads **model weight memory**, as these models typically do not use KV caches.
+
+This means:
+
+- Diffusion models can now enter Level 1 sleep.
+- Pipeline states (e.g., noise schedulers, buffers) remain intact after waking.
+- Useful for releasing VRAM between image generation or training cycles.
+
+---
+
+## Enable sleep mode
+To enable sleep mode, set the `enable_sleep_mode` in `engine_args` to `True`
+
+
+Example:
+```python
+omni = Omni(model=...,enable_sleep_mode=True)
+```

docs/getting_started/installation/gpu.md

@@ -49,15 +49,18 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ## Set up using Docker
 
+### Pre-built images
+
+=== "NVIDIA CUDA"
+
+    --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-images"
+
+=== "AMD ROCm"
+
+    --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:pre-built-images"
 
 ### Build your own docker image
 
 === "AMD ROCm"
 
     --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-docker"
-
-### Pre-built images
-
-=== "NVIDIA CUDA"
-
-    --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-images"

docs/getting_started/installation/gpu/cuda.inc.md

@@ -17,12 +17,12 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 # --8<-- [start:pre-built-wheels]
 
 #### Installation of vLLM
-Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1 and 0.12.0rc1. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.11.0 --torch-backend=auto
+uv pip install vllm==0.12.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -93,7 +93,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.11.0rc1 \
+    vllm/vllm-omni:v0.12.0rc1 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

docs/getting_started/installation/gpu/rocm.inc.md

@@ -52,4 +52,42 @@ bash
 
 # --8<-- [start:pre-built-images]
 
+vLLM-Omni offers an official docker image for deployment. These images are built on top of vLLM docker images and available on Docker Hub as [vllm/vllm-omni-rocm](https://hub.docker.com/r/vllm/vllm-omni-rocm/tags). The version of vLLM-Omni indicates which release of vLLM it is based on.
+
+#### Launch vLLM-Omni Server
+Here's an example deployment command that has been verified on 2 x MI300's:
+```bash
+docker run -it \
+  --network=host \
+  --group-add=video \
+  --ipc=host \
+  --cap-add=SYS_PTRACE \
+  --security-opt seccomp=unconfined \
+  --device /dev/kfd \
+  --device /dev/dri \
+  -v <path/to/model>:/app/model \
+  -v ~/.cache/huggingface:/root/.cache/huggingface \
+  --env "HF_TOKEN=$HF_TOKEN" \
+  vllm/vllm-omni-rocm:v0.12.0rc1 \
+  vllm serve --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+```
+
+#### Launch an interactive terminal with prebuilt docker image.
+If you want to run in dev environment you can launch the docker image as follows:
+```bash
+docker run -it \
+  --network=host \
+  --group-add=video \
+  --ipc=host \
+  --cap-add=SYS_PTRACE \
+  --security-opt seccomp=unconfined \
+  --device /dev/kfd \
+  --device /dev/dri \
+  -v <path/to/model>:/app/model \
+  -v ~/.cache/huggingface:/root/.cache/huggingface \
+  --env "HF_TOKEN=$HF_TOKEN" \
+  vllm/vllm-omni-rocm:v0.12.0rc1 \
+  bash
+```
+
 # --8<-- [end:pre-built-images]

docs/getting_started/installation/npu/npu.inc.md

@@ -13,10 +13,10 @@ export DEVICE0=/dev/davinci0
 export DEVICE1=/dev/davinci1
 # Update the vllm-ascend image
 # Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1-a3
+export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \
@@ -34,12 +34,22 @@ docker run --rm \
     -p 8000:8000 \
     -it $IMAGE bash
 
+# Install the missing dependency of mooncake in the origin image.
+apt update
+apt install libjemalloc2
+echo "export LD_PRELOAD=/usr/lib/$(uname -m)-linux-gnu/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc
+source ~/.bashrc
+
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.11.0rc1 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.12.0rc1 https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 pip install -v -e .
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+# (Optional) Disable mooncake for stable capability
+mv /usr/local/Ascend/ascend-toolkit/latest/python/site-packages/mooncake \
+   /usr/local/Ascend/ascend-toolkit/latest/python/site-packages/mooncake.disabled
 ```
 
 The default workdir is `/workspace`, with vLLM, vLLM-Ascend and vLLM-Omni code placed in `/vllm-workspace` installed in development mode.

docs/models/supported_models.md

@@ -21,6 +21,7 @@ th {
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
+| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
 | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
@@ -31,6 +32,9 @@ th {
 |`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
 |`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
 |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
+|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
+|`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` |
+|`LTX2Pipeline` | LTX-2 | `Lightricks/LTX-2` |
 
 
 ## List of Supported Models for NPU

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -0,0 +1,29 @@
+# CPU Offloading for Diffusion Model
+
+## Overview
+CPU offload lets the diffusion worker move large model components between GPU and CPU memory on demand. It keeps the DiT transformer resident on GPU only while it is actively running, and swaps it out when encoders modules need the device. This reduces peak VRAM usage so bigger checkpoints run on smaller GPUs, or multiple requests can share the same GPU.
+
+## Execution Model
+1. Text encoders run on GPU while the DiT transformer is offloaded to CPU.
+2. Before denoising, weights are prefetched back to GPU, honoring pinned-memory copies for speed.
+3. After the diffusion step, the transformer returns to CPU and the process repeats as needed.
+
+Transfers use pinned host buffers, and the worker coordinates swaps via mutex-style hooks so components never compete for memory.
+
+## Configuration
+You can enable CPU offload in two ways:
+
+- **Python API**: set `enable_cpu_offload=True`.
+
+```python
+from vllm_omni import Omni
+
+if __name__ == "__main__":
+
+    m = Omni(model="Qwen/Qwen-Image",enable_cpu_offload=True)
+```
+
+- **CLI**: pass `--enable-cpu-offload` to the diffusion service entrypoint.
+
+## Known Limitations
+- Cold start latency increases for over one minute for some models(e.g., Qwen-Image)

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -10,27 +10,58 @@ The following parallelism methods are currently supported in vLLM-Omni:
 
 2. [Ring-Attention](#ring-attention) - splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results, keeping the sequence dimension sharded
 
+3. Classifier-Free-Guidance Parallel (CFG-Parallel): CFG-Parallel runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step.
+
+4. [Tensor Parallelism](#tensor-parallelism): Tensor parallelism shards model weights across devices. This can reduce per-GPU memory usage. Note that for diffusion models we currently shard the majority of layers within the DiT.
 
 The following table shows which models are currently supported by parallelism method:
 
 ### ImageGen
 
-| Model | Model Identifier | Ulysses-SP | Ring-SP |
-|-------|------------------|-----------|---------|
-| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ❌ |
-| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ❌ |
-| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ |
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ |
-| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ |
-| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ❌ |
+| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
+|-------|------------------|------------|---------|--------------|--------------------------|
+| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ❌ |
+| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ❌ |
+| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ |
+| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ❌ |
+| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ❌ |
+| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ❌ |
+| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ❌ |
+| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ❌ | ❌ | ✅ (TP=2 only) |
+| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
+
+!!! note "Why Z-Image is TP=2 only"
+    Z-Image Turbo is currently limited to `tensor_parallel_size` of **1 or 2** due to model shape divisibility constraints.
+    For example, the model has `n_heads=30` and a final projection out dimension of `64`, so valid TP sizes must divide both 30 and 64; the only common divisors are **1 and 2**.
 
 ### VideoGen
 
-| Model | Model Identifier | Ulysses-SP | Ring-SP |
-|-------|------------------|-----------|---------|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ❌ |
+| Model | Model Identifier | Ulysses-SP | Ring-SP | Tensor-Parallel |
+|-------|------------------|------------|---------|--------------------------|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ❌ | ❌ |
+
+### Tensor Parallelism
+
+Tensor parallelism splits model parameters across GPUs. In vLLM-Omni, tensor parallelism is configured via `DiffusionParallelConfig.tensor_parallel_size`.
+
+#### Offline Inference
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
+)
+
+outputs = omni.generate(
+    prompt="a cat reading a book",
+    num_inference_steps=9,
+    width=512,
+    height=512,
+)
+```
 
 ### Sequence Parallelism
 
@@ -181,6 +212,7 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m
 
 If a diffusion model has been deployed in vLLM-Omni and supports single-card inference, you can refer to the following instructions to parallelize it with [Ulysses-SP](https://arxiv.org/pdf/2309.14509).
 
+
 This section uses **Qwen-Image** (`QwenImageTransformer2DModel`) as the reference implementation. Qwen-Image is a **dual-stream** transformer (text + image) that performs **joint attention** across the concatenated sequences. Because of that, when enabling sequence parallel you typically:
 
 - Chunk **image tokens** (`hidden_states`) across SP ranks along the **sequence dimension**.
@@ -272,3 +304,101 @@ omni = Omni(
 
 outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
 ```
+
+
+### CFG-Parallel
+
+##### Offline Inference
+
+CFG-Parallel is enabled through `DiffusionParallelConfig(cfg_parallel_size=...)`. The recommended configuration is `cfg_parallel_size=2` (one rank for the positive branch and one rank for the negative branch).
+
+An example of offline inference using CFG-Parallel (image-to-image) is shown below:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Qwen/Qwen-Image-Edit",
+    parallel_config=DiffusionParallelConfig(cfg_parallel_size=2),
+)
+
+outputs = omni.generate(
+    prompt="turn this cat to a dog",
+    negative_prompt="low quality, blurry",
+    true_cfg_scale=4.0,
+    pil_image=input_image,
+    num_inference_steps=50,
+)
+```
+
+Notes:
+
+- CFG-Parallel is only effective when **true CFG** is enabled (i.e., `true_cfg_scale > 1` and a `negative_prompt` is provided).
+
+#### How to parallelize a pipeline
+
+This section describes how to add CFG-Parallel to a diffusion **pipeline**. We use the Qwen-Image pipeline (`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`) as the reference implementation.
+
+In `QwenImagePipeline`, each diffusion step runs two denoiser forward passes sequentially:
+
+- positive (prompt-conditioned)
+- negative (negative-prompt-conditioned)
+
+CFG-Parallel assigns these two branches to different ranks in the **CFG group** and synchronizes the results.
+
+Below is an example of CFG-Parallel implementation:
+
+```python
+def diffuse(
+        self,
+        ...
+        ):
+    # Enable CFG-parallel: rank0 computes positive, rank1 computes negative.
+    cfg_parallel_ready = do_true_cfg and get_classifier_free_guidance_world_size() > 1
+
+    self.transformer.do_true_cfg = do_true_cfg
+
+    if cfg_parallel_ready:
+        cfg_group = get_cfg_group()
+        cfg_rank = get_classifier_free_guidance_rank()
+
+        if cfg_rank == 0:
+            local_pred = self.transformer(
+                hidden_states=latents,
+                timestep=timestep / 1000,
+                guidance=guidance,
+                encoder_hidden_states_mask=prompt_embeds_mask,
+                encoder_hidden_states=prompt_embeds,
+                img_shapes=img_shapes,
+                txt_seq_lens=txt_seq_lens,
+                attention_kwargs=self.attention_kwargs,
+                return_dict=False,
+            )[0]
+        else:
+            local_pred = self.transformer(
+                hidden_states=latents,
+                timestep=timestep / 1000,
+                guidance=guidance,
+                encoder_hidden_states_mask=negative_prompt_embeds_mask,
+                encoder_hidden_states=negative_prompt_embeds,
+                img_shapes=img_shapes,
+                txt_seq_lens=negative_txt_seq_lens,
+                attention_kwargs=self.attention_kwargs,
+                return_dict=False,
+            )[0]
+
+        gathered = cfg_group.all_gather(local_pred, separate_tensors=True)
+        if cfg_rank == 0:
+            noise_pred = gathered[0]
+            neg_noise_pred = gathered[1]
+            comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+            cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
+            noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
+            noise_pred = comb_pred * (cond_norm / noise_norm)
+            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+        cfg_group.broadcast(latents, src=0)
+    else:
+        # fallback: run positive then negative sequentially on one rank
+        ...
+```

docs/user_guide/diffusion_acceleration.md

@@ -6,18 +6,21 @@ vLLM-Omni supports various cache acceleration methods to speed up diffusion mode
 
 vLLM-Omni currently supports two main cache acceleration backends:
 
-1. **[TeaCache](acceleration/teacache.md)** - Hook-based adaptive caching that caches transformer computations when consecutive timesteps are similar
-2. **[Cache-DiT](acceleration/cache_dit_acceleration.md)** - Library-based acceleration using multiple techniques:
+1. **[TeaCache](diffusion/teacache.md)** - Hook-based adaptive caching that caches transformer computations when consecutive timesteps are similar
+2. **[Cache-DiT](diffusion/cache_dit_acceleration.md)** - Library-based acceleration using multiple techniques:
     - **DBCache** (Dual Block Cache): Caches intermediate transformer block outputs based on residual differences
     - **TaylorSeer**: Uses Taylor expansion-based forecasting for faster inference
     - **SCM** (Step Computation Masking): Selectively computes steps based on adaptive masking
 
 Both methods can provide significant speedups (typically **1.5x-2.0x**) while maintaining high output quality.
 
-vLLM-Omni also supports the sequence parallelism (SP) for the diffusion model, that includes:
+vLLM-Omni also supports parallelism methods for diffusion models, including:
 
-1. [Ulysses-SP](acceleration/parallelism_acceleration.md#ulysses-sp) - splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
-2. [Ring-Attention](acceleration/parallelism_acceleration.md#ring-attention) - splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results, keeping the sequence dimension sharded.
+1. [Ulysses-SP](diffusion/parallelism_acceleration.md#ulysses-sp) - splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
+
+2. [Ring-Attention](diffusion/parallelism_acceleration.md#ring-attention) - splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results, keeping the sequence dimension sharded.
+
+3. [CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel) - runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step.
 
 ## Quick Comparison
 
@@ -34,26 +37,25 @@ The following table shows which models are currently supported by each accelerat
 
 ### ImageGen
 
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention |
-|-------|------------------|----------|-----------|-----------|----------------|
-| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ |
-| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ |
-| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ |
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ❌ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ✅ |❌ | ❌ |
-| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ |
-| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ |
-| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ |
-| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ |
+| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel |
+|-------|------------------|:----------:|:-----------:|:-----------:|:----------------:|:----------------:|
+| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ |
+| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ✅ |❌ | ❌ | ❌ |
+| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ❌ | ✅ | ❌ | ❌ | ❌ |
 
 ### VideoGen
 
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP |
-|-------|------------------|----------|-----------|-----------|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ❌ |
+| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention |
+|-------|------------------|:--------:|:---------:|:----------:|:--------------:|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ❌ | ❌ |
 
 
 ## Performance Benchmarks
@@ -134,7 +136,7 @@ ulysses_degree = 2
 
 omni = Omni(
     model="Qwen/Qwen-Image",
-    parallel_config=DiffusionParallelConfig(ulysses_degree=2)
+    parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
 )
 
 outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
@@ -149,7 +151,7 @@ ulysses_degree = 2
 
 omni = Omni(
     model="Qwen/Qwen-Image-Edit",
-    parallel_config=DiffusionParallelConfig(ulysses_degree=2)
+    parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
 )
 
 outputs = omni.generate(prompt="turn this cat to a dog",
@@ -172,10 +174,31 @@ omni = Omni(
 outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
 ```
 
+### Using CFG-Parallel
+
+Run image-to-image:
+
+CFG-Parallel splits the CFG positive/negative branches across GPUs. Use it when you set a non-trivial `true_cfg_scale`.
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+cfg_parallel_size = 2
+
+omni = Omni(
+    model="Qwen/Qwen-Image-Edit",
+    parallel_config=DiffusionParallelConfig(cfg_parallel_size=cfg_parallel_size)
+)
+
+outputs = omni.generate(prompt="turn this cat to a dog",
+        pil_image=input_image, num_inference_steps=50, true_cfg_scale=4.0)
+```
+
 ## Documentation
 
 For detailed information on each acceleration method:
 
-- **[TeaCache Guide](acceleration/teacache.md)** - Complete TeaCache documentation, configuration options, and best practices
-- **[Cache-DiT Acceleration Guide](acceleration/cache_dit_acceleration.md)** - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters
-- **[Sequence Parallelism](acceleration/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
+- **[TeaCache Guide](diffusion/teacache.md)** - Complete TeaCache documentation, configuration options, and best practices
+- **[Cache-DiT Acceleration Guide](diffusion/cache_dit_acceleration.md)** - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters
+- **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
+- **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.

docs/user_guide/examples/offline_inference/image_to_video.md

@@ -0,0 +1,66 @@
+# Image-To-Video
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video>.
+
+
+This example demonstrates how to generate videos from images using Wan2.2 Image-to-Video models with vLLM-Omni's offline inference API.
+
+## Local CLI Usage
+
+### Wan2.2-I2V-A14B-Diffusers (MoE)
+```bash
+python image_to_video.py \
+  --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
+  --image input.png \
+  --prompt "A cat playing with yarn, smooth motion" \
+  --negative_prompt "<optional quality filter>" \
+  --height 480 \
+  --width 832 \
+  --num_frames 48 \
+  --guidance_scale 5.0 \
+  --guidance_scale_high 6.0 \
+  --num_inference_steps 40 \
+  --boundary_ratio 0.875 \
+  --flow_shift 12.0 \
+  --fps 16 \
+  --output i2v_output.mp4
+```
+
+### Wan2.2-TI2V-5B-Diffusers (Unified)
+```bash
+python image_to_video.py \
+  --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+  --image input.png \
+  --prompt "A cat playing with yarn, smooth motion" \
+  --negative_prompt "<optional quality filter>" \
+  --height 480 \
+  --width 832 \
+  --num_frames 48 \
+  --guidance_scale 4.0 \
+  --num_inference_steps 40 \
+  --flow_shift 12.0 \
+  --fps 16 \
+  --output i2v_output.mp4
+```
+
+Key arguments:
+
+- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V).
+- `--image`: Path to input image (required).
+- `--prompt`: Text description of desired motion/animation.
+- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16.
+- `--num_frames`: Number of frames (default 81).
+- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high-noise stages for MoE).
+- `--negative_prompt`: Optional list of artifacts to suppress.
+- `--boundary_ratio`: Boundary split ratio for two-stage MoE models.
+- `--flow_shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
+- `--num_inference_steps`: Number of denoising steps (default 50).
+- `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
+- `--output`: Path to save the generated video.
+
+## Example materials
+
+??? abstract "image_to_video.py"
+    ``````py
+    --8<-- "examples/offline_inference/image_to_video/image_to_video.py"
+    ``````

docs/user_guide/examples/offline_inference/ltx2.md

@@ -0,0 +1,42 @@
+# LTX-2 Text-To-Video
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/ltx2>.
+
+The `Lightricks/LTX-2` pipeline generates high-quality videos from text prompts using a 19B parameter DiT architecture. It supports up to 4K video generation with synchronized audio.
+
+## Local CLI Usage
+
+```bash
+python text_to_video.py \
+  --prompt "A panda riding a bicycle through a forest, cinematic lighting" \
+  --height 512 \
+  --width 768 \
+  --num_frames 121 \
+  --num_inference_steps 40 \
+  --guidance_scale 4.0 \
+  --output ltx2_output.mp4
+```
+
+Key arguments:
+
+- `--prompt`: Text description of the video to generate (string).
+- `--height/--width`: Output resolution (defaults 512x768). Dimensions must be divisible by 32.
+- `--num_frames`: Number of frames (should be 8*n+1, e.g., 25, 81, 121). Default is 121.
+- `--guidance_scale`: Classifier-free guidance scale (default: 4.0). Range 3.0-5.0 recommended.
+- `--negative_prompt`: Optional text describing what to avoid in the video.
+- `--num_inference_steps`: Number of denoising steps (default: 40).
+- `--seed`: Random seed for reproducibility (default: 42).
+- `--fps`: Frames per second for the saved MP4 (default: 24).
+- `--output`: Path to save the generated video.
+
+## Example materials
+
+??? abstract "text_to_video.py"
+    ``````py
+    --8<-- "examples/offline_inference/ltx2/text_to_video.py"
+    ``````
+
+??? abstract "text_to_video.md"
+    ``````md
+    --8<-- "examples/offline_inference/ltx2/text_to_video.md"
+    ``````

docs/user_guide/examples/offline_inference/qwen2_5_omni.md

@@ -2,6 +2,7 @@
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen2_5_omni>.
 
+
 ## Setup
 Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
 

docs/user_guide/examples/offline_inference/qwen3_omni.md

@@ -2,6 +2,7 @@
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen3_omni>.
 
+
 ## Setup
 Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
 

docs/user_guide/examples/online_serving/gradio_demo.md

@@ -1,7 +0,0 @@
-# Gradio Demo
-
-Source <https://github.com/vllm-project/vllm-omni/blob/main/examples/online_serving/gradio_demo.py>.
-
-``````py
---8<-- "examples/online_serving/gradio_demo.py"
-``````

docs/user_guide/examples/online_serving/image_to_image.md

@@ -230,6 +230,10 @@ Provide multiple images in `content` (order matters):
     ``````py
     --8<-- "examples/online_serving/image_to_image/openai_chat_client.py"
     ``````
+??? abstract "run_curl_image_edit.sh"
+    ``````sh
+    --8<-- "examples/online_serving/image_to_image/run_curl_image_edit.sh"
+    ``````
 ??? abstract "run_server.sh"
     ``````sh
     --8<-- "examples/online_serving/image_to_image/run_server.sh"

docs/user_guide/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.md

@@ -1,7 +0,0 @@
-# OpenAI Chat Completion Client For Multimodal Generation
-
-Source <https://github.com/vllm-project/vllm-omni/blob/main/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py>.
-
-``````py
---8<-- "examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py"
-``````

docs/user_guide/examples/online_serving/qwen2_5_omni.md

@@ -137,6 +137,7 @@ response = client.chat.completions.create(
 # Response contains two choices: one with text, one with audio
 print(response.choices[0].message.content)  # Text response
 print(response.choices[1].message.audio)    # Audio response
+```
 
 ## Streaming Output
 If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -145,6 +145,7 @@ response = client.chat.completions.create(
 # Response contains two choices: one with text, one with audio
 print(response.choices[0].message.content)  # Text response
 print(response.choices[1].message.audio)    # Audio response
+```
 
 ## Streaming Output
 If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.

docs/user_guide/examples/online_serving/text_to_image.md

@@ -12,6 +12,8 @@ This example demonstrates how to deploy Qwen-Image model for online image genera
 ```bash
 vllm serve Qwen/Qwen-Image --omni --port 8091
 ```
+!!! note
+    If you encounter Out-of-Memory (OOM) issues or have limited GPU memory, you can enable VAE slicing and tiling to reduce memory usage, --vae-use-slicing --vae-use-tiling
 
 ### Start with Parameters
 
@@ -43,7 +45,7 @@ curl -s http://localhost:8091/v1/chat/completions \
       "true_cfg_scale": 4.0,
       "seed": 42
     }
-  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
 ```
 
 ### Method 2: Using Python Client
@@ -146,7 +148,7 @@ Use `extra_body` to pass generation parameters:
 
 ```bash
 # Extract base64 from response and decode to image
-cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png
+cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
 ```
 
 ## File Description

examples/offline_inference/image_to_image/image_edit.py

@@ -49,6 +49,22 @@ Usage (layered):
         --layers 4 \
         --color-format "RGBA"
 
+Usage (with CFG Parallel):
+    python image_edit.py \
+        --image input.png \
+        --prompt "Edit description" \
+        --cfg_parallel_size 2 \
+        --num_inference_steps 50 \
+        --cfg_scale 4.0
+
+Usage (disable torch.compile):
+    python image_edit.py \
+        --image input.png \
+        --prompt "Edit description" \
+        --enforce_eager \
+        --num_inference_steps 50 \
+        --cfg_scale 4.0
+
 For more options, run:
     python image_edit.py --help
 """
@@ -245,7 +261,18 @@ def parse_args() -> argparse.Namespace:
         default=0.2,
         help="[tea_cache] Threshold for accumulated relative L1 distance.",
     )
-
+    parser.add_argument(
+        "--cfg_parallel_size",
+        type=int,
+        default=1,
+        choices=[1, 2],
+        help="Number of GPUs used for classifier free guidance parallel size.",
+    )
+    parser.add_argument(
+        "--enforce_eager",
+        action="store_true",
+        help="Disable torch.compile and force eager execution.",
+    )
     return parser.parse_args()
 
 
@@ -273,7 +300,10 @@ def main():
     # Enable VAE memory optimizations on NPU
     vae_use_slicing = is_npu()
     vae_use_tiling = is_npu()
-    parallel_config = DiffusionParallelConfig(ulysses_degree=args.ulysses_degree, ring_degree=args.ring_degree)
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree, ring_degree=args.ring_degree, cfg_parallel_size=args.cfg_parallel_size
+    )
+
     # Configure cache based on backend type
     cache_config = None
     if args.cache_backend == "cache_dit":
@@ -304,6 +334,7 @@ def main():
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         parallel_config=parallel_config,
+        enforce_eager=args.enforce_eager,
     )
     print("Pipeline loaded")
 
@@ -319,7 +350,9 @@ def main():
             print(f"    Image {idx + 1} size: {img.size}")
     else:
         print(f"  Input image size: {input_image.size}")
-    print(f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}")
+    print(
+        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}"
+    )
     print(f"{'=' * 60}\n")
 
     generation_start = time.perf_counter()

examples/offline_inference/ltx2/text_to_video.md

@@ -0,0 +1,77 @@
+# LTX-2 Text-to-Video Generation
+
+This example demonstrates how to use Lightricks LTX-2 model for text-to-video generation with vLLM-Omni.
+
+## Model Information
+
+LTX-2 is a state-of-the-art text-to-video model by Lightricks that can generate high-quality videos from text prompts. It supports:
+- Text-to-video generation
+- Image-to-video generation (conditioned on an input image)
+- High-resolution video generation (up to 4K)
+- Synchronized audio-video generation
+
+For more information, see:
+- [Lightricks/LTX-2 on HuggingFace](https://huggingface.co/Lightricks/LTX-2)
+- [LTX-2 GitHub Repository](https://github.com/Lightricks/LTX-2)
+
+## Requirements
+
+- Python >= 3.10
+- diffusers >= 0.36.0 (already included in vllm-omni dependencies)
+- GPU with sufficient VRAM (recommended: >= 24GB for high-quality generation)
+
+## Usage
+
+### Basic Text-to-Video Generation
+
+```bash
+python text_to_video.py \
+    --model Lightricks/LTX-2 \
+    --prompt "A panda riding a bicycle through a forest, cinematic lighting" \
+    --height 512 \
+    --width 768 \
+    --num_frames 121 \
+    --num_inference_steps 40 \
+    --guidance_scale 4.0 \
+    --output output.mp4
+```
+
+### Parameters
+
+- `--model`: Model ID or local path (default: `Lightricks/LTX-2`)
+- `--prompt`: Text description of the video to generate
+- `--negative_prompt`: Text describing what to avoid in the video
+- `--height`: Video height in pixels, must be divisible by 32 (default: 512)
+- `--width`: Video width in pixels, must be divisible by 32 (default: 768)
+- `--num_frames`: Number of frames, should be 8*n+1 (e.g., 25, 81, 121) (default: 121)
+- `--num_inference_steps`: Number of denoising steps (default: 40)
+- `--guidance_scale`: Classifier-free guidance scale (default: 4.0)
+- `--seed`: Random seed for reproducibility (default: 42)
+- `--fps`: Output video frames per second (default: 24)
+- `--output`: Output video file path (default: ltx2_output.mp4)
+
+### Example Prompts
+
+```bash
+# Cinematic scene
+python text_to_video.py --prompt "A serene lakeside sunrise with mist over the water, cinematic"
+
+# Action scene
+python text_to_video.py --prompt "A futuristic cityscape with flying cars, neon lights, cyberpunk style"
+
+# Nature scene
+python text_to_video.py --prompt "A waterfall in a lush rainforest, birds flying overhead"
+```
+
+## Performance Tips
+
+1. **Resolution**: Start with lower resolutions (512x768) for faster generation. Higher resolutions require more VRAM.
+2. **Frames**: More frames take longer to generate. Start with 25 or 81 frames for testing.
+3. **Inference Steps**: 40 steps provide good quality. Fewer steps are faster but may reduce quality.
+4. **Guidance Scale**: Values between 3.0-5.0 typically work well. Higher values follow the prompt more strictly.
+
+## Notes
+
+- The first run will download the model from HuggingFace (approximately 20-30GB).
+- Generation time depends on your GPU and the parameters chosen.
+- For best results, use descriptive prompts with style keywords (e.g., "cinematic", "high quality", "detailed").

examples/offline_inference/ltx2/text_to_video.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.utils.platform_utils import detect_device_type
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate a video with LTX-2 T2V.")
+    parser.add_argument(
+        "--model",
+        default="Lightricks/LTX-2",
+        help="LTX-2 model ID or local path.",
+    )
+    parser.add_argument(
+        "--prompt",
+        default="A serene lakeside sunrise with mist over the water.",
+        help="Text prompt for video generation.",
+    )
+    parser.add_argument("--negative_prompt", default="", help="Negative prompt.")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
+    parser.add_argument("--guidance_scale", type=float, default=4.0, help="CFG scale.")
+    parser.add_argument("--height", type=int, default=512, help="Video height (must be divisible by 32).")
+    parser.add_argument("--width", type=int, default=768, help="Video width (must be divisible by 32).")
+    parser.add_argument(
+        "--num_frames", type=int, default=121, help="Number of frames (should be 8*n+1, e.g., 25, 81, 121)."
+    )
+    parser.add_argument("--num_inference_steps", type=int, default=40, help="Sampling steps.")
+    parser.add_argument("--output", type=str, default="ltx2_output.mp4", help="Path to save the video (mp4).")
+    parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    device = detect_device_type()
+    generator = torch.Generator(device=device).manual_seed(args.seed)
+
+    omni = Omni(
+        model=args.model,
+    )
+
+    frames = omni.generate(
+        args.prompt,
+        negative_prompt=args.negative_prompt,
+        height=args.height,
+        width=args.width,
+        generator=generator,
+        guidance_scale=args.guidance_scale,
+        num_inference_steps=args.num_inference_steps,
+        num_frames=args.num_frames,
+    )
+
+    # Extract video frames from OmniRequestOutput
+    if isinstance(frames, list) and len(frames) > 0:
+        first_item = frames[0]
+
+        # Check if it's an OmniRequestOutput
+        if hasattr(first_item, "final_output_type"):
+            if first_item.final_output_type != "image":
+                raise ValueError(
+                    f"Unexpected output type '{first_item.final_output_type}', "
+                    f"expected 'image' for video frames output."
+                )
+
+            # Pipeline mode: extract from nested request_output
+            if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
+                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
+                    inner_output = first_item.request_output[0]
+                    if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
+                        frames = inner_output.images[0] if inner_output.images else None
+                        if frames is None:
+                            raise ValueError("No video frames found in output.")
+            # Diffusion mode: use direct images field
+            elif hasattr(first_item, "images") and first_item.images:
+                frames = first_item.images
+            else:
+                raise ValueError("No video frames found in OmniRequestOutput.")
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        from diffusers.utils import export_to_video
+    except ImportError:
+        raise ImportError("diffusers is required for export_to_video.")
+
+    # frames may be np.ndarray (preferred) or torch.Tensor
+    # export_to_video expects a list of frames with values in [0, 1]
+    if isinstance(frames, torch.Tensor):
+        video_tensor = frames.detach().cpu()
+        if video_tensor.dim() == 5:
+            # [B, C, F, H, W] or [B, F, H, W, C]
+            if video_tensor.shape[1] in (3, 4):
+                video_tensor = video_tensor[0].permute(1, 2, 3, 0)
+            else:
+                video_tensor = video_tensor[0]
+        elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
+            video_tensor = video_tensor.permute(1, 2, 3, 0)
+        # If float, assume [-1,1] and normalize to [0,1]
+        if video_tensor.is_floating_point():
+            video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
+        video_array = video_tensor.float().numpy()
+    else:
+        video_array = frames
+        if hasattr(video_array, "shape") and video_array.ndim == 5:
+            video_array = video_array[0]
+
+    # Convert 4D array (frames, H, W, C) to list of frames for export_to_video
+    if isinstance(video_array, np.ndarray) and video_array.ndim == 4:
+        video_array = list(video_array)
+
+    export_to_video(video_array, str(output_path), fps=args.fps)
+    print(f"Saved generated video to {output_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Example script for text-to-audio generation using Stable Audio Open.
+
+This script demonstrates how to generate audio from text prompts using
+the Stable Audio Open model with vLLM-Omni.
+
+Usage:
+    python text_to_audio.py --prompt "The sound of a dog barking"
+    python text_to_audio.py --prompt "A piano playing a gentle melody" --audio_length 10.0
+    python text_to_audio.py --prompt "Thunder and rain sounds" --negative_prompt "Low quality"
+"""
+
+import argparse
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.utils.platform_utils import detect_device_type
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate audio with Stable Audio Open.")
+    parser.add_argument(
+        "--model",
+        default="stabilityai/stable-audio-open-1.0",
+        help="Stable Audio model name or local path.",
+    )
+    parser.add_argument(
+        "--prompt",
+        default="The sound of a hammer hitting a wooden surface.",
+        help="Text prompt for audio generation.",
+    )
+    parser.add_argument(
+        "--negative_prompt",
+        default="Low quality.",
+        help="Negative prompt for classifier-free guidance.",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed for deterministic results.",
+    )
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=7.0,
+        help="Classifier-free guidance scale.",
+    )
+    parser.add_argument(
+        "--audio_start",
+        type=float,
+        default=0.0,
+        help="Audio start time in seconds.",
+    )
+    parser.add_argument(
+        "--audio_length",
+        type=float,
+        default=10.0,
+        help="Audio length in seconds (max ~47s for stable-audio-open-1.0).",
+    )
+    parser.add_argument(
+        "--num_inference_steps",
+        type=int,
+        default=100,
+        help="Number of denoising steps for the diffusion sampler.",
+    )
+    parser.add_argument(
+        "--num_waveforms",
+        type=int,
+        default=1,
+        help="Number of audio waveforms to generate for the given prompt.",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="stable_audio_output.wav",
+        help="Path to save the generated audio (WAV format).",
+    )
+    parser.add_argument(
+        "--sample_rate",
+        type=int,
+        default=44100,
+        help="Sample rate for output audio (Stable Audio uses 44100 Hz).",
+    )
+    return parser.parse_args()
+
+
+def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 44100):
+    """Save audio data to a WAV file."""
+    try:
+        import soundfile as sf
+
+        sf.write(output_path, audio_data, sample_rate)
+    except ImportError:
+        try:
+            import scipy.io.wavfile as wav
+
+            # Ensure audio is in the correct format for scipy
+            if audio_data.dtype == np.float32 or audio_data.dtype == np.float64:
+                # Normalize to int16 range
+                audio_data = np.clip(audio_data, -1.0, 1.0)
+                audio_data = (audio_data * 32767).astype(np.int16)
+            wav.write(output_path, sample_rate, audio_data)
+        except ImportError:
+            raise ImportError(
+                "Either 'soundfile' or 'scipy' is required to save audio files. "
+                "Install with: pip install soundfile or pip install scipy"
+            )
+
+
+def main():
+    args = parse_args()
+    device = detect_device_type()
+    generator = torch.Generator(device=device).manual_seed(args.seed)
+
+    print(f"\n{'=' * 60}")
+    print("Stable Audio Open - Text-to-Audio Generation")
+    print(f"{'=' * 60}")
+    print(f"  Model: {args.model}")
+    print(f"  Prompt: {args.prompt}")
+    print(f"  Negative prompt: {args.negative_prompt}")
+    print(f"  Audio length: {args.audio_length}s")
+    print(f"  Inference steps: {args.num_inference_steps}")
+    print(f"  Guidance scale: {args.guidance_scale}")
+    print(f"  Seed: {args.seed}")
+    print(f"{'=' * 60}\n")
+
+    # Initialize Omni with Stable Audio model
+    omni = Omni(model=args.model)
+
+    # Calculate audio end time
+    audio_end_in_s = args.audio_start + args.audio_length
+
+    # Time profiling for generation
+    generation_start = time.perf_counter()
+
+    # Generate audio
+    audio = omni.generate(
+        args.prompt,
+        negative_prompt=args.negative_prompt,
+        generator=generator,
+        guidance_scale=args.guidance_scale,
+        num_inference_steps=args.num_inference_steps,
+        num_outputs_per_prompt=args.num_waveforms,
+        extra={
+            "audio_start_in_s": args.audio_start,
+            "audio_end_in_s": audio_end_in_s,
+        },
+    )
+
+    generation_end = time.perf_counter()
+    generation_time = generation_end - generation_start
+
+    print(f"Total generation time: {generation_time:.2f} seconds")
+
+    # Process and save audio
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    suffix = output_path.suffix or ".wav"
+    stem = output_path.stem or "stable_audio_output"
+
+    # Handle different output formats
+    if isinstance(audio, torch.Tensor):
+        audio = audio.cpu().float().numpy()
+
+    # Audio shape is typically [batch, channels, samples] or [channels, samples]
+    if audio.ndim == 3:
+        # [batch, channels, samples]
+        if args.num_waveforms <= 1:
+            audio_data = audio[0].T  # [samples, channels]
+            save_audio(audio_data, str(output_path), args.sample_rate)
+            print(f"Saved generated audio to {output_path}")
+        else:
+            for idx in range(audio.shape[0]):
+                audio_data = audio[idx].T  # [samples, channels]
+                save_path = output_path.parent / f"{stem}_{idx}{suffix}"
+                save_audio(audio_data, str(save_path), args.sample_rate)
+                print(f"Saved generated audio to {save_path}")
+    elif audio.ndim == 2:
+        # [channels, samples]
+        audio_data = audio.T  # [samples, channels]
+        save_audio(audio_data, str(output_path), args.sample_rate)
+        print(f"Saved generated audio to {output_path}")
+    else:
+        # [samples] - mono audio
+        save_audio(audio, str(output_path), args.sample_rate)
+        print(f"Saved generated audio to {output_path}")
+
+    print(f"\nGenerated {args.audio_length}s of audio at {args.sample_rate} Hz")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/text_to_image/text_to_image.py

@@ -81,6 +81,24 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for ring sequence parallelism.",
     )
+    parser.add_argument(
+        "--cfg_parallel_size",
+        type=int,
+        default=1,
+        choices=[1, 2],
+        help="Number of GPUs used for classifier free guidance parallel size.",
+    )
+    parser.add_argument(
+        "--enforce_eager",
+        action="store_true",
+        help="Disable torch.compile and force eager execution.",
+    )
+    parser.add_argument(
+        "--tensor_parallel_size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
+    )
     return parser.parse_args()
 
 
@@ -123,7 +141,13 @@ def main():
         }
 
     # assert args.ring_degree == 1, "Ring attention is not supported yet"
-    parallel_config = DiffusionParallelConfig(ulysses_degree=args.ulysses_degree, ring_degree=args.ring_degree)
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+    )
+
     omni = Omni(
         model=args.model,
         vae_use_slicing=vae_use_slicing,
@@ -131,6 +155,7 @@ def main():
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         parallel_config=parallel_config,
+        enforce_eager=args.enforce_eager,
     )
 
     # Time profiling for generation
@@ -139,7 +164,10 @@ def main():
     print(f"  Model: {args.model}")
     print(f"  Inference steps: {args.num_inference_steps}")
     print(f"  Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}")
-    print(f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}")
+    print(
+        f"  Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, "
+        f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}"
+    )
     print(f"  Image size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")
 

examples/online_serving/image_to_image/openai_chat_client.py

@@ -127,7 +127,7 @@ def main():
     parser.add_argument("--width", type=int, default=1024, help="Output image width")
     parser.add_argument("--steps", type=int, default=50, help="Inference steps")
     parser.add_argument("--guidance", type=float, default=7.5, help="Guidance scale")
-    parser.add_argument("--seed", type=int, help="Random seed")
+    parser.add_argument("--seed", type=int, default=0, help="Random seed")
     parser.add_argument("--negative", help="Negative prompt")
 
     args = parser.parse_args()

examples/online_serving/qwen2_5_omni/README.md

@@ -134,6 +134,7 @@ response = client.chat.completions.create(
 # Response contains two choices: one with text, one with audio
 print(response.choices[0].message.content)  # Text response
 print(response.choices[1].message.audio)    # Audio response
+```
 
 ## Streaming Output
 If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.

examples/online_serving/qwen3_omni/README.md

@@ -142,6 +142,7 @@ response = client.chat.completions.create(
 # Response contains two choices: one with text, one with audio
 print(response.choices[0].message.content)  # Text response
 print(response.choices[1].message.audio)    # Audio response
+```
 
 ## Streaming Output
 If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.

examples/online_serving/text_to_image/README.md

@@ -9,6 +9,8 @@ This example demonstrates how to deploy Qwen-Image model for online image genera
 ```bash
 vllm serve Qwen/Qwen-Image --omni --port 8091
 ```
+!!! note
+    If you encounter Out-of-Memory (OOM) issues or have limited GPU memory, you can enable VAE slicing and tiling to reduce memory usage, --vae-use-slicing --vae-use-tiling
 
 ### Start with Parameters
 
@@ -40,7 +42,7 @@ curl -s http://localhost:8091/v1/chat/completions \
       "true_cfg_scale": 4.0,
       "seed": 42
     }
-  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
 ```
 
 ### Method 2: Using Python Client
@@ -143,7 +145,7 @@ Use `extra_body` to pass generation parameters:
 
 ```bash
 # Extract base64 from response and decode to image
-cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png
+cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
 ```
 
 ## File Description

examples/online_serving/text_to_image/openai_chat_client.py

@@ -100,7 +100,7 @@ def main():
     parser.add_argument("--width", type=int, default=1024, help="Image width")
     parser.add_argument("--steps", type=int, default=50, help="Inference steps")
     parser.add_argument("--cfg-scale", type=float, default=4.0, help="True CFG scale")
-    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    parser.add_argument("--seed", type=int, default=0, help="Random seed")
     parser.add_argument("--negative", help="Negative prompt")
 
     args = parser.parse_args()

mkdocs.yml

@@ -60,6 +60,10 @@ hooks:
   - docs/mkdocs/hooks/url_schemes.py
   - docs/mkdocs/hooks/generate_examples.py
 
+# Exclude include files from navigation warnings
+exclude_docs: |
+  **/*.inc.md
+
 # Plugins
 plugins:
   - meta

pyproject.toml

@@ -38,6 +38,7 @@ dependencies = [
     "soundfile>=0.13.1",
     "cache-dit==1.1.8",
     "tqdm>=4.66.0",
+    "torchsde>=0.2.6",  # Required for Stable Audio scheduler
     # "vllm==0.12.0",  # TODO: fix the entrypoints overwrite problem
 ]
 
@@ -49,6 +50,8 @@ dev = [
     "pytest-cov>=4.0.0",
     "mypy==1.11.1",
     "pre-commit==4.0.1",
+    "openai-whisper>=20250625",
+    "psutil>=7.2.0"
 ]
 
 docs = [
@@ -150,11 +153,34 @@ addopts = [
     "--cov-report=xml",
 ]
 markers = [
-    "unit: Unit tests",
-    "integration: Integration tests",
+    # ci/cd required
+    "core_model: Core model tests (run in each PR)",
+    # function module markers
+    "diffusion: Diffusion model tests",
+    "omni: Omni model tests",
+    "cache: Cache backend tests",
+    "parallel: Parallelism/distributed tests",
+    # platform markers
+    "cpu: Tests that run on CPU",
+    "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)",
+    "npu: Tests that run on NPU/Ascend (auto-added)",
+    # specified computation resources marks (auto-added)
+    "H100: Tests that require H100 GPU",
+    "L4: Tests that require L4 GPU",
+    "MI325: Tests that require MI325 GPU (AMD/ROCm)",
+    "A2: Tests that require A2 NPU",
+    "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_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_npu: Skip if the num of NPU cards is less than the required",
+    # more detailed markers
+    "slow: Slow tests (may skip in quick CI)",
     "benchmark: Benchmark tests",
-    "slow: Slow tests",
-    "core_model: enable this model test in each PR instead of only nightly",
 ]
 
 [tool.typos.default]

pytest.ini

@@ -1,3 +0,0 @@
-[pytest]
-markers =
-    gpu_mem_high: needs high VRAM

tests/conftest.py

@@ -1,8 +1,19 @@
+import base64
 import os
+import socket
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Any
 
+import psutil
 import pytest
 import torch
+import whisper
+import yaml
 from vllm.logger import init_logger
+from vllm.utils import get_open_port
 
 logger = init_logger(__name__)
 
@@ -34,3 +45,286 @@ def clean_gpu_memory_between_tests():
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         gc.collect()
+
+
+def dummy_messages_from_mix_data(
+    system_prompt: dict[str, Any] = None,
+    video_data_url: Any = None,
+    audio_data_url: Any = None,
+    image_data_url: Any = None,
+    content_text: str = None,
+):
+    """Create messages with video、image、audio data URL for OpenAI API."""
+
+    if content_text is not None:
+        content = [{"type": "text", "text": content_text}]
+    else:
+        content = []
+
+    media_items = []
+    if isinstance(video_data_url, list):
+        for video_url in video_data_url:
+            media_items.append((video_url, "video"))
+    else:
+        media_items.append((video_data_url, "video"))
+
+    if isinstance(image_data_url, list):
+        for url in image_data_url:
+            media_items.append((url, "image"))
+    else:
+        media_items.append((image_data_url, "image"))
+
+    if isinstance(audio_data_url, list):
+        for url in audio_data_url:
+            media_items.append((url, "audio"))
+    else:
+        media_items.append((audio_data_url, "audio"))
+
+    content.extend(
+        {"type": f"{media_type}_url", f"{media_type}_url": {"url": url}}
+        for url, media_type in media_items
+        if url is not None
+    )
+    messages = [{"role": "user", "content": content}]
+    if system_prompt is not None:
+        messages = [system_prompt] + messages
+    return messages
+
+
+def cosine_similarity_text(s1, s2):
+    """
+        Calculate cosine similarity between two text strings.
+        Notes:
+    ------
+    - Higher score means more similar texts
+    - Score of 1.0 means identical word composition (bag-of-words)
+    - Score of 0.0 means completely different vocabulary
+    """
+    from sklearn.feature_extraction.text import CountVectorizer
+    from sklearn.metrics.pairwise import cosine_similarity
+
+    vectorizer = CountVectorizer().fit_transform([s1, s2])
+    vectors = vectorizer.toarray()
+    return cosine_similarity([vectors[0]], [vectors[1]])[0][0]
+
+
+def convert_audio_to_text(audio_data):
+    """
+    Convert base64 encoded audio data to text using speech recognition.
+    """
+
+    audio_data = base64.b64decode(audio_data)
+    output_path = f"./test_{int(time.time())}"
+    with open(output_path, "wb") as audio_file:
+        audio_file.write(audio_data)
+
+    print(f"audio data is saved: {output_path}")
+    model = whisper.load_model("base")
+    text = model.transcribe(output_path)["text"]
+    if text:
+        return text
+    else:
+        return ""
+
+
+def modify_stage_config(
+    yaml_path: str,
+    stage_updates: dict[int, dict[str, Any]],
+) -> str:
+    """
+    Batch modify configurations for multiple stages in a YAML file.
+
+    Args:
+        yaml_path: Path to the YAML configuration file.
+        stage_updates: Dictionary where keys are stage IDs and values are dictionaries of
+                      modifications for that stage. Each modification dictionary uses
+                      dot-separated paths as keys and new configuration values as values.
+                      Example: {
+                          0: {'engine_args.max_model_len': 5800},
+                          1: {'runtime.max_batch_size': 2}
+                      }
+
+    Returns:
+        str: Path to the newly created modified YAML file with timestamp suffix.
+
+    Example:
+        >>> output_file = modify_stage_config(
+        ...     'config.yaml',
+        ...     {
+        ...         0: {'engine_args.max_model_len': 5800},
+        ...         1: {'runtime.max_batch_size': 2}
+        ...     }
+        ... )
+        >>> print(f"Modified configuration saved to: {output_file}")
+        Modified configuration saved to: config_1698765432.yaml
+    """
+    path = Path(yaml_path)
+    if not path.exists():
+        raise FileNotFoundError(f"yaml does not exist: {path}")
+    try:
+        with open(yaml_path, encoding="utf-8") as f:
+            config = yaml.safe_load(f) or {}
+    except Exception as e:
+        raise ValueError(f"Cannot parse YAML file: {e}")
+
+    stage_args = config.get("stage_args", [])
+    if not stage_args:
+        raise ValueError("the stage_args does not exist")
+
+    for stage_id, config_dict in stage_updates.items():
+        target_stage = None
+        for stage in stage_args:
+            if stage.get("stage_id") == stage_id:
+                target_stage = stage
+                break
+
+        if target_stage is None:
+            available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s]
+            raise KeyError(f"Stage ID {stage_id} is not exist, available IDs: {available_ids}")
+
+        for key_path, value in config_dict.items():
+            current = target_stage
+            keys = key_path.split(".")
+            for i in range(len(keys) - 1):
+                key = keys[i]
+                if key not in current:
+                    raise KeyError(f"the {'.'.join(keys[: i + 1])} does not exist")
+
+                elif not isinstance(current[key], dict) and i < len(keys) - 2:
+                    raise ValueError(f"{'.'.join(keys[: i + 1])}' cannot continue deeper because it's not a dict")
+                current = current[key]
+            current[keys[-1]] = value
+
+    output_path = f"{yaml_path.split('.')[0]}_{int(time.time())}.yaml"
+    with open(output_path, "w", encoding="utf-8") as f:
+        yaml.dump(config, f, default_flow_style=False, sort_keys=False, allow_unicode=True, indent=2)
+
+    return output_path
+
+
+class OmniServer:
+    """Omniserver for vLLM-Omni tests."""
+
+    def __init__(
+        self,
+        model: str,
+        serve_args: list[str],
+        *,
+        env_dict: dict[str, str] | None = None,
+    ) -> None:
+        self.model = model
+        self.serve_args = serve_args
+        self.env_dict = env_dict
+        self.proc: subprocess.Popen | None = None
+        self.host = "127.0.0.1"
+        self.port = get_open_port()
+
+    def _start_server(self) -> None:
+        """Start the vLLM-Omni server subprocess."""
+        env = os.environ.copy()
+        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+        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),
+        ] + self.serve_args
+
+        print(f"Launching OmniServer with: {' '.join(cmd)}")
+        self.proc = subprocess.Popen(
+            cmd,
+            env=env,
+            cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),  # Set working directory to vllm-omni root
+        )
+
+        # Wait for server to be ready
+        max_wait = 600  # 10 minutes
+        start_time = time.time()
+        while time.time() - start_time < max_wait:
+            try:
+                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+                    sock.settimeout(1)
+                    result = sock.connect_ex((self.host, self.port))
+                    if result == 0:
+                        print(f"Server ready on {self.host}:{self.port}")
+                        return
+            except Exception:
+                pass
+            time.sleep(2)
+
+        raise RuntimeError(f"Server failed to start within {max_wait} seconds")
+
+    def _kill_process_tree(self, pid):
+        """kill process and its children"""
+        try:
+            parent = psutil.Process(pid)
+            children = parent.children(recursive=True)
+            for child in children:
+                try:
+                    child.terminate()
+                except psutil.NoSuchProcess:
+                    pass
+
+            gone, still_alive = psutil.wait_procs(children, timeout=10)
+
+            for child in still_alive:
+                try:
+                    child.kill()
+                except psutil.NoSuchProcess:
+                    pass
+
+            try:
+                parent.terminate()
+                parent.wait(timeout=10)
+            except (psutil.NoSuchProcess, psutil.TimeoutExpired):
+                try:
+                    parent.kill()
+                except psutil.NoSuchProcess:
+                    pass
+
+        except psutil.NoSuchProcess:
+            pass
+
+    def __enter__(self):
+        self._start_server()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.proc:
+            try:
+                parent = psutil.Process(self.proc.pid)
+                children = parent.children(recursive=True)
+                for child in children:
+                    try:
+                        child.terminate()
+                    except psutil.NoSuchProcess:
+                        pass
+
+                gone, still_alive = psutil.wait_procs(children, timeout=10)
+
+                for child in still_alive:
+                    try:
+                        child.kill()
+                    except psutil.NoSuchProcess:
+                        pass
+
+                try:
+                    parent.terminate()
+                    parent.wait(timeout=10)
+                except (psutil.NoSuchProcess, psutil.TimeoutExpired):
+                    try:
+                        parent.kill()
+                    except psutil.NoSuchProcess:
+                        pass
+
+            except psutil.NoSuchProcess:
+                pass

tests/diffusion/attention/test_flash_attn.py

@@ -0,0 +1,290 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Test script for FlashAttention backend with padding handling.
+
+This script tests two main scenarios:
+1. Case 1: Comparing padded vs unpadded inputs for batch_size=1
+2. Case 2: Comparing FlashAttention and SDPA backends for batch_size=2 with padding
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
+from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+
+
+def create_attention_mask(batch_size: int, seq_len: int, valid_len: int, device: torch.device) -> torch.Tensor:
+    """
+    Create attention mask where first valid_len tokens are valid (1) and rest are padding (0).
+
+    Args:
+        batch_size: Batch size
+        seq_len: Total sequence length (including padding)
+        valid_len: Number of valid (non-padded) tokens
+
+    Returns:
+        Attention mask of shape (batch_size, seq_len)
+    """
+    mask = torch.zeros(batch_size, seq_len, dtype=torch.bool, device=device)
+    mask[:, :valid_len] = True
+    return mask
+
+
+def pad_tensor(tensor: torch.Tensor, target_seq_len: int, pad_value: float = 0.0) -> torch.Tensor:
+    """
+    Pad tensor along sequence dimension (dim=1).
+
+    Args:
+        tensor: Input tensor of shape (batch_size, seq_len, num_heads, head_dim)
+        target_seq_len: Target sequence length after padding
+        pad_value: Value to use for padding
+
+    Returns:
+        Padded tensor of shape (batch_size, target_seq_len, num_heads, head_dim)
+    """
+    batch_size, seq_len, num_heads, head_dim = tensor.shape
+    if target_seq_len <= seq_len:
+        return tensor
+
+    padding = torch.full(
+        (batch_size, target_seq_len - seq_len, num_heads, head_dim), pad_value, dtype=tensor.dtype, device=tensor.device
+    )
+    return torch.cat([tensor, padding], dim=1)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="FlashAttention requires CUDA")
+def test_padding_equivalence():
+    """
+    Case 1: Test that padded and unpadded inputs produce similar outputs.
+
+    - Input A: batch_size=1, hidden_states (1, 48), encoder_hidden_states (1, 16)
+      Concatenated length: 64, NO attention_mask
+    - Input B: Same data but padded: hidden_states (1, 58), encoder_hidden_states (1, 26)
+      Concatenated length: 84, WITH attention_mask
+
+    Expected: Output A and Output B should be very close.
+    """
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+
+    # Configuration
+    batch_size = 1
+    hidden_seq_len = 48
+    encoder_seq_len = 16
+    pad_length = 10
+    num_heads = 8
+    head_dim = 64
+
+    # Initialize FlashAttention
+    fa_impl = FlashAttentionImpl(
+        num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
+    )
+
+    # Create base tensors with random values (same for both A and B)
+    torch.manual_seed(42)
+    hidden_states_base = torch.randn(batch_size, hidden_seq_len, num_heads, head_dim, device=device, dtype=dtype)
+    encoder_hidden_states_base = torch.randn(
+        batch_size, encoder_seq_len, num_heads, head_dim, device=device, dtype=dtype
+    )
+
+    # ========== Input A: Unpadded, no attention mask ==========
+    query_a = torch.cat([hidden_states_base, encoder_hidden_states_base], dim=1)
+    key_a = query_a.clone()
+    value_a = query_a.clone()
+
+    attn_metadata_a = AttentionMetadata(attn_mask=None)
+
+    output_a = fa_impl.forward(query=query_a, key=key_a, value=value_a, attn_metadata=attn_metadata_a)
+
+    # ========== Input B: Padded with attention mask ==========
+    hidden_states_padded = pad_tensor(hidden_states_base, hidden_seq_len + pad_length)
+    encoder_hidden_states_padded = pad_tensor(encoder_hidden_states_base, encoder_seq_len + pad_length)
+
+    query_b = torch.cat([hidden_states_padded, encoder_hidden_states_padded], dim=1)
+    key_b = query_b.clone()
+    value_b = query_b.clone()
+
+    # Create attention mask
+    attn_mask_b = torch.cat(
+        [
+            create_attention_mask(batch_size, hidden_seq_len + pad_length, hidden_seq_len, device),
+            create_attention_mask(batch_size, encoder_seq_len + pad_length, encoder_seq_len, device),
+        ],
+        dim=1,
+    )
+
+    attn_metadata_b = AttentionMetadata(attn_mask=attn_mask_b)
+
+    output_b = fa_impl.forward(query=query_b, key=key_b, value=value_b, attn_metadata=attn_metadata_b)
+
+    # Extract non-padded portion from output_b
+    output_b_unpadded = torch.cat(
+        [
+            output_b[:, :hidden_seq_len, :, :],
+            output_b[:, hidden_seq_len + pad_length : hidden_seq_len + pad_length + encoder_seq_len, :, :],
+        ],
+        dim=1,
+    )
+
+    # Compare outputs
+    max_diff = torch.max(torch.abs(output_a - output_b_unpadded)).item()
+    mean_diff = torch.mean(torch.abs(output_a - output_b_unpadded)).item()
+
+    print("\n=== Case 1: Padding Equivalence Test ===")
+    print(f"Output A shape: {output_a.shape}")
+    print(f"Output B shape: {output_b.shape}")
+    print(f"Output B unpadded shape: {output_b_unpadded.shape}")
+    print(f"Max absolute difference: {max_diff:.6f}")
+    print(f"Mean absolute difference: {mean_diff:.6f}")
+
+    # Assert that outputs are close
+    # Using higher tolerance for bfloat16
+    assert max_diff < 0.1, f"Max difference {max_diff} exceeds threshold 0.1"
+    assert mean_diff < 0.01, f"Mean difference {mean_diff} exceeds threshold 0.01"
+
+    print("✓ Case 1 PASSED: Padded and unpadded outputs are very close!")
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="FlashAttention requires CUDA")
+def test_fa_vs_sdpa():
+    """
+    Case 2: Compare FlashAttention and SDPA backends with padding.
+
+    - batch_size=2
+    - hidden_states: (2, 48) padded to (2, 58)
+    - encoder_hidden_states: (2, 16) padded to (2, 26)
+    - Concatenated length: 84
+    - Compare FA and SDPA outputs
+
+    Expected: FA and SDPA outputs should be very close.
+    """
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+
+    # Configuration
+    batch_size = 2
+    hidden_seq_len = 48
+    encoder_seq_len = 16
+    pad_length = 10
+    num_heads = 8
+    head_dim = 64
+
+    # Initialize both backends
+    fa_impl = FlashAttentionImpl(
+        num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
+    )
+
+    sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
+
+    # Create base tensors
+    torch.manual_seed(123)
+    hidden_states_base = torch.randn(batch_size, hidden_seq_len, num_heads, head_dim, device=device, dtype=dtype)
+    encoder_hidden_states_base = torch.randn(
+        batch_size, encoder_seq_len, num_heads, head_dim, device=device, dtype=dtype
+    )
+
+    # Pad tensors
+    hidden_states_padded = pad_tensor(hidden_states_base, hidden_seq_len + pad_length)
+    encoder_hidden_states_padded = pad_tensor(encoder_hidden_states_base, encoder_seq_len + pad_length)
+
+    # Concatenate
+    query = torch.cat([hidden_states_padded, encoder_hidden_states_padded], dim=1)
+    key = query.clone()
+    value = query.clone()
+
+    # Create attention mask
+    attn_mask = torch.cat(
+        [
+            create_attention_mask(batch_size, hidden_seq_len + pad_length, hidden_seq_len, device),
+            create_attention_mask(batch_size, encoder_seq_len + pad_length, encoder_seq_len, device),
+        ],
+        dim=1,
+    )
+
+    attn_metadata = AttentionMetadata(attn_mask=attn_mask)
+
+    # Run FlashAttention
+    output_fa = fa_impl.forward(query=query.clone(), key=key.clone(), value=value.clone(), attn_metadata=attn_metadata)
+
+    # Run SDPA
+    # SDPA expects 4D attention mask: (batch_size, 1, seq_len, seq_len) or (batch_size, seq_len)
+    # For causal=False, we need to convert 2D mask to 4D
+    if attn_mask is not None:
+        # Expand mask for SDPA: (batch_size, seq_len) -> (batch_size, 1, 1, seq_len)
+        attn_mask_4d = attn_mask.unsqueeze(1).unsqueeze(2)
+        # Convert bool to float: True -> 0.0, False -> -inf
+        attn_mask_float = torch.zeros_like(attn_mask_4d, dtype=dtype)
+        attn_mask_float.masked_fill_(~attn_mask_4d, float("-inf"))
+        attn_metadata_sdpa = AttentionMetadata(attn_mask=attn_mask_float)
+    else:
+        attn_metadata_sdpa = AttentionMetadata(attn_mask=None)
+
+    output_sdpa = sdpa_impl.forward(
+        query=query.clone(), key=key.clone(), value=value.clone(), attn_metadata=attn_metadata_sdpa
+    )
+
+    # Compare outputs (only compare valid regions)
+    output_fa_valid = torch.cat(
+        [
+            output_fa[:, :hidden_seq_len, :, :],
+            output_fa[:, hidden_seq_len + pad_length : hidden_seq_len + pad_length + encoder_seq_len, :, :],
+        ],
+        dim=1,
+    )
+    output_sdpa_valid = torch.cat(
+        [
+            output_sdpa[:, :hidden_seq_len, :, :],
+            output_sdpa[:, hidden_seq_len + pad_length : hidden_seq_len + pad_length + encoder_seq_len, :, :],
+        ],
+        dim=1,
+    )
+
+    max_diff = torch.max(torch.abs(output_fa_valid - output_sdpa_valid)).item()
+    mean_diff = torch.mean(torch.abs(output_fa_valid - output_sdpa_valid)).item()
+
+    print("\n=== Case 2: FA vs SDPA Comparison ===")
+    print(f"Batch size: {batch_size}")
+    print(f"FA output shape: {output_fa.shape}")
+    print(f"SDPA output shape: {output_sdpa.shape}")
+    print(f"Max absolute difference (valid region): {max_diff:.6f}")
+    print(f"Mean absolute difference (valid region): {mean_diff:.6f}")
+
+    # Assert that outputs are close
+    # Using higher tolerance for bfloat16 and different implementations
+    assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
+    assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
+
+    print("✓ Case 2 PASSED: FA and SDPA outputs are very close!")
+
+
+if __name__ == "__main__":
+    print("Running FlashAttention Padding Tests...")
+    print("=" * 60)
+
+    # Try to run CUDA tests
+    if torch.cuda.is_available():
+        try:
+            print("\n[Running Case 1: Padding Equivalence for FA]")
+            test_padding_equivalence()
+        except Exception as e:
+            print(f"✗ Case 1 failed: {e}")
+            import traceback
+
+            traceback.print_exc()
+
+        try:
+            print("\n[Running Case 2: FA vs SDPA]")
+            test_fa_vs_sdpa()
+        except Exception as e:
+            print(f"✗ Case 2 failed: {e}")
+            import traceback
+
+            traceback.print_exc()
+    else:
+        raise RuntimeError("CUDA is not available")
+    print("\n" + "=" * 60)
+    print("Test suite completed!")

tests/diffusion/attention/test_sequence_parallel.py

@@ -12,13 +12,13 @@ from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import (
     DiffusionParallelConfig,
     OmniDiffusionConfig,
-    set_current_omni_diffusion_config,
 )
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
     init_distributed_environment,
     initialize_model_parallel,
 )
+from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.utils.platform_utils import detect_device_type
 
 device_type = detect_device_type()
@@ -385,7 +385,7 @@ def ulysses_attention_on_test_model(
     )
 
     # Set the config so Attention can access it
-    with set_current_omni_diffusion_config(od_config):
+    with set_forward_context(omni_diffusion_config=od_config):
         # Create model
         hidden_size = num_heads * head_size
 

tests/diffusion/models/z_image/test_zimage_tp_constraints.py

@@ -0,0 +1,44 @@
+import pytest
+
+from vllm_omni.diffusion.models.z_image.z_image_transformer import validate_zimage_tp_constraints
+
+
+def test_validate_zimage_tp_constraints_tp2_ok():
+    ffn_hidden_dim, final_out_dims, supported_tp = validate_zimage_tp_constraints(
+        dim=3840,
+        n_heads=30,
+        n_kv_heads=30,
+        in_channels=16,
+        all_patch_size=(2,),
+        all_f_patch_size=(1,),
+        tensor_parallel_size=2,
+    )
+    assert ffn_hidden_dim == 10240
+    assert final_out_dims == [64]
+    assert supported_tp == [1, 2]
+
+
+def test_validate_zimage_tp_constraints_tp4_fails_on_heads():
+    with pytest.raises(ValueError, match=r"n_heads % tensor_parallel_size"):
+        validate_zimage_tp_constraints(
+            dim=3840,
+            n_heads=30,
+            n_kv_heads=30,
+            in_channels=16,
+            all_patch_size=(2,),
+            all_f_patch_size=(1,),
+            tensor_parallel_size=4,
+        )
+
+
+def test_validate_zimage_tp_constraints_tp3_fails_on_ffn_hidden_dim():
+    with pytest.raises(ValueError, match=r"ffn_hidden_dim % tensor_parallel_size"):
+        validate_zimage_tp_constraints(
+            dim=3840,
+            n_heads=30,
+            n_kv_heads=30,
+            in_channels=16,
+            all_patch_size=(2,),
+            all_f_patch_size=(1,),
+            tensor_parallel_size=3,
+        )

tests/diffusion/test_gpu_diffusion_worker.py

@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Unit tests for GPUDiffusionWorker class.
+
+This module tests the GPUDiffusionWorker implementation:
+- load_weights: Loading model weights
+- sleep: Putting worker into sleep mode (levels 1 and 2)
+- wake_up: Waking worker from sleep mode
+"""
+
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.worker.gpu_diffusion_worker import GPUDiffusionWorker
+
+
+@pytest.fixture
+def mock_od_config():
+    """Create a mock OmniDiffusionConfig."""
+    config = Mock()
+    config.num_gpus = 1
+    config.master_port = 12345
+    config.enable_sleep_mode = False
+    config.cache_backend = None
+    config.cache_config = None
+    config.model = "test-model"
+    return config
+
+
+@pytest.fixture
+def mock_gpu_worker(mock_od_config):
+    """Create a GPUDiffusionWorker with mocked initialization."""
+    with patch.object(GPUDiffusionWorker, "init_device"):
+        worker = GPUDiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
+        # Mock the model_runner with pipeline
+        worker.model_runner = Mock()
+        worker.model_runner.pipeline = Mock()
+        worker._sleep_saved_buffers = {}
+        return worker
+
+
+class TestGPUDiffusionWorkerLoadWeights:
+    """Test GPUDiffusionWorker.load_weights method."""
+
+    def test_load_weights_calls_pipeline(self, mock_gpu_worker):
+        """Test that load_weights delegates to model_runner.load_weights."""
+        # Setup mock weights
+        mock_weights = [
+            ("layer1.weight", torch.randn(10, 10)),
+            ("layer2.weight", torch.randn(20, 20)),
+        ]
+        expected_loaded = {"layer1.weight", "layer2.weight"}
+
+        # Configure model_runner mock
+        mock_gpu_worker.model_runner.load_weights = Mock(return_value=expected_loaded)
+
+        # Call load_weights
+        result = mock_gpu_worker.load_weights(mock_weights)
+
+        # Verify model_runner.load_weights was called with the weights
+        mock_gpu_worker.model_runner.load_weights.assert_called_once_with(mock_weights)
+        assert result == expected_loaded
+
+    def test_load_weights_empty_iterable(self, mock_gpu_worker):
+        """Test load_weights with empty weights iterable."""
+        mock_gpu_worker.model_runner.load_weights = Mock(return_value=set())
+
+        result = mock_gpu_worker.load_weights([])
+
+        mock_gpu_worker.model_runner.load_weights.assert_called_once_with([])
+        assert result == set()
+
+
+class TestGPUDiffusionWorkerSleep:
+    """Test GPUDiffusionWorker.sleep method."""
+
+    @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
+    @patch("vllm.device_allocator.cumem.CuMemAllocator")
+    def test_sleep_level_1(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+        """Test sleep mode level 1 (offload weights only)."""
+        # Setup memory info mocks
+        # Before sleep: 1GB free, 8GB total
+        # After sleep: 3GB free, 8GB total (freed 2GB)
+        mock_mem_info.side_effect = [
+            (1 * 1024**3, 8 * 1024**3),  # Before sleep
+            (3 * 1024**3, 8 * 1024**3),  # After sleep
+        ]
+
+        # Setup allocator mock
+        mock_allocator = Mock()
+        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
+        mock_allocator.sleep = Mock()
+
+        # Call sleep with level 1
+        result = mock_gpu_worker.sleep(level=1)
+
+        # Verify sleep was called with correct tags
+        mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",))
+        assert result is True
+        # Verify buffers were NOT saved (level 1 doesn't save buffers)
+        assert len(mock_gpu_worker._sleep_saved_buffers) == 0
+
+    @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
+    @patch("vllm.device_allocator.cumem.CuMemAllocator")
+    def test_sleep_level_2(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+        """Test sleep mode level 2 (offload all, save buffers)."""
+        # Setup memory info mocks
+        mock_mem_info.side_effect = [
+            (1 * 1024**3, 8 * 1024**3),  # Before sleep
+            (5 * 1024**3, 8 * 1024**3),  # After sleep (freed 4GB)
+        ]
+
+        # Setup allocator mock
+        mock_allocator = Mock()
+        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
+        mock_allocator.sleep = Mock()
+
+        # Mock pipeline buffers
+        mock_buffer1 = torch.randn(10, 10)
+        mock_buffer2 = torch.randn(20, 20)
+        mock_gpu_worker.model_runner.pipeline.named_buffers = Mock(
+            return_value=[
+                ("buffer1", mock_buffer1),
+                ("buffer2", mock_buffer2),
+            ]
+        )
+
+        # Call sleep with level 2
+        result = mock_gpu_worker.sleep(level=2)
+
+        # Verify sleep was called with empty tags (offload all)
+        mock_allocator.sleep.assert_called_once_with(offload_tags=tuple())
+        assert result is True
+
+        # Verify buffers were saved
+        assert len(mock_gpu_worker._sleep_saved_buffers) == 2
+        assert "buffer1" in mock_gpu_worker._sleep_saved_buffers
+        assert "buffer2" in mock_gpu_worker._sleep_saved_buffers
+
+    @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
+    @patch("vllm.device_allocator.cumem.CuMemAllocator")
+    def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+        """Test that sleep validates memory was actually freed."""
+        # Simulate memory increase (should trigger assertion error)
+        mock_mem_info.side_effect = [
+            (3 * 1024**3, 8 * 1024**3),  # Before sleep: 3GB free
+            (1 * 1024**3, 8 * 1024**3),  # After sleep: 1GB free (negative freed!)
+        ]
+
+        mock_allocator = Mock()
+        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
+        mock_allocator.sleep = Mock()
+
+        # This should raise an assertion error
+        with pytest.raises(AssertionError, match="Memory usage increased after sleeping"):
+            mock_gpu_worker.sleep(level=1)
+
+
+class TestGPUDiffusionWorkerWakeUp:
+    """Test GPUDiffusionWorker.wake_up method."""
+
+    @patch("vllm.device_allocator.cumem.CuMemAllocator")
+    def test_wake_up_without_buffers(self, mock_allocator_class, mock_gpu_worker):
+        """Test wake_up without saved buffers (level 1 sleep)."""
+        # Setup allocator mock
+        mock_allocator = Mock()
+        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
+        mock_allocator.wake_up = Mock()
+
+        # Ensure no saved buffers
+        mock_gpu_worker._sleep_saved_buffers = {}
+
+        # Call wake_up
+        result = mock_gpu_worker.wake_up(tags=["weights"])
+
+        # Verify allocator.wake_up was called
+        mock_allocator.wake_up.assert_called_once_with(["weights"])
+        assert result is True
+
+    @patch("vllm.device_allocator.cumem.CuMemAllocator")
+    def test_wake_up_with_buffers(self, mock_allocator_class, mock_gpu_worker):
+        """Test wake_up with saved buffers (level 2 sleep)."""
+        # Setup allocator mock
+        mock_allocator = Mock()
+        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
+        mock_allocator.wake_up = Mock()
+
+        # Create saved buffers
+        saved_buffer1 = torch.randn(10, 10)
+        saved_buffer2 = torch.randn(20, 20)
+        mock_gpu_worker._sleep_saved_buffers = {
+            "buffer1": saved_buffer1,
+            "buffer2": saved_buffer2,
+        }
+
+        # Mock pipeline buffers (these will be restored)
+        mock_buffer1 = Mock()
+        mock_buffer1.data = Mock()
+        mock_buffer2 = Mock()
+        mock_buffer2.data = Mock()
+
+        mock_gpu_worker.model_runner.pipeline.named_buffers = Mock(
+            return_value=[
+                ("buffer1", mock_buffer1),
+                ("buffer2", mock_buffer2),
+            ]
+        )
+
+        # Call wake_up
+        result = mock_gpu_worker.wake_up(tags=None)
+
+        # Verify allocator.wake_up was called
+        mock_allocator.wake_up.assert_called_once_with(None)
+
+        # Verify buffers were restored
+        mock_buffer1.data.copy_.assert_called_once()
+        mock_buffer2.data.copy_.assert_called_once()
+
+        # Verify saved buffers were cleared
+        assert len(mock_gpu_worker._sleep_saved_buffers) == 0
+        assert result is True
+
+    @patch("vllm.device_allocator.cumem.CuMemAllocator")
+    def test_wake_up_partial_buffer_restore(self, mock_allocator_class, mock_gpu_worker):
+        """Test wake_up only restores buffers that were saved."""
+        # Setup allocator mock
+        mock_allocator = Mock()
+        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
+        mock_allocator.wake_up = Mock()
+
+        # Only save buffer1, not buffer2
+        saved_buffer1 = torch.randn(10, 10)
+        mock_gpu_worker._sleep_saved_buffers = {
+            "buffer1": saved_buffer1,
+        }
+
+        # Mock pipeline has both buffers
+        mock_buffer1 = Mock()
+        mock_buffer1.data = Mock()
+        mock_buffer2 = Mock()
+        mock_buffer2.data = Mock()
+
+        mock_gpu_worker.model_runner.pipeline.named_buffers = Mock(
+            return_value=[
+                ("buffer1", mock_buffer1),
+                ("buffer2", mock_buffer2),
+            ]
+        )
+
+        # Call wake_up
+        result = mock_gpu_worker.wake_up()
+
+        # Verify only buffer1 was restored
+        mock_buffer1.data.copy_.assert_called_once()
+        # buffer2 should NOT be restored since it wasn't saved
+        mock_buffer2.data.copy_.assert_not_called()
+
+        assert result is True

tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml

@@ -45,7 +45,7 @@ stage_args:
        worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.3
-       enforce_eager: true
+       enforce_eager: false
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
       #  tensor_parallel_size: 2

tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml

@@ -0,0 +1,98 @@
+# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
+# Stage 0: Thinker (multimodal understanding + text generation)
+# Stage 1: Talker (text embeddings → 16-layer RVQ codec codes)
+# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
+
+# The following config has been verified on 2x H100-80G GPUs.
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0,1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.6
+      enforce_eager: false
+      trust_remote_code: true
+      engine_output_type: latent  # Output hidden states for talker
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      hf_config_name: thinker_config
+      tensor_parallel_size: 2
+      load_format: dummy
+    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
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    runtime:
+       devices: "1"
+       max_batch_size: 1
+    engine_args:
+       model_stage: talker
+       model_arch: Qwen3OmniMoeForConditionalGeneration
+       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+       gpu_memory_utilization: 0.3
+       enforce_eager: true
+       trust_remote_code: true
+       engine_output_type: latent  # Output codec codes for code2wav
+      #  tensor_parallel_size: 2
+       enable_prefix_caching: false
+       distributed_executor_backend: "mp"
+       hf_config_name: talker_config
+       load_format: dummy
+    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: 100
+      seed: 42
+      detokenize: False
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 2
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      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.1
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 1000000
+      hf_config_name: thinker_config
+      load_format: dummy
+    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: 200
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -0,0 +1,90 @@
+import sys
+import threading
+import time
+from pathlib import Path
+
+import pytest
+import torch
+
+from vllm_omni.utils.platform_utils import is_npu, is_rocm
+
+# ruff: noqa: E402
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni import Omni
+
+
+class GPUMemoryMonitor:
+    """Poll global device memory usage via CUDA APIs."""
+
+    def __init__(self, device_index: int, interval: float = 0.05):
+        self.device_index = device_index
+        self.interval = interval
+        self.peak_used_mb = 0.0
+        self._stop_event = threading.Event()
+        self._thread: threading.Thread | None = None
+
+    def start(self) -> None:
+        def monitor_loop() -> None:
+            while not self._stop_event.is_set():
+                try:
+                    with torch.cuda.device(self.device_index):
+                        free_bytes, total_bytes = torch.cuda.mem_get_info()
+                    used_mb = (total_bytes - free_bytes) / (1024**2)
+                    self.peak_used_mb = max(self.peak_used_mb, used_mb)
+                except Exception:
+                    pass
+                time.sleep(self.interval)
+
+        self._thread = threading.Thread(target=monitor_loop, daemon=True)
+        self._thread.start()
+
+    def stop(self) -> None:
+        if self._thread is None:
+            return
+        self._stop_event.set()
+        self._thread.join(timeout=2.0)
+
+
+models = ["riverclouds/qwen_image_random"]
+
+
+@pytest.mark.skipif(is_npu() or is_rocm(), reason="Hardware not supported")
+@pytest.mark.parametrize("model_name", models)
+def test_cpu_offload_diffusion_model(model_name: str):
+    def inference(offload: bool = True):
+        torch.cuda.empty_cache()
+        device_index = torch.cuda.current_device()
+        monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
+        monitor.start()
+        m = Omni(model=model_name, enable_cpu_offload=offload)
+        torch.cuda.reset_peak_memory_stats(device=device_index)
+        height = 256
+        width = 256
+
+        m.generate(
+            "a photo of a cat sitting on a laptop keyboard",
+            height=height,
+            width=width,
+            num_inference_steps=9,
+            guidance_scale=0.0,
+            generator=torch.Generator("cuda").manual_seed(42),
+        )
+
+        monitor.stop()
+        torch.cuda.synchronize(device_index)
+        fallback_alloc = torch.cuda.max_memory_allocated(device=device_index) / (1024**2)
+        fallback_reserved = torch.cuda.max_memory_reserved(device=device_index) / (1024**2)
+        peak_memory_mb = max(monitor.peak_used_mb, fallback_alloc, fallback_reserved)
+
+        return peak_memory_mb
+
+    offload_peak_memory = inference(offload=True)
+    no_offload_peak_memory = inference(offload=False)
+    print(f"Offload peak memory: {offload_peak_memory} MB")
+    print(f"No offload peak memory: {no_offload_peak_memory} MB")
+    assert offload_peak_memory + 2500 < no_offload_peak_memory, (
+        f"Offload peak memory {offload_peak_memory} MB should be less than no offload peak memory {no_offload_peak_memory} MB"
+    )

tests/e2e/offline_inference/test_qwen3_omni.py

@@ -10,14 +10,20 @@ from pathlib import Path
 import pytest
 from vllm.assets.video import VideoAsset
 
+from vllm_omni.utils import is_rocm
+
 from .conftest import OmniRunner
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
-# CI stage config for 2xH100-80G GPUs
-stage_configs = [str(Path(__file__).parent / "stage_configs" / "qwen3_omni_ci.yaml")]
+# CI stage config for 2xH100-80G GPUs or AMD GPU MI325
+if is_rocm():
+    # ROCm stage config optimized for MI325 GPU
+    stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
+else:
+    stage_configs = [str(Path(__file__).parent / "stage_configs" / "qwen3_omni_ci.yaml")]
 
 # Create parameter combinations for model and stage config
 test_params = [(model, stage_config) for model in models for stage_config in stage_configs]

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -202,3 +202,114 @@ def test_sequence_parallel(
         f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}); "
         f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
     )
+
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_sequence_parallel_ulysses_sp_only(
+    model_name: str,
+    dtype: torch.dtype,
+    attn_backend: str,
+):
+    """Test sequence parallel with ulysses_degree=4, ring_degree=1, and the image size (332x332) where the sequence length is NOT divisible by sp_size."""
+    ulysses_degree = 4
+    ring_degree = 1
+
+    # Skip if not enough GPUs available for SP run
+    if device_count() < ulysses_degree * ring_degree:
+        pytest.skip(f"Test requires {ulysses_degree * ring_degree} GPUs but only {device_count()} available")
+
+    # (272/8) * (272/8) = 17 * 17 = 289 Not divisible by sp_size=4
+    height = 272
+    width = 272
+    num_inference_steps = 4  # Minimal steps for fast test
+    seed = 42
+
+    # Step 1: Baseline (no Ulysses sequence parallel)
+    baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1)
+    baseline = Omni(
+        model=model_name,
+        parallel_config=baseline_parallel_config,
+        dtype=dtype,
+        attention_backend=attn_backend,
+    )
+    try:
+        outputs = baseline.generate(
+            PROMPT,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=0.0,
+            generator=torch.Generator(get_device_name()).manual_seed(seed),
+            num_outputs_per_prompt=1,
+        )
+        baseline_images = outputs[0].request_output[0].images
+    finally:
+        baseline.close()
+        if dist.is_initialized():
+            dist.destroy_process_group()
+        for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
+            os.environ.pop(key, None)
+        time.sleep(2)  # Wait for resources to release
+
+    assert baseline_images is not None
+    assert len(baseline_images) == 1
+    assert baseline_images[0].width == width
+    assert baseline_images[0].height == height
+
+    # Step 2: SP (Ulysses-SP + Ring-SP)
+    sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
+    sp = Omni(
+        model=model_name,
+        parallel_config=sp_parallel_config,
+        dtype=dtype,
+        attention_backend=attn_backend,
+    )
+    try:
+        outputs = sp.generate(
+            PROMPT,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=0.0,
+            generator=torch.Generator(get_device_name()).manual_seed(seed),
+            num_outputs_per_prompt=1,
+        )
+        sp_images = outputs[0].request_output[0].images
+    finally:
+        sp.close()
+        if dist.is_initialized():
+            dist.destroy_process_group()
+        for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
+            os.environ.pop(key, None)
+        time.sleep(2)
+
+    assert sp_images is not None
+    assert len(sp_images) == 1
+    assert sp_images[0].width == width
+    assert sp_images[0].height == height
+
+    # Step 3: Compare outputs
+    mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0])
+
+    # FP16/BF16 may differ slightly due to different computation order under parallelism.
+    if dtype in (torch.float16, torch.bfloat16):
+        mean_threshold = 2e-2
+        max_threshold = 2e-1
+    else:
+        mean_threshold = 1e-2
+        max_threshold = 1e-1
+
+    print(
+        "Image diff stats (baseline ulysses_degree*ring_degree=1 vs SP): "
+        f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; "
+        f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
+        f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
+    )
+
+    assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
+        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} "
+        f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}); "
+        f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}"
+    )

tests/e2e/offline_inference/test_stable_audio_model.py

@@ -0,0 +1,66 @@
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.outputs import OmniRequestOutput
+
+# ruff: noqa: E402
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni import Omni
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+# Use random weights model for CI testing (small, no authentication required)
+models = ["linyueqian/stable_audio_random"]
+
+
+@pytest.mark.parametrize("model_name", models)
+def test_stable_audio_model(model_name: str):
+    m = Omni(model=model_name)
+
+    # Use minimal settings for testing
+    # Generate a short 2-second audio clip with minimal inference steps
+    audio_start_in_s = 0.0
+    audio_end_in_s = 2.0  # Short duration for fast testing
+    sample_rate = 44100  # Stable Audio uses 44100 Hz
+
+    outputs = m.generate(
+        "The sound of a dog barking",
+        negative_prompt="Low quality.",
+        num_inference_steps=4,  # Minimal steps for speed
+        guidance_scale=7.0,
+        generator=torch.Generator("cuda").manual_seed(42),
+        num_outputs_per_prompt=1,
+        extra={
+            "audio_start_in_s": audio_start_in_s,
+            "audio_end_in_s": audio_end_in_s,
+        },
+    )
+
+    # Extract audio from OmniRequestOutput
+    assert outputs is not None
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"  # Generic output type
+    assert hasattr(first_output, "request_output") and first_output.request_output
+
+    req_out = first_output.request_output[0]
+    assert isinstance(req_out, OmniRequestOutput)
+    assert hasattr(req_out, "images") and len(req_out.images) >= 1
+
+    # For stable audio, the "images" field contains audio numpy arrays
+    audio = req_out.images[0]
+    assert isinstance(audio, np.ndarray)
+    # audio shape: (batch, channels, samples)
+    # For stable-audio-open-1.0: sample_rate=44100, so 2 seconds = 88200 samples
+    assert audio.ndim == 3
+    assert audio.shape[0] == 1  # batch size
+    assert audio.shape[1] == 2  # stereo channels
+    expected_samples = int((audio_end_in_s - audio_start_in_s) * sample_rate)
+    assert audio.shape[2] == expected_samples  # 88200 samples for 2 seconds

tests/e2e/offline_inference/test_zimage_tensor_parallel.py

@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
+
+# ruff: noqa: E402
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.utils.platform_utils import is_npu, is_rocm
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+PROMPT = "a photo of a cat sitting on a laptop keyboard"
+
+
+def _get_zimage_model() -> str:
+    # Allow overriding the model for local/offline environments.
+    # Can be either a HuggingFace repo id or a local path.
+    return os.environ.get("VLLM_TEST_ZIMAGE_MODEL", "Tongyi-MAI/Z-Image-Turbo")
+
+
+def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor:
+    """Convert PIL image to float32 RGB tensor in [0, 1] with shape [H, W, 3]."""
+    arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
+    return torch.from_numpy(arr)
+
+
+def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
+    """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1]."""
+    ta = _pil_to_float_rgb_tensor(a)
+    tb = _pil_to_float_rgb_tensor(b)
+    assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
+    abs_diff = torch.abs(ta - tb)
+    return abs_diff.mean().item(), abs_diff.max().item()
+
+
+def _extract_single_image(outputs) -> Image.Image:
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    if not hasattr(first_output, "request_output") or not first_output.request_output:
+        raise ValueError("No request_output found in OmniRequestOutput")
+
+    req_out = first_output.request_output[0]
+    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
+        raise ValueError("Invalid request_output structure or missing 'images' key")
+
+    images = req_out.images
+    if images is None or len(images) != 1:
+        raise ValueError(f"Expected 1 image, got {0 if images is None else len(images)}")
+    return images[0]
+
+
+def _run_zimage_generate(
+    *, tp_size: int, height: int, width: int, num_inference_steps: int, seed: int
+) -> tuple[Image.Image, float]:
+    m = Omni(
+        model=_get_zimage_model(),
+        parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp_size),
+    )
+    try:
+        # NOTE: Omni closes itself when a generate() call is exhausted.
+        # To avoid measuring teardown time (process shutdown, memory cleanup),
+        # we measure the latency to produce *subsequent* outputs within a single
+        # generator run.
+        #
+        # This also serves as a warmup: the first output may include extra
+        # compilation/caching overhead, while later outputs are closer to
+        # steady-state inference.
+        num_requests = 4  # 1 warmup + 3 timed
+        gen = m.generate(
+            [PROMPT] * num_requests,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=0.0,
+            seed=seed,
+            num_outputs_per_prompt=1,
+            py_generator=True,
+        )
+
+        warmup_output = next(gen)
+        t_prev = time.perf_counter()
+        per_request_times_s: list[float] = []
+        last_output = warmup_output
+        for _ in range(num_requests - 1):
+            last_output = next(gen)
+            t_now = time.perf_counter()
+            per_request_times_s.append(t_now - t_prev)
+            t_prev = t_now
+
+        # Ensure the generator is fully consumed so it can clean up.
+        for _ in gen:
+            pass
+
+        median_time_s = float(np.median(per_request_times_s))
+        return _extract_single_image([last_output]), median_time_s
+    finally:
+        m.close()
+        cleanup_dist_env_and_memory()
+
+
+@pytest.mark.integration
+def test_zimage_tensor_parallel_tp2(tmp_path: Path):
+    if is_npu() or is_rocm():
+        pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.")
+    if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
+        pytest.skip("Z-Image TP=2 requires >= 2 CUDA devices.")
+
+    height = 256
+    width = 256
+    num_inference_steps = 2
+    seed = 42
+
+    tp1_img, tp1_time_s = _run_zimage_generate(
+        tp_size=1,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+    )
+    tp2_img, tp2_time_s = _run_zimage_generate(
+        tp_size=2,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+    )
+
+    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
+
+    mean_abs_diff, max_abs_diff = _diff_metrics(tp1_img, tp2_img)
+    mean_threshold = 3e-2
+    max_threshold = 3.5e-1
+    print(
+        "Z-Image TP image diff stats (TP=1 vs TP=2): "
+        f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; "
+        f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
+        f"tp1_img={tp1_path}, tp2_img={tp2_path}"
+    )
+    assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
+        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} "
+        f"(thresholds: mean<={mean_threshold:.6e}, max<={max_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})"

tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml

@@ -45,7 +45,7 @@ stage_args:
        worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
        gpu_memory_utilization: 0.3
-       enforce_eager: true
+       enforce_eager: false
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
       #  tensor_parallel_size: 2

tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml

@@ -0,0 +1,98 @@
+# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
+# Stage 0: Thinker (multimodal understanding + text generation)
+# Stage 1: Talker (text embeddings → 16-layer RVQ codec codes)
+# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
+
+# The following config has been verified on 2x H100-80G GPUs.
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0,1"
+      max_batch_size: 5
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.6
+      enforce_eager: false
+      trust_remote_code: true
+      engine_output_type: latent  # Output hidden states for talker
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      hf_config_name: thinker_config
+      tensor_parallel_size: 2
+      load_format: dummy
+    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
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    runtime:
+       devices: "1"
+       max_batch_size: 5
+    engine_args:
+       model_stage: talker
+       model_arch: Qwen3OmniMoeForConditionalGeneration
+       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+       gpu_memory_utilization: 0.3
+       enforce_eager: true
+       trust_remote_code: true
+       engine_output_type: latent  # Output codec codes for code2wav
+      #  tensor_parallel_size: 2
+       enable_prefix_caching: false
+       distributed_executor_backend: "mp"
+       hf_config_name: talker_config
+       load_format: dummy
+    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: 100
+      seed: 42
+      detokenize: False
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 2
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      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.1
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 1000000
+      hf_config_name: thinker_config
+      load_format: dummy
+    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: 200
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1

tests/e2e/online_serving/test_async_omni.py

@@ -5,7 +5,7 @@ from pathlib import Path
 import pytest
 from vllm.inputs import PromptType
 
-from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.async_omni import AsyncOmni, ClientRequestState
 
 SEED = 42
 
@@ -98,3 +98,46 @@ async def test_abort():
         num_generated_tokens, request_id = await task
         assert num_generated_tokens == NUM_EXPECTED_TOKENS
     await asyncio.sleep(5)
+
+
+@pytest.mark.asyncio
+async def test_build_and_log_summary(monkeypatch):
+    from vllm_omni.entrypoints.utils import get_final_stage_id_for_e2e
+
+    RealCRS = ClientRequestState
+    capture_metrics = {}
+
+    class MockCRS(RealCRS):
+        def __init__(self, request_id: str):
+            super().__init__(request_id)
+            capture_metrics[request_id] = self
+
+    monkeypatch.setattr("vllm_omni.entrypoints.async_omni.ClientRequestState", MockCRS)
+    monkeypatch.setattr("vllm_omni.entrypoints.client_request_state.ClientRequestState", MockCRS)
+
+    with ExitStack() as after:
+        engine = AsyncOmni(model=model, stage_configs_path=stage_config)
+        after.callback(engine.shutdown)
+        prompt = "Hello my name is Robert and "
+        NUM_EXPECTED_TOKENS = 100
+        NUM_REQUESTS = 5
+        request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]
+
+        # Create concurrent requests.
+        tasks: list[asyncio.Task] = []
+        for idx, request_id in enumerate(request_ids):
+            tasks.append(asyncio.create_task(generate(engine, request_id, prompt, NUM_EXPECTED_TOKENS)))
+
+        # Confirm the requests are okay.
+        for idx, task in enumerate(tasks):
+            await task
+            output_modalities = ["text"]
+            final_stage_id_for_e2e = get_final_stage_id_for_e2e(
+                output_modalities, engine.output_modalities, engine.stage_list
+            )
+            summary = capture_metrics[request_ids[idx]].metrics.build_and_log_summary(final_stage_id_for_e2e)
+
+            # Check that total tokens matches sum of stage tokens.
+            assert summary["e2e_total_tokens"] == sum(stage["tokens"] for stage in summary["stages"])
+            # Check that total time matches sum of stage times.
+            assert summary["e2e_total_time_ms"] >= sum(stage["total_time_ms"] for stage in summary["stages"])

tests/e2e/online_serving/test_qwen3_omni.py

@@ -20,12 +20,18 @@ import pytest
 from vllm.assets.video import VideoAsset
 from vllm.utils import get_open_port
 
+from vllm_omni.utils import is_rocm
+
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
-# CI stage config for 2*H100-80G GPUs
-stage_configs = [str(Path(__file__).parent / "stage_configs" / "qwen3_omni_ci.yaml")]
+# CI stage config for 2xH100-80G GPUs or AMD GPU MI325
+if is_rocm():
+    # ROCm stage config optimized for MI325 GPU
+    stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
+else:
+    stage_configs = [str(Path(__file__).parent / "stage_configs" / "qwen3_omni_ci.yaml")]
 
 # Create parameter combinations for model and stage config
 test_params = [(model, stage_config) for model in models for stage_config in stage_configs]

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -0,0 +1,158 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E Online tests for Qwen3-Omni model.
+"""
+
+import concurrent.futures
+import os
+import time
+from pathlib import Path
+
+import openai
+import pytest
+
+from tests.conftest import (
+    OmniServer,
+    convert_audio_to_text,
+    cosine_similarity_text,
+    dummy_messages_from_mix_data,
+    modify_stage_config,
+)
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
+
+# CI stage config for 2*H100-80G GPUs
+stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")]
+
+# Create parameter combinations for model and stage config
+test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+
+
+def client(omni_server):
+    """OpenAI client for the running vLLM-Omni server."""
+    return openai.OpenAI(
+        base_url=f"http://{omni_server.host}:{omni_server.port}/v1",
+        api_key="EMPTY",
+    )
+
+
+def get_system_prompt():
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, "
+                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
+                    "as well as generating text and speech."
+                ),
+            }
+        ],
+    }
+
+
+def get_prompt(prompt_type="text_only"):
+    prompts = {
+        "text_only": "What is the capital of China?",
+        "mix": "What is recited in the audio? What is in this image? Describe the video briefly.",
+    }
+    return prompts.get(prompt_type, prompts["text_only"])
+
+
+def get_max_batch_size(size_type="few"):
+    batch_sizes = {"few": 5, "medium": 100, "large": 256}
+    return batch_sizes.get(size_type, 5)
+
+
+@pytest.mark.parametrize("test_config", test_params)
+def test_text_to_text_001(test_config: tuple[str, str]) -> None:
+    """Test processing text, generating text output via OpenAI API."""
+    model, stage_config_path = test_config
+    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
+        messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
+
+        # Test single completion
+        api_client = client(server)
+        start_time = time.perf_counter()
+        chat_completion = api_client.chat.completions.create(
+            model=server.model, messages=messages, max_tokens=20, modalities=["text"]
+        )
+        # Verify E2E
+        print(f"the request e2e is: {time.perf_counter() - start_time}")
+        # TODO: Verify the E2E latency after confirmation baseline.
+
+        # Verify only output text
+        assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
+
+        # Verify text output success
+        text_choice = chat_completion.choices[0]
+        assert text_choice.message.content is not None, "No text output is generated"
+        assert chat_completion.usage.completion_tokens <= 20, "The output length more than the requested max_tokens."
+        assert "beijing" in text_choice.message.content.lower(), "The output do not contain keywords."
+
+
+@pytest.mark.parametrize("test_config", test_params)
+def test_text_to_text_audio_001(test_config: tuple[str, str]) -> None:
+    """Test processing text, generating text and audio output via OpenAI API."""
+
+    model, stage_config_path = test_config
+    num_concurrent_requests = get_max_batch_size()
+    stage_config_path = modify_stage_config(
+        stage_config_path,
+        {
+            0: {"runtime.max_batch_size": num_concurrent_requests},
+            1: {"runtime.max_batch_size": num_concurrent_requests},
+        },
+    )
+    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
+        messages = dummy_messages_from_mix_data(
+            system_prompt=get_system_prompt(), content_text="What is the capital of China?"
+        )
+
+        # Test single completion
+        api_client = client(server)
+        e2e_list = list()
+        with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
+            # Submit multiple completion requests concurrently
+            futures = [
+                executor.submit(api_client.chat.completions.create, model=server.model, messages=messages)
+                for _ in range(num_concurrent_requests)
+            ]
+            start_time = time.perf_counter()
+            # Wait for all requests to complete and collect results
+            chat_completions = list()
+            for future in concurrent.futures.as_completed(futures):
+                chat_completions.append(future.result())
+                # Verify E2E
+                current_e2e = time.perf_counter() - start_time
+                print(f"the request e2e is: {current_e2e}")
+                # TODO: Verify the E2E latency after confirmation baseline.
+                e2e_list.append(current_e2e)
+
+        print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
+        # Verify all completions succeeded
+        assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded."
+        for chat_completion in chat_completions:
+            # Verify audio output success
+            audio_message = chat_completion.choices[1].message
+            audio_data = audio_message.audio.data
+            assert audio_data is not None, "No audio output is generated"
+            assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
+
+            # Verify text output success
+            text_choice = chat_completion.choices[0]
+            text_content = text_choice.message.content
+            assert text_choice.message.content is not None, "No text output is generated"
+            assert "beijing" in text_choice.message.content.lower(), "The output do not contain keywords."
+
+            # Verify text output same as audio output
+            audio_content = convert_audio_to_text(audio_data)
+            print(f"text content is: {text_content}")
+            print(f"audio content is: {audio_content}")
+            assert cosine_similarity_text(audio_content.lower(), text_content.lower()) > 0.9, (
+                "The audio content is not same as the text"
+            )

tests/e2e/stage_configs/qwen3_omni_ci.yaml

@@ -0,0 +1,95 @@
+# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
+# Stage 0: Thinker (multimodal understanding + text generation)
+# Stage 1: Talker (text embeddings → 16-layer RVQ codec codes)
+# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
+
+# The following config has been verified on 2x H100-80G GPUs.
+stage_args:
+- stage_id: 0
+  runtime:
+    devices: "0,1"
+    max_batch_size: 1
+  engine_args:
+    model_stage: thinker
+    model_arch: Qwen3OmniMoeForConditionalGeneration
+    worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+    gpu_memory_utilization: 0.6
+    enforce_eager: true
+    trust_remote_code: true
+    engine_output_type: latent  # Output hidden states for talker
+    distributed_executor_backend: "mp"
+    max_num_batched_tokens: 32768
+    enable_prefix_caching: false
+    hf_config_name: thinker_config
+    tensor_parallel_size: 2
+  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
+  runtime:
+    devices: "1"
+    max_batch_size: 1
+  engine_args:
+    model_stage: talker
+    model_arch: Qwen3OmniMoeForConditionalGeneration
+    worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+    gpu_memory_utilization: 0.3
+    enforce_eager: true
+    trust_remote_code: true
+    engine_output_type: latent  # Output codec codes for code2wav
+    enable_prefix_caching: false
+    max_num_batched_tokens: 32768
+    distributed_executor_backend: "mp"
+    hf_config_name: talker_config
+  engine_input_source: [0]
+  custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+  default_sampling_params:
+    temperature: 0.9
+    top_k: 50
+    max_tokens: 100
+    seed: 42
+    detokenize: False
+    repetition_penalty: 1.05
+    stop_token_ids: [2150]
+
+- stage_id: 2
+  runtime:
+    devices: "0"
+    max_batch_size: 1
+  engine_args:
+    model_stage: code2wav
+    model_arch: Qwen3OmniMoeForConditionalGeneration
+    worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+    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.1
+    distributed_executor_backend: "mp"
+    max_num_batched_tokens: 1000000
+    hf_config_name: thinker_config
+  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: 200
+    seed: 42
+    detokenize: True
+    repetition_penalty: 1.1

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -7,7 +7,7 @@ from vllm_omni.entrypoints.async_omni import AsyncOmni
 
 def test_default_stage_config_includes_cache_backend(monkeypatch):
     """Ensure cache_backend/cache_config are preserved in default diffusion stage."""
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model: [])
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
     monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
     monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
     monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
@@ -37,7 +37,7 @@ def test_default_stage_config_includes_cache_backend(monkeypatch):
 
 def test_default_cache_config_used_when_missing(monkeypatch):
     """Ensure default cache_config is applied when cache_backend is set."""
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model: [])
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
     monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
     monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
     monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
@@ -55,7 +55,7 @@ def test_default_cache_config_used_when_missing(monkeypatch):
 
 def test_default_stage_devices_from_sequence_parallel(monkeypatch):
     """Ensure devices list reflects sequence parallel size when no parallel_config is provided."""
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model: [])
+    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
     monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
     monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
     monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)

tests/entrypoints/test_omni_llm.py

@@ -445,7 +445,7 @@ def mock_get_config(monkeypatch):
 def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_config):
     """Test that stage configs are auto-loaded when stage_configs_path is None."""
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [
             _FakeStageConfig(fake_stage_config),
             _FakeStageConfig(fake_stage_config),
@@ -510,7 +510,7 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
 def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
     """Test that generate raises ValueError when sampling_params_list length doesn't match."""
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [_FakeStageConfig(fake_stage_config)]
 
     import sys
@@ -562,7 +562,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
     stage_cfg1 = dict(fake_stage_config)
     stage_cfg1["processed_input"] = ["processed-for-stage-1"]
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
 
     import sys
@@ -660,7 +660,7 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
     stage_cfg0["final_output"] = False
     stage_cfg1["final_output"] = False
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
 
     import sys
@@ -734,7 +734,7 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
     stage_cfg0["final_output"] = False
     stage_cfg1["final_output"] = False
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
 
     import sys
@@ -802,7 +802,7 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
 def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
     """Test that _wait_for_stages_ready handles timeout correctly."""
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [_FakeStageConfig(fake_stage_config)]
 
     import sys
@@ -858,7 +858,7 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
 def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
     """Test that generate handles error messages from stages correctly."""
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [_FakeStageConfig(fake_stage_config)]
 
     import sys
@@ -934,7 +934,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
 def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
     """Test that close() sends shutdown signal to all input queues."""
 
-    def _fake_loader(model: str):
+    def _fake_loader(model: str, base_engine_args=None):
         return [_FakeStageConfig(fake_stage_config)]
 
     import sys

tests/entrypoints/test_stage_utils.py

@@ -55,4 +55,3 @@ def test_set_stage_devices_respects_logical_ids(monkeypatch):
     set_stage_devices(stage_id=0, devices="0,1")
 
     assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
-    assert call_log and call_log[0] == 0  # current device set after remap

tests/utils.py

@@ -1,11 +1,24 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
+# Some functions are copied from vllm/tests/utils.py
+import functools
 import os
+import signal
+import subprocess
+import sys
+import tempfile
 import time
-from contextlib import contextmanager
+from collections.abc import Callable
+from contextlib import ExitStack, contextmanager, suppress
+from typing import Any, Literal
 
+import cloudpickle
+import pytest
+from typing_extensions import ParamSpec
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import cuda_device_count_stateless
+
+_P = ParamSpec("_P")
 
 if current_platform.is_rocm():
     from amdsmi import (
@@ -90,10 +103,16 @@ def wait_for_gpu_memory_to_clear(
         print("")
 
         if threshold_bytes is not None:
-            is_free = lambda used, total: used <= threshold_bytes / 2**30  # noqa E731
+
+            def is_free(used, total):
+                return used <= threshold_bytes / 2**30  # noqa E731
+
             threshold = f"{threshold_bytes / 2**30} GiB"
         else:
-            is_free = lambda used, total: used / total <= threshold_ratio  # noqa E731
+
+            def is_free(used, total):
+                return used / total <= threshold_ratio  # noqa E731
+
             threshold = f"{threshold_ratio:.2f}"
 
         dur_s = time.time() - start_time
@@ -105,3 +124,353 @@ def wait_for_gpu_memory_to_clear(
             raise ValueError(f"Memory of devices {devices=} not free after {dur_s=:.02f} ({threshold=})")
 
         time.sleep(5)
+
+
+def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, None]:
+    """Decorator to fork a new process for each test function.
+    See https://github.com/vllm-project/vllm/issues/7053 for more details.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
+        # Make the process the leader of its own process group
+        # to avoid sending SIGTERM to the parent process
+        os.setpgrp()
+        from _pytest.outcomes import Skipped
+
+        # Create a unique temporary file to store exception info from child
+        # process. Use test function name and process ID to avoid collisions.
+        with (
+            tempfile.NamedTemporaryFile(
+                delete=False, mode="w+b", prefix=f"vllm_test_{func.__name__}_{os.getpid()}_", suffix=".exc"
+            ) as exc_file,
+            ExitStack() as delete_after,
+        ):
+            exc_file_path = exc_file.name
+            delete_after.callback(os.remove, exc_file_path)
+
+            pid = os.fork()
+            print(f"Fork a new process to run a test {pid}")
+            if pid == 0:
+                # Parent process responsible for deleting, don't delete
+                # in child.
+                delete_after.pop_all()
+                try:
+                    func(*args, **kwargs)
+                except Skipped as e:
+                    # convert Skipped to exit code 0
+                    print(str(e))
+                    os._exit(0)
+                except Exception as e:
+                    import traceback
+
+                    tb_string = traceback.format_exc()
+
+                    # Try to serialize the exception object first
+                    exc_to_serialize: dict[str, Any]
+                    try:
+                        # First, try to pickle the actual exception with
+                        # its traceback.
+                        exc_to_serialize = {"pickled_exception": e}
+                        # Test if it can be pickled
+                        cloudpickle.dumps(exc_to_serialize)
+                    except (Exception, KeyboardInterrupt):
+                        # Fall back to string-based approach.
+                        exc_to_serialize = {
+                            "exception_type": type(e).__name__,
+                            "exception_msg": str(e),
+                            "traceback": tb_string,
+                        }
+                    try:
+                        with open(exc_file_path, "wb") as f:
+                            cloudpickle.dump(exc_to_serialize, f)
+                    except Exception:
+                        # Fallback: just print the traceback.
+                        print(tb_string)
+                    os._exit(1)
+                else:
+                    os._exit(0)
+            else:
+                pgid = os.getpgid(pid)
+                _pid, _exitcode = os.waitpid(pid, 0)
+                # ignore SIGTERM signal itself
+                old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN)
+                # kill all child processes
+                os.killpg(pgid, signal.SIGTERM)
+                # restore the signal handler
+                signal.signal(signal.SIGTERM, old_signal_handler)
+                if _exitcode != 0:
+                    # Try to read the exception from the child process
+                    exc_info = {}
+                    if os.path.exists(exc_file_path):
+                        with suppress(Exception), open(exc_file_path, "rb") as f:
+                            exc_info = cloudpickle.load(f)
+
+                    if (original_exception := exc_info.get("pickled_exception")) is not None:
+                        # Re-raise the actual exception object if it was
+                        # successfully pickled.
+                        assert isinstance(original_exception, Exception)
+                        raise original_exception
+
+                    if (original_tb := exc_info.get("traceback")) is not None:
+                        # Use string-based traceback for fallback case
+                        raise AssertionError(
+                            f"Test {func.__name__} failed when called with"
+                            f" args {args} and kwargs {kwargs}"
+                            f" (exit code: {_exitcode}):\n{original_tb}"
+                        ) from None
+
+                    # Fallback to the original generic error
+                    raise AssertionError(
+                        f"function {func.__name__} failed when called with"
+                        f" args {args} and kwargs {kwargs}"
+                        f" (exit code: {_exitcode})"
+                    ) from None
+
+    return wrapper
+
+
+def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]:
+    """Decorator to spawn a new process for each test function."""
+
+    @functools.wraps(f)
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
+        # Check if we're already in a subprocess
+        if os.environ.get("RUNNING_IN_SUBPROCESS") == "1":
+            # If we are, just run the function directly
+            return f(*args, **kwargs)
+
+        import torch.multiprocessing as mp
+
+        with suppress(RuntimeError):
+            mp.set_start_method("spawn")
+
+        # Get the module
+        module_name = f.__module__
+
+        # Create a process with environment variable set
+        env = os.environ.copy()
+        env["RUNNING_IN_SUBPROCESS"] = "1"
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            output_filepath = os.path.join(tempdir, "new_process.tmp")
+
+            # `cloudpickle` allows pickling complex functions directly
+            input_bytes = cloudpickle.dumps((f, output_filepath))
+
+            cmd = [sys.executable, "-m", f"{module_name}"]
+
+            returned = subprocess.run(cmd, input=input_bytes, capture_output=True, env=env)
+
+            # check if the subprocess is successful
+            try:
+                returned.check_returncode()
+            except Exception as e:
+                # wrap raised exception to provide more information
+                raise RuntimeError(f"Error raised in subprocess:\n{returned.stderr.decode()}") from e
+
+    return wrapper
+
+
+def create_new_process_for_each_test(
+    method: Literal["spawn", "fork"] | None = None,
+) -> Callable[[Callable[_P, None]], Callable[_P, None]]:
+    """Creates a decorator that runs each test function in a new process.
+
+    Args:
+        method: The process creation method. Can be either "spawn" or "fork".
+               If not specified, it defaults to "spawn" on ROCm and XPU
+               platforms and "fork" otherwise.
+
+    Returns:
+        A decorator to run test functions in separate processes.
+    """
+    if method is None:
+        # TODO: Spawn is not working correctly on ROCm
+        # The test content will not run and tests passed immediately.
+        # For now, using `fork` for ROCm as it can run with `fork`
+        # and tests are running correctly.
+        use_spawn = current_platform.is_xpu()
+        method = "spawn" if use_spawn else "fork"
+
+    assert method in ["spawn", "fork"], "Method must be either 'spawn' or 'fork'"
+
+    if method == "fork":
+        return fork_new_process_for_each_test
+
+    return spawn_new_process_for_each_test
+
+
+def cuda_marks(*, res: str, num_cards: int):
+    """
+    Get a collection of pytest marks to apply for `@cuda_test`.
+
+    Args:
+        res: Resource type, e.g., "L4" or "H100".
+        num_cards: Number of GPU cards required.
+
+    Returns:
+        List of pytest marks to apply.
+    """
+    test_platform_detail = pytest.mark.cuda
+
+    if res == "L4":
+        test_resource = pytest.mark.L4
+    elif res == "H100":
+        test_resource = pytest.mark.H100
+    else:
+        raise ValueError(f"Invalid CUDA resource type: {res}. Supported: L4, H100")
+
+    marks = [test_resource, test_platform_detail]
+
+    if num_cards == 1:
+        return marks
+    else:
+        test_distributed = pytest.mark.distributed_cuda(num_cards=num_cards)
+        test_skipif = pytest.mark.skipif_cuda(
+            cuda_device_count_stateless() < num_cards,
+            reason=f"Need at least {num_cards} CUDA GPUs to run the test.",
+        )
+        return marks + [test_distributed, test_skipif]
+
+
+def rocm_marks(*, res: str, num_cards: int):
+    """
+    Get a collection of pytest marks to apply for `@rocm_test`.
+
+    Args:
+        res: Resource type, e.g., "MI325".
+        num_cards: Number of GPU cards required.
+
+    Returns:
+        List of pytest marks to apply.
+    """
+    test_platform_detail = pytest.mark.rocm
+
+    if res == "MI325":
+        test_resource = pytest.mark.MI325
+    else:
+        raise ValueError(f"Invalid ROCm resource type: {res}. Supported: MI325")
+
+    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 ROCm support for `skipif_rocm` 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.
+        num_cards: Number of GPU cards required.
+
+    Returns:
+        List of pytest marks to apply.
+    """
+    test_platform = pytest.mark.gpu
+    if res in ("L4", "H100"):
+        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)
+    raise ValueError(f"Invalid resource type: {res}. Supported: L4, H100, MI325")
+
+
+def npu_marks(*, res: str, num_cards: int):
+    """Get a collection of pytest marks to apply for `@npu_test`."""
+    test_platform = pytest.mark.npu
+    if res == "A2":
+        test_resource = pytest.mark.A2
+    elif res == "A3":
+        test_resource = pytest.mark.A3
+    else:
+        # TODO: Currently we don't have various NPU card types defined
+        # Use None to skip resource-specific marking for unknown types
+        test_resource = None
+
+    if num_cards == 1:
+        return [mark for mark in [test_platform, test_resource] if mark is not None]
+    else:
+        # Multiple cards scenario needs distributed_npu mark
+        test_distributed = pytest.mark.distributed_npu(num_cards=num_cards)
+        # TODO: add NPU support for `skipif_npu` marker
+        return [mark for mark in [test_platform, test_resource, test_distributed] if mark is not None]
+
+
+def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
+    """
+    Decorate a test for multiple hardware platforms with a single call.
+    Automatically wraps the test with @create_new_process_for_each_test() for distributed tests.
+
+    Args:
+        res: Mapping from platform to resource type. Supported platforms/resources:
+            - cuda: L4, H100
+            - rocm: MI325
+            - npu: A2, A3
+        num_cards: Number of cards required. Can be:
+            - int: same card count for all platforms (default: 1)
+            - dict: per-platform card count, e.g., {"cuda": 2, "rocm": 2}
+
+    Example:
+        @hardware_test(
+            res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+            num_cards={"cuda": 2, "rocm": 2, "npu": 2},
+        )
+        def test_multi_platform():
+            ...
+    """
+    # Validate platforms
+    # Don't validate platform details in this decorator
+    for platform, _ in res.items():
+        if platform not in ("cuda", "rocm", "npu"):
+            raise ValueError(f"Unsupported platform: {platform}")
+
+    # Normalize num_cards
+    if isinstance(num_cards, int):
+        num_cards_dict = {platform: num_cards for platform in res.keys()}
+    else:
+        num_cards_dict = num_cards
+        for platform in num_cards_dict.keys():
+            if platform not in res:
+                raise ValueError(
+                    f"Platform '{platform}' in num_cards but not in res. Available platforms: {list(res.keys())}"
+                )
+        for platform in res.keys():
+            if platform not in num_cards_dict:
+                num_cards_dict[platform] = 1
+
+    # Collect marks from all platforms
+    all_marks: list[Callable[[Callable[_P, None]], Callable[_P, None]]] = []
+    for platform, resource in res.items():
+        cards = num_cards_dict[platform]
+        if platform == "cuda" or platform == "rocm":
+            marks = gpu_marks(res=resource, num_cards=cards)
+        elif platform == "npu":
+            marks = npu_marks(res=resource, num_cards=cards)
+        else:
+            raise ValueError(f"Unsupported platform: {platform}")
+        all_marks.extend(marks)
+
+    create_new_process_flag = False
+    for cards in num_cards_dict.values():
+        if cards > 1:
+            create_new_process_flag = True
+            break
+
+    def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
+        if create_new_process_flag:
+            # only for distributed tests
+            func = create_new_process_for_each_test()(f)
+        else:
+            func = f
+        for mark in reversed(all_marks):
+            func = mark(func)
+        return func
+
+    return wrapper

tests/worker/test_omni_gpu_model_runner.py

@@ -0,0 +1,123 @@
+from contextlib import contextmanager
+
+import torch
+
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+
+
+class DummyBuffer:
+    """A minimal buffer wrapper that exposes the `.gpu` attribute."""
+
+    def __init__(self, t: torch.Tensor):
+        self.gpu = t
+
+
+class DummyInputBatch:
+    """A minimal input batch that only provides `req_ids`."""
+
+    def __init__(self, req_ids):
+        self.req_ids = req_ids
+
+
+class DummyReqState:
+    """A minimal request state container."""
+
+    pass
+
+
+class DummyTalkerMTP(torch.nn.Module):
+    """A fake talker_mtp module for deterministic CPU testing."""
+
+    def forward(self, req_input_ids, req_embeds, last_talker_hidden, text_step):
+        # Deterministic behavior:
+        # - output embeds = input embeds + 1
+        # - output codes = [[0], [1], ...]
+        bsz = req_embeds.shape[0]
+        new_embeds = req_embeds + 1.0
+        codes = torch.arange(bsz, dtype=torch.int64).view(bsz, 1)
+        return new_embeds, codes
+
+
+@contextmanager
+def _noop_forward_context(*args, **kwargs):
+    """A no-op context manager to replace vLLM forward context in CPU tests."""
+    yield
+
+
+def _make_runner(req_ids=("r1", "r2"), hidden_size=4):
+    # Create an instance without calling OmniGPUModelRunner.__init__
+    runner = object.__new__(OmniGPUModelRunner)
+
+    # Minimal attributes used by OmniGPUModelRunner._talker_mtp_forward
+    runner.input_batch = DummyInputBatch(list(req_ids))
+    runner.requests = {rid: DummyReqState() for rid in req_ids}
+
+    # query_start_loc.cpu[req_index] is used to locate the token position
+    # in the flattened `inputs_embeds`.
+    runner.query_start_loc = type("QSL", (), {})()
+    # Map: r1 -> offset 0, r2 -> offset 3
+    runner.query_start_loc.cpu = torch.tensor([0, 3], dtype=torch.int32)
+
+    bsz = len(req_ids)
+    runner.talker_mtp_input_ids = DummyBuffer(torch.zeros((bsz,), dtype=torch.int64))
+    runner.talker_mtp_inputs_embeds = DummyBuffer(torch.zeros((bsz, hidden_size), dtype=torch.float32))
+    runner.last_talker_hidden = DummyBuffer(torch.zeros((bsz, hidden_size), dtype=torch.float32))
+    runner.text_step = DummyBuffer(torch.zeros((bsz, hidden_size), dtype=torch.float32))
+
+    runner.talker_mtp = DummyTalkerMTP()
+    runner.vllm_config = object()
+
+    # Provide a minimal implementation that returns the expected 4-tuple.
+    def _determine_batch_execution_and_padding(**kwargs):
+        return None, object(), None, None
+
+    runner._determine_batch_execution_and_padding = _determine_batch_execution_and_padding
+
+    # Use the real merge method from OmniGPUModelRunner.
+    return runner
+
+
+def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch):
+    # Patch the module-level `set_forward_context` symbol used inside
+    # OmniGPUModelRunner._talker_mtp_forward.
+    import vllm_omni.worker.gpu_model_runner as mod  # Must be the same module that defines OmniGPUModelRunner
+
+    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
+
+    runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
+
+    # Initialize per-request embeds (batch-major inside talker_mtp_inputs_embeds)
+    runner.talker_mtp_inputs_embeds.gpu[0] = torch.tensor([1.0, 2.0, 3.0, 4.0])
+    runner.talker_mtp_inputs_embeds.gpu[1] = torch.tensor([10.0, 20.0, 30.0, 40.0])
+
+    # Flattened `inputs_embeds`: offsets 0 and 3 will be overwritten
+    inputs_embeds = torch.zeros((6, 4), dtype=torch.float32)
+
+    # Call the original implementation from OmniGPUModelRunner (no re-implementation)
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
+
+    # Validate embeds were written back (+1)
+    assert torch.allclose(inputs_embeds[0], torch.tensor([2.0, 3.0, 4.0, 5.0]))
+    assert torch.allclose(inputs_embeds[3], torch.tensor([11.0, 21.0, 31.0, 41.0]))
+
+    # Validate per-request additional_information_cpu was updated
+    info_r1 = runner.requests["r1"].additional_information_cpu
+    info_r2 = runner.requests["r2"].additional_information_cpu
+    assert int(info_r1["code_predictor_codes"][0, 0]) == 0
+    assert int(info_r2["code_predictor_codes"][0, 0]) == 1
+
+
+def test_talker_mtp_forward_cpu_empty_batch_noop(monkeypatch):
+    import vllm_omni.worker.gpu_model_runner as mod
+
+    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
+
+    runner = _make_runner(req_ids=("r1",), hidden_size=4)
+
+    inputs_embeds = torch.randn((2, 4))
+    before = inputs_embeds.clone()
+
+    OmniGPUModelRunner._talker_mtp_forward(runner, [], inputs_embeds)
+
+    # Ensure no changes were made
+    assert torch.allclose(inputs_embeds, before)

tools/pre_commit/check_pickle_imports.py

@@ -18,6 +18,7 @@ import regex as re
 ALLOWED_FILES = {
     "vllm_omni/entrypoints/omni_llm.py",
     "tests/e2e/offline_inference/utils.py",
+    "tests/utils.py",
     "vllm_omni/diffusion/distributed/group_coordinator.py",
     "tests/diffusion/attention/test_sequence_parallel.py",
 }

vllm_omni/diffusion/attention/backends/abstract.py

@@ -48,6 +48,8 @@ class AttentionBackend(ABC):
 @dataclass
 class AttentionMetadata:
     attn_mask: torch.Tensor | None = None
+    joint_attn_mask: torch.Tensor | None = None
+    # a joint mask for the joint query, key, and value, depends the joint_strategy
     joint_query: torch.Tensor | None = None
     # a replicated tensor among processes appended to the front or rear of query, depends the joint_strategy
     joint_key: torch.Tensor | None = None

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -9,13 +9,14 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionImpl,
     AttentionMetadata,
 )
+from vllm_omni.diffusion.attention.backends.utils.fa import _pad_input, _unpad_input, _upad_input
 
 logger = init_logger(__name__)
 
 try:
     # only tested with flash_attn v3
     # from flash_attn_interface import flash_attn_func as flash_attn_3_func  # not available in flash-attn 2.8.1
-    from flash_attn import flash_attn_func  # can be FA2 or FA3
+    from flash_attn import flash_attn_func, flash_attn_varlen_func  # can be FA2 or FA3
 except ImportError:
     logger.warning(
         "FlashAttentionBackend is not available. You may install flash-attn "
@@ -63,12 +64,51 @@ class FlashAttentionImpl(AttentionImpl):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata = None,
     ) -> torch.Tensor:
-        # TODO: flash_attn_func does not support attn_mask.
-        out: torch.Tensor = flash_attn_func(
-            query,
-            key,
-            value,
-            causal=self.causal,
-            softmax_scale=self.softmax_scale,
-        )
+        """
+        Flash attention implementation.
+
+        Args:
+            query: (batch_size, seq_len, num_heads, head_dim)
+            key: (batch_size, seq_len, num_heads, head_dim)
+            value: (batch_size, seq_len, num_heads, head_dim)
+            attn_metadata: AttentionMetadata. Attention mask is supported as attn_metadata.attn_mask
+
+        Returns:
+            (batch_size, seq_len, num_heads, head_dim)
+        """
+        query_length = query.size(1)
+        attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
+        #  Contains at least one padding token in the sequence
+        if attention_mask is not None and torch.any(~attention_mask):
+            assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)"
+            q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input(
+                query, key, value, attention_mask, query_length, _unpad_input
+            )
+
+            out_unpad = flash_attn_varlen_func(
+                q,
+                k,
+                v,
+                cu_seqlens_q=cu_seq_lens_q,
+                cu_seqlens_k=cu_seq_lens_k,
+                max_seqlen_q=max_length_q,
+                max_seqlen_k=max_length_k,
+                **{
+                    "causal": self.causal,
+                    "softmax_scale": self.softmax_scale,
+                },
+            )
+            if isinstance(out_unpad, tuple):
+                out_unpad = out_unpad[0]
+
+            out = _pad_input(out_unpad, indices_q, query.size(0), query_length)
+
+        else:
+            out: torch.Tensor = flash_attn_func(
+                query,
+                key,
+                value,
+                causal=self.causal,
+                softmax_scale=self.softmax_scale,
+            )
         return out

vllm_omni/diffusion/attention/backends/ring_flash_attn.py

@@ -28,6 +28,20 @@ def ring_flash_attn_forward(
     joint_tensor_value=None,
     joint_strategy="front",
 ):
+    # Validate causal + joint_strategy combination
+    # When causal=True and joint_strategy="rear", the causal mask would incorrectly
+    # prevent local query tokens from attending to joint key tokens (which are
+    # concatenated at the end). This breaks the semantics where joint tokens
+    # (e.g., text conditioning) should be visible to all local tokens.
+    if causal and joint_tensor_key is not None and joint_strategy == "rear":
+        raise ValueError(
+            "joint_strategy='rear' is not compatible with causal=True in Ring Attention. "
+            "When using causal attention with joint tokens, use joint_strategy='front' "
+            "to ensure joint tokens act as a visible prefix for all local tokens. "
+            "With 'rear' strategy, the causal mask would incorrectly block local tokens "
+            "from seeing the joint tokens."
+        )
+
     comm = RingComm(process_group)
 
     out = None

vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py

@@ -62,6 +62,20 @@ class RingAttentionFunc(torch.autograd.Function):
         joint_tensor_value=None,
         joint_strategy="front",
     ):
+        # Validate causal + joint_strategy combination
+        # When causal=True and joint_strategy="rear", the causal mask would incorrectly
+        # prevent local query tokens from attending to joint key tokens (which are
+        # concatenated at the end). This breaks the semantics where joint tokens
+        # (e.g., text conditioning) should be visible to all local tokens.
+        if is_causal and joint_tensor_key is not None and joint_strategy == "rear":
+            raise ValueError(
+                "joint_strategy='rear' is not compatible with causal=True in Ring Attention. "
+                "When using causal attention with joint tokens, use joint_strategy='front' "
+                "to ensure joint tokens act as a visible prefix for all local tokens. "
+                "With 'rear' strategy, the causal mask would incorrectly block local tokens "
+                "from seeing the joint tokens."
+            )
+
         comm = RingComm(group)
         # Ensure tensors are contiguous for P2P communication
         q = q.contiguous()

vllm_omni/diffusion/attention/backends/utils/__init__.py

@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Utils for attention backends.
+"""
+
+from vllm_omni.diffusion.attention.backends.utils.fa import _pad_input, _unpad_input, _upad_input
+
+__all__ = [
+    "_pad_input",
+    "_unpad_input",
+    "_upad_input",
+]

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -0,0 +1,209 @@
+# Copyright 2025 The Fairseq Authors and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Adapted from https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_flash_attention_utils.py
+import torch
+import torch.nn.functional as F
+
+
+def _index_first_axis(tensor, indices):
+    """
+    A local implementation of the PyTorch indexing operation `tensor[indices]` on the first axis,
+    after flattening the first two dimensions of the tensor. This is functionally equivalent to
+    FA2's `index_first_axis` and replaces the need to import it.
+    """
+    # The input tensor is expected to be of shape (batch, seq_len, ...). We flatten the first
+    # two dimensions to get (total_tokens, ...) before indexing.
+    reshaped_tensor = tensor.reshape(-1, *tensor.shape[2:])
+    return reshaped_tensor[indices]
+
+
+def _unpad_input(hidden_states, attention_mask, unused_mask=None):
+    """
+    unpad_input function for flash attention variants that do not have them within their pkg themselves, e.g. fa3.
+
+    Arguments:
+        hidden_states: (batch, seqlen, ...)
+        attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
+        unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused.
+
+    Return:
+        hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask.
+        indices: (total_nnz), the indices of masked tokens from the flattened input sequence.
+        cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states.
+        max_seqlen_in_batch: int
+        seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask.
+    """
+    all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask
+    seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32)
+    used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+
+    return (
+        _index_first_axis(hidden_states, indices),
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+        used_seqlens_in_batch,
+    )
+
+
+def _pad_input(hidden_states, indices, batch, seqlen):
+    """
+    pad_input function for flash attention variants that do not have them within their pkg themselves, e.g. fa3.
+
+    Arguments:
+        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
+        indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence.
+        batch: int, batch size for the padded sequence.
+        seqlen: int, maximum sequence length for the padded sequence.
+
+    Return:
+        hidden_states: (batch, seqlen, ...)
+    """
+    dim = hidden_states.shape[1:]
+    output = torch.zeros((batch * seqlen), *dim, device=hidden_states.device, dtype=hidden_states.dtype)
+    output[indices] = hidden_states
+    return output.view(batch, seqlen, *dim)
+
+
+def _get_unpad_data(attention_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]:
+    """
+    Retrieves indexing data required to repad unpadded (ragged) tensors.
+
+    Arguments:
+        attention_mask (`torch.Tensor`):
+            Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid.
+
+    Return:
+        indices (`torch.Tensor`):
+            The indices of non-masked tokens from the flattened input sequence.
+        cu_seqlens (`torch.Tensor`):
+            The cumulative sequence lengths, used to index into ragged (unpadded) tensors.
+             `cu_seqlens` shape is (batch_size + 1,).
+        max_seqlen_in_batch (`int`):
+            Maximum sequence length in batch.
+    """
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    # NOTE: Similar to the `.item()` in prepare_fa2_from_position_ids, with torch compile,
+    # this might cause a graph break
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+    return (
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+def _upad_input(
+    query_layer: torch.Tensor,
+    key_layer: torch.Tensor,
+    value_layer: torch.Tensor,
+    attention_mask: torch.Tensor,
+    query_length: int,
+    unpad_input_func,
+):
+    """
+    Unpads query, key, and values tensors, using a single dimension for all tokens even though they belong
+    to different batches. This function is used instead of `flash_attn.bert_padding.unpad_input` in
+    order to avoid the recomputation of the same intermediary tensors for query, key, value tensors.
+
+    Arguments:
+        query_layer (`torch.Tensor`):
+            Query state with padding. Shape: (batch_size, query_length, num_heads, head_dim).
+        key_layer (`torch.Tensor`):
+            Key state with padding. Shape: (batch_size, kv_seq_len, num_key_value_heads, head_dim).
+        value_layer (`torch.Tensor`):
+            Value state with padding. Shape: (batch_size, kv_seq_len, num_key_value_heads, head_dim).
+        attention_mask (`torch.Tensor`):
+            Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid.
+        query_length (`int`):
+            Target length.
+        unpad_input_func:
+            The function to use for unpadding the input tensors.
+
+    Return:
+        query_layer (`torch.Tensor`):
+            Query state without padding. Shape: (total_target_length, num_heads, head_dim).
+        key_layer (`torch.Tensor`):
+            Key state with padding. Shape: (total_source_length, num_key_value_heads, head_dim).
+        value_layer (`torch.Tensor`):
+            Value state with padding. Shape: (total_source_length, num_key_value_heads, head_dim).
+        indices_q (`torch.Tensor`):
+            The indices of non-masked tokens from the flattened input target sequence.
+        (cu_seqlens_q, cu_seqlens_k) (`tuple[int]`):
+            The cumulative sequence lengths for the target (query) and source (key, value), used to index into
+             ragged (unpadded) tensors. `cu_seqlens` shape is (batch_size + 1,).
+        (max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`tuple[int]`):
+            Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence i.e. query,
+            `max_seqlen_in_batch_k` for the source sequence i.e. key/value).
+    """
+    if torch.compiler.is_compiling():
+        # allow PyTorch compiler to include operations that return scalar values (like .item()
+        torch._dynamo.config.capture_scalar_outputs = True
+    indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+
+    # With static caches, the k/v states may be larger than the mask ->
+    # we need to slice them to avoid generating garbage
+    # It's a bit of an anti-pattern, but otherwise we silently compute wrong attentions scores
+    if key_layer.shape[1] > (seq_len := attention_mask.shape[-1]):
+        key_layer, value_layer = key_layer[:, :seq_len, :, :], value_layer[:, :seq_len, :, :]
+
+    batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+
+    key_layer = _index_first_axis(key_layer, indices_k)
+    value_layer = _index_first_axis(value_layer, indices_k)
+    if query_length == kv_seq_len:
+        query_layer = _index_first_axis(query_layer, indices_k)
+        cu_seqlens_q = cu_seqlens_k
+        max_seqlen_in_batch_q = max_seqlen_in_batch_k
+        indices_q = indices_k
+    elif query_length == 1:
+        max_seqlen_in_batch_q = 1
+        cu_seqlens_q = torch.arange(
+            batch_size + 1, dtype=torch.int32, device=query_layer.device
+        )  # There is a memcpy here, that is very bad.
+        indices_q = cu_seqlens_q[:-1]
+        query_layer = query_layer.squeeze(1)
+    else:
+        # The -q_len: slice assumes left padding.
+        attention_mask = attention_mask[:, -query_length:]
+        query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q, *_ = unpad_input_func(query_layer, attention_mask)
+
+    return (
+        query_layer,
+        key_layer,
+        value_layer,
+        indices_q,
+        (cu_seqlens_q, cu_seqlens_k),
+        (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+    )
+
+
+def _is_packed_sequence(position_ids, batch_size):
+    """
+    Check the position ids whether packed sequences are indicated or not
+        1. Position ids exist
+        2. Flattened sequences only are supported
+        3. Compile-friendly `not (torch.diff(position_ids, dim=-1) >= 0).all()`, i.e.
+        we have multiple increasing sequences
+    """
+    if position_ids is None:
+        return False
+
+    increasing_position_sequences = torch.arange(position_ids.shape[1], device=position_ids.device) + position_ids.min()
+    return batch_size == 1 and (increasing_position_sequences - position_ids).abs().sum().bool()

vllm_omni/diffusion/attention/layer.py

@@ -15,8 +15,8 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy
 from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
 from vllm_omni.diffusion.attention.selector import get_attn_backend
-from vllm_omni.diffusion.data import get_current_omni_diffusion_config
 from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
+from vllm_omni.diffusion.forward_context import get_forward_context
 
 logger = init_logger(__name__)
 
@@ -58,7 +58,7 @@ class Attention(nn.Module):
         self.ring_runner = None
 
         try:
-            config = get_current_omni_diffusion_config()
+            config = get_forward_context().omni_diffusion_config
             self.backend_pref = config.attention_backend
             if config.parallel_config.ring_degree > 1:
                 self.use_ring = True

vllm_omni/diffusion/attention/parallel/factory.py

@@ -6,8 +6,8 @@ from __future__ import annotations
 from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention, ParallelAttentionStrategy
 from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
 from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
-from vllm_omni.diffusion.data import get_current_omni_diffusion_config
 from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size, get_sp_group
+from vllm_omni.diffusion.forward_context import get_forward_context
 
 
 def build_parallel_attention_strategy(
@@ -24,7 +24,7 @@ def build_parallel_attention_strategy(
       and initialized process groups.
     """
     try:
-        cfg = get_current_omni_diffusion_config()
+        cfg = get_forward_context().omni_diffusion_config
         p = cfg.parallel_config
     except Exception:
         return NoParallelAttention()

vllm_omni/diffusion/attention/parallel/ring.py

@@ -16,10 +16,10 @@ from vllm_omni.diffusion.attention.parallel.base import (
     ParallelAttentionContext,
     # ParallelAttentionStrategy, # Not used in type hint below currently
 )
+from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGroupCoordinator
 
 # from vllm_omni.diffusion.attention.backends.ring_selector import AttnType # Already imported above
-from vllm_omni.diffusion.data import get_current_omni_diffusion_config
-from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGroupCoordinator
+from vllm_omni.diffusion.forward_context import get_forward_context
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
@@ -110,7 +110,7 @@ class RingParallelAttention:
         backend_pref = self.attn_backend_pref
         if backend_pref is None:
             try:
-                config = get_current_omni_diffusion_config()
+                config = get_forward_context().omni_diffusion_config
                 # config might not have attention_backend attribute if not updated
                 backend_pref = getattr(config, "attention_backend", None)
             except Exception:

vllm_omni/diffusion/attention/parallel/ulysses.py

@@ -163,6 +163,36 @@ class UlyssesParallelAttention:
             joint_len=joint_len,
             joint_strategy=joint_strategy,
         )
+
+        if attn_metadata is not None:
+            if is_joint:
+                if attn_metadata.joint_attn_mask is None and attn_metadata.attn_mask is None:
+                    attn_metadata.attn_mask = None
+                else:
+                    if attn_metadata.attn_mask is None:
+                        attn_metadata.attn_mask = torch.ones(
+                            [query.shape[0], query.shape[1] - attn_metadata.joint_attn_mask.shape[1]],
+                            dtype=torch.bool,
+                            device=query.device,
+                        )
+                    elif attn_metadata.joint_attn_mask is None:
+                        attn_metadata.joint_attn_mask = torch.ones(
+                            [query.shape[0], query.shape[1] - attn_metadata.attn_mask.shape[1]],
+                            dtype=torch.bool,
+                            device=query.device,
+                        )
+                    attn_metadata.attn_mask = (
+                        torch.cat([attn_metadata.joint_attn_mask, attn_metadata.attn_mask], dim=1)
+                        if joint_strategy == "front"
+                        else torch.cat([attn_metadata.attn_mask, attn_metadata.joint_attn_mask], dim=1)
+                    )
+
+            if attn_metadata.attn_mask is not None:
+                # the final attn_mask is ready, the length should be aligedn with query length
+                assert attn_metadata.attn_mask.shape[1] == query.shape[1], (
+                    f"attn_mask length: {attn_metadata.attn_mask.shape[1]} != query length: {query.shape[1]}"
+                )
+                attn_metadata.attn_mask = attn_metadata.attn_mask.bool().contiguous()
         return query, key, value, attn_metadata, ctx
 
     def post_attention(self, attn_output: torch.Tensor, ctx: ParallelAttentionContext | None) -> torch.Tensor:

vllm_omni/diffusion/attention/selector.py

@@ -31,6 +31,8 @@ _BACKEND_CONFIG = {
     "ASCEND": {"module": "vllm_omni.diffusion.attention.backends.ascend_attn", "class": "AscendAttentionBackend"},
 }
 
+_BACKENDS_SUPPORT_ATTENTION_MASK = ["SDPA", "ASCEND", "FLASH_ATTN"]
+
 
 def load_backend(backend_name: str) -> type[AttentionBackend]:
     config = _BACKEND_CONFIG[backend_name]

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -7,11 +7,19 @@ This module provides a CacheDiTBackend class to enable cache-dit acceleration on
 pipelines in vllm-omni, supporting both single and dual-transformer architectures.
 """
 
+import functools
 from collections.abc import Callable
+from contextlib import ExitStack
 from typing import Any, Optional
 
 import cache_dit
+import torch
 from cache_dit import BlockAdapter, DBCacheConfig, ForwardPattern, ParamsModifier, TaylorSeerCalibratorConfig
+from cache_dit.caching.block_adapters import FakeDiffusionPipeline
+from cache_dit.caching.cache_adapters.cache_adapter import CachedAdapter
+from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
+from cache_dit.caching.cache_contexts import BasicCacheConfig
+from cache_dit.caching.cache_contexts.cache_manager import CachedContextManager
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.cache.base import CacheBackend
@@ -401,6 +409,384 @@ def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], No
     return refresh_cache_context
 
 
+class BagelCachedContextManager(CachedContextManager):
+    """
+    Custom CachedContextManager for Bagel that safely handles NaiveCache objects
+    (mapped to encoder_hidden_states) by skipping tensor operations on them.
+    """
+
+    @torch.compiler.disable
+    def apply_cache(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        prefix: str = "Bn",
+        encoder_prefix: str = "Bn_encoder",
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        # Allow Bn and Fn prefix to be used for residual cache.
+        if "Bn" in prefix:
+            hidden_states_prev = self.get_Bn_buffer(prefix)
+        else:
+            hidden_states_prev = self.get_Fn_buffer(prefix)
+
+        assert hidden_states_prev is not None, f"{prefix}_buffer must be set before"
+
+        if self.is_cache_residual():
+            hidden_states = hidden_states_prev + hidden_states
+        else:
+            # If cache is not residual, we use the hidden states directly
+            hidden_states = hidden_states_prev
+
+        hidden_states = hidden_states.contiguous()
+
+        if encoder_hidden_states is not None:
+            if "Bn" in encoder_prefix:
+                encoder_hidden_states_prev = self.get_Bn_encoder_buffer(encoder_prefix)
+            else:
+                encoder_hidden_states_prev = self.get_Fn_encoder_buffer(encoder_prefix)
+
+            if encoder_hidden_states_prev is not None:
+                if self.is_encoder_cache_residual():
+                    # FIX: Check if encoder_hidden_states is a tensor before adding
+                    if isinstance(encoder_hidden_states, torch.Tensor) and isinstance(
+                        encoder_hidden_states_prev, torch.Tensor
+                    ):
+                        encoder_hidden_states = encoder_hidden_states_prev + encoder_hidden_states
+                else:
+                    # If encoder cache is not residual, we use the encoder hidden states directly
+                    encoder_hidden_states = encoder_hidden_states_prev
+
+            # FIX: Check if encoder_hidden_states is a tensor before calling contiguous
+            if isinstance(encoder_hidden_states, torch.Tensor):
+                encoder_hidden_states = encoder_hidden_states.contiguous()
+
+        return hidden_states, encoder_hidden_states
+
+
+class BagelCachedBlocks(CachedBlocks_Pattern_0_1_2):
+    """
+    Custom CachedBlocks for Bagel that safely handles NaiveCache objects
+    by adding isinstance checks in call_Mn_blocks and compute_or_prune.
+    """
+
+    def call_Mn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        original_hidden_states = hidden_states
+        original_encoder_hidden_states = encoder_hidden_states
+        for block in self._Mn_blocks():
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+            hidden_states, encoder_hidden_states = self._process_block_outputs(hidden_states, encoder_hidden_states)
+
+        # compute hidden_states residual
+        hidden_states = hidden_states.contiguous()
+
+        hidden_states_residual = hidden_states - original_hidden_states
+
+        if (
+            encoder_hidden_states is not None
+            and original_encoder_hidden_states is not None
+            and isinstance(encoder_hidden_states, torch.Tensor)  # FIX: Added Check
+        ):
+            encoder_hidden_states = encoder_hidden_states.contiguous()
+            encoder_hidden_states_residual = encoder_hidden_states - original_encoder_hidden_states
+        else:
+            encoder_hidden_states_residual = None
+
+        return (
+            hidden_states,
+            encoder_hidden_states,
+            hidden_states_residual,
+            encoder_hidden_states_residual,
+        )
+
+    def compute_or_prune(
+        self,
+        block_id: int,  # Block index in the transformer blocks
+        # Below are the inputs to the block
+        block,  # The transformer block to be executed
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        # NOTE: Although Bagel likely won't use pruning, implementing safe version just in case.
+        # Copy-pasted from original but adding checks.
+
+        original_hidden_states = hidden_states
+        original_encoder_hidden_states = encoder_hidden_states
+
+        can_use_prune = self._maybe_prune(
+            block_id,
+            hidden_states,
+            prefix=f"{self.cache_prefix}_{block_id}_Fn_original",
+        )
+
+        torch._dynamo.graph_break()
+        if can_use_prune:
+            self.context_manager.add_pruned_step()
+            hidden_states, encoder_hidden_states = self.context_manager.apply_prune(
+                hidden_states,
+                encoder_hidden_states,
+                prefix=(
+                    f"{self.cache_prefix}_{block_id}_Bn_residual"
+                    if self.context_manager.is_cache_residual()
+                    else f"{self.cache_prefix}_Bn_hidden_states"
+                ),
+                encoder_prefix=(
+                    f"{self.cache_prefix}_{block_id}_Bn_encoder_residual"
+                    if self.context_manager.is_encoder_cache_residual()
+                    else f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states"
+                ),
+            )
+            torch._dynamo.graph_break()
+        else:
+            # Normal steps: Compute the block and cache the residuals.
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                *args,
+                **kwargs,
+            )
+            hidden_states, encoder_hidden_states = self._process_block_outputs(hidden_states, encoder_hidden_states)
+            if not self._skip_prune(block_id):
+                hidden_states = hidden_states.contiguous()
+                hidden_states_residual = hidden_states - original_hidden_states
+
+                if (
+                    encoder_hidden_states is not None
+                    and original_encoder_hidden_states is not None
+                    and isinstance(encoder_hidden_states, torch.Tensor)  # FIX: Added Check
+                ):
+                    encoder_hidden_states = encoder_hidden_states.contiguous()
+                    encoder_hidden_states_residual = encoder_hidden_states - original_encoder_hidden_states
+                else:
+                    encoder_hidden_states_residual = None
+
+                self.context_manager.set_Fn_buffer(
+                    original_hidden_states,
+                    prefix=f"{self.cache_prefix}_{block_id}_Fn_original",
+                )
+                if self.context_manager.is_cache_residual():
+                    self.context_manager.set_Bn_buffer(
+                        hidden_states_residual,
+                        prefix=f"{self.cache_prefix}_{block_id}_Bn_residual",
+                    )
+                else:
+                    self.context_manager.set_Bn_buffer(
+                        hidden_states,
+                        prefix=f"{self.cache_prefix}_{block_id}_Bn_hidden_states",
+                    )
+                if encoder_hidden_states_residual is not None:
+                    if self.context_manager.is_encoder_cache_residual():
+                        self.context_manager.set_Bn_encoder_buffer(
+                            encoder_hidden_states_residual,
+                            prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_residual",
+                        )
+                    else:
+                        self.context_manager.set_Bn_encoder_buffer(
+                            encoder_hidden_states_residual,
+                            prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states",
+                        )
+            torch._dynamo.graph_break()
+
+        return hidden_states, encoder_hidden_states
+
+
+class BagelCachedAdapter(CachedAdapter):
+    """
+    Custom CachedAdapter for Bagel that uses BagelCachedContextManager and BagelCachedBlocks.
+    """
+
+    @classmethod
+    def create_context(
+        cls,
+        block_adapter: BlockAdapter,
+        **context_kwargs,
+    ) -> tuple[list[str], list[dict[str, Any]]]:
+        # Override to use BagelCachedContextManager
+
+        BlockAdapter.assert_normalized(block_adapter)
+
+        if BlockAdapter.is_cached(block_adapter.pipe):
+            return block_adapter.pipe
+
+        # Check context_kwargs
+        context_kwargs = cls.check_context_kwargs(block_adapter, **context_kwargs)
+
+        # Each Pipeline should have it's own context manager instance.
+        cache_config: BasicCacheConfig = context_kwargs.get("cache_config", None)
+        assert cache_config is not None, "cache_config can not be None."
+
+        # Apply cache on pipeline: wrap cache context
+        pipe_cls_name = block_adapter.pipe.__class__.__name__
+
+        # USE CUSTOM CONTEXT MANAGER
+        context_manager = BagelCachedContextManager(
+            name=f"{pipe_cls_name}_{hash(id(block_adapter.pipe))}",
+            persistent_context=isinstance(block_adapter.pipe, FakeDiffusionPipeline),
+        )
+
+        flatten_contexts, contexts_kwargs = cls.modify_context_params(block_adapter, **context_kwargs)
+
+        block_adapter.pipe._context_manager = context_manager  # instance level
+
+        if not context_manager.persistent_context:
+            original_call = block_adapter.pipe.__class__.__call__
+
+            @functools.wraps(original_call)
+            def new_call(self, *args, **kwargs):
+                with ExitStack() as stack:
+                    # cache context will be reset for each pipe inference
+                    for context_name, context_kwargs in zip(flatten_contexts, contexts_kwargs):
+                        stack.enter_context(
+                            context_manager.enter_context(
+                                context_manager.reset_context(
+                                    context_name,
+                                    **context_kwargs,
+                                ),
+                            )
+                        )
+                    outputs = original_call(self, *args, **kwargs)
+                    cls.apply_stats_hooks(block_adapter)
+                    return outputs
+
+            block_adapter.pipe.__class__.__call__ = new_call
+            block_adapter.pipe.__class__._original_call = original_call
+
+        else:
+            # Init persistent cache context for transformer
+            for context_name, context_kwargs in zip(flatten_contexts, contexts_kwargs):
+                context_manager.reset_context(
+                    context_name,
+                    **context_kwargs,
+                )
+
+        block_adapter.pipe.__class__._is_cached = True
+
+        cls.apply_params_hooks(block_adapter, contexts_kwargs)
+
+        return flatten_contexts, contexts_kwargs
+
+    @classmethod
+    def collect_unified_blocks(
+        cls,
+        block_adapter: BlockAdapter,
+        contexts_kwargs: list[dict],
+    ) -> list[dict[str, torch.nn.ModuleList]]:
+        # Override to use BagelCachedBlocks
+
+        BlockAdapter.assert_normalized(block_adapter)
+
+        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
+        assert hasattr(block_adapter.pipe, "_context_manager")
+        # Skipping isinstance check for ContextManager._supported_managers to avoid import issues
+
+        for i in range(len(block_adapter.transformer)):
+            unified_blocks_bind_context = {}
+            for j in range(len(block_adapter.blocks[i])):
+                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
+
+                # Directly instantiate BagelCachedBlocks
+                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
+                    [
+                        BagelCachedBlocks(
+                            # 0. Transformer blocks configuration
+                            block_adapter.blocks[i][j],
+                            transformer=block_adapter.transformer[i],
+                            forward_pattern=block_adapter.forward_pattern[i][j],
+                            check_forward_pattern=block_adapter.check_forward_pattern,
+                            check_num_outputs=block_adapter.check_num_outputs,
+                            # 1. Cache/Prune context configuration
+                            cache_prefix=block_adapter.blocks_name[i][j],
+                            cache_context=block_adapter.unique_blocks_name[i][j],
+                            context_manager=block_adapter.pipe._context_manager,
+                            cache_type=cache_config.cache_type,
+                        )
+                    ]
+                )
+
+            total_cached_blocks.append(unified_blocks_bind_context)
+
+        return total_cached_blocks
+
+
+def enable_cache_for_bagel(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for Bagel model (via OmniDiffusion pipeline).
+
+    Args:
+        pipeline: The OmniDiffusion pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+
+    Returns:
+        A refresh function that can be called to update cache context with new num_inference_steps.
+    """
+    # Build DBCacheConfig
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    # Build calibrator config if TaylorSeer is enabled
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    # Access the transformer: BagelPipeline -> Qwen2MoTForCausalLM -> Qwen2MoTModel
+    # BagelPipeline has self.language_model which is Qwen2MoTForCausalLM
+    # Qwen2MoTForCausalLM has self.model which is Qwen2MoTModel
+    transformer = pipeline.language_model.model
+
+    logger.info(
+        f"Enabling cache-dit on Bagel transformer: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    # Enable cache-dit on the transformer
+    # Pattern_0 corresponds to (hidden_states, encoder_hidden_states) input, output
+    # Custom adapter for Bagel to handle NaiveCache correctly
+    # from vllm_omni.diffusion.cache.bagel_cache_adapter import BagelCachedAdapter # No longer needed
+    BagelCachedAdapter.apply(
+        BlockAdapter(
+            transformer=transformer,
+            blocks=transformer.layers,
+            forward_pattern=ForwardPattern.Pattern_0,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        transformer = pipeline.language_model.model
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
 # Register custom cache-dit enablers after function definitions
 CUSTOM_DIT_ENABLERS.update(
     {
@@ -409,6 +795,7 @@ CUSTOM_DIT_ENABLERS.update(
         "LongCatImagePipeline": enable_cache_for_longcat_image,
         "LongCatImageEditPipeline": enable_cache_for_longcat_image,
         "StableDiffusion3Pipeline": enable_cache_for_sd3,
+        "BagelPipeline": enable_cache_for_bagel,
     }
 )
 

vllm_omni/diffusion/cache/teacache/hook.py

@@ -20,6 +20,10 @@ import torch
 from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
 from vllm_omni.diffusion.cache.teacache.extractors import get_extractor
 from vllm_omni.diffusion.cache.teacache.state import TeaCacheState
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_classifier_free_guidance_rank,
+    get_classifier_free_guidance_world_size,
+)
 from vllm_omni.diffusion.hooks import HookRegistry, ModelHook, StateManager
 
 
@@ -34,6 +38,7 @@ class TeaCacheHook(ModelHook):
     Key features:
     - Zero changes to model code
     - CFG-aware with separate states for positive/negative branches
+    - CFG-parallel compatible: properly detects branch identity across ranks
     - Model-specific polynomial rescaling
     - Auto-detection of model types
 
@@ -113,8 +118,18 @@ class TeaCacheHook(ModelHook):
         # GENERIC CACHING LOGIC (works for all models)
         # ============================================================================
         # Set context based on CFG branch for separate state tracking
-        if module.do_true_cfg and self._forward_cnt % 2 == 1:
-            cache_branch = "negative"
+        # With CFG-parallel, each rank processes only one branch:
+        #   - cfg_rank 0: positive branch
+        #   - cfg_rank > 0: negative branch
+        # Without CFG-parallel, branches alternate within a single rank
+        if module.do_true_cfg:
+            cfg_parallel_size = get_classifier_free_guidance_world_size()
+            if cfg_parallel_size > 1:
+                cfg_rank = get_classifier_free_guidance_rank()
+                cache_branch = "negative" if cfg_rank > 0 else "positive"
+            else:
+                # No CFG-parallel: use forward counter to alternate branches
+                cache_branch = "negative" if self._forward_cnt % 2 == 1 else "positive"
         else:
             cache_branch = "positive"
 

vllm_omni/diffusion/compile.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Any
+
+import torch.nn as nn
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def regionally_compile(model: nn.Module, *compile_args: Any, **compile_kwargs: Any) -> nn.Module:
+    """
+    Apply regional compilation to a PyTorch model.
+
+    Args:
+        model: The PyTorch model instance to compile
+        *compile_args: Positional arguments forwarded to torch.compile
+        **compile_kwargs: Keyword arguments forwarded to torch.compile
+
+    Returns:
+        The same model instance (modified in-place)
+    """
+    # Get the list of repeated blocks from the model
+    repeated_blocks = getattr(model, "_repeated_blocks", None)
+
+    if not repeated_blocks:
+        logger.warning("Regional compilation skipped because the model does not define `_repeated_blocks`.")
+        return model
+
+    # Check if we have modules with the specified class names
+    has_compiled_region = False
+    for submod in model.modules():
+        if submod.__class__.__name__ in repeated_blocks:
+            # Compile this submodule
+            submod.compile(*compile_args, **compile_kwargs)
+            has_compiled_region = True
+
+    if not has_compiled_region:
+        logger.warning(f"Regional compilation skipped because {repeated_blocks} classes are not found in the model.")
+
+    return model