Voxtral TTS: drop hardcoded CUDA in audio tokenizer; add XPU stage config (#2428)

Signed-off-by: Joshna Medisetty <joshna.medisetty@intel.com>
Signed-off-by: Joshna-Medisetty <joshna.medisetty@intel.com>

.buildkite/bootstrap-amd-omni.sh

@@ -90,10 +90,18 @@ upload_pipeline() {
     FAIL_FAST=$(fail_fast)
 
     cd .buildkite
+
+    # Select test definition file: merge suite for main, ready suite for PRs
+    if [[ $BUILDKITE_BRANCH == "main" ]]; then
+        TEST_YAML="test-amd-merge.yml"
+    else
+        TEST_YAML="test-amd-ready.yaml"
+    fi
+
     (
         set -x
         # Output pipeline.yaml with all blank lines removed
-        minijinja-cli test-template.j2 test-amd.yaml \
+        minijinja-cli test-template.j2 "$TEST_YAML" \
             -D branch="$BUILDKITE_BRANCH" \
             -D list_file_diff="$LIST_FILE_DIFF" \
             -D run_all="$RUN_ALL" \

.buildkite/bootstrap-intel-omni.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+# vllm-omni Intel bootstrap
+# Uses static pipeline-intel.yaml for Intel XPU tests
+# Based on: bootstrap-amd-omni.sh
+
+set -euo pipefail
+
+if [[ -z "${DOCS_ONLY_DISABLE:-}" ]]; then
+    DOCS_ONLY_DISABLE=0
+fi
+
+upload_pipeline() {
+    echo "--- 🛠 Preparing Intel pipeline"
+    buildkite-agent pipeline upload .buildkite/pipeline-intel.yaml
+}
+
+get_diff() {
+    $(git add .)
+    echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
+}
+
+get_diff_main() {
+    $(git add .)
+    echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
+}
+
+file_diff=$(get_diff)
+if [[ $BUILDKITE_BRANCH == "main" ]]; then
+    file_diff=$(get_diff_main)
+fi
+
+# ----------------------------------------------------------------------
+# Early exit: skip pipeline if all changed files are under docs/
+# ----------------------------------------------------------------------
+if [[ "${DOCS_ONLY_DISABLE}" != "1" ]] && [[ -n "${file_diff:-}" ]]; then
+    docs_only=1
+    while IFS= read -r f; do
+        [[ -z "$f" ]] && continue
+        if [[ "$f" != docs/* ]]; then
+            docs_only=0
+            break
+        fi
+    done < <(printf '%s\n' "$file_diff" | tr ' ' '\n' | tr -d '\r')
+
+    if [[ "$docs_only" -eq 1 ]]; then
+        buildkite-agent annotate ":memo: CI skipped — docs only" --style "info"
+        exit 0
+    fi
+fi
+
+upload_pipeline

.buildkite/nightly-release-pipeline.yaml

@@ -0,0 +1,20 @@
+steps:
+  - label: "Build and upload wheel"
+    key: "build-wheel"
+    agents:
+      queue: cpu_queue_release
+    commands:
+      - "curl -LsSf https://astral.sh/uv/install.sh | sh"
+      - 'export PATH="$HOME/.local/bin:$PATH"'
+      - "uv venv --python=3.12 && source .venv/bin/activate"
+      - "uv pip install --upgrade build"
+      - "python3 -m build"
+      - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+
+  - label: "Generate and upload wheel indices"
+    depends_on: "build-wheel"
+    allow_dependency_failure: true
+    agents:
+      queue: small_cpu_queue_release
+    commands:
+      - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"

.buildkite/pipeline-intel.yaml

@@ -0,0 +1,16 @@
+# .buildkite/pipeline-intel.yaml
+steps:
+  - group: "INTEL Tests"
+    steps:
+      - label: "Intel: XPU Test"
+        command: "bash .buildkite/scripts/hardware_ci/run-xpu-test.sh"
+        agents:
+          queue: "intel-gpu-omni"
+        env:
+          DOCKER_BUILDKIT: "1"
+          # Buildkite will automatically replace this with the actual commit hash
+          VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
+          VLLM_VERSION: "v0.19.0"
+        priority: 100
+        timeout_in_minutes: 60
+        soft_fail: true

.buildkite/pipeline.yml

@@ -13,7 +13,7 @@ steps:
   - label: "Upload Ready Pipeline"
     depends_on: image-build
     key: upload-ready-pipeline
-    if: build.branch != "main"
+    if: build.branch != "main" && build.pull_request.labels includes "ready"
     commands:
       - buildkite-agent pipeline upload .buildkite/test-ready.yml
     agents:
@@ -23,17 +23,17 @@ steps:
   - label: "Upload Merge Pipeline"
     depends_on: image-build
     key: upload-merge-pipeline
-    if: build.branch == "main"
+    if: build.branch == "main" && build.env("NIGHTLY") != "1"
     commands:
       - buildkite-agent pipeline upload .buildkite/test-merge.yml
     agents:
       queue: "cpu_queue_premerge"
 
-  # L4 Test
+  # L4 Test — main+NIGHTLY=1 (scheduled), or PR with label nightly-test (e.g. add label then Rebuild)
   - label: "Upload Nightly Pipeline"
     depends_on: image-build
     key: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1"
+    if: '(build.branch == "main" && build.env("NIGHTLY") == "1") || (build.branch != "main" && build.pull_request.labels includes "nightly-test")'
     commands:
       - buildkite-agent pipeline upload .buildkite/test-nightly.yml
     agents:

.buildkite/scripts/generate-and-upload-nightly-index.sh

@@ -0,0 +1,87 @@
+#!/usr/bin/env bash
+
+set -ex
+
+# Generate and upload wheel indices for all vllm-omni wheels in the commit directory.
+# This script should run once after all wheels have been built and uploaded.
+# All paths are under the omni/ prefix in the vllm-wheels S3 bucket.
+
+# ======== setup ========
+
+BUCKET="vllm-wheels"
+INDICES_OUTPUT_DIR="indices"
+PYTHON="${PYTHON_PROG:-python3}"
+SUBPATH="omni/$BUILDKITE_COMMIT"
+S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
+
+# detect if python3.12+ is available
+has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
+if [[ "$has_new_python" -eq 0 ]]; then
+    # use new python from docker
+    docker pull python:3-slim
+    PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3"
+fi
+
+echo "Using python interpreter: $PYTHON"
+echo "Python version: $($PYTHON --version)"
+
+# ======== generate and upload indices ========
+
+# list all wheels in the commit directory
+echo "Existing wheels on S3:"
+aws s3 ls "$S3_COMMIT_PREFIX" || echo "(no objects found)"
+obj_json="objects.json"
+aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$SUBPATH/" --delimiter / --output json > "$obj_json"
+mkdir -p "$INDICES_OUTPUT_DIR"
+
+# HACK: we do not need regex module here, but it is required by pre-commit hook
+# To avoid any external dependency, we simply replace it back to the stdlib re module
+sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
+
+# Generate indices -- the version is just the commit hash (not omni/{commit})
+# because relative paths are computed between the index and wheel directories,
+# both of which live under the omni/ prefix in S3.
+$PYTHON .buildkite/scripts/generate-nightly-index.py \
+    --version "$BUILDKITE_COMMIT" \
+    --current-objects "$obj_json" \
+    --output-dir "$INDICES_OUTPUT_DIR" \
+    --comment "commit $BUILDKITE_COMMIT"
+
+# copy indices to /omni/{commit}/ unconditionally
+echo "Uploading indices to $S3_COMMIT_PREFIX"
+aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
+
+# copy to /omni/nightly/ when NIGHTLY=1
+if [[ "${NIGHTLY:-}" == "1" ]]; then
+    echo "Uploading indices to overwrite /omni/nightly/"
+    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/nightly/"
+fi
+
+# detect version from any wheel in the commit directory
+first_wheel_key=$($PYTHON -c "import json; obj=json.load(open('$obj_json')); print(next((c['Key'] for c in obj.get('Contents', []) if c['Key'].endswith('.whl')), ''))")
+if [[ -z "$first_wheel_key" ]]; then
+    echo "Error: No wheels found in $S3_COMMIT_PREFIX"
+    exit 1
+fi
+first_wheel=$(basename "$first_wheel_key")
+aws s3 cp "s3://$BUCKET/${first_wheel_key}" "/tmp/${first_wheel}"
+version=$(unzip -p "/tmp/${first_wheel}" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
+rm -f "/tmp/${first_wheel}"
+echo "Version in wheel: $version"
+pure_version="${version%%+*}"
+echo "Pure version (without variant): $pure_version"
+
+# re-generate and copy to /omni/{version}/ only if it does not have "dev" in the version
+if [[ "$version" != *"dev"* ]]; then
+    echo "Re-generating indices for /omni/$pure_version/"
+    rm -rf "${INDICES_OUTPUT_DIR:?}"
+    mkdir -p "$INDICES_OUTPUT_DIR"
+    # wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
+    $PYTHON .buildkite/scripts/generate-nightly-index.py \
+        --version "$pure_version" \
+        --wheel-dir "$BUILDKITE_COMMIT" \
+        --current-objects "$obj_json" \
+        --output-dir "$INDICES_OUTPUT_DIR" \
+        --comment "version $pure_version"
+    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/$pure_version/"
+fi

.buildkite/scripts/generate-nightly-index.py

@@ -0,0 +1,193 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import json
+import sys
+from dataclasses import asdict, dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+from urllib.parse import quote
+
+import regex as re
+
+
+def normalize_package_name(name: str) -> str:
+    """Normalize package name per PEP 503."""
+    return re.sub(r"[-_.]+", "-", name).lower()
+
+
+if not sys.version_info >= (3, 12):
+    raise RuntimeError("This script requires Python 3.12 or higher.")
+
+INDEX_HTML_TEMPLATE = """<!DOCTYPE html>
+<html>
+  <!-- {comment} -->
+  <meta name="pypi:repository-version" content="1.0">
+  <body>
+{items}
+  </body>
+</html>
+"""
+
+
+@dataclass
+class WheelFileInfo:
+    package_name: str
+    version: str
+    build_tag: str | None
+    python_tag: str
+    abi_tag: str
+    platform_tag: str
+    filename: str
+
+
+def parse_from_filename(file: str) -> WheelFileInfo:
+    """
+    Parse wheel filename per PEP 427:
+        {package_name}-{version}(-{build_tag})?-{python_tag}-{abi_tag}-{platform_tag}.whl
+    """
+    wheel_file_re = re.compile(
+        r"^(?P<package_name>.+)-(?P<version>[^-]+?)(-(?P<build_tag>[^-]+))?-(?P<python_tag>[^-]+)-(?P<abi_tag>[^-]+)-(?P<platform_tag>[^-]+)\.whl$"
+    )
+    match = wheel_file_re.match(file)
+    if not match:
+        raise ValueError(f"Invalid wheel file name: {file}")
+
+    return WheelFileInfo(
+        package_name=match.group("package_name"),
+        version=match.group("version"),
+        build_tag=match.group("build_tag"),
+        python_tag=match.group("python_tag"),
+        abi_tag=match.group("abi_tag"),
+        platform_tag=match.group("platform_tag"),
+        filename=file,
+    )
+
+
+def generate_project_list(package_names: list[str], comment: str = "") -> str:
+    """Generate top-level PEP 503 project list HTML."""
+    href_tags = []
+    for name in sorted(package_names):
+        href_tags.append(f'    <a href="{name}/">{name}/</a><br/>')
+    return INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
+
+
+def generate_package_index(
+    wheel_files: list[WheelFileInfo],
+    wheel_base_dir: Path,
+    index_base_dir: Path,
+    comment: str = "",
+) -> tuple[str, str]:
+    """Generate package index HTML and metadata JSON linking to wheel files."""
+    href_tags = []
+    metadata = []
+    for file in sorted(wheel_files, key=lambda x: x.filename):
+        relative_path = wheel_base_dir.relative_to(index_base_dir, walk_up=True) / file.filename
+        # handle '+' in URL; avoid double-encoding '/' and '%2B' (AWS S3 behavior)
+        file_path_quoted = quote(relative_path.as_posix(), safe=":%/")
+        href_tags.append(f'    <a href="{file_path_quoted}">{file.filename}</a><br/>')
+        file_meta = asdict(file)
+        file_meta["path"] = file_path_quoted
+        metadata.append(file_meta)
+    index_str = INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
+    metadata_str = json.dumps(metadata, indent=2)
+    return index_str, metadata_str
+
+
+def generate_index(
+    whl_files: list[str],
+    wheel_base_dir: Path,
+    index_base_dir: Path,
+    comment: str = "",
+):
+    """
+    Generate PEP 503 index for all wheel files.
+
+    Output structure:
+        index_base_dir/
+            index.html          # project list linking to vllm-omni/
+            vllm-omni/
+                index.html      # package index linking to wheel files
+                metadata.json   # machine-readable metadata
+    """
+    parsed_files = [parse_from_filename(f) for f in whl_files]
+
+    if not parsed_files:
+        print("No wheel files found, skipping index generation.")
+        return
+
+    comment_str = f" ({comment})" if comment else ""
+    comment_tmpl = f"Generated on {datetime.now().isoformat()}{comment_str}"
+
+    # Group by normalized package name
+    packages: dict[str, list[WheelFileInfo]] = {}
+    for file in parsed_files:
+        name = normalize_package_name(file.package_name)
+        packages.setdefault(name, []).append(file)
+
+    print(f"Found packages: {list(packages.keys())}")
+
+    # Generate per-package index
+    for package, files in packages.items():
+        package_dir = index_base_dir / package
+        package_dir.mkdir(parents=True, exist_ok=True)
+        index_str, metadata_str = generate_package_index(files, wheel_base_dir, package_dir, comment)
+        with open(package_dir / "index.html", "w") as f:
+            f.write(index_str)
+        with open(package_dir / "metadata.json", "w") as f:
+            f.write(metadata_str)
+
+    # Generate top-level project list
+    project_list_str = generate_project_list(sorted(packages.keys()), comment_tmpl)
+    with open(index_base_dir / "index.html", "w") as f:
+        f.write(project_list_str)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Generate PEP 503 wheel index from S3 object listing.")
+    parser.add_argument("--version", type=str, required=True, help="Version string (e.g., commit hash)")
+    parser.add_argument("--current-objects", type=str, required=True, help="Path to JSON from S3 list-objects-v2")
+    parser.add_argument("--output-dir", type=str, required=True, help="Directory to write index files")
+    parser.add_argument("--wheel-dir", type=str, default=None, help="Wheel directory (defaults to --version)")
+    parser.add_argument("--comment", type=str, default="", help="Comment for generated HTML")
+
+    args = parser.parse_args()
+
+    version = args.version
+    if "\\" in version or "/" in version:
+        raise ValueError("Version string must not contain slashes or backslashes.")
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    with open(args.current_objects) as f:
+        current_objects: dict[str, list[dict[str, Any]]] = json.load(f)
+
+    wheel_files = [
+        item["Key"].split("/")[-1] for item in current_objects.get("Contents", []) if item["Key"].endswith(".whl")
+    ]
+
+    print(f"Found {len(wheel_files)} wheel files for version {version}: {wheel_files}")
+
+    # For release versions, filter to only matching non-dev wheels
+    PY_VERSION_RE = re.compile(r"^\d+\.\d+\.\d+([a-zA-Z0-9.+-]*)?$")
+    if PY_VERSION_RE.match(version):
+        wheel_files = [f for f in wheel_files if version in f and "dev" not in f]
+        print(f"Non-nightly version detected, wheel files used: {wheel_files}")
+    else:
+        print("Nightly version detected, keeping all wheel files.")
+
+    wheel_dir = (args.wheel_dir or version).strip().rstrip("/")
+    wheel_base_dir = Path(output_dir).parent / wheel_dir
+    index_base_dir = Path(output_dir)
+
+    generate_index(
+        whl_files=wheel_files,
+        wheel_base_dir=wheel_base_dir,
+        index_base_dir=index_base_dir,
+        comment=args.comment.strip(),
+    )
+    print(f"Successfully generated index in {output_dir}")

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

@@ -64,15 +64,22 @@ while true; do
 done
 
 echo "--- Pulling container"
-image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
+## Temporary change to use AMD Docker Hub to store the vllm-ci image
+# to bypass the rate limit issue with ECR Public Gallery.
+# TODO: @tjtanaa point back to ECR Public Gallery
+# once the amd agents are configured to use ECR Public Gallery.
+# image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
+image_name="rocm/vllm-ci:${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
+# TODO: @tjtanaa uncomment this once the amd agents are configured to use ECR Public Gallery.
+# # 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 AMD Docker Hub
 
 docker pull "${image_name}"
 
@@ -87,7 +94,11 @@ HF_CACHE="$(realpath ~)/huggingface"
 mkdir -p "${HF_CACHE}"
 HF_MOUNT="/root/.cache/huggingface"
 
-commands=$@
+if [[ -n "${TEST_COMMAND:-}" ]]; then
+    commands="$TEST_COMMAND"
+else
+    commands="$@"
+fi
 echo "Commands:$commands"
 
 PARALLEL_JOB_COUNT=8
@@ -102,6 +113,7 @@ if [[ -z "$render_gid" ]]; then
 fi
 
 # check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs.
+# TODO: @tjtanaa reenable to run VLLM_ROCM_USE_AITER=1 when AITER is shipped with prebuilt kernels.
 if [[ $commands == *"--shard-id="* ]]; then
   # assign job count as the number of shards used
   commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g')
@@ -118,11 +130,9 @@ if [[ $commands == *"--shard-id="* ]]; then
         --rm \
         -e MIOPEN_DEBUG_CONV_DIRECT=0 \
         -e MIOPEN_DEBUG_CONV_GEMM=0 \
-        -e VLLM_ROCM_USE_AITER=1 \
+        -e VLLM_ROCM_USE_AITER=0 \
         -e HIP_VISIBLE_DEVICES="${GPU}" \
         -e HF_TOKEN \
-        -e AWS_ACCESS_KEY_ID \
-        -e AWS_SECRET_ACCESS_KEY \
         -v "${HF_CACHE}:${HF_MOUNT}" \
         -e "HF_HOME=${HF_MOUNT}" \
         -e "PYTHONPATH=${MYPYTHONPATH}" \
@@ -153,10 +163,8 @@ else
           --rm \
           -e MIOPEN_DEBUG_CONV_DIRECT=0 \
           -e MIOPEN_DEBUG_CONV_GEMM=0 \
-          -e VLLM_ROCM_USE_AITER=1 \
+          -e VLLM_ROCM_USE_AITER=0 \
           -e HF_TOKEN \
-          -e AWS_ACCESS_KEY_ID \
-          -e AWS_SECRET_ACCESS_KEY \
           -v "${HF_CACHE}:${HF_MOUNT}" \
           -e "HF_HOME=${HF_MOUNT}" \
           -e "PYTHONPATH=${MYPYTHONPATH}" \

.buildkite/scripts/upload-nightly-wheels.sh

@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+
+set -ex
+
+# Upload a single wheel to S3 under the omni/ prefix.
+# Index generation is handled separately by generate-and-upload-nightly-index.sh.
+
+BUCKET="vllm-wheels"
+SUBPATH="omni/$BUILDKITE_COMMIT"
+S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
+
+# ========= collect & upload the wheel ==========
+
+# python3 -m build outputs to dist/ by default
+wheel_files=(dist/*.whl)
+
+# Check that exactly one wheel is found
+if [[ ${#wheel_files[@]} -ne 1 ]]; then
+  echo "Error: Expected exactly one wheel file in dist/, but found ${#wheel_files[@]}"
+  exit 1
+fi
+wheel="${wheel_files[0]}"
+
+echo "Uploading wheel: $wheel"
+
+# Extract the version from the wheel
+version=$(unzip -p "$wheel" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
+echo "Version in wheel: $version"
+
+# Upload wheel to S3
+aws s3 cp "$wheel" "$S3_COMMIT_PREFIX"
+
+echo "Wheel uploaded to $S3_COMMIT_PREFIX. Index generation is handled by a separate step."

.buildkite/test-amd-merge.yml

@@ -0,0 +1,215 @@
+steps:
+
+- label: "Simple Unit Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_ROCM_USE_AITER=0
+    - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+
+- label: "Diffusion Model Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
+
+- label: "Diffusion Images API LoRA E2E"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
+
+- label: "Diffusion Model CPU offloading Test"
+  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
+    - |
+      timeout 20m bash -c '
+        set +e
+        pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+        EXIT1=\$?
+        pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+        EXIT2=\$?
+        exit \$((EXIT1 | EXIT2))
+      '
+
+## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
+# - label: "Audio Generation Model Test"
+#   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
+#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
+
+- label: "Diffusion Cache Backend Test"
+  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
+    - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
+
+- label: "Diffusion Sequence Parallelism Test"
+  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
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+
+# merge-only tests
+- label: "Diffusion Tensor Parallelism Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
+
+- label: "Diffusion GPU Worker Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
+
+- label: "Benchmark & Engine Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - |
+      timeout 20m bash -c '
+        set +e
+        pytest -s -v tests/benchmarks/test_serve_cli.py
+        EXIT1=\$?
+        pytest -s -v tests/engine/test_async_omni_engine_abort.py
+        EXIT2=\$?
+        exit \$((EXIT1 | EXIT2))
+      '
+
+- label: "Omni Model Test Qwen2-5-Omni"
+  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
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+
+- label: "Omni Model Test Qwen3-Omni"
+  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 VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
+
+- label: "Qwen3-TTS CustomVoice E2E Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - |
+      timeout 20m bash -c '
+        export VLLM_LOGGING_LEVEL=DEBUG
+        export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
+      '
+
+- label: "Qwen3-TTS Base E2E Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - |
+      timeout 20m bash -c '
+        export VLLM_LOGGING_LEVEL=DEBUG
+        export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_base.py
+      '
+
+- label: "Diffusion Image Edit Test"
+  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
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+
+# split Bagel Model Test with H100 (Real Weights) into three tests
+- label: "Bagel Text2Img Model Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
+
+- label: "Bagel Img2Img Model Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
+
+- label: "Bagel Online Serving Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_IMAGE_FETCH_TIMEOUT=60
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model" -k "rocm"

.buildkite/test-amd-ready.yaml

@@ -0,0 +1,179 @@
+steps:
+
+- label: "Simple Unit Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_ROCM_USE_AITER=0
+    - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+
+- label: "Diffusion Model Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+
+- label: "Diffusion Model CPU offloading Test"
+  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
+    - |
+      timeout 20m bash -c '
+        set +e
+        pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+        EXIT1=\$?
+        pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+        EXIT2=\$?
+        exit \$((EXIT1 | EXIT2))
+      '
+
+## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
+# - label: "Audio Generation Model Test"
+#   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
+#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
+
+- label: "Diffusion Cache Backend Test"
+  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
+    - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
+
+- label: "Diffusion Sequence Parallelism Test"
+  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
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
+
+- label: "Diffusion GPU Worker Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
+
+- label: "Benchmark & Engine Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - |
+      timeout 30m bash -c '
+        set +e
+        pytest -s -v tests/benchmarks/test_serve_cli.py
+        EXIT1=\$?
+        pytest -s -v tests/engine/test_async_omni_engine_abort.py
+        EXIT2=\$?
+        exit \$((EXIT1 | EXIT2))
+      '
+
+- label: "Omni Model Test Qwen2-5-Omni"
+  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
+    - timeout 17m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+
+- label: "Omni Model Test Qwen3-Omni"
+  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 VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+    - timeout 10m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
+
+- label: "Qwen3-TTS E2E Test"
+  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 VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
+
+- label: "Diffusion Image Edit Test"
+  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
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+
+- label: "Bagel Text2Img Model Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model" -k "rocm"
+
+- label: "Bagel Img2Img Model Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model" -k "rocm"
+
+- label: "Bagel Online Serving Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_IMAGE_FETCH_TIMEOUT=60
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model" -k "rocm"

.buildkite/test-amd.yaml

@@ -1,7 +1,7 @@
 steps:
 
 - label: "Diffusion Model Test"
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
   agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
@@ -11,7 +11,7 @@ steps:
     - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
 
 - label: "Diffusion Images API LoRA E2E"
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
   agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
@@ -101,6 +101,7 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+    - export VLLM_ROCM_USE_AITER=0
     - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
 

.buildkite/test-merge.yml

@@ -12,11 +12,12 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Model Test"
-    timeout_in_minutes: 20
+    timeout_in_minutes: 30
     depends_on: upload-merge-pipeline
     commands:
       - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
@@ -29,11 +30,12 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Images API LoRA E2E"
-    timeout_in_minutes: 20
+    timeout_in_minutes: 30
     depends_on: upload-merge-pipeline
     commands:
       - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
@@ -46,6 +48,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -53,8 +56,8 @@ steps:
     timeout_in_minutes: 20
     depends_on: upload-merge-pipeline
     commands:
-      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
-      - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+      # Single pytest session for one combined summary at end of log.
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
     agents:
       queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
@@ -64,6 +67,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -81,6 +85,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -98,6 +103,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -105,7 +111,7 @@ steps:
     timeout_in_minutes: 20
     depends_on: upload-merge-pipeline
     commands:
-      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py tests/diffusion/distributed/test_ulysses_uaa_perf.py
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
@@ -116,6 +122,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -134,6 +141,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -152,24 +160,64 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Benchmark & Engine Test"
+  - label: "Engine Test"
     depends_on: upload-merge-pipeline
     commands:
       - |
         timeout 15m bash -c '
                 export VLLM_WORKER_MULTIPROC_METHOD=spawn
-                set +e
-                pytest -s -v tests/benchmarks/test_serve_cli.py
-                EXIT1=$$?
                 pytest -s -v tests/engine/test_async_omni_engine_abort.py
-                EXIT2=$$?
-                exit $$((EXIT1 | EXIT2))
         '
+    agents:
+      queue: "gpu_1_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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Omni Model Test"
+    timeout_in_minutes: 30
+    depends_on: upload-merge-pipeline
+    commands:
+      - export VLLM_LOGGING_LEVEL=DEBUG
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
     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
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Qwen3-TTS CustomVoice E2E Test"
+    depends_on: upload-merge-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
+        '
+    agents:
+      queue: "gpu_1_queue"
     plugins:
       - docker#v5.2.0:
           image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -181,20 +229,24 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Omni Model Test"
-    timeout_in_minutes: 15
+  - label: "Qwen3-TTS Base E2E Test"
     depends_on: upload-merge-pipeline
     commands:
-      - export VLLM_LOGGING_LEVEL=DEBUG
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+      - |
+        timeout 20m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
+        '
     agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+      queue: "gpu_1_queue"
     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:
@@ -206,8 +258,7 @@ steps:
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
       - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
+      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -226,6 +277,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -261,6 +317,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -272,37 +333,103 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  # - label: "Bagel Text2Img Model Test with H100"
-  #   timeout_in_minutes: 30
-  #   depends_on: upload-merge-pipeline
-  #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #               resources:
-  #                 limits:
-  #                   nvidia.com/gpu: 1
-  #               volumeMounts:
-  #                 - name: devshm
-  #                   mountPath: /dev/shm
-  #                 - name: hf-cache
-  #                   mountPath: /root/.cache/huggingface
-  #               env:
-  #                 - name: HF_HOME
-  #                   value: /root/.cache/huggingface
-  #           nodeSelector:
-  #             node.kubernetes.io/instance-type: gpu-h100-sxm
-  #           volumes:
-  #             - name: devshm
-  #               emptyDir:
-  #                 medium: Memory
-  #             - name: hf-cache
-  #               hostPath:
-  #                 path: /mnt/hf-cache
-  #                 type: DirectoryOrCreate
+  - label: "Bagel Model Test with H100 (Real Weights)"
+    timeout_in_minutes: 60
+    depends_on: upload-merge-pipeline
+    commands:
+      - |
+        timeout 55m bash -c '
+          set -e
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          export VLLM_IMAGE_FETCH_TIMEOUT=60
+          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory"
+          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: "Voxtral-TTS E2E Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+        '
+
+  - label: "CosyVoice3-TTS E2E Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts.py -m "advanced_model" --run-level "advanced_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate

.buildkite/test-nightly.yml

@@ -2,16 +2,14 @@ steps:
   - label: ":full_moon: Omni Model Test with H100"
     timeout_in_minutes: 90
     depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1"
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
       - |
-        set +e
-        pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "advanced_model" --run-level "advanced_model"
-        EXIT1=$$?
-        pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
-        EXIT2=$$?
-        exit $$((EXIT1 | EXIT2))
+        pytest -s -v \
+          tests/examples/ \
+          tests/e2e/online_serving/test_*_expansion.py \
+          -m "advanced_model and H100 and omni" --run-level "advanced_model"
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -30,6 +28,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -41,13 +44,14 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: ":full_moon: Omni Model Test"
-    timeout_in_minutes: 60
+  - label: ":full_moon: Omni Model Test with L4"
+    timeout_in_minutes: 90
     depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1"
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
+      - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+      - pytest -s -v tests/examples/ tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
@@ -58,16 +62,17 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: ":full_moon: Diffusion Model Test with H100"
-    timeout_in_minutes: 60
+    timeout_in_minutes: 120
     depends_on: upload-nightly-pipeline
-    # if: build.env("NIGHTLY") == "1"
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
+      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -86,6 +91,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -97,17 +107,56 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
-    timeout_in_minutes: 30
+  - label: ":full_moon: Diffusion Model (Wan2.2) Test with H100"
+    timeout_in_minutes: 90
     depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1"
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
     commands:
-      - |
-        huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
-        export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -m "advanced_model" --run-level "advanced_model"
     agents:
-      queue: "gpu_4_queue"
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: ":full_moon: Diffusion Model Test"
+    timeout_in_minutes: 60
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model"
+    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
@@ -116,21 +165,19 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: ":full_moon: Omni Model Perf Test & Test Case Statistics"
-    key: nightly-performance
-    timeout_in_minutes: 180
+
+  - label: ":full_moon: Doc Example Code Test with H100"
+    timeout_in_minutes: 60
     depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1"
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export BENCHMARK_DIR=tests
-      - pytest -s -v tests/perf/scripts/run_benchmark.py
-      - buildkite-agent artifact upload "tests/*.json"
-      - python tools/nightly/buildkite_testcase_statistics.py -o tests/buildkite_testcase_statistics.html
-      - buildkite-agent artifact upload "tests/*.html"
+      - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+      - pytest -s -v tests/examples/online_serving/test_text_to_image.py tests/examples/offline_inference/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -149,6 +196,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -160,19 +212,248 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
+  - label: ":full_moon: Omni Model Perf Test & Testcase Statistics with H100"
+    key: nightly-omni-performance
+    timeout_in_minutes: 180
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export BENCHMARK_DIR=tests/dfx/perf/results
+      - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+      - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
+      - buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+      - python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
+      - buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: ":full_moon: GEBench Accuracy Test with H100"
+    key: nightly-gebench-accuracy
+    timeout_in_minutes: 60
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
+      - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: ":full_moon: GEdit-Bench Accuracy Test with H100"
+    key: nightly-gedit-bench-accuracy
+    timeout_in_minutes: 60
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
+      - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
+      - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
+                    value: "120"
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: ":full_moon: Wan22 I2V Accuracy Test with H100"
+    key: nightly-wan22-i2v-accuracy
+    timeout_in_minutes: 180
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --run-level advanced_model
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: ":full_moon: Diffusion Perf Test with H100"
+    key: nightly-qwen-image-performance
+    timeout_in_minutes: 180
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+      - export CACHE_DIT_VERSION=1.3.0
+      - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
+      - buildkite-agent artifact upload "tests/dfx/perf/results/benchmark_results_*.json"
+      - buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 4
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  # No need to run this step for PRs with label nightly-test
   - label: ":email: Nightly Collection & Email"
     key: nightly-perf-distribution
     depends_on:
-    - nightly-performance
+      - nightly-omni-performance
+      - nightly-qwen-image-performance
     if: build.env("NIGHTLY") == "1"
     commands:
       - pip install openpyxl
-      - export DEFAULT_INPUT_DIR=tests
-      - export DEFAULT_OUTPUT_DIR=tests
-      - buildkite-agent artifact download "tests/*.json" . --step nightly-performance
-      - buildkite-agent artifact download "tests/*.html" . --step nightly-performance
+      - export DEFAULT_INPUT_DIR=tests/dfx/perf/results
+      - export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-qwen-image-performance
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-omni-performance
       - python tools/nightly/generate_nightly_perf_excel.py
-      - python tools/nightly/send_nightly_email.py --report-file "tests/*.xlsx, tests/*.html"
-      - buildkite-agent artifact upload "tests/*.xlsx"
+      - python tools/nightly/generate_nightly_perf_html.py
+      - python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html"
+      - buildkite-agent artifact upload "tests/dfx/perf/results/*.xlsx"
+      - buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
     agents:
       queue: "cpu_queue_premerge"

.buildkite/test-ready.yml

@@ -5,6 +5,24 @@ steps:
       - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
     agents:
       queue: "gpu_1_queue"
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Voxtral TTS CUDA Unit Test"
+    timeout_in_minutes: 10
+    depends_on: upload-ready-pipeline
+    commands:
+      - "timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py"
+    agents:
+      queue: "gpu_1_queue"
     plugins:
       - docker#v5.2.0:
           image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -18,7 +36,41 @@ steps:
   - label: "Diffusion Model Test"
     depends_on: upload-ready-pipeline
     commands:
-      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+      - timeout 30m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+    agents:
+      queue: "gpu_1_queue"
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Batching Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Custom Pipeline Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
     agents:
       queue: "gpu_1_queue"
     plugins:
@@ -44,6 +96,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -60,6 +113,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -76,6 +130,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -93,6 +148,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -110,25 +166,21 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
 
-  - label: "Benchmark & Engine Test"
+  - label: "Engine Test"
     depends_on: upload-ready-pipeline
     commands:
       - |
         timeout 15m bash -c '
                 export VLLM_WORKER_MULTIPROC_METHOD=spawn
-                set +e
-                pytest -s -v tests/benchmarks/test_serve_cli.py
-                EXIT1=$$?
                 pytest -s -v tests/engine/test_async_omni_engine_abort.py
-                EXIT2=$$?
-                exit $$((EXIT1 | EXIT2))
         '
     agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+      queue: "gpu_1_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
@@ -137,6 +189,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -148,7 +201,7 @@ steps:
         timeout 17m bash -c '
           export VLLM_LOGGING_LEVEL=DEBUG
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+          pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "core_model" --run-level "core_model"
         '
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
@@ -159,60 +212,126 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  # - label: "Omni Model Test with H100"
-  #   depends_on: upload-ready-pipeline
-  #   commands:
-  #     - |
-  #       timeout 20m bash -c '
-  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #         export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-  #         # - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-  #         pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
-  #       '
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #               resources:
-  #                 limits:
-  #                   nvidia.com/gpu: 2
-  #               volumeMounts:
-  #                 - name: devshm
-  #                   mountPath: /dev/shm
-  #                 - name: hf-cache
-  #                   mountPath: /root/.cache/huggingface
-  #               env:
-  #                 - name: HF_HOME
-  #                   value: /root/.cache/huggingface
-  #           nodeSelector:
-  #             node.kubernetes.io/instance-type: gpu-h100-sxm
-  #           volumes:
-  #             - name: devshm
-  #               emptyDir:
-  #                 medium: Memory
-  #             - name: hf-cache
-  #               hostPath:
-  #                 path: /mnt/hf-cache
-  #                 type: DirectoryOrCreate
+  - label: "Omni Model Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: "MiMo-Audio E2E Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 30m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
 
   - label: "Qwen3-TTS E2E Test"
     depends_on: upload-ready-pipeline
     commands:
       - |
         timeout 20m bash -c '
-          huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
           export VLLM_LOGGING_LEVEL=DEBUG
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py -k "not NoAsyncChunk"
+          export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
         '
     agents:
-      queue: "gpu_4_queue"
+      queue: "gpu_1_queue"
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "OmniVoice E2E Test"
+    timeout_in_minutes: 20
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_omnivoice.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "gpu_1_queue"
     plugins:
       - docker#v5.2.0:
           image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -224,6 +343,45 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+  - label: "Voxtral-TTS E2E Test"
+    timeout_in_minutes: 20
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
   # - label: "Diffusion Image Edit Test with H100 (1 GPU)"
   #   depends_on: upload-ready-pipeline
   #   commands:
@@ -261,39 +419,175 @@ steps:
   #                 path: /mnt/hf-cache
   #                 type: DirectoryOrCreate
 
-  # - label: "Bagel Text2Img Model Test with H100"
-  #   depends_on: upload-ready-pipeline
-  #   commands:
-  #     - |
-  #       timeout 30m bash -c '
-  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #         pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
-  #       '
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #               resources:
-  #                 limits:
-  #                   nvidia.com/gpu: 1
-  #               volumeMounts:
-  #                 - name: devshm
-  #                   mountPath: /dev/shm
-  #                 - name: hf-cache
-  #                   mountPath: /root/.cache/huggingface
-  #               env:
-  #                 - name: HF_HOME
-  #                   value: /root/.cache/huggingface
-  #           nodeSelector:
-  #             node.kubernetes.io/instance-type: gpu-h100-sxm
-  #           volumes:
-  #             - name: devshm
-  #               emptyDir:
-  #                 medium: Memory
-  #             - name: hf-cache
-  #               hostPath:
-  #                 path: /mnt/hf-cache
-  #                 type: DirectoryOrCreate
+  - label: "Bagel Text2Img Model Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 30m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: "Bagel Img2Img Model Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 30m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: "Bagel Online Serving Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 40m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          export VLLM_IMAGE_FETCH_TIMEOUT=60
+          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: "CosyVoice3-TTS E2E Test"
+    timeout_in_minutes: 20
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate

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

@@ -3,7 +3,7 @@
    Last synced: 2025-12-15
    Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests
 #}
-{% set docker_image_amd = "public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT-rocm-omni" %}
+{% set docker_image_amd = "rocm/vllm-ci:$BUILDKITE_COMMIT-rocm-omni" %}
 {% set default_working_dir = "/app/vllm-omni" %}
 
   - group: "AMD Tests"
@@ -13,7 +13,6 @@
         depends_on: ~
         soft_fail: false
         commands:
-          - "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
           - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target test --build-arg ARG_PYTORCH_ROCM_ARCH=gfx942 --progress plain ."
           - "docker push {{ docker_image_amd }}"
         key: "amd-build"
@@ -22,13 +21,15 @@
         retry:
           automatic:
             - exit_status: -1  # Agent was lost
-              limit: 1
+              limit: 2
             - exit_status: -10  # Agent was lost
-              limit: 1
+              limit: 2
+            - exit_status: 128  # Git connectivity issues
+              limit: 2
             - exit_status: 1  # Machine occasionally fail
               limit: 1
         agents:
-          queue: cpu_queue_premerge
+          queue: amd-cpu
 
     {% for step in steps %}
     {% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
@@ -40,9 +41,14 @@
           {% else %}
           queue: amd_mi325_1
           {% endif %}
-        command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh "(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe  }} ; {{ step.command  or (step.commands | join(" && ")) | safe }}"
+{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
+{% set indented_cmd = cmd_body | replace("\n", "\n            ") %}
+        command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
         env:
           DOCKER_BUILDKIT: "1"
+          TEST_COMMAND: |-
+            (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
+            {{ indented_cmd | safe }}
         priority: 100
         {% if step.grade and step.grade == "Blocking" %}
         soft_fail: false

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

@@ -74,28 +74,21 @@ body:
       If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. For example:
 
       ```python
-      from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config
+      from vllm_omni.entrypoints.omni import Omni
+      from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+      from vllm import SamplingParams
 
-      # Create stage configurations
-      ar_config = create_ar_stage_config(
-          stage_id=0,
-          model_path="Qwen/Qwen3-0.6B",
-          input_modalities=["text"],
-          output_modalities=["text"]
+      omni = Omni(
+          model="Qwen/Qwen-Image",
+          stage_configs_path="/path/to/stage_configs.yaml",
       )
 
-      dit_config = create_dit_stage_config(
-          stage_id=1,
-          model_path="stabilityai/stable-diffusion-2-1",
-          input_modalities=["text"],
-          output_modalities=["image"]
-      )
-
-      # Initialize OmniLLM
-      omni = OmniLLM([ar_config, dit_config])
-
-      # Generate
-      outputs = omni.generate(prompt="A scenic watercolor painting of a lighthouse at sunset")
+      prompts = [{"prompt": "A scenic watercolor painting of a lighthouse at sunset"}]
+      sampling_params_list = [
+          SamplingParams(max_tokens=1),
+          OmniDiffusionSamplingParams(num_outputs_per_prompt=1),
+      ]
+      outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
       ```
 
       If the code is too long (hopefully, it isn't), feel free to put it in a public gist and link it in the issue: https://gist.github.com.

.github/ISSUE_TEMPLATE/750-RFC.yml

@@ -7,7 +7,7 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Please take a look at previous [RFCs](https://github.com/vllm-project/vllm-omni/issues?q=label%3ARFC+sort%3Aupdated-desc) for reference.
+      #### Please take a look at previous [RFCs](https://github.com/vllm-project/vllm-omni/issues?q=in%3Atitle%20RFC%20sort%3Aupdated-desc) for reference.
 - type: textarea
   attributes:
     label: Motivation.
@@ -21,7 +21,7 @@ body:
     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).
+      Please provide the detailed design document of the RFC using the [template](https://docs.google.com/document/d/1jcgR3cDaUQH3VczD4ZcKaJAoYWHjCmnYzHkCNyz-9fk/edit?usp=sharing).
   validations:
     required: true
 - type: textarea

.github/workflows/build_wheel.yml

@@ -39,10 +39,10 @@ jobs:
         python-version: ["3.11", "3.12"]
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
 
@@ -57,7 +57,7 @@ jobs:
           bash scripts/build_wheel.sh --python python
 
       - name: Upload Wheel Artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v7
         with:
           name: vllm-omni-wheel-py${{ matrix.python-version }}
           path: dist/*.whl

.github/workflows/pr-reviewer.yml

@@ -24,12 +24,12 @@ jobs:
        github.event.comment.author_association == 'OWNER')
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4.2.2
+        uses: actions/checkout@v6.0.2
         with:
           fetch-depth: 0
 
       - name: Set up Python
-        uses: actions/setup-python@v5.3.0
+        uses: actions/setup-python@v6.2.0
         with:
           python-version: '3.11'
           cache: 'pip'
@@ -54,7 +54,7 @@ jobs:
 
       - name: Upload review logs
         if: always()
-        uses: actions/upload-artifact@v4.4.3
+        uses: actions/upload-artifact@v7.0.0
         with:
           name: pr-review-logs-${{ github.event.issue.number || github.event.pull_request.number }}
           path: /tmp/pr_review_*.log

.github/workflows/pre-commit.yml

@@ -18,8 +18,8 @@ jobs:
   pre-commit:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
-      - uses: actions/setup-python@v5
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
         with:
           python-version: "3.12"
           cache: 'pip'

README.md

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

benchmarks/README.md

@@ -15,6 +15,8 @@ This document explains the benchmark architecture across all benchmark assets in
 ## Directory layout
 - `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
 - `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
+- `benchmarks/accuracy/` — accuracy benchmark integrations that adapt external
+  benchmark suites to vLLM-Omni serving and evaluation flows.
 - Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
 
 ## Reference workflows

benchmarks/__init__.py

@@ -0,0 +1 @@
+"""Benchmark helpers and runnable benchmark entrypoints."""

benchmarks/accuracy/README.md

@@ -0,0 +1,27 @@
+# Accuracy Benchmarks
+
+This directory hosts accuracy benchmark integrations that run entirely through a
+local `vllm-omni serve` deployment.
+
+Current integrations:
+
+- `text_to_image/`: GEBench generation + local judge scoring flow.
+- `image_to_image/`: GEdit-Bench generation + local VIEScore-style scoring flow.
+
+Design notes:
+
+- Generation is executed through the OpenAI-compatible endpoints exposed by
+  `vllm-omni serve`.
+- Evaluation is also executed through a local OpenAI-compatible judge model
+  served by `vllm-omni`.
+- Both generation and judge requests accept either `http://host:port` or
+  `http://host:port/v1`.
+- Output directory layout intentionally stays close to the upstream repos.
+
+Test guidance:
+
+- Local static/self-checks live in `tests/benchmarks/test_accuracy_bench_utils.py`.
+- End-to-end generation/evaluation should be validated in a remote GPU
+  environment. In the current repo marker system there is `L4` but no `L5`
+  marker, so benchmark smoke tests should be wired as `advanced_model +
+  benchmark + L4` when GPU capacity is available.

benchmarks/accuracy/__init__.py

@@ -0,0 +1 @@
+"""Accuracy benchmark integrations for vLLM-Omni."""

benchmarks/accuracy/common.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+import base64
+import io
+import json
+from pathlib import Path
+from typing import Any
+
+import requests
+from PIL import Image
+
+
+def ensure_dir(path: Path) -> Path:
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def load_json(path: Path) -> dict[str, Any]:
+    with path.open("r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def write_json(path: Path, payload: dict[str, Any]) -> None:
+    ensure_dir(path.parent)
+    with path.open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2, ensure_ascii=False)
+
+
+def save_image(path: Path, image: Image.Image) -> None:
+    ensure_dir(path.parent)
+    image.save(path)
+
+
+def find_first_image(folder: Path, stem: str | None = None) -> Path | None:
+    patterns = [f"{stem}.*"] if stem else ["*.png", "*.jpg", "*.jpeg", "*.webp"]
+    for pattern in patterns:
+        for candidate in sorted(folder.glob(pattern)):
+            if candidate.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"}:
+                return candidate
+    return None
+
+
+def extract_json_object(raw_text: str) -> dict[str, Any]:
+    raw_text = raw_text.strip()
+    delimiter = "||V^=^V||"
+    if raw_text.count(delimiter) >= 2:
+        start = raw_text.find(delimiter) + len(delimiter)
+        end = raw_text.rfind(delimiter)
+        raw_text = raw_text[start:end].strip()
+
+    start = raw_text.find("{")
+    end = raw_text.rfind("}")
+    if start == -1 or end == -1 or end < start:
+        raise ValueError(f"Could not find JSON object in: {raw_text[:200]}")
+    return json.loads(raw_text[start : end + 1])
+
+
+def build_openai_url(base_url: str, api_path: str) -> str:
+    base = base_url.rstrip("/")
+    normalized_path = api_path if api_path.startswith("/") else f"/{api_path}"
+    if base.endswith(normalized_path):
+        return base
+    if base.endswith("/v1"):
+        return f"{base}{normalized_path}"
+    return f"{base}/v1{normalized_path}"
+
+
+def pil_to_base64(image: Image.Image, image_format: str = "PNG") -> str:
+    buffer = io.BytesIO()
+    image.save(buffer, format=image_format)
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
+def pil_to_data_url(image: Image.Image, image_format: str = "PNG") -> str:
+    return f"data:image/{image_format.lower()};base64,{pil_to_base64(image, image_format=image_format)}"
+
+
+def decode_base64_image(encoded: str) -> Image.Image:
+    image = Image.open(io.BytesIO(base64.b64decode(encoded)))
+    image.load()
+    return image.convert("RGB")
+
+
+def pil_to_png_bytes(image: Image.Image) -> bytes:
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    return buffer.getvalue()
+
+
+class VllmOmniImageClient:
+    """Thin OpenAI-compatible image client for vLLM-Omni serving."""
+
+    def __init__(self, base_url: str, api_key: str = "EMPTY", timeout: int = 600):
+        self.base_url = base_url.rstrip("/")
+        self.api_key = api_key
+        self.timeout = timeout
+
+    @property
+    def _headers(self) -> dict[str, str]:
+        return {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+
+    def generate_text_to_image(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        num_inference_steps: int = 20,
+        guidance_scale: float | None = None,
+        seed: int | None = None,
+        output_compression: int | None = None,
+    ) -> Image.Image:
+        payload: dict[str, Any] = {
+            "model": model,
+            "prompt": prompt,
+            "n": 1,
+            "size": f"{width}x{height}",
+            "response_format": "b64_json",
+            "num_inference_steps": num_inference_steps,
+        }
+        if guidance_scale is not None:
+            payload["guidance_scale"] = guidance_scale
+        if seed is not None:
+            payload["seed"] = seed
+        if output_compression is not None:
+            payload["output_compression"] = output_compression
+
+        response = requests.post(
+            build_openai_url(self.base_url, "/images/generations"),
+            json=payload,
+            headers=self._headers,
+            timeout=self.timeout,
+        )
+        response.raise_for_status()
+        return decode_base64_image(response.json()["data"][0]["b64_json"])
+
+    def generate_image_edit(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        images: Image.Image | list[Image.Image],
+        width: int,
+        height: int,
+        num_inference_steps: int = 20,
+        guidance_scale: float | None = None,
+        seed: int | None = None,
+        negative_prompt: str | None = None,
+        output_compression: int | None = None,
+    ) -> Image.Image:
+        if not isinstance(images, list):
+            images = [images]
+        data: dict[str, Any] = {
+            "model": model,
+            "prompt": prompt,
+            "n": 1,
+            "size": f"{width}x{height}",
+            "response_format": "b64_json",
+            "num_inference_steps": str(num_inference_steps),
+        }
+        if guidance_scale is not None:
+            data["guidance_scale"] = str(guidance_scale)
+        if seed is not None:
+            data["seed"] = str(seed)
+        if negative_prompt:
+            data["negative_prompt"] = negative_prompt
+        if output_compression is not None:
+            data["output_compression"] = str(output_compression)
+
+        files = [
+            (
+                "image[]" if len(images) > 1 else "image",
+                (f"image_{index}.png", pil_to_png_bytes(image), "image/png"),
+            )
+            for index, image in enumerate(images)
+        ]
+
+        edit_paths = ["/images/edits", "/images/edit"]
+        last_response: requests.Response | None = None
+        for api_path in edit_paths:
+            response = requests.post(
+                build_openai_url(self.base_url, api_path),
+                data=data,
+                files=files,
+                headers={"Authorization": f"Bearer {self.api_key}"},
+                timeout=self.timeout,
+            )
+            last_response = response
+            if response.status_code != 404:
+                response.raise_for_status()
+                return decode_base64_image(response.json()["data"][0]["b64_json"])
+
+        assert last_response is not None
+        last_response.raise_for_status()
+        raise ValueError("No image payload returned from image edit endpoint")

benchmarks/accuracy/image_to_image/README.md

@@ -0,0 +1,103 @@
+# GEdit-Bench on vLLM-Omni
+
+This integration adapts the upstream `stepfun-ai/Step1X-Edit/GEdit-Bench`
+evaluation flow into `vllm-omni/benchmarks/accuracy/image_to_image`.
+
+Upstream mapping:
+
+- `run_gedit_score.py` -> `run_gedit_bench.py evaluate`
+- `calculate_statistics.py` -> `run_gedit_bench.py summarize`
+- upstream output layout under `results/<method>/fullset/...` is preserved
+
+What changed:
+
+- The upstream repo mainly ships evaluation scripts. This integration adds a
+  generation runner that uses the local `vllm-omni` OpenAI-compatible
+  `/v1/images/edits` endpoint to produce benchmark outputs in the expected
+  directory structure.
+- The evaluator keeps the same VIEScore-style decomposition:
+  - `semantics_score`
+  - `quality_score`
+  - `overall_score = sqrt(semantics_score * quality_score)`
+- Judge calls are routed to a local OpenAI-compatible model served by
+  `vllm-omni`, not a remote provider.
+
+Dataset:
+
+- Default `--dataset-ref` is `stepfun-ai/GEdit-Bench`
+- You can also pass a local dataset directory previously saved with
+  Hugging Face `datasets`
+
+Example usage:
+
+```bash
+python benchmarks/accuracy/image_to_image/run_gedit_bench.py generate \
+  --output-root benchmarks/accuracy/image_to_image/results \
+  --base-url http://127.0.0.1:8000 \
+  --model Qwen/Qwen-Image-Edit \
+  --model-name qwen_image_edit \
+  --dataset-ref stepfun-ai/GEdit-Bench \
+  --task-type all \
+  --instruction-language en
+```
+
+```bash
+python benchmarks/accuracy/image_to_image/run_gedit_bench.py evaluate \
+  --output-root benchmarks/accuracy/image_to_image/results \
+  --model-name qwen_image_edit \
+  --save-dir benchmarks/accuracy/image_to_image/scores \
+  --dataset-ref stepfun-ai/GEdit-Bench \
+  --judge-base-url http://127.0.0.1:8000 \
+  --judge-model Qwen/Qwen2.5-VL-7B-Instruct \
+  --judge-api-key EMPTY
+```
+
+```bash
+python benchmarks/accuracy/image_to_image/run_gedit_bench.py summarize \
+  --csv-path benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_vie_score.csv \
+  --language en
+```
+
+Example summary output:
+
+```json
+{
+  "language": "all",
+  "languages": {
+    "en": {
+      "overall": {"count": 110, "Q_SC": 6.58, "Q_PQ": 5.89, "Q_O": 5.86},
+      "intersection": {"count": 78, "Q_SC": 6.50, "Q_PQ": 5.66, "Q_O": 5.65}
+    },
+    "cn": {
+      "overall": {"count": 110, "Q_SC": 6.90, "Q_PQ": 5.78, "Q_O": 6.11},
+      "intersection": {"count": 63, "Q_SC": 7.22, "Q_PQ": 5.59, "Q_O": 6.28}
+    }
+  }
+}
+```
+
+Example generated images to inspect:
+
+- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/background_change/en/<sample>.png`
+- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/text_change/en/<sample>.png`
+- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/subject-replace/cn/<sample>.png`
+
+Example score artifacts to inspect together with the images:
+
+- `benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_vie_score.csv`
+- `benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_summary.json`
+
+What to expect:
+
+- `Q_SC` measures instruction following and content preservation.
+- `Q_PQ` measures image naturalness and artifact quality.
+- `Q_O` is the combined overall score; higher is better.
+- `overall.count` is the number of evaluated samples for that language, while `intersection.count` is the subset with `intersection_exist == True`.
+
+Notes:
+
+- This flow requires the optional Hugging Face `datasets` package.
+- `generate` writes `generation_manifest.json` with local output coverage.
+- The current repo marker set exposes `L4` but not `L5`, so if you promote an
+  end-to-end smoke test into CI, use the existing `advanced_model`, `benchmark`,
+  and `L4` markers or introduce a new repo-wide marker explicitly first.

benchmarks/accuracy/image_to_image/__init__.py

@@ -0,0 +1 @@
+"""GEdit-Bench integration for vLLM-Omni."""

benchmarks/accuracy/image_to_image/gedit_bench.py

@@ -0,0 +1,718 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import logging
+import math
+import statistics
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from datetime import datetime, timezone
+from io import BytesIO
+from pathlib import Path
+from typing import Any
+
+import requests
+from PIL import Image
+from tqdm.auto import tqdm
+
+from benchmarks.accuracy.common import (
+    VllmOmniImageClient,
+    build_openai_url,
+    extract_json_object,
+    write_json,
+)
+
+GROUPS = [
+    "background_change",
+    "color_alter",
+    "material_alter",
+    "motion_change",
+    "ps_human",
+    "style_change",
+    "subject-add",
+    "subject-remove",
+    "subject-replace",
+    "text_change",
+    "tone_transfer",
+]
+DEFAULT_SAMPLES_PER_GROUP = 10
+logger = logging.getLogger(__name__)
+
+
+def infer_model_name(model: str) -> str:
+    normalized = model.rstrip("/\\")
+    name = Path(normalized).name
+    return name or normalized
+
+
+def resolve_model_name(*, model_name: str | None, model: str | None = None, output_root: Path | None = None) -> str:
+    if model_name:
+        return model_name
+    if model:
+        return infer_model_name(model)
+    if output_root is not None:
+        candidates = sorted(path.name for path in output_root.iterdir() if path.is_dir())
+        if len(candidates) == 1:
+            return candidates[0]
+        if not candidates:
+            raise ValueError(f"Could not infer model-name from empty output root: {output_root}")
+        raise ValueError(
+            f"Could not infer model-name from output root {output_root}; multiple candidates found: {candidates}"
+        )
+    raise ValueError("model-name is required when it cannot be inferred from model or output-root")
+
+
+def parse_score_payload(raw_text: str) -> dict[str, Any]:
+    try:
+        parsed = extract_json_object(raw_text)
+    except Exception:
+        stripped = raw_text.strip()
+        if stripped.startswith("[") and stripped.endswith("]"):
+            parsed = {"score": json.loads(stripped), "reasoning": ""}
+        else:
+            raise
+    score = parsed.get("score", [])
+    if not isinstance(score, list):
+        score = [score]
+    parsed["score"] = [int(value) for value in score]
+    return parsed
+
+
+def _utc_timestamp() -> str:
+    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+
+
+def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
+    by_task: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    by_language: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    for record in records:
+        by_task[record["task_type"]].append(record)
+        by_language[record["instruction_language"]].append(record)
+
+    return {
+        "count": len(records),
+        "by_task": {
+            task_type: {
+                "count": len(rows),
+                "samples": sorted(row["key"] for row in rows),
+            }
+            for task_type, rows in sorted(by_task.items())
+        },
+        "by_language": {
+            language: {
+                "count": len(rows),
+                "samples": sorted(row["key"] for row in rows),
+            }
+            for language, rows in sorted(by_language.items())
+        },
+    }
+
+
+def select_balanced_gedit_rows(
+    rows: list[dict[str, Any]],
+    *,
+    task_type: str = "all",
+    instruction_language: str = "all",
+    samples_per_group: int | None,
+) -> list[dict[str, Any]]:
+    filtered_rows = []
+    for row in rows:
+        if task_type != "all" and row["task_type"] != task_type:
+            continue
+        if instruction_language != "all" and row["instruction_language"] != instruction_language:
+            continue
+        filtered_rows.append(row)
+
+    if samples_per_group is None:
+        return filtered_rows
+
+    def _select_rows_for_group(group_rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        if instruction_language != "all":
+            return group_rows[:samples_per_group]
+
+        rows_by_language: dict[str, list[dict[str, Any]]] = defaultdict(list)
+        for row in group_rows:
+            rows_by_language[str(row.get("instruction_language", "")).strip()].append(row)
+
+        ordered_languages = ["en", "cn"]
+        per_language_quota = samples_per_group // max(len(ordered_languages), 1)
+        remainder = samples_per_group % max(len(ordered_languages), 1)
+
+        selected_rows: list[dict[str, Any]] = []
+        leftovers: list[dict[str, Any]] = []
+        for index, language in enumerate(ordered_languages):
+            language_rows = rows_by_language.get(language, [])
+            quota = per_language_quota + (1 if index < remainder else 0)
+            selected_rows.extend(language_rows[:quota])
+            leftovers.extend(language_rows[quota:])
+
+        if len(selected_rows) < samples_per_group:
+            selected_rows.extend(leftovers[: samples_per_group - len(selected_rows)])
+
+        return selected_rows
+
+    if task_type != "all":
+        return _select_rows_for_group(filtered_rows)
+
+    grouped_rows: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    for row in filtered_rows:
+        grouped_rows[row["task_type"]].append(row)
+
+    selected: list[dict[str, Any]] = []
+    for group in GROUPS:
+        selected.extend(_select_rows_for_group(grouped_rows.get(group, [])))
+    return selected
+
+
+def summarize_gedit_rows(rows: list[dict[str, Any]], language: str = "all") -> dict[str, Any]:
+    return summarize_gedit_rows_with_backbone(rows, language=language)
+
+
+def _mean_or_none(values: list[float]) -> float | None:
+    return statistics.fmean(values) if values else None
+
+
+def _to_q_metrics(section: dict[str, float | None]) -> dict[str, float | None]:
+    return {
+        "count": section.get("count"),
+        "Q_SC": section.get("avg_semantics"),
+        "Q_PQ": section.get("avg_quality"),
+        "Q_O": section.get("avg_overall"),
+    }
+
+
+def _summarize_gedit_rows_single_language(
+    rows: list[dict[str, Any]],
+    *,
+    language: str,
+) -> dict[str, Any]:
+    filtered_rows = [row for row in rows if str(row.get("instruction_language", "")).strip() == language]
+
+    def _to_bool(value: Any) -> bool:
+        if isinstance(value, bool):
+            return value
+        return str(value).strip().lower() in {"1", "true", "yes"}
+
+    per_group: dict[str, dict[str, float | None]] = {}
+    intersection_group: dict[str, dict[str, float | None]] = {}
+
+    for group in GROUPS:
+        group_rows = [row for row in filtered_rows if row.get("task_type") == group]
+        semantics = [float(row["semantics_score"]) for row in group_rows]
+        quality = [float(row["quality_score"]) for row in group_rows]
+        overall = [math.sqrt(float(row["semantics_score"]) * float(row["quality_score"])) for row in group_rows]
+
+        per_group[group] = {
+            "count": len(group_rows),
+            "avg_semantics": _mean_or_none(semantics),
+            "avg_quality": _mean_or_none(quality),
+            "avg_overall": _mean_or_none(overall),
+        }
+
+        intersection_rows = [row for row in group_rows if _to_bool(row.get("intersection_exist", False))]
+        intersection_semantics = [float(row["semantics_score"]) for row in intersection_rows]
+        intersection_quality = [float(row["quality_score"]) for row in intersection_rows]
+        intersection_overall = [
+            math.sqrt(float(row["semantics_score"]) * float(row["quality_score"])) for row in intersection_rows
+        ]
+        intersection_group[group] = {
+            "count": len(intersection_rows),
+            "avg_semantics": _mean_or_none(intersection_semantics),
+            "avg_quality": _mean_or_none(intersection_quality),
+            "avg_overall": _mean_or_none(intersection_overall),
+        }
+
+    overall_section = {
+        "count": len(filtered_rows),
+        "avg_semantics": _mean_or_none(
+            [score["avg_semantics"] for score in per_group.values() if score["avg_semantics"] is not None]
+        ),
+        "avg_quality": _mean_or_none(
+            [score["avg_quality"] for score in per_group.values() if score["avg_quality"] is not None]
+        ),
+        "avg_overall": _mean_or_none(
+            [score["avg_overall"] for score in per_group.values() if score["avg_overall"] is not None]
+        ),
+    }
+    intersection_section = {
+        "count": sum(score["count"] for score in intersection_group.values()),
+        "avg_semantics": _mean_or_none(
+            [score["avg_semantics"] for score in intersection_group.values() if score["avg_semantics"] is not None]
+        ),
+        "avg_quality": _mean_or_none(
+            [score["avg_quality"] for score in intersection_group.values() if score["avg_quality"] is not None]
+        ),
+        "avg_overall": _mean_or_none(
+            [score["avg_overall"] for score in intersection_group.values() if score["avg_overall"] is not None]
+        ),
+    }
+
+    return {
+        "language": language,
+        "by_group": {group: _to_q_metrics(section) for group, section in per_group.items()},
+        "overall": _to_q_metrics(overall_section),
+        "intersection": _to_q_metrics(intersection_section),
+    }
+
+
+def summarize_gedit_rows_with_backbone(
+    rows: list[dict[str, Any]],
+    *,
+    language: str = "all",
+) -> dict[str, Any]:
+    if language == "all":
+        return {
+            "language": "all",
+            "languages": {
+                single_language: _summarize_gedit_rows_single_language(
+                    rows,
+                    language=single_language,
+                )
+                for single_language in ["en", "cn"]
+            },
+        }
+
+    return _summarize_gedit_rows_single_language(rows, language=language)
+
+
+def _require_datasets():
+    try:
+        from datasets import load_dataset, load_from_disk
+    except ImportError as exc:
+        raise ImportError("GEdit-Bench requires the optional `datasets` package.") from exc
+    return load_dataset, load_from_disk
+
+
+def _load_gedit_dataset(dataset_ref: str):
+    load_dataset, load_from_disk = _require_datasets()
+    dataset_path = Path(dataset_ref)
+    if dataset_path.exists():
+        if (dataset_path / "state.json").exists() and (dataset_path / "dataset_info.json").exists():
+            return load_from_disk(str(dataset_path))
+        return load_dataset(str(dataset_path))
+    return load_dataset(dataset_ref)
+
+
+def _resolve_gedit_split(dataset_obj: Any) -> Any:
+    if isinstance(dataset_obj, dict):
+        if "train" in dataset_obj:
+            return dataset_obj["train"]
+        return dataset_obj
+    try:
+        return dataset_obj["train"]
+    except Exception:
+        return dataset_obj
+
+
+def _to_pil_image(value: Any) -> Image.Image:
+    if isinstance(value, Image.Image):
+        return value.convert("RGB")
+    if isinstance(value, dict) and "bytes" in value:
+        image = Image.open(BytesIO(value["bytes"]))
+        image.load()
+        return image.convert("RGB")
+    raise TypeError(f"Unsupported image payload type: {type(value)!r}")
+
+
+class LocalVIEScorer:
+    def __init__(self, *, base_url: str, api_key: str, model: str, timeout: int = 600):
+        self.base_url = base_url.rstrip("/")
+        self.api_key = api_key
+        self.model = model
+        self.timeout = timeout
+        self.sc_prompt = (
+            "You are evaluating image editing quality.\n"
+            "Two images are provided: the source image and the edited image.\n"
+            'Return JSON only in the format {"score": [edit_success, content_preservation], '
+            '"reasoning": "..."}.\n'
+            "Each score must be an integer from 0 to 10.\n"
+            "Editing instruction: <instruction>"
+        )
+        self.pq_prompt = (
+            "You are evaluating image quality.\n"
+            'Return JSON only in the format {"score": [naturalness, artifact_free], "reasoning": "..."}.\n'
+            "Each score must be an integer from 0 to 10."
+        )
+
+    def _request(self, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
+        from benchmarks.accuracy.common import pil_to_data_url
+
+        content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
+        for image in images:
+            content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
+
+        response = requests.post(
+            build_openai_url(self.base_url, "/chat/completions"),
+            json={
+                "model": self.model,
+                "messages": [{"role": "user", "content": content}],
+                "temperature": 0,
+            },
+            headers={
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json",
+            },
+            timeout=self.timeout,
+        )
+        response.raise_for_status()
+        message_content = response.json()["choices"][0]["message"]["content"]
+        if isinstance(message_content, list):
+            text = "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
+        else:
+            text = str(message_content)
+        return parse_score_payload(text)
+
+    def evaluate(self, source_image: Image.Image, edited_image: Image.Image, instruction: str) -> dict[str, float]:
+        sc_payload = self._request(self.sc_prompt.replace("<instruction>", instruction), [source_image, edited_image])
+        pq_payload = self._request(self.pq_prompt, [edited_image])
+        semantics = float(min(sc_payload["score"])) if sc_payload["score"] else 0.0
+        quality = float(min(pq_payload["score"])) if pq_payload["score"] else 0.0
+        overall = math.sqrt(semantics * quality)
+        return {
+            "semantics_score": semantics,
+            "quality_score": quality,
+            "overall_score": overall,
+        }
+
+
+class GEditBenchRunner:
+    def __init__(
+        self,
+        *,
+        dataset_ref: str,
+        output_root: Path,
+        base_url: str,
+        model: str,
+        api_key: str = "EMPTY",
+        width: int = 512,
+        height: int = 512,
+        num_inference_steps: int = 20,
+        guidance_scale: float | None = None,
+        seed: int | None = 42,
+    ):
+        self.dataset_ref = dataset_ref
+        self.output_root = output_root
+        self.model = model
+        self.width = width
+        self.height = height
+        self.num_inference_steps = num_inference_steps
+        self.guidance_scale = guidance_scale
+        self.seed = seed
+        self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
+
+    def generate(
+        self,
+        *,
+        model_name: str,
+        task_type: str = "all",
+        instruction_language: str = "all",
+        workers: int = 1,
+        max_samples: int | None = None,
+        samples_per_group: int | None = None,
+    ) -> list[dict[str, Any]]:
+        dataset = _resolve_gedit_split(_load_gedit_dataset(self.dataset_ref))
+        rows = select_balanced_gedit_rows(
+            list(dataset),
+            task_type=task_type,
+            instruction_language=instruction_language,
+            samples_per_group=samples_per_group,
+        )
+        if max_samples is not None:
+            rows = rows[:max_samples]
+
+        outputs: list[dict[str, Any]] = []
+        total = len(rows)
+        if workers <= 1:
+            with tqdm(total=total, desc="GEdit generate", unit="sample") as progress:
+                for item in rows:
+                    result = self._safe_generate_one(model_name, item)
+                    if result:
+                        outputs.append(result)
+                    progress.update(1)
+            return outputs
+
+        with tqdm(total=total, desc="GEdit generate", unit="sample") as progress:
+            with ThreadPoolExecutor(max_workers=workers) as executor:
+                futures = [executor.submit(self._safe_generate_one, model_name, item) for item in rows]
+                for future in as_completed(futures):
+                    result = future.result()
+                    if result:
+                        outputs.append(result)
+                    progress.update(1)
+        return outputs
+
+    def _safe_generate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
+        try:
+            return self._generate_one(model_name, item)
+        except Exception:
+            logger.exception("Failed to generate GEdit-Bench sample %s", item.get("key", "<unknown>"))
+            return None
+
+    def _generate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
+        output_path = (
+            self.output_root
+            / model_name
+            / "fullset"
+            / item["task_type"]
+            / item["instruction_language"]
+            / f"{item['key']}.png"
+        )
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        if output_path.exists():
+            return {
+                "task_type": item["task_type"],
+                "instruction_language": item["instruction_language"],
+                "key": item["key"],
+                "output_path": str(output_path),
+            }
+
+        source_image = _to_pil_image(item["input_image_raw"])
+        edited_image = self.client.generate_image_edit(
+            model=self.model,
+            prompt=item["instruction"],
+            images=source_image,
+            width=self.width,
+            height=self.height,
+            num_inference_steps=self.num_inference_steps,
+            guidance_scale=self.guidance_scale,
+            seed=self.seed,
+        )
+        edited_image.save(output_path)
+        return {
+            "task_type": item["task_type"],
+            "instruction_language": item["instruction_language"],
+            "key": item["key"],
+            "output_path": str(output_path),
+        }
+
+
+class GEditBenchEvaluator:
+    def __init__(self, *, dataset_ref: str, output_root: Path, scorer: LocalVIEScorer):
+        self.dataset_ref = dataset_ref
+        self.output_root = output_root
+        self.scorer = scorer
+
+    def evaluate(
+        self,
+        *,
+        model_name: str,
+        save_dir: Path,
+        task_type: str = "all",
+        instruction_language: str = "all",
+        workers: int = 1,
+        max_samples: int | None = None,
+        samples_per_group: int | None = None,
+    ) -> dict[str, Any]:
+        dataset = _resolve_gedit_split(_load_gedit_dataset(self.dataset_ref))
+        rows = select_balanced_gedit_rows(
+            list(dataset),
+            task_type=task_type,
+            instruction_language=instruction_language,
+            samples_per_group=samples_per_group,
+        )
+        if max_samples is not None:
+            rows = rows[:max_samples]
+
+        results: list[dict[str, Any]] = []
+        total = len(rows)
+        if workers <= 1:
+            with tqdm(total=total, desc="GEdit evaluate", unit="sample") as progress:
+                for item in rows:
+                    result = self._safe_evaluate_one(model_name, item)
+                    if result:
+                        results.append(result)
+                    progress.update(1)
+        else:
+            with tqdm(total=total, desc="GEdit evaluate", unit="sample") as progress:
+                with ThreadPoolExecutor(max_workers=workers) as executor:
+                    futures = [executor.submit(self._safe_evaluate_one, model_name, item) for item in rows]
+                    for future in as_completed(futures):
+                        result = future.result()
+                        if result:
+                            results.append(result)
+                        progress.update(1)
+
+        save_dir.mkdir(parents=True, exist_ok=True)
+        base_name = f"{model_name}_{task_type}_{instruction_language}"
+        timestamp = _utc_timestamp()
+        csv_path = save_dir / f"{base_name}_vie_score.csv"
+        timestamped_csv_path = save_dir / f"{base_name}_vie_score_{timestamp}.csv"
+
+        def _write_csv(path: Path) -> None:
+            with path.open("w", encoding="utf-8", newline="") as handle:
+                writer = csv.DictWriter(
+                    handle,
+                    fieldnames=[
+                        "key",
+                        "task_type",
+                        "edited_image",
+                        "instruction",
+                        "semantics_score",
+                        "quality_score",
+                        "overall_score",
+                        "intersection_exist",
+                        "instruction_language",
+                    ],
+                )
+                writer.writeheader()
+                for row in results:
+                    writer.writerow(row)
+
+        _write_csv(csv_path)
+        _write_csv(timestamped_csv_path)
+
+        summary = summarize_gedit_rows_with_backbone(
+            results,
+            language=instruction_language,
+        )
+        summary_path = save_dir / f"{base_name}_summary.json"
+        timestamped_summary_path = save_dir / f"{base_name}_summary_{timestamp}.json"
+        write_json(summary_path, summary)
+        write_json(timestamped_summary_path, summary)
+        return {
+            "results": results,
+            "summary": summary,
+            "csv_path": str(csv_path),
+            "summary_path": str(summary_path),
+            "timestamped_csv_path": str(timestamped_csv_path),
+            "timestamped_summary_path": str(timestamped_summary_path),
+        }
+
+    def _safe_evaluate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
+        try:
+            return self._evaluate_one(model_name, item)
+        except Exception:
+            logger.exception("Failed to evaluate GEdit-Bench sample %s", item.get("key", "<unknown>"))
+            return None
+
+    def _evaluate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
+        edited_image_path = (
+            self.output_root
+            / model_name
+            / "fullset"
+            / item["task_type"]
+            / item["instruction_language"]
+            / f"{item['key']}.png"
+        )
+        if not edited_image_path.exists():
+            return None
+
+        source_image = _to_pil_image(item["input_image_raw"])
+        edited_image = Image.open(edited_image_path).convert("RGB")
+        scores = self.scorer.evaluate(source_image, edited_image, item["instruction"])
+        return {
+            "key": item["key"],
+            "task_type": item["task_type"],
+            "edited_image": str(edited_image_path),
+            "instruction": item["instruction"],
+            "semantics_score": scores["semantics_score"],
+            "quality_score": scores["quality_score"],
+            "overall_score": scores["overall_score"],
+            "intersection_exist": item.get("Intersection_exist", False),
+            "instruction_language": item["instruction_language"],
+        }
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Run the GEdit-Bench integration against a local vLLM-Omni server.")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    generate = subparsers.add_parser("generate")
+    generate.add_argument("--dataset-ref", type=str, default="stepfun-ai/GEdit-Bench")
+    generate.add_argument("--output-root", type=Path, required=True)
+    generate.add_argument("--base-url", type=str, required=True)
+    generate.add_argument("--model", type=str, required=True)
+    generate.add_argument("--model-name", type=str, default=None)
+    generate.add_argument("--api-key", type=str, default="EMPTY")
+    generate.add_argument("--task-type", choices=["all", *GROUPS], default="all")
+    generate.add_argument("--instruction-language", choices=["all", "en", "cn"], default="all")
+    generate.add_argument("--width", type=int, default=512)
+    generate.add_argument("--height", type=int, default=512)
+    generate.add_argument("--num-inference-steps", type=int, default=20)
+    generate.add_argument("--guidance-scale", type=float, default=None)
+    generate.add_argument("--seed", type=int, default=42)
+    generate.add_argument("--workers", type=int, default=1)
+    generate.add_argument("--max-samples", type=int, default=None)
+    generate.add_argument("--samples-per-group", type=int, default=None)
+
+    evaluate = subparsers.add_parser("evaluate")
+    evaluate.add_argument("--dataset-ref", type=str, default="stepfun-ai/GEdit-Bench")
+    evaluate.add_argument("--output-root", type=Path, required=True)
+    evaluate.add_argument("--model-name", type=str, default=None)
+    evaluate.add_argument("--save-dir", type=Path, required=True)
+    evaluate.add_argument("--task-type", choices=["all", *GROUPS], default="all")
+    evaluate.add_argument("--instruction-language", choices=["all", "en", "cn"], default="all")
+    evaluate.add_argument("--judge-base-url", type=str, required=True)
+    evaluate.add_argument("--judge-model", type=str, required=True)
+    evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
+    evaluate.add_argument("--workers", type=int, default=1)
+    evaluate.add_argument("--max-samples", type=int, default=None)
+    evaluate.add_argument("--samples-per-group", type=int, default=None)
+
+    summarize = subparsers.add_parser("summarize")
+    summarize.add_argument("--csv-path", type=Path, required=True)
+    summarize.add_argument("--language", choices=["all", "en", "cn"], default="all")
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    if args.command == "generate":
+        model_name = resolve_model_name(model_name=args.model_name, model=args.model)
+        runner = GEditBenchRunner(
+            dataset_ref=args.dataset_ref,
+            output_root=args.output_root,
+            base_url=args.base_url,
+            model=args.model,
+            api_key=args.api_key,
+            width=args.width,
+            height=args.height,
+            num_inference_steps=args.num_inference_steps,
+            guidance_scale=args.guidance_scale,
+            seed=args.seed,
+        )
+        records = runner.generate(
+            model_name=model_name,
+            task_type=args.task_type,
+            instruction_language=args.instruction_language,
+            workers=args.workers,
+            max_samples=args.max_samples,
+            samples_per_group=args.samples_per_group,
+        )
+        payload = {"records": records, "summary": summarize_generated_records(records)}
+        write_json(args.output_root / model_name / "generation_manifest.json", payload)
+        return 0
+
+    if args.command == "evaluate":
+        model_name = resolve_model_name(model_name=args.model_name, output_root=args.output_root)
+        scorer = LocalVIEScorer(
+            base_url=args.judge_base_url,
+            api_key=args.judge_api_key,
+            model=args.judge_model,
+        )
+        evaluator = GEditBenchEvaluator(dataset_ref=args.dataset_ref, output_root=args.output_root, scorer=scorer)
+        evaluator.evaluate(
+            model_name=model_name,
+            save_dir=args.save_dir,
+            task_type=args.task_type,
+            instruction_language=args.instruction_language,
+            workers=args.workers,
+            max_samples=args.max_samples,
+            samples_per_group=args.samples_per_group,
+        )
+        return 0
+
+    if args.command == "summarize":
+        with args.csv_path.open("r", encoding="utf-8", newline="") as handle:
+            rows = list(csv.DictReader(handle))
+        summary = summarize_gedit_rows_with_backbone(rows, language=args.language)
+        print(json.dumps(summary, indent=2, ensure_ascii=False))
+        return 0
+
+    parser.error(f"Unknown command: {args.command}")
+    return 1

benchmarks/accuracy/image_to_image/run_gedit_bench.py

@@ -0,0 +1,13 @@
+# ruff: noqa: E402, I001
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from benchmarks.accuracy.image_to_image.gedit_bench import main
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

benchmarks/accuracy/text_to_image/README.md

@@ -0,0 +1,103 @@
+# GEBench on vLLM-Omni
+
+This integration adapts the upstream `stepfun-ai/GEBench` scripts into
+`vllm-omni/benchmarks/accuracy/text_to_image`.
+
+Upstream mapping:
+
+- `scripts/generate.py` -> `run_gebench.py generate`
+- `scripts/evaluate.py` -> `run_gebench.py evaluate`
+- upstream prompt / judge logic -> `gbench.py`
+
+What changed:
+
+- Generation calls the local OpenAI-compatible `vllm-omni` endpoints:
+  - `/v1/images/generations` for text-only frame generation
+  - `/v1/images/edits` for image-conditioned GUI transition generation
+- Evaluation still keeps the GEBench scoring dimensions:
+  - `goal`
+  - `logic`
+  - `cons`
+  - `ui`
+  - `qual`
+- Judge calls are also routed to a local OpenAI-compatible model served by
+  `vllm-omni` instead of a remote service.
+
+Expected dataset layout:
+
+- Clone the benchmark dataset from Hugging Face into a local directory:
+
+```bash
+git clone https://huggingface.co/datasets/stepfun-ai/GEBench /path/to/GEBench
+```
+
+Example usage:
+
+```bash
+python benchmarks/accuracy/text_to_image/run_gebench.py generate \
+  --dataset-root /path/to/GEBench \
+  --output-root benchmarks/accuracy/text_to_image/outputs \
+  --base-url http://127.0.0.1:8000 \
+  --model Tongyi-MAI/Z-Image-Turbo \
+  --data-type type3
+```
+
+```bash
+python benchmarks/accuracy/text_to_image/run_gebench.py evaluate \
+  --dataset-root /path/to/GEBench \
+  --output-root benchmarks/accuracy/text_to_image/outputs \
+  --data-type type3 \
+  --judge-base-url http://127.0.0.1:8000 \
+  --judge-model Qwen/Qwen2.5-VL-7B-Instruct \
+  --judge-api-key EMPTY
+```
+
+```bash
+python benchmarks/accuracy/text_to_image/run_gebench.py summarize \
+  --output-root benchmarks/accuracy/text_to_image/outputs
+```
+
+Example summary output:
+
+```json
+{
+  "generation": {
+    "count": 20,
+    "by_type": {
+      "type3": {"count": 10},
+      "type4": {"count": 10}
+    }
+  },
+  "evaluation": {
+    "count": 20,
+    "overall_mean": 0.52,
+    "by_type": {
+      "type3": {"count": 10, "overall_mean": 0.50, "overall_mean_100": 50.0},
+      "type4": {"count": 10, "overall_mean": 0.54, "overall_mean_100": 54.0}
+    }
+  }
+}
+```
+
+Example generated images to inspect:
+
+- `benchmarks/accuracy/text_to_image/outputs/03_trajectory_text_fictionalapp/<lang_device>/<sample>/frame0.png`
+- `benchmarks/accuracy/text_to_image/outputs/03_trajectory_text_fictionalapp/<lang_device>/<sample>/frame5.png`
+- `benchmarks/accuracy/text_to_image/outputs/04_trajectory_text_realapp/<lang_device>/<sample>/frame0.png`
+- `benchmarks/accuracy/text_to_image/outputs/04_trajectory_text_realapp/<lang_device>/<sample>/frame5.png`
+
+What to expect:
+
+- `overall_mean` is normalized to `0.0 ~ 1.0`; higher is better.
+- `frame0.png` is the initial GUI frame; `frame5.png` is the final trajectory frame most often used for quick inspection.
+- For full debugging, inspect the whole `frame0.png` ... `frame5.png` sequence for one sample directory.
+
+Notes:
+
+- GEBench upstream leaves type3/type4 generation unfinished. This integration
+  fills that gap with a trajectory runner that generates `frame0.png` followed
+  by `frame1.png` ... `frame5.png`.
+- Type1/2/5 require an image-edit capable model exposed through
+  `vllm-omni serve`.
+- `summarize` will report both generated coverage and any existing evaluation
+  summary files.

benchmarks/accuracy/text_to_image/__init__.py

@@ -0,0 +1 @@
+"""GEBench integration for vLLM-Omni."""

benchmarks/accuracy/text_to_image/gbench.py

@@ -0,0 +1,927 @@
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import requests
+from PIL import Image
+
+from benchmarks.accuracy.common import (
+    VllmOmniImageClient,
+    build_openai_url,
+    ensure_dir,
+    extract_json_object,
+    find_first_image,
+    load_json,
+    pil_to_data_url,
+    save_image,
+    write_json,
+)
+
+TYPE_TO_FOLDER = {
+    "type1": "01_single_step",
+    "type2": "02_multi_step",
+    "type3": "03_trajectory_text_fictionalapp",
+    "type4": "04_trajectory_text_realapp",
+    "type5": "05_grounding_data",
+}
+SCORE_KEYS = ("goal", "logic", "cons", "ui", "qual")
+DEFAULT_SAMPLES_PER_TYPE = 10
+
+
+def _utc_timestamp() -> str:
+    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+
+
+def _write_json_with_timestamp(path: Path, payload: dict[str, Any]) -> Path:
+    write_json(path, payload)
+    timestamped_path = path.with_name(f"{path.stem}_{_utc_timestamp()}{path.suffix}")
+    write_json(timestamped_path, payload)
+    return timestamped_path
+
+
+@dataclass(frozen=True)
+class GEBenchSampleSpec:
+    sample_path: Path
+    metadata: dict[str, Any]
+    sample_name: str
+    lang_device: str
+
+
+def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
+    by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    for record in records:
+        by_type[record["data_type"]].append(record)
+
+    return {
+        "count": len(records),
+        "by_type": {
+            data_type: {
+                "count": len(rows),
+            }
+            for data_type, rows in sorted(by_type.items())
+        },
+    }
+
+
+def summarize_gebench_results(results: list[dict[str, Any]]) -> dict[str, Any]:
+    by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    for result in results:
+        by_type[result["data_type"]].append(result)
+
+    summary: dict[str, Any] = {
+        "count": len(results),
+        "overall_mean": statistics.fmean(r["overall"] for r in results) if results else 0.0,
+        "by_type": {},
+    }
+    for data_type, rows in by_type.items():
+        score_means: dict[str, float] = {}
+        all_score_keys = {key for row in rows for key in row.get("scores", {}).keys()}
+        for score_key in all_score_keys:
+            values = [row["scores"][score_key] for row in rows if score_key in row.get("scores", {})]
+            score_means[score_key] = statistics.fmean(values) if values else 0.0
+        overall_mean = statistics.fmean(row["overall"] for row in rows)
+        summary["by_type"][data_type] = {
+            "count": len(rows),
+            "overall_mean": overall_mean,
+            "overall_mean_100": overall_mean * 100.0,
+            "score_means": score_means,
+        }
+    return summary
+
+
+def select_balanced_gebench_samples(
+    sample_paths_by_type: dict[str, list[Any]],
+    *,
+    samples_per_type: int | None,
+) -> dict[str, list[Any]]:
+    if samples_per_type is None:
+        return {data_type: list(paths) for data_type, paths in sample_paths_by_type.items()}
+    return {data_type: list(paths)[:samples_per_type] for data_type, paths in sample_paths_by_type.items()}
+
+
+def collect_gebench_generation_summary(output_root: Path) -> dict[str, Any]:
+    records: list[dict[str, Any]] = []
+    for data_type, folder_name in TYPE_TO_FOLDER.items():
+        type_root = output_root / folder_name
+        if not type_root.exists():
+            continue
+        for lang_dir in sorted(path for path in type_root.iterdir() if path.is_dir()):
+            for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir()):
+                expected = sample_dir / "frame5.png" if data_type in {"type2", "type3", "type4"} else None
+                if expected is None:
+                    expected = find_first_image(sample_dir)
+                elif not expected.exists():
+                    expected = None
+                if expected is None:
+                    continue
+                records.append(
+                    {
+                        "data_type": data_type,
+                        "sample_name": f"{lang_dir.name}/{sample_dir.name}",
+                        "output_path": str(expected),
+                    }
+                )
+    return summarize_generated_records(records)
+
+
+def _normalize_score_key(key: str) -> str:
+    mapping = {
+        "goal": "goal",
+        "logic": "logic",
+        "cons": "cons",
+        "consistency": "cons",
+        "ui": "ui",
+        "qual": "qual",
+        "quality": "qual",
+    }
+    return mapping.get(key.lower(), key.lower())
+
+
+def _normalize_scores(raw_scores: dict[str, Any]) -> dict[str, int]:
+    scores: dict[str, int] = {}
+    for key, value in raw_scores.items():
+        normalized = _normalize_score_key(key)
+        if normalized not in SCORE_KEYS:
+            continue
+        scalar = value.get("s", 0) if isinstance(value, dict) else value
+        try:
+            scores[normalized] = int(scalar)
+        except (TypeError, ValueError):
+            scores[normalized] = 0
+    for key in SCORE_KEYS:
+        scores.setdefault(key, 0)
+    return scores
+
+
+def _compute_overall(scores: dict[str, int]) -> float:
+    return sum(scores.values()) / (len(SCORE_KEYS) * 5.0)
+
+
+def _iter_sample_paths(dataset_root: Path, data_type: str) -> list[Path]:
+    data_dir = dataset_root / TYPE_TO_FOLDER[data_type]
+    if not data_dir.exists():
+        raise FileNotFoundError(f"Dataset path not found: {data_dir}")
+
+    samples: list[Path] = []
+    for lang_dir in sorted(path for path in data_dir.iterdir() if path.is_dir()):
+        for child in sorted(lang_dir.iterdir()):
+            if child.is_dir():
+                samples.append(child)
+            elif child.suffix.lower() == ".json":
+                samples.append(child)
+    return samples
+
+
+def _sample_name_from_metadata(metadata: dict[str, Any], sample_path: Path, item_index: int | None = None) -> str:
+    sample_id = metadata.get("id") or metadata.get("sample_id") or metadata.get("name")
+    if sample_id:
+        return str(sample_id)
+    if item_index is not None:
+        return f"{sample_path.stem}_{item_index:04d}"
+    return sample_path.stem if sample_path.is_file() else sample_path.name
+
+
+def _expand_sample_path(sample_path: Path) -> list[GEBenchSampleSpec]:
+    if sample_path.is_dir():
+        metadata = _load_metadata(sample_path)
+        return [
+            GEBenchSampleSpec(
+                sample_path=sample_path,
+                metadata=metadata,
+                sample_name=_sample_name(sample_path),
+                lang_device=_lang_device(sample_path, metadata),
+            )
+        ]
+
+    payload = load_json(sample_path)
+    if isinstance(payload, dict):
+        return [
+            GEBenchSampleSpec(
+                sample_path=sample_path,
+                metadata=payload,
+                sample_name=_sample_name_from_metadata(payload, sample_path),
+                lang_device=_lang_device(sample_path, payload),
+            )
+        ]
+
+    if isinstance(payload, list):
+        specs: list[GEBenchSampleSpec] = []
+        for item_index, item in enumerate(payload):
+            if not isinstance(item, dict):
+                continue
+            specs.append(
+                GEBenchSampleSpec(
+                    sample_path=sample_path,
+                    metadata=item,
+                    sample_name=_sample_name_from_metadata(item, sample_path, item_index),
+                    lang_device=_lang_device(sample_path, item),
+                )
+            )
+        return specs
+
+    raise TypeError(f"Unsupported metadata payload type for sample {sample_path}: {type(payload)!r}")
+
+
+def _iter_sample_specs(dataset_root: Path, data_type: str) -> list[GEBenchSampleSpec]:
+    sample_specs: list[GEBenchSampleSpec] = []
+    for sample_path in _iter_sample_paths(dataset_root, data_type):
+        sample_specs.extend(_expand_sample_path(sample_path))
+    return sample_specs
+
+
+def _load_metadata(sample_path: Path) -> dict[str, Any]:
+    if sample_path.is_file():
+        return load_json(sample_path)
+    for candidate in ("meta_data.json", "metadata.json"):
+        meta_path = sample_path / candidate
+        if meta_path.exists():
+            return load_json(meta_path)
+    raise FileNotFoundError(f"Metadata not found for sample: {sample_path}")
+
+
+def _sample_name(sample_path: Path) -> str:
+    return sample_path.stem if sample_path.is_file() else sample_path.name
+
+
+def _lang_device(sample_path: Path, metadata: dict[str, Any]) -> str:
+    return str(metadata.get("lang_device") or sample_path.parent.name)
+
+
+def _resolve_referenced_image(
+    *,
+    metadata: dict[str, Any],
+    sample_path: Path,
+    dataset_root: Path,
+    data_type: str,
+) -> Image.Image | None:
+    for key in ("image", "input_image", "initial_image", "reference_image"):
+        image_ref = metadata.get(key)
+        if not image_ref:
+            continue
+        candidate = dataset_root / TYPE_TO_FOLDER[data_type] / str(image_ref)
+        if candidate.exists():
+            image = Image.open(candidate)
+            image.load()
+            return image.convert("RGB")
+    if sample_path.is_dir():
+        local_image = find_first_image(sample_path)
+        if local_image:
+            image = Image.open(local_image)
+            image.load()
+            return image.convert("RGB")
+    return None
+
+
+def _trajectory_steps(metadata: dict[str, Any]) -> list[dict[str, Any]]:
+    for key in ("trajectory", "steps", "frames"):
+        value = metadata.get(key)
+        if isinstance(value, list):
+            return [step for step in value if isinstance(step, dict)]
+    extracted: list[dict[str, Any]] = []
+    for index in range(1, 6):
+        value = metadata.get(f"step{index}") or metadata.get(str(index))
+        if isinstance(value, dict):
+            extracted.append(value)
+    return extracted
+
+
+def _text_or_default(value: Any, default: str = "") -> str:
+    return str(value).strip() if value is not None else default
+
+
+def _type1_prompt(metadata: dict[str, Any]) -> str:
+    caption = _text_or_default(metadata.get("caption") or metadata.get("instruction"), "Transform the reference GUI.")
+    return (
+        "Using the reference GUI screenshot, generate the next GUI state after the requested interaction.\n\n"
+        f"Requested change:\n{caption}\n\n"
+        "Requirements:\n"
+        "- Preserve layout, visual identity, and unrelated regions.\n"
+        "- Only apply the requested state change.\n"
+        "- Keep all text and controls readable.\n"
+    )
+
+
+def _type2_prompt(goal: str, step_num: int) -> str:
+    return (
+        "Generate the next GUI state for a multi-step task.\n\n"
+        f"Overall goal: {goal}\n"
+        f"Current progress step: {step_num}/5\n\n"
+        "Requirements:\n"
+        "- The change should be incremental and plausible.\n"
+        "- Preserve layout and visual identity.\n"
+        "- Make text/buttons readable.\n"
+    )
+
+
+def _type34_initial_prompt(metadata: dict[str, Any], first_step: dict[str, Any]) -> str:
+    app_name = _text_or_default(metadata.get("app_name"), "App")
+    final_goal = _text_or_default(metadata.get("final_goal") or metadata.get("instruction"), "Complete the task.")
+    visual_description = _text_or_default(
+        metadata.get("visual_description") or first_step.get("visual_description") or first_step.get("description"),
+        "A clean product-quality app home screen.",
+    )
+    return (
+        "Generate the first GUI frame for a task trajectory.\n\n"
+        f"App name: {app_name}\n"
+        f"Final goal: {final_goal}\n"
+        f"Visual description:\n{visual_description}\n\n"
+        "Requirements:\n"
+        "- Generate a production-looking UI screenshot only.\n"
+        "- Keep the layout coherent and readable.\n"
+    )
+
+
+def _type34_next_prompt(step_num: int, step_info: dict[str, Any]) -> str:
+    action = _text_or_default(step_info.get("action") or step_info.get("instruction"), "Continue the task.")
+    visual_description = _text_or_default(
+        step_info.get("visual_description") or step_info.get("description"),
+        "Reflect the expected next GUI state.",
+    )
+    return (
+        "Using the previous frame as reference, generate the next GUI frame.\n\n"
+        f"Step {step_num} action: {action}\n"
+        f"Expected visual state:\n{visual_description}\n\n"
+        "Requirements:\n"
+        "- Only change UI regions affected by this action.\n"
+        "- Preserve persistent bars, layout, and style.\n"
+        "- Keep text and icons readable.\n"
+    )
+
+
+def _type5_prompt(metadata: dict[str, Any]) -> str:
+    grounding = metadata.get("grounding") or {}
+    explanation = _text_or_default(
+        metadata.get("grounding_explanation") or grounding.get("effect") or grounding.get("description"),
+        "Predict the immediate GUI reaction to the indicated target.",
+    )
+    return (
+        "Using the reference GUI screenshot, predict the immediate GUI state after the grounded interaction.\n\n"
+        f"Expected effect: {explanation}\n"
+        f"Grounding metadata: {json.dumps(grounding, ensure_ascii=False)}\n\n"
+        "Requirements:\n"
+        "- Apply only the interaction-triggered change.\n"
+        "- Preserve unrelated regions.\n"
+        "- Keep the UI realistic and readable.\n"
+    )
+
+
+def _make_storyboard_image(
+    frames: list[Image.Image],
+    *,
+    columns: int = 3,
+    background_color: tuple[int, int, int] = (255, 255, 255),
+) -> Image.Image:
+    if not frames:
+        raise ValueError("Expected at least one frame to build a storyboard image.")
+
+    normalized = [frame.convert("RGB") for frame in frames]
+    frame_width = max(frame.width for frame in normalized)
+    frame_height = max(frame.height for frame in normalized)
+    rows = (len(normalized) + columns - 1) // columns
+    storyboard = Image.new("RGB", (frame_width * columns, frame_height * rows), color=background_color)
+
+    for index, frame in enumerate(normalized):
+        x_offset = (index % columns) * frame_width
+        y_offset = (index // columns) * frame_height
+        if frame.size != (frame_width, frame_height):
+            frame = frame.resize((frame_width, frame_height))
+        storyboard.paste(frame, (x_offset, y_offset))
+    return storyboard
+
+
+def _trajectory_judge_payload(frames: list[Image.Image]) -> tuple[str, list[Image.Image]]:
+    storyboard = _make_storyboard_image(frames, columns=3)
+    prompt_suffix = (
+        "The attached image is a storyboard containing six frames arranged left-to-right, "
+        "top-to-bottom as frame0, frame1, frame2, frame3, frame4, frame5."
+    )
+    return prompt_suffix, [storyboard]
+
+
+class LocalJudgeClient:
+    def __init__(self, base_url: str, api_key: str, model: str, timeout: int = 600):
+        self.base_url = base_url.rstrip("/")
+        self.api_key = api_key
+        self.model = model
+        self.timeout = timeout
+
+    def _build_scoring_prompt(self, task_prompt: str) -> str:
+        return (
+            "You are an expert evaluator for GUI image editing and GUI trajectory generation.\n"
+            "Evaluate whether the generated image(s) satisfy the task.\n\n"
+            "Score these five dimensions from 0 to 5:\n"
+            "- goal: whether the user goal is completed correctly\n"
+            "- logic: whether the transition/state change is logically correct\n"
+            "- cons: whether unrelated regions remain consistent\n"
+            "- ui: whether the UI layout/components remain realistic and coherent\n"
+            "- qual: whether the images are visually clear and artifact-free\n\n"
+            "Return JSON only. Do not add any prose outside JSON.\n"
+            "Use exactly this schema:\n"
+            "{\n"
+            '  "goal": 0,\n'
+            '  "logic": 0,\n'
+            '  "cons": 0,\n'
+            '  "ui": 0,\n'
+            '  "qual": 0,\n'
+            '  "reasoning": "short explanation"\n'
+            "}\n\n"
+            "Scoring task:\n"
+            f"{task_prompt}"
+        )
+
+    def _request_text(self, prompt: str, images: list[Image.Image]) -> str:
+        content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
+        for image in images:
+            content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
+
+        response = requests.post(
+            build_openai_url(self.base_url, "/chat/completions"),
+            json={
+                "model": self.model,
+                "messages": [{"role": "user", "content": content}],
+                "temperature": 0,
+            },
+            headers={
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json",
+            },
+            timeout=self.timeout,
+        )
+        response.raise_for_status()
+        message_content = response.json()["choices"][0]["message"]["content"]
+        if isinstance(message_content, list):
+            return "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
+        return str(message_content)
+
+    def evaluate(self, *, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
+        primary_prompt = self._build_scoring_prompt(prompt)
+        text = self._request_text(primary_prompt, images)
+        try:
+            return extract_json_object(text)
+        except ValueError:
+            retry_prompt = (
+                self._build_scoring_prompt(prompt) + "\n\nYour previous response was not valid JSON. "
+                "Return only the JSON object with integer scores."
+            )
+            retry_text = self._request_text(retry_prompt, images)
+            try:
+                return extract_json_object(retry_text)
+            except ValueError:
+                return {
+                    "goal": 0,
+                    "logic": 0,
+                    "cons": 0,
+                    "ui": 0,
+                    "qual": 0,
+                    "reasoning": retry_text.strip() or text.strip() or "Judge response was not valid JSON.",
+                }
+
+
+class GEBenchRunner:
+    def __init__(
+        self,
+        *,
+        dataset_root: Path,
+        output_root: Path,
+        base_url: str,
+        model: str,
+        api_key: str = "EMPTY",
+        width: int = 768,
+        height: int = 576,
+        num_inference_steps: int = 8,
+        output_compression: int | None = 98,
+        guidance_scale: float | None = None,
+        seed: int | None = 42,
+    ):
+        self.dataset_root = dataset_root
+        self.output_root = output_root
+        self.model = model
+        self.width = width
+        self.height = height
+        self.num_inference_steps = num_inference_steps
+        self.output_compression = output_compression
+        self.guidance_scale = guidance_scale
+        self.seed = seed
+        self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
+
+    def generate(
+        self,
+        *,
+        data_type: str,
+        workers: int = 1,
+        max_samples: int | None = None,
+        samples_per_type: int | None = None,
+    ) -> list[dict[str, Any]]:
+        sample_specs = select_balanced_gebench_samples(
+            {data_type: _iter_sample_specs(self.dataset_root, data_type)},
+            samples_per_type=samples_per_type,
+        )[data_type]
+        if max_samples is not None:
+            sample_specs = sample_specs[:max_samples]
+
+        results: list[dict[str, Any]] = []
+        if workers <= 1:
+            for sample_spec in sample_specs:
+                result = self._generate_one(data_type, sample_spec)
+                if result:
+                    results.append(result)
+            return results
+
+        with ThreadPoolExecutor(max_workers=workers) as executor:
+            futures = [executor.submit(self._generate_one, data_type, sample_spec) for sample_spec in sample_specs]
+            for future in as_completed(futures):
+                result = future.result()
+                if result:
+                    results.append(result)
+        return results
+
+    def _generate_one(self, data_type: str, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
+        sample_path = sample_spec.sample_path
+        metadata = sample_spec.metadata
+        lang_device = sample_spec.lang_device
+        sample_name = sample_spec.sample_name
+        output_dir = ensure_dir(self.output_root / TYPE_TO_FOLDER[data_type] / lang_device / sample_name)
+
+        if data_type == "type1":
+            output_path = output_dir / "generated.png"
+            if output_path.exists():
+                return {
+                    "data_type": data_type,
+                    "sample_name": f"{lang_device}/{sample_name}",
+                    "output_path": str(output_path),
+                }
+            source = _resolve_referenced_image(
+                metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
+            )
+            if source is None:
+                return None
+            generated = self.client.generate_image_edit(
+                model=self.model,
+                prompt=_type1_prompt(metadata),
+                images=source,
+                width=self.width,
+                height=self.height,
+                num_inference_steps=self.num_inference_steps,
+                output_compression=self.output_compression,
+                guidance_scale=self.guidance_scale,
+                seed=self.seed,
+            )
+            save_image(output_path, generated)
+            return {
+                "data_type": data_type,
+                "sample_name": f"{lang_device}/{sample_name}",
+                "output_path": str(output_path),
+            }
+
+        if data_type == "type2":
+            goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
+            source = _resolve_referenced_image(
+                metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
+            )
+            if source is None:
+                return None
+            frame0_path = output_dir / "frame0.png"
+            if not frame0_path.exists():
+                save_image(frame0_path, source)
+            previous = source
+            for step_num in range(1, 6):
+                frame_path = output_dir / f"frame{step_num}.png"
+                if frame_path.exists():
+                    previous = Image.open(frame_path).convert("RGB")
+                    continue
+                generated = self.client.generate_image_edit(
+                    model=self.model,
+                    prompt=_type2_prompt(goal, step_num),
+                    images=previous,
+                    width=self.width,
+                    height=self.height,
+                    num_inference_steps=self.num_inference_steps,
+                    output_compression=self.output_compression,
+                    guidance_scale=self.guidance_scale,
+                    seed=self.seed,
+                )
+                save_image(frame_path, generated)
+                previous = generated
+            output_path = output_dir / "frame5.png"
+            return {
+                "data_type": data_type,
+                "sample_name": f"{lang_device}/{sample_name}",
+                "output_path": str(output_path),
+            }
+
+        if data_type in {"type3", "type4"}:
+            steps = _trajectory_steps(metadata)
+            frame0_path = output_dir / "frame0.png"
+            if frame0_path.exists():
+                previous = Image.open(frame0_path).convert("RGB")
+            else:
+                previous = self.client.generate_text_to_image(
+                    model=self.model,
+                    prompt=_type34_initial_prompt(metadata, steps[0] if steps else {}),
+                    width=self.width,
+                    height=self.height,
+                    num_inference_steps=self.num_inference_steps,
+                    output_compression=self.output_compression,
+                    guidance_scale=self.guidance_scale,
+                    seed=self.seed,
+                )
+                save_image(frame0_path, previous)
+
+            for step_num in range(1, 6):
+                frame_path = output_dir / f"frame{step_num}.png"
+                if frame_path.exists():
+                    previous = Image.open(frame_path).convert("RGB")
+                    continue
+                step_info = steps[step_num - 1] if step_num - 1 < len(steps) else {}
+                generated = self.client.generate_image_edit(
+                    model=self.model,
+                    prompt=_type34_next_prompt(step_num, step_info),
+                    images=previous,
+                    width=self.width,
+                    height=self.height,
+                    num_inference_steps=self.num_inference_steps,
+                    output_compression=self.output_compression,
+                    guidance_scale=self.guidance_scale,
+                    seed=self.seed,
+                )
+                save_image(frame_path, generated)
+                previous = generated
+            output_path = output_dir / "frame5.png"
+            return {
+                "data_type": data_type,
+                "sample_name": f"{lang_device}/{sample_name}",
+                "output_path": str(output_path),
+            }
+
+        if data_type == "type5":
+            output_path = output_dir / "generated.png"
+            if output_path.exists():
+                return {
+                    "data_type": data_type,
+                    "sample_name": f"{lang_device}/{sample_name}",
+                    "output_path": str(output_path),
+                }
+            source = _resolve_referenced_image(
+                metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
+            )
+            if source is None:
+                return None
+            generated = self.client.generate_image_edit(
+                model=self.model,
+                prompt=_type5_prompt(metadata),
+                images=source,
+                width=self.width,
+                height=self.height,
+                num_inference_steps=self.num_inference_steps,
+                output_compression=self.output_compression,
+                guidance_scale=self.guidance_scale,
+                seed=self.seed,
+            )
+            save_image(output_path, generated)
+            return {
+                "data_type": data_type,
+                "sample_name": f"{lang_device}/{sample_name}",
+                "output_path": str(output_path),
+            }
+
+        raise ValueError(f"Unsupported data type: {data_type}")
+
+
+class GEBenchEvaluator:
+    def __init__(self, *, dataset_root: Path, output_root: Path, judge: LocalJudgeClient):
+        self.dataset_root = dataset_root
+        self.output_root = output_root
+        self.judge = judge
+
+    def evaluate(
+        self,
+        *,
+        data_type: str,
+        workers: int = 1,
+        max_samples: int | None = None,
+        samples_per_type: int | None = None,
+    ) -> dict[str, Any]:
+        output_type_dir = self.output_root / TYPE_TO_FOLDER[data_type]
+        sample_specs_by_name = {
+            (spec.lang_device, spec.sample_name): spec for spec in _iter_sample_specs(self.dataset_root, data_type)
+        }
+        if not output_type_dir.exists():
+            payload = {"data_type": data_type, "results": [], "summary": summarize_gebench_results([])}
+            write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
+            return payload
+        sample_dirs = [
+            sample_dir
+            for lang_dir in sorted(path for path in output_type_dir.iterdir() if path.is_dir())
+            for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir())
+            if (lang_dir.name, sample_dir.name) in sample_specs_by_name
+        ]
+        sample_dirs = select_balanced_gebench_samples(
+            {data_type: sample_dirs},
+            samples_per_type=samples_per_type,
+        )[data_type]
+        if max_samples is not None:
+            sample_dirs = sample_dirs[:max_samples]
+        results: list[dict[str, Any]] = []
+        if workers <= 1:
+            for sample_dir in sample_dirs:
+                result = self._evaluate_one(
+                    data_type,
+                    sample_dir,
+                    sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
+                )
+                if result:
+                    results.append(result)
+        else:
+            with ThreadPoolExecutor(max_workers=workers) as executor:
+                futures = [
+                    executor.submit(
+                        self._evaluate_one,
+                        data_type,
+                        sample_dir,
+                        sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
+                    )
+                    for sample_dir in sample_dirs
+                ]
+                for future in as_completed(futures):
+                    result = future.result()
+                    if result:
+                        results.append(result)
+
+        payload = {"data_type": data_type, "results": results, "summary": summarize_gebench_results(results)}
+        write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
+        return payload
+
+    def _evaluate_one(self, data_type: str, sample_dir: Path, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
+        lang_device = sample_dir.parent.name
+        sample_name = sample_dir.name
+        dataset_sample = sample_spec.sample_path
+        metadata = sample_spec.metadata
+
+        if data_type == "type1":
+            source = _resolve_referenced_image(
+                metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
+            )
+            generated_path = find_first_image(sample_dir)
+            if source is None or generated_path is None:
+                return None
+            generated = Image.open(generated_path).convert("RGB")
+            raw_scores = self.judge.evaluate(prompt=_type1_prompt(metadata), images=[source, generated])
+        elif data_type == "type2":
+            frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
+            goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
+            prompt_suffix, judge_images = _trajectory_judge_payload(frames)
+            raw_scores = self.judge.evaluate(
+                prompt=f"Evaluate a six-frame GUI trajectory.\nTask: {goal}\n{prompt_suffix}",
+                images=judge_images,
+            )
+        elif data_type in {"type3", "type4"}:
+            frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
+            instruction = _text_or_default(metadata.get("instruction") or metadata.get("caption"), "Complete the task.")
+            prompt_suffix, judge_images = _trajectory_judge_payload(frames)
+            raw_scores = self.judge.evaluate(
+                prompt=f"Evaluate a six-frame GUI trajectory.\nInstruction: {instruction}\n{prompt_suffix}",
+                images=judge_images,
+            )
+        elif data_type == "type5":
+            source = _resolve_referenced_image(
+                metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
+            )
+            generated_path = find_first_image(sample_dir)
+            if source is None or generated_path is None:
+                return None
+            generated = Image.open(generated_path).convert("RGB")
+            raw_scores = self.judge.evaluate(prompt=_type5_prompt(metadata), images=[source, generated])
+        else:
+            raise ValueError(f"Unsupported data type: {data_type}")
+
+        scores = _normalize_scores(raw_scores)
+        return {
+            "sample_name": f"{lang_device}/{sample_name}",
+            "data_type": data_type,
+            "scores": scores,
+            "overall": _compute_overall(scores),
+            "raw_scores": raw_scores,
+        }
+
+
+def _data_types_arg(value: str) -> list[str]:
+    return list(TYPE_TO_FOLDER.keys()) if value == "all" else [value]
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Run local GEBench generation and scoring against vLLM-Omni.")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    generate = subparsers.add_parser("generate")
+    generate.add_argument("--dataset-root", type=Path, required=True)
+    generate.add_argument("--output-root", type=Path, required=True)
+    generate.add_argument("--base-url", type=str, required=True)
+    generate.add_argument("--model", type=str, required=True)
+    generate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
+    generate.add_argument("--api-key", type=str, default="EMPTY")
+    generate.add_argument("--width", type=int, default=768)
+    generate.add_argument("--height", type=int, default=576)
+    generate.add_argument("--num-inference-steps", type=int, default=8)
+    generate.add_argument("--output-compression", type=int, default=98)
+    generate.add_argument("--guidance-scale", type=float, default=None)
+    generate.add_argument("--seed", type=int, default=42)
+    generate.add_argument("--workers", type=int, default=1)
+    generate.add_argument("--max-samples", type=int, default=None)
+    generate.add_argument("--samples-per-type", type=int, default=None)
+
+    evaluate = subparsers.add_parser("evaluate")
+    evaluate.add_argument("--dataset-root", type=Path, required=True)
+    evaluate.add_argument("--output-root", type=Path, required=True)
+    evaluate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
+    evaluate.add_argument("--judge-base-url", type=str, required=True)
+    evaluate.add_argument("--judge-model", type=str, required=True)
+    evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
+    evaluate.add_argument("--workers", type=int, default=1)
+    evaluate.add_argument("--max-samples", type=int, default=None)
+    evaluate.add_argument("--samples-per-type", type=int, default=None)
+
+    summarize = subparsers.add_parser("summarize")
+    summarize.add_argument("--output-root", type=Path, required=True)
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    if args.command == "generate":
+        runner = GEBenchRunner(
+            dataset_root=args.dataset_root,
+            output_root=args.output_root,
+            base_url=args.base_url,
+            model=args.model,
+            api_key=args.api_key,
+            width=args.width,
+            height=args.height,
+            num_inference_steps=args.num_inference_steps,
+            output_compression=args.output_compression,
+            guidance_scale=args.guidance_scale,
+            seed=args.seed,
+        )
+        records: list[dict[str, Any]] = []
+        for data_type in _data_types_arg(args.data_type):
+            records.extend(
+                runner.generate(
+                    data_type=data_type,
+                    workers=args.workers,
+                    max_samples=args.max_samples,
+                    samples_per_type=args.samples_per_type,
+                )
+            )
+        payload = {"records": records, "summary": summarize_generated_records(records)}
+        write_json(args.output_root / "generation_manifest.json", payload)
+        return 0
+
+    if args.command == "evaluate":
+        judge = LocalJudgeClient(
+            base_url=args.judge_base_url,
+            api_key=args.judge_api_key,
+            model=args.judge_model,
+        )
+        evaluator = GEBenchEvaluator(dataset_root=args.dataset_root, output_root=args.output_root, judge=judge)
+        combined_results: list[dict[str, Any]] = []
+        for data_type in _data_types_arg(args.data_type):
+            payload = evaluator.evaluate(
+                data_type=data_type,
+                workers=args.workers,
+                max_samples=args.max_samples,
+                samples_per_type=args.samples_per_type,
+            )
+            combined_results.extend(payload["results"])
+        _write_json_with_timestamp(
+            args.output_root / "evaluations" / "summary.json",
+            {"summary": summarize_gebench_results(combined_results)},
+        )
+        return 0
+
+    if args.command == "summarize":
+        generation_summary = collect_gebench_generation_summary(args.output_root)
+        evaluation_dir = args.output_root / "evaluations"
+        result_records: list[dict[str, Any]] = []
+        if evaluation_dir.exists():
+            for file_path in sorted(evaluation_dir.glob("type*.json")):
+                payload = load_json(file_path)
+                result_records.extend(payload.get("results", []))
+        payload: dict[str, Any] = {"generation": generation_summary}
+        if result_records:
+            payload["evaluation"] = summarize_gebench_results(result_records)
+        _write_json_with_timestamp(args.output_root / "summary.json", payload)
+        print(json.dumps(payload, indent=2, ensure_ascii=False))
+        return 0
+
+    parser.error(f"Unknown command: {args.command}")
+    return 1

benchmarks/accuracy/text_to_image/run_gebench.py

@@ -0,0 +1,13 @@
+# ruff: noqa: E402, I001
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from benchmarks.accuracy.text_to_image.gbench import main
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

benchmarks/diffusion/backends.py

@@ -1,3 +1,4 @@
+import asyncio
 import base64
 import mimetypes
 import os
@@ -35,6 +36,7 @@ class RequestFuncOutput:
     error: str = ""
     start_time: float = 0.0
     response_body: dict[str, Any] = field(default_factory=dict)
+    stage_durations: dict[str, float] = field(default_factory=dict)
     peak_memory_mb: float = 0.0
     slo_achieved: bool | None = None
 
@@ -55,6 +57,7 @@ async def async_request_chat_completions(
     input: RequestFuncInput,
     session: aiohttp.ClientSession,
     pbar: tqdm | None = None,
+    enable_diffusion_pipeline_profiler: bool = False,
 ) -> RequestFuncOutput:
     output = RequestFuncOutput()
     output.start_time = time.perf_counter()
@@ -106,8 +109,19 @@ async def async_request_chat_completions(
                 resp_json = await response.json()
                 output.response_body = resp_json
                 output.success = True
-                if "peak_memory_mb" in resp_json:
-                    output.peak_memory_mb = resp_json["peak_memory_mb"]
+                try:
+                    choices = resp_json.get("choices", [])
+                    if choices and isinstance(choices, list):
+                        msg = choices[0].get("message", {})
+                        if isinstance(msg, dict):
+                            content = msg.get("content", [])
+                            if content and isinstance(content, list) and len(content) > 0:
+                                first_item = content[0]
+                                if isinstance(first_item, dict):
+                                    output.stage_durations = first_item.get("stage_durations") or {}
+                                    output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
+                except (IndexError, TypeError, AttributeError):
+                    pass
             else:
                 output.error = f"HTTP {response.status}: {await response.text()}"
                 output.success = False
@@ -192,7 +206,140 @@ async def async_request_openai_images(
     return output
 
 
+async def async_request_v1_videos(
+    input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput:
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    files = dict(input.extra_body)
+    if input.prompt:
+        files.setdefault("prompt", input.prompt)
+    if input.width and input.height:
+        files.setdefault("height", input.height)
+        files.setdefault("width", input.width)
+    if input.num_frames:
+        files.setdefault("num_frames", input.num_frames)
+    if input.num_inference_steps:
+        files.setdefault("num_inference_steps", input.num_inference_steps)
+    if input.seed is not None:
+        files.setdefault("seed", input.seed)
+    if input.fps:
+        files.setdefault("fps", input.fps)
+
+    form = aiohttp.FormData()
+    for k, v in files.items():
+        form.add_field(k, str(v))
+
+    image_file = None
+    if input.image_paths and len(input.image_paths) > 0:
+        image_path = input.image_paths[0]
+        image_file = open(image_path, "rb")
+        form.add_field(
+            "input_reference",
+            image_file,
+            filename=os.path.basename(image_path),
+            content_type="application/octet-stream",
+        )
+
+    job_id = None
+    job_status = None
+    poll_json = {}
+    resp_json = {}
+
+    try:
+        # invoke a post request (POST /v1/videos)
+        async with session.post(input.api_url, data=form) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                job_id = resp_json.get("id")
+                job_status = resp_json.get("status")
+                if not job_id or not job_status:
+                    output.error = "API response missing job 'id' or 'status' field."
+                    output.success = False
+                    return output
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+                return output
+
+        # invoke a poll request (GET /v1/videos/{video_id})
+        poll_interval = 2.0  # Unit(s)
+        timeout_seconds = 600.0
+        deadline = time.perf_counter() + timeout_seconds
+        job_url = f"{input.api_url}/{job_id}"
+
+        while job_status not in {"completed", "failed"}:
+            await asyncio.sleep(poll_interval)
+
+            async with session.get(job_url) as poll_response:
+                if poll_response.status != 200:
+                    output.error = f"Polling failed HTTP {poll_response.status}: {await poll_response.text()}"
+                    output.success = False
+                    return output
+
+                poll_json = await poll_response.json()
+                job_status = poll_json.get("status")
+
+                if time.perf_counter() >= deadline:
+                    output.error = f"Timed out waiting for video job {job_id} to complete."
+                    output.success = False
+                    return output
+
+        if job_status == "failed":
+            output.error = f"Video job failed: {poll_json}"
+            output.success = False
+            return output
+
+        # invoke a get request (GET /v1/videos/{video_id}/content)
+        content_url = f"{job_url}/content"
+        async with session.get(content_url) as content_response:
+            if content_response.status != 200:
+                output.error = (
+                    f"Content retrieval failed HTTP {content_response.status}: {await content_response.text()}"
+                )
+                output.success = False
+                return output
+
+            video_bytes = await content_response.read()
+            output.response_body = video_bytes
+            output.success = True
+            if "peak_memory_mb" in poll_json:
+                output.peak_memory_mb = poll_json["peak_memory_mb"]
+            elif "peak_memory_mb" in resp_json:
+                output.peak_memory_mb = resp_json["peak_memory_mb"]
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+    finally:
+        if image_file is not None:
+            image_file.close()
+
+        if job_id is not None:
+            try:
+                async with session.delete(f"{input.api_url}/{job_id}") as _:
+                    pass
+            except Exception as e:
+                print(f"Failed to clean up video job {job_id}: {e}")
+
+    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
+
+
 backends_function_mapping = {
-    "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
-    "openai": (async_request_openai_images, "/v1/images/generations"),
+    "2i": {
+        "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
+        "openai": (async_request_openai_images, "/v1/images/generations"),
+    },
+    "2v": {
+        "v1/videos": (async_request_v1_videos, "/v1/videos"),
+    },
 }

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -1,4 +1,4 @@
-# adapted from sglang and fastvideo
+# adapted from fastvideo
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
@@ -9,6 +9,7 @@ If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
 Supports multiple backends:
     - vllm-omni: Uses /v1/chat/completions endpoint (default)
     - openai: Uses /v1/images/generations endpoint
+    - v1/videos: Use /v1/videos endpoint
 
 Usage:
     # Video (vllm-omni backend)
@@ -28,6 +29,16 @@ Usage:
         --backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
         --height 1024 --width 1024
 
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+        --backend vllm-omni --dataset random --task t2i --num-prompts 1 \
+        --max-concurrency 1 --enable-negative-prompt \
+        --random-request-config '[
+            {"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
+            {"width":768,"height":768,"num_inference_steps":20,"weight":0.25},
+            {"width":1024,"height":1024,"num_inference_steps":25,"weight":0.45},
+            {"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
+        ]'
+
     i2i:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
         --backend vllm-omni --dataset vbench --task i2i --num-prompts 10
@@ -38,6 +49,16 @@ Usage:
         --backend openai --dataset vbench --task t2i --num-prompts 10 \
         --height 1024 --width 1024 --port 3000
 
+    # Video (v1/vedeos)
+    t2v:
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+        --backend v1/videos --dataset random --task t2v --num-prompts 1 \
+        --max-concurrency 1 --enable-negative-prompt \
+        --random-request-config '[
+            {"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1}
+        ]'
+
+
 """
 
 import argparse
@@ -45,7 +66,10 @@ import ast
 import asyncio
 import glob
 import json
+import logging
 import os
+import random
+import tempfile
 import time
 import uuid
 from abc import ABC, abstractmethod
@@ -57,8 +81,11 @@ import aiohttp
 import numpy as np
 import requests
 from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
+from PIL import Image
 from tqdm.asyncio import tqdm
 
+logger = logging.getLogger(__name__)
+
 
 class BaseDataset(ABC):
     def __init__(self, args, api_url: str, model: str):
@@ -262,6 +289,9 @@ class VBenchDataset(BaseDataset):
 
     def _resize_data(self, data: list[dict[str, Any]]) -> list[dict[str, Any]]:
         """Resize data to match num_prompts."""
+        if not data:
+            raise ValueError("No benchmark data available. Install Pillow or provide --dataset-path.")
+
         if not self.args.num_prompts:
             return data
 
@@ -485,6 +515,14 @@ class TraceDataset(BaseDataset):
         timestamp = self._coerce_float(row.get("timestamp"))
         slo_ms = self._coerce_float(row.get("slo_ms"))
         image_paths = row.get("image_paths")
+        if not image_paths:
+            single = row.get("image_path")
+            image_paths = [single] if single else None
+
+        if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i"]:
+            raise ValueError(
+                f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row."
+            )
 
         override_w = self.args.width
         override_h = self.args.height
@@ -516,26 +554,66 @@ class TraceDataset(BaseDataset):
 
 
 class RandomDataset(BaseDataset):
-    def __init__(self, args, api_url: str, model: str):
-        self.args = args
-        self.api_url = api_url
-        self.model = model
+    def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool = False):
+        super().__init__(args, api_url, model)
         self.num_prompts = args.num_prompts
+        self.enable_negative_prompt = enable_negative_prompt
+        self.random_request_config = getattr(args, "random_request_config", None)
+        if self.random_request_config:
+            self.random_request_config = json.loads(self.random_request_config)
+            self._weights = [p["weight"] for p in self.random_request_config]
+
+            self.random_request_config = [
+                {k: v for k, v in p.items() if k != "weight"} for p in self.random_request_config
+            ]
+
+            seed = getattr(args, "random_request_seed", 42)
+            self._rng = random.Random(seed)
+
+            self._sampled_requests = self._rng.choices(
+                self.random_request_config,
+                weights=self._weights,
+                k=self.num_prompts,
+            )
+        else:
+            self._sampled_requests = None
+
+        # Random image generate
+        if self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
+            img = Image.new("RGB", (512, 512), (255, 255, 255))
+
+            image_path = os.path.join(tempfile.gettempdir(), "diffusion_benchmark_random_image.png")
+            self._random_image_path = [image_path]
+            img.save(image_path)
+        else:
+            self._random_image_path = None
 
     def __len__(self) -> int:
         return self.num_prompts
 
     def __getitem__(self, idx: int) -> RequestFuncInput:
+        extra_body = {}
+        if self.enable_negative_prompt:
+            extra_body["negative_prompt"] = f"Negative prompt {idx} for benchmarking diffusion models"
+
+        params = {
+            "width": self.args.width,
+            "height": self.args.height,
+            "num_frames": self.args.num_frames,
+            "num_inference_steps": self.args.num_inference_steps,
+            "fps": self.args.fps,
+        }
+        if self._sampled_requests:
+            profile = self._sampled_requests[idx]
+            params.update(profile)
         return RequestFuncInput(
             prompt=f"Random prompt {idx} for benchmarking diffusion models",
             api_url=self.api_url,
             model=self.model,
-            width=self.args.width,
-            height=self.args.height,
-            num_frames=self.args.num_frames,
-            num_inference_steps=self.args.num_inference_steps,
             seed=self.args.seed,
-            fps=self.args.fps,
+            extra_body=extra_body,
+            image_paths=self._random_image_path,
+            **params,
         )
 
     def get_requests(self) -> list[RequestFuncInput]:
@@ -647,19 +725,19 @@ 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.
+    """Yield requests using a Poisson process 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.
+    - Otherwise, inter-arrival times follow an exponential distribution.
     """
 
     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:
+            interval_s = random.expovariate(request_rate)
             await asyncio.sleep(interval_s)
         yield req
 
@@ -678,6 +756,15 @@ def calculate_metrics(
     latencies = [o.latency for o in success_outputs]
     peak_memories = [o.peak_memory_mb for o in success_outputs if o.peak_memory_mb > 0]
 
+    # Aggregate per-stage durations across all successful requests that reported them.
+    stage_duration_lists: dict[str, list[float]] = {}
+    for o in success_outputs:
+        for stage, duration in (o.stage_durations or {}).items():
+            stage_duration_lists.setdefault(stage, []).append(duration)
+    stage_durations_mean = {s: float(np.mean(v)) for s, v in stage_duration_lists.items()}
+    stage_durations_p50 = {s: float(np.percentile(v, 50)) for s, v in stage_duration_lists.items()}
+    stage_durations_p99 = {s: float(np.percentile(v, 99)) for s, v in stage_duration_lists.items()}
+
     metrics = {
         "duration": total_duration,
         "completed_requests": num_success,
@@ -686,10 +773,14 @@ def calculate_metrics(
         "latency_mean": np.mean(latencies) if latencies else 0,
         "latency_median": np.median(latencies) if latencies else 0,
         "latency_p99": np.percentile(latencies, 99) if latencies else 0,
+        "latency_p95": np.percentile(latencies, 95) if latencies else 0,
         "latency_p50": np.percentile(latencies, 50) if latencies else 0,
         "peak_memory_mb_max": max(peak_memories) if peak_memories else 0,
         "peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0,
         "peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0,
+        "stage_durations_mean": stage_durations_mean,
+        "stage_durations_p50": stage_durations_p50,
+        "stage_durations_p99": stage_durations_p99,
     }
 
     if slo_enabled:
@@ -742,8 +833,32 @@ async def benchmark(args):
     if args.base_url is None:
         args.base_url = f"http://{args.host}:{args.port}"
 
+    VIDEO_TASKS = {"t2v", "i2v", "ti2v"}
+    IMAGE_TASKS = {"t2i", "i2i", "ti2i"}
+
+    if args.task in VIDEO_TASKS:
+        task_type = "2v"
+    elif args.task in IMAGE_TASKS:
+        task_type = "2i"
+    else:
+        raise ValueError(
+            f"Unsupported task: '{args.task}'. "
+            f"Valid video tasks: {sorted(VIDEO_TASKS)}, "
+            f"Valid image tasks: {sorted(IMAGE_TASKS)}"
+        )
+
+    valid_backends = sorted(backends_function_mapping[task_type].keys())
+
+    if args.backend not in valid_backends:
+        logger.error(
+            f"Invalid backend '{args.backend}' for task '{args.task}' (task type: '{task_type}').\n"
+            f"Valid backends for this task type: {valid_backends}\n"
+            f"Example usage: --task {args.task} --backend {valid_backends[0]}"
+        )
+        raise ValueError("Backend validation failed. See log above for valid options.")
+
     # Setup API URL and request function based on backend
-    request_func, api_url = backends_function_mapping[args.backend]
+    request_func, api_url = backends_function_mapping[task_type][args.backend]
     api_url = f"{args.base_url}{api_url}"
 
     if args.dataset == "vbench":
@@ -751,7 +866,7 @@ async def benchmark(args):
     elif args.dataset == "trace":
         dataset = TraceDataset(args, api_url, args.model)
     elif args.dataset == "random":
-        dataset = RandomDataset(args, api_url, args.model)
+        dataset = RandomDataset(args, api_url, args.model, args.enable_negative_prompt)
     else:
         raise ValueError(f"Unknown dataset: {args.dataset}")
 
@@ -789,6 +904,8 @@ async def benchmark(args):
                         warm_req,
                         num_inference_steps=args.warmup_num_inference_steps,
                     )
+                if args.task == "t2v":
+                    warm_req = replace(warm_req, num_frames=1)
                 warm_out = await limited_request_func(warm_req, session, None)
                 warmup_pairs.append((warm_req, warm_out))
 
@@ -847,6 +964,7 @@ async def benchmark(args):
     print("{:<40} {:<15.4f}".format("Latency Mean (s):", metrics["latency_mean"]))
     print("{:<40} {:<15.4f}".format("Latency Median (s):", metrics["latency_median"]))
     print("{:<40} {:<15.4f}".format("Latency P99 (s):", metrics["latency_p99"]))
+    print("{:<40} {:<15.4f}".format("Latency P95 (s):", metrics["latency_p95"]))
 
     if args.slo:
         print(f"{'-' * 50}")
@@ -860,6 +978,12 @@ async def benchmark(args):
         print("{:<40} {:<15.2f}".format("Peak Memory Mean (MB):", metrics["peak_memory_mb_mean"]))
         print("{:<40} {:<15.2f}".format("Peak Memory Median (MB):", metrics["peak_memory_mb_median"]))
 
+    if metrics["stage_durations_mean"]:
+        print(f"{'-' * 50}")
+        print("Stage Durations Mean (s):")
+        for stage, val in metrics["stage_durations_mean"].items():
+            print("{:<40} {:<15.4f}".format(f"  {stage}:", val))
+
     print("\n" + "=" * 60)
 
     if args.output_file:
@@ -883,7 +1007,7 @@ if __name__ == "__main__":
         "--backend",
         type=str,
         default="vllm-omni",
-        choices=["vllm-omni", "openai"],
+        choices=["vllm-omni", "openai", "v1/videos"],
         help="Backend to target the benchmark to.",
     )
     parser.add_argument(
@@ -972,6 +1096,25 @@ if __name__ == "__main__":
         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.")
+    parser.add_argument(
+        "--enable-negative-prompt",
+        action="store_true",
+        default=False,
+        help="Generate negative prompts when using the random dataset.",
+    )
+    parser.add_argument(
+        "--random-request-config",
+        type=str,
+        default=None,
+        help=(
+            "JSON string defining random request profiles. "
+            "Each profile may contain: width, height, num_inference_steps, etc. "
+            "The 'weight' field controls sampling probability (relative weight). "
+            "Example: "
+            '[{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},'
+            '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
+        ),
+    )
 
     args = parser.parse_args()
 

benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md

@@ -0,0 +1,169 @@
+# Qwen-Image Serving Performance Dashboard
+
+This document describes how to deploy and benchmark **Qwen-Image** using vLLM-Omni. It includes service startup configuration, acceleration-related options, benchmark methodology, dataset settings, and performance results.
+
+---
+
+# 1. Overview
+
+Qwen-Image is a multimodal text-to-image generation model served through the vLLM-Omni infrastructure.
+
+This document covers:
+
+* Service launch configuration (including acceleration options)
+* Benchmark scripts and usage
+* Dataset and workload settings
+* Performance measurement results
+* Reproducibility guidelines
+
+---
+
+# 2. Test Environment
+| Component | Specification |
+|------------|----------------|
+| GPU | NVIDIA A100-SXM4-80GB |
+| Diffusion Attention Backend | FlashAttention |
+
+# 3. Service Launch Configuration
+
+## 3.1 Basic Serving Command
+
+```bash
+vllm serve Qwen/Qwen-Image --omni \
+    --port 8091
+```
+
+## 3.2 Key Parameters
+
+| Parameter             | Description              |
+| --------------------- | ------------------------ |
+| `--cfg-parallel-size` | CFG parallelism degree   |
+| `--ulysses-degree`    | Ulysses parallel degree  |
+| `--vae-patch-parallel-size`    | VAE parallel degree  |
+| `--tensor-parallel-size` | Tensor parallelism degree |
+
+Record these parameters when reporting performance results.
+
+---
+
+# 4. Benchmark Script
+
+## 4.1 Benchmark Entry
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --backend vllm-omni \
+    --dataset <DATASET_NAME> \
+    --task t2i \
+    --num-prompts <N> \
+    --max-concurrency <C> \
+    --enable-negative-prompt \
+    --random-request-config <CFG>
+```
+
+## 4.2 Key Benchmark Arguments
+
+| Parameter              | Description                       |
+| ---------------------- | --------------------------------- |
+| `--backend`            | Serving backend (use `vllm-omni`) |
+| `--dataset`            | Dataset name (`random` or custom) |
+| `--task`               | Task type (e.g., `t2i`)           |
+| `--num-prompts`        | Total number of requests          |
+| `--max-concurrency`    | Client-side concurrency           |
+| `--random-request-config`| JSON string defining random request |
+
+---
+
+# 5. Dataset & Workload Settings
+
+## 5.1 Recommended Evaluation Configurations
+
+### Dataset A ( 512 Resolution)
+
+* Dataset: `random`
+* Task: t2i
+* Concurrency: 1
+* Mix Resolution
+```
+[
+ {"width":512,"height":512,"num_inference_steps":20,"weight":1}
+]
+```
+
+### Dataset B (1536 Resolution)
+
+* Dataset: `random`
+* Task: t2i
+* Concurrency: 1
+* Mix Resolution
+```
+[
+ {"width":1536,"height":1536,"num_inference_steps":35,"weight":1}
+]
+```
+
+### Dataset C (Mix Resolution)
+
+* Dataset: `random`
+* Task: t2i
+* Concurrency: 1
+* Mix Resolution
+```
+[
+ {"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
+ {"width":768,"height":768,"num_inference_steps":20,"weight":0.25},
+ {"width":1024,"height":1024,"num_inference_steps":25,"weight":0.45},
+ {"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
+]
+```
+---
+
+## 5.2 Example Benchmark Command
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --backend vllm-omni \
+    --dataset random \
+    --task t2i \
+    --num-prompts 1 \
+    --max-concurrency 1 \
+    --enable-negative-prompt \
+    --random-request-config '[
+        {"width":512,"height":512,"num_inference_steps":20,"weight":1}
+    ]'
+```
+
+---
+
+# 6. Performance Metrics
+
+The following metrics are collected during benchmarking:
+
+| Metric             | Description                   | Unit    |
+| ------------------ | ----------------------------- | ------- |
+| Mean Latency        | Mean of latency       | seconds |
+| P99 Latency        | P99 of latency             | seconds |
+
+---
+
+# 7. Performance Results
+
+| Dataset Configuration | Max Concur. | CFG | Usp | Tp | VAE Parallel | Mean Latency (s) | P99 Latency (s) |
+|-----------------------|-----|-----|-----|----|--------------|------------------|------------------|
+| Dataset A | 1 | 2 | 2 | Off | Off          | 2.2087          | 2.2087          |
+| Dataset B | 1 | 2 | 2 | Off | Off          | 19.6739          | 19.6739          |
+| Dataset C | 1 | 2 | 2 | Off | Off          | 5.67259        | 18.6234 |
+---
+
+# 8. Reproducibility Checklist
+
+To ensure consistent and comparable benchmark results:
+
+* Record GPU type
+* Record parallel configuration
+* Record benchmark parameters (resolution, concurrency, number of prompts)
+* Ensure no background workload on GPUs during testing
+
+---
+
+This document serves as the official Qwen-Image serving performance reference under vLLM-Omni.

benchmarks/diffusion/performance_dashboard/wan_2_2_serving_performance.md

@@ -0,0 +1,170 @@
+# Wan2.2 Serving Performance Dashboard
+
+This document describes how to deploy and benchmark **Wan-AI/Wan2.2-T2V-A14B-Diffusers** using vLLM-Omni. It includes service startup configuration, acceleration-related options, benchmark methodology, dataset settings, and performance results.
+
+---
+
+# 1. Overview
+
+Wan-AI/Wan2.2-T2V-A14B-Diffusers is a multimodal text-to-video generation model served through the vLLM-Omni infrastructure.
+
+This document covers:
+
+* Service launch configuration (including acceleration options)
+* Benchmark scripts and usage
+* Dataset and workload settings
+* Performance measurement results
+* Reproducibility guidelines
+
+---
+
+# 2. Test Environment
+| Component | Specification |
+|------------|----------------|
+| GPU | NVIDIA A100-SXM4-80GB |
+| Diffusion Attention Backend | FlashAttention |
+
+# 3. Service Launch Configuration
+
+## 3.1 Basic Serving Command
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni \
+    --port 8091
+```
+
+## 3.2 Key Parameters
+
+| Parameter             | Description              |
+| --------------------- | ------------------------ |
+| `--cfg-parallel-size` | CFG parallelism degree   |
+| `--ulysses-degree`    | Ulysses parallel degree  |
+| `--vae-patch-parallel-size`    | VAE parallel degree  |
+| `--tensor-parallel-size` | Tensor parallelism degree |
+| `--use-hsdp` | Enable Hybrid Sharded Data Parallel to shard model weights across GPUs |
+
+Record these parameters when reporting performance results.
+
+---
+
+# 4. Benchmark Script
+
+## 4.1 Benchmark Entry
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --backend v1/videos \
+    --dataset <DATASET_NAME> \
+    --task t2v \
+    --num-prompts <N> \
+    --max-concurrency <C> \
+    --enable-negative-prompt \
+    --random-request-config <CFG>
+```
+
+## 4.2 Key Benchmark Arguments
+
+| Parameter              | Description                       |
+| ---------------------- | --------------------------------- |
+| `--backend`            | Serving backend (use `v1/videos`) |
+| `--dataset`            | Dataset name (`random` or custom) |
+| `--task`               | Task type (e.g., `t2v`)           |
+| `--num-prompts`        | Total number of requests          |
+| `--max-concurrency`    | Client-side concurrency           |
+| `--random-request-config`| JSON string defining random request |
+
+---
+
+# 5. Dataset & Workload Settings
+
+## 5.1 Recommended Evaluation Configurations
+
+### Dataset A (480p)
+
+* Dataset: `random`
+* Task: t2v
+* Concurrency: 1
+* Mix Resolution
+```
+[
+    {"width":854,"height":480,"num_inference_steps":3,"num_frames":80,"fps":16,"weight":1}
+]
+```
+### Dataset B (720p)
+
+* Dataset: `random`
+* Task: t2v
+* Concurrency: 1
+* Mix Resolution
+```
+[
+    {"width":1280,"height":720,"num_inference_steps":6,"num_frames":80,"fps":16,"weight":1}
+]
+```
+### Dataset C (Mix Resolution)
+
+* Dataset: `random`
+* Task: t2v
+* Concurrency: 1
+* Mix Resolution
+```
+[
+ {"width":854,"height":480,"num_inference_steps":3,"num_frames":80,"fps":16,"weight":0.15},
+ {"width":854,"height":480,"num_inference_steps":4,"num_frames":120,"fps":24,"weight":0.25},
+ {"width":1280,"height":720,"num_inference_steps":6,"num_frames":80,"fps":16,"weight":0.6}
+]
+```
+---
+
+## 5.2 Example Benchmark Command
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --backend v1/videos \
+    --dataset random \
+    --task t2v \
+    --num-prompts 1 \
+    --max-concurrency 1 \
+    --enable-negative-prompt \
+    --random-request-config '[
+        {"width":854,"height":480,"num_inference_steps":18,"num_frames": 33,"fps":16",weight":1}
+    ]'
+```
+
+---
+
+# 6. Performance Metrics
+
+The following metrics are collected during benchmarking:
+
+| Metric             | Description                   | Unit    |
+| ------------------ | ----------------------------- | ------- |
+| Mean Latency        | Mean of latency       | seconds |
+| P99 Latency        | P99 of latency             | seconds |
+
+---
+
+# 7. Performance Results
+
+| Dataset Configuration | Max Concur. | CFG | Usp | Tp | Hsdp | VAE Parallel | Mean Latency (s) | P99 Latency (s) |
+|-----------------------|-----|-----|-----|-----|----|--------------|------------------|------------------|
+| Dataset A | 1 | 2 | 2 | 1 | On | 1          | 24.6766          | 24.6766          |
+| Dataset A | 1 | 2 | 2 | 1 | On | 4          | 21.6810          | 21.6810          |
+| Dataset B | 1 | 2 | 2 | 1 | On | 1          | 124.6639         | 124.6639          |
+| Dataset B | 1 | 2 | 2 | 1 | On | 4          | 117.44          | 117.44          |
+| Dataset C | 1 | 2 | 2 | 1 | On | 1          | 79.2175        | 124.2565 |
+| Dataset C | 1 | 2 | 2 | 1 | On | 4          | 74.4977        | 117.710 |
+---
+
+# 8. Reproducibility Checklist
+
+To ensure consistent and comparable benchmark results:
+
+* Record GPU type
+* Record parallel configuration
+* Record benchmark parameters (resolution, concurrency, number of prompts)
+* Ensure no background workload on GPUs during testing
+
+---
+
+This document serves as the official Wan2.2 serving performance reference under vLLM-Omni.

benchmarks/diffusion/quantization_quality.py

@@ -0,0 +1,460 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Benchmark quantization quality loss for diffusion models (image & video).
+
+Generates outputs with BF16 (baseline) and a quantized config using the same
+seed, then computes LPIPS perceptual distance between them. Results are printed
+as a Markdown table ready to paste into a PR description.
+
+Requirements:
+    pip install lpips Pillow numpy
+
+Image example (text-to-image):
+    python benchmarks/diffusion/quantization_quality.py \
+        --model Tongyi-MAI/Z-Image-Turbo \
+        --task t2i \
+        --quantization fp8 \
+        --prompts \
+            "an aerial view of a coral reef with crystal clear turquoise water" \
+            "a campfire in a dark forest with sparks rising into a starry sky" \
+            "a gourmet dessert plate with chocolate mousse and gold leaf" \
+        --height 1024 --width 1024 \
+        --num-inference-steps 50 --seed 42
+
+Video example (text-to-video):
+    python benchmarks/diffusion/quantization_quality.py \
+        --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+        --task t2v \
+        --quantization fp8 \
+        --prompts \
+            "A serene lakeside sunrise with mist over the water" \
+            "A cat walking across a wooden bridge in autumn" \
+        --height 720 --width 1280 \
+        --num-frames 81 --num-inference-steps 40 --seed 42
+
+Multiple quantization methods:
+    python benchmarks/diffusion/quantization_quality.py \
+        --model Tongyi-MAI/Z-Image-Turbo \
+        --task t2i \
+        --quantization fp8 int8 bitsandbytes \
+        --prompts "a cup of coffee on the table" \
+        --height 1024 --width 1024 \
+        --num-inference-steps 50 --seed 42
+
+Output directory structure (--output-dir, default: ./quant_bench_output):
+    quant_bench_output/
+        baseline/           # BF16 outputs
+        <method>/           # Quantized outputs per method
+        results.md          # Markdown table
+"""
+
+import argparse
+import gc
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+def compute_lpips_images(
+    baseline_images: list,
+    quantized_images: list,
+    net: str = "alex",
+) -> list[float]:
+    """Compute LPIPS between paired lists of PIL images."""
+    import lpips
+    from torchvision import transforms
+
+    loss_fn = lpips.LPIPS(net=net).eval()
+    if torch.cuda.is_available():
+        loss_fn = loss_fn.cuda()
+
+    transform = transforms.Compose(
+        [
+            transforms.Resize((256, 256)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ]
+    )
+
+    scores = []
+    for img_bl, img_qt in zip(baseline_images, quantized_images):
+        t_bl = transform(img_bl.convert("RGB")).unsqueeze(0)
+        t_qt = transform(img_qt.convert("RGB")).unsqueeze(0)
+        if torch.cuda.is_available():
+            t_bl, t_qt = t_bl.cuda(), t_qt.cuda()
+        with torch.no_grad():
+            score = loss_fn(t_bl, t_qt).item()
+        scores.append(score)
+    return scores
+
+
+def compute_lpips_video(
+    baseline_frames: np.ndarray,
+    quantized_frames: np.ndarray,
+    net: str = "alex",
+) -> float:
+    """Compute mean per-frame LPIPS for a video pair.
+
+    Args:
+        baseline_frames: (F, H, W, C) float array in [0, 1].
+        quantized_frames: same shape.
+
+    Returns:
+        Mean LPIPS across all frames.
+    """
+    import lpips
+
+    loss_fn = lpips.LPIPS(net=net).eval()
+    if torch.cuda.is_available():
+        loss_fn = loss_fn.cuda()
+
+    num_frames = min(len(baseline_frames), len(quantized_frames))
+    scores = []
+    for i in range(num_frames):
+        # Convert (H, W, C) float [0,1] -> (1, C, H, W) float [-1, 1]
+        f_bl = torch.from_numpy(baseline_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
+        f_qt = torch.from_numpy(quantized_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
+        if torch.cuda.is_available():
+            f_bl, f_qt = f_bl.cuda(), f_qt.cuda()
+        with torch.no_grad():
+            score = loss_fn(f_bl, f_qt).item()
+        scores.append(score)
+    return float(np.mean(scores))
+
+
+def _build_omni_kwargs(args, quantization=None):
+    """Build kwargs dict for Omni() constructor."""
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        tensor_parallel_size=args.tensor_parallel_size,
+    )
+    kwargs = {
+        "model": args.model,
+        "parallel_config": parallel_config,
+        "enforce_eager": args.enforce_eager,
+    }
+    if quantization:
+        kwargs["quantization_config"] = quantization
+    return kwargs
+
+
+def _generate_image(omni, args, prompt, seed):
+    """Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.platforms import current_omni_platform
+
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
+    torch.cuda.reset_peak_memory_stats()
+    start = time.perf_counter()
+    outputs = omni.generate(
+        {"prompt": prompt},
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            num_inference_steps=args.num_inference_steps,
+        ),
+    )
+    elapsed = time.perf_counter() - start
+    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+
+    first = outputs[0]
+    req_out = first.request_output[0] if hasattr(first, "request_output") else first
+    img = req_out.images[0]
+    return img, elapsed, peak_mem
+
+
+def _generate_video(omni, args, prompt, seed):
+    """Generate a video and return (np.ndarray [F,H,W,C], time_seconds, memory_gib)."""
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.outputs import OmniRequestOutput
+    from vllm_omni.platforms import current_omni_platform
+
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
+    torch.cuda.reset_peak_memory_stats()
+    start = time.perf_counter()
+    outputs = omni.generate(
+        {"prompt": prompt, "negative_prompt": ""},
+        OmniDiffusionSamplingParams(
+            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,
+        ),
+    )
+    elapsed = time.perf_counter() - start
+    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+
+    first = outputs[0]
+    if hasattr(first, "request_output") and isinstance(first.request_output, list):
+        inner = first.request_output[0]
+        if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
+            frames = inner.images[0] if inner.images else None
+        else:
+            frames = inner
+    elif hasattr(first, "images") and first.images:
+        frames = first.images
+    else:
+        raise ValueError("Could not extract video frames from output.")
+
+    if isinstance(frames, torch.Tensor):
+        video = frames.detach().cpu()
+        if video.dim() == 5:
+            video = video[0].permute(1, 2, 3, 0) if video.shape[1] in (3, 4) else video[0]
+        elif video.dim() == 4 and video.shape[0] in (3, 4):
+            video = video.permute(1, 2, 3, 0)
+        if video.is_floating_point():
+            video = video.clamp(-1, 1) * 0.5 + 0.5
+        frames_array = video.float().numpy()
+    else:
+        frames_array = np.asarray(frames)
+        if frames_array.ndim == 5:
+            frames_array = frames_array[0]
+
+    return frames_array, elapsed, peak_mem
+
+
+def _unload_omni(omni):
+    """Delete Omni instance and free GPU memory."""
+    del omni
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+
+def run_benchmark(args):
+    from vllm_omni.entrypoints.omni import Omni
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    is_video = args.task == "t2v"
+    prompts = args.prompts
+    seed = args.seed
+
+    # Determine configs to benchmark
+    configs = []  # list of (label, quantization_method)
+    for method in args.quantization:
+        configs.append((method, method))
+
+    # --- Baseline run ---
+    print("\n" + "=" * 60)
+    print("Running BF16 baseline...")
+    print("=" * 60)
+    bl_kwargs = _build_omni_kwargs(args, quantization=None)
+    omni_bl = Omni(**bl_kwargs)
+
+    baseline_outputs = {}  # prompt -> (output, time, mem)
+    for prompt in prompts:
+        print(f"  Generating: {prompt[:60]}...")
+        if is_video:
+            out, t, mem = _generate_video(omni_bl, args, prompt, seed)
+        else:
+            out, t, mem = _generate_image(omni_bl, args, prompt, seed)
+        baseline_outputs[prompt] = (out, t, mem)
+
+    bl_avg_time = np.mean([v[1] for v in baseline_outputs.values()])
+    bl_mem = baseline_outputs[prompts[0]][2]  # use first prompt's memory
+    _unload_omni(omni_bl)
+
+    # Save baseline outputs
+    bl_dir = output_dir / "baseline"
+    bl_dir.mkdir(parents=True, exist_ok=True)
+    for i, prompt in enumerate(prompts):
+        out = baseline_outputs[prompt][0]
+        if is_video:
+            try:
+                from diffusers.utils import export_to_video
+
+                frames_list = list(out) if isinstance(out, np.ndarray) and out.ndim == 4 else out
+                export_to_video(frames_list, str(bl_dir / f"prompt_{i}.mp4"), fps=args.fps)
+            except ImportError:
+                np.save(bl_dir / f"prompt_{i}.npy", out)
+        else:
+            out.save(bl_dir / f"prompt_{i}.png")
+
+    # --- Quantized runs ---
+    all_results = []  # list of dicts
+
+    for config_label, quant_method in configs:
+        print(f"\n{'=' * 60}")
+        print(f"Running: {config_label}...")
+        print("=" * 60)
+
+        qt_kwargs = _build_omni_kwargs(args, quantization=quant_method)
+        omni_qt = Omni(**qt_kwargs)
+
+        qt_outputs = {}
+        for prompt in prompts:
+            print(f"  Generating: {prompt[:60]}...")
+            if is_video:
+                out, t, mem = _generate_video(omni_qt, args, prompt, seed)
+            else:
+                out, t, mem = _generate_image(omni_qt, args, prompt, seed)
+            qt_outputs[prompt] = (out, t, mem)
+
+        qt_avg_time = np.mean([v[1] for v in qt_outputs.values()])
+        qt_mem = qt_outputs[prompts[0]][2]
+        _unload_omni(omni_qt)
+
+        # Save quantized outputs
+        qt_dir = output_dir / config_label.replace(" ", "_")
+        qt_dir.mkdir(parents=True, exist_ok=True)
+
+        # Compute LPIPS per prompt
+        per_prompt = []
+        for i, prompt in enumerate(prompts):
+            bl_out = baseline_outputs[prompt][0]
+            qt_out = qt_outputs[prompt][0]
+            if is_video:
+                lpips_score = compute_lpips_video(bl_out, qt_out, net=args.lpips_net)
+                try:
+                    from diffusers.utils import export_to_video
+
+                    frames_list = list(qt_out) if isinstance(qt_out, np.ndarray) and qt_out.ndim == 4 else qt_out
+                    export_to_video(frames_list, str(qt_dir / f"prompt_{i}.mp4"), fps=args.fps)
+                except ImportError:
+                    np.save(qt_dir / f"prompt_{i}.npy", qt_out)
+            else:
+                lpips_score = compute_lpips_images([bl_out], [qt_out], net=args.lpips_net)[0]
+                qt_out.save(qt_dir / f"prompt_{i}.png")
+            per_prompt.append({"prompt": prompt, "lpips": lpips_score})
+
+        mean_lpips = np.mean([p["lpips"] for p in per_prompt])
+        speedup = bl_avg_time / qt_avg_time if qt_avg_time > 0 else float("inf")
+        mem_reduction = (bl_mem - qt_mem) / bl_mem * 100
+
+        all_results.append(
+            {
+                "config": config_label,
+                "avg_time": qt_avg_time,
+                "speedup": speedup,
+                "memory_gib": qt_mem,
+                "mem_reduction_pct": mem_reduction,
+                "mean_lpips": mean_lpips,
+                "per_prompt": per_prompt,
+            }
+        )
+
+    # --- Print results ---
+    print("\n\n")
+    print("=" * 80)
+    print("RESULTS")
+    print("=" * 80)
+
+    # Summary table
+    lines = []
+    lines.append(f"## Quantization Quality Benchmark — {args.model.split('/')[-1]}")
+    lines.append(
+        f"Setup: {args.height}x{args.width}, {args.num_inference_steps} steps, "
+        f"seed={args.seed}, LPIPS ({args.lpips_net})"
+    )
+    if is_video:
+        lines.append(f"Video: {args.num_frames} frames")
+    lines.append("")
+    lines.append("### Summary")
+    lines.append("")
+    lines.append("| Config | Avg Time | Speedup | Memory (GiB) | Mem Reduction | Mean LPIPS |")
+    lines.append("|--------|----------|---------|--------------|---------------|------------|")
+    lines.append(f"| BF16 baseline | {bl_avg_time:.2f}s | 1.00x | {bl_mem:.2f} | — | (ref) |")
+    for r in all_results:
+        lines.append(
+            f"| {r['config']} | {r['avg_time']:.2f}s | {r['speedup']:.2f}x "
+            f"| {r['memory_gib']:.2f} | {r['mem_reduction_pct']:.0f}% "
+            f"| {r['mean_lpips']:.4f} |"
+        )
+    lines.append("")
+    lines.append("> LPIPS < 0.01 = imperceptible, > 0.1 = clearly noticeable.")
+    lines.append("")
+
+    # Per-prompt table
+    if len(prompts) > 1:
+        lines.append("### Per-Prompt LPIPS")
+        lines.append("")
+        header = "| Prompt |"
+        sep = "|--------|"
+        for r in all_results:
+            header += f" {r['config']} |"
+            sep += "--------|"
+        lines.append(header)
+        lines.append(sep)
+        for i, prompt in enumerate(prompts):
+            short = prompt[:50] + "..." if len(prompt) > 50 else prompt
+            row = f"| {short} |"
+            for r in all_results:
+                row += f" {r['per_prompt'][i]['lpips']:.4f} |"
+            lines.append(row)
+        lines.append("")
+
+    md = "\n".join(lines)
+    print(md)
+
+    # Save markdown
+    results_path = output_dir / "results.md"
+    results_path.write_text(md, encoding="utf-8")
+    print(f"\nResults saved to {results_path}")
+    print(f"Baseline outputs in {bl_dir}")
+    for r in all_results:
+        qt_dir = output_dir / r["config"].replace(" ", "_")
+        print(f"Quantized outputs in {qt_dir}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Benchmark quantization quality loss for diffusion models.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--model", required=True, help="Model name or local path.")
+    parser.add_argument(
+        "--task",
+        default="t2i",
+        choices=["t2i", "t2v"],
+        help="Task type: t2i (text-to-image) or t2v (text-to-video).",
+    )
+    parser.add_argument(
+        "--quantization",
+        nargs="+",
+        required=True,
+        help="One or more quantization methods to benchmark (e.g. fp8 int8 bitsandbytes).",
+    )
+    parser.add_argument(
+        "--prompts",
+        nargs="+",
+        default=["a cup of coffee on the table"],
+        help="One or more prompts to generate.",
+    )
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--num-inference-steps", type=int, default=50)
+    parser.add_argument("--num-frames", type=int, default=81, help="Number of video frames (t2v only).")
+    parser.add_argument("--fps", type=int, default=24, help="Video FPS for saving (t2v only).")
+    parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (used for video).")
+    parser.add_argument("--output-dir", type=str, default="./quant_bench_output", help="Directory to save outputs.")
+    parser.add_argument(
+        "--lpips-net",
+        type=str,
+        default="alex",
+        choices=["alex", "vgg", "squeeze"],
+        help="LPIPS backbone network.",
+    )
+    parser.add_argument("--ulysses-degree", type=int, default=1)
+    parser.add_argument("--ring-degree", type=int, default=1)
+    parser.add_argument("--tensor-parallel-size", type=int, default=1)
+    parser.add_argument("--enforce-eager", action="store_true")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_benchmark(args)

benchmarks/qwen3-omni/README.md

@@ -56,9 +56,9 @@ What it does:
 - Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
 - Uses `benchmarks/build_dataset/top100.txt` and writes to:
   - Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
-    - `omni_llm_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
-    - `omni_llm_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
-    - `omni_llm_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
+    - `omni_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
+    - `omni_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
+    - `omni_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
   - Outputs: `benchmarks/qwen3-omni/vllm_omni/outputs/` — ~100 text and `.wav` files.
 
 Key checks:

benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh

@@ -26,12 +26,12 @@ else
                       --log-stats \
                       --log-dir $log_dir
     echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
-    echo "  - omni_llm_pipeline_text                       run dir/base name"
-    echo "  - omni_llm_pipeline_text.orchestrator.stats.jsonl  orchestrator-stage latency stats"
-    echo "  - omni_llm_pipeline_text.overall.stats.jsonl       overall latency/TPS stats"
-    echo "  - omni_llm_pipeline_text.stage0.log                per-stage detailed logs"
-    echo "  - omni_llm_pipeline_text.stage1.log"
-    echo "  - omni_llm_pipeline_text.stage2.log"
+    echo "  - omni_pipeline_text                       run dir/base name"
+    echo "  - omni_pipeline_text.orchestrator.stats.jsonl  orchestrator-stage latency stats"
+    echo "  - omni_pipeline_text.overall.stats.jsonl       overall latency/TPS stats"
+    echo "  - omni_pipeline_text.stage0.log                per-stage detailed logs"
+    echo "  - omni_pipeline_text.stage1.log"
+    echo "  - omni_pipeline_text.stage2.log"
     echo "Key checks: overall.stats.jsonl for end-to-end latency/TPS; orchestrator.stats.jsonl for stable per-stage latency; stage*.log for errors or long tails."
     echo "  - outputs/             Generated txt and wav files, there should be 100 text and wav files generated respectively"
 fi

benchmarks/qwen3-tts/README.md

@@ -32,8 +32,11 @@ bash run_benchmark.sh --hf-only
 # Use a different model (e.g. 1.7B)
 MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
 
-# Use batch_size=4 config for higher throughput
-STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
+# Use a Voice Clone model
+MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-Base TASK_TYPE=Base bash run_benchmark.sh --async-only
+
+# Use bs16 config for higher throughput
+STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs16.yaml bash run_benchmark.sh --async-only
 
 # Custom GPU, prompt count, concurrency levels
 GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
@@ -54,7 +57,7 @@ CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
 ### 2) Run online serving benchmark
 
 ```bash
-python vllm_omni/bench_tts_serve.py \
+python benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py \
     --port 8000 \
     --num-prompts 50 \
     --max-concurrency 1 4 10 \
@@ -65,7 +68,7 @@ python vllm_omni/bench_tts_serve.py \
 ### 3) Run HuggingFace baseline
 
 ```bash
-python transformers/bench_tts_hf.py \
+python benchmarks/qwen3-tts/transformers/bench_tts_hf.py \
     --model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
     --num-prompts 50 \
     --gpu-device 0 \
@@ -75,7 +78,7 @@ python transformers/bench_tts_hf.py \
 ### 4) Generate comparison plots
 
 ```bash
-python plot_results.py \
+python benchmarks/qwen3-tts/plot_results.py \
     --results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
     --labels "vllm-omni" "hf_transformers" \
     --output results/comparison.png
@@ -83,14 +86,14 @@ python plot_results.py \
 
 ## Stage Configs
 
-| Config | Batch Size | Description |
-|--------|:----------:|-------------|
-| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lower latency) |
-| `vllm_omni/configs/qwen3_tts_bs4.yaml` | 4 | Concurrent request processing (higher throughput) |
+| Config | max_num_seqs | Description |
+|--------|:------------:|-------------|
+| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lowest latency) |
+| `vllm_omni/configs/qwen3_tts_bs16.yaml` | 16 | High-throughput concurrent processing |
 
-Both configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
+All configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
 
-The model is specified via the CLI `--model` flag (or `MODEL` env var in `run_benchmark.sh`), so the same configs work for both the 0.6B and 1.7B model variants.
+The model is specified via the CLI `--model` flag (or `MODEL` env var), so the same configs work for both the 0.6B and 1.7B model variants.
 
 ## Metrics
 

benchmarks/qwen3-tts/run_benchmark.sh

@@ -23,6 +23,9 @@
 #   # Use 1.7B model:
 #   MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
 #
+#   # Use Voice Clone model
+#   MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-Base TASK_TYPE=Base bash run_benchmark.sh --async-only
+#
 #   # Use batch_size=4 config:
 #   STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
 #
@@ -35,6 +38,7 @@
 #   GPU_MEM_TALKER   - gpu_memory_utilization for talker stage (default: 0.3)
 #   GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.2)
 #   STAGE_CONFIG     - Path to stage config YAML (default: configs/qwen3_tts_bs1.yaml)
+#   TASK_TYPE        - Task type: CustomVoice, VoiceDesign, Base (default: CustomVoice)
 
 set -euo pipefail
 
@@ -53,6 +57,7 @@ NUM_WARMUPS="${NUM_WARMUPS:-3}"
 STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/configs/qwen3_tts_bs1.yaml}"
 RESULT_DIR="${SCRIPT_DIR}/results"
 TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+TASK_TYPE="${TASK_TYPE:-CustomVoice}"
 
 # Parse args
 RUN_ASYNC=true
@@ -77,6 +82,7 @@ echo " Concurrency:  ${CONCURRENCY}"
 echo " Port:         ${PORT}"
 echo " Stage config: ${STAGE_CONFIG}"
 echo " Results:      ${RESULT_DIR}"
+echo " Task type:    ${TASK_TYPE}"
 echo "============================================================"
 
 # Prepare stage config with correct GPU device and memory settings
@@ -195,7 +201,8 @@ run_bench() {
         --max-concurrency ${conc_args} \
         --num-warmups "${NUM_WARMUPS}" \
         --config-name "${config_name}" \
-        --result-dir "${RESULT_DIR}"
+        --result-dir "${RESULT_DIR}" \
+        --task-type "${TASK_TYPE}"
 
     stop_server
 
@@ -222,7 +229,8 @@ if [ "${RUN_HF}" = true ]; then
         --num-warmups "${NUM_WARMUPS}" \
         --gpu-device "${GPU_DEVICE}" \
         --config-name "hf_transformers" \
-        --result-dir "${RESULT_DIR}"
+        --result-dir "${RESULT_DIR}" \
+        --task-type "${TASK_TYPE}"
 
     # Allow GPU memory to settle
     sleep 5

benchmarks/qwen3-tts/transformers/bench_tts_hf.py

@@ -38,6 +38,10 @@ PROMPTS = [
     "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
 ]
 
+REF_AUDIO = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
+REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+INSTRUCT = "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
+
 
 @dataclass
 class BenchmarkResult:
@@ -75,6 +79,29 @@ class BenchmarkResult:
     per_request: list = field(default_factory=list)
 
 
+def generate_audio(model, prompt: str, args):
+    if args.task_type == "Base":
+        return model.generate_voice_clone(
+            text=prompt,
+            language=args.language,
+            ref_audio=REF_AUDIO,
+            ref_text=REF_TEXT,
+        )
+
+    if args.task_type == "VoiceDesign":
+        return model.generate_voice_design(
+            text=prompt,
+            language=args.language,
+            instruct=INSTRUCT,
+        )
+
+    return model.generate_custom_voice(
+        text=prompt,
+        language=args.language,
+        speaker=args.voice,
+    )
+
+
 def run_benchmark(args):
     from qwen_tts import Qwen3TTSModel
 
@@ -95,11 +122,7 @@ def run_benchmark(args):
         print(f"Warming up with {args.num_warmups} requests...")
         for i in range(args.num_warmups):
             p = PROMPTS[i % len(PROMPTS)]
-            wavs, sr = model.generate_custom_voice(
-                text=p,
-                language=args.language,
-                speaker=args.voice,
-            )
+            wavs, sr = generate_audio(model, p, args)
         # Sync GPU
         torch.cuda.synchronize(device)
         print("Warmup done.")
@@ -124,11 +147,7 @@ def run_benchmark(args):
             torch.cuda.synchronize(device)
             st = time.perf_counter()
 
-            wavs, sr = model.generate_custom_voice(
-                text=prompt,
-                language=args.language,
-                speaker=args.voice,
-            )
+            wavs, sr = generate_audio(model, prompt, args)
 
             torch.cuda.synchronize(device)
             elapsed = time.perf_counter() - st
@@ -268,6 +287,7 @@ def parse_args():
     parser.add_argument("--gpu-device", type=int, default=0)
     parser.add_argument("--voice", type=str, default="Vivian")
     parser.add_argument("--language", type=str, default="English")
+    parser.add_argument("--task-type", type=str, default="CustomVoice", choices=["CustomVoice", "VoiceDesign", "Base"])
     parser.add_argument(
         "--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
     )

benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py

@@ -37,6 +37,9 @@ PROMPTS = [
     "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
     "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
 ]
+REF_AUDIO = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
+REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+INSTRUCT = "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
 
 
 @dataclass
@@ -93,22 +96,39 @@ def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width
     return num_samples / sample_rate
 
 
+def create_payload(
+    prompt: str, task_type: str = "CustomVoice", voice: str = "vivian", language: str = "English"
+) -> dict:
+    payload = {
+        "input": prompt,
+        "language": language,
+        "stream": True,
+        "response_format": "pcm",
+        "task_type": task_type,
+    }
+
+    if task_type == "Base":
+        payload["ref_audio"] = REF_AUDIO
+        payload["ref_text"] = REF_TEXT
+    elif task_type == "CustomVoice":
+        payload["voice"] = voice
+    elif task_type == "VoiceDesign":
+        payload["instructions"] = INSTRUCT
+
+    return payload
+
+
 async def send_tts_request(
     session: aiohttp.ClientSession,
     api_url: str,
     prompt: str,
+    task_type: str = "CustomVoice",
     voice: str = "vivian",
     language: str = "English",
     pbar: tqdm | None = None,
 ) -> RequestResult:
     """Send a streaming TTS request and measure latency metrics."""
-    payload = {
-        "input": prompt,
-        "voice": voice,
-        "language": language,
-        "stream": True,
-        "response_format": "pcm",
-    }
+    payload = create_payload(prompt, task_type, voice, language)
 
     result = RequestResult(prompt=prompt)
     st = time.perf_counter()
@@ -153,6 +173,7 @@ async def run_benchmark(
     num_prompts: int,
     max_concurrency: int,
     num_warmups: int = 3,
+    task_type: str = "CustomVoice",
     voice: str = "vivian",
     language: str = "English",
 ) -> BenchmarkResult:
@@ -175,7 +196,7 @@ async def run_benchmark(
         warmup_tasks = []
         for i in range(num_warmups):
             prompt = PROMPTS[i % len(PROMPTS)]
-            warmup_tasks.append(send_tts_request(session, api_url, prompt, voice, language))
+            warmup_tasks.append(send_tts_request(session, api_url, prompt, task_type, voice, language))
         await asyncio.gather(*warmup_tasks)
         print("  Warmup done.")
 
@@ -189,7 +210,7 @@ async def run_benchmark(
 
     async def limited_request(prompt):
         async with semaphore:
-            return await send_tts_request(session, api_url, prompt, voice, language, pbar)
+            return await send_tts_request(session, api_url, prompt, task_type, voice, language, pbar)
 
     start_time = time.perf_counter()
     tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
@@ -306,6 +327,7 @@ async def main(args):
             num_prompts=args.num_prompts,
             max_concurrency=concurrency,
             num_warmups=args.num_warmups,
+            task_type=args.task_type,
             voice=args.voice,
             language=args.language,
         )
@@ -334,6 +356,7 @@ def parse_args():
         "--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
     )
     parser.add_argument("--num-warmups", type=int, default=3)
+    parser.add_argument("--task-type", type=str, default="CustomVoice", choices=["CustomVoice", "VoiceDesign", "Base"])
     parser.add_argument("--voice", type=str, default="vivian")
     parser.add_argument("--language", type=str, default="English")
     parser.add_argument(

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml

@@ -4,19 +4,18 @@ async_chunk: true
 stage_args:
   - stage_id: 0
     stage_type: llm
+    is_comprehension: true
     runtime:
       devices: "0"
-      max_batch_size: 1
     engine_args:
+      max_num_seqs: 1
       model_stage: qwen3_tts
       model_arch: Qwen3TTSTalkerForConditionalGeneration
-      hf_overrides:
-        architectures: [Qwen3TTSTalkerForConditionalGeneration]
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       enforce_eager: false
       trust_remote_code: true
-      async_scheduling: false
+      async_scheduling: true
       enable_prefix_caching: false
       engine_output_type: latent
       gpu_memory_utilization: 0.3
@@ -39,20 +38,18 @@ stage_args:
     stage_type: llm
     runtime:
       devices: "0"
-      max_batch_size: 1
     engine_args:
+      max_num_seqs: 1
       model_stage: code2wav
       model_arch: Qwen3TTSCode2Wav
-      hf_overrides:
-        architectures: [Qwen3TTSCode2Wav]
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
-      async_scheduling: false
+      async_scheduling: true
       enable_prefix_caching: false
       engine_output_type: audio
-      gpu_memory_utilization: 0.2
+      gpu_memory_utilization: 0.3
       distributed_executor_backend: "mp"
       max_num_batched_tokens: 8192
       max_model_len: 32768

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs16.yaml

@@ -0,0 +1,94 @@
+# Qwen3-TTS max_num_seqs=16 config (streaming with async_chunk)
+# High-throughput concurrent request processing
+# 2-stage pipeline: Talker -> Code2Wav
+async_chunk: true
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    is_comprehension: true
+    runtime:
+      devices: "0"
+    engine_args:
+      max_num_seqs: 16
+      model_stage: qwen3_tts
+      model_arch: Qwen3TTSTalkerForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      enforce_eager: false
+      trust_remote_code: true
+      async_scheduling: true
+      enable_prefix_caching: false
+      engine_output_type: latent
+      gpu_memory_utilization: 0.3
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 4096
+      max_model_len: 4096
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: false
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+    engine_args:
+      max_num_seqs: 16
+      model_stage: code2wav
+      model_arch: Qwen3TTSCode2Wav
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      async_scheduling: true
+      enable_prefix_caching: false
+      engine_output_type: audio
+      gpu_memory_utilization: 0.2
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 16384
+      max_model_len: 32768
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: audio
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    tts_args:
+      max_instructions_length: 500
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.0
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 16
+
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+        codec_streaming: true
+        connector_get_sleep_s: 0.01
+        connector_get_max_wait_first_chunk: 3000
+        connector_get_max_wait: 300
+        codec_chunk_frames: 25
+        codec_left_context_frames: 25
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml

@@ -1,23 +1,22 @@
 # Qwen3-TTS batch_size=4 config (streaming with async_chunk)
-# Enables concurrent request processing with max_inflight=4
+# Enables concurrent request processing
 # 2-stage pipeline: Talker -> Code2Wav
 async_chunk: true
 stage_args:
   - stage_id: 0
     stage_type: llm
+    is_comprehension: true
     runtime:
       devices: "0"
-      max_batch_size: 4
     engine_args:
+      max_num_seqs: 4
       model_stage: qwen3_tts
       model_arch: Qwen3TTSTalkerForConditionalGeneration
-      hf_overrides:
-        architectures: [Qwen3TTSTalkerForConditionalGeneration]
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       enforce_eager: false
       trust_remote_code: true
-      async_scheduling: false
+      async_scheduling: true
       enable_prefix_caching: false
       engine_output_type: latent
       gpu_memory_utilization: 0.3
@@ -40,17 +39,15 @@ stage_args:
     stage_type: llm
     runtime:
       devices: "0"
-      max_batch_size: 4
     engine_args:
+      max_num_seqs: 4
       model_stage: code2wav
       model_arch: Qwen3TTSCode2Wav
-      hf_overrides:
-        architectures: [Qwen3TTSCode2Wav]
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
-      async_scheduling: false
+      async_scheduling: true
       enable_prefix_caching: false
       engine_output_type: audio
       gpu_memory_utilization: 0.2

docker/Dockerfile.ci

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.17.0
+ARG VLLM_BASE_TAG=v0.19.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -11,9 +11,7 @@ RUN apt-get update && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-# Install vllm-omni into the same uv-managed Python environment used by the base image.
-# Use bash -c so that $(python3 -c ...) is expanded inside the container.
-RUN uv pip install --system --no-cache-dir ".[dev]"
+RUN uv pip install --system ".[dev]"
 
 RUN ln -sf /usr/bin/python3 /usr/bin/python
 

docker/Dockerfile.npu

@@ -1,17 +1,30 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.14.0rc1
+ARG VLLM_ASCEND_TAG=v0.18.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-WORKDIR /vllm-workspace/vllm-ascend
-RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
-RUN pip install -v -e .
+# WORKDIR /vllm-workspace/vllm
+# RUN git fetch origin --tags && git checkout v0.18.0
+
+# WORKDIR /vllm-workspace/vllm-ascend
+# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+# # Install vllm-ascend
+# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
+# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+#     source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+#     source /usr/local/Ascend/nnal/atb/set_env.sh && \
+#     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+#     python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
+#     python3 -m pip cache purge
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
-
 COPY . .
 
-RUN pip install -v -e .
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

docker/Dockerfile.npu.a3

@@ -1,17 +1,30 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
+ARG VLLM_ASCEND_TAG=v0.18.0rc1-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-WORKDIR /vllm-workspace/vllm-ascend
-RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
-RUN pip install -v -e .
+# WORKDIR /vllm-workspace/vllm
+# RUN git fetch origin --tags && git checkout v0.18.0
+
+# WORKDIR /vllm-workspace/vllm-ascend
+# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+# # Install vllm-ascend
+# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
+# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+#     source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+#     source /usr/local/Ascend/nnal/atb/set_env.sh && \
+#     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+#     python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
+#     python3 -m pip cache purge
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
-
 COPY . .
 
-RUN pip install -v -e .
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.17.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.19.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
 #    we should swap over to use stable release ASAP.
 #    We should avoid relying on custom commits.
 ARG USE_NIGHTLY_BUILD=0
-ARG VLLM_VERSION_OR_COMMIT_HASH=2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
+ARG VLLM_VERSION_OR_COMMIT_HASH=89138b21cc246ae944c741d5c399c148e2b770ab
 ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
 

docker/Dockerfile.xpu

@@ -76,7 +76,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"
 
-ARG VLLM_VERSION=v0.17.0
+ARG VLLM_VERSION=v0.19.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm
 

docs/.nav.yml

@@ -6,6 +6,7 @@ nav:
     - getting_started/installation/*
   - Serving:
     - OpenAI-Compatible API:
+      - Diffusion Chat API: serving/diffusion_chat_api.md
       - Image Generation: serving/image_generation_api.md
       - Image Edit: serving/image_edit_api.md
       - Text to Speech: serving/speech_api.md
@@ -28,6 +29,7 @@ nav:
       - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
       - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
       - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
+      - Voxtral TTS Offline Inference: user_guide/examples/offline_inference/voxtral_tts.md
       - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
     - Online Serving:
       - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
@@ -52,18 +54,28 @@ nav:
   - Features:
     - Sleep Mode: features/sleep_mode.md
     - Diffusion Features:
-      - Acceleration Overview: user_guide/diffusion_acceleration.md
-      - TeaCache: user_guide/diffusion/teacache.md
-      - Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
+      - Overview: user_guide/diffusion_features.md
+      - Feature Compatibility: user_guide/feature_compatibility.md
+      - Cache Acceleration:
+        - TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
+        - Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
       - Quantization:
         - Overview: user_guide/diffusion/quantization/overview.md
         - FP8: user_guide/diffusion/quantization/fp8.md
+        - Int8: user_guide/diffusion/quantization/int8.md
         - GGUF: user_guide/diffusion/quantization/gguf.md
-      - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
+      - Parallelism:
+        - Overview: user_guide/diffusion/parallelism/overview.md
+        - CFG Parallel: user_guide/diffusion/parallelism/cfg_parallel.md
+        - Expert Parallel: user_guide/diffusion/parallelism/expert_parallel.md
+        - Hybrid Sharded Data Parallel: user_guide/diffusion/parallelism/hsdp.md
+        - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
+        - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
+        - VAE Patch Parallel: user_guide/diffusion/parallelism/vae_patch_parallel.md
       - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
       - LoRA: user_guide/diffusion/lora.md
-      - Hybrid Sharded Data Parallel: design/feature/hsdp.md
       - Custom Pipeline: features/custom_pipeline.md
+      - Step Execution: user_guide/diffusion/step_execution.md
     - ComfyUI: features/comfyui.md
 - Developer Guide:
   - General:
@@ -76,6 +88,7 @@ nav:
     - contributing/model/adding_omni_model.md
     - contributing/model/adding_tts_model.md
     - contributing/model/adding_diffusion_model.md
+    - contributing/model/adding_quantization_model.md
   - CI: contributing/ci
   - Design Documents:
     - design/index.md
@@ -87,19 +100,25 @@ nav:
       - design/feature/cfg_parallel.md
       - design/feature/sequence_parallel.md
       - design/feature/tensor_parallel.md
+      - design/feature/vae_parallel.md
+      - design/feature/hsdp.md
       - design/feature/cache_dit.md
       - design/feature/teacache.md
       - design/feature/async_chunk_design.md
+      - design/feature/vae_parallel.md
+      - design/feature/diffusion_step_execution.md
     - Module Design:
       - design/module/ar_module.md
       - design/module/dit_module.md
       - design/module/entrypoint_module.md
+      - design/module/async_omni_architecture.md
   - Docs Guide: contributing/DOCS_GUIDE.md
 - API Reference:
   - api/README.md
   - api/vllm_omni
 - CLI Reference: cli
 - Community:
+  - Governance: community/governance.md
   - community/*
   - Slack: https://slack.vllm.ai
   - Blog: https://blog.vllm.ai

docs/api/README.md

@@ -5,25 +5,21 @@
 Main entry points for vLLM-Omni inference and serving.
 
 - [vllm_omni.entrypoints.async_omni.AsyncOmni][]
-- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
-- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
+- [vllm_omni.entrypoints.cfg_companion_tracker.CfgCompanionTracker][]
 - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
 - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
 - [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
 - [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
 - [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
 - [vllm_omni.entrypoints.omni.Omni][]
-- [vllm_omni.entrypoints.omni.OmniBase][]
-- [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]
-- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
-- [vllm_omni.entrypoints.omni_stage.OmniStage][]
-- [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
-- [vllm_omni.entrypoints.zmq_utils.ZmqQueue][]
+- [vllm_omni.entrypoints.omni_base.OmniBase][]
+- [vllm_omni.entrypoints.pd_utils.PDDisaggregationMixin][]
 
 ## Inputs
 
 Input data structures for multi-modal inputs.
 
+- [vllm_omni.inputs.data.OmniCustomPrompt][]
 - [vllm_omni.inputs.data.OmniDiffusionSamplingParams][]
 - [vllm_omni.inputs.data.OmniEmbedsPrompt][]
 - [vllm_omni.inputs.data.OmniTextPrompt][]
@@ -48,11 +44,19 @@ Engine classes for offline and online inference.
 - [vllm_omni.engine.OmniEngineCoreOutputs][]
 - [vllm_omni.engine.OmniEngineCoreRequest][]
 - [vllm_omni.engine.PromptEmbedsPayload][]
-- [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
 - [vllm_omni.engine.arg_utils.OmniEngineArgs][]
-- [vllm_omni.engine.input_processor.OmniInputProcessor][]
+- [vllm_omni.engine.async_omni_engine.AsyncOmniEngine][]
+- [vllm_omni.engine.mm_outputs.MultimodalCompletionOutput][]
+- [vllm_omni.engine.mm_outputs.MultimodalPayload][]
+- [vllm_omni.engine.orchestrator.Orchestrator][]
+- [vllm_omni.engine.orchestrator.OrchestratorRequestState][]
+- [vllm_omni.engine.output_modality.OutputModality][]
+- [vllm_omni.engine.output_modality.TensorAccumulationStrategy][]
 - [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
 - [vllm_omni.engine.output_processor.OmniRequestState][]
+- [vllm_omni.engine.stage_engine_core_client.StageEngineCoreClient][]
+- [vllm_omni.engine.stage_init_utils.StageMetadata][]
+- [vllm_omni.engine.stage_init_utils.StartedLlmStage][]
 
 ## Core
 
@@ -64,6 +68,19 @@ Core scheduling and caching components.
 - [vllm_omni.core.sched.output.OmniCachedRequestData][]
 - [vllm_omni.core.sched.output.OmniNewRequestData][]
 - [vllm_omni.core.sched.output.OmniSchedulerOutput][]
+- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.BASECFM][]
+- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.CausalConditionalCFM][]
+- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.CausalMaskedDiffWithDiT][]
+- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.ConditionalCFM][]
+- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConv1d][]
+- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConv1dUpsample][]
+- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConvRNNF0Predictor][]
+- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalHiFTGenerator][]
+- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.HiFTGenerator][]
+- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SineGen][]
+- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SineGen2][]
+- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.Snake][]
+- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SourceModuleHnNSF][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.EuclideanCodebook][]
@@ -75,9 +92,17 @@ Core scheduling and caching components.
 Configuration classes.
 
 - [vllm_omni.config.model.OmniModelConfig][]
+- [vllm_omni.config.stage_config.ModelPipeline][]
+- [vllm_omni.config.stage_config.StageConfig][]
+- [vllm_omni.config.stage_config.StageConfigFactory][]
+- [vllm_omni.config.stage_config.StageType][]
 - [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
 - [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
 - [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
+- [vllm_omni.model_executor.models.cosyvoice3.config.CosyVoice3Config][]
+- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechConfig][]
+- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechFastARConfig][]
+- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechSlowARConfig][]
 - [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioConfig][]
 - [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioTokenizerConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSConfig][]
@@ -91,6 +116,10 @@ Configuration classes.
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
+- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
+- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
+- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
+- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
 
 ## Workers
 
@@ -101,6 +130,8 @@ Worker classes and model runners for distributed inference.
 - [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
 - [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
 - [vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase][]
+- [vllm_omni.diffusion.worker.utils.DiffusionRequestState][]
+- [vllm_omni.diffusion.worker.utils.RunnerOutput][]
 - [vllm_omni.platforms.npu.worker.npu_ar_model_runner.ExecuteModelState][]
 - [vllm_omni.platforms.npu.worker.npu_ar_model_runner.NPUARModelRunner][]
 - [vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker][]
@@ -118,5 +149,6 @@ Worker classes and model runners for distributed inference.
 - [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
 - [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
 - [vllm_omni.worker.gpu_memory_utils.parse_cuda_visible_devices][]
+- [vllm_omni.worker.gpu_model_runner.CUDAGraphWrapper][]
 - [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
 - [vllm_omni.worker.mixins.OmniWorkerMixin][]

docs/community/governance.md

@@ -0,0 +1,54 @@
+# Governance
+
+vLLM-Omni's governance is inspired by the [vLLM governance process](https://docs.vllm.ai/en/latest/governance/process/). We share the same commitment to open source community and meritocratic norms.
+
+## Values
+
+vLLM-Omni aims to be the fastest and easiest-to-use omni-modality inference and serving engine. Our values are aligned with [vLLM's values](https://docs.vllm.ai/en/latest/governance/process/#values):
+
+### Design Values
+
+1. **Top performance**: System performance is our top priority. We monitor overheads, optimize kernels, and publish benchmarks. We never leave performance on the table.
+2. **Ease of use**: vLLM-Omni must be simple to install, configure, and operate. We provide clear documentation, fast startup, clean logs, helpful error messages, and monitoring guides. Many users fork our code or study it deeply, so we keep it readable and modular.
+3. **Wide coverage**: vLLM-Omni supports frontier models and high-performance accelerators. We make it easy to add new models and hardware. vLLM-Omni + PyTorch form a simple interface that avoids complexity.
+4. **Production ready**: vLLM-Omni runs 24/7 in production. It must be easy to operate and monitor for health issues.
+5. **Extensibility**: vLLM-Omni serves as fundamental omni-modality infrastructure. Our codebase cannot cover every use case, so we design for easy forking and customization.
+
+### Collaboration Values
+
+1. **Tightly Knit and Fast-Moving**: Our maintainer team is aligned on vision, philosophy, and roadmap. We work closely to unblock each other and move quickly.
+2. **Individual Merit**: No one buys their way into governance. Committer status belongs to individuals, not companies. We reward contribution, maintenance, and project stewardship.
+
+## Project Maintainers
+
+### Lead Maintainers
+
+Lead maintainers are responsible for the overall direction and strategy of the project:
+
+- [@Gaohan123](https://github.com/Gaohan123)
+- [@hsliuustc0106](https://github.com/hsliuustc0106)
+- [@ywang96](https://github.com/ywang96)
+
+### Active Committers
+
+Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
+
+- [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
+- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
+- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
+- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
+- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
+- [@princepride](https://github.com/princepride): Diffusion and Omni Support
+- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
+- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
+- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism
+- [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
+- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
+
+## Meetings
+
+Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
+
+## Committer Nomination Process
+
+Every month, any active committer can nominate new committer(s) to the project. Up to **two new committers** will be admitted per month based on the quality and impact of their contributions.

docs/configuration/README.md

@@ -14,9 +14,10 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
 
 - **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
 
+## Multi-Stage Recipes
+
+- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
+
 ## Optimization Features
 
-- **[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
-- **[CPU Offloading](../user_guide/diffusion/cpu_offload_diffusion.md)** - Enable CPU offloading (model-level and layerwise) for for DiT models
+- **[Diffusion Features Overview](../user_guide/diffusion_features.md)** - Complete overview of all diffusion model features and supported models

docs/configuration/gpu_memory_utilization.md

@@ -158,7 +158,7 @@ The stage initialized but ran out of memory during processing.
 
 **Solutions:**
 1. Reduce `max_num_batched_tokens`
-2. Reduce `max_batch_size` in runtime config
+2. Reduce `max_num_seqs` in engine_args
 3. Lower `gpu_memory_utilization` slightly
 4. Enable quantization if supported
 

docs/configuration/pd_disaggregation.md

@@ -0,0 +1,177 @@
+# Prefill-Decode (PD) Disaggregation
+
+PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
+stages so prompt processing and token generation can run on different workers.
+
+This is documented as a stage-config recipe instead of a bundled YAML because the
+deployment-specific values usually change per environment:
+
+- GPU placement
+- `tensor_parallel_size`
+- connector backend and connector ports
+- connector IPs or bootstrap addresses
+
+Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml)
+and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
+changes below.
+
+## Requirements
+
+- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
+- A KV connector supported by vLLM, such as `MooncakeConnector`
+- Matching `tensor_parallel_size` on the prefill and decode thinker stages
+
+## 1. Split the thinker into prefill and decode stages
+
+Replace the original thinker stage with two stages:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    is_prefill_only: true
+    runtime:
+      devices: "0"
+    engine_args:
+      max_num_seqs: 16
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 1
+      kv_transfer_config:
+        kv_connector: "MooncakeConnector"
+        kv_role: "kv_producer"
+        kv_rank: 0
+        kv_parallel_size: 2
+        kv_connector_extra_config:
+          mooncake_bootstrap_port: 25201
+    final_output: false
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    stage_type: llm
+    is_decode_only: true
+    runtime:
+      devices: "1"
+    engine_args:
+      max_num_seqs: 64
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 1
+      kv_transfer_config:
+        kv_connector: "MooncakeConnector"
+        kv_role: "kv_consumer"
+        kv_rank: 1
+        kv_parallel_size: 2
+        kv_connector_extra_config:
+          mooncake_bootstrap_port: 25202
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+```
+
+Notes:
+
+- `is_prefill_only: true` marks the thinker stage that only saves KV.
+- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
+- `kv_transfer_config` is required on both stages.
+- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
+  prefill side only processes the prompt and exports KV.
+
+## 2. Shift the downstream stages by one index
+
+After inserting the extra thinker stage, renumber the remaining stages:
+
+```yaml
+  - stage_id: 2
+    runtime:
+      devices: "2"
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+
+  - stage_id: 3
+    runtime:
+      devices: "2"
+    engine_args:
+      max_num_seqs: 1
+    engine_input_source: [2]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+```
+
+Compared with the default Qwen3-Omni config:
+
+- the talker becomes stage `2` instead of stage `1`
+- the code2wav stage becomes stage `3` instead of stage `2`
+- the talker now reads from decode stage `1`
+
+## 3. Add runtime edges for the four-stage pipeline
+
+```yaml
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
+    - from: 1
+      to: 2
+      window_size: -1
+    - from: 2
+      to: 3
+      window_size: -1
+```
+
+## 4. Launch with your custom config
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+    --stage-configs-path /path/to/qwen3_omni_pd.yaml
+```
+
+## Operational Notes
+
+- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
+  pair. Keep prefill and decode at the same `tensor_parallel_size`.
+- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
+  separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
+- Choose connector ports and addresses that match your deployment. The values
+  shown above are examples only.

docs/configuration/stage_configs.md

@@ -18,7 +18,7 @@ If users want to modify some part of it. The custom stage_configs file can be in
 For offline (Assume necessary dependencies have ben imported):
 ```python
 model_name = "Qwen/Qwen2.5-Omni-7B"
-omni_llm = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
+omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
 ```
 
 For online serving:
@@ -30,15 +30,15 @@ vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to
 
 Below is a specific example of stage_configs.yaml in Qwen2.5-omni.
 ```python
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
 stage_args:
   - stage_id: 0 # mark the unique id for each stage
     runtime: # The disaggregated configuration
       process: true  # Run this stage in a separate process
       devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
-      max_batch_size: 1 # the batch_size for offline inference
     engine_args: # Engine arguments for a certain engine
       model_stage: thinker
+      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
       worker_type: ar # The specific worker used
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used
@@ -62,9 +62,9 @@ stage_args:
     runtime:
       process: true
       devices: "1"
-      max_batch_size: 3
     engine_args:
       model_stage: talker
+      max_num_seqs: 3
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -88,9 +88,9 @@ stage_args:
     runtime:
       process: true
       devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
-      max_batch_size: 1
     engine_args:
       model_stage: code2wav
+      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
@@ -159,9 +159,9 @@ Visible devices for this stage, specified as a string. This controls which GPU d
 
 Default: `"0"`
 
-#### `runtime.max_batch_size`
+#### `engine_args.max_num_seqs`
 
-The maximum batch size for offline inference in this stage. This limits how many sequences can be processed together in a single batch during offline inference operations.
+The maximum number of sequences for concurrent processing in this stage. For LLM stages, this controls the vLLM scheduler's maximum concurrent sequences. For all stage types, this also controls how many tasks can be batched together in the task processing loop.
 
 Default: `1`
 

docs/configuration/stage_configs/qwen2_5_omni.yaml

@@ -1,12 +1,12 @@
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
 stage_args:
   - stage_id: 0
     runtime:
       process: true            # Run this stage in a separate process
       devices: "0"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
-      max_batch_size: 1
     engine_args:
       model_stage: thinker
+      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -30,9 +30,9 @@ stage_args:
     runtime:
       process: true
       devices: "1"
-      max_batch_size: 1
     engine_args:
       model_stage: talker
+      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -56,9 +56,9 @@ stage_args:
     runtime:
       process: true
       devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
-      max_batch_size: 1
     engine_args:
       model_stage: code2wav
+      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler

docs/contributing/README.md

@@ -107,7 +107,7 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
 - `[CI/Build]` for build or continuous integration improvements.
 - `[Doc]` for documentation fixes and improvements.
 - `[Model]` for adding a new model or improving an existing model. Model name should appear in the title.
-- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `OmniLLM` class, etc.)
+- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `Omni`/`AsyncOmni`, etc.)
 - `[Kernel]` for changes affecting CUDA kernels or other compute kernels.
 - `[Core]` for changes in the core vLLM-Omni logic (e.g., `OmniProcessor`, `OmniARScheduler`, etc.)
 - `[Hardware][Vendor]` for hardware-specific changes. Vendor name should appear in the prefix, such as [Ascend] for Ascend NPUs.

docs/contributing/ci/.nav.yaml

@@ -0,0 +1,6 @@
+nav:
+  - CI_5levels.md
+  - failures.md
+  - test_guide.md
+  - test_markers.md
+  - test_style.md

docs/contributing/ci/CI_5levels.md

@@ -86,7 +86,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
         /tests/e2e/online_serving/test_{model_name}_expansion.py<br>
         /tests/e2e/offline_inference/test_{model_name}_expansion.py<br>
         <strong>Performance:</strong><br>
-        /tests/e2e/perf/nightly.json<br>
+        /tests/dfx/perf/tests/test.json<br>
         <strong>Doc Test:</strong><br>
         tests/example/online_serving/test_{model_name}.py<br>
         tests/example/offline_inference/test_{model_name}.py
@@ -104,7 +104,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
       <td> Depends on reality </td>
       <td>
         <strong>Stability:</strong><br>
-        tests/e2e/stability/weekly.json<br>
+        /tests/dfx/stability/tests/test.json<br>
         <strong>Reliability:</strong><br>
         tests/e2e/reliability/test_{model_name}.py
       </td>
@@ -168,12 +168,6 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
-│   ├── omni.py                                 │   ├── test_omni.py                    ⬜  (E2E covered by e2e/offline, e2e/online)
-│   ├── omni_llm.py                             │   ├── test_omni_llm.py                ✅
-│   ├── omni_stage.py                            │   ├── test_omni_stage.py              ⬜  (partial in test_omni_stage_diffusion_config.py)
-│   ├── omni_diffusion.py                       │   ├── test_omni_diffusion.py          ✅
-│   ├── async_omni.py                            │   ├── test_async_omni.py              ✅ actually in e2e/online_serving/test_async_omni.py
-│   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
 │   │   └── ...                                  │   │   └── test_*.py                  ⬜
@@ -230,9 +224,11 @@ vllm_omni/                                    tests/
 │
 └── e2e (tests side only)               →     ├── e2e/
                                                ├── online_serving/                     ✅ non-empty
+                                               │   ├── test_qwen2_5_omni.py
                                                │   ├── test_async_omni.py
                                                │   ├── test_qwen3_omni.py
                                                │   ├── test_qwen3_omni_expansion.py
+                                               │   ├── test_mimo_audio.py
                                                │   ├── test_image_gen_edit.py
                                                │   ├── test_images_generations_lora.py
                                                │   └── stage_configs/
@@ -534,13 +530,13 @@ L4 level testing is a comprehensive quality audit before a version release. It e
 ### 3.2 Testing Content and Scope
 
 -   ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling.
--   ***Performance Testing***: Uses the `/tests/e2e/perf/nightly.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
+-   ***Performance Testing***: Uses the `tests/dfx/perf/tests/test.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
 -   ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description.
 
 ### 3.3 Test Directory and Execution Files
 
 -   ***Functional Testing***: Same directories as L3.
--   ***Performance Test Configuration***: `/tests/e2e/perf/nightly.json`
+-   ***Performance Test Configuration***: `tests/dfx/perf/tests/test.json`
 -   ***Documentation Example Tests***:
 -   -   `tests/example/online_serving/test_{model_name}.py`
     -   `tests/example/offline_inference/test_{model_name}.py`
@@ -550,97 +546,20 @@ L4 level testing is a comprehensive quality audit before a version release. It e
 -   ***Trigger Timing***: **`Nightly`**, automatically executed every night.
 -   ***Execution Environment***: ***GPU*** server clusters to meet the resource demands of performance testing.
 -   ***Script Example***:
-<details>
-<summary> Test Examples</summary>
-When you want to add L4-level performance test cases, you can refer to the following format for case addition in tests/perf/tests/test.json:
 
-```JSON
-{
-    "test_name": "test_qwen3_omni",
-    "server_params": {
-        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-        "stage_config_name": "qwen3_omni.yaml"
-    },
-    "benchmark_params": [
-        {
-            "dataset_name": "random",
-            "num_prompts": [10, 20],
-            "request_rate": [0.5, 1],
-            "random_input_len": 2500,
-            "random_output_len": 900,
-            "ignore_eos": true,
-            "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-            "baseline": {
-                "mean_ttft_ms": 100000,
-                "mean_audio_ttfp_ms": 100000,
-                "mean_audio_rtf": 100000
-            }
-        }
-    ]
-}
-```
+??? example "Test Examples: Documentation Example Tests"
 
-#### Parameter Explanation
+    --8<-- "docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md"
 
-***Overview***
+??? example "Test Examples: Performance Tests"
 
-| Field            | Required | Description                                                     |
-| ---------------- | -------- | --------------------------------------------------------------- |
-| test_name        | Yes      | Unique identifier for the test case                             |
-| server_params    | Yes      | Server-side configuration parameters                            |
-| benchmark_params | Yes      | Benchmark running parameters (supports multiple configurations) |
+    --8<-- "docs/contributing/ci/test_examples/l4_performance_tests.inc.md"
 
-#### server_params Configuration
+??? example "Test Examples: Functionality Tests"
 
-##### Basic Parameters
+    --8<-- "docs/contributing/ci/test_examples/l4_functionality_tests.inc.md"
 
-| Parameter         | Required | Example                            | Description                   |
-| ----------------- | -------- | ---------------------------------- | ----------------------------- |
-| model             | Yes      | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path            |
-| stage_config_name | Yes      | "qwen3_omni.yaml"                  | Stage configuration file name |
-
-##### Dynamic Configuration (update/delete)
-
-Supports incremental modifications based on the basic configuration:
-
-| Operation | Description                          |
-| --------- | ------------------------------------ |
-| update    | Update or add configuration items    |
-| delete    | Delete specified configuration items |
-
-***Example***:
-```
-"update": {
-    "async_chunk": true,  // Enable asynchronous chunk processing
-    "stage_args": {
-        "0": {
-            "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
-        }
-    }
-},
-"delete": {
-    "stage_args": {
-        "2": ["custom_process_input_func"]  // Delete this configuration for stage 2
-    }
-}
-```
-#### benchmark_params Configuration
-
-You can add any benchmark running parameters you need here. For all optional parameters, refer to the [benchmark documentation](https://github.com/vllm-project/vllm-omni/blob/main/docs/cli/bench/serve.md). General modifications are as follows:
-
-1.  Change the ---xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
-2.  For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
-3.  Add the baseline parameter to specify the required validation values, ensuring the validation metric names match those in the result.json generated by the benchmark.
-4.  The qps and concurrency modes are mutually exclusive. For detailed explanations, see the table below:
-
-| Parameter       | Type        | Required | Example/Values  | Description                                                                                                                                                                                                                                                          |
-| --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| num_prompts     | int / array | Yes      | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
-| request_rate    | int / array | No       | 1, [1, 2, 3]    | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
-| max_concurrency | int / array | No       | 1, [1, 2, 3]    | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
-</details>
-
--   -   ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).
+-   ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).
 
 ## Chapter 4: L5 Level Testing - Stability and Reliability Testing
 
@@ -652,12 +571,12 @@ L5 level testing focuses on the performance of model services under ***long-runn
 
 ### 4.2 Testing Content and Scope
 
--   ***Long-term Stability (Stability) Testing***: Uses the `tests/e2e/stability/weekly.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
+-   ***Long-term Stability (Stability) Testing***: Uses the `tests/dfx/stability/tests/test.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
 -   ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
 
 ### 4.3 Test Directory and Execution Files
 
--   ***Stability Test Configuration***: `tests/e2e/stability/weekly.json`
+-   ***Stability Test Configuration***: `tests/dfx/stability/tests/test.json`
 -   ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
 
 ### 4.4 Execution Method and Example
@@ -667,12 +586,79 @@ L5 level testing focuses on the performance of model services under ***long-runn
 -   ***Script Example***:
 <details>
 <summary> Test Examples</summary>
-```python
-# WIP
+
+When you want to add L5-level stability test cases, you can refer to the following format for case addition in `tests/dfx/stability/tests/test.json`:
+
+```json
+{
+    "test_name": "test_qwen3_omni_stability",
+    "server_params": {
+        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "stage_config_name": "qwen3_omni.yaml"
+    },
+    "benchmark_params": [
+        {
+            "dataset_name": "random",
+            "backend": "openai-chat-omni",
+            "endpoint": "/v1/chat/completions",
+            "duration_sec": 43200,
+            "request_rate": 0.5,
+            "num_prompts_per_batch": 20,
+            "random_input_len": 2500,
+            "random_output_len": 900,
+            "ignore_eos": true,
+            "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration"
+        }
+    ]
+}
 ```
+
+#### Parameter Explanation
+
+***Overview***
+
+| Field            | Required | Description                                                                 |
+| ---------------- | -------- | --------------------------------------------------------------------------- |
+| test_name        | Yes      | Unique identifier for the stability test case                               |
+| server_params    | Yes      | Server-side configuration parameters (model, stage configuration, etc.)     |
+| benchmark_params | Yes      | Stability benchmark running parameters (supports multiple configurations)   |
+
+#### server_params Configuration
+
+##### Basic Parameters
+
+| Parameter         | Required | Example                            | Description                         |
+| ----------------- | -------- | ---------------------------------- | ----------------------------------- |
+| model             | Yes      | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path                  |
+| stage_config_name | Yes      | "qwen3_omni.yaml"                  | Stage configuration file name       |
+
+##### Dynamic Configuration (update/delete)
+
+Supports incremental modifications based on the basic configuration:
+
+| Operation | Description                          |
+| --------- | ------------------------------------ |
+| update    | Update or add configuration items    |
+| delete    | Delete specified configuration items |
+
+***Example***:
+You can refer to Test Examples in Chapter 3.4
+
+#### benchmark_params Configuration
+
+For stability testing, the key parameters are:
+
+-   **duration_sec**: Total duration (in seconds) during which benchmark traffic is sent. The stability benchmark will keep sending batches until this duration is reached.
+-   **request_rate** / **max_concurrency**: Exactly one of them must be specified. They control how the traffic is generated for each batch:
+    -   `request_rate`: Number of requests per second. The benchmark will send `num_prompts_per_batch` requests at the given rate.
+    -   `max_concurrency`: Maximum number of concurrent requests. When this is used, `request_rate` is set to `inf` internally.
+-   **num_prompts_per_batch**: Number of prompts sent in each batch. Multiple batches will be executed sequentially within `duration_sec`.
+
+All other optional parameters follow the same rules as the in Chapter 3.4.
+
 </details>
 
--   -   ***Stability***: (Execution would be driven by the configuration in `weekly.json`).
+-   -   ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_{model_name}.py`
     -   ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py`
 
 ## Summary

docs/contributing/ci/failures.md

@@ -1,10 +1,106 @@
 # CI Failures
 
-What should I do when a CI job fails on my PR, but I don't think my PR caused the failure?
+## Overview of CI Checks
 
+When you open a PR against vLLM-Omni, several CI checks run automatically:
 
-## Common Case 1:
+| Check | Platform | What it does |
+| ----- | -------- | ------------ |
+| **pre-commit** | GitHub Actions | Runs linting (Ruff), formatting, spell-checking (typos), and YAML validation. |
+| **Build Wheel** | GitHub Actions | Builds Python wheels for Python 3.11 and 3.12 on Ubuntu. Skipped for docs-only or Markdown-only changes (controlled by `paths-ignore` in the workflow). |
+| **DCO** | GitHub | Verifies every commit has a `Signed-off-by` line. |
+| **docs/readthedocs.org:vllm-omni** | Read the Docs | Builds the MkDocs documentation site. |
+| **buildkite/vllm-omni** | Buildkite | Runs GPU-based tests on NVIDIA CUDA hardware (L4, H100). |
+| **buildkite/vllm-omni-amd** | Buildkite | Runs GPU-based tests on AMD ROCm hardware (MI325). |
+| **buildkite/vllm-omni-intel** | Buildkite | Runs GPU-based tests on Intel XPU hardware (Intel Arc BMG). |
 
-🚨 Error: The command was interrupted by a signal: signal: terminated
+## Step 1: Identify the Failing Check
 
-Reason: The test is terminated due to exceed the time limits. Sometimes rebuild the test will pass.
+Click the **Details** link next to the failing check on your PR to open the build log. The most common failures fall into these categories:
+
+### pre-commit failures
+
+These are typically formatting or linting issues introduced by your PR. Fix them locally:
+
+```bash
+uv pip install pre-commit
+pre-commit run --all-files
+```
+
+Then commit the fixes and push.
+
+### DCO failures
+
+Every commit must include a `Signed-off-by` line. If you forgot, amend your commits:
+
+```bash
+git commit --amend -s
+git push --force-with-lease
+```
+
+For multiple commits, use an interactive rebase to add the sign-off to each one.
+
+### Read the Docs failures
+
+The documentation build uses MkDocs with `fail_on_warning: true`, so even a minor warning (not just errors) will cause the build to fail. To reproduce locally:
+
+```bash
+uv pip install -e ".[docs]"
+mkdocs build --strict
+```
+
+Common causes include broken cross-references, invalid admonition syntax, or missing files referenced by `--8<--` includes.
+
+### Buildkite failures
+
+Buildkite runs GPU tests in Docker containers. These are the most complex checks and can fail for reasons unrelated to your PR (infrastructure issues, flaky tests, etc.). See the sections below for how to investigate.
+
+## Step 2: Check if the Failure Is a Known Issue
+
+Before spending time debugging, check whether the failure already exists on the `main` branch:
+
+1. **Look at the Buildkite build log** — the test name and error message are usually enough to identify the issue.
+2. **Check recent CI runs on `main`** — if the same test is failing there, the failure is not caused by your PR.
+3. **Search existing issues** — look for open issues in the [vllm-omni issue tracker](https://github.com/vllm-project/vllm-omni/issues) with the test name or error message.
+
+If the failure is already tracked, leave a comment on your PR noting that the failure is pre-existing and link the issue.
+
+## Step 3: Investigate the Failure
+
+If the failure appears to be new, investigate whether your changes caused it.
+
+### Reading Buildkite Logs
+
+1. Click **Details** next to the Buildkite check on your PR.
+2. Find the failing step in the pipeline (e.g., "Diffusion Model Test", "Simple Unit Test").
+3. Expand the step to see the full test output with the traceback.
+
+### Running Tests Locally
+
+For instructions on running tests locally (including specific test files, functions, and markers), see the [Running Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_guide/#running-tests) section in the Test Guide.
+
+## Step 4: Raise an Issue or Fix It
+
+### If the failure is pre-existing (not caused by your PR)
+
+1. **Raise a new issue** if one doesn't already exist, using the title format:
+   `[CI Failure]: [job-name] - [test-path]`
+2. Include the error message, relevant log excerpts, and the commit hash where the failure occurs (e.g., "Still failing on main as of commit `abc1234`").
+3. Leave a comment on your PR linking to the issue and noting that the failure is unrelated to your changes.
+
+### If the failure is caused by your PR
+
+1. Fix the issue in your branch and push the update.
+2. If the fix is non-trivial, consider adding a test to prevent regression.
+
+## Common Failure Patterns
+
+| Symptom | Likely Cause | Fix |
+| ------- | ------------ | --- |
+| `ruff` or formatting errors | Code style violation | Run `pre-commit run --all-files` |
+| `Signed-off-by` missing | DCO check | Amend commits with `git commit --amend -s` |
+| MkDocs build warning | Broken docs reference | Run `mkdocs build --strict` locally |
+| `OOM` or `CUDA out of memory` | Test exceeds GPU memory | Check if your changes increased memory usage; use `--vae_use_slicing` / `--vae_use_tiling` for diffusion tests |
+| Import errors | Missing or changed dependency | Check `pyproject.toml` and make sure dependencies are correct |
+| Timeout (step exceeded N minutes) | Test is too slow or hangs | Profile the test; check for infinite loops or deadlocks |
+| `Agent lost` in Buildkite | Infrastructure issue (not your fault) | Re-trigger the build; comment on your PR |

docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md

@@ -0,0 +1,49 @@
+**Preferred Test Strategy**
+
+Use one of the following patterns depending on page type:
+
+- **Dynamic code-block extraction (preferred for offline docs)**
+    - Extract Python/Bash code blocks from markdown AST analyzer, then execute them directly in tests.
+    - Benefit: test logic stays automatically aligned with docs.
+    - Basic idea: Use `ReadmeSnippet.extract_readme_snippets` to extract a list of code blocks as a global variable in file,
+    use this list as `pytest.mark.parametrize` parameters, and pass each snippet item to `example_runner.run` inside the parametrized test.
+    Additionally pass an `output_subfolder` argument for the 2nd-level output folder explained in **Output Directory Structure** below.
+    If any extra environment variable is need for a test (e.g., the example script reads it), `example_runner.run` also accepts a 3rd `env` parameter.
+    - See [tests/examples/offline_inference/test_text_to_image.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/examples/offline_inference/test_text_to_image.py) for reference implementation.
+
+- **Explicit copied scripts (used by online docs for now until further update)**
+    - For online serving pages, it is acceptable to copy code from docs into dedicated test functions, because only client-side, request-sending scripts are tested.
+    - Benefit: dynamic extraction is overly complex: need to tell server-launch and client-request scripts.
+    - Requirement: copied test code must be kept in sync with doc updates.
+
+**Test Case Naming Convention**
+
+- Dynamic code extraction (auto-generated internally):
+    - `test_{single_function_name_matching_file_name}[h2_heading_00X]`
+    - Example: `test_text_to_image[basic_usage_001]`
+- Explicit copied scripts:
+    - `test_{h2_heading_00X}[{dummy_param_id_for_omni_server}]`
+    - Example: `test_api_calls_001[omni_server0]`
+
+**Runtime Configuration**
+
+In the example code tests, do **not** reduce `num_inference_steps` just to speed up the tests  unless there is a strong CI reliability reason to do otherwise.
+
+**Skipping Rules**
+
+You may skip examples falling in the following categories using `pytest.mark.skip` or `pytest.skip`:
+
+- Gradio UI scripts
+- Scenarios that significantly overlap with existing tests and add little new coverage.
+
+**Output Directory Structure**
+
+Use a three-layer output structure to store output artifacts:
+
+1. Root output directory
+    - Auto-detected from `OUTPUT_DIR` env var or auto-generated under `/tmp`.
+2. Doc-page directory
+    - Define and use a clear page-level folder name in each `test_*.py` yourself (abbreviations are acceptable, e.g., `example_offline_t2i`).
+3. Test-case directory
+    - Must match the case identifier (e.g., `basic_usage_001`).
+    - Auto-generated for dynamic extracted tests.

docs/contributing/ci/test_examples/l4_functionality_tests.inc.md

@@ -0,0 +1,46 @@
+**Scope**
+
+For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
+
+**Test Case Design**
+
+For a *high priority* model (currently listed in [issue #1832](https://github.com/vllm-project/vllm-omni/issues/1832)), we use several test cases, each with multiple features turned on, so that each supported feature is tested in at least one test case. This is to relieve the GPU workload on the CI machine. The suggested test case combination is as follows:
+
+- If the model can fit into 4 L4 GPU (with quantization and tensor parallel always on) (20GB RAM each)
+    - (1 GPU) TeaCache + Layerwise CPU offloading + GGUF
+    - (4 GPUs) CacheDiT + Ulysses=2 + TP=2 + VAE=2 + FP8
+    - (4 GPUs) CacheDiT + Ring=2 + HSDP=2 + VAE=2 + GGUF
+    - (4 GPUs) TeaCache + CFG=2 + TP=2 + VAE=2 + FP8
+- Otherwise, consider 2 H100 GPU environment (80GB RAM each) with the following tests
+    - (1 GPU) TeaCache + Layerwise CPU offloading + GGUF
+    - (2 GPUs) CacheDiT + Ulysses=2 + FP8
+    - (2 GPUs) CacheDiT + Ring=2 + GGUF
+    - (2 GPUs) TeaCache + CFG=2 + FP8
+    - (2 GPUs) CacheDiT + TP=2 + VAE=2 + FP8
+    - (2 GPUs) CacheDiT + HSDP=2 + VAE=2 + GGUF
+- If 2 H100 GPU cannot handle the model either (e.g., HunyuanImage 3.0)
+    - Still design tests and feature combinations that can best fit real-world scenario.
+    - Do not include it in CI (or exclude it from the CI's filtering criteria). Instead, relevant PR authors are suggested to run these tests locally.
+- Fallback plan
+    - If the model does not support layerwise CPU offloading, replace the corresponding test case with module-wise offloading
+    - If the model only supports specific or no caching feature, use this option or remove this option in all test cases.
+    - If the model only supports specific or no quantization feature, use this option or remove this option in all test cases.
+    - If the model does not support certain other features, remove this option from that test case. If, consequently, the coverage of this modified test case completely overlaps with others, remove this test case.
+
+For a *normal priority* model, further reduce the number of test cases.
+
+- Only write one or two test cases for the most common feature combinations.
+- The author can explore themselves to see which feature combination balances output quality and performance. Alternatively, the author can refer to any example code in the PR that adds the model, or the example code in the PR that adds a feature (if the code involves this model of interest).
+
+Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`.
+
+**Code Style**
+
+- Validation: test that the multimodal output files of your model have the correct shapes. `OpenAIClientHandler.send_diffusion_request` should have taken care of this.
+- Test marks: always add `advanced_model` and `diffusion`. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/).
+- To maximize code reuse, you may refer to
+    - `tests/conftest.py` for `omni_server` (running server in subprocess) and `openai_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper.
+    - `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
+    - [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
+- Doc: add a concise docstring for each test function.
+- Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).

docs/contributing/ci/test_examples/l4_performance_tests.inc.md

@@ -0,0 +1,89 @@
+When you want to add L4-level ***performance test*** cases, you can refer to the following format for case addition in tests/dfx/perf/tests/test.json:
+
+```JSON
+{
+    "test_name": "test_qwen3_omni",
+    "server_params": {
+        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "stage_config_name": "qwen3_omni.yaml"
+    },
+    "benchmark_params": [
+        {
+            "dataset_name": "random",
+            "num_prompts": [10, 20],
+            "max_concurrency": [1, 4],
+            "random_input_len": 2500,
+            "random_output_len": 900,
+            "ignore_eos": true,
+            "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+            "baseline": {
+                "mean_ttft_ms": [500, 800],
+                "mean_audio_ttfp_ms": [2000, 3500],
+                "mean_audio_rtf": [0.25, 0.35]
+            }
+        }
+    ]
+}
+```
+
+**Parameter Explanation**
+
+*Overview*
+
+| Field            | Required | Description                                                     |
+| ---------------- | -------- | --------------------------------------------------------------- |
+| test_name        | Yes      | Unique identifier for the test case                             |
+| server_params    | Yes      | Server-side configuration parameters                            |
+| benchmark_params | Yes      | Benchmark running parameters (supports multiple configurations) |
+
+**`server_params` Configuration**
+
+*Basic Parameters*
+
+| Parameter         | Required | Example                            | Description                   |
+| ----------------- | -------- | ---------------------------------- | ----------------------------- |
+| model             | Yes      | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path            |
+| stage_config_name | Yes      | "qwen3_omni.yaml"                  | Stage configuration file name |
+
+*Dynamic Configuration (update/delete)*
+
+Supports incremental modifications based on the basic configuration:
+
+| Operation | Description                          |
+| --------- | ------------------------------------ |
+| update    | Update or add configuration items    |
+| delete    | Delete specified configuration items |
+
+**Example**:
+
+```
+"update": {
+    "async_chunk": true,  // Enable asynchronous chunk processing
+    "stage_args": {
+        "0": {
+            "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
+        }
+    }
+},
+"delete": {
+    "stage_args": {
+        "2": ["custom_process_input_func"]  // Delete this configuration for stage 2
+    }
+}
+```
+
+**`benchmark_params` Configuration**
+
+You can add any benchmark running parameters you need here. For all optional parameters, refer to the [benchmark documentation](https://github.com/vllm-project/vllm-omni/blob/main/docs/cli/bench/serve.md). General modifications are as follows:
+
+1.  Change the --xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
+2.  For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
+3.  Optionally add a `baseline` object (see **Baseline thresholds** below). If you omit `baseline` or leave it empty, the performance test still runs but does not assert metric thresholds from this field.
+4.  The qps and concurrency modes are recommended to be mutually exclusive. For detailed explanations, see the table below:
+
+| Parameter       | Type        | Required | Example/Values  | Description                                                                                                                                                                                                                                                          |
+| --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| num_prompts     | int / array | Yes      | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
+| request_rate    | float / array | No  | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
+| max_concurrency | int / array | No       | 1, [1, 2, 3]    | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode).                                                                                                                                                                                                             |
+| baseline        | object      | No       | see above       | Optional per-metric thresholds; keys must match benchmark output fields. Scalar, list (per sweep step), or object (keyed by concurrency or QPS string).  

docs/contributing/ci/tests_style.md

@@ -24,8 +24,6 @@ End-to-end tests verify the complete functionality of a system or component. For
 
 - **`tests/e2e/online_serving/`**: Tests for online serving scenarios (e.g., API server tests)
 
-**Example:** The test file for `vllm_omni/entrypoints/omni_llm.py` should be located at `tests/entrypoints/test_omni_llm.py`.
-
 ## Test Directory Structure
 
 The ideal directory structure mirrors the source code organization. Legend: `✅` = test exists, `⬜` = suggested to add.
@@ -75,12 +73,6 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
-│   ├── omni.py                                 │   ├── test_omni.py                    ⬜  (E2E covered by e2e/offline, e2e/online)
-│   ├── omni_llm.py                             │   ├── test_omni_llm.py                ✅
-│   ├── omni_stage.py                            │   ├── test_omni_stage.py              ⬜  (partial in test_omni_stage_diffusion_config.py)
-│   ├── omni_diffusion.py                       │   ├── test_omni_diffusion.py          ✅
-│   ├── async_omni.py                            │   ├── test_async_omni.py              ✅ actually in e2e/online_serving/test_async_omni.py
-│   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
 │   │   └── ...                                  │   │   └── test_*.py                  ⬜
@@ -137,9 +129,11 @@ vllm_omni/                                    tests/
 │
 └── e2e (tests side only)               →     ├── e2e/
                                                ├── online_serving/                     ✅ non-empty
+                                               │   ├── test_qwen2_5_omni.py
                                                │   ├── test_async_omni.py
                                                │   ├── test_qwen3_omni.py
                                                │   ├── test_qwen3_omni_expansion.py
+                                               │   ├── test_mimo_audio.py
                                                │   ├── test_image_gen_edit.py
                                                │   ├── test_images_generations_lora.py
                                                │   └── stage_configs/
@@ -164,13 +158,19 @@ vllm_omni/                                    tests/
                                                        ├── qwen3_omni_ci.yaml
                                                        ├── bagel_*.yaml
                                                        └── npu/, rocm/, etc.
+examples/                                     tests
+│                                             └── examples
+├── online_serving/                     →         ├── online_serving/
+│   └── {doc_page_title}/README.md                │   └── test_{doc_page_title}.py  ⬜
+└── offline_inference/                  →         └── offline_inference/
+    └── {doc_page_title}/README.md                    └── test_{doc_page_title}.py  ⬜
 ```
 
 
 
 ### Naming Conventions
 
-- **Unit Tests**: Use `test_<module_name>.py` format. Example: `omni_llm.py` → `test_omni_llm.py`
+- **Unit Tests**: Use `test_<module_name>.py` format. Example: `stage_utils.py` → `test_stage_utils.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`
 

docs/contributing/model/adding_diffusion_model.md

@@ -653,25 +653,7 @@ For a fair comparison, keep the same **prompt**, **seed**, **resolution**, **num
 
 To ensure project maintainability and sustainable development, please submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
 
-For comprehensive testing guidelines and the definition of test levels (L1-L5), please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
-The following tests are required to add:
-
-- L4 test of the model's full *functionality* (i.e., all the *diffusion features* that are supported by this model), including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
-  - Test cases: Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`. The following test cases shall cover all features:
-    - 1 GPU: TeaCache & GGUF (or fallback to FP8, or disable it) & Layer-wise CPU offloading (or fallback to Module-wise)
-    - 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Ulysses = 2
-    - 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & Ring = 2
-    - 2 GPUs: TeaCache & FP8 (or fallback to GGUF, or disable it) & CFG Parallel = 2
-    - 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Tensor Parallel = 2 & VAE Patch Parallel = 2
-    - 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & HSDP = 2 & VAE Patch Parallel = 2
-  - Validation: test that the multimodal output files of your model have the correct shapes.
-  - Test marks: always add `advanced_model` and `diffusion`. Add `parallel` and GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/)
-  - To maximize code reuse, you may refer to
-    - `tests/conftest.py` for `omni_server` and `openai_client` fixtures, `generate_synthetic_image` and `assert_XXX_valid` helper.
-    - `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
-    - [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
-  - Doc: add a concise dostring for each test function.
-  - Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).
+For comprehensive testing guidelines and the definition of test levels (L1-L5), please refer to the [Multi-Level Automated Testing System Documentation](../ci/CI_5levels.md). You are at least required to add an L4 *functionality* test described in that document.
 
 ---
 
@@ -739,6 +721,22 @@ See detailed guide: [How to add Sequence Parallel support](../../design/feature/
 omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2)
 ```
 
+### Step Execution
+
+See detailed design guide: [How to add step execution support](../../design/feature/diffusion_step_execution.md)
+
+Use this only when your pipeline can be split into stable request-scoped and
+step-scoped phases. The reference implementation is
+`QwenImagePipeline`, which maps its request-level `forward()` into:
+
+1. `prepare_encode()` for prompt encoding, latent init, timestep prep, and per-request scheduler setup.
+2. `denoise_step()` for one transformer/noise prediction.
+3. `step_scheduler()` for one scheduler update and `step_index` advance.
+4. `post_decode()` for the final VAE decode.
+
+Do not enable `step_execution=True` until those four methods are implemented
+and validated against the request-level path.
+
 ### Cache Acceleration
 
 #### TeaCache
@@ -815,6 +813,107 @@ class WanTransformer3DModel(nn.Module):
 
 ---
 
+### Diffusion Timing (Performance Profiling)
+When adapting a new diffusion model, it is often useful to analyze the latency of key components such as text encoding, diffusion denoising, and VAE decoding.
+vLLM-Omni provides a timing utility via `DiffusionPipelineProfilerMixin` to help developers quickly identify performance bottlenecks.
+
+!!! info
+      `DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md#3-profiling-diffusion-models). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
+
+This tool automatically measures the execution time of selected pipeline modules and prints the results in the logs.
+
+**Enabling Diffusion Timing**
+
+
+Enable timing by setting:
+```
+vllm serve Qwen/Qwen-Image --omni --port 8091 --enable-diffusion-pipeline-profiler
+```
+You can optionally specify which modules to profile:
+```
+class YourPipeline(xxx, DiffusionPipelineProfilerMixin):
+    def __init__(self, xxx):
+        ...
+        self.setup_diffusion_pipeline_profiler(profiler_targets=["diffuse"], enable_diffusion_pipeline_profiler)
+```
+If not specified, the default targets are used:
+```
+["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"]
+```
+**Adding DiffusionPipelineProfilerMixin to a Pipeline**
+To enable timing support in your pipeline, inherit from DiffusionPipelineProfilerMixin.
+```python
+from vllm_omni.diffusion.utils.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+
+class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
+    # Optional: Specify custom timing targets
+    _PROFILER_TARGETS = ["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"]
+
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+        self.parallel_config = od_config.parallel_config
+        # initialize pipeline components
+        ...
+
+        # initialize timing profiler
+        self.setup_diffusion_pipeline_profiler(enable_diffusion_pipeline_profiler)
+```
+The mixin dynamically wraps selected methods and records their execution time during inference.
+
+If you need to fetch the execution time of different modules, you will need to pass `self.stage_durations` to `DiffusionOutput`, as shown below:
+
+```diff
+-   return DiffusionOutput(output=img)
++   return DiffusionOutput(
+            output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
+        )
+```
+
+**Pipeline Design for Timing**
+The current diffusion timing utility is function-based, meaning it measures the execution time of individual methods.
+
+When implementing a new pipeline, avoid putting all logic inside a single function (e.g., forward). Instead, structure the pipeline in a modular way by separating key stages into independent methods, such as the diffusion loop.
+
+For example:
+```
+def forward(self, req: OmniDiffusionRequest):
+    prompt_embeds = self.encode_prompt(req)
+    latents = self.diffuse(prompt_embeds, req)
+    images = self.vae.decode(latents)
+    return DiffusionOutput(output=images)
+```
+This allows the timing utility to measure each stage (e.g., encode_prompt, diffuse, vae.decode) separately and helps identify performance bottlenecks more easily.
+
+
+**Default Profiled Modules**
+
+By default, the following pipeline modules are timed:
+```
+vae.encode
+vae.decode
+diffuse
+text_encoder.forward
+tokenizer.forward
+```
+
+**Example Output**
+
+When enabled, timing logs appear like this:
+```
+[DiffusionTiming] text_encoder.forward took 0.018s
+[DiffusionTiming] diffuse took 2.412s
+[DiffusionTiming] vae.decode took 0.063s
+```
+These measurements help identify bottlenecks during model adaptation and optimization
+
+
+
 ## Troubleshooting
 
 
@@ -866,7 +965,7 @@ hidden_states = hidden_states.reshape(batch_size, seq_len, -1)
 
 1. **Reduce batch size:**
    ```python
-   omni.generate(prompts=[...], max_batch_size=2)
+   omni.generate(prompts=[...], max_num_seqs=2)
    ```
 
 2. **Use smaller image size:**

docs/contributing/model/adding_omni_model.md

@@ -330,54 +330,46 @@ Stage transitions are the mechanism by which outputs from one stage are converte
 
 ### Where Stage Transitions Are Called
 
-Stage transitions happen automatically in the orchestrator (`OmniLLM` class) during the generation loop. Here's the detailed flow:
+Stage transitions happen automatically in the runtime orchestrator. Here's the detailed flow:
 
-1. **Location**: `vllm_omni/entrypoints/omni_llm.py` in the `_run_generation()` method
+1. **Location**: `vllm_omni/engine/orchestrator.py` in `_forward_to_next_stage()`
 2. **Trigger**: When a stage completes processing and produces outputs
 3. **Execution Flow**:
    ```python
-   # In omni_llm.py, _run_generation() method (around line 345-460)
+   # In orchestrator.py
+   next_stage_id = stage_id + 1
+   next_client = self.stage_clients[next_stage_id]
+   params = req_state.sampling_params_list[next_stage_id]
 
-   # Main orchestrator loop polls each stage for completed requests
-   for stage_id, stage in enumerate(self.stage_list):
-       result = stage.try_collect()  # Get completed request
-       if result is None:
-           continue
+   # Save current stage outputs so stage_input_processors can consume them.
+   self.stage_clients[stage_id].set_engine_outputs([output])
 
-       # Store outputs from this stage
-       engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
-       stage.set_engine_outputs(engine_outputs)
+   # THIS IS WHERE STAGE TRANSITION HAPPENS
+   next_inputs = next_client.process_engine_inputs(
+       stage_list=self.stage_clients,
+       prompt=req_state.prompt,
+   )
 
-       # Check if there's a next stage to forward to
-       next_stage_id = stage_id + 1
-       if next_stage_id < len(self.stage_list):
-           next_stage: OmniStage = self.stage_list[next_stage_id]
-
-           # THIS IS WHERE STAGE TRANSITION HAPPENS
-           next_inputs = next_stage.process_engine_inputs(
-               self.stage_list,
-               [request_id_to_prompt[req_id]]
-           )
-
-           # Submit to next stage
-           task = {
-               "type": OmniStageTaskType.GENERATE,
-               "request_id": req_id,
-               "engine_inputs": next_inputs[0],
-               "sampling_params": sampling_params_list[next_stage_id],
-           }
-           next_stage.submit(task)
+   # Build and submit request(s) to the next stage.
+   for next_input in next_inputs:
+       request = build_engine_core_request_from_tokens(
+           request_id=req_id,
+           prompt=next_input,
+           params=params,
+           model_config=self.stage_vllm_configs[next_stage_id].model_config,
+       )
+       await next_client.add_request_async(request)
    ```
 
 ### How Stage Transitions Work
 
 The stage transition process follows these steps:
 
-1. **Stage Completion**: When a stage finishes processing a request, it stores outputs via `stage.set_engine_outputs(engine_outputs)`
+1. **Stage Completion**: When a stage finishes processing a request, the orchestrator stores outputs via `stage_client.set_engine_outputs(...)`
 
 2. **Transition Detection**: The orchestrator checks if there's a next stage and calls `process_engine_inputs()` on it
 
-3. **Input Processing**: The `process_engine_inputs()` method in `OmniStage` (`omni_stage.py`) handles the transition:
+3. **Input Processing**: The stage input processor configured in stage YAML (under `vllm_omni/model_executor/stage_input_processors/`) handles the transition:
    ```python
    def process_engine_inputs(
        self, stage_list: list[Any], prompt: OmniTokensPrompt | TextPrompt = None

docs/contributing/model/adding_quantization_model.md

@@ -0,0 +1,478 @@
+# Support FP8 Quantization
+
+This section describes how to add FP8 quantization support to a diffusion model. We use the Qwen-Image pipeline as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Troubleshooting](#troubleshooting)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is FP8 Quantization?
+
+FP8 quantization reduces model memory footprint by storing weights in 8-bit floating point (`float8_e4m3fn`) instead of 16-bit (`bfloat16`/`float16`). This saves ~50% of weight memory.
+
+A diffusion pipeline typically has three quantizable components:
+
+| Component | Layer Types | Quantization Approach |
+|-----------|-----------|----------------------|
+| **DiT (transformer)** | `nn.Linear` | Replace with vLLM quantized linear layers |
+| **Text encoder** | `nn.Linear` | FP8 weight storage via hooks |
+| **VAE** | `nn.Conv2d`, `nn.Conv3d` | FP8 weight storage via hooks |
+
+### Architecture
+
+The quantization system has two mechanisms:
+
+**1. vLLM Quantized Linear Layers (for DiT)**
+
+vLLM provides drop-in replacements for `nn.Linear` that perform FP8 W8A8 computation on supported hardware. These are configured via `quant_config` passed to layer constructors.
+
+Key classes:
+
+| Class | Location | Purpose |
+|-------|----------|---------|
+| `build_quant_config` | `vllm_omni/quantization/__init__.py` | Build quantization config from string/dict/per-component spec |
+| `ComponentQuantizationConfig` | `vllm_omni/quantization/component_config.py` | Per-component routing by layer prefix |
+
+**2. FP8 Weight Storage via Hooks (for text encoder and VAE)**
+
+For models loaded via `from_pretrained()` (outside vLLM's weight pipeline), we use forward hooks to store weights in FP8 and dequantize on-the-fly:
+
+```
+At rest:     weight.data = fp8_tensor (half memory)
+Pre-hook:    weight.data = fp8_weight.to(bf16) * scale
+Forward:     Normal computation in BF16
+Post-hook:   weight.data = fp8_weight (back to FP8)
+```
+
+Key function:
+
+| Function | Location | Purpose |
+|----------|----------|---------|
+| `apply_fp8_weight_storage` | `vllm_omni/diffusion/models/utils.py` | Apply FP8 hooks to Linear/Conv layers (introduced in [#1412](https://github.com/vllm-project/vllm-omni/pull/1412)) |
+
+---
+
+## Step-by-Step Implementation
+
+### Step 1: Add `quant_config` to the Transformer
+
+Pass `quant_config` to vLLM parallel linear layers in the transformer. This enables FP8 W8A8 for DiT linear layers.
+
+#### 1.1: Accept `quant_config` in Constructor
+
+```python
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+class YourTransformer2DModel(nn.Module):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        quant_config: QuantizationConfig | None = None,
+        # ... other params
+    ):
+        self.quant_config = quant_config
+        # Pass to sub-modules
+        self.blocks = nn.ModuleList([
+            YourTransformerBlock(quant_config=quant_config, ...)
+            for _ in range(num_layers)
+        ])
+```
+
+#### 1.2: Pass `quant_config` to vLLM Linear Layers
+
+Replace `nn.Linear` with vLLM parallel linear layers that accept `quant_config`:
+
+```python
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+
+class YourAttentionBlock(nn.Module):
+    def __init__(self, hidden_size, num_heads, quant_config=None, prefix=""):
+        # QKV projection
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=hidden_size // num_heads,
+            total_num_heads=num_heads,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_qkv",
+        )
+        # Output projection
+        self.to_out = RowParallelLinear(
+            input_size=hidden_size,
+            output_size=hidden_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_out",
+        )
+
+class YourFeedForward(nn.Module):
+    def __init__(self, hidden_size, intermediate_size, quant_config=None, prefix=""):
+        self.gate_up_proj = ColumnParallelLinear(
+            input_size=hidden_size,
+            output_size=intermediate_size * 2,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+```
+
+!!! important
+    Always pass the `prefix` parameter. It is used by `ignored_layers` matching to skip sensitive layers.
+
+#### 1.3: Identify Sensitive Layers
+
+Some layers degrade output quality when quantized. Common sensitive layers:
+
+| Layer | Why Sensitive | Typical Models | Quant Methods Affected |
+|-------|-------------|---------------|----------------------|
+| `img_mlp` | Processes denoising latents with shifting dynamic range | Qwen-Image | FP8, Int8 |
+| `feed_forward` | FFN layers in DiT blocks, large dynamic range | Z-Image | Int8 |
+| `proj_out` | Final output projection, small quantization errors amplified | Various | FP8, Int8 |
+| `lm_head` | Output vocabulary projection, precision-critical | Omni (thinker) | NVFP4 |
+| `mlp.gate` | MoE router gates, routing decisions are precision-critical | MoE models | NVFP4 |
+
+**Quantization quality benchmarks** (LPIPS: lower = better, <0.01 imperceptible, >0.1 noticeable):
+
+| Model | Method | All Layers | With `ignored_layers` | Recommendation |
+|-------|--------|-----------|----------------------|----------------|
+| Qwen-Image-2512 | Int8 | LPIPS 0.0197 | 0.0027 (skip `img_mlp`) | Skip `img_mlp` |
+| Z-Image-Turbo | Int8 | LPIPS 0.1597 | 0.0290 (skip `feed_forward`) | Skip `feed_forward` |
+| Z-Image-Turbo | FP8 | LPIPS ~0.005 | — | All layers OK |
+
+To identify sensitive layers for a new model:
+
+1. Run inference with all layers quantized
+2. Compare output quality with BF16 baseline using LPIPS (see `benchmarks/diffusion/quantization_quality.py`)
+3. Selectively disable quantization per layer (via `ignored_layers`) until quality matches
+4. Document the recommended `ignored_layers` in the model's supported models table
+
+### Step 2: Wire Up Quantization in the Pipeline
+
+#### 2.1: Create `quant_config` and Pass to Transformer
+
+In the pipeline's `__init__`, extract the vLLM quant config and pass it to the transformer:
+
+```python
+from vllm_omni.quantization import build_quant_config
+
+class YourPipeline(nn.Module):
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__()
+        # ... load scheduler, text_encoder, vae ...
+
+        # Build quant config from OmniDiffusionConfig
+        quant_config = build_quant_config(od_config.quantization_config)
+
+        # Pass to transformer
+        self.transformer = YourTransformer2DModel(
+            od_config=od_config,
+            quant_config=quant_config,
+            # ... other kwargs
+        )
+```
+
+#### 2.2: Apply FP8 Weight Storage to Text Encoder and VAE
+
+For models loaded via `from_pretrained()`, apply hook-based FP8 weight storage.
+
+!!! note
+    The `apply_fp8_weight_storage` function is provided by `vllm_omni/diffusion/models/utils.py`, introduced in [PR #1412](https://github.com/vllm-project/vllm-omni/pull/1412).
+
+```python
+from vllm_omni.diffusion.models.utils import apply_fp8_weight_storage
+
+class YourPipeline(nn.Module):
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__()
+
+        self.text_encoder = SomeTextEncoder.from_pretrained(...)
+        self.vae = SomeVAE.from_pretrained(...).to(self.device)
+
+        # Apply FP8 weight storage to text encoder and VAE
+        if (
+            od_config.quantization_config is not None
+            and getattr(od_config.quantization_config, "quant_method", None) == "fp8"
+        ):
+            apply_fp8_weight_storage(self.vae)
+            apply_fp8_weight_storage(self.text_encoder)
+
+        # ... rest of init
+```
+
+`apply_fp8_weight_storage` walks all `nn.Linear`, `nn.Conv2d`, and `nn.Conv3d` modules in the model and applies FP8 hooks to each.
+
+#### 2.3: Fix `load_weights` to Include Pre-loaded Components
+
+Since text encoders and VAEs are loaded via `from_pretrained()` (not through vLLM's weight pipeline), their parameters won't appear in the `loaded_weights` set. Mark them explicitly:
+
+```python
+def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+    loader = AutoWeightsLoader(self)
+    loaded_weights = loader.load_weights(weights)
+    # Mark pre-loaded component weights as loaded
+    loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
+    loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
+    return loaded_weights
+```
+
+### Step 3: Add to Supported Models Documentation
+
+Update the FP8 supported models table in `docs/user_guide/diffusion/quantization/fp8.md`:
+
+```markdown
+| Your-Model | `org/Your-Model` | ✅ | ✅ | ✅ | `sensitive_layer` |
+```
+
+---
+
+## Testing
+
+### Quick Verification
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+# Test with FP8 enabled
+omni = Omni(model="your-model", quantization="fp8")
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+### Quality Gate Test (LPIPS)
+
+We provide a pytest-based quality gate at `tests/diffusion/quantization/test_quantization_quality.py`.
+It generates outputs with both BF16 and your quantized method using the same seed, computes
+[LPIPS](https://github.com/richzhang/PerceptualSimilarity) perceptual distance, and **fails if it
+exceeds a threshold**. This is the recommended way to validate that a quantization method does not
+unacceptably degrade output quality.
+
+**Step 1: Add your test config** — edit `QUALITY_CONFIGS` in the test file:
+
+```python
+QualityTestConfig(
+    id="int8_z_image",              # pytest ID
+    model="Tongyi-MAI/Z-Image-Turbo",
+    quantization="int8",            # your quantization method
+    task="t2i",                     # "t2i" or "t2v"
+    prompt="a cup of coffee on a wooden table, morning light",
+    max_lpips=0.15,                 # fail threshold (see guide below)
+    num_inference_steps=20,
+),
+```
+
+**Step 2: Run the test:**
+
+```bash
+# Install dependency
+pip install lpips
+
+# Run all quality tests
+pytest tests/diffusion/quantization/test_quantization_quality.py -v -m ""
+
+# Run only your method
+pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "int8"
+
+# Run a specific model
+pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image"
+```
+
+**Step 3: Read the output:**
+
+```
+Quantization Quality: int8_z_image
+  Model:         Tongyi-MAI/Z-Image-Turbo
+  Method:        int8
+  LPIPS:         0.0290  (threshold: 0.15)
+  BF16 memory:   19.15 GiB
+  Quant memory:  14.80 GiB  (23% reduction)
+  Result:        PASS
+```
+
+**LPIPS threshold guide:**
+
+| LPIPS Range | Meaning | Recommendation |
+|-------------|---------|----------------|
+| < 0.01 | Imperceptible | Ideal — no visible difference |
+| 0.01 – 0.05 | Minor differences | Good — acceptable for most use cases |
+| 0.05 – 0.15 | Noticeable on close inspection | Acceptable with `ignored_layers` tuning |
+| 0.15 – 0.30 | Clearly visible | Investigate — likely needs `ignored_layers` |
+| > 0.30 | Significant degradation | Quantization method may not suit this model |
+
+!!! tip
+    If your test fails, don't raise the threshold — instead, identify sensitive layers and add
+    them to `ignored_layers` (see next section). Then re-run the test with the lower LPIPS.
+
+### Quality Benchmark (Detailed)
+
+For a full multi-prompt benchmark with Markdown output (suitable for PR descriptions), use the
+standalone benchmark script:
+
+```bash
+python benchmarks/diffusion/quantization_quality.py \
+    --model Tongyi-MAI/Z-Image-Turbo \
+    --quantization fp8 \
+    --prompts \
+        "a cup of coffee on the table" \
+        "an aerial view of a coral reef" \
+    --height 1024 --width 1024 \
+    --num-inference-steps 50 --seed 42
+```
+
+This generates BF16 and quantized outputs, computes per-prompt LPIPS, and prints a Markdown
+summary table ready to paste into your PR.
+
+### Identifying Sensitive Layers with `ignored_layers`
+
+Not all layers tolerate quantization equally. **`ignored_layers`** lets you keep specific layers
+in BF16 while quantizing everything else. This is critical for achieving good quality with
+aggressive quantization methods (Int8, NVFP4).
+
+**Why some layers are sensitive:**
+
+| Layer Type | Why Sensitive | Typical Impact |
+|------------|--------------|----------------|
+| `img_mlp` | Processes denoising latents with shifting dynamic range across timesteps. No built-in normalization to absorb quantization error. | Color shifts, blurriness |
+| `feed_forward` | FFN layers with large dynamic range in DiT blocks. | Artifacts, loss of detail |
+| `proj_out` | Final output projection — small errors get amplified since there are no subsequent layers to correct them. | Overall quality loss |
+| `lm_head` | Vocabulary projection in LLMs — precision-critical for token selection. | Garbage text output |
+| `mlp.gate` | MoE router gates — routing decisions are binary and precision-critical. | Wrong expert selection, quality collapse |
+
+**How to find sensitive layers for your model:**
+
+1. Run the quality gate test with all layers quantized:
+    ```bash
+    pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "your_model"
+    ```
+
+2. If LPIPS exceeds your threshold, try skipping common sensitive layers one at a time:
+    ```python
+    QualityTestConfig(
+        id="fp8_your_model_skip_mlp",
+        model="your-org/Your-Model",
+        quantization="fp8",
+        task="t2i",
+        prompt="a cup of coffee on a wooden table",
+        max_lpips=0.05,
+    ),
+    ```
+
+    And in the test, modify the Omni call:
+    ```python
+    omni_qt = Omni(
+        model=config.model,
+        quantization_config={
+            "method": config.quantization,
+            "ignored_layers": ["img_mlp"],
+        },
+    )
+    ```
+
+3. Compare LPIPS with and without `ignored_layers` to isolate the problematic layers.
+
+4. Document the recommended `ignored_layers` in the supported models table.
+
+**Real-world examples:**
+
+| Model | Method | All Layers (LPIPS) | With `ignored_layers` | Which Layers to Skip |
+|-------|--------|-------------------|----------------------|---------------------|
+| Qwen-Image-2512 | Int8 | 0.0197 | 0.0027 | `img_mlp` |
+| Z-Image-Turbo | Int8 | 0.1597 | 0.0290 | `feed_forward` |
+| Z-Image-Turbo | FP8 | ~0.005 | — | None needed |
+
+### What to Verify (Checklist)
+
+1. **Model loads without errors** — Check logs for quantization initialization messages
+2. **Quality gate passes** — LPIPS within threshold vs BF16 baseline
+3. **Memory reduction** — Expect ~30-50% savings on quantized components
+4. **Sensitive layers identified** — Document which layers (if any) need `ignored_layers`
+5. **`ignored_layers` pattern works** — Verify that the layer name prefix matches your transformer's module names
+
+---
+
+## Troubleshooting
+
+### Issue: Output Quality Degradation
+
+**Symptoms:** Generated images have visible artifacts, color shifts, or blurriness compared to BF16 baseline.
+
+**Causes & Solutions:**
+
+- **Sensitive layers being quantized:**
+
+    **Solution:** Identify and skip sensitive layers:
+    ```python
+    omni = Omni(
+        model="your-model",
+        quantization_config={
+            "method": "fp8",
+            "ignored_layers": ["img_mlp"],
+        },
+    )
+    ```
+
+- **All layers sensitive (rare):**
+
+    **Solution:** The model may not be suitable for FP8 quantization. Consider weight-only quantization or other methods.
+
+### Issue: Missing Weights Warning
+
+**Symptoms:** Warnings about unloaded weights for `vae.*` or `text_encoder.*` parameters.
+
+**Solution:** Ensure `load_weights` marks pre-loaded component weights. See [Step 2.3](#23-fix-load_weights-to-include-pre-loaded-components).
+
+### Issue: CUDA Out of Memory
+
+**Symptoms:** OOM errors during model loading despite FP8 being enabled.
+
+**Causes & Solutions:**
+
+- **Text encoder/VAE loaded in BF16 before FP8 conversion:**
+
+    FP8 weight storage is applied after `from_pretrained()`, so peak load memory is still BF16. This is a one-time cost during initialization.
+
+---
+
+## Reference Implementations
+
+| Model | Pipeline File | Transformer File | Notes |
+|-------|-------------|-----------------|-------|
+| **Z-Image** | `models/z_image/pipeline_z_image.py` | `models/z_image/z_image_transformer.py` | DiT + text encoder FP8 |
+| **Qwen-Image** | `models/qwen_image/pipeline_qwen_image.py` | `models/qwen_image/qwen_image_transformer.py` | DiT + text encoder + VAE FP8 |
+
+All files are under `vllm_omni/diffusion/`.
+
+---
+
+## Summary
+
+Adding FP8 quantization support to a new model:
+
+1. ✅ **Add `quant_config` to transformer** — Accept and pass to vLLM parallel linear layers with proper `prefix`
+2. ✅ **Identify sensitive layers** — Profile which layers need `ignored_layers` for quality
+3. ✅ **Apply FP8 weight storage** — Call `apply_fp8_weight_storage()` on text encoder and VAE in pipeline `__init__`
+4. ✅ **Fix `load_weights`** — Mark pre-loaded component weights as loaded
+5. ✅ **Update docs** — Add model to the supported models table in `docs/user_guide/diffusion/quantization/fp8.md`
+6. ✅ **Test** — Verify quality, memory savings, and sensitive layer behavior

docs/contributing/model/adding_tts_model.md

@@ -361,9 +361,9 @@ stage_args:
     stage_type: llm
     runtime:
       devices: "0"
-      max_batch_size: 64
     engine_args:
       model_stage: ar_stage
+      max_num_seqs: 64
       model_arch: YourTTSModelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -403,9 +403,9 @@ stage_args:
     stage_type: llm
     runtime:
       devices: "0"
-      max_batch_size: 64
     engine_args:
       model_stage: ar_stage
+      max_num_seqs: 64
       model_arch: YourTTSModelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler

docs/contributing/profiling.md

@@ -4,49 +4,57 @@
 
 vLLM-Omni uses the PyTorch Profiler to analyze performance across both **multi-stage omni-modality models** and **diffusion models**.
 
-### 1. Set the Output Directory
-Before running any script, set this environment variable. The system detects this and automatically saves traces here.
+### 1. Configure Profiling in the Stage YAML
 
-```bash
-export VLLM_TORCH_PROFILER_DIR=./profiles
+Enable profiling by adding `profiler_config` under `engine_args` for the stage(s) you want to profile in your stage config YAML:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    engine_args:
+      # ... other engine args ...
+      profiler_config:
+        profiler: torch
+        torch_profiler_dir: ./perf
 ```
 
+| Field | Description |
+|---|---|
+| `profiler` | Profiler backend to use. Currently supports `torch`. |
+| `torch_profiler_dir` | Directory where trace files are saved. Created automatically if it doesn't exist. |
+
+> **Tip:** Only enable `profiler_config` on stages you actually need to profile. Stages without it will not start a profiler, keeping overhead minimal.
+
 ### 2. Profiling Omni-Modality Models
 
-It is best to limit profiling to one iteration to keep trace files manageable.
-
-```bash
-export VLLM_PROFILER_MAX_ITERS=1
-```
-
 **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:
+
+It is highly recommended to profile specific stages to prevent producing overly large trace files:
+
 ```python
 # Profile all stages
 omni_llm.start_profile()
 
 # 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])
 ```
 
+> **Important:** Always pass the same `stages` list to both `start_profile()` and `stop_profile()`. If you omit `stages` from `stop_profile()`, it defaults to stopping all stages — including ones that were never started — which will produce errors.
+
 **Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
 
 ```python
-from vllm_omni import omni_llm
+profiler_stages = [0]  # Only profile the stages you need
 
-profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
-
-# 1. Start profiling if enabled
-if profiler_enabled:
-    omni_llm.start_profile(stages=[0])
+# 1. Start profiling
+omni.start_profile(stages=profiler_stages)
 
 # Initialize generator
-omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
 total_requests = len(prompts)
 processed_count = 0
@@ -64,7 +72,8 @@ for stage_outputs in omni_generator:
         print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
 
         # Stop the profiler while workers are still active
-        omni_llm.stop_profile()
+        # Pass the same stages list used in start_profile()
+        omni_llm.stop_profile(stages=profiler_stages)
 
         # Wait for traces to flush to disk
         print("[Info] Waiting 30s for workers to write trace files to disk...")
@@ -75,24 +84,38 @@ omni_llm.close()
 ```
 
 
+**CLI Usage** (using `end2end.py`):
+```bash
+# Profile only Stage 0 (Thinker)
+python end2end.py --output-wav output_audio \
+    --query-type text --enable-profiler --profiler-stages 0
+
+# Profile Stage 0 and Stage 2
+python end2end.py --output-wav output_audio \
+    --query-type text --enable-profiler --profiler-stages 0 2
+
+# Profile all stages (omit --profiler-stages)
+python end2end.py --output-wav output_audio \
+    --query-type text --enable-profiler
+```
+
 **Examples**:
 
 1. **Qwen2.5-Omni**:  [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. **Qwen3-Omni**:   [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. Profiling diffusion models
 
-Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding.
+Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding. Standalone diffusion scripts use `--profiler-dir` to enable profiling.
 
 **CLI Usage:**
-```python
-
+```bash
 python image_to_video.py \
     --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
     --image qwen-bear.png \
     --prompt "A cat playing with yarn, smooth motion" \
+    --profiler-dir \
     \
     # Minimize Spatial Dimensions (Optional but helpful):
     #    Drastically reduces memory usage so the profiler doesn't
@@ -122,25 +145,72 @@ python image_to_video.py \
     --flow-shift 12.0 \
     --fps 16 \
     --output i2v_output.mp4
-
 ```
 
+> **Note:** For diffusion stages within a multi-stage omni pipeline, use `profiler_config` in the stage YAML instead (see Section 1).
+
 **Examples**:
 
 1. **Qwen image edit**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
 
 2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**:   [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
 
-### 4. Analyzing Omni Traces
+### 4. Profiling Online Serving
 
-Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
+When `profiler_config` is set in the stage YAML, the server automatically exposes `/start_profile` and `/stop_profile` HTTP endpoints.
+
+**1. Start the server** with a stage YAML that has `profiler_config` enabled:
+```bash
+vllm serve Qwen/Qwen2.5-Omni-7B \
+    --omni \
+    --stage-configs-path qwen2_5_omni.yaml \
+    --port 8091
+```
+
+Or for one stage diffusion models:
+
+```bash
+vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --port 8091 --profiler-config '{"profiler": "torch", "torch_profiler_dir": "./vllm_profile"}'
+```
+
+**2. Start profiling** by sending a POST request:
+```bash
+# Profile all stages that have profiler_config set
+curl -X POST http://localhost:8091/start_profile
+
+# Profile specific stages only
+curl -X POST http://localhost:8091/start_profile \
+    -H "Content-Type: application/json" \
+    -d '{"stages": [0]}'
+```
+
+**3. Send your inference requests** as normal while the profiler is running.
+
+**4. Stop profiling** and collect traces:
+```bash
+# Stop all stages
+curl -X POST http://localhost:8091/stop_profile
+
+# Stop specific stages (must match the stages you started)
+curl -X POST http://localhost:8091/stop_profile \
+    -H "Content-Type: application/json" \
+    -d '{"stages": [0]}'
+```
+
+Trace files are written to the `torch_profiler_dir` specified in your stage YAML.
+
+> **Important:** Always stop the same stages you started. Stopping a stage that was never started will produce errors.
+
+### 5. Analyzing Traces
+
+Output files are saved to the `torch_profiler_dir` specified in your stage YAML config.
 
 **Output**
-**Chrome Trace** (```.json.gz```): Visual timeline of kernels and stages. Open in Perfetto UI.
+**Chrome Trace** (`.json.gz`): Visual timeline of kernels and stages. Open in Perfetto UI.
 
 **Viewing Tools:**
 
-- [Perfetto](https://ui.perfetto.dev/)(recommended)
-- ```chrome://tracing```(Chrome only)
+- [Perfetto](https://ui.perfetto.dev/) (recommended)
+- `chrome://tracing` (Chrome only)
 
 **Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation:  [vLLM Profiling Guide](https://docs.vllm.ai/en/stable/contributing/profiling/)

docs/design/architecture_overview.md

@@ -67,12 +67,12 @@ According to analysis for current popular open-source models, most of them have
 | 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 |
+| **EntryPoints**   | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni), while `AsyncOmniEngine` and `Orchestrator` coordinate multi-stage AR/DiT execution |
 | **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.
+Disaggregated stages are managed through stage configuration. In Qwen3-Omni, Thinker/Talker/Code2wav are declared as separate configured stages, and runtime routing is handled by `Orchestrator` over `StageEngineCoreClient` / `StageDiffusionClient`.
 
 ## Main features
 
@@ -127,10 +127,10 @@ Taking **Qwen3-Omni** as an example:
 The **Omni** class provides a Python interface for offline batched inference. Users initialize the Omni class with a Hugging Face model name and use the generate method, passing inputs that include both text prompts and multi-modal data:
 
 ```
-# Create an omni_lm with HF model name.
+# Create an omni runtime with HF model name.
 from vllm_omni.entrypoints.omni import Omni
 
-omni_lm = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
+omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
 
 # Example prompts.
 om_inputs = {"prompt": prompt,
@@ -140,7 +140,7 @@ om_inputs = {"prompt": prompt,
              }}
 
 # Generate texts and audio from the multi-modality inputs.
-outputs = omni_lm.generate(om_inputs, sampling_params_list)
+outputs = omni.generate(om_inputs, sampling_params_list)
 ```
 
 ## Online Serving

docs/design/feature/async_chunk_design.md

@@ -171,7 +171,7 @@ stage_args:
 - `custom_process_next_stage_input_func: str`: Path to custom chunk processing function; receives `(transfer_manager, pooling_output, request)`. For qwen3-omni: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
 - `stage_connector_config: dict`: Connector configuration
 - `worker_type: str`: Model type, e.g. `"ar"` or `"generation"` (used by OmniChunkTransferAdapter for mode-specific payload handling)
-- `max_batch_size: int`: Maximum batch size for the stage
+- `max_num_seqs: int`: Maximum number of sequences for concurrent processing in the stage
 
 
 ### Connector Configuration
@@ -195,9 +195,9 @@ stage_args:
   - stage_id: 2  # code2wav stage
     runtime:
       devices: "1"
-      max_batch_size: 64  # Enables batched audio generation
     engine_args:
       model_stage: code2wav
+      max_num_seqs: 64  # Enables batched audio generation
 ```
 
 ## Related Files

docs/design/feature/cache_dit.md

@@ -113,13 +113,13 @@ def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], No
 
 Some models require custom handling:
 
-- **Dual-transformer:** Models with separate high-noise and low-noise transformers (e.g., Wan2.2)
+- **Single or dual-transformer:** Models that may use one or two transformers (e.g., Wan2.2)
 - **Multi-block-list:** Models with multiple block lists in one transformer (e.g., LongCatImage with `transformer_blocks` + `single_transformer_blocks`)
 - **Special forward patterns:** Models with non-standard block execution patterns
 
-### Example 1: Dual-Transformer Model (Wan2.2)
+### Example 1: Single or Dual-Transformer Model (Wan2.2)
 
-Wan2.2 uses two transformers: one for high-noise steps and one for low-noise steps.
+Wan2.2 can use either a single transformer or two transformers (one for high-noise steps and one for low-noise steps). The implementation automatically detects the mode based on the presence of `transformer_2`.
 
 **Key difference:** Use `BlockAdapter` to wrap multiple transformers with separate configurations.
 
@@ -256,7 +256,7 @@ cache_config={
 }
 ```
 
-Check the [user guide for cache_dit](../../user_guide/diffusion/cache_dit_acceleration.md) for more adjustable parameters.
+Check the [user guide for cache_dit](../../user_guide/diffusion/cache_acceleration/cache_dit.md) for more adjustable parameters.
 
 ---
 
@@ -267,7 +267,7 @@ Complete examples in the codebase:
 | Model | Path | Pattern | Notes |
 |-------|------|---------|-------|
 | **Standard DiT** | [`cache_dit_backend.py::enable_cache_for_dit`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Default enabler | Single transformer, automatic |
-| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Dual-transformer | Separate high/low noise transformers |
+| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Single or dual-transformer | Auto-detects mode based on transformer_2 presence |
 | **LongCat** | [`cache_dit_backend.py::enable_cache_for_longcat_image`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_longcat_image) | Multi-block-list | Two block lists in one transformer |
 | **BAGEL** | [`cache_dit_backend.py::enable_cache_for_bagel`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_bagel) | Omni model | Complex architecture |
 

docs/design/feature/cfg_parallel.md

@@ -34,7 +34,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
 | Method | Purpose | Automatic Behavior |
 |--------|---------|-------------------|
 | [`predict_noise_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with CFG | Detects parallel mode, distributes computation, gathers results |
-| [`scheduler_step_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Step scheduler with sync | Rank 0 steps, broadcasts latents to all ranks |
+| [`scheduler_step_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Step scheduler | All ranks step locally (no broadcast needed) |
 | [`combine_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine positive/negative | Applies CFG formula with optional normalization |
 | [`predict_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Forward pass wrapper | Override for custom transformer calls |
 | [`cfg_normalize_function()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Normalize CFG output | Override for custom normalization |
@@ -47,7 +47,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
   - Rank 0 computes positive prompt prediction
   - Rank 1 computes negative prompt prediction
   - Results are gathered via `all_gather()`
-  - Combined on rank 0 using CFG formula
+  - All ranks compute CFG combine locally (deterministic, identical results)
 
 - **Sequential mode** (when `cfg_world_size == 1`):
   - Single rank computes both positive and negative predictions
@@ -55,14 +55,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
 
 `scheduler_step_maybe_with_cfg()` ensures consistent latent states across all ranks:
 
-- **CFG-Parallel mode**:
-  - Only rank 0 performs the scheduler step (applies noise prediction to update latents)
-  - Updated latents are broadcast to all other ranks via `broadcast()`
-  - All ranks maintain synchronized latent states for the next iteration
-
-- **Sequential mode**:
-  - Single rank directly performs the scheduler step
-  - No synchronization needed
+- All ranks compute the scheduler step locally — no broadcast needed because `predict_noise_maybe_with_cfg` already ensures all ranks have identical noise predictions after `all_gather` + local combine.
 
 ---
 
@@ -177,6 +170,70 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin):
         # return noise_pred
 ```
 
+
+### Override `combine_cfg_noise()` for Multi-Output Models
+
+When `predict_noise()` returns a tuple (e.g., video + audio), the default `combine_cfg_noise()` applies CFG to every element. Override it to apply different logic per element — for example, CFG on video but positive-only on audio:
+
+```python
+class MyVideoAudioPipeline(nn.Module, CFGParallelMixin):
+    def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, scale, normalize):
+        (video_pos, audio_pos) = positive_noise_pred
+        (video_neg, audio_neg) = negative_noise_pred
+        video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize)
+        return (video_combined, audio_pos)  # audio: positive only, no CFG
+```
+
+This also requires `predict_noise()` to return a tuple (see [Override predict_noise](#override-predict_noise-for-custom-transformer-calls) above).
+
+### Implement a Composite Scheduler for Multi-Output Models
+
+When each output has its own denoising schedule, implement a composite scheduler that dispatches to per-output schedulers. Assign it to `self.scheduler` so the default `scheduler_step()` works without override.
+
+**Complete example (video + audio with separate schedulers and diffuse loop):**
+
+```python
+class VideoAudioScheduler:
+    """Composite scheduler dispatching to video and audio schedulers."""
+    def __init__(self, video_scheduler, audio_scheduler):
+        self.video_scheduler = video_scheduler
+        self.audio_scheduler = audio_scheduler
+
+    def step(self, noise_pred, t, latents, return_dict=False, generator=None):
+        video_out = self.video_scheduler.step(noise_pred[0], t[0], latents[0], return_dict=False, generator=generator)[0]
+        audio_out = self.audio_scheduler.step(noise_pred[1], t[1], latents[1], return_dict=False, generator=generator)[0]
+        return ((video_out, audio_out),)
+
+class MyVideoAudioPipeline(nn.Module, CFGParallelMixin):
+    def __init__(self, ...):
+        self.scheduler = VideoAudioScheduler(video_sched, audio_sched)
+
+    def predict_noise(self, **kwargs):
+        video_pred, audio_pred = self.transformer(**kwargs)
+        return (video_pred, audio_pred)
+
+    def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, scale, normalize):
+        # ... (as above)
+
+    def diffuse(self, video_latents, audio_latents, timesteps_video, timesteps_audio, ...):
+        for t_v, t_a in zip(timesteps_video, timesteps_audio):
+            positive_kwargs = {...}  
+            negative_kwargs = {...} if do_true_cfg else None
+
+            video_pred, audio_pred = self.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_true_cfg, true_cfg_scale=self.guidance_scale,
+                positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs,
+            )
+            video_latents, audio_latents = self.scheduler_step_maybe_with_cfg(
+                (video_pred, audio_pred), (t_v, t_a),
+                (video_latents, audio_latents), do_true_cfg=do_true_cfg,
+                generator=generator,
+            )
+        return video_latents, audio_latents
+```
+
+> **Note:** If you use a non-deterministic scheduler, e.g., DDPM, please set `self.scheduler_step_maybe_with_cfg(..., generator=torch.Generator(device).manual_seed(seed))` explicitly to control the randomness of scheduler step among ranks.
+
 ---
 
 ## Testing

docs/design/feature/diffusion_step_execution.md

@@ -0,0 +1,121 @@
+# Diffusion Step Execution
+
+This guide documents vLLM-Omni's stepwise diffusion contract for model authors
+and contributors implementing `step_execution=True` support for a diffusion
+pipeline.
+
+For end-user enablement, supported models, and current limitations, see
+[Step Execution](../../user_guide/diffusion/step_execution.md).
+
+## Current Support Scope
+
+`step_execution` is **not** a generic diffusion toggle. It only works for
+pipelines that implement the segmented stateful contract in
+[`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py).
+
+Current in-tree support:
+
+| Pipeline | Example models | Step execution |
+|----------|----------------|----------------|
+| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
+| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
+
+Current engine/runtime limitations:
+
+- `StepScheduler` only schedules `batch_size=1`.
+- `cache_backend` is not supported in step mode.
+- Request-mode extras such as KV transfer are not wired into step mode yet.
+- Unsupported pipelines now fail early during model loading instead of failing on the first request.
+
+## Execution Contract
+
+Step mode is driven by four pipeline methods plus the shared mutable request
+state object:
+
+- `prepare_encode(state)`: one-time request preparation.
+- `denoise_step(state)`: compute the noise prediction for the current step.
+- `step_scheduler(state, noise_pred)`: mutate latents and advance step state.
+- `post_decode(state)`: decode the final output after denoising is complete.
+
+The state lives in
+[`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
+as `DiffusionRequestState`. Store request-scoped tensors there, or use
+`state.extra` for model-specific fields that do not justify extending the core
+dataclass.
+
+The worker-side step loop lives in
+[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py):
+
+1. `prepare_encode()` runs once for a new request.
+2. `denoise_step()` runs every scheduler tick.
+3. `step_scheduler()` mutates `state.latents` and advances `state.step_index`.
+4. `post_decode()` runs exactly once after `state.denoise_completed` becomes true.
+
+## Recommended Split
+
+When converting an existing request-level `forward()` pipeline, keep the split
+strict and mechanical:
+
+| Request-level phase | Stepwise method | What belongs there |
+|---------------------|-----------------|--------------------|
+| Input validation, prompt encoding, latent init, timestep prep, per-request scheduler creation | `prepare_encode()` | Anything that should happen once per request |
+| Transformer forward / noise prediction | `denoise_step()` | Pure denoise computation for the current timestep |
+| `scheduler.step(...)` and `step_index += 1` | `step_scheduler()` | Only latent/state mutation for one step |
+| VAE decode / postprocess | `post_decode()` | Final decode only |
+
+Keep the stepwise path reusing the same helpers as the request-level path
+whenever possible. Reimplementing the denoise loop from scratch is the easiest
+way to introduce behavioral drift.
+
+## Qwen-Image Reference
+
+[`pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)
+is the reference implementation and is split correctly for the current
+contract:
+
+- `prepare_encode()` reuses `_prepare_generation_context()` so prompt encoding,
+  latent init, timestep creation, CFG setup, and shape bookkeeping stay aligned
+  with `forward()`.
+- `prepare_encode()` deep-copies `self.scheduler` **after**
+  `prepare_timesteps()` so request-specific scheduler state is isolated.
+- `denoise_step()` reuses `_build_denoise_kwargs()` plus
+  `predict_noise_maybe_with_cfg()`, so sequential CFG, CFG-parallel, and
+  non-CFG behavior stay identical to the request-level path.
+- `step_scheduler()` only calls
+  `scheduler_step_maybe_with_cfg(..., per_request_scheduler=state.scheduler)`
+  and increments `state.step_index`.
+- `post_decode()` reuses `_decode_latents()`, so the final image decode matches
+  the normal `forward()` path.
+
+That decomposition is the target pattern for future models.
+
+## Rules For New Pipelines
+
+- Do not keep request-scoped scheduler state on `self.scheduler`. Copy it into
+  `state.scheduler` during `prepare_encode()`.
+- Do not mutate `state.step_index` inside `denoise_step()`. Only
+  `step_scheduler()` should advance the step.
+- Do not decode partial outputs in `denoise_step()` or `step_scheduler()`.
+- If the request-level pipeline has condition latents, masks, or edit-specific
+  tensors, store them in `state` or `state.extra`, not in global pipeline
+  attributes.
+- Preserve CFG behavior by sharing the same helper path used by `forward()`.
+- Keep `post_decode()` equivalent to the tail of `forward()`.
+
+## Validation Checklist
+
+Before marking a pipeline as `supports_step_execution = True`, verify:
+
+- Stepwise output matches request-level output for the same seed and sampling params.
+- Per-request scheduler state is isolated across concurrent requests.
+- Abort during denoise does not leak cached state.
+- `step_index` reported by `RunnerOutput` matches the scheduler progress.
+- CFG-parallel and non-CFG paths both work if the request-level pipeline supports them.
+
+## Related Files
+
+- Contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py)
+- State: [`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
+- Runner loop: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
+- Scheduler transport: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
+- Reference pipeline: [`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)

docs/design/feature/sequence_parallel.md

@@ -7,6 +7,7 @@ This section describes how to add Sequence Parallel (SP) to a diffusion transfor
 ## Table of Contents
 
 - [Overview](#overview)
+- [UAA Mode (Experimental)](#uaa-mode-experimental)
 - [Approach 1: Non-Intrusive `_sp_plan` (Recommended)](#approach-1-non-intrusive-_sp_plan-recommended)
 - [Approach 2: Intrusive Modification (For Complex Cases)](#approach-2-intrusive-modification-for-complex-cases)
 - [Testing](#testing)
@@ -46,6 +47,41 @@ from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
 
 ---
 
+## UAA Mode (Experimental)
+
+`ulysses_mode="advanced_uaa"` enables the experimental UAA ("Ulysses Anything Attention") feature, which lets Ulysses attention handle arbitrary sequence lengths and arbitrary attention head counts. The same idea is also supported by [Cache-DiT](https://cache-dit.readthedocs.io/en/latest/user_guide/CONTEXT_PARALLEL/#uaa-ulysses-anything-attention).
+
+Use it when plain Ulysses-SP would otherwise fail because:
+
+- the local sequence shards are not evenly divisible after split hooks, or
+- the attention head count is not divisible by `ulysses_degree`.
+
+### Design Summary
+
+1. **Strict mode stays unchanged.**
+   `ulysses_mode="strict"` keeps the original fast path and still requires divisible sequence/head shapes.
+
+2. **UAA uses variable all-to-all split sizes for sequence shards.**
+   Before the Ulysses Q/K/V exchange, each rank all-gathers its local sequence length and uses those lengths as `all_to_all_single(..., output_split_sizes=seq_lens)`. This lets Ulysses gather the full sequence even when each rank started with a different local shard length.
+
+3. **UAA pads heads only inside the Ulysses exchange.**
+   If `head_cnt % ulysses_degree != 0`, UAA pads the head dimension up to the next multiple of `ulysses_degree`, performs the forward/reverse all-to-all, then slices the temporary head padding away after the reverse exchange. The same rule is applied to joint attention tensors.
+
+4. **Hybrid Ulysses + Ring is still shape-constrained.**
+   Ring attention expects every rank in a ring group to exchange exactly the same post-Ulysses sequence shape. UAA therefore validates those shapes before entering the ring path and raises a clear error if ring peers disagree on `S_global`.
+
+5. **Tiny scalar gathers stay out of TorchDynamo tracing.**
+   `_all_gather_int()` is marked with `@torch.compiler.disable` so the scalar `item()` conversions used by UAA metadata collection do not get traced into `torch.compile`.
+
+### UAA vs `auto_pad`
+
+- `auto_pad=True` pads sequence tokens in `_sp_plan` and requires attention backends that support `attention_mask`.
+- `advanced_uaa` does not depend on mask-based token padding inside Ulysses attention. It is therefore a better fit for non-divisible head counts and uneven Ulysses shard sizes.
+- `auto_pad=True` remains incompatible with Ring attention because the ring backend does not consume `attention_mask`.
+- `advanced_uaa` is still experimental and hybrid mode remains limited by Ring's equal-shape requirement.
+
+---
+
 ## Approach 1: Non-Intrusive `_sp_plan` (Recommended)
 
 The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries.
@@ -364,9 +400,14 @@ _sp_plan = {
 
 - **Sequence Length not divisible by sp_size:**
 
-**Problem:** `SequenceParallelInput(auto_pad=False)` - auto_pad should be True to enable automatic sequence padding.
+**Problem:** strict Ulysses sequence parallel requires divisible shapes. If the shard length is uneven, or if the model head count is not divisible by `ulysses_degree`, the strict path will raise an error.
 
-**Solution:** In `SequenceParallelInput`, set `auto_pad=True` and add attention mask support.
+**Solutions:**
+
+1. Use `ulysses_mode="advanced_uaa"` for Ulysses-SP when you want the experimental uneven-shape path without relying on attention-mask padding.
+2. If the model already uses `_sp_plan` token padding and the attention backend supports `attention_mask`, set `auto_pad=True` and add attention-mask plumbing.
+
+> **Experimental Feature:** `ulysses_mode="advanced_uaa"` is experimental. It is intended to relax Ulysses divisibility constraints, but hybrid Ulysses + Ring still requires equal post-Ulysses sequence lengths inside each ring group.
 
 > **Experimental Feature:** `auto_pad=True` is an experimental feature and may be changed in the future. We plan to improve this solution to involve minimal changes to model files. More details are [here](https://github.com/vllm-project/vllm-omni/issues/1324).
 

docs/design/feature/teacache.md

@@ -369,7 +369,7 @@ images = omni.generate(
 2. **Compare performance** - Measure speedup vs baseline (expect 1.5x-2.0x)
 3. **Verify output quality** - Visually compare cached vs uncached outputs (should be nearly identical)
 
-See more detailed examples in [user guide for teacache](../../user_guide/diffusion/teacache.md).
+See more detailed examples in [user guide for teacache](../../user_guide/diffusion/cache_acceleration/teacache.md).
 
 ---
 

docs/design/module/ar_module.md

@@ -85,9 +85,6 @@ classDiagram
     class InputProcessor {
         +process_inputs() EngineCoreRequest
     }
-    class OmniInputProcessor {
-        +process_inputs() OmniEngineCoreRequest
-    }
 
     class VLLMOutputProcessor {
         +process_outputs() OutputProcessorOutput
@@ -96,7 +93,6 @@ classDiagram
         +process_outputs() OutputProcessorOutput
         +_route_and_normalize()
     }
-    InputProcessor <|-- OmniInputProcessor
     VLLMOutputProcessor <|-- MultimodalOutputProcessor
 ```
 
@@ -105,7 +101,7 @@ classDiagram
 - **Scheduler**: `OmniARScheduler` extends `vllm.v1.core.sched.scheduler.Scheduler` to enrich scheduled requests with omni-specific payloads
 - **Worker**: `GPUARWorker` extends `vllm.v1.worker.gpu_worker.Worker` to initialize AR-specific model runners
 - **ModelRunner**: `GPUARModelRunner` extends `OmniGPUModelRunner` → `vllm.v1.worker.gpu_model_runner.GPUModelRunner` to expose hidden states and handle multimodal outputs
-- **InputProcessor**: `OmniInputProcessor` extends `vllm.v1.engine.input_processor.InputProcessor` to serialize prompt embeddings and additional information
+- **InputProcessor**: Stage-0 uses upstream `vllm.v1.engine.input_processor.InputProcessor`; `AsyncOmniEngine` then restores omni-specific payloads (for example `additional_information` and `prompt_embeds`) when building `OmniEngineCoreRequest`
 - **OutputProcessor**: `MultimodalOutputProcessor` extends `vllm.v1.engine.output_processor.OutputProcessor` to route and accumulate multimodal outputs
 
 ## 3. Scheduler Design
@@ -118,7 +114,7 @@ The following diagram illustrates the request flow through the AR module compone
 
 ```mermaid
 flowchart TD
-    A[OmniInputProcessor] -->|OmniEngineCoreRequest| B[OmniARScheduler]
+    A[InputProcessor stage-0 in AsyncOmniEngine] -->|EngineCoreRequest then upgraded to OmniEngineCoreRequest| B[OmniARScheduler]
     B -->|schedule: OmniNewRequestData| C[GPUARWorker]
     C -->|SchedulerOutput| D[GPUARModelRunner]
     D -->|execute_model: None| E[Model Forward Pass]
@@ -297,15 +293,17 @@ The input/output processing pipeline handles serialization, routing, and accumul
 ```mermaid
 sequenceDiagram
     participant Client
-    participant OmniInputProcessor
+    participant AsyncOmniEngine
+    participant InputProcessor
     participant Scheduler
     participant ModelRunner
     participant MultimodalOutputProcessor
     participant Client
 
-    Client->>OmniInputProcessor: prompt + prompt_embeds + additional_info
-    OmniInputProcessor->>OmniInputProcessor: Serialize tensors to payloads
-    OmniInputProcessor->>Scheduler: OmniEngineCoreRequest (with payloads)
+    Client->>AsyncOmniEngine: prompt + prompt_embeds + additional_info
+    AsyncOmniEngine->>InputProcessor: process_inputs()
+    InputProcessor->>Scheduler: EngineCoreRequest
+    AsyncOmniEngine->>AsyncOmniEngine: _upgrade_to_omni_request() + serialize_additional_information()
     Scheduler->>ModelRunner: OmniNewRequestData (with payloads)
     ModelRunner->>ModelRunner: Decode payloads → CPU tensors
     ModelRunner->>ModelRunner: Overlay prompt_embeds on inputs_embeds
@@ -318,43 +316,18 @@ sequenceDiagram
     MultimodalOutputProcessor->>Client: RequestOutput (with multimodal_output)
 ```
 
-### OmniInputProcessor
+### Stage-0 Input Processing
 
-`OmniInputProcessor` extends the base input processor to serialize prompt embeddings and additional information for inter-stage transfer.
-
-#### Payload Serialization
-
-Converts PyTorch tensors to serialized payloads:
+Stage-0 now uses upstream `InputProcessor` directly, and `AsyncOmniEngine` upgrades the request to `OmniEngineCoreRequest` while restoring omni-specific payloads.
 
 ```python
-def process_inputs(self, ...) -> OmniEngineCoreRequest:
-    # Serialize prompt_embeds
-    if "prompt_embeds" in decoder_inputs:
-        pe_cpu = decoder_inputs["prompt_embeds"].detach().to("cpu").contiguous()
-        prompt_embeds_payload = PromptEmbedsPayload(
-            data=pe_cpu.numpy().tobytes(),
-            shape=[seq_len, hidden_size],
-            dtype=dtype_str,
-        )
-
-    # Serialize additional_information
-    if "additional_information" in decoder_inputs:
-        entries = {}
-        for key, value in raw_info.items():
-            if isinstance(value, torch.Tensor):
-                entry = AdditionalInformationEntry(
-                    tensor_data=value.numpy().tobytes(),
-                    tensor_shape=list(value.shape),
-                    tensor_dtype=dtype_str,
-                )
-            entries[key] = entry
-        additional_information_payload = AdditionalInformationPayload(entries=entries)
-
-    return OmniEngineCoreRequest(
-        # ... standard fields ...
-        prompt_embeds=prompt_embeds_payload,
-        additional_information=additional_information_payload,
-    )
+request = self.input_processor.process_inputs(
+    request_id=request_id,
+    prompt=prompt,
+    params=params,
+    supported_tasks=self.supported_tasks,
+)
+request = _upgrade_to_omni_request(request, prompt)
 ```
 
 ### MultimodalOutputProcessor
@@ -400,13 +373,13 @@ The AR module of vLLM-Omni extends vLLM through strategic inheritance and minima
 ### Key Design Patterns
 
 1. **Inheritance over composition**: Extends vLLM classes to preserve compatibility with existing scheduling, batching, and execution mechanisms
-2. **Payload serialization**: Uses `PromptEmbedsPayload` and `AdditionalInformationPayload` for efficient inter-stage data transfer
+2. **Payload serialization**: Uses serialized `additional_information` payloads together with prompt-embedding handoff for efficient inter-stage transfer
 3. **Two-phase execution**: Maintains vLLM's execute/sample separation for AR models while supporting single-phase execution for generation models
 4. **Multimodal routing**: Routes outputs by `output_type` and accumulates tensors incrementally to support streaming
 
 ### Differences from vLLM
 
-- **Payload support**: Serialized prompt embeddings and additional information enable direct transfer between pipeline stages
+- **Payload support**: Serialized additional information and prompt embeddings enable direct transfer between pipeline stages
 - **Multimodal handling**: Extended input/output processors support images, audio, and other modalities alongside text
 - **Hidden state exposure**: AR model runners expose per-request hidden states via `pooler_output` for downstream consumption
 - **Generation scheduler**: Fast-path scheduling for basic heterogeneous architectures that complete in one step

docs/design/module/async_omni_architecture.md

@@ -0,0 +1,203 @@
+# AsyncOmni Architecture (Qwen3-Omni Example)
+
+## 1. System Architecture
+
+```text
+• ┌─────────────────────────────────────────────────────────────────────────────────┐
+  │                                    API Layer                                    │
+  │  ┌─────────────────────────────────────┐  ┌──────────────────────────────────┐  │
+  │  │ AsyncOmni (EngineClient)            │  │ Omni                             │  │
+  │  │ • generate() / abort() / shutdown() │  │ • generate()                     │  │
+  │  │ • _final_output_handler()           │  │                                  |  │
+  │  └─────────────────────────────────────┘  └──────────────────────────────────┘  │
+  ├─────────────────────────────────────────────────────────────────────────────────┤
+  │                              Engine Layer (Proxy)                               │
+  │  ┌───────────────────────────────────────────────────────────────────────────┐  │
+  │  │ AsyncOmniEngine                                                           │  │
+  │  │ • _bootstrap_orchestrator() & _initialize_stages()                        │  │
+  │  │ • add_request() / add_request_async() -> input_processor.process_inputs() │  │
+  │  │ • try_get_output() / try_get_output_async()                               │  │
+  │  └───────────────────┬─────────────────────────────────▲─────────────────────┘  │
+  │         request_queue (janus.Queue)        output_queue (janus.Queue)           │
+  ├──────────────────────┼─────────────────────────────────┼────────────────────────┤
+  │                      ▼        Orchestration Layer      │                        │
+  │  ┌───────────────────────────────────────────────────────────────────────────┐  │
+  │  │ Orchestrator [background thread]                                          │  │
+  │  │ • _request_handler()                                                      │  │
+  │  │     -  stage_client.add_request_async() & _prewarm_async_chunk_stages()   │  │
+  │  │ • _orchestration_output_handler()                                         │  │
+  │  │     -  _process_stage_outputs() -> output_processors[i].process_outputs() │  │
+  │  │     -  _route_output() & _forward_to_next_stage()                         │  │
+  │  └──────────┬─────────────────────────┬────────────────────────┬─────────────┘  │
+  ├─────────────┼─────────────────────────┼────────────────────────┼────────────────┤
+  │             │                 Communication Layer              │                │
+  │  ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐  │
+  │  │ StageEngineCoreClient │ │ StageEngineCoreClient │ │ StageDiffusionClient  │  │
+  │  │ • ZMQ ROUTER / PULL   │ │ • ZMQ ROUTER / PULL   │ │ • ZMQ ROUTER / PULL   │  │
+  │  │ • Msgpack codec       │ │ • Msgpack codec       │ │ • Msgpack codec       │  │
+  │  └──────────┬────────────┘ └──────────┬────────────┘ └──────────┬────────────┘  │
+  │             ▼ ZMQ IPC                 ▼ ZMQ IPC                 ▼ ZMQ IPC       │
+  ├─────────────────────────────────────────────────────────────────────────────────┤
+  │                                 Execution Layer                                 │
+  │  ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐  │
+  │  │ StageCoreProc         │ │ StageCoreProc         │ │ DiffusionEngine       │  │
+  │  │ [background process]  │ │ [background process]  │ │ [background process]  │  │
+  │  └───────────────────────┘ └───────────────────────┘ └───────────────────────┘  │
+  └─────────────────────────────────────────────────────────────────────────────────┘
+```
+
+## 2. Execution Flow (Arrow Steps, one generate request)
+
+```text
+[1] App
+    -> AsyncOmni.generate(prompt, request_id)
+
+[2] AsyncOmni
+    -> _final_output_handler()   (started on first request)
+    -> AsyncOmniEngine.add_request(stage_id=0, ...)
+
+[3] AsyncOmniEngine.add_request
+    -> (if stage-0 is llm and input is not EngineCoreRequest)
+       InputProcessor.process_inputs()
+       OutputProcessor[0].add_request()
+    -> request_queue.put(add_request_msg)
+
+[4] Orchestrator._request_handler
+    -> _handle_add_request(msg)
+    -> stage_clients[0].add_request_async(...)
+
+[5] Orchestrator._orchestration_loop (loop)
+    -> poll stage output
+       - llm stage: await get_output_async()
+       - diffusion stage: get_diffusion_output_nowait()
+    -> (llm stage) output_processors[i].process_outputs(...)
+    -> _route_output(...)
+    -> if finished and not final_stage and non-async-chunk:
+         _forward_to_next_stage(...)
+         -> next_stage.add_request_async(...)
+    -> output_queue.put(output)
+
+[6] AsyncOmni._final_output_loop (background coroutine)
+    -> AsyncOmniEngine.try_get_output_async()
+    -> route by request_id to ClientRequestState.queue
+
+[7] AsyncOmni._process_orchestrator_results
+    -> read from ClientRequestState.queue
+    -> _process_single_result(...)
+    -> yield OmniRequestOutput
+
+[8] Exit condition
+    -> receive result["finished"] == True
+    -> generate() ends
+```
+
+## 3. Runtime Sequence (one generate request)
+
+```mermaid
+sequenceDiagram
+    participant APP as App
+    participant AO as AsyncOmni
+    participant ENG as AsyncOmniEngine
+    participant ORCH as Orchestrator
+    participant S0 as Stage-0 Client
+    participant SN as Next Stage Client
+
+    APP->>AO: generate
+    AO->>AO: start output_handler once
+    AO->>ENG: add_request(stage_id=0, ...)
+    ENG->>ENG: input_processor.process_inputs()
+    ENG->>ORCH: request_queue.put(add_request)
+
+    ORCH->>ORCH: _handle_add_request
+    ORCH->>S0: add_request_async
+
+    loop poll route forward
+        ORCH->>S0: get_output_async / get_diffusion_output_nowait
+        ORCH->>ORCH: _route_output
+        alt need forward to next stage
+            ORCH->>SN: add_request_async
+        end
+        ORCH-->>ENG: output_queue.put
+    end
+
+    AO->>ENG: try_get_output_async
+    ENG-->>AO: message
+    AO-->>APP: yield OmniRequestOutput
+```
+
+## 4. Comparison
+
+Previous topology (reference):
+
+```text
+┌────────────────────────────────────────────────────────────────────────────┐
+│ Main Process                                                               │
+│  ┌──────────────────────┐   ┌────────────────────────────────────────────┐ │
+│  │ generate()           │   │ final_output_handler()                     │ │
+│  └──────────────────────┘   └────────────────────────────────────────────┘ │
+└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
+  mp.Queue (in_q/out_q)    mp.Queue (in_q/out_q)    mp.Queue (in_q/out_q)
+           ▼▲                        ▼▲                        ▼▲
+┌───────────────────────┐  ┌───────────────────────┐  ┌──────────────────────┐
+│ Worker Proc-0         │  │ Worker Proc-1         │  │ Worker Proc-2        │
+│ (Thinker LLM)         │  │ (Talker LLM)          │  │ (Vocoder)            │
+│  ┌────────────────┐   │  │  ┌────────────────┐   │  │  ┌────────────────┐  │
+│  │_stage_worker   │   │  │  │_stage_worker   │   │  │  │_stage_worker   │  │
+│  │_async()        │   │  │  │_async()        │   │  │  │_async()        │  │
+│  └────────────────┘   │  │  └────────────────┘   │  │  └────────────────┘  │
+│  ┌────────────────┐   │  │  ┌────────────────┐   │  │  ┌────────────────┐  │
+│  │output_handler()│   │  │  │output_handler()│   │  │  │output_handler()│  │
+│  └────────────────┘   │  │  └────────────────┘   │  │  └────────────────┘  │
+└──────────┬────────────┘  └──────────┬────────────┘  └──────────┬───────────┘
+       ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ
+┌──────────────────────┐   ┌──────────────────────┐   ┌──────────────────────┐
+│ EngineCore Proc-0    │   │ EngineCore Proc-1    │   │ EngineCore Proc-2    │
+│ (Thinker)            │   │ (Talker)             │   │ (Vocoder)            │
+└──────────────────────┘   └──────────────────────┘   └──────────────────────┘
+```
+
+Current topology:
+
+```text
+┌────────────────────────────────────────────────────────────────────────────┐
+│ Main Process                                                               │
+│  ┌──────────────────────────────────────────────────────────────────────┐  │
+│  │ Main Thread                                                          │  │
+│  │  ┌──────────────────────┐   ┌─────────────────────────────────────┐  │  │
+│  │  │ generate()           │   │ final_output_handler()              │  │  │
+│  │  └──────────────────────┘   └─────────────────────────────────────┘  │  │
+│  └──────────────────────────────────────────────────────────────────────┘  │
+│         janus.Queue (request_queue) ▼  ▲ janus.Queue (output_queue)        │
+│  ┌──────────────────────────────────────────────────────────────────────┐  │
+│  │ Orchestrator Thread                                                  │  │
+│  │  ┌──────────────────────┐  ┌──────────────────────────────────────┐  │  │
+│  │  │ _request_handler()   │  │ _orchestration_output_handler()      │  │  │
+│  │  └──────────────────────┘  └──────────────────────────────────────┘  │  │
+│  │  ┌────────────────────────────────────────────────────────────────┐  │  │
+│  │  │ _orchestration_loop(): poll/process/route outputs for all stages│  │  │
+│  │  └────────────────────────────────────────────────────────────────┘  │  │
+│  └───────┬─────────────────────────┬─────────────────────────┬──────────┘  │
+└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
+       ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ  
+  ┌──────────────────────┐  ┌──────────────────────┐  ┌──────────────────────┐
+  │ EngineCore Proc-0    │  │ EngineCore Proc-1    │  │ EngineCore Proc-2    │
+  │ (Thinker)            │  │ (Talker)             │  │ (Vocoder)            │
+  └──────────────────────┘  └──────────────────────┘  └──────────────────────┘
+```
+
+
+Test scripts:
+```bash
+# enter offline inference folder.
+cd examples/offline_inference/qwen2_5_omni
+python end2end.py --output-dir output_audio --query-type use_mixed_modalities
+
+cd ../qwen3_omni
+python end2end.py --output-dir output_audio --query-type text --async-chunk --enable-stats
+
+cd ../bagel
+python end2end.py --prompts "A cute cat"
+
+cd ../text_to_image
+python text_to_image.py --prompt "a cup of coffee on the table" --output output.png
+```

docs/design/module/dit_module.md

@@ -127,80 +127,117 @@ def step(self, requests: list[OmniDiffusionRequest]):
 
 ## 2. Scheduler
 
-**Location**: `vllm_omni/diffusion/scheduler.py`
+**Location**: `vllm_omni/diffusion/sched/`
 
 ### Architecture
 
-The `Scheduler` is implemented as a **Singleton** pattern to ensure a single coordination point across the system, i.e., only one scheduler instance exists for coordination.
+The scheduler is a **request-state scheduler**. It owns request lifecycle management and scheduling decisions, while execution stays in `DiffusionEngine` and the executor.
 
 ### Key Components
 
-#### 2.1 Message Queue System
+#### 2.1 Scheduler Interface
 
 ```python
-class Scheduler:
-    def initialize(self, od_config: OmniDiffusionConfig):
-        # Broadcast queue: scheduler -> all workers
-        self.mq = MessageQueue(
-            n_reader=od_config.num_gpus,
-            n_local_reader=od_config.num_gpus,
-            local_reader_ranks=list(range(od_config.num_gpus)),
-        )
-
-        # Result queue: rank 0 worker -> scheduler
-        self.result_mq = None  # Initialized later
+class SchedulerInterface(ABC):
+    def add_request(self, request: OmniDiffusionRequest) -> str: ...
+    def schedule(self) -> DiffusionSchedulerOutput: ...
+    def update_from_output(
+        self,
+        sched_output: DiffusionSchedulerOutput,
+        output: DiffusionOutput,
+    ) -> set[str]: ...
 ```
 
-**Communication Pattern**:
+**Responsibilities**:
 
-- **Broadcast Queue**: One-to-many communication (scheduler → all workers)
+- **Lifecycle contract**: Defines how the engine adds requests, triggers one scheduling cycle, and feeds executor results back.
 
-- **Result Queue**: One-to-one communication (rank 0 → scheduler)
+- **Stable boundary**: `DiffusionSchedulerOutput` is the only scheduling result consumed by `DiffusionEngine`.
 
-- **Shared Memory**: Uses `MessageQueue` (ZMQ-based) for efficient IPC
+- **Pluggability**: Different scheduler policies can reuse the same engine integration path.
 
-#### 2.2 Request Distribution
+#### 2.2 Request State Model
 
 ```python
-def add_req(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
-    # Broadcast request to all workers
-    self.mq.enqueue(requests)
+class DiffusionRequestStatus(enum.IntEnum):
+    WAITING = ...
+    RUNNING = ...
+    PREEMPTED = ...
+    FINISHED_COMPLETED = ...
+    FINISHED_ABORTED = ...
+    FINISHED_ERROR = ...
 
-    # Wait for result from Rank 0
-    output = self.result_mq.dequeue()
-    return output
+@dataclass
+class DiffusionRequestState:
+    sched_req_id: str
+    req: OmniDiffusionRequest
+    status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
 ```
 
 **Design Features**:
 
-- **Broadcast Model**: All workers receive the same request (for tensor parallelism)
+- **Scheduler-owned ID**: Each `OmniDiffusionRequest` is tracked by an internal `sched_req_id`, separated from public `request_id` values.
 
-- **Single Response**: Only rank 0 sends results back (avoids duplicate outputs)
+- **Explicit lifecycle**: Requests move through waiting, running, optional preemption, and terminal states.
 
-- **Synchronous**: Blocks until result is received (can be made async)
+- **Centralized error handling**: Completion, abort, and error states are all normalized in the scheduler layer.
 
-#### 2.3 Singleton Pattern
+#### 2.3 Shared Bookkeeping in `_BaseScheduler`
 
 ```python
-class Scheduler:
-    _instance = None
-
-    def __new__(cls, *args, **kwargs):
-        if not cls._instance:
-            cls._instance = super().__new__(cls)
-        return cls._instance
-
-# Global singleton instance
-scheduler = Scheduler()
+class _BaseScheduler(SchedulerInterface):
+    def __init__(self) -> None:
+        self._request_states = {}
+        self._request_id_to_sched_req_id = {}
+        self._waiting = deque()
+        self._running = []
+        self._finished_req_ids = set()
+        self.max_num_running_reqs = 1
 ```
 
-**Benefits**:
+**Design Features**:
 
-- **Single Point of Control**: Ensures consistent state
+- **Common state storage**: Shared request maps and waiting/running sets live in the base class.
 
-- **Easy Access**: Global `scheduler` instance accessible everywhere
+- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
 
-- **Resource Management**: Centralized queue management
+- **Current constraint**: `max_num_running_reqs` remains `1` because the current engine path is still synchronous request-mode execution.
+
+#### 2.4 Current `RequestScheduler` Policy
+
+```python
+class RequestScheduler(_BaseScheduler):
+    def schedule(self) -> DiffusionSchedulerOutput:
+        # 1. keep existing RUNNING requests in the scheduling result
+        # 2. pull WAITING requests while capacity remains
+        # 3. move newly admitted requests into RUNNING
+```
+
+**Behavior**:
+
+- **FIFO request scheduling**: Waiting requests are promoted in queue order.
+
+- **Single-request admission**: The current policy only admits one active request at a time.
+
+- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished scheduler ids.
+
+#### 2.5 Engine-Driven Execution Loop
+
+```python
+sched_req_id = scheduler.add_request(request)
+while True:
+    sched_output = scheduler.schedule()
+    output = executor.add_req(req)
+    finished_req_ids = scheduler.update_from_output(sched_output, output)
+```
+
+**Design Decisions**:
+
+- **Separation of concerns**: Scheduler manages state and policy; executor handles runtime execution.
+
+- **No scheduler-owned IPC**: Scheduler no longer talks to workers directly.
+
+- **Conservative concurrency**: The current request-mode implementation still allows only one active request at a time.
 
 ---
 
@@ -880,8 +917,9 @@ def initialize_model_parallel(
        └─> Model-specific transformations
 
 3. Scheduling
-   └─> scheduler.add_req(requests)
-       └─> Broadcast via MessageQueue to all workers
+   └─> scheduler.add_request(request)
+       └─> scheduler.schedule()
+           └─> DiffusionEngine submits scheduled request to executor.add_req(req)
 
 4. Worker Execution
    └─> WorkerProc.worker_busy_loop()
@@ -895,8 +933,9 @@ def initialize_model_parallel(
                └─> vae.decode()
 
 5. Result Collection
-   └─> Rank 0 sends DiffusionOutput via result queue
-       └─> Scheduler receives and returns
+   └─> Executor returns DiffusionOutput
+       └─> scheduler.update_from_output(...)
+           └─> DiffusionEngine pops finished request state
 
 6. Post-processing
    └─> post_process_func(output)

docs/features/custom_pipeline.md

@@ -97,7 +97,7 @@ outputs = omni.generate(
 )
 
 # Access custom trajectory data
-output = outputs[0].request_output[0]
+output = outputs[0].request_output
 print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
 print(f"Trajectory latents shape: {output.latents.shape}")
 ```

docs/getting_started/installation/gpu.md

@@ -30,7 +30,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0. 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 available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, and 0.18.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

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

@@ -29,18 +29,23 @@ uv pip install vllm --torch-backend=auto
 uv pip install vllm-omni
 ```
 
+To run Gradio demos, also install the optional extras:
+```bash
+uv pip install 'vllm-omni[demo]'
+```
+
 # --8<-- [end:pre-built-wheels]
 
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library
 
 ```bash
-uv pip install vllm==0.17.0 --torch-backend=auto
+uv pip install vllm==0.19.0 --torch-backend=auto
 ```
 
-The release 0.17.0 of vLLM is based on PyTorch 2.10.0 which requires CUDA 12.9 environment.
+The 0.19.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -50,17 +55,23 @@ cd vllm-omni
 uv pip install -e .
 ```
 
+To run Gradio demos, install with optional extras:
+```bash
+uv pip install -e '.[demo]'
+```
+
 <details><summary>(Optional) Installation of vLLM from source</summary>
 If you want to check, modify or debug with source code of vLLM, install the library from source with the following instructions:
 
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.17.0
+git checkout v0.19.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_x86_64.whl
+#For CUDA 13.0
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.19.0/vllm-0.19.0+cu130-cp38-abi3-manylinux_2_35_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -91,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.16.0 \
+    vllm/vllm-omni:v0.18.0 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

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

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
+uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library
 
 ```bash
-uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
+uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
 ```
 
-The release 0.17.0 of vLLM requires ROCm 7.0 environment.
+The pre-built 0.19.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.17.0
+git checkout v0.19.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.16.0 \
+  vllm/vllm-omni-rocm:v0.19.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -149,7 +149,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.16.0
+  vllm/vllm-omni-rocm:v0.19.0
 ```
 
 # --8<-- [end:pre-built-images]

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

@@ -10,10 +10,10 @@ The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built
 ```bash
 # Update the vllm-ascend image
 # Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1-a3
+export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \
@@ -33,18 +33,13 @@ docker run --rm \
     -p 8000:8000 \
     -it $IMAGE bash
 
-# Because vllm-ascend will release v0.16.0rc1 after vllm-omni 0.16.0,
-# we have to pin vllm-ascend at the current commit.
-cd /vllm-workspace/vllm-ascend
-git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
-pip install -v -e .
-
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.16.0 https://github.com/vllm-project/vllm-omni.git
-
+git clone -b v0.18.0 https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
-pip install -v -e .
+pip install -v -e . --no-build-isolation
+# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
+
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
 ```
 
@@ -61,23 +56,16 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
 You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
 
 ```bash
-# Pin vLLM version to 0.16.0
-cd /vllm-workspace/vllm
-git pull origin main
-git fetch origin --tags
-git checkout v0.16.0
+# Pin vLLM version to 0.18.0
+git clone -b v0.18.0 https://github.com/vllm-project/vllm.git
 VLLM_TARGET_DEVICE=empty pip install -v -e .
 
-# Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
-cd ../vllm-ascend
-git pull origin main
-git fetch origin --tags
-git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+git clone -b v0.18.0rc1 https://github.com/vllm-project/vllm-ascend.git
 pip install -v -e .
 
 # Install vLLM-Omni from the latest main branch
-cd ../vllm-omni
 git clone https://github.com/vllm-project/vllm-omni.git
+cd /vllm-workspace/vllm-omni
 pip install -v -e . --no-build-isolation
 # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
@@ -100,10 +88,10 @@ Here's an example deployment command that has been verified on 4 x NPUs:
 
 ```bash
 # Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-omni:v0.16.0
+# export IMAGE=quay.io/ascend/vllm-omni:v0.18.0
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-omni:v0.16.0-a3
-export IMAGE=quay.io/ascend/vllm-omni:v0.16.0
+# export IMAGE=quay.io/ascend/vllm-omni:v0.18.0-a3
+export IMAGE=quay.io/ascend/vllm-omni:v0.18.0
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \

docs/getting_started/quickstart.md

@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.17.0 --torch-backend=auto
+uv pip install vllm==0.19.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.17.0 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
+uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -42,7 +42,7 @@ if __name__ == "__main__":
     omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = outputs[0].request_output[0].images
+    images = outputs[0].request_output.images
     images[0].save("coffee.png")
 ```
 
@@ -70,7 +70,7 @@ if __name__ == "__main__":
     ]
     omni_outputs = omni.generate(prompts)
     for i_prompt, prompt_output in enumerate(omni_outputs):
-        this_request_output = prompt_output.request_output[0]
+        this_request_output = prompt_output.request_output
         this_images = this_request_output.images
         for i_image, image in enumerate(this_images):
             image.save(f"p{i_prompt}-img{i_image}.jpg")
@@ -82,7 +82,7 @@ if __name__ == "__main__":
 
 !!! info
 
-    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
     list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
     batched inputs, you can [modify this configuration](../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
 

docs/mkdocs/hooks/generate_api_readme.py

@@ -139,8 +139,11 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]:
             # Skip excluded modules (avoid importing vllm during docs build)
             excluded_prefixes = [
                 "vllm_omni.diffusion.models.qwen_image",
+                "vllm_omni.diffusion.quantization",
+                "vllm_omni.quantization",
                 "vllm_omni.entrypoints.async_diffusion",
                 "vllm_omni.entrypoints.openai",
+                "vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts",
             ]
             if any(module_path.startswith(prefix) for prefix in excluded_prefixes):
                 continue

docs/mkdocs/hooks/generate_argparse.py

@@ -120,6 +120,7 @@ def extract_omni_serve_subparser_init():
                         "_FlexibleArgumentParser": _FlexibleArgumentParser,
                         "FlexibleArgumentParser": _FlexibleArgumentParser,
                         "make_arg_parser": lambda parser: parser,  # no-op for doc
+                        "_ensure_vllm_platform": lambda: None,  # no-op for doc
                         "VLLM_SUBCMD_PARSER_EPILOG": "",
                         "logger": logger,
                         "DummySubparsers": DummySubparsers,

docs/models/supported_models.md

@@ -6,7 +6,7 @@ vLLM-Omni supports unified multimodal comprehension and generation models across
 
 If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm_omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.
 
-## List of Supported Models for Nvidia GPU / AMD GPU
+## List of Supported Models
 
 <style>
 th {
@@ -15,71 +15,49 @@ th {
 }
 </style>
 
-| Architecture | Models | Example HF Models |
-|--------------|--------|-------------------|
-| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
-| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` |
-| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |
-| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |
-| `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` |
-| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` |
-|`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
-| `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
-| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
-| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` |
-| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` |
-| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` |
-| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` |
-|`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` |
-|`CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` |
-|`MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` |
-|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
-|`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
-|`OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` |
-|`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` |
-|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
-|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
-|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
-|`GlmImageForConditionalGeneration` | GLM-Image | `zai-org/GLM-Image` |
-|`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` |
-|`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` |
-|`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
-|`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` |
-|`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` |
+| Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU |
+|--------------|--------|-------------------|------------|---------|-----|-----------|
+| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
+| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | |
+| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
+| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
+| `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
+| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
+| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
+| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
+| `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` | ✅︎ | ✅︎ | | ✅︎ |
+| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | |
+| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
+| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
+| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ |
+| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
+| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
+| `MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
+| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | |
+| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
+| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
+| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
+| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
+| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
 
-
-## List of Supported Models for NPU
-
-<style>
-th {
-  white-space: nowrap;
-  min-width: 0 !important;
-}
-</style>
-
-| Architecture | Models | Example HF Models |
-|--------------|--------|-------------------|
-| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
-| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
-| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |
-| `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` |
-| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` |
-|`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
-| `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
-| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
-|`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
-|`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
-|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
-|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
-|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
-|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
+✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.

docs/serving/diffusion_chat_api.md

@@ -0,0 +1,78 @@
+# Diffusion Chat Completions API
+
+vLLM-Omni supports generating and editing images via the `/v1/chat/completions`
+endpoint using diffusion models. This page explains how to pass generation
+parameters (such as `num_inference_steps`, `height`, `width`) to diffusion
+models through this endpoint.
+
+!!! tip
+    For dedicated endpoints that accept generation parameters as top-level
+    fields, see [Image Generation API](image_generation_api.md) and
+    [Image Edit API](image_edit_api.md).
+
+## Passing Generation Parameters
+
+The `/v1/chat/completions` endpoint follows the OpenAI Chat API schema, which
+does not natively include diffusion-specific fields like `num_inference_steps`
+or `height`. How you pass these extra fields depends on your client.
+
+### curl / Python `requests`
+
+Wrap generation parameters inside an `"extra_body"` key in the JSON body:
+
+```bash
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      {"role": "user", "content": "A beautiful landscape painting"}
+    ],
+    "extra_body": {
+      "num_inference_steps": 50,
+      "seed": 42
+    }
+  }'
+```
+
+### OpenAI Python SDK
+
+Use the `extra_body` **keyword argument**. The SDK automatically merges these
+fields into the top-level request body:
+
+```python
+response = client.chat.completions.create(
+    model="Qwen/Qwen-Image",
+    messages=[{"role": "user", "content": "A beautiful landscape painting"}],
+    extra_body={
+        "num_inference_steps": 50,
+        "seed": 42,
+    },
+)
+```
+
+!!! note "SDK `extra_body` vs. JSON `extra_body`"
+    These two `extra_body` usages look similar but work differently under the
+    hood. The SDK flattens the dict into the top-level request JSON, while the
+    curl/requests approach sends it as a nested `"extra_body"` key. Both are
+    handled correctly by the server.
+
+!!! note "About the `ignored fields` warning"
+    You may see a log message like:
+
+    ```
+    WARNING: The following fields were present in the request but ignored: {'height', 'width', ...}
+    ```
+
+    This is **harmless**. It is emitted by vLLM's request validation layer
+    because these fields are not part of the standard OpenAI
+    `ChatCompletionRequest` schema. The fields are still stored internally
+    and correctly forwarded to the diffusion pipeline.
+
+## Model-Specific Examples
+
+For complete examples with full request/response details, see the model-specific
+guides:
+
+- [Text-to-Image (Qwen-Image)](../user_guide/examples/online_serving/text_to_image.md)
+- [Image-to-Image (Qwen-Image-Edit, Qwen-Image-Layered)](../user_guide/examples/online_serving/image_to_image.md)
+- [GLM-Image](../user_guide/examples/online_serving/glm_image.md)

docs/serving/speech_api.md

@@ -4,6 +4,7 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
 
 - **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
 - **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
+- **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
 
 Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
 
@@ -28,6 +29,14 @@ vllm-omni serve fishaudio/s2-pro \
     --trust-remote-code \
     --enforce-eager \
     --gpu-memory-utilization 0.9
+
+# Voxtral TTS
+vllm serve mistralai/Voxtral-4B-TTS-2603 \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
 ```
 
 ### Generate Speech
@@ -109,6 +118,7 @@ Content-Type: application/json
 | `instructions` | string | "" | Voice style/emotion instructions |
 | `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
 | `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
+| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
 
 **Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
 
@@ -134,9 +144,23 @@ Lists available voices for the loaded model.
 
 ```json
 {
-    "voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
+    "voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian", "custom_voice_1"],
+    "uploaded_voices": [
+        {
+            "name": "custom_voice_1",
+            "consent": "user_consent_id",
+            "created_at": 1738660000,
+            "file_size": 1024000,
+            "mime_type": "audio/wav",
+            "ref_text": "The exact transcript of the audio sample.",
+            "speaker_description": "warm narrator"
+        }
+    ]
 }
 ```
+
+`uploaded_voices` is always present (empty list when no custom voices have been uploaded). Fields `ref_text` and `speaker_description` are omitted per-entry when not provided at upload time.
+
 ```
 POST /v1/audio/voices
 Content-Type: multipart/form-data
@@ -151,6 +175,8 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
 | `audio_sample` | file | Yes | Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4) |
 | `consent` | string | Yes | Consent recording ID |
 | `name` | string | Yes | Name for the new voice |
+| `ref_text` | string | No | Transcript of the audio. When provided, enables in-context voice cloning (higher quality). Without it, only the speaker embedding is extracted. |
+| `speaker_description` | string | No | Free-form description of the voice (e.g. "warm narrator", "energetic presenter"). Stored as metadata and returned in `GET /v1/audio/voices`. |
 
 **Response Example:**
 
@@ -162,18 +188,24 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
     "consent": "user_consent_id",
     "created_at": 1738660000,
     "mime_type": "audio/wav",
-    "file_size": 1024000
+    "file_size": 1024000,
+    "ref_text": "The exact transcript of the audio sample.",
+    "speaker_description": "warm narrator"
   }
 }
 ```
 
+Fields `ref_text` and `speaker_description` are omitted when not provided at upload time.
+
 **Usage Example:**
 
 ```bash
 curl -X POST http://localhost:8091/v1/audio/voices \
   -F "audio_sample=@/path/to/voice_sample.wav" \
   -F "consent=user_consent_id" \
-  -F "name=custom_voice_1"
+  -F "name=custom_voice_1" \
+  -F "ref_text=The exact transcript of the audio sample." \
+  -F "speaker_description=warm narrator"
 ```
 
 ## Streaming Text Input (WebSocket)
@@ -308,7 +340,9 @@ curl -X POST http://localhost:8091/v1/audio/speech \
     }' --output cloned.wav
 ```
 
-upload voice
+### Upload Voice
+
+Upload voice (speaker embedding only):
 ```bash
 curl -X POST http://localhost:8091/v1/audio/voices \
   -F "audio_sample=@/path/to/voice_sample.wav" \
@@ -316,17 +350,183 @@ curl -X POST http://localhost:8091/v1/audio/voices \
   -F "name=custom_voice_1"
 ```
 
-use upload voice
+Upload voice with transcript (in-context cloning, higher quality):
+```bash
+curl -X POST http://localhost:8091/v1/audio/voices \
+  -F "audio_sample=@/path/to/voice_sample.wav" \
+  -F "consent=user_consent_id" \
+  -F "name=custom_voice_2" \
+  -F "ref_text=The exact transcript of the audio sample."
+```
+
+### Use Uploaded Voice
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \
     -H "Content-Type: application/json" \
     -d '{
         "input": "Hello, this is a cloned voice",
-        "task_type": "Base",
         "voice": "custom_voice_1"
     }' --output cloned.wav
 ```
 
+## Batch Speech Generation
+
+The batch endpoint synthesizes multiple texts in a single request, returning all results as JSON with base64-encoded audio.
+
+### Endpoint
+
+```
+POST /v1/audio/speech/batch
+Content-Type: application/json
+```
+
+### Request Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `items` | array | **required** | List of items to synthesize (1–32) |
+| `model` | string | server's model | Model to use |
+| `voice` | string | null | Default voice for all items |
+| `response_format` | string | "wav" | Default audio format for all items |
+| `speed` | float | 1.0 | Default playback speed (0.25–4.0) |
+| `task_type` | string | null | Default TTS task type |
+| `language` | string | null | Default language |
+| `instructions` | string | null | Default voice style instructions |
+| `ref_audio` | string | null | Default reference audio (Base task) |
+| `ref_text` | string | null | Default reference transcript (Base task) |
+| `max_new_tokens` | integer | null | Default max tokens |
+
+Each item in the `items` array requires only `input` (the text). All other fields are optional and override the batch-level defaults when set:
+
+| Field | Type | Description |
+|-------|------|-------------|
+| `input` | string | **required** — text to synthesize |
+| `voice` | string | Override voice for this item |
+| `response_format` | string | Override format for this item |
+| `speed` | float | Override speed for this item |
+| `task_type` | string | Override task type |
+| `language` | string | Override language |
+| `instructions` | string | Override instructions |
+| `ref_audio` | string | Override reference audio |
+| `ref_text` | string | Override reference transcript |
+| `max_new_tokens` | integer | Override max tokens |
+
+### Response Format
+
+```json
+{
+    "id": "speech-batch-abc123",
+    "results": [
+        {
+            "index": 0,
+            "status": "success",
+            "audio_data": "<base64-encoded audio>",
+            "media_type": "audio/wav"
+        },
+        {
+            "index": 1,
+            "status": "error",
+            "error": "Input text cannot be empty"
+        }
+    ],
+    "total": 2,
+    "succeeded": 1,
+    "failed": 1
+}
+```
+
+### Examples
+
+**Basic batch with shared defaults:**
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech/batch \
+    -H "Content-Type: application/json" \
+    -d '{
+        "items": [
+            {"input": "Hello, how are you?"},
+            {"input": "Goodbye, see you later!"}
+        ],
+        "voice": "vivian",
+        "language": "English"
+    }'
+```
+
+**Per-item overrides (different voices and formats):**
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech/batch \
+    -H "Content-Type: application/json" \
+    -d '{
+        "items": [
+            {"input": "Hello!", "voice": "vivian", "response_format": "mp3"},
+            {"input": "你好！", "voice": "ryan", "language": "Chinese"}
+        ],
+        "response_format": "wav"
+    }'
+```
+
+**Voice cloning with shared reference audio (Base task):**
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech/batch \
+    -H "Content-Type: application/json" \
+    -d '{
+        "items": [
+            {"input": "First sentence in the cloned voice."},
+            {"input": "Second sentence in the cloned voice."}
+        ],
+        "task_type": "Base",
+        "ref_audio": "https://example.com/reference.wav",
+        "ref_text": "Transcript of the reference audio"
+    }'
+```
+
+Setting `ref_audio` at the batch level applies it to all items, avoiding the need to repeat it per item.
+
+**Decoding the response in Python:**
+
+```python
+import base64
+import httpx
+
+response = httpx.post(
+    "http://localhost:8091/v1/audio/speech/batch",
+    json={
+        "items": [
+            {"input": "First sentence."},
+            {"input": "Second sentence."},
+        ],
+        "voice": "vivian",
+    },
+    timeout=300.0,
+)
+
+for result in response.json()["results"]:
+    if result["status"] == "success":
+        audio_bytes = base64.b64decode(result["audio_data"])
+        with open(f"output_{result['index']}.wav", "wb") as f:
+            f.write(audio_bytes)
+```
+
+### Configuration
+
+| Parameter | Source | Default | Description |
+|-----------|--------|---------|-------------|
+| `tts_batch_max_items` | engine kwarg | 32 | Maximum number of items per batch request |
+
+All items are fanned out to `generate()` concurrently. The engine's stage worker automatically batches them up to the configured `max_batch_size` and queues the rest — no client-side throttling needed.
+
+For best throughput, use a batch-optimized stage config with `max_batch_size > 1`:
+
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml \
+    --omni --port 8091 --trust-remote-code --enforce-eager
+```
+
+The default `qwen3_tts.yaml` uses `max_batch_size: 1` (single request). The `qwen3_tts_batch.yaml` config sets `max_batch_size: 4` for ~4x throughput.
+
 ## Supported Models
 
 ### Qwen3-TTS
@@ -347,6 +547,12 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 
 Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech online serving example](../user_guide/examples/online_serving/fish_speech.md) for details.
 
+### Voxtral TTS
+
+| Model | Description |
+|-------|-------------|
+| `mistralai/Voxtral-4B-TTS-2603` | 3B AR + FlowMatching TTS. Supports text-to-speech with preset voices. |
+
 ## Error Responses
 
 ### 400 Bad Request

docs/user_guide/diffusion/cache_acceleration/cache_dit.md

@@ -0,0 +1,285 @@
+# Cache-DiT Guide
+
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Example Script](#example-script)
+- [Acceleration Methods](#acceleration-methods)
+- [Configuration Parameters](#configuration-parameters)
+- [Best Practices](#best-practices)
+- [Troubleshooting](#troubleshooting)
+- [Summary](#summary)
+- [Additional Resources](#additional-resources)
+
+---
+
+## Overview
+
+Cache-DiT accelerates diffusion transformer models through intelligent caching mechanisms, providing significant speedup with minimal quality loss. It supports multiple acceleration techniques that can be combined for optimal performance:
+
+- **DBCache**: Dual Block Cache for reducing redundant computations
+- **TaylorSeer**: Taylor expansion-based forecasting for faster inference
+- **SCM**: Step Computation Masking for selective step computation
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+---
+
+## Quick Start
+
+### Basic Usage
+
+Enable cache-dit acceleration by simply setting `cache_backend="cache_dit"`:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_backend="cache_dit",  # Enable Cache-DiT with defaults
+)
+
+outputs = omni.generate(
+    "a beautiful landscape",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Note**: When `cache_config` is not provided, Cache-DiT uses optimized default values. See the [Configuration Parameters](#configuration-parameters) section for details.
+
+### Custom Configuration
+
+To customize cache-dit settings, provide a `cache_config` dictionary, for example:
+
+```python
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1,
+        "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4,
+        "residual_diff_threshold": 0.12,
+    },
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use the example script under `examples/offline_inference/text_to_image`:
+
+```bash
+cd examples/offline_inference/text_to_image
+python text_to_image.py \
+    --model Qwen/Qwen-Image \
+    --prompt "a cup of coffee on the table" \
+    --cache-backend cache_dit \
+    --num-inference-steps 50
+```
+
+See the [text_to_image.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/text_to_image/text_to_image.py) for detailed configuration options.
+
+The script uses cache-dit acceleration with a hybrid configuration combining DBCache, SCM, and TaylorSeer:
+
+```python
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_backend="cache_dit",
+    cache_config={
+        # Scheme: Hybrid DBCache + SCM + TaylorSeer
+        "Fn_compute_blocks": 1,  # Optimized for single-transformer models
+        "Bn_compute_blocks": 0,  # Number of backward compute blocks
+        "max_warmup_steps": 4,  # Maximum warmup steps (works for few-step models)
+        "residual_diff_threshold": 0.24,  # Higher threshold for more aggressive caching
+        "max_continuous_cached_steps": 3,  # Limit to prevent precision degradation
+        # TaylorSeer parameters [cache-dit only]
+        "enable_taylorseer": False,  # Disabled by default (not suitable for few-step models)
+        "taylorseer_order": 1,  # TaylorSeer polynomial order
+        # SCM (Step Computation Masking) parameters [cache-dit only]
+        "scm_steps_mask_policy": None,  # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
+        "scm_steps_policy": "dynamic",  # SCM steps policy: "dynamic" or "static"
+    }
+)
+```
+
+You can customize the configuration by modifying the `cache_config` dictionary to use only specific methods (e.g., DBCache only, DBCache + SCM, etc.) based on your quality and speed requirements.
+
+For image-to-image tasks, use the example script under `examples/offline_inference/image_to_image`:
+
+```bash
+cd examples/offline_inference/image_to_image
+python image_edit.py \
+    --model Qwen/Qwen-Image-Edit \
+    --prompt "make the sky more colorful" \
+    --image path/to/input/image.jpg \
+    --cache-backend cache_dit \
+    --num-inference-steps 50 \
+    --cache-dit-max-continuous-cached-steps 3 \
+    --cache-dit-residual-diff-threshold 0.24 \
+    --cache-dit-enable-taylorseer
+```
+
+See the [image_edit.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py) for detailed configuration options.
+
+### Online Serving
+
+```bash
+# Default configuration (recommended)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --cache-backend cache_dit
+
+# Custom configuration
+vllm serve Qwen/Qwen-Image --omni --port 8091 \
+  --cache-backend cache_dit \
+  --cache-config '{"Fn_compute_blocks": 1, "residual_diff_threshold": 0.12}'
+```
+
+---
+
+## Acceleration Methods
+
+For comprehensive illustration, please view Cache-DiT [User Guide](https://cache-dit.readthedocs.io/en/latest/user_guide/OVERVIEWS/).
+
+### 1. DBCache (Dual Block Cache)
+
+DBCache intelligently caches intermediate transformer block outputs when the residual differences between consecutive steps are small, reducing redundant computations without sacrificing quality.
+
+**Example Configuration**:
+
+```python
+cache_config={
+    "Fn_compute_blocks": 8,           # Use first 8 blocks for difference computation
+    "Bn_compute_blocks": 0,           # No additional fusion blocks
+    "max_warmup_steps": 8,            # Cache after 8 warmup steps
+    "residual_diff_threshold": 0.12,  # Lower threshold for faster inference
+    "max_cached_steps": -1,           # No limit on cached steps
+}
+```
+
+**Performance Tips**:
+
+- Default `Fn_compute_blocks=1` works well for most cases. Some models (e.g., FLUX.2-klein) use a larger value for `Fn_compute_blocks` for a balanced performance.
+- Increase `residual_diff_threshold` (e.g., 0.12-0.15) for faster inference with slight quality trade-off, or decrease from default 0.24 for higher quality.
+- Default `max_warmup_steps=4` is optimized for few-step models. Increase to 6-8 for more steps if needed.
+
+### 2. TaylorSeer
+
+TaylorSeer uses Taylor expansion to forecast future hidden states, allowing the model to skip some computation steps while maintaining quality.
+
+**Example Configuration**:
+
+```python
+cache_config={
+    "enable_taylorseer": True,
+    "taylorseer_order": 1,  # First-order Taylor expansion
+}
+```
+
+**Performance Tips**:
+
+- TaylorSeer is **not suitable for few-step distilled models**.
+- Use `taylorseer_order=1` for most cases (good balance of speed and quality).
+- Combine with DBCache for maximum acceleration.
+- Higher orders (2-3) may improve quality but reduce speed gains.
+
+### 3. SCM (Step Computation Masking)
+
+SCM allows you to specify which steps must be computed and which can use cached results, similar to LeMiCa/EasyCache style acceleration.
+
+`scm_steps_mask_policy` options (number of compute steps out of 28):
+
+| Policy | Compute Steps | Speed | Quality |
+|--------|--------------|-------|---------|
+| `None` (default) | All | Baseline | Best |
+| `"slow"` | 18 / 28 | Moderate | High |
+| `"medium"` | 15 / 28 | Balanced | Good |
+| `"fast"` | 11 / 28 | Fast | Moderate |
+| `"ultra"` | 8 / 28 | Fastest | Lower |
+
+**Example Configuration**:
+
+```python
+cache_config={
+    "scm_steps_mask_policy": "medium",  # Balanced speed/quality
+    "scm_steps_policy": "dynamic",      # Use dynamic cache
+}
+```
+
+**Performance Tips**:
+
+- SCM is disabled by default. Enable it by setting a policy value if you need additional acceleration.
+- Start with `"medium"` policy and adjust based on quality requirements.
+- Use `"fast"` or `"ultra"` for maximum speed when quality can be slightly compromised.
+- `"dynamic"` policy generally provides better quality than `"static"`.
+- SCM mask is automatically regenerated when `num_inference_steps` changes during inference.
+
+---
+
+## Configuration Parameters
+
+In `cache_config` passed to `Omni` constructor, it accepts the arguments of `DBCacheConfig` ([Cache-DiT API Reference](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/)). Key parameters are listed below:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `Fn_compute_blocks` | int | 1 | First n blocks for difference computation (optimized for single-transformer models) |
+| `Bn_compute_blocks` | int | 0 | Last n blocks for fusion |
+| `max_warmup_steps` | int | 4 | Steps before caching starts (optimized for few-step distilled models) |
+| `max_cached_steps` | int | -1 | Max cached steps (-1 = unlimited) |
+| `max_continuous_cached_steps` | int | 3 | Max consecutive cached steps (prevents precision degradation) |
+| `residual_diff_threshold` | float | 0.24 | Residual difference threshold (higher for more aggressive caching) |
+| `num_inference_steps` | int \| None | None | Initial inference steps for SCM mask generation (optional, auto-refreshed during inference) |
+| `enable_taylorseer` | bool | False | Enable TaylorSeer acceleration (not suitable for few-step distilled models) |
+| `taylorseer_order` | int | 1 | Taylor expansion order |
+| `scm_steps_mask_policy` | str \| None | None | SCM mask policy (None, "slow", "medium", "fast", "ultra") |
+| `scm_steps_policy` | str | "dynamic" | SCM computation policy ("dynamic" or "static") |
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- Production deployments requiring fast inference
+- Diffusion transformer models (DiT architecture)
+- Scenarios where 1.5x-3x speedup is valuable
+
+**Not for:**
+
+- Non-DiT architectures (use model-specific acceleration instead)
+- Models already using few-step distillation (< 10 steps)
+
+---
+
+## Troubleshooting
+
+### Common Issue 1: Quality Degradation
+
+**Symptoms**: Generated images have visible artifacts or lower quality
+
+**Solution**:
+```python
+# Reduce aggressiveness - use more conservative settings
+cache_config={
+    "residual_diff_threshold": 0.20,  # Lower threshold (closer to default 0.24)
+    "Fn_compute_blocks": 8,            # Use more blocks for better decisions
+    "max_warmup_steps": 6,             # Longer warmup
+    "scm_steps_mask_policy": "slow",   # More compute steps
+}
+```
+
+---
+
+## Summary
+
+Using Cache-DiT acceleration:
+
+1. ✅ **Enable Cache-DiT** - Set `cache_backend="cache_dit"` to get 1.5x-3x speedup with optimized defaults
+2. ✅ **(Optional) Customize** - Adjust `cache_config` parameters for specific speed/quality trade-offs

docs/user_guide/diffusion/cache_acceleration/teacache.md

@@ -0,0 +1,194 @@
+# TeaCache Guide
+
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Example Script](#example-script)
+- [Configuration Parameters](#configuration-parameters)
+- [Best Practices](#best-practices)
+- [Troubleshooting](#troubleshooting)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+TeaCache accelerates diffusion model inference by caching transformer computations when consecutive timesteps are similar, providing **1.5x-2.0x speedup** with minimal quality loss. It dynamically decides whether to reuse cached outputs based on input similarity, making it ideal for production deployments where inference speed matters without sacrificing generation quality.
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+---
+
+## Quick Start
+
+
+
+### Basic Usage
+
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_backend="tea_cache",
+)
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+### Custom Configuration
+
+```python
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_backend="tea_cache",
+    cache_config={
+        "rel_l1_thresh": 0.2,  # Controls speed/quality tradeoff
+    },
+)
+```
+
+### Using Environment Variable
+
+You can also enable TeaCache via environment variable:
+
+```bash
+export DIFFUSION_CACHE_BACKEND=tea_cache
+```
+
+Then initialize without explicitly setting `cache_backend`:
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_config={"rel_l1_thresh": 0.2}
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use python script under `examples/offline_inference/text_to_image/` or `examples/offline_inference/image_to_image/` with CLI:
+
+```bash
+# Text-to-image example
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image \
+  --cache-backend tea_cache
+
+# Image-to-image example
+python examples/offline_inference/image_to_image/image_edit.py \
+  --model Qwen/Qwen-Image-Edit \
+  --image input.png \
+  --prompt "Edit description" \
+  --cache-backend tea_cache \
+  --tea-cache-rel-l1-thresh 0.25
+```
+
+See the [text_to_image.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/text_to_image/text_to_image.py) or [image_edit.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py) for detailed configuration options.
+
+### Online Serving
+
+```bash
+# Default configuration
+vllm serve Qwen/Qwen-Image --omni --port 8091 --cache-backend tea_cache
+
+# Custom configuration
+vllm serve Qwen/Qwen-Image --omni --port 8091 \
+  --cache-backend tea_cache \
+  --cache-config '{"rel_l1_thresh": 0.2}'
+```
+
+---
+
+## Configuration Parameters
+
+In `OmniDiffusionConfig`
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `rel_l1_thresh` | float | `0.2` | Similarity threshold for cache reuse. Lower values prioritize quality (less caching), higher values prioritize speed (more caching). Suggested range: 0.1-0.8 |
+| `coefficients` | list[float] \| None | `None` | Polynomial coefficients for rescaling L1 distance. Must contain exactly 5 elements if provided. If `None`, uses model-specific defaults based on transformer type. |
+
+Users can find the default model coefficients in [`vllm_omni/diffusion/cache/teacache/config.py`](https://github.com/vllm-project/vllm-omni/blob/main/vllm_omni/diffusion/cache/teacache/config.py), for example:
+
+```python
+_MODEL_COEFFICIENTS = {
+    # Qwen-Image transformer coefficients from ComfyUI-TeaCache
+    # Tuned specifically for Qwen's dual-stream transformer architecture
+    # Used for all Qwen-Image Family pipelines, in general
+    "QwenImageTransformer2DModel": [
+        -4.50000000e02,
+        2.80000000e02,
+        -4.50000000e01,
+        3.20000000e00,
+        -2.00000000e-02,
+    ],
+    ...
+}
+```
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- Production deployments requiring faster inference, tolerant of minimal quality loss
+- Scenarios where 1.5-2x speedup is valuable
+- Useful for single-card acceleration
+
+**Not for:**
+
+- Maximum quality requirements where no degradation is acceptable
+- Very short inference runs (< 20 steps) where caching overhead may outweigh benefits
+
+
+---
+
+## Troubleshooting
+
+### Common Issue 1: Quality Degradation
+
+**Symptoms**: Generated images show artifacts, reduced detail, or inconsistent quality compared to non-cached results
+
+**Solution**:
+
+```python
+# Lower the threshold for more conservative caching
+cache_config={"rel_l1_thresh": 0.1}
+```
+
+### Common Issue 2: Limited Speedup
+
+**Symptoms**: Actual speedup is less than expected (< 1.3x)
+
+**Solutions**:
+1. Increase the threshold to enable more aggressive caching:
+   ```python
+   cache_config={"rel_l1_thresh": 0.8}
+   ```
+2. Ensure you're using sufficient inference steps (35+ recommended)
+3. Check that your model architecture is supported (see Supported Models section)
+
+---
+
+
+## Summary
+
+1. ✅ **Enable TeaCache** - Set `cache_backend="tea_cache"` to get 1.5x-2.0x speedup with optimized defaults
+2. ✅ **(Optional) Customize** - Adjust thresholds and polynomial coefficients for specific speed/quality trade-offs

docs/user_guide/diffusion/cache_dit_acceleration.md

@@ -1,228 +0,0 @@
-# Cache-DiT Acceleration Guide
-
-This guide explains how to use cache-dit acceleration in vLLM-Omni to speed up diffusion model inference.
-
-## Overview
-
-Cache-dit is a library that accelerates diffusion transformer models through intelligent caching mechanisms. It supports multiple acceleration techniques that can be combined for optimal performance:
-
-- **DBCache**: Dual Block Cache for reducing redundant computations
-- **TaylorSeer**: Taylor expansion-based forecasting for faster inference
-- **SCM**: Step Computation Masking for selective step computation
-
-## Quick Start
-
-### Basic Usage
-
-Enable cache-dit acceleration by simply setting `cache_backend="cache_dit"`. Cache-dit will use its recommended default parameters:
-
-```python
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-# Simplest way: just enable cache-dit with default parameters
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_backend="cache_dit",
-)
-
-images = omni.generate(
-    "a beautiful landscape",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-**Default Parameters**: When `cache_config` is not provided, cache-dit uses optimized default values. See the [Configuration Reference](#configuration-reference) section for a complete list of all parameters and their default values.
-
-### Custom Configuration
-
-To customize cache-dit settings, provide a `cache_config` dictionary, for example:
-
-```python
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_backend="cache_dit",
-    cache_config={
-        "Fn_compute_blocks": 1,
-        "Bn_compute_blocks": 0,
-        "max_warmup_steps": 4,
-        "residual_diff_threshold": 0.12,
-    },
-)
-```
-
-## Online Serving (OpenAI-Compatible)
-
-Enable Cache-DiT for online serving by passing `--cache-backend cache_dit` when starting the server:
-
-```bash
-# Use Cache-DiT default (recommended) parameters
-vllm serve Qwen/Qwen-Image --omni --port 8091 --cache-backend cache_dit
-```
-
-To customize Cache-DiT settings for online serving, pass a JSON string via `--cache-config`:
-
-```bash
-vllm serve Qwen/Qwen-Image --omni --port 8091 \
-  --cache-backend cache_dit \
-  --cache-config '{"Fn_compute_blocks": 1, "Bn_compute_blocks": 0, "max_warmup_steps": 4, "residual_diff_threshold": 0.12}'
-```
-
-## Acceleration Methods
-
-For comprehensive illustration, please view cache-dit [User_Guide](https://cache-dit.readthedocs.io/en/latest/user_guide/OVERVIEWS/)
-
-### 1. DBCache (Dual Block Cache)
-
-DBCache intelligently caches intermediate transformer block outputs when the residual differences between consecutive steps are small, reducing redundant computations without sacrificing quality.
-
-**Key Parameters**:
-
-- `Fn_compute_blocks` (int, default: 1): Number of **first n** transformer blocks used to compute stable feature differences. Higher values provide more accurate caching decisions but increase computation.
-- `Bn_compute_blocks` (int, default: 0): Number of **last n** transformer blocks used for additional fusion. These blocks act as an auto-scaler for approximate hidden states.
-- `max_warmup_steps` (int, default: 4): Number of initial steps where caching is disabled to ensure the model learns sufficient features before caching begins. Optimized for few-step distilled models.
-- `residual_diff_threshold` (float, default: 0.24): Threshold for residual difference. Higher values lead to faster performance but may reduce precision. Default uses a relatively higher threshold for more aggressive caching.
-- `max_cached_steps` (int, default: -1): Maximum number of cached steps. Set to -1 for unlimited caching.
-- `max_continuous_cached_steps` (int, default: 3): Maximum number of consecutive cached steps. Limits consecutive caching to prevent precision degradation.
-
-**Example Configuration**:
-
-```python
-cache_config={
-    "Fn_compute_blocks": 8,      # Use first 8 blocks for difference computation
-    "Bn_compute_blocks": 0,       # No additional fusion blocks
-    "max_warmup_steps": 8,        # Cache after 8 warmup steps
-    "residual_diff_threshold": 0.12,  # Higher threshold for faster inference
-    "max_cached_steps": -1,        # No limit on cached steps
-}
-```
-
-**Performance Tips**:
-
-- Default `Fn_compute_blocks=1` works well for most cases. Increase to 8-12 for larger models or when more accuracy is needed
-- Increase `residual_diff_threshold` (e.g., 0.12-0.15) for faster inference with slight quality trade-off, or decrease from default 0.24 for higher quality
-- Default `max_warmup_steps=4` is optimized for few-step models. Increase to 6-8 for more steps if needed
-
-### 2. TaylorSeer
-
-TaylorSeer uses Taylor expansion to forecast future hidden states, allowing the model to skip some computation steps while maintaining quality.
-
-**Key Parameters**:
-
-- `enable_taylorseer` (bool, default: False): Enable TaylorSeer acceleration
-- `taylorseer_order` (int, default: 1): Order of Taylor expansion. Higher orders provide better accuracy but require more computation.
-
-**Example Configuration**:
-
-```python
-cache_config={
-    "enable_taylorseer": True,
-    "taylorseer_order": 1,  # First-order Taylor expansion
-}
-```
-
-**Performance Tips**:
-
-- Use `taylorseer_order=1` for most cases (good balance of speed and quality)
-- Combine with DBCache for maximum acceleration
-- Higher orders (2-3) may improve quality but reduce speed gains
-
-### 3. SCM (Step Computation Masking)
-
-SCM allows you to specify which steps must be computed and which can use cached results, similar to LeMiCa/EasyCache style acceleration.
-
-**Key Parameters**:
-
-- `scm_steps_mask_policy` (str | None, default: None): Predefined mask policy. Options:
-  - `None`: SCM disabled (default)
-  - `"slow"`: More compute steps, higher quality (18 compute steps out of 28)
-  - `"medium"`: Balanced (15 compute steps out of 28)
-  - `"fast"`: More cache steps, faster inference (11 compute steps out of 28)
-  - `"ultra"`: Maximum speed (8 compute steps out of 28)
-- `scm_steps_policy` (str, default: "dynamic"): Policy for cached steps:
-  - `"dynamic"`: Use dynamic cache for masked steps (recommended)
-  - `"static"`: Use static cache for masked steps
-
-**Example Configuration**:
-
-```python
-cache_config={
-    "scm_steps_mask_policy": "medium",  # Balanced speed/quality
-    "scm_steps_policy": "dynamic",      # Use dynamic cache
-}
-```
-
-**Performance Tips**:
-
-- SCM is disabled by default (`scm_steps_mask_policy=None`). Enable it by setting a policy value if you need additional acceleration
-- Start with `"medium"` policy and adjust based on quality requirements
-- Use `"fast"` or `"ultra"` for maximum speed when quality can be slightly compromised
-- `"dynamic"` policy generally provides better quality than `"static"`
-- SCM mask is automatically regenerated when `num_inference_steps` changes during inference
-
-## Configuration Reference
-
-### DiffusionCacheConfig Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `Fn_compute_blocks` | int | 1 | First n blocks for difference computation (optimized for single-transformer models) |
-| `Bn_compute_blocks` | int | 0 | Last n blocks for fusion |
-| `max_warmup_steps` | int | 4 | Steps before caching starts (optimized for few-step distilled models) |
-| `max_cached_steps` | int | -1 | Max cached steps (-1 = unlimited) |
-| `max_continuous_cached_steps` | int | 3 | Max consecutive cached steps (prevents precision degradation) |
-| `residual_diff_threshold` | float | 0.24 | Residual difference threshold (higher for more aggressive caching) |
-| `num_inference_steps` | int \| None | None | Initial inference steps for SCM mask generation (optional, auto-refreshed during inference) |
-| `enable_taylorseer` | bool | False | Enable TaylorSeer acceleration (not suitable for few-step distilled models) |
-| `taylorseer_order` | int | 1 | Taylor expansion order |
-| `scm_steps_mask_policy` | str \| None | None | SCM mask policy (None, "slow", "medium", "fast", "ultra") |
-| `scm_steps_policy` | str | "dynamic" | SCM computation policy ("dynamic" or "static") |
-
-## Example: Accelerate Text-to-Image Generation with CacheDiT
-
-See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example with cache-dit acceleration.
-
-```bash
-# Enable cache-dit with hybrid acceleration
-cd examples/offline_inference/text_to_image
-python text_to_image.py \
-    --model Qwen/Qwen-Image \
-    --prompt "a cup of coffee on the table" \
-    --cache-backend cache_dit \
-    --num-inference-steps 50
-```
-
-
-The script uses cache-dit acceleration with a hybrid configuration combining DBCache, SCM, and TaylorSeer:
-
-```python
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_backend="cache_dit",
-    cache_config={
-        # Scheme: Hybrid DBCache + SCM + TaylorSeer
-        # DBCache
-        "Fn_compute_blocks": 8,
-        "Bn_compute_blocks": 0,
-        "max_warmup_steps": 4,
-        "residual_diff_threshold": 0.12,
-        # TaylorSeer
-        "enable_taylorseer": True,
-        "taylorseer_order": 1,
-        # SCM
-        "scm_steps_mask_policy": "fast",  # Set to None to disable SCM
-        "scm_steps_policy": "dynamic",
-    },
-)
-```
-
-You can customize the configuration by modifying the `cache_config` dictionary to use only specific methods (e.g., DBCache only, DBCache + SCM, etc.) based on your quality and speed requirements.
-
-To test another model, you can modify `--model` with the target model identifier like `Tongyi-MAI/Z-Image-Turbo` and update `cache_config` according the model architecture (e.g., number of transformer blocks).
-
-
-## Additional Resources
-
-- [Cache-DiT User Guide](https://cache-dit.readthedocs.io/en/latest/user_guide/OVERVIEWS/)
-- [Cache-DiT Benchmark](https://cache-dit.readthedocs.io/en/latest/benchmark/HYBRID_CACHE/)
-- [DBCache Technical Details](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/)