lint

Signed-off-by: khluu <khluu000@gmail.com>

.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

@@ -10,7 +10,7 @@ steps:
           DOCKER_BUILDKIT: "1"
           # Buildkite will automatically replace this with the actual commit hash
           VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.18.0"
+          VLLM_VERSION: "v0.19.0"
         priority: 100
         timeout_in_minutes: 60
         soft_fail: true

.buildkite/pipeline.yml

@@ -1,6 +1,21 @@
+# Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci
+# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs if nightly-test
+# / main NIGHTLY so upload-nightly is not skipped together with test-ready/test-merge.
+#
+# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite).
+steps:
+  - label: ":github: Resolve skip-ci & upload pipeline"
+    key: upload-ci-pipeline
+    commands:
+      - "bash .buildkite/scripts/upload_pipeline_with_skip_ci.sh"
+    agents:
+      queue: "cpu_queue_premerge"
+
+---
 steps:
   - label: ":docker: Build image"
     key: image-build
+    if: __IMAGE_BUILD_IF__
     commands:
       - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
       - "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
@@ -13,7 +28,7 @@ steps:
   - label: "Upload Ready Pipeline"
     depends_on: image-build
     key: upload-ready-pipeline
-    if: build.branch != "main" && build.pull_request.labels includes "ready"
+    if: __UPLOAD_READY_IF__
     commands:
       - buildkite-agent pipeline upload .buildkite/test-ready.yml
     agents:
@@ -23,7 +38,7 @@ steps:
   - label: "Upload Merge Pipeline"
     depends_on: image-build
     key: upload-merge-pipeline
-    if: build.branch == "main" && build.env("NIGHTLY") != "1"
+    if: __UPLOAD_MERGE_IF__
     commands:
       - buildkite-agent pipeline upload .buildkite/test-merge.yml
     agents:

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

@@ -0,0 +1,88 @@
+#!/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 --user $(id -u):$(id -g) -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.bak 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py && rm -f .buildkite/scripts/generate-nightly-index.py.bak
+
+# 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
+    s3_version="v$pure_version"
+    echo "Re-generating indices for /omni/$s3_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 "$s3_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/$s3_version/"
+fi

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

@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import json
+import re
+import sys
+from dataclasses import asdict, dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+from urllib.parse import quote
+
+
+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}"
 
@@ -126,8 +133,6 @@ if [[ $commands == *"--shard-id="* ]]; then
         -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}" \
@@ -160,8 +165,6 @@ else
           -e MIOPEN_DEBUG_CONV_GEMM=0 \
           -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/scripts/upload_pipeline_with_skip_ci.sh

@@ -0,0 +1,137 @@
+#!/usr/bin/env bash
+# Evaluate docs-only skip-ci and upload continuation steps from the same `.buildkite/pipeline.yml`
+# (YAML document after the first `---`). Buildkite `if` is evaluated at upload time.
+set -euo pipefail
+
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+PIPELINE_YML="${ROOT}/.buildkite/pipeline.yml"
+
+# Prints a single digit to stdout: 1 = skip image CI, 0 = run. Logs go to stderr.
+is_docs_only_change() {
+  local file_path
+  local has_any=0
+
+  while IFS= read -r file_path; do
+    [[ -z "${file_path}" ]] && continue
+    has_any=1
+
+    if [[ "${file_path}" == docs/* ]]; then
+      continue
+    fi
+    if [[ "${file_path}" == *.md ]]; then
+      continue
+    fi
+    if [[ "${file_path}" == "mkdocs.yaml" ]]; then
+      continue
+    fi
+    return 1
+  done
+
+  [[ "${has_any}" -eq 1 ]]
+}
+
+resolve_skip_ci() {
+  local is_pr_build=0
+  local files
+  local base_branch base_ref
+
+  if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" && -n "${BUILDKITE_PULL_REQUEST:-}" ]]; then
+    is_pr_build=1
+  fi
+
+  if [[ "${is_pr_build}" -eq 1 ]]; then
+    base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
+    if ! git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
+      echo "resolve_skip_ci: origin/${base_branch} not found locally; trying fetch" >&2
+      git fetch --depth=200 origin "${base_branch}" >/dev/null 2>&1 || true
+    fi
+
+    base_ref=""
+    if git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
+      base_ref="origin/${base_branch}"
+    elif git rev-parse --verify "${base_branch}" >/dev/null 2>&1; then
+      base_ref="${base_branch}"
+    else
+      echo "resolve_skip_ci: cannot resolve PR base ${base_branch}; skip-ci=0" >&2
+      echo -n 0
+      return 0
+    fi
+
+    if ! files="$(git diff --name-only "${base_ref}...${BUILDKITE_COMMIT}" 2>/dev/null)"; then
+      echo "resolve_skip_ci: failed to compute PR changed files; skip-ci=0" >&2
+      echo -n 0
+      return 0
+    fi
+  elif [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then
+    if ! git rev-parse --verify "${BUILDKITE_COMMIT}^" >/dev/null 2>&1; then
+      echo "resolve_skip_ci: commit has no parent on main; skip-ci=0" >&2
+      echo -n 0
+      return 0
+    fi
+    if ! files="$(git diff --name-only "${BUILDKITE_COMMIT}^..${BUILDKITE_COMMIT}" 2>/dev/null)"; then
+      echo "resolve_skip_ci: failed to compute main changed files; skip-ci=0" >&2
+      echo -n 0
+      return 0
+    fi
+  else
+    echo "resolve_skip_ci: not PR/main build; skip-ci=0" >&2
+    echo -n 0
+    return 0
+  fi
+
+  if is_docs_only_change <<< "${files}"; then
+    echo "resolve_skip_ci: docs-only change detected; skip-ci=1" >&2
+    echo -n 1
+    return 0
+  fi
+
+  echo "resolve_skip_ci: non-doc changes detected; skip-ci=0" >&2
+  echo -n 0
+}
+
+SKIP_CI="$(resolve_skip_ci)"
+
+if [[ ! -f "${PIPELINE_YML}" ]]; then
+  echo "upload_pipeline_with_skip_ci: missing ${PIPELINE_YML}" >&2
+  exit 1
+fi
+
+export ROOT SKIP_CI PIPELINE_YML
+python3 <<'PY' | buildkite-agent pipeline upload
+import os
+import pathlib
+
+path = pathlib.Path(os.environ["PIPELINE_YML"])
+text = path.read_text(encoding="utf-8")
+sep = "\n---\n"
+if sep not in text:
+    raise SystemExit(
+        "upload_pipeline_with_skip_ci: .buildkite/pipeline.yml must contain a '\\n---\\n' separator "
+        "(document 1 = bootstrap, document 2 = uploaded steps)"
+    )
+_, continuation = text.split(sep, 1)
+
+skip = os.environ.get("SKIP_CI") == "1"
+# When docs-only skip-ci: skip default CI image, but still build for L4 nightly (PR label nightly-test or
+# main NIGHTLY=1), otherwise upload-nightly (depends_on image-build) would be skipped too.
+nightly_only = (
+    '(build.pull_request.labels includes "nightly-test") '
+    '|| (build.branch == "main" && build.env("NIGHTLY") == "1")'
+)
+# Placeholder in pipeline.yml is `if: __IMAGE_BUILD_IF__` (valid YAML); replace value only.
+if skip:
+    rep = f"'{nightly_only}'"
+    ready_rep = "'false'"
+    merge_rep = "'false'"
+else:
+    rep = "'true'"
+    ready_rep = "'build.branch != \"main\" && build.pull_request.labels includes \"ready\"'"
+    merge_rep = "'build.branch == \"main\" && build.env(\"NIGHTLY\") != \"1\"'"
+rendered = (
+    continuation
+    .replace("__IMAGE_BUILD_IF__", rep)
+    .replace("__UPLOAD_READY_IF__", ready_rep)
+    .replace("__UPLOAD_MERGE_IF__", merge_rep)
+)
+print(rendered, end="")
+PY

.buildkite/test-amd-merge.yml

@@ -134,15 +134,33 @@ steps:
     - 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 E2E Test"
+- label: "Qwen3-TTS CustomVoice 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
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.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_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

.buildkite/test-amd-ready.yaml

@@ -124,7 +124,8 @@ steps:
   commands:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
+    - 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

.buildkite/test-merge.yml

@@ -56,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:
@@ -111,8 +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/diffusion/distributed/test_ulysses_uaa_perf.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:
@@ -174,7 +173,7 @@ steps:
                 pytest -s -v tests/engine/test_async_omni_engine_abort.py
         '
     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
@@ -193,8 +192,7 @@ steps:
     commands:
       - 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 "advanced_model" --run-level "advanced_model"
+      - 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:
@@ -208,15 +206,59 @@ steps:
           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
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Qwen3-TTS Base 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_base.py tests/e2e/offline_inference/test_qwen3_tts_base.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
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Omni Model Test with H100"
     timeout_in_minutes: 30
     depends_on: upload-merge-pipeline
     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/online_serving/test_mimo_audio.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:
@@ -302,8 +344,7 @@ steps:
           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 -m "advanced_model" --run-level "advanced_model"
-          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
+          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"
@@ -347,8 +388,7 @@ steps:
         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 "advanced_model" --run-level "advanced_model"
-          pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+          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"
         '
     agents:
       queue: "mithril-h100-pool"
@@ -368,6 +408,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:

.buildkite/test-nightly.yml

@@ -6,16 +6,10 @@ steps:
     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=$$?
-        pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
-        EXIT3=$$?
-        pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
-        EXIT4=$$?
-        exit $$((EXIT1 | EXIT2 | EXIT3 | EXIT4))
+        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:
@@ -50,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" || 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:
@@ -77,7 +72,7 @@ steps:
     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 -k "not test_wan22_expansion" -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 and not test_wan_2_1_vace_expansion" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -112,13 +107,13 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: ":full_moon: Diffusion Model Wan22 completed Test with H100"
+  - label: ":full_moon: Diffusion Model (Wan) Test with H100"
     timeout_in_minutes: 90
     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_wan22_expansion.py -m "advanced_model" --run-level "advanced_model"
+      - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py -m "advanced_model" --run-level "advanced_model"
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -153,7 +148,7 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: ":full_moon: Diffusion Model Test with L4"
+  - 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"
@@ -175,7 +170,7 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
 
-  - label: ":full_moon: Documentation Example Code Test with H100"
+  - label: ":full_moon: Doc Example Code Test with H100"
     timeout_in_minutes: 60
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
@@ -217,30 +212,7 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
-    timeout_in_minutes: 30
-    depends_on: upload-nightly-pipeline
-    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
-    agents:
-      queue: "gpu_4_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: ":full_moon: Omni Model Perf Test & Test Case Statistics"
+  - label: ":full_moon: Omni Model Perf Test & Testcase Statistics with H100"
     key: nightly-omni-performance
     timeout_in_minutes: 180
     depends_on: upload-nightly-pipeline
@@ -248,6 +220,7 @@ steps:
     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
@@ -375,9 +348,51 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: ":full_moon: Qwen-Image Diffusion Perf Test with H100"
+  - 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: 300
+    timeout_in_minutes: 180
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
     commands:

.buildkite/test-ready.yml

@@ -180,7 +180,7 @@ steps:
                 pytest -s -v tests/engine/test_async_omni_engine_abort.py
         '
     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
@@ -216,44 +216,42 @@ steps:
           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
@@ -273,7 +271,7 @@ steps:
               - 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
+                    nvidia.com/gpu: 1
                 volumeMounts:
                   - name: devshm
                     mountPath: /dev/shm
@@ -303,13 +301,13 @@ steps:
     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
@@ -322,6 +320,29 @@ steps:
           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
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Voxtral-TTS E2E Test"
     timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
@@ -330,8 +351,7 @@ steps:
         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 "advanced_model" --run-level "advanced_model"
-          pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "core_model" --run-level "core_model"
         '
     agents:
       queue: "mithril-h100-pool"
@@ -528,3 +548,46 @@ steps:
                 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 %}

.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

.gitignore

@@ -83,6 +83,9 @@ target/
 profile_default/
 ipython_config.py
 
+# uv
+uv.lock
+
 # pyenv
 #   For a library or package, you might want to ignore these files since the code is
 #   intended to run in multiple environments; otherwise, check them in:

README.md

@@ -16,7 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---
 
 *Latest News* 🔥
-
+- [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.

benchmarks/qwen3-tts/README.md

@@ -32,6 +32,9 @@ 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 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
 

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(

docker/Dockerfile.ci

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.18.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.cuda

@@ -0,0 +1,22 @@
+ARG BASE_IMAGE=vllm/vllm-openai:v0.19.0
+FROM ${BASE_IMAGE}
+
+ARG COMMON_WORKDIR=/app
+
+WORKDIR ${COMMON_WORKDIR}
+
+# Step 1: Setup - Install system dependencies
+RUN apt-get update && \
+    apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
+
+# Step 2: Copy vllm-omni code and install
+COPY . ${COMMON_WORKDIR}/vllm-omni
+RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir "."
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+ENTRYPOINT []

docker/Dockerfile.npu

@@ -1,20 +1,20 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.17.0rc1
+ARG VLLM_ASCEND_TAG=v0.18.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-WORKDIR /vllm-workspace/vllm
-RUN git fetch origin --tags && git checkout v0.18.0
+# 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
+# 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}

docker/Dockerfile.npu.a3

@@ -1,20 +1,20 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
+ARG VLLM_ASCEND_TAG=v0.18.0rc1-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-WORKDIR /vllm-workspace/vllm
-RUN git fetch origin --tags && git checkout v0.18.0
+# 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
+# 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}

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.18.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.

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.18.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

@@ -54,20 +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
-      - Step Execution: user_guide/diffusion/step_execution.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:
@@ -92,6 +100,8 @@ 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

docs/api/README.md

@@ -5,7 +5,6 @@
 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.cfg_companion_tracker.CfgCompanionTracker][]
 - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
 - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]

docs/community/governance.md

@@ -37,10 +37,12 @@ Committers have write access and merge rights. They typically have deep expertis
 - [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
 - [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
 - [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
+- [@lishunyang12](https://github.com/lishunyang12): 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):
+- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism
+- [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
 - [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
 
 ## Meetings

docs/configuration/README.md

@@ -20,7 +20,4 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
 
 ## 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/contributing/ci/CI_5levels.md

@@ -168,7 +168,6 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
-│   ├── 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                  ⬜
@@ -272,7 +271,7 @@ Before entering specific testing levels, the project establishes two common spec
 
 L1 and L2 level testing form the foundation of the quality assurance system. L1 level testing focuses on verifying the internal logic correctness of code units (e.g., functions, classes), ensuring each independent component behaves as designed.
 
-L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle  .
+L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle.
 
 
 
@@ -548,103 +547,19 @@ L4 level testing is a comprehensive quality audit before a version release. It e
 -   ***Execution Environment***: ***GPU*** server clusters to meet the resource demands of performance testing.
 -   ***Script Example***:
 
-???+ example "Test Examples"
+??? example "Test Examples: Documentation Example Tests"
 
-    When adding L4-level ***documentation example Tests***, please pay attention to the following guides.
+    --8<-- "docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md"
 
-    --8<-- "docs/contributing/ci/test_examples/doc_example_tests.inc.md"
+??? example "Test Examples: Performance Tests"
 
-    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:
+    --8<-- "docs/contributing/ci/test_examples/l4_performance_tests.inc.md"
 
-    ```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: Functionality Tests"
 
-    **Parameter Explanation**
+    --8<-- "docs/contributing/ci/test_examples/l4_functionality_tests.inc.md"
 
-    *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.  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
 

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_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

@@ -73,7 +73,6 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
-│   ├── 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                  ⬜

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.
 
 ---
 
@@ -820,7 +802,7 @@ omni = Omni(model="your-model", enable_layerwise_offload=True)
 
 ```python
 class WanTransformer3DModel(nn.Module):
-    _layerwise_offload_blocks_attr = "blocks"  # Attribute name containing transformer blocks
+    _layerwise_offload_blocks_attrs = ["blocks"]  # Attribute name containing transformer blocks
 
     def __init__(self):
         self.blocks = nn.ModuleList([...])  # Transformer blocks

docs/design/feature/cache_dit.md

@@ -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.
 
 ---
 

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/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/feature/vae_parallel.md

@@ -1,14 +1,15 @@
 # VAE Patch Parallelism
 
 This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
-We use **Qwen-Image** as the reference implementation.
+We use **Qwen-Image** as the reference implementation for decode parallel, and **Wan2.2** for encode parallel.
 
 ---
 
 ## Table of Contents
 
 - [Overview](#overview)
-- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Step-by-Step Implementation (Decode)](#step-by-step-implementation-decode)
+- [Encode Parallel Implementation](#encode-parallel-implementation)
 - [Testing](#testing)
 - [Reference Implementations](#reference-implementations)
 - [Summary](#summary)
@@ -19,13 +20,13 @@ We use **Qwen-Image** as the reference implementation.
 
 ### What is Vae Patch parallel?
 
-**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is:
+**VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is:
 
 + Split into multiple spatial tiles
 
 + Distributed across multiple ranks
 
-+ Decoded in parallel
++ Encoded/Decoded in parallel
 
 + Merged to reconstruct the final output
 
@@ -35,10 +36,17 @@ This approach:
 
 + Reduces peak memory usage per device
 
-+ Accelerates decoding latency
++ Accelerates encoding/decoding latency
+
+### When to Use Encode vs Decode Parallel
+
+| Operation | Use Case | Example |
+|-----------|----------|---------|
+| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
+| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
 
 ### Architecture
-We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding.
+We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding.
 
 The executor is model-agnostic and accepts three function parameters:
 
@@ -84,7 +92,7 @@ Therefore:
 
 + Merge must perform blending to avoid seams
 
-## Step-by-Step Implementation
+## Step-by-Step Implementation (Decode)
 
 ### Step 1: Implement DistributedAutoencoderKLQwenImage
 `QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
@@ -205,14 +213,14 @@ def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid
 We need to override tiled_decode, the main logic is:
 + check distributed is enabled
 + select split/exec/merge
-+ Invoke self.distributed_decoder.execute to decode
++ Invoke self.distributed_executor.execute to decode
 ```
 def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
     if not self.is_distributed_enabled():
         return super().tiled_decode(z, return_dict=return_dict)
 
     logger.info("Decode run with distributed executor")
-    result = self.distributed_decoder.execute(
+    result = self.distributed_executor.execute(
         z,
         DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
         broadcast_result=True,
@@ -243,6 +251,166 @@ class YourModelPipeline(nn.Module):
 +       ).to(self.device)
 ```
 
+## Encode Parallel Implementation
+
+For models that require VAE encoding (e.g., Image-to-Video), you can also parallelize the encode operation. We use **Wan2.2** as the reference implementation.
+
+### Step 1: Implement encode_tile_split
+
+Similar to decode, split the input tensor into tiles. Key considerations:
+
++ **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly
++ **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x)
+
+```python
+def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+    _, _, num_frames, height, width = x.shape
+    encode_spatial_compression_ratio = self.spatial_compression_ratio
+
+    # Scale tile parameters for patchified coordinate system
+    tile_sample_min_height = self.tile_sample_min_height
+    tile_sample_min_width = self.tile_sample_min_width
+    tile_sample_stride_height = self.tile_sample_stride_height
+    tile_sample_stride_width = self.tile_sample_stride_width
+
+    if self.config.patch_size is not None:
+        # When input is patchified, scale tile parameters accordingly
+        encode_spatial_compression_ratio = self.spatial_compression_ratio // self.config.patch_size
+        tile_sample_min_height = tile_sample_min_height // self.config.patch_size
+        tile_sample_min_width = tile_sample_min_width // self.config.patch_size
+        tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
+        tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
+
+    latent_height = height // encode_spatial_compression_ratio
+    latent_width = width // encode_spatial_compression_ratio
+
+    tile_latent_min_height = tile_sample_min_height // encode_spatial_compression_ratio
+    tile_latent_min_width = tile_sample_min_width // encode_spatial_compression_ratio
+    tile_latent_stride_height = tile_sample_stride_height // encode_spatial_compression_ratio
+    tile_latent_stride_width = tile_sample_stride_width // encode_spatial_compression_ratio
+
+    blend_height = tile_latent_min_height - tile_latent_stride_height
+    blend_width = tile_latent_min_width - tile_latent_stride_width
+
+    tiletask_list = []
+    # Use temporal compression ratio from config instead of hardcoding
+    temporal_compression = self.config.scale_factor_temporal
+
+    for i in range(0, height, tile_sample_stride_height):
+        for j in range(0, width, tile_sample_stride_width):
+            time_list = []
+            frame_range = 1 + (num_frames - 1) // temporal_compression
+            for k in range(frame_range):
+                if k == 0:
+                    tile = x[:, :, :1, i : i + tile_sample_min_height, j : j + tile_sample_min_width]
+                else:
+                    tile = x[
+                        :, :,
+                        1 + temporal_compression * (k - 1) : 1 + temporal_compression * k,
+                        i : i + tile_sample_min_height,
+                        j : j + tile_sample_min_width,
+                    ]
+                time_list.append(tile)
+            tiletask_list.append(
+                TileTask(len(tiletask_list), (i // tile_sample_stride_height, j // tile_sample_stride_width),
+                         time_list, workload=time_list[0].shape[3] * time_list[0].shape[4])
+            )
+
+    grid_spec = GridSpec(
+        split_dims=(3, 4),
+        grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+        tile_spec={
+            "latent_height": latent_height, "latent_width": latent_width,
+            "blend_height": blend_height, "blend_width": blend_width,
+            "tile_latent_stride_height": tile_latent_stride_height,
+            "tile_latent_stride_width": tile_latent_stride_width,
+        },
+        output_dtype=self.dtype,
+    )
+    return tiletask_list, grid_spec
+```
+
+### Step 2: Implement encode_tile_exec
+
+```python
+def encode_tile_exec(self, task: TileTask) -> torch.Tensor:
+    """Encode a single sample tile into latent space."""
+    self.clear_cache()
+    time = []
+    for k, tile in enumerate(task.tensor):
+        self._enc_conv_idx = [0]
+        encoded = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+        encoded = self.quant_conv(encoded)
+        time.append(encoded)
+    result = torch.cat(time, dim=2)
+    self.clear_cache()
+    return result
+```
+
+### Step 3: Implement encode_tile_merge
+
+```python
+def encode_tile_merge(
+    self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec
+) -> torch.Tensor:
+    """Merge encoded tiles into a full latent tensor."""
+    grid_h, grid_w = grid_spec.grid_shape
+    result_rows = []
+    for i in range(grid_h):
+        result_row = []
+        for j in range(grid_w):
+            tile = coord_tensor_map[(i, j)]
+            if i > 0:
+                tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+            if j > 0:
+                tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+            result_row.append(tile[:, :, :,
+                : grid_spec.tile_spec["tile_latent_stride_height"],
+                : grid_spec.tile_spec["tile_latent_stride_width"]])
+        result_rows.append(torch.cat(result_row, dim=-1))
+
+    enc = torch.cat(result_rows, dim=3)[
+        :, :, :, : grid_spec.tile_spec["latent_height"], : grid_spec.tile_spec["latent_width"]
+    ]
+    return enc
+```
+
+### Step 4: Override tiled_encode method
+
+Override `tiled_encode` instead of `encode`. The parent's `_encode()` handles patchify before calling `tiled_encode()`, so input `x` is already patchified.
+
+```python
+def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+    """
+    Encode using distributed VAE executor.
+
+    Note: x is already patchified by parent's _encode() before calling this method.
+    """
+    if not self.is_distributed_enabled():
+        return super().tiled_encode(x)
+
+    self.clear_cache()
+    result = self.distributed_executor.execute(
+        x,
+        DistributedOperator(
+            split=self.encode_tile_split,
+            exec=self.encode_tile_exec,
+            merge=self.encode_tile_merge,
+        ),
+        broadcast_result=True,  # Latents needed by all ranks for diffusion
+    )
+    self.clear_cache()
+    return result
+```
+
+**Key differences from decode parallel:**
+
+| Aspect | Decode Parallel | Encode Parallel |
+|--------|-----------------|-----------------|
+| `broadcast_result` | Often `False` (only rank 0 needs output) | `True` (all ranks need latents for diffusion) |
+| Patchify | Applied in merge (unpatchify) | Handled by parent `_encode()` before `tiled_encode()` |
+| Temporal chunking | Frame-by-frame | Chunk-based (e.g., 1 + 4n frames) |
+
 ## Testing
 Verify numerical consistency between:
 + vae_patch_parallel_size = 1
@@ -272,18 +440,20 @@ When vae_patch_parallel_size is larger than the DiT world size, it will automati
 
 Complete examples in the codebase:
 
-| Model | Path | Notes |
-|-------|------|-------|
-| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
-| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
-| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage |
+| Model | Path | Decode Parallel | Encode Parallel |
+|-------|------|-----------------|-----------------|
+| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | ✅ | ❌ |
+| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | ✅ | ✅ |
+| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | ✅ | ❌ |
 
 ---
 
 ## Summary
 
-Adding Vae Patch Parallel support to diffusion model:
+Adding VAE Patch Parallel support to diffusion model:
 
-1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
-2. **Change vae model in pipeline to Distributed Vae**
-3. **Test** - Verify with `tensor_parallel_size=N` quality
+1. **Implement Distributed VAE** - Inherit from base VAE class and `DistributedVaeMixin`
+2. **Decode Parallel** - Refactor `tiled_decode` into `tile_split`/`tile_exec`/`tile_merge`
+3. **Encode Parallel** (optional) - Implement `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models
+4. **Change VAE model in pipeline** - Use the distributed version
+5. **Test** - Verify numerical consistency with `vae_patch_parallel_size=1` vs `N`

docs/design/module/async_omni_architecture.md

@@ -69,7 +69,7 @@
 [5] Orchestrator._orchestration_loop (loop)
     -> poll stage output
        - llm stage: await get_output_async()
-       - diffusion stage: get_diffusion_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:
@@ -112,7 +112,7 @@ sequenceDiagram
     ORCH->>S0: add_request_async
 
     loop poll route forward
-        ORCH->>S0: get_output_async / get_diffusion_output_async
+        ORCH->>S0: get_output_async / get_diffusion_output_nowait
         ORCH->>ORCH: _route_output
         alt need forward to next stage
             ORCH->>SN: add_request_async

docs/design/module/dit_module.md

@@ -192,7 +192,7 @@ class _BaseScheduler(SchedulerInterface):
         self._waiting = deque()
         self._running = []
         self._finished_req_ids = set()
-        self._max_batch_size = 1
+        self.max_num_running_reqs = 1
 ```
 
 **Design Features**:
@@ -201,7 +201,7 @@ class _BaseScheduler(SchedulerInterface):
 
 - **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
 
-- **Current constraint**: `_max_batch_size` remains `1` because the current engine path is still synchronous request-mode execution.
+- **Current constraint**: `max_num_running_reqs` remains `1` because the current engine path is still synchronous request-mode execution.
 
 #### 2.4 Current `RequestScheduler` Policy
 

docs/getting_started/installation/README.md

@@ -6,4 +6,5 @@ vLLM-Omni supports the following hardware platforms:
     - [NVIDIA CUDA](gpu.md)
     - [AMD ROCm](gpu.md)
     - [Intel XPU](gpu.md)
+    - [MThreads MUSA](gpu.md)
 - [NPU](npu.md)

docs/getting_started/installation/gpu.md

@@ -22,6 +22,10 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:requirements"
 
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:requirements"
+
 ## Set up using Python
 
 ### Create a new Python environment
@@ -44,6 +48,10 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:pre-built-wheels"
 
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:pre-built-wheels"
+
 [](){ #build-from-source }
 
 ### Build wheel from source
@@ -60,6 +68,10 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-wheel-from-source"
 
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:build-wheel-from-source"
+
 ## Set up using Docker
 
 ### Pre-built images
@@ -76,6 +88,10 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:pre-built-images"
 
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:pre-built-images"
+
 ### Build your own docker image
 
 === "AMD ROCm"
@@ -85,3 +101,7 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 === "Intel XPU"
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-docker"
+
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:build-docker"

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.18.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.18.0 --torch-backend=auto
+uv pip install vllm==0.19.0 --torch-backend=auto
 ```
 
-The 0.18.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+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.18.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.18.0/vllm-0.18.0+cu129-cp38-abi3-manylinux_2_35_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

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

@@ -0,0 +1,65 @@
+# --8<-- [start:requirements]
+
+- GPU: Moore Threads GPU with MUSA SDK installed (validated on MTT S5000)
+
+# --8<-- [end:requirements]
+# --8<-- [start:set-up-using-python]
+
+vLLM-Omni for MUSA requires building from source. Pre-built wheels are not currently available.
+
+!!! note
+    MUSA platform requires vLLM-MUSA to be installed first.
+
+# --8<-- [start:pre-built-wheels]
+
+# --8<-- [end:pre-built-wheels]
+
+# --8<-- [start:build-wheel-from-source]
+
+#### Prerequisites
+
+- **MUSA SDK**: Download from [MUSA SDK Download](https://developer.mthreads.com/sdk/download/musa)
+- **torchada**: CUDA→MUSA compatibility layer for PyTorch (`pip install torchada`)
+- **mthreads-ml-py**: MTML Python bindings (`pip install mthreads-ml-py`)
+- **MATE**: MUSA AI Tensor Engine ([GitHub](https://github.com/MooreThreads/mate))
+
+#### Installation of vLLM-MUSA
+
+```bash
+git clone https://github.com/MooreThreads/vllm-musa.git
+cd vllm-musa
+git checkout v0.18.0-dev
+pip install . --no-build-isolation -v
+```
+
+#### Installation of vLLM-Omni
+
+```bash
+git clone https://github.com/vllm-project/vllm-omni.git
+cd vllm-omni
+VLLM_OMNI_TARGET_DEVICE=musa pip install -e . --no-build-isolation
+```
+
+For Gradio demos:
+
+```bash
+pip install -e '.[demo]' --no-build-isolation
+```
+
+#### Environment Variables
+
+```bash
+export MUSA_VISIBLE_DEVICES=0,1
+export VLLM_WORKER_MULTIPROC_METHOD=spawn
+export VLLM_MUSA_CUSTOM_OP_USE_NATIVE=false
+```
+
+# --8<-- [end:build-wheel-from-source]
+
+# --8<-- [start:build-docker]
+
+# --8<-- [end:build-docker]
+
+# --8<-- [start:pre-built-images]
+
+# --8<-- [end:pre-built-images]

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.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.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.18.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.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
+uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
 ```
 
-The pre-built 0.18.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.
+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.18.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.18.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.18.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.17.0rc1
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.17.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,17 +33,6 @@ docker run --rm \
     -p 8000:8000 \
     -it $IMAGE bash
 
-cd /vllm-workspace/vllm
-git fetch origin --tags
-git checkout v0.18.0
-
-# Because vllm-ascend will release v0.18.0rc1 after vllm-omni 0.16.0,
-# we have to pin vllm-ascend at the current commit.
-cd /vllm-workspace/vllm-ascend
-git pull origin main
-git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
-pip install -v -e .
-
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
 git clone -b v0.18.0 https://github.com/vllm-project/vllm-omni.git
@@ -68,15 +57,10 @@ You can also build vLLM-Omni from the latest main branch if you want to use the
 
 ```bash
 # Pin vLLM version to 0.18.0
-cd /vllm-workspace/vllm
-git fetch origin --tags
-git checkout v0.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-workspace/vllm-ascend
-git pull origin main
-git fetch origin --tags
-git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+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

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.18.0 --torch-backend=auto
+uv pip install vllm==0.19.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.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

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

@@ -31,9 +31,13 @@ th {
 | `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` | ✅︎ | ✅︎ | ✅︎ | |
+| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
 | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |

docs/serving/speech_api.md

@@ -118,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
 
@@ -143,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
@@ -161,6 +176,7 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
 | `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:**
 
@@ -172,11 +188,15 @@ 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
@@ -184,7 +204,8 @@ 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 "ref_text=The exact transcript of the audio sample."
+  -F "ref_text=The exact transcript of the audio sample." \
+  -F "speaker_description=warm narrator"
 ```
 
 ## Streaming Text Input (WebSocket)

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/)

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -91,12 +91,19 @@ Models must define the blocks attribute name for layerwise offloading:
 
 ```python
 class WanTransformer3DModel(nn.Module):
-    _layerwise_offload_blocks_attr = "blocks"  # Attribute name containing transformer blocks
+    _layerwise_offload_blocks_attrs = ["blocks"]  # Attribute names containing transformer blocks
 
     def __init__(self):
         self.blocks = nn.ModuleList([...])  # Transformer blocks
 ```
 
+For models with multiple block types:
+
+```python
+class Flux2Transformer2DModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
+```
+
 ### Limitations
 - Cold start latency increases because of
     1) components are loaded to CPU first at the very first during initialization,
@@ -140,4 +147,4 @@ Factory function `get_offload_backend()` selects the appropriate backend based o
 
 **Notes:**
 - Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
-- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attr` pointing to transformer blocks
+- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks

docs/user_guide/diffusion/parallelism/cfg_parallel.md

@@ -0,0 +1,169 @@
+# CFG-Parallel 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
+
+CFG-Parallel accelerates diffusion models by distributing positive and negative classifier-free guidance (CFG) passes across different GPUs, providing ~1.8x speedup when CFG is enabled. It's ideal for image editing tasks that require guidance scales greater than 1.0.
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+---
+
+## Quick Start
+
+### Basic Usage
+
+Simplest working example:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from PIL import Image
+
+omni = Omni(
+    model="Qwen/Qwen-Image-Edit",
+    parallel_config=DiffusionParallelConfig(cfg_parallel_size=2),  # Enable CFG-Parallel
+)
+
+input_image = Image.open("input.png").convert("RGB")
+outputs = omni.generate(
+    {
+        "prompt": "turn this cat to a dog",
+        "negative_prompt": "low quality, blurry",
+        "multi_modal_data": {"image": input_image},
+    },
+    OmniDiffusionSamplingParams(
+        true_cfg_scale=4.0,
+        num_inference_steps=50,
+    ),
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use python script under `examples/offline_inference/image_to_image/image_edit.py`:
+
+```bash
+cd examples/offline_inference/image_to_image/
+python image_edit.py \
+  --model "Qwen/Qwen-Image-Edit" \
+  --image "input.png" \
+  --prompt "turn this cat to a dog" \
+  --negative-prompt "low quality, blurry" \
+  --cfg-scale 4.0 \
+  --output "edited_image.png" \
+  --cfg-parallel-size 2
+```
+
+### Online Serving
+
+Enable CFG-Parallel in online serving:
+
+```bash
+# Default configuration
+vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2
+
+```
+
+---
+
+## Configuration Parameters
+
+In `DiffusionParallelConfig`
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `cfg_parallel_size` | int | 1 | Number of GPUs for CFG parallelism. Set to 2 to enable CFG-Parallel (rank 0 for positive, rank 1 for negative branch) |
+
+
+!!! info
+    Most models support `cfg_parallel_size=2` (positive branch on rank 0, negative branch on rank 1). **Bagel** is an exception: it supports `cfg_parallel_size=3`, which adds a third branch on rank 2 for full three-way CFG parallelism.
+
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- Tasks requiring classifier-free guidance
+- Multi-GPU setups (at least 2 GPUs available)
+- Combining with other parallelism methods (sequence/tensor parallel)
+
+**Not for:**
+
+- Single GPU setups
+- Models that don't support CFG-Parallel (check [supported models](../../diffusion_features.md#supported-models))
+- Workloads without negative prompts or classifier-free guidance
+- Very short inference runs (< 10 steps) where parallelism overhead may outweigh benefits
+
+### Expected Performance
+
+| Configuration | Speedup | Quality | Use Case |
+|--------------|---------|---------|----------|
+| CFG-Parallel (2 GPUs) | 1.5~1.8x | No degradation | Large model, VRAM limited |
+
+---
+
+## Troubleshooting
+
+### Common Issue 1: No Speedup with CFG-Parallel
+
+**Symptoms**: CFG-Parallel enabled but no performance improvement
+
+**Solutions**:
+
+1. **Ensure CFG scale is set correctly:**
+```python
+# Bad: No CFG effect
+sampling_params = OmniDiffusionSamplingParams(num_inference_steps=50)
+
+# Good: CFG-Parallel will work
+sampling_params = OmniDiffusionSamplingParams(
+    num_inference_steps=50,
+    true_cfg_scale=4.0  # Must be > 1.0
+)
+```
+
+2. **Add negative prompt:**
+```python
+outputs = omni.generate(
+    {
+        "prompt": "beautiful landscape",
+        "negative_prompt": "low quality, blurry",  # Required for best results
+        "multi_modal_data": {"image": input_image}
+    },
+    sampling_params
+)
+```
+
+3. **Check model support:**
+   - Verify your model in [supported models](../../diffusion_features.md#supported-models)
+   - Some models don't support CFG-Parallel
+
+---
+
+## Summary
+
+1. ✅ **Enable CFG-Parallel** - Set `cfg_parallel_size=2` in `DiffusionParallelConfig` to get speedup when using CFG
+2. ✅ **Set CFG Scale** - Ensure `true_cfg_scale > 1.0` in `OmniDiffusionSamplingParams` for CFG-Parallel to take effect
+3. ✅ **Check Model Support** - Verify your model supports CFG-Parallel in [supported models](../../diffusion_features.md#supported-models)

docs/user_guide/diffusion/parallelism/expert_parallel.md

@@ -0,0 +1,87 @@
+# Expert Parallelism Guide
+
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Configuration Parameters](#configuration-parameters)
+- [Best Practices](#best-practices)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+Unlike Tensor Parallelism which shards every layer's weights, Expert Parallelism (EP) only shards the MoE expert MLP blocks. This significantly reduces the memory footprint of MoE models (e.g., HunyuanImage3.0) while maintaining constant dense-equivalent compute efficiency.
+
+During the forward pass, a gating mechanism routes tokens to their designated experts, requiring all-to-all communication to dispatch tokens to the correct ranks and combine results.
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+!!! note "EP Size Constraint"
+    The effective EP size equals `tp × sp × cfg × dp`. At least one of TP/SP/CFG/DP must be set when EP is enabled.
+
+---
+
+## Quick Start
+
+### Basic Usage
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="tencent/HunyuanImage-3.0",
+    parallel_config=DiffusionParallelConfig(
+        tensor_parallel_size=8,
+        enable_expert_parallel=True,
+    ),
+)
+
+outputs = omni.generate(
+    "A brown and white dog is running on the grass",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+        width=1024,
+        height=1024,
+    ),
+)
+```
+
+---
+
+## Configuration Parameters
+
+In `DiffusionParallelConfig`:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `enable_expert_parallel` | bool | False | Enable Expert Parallelism for MoE models |
+
+EP size is derived automatically as `tp × sp × cfg × dp` — configure at least one of those to set the EP degree.
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- MoE models (e.g., HunyuanImage3.0) with numbers of experts
+- Memory-constrained multi-GPU setups where only expert blocks need sharding
+
+**Not for:**
+
+- Dense models (no MoE layers) — EP has no effect
+- Single GPU setups
+
+---
+
+## Summary
+
+1. ✅ **Enable EP** - Set `enable_expert_parallel=True` in `DiffusionParallelConfig` for MoE models
+2. ✅ **Set parallelism degree** - At least one of `tensor_parallel_size` / `ulysses_degree` / `cfg_parallel_size` must be > 1 to define the EP size

docs/user_guide/diffusion/parallelism/hsdp.md

@@ -0,0 +1,149 @@
+# HSDP Guide
+
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Example Script](#example-script)
+- [Configuration Parameters](#configuration-parameters)
+- [Best Practices](#best-practices)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
+
+Unlike Tensor Parallelism which splits computation, HSDP uses PyTorch's FSDP2 to shard and redistribute weights at runtime. Each GPU only holds a fraction of the model weights, and weights are gathered on-demand during forward passes.
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+**Operating Modes:**
+
+- **Standalone Mode**: HSDP alone without other parallelism. Must specify `hsdp_shard_size` explicitly.
+- **Combined Mode**: HSDP overlays on top of other parallelism (Ulysses-SP, CFG-Parallel). HSDP dimensions must match world_size.
+
+---
+
+## Quick Start
+
+### Basic Usage
+
+Simplest working example (standalone HSDP, shard across 4 GPUs):
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        use_hsdp=True,
+        hsdp_shard_size=4,  # Shard across 4 GPUs
+    ),
+)
+
+outputs = omni.generate(
+    "A cat playing piano",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+### Combined with Sequence Parallel
+
+```python
+omni = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        ulysses_degree=4,  # Sequence parallel
+        use_hsdp=True,     # HSDP overlays on SP
+    ),
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use Python script under `examples/offline_inference/image_to_video/`:
+
+```bash
+# Standalone HSDP: shard across 4 GPUs
+python examples/offline_inference/image_to_video/image_to_video.py \
+    --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+    --use-hsdp \
+    --hsdp-shard-size 4
+
+# Combined HSDP + Sequence Parallel
+python examples/offline_inference/image_to_video/image_to_video.py \
+    --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+    --ulysses-degree 4 \
+    --use-hsdp
+```
+
+### Online Serving
+
+**Standalone HSDP** (shard model across 4 GPUs):
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 \
+    --use-hsdp --hsdp-shard-size 4
+```
+
+**Combined with Sequence Parallel**:
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 \
+    --use-hsdp --usp 4
+```
+
+---
+
+## Configuration Parameters
+
+In `DiffusionParallelConfig`:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `use_hsdp` | bool | False | Enable HSDP |
+| `hsdp_shard_size` | int | -1 | Number of GPUs to shard weights across. `-1` = auto (requires other parallelism > 1) |
+| `hsdp_replicate_size` | int | 1 | Number of replica groups. Each group holds a full sharded copy |
+
+**Constraints:**
+
+- `hsdp_replicate_size × hsdp_shard_size == world_size`
+- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- Very large models (e.g., Wan2.2 14B)
+- Multi-GPU setups where memory reduction is the primary goal
+- Combining with Sequence Parallelism for large video models
+
+**Not for:**
+
+- Models that fit comfortably in single-GPU memory
+- Use cases requiring Tensor Parallelism (HSDP and TP are mutually exclusive)
+
+### Adding HSDP Support to New Models
+
+For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../../design/feature/hsdp.md).
+
+---
+
+## Summary
+
+1. ✅ **Enable HSDP** - Set `use_hsdp=True` and `hsdp_shard_size` to reduce per-GPU memory for large models
+2. ✅ **Combine with SP** - Use together with `ulysses_degree` for video models requiring both memory reduction and sequence parallelism
+3. ⚠️ **Incompatible with TP** - `tensor_parallel_size` must be 1 when HSDP is enabled

docs/user_guide/diffusion/parallelism/overview.md

@@ -0,0 +1,16 @@
+# Parallelism Acceleration Guide
+
+This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion model inference and reducing per-device memory requirements.
+
+## Supported Methods
+
+| Method | Description |
+|--------|-------------|
+| **[Tensor Parallelism](tensor_parallel.md)** | Shards DiT weights across GPUs to reduce per-GPU memory |
+| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
+| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation |
+| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
+| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs |
+| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) |
+
+See [Supported Models](../../diffusion_features.md#supported-models) for per-model compatibility.

docs/user_guide/diffusion/parallelism/sequence_parallel.md

@@ -0,0 +1,233 @@
+# Sequence Parallelism Guide
+
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Example Script](#example-script)
+- [Best Practices](#best-practices)
+- [Troubleshooting](#troubleshooting)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+Sequence parallelism splits the input along the sequence dimension across multiple GPUs, allowing each device to process only a portion of the sequence. vLLM-Omni provides 1.5x-3.6x speedup for large images and videos using DeepSpeed Ulysses, Ring-Attention, or hybrid approaches. Use sequence parallelism when generating high-resolution images/videos that don't fit on a single GPU or require faster inference.
+
+See supported models list in [Diffusion Features - Supported Models](../../diffusion_features.md#supported-models).
+
+**Supported Methods:**
+
+- **DeepSpeed Ulysses Sequence Parallel (Ulysses-SP)** ([paper](https://arxiv.org/pdf/2309.14509)): Uses all-to-all communication for subset of attention heads per device
+- **Ring-Attention** ([paper](https://arxiv.org/abs/2310.01889)): Uses ring-based P2P communication with sharded sequence dimension throughout
+- **Hybrid Ulysses + Ring**: Combines both for larger scale parallelism (`ulysses_degree × ring_degree`)
+
+---
+
+## Quick Start
+
+### Basic Usage - Ulysses-SP
+
+Simplest working example with Ulysses Sequence Parallel:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    parallel_config=DiffusionParallelConfig(ulysses_degree=2)  # Enable Ulysses-SP
+)
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=1024, height=1024),
+)
+```
+
+!!! note "Experimental UAA mode"
+    `ulysses_mode="advanced_uaa"` is an experimental extension to Ulysses-SP. It lets Ulysses attention handle arbitrary sequence lengths and arbitrary attention head counts without relying on `attention_mask`-based token padding.
+
+    In hybrid Ulysses + Ring mode, Ring still requires every rank in the same ring group to observe the same post-Ulysses sequence length. If that condition is not met, vLLM-Omni raises a validation error instead of entering the ring kernel with inconsistent shapes.
+
+To enable the experimental UAA mode, use a model/configuration that requires it. For example, `Tongyi-MAI/Z-Image-Turbo` has 30 attention heads, so `ulysses_degree=4` requires UAA because 30 is not divisible by 4:
+
+```python
+omni = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    parallel_config=DiffusionParallelConfig(
+        ulysses_degree=4,
+        ulysses_mode="advanced_uaa",
+    ),
+)
+```
+
+### Alternative Methods
+
+**Ring-Attention** (better for very long sequences):
+
+```python
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    parallel_config=DiffusionParallelConfig(ring_degree=2)  # Enable Ring-Attention
+)
+```
+
+**Hybrid Ulysses + Ring** (for larger scale):
+
+```python
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    parallel_config=DiffusionParallelConfig(ulysses_degree=2, ring_degree=2)  # 4 GPUs total
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use Python script under `examples/offline_inference/text_to_image/text_to_image.py`:
+
+**Ulysses-SP:**
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model Qwen/Qwen-Image \
+    --prompt "A cat sitting on a windowsill" \
+    --ulysses-degree 2 \
+    --width 1024 --height 1024
+```
+
+**Ring-Attention:**
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model Qwen/Qwen-Image \
+    --prompt "A cat sitting on a windowsill" \
+    --ring-degree 2 \
+    --width 1024 --height 1024
+```
+
+**Hybrid Ulysses + Ring:**
+
+```bash
+# Hybrid: 2 Ulysses × 2 Ring = 4 GPUs total
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model Qwen/Qwen-Image \
+    --prompt "A cat sitting on a windowsill" \
+    --ulysses-degree 2 --ring-degree 2 \
+    --width 1024 --height 1024
+```
+
+### Online Serving
+
+**Ulysses-SP:**
+
+```bash
+# Text-to-image (requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
+```
+
+**Ulysses-SP with UAA mode** (for models with non-divisible head counts):
+
+```bash
+vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 --usp 4 --ulysses-mode advanced_uaa
+```
+
+**Ring-Attention:**
+
+```bash
+# Text-to-image (requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2
+```
+
+**Hybrid Ulysses + Ring:**
+
+```bash
+# Text-to-image (requires >= 4 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
+```
+
+---
+
+## Configuration Parameters
+
+In `DiffusionParallelConfig`:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `ulysses_degree` | int | 1 | Number of GPUs for Ulysses-SP. Uses all-to-all communication. |
+| `ring_degree` | int | 1 | Number of GPUs for Ring-Attention. Uses P2P ring communication. |
+| `ulysses_mode` | str | `"default"` | Ulysses attention mode. Set to `"advanced_uaa"` to handle arbitrary sequence lengths and head counts without padding. |
+
+**Notes:**
+- Total sequence parallel size equals to `ulysses_degree × ring_degree`
+- Degrees must evenly divide the sequence length for optimal performance (or use `ulysses_mode="advanced_uaa"` for Ulysses-SP)
+
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- Large images (1024x1024 or higher) or videos
+- Fast inter-GPU communication, larger bandwidth (e.g., NVLink)
+
+**Not for:**
+
+- Small images (<1024px) - overhead exceeds benefit, use single GPU with cache instead
+
+
+---
+
+## Troubleshooting
+
+### Common Issue 1: Performance Not Scaling
+
+**Symptoms**: Adding GPUs doesn't improve speed proportionally, or higher parallelism degree is slower
+
+**Diagnosis:**
+```bash
+# Check GPU topology
+nvidia-smi topo -m
+
+```
+
+**Solutions:**
+
+1. Check inter-GPU communication - NVLink is better than PCIe
+2. Reduce parallelism degree if over-parallelized:
+```python
+# If 4 GPUs is slower than 2
+parallel_config=DiffusionParallelConfig(ulysses_degree=2)
+```
+3. Try to switch between Ring-Attention and Ulysses-SP
+
+- Ring-Attention has advantages, like communication-computation overlap, but the block-wise loop overhead is relatively higher, especially for short sequences
+- Ulysses-SP: can benefit from larger bandwidth (such as NVLink), with two major constraints, the sequence length should be divisible by usp size, and the number of heads should be divisible by usp size (or use `ulysses_mode="advanced_uaa"`)
+
+
+### Common Issue 2: Out of Memory (OOM)
+
+**Symptoms**: CUDA OOM errors or process crashes with memory errors
+
+**Solutions:**
+
+1. Increase parallelism degree to split sequence more:
+```python
+parallel_config=DiffusionParallelConfig(ulysses_degree=4)  # From 2
+```
+2. Combine with other parallelism method, e.g., tensor parallel, and memory optimization methods, e.g., cpu offloading.
+
+
+## Summary
+
+1. ✅ **Enable Sequence Parallelism** - Set `ulysses_degree` or `ring_degree` for long sequence generation
+2. ✅ **UAA mode** - Use `ulysses_mode="advanced_uaa"` when head count is not divisible by `ulysses_degree`
+3. ✅ **Troubleshooting** - Check GPU topology with `nvidia-smi topo -m`, reduce degree if performance doesn't scale

docs/user_guide/diffusion/parallelism/tensor_parallel.md

@@ -0,0 +1,151 @@
+# Tensor Parallelism 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
+
+Tensor Parallelism (TP) shards some model weights across multiple GPUs, usually the Linear layers. This enables running large models that don't fit on a single GPU. It's essential for memory-constrained setups or very large models.
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+!!! note "TP Limitations for Diffusion Models"
+    We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
+
+    - Good news: The text_encoder typically has minimal impact on overall inference performance.
+    - Bad news: When TP is enabled, every TP process retains a full copy of the text_encoder weights, leading to significant GPU memory waste.
+
+    We are actively refactoring this design to address this. For details and progress, please refer to [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
+
+---
+
+## Quick Start
+
+
+### Basic Usage
+
+Simplest working example:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),  # Enable TP
+)
+
+outputs = omni.generate(
+    "a cat reading a book",
+    OmniDiffusionSamplingParams(num_inference_steps=9),
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use Python script under `examples/offline_inference`, and enable TP:
+
+```bash
+# Text-to-Image with Qwen-Image
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model Qwen/Qwen-Image \
+    --tensor-parallel-size 2
+
+# Image Editing with Qwen-Image-Edit
+python examples/offline_inference/image_to_image/image_edit.py \
+    --model Qwen/Qwen-Image-Edit \
+    --image input.png \
+    --prompt "Edit description" \
+    --tensor-parallel-size 2
+```
+
+### Online Serving
+
+You can enable tensor parallelism in online serving via `--tensor-parallel-size`:
+
+```bash
+# Text-to-Image with Qwen-Image on 2 GPUs
+vllm serve Qwen/Qwen-Image --omni --port 8091 \
+    --tensor-parallel-size 2
+
+# Text-to-Image with Z-Image (TP=2 only)
+vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 \
+    --tensor-parallel-size 2
+```
+
+---
+
+## Configuration Parameters
+
+In `DiffusionParallelConfig`:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `tensor_parallel_size` | int | 1 | Number of GPUs to shard model weights across. Must divide number of heads. |
+
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- Large models that don't fit on a single GPU, especially for models with large DiT blocks (transformer layers)
+- Memory-constrained environments
+
+**Not for:**
+
+- When maximum throughput is needed and memory is sufficient
+- Models with incompatible dimensions (e.g., Z-Image `num_heads=30`, which now supports `tensor_parallel_size=2`)
+
+
+## Troubleshooting
+
+### Common Issue 1: Out of Memory (OOM)
+
+**Symptoms**: CUDA OOM errors during model loading or inference, process crashes with memory errors
+
+**Solution**:
+```python
+# Step 1: Enable TP with smallest degree
+parallel_config=DiffusionParallelConfig(tensor_parallel_size=2)
+
+# Step 2: If still OOM, increase TP degree
+parallel_config=DiffusionParallelConfig(tensor_parallel_size=4)
+
+```
+
+### Common Issue 2: Divisibility Error
+
+**Symptoms**: Error like "Model dimension X not divisible by tensor_parallel_size Y"
+
+**Solutions**:
+1. Check model-specific constraints (e.g., Z-Image only supports TP=2)
+2. Use a smaller TP size that divides model dimensions
+3. Consult [Supported Models](../../diffusion_features.md#supported-models) for compatible TP sizes
+
+
+---
+
+## Summary
+
+1. ✅ **Enable TP** - Set `--tensor-parallel-size` to reduce per-GPU memory
+2. ✅ **Increase TP size** - Only increase if OOM persists
+3. ⚠️ **Text encoder not sharded** - Known limitation

docs/user_guide/diffusion/parallelism/vae_patch_parallel.md

@@ -0,0 +1,200 @@
+# VAE Patch Parallelism 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
+
+VAE Patch Parallelism distributes the VAE (Variational AutoEncoder) decode/encode computation across multiple GPUs by splitting the latent space into spatial tiles or patches. Each GPU processes a subset of tiles in parallel, significantly reducing peak memory consumption during the VAE decode stage while maintaining output quality.
+
+This is particularly useful for:
+- **High-resolution image generation** where VAE decode can become a memory bottleneck
+- **Memory-constrained environments** where the VAE decode activation peak exceeds available VRAM
+- **Multi-GPU setups** where you want to leverage distributed resources for the VAE stage
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+
+VAE Patch Parallelism uses two strategies based on image size:
+
+| Strategy | Use Case | How It Works | Overlap Handling | Output Quality |
+|----------|----------|--------------|------------------|----------------|
+| **Tiled Decode** | Large images (triggers VAE tiling) | Distributes existing VAE tiling computation across ranks. Each rank decodes a subset of overlapping tiles. | Uses VAE's native `blend_v` and `blend_h` functions to seamlessly merge overlapping regions | Bit-identical (same logic as single-GPU tiling) |
+| **Patch Decode** | Small images (no VAE tiling) | Splits latent into spatial patches with halos. Each rank decodes one patch with boundary context. | Halo regions provide edge context; core regions are directly stitched without blending | Near-identical (diff < 0.5%, visually imperceptible) |
+
+
+VAE Patch Parallelism **reuses the DiT process group** (`dit_group`) and does not initialize a separate ProcessGroup. This means:
+
+- **Shared ranks**: VAE patch parallelism uses the same GPU ranks as DiT parallelism (Tensor Parallel, Sequence Parallel, etc.)
+- **Combined usage**: VAE patch parallelism is typically used together with other parallelism methods
+- **Configuration alignment**: The `vae_patch_parallel_size` should be no greater than the size of your DiT process group
+
+---
+
+## Quick Start
+
+### Basic Usage
+
+Simplest working example:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+# TP=2 for DiT, VAE patch parallel also uses these 2 GPUs
+omni = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    parallel_config=DiffusionParallelConfig(
+        tensor_parallel_size=2,          # Enable tensor parallelism for DiT
+        vae_patch_parallel_size=2,       # Enable VAE patch parallelism
+    ),
+    vae_use_tiling=True,  # Required for VAE patch parallelism
+)
+
+outputs = omni.generate(
+    "a futuristic city at sunset, high resolution, 8k",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=9,
+        height=1152,  # High resolution benefits from VAE patch parallel
+        width=1152,
+    ),
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use Python script under `examples/offline_inference/text_to_image/`:
+
+```bash
+# Text-to-Image with Z-Image
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model Tongyi-MAI/Z-Image-Turbo \
+    --prompt "a futuristic city at sunset" \
+    --height 1152 \
+    --width 1152 \
+    --tensor-parallel-size 2 \
+    --vae-patch-parallel-size 2 \
+    --vae-use-tiling
+```
+
+### Online Serving
+
+You can enable VAE patch parallelism in online serving via `--vae-patch-parallel-size`:
+
+```bash
+# Text-to-Image with Z-Image, TP=2 + VAE patch parallel=2
+vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 \
+    --tensor-parallel-size 2 \
+    --vae-patch-parallel-size 2 \
+    --vae-use-tiling
+```
+
+---
+
+## Configuration Parameters
+
+In `DiffusionParallelConfig`:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. |
+
+Additional requirements:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `vae_use_tiling` | bool | False | Must be set to `True` when using VAE patch parallelism. |
+
+!!! note "Automatic VAE Tiling"
+    When `vae_patch_parallel_size > 1` and the model has a distributed VAE (`DistributedVaeMixin`), the system automatically sets `vae_use_tiling=True` if not already enabled.
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- High-resolution image generation and long video generation
+- Memory-constrained setups where VAE decode causes OOM
+- Multi-GPU environments
+
+**Not for:**
+
+- Low-resolution images/videos where VAE decode is not a bottleneck
+- Single GPU setups should use vae tiling decode, but not parallel vae tiling decode
+- Models that do not support vae patch parallel
+
+---
+
+## Troubleshooting
+
+### Common Issue 1: Model Not Support VAE Patch Parallel
+
+**Symptoms**:
+```
+WARNING: vae_patch_parallel_size=2 is set but VAE patch parallelism is NOT enabled for xxxPipeline; ignoring.
+```
+
+**Root Cause**: VAE Patch Parallelism requires the model's VAE to implement `DistributedVaeMixin`. At startup, `vllm_omni/diffusion/registry.py` checks whether the instantiated pipeline has a `.vae` attribute that is an instance of `DistributedVaeMixin`. If it does not, the setting is silently ignored:
+
+```python
+vae_pp_size = od_config.parallel_config.vae_patch_parallel_size
+is_distributed_vae = hasattr(model, "vae") and isinstance(model.vae, DistributedVaeMixin)
+if vae_pp_size > 1 and not is_distributed_vae:
+    logger.warning(
+        "vae_patch_parallel_size=%d is set but VAE patch parallelism is NOT enabled for %s; ignoring.",
+        vae_pp_size,
+        od_config.model_class_name,
+    )
+```
+
+**Solutions**:
+
+1. **Use a supported model** (recommended): check [Supported Models](../../diffusion_features.md#supported-models) for the VAE-Patch-Parallel column.
+
+2. To add support for a new model, implement `DistributedVaeMixin` on its VAE class (contributions are welcome).
+
+
+### Common Issue 2: `vae_patch_parallel_size` Exceeds DiT Process Group Size
+
+**Symptoms**: Shows warning message, and vae patch parallel size is resized to DiT process group size
+
+**Root Cause**: VAE Patch Parallelism reuses the DiT process group.
+
+**Recommendation**: Always set `vae_patch_parallel_size` to be no greater than your DiT process group size.
+
+Note that the size of DiT process group size equals to:
+```text
+dit_parallel_size = data_parallel_size
+                  × cfg_parallel_size
+                  × sequence_parallel_size
+                  × pipeline_parallel_size
+                  × tensor_parallel_size
+
+```
+_sequence_parallel_size = ulysses_degree × ring_degree_
+
+---
+
+## Summary
+
+1. ✅ **Enable VAE Patch Parallelism** - Set `vae_patch_parallel_size`， `vae_use_tiling=True` in `DiffusionParallelConfig` to reduce VAE decode peak memory
+2. ✅ **Use Long Sequence** - VAE patch parallelism benefits are most apparent at long sequence decoding
+3. ✅ **Combine with other parallelism methods** - Suggest to use together with Tensor Parallel or CFG-Parallel for maximum memory savings

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -1,477 +0,0 @@
-# Parallelism Acceleration Guide
-
-This guide includes how to use parallelism methods in vLLM-Omni to speed up diffusion model inference as well as reduce the memory requirement on each device.
-
-## Overview
-
-The following parallelism methods are currently supported in vLLM-Omni:
-
-1. DeepSpeed Ulysses Sequence Parallel (DeepSpeed Ulysses-SP) ([arxiv paper](https://arxiv.org/pdf/2309.14509)): Ulysses-SP splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
-
-2. [Ring-Attention](#ring-attention) - splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results, keeping the sequence dimension sharded
-
-3. Classifier-Free-Guidance Parallel (CFG-Parallel): CFG-Parallel runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step.
-
-4. [Tensor Parallelism](#tensor-parallelism): Tensor parallelism shards model weights across devices. This can reduce per-GPU memory usage. Note that for diffusion models we currently shard the majority of layers within the DiT.
-
-5. [VAE Patch Parallelism](#vae-patch-parallelism): VAE patch parallelism shards VAE decode spatially across ranks. This can reduce the peak memory of VAE decode and (depending on resolution and communication overhead) speed up VAE decode.
-
-6. [HSDP](#hsdp): Hybrid Sharded Data Parallel shards model weights across GPUs using PyTorch FSDP2. This reduces per-GPU memory usage, enabling inference of large models on GPUs with limited memory.
-
-7. [Expert Parallel](#expert-parallelism): Expert Parallelism shards the Experts of a Mixture-of-Experts (MoE) layer across multiple devices. During the forward, a gating mechanism routes tokens to their designated experts, necessitating cross-cards communication(all-to-all) to dispatch tokens to the correct ranks and combine the results. This parallelism allows for massive scaling of model parameters without a proportional increase in the computational load per device.
-
-The following table shows which models are currently supported by parallelism method:
-
-### ImageGen
-
-| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel | HSDP |
-|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|:----:|
-| **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |      N/A        | ❌   |
-| **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |      N/A        | ❌   |
-| **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |         ❌          |      N/A        | ❌   |
-| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |         ✅          |      N/A        | ❌   |
-| **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |      N/A        | ❌   |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |      N/A        | ❌   |
-| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |      N/A        | ❌   |
-| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |         ✅          |      N/A        | ❌   |
-| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ✅        |         ✅          |      N/A        | ❌   |
-| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |      N/A        | ✅   |
-| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ✅       |        ✅        |         ❌          |      N/A        | ✅   |
-| **FLUX.2-dev**           | `black-forest-labs/FLUX.2-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |      N/A        | ✅   |
-| **HunyuanImage3.0**      | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |      ✅        | ❌   |
-| **Bagel**                  | `ByteDance-Seed/BAGEL-7B-MoT` |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |      N/A        | ❌   |
-| **DreamID-Omni**           | `XuGuo699/DreamID-Omni`       |     ❌      |    ❌    |      ✅       |        ❌        |         ❌          |      N/A        | ❌   |
-| **FLUX.1-Kontext-dev**   | `black-forest-labs/FLUX.1-Kontext-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |      N/A        | ✅   |
-| **OmniGen2**               | `OmniGen2/OmniGen2`           |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |      N/A        | ❌   |
-
-!!! note "TP Limitations for Diffusion Models"
-    We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
-
-    - Good news: The text_encoder typically has minimal impact on overall inference performance.
-    - Bad news: When TP is enabled, every TP process retains a full copy of the text_encoder weights, leading to significant GPU memory waste.
-
-    We are actively refactoring this design to address this. For details and progress, please refer to [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
-
-
-!!! note "Why Z-Image is TP=2 only"
-    Z-Image Turbo is currently limited to `tensor_parallel_size` of **1 or 2** due to model shape divisibility constraints.
-    For example, the model has `n_heads=30` and a final projection out dimension of `64`, so valid TP sizes must divide both 30 and 64; the only common divisors are **1 and 2**.
-
-### VideoGen
-
-| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | VAE-Patch-Parallel |
-|-------|------------------|:----------:|:--------------:|:---------------:|:----:| :----:|
-| **Wan2.1** | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Wan2.1** | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **LTX-2** | `Lightricks/LTX-2` | ✅ | ✅ | ✅ | ❌ | ❌ |
-
-### Tensor Parallelism
-
-Tensor parallelism splits model parameters across GPUs. In vLLM-Omni, tensor parallelism is configured via `DiffusionParallelConfig.tensor_parallel_size`.
-
-#### Offline Inference
-
-```python
-from vllm_omni import Omni
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="Tongyi-MAI/Z-Image-Turbo",
-    parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
-)
-
-outputs = omni.generate(
-    "a cat reading a book",
-    OmniDiffusionSamplingParams(
-        num_inference_steps=9,
-        width=512,
-        height=512,
-    ),
-)
-```
-
-### VAE Patch Parallelism
-
-VAE patch parallelism distributes the VAE decode workload across multiple ranks by splitting the latent spatially. It is configured via `DiffusionParallelConfig.vae_patch_parallel_size` and can be combined with other parallelism methods (e.g., TP).
-
-!!! note "Enablement and feature gate"
-    - VAE patch parallelism is currently **enabled only for validated pipelines** (check [ImageGen](#imagegen) and [VideoGen](#videogen) for more information).
-    - If `vae_patch_parallel_size > 1` is set for a validated pipeline, vLLM-Omni will automatically enable `vae_use_tiling` as a safety gate. (We use `vae_use_tiling` because it indicates the VAE supports diffusers tiling parameters like `tile_latent_min_size` and `tile_overlap_factor`.)
-
-#### Offline Inference
-
-```python
-from vllm_omni import Omni
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="Tongyi-MAI/Z-Image-Turbo",
-    parallel_config=DiffusionParallelConfig(
-        tensor_parallel_size=2,
-        vae_patch_parallel_size=2,
-    ),
-    vae_use_tiling=True,
-)
-
-outputs = omni.generate(
-    prompt="a cat reading a book",
-    num_inference_steps=9,
-    width=1024,
-    height=1024,
-)
-```
-
-#### How it works (method selection)
-
-VAE patch parallelism automatically selects between two internal decode methods based on whether diffusers tiling would kick in:
-
-- `_distributed_tiled_decode`: Used when the latent spatial size exceeds `vae.tile_latent_min_size` (i.e., diffusers tiled decode). Each rank decodes a subset of tiles; rank0 gathers and runs the same overlap+blend+stitch logic as diffusers. This matches the single-rank diffusers tiled output.
-
-- `_distributed_patch_decode`: Used when diffusers tiling would not kick in. Each rank decodes a grid patch expanded with a latent-space halo; then rank0 gathers the cropped core patches and stitches them into the full image. This path has no blending and can introduce small numerical differences compared to the non-parallel decode.
-
-### Sequence Parallelism
-
-#### Ulysses-SP
-
-!!! note "Experimental UAA mode"
-    `ulysses_mode="advanced_uaa"` is an experimental extension to Ulysses-SP. It lets Ulysses attention handle arbitrary sequence lengths and arbitrary attention head counts without relying on `attention_mask`-based token padding.
-
-    In hybrid Ulysses + Ring mode, Ring still requires every rank in the same ring group to observe the same post-Ulysses sequence length. If that condition is not met, vLLM-Omni raises a validation error instead of entering the ring kernel with inconsistent shapes.
-
-##### Offline Inference
-
-An example of offline inference script using [Ulysses-SP](https://arxiv.org/pdf/2309.14509) is shown below:
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-ulysses_degree = 2
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    parallel_config=DiffusionParallelConfig(ulysses_degree=2)
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
-)
-```
-
-See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
-
-To enable the experimental UAA mode explicitly, use a model/configuration that actually requires it. For example, `Tongyi-MAI/Z-Image-Turbo` has 30 attention heads, so `ulysses_degree=4` requires UAA because 30 is not divisible by 4:
-
-```python
-omni = Omni(
-    model="Tongyi-MAI/Z-Image-Turbo",
-    parallel_config=DiffusionParallelConfig(
-        ulysses_degree=4,
-        ulysses_mode="advanced_uaa",
-    ),
-)
-```
-
-##### Online Serving
-
-You can enable Ulysses-SP in online serving for diffusion models via `--usp`:
-
-```bash
-# Text-to-image (requires >= 2 GPUs)
-vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
-
-# Experimental UAA mode for a model with 30 attention heads
-vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 --usp 4 --ulysses-mode advanced_uaa
-```
-
-##### Benchmarks
-!!! note "Benchmark Disclaimer"
-    These benchmarks are provided for **general reference only**. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:
-
-    - Specific model and use case
-    - Hardware configuration
-    - Careful parameter tuning
-    - Different inference settings (e.g., number of steps, image resolution)
-
-
-To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** model generating images (**2048x2048** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA H800 GPUs. `sdpa` is the attention backends.
-
-| Configuration | Ulysses degree |Generation Time | Speedup |
-|---------------|----------------|---------|---------|
-| **Baseline (diffusers)** | - | 112.5s | 1.0x |
-| Ulysses-SP  |  2  |  65.2s | 1.73x |
-| Ulysses-SP  |  4  | 39.6s | 2.84x |
-| Ulysses-SP  |  8  | 30.8s | 3.65x |
-
-#### Ring-Attention
-
-Ring-Attention ([arxiv paper](https://arxiv.org/abs/2310.01889)) splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results. Unlike Ulysses-SP which uses all-to-all communication, Ring-Attention keeps the sequence dimension sharded throughout the computation and circulates Key/Value blocks through a ring topology.
-
-##### Offline Inference
-
-An example of offline inference script using Ring-Attention is shown below:
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-ring_degree = 2
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    parallel_config=DiffusionParallelConfig(ring_degree=2)
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
-)
-```
-
-See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
-
-
-##### Online Serving
-
-You can enable Ring-Attention in online serving for diffusion models via `--ring`:
-
-```bash
-# Text-to-image (requires >= 2 GPUs)
-vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2
-```
-
-##### Benchmarks
-!!! note "Benchmark Disclaimer"
-    These benchmarks are provided for **general reference only**. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:
-
-    - Specific model and use case
-    - Hardware configuration
-    - Careful parameter tuning
-    - Different inference settings (e.g., number of steps, image resolution)
-
-
-To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** model generating images (**1024x1024** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA A100 GPUs. `flash_attn` is the attention backends.
-
-| Configuration | Ring degree |Generation Time | Speedup |
-|---------------|----------------|---------|---------|
-| **Baseline (diffusers)** | - | 45.2s | 1.0x |
-| Ring-Attention  |  2  |  29.9s | 1.51x |
-| Ring-Attention  |  4  | 23.3s | 1.94x |
-
-
-#### Hybrid Ulysses + Ring
-
-You can combine both Ulysses-SP and Ring-Attention for larger scale parallelism. The total sequence parallel size equals `ulysses_degree × ring_degree`.
-
-!!! note "Experimental UAA in hybrid mode"
-    `ulysses_mode="advanced_uaa"` can also be used with hybrid Ulysses + Ring, but this does not remove Ring's shape requirement. Every rank in the same ring group must still have the same post-Ulysses sequence length.
-
-##### Offline Inference
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-# Hybrid: 2 Ulysses × 2 Ring = 4 GPUs total
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    parallel_config=DiffusionParallelConfig(
-        ulysses_degree=2,
-        ring_degree=2,
-    )
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
-)
-```
-
-##### Online Serving
-
-```bash
-# Text-to-image (requires >= 4 GPUs)
-vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
-```
-
-##### Benchmarks
-!!! note "Benchmark Disclaimer"
-    These benchmarks are provided for **general reference only**. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:
-
-    - Specific model and use case
-    - Hardware configuration
-    - Careful parameter tuning
-    - Different inference settings (e.g., number of steps, image resolution)
-
-
-To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** model generating images (**1024x1024** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA A100 GPUs. `flash_attn` is the attention backends.
-
-| Configuration | Ulysses degree | Ring degree | Generation Time | Speedup |
-|---------------|----------------|-------------|-----------------|---------|
-| **Baseline (diffusers)** | - | - | 45.2s | 1.0x |
-| Hybrid Ulysses + Ring  |  2  |  2  |  24.3s | 1.87x |
-
-
-### CFG-Parallel
-
-#### Offline Inference
-
-CFG-Parallel is enabled through `DiffusionParallelConfig(cfg_parallel_size=2)`, which runs one rank for the positive branch and one rank for the negative branch.
-
-An example of offline inference using CFG-Parallel (image-to-image) is shown below:
-
-```python
-from vllm_omni import Omni
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-image_path = "path_to_image.png"
-omni = Omni(
-    model="Qwen/Qwen-Image-Edit",
-    parallel_config=DiffusionParallelConfig(cfg_parallel_size=2),
-)
-input_image = Image.open(image_path).convert("RGB")
-
-outputs = omni.generate(
-    {
-        "prompt": "turn this cat to a dog",
-        "negative_prompt": "low quality, blurry",
-        "multi_modal_data": {"image": input_image},
-    },
-    OmniDiffusionSamplingParams(
-        true_cfg_scale=4.0,
-        num_inference_steps=50,
-    ),
-)
-```
-
-Notes:
-
-- CFG-Parallel is only effective when a `negative_prompt` is provided AND a guidance scale (or `cfg_scale`) is greater than 1.
-
-See `examples/offline_inference/image_to_image/image_edit.py` for a complete working example.
-```bash
-cd examples/offline_inference/image_to_image/
-python image_edit.py \
-  --model "Qwen/Qwen-Image-Edit" \
-  --image "qwen_image_output.png" \
-  --prompt "turn this cat to a dog" \
-  --negative-prompt "low quality, blurry" \
-  --cfg-scale 4.0 \
-  --output "edited_image.png" \
-  --cfg-parallel-size 2
-```
-
-#### Online Serving
-
-You can enable CFG-Parallel in online serving for diffusion models via `--cfg-parallel-size`:
-
-```bash
-vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2
-```
-
-### HSDP
-
-HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
-
-Unlike Tensor Parallelism which splits computation, HSDP uses PyTorch's FSDP2 to shard and redistribute weights at runtime. Each GPU only holds a fraction of the model weights, and weights are gathered on-demand during forward passes.
-
-#### Configuration
-
-HSDP is configured via `DiffusionParallelConfig`:
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `use_hsdp` | bool | False | Enable HSDP |
-| `hsdp_shard_size` | int | -1 | Number of GPUs to shard weights across. -1 = auto (requires other parallelism > 1) |
-| `hsdp_replicate_size` | int | 1 | Number of replica groups. Each group holds a full sharded copy |
-
-**Constraints:**
-
-- `hsdp_replicate_size × hsdp_shard_size == world_size`
-- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
-
-#### Operating Modes
-
-HSDP can work in two modes:
-
-- **Standalone Mode**: HSDP alone without other parallelism. Must specify `hsdp_shard_size` explicitly.
-- **Combined Mode**: HSDP overlays on top of other parallelism (Ulysses Sequence Parallel, CFG Parallel). HSDP dimensions must match world_size.
-
-#### Offline Inference
-
-**Standalone HSDP** (shard across 4 GPUs, no other parallelism):
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-    parallel_config=DiffusionParallelConfig(
-        use_hsdp=True,
-        hsdp_shard_size=4,  # Shard across 4 GPUs
-    ),
-)
-
-outputs = omni.generate(
-    "A cat playing piano",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-**Combined HSDP + Sequence Parallel**:
-
-```python
-omni = Omni(
-    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-    parallel_config=DiffusionParallelConfig(
-        ulysses_degree=4,  # Sequence parallel
-        use_hsdp=True,     # HSDP overlays on SP
-    ),
-)
-```
-
-#### Online Serving
-
-**Standalone HSDP** (shard model across 4 GPUs):
-
-```bash
-vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --hsdp-shard-size 4
-```
-
-**Combined with Sequence Parallel**:
-
-```bash
-vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --usp 4
-```
-
-#### Adding HSDP Support to New Models
-
-For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../design/feature/hsdp.md).
-
-### Expert Parallelism
-
-Unlike Tensor Parallelism which shards every layer's weights, EP only shards the MoE expert MLP blocks. This significantly reduces the memory footprint of MoE models (e.g., HunyuanImage3.0) while maintaining constant dense-equivalent compute efficiency. Expert Parallelism is enabled via `DiffusionParallelConfig.enable_expert_parallel`. And `self.ep = tp * sp * cfg * dp` for now, so at least one of TP/SP/CFG/DP should set when EP enabled.
-
-#### Offline Inference
-
-```python
-from vllm_omni import Omni
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="tencent/HunyuanImage-3.0",
-    parallel_config=DiffusionParallelConfig(tensor_parallel_size=8, enable_expert_parallel=True),
-)
-
-outputs = omni.generate(
-    "A brown and white dog is running on the grass",
-    OmniDiffusionSamplingParams(
-        num_inference_steps=50,
-        width=1024,
-        height=1024,
-    ),
-)
-```

docs/user_guide/diffusion/quantization/autoround.md

@@ -0,0 +1,91 @@
+# AutoRound Quantization
+
+## Overview
+
+[AutoRound](https://github.com/intel/auto-round) is an advanced quantization toolkit designed for Large Language Models (LLMs), Vision-Language Models (VLMs), and diffusion models. It achieves high accuracy at ultra-low bit widths (2–4 bits) with minimal tuning by leveraging sign-gradient descent, while providing broad hardware compatibility with multi-datatype support.
+
+The quantization config is auto-detected from the checkpoint's `config.json` (`quantization_config.quant_method = "auto-round"`). No extra CLI flags are needed.
+
+### Supported Schemes
+
+| Scheme | Bits | Status |
+|--------|------|--------|
+| W4A16 | 4 | ✅ Supported |
+| W8A16 | 8 | Planned |
+
+W4A16 is the first supported scheme. Additional schemes will be added in future releases.
+
+## Configuration
+
+1. **Python API**: point `model` at a pre-quantized checkpoint. The quantization is detected automatically.
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(model="vllm-project-org/FLUX.1-dev-AutoRound-w4a16")
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=28),
+)
+outputs[0].save_images("output.png")
+```
+
+2. **CLI**: pass the quantized model path directly.
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model vllm-project-org/FLUX.1-dev-AutoRound-w4a16 \
+  --prompt "A cat sitting on a windowsill" \
+  --num-inference-steps 28 \
+  --output outputs/flux_w4a16.png
+```
+
+No `--quantization` flag is needed — the quantization method is read from the checkpoint.
+
+## How It Works
+
+The checkpoint's `config.json` contains:
+
+```json
+{
+  "quantization_config": {
+    "quant_method": "auto-round",
+    "bits": 4,
+    "group_size": 128,
+    "sym": true,
+    "packing_format": "auto_round:auto_gptq",
+    "block_name_to_quantize": "transformer_blocks,single_transformer_blocks"
+  }
+}
+```
+
+At load time:
+
+1. `TransformerConfig.from_dict()` parses the `quantization_config` section and builds a vLLM `INCConfig` via `build_quant_config("auto-round", ...)`.
+2. `OmniDiffusionConfig.set_tf_model_config()` propagates the detected config to the engine.
+3. The appropriate compute kernel (e.g. GPTQ-Marlin for W4A16) is selected automatically based on the checkpoint's bit-width and packing format.
+
+## Supported Models
+
+| Model | HF Checkpoint | Scheme | Group Size | Backend |
+|-------|--------------|--------|------------|---------|
+| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | W4A16 | 128 | GPTQ-Marlin |
+
+## Creating a Quantized Checkpoint
+
+Use [AutoRound](https://github.com/intel/auto-round) to quantize a BF16 model. The `--scheme` flag selects the quantization scheme:
+
+```bash
+# W4A16 (4-bit weight, 16-bit activation)
+auto-round \
+    --model black-forest-labs/FLUX.1-dev \
+    --scheme W4A16 \
+    --batch_size 1 \
+    --disable_opt_rtn \
+    --dataset coco2014 \
+    --iters 0
+```
+
+The output directory can be used directly as the `model` argument. See the [AutoRound documentation](https://github.com/intel/auto-round) for all available schemes and options.

docs/user_guide/diffusion/quantization/fp8.md

@@ -65,6 +65,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
 | Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
 | HunyuanImage-3 | `tencent/HunyuanImage3` | All layers | None |
 | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None |
+| Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | All layers | None |
 
 ## Combining with Other Features
 

docs/user_guide/diffusion/quantization/overview.md

@@ -11,6 +11,7 @@ vLLM-Omni provides a unified quantization framework that supports both diffusion
 | FP8 | [FP8](fp8.md) | FP8 W8A8, dynamic or static | Z-Image, Qwen-Image, Flux, Bagel | SM 89 (Ada) |
 | Int8 | [Int8](int8.md) | Int8 W8A8 | Z-Image, Qwen-Image | SM 89 (Ada) / Ascend NPU |
 | GGUF | [GGUF](gguf.md) | GGUF format, dequant+GEMM for N-D tensors | Z-Image, Flux | SM 60 |
+| AutoRound | [AutoRound](autoround.md) | W4A16 (pre-quantized) | Flux | SM 80 (Ampere) |
 
 ### Multi-stage Omni Models (Pre-quantized Checkpoints)
 
@@ -102,6 +103,9 @@ config = build_quant_config("fp8")
 # Dict with parameters
 config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
 
+# AutoRound / INC (auto-detected from checkpoint, or explicit)
+config = build_quant_config("auto-round", bits=4, group_size=128)
+
 # Per-component dict
 config = build_quant_config({
     "transformer": {"method": "fp8"},

docs/user_guide/diffusion/step_execution.md

@@ -46,7 +46,7 @@ its stepwise request state machine. For normal diffusion inference, leave it
 disabled unless your workflow depends on this mode.
 
 If you are looking for general diffusion speedups, see
-[Diffusion Acceleration Overview](../diffusion_acceleration.md).
+[Diffusion Features Overview](../diffusion_features.md).
 
 ## Troubleshooting
 

docs/user_guide/diffusion/teacache.md

@@ -1,145 +0,0 @@
-# TeaCache Configuration Guide
-
-TeaCache speeds up diffusion model inference by caching transformer computations when consecutive timesteps are similar. This typically provides **1.5x-2.0x speedup** with minimal quality loss.
-
-## Quick Start
-
-Enable TeaCache by setting `cache_backend` to `"tea_cache"`:
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-# Simple configuration - model_type is automatically extracted from pipeline.__class__.__name__
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_backend="tea_cache",
-    cache_config={
-        "rel_l1_thresh": 0.2  # Optional, defaults to 0.2
-    }
-)
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(
-        num_inference_steps=50,
-    ),
-)
-```
-
-### 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}  # Optional
-)
-```
-
-## Online Serving (OpenAI-Compatible)
-
-Enable TeaCache for online serving by passing `--cache-backend tea_cache` when starting the server:
-
-```bash
-vllm serve Qwen/Qwen-Image --omni --port 8091 \
-  --cache-backend tea_cache \
-  --cache-config '{"rel_l1_thresh": 0.2}'
-```
-
-## Configuration Parameters
-
-### `rel_l1_thresh` (float, default: `0.2`)
-
-Controls the balance between speed and quality. Lower values prioritize quality, higher values prioritize speed.
-
-**Recommended values:**
-
-- `0.2` - **~1.5x speedup** with minimal quality loss (recommended)
-- `0.4` - **~1.8x speedup** with slight quality loss
-- `0.6` - **~2.0x speedup** with noticeable quality loss
-- `0.8` - **~2.25x speedup** with significant quality loss
-
-## Examples
-
-### Python API
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_backend="tea_cache",
-    cache_config={"rel_l1_thresh": 0.2}
-)
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(
-        num_inference_steps=50,
-    ),
-)
-```
-
-## Performance Tuning
-
-Start with the default `rel_l1_thresh=0.2` and adjust based on your needs:
-
-- **Maximum quality**: Use `0.1-0.2`
-- **Balanced**: Use `0.2-0.4` (recommended)
-- **Maximum speed**: Use `0.6-0.8` (may reduce quality)
-
-## Troubleshooting
-
-### Quality Degradation
-
-If you notice quality issues, lower the threshold:
-
-```python
-cache_config={"rel_l1_thresh": 0.1}  # More conservative caching
-```
-
-## Supported Models
-
-### ImageGen
-
-<style>
-th {
-  white-space: nowrap;
-  min-width: 0 !important;
-}
-</style>
-
-| Architecture | Models | Example HF Models |
-|--------------|--------|-------------------|
-| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
-| `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` |
-| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |
-
-### VideoGen
-
-No VideoGen models are supported by TeaCache yet.
-
-### Coming Soon
-
-<style>
-th {
-  white-space: nowrap;
-  min-width: 0 !important;
-}
-</style>
-
-| Architecture | Models | Example HF Models |
-|--------------|--------|-------------------|
-| `FluxPipeline` | Flux | - |
-| `CogVideoXPipeline` | CogVideoX | - |

docs/user_guide/diffusion_acceleration.md

@@ -1,380 +0,0 @@
-# Diffusion Acceleration Overview
-
-vLLM-Omni supports various acceleration methods to speed up diffusion model inference with minimal quality degradation. These include **cache methods** that intelligently cache intermediate computations to avoid redundant work across diffusion timesteps, **parallelism methods** that distribute the computation across multiple devices, and **quantization methods** that reduce memory footprint while preserving accuracy.
-
-## Supported Acceleration Methods
-
-vLLM-Omni currently supports two main cache acceleration backends:
-
-1. **[TeaCache](diffusion/teacache.md)** - Hook-based adaptive caching that caches transformer computations when consecutive timesteps are similar
-2. **[Cache-DiT](diffusion/cache_dit_acceleration.md)** - Library-based acceleration using multiple techniques:
-    - **DBCache** (Dual Block Cache): Caches intermediate transformer block outputs based on residual differences
-    - **TaylorSeer**: Uses Taylor expansion-based forecasting for faster inference
-    - **SCM** (Step Computation Masking): Selectively computes steps based on adaptive masking
-
-Both methods can provide significant speedups (typically **1.5x-2.0x**) while maintaining high output quality.
-
-vLLM-Omni also supports quantization methods:
-
-3. **[Quantization](diffusion/quantization/overview.md)** - Reduces DiT linear layers from BF16 to FP8 or Int8, providing ~1.28x speedup with minimal quality loss. Supports per-layer skip for sensitive layers.
-
-vLLM-Omni also supports parallelism methods for diffusion models, including:
-
-1. [Ulysses-SP](diffusion/parallelism_acceleration.md#ulysses-sp) - splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
-
-2. [Ring-Attention](diffusion/parallelism_acceleration.md#ring-attention) - splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results, keeping the sequence dimension sharded.
-
-3. [CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel) - runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step.
-
-4. [Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism) - shards DiT weights across devices to reduce per-GPU memory usage.
-
-5. [VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism) - shards VAE decode/encode spatially across ranks to reduce VAE peak memory (and can speed up VAE decode).
-
-6. [HSDP](diffusion/parallelism_acceleration.md#hsdp) - Hybrid Sharded Data Parallel shards model weights across GPUs to reduce per-GPU memory usage, enabling inference of large models on limited GPU memory.
-
-## Quick Comparison
-
-### Cache Methods
-
-| Method | Configuration | Description | Best For |
-|--------|--------------|-------------|----------|
-| **TeaCache** | `cache_backend="tea_cache"` | Simple, adaptive caching with minimal configuration | Quick setup, balanced speed/quality |
-| **Cache-DiT** | `cache_backend="cache_dit"` | Advanced caching with multiple techniques (DBCache, TaylorSeer, SCM) | Maximum acceleration, fine-grained control |
-
-### Quantization Methods
-
-| Method | Configuration | Description | Best For |
-|--------|--------------|-------------|----------|
-| **FP8** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup |
-| **Int8** | `quantization="int8"` | Int8 W8A8 | Memory reduction, inference speedup |
-
-## Supported Models
-
-The following table shows which models are currently supported by each acceleration method:
-
-### ImageGen
-
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel |
-|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:---------------:|:------------------:|
-| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
-| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
-| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
-| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ (TP=2 only) | ✅ |
-| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
-| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
-| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
-| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
-| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
-| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
-| **FLUX.1-Kontext-dev** | `black-forest-labs/FLUX.1-Kontext-dev` | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
-| **GLM-Image** | `zai-org/GLM-Image` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
-
-### VideoGen
-
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | VAE-Patch-Parallel |
-|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
-| **Wan2.1-T2V** | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Wan2.1-T2V** | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ |
-
-### Quantization
-
-| Model | Model Identifier | FP8 | Int8 |
-|-------|------------------|:---:|:---:|
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ |
-| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ |
-
-
-### AudioGen
-
-| Model                  | Model Identifier                         | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel |
-|------------------------|------------------------------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|
-| **Stable-Audio-Open**  | `stabilityai/stable-audio-open-1.0`      |    ✅    |           |            |                |              |
-
-## Performance Benchmarks
-
-The following benchmarks were measured on **Qwen/Qwen-Image** and **Qwen/Qwen-Image-Edit** models generating 1024x1024 images with 50 inference steps:
-
-!!! note "Benchmark Disclaimer"
-    These benchmarks are provided for **general reference only**. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:
-
-    - Specific model and use case
-    - Hardware configuration
-    - Careful parameter tuning
-    - Different inference settings (e.g., number of steps, image resolution)
-
-    For optimal performance in your specific scenario, we recommend experimenting with different parameter configurations as described in the detailed guides below.
-
-| Model | Cache Backend | Cache Config | Generation Time | Speedup | Notes |
-|-------|---------------|--------------|----------------|---------|-------|
-| **Qwen/Qwen-Image** | None  | None | 20.0s | 1.0x | Baseline (diffusers) |
-| **Qwen/Qwen-Image** | TeaCache | `rel_l1_thresh=0.2` | 10.47s | **1.91x** | Recommended default setting |
-| **Qwen/Qwen-Image** | Cache-DiT | DBCache + TaylorSeer (Fn=1, Bn=0, W=8, TaylorSeer order=1) | 10.8s | **1.85x** | - |
-| **Qwen/Qwen-Image** | Cache-DiT | DBCache + TaylorSeer + SCM (Fn=8, Bn=0, W=4, TaylorSeer order=1, SCM fast) | 14.0s | **1.43x** | - |
-| **Qwen/Qwen-Image-Edit** | None | No acceleration | 51.5s | 1.0x | Baseline (diffusers) |
-| **Qwen/Qwen-Image-Edit** | Cache-DiT | Default (Fn=1, Bn=0, W=4, TaylorSeer disabled, SCM disabled) | 21.6s | **2.38x** | - |
-
-To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** model generating images (**2048x2048** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA H800 GPUs. `sdpa` is the attention backends.
-
-| Configuration | Ulysses degree |Generation Time | Speedup |
-|---------------|----------------|---------|---------|
-| **Baseline (diffusers)** | - | 112.5s | 1.0x |
-| Ulysses-SP  |  2  |  65.2s | 1.73x |
-| Ulysses-SP  |  4  | 39.6s | 2.84x |
-| Ulysses-SP  |  8  | 30.8s | 3.65x |
-
-## Quick Start
-
-### Using TeaCache
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_backend="tea_cache",
-    cache_config={"rel_l1_thresh": 0.2}  # Optional, defaults to 0.2
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(
-        num_inference_steps=50,
-    ),
-)
-```
-
-### Using 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",
-    cache_config={
-        "Fn_compute_blocks": 1,
-        "Bn_compute_blocks": 0,
-        "max_warmup_steps": 8,
-        "enable_taylorseer": True,
-        "taylorseer_order": 1,
-    }
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(
-        num_inference_steps=50,
-    ),
-)
-```
-
-### Using Ulysses-SP
-
-Run text-to-image:
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-ulysses_degree = 2
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
-)
-```
-
-
-Run image-to-image:
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-ulysses_degree = 2
-
-omni = Omni(
-    model="Qwen/Qwen-Image-Edit",
-    parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
-)
-
-outputs = omni.generate(
-    {
-        "prompt": "turn this cat to a dog",
-        "multi_modal_data": {"image": input_image}
-    },
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-### Using Ring-Attention
-
-Run text-to-image:
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-ring_degree = 2
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    parallel_config=DiffusionParallelConfig(ring_degree=2)
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
-)
-```
-
-### Using Tensor Parallelism
-
-```python
-from vllm_omni import Omni
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="Tongyi-MAI/Z-Image-Turbo",
-    parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
-)
-
-outputs = omni.generate(
-    prompt="a cat reading a book",
-    num_inference_steps=9,
-    width=512,
-    height=512,
-)
-```
-
-### Using VAE Patch Parallelism
-
-```python
-from vllm_omni import Omni
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="Tongyi-MAI/Z-Image-Turbo",
-    parallel_config=DiffusionParallelConfig(
-        tensor_parallel_size=2,
-        vae_patch_parallel_size=2,
-    ),
-)
-
-outputs = omni.generate(
-    prompt="a cat reading a book",
-    num_inference_steps=9,
-    width=1024,
-    height=1024,
-)
-```
-
-### Using HSDP
-
-HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-    parallel_config=DiffusionParallelConfig(
-        use_hsdp=True,
-        hsdp_replicate_size=1,  # Number of replica groups
-        hsdp_shard_size=8,      # Number of GPUs to shard across (or -1 for auto)
-    ),
-)
-
-outputs = omni.generate(
-    "A cat playing piano",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-### Using CFG-Parallel
-
-Run image-to-image:
-
-CFG-Parallel splits the CFG positive/negative branches across GPUs. Use it when you set a non-trivial `true_cfg_scale`.
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-cfg_parallel_size = 2
-
-omni = Omni(
-    model="Qwen/Qwen-Image-Edit",
-    parallel_config=DiffusionParallelConfig(cfg_parallel_size=cfg_parallel_size)
-)
-
-outputs = omni.generate(
-    {
-        "prompt": "turn this cat to a dog",
-        "multi_modal_data": {"image": input_image}
-    },
-    OmniDiffusionSamplingParams(num_inference_steps=50, true_cfg_scale=4.0),
-)
-```
-
-### Using FP8 Quantization
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(
-    model="<your-model>",
-    quantization="fp8",
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-### Using Int8 Quantization
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(
-    model="<your-model>",
-    quantization="int8",
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-## Documentation
-
-For detailed information on each acceleration method:
-
-- **[TeaCache Guide](diffusion/teacache.md)** - Complete TeaCache documentation, configuration options, and best practices
-- **[Cache-DiT Acceleration Guide](diffusion/cache_dit_acceleration.md)** - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters
-- **[Quantization Guide](diffusion/quantization/overview.md)** - Quantization for DiT models with per-layer control
-- **[Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism)** - Guidance on how to enable TP for diffusion models.
-- **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
-- **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.
-- **[VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism)** - Guidance on how to reduce VAE memory via patch/tile parallelism.
-- **[HSDP](diffusion/parallelism_acceleration.md#hsdp)** - Hybrid Sharded Data Parallel for memory-efficient inference of large models.

docs/user_guide/diffusion_features.md

@@ -0,0 +1,209 @@
+# Diffusion Advanced Features
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Supported Features](#supported-features)
+- [Supported Models](#supported-models)
+- [Feature Compatibility](#feature-compatibility)
+- [Learn More](#learn-more)
+
+## Overview
+
+vLLM-Omni supports various advanced features for diffusion models:
+
+- Acceleration: **cache methods**, **parallelism methods**
+- Memory optimization: **cpu offloading**, **quantization**
+- Extensions: **LoRA inference**
+- Execution modes: **step execution**
+
+## Supported Features
+
+### Acceleration
+
+#### Lossy Acceleration
+
+Cache methods trade minimal quality for significant speedup. Quality loss is typically imperceptible with proper tuning.
+
+| Method | Description | Best For |
+|--------|-------------|----------|
+| **[TeaCache](diffusion/cache_acceleration/teacache.md)** | Adaptive caching using modulated inputs | Quick setup, balanced quality/speed on single GPU |
+| **[Cache-DiT](diffusion/cache_acceleration/cache_dit.md)** | Multiple caching techniques: DBCache, TaylorSeer, SCM | Fine-grained control, tunable quality-speed tradeoff |
+
+
+#### Lossless Acceleration
+
+Parallelism methods distribute computation across GPUs without quality loss (mathematically equivalent to single-GPU).
+
+| Method | Description | Best For |
+|--------|-------------|----------|
+| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via all-to-all communication | High-resolution images (>1536px) or long videos with 2-8 GPUs |
+| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via ring-based communication | Videos, very long sequences, memory-constrained, with 2-8 GPUs |
+| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)** | Splits CFG positive/negative branches across devices | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs |
+| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs |
+| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
+| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
+
+**Note:** Some acceleration methods can be combined together for optimized performance. See [Feature Compatibility Table](#feature-compatibility) and [Feature Compatibility Tutorial](feature_compatibility.md) for detailed configuration examples.
+
+### Memory Optimization
+
+Memory optimization methods help reduce GPU memory usage, enabling inference on resource-constrained hardware or larger models.
+
+| Method | Description | Best For |
+|--------|-------------|----------|
+| **[CPU Offload](diffusion/cpu_offload_diffusion.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
+| **[Quantization](diffusion/quantization/overview.md)** | Reduces DiT layers from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss    |
+| **[VAE Patch Parallelism](diffusion/parallelism/vae_patch_parallel.md)** | Distributes VAE decode tiling across GPUs | High-resolution generation with reduced VAE memory peak |
+
+### Extensions
+
+Extension methods add specialized capabilities to diffusion models beyond standard inference.
+
+| Method | Description | Best For |
+|--------|-------------|----------|
+| **[LoRA Inference](diffusion/lora.md)** | Enables inference with Low-Rank Adaptation (LoRA) adapters weights | Reinforcement learning extensions |
+
+
+### Execution Modes
+
+Execution modes control how the diffusion pipeline processes denoise steps.
+
+| Method | Description | Best For |
+|--------|-------------|----------|
+| **[Step Execution](diffusion/step_execution.md)** | Per-step denoise execution with mid-request abort support | Request cancellation between denoise steps, fine-grained execution control |
+
+**Note:** Step execution is currently supported by QwenImagePipeline only. See [Supported Models](#supported-models) for details.
+
+### Quantization Methods
+
+| Method | Configuration | Description | Best For |
+|--------|--------------|-------------|----------|
+| **[FP8](diffusion/quantization/fp8.md)** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup |
+| **[INT8](diffusion/quantization/int8.md)** | `quantization="int8"` | INT8 weight-only, no calibration or pre-quantized checkpoint needed | Memory reduction, broad GPU compatibility |
+| **[GGUF](diffusion/quantization/gguf.md)** | `quantization="gguf"` | Native GGUF transformer-only weights (Q4, Q8, etc.) | Memory reduction on consumer GPUs |
+
+## Supported Models
+
+The following tables show which models support each feature:
+
+- **🔀SP (Ulysses & Ring)**: Includes both Ulysses-SP and Ring-Attention methods
+- ✅ = Fully supported
+- ❌ = Not supported
+
+> Notes:
+
+> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only.
+> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](diffusion/quantization/overview.md) for per-method (FP8, GGUF, …) and per-model support details.
+
+### ImageGen
+
+| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
+| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
+| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
+| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
+| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
+| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
+| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ (decode) | ❌ | ❌ |
+| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
+
+> Notes:
+> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
+> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
+
+### VideoGen
+
+| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
+| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
+| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+
+### AudioGen
+
+| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Stable-Audio-Open** | ❌ | ❌ | ❓ | ❓ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
+
+
+## Feature Compatibility
+
+**Legend:**
+
+- ✅: Functionality is supported
+- ❌: No support plan
+- ❓: Not verified yet and Not Recommended
+
+|  | ⚡TeaCache | ⚡Cache-DiT | 🔀Ulysses-SP | 🔀Ring-Attn | 🔀CFG-Parallel | 🔀Tensor Parallel | 🔀HSDP | 🔀Expert Parallel | 💾CPU Offloading (Layerwise) | 💾CPU Offloading (Module-wise) | 💾VAE Patch Parallel | 💾FP8 Quant | 🔧LoRA Inference | 🔄Step Execution |
+|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| **⚡TeaCache** | | | | | | | | | | | | | | |
+| **⚡Cache-DiT** | ❌ | | | | | | | | | | | | | |
+| **🔀Ulysses-SP** | ✅ | ✅ | | | | | | | | | | | | |
+| **🔀Ring-Attn** | ✅ | ✅ | ✅ | | | | | | | | | | | |
+| **🔀CFG-Parallel** | ✅ | ✅ | ✅ | ✅ | | | | | | | | | | |
+| **🔀Tensor Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | |
+| **🔀HSDP** | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | | | | | | | |
+| **🔀Expert Parallel** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | | | | | | |
+| **💾CPU Offloading (Layerwise)** | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | | | | | |
+| **💾CPU Offloading (Module-wise)** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ❌ | | | | | |
+| **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | | |
+| **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | | |
+| **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | |
+| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ❌ | |
+
+!!! info
+
+    1. Tensor Parallel and HSDP are not compatible.
+    2. TeaCache and Cache-DiT are not compatible.
+    3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
+    4. CPU Offloading (Layerwise) supports single-card for now.
+    5. Using FP8-Quant as an example of qunatization methods.
+    6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT) or LoRA.
+
+
+## Learn More
+
+**Cache Acceleration:**
+
+- **[TeaCache Configuration Guide](diffusion/cache_acceleration/teacache.md)** - Parameter tuning, performance tips, troubleshooting
+- **[Cache-DiT Advanced Guide](diffusion/cache_acceleration/cache_dit.md)** - DBCache, TaylorSeer, SCM techniques and optimization
+
+**Parallelism Methods:**
+
+- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, HSDP, and Expert Parallelism
+
+**Memory Optimization:**
+
+- **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage
+- **[VAE Patch Parallelism Guide](diffusion/parallelism/vae_patch_parallel.md)** - Distribute VAE decode tiling across GPUs for high-resolution images
+- **[Quantization Overview](diffusion/quantization/overview.md)** - Overview of quantization methods for diffusion models
+
+**Extensions:**
+
+- **[LoRA Inference Guide](diffusion/lora.md)** - Low-Rank Adaptation for style customization and fine-tuning
+
+**Execution Modes:**
+
+- **[Step Execution Guide](diffusion/step_execution.md)** - Per-step denoise execution with mid-request abort support
+
+**Advanced Topics:**
+
+- **[Feature Compatibility](feature_compatibility.md)** - How to combine multiple features for maximum performance

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -256,6 +256,15 @@ lora_adapter/
 
 ## Web UI Demo
 
+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 Launch the Gradio demo:
 
 ```bash

docs/user_guide/examples/online_serving/fish_speech.md

@@ -11,6 +11,15 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serv
 
 ## Gradio Demo
 
+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support.
 
 ```bash

docs/user_guide/examples/online_serving/image_to_image.md

@@ -122,6 +122,15 @@ python openai_chat_client.py --input input1.png input2.png --prompt "Combine the
 
 ### Method 4: Using Gradio Demo
 
+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 ```bash
 python gradio_demo.py
 # Visit http://localhost:7861

docs/user_guide/examples/online_serving/image_to_video.md

@@ -37,12 +37,45 @@ artifact, poll the job status and then download the completed file from the
 content endpoint.
 
 The main endpoints are:
-- `POST /v1/videos`: create a video generation job
+- `POST /v1/videos`: create a video generation job (async)
+- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
 - `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
 - `GET /v1/videos`: list stored video jobs
 - `GET /v1/videos/{video_id}/content`: download the generated video file
 - `DELETE /v1/videos/{video_id}`: delete the job and any stored output
 
+## Sync API (Benchmark / Testing)
+
+`POST /v1/videos/sync` is a synchronous alternative that blocks until generation
+completes and returns the raw video bytes (`video/mp4`) directly in the response
+body. It is designed for benchmark and testing scenarios where one-shot
+request/response latency measurement is needed.
+
+The sync endpoint accepts the same form parameters as `POST /v1/videos`. It does
+not create any stored job record — the response is purely the generated video
+file. Metadata is returned via response headers:
+
+- `X-Request-Id`: unique identifier for this generation request
+- `X-Model`: model name used for generation
+- `X-Inference-Time-S`: wall-clock inference time in seconds
+
+```bash
+curl -X POST http://localhost:8091/v1/videos/sync \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F "input_reference=@/path/to/input.png" \
+  -F "size=832x480" \
+  -F "num_frames=33" \
+  -F "fps=16" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=1.0" \
+  -F "guidance_scale_2=1.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=12.0" \
+  -F "seed=42" \
+  -o sync_i2v_output.mp4
+```
+
 ## Storage
 
 Generated video files are stored on local disk by the async video API.

docs/user_guide/examples/online_serving/qwen2_5_omni.md

@@ -149,6 +149,15 @@ python openai_chat_completion_client_for_multimodal_generation.py \
 
 ## Run Local Web UI Demo
 
+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 This Web UI demo allows users to interact with the model through a web browser.
 
 ### Running Gradio Demo

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -154,7 +154,7 @@ print(response.choices[1].message.audio)    # Audio response
 ```
 
 ## Streaming Output
-If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.
+If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. We support both text streaming output and audio streaming output. Other modalities can output normally.
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py \
     --query-type use_image \
@@ -163,6 +163,15 @@ python openai_chat_completion_client_for_multimodal_generation.py \
 
 ## Run Local Web UI Demo
 
+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 This Web UI demo allows users to interact with the model through a web browser.
 
 ### Running Gradio Demo

docs/user_guide/examples/online_serving/qwen3_tts.md

@@ -19,6 +19,15 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
 
 ## Gradio Demo
 
+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 Two interactive Gradio demos are available, both supporting all 3 task types:
 
 | Demo     | File                     | Transport    | Streaming Quality                                  |
@@ -150,7 +159,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
     -H "Content-Type: application/json" \
     -d '{
         "input": "Hello, how are you?",
-        "speaker": "vivian",
+        "voice": "vivian",
         "language": "English"
     }' --output output.wav
 
@@ -159,7 +168,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
     -H "Content-Type: application/json" \
     -d '{
         "input": "I am so excited!",
-        "speaker": "vivian",
+        "voice": "vivian",
         "instructions": "Speak with great enthusiasm"
     }' --output excited.wav
 
@@ -176,7 +185,7 @@ client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
 
 response = client.audio.speech.create(
     model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
-    speaker="vivian",
+    voice="vivian",
     input="Hello, how are you?",
 )
 
@@ -192,7 +201,7 @@ response = httpx.post(
     "http://localhost:8091/v1/audio/speech",
     json={
         "input": "Hello, how are you?",
-        "speaker": "vivian",
+        "voice": "vivian",
         "language": "English",
     },
     timeout=300.0,
@@ -228,12 +237,16 @@ List all available voices/speakers from the loaded model, including both built-i
       "consent": "user_consent_id",
       "created_at": 1738660000,
       "file_size": 1024000,
-      "mime_type": "audio/wav"
+      "mime_type": "audio/wav",
+      "ref_text": "The exact transcript of the audio sample.",
+      "speaker_description": "warm narrator"
     }
   ]
 }
 ```
 
+Fields `ref_text` and `speaker_description` are omitted per-entry when not provided at upload time.
+
 #### POST /v1/audio/voices
 
 Upload a new voice sample for voice cloning in Base task TTS requests.
@@ -243,6 +256,7 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
 - `consent` (required): Consent recording ID
 - `name` (required): Name for the new voice
 - `ref_text` (optional): Transcript of the audio. Enables in-context voice cloning (higher quality).
+- `speaker_description` (optional): Free-form description of the voice (e.g. "warm narrator", "energetic presenter"). Stored as metadata.
 
 **Response Example:**
 ```json
@@ -253,18 +267,23 @@ 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:8000/v1/audio/voices \
+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 "ref_text=The exact transcript of the audio sample."
+  -F "ref_text=The exact transcript of the audio sample." \
+  -F "speaker_description=warm narrator"
 ```
 
 ### Endpoint
@@ -281,7 +300,7 @@ This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/
 ```json
 {
     "input": "Text to synthesize",
-    "speaker": "vivian",
+    "voice": "vivian",
     "response_format": "wav",
     "task_type": "CustomVoice",
     "language": "Auto",
@@ -301,7 +320,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
 
 ### Voice and language (summary)
 
-- **Speaker**: Use the `speaker` request field to select the speaker (e.g., `vivian`, `ryan`, `aiden`). List available speakers with `GET /v1/audio/voices`.
+- **Speaker**: Use the `voice` request field to select the speaker (e.g., `vivian`, `ryan`, `aiden`). List available speakers with `GET /v1/audio/voices`.
 - **Language**: Use the `language` field for the codec language tag (`Auto`, `Chinese`, `English`, etc.). Default is `Auto` for automatic detection.
 - **CustomVoice**: Requires a valid `voice` from the model’s speaker set. **VoiceDesign**: Use `instructions` to describe the voice. **Base**: Use `ref_audio` and `ref_text` for voice cloning.
 
@@ -313,7 +332,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
 | ----------------- | ------ | -------------- | ----------------------------------------------------------- |
 | `input`           | string | **required**   | Text to synthesize                                          |
 | `model`           | string | server's model | Model to use (optional, should match server if specified)   |
-| `speaker`         | string | "vivian"       | Speaker name (e.g., vivian, ryan, aiden)                    |
+| `voice`           | string | "vivian"       | Speaker name (e.g., vivian, ryan, aiden)                    |
 | `response_format` | string | "wav"          | Audio format: wav, mp3, flac, pcm, aac, opus                |
 | `speed`           | float  | 1.0            | Playback speed (0.25-4.0, not supported with `stream=true`) |
 
@@ -348,7 +367,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
     -H "Content-Type: application/json" \
     -d '{
         "input": "Hello, how are you?",
-        "speaker": "vivian",
+        "voice": "vivian",
         "language": "English",
         "stream": true,
         "response_format": "pcm"

docs/user_guide/examples/online_serving/text_to_image.md

@@ -105,6 +105,15 @@ python openai_chat_client.py --prompt "A beautiful landscape painting" --output
 
 ### Method 4: Using Gradio Demo
 
+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 ```bash
 python gradio_demo.py
 # Visit http://localhost:7860

docs/user_guide/examples/online_serving/text_to_video.md

@@ -37,12 +37,43 @@ artifact, poll the job status and then download the completed file from the
 content endpoint.
 
 The main endpoints are:
-- `POST /v1/videos`: create a video generation job
+- `POST /v1/videos`: create a video generation job (async)
+- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
 - `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
 - `GET /v1/videos`: list stored video jobs
 - `GET /v1/videos/{video_id}/content`: download the generated video file
 - `DELETE /v1/videos/{video_id}`: delete the job and any stored output
 
+## Sync API (Benchmark / Testing)
+
+`POST /v1/videos/sync` is a synchronous alternative that blocks until generation
+completes and returns the raw video bytes (`video/mp4`) directly in the response
+body. It is designed for benchmark and testing scenarios where one-shot
+request/response latency measurement is needed.
+
+The sync endpoint accepts the same form parameters as `POST /v1/videos`. It does
+not create any stored job record — the response is purely the generated video
+file. Metadata is returned via response headers:
+
+- `X-Request-Id`: unique identifier for this generation request
+- `X-Model`: model name used for generation
+- `X-Inference-Time-S`: wall-clock inference time in seconds
+
+```bash
+curl -X POST http://localhost:8091/v1/videos/sync \
+  -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
+  -F "size=832x480" \
+  -F "num_frames=33" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=4.0" \
+  -F "guidance_scale_2=4.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=5.0" \
+  -F "seed=42" \
+  -o sync_t2v_output.mp4
+```
+
 ## Storage
 
 Generated video files are stored on local disk by the async video API.

docs/user_guide/feature_compatibility.md

@@ -0,0 +1,218 @@
+# Feature Compatibility
+
+This guide explains the compatibility matrix of different diffusion features in vLLM-Omni. You can use cache methods together with parallelism methods and other features to achieve optimal speed and efficiency.
+
+## Overview
+
+vLLM-Omni supports combining:
+
+- **Cache methods** (TeaCache, Cache-DiT) with **Parallelism methods** (Ulysses-SP, Ring-Attention, CFG-Parallel, Tensor Parallelism)
+- **Multiple parallelism methods** together (e.g., Ulysses-SP + Ring-Attention, CFG-Parallel + Sequence Parallelism)
+- **LoRA adapters** with most acceleration features
+- **CPU offloading** with other memory optimization features
+
+See the feature compatibility matrix in [Table](diffusion_features.md#feature-compatibility)
+
+## Common Combinations
+
+### 1. Cache + Sequence Parallelism (Recommended)
+
+Best for: **Large images (>1536px) or videos**
+
+Combines cache acceleration with sequence parallelism for maximum speedup on single-device-challenging workloads.
+
+**Using TeaCache + Ulysses-SP:**
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image \
+  --prompt "A beautiful mountain landscape" \
+  --cache-backend tea_cache \
+  --ulysses-degree 2
+```
+
+**Using Cache-DiT + Ring-Attention:**
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image \
+  --prompt "A futuristic city" \
+  --cache-backend cache_dit \
+  --ring-degree 2
+```
+
+### 2. Cache + CFG-Parallel
+
+Best for: **Image editing with Classifier-Free Guidance**
+
+Accelerates both the diffusion process and CFG computation.
+
+```bash
+python examples/offline_inference/image_to_image/image_edit.py \
+  --model Qwen/Qwen-Image-Edit \
+  --prompt "make it sunset" \
+  --negative-prompt "low quality, blurry" \
+  --image input.png \
+  --cache-backend cache_dit \
+  --cfg-parallel-size 2 \
+  --cfg-scale 4.0
+```
+
+### 3. CFG-Parallel + Sequence Parallelism
+
+Best for: **Large resolution image editing with CFG**
+
+Combines both CFG branch splitting and sequence parallelism for maximum GPU utilization.
+
+**CFG-Parallel + Ulysses-SP:**
+
+```bash
+python examples/offline_inference/image_to_image/image_edit.py \
+  --model Qwen/Qwen-Image-Edit \
+  --prompt "transform into autumn scene" \
+  --negative-prompt "low quality" \
+  --image input.png \
+  --cache-backend cache_dit \
+  --cfg-parallel-size 2 \
+  --ulysses-degree 2 \
+  --cfg-scale 4.0
+```
+
+### 4. Hybrid Ulysses + Ring + Vae tiling
+
+Best for: **Very large images or videos on multiple devices**
+
+Combines Ulysses-SP (all-to-all) with Ring-Attention (ring P2P) for scalable parallelism.
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image \
+  --prompt "Epic fantasy landscape" \
+  --cache-backend cache_dit \
+  --ulysses-degree 2 \
+  --ring-degree 2 \
+  --num-inference-steps 50 \
+  --width 2048 \
+  --height 2048 \
+  --vae-use-tiling
+```
+
+### 5. Cache + Tensor Parallelism
+
+Best for: **Large models that don't fit in single GPU memory**
+
+Reduces per-GPU memory usage while maintaining cache acceleration.
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Tongyi-MAI/Z-Image-Turbo \
+  --prompt "A cat reading a book" \
+  --cache-backend tea_cache \
+  --tensor-parallel-size 2 \
+  --num-inference-steps 9 \
+```
+
+## Online Serving
+
+### Cache + Sequence Parallelism
+
+```bash
+# TeaCache + Ulysses-SP
+vllm serve Qwen/Qwen-Image --omni --port 8091 \
+  --cache-backend tea_cache \
+  --cache-config '{"rel_l1_thresh": 0.2}' \
+  --usp 2
+
+# Cache-DiT + Ring-Attention
+vllm serve Qwen/Qwen-Image --omni --port 8091 \
+  --cache-backend cache_dit \
+  --cache-config '{"Fn_compute_blocks": 1, "max_warmup_steps": 8}' \
+  --ring 2
+```
+
+### Cache + CFG-Parallel
+
+```bash
+vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 \
+  --cache-backend cache_dit \
+  --cfg-parallel-size 2
+```
+
+### Multiple Parallelism Methods
+
+```bash
+# CFG-Parallel + Ulysses-SP (4 GPUs total)
+vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 \
+  --cache-backend cache_dit \
+  --cfg-parallel-size 2 \
+  --usp 2
+
+# Hybrid Ulysses + Ring (4 GPUs total)
+vllm serve Qwen/Qwen-Image --omni --port 8091 \
+  --cache-backend cache_dit \
+  --usp 2 \
+  --ring 2
+```
+
+
+## Limitations
+
+### Incompatibilities
+
+- **TeaCache + Cache-DiT**: These two cache methods cannot be used together. Only one cache backend can be active at a time. Attempting to enable both will result in an error.
+
+### Partial Support
+
+- **Tensor Parallelism — Text Encoder Not Sharded**: TP currently only shards the DiT blocks. Each TP rank retains a **full copy of the text encoder weights**, leading to significant GPU memory overhead proportional to TP degree. Tracked in [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
+
+- **CPU Offloading — Two Modes Are Mutually Exclusive**: Model-level offload (`enable_cpu_offload`) and layerwise offload (`enable_layerwise_offload`) cannot be used simultaneously. If both are set, layerwise takes priority and model-level is silently ignored.
+
+- **CPU Offloading — VAE stays on GPU**: Both offloading strategies keep the VAE on GPU at all times. For high-resolution generation, VAE decode can still cause OOM. Mitigate by combining with `vae_use_tiling=True` or VAE Patch Parallelism.
+
+- **VAE Patch Parallelism — DistributedVaeExecutor Required**: VAE Patch Parallelism is only enabled for models that have `DistributedVaeExecutor`. Unsupported models will silently ignore `vae_patch_parallel_size`, and use sequential vae tiling instead.
+
+### Configuration Constraints
+
+- **GPU Count Must Match Parallel Degrees**: Total GPU count must satisfy:
+  ```
+  total_gpus = ulysses_degree × ring_degree × cfg_parallel_size × tensor_parallel_size
+  ```
+  Any mismatch will cause a configuration error at startup.
+
+- **VAE Patch Parallel Size ≤ DiT Process Group Size**: `vae_patch_parallel_size` reuses the DiT process group and cannot exceed it. Larger values are automatically clamped with a warning.
+
+- **Model-Specific TP Constraints**: Some models impose divisibility constraints on TP size. For example, Z-Image Turbo (`num_heads=30`) only supports `tensor_parallel_size=2`. Check [Supported Models](diffusion_features.md#supported-models) for per-model constraints.
+
+## Troubleshooting
+
+### Performance Not Scaling
+
+**Symptoms:** Adding more GPUs doesn't improve speed proportionally
+
+**Solutions:**
+1. Check GPU communication bandwidth (use `nvidia-smi topo -m`)
+2. Reduce parallelism degree if communication overhead is high
+3. For very long sequences, prefer Ring-Attention over Ulysses-SP
+4. Ensure batch size is large enough to saturate GPUs
+
+### Out of Memory with Parallelism
+
+**Symptoms:** OOM errors when combining methods
+
+**Solutions:**
+1. Enable Tensor Parallelism to shard weights
+2. Reduce resolution or batch size
+3. Combine with memory efficient methods, such as cpu offloading
+
+### Configuration Errors
+
+**Symptoms:** Errors about invalid parallel configuration
+
+**Solutions:**
+1. Verify total GPU count matches: `ulysses × ring × cfg × tp`
+2. Check model supports all enabled methods
+3. Ensure divisibility constraints (e.g., Z-Image TP=1 or 2 only)
+
+## See Also
+
+- [Diffusion Acceleration Overview](diffusion_features.md) - Main acceleration guide

examples/offline_inference/bagel/end2end.py

@@ -2,6 +2,10 @@ import argparse
 import os
 
 from vllm_omni.inputs.data import OmniPromptType
+from vllm_omni.model_executor.stage_input_processors.bagel import (
+    GEN_THINK_SYSTEM_PROMPT,
+    VLM_THINK_SYSTEM_PROMPT,
+)
 
 
 def parse_args():
@@ -65,6 +69,17 @@ def parse_args():
         help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
     )
     parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.")
+    parser.add_argument(
+        "--cfg-interval",
+        type=float,
+        nargs=2,
+        default=None,
+        help="CFG interval [start, end] (default: pipeline default)",
+    )
+    parser.add_argument(
+        "--cfg-renorm-type", type=str, default=None, help="CFG renorm type: global, text_channel, channel"
+    )
+    parser.add_argument("--cfg-renorm-min", type=float, default=None, help="CFG renorm min")
     parser.add_argument(
         "--enable-diffusion-pipeline-profiler",
         action="store_true",
@@ -76,6 +91,12 @@ def parse_args():
         default=None,
         help="Quantization method (e.g. 'fp8').",
     )
+    parser.add_argument(
+        "--think",
+        action="store_true",
+        default=False,
+        help="Enable thinking mode: AR stage decodes <think>...</think> planning tokens before image generation.",
+    )
 
     args = parser.parse_args()
     return args
@@ -101,7 +122,7 @@ def main():
 
     if not prompts:
         # Default prompt for text2img test if none provided
-        prompts = ["<|im_start|>A cute cat<|im_end|>"]
+        prompts = ["A cute cat"]
         print(f"[Info] No prompts provided, using default: {prompts}")
     omni_outputs = []
 
@@ -110,8 +131,12 @@ def main():
     from vllm_omni.entrypoints.omni import Omni
 
     omni_kwargs = {}
-    if args.stage_configs_path:
-        omni_kwargs["stage_configs_path"] = args.stage_configs_path
+    stage_configs_path = args.stage_configs_path
+    if args.think and stage_configs_path is None:
+        stage_configs_path = "vllm_omni/model_executor/stage_configs/bagel_think.yaml"
+        print(f"[Info] Think mode enabled, using stage config: {stage_configs_path}")
+    if stage_configs_path:
+        omni_kwargs["stage_configs_path"] = stage_configs_path
 
     omni_kwargs.update(
         {
@@ -136,7 +161,8 @@ def main():
             if not args.image_path or not os.path.exists(args.image_path):
                 raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}")
             loaded_image = Image.open(args.image_path).convert("RGB")
-            final_prompt_text = f"<|fim_middle|><|im_start|>{p}<|im_end|>"
+            think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
+            final_prompt_text = f"{think_prefix}<|fim_middle|><|im_start|>{p}<|im_end|>"
             prompt_dict = {
                 "prompt": final_prompt_text,
                 "multi_modal_data": {"img2img": loaded_image},
@@ -148,7 +174,10 @@ def main():
         elif args.modality == "img2text":
             if args.image_path:
                 loaded_image = Image.open(args.image_path).convert("RGB")
-                final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                think_prefix = f"<|im_start|>system\n{VLM_THINK_SYSTEM_PROMPT}<|im_end|>\n" if args.think else ""
+                final_prompt_text = (
+                    f"{think_prefix}<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                )
                 prompt_dict = {
                     "prompt": final_prompt_text,
                     "multi_modal_data": {"image": loaded_image},
@@ -156,11 +185,13 @@ def main():
                 }
                 formatted_prompts.append(prompt_dict)
         elif args.modality == "text2text":
-            final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
+            think_prefix = f"<|im_start|>{VLM_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
+            final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|><|im_start|>"
             prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
             formatted_prompts.append(prompt_dict)
         else:
-            final_prompt_text = f"<|im_start|>{p}<|im_end|>"
+            think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
+            final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|>"
             prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
             if args.negative_prompt is not None:
                 prompt_dict["negative_prompt"] = args.negative_prompt
@@ -168,7 +199,6 @@ def main():
 
     params_list = omni.default_sampling_params_list
     if args.modality in ("text2img", "img2img"):
-        params_list[0].max_tokens = 1  # type: ignore
         if len(params_list) > 1:
             diffusion_params = params_list[1]
             diffusion_params.num_inference_steps = args.steps  # type: ignore
@@ -179,6 +209,12 @@ def main():
                 "cfg_text_scale": args.cfg_text_scale,
                 "cfg_img_scale": args.cfg_img_scale,
             }
+            if args.cfg_interval is not None:
+                extra["cfg_interval"] = tuple(args.cfg_interval)
+            if args.cfg_renorm_type is not None:
+                extra["cfg_renorm_type"] = args.cfg_renorm_type
+            if args.cfg_renorm_min is not None:
+                extra["cfg_renorm_min"] = args.cfg_renorm_min
             if args.negative_prompt is not None:
                 extra["negative_prompt"] = args.negative_prompt
             diffusion_params.extra_args = extra  # type: ignore
@@ -187,6 +223,13 @@ def main():
 
     img_idx = 0
     for req_output in omni_outputs:
+        if args.think:
+            ro = getattr(req_output, "request_output", None)
+            if ro and getattr(ro, "outputs", None):
+                txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
+                if txt:
+                    print(txt)
+
         images = getattr(req_output, "images", None)
 
         if not images:
@@ -195,6 +238,7 @@ def main():
         for j, img in enumerate(images):
             save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
             img.save(save_path)
+            print(f"[Output] Saved image to {save_path}")
         img_idx += 1
 
     print(omni_outputs)

examples/offline_inference/glm_image/README.md

@@ -129,8 +129,17 @@ The first run loads model weights. Subsequent runs are faster:
 --stage-init-timeout 900  # Increase timeout for slow storage
 ```
 
+### `Transformers does not recognize this architecture` Error
+
+Your have to upgrade `transformers` package to `5.3.0` or above:
+
+```
+pip install --upgrade transformers
+```
+
 ## Requirements
 
 - vLLM-Omni with GLM-Image support
 - CUDA-capable GPU (recommended: H100/A100 with 80GB)
 - GLM-Image model weights
+- `transformers` v5.3.0 or above

examples/offline_inference/helios/end2end.py

@@ -178,6 +178,12 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        help="Quantization method (e.g. 'fp8').",
+    )
 
     # Memory & parallelism
     parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
@@ -213,6 +219,7 @@ def main():
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+        quantization=args.quantization,
     )
 
     # Validate I2V / V2V arguments

examples/offline_inference/image_to_video/image_to_video.py

@@ -490,10 +490,6 @@ def main():
         if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
             frames_np = frames_np[..., :3]
 
-        frames_np = np.clip(frames_np, 0.0, 1.0)
-        frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
-        video_tensor = torch.from_numpy(frames_u8)
-
         audio_out = None
         if audio is not None:
             if isinstance(audio, list):
@@ -507,7 +503,7 @@ def main():
                 audio_out = audio_out.float().cpu()
 
         encode_video(
-            video_tensor,
+            frames_np,
             fps=fps,
             audio=audio_out,
             audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,

examples/offline_inference/magi_human/README.md

@@ -0,0 +1,72 @@
+# MagiHuman Generation
+
+MagiHuman is an advanced, omni-modality model that generates both high-quality video and lip-synced audio from a text prompt.
+
+Because MagiHuman is a very large model featuring a powerful DiT MoE backbone and a ~9B parameter T5Gemma text encoder, it natively supports **Tensor Parallelism (TP)** in vLLM-Omni to run efficiently across multi-GPU setups, reducing device memory bottlenecks.
+
+## Setup
+
+### Install MagiCompiler (recommended)
+
+MagiHuman relies on [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for custom-op registration used by the DiT attention kernels. While the pipeline can fall back to stub implementations, installing MagiCompiler is **strongly recommended** for correct behaviour.
+
+```bash
+# Clone the repo
+git clone https://github.com/SandAI-org/MagiCompiler.git
+cd MagiCompiler
+
+# System dependencies (optional, for FX graph visualization; Debian/Ubuntu)
+sudo apt update && sudo apt install -y graphviz
+
+# Python dependencies
+pip install -r requirements.txt
+
+# Install MagiCompiler
+pip install .          # end users (recommended)
+# pip install -e .     # developers (editable install)
+```
+
+### Hardware requirements
+
+Ensure your hardware has enough VRAM. For a standard node with 80GB GPUs, running with `--tensor-parallel-size 4` is recommended to shard both the MoE weights and the T5Gemma text encoder across 4 GPUs, reducing the per-GPU peak VRAM overhead significantly (by roughly ~13.5GB per GPU compared to single-device inference).
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for further details on allocating memory.
+
+## Run Examples
+
+Get into the example folder:
+```bash
+cd examples/offline_inference/magi_human
+```
+
+### End-to-End Generation (Text to Video+Audio)
+
+Generate a video with synchronized speech natively generated by the model.
+
+```bash
+python end2end.py \
+    --model /proj-tango-pvc/users/zhipeng.wang/workspace/models/daVinci-MagiHuman \
+    --prompt "A young woman with long, wavy golden blonde hair..." \
+    --tensor-parallel-size 4 \
+    --output output_magihuman.mp4
+```
+
+## Common Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--model` | *(Required)* | Local model path or HuggingFace ID |
+| `--prompt` | *(built-in demo prompt)* | Highly detailed text prompt dictating visual look and dialogue text |
+| `--tensor-parallel-size` | `4` | Tensor parallelism size (Number of GPUs) |
+| `--height` | `256` | Initial resolution height |
+| `--width` | `448` | Initial resolution width |
+| `--num-inference-steps` | `8` | Denoising steps |
+| `--seed` | `52` | Random seed |
+| `--output` | `output_magihuman.mp4` | Output video with audio path |
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/magi_human/end2end.py"
+    ``````

examples/offline_inference/magi_human/end2end.py

@@ -0,0 +1,117 @@
+import argparse
+
+from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="End-to-end inference script for MagiHuman.")
+    parser.add_argument("--model", type=str, required=True, help="Path or ID of the MagiHuman model.")
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default="",
+        help="Text prompt containing visual description, dialogue, and background sound.",
+    )
+    parser.add_argument(
+        "--tensor-parallel-size", "-tp", type=int, default=4, help="Tensor parallel size (number of GPUs)."
+    )
+    parser.add_argument(
+        "--output", type=str, default="output_magihuman.mp4", help="Path to save the generated mp4 file."
+    )
+    parser.add_argument("--height", type=int, default=256, help="Video height.")
+    parser.add_argument("--width", type=int, default=448, help="Video width.")
+    parser.add_argument("--num-inference-steps", type=int, default=8, help="Number of denoising steps.")
+    parser.add_argument("--seed", type=int, default=52, help="Random seed for generation.")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    print(f"Initializing MagiHuman pipeline with TP={args.tensor_parallel_size}...")
+    omni = Omni(
+        model=args.model,
+        init_timeout=1200,
+        tensor_parallel_size=args.tensor_parallel_size,
+        devices=list(range(args.tensor_parallel_size)),
+    )
+
+    prompt = args.prompt
+    if not prompt:
+        prompt = (
+            "A young woman with long, wavy golden blonde hair and bright blue eyes, "
+            "wearing a fitted ivory silk blouse with a delicate lace collar, sits "
+            "stationary in front of a softly lit, blurred warm-toned interior. Her "
+            "overall disposition is warm, composed, and gently confident. The camera "
+            "holds a static medium close-up, framing her from the shoulders up, "
+            "with shallow depth of field keeping her face in sharp focus. Soft "
+            "directional key light falls from the upper left, casting a gentle "
+            "highlight along her cheekbone and nose bridge. She draws a quiet breath, "
+            "the levator labii superiors relaxing as her lips part. She speaks in "
+            "clear, warm, unhurried American English: "
+            "\"The most beautiful things in life aren't things at all — "
+            "they're moments, feelings, and the people who make you feel truly alive.\" "
+            "Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
+            "shapes each vowel with precision. A faint, genuine smile engages the "
+            "zygomaticus major, lifting her lip corners fractionally. Her brows rest "
+            "in a soft, neutral arch throughout. She maintains steady, forward-facing "
+            "eye contact. Head position remains level; no torso displacement occurs.\n\n"
+            "Dialogue:\n"
+            "<Young blonde woman, American English>: "
+            "\"The most beautiful things in life aren't things at all — "
+            "they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
+            "Background Sound:\n"
+            "<Soft, warm indoor ambience with a faint distant piano melody>"
+        )
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.num_inference_steps,
+        seed=args.seed,
+        extra_args={
+            "seconds": 5,
+            "sr_height": 1080,
+            "sr_width": 1920,
+            "sr_num_inference_steps": 5,
+        },
+    )
+
+    print(f"Generating with prompt: {prompt[:80]}...")
+    outputs = omni.generate(
+        prompts=[prompt],
+        sampling_params_list=[sampling_params],
+    )
+
+    print(f"Generation complete. Output type: {type(outputs)}")
+    if outputs:
+        first = outputs[0]
+
+        if hasattr(first, "images") and first.images:
+            video_frames = first.images[0]
+            print(f"Video frames: shape={video_frames.shape}, dtype={video_frames.dtype}")
+
+            audio_waveform = None
+            if hasattr(first, "multimodal_output") and first.multimodal_output:
+                audio_waveform = first.multimodal_output.get("audio")
+                if audio_waveform is not None:
+                    print(f"Audio waveform: shape={audio_waveform.shape}, dtype={audio_waveform.dtype}")
+
+            video_bytes = mux_video_audio_bytes(
+                video_frames,
+                audio_waveform,
+                fps=25.0,
+                audio_sample_rate=44100,
+            )
+            with open(args.output, "wb") as f:
+                f.write(video_bytes)
+            print(f"Saved MP4 ({len(video_bytes)} bytes) to {args.output}")
+        print("SUCCESS: MagiHuman pipeline generation completed.")
+    else:
+        print("WARNING: No outputs returned.")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/omnivoice/README.md

@@ -0,0 +1,73 @@
+# OmniVoice
+
+This directory contains an offline demo for running OmniVoice TTS models with vLLM Omni. It generates speech from text and saves WAV files locally.
+
+## Model Overview
+
+[OmniVoice](https://huggingface.co/k2-fsa/OmniVoice) is a zero-shot multilingual TTS model supporting 600+ languages. It uses a diffusion language model (Qwen3-0.6B backbone) with iterative masked unmasking to generate speech.
+
+Three inference modes are supported:
+
+- **Auto Voice**: Generate speech without any reference — the model picks a voice automatically.
+- **Voice Clone**: Clone a voice from a reference audio + transcription.
+- **Voice Design**: Control voice style via natural language instruction (e.g., "female, low pitch, british accent").
+
+## Setup
+
+Ensure the model is downloaded:
+
+```bash
+huggingface-cli download k2-fsa/OmniVoice
+```
+
+> **Note:** Voice cloning requires `transformers>=5.3.0` for `HiggsAudioV2TokenizerModel`. Auto voice and voice design modes work with `transformers>=4.57.0`.
+
+## Quick Start
+
+Auto voice (text only):
+
+```bash
+python end2end.py --model k2-fsa/OmniVoice --text "Hello, this is a test."
+```
+
+Voice design (with style instruction):
+
+```bash
+python end2end.py --model k2-fsa/OmniVoice \
+    --text "Hello, this is a test." \
+    --instruct "female, low pitch, british accent"
+```
+
+Voice clone (with reference audio):
+
+```bash
+python end2end.py --model k2-fsa/OmniVoice \
+    --text "Hello, this is a test." \
+    --ref-audio ref.wav \
+    --ref-text "This is the reference transcription."
+```
+
+## Language Support
+
+Specify a language for improved quality:
+
+```bash
+python end2end.py --model k2-fsa/OmniVoice \
+    --text "你好，这是一个测试。" \
+    --lang zh
+```
+
+## Architecture
+
+OmniVoice uses a two-stage pipeline:
+
+- **Stage 0 (Generator)**: Qwen3-0.6B transformer with 32-step iterative unmasking and classifier-free guidance. Generates 8-codebook audio tokens from text.
+- **Stage 1 (Decoder)**: HiggsAudioV2 RVQ quantizer + DAC acoustic decoder. Converts tokens to 24kHz waveform.
+
+Both stages use `GPUGenerationWorker` with `OmniGenerationScheduler`.
+
+## Notes
+
+- Output audio is saved to `output.wav` by default. Use `--output` to change the path.
+- The model estimates duration from text automatically via `RuleDurationEstimator`.
+- Use `--stage-init-timeout` to increase the stage initialization timeout for first-time model downloads.

examples/offline_inference/omnivoice/end2end.py

@@ -0,0 +1,164 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end OmniVoice TTS inference via vLLM-Omni.
+
+Supports:
+- Auto voice mode: text only → generated speech
+- Voice cloning mode: text + reference audio → cloned voice speech
+
+Usage:
+    # Auto voice
+    python end2end.py --model k2-fsa/OmniVoice --text "Hello world"
+
+    # Voice cloning
+    python end2end.py --model k2-fsa/OmniVoice --text "Hello" \
+        --ref-audio ref.wav --ref-text "reference transcription"
+"""
+
+import argparse
+import os
+
+import numpy as np
+import soundfile as sf
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+
+def run_e2e():
+    parser = argparse.ArgumentParser(description="OmniVoice E2E TTS inference")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="k2-fsa/OmniVoice",
+        help="Model name or path (HuggingFace or local)",
+    )
+    parser.add_argument(
+        "--stage-config",
+        type=str,
+        default="vllm_omni/model_executor/stage_configs/omnivoice.yaml",
+    )
+    parser.add_argument(
+        "--text",
+        type=str,
+        default="Hello, this is a test of the OmniVoice text to speech system.",
+    )
+    parser.add_argument(
+        "--ref-audio",
+        type=str,
+        default=None,
+        help="Reference audio for voice cloning (WAV file)",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Transcription of reference audio",
+    )
+    parser.add_argument(
+        "--lang",
+        type=str,
+        default=None,
+        help="Language code (e.g., 'en', 'zh')",
+    )
+    parser.add_argument(
+        "--instruct",
+        type=str,
+        default=None,
+        help="Voice design instruction (e.g., 'female, low pitch, british accent')",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="output.wav",
+        help="Output audio file path",
+    )
+    parser.add_argument(
+        "--stage-init-timeout",
+        type=int,
+        default=600,
+        help="Stage initialization timeout in seconds",
+    )
+    args = parser.parse_args()
+
+    if not os.path.exists(args.stage_config):
+        raise FileNotFoundError(f"Stage config not found: {args.stage_config}")
+
+    print(f"Initializing OmniVoice with model={args.model}")
+
+    omni = Omni(
+        model=args.model,
+        stage_configs_path=args.stage_config,
+        trust_remote_code=True,
+        log_stats=True,
+    )
+
+    print("Model initialized. Preparing inputs...")
+
+    # Build prompt
+    mm_processor_kwargs = {}
+    multi_modal_data = {}
+
+    if args.ref_audio:
+        if not os.path.exists(args.ref_audio):
+            raise FileNotFoundError(f"Reference audio not found: {args.ref_audio}")
+
+        import librosa
+
+        audio_signal, sr = librosa.load(args.ref_audio, sr=None)
+        multi_modal_data["audio"] = (audio_signal.astype(np.float32), sr)
+        mm_processor_kwargs["ref_text"] = args.ref_text or ""
+        mm_processor_kwargs["sample_rate"] = sr
+
+    if args.lang:
+        mm_processor_kwargs["lang"] = args.lang
+    if args.instruct:
+        mm_processor_kwargs["instruct"] = args.instruct
+
+    prompts = {"prompt": args.text}
+    if multi_modal_data:
+        prompts["multi_modal_data"] = multi_modal_data
+    if mm_processor_kwargs:
+        prompts["mm_processor_kwargs"] = mm_processor_kwargs
+
+    sampling_params_list = [OmniDiffusionSamplingParams()]
+
+    print(f"Generating speech for: {args.text}")
+
+    outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list))
+
+    print(f"Received {len(outputs)} outputs.")
+    for i, output in enumerate(outputs):
+        try:
+            ro = output.request_output
+            if ro is None:
+                print("No request_output found.")
+                continue
+
+            mm = getattr(ro, "multimodal_output", None)
+            if not mm and ro.outputs:
+                mm = getattr(ro.outputs[0], "multimodal_output", None)
+
+            if mm:
+                print(f"Multimodal output keys: {mm.keys()}")
+                if "audio" in mm:
+                    audio_out = mm["audio"]
+                    sr = mm.get("sr", 24000)
+                    if isinstance(audio_out, np.ndarray):
+                        audio_np = audio_out
+                    else:
+                        audio_np = audio_out.cpu().numpy().squeeze()
+                    out_path = args.output if i == 0 else f"output_{i}.wav"
+                    sf.write(out_path, audio_np, sr)
+                    print(f"Saved audio to {out_path} ({sr}Hz, {len(audio_np) / sr:.2f}s)")
+            else:
+                print("No multimodal output found.")
+        except Exception as e:
+            print(f"Error inspecting output: {e}")
+
+    omni.close()
+    print("Done.")
+
+
+if __name__ == "__main__":
+    run_e2e()

examples/offline_inference/text_to_image/README.md

@@ -33,6 +33,7 @@ This folder provides several entrypoints for experimenting with text-to-image di
 | `black-forest-labs/FLUX.2-klein-4B` | 1024 x 1024 | 72.7 | 14.9 |
 | `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 |
 | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
+| `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3)  | 160 |
 
 !!! info
 *Peak VRAM:  based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU.
@@ -90,6 +91,8 @@ python text_to_image.py \
 | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
 | `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
 | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
+| `--use-system-prompt` | str | `None` | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text. Recommended: `en_unified`. Only for HunyuanImage-3.0.|
+| `--system-prompt` | str | `None` | Custom system prompt text. Only used when `--use-system-prompt` is set to `custom`. Only for HunyuanImage-3.0.|
 
 **NextStep-1.1 specific arguments:**
 

examples/offline_inference/text_to_image/gradio_demo.py

@@ -1,7 +1,10 @@
 import argparse
 from functools import lru_cache
 
-import gradio as gr
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
 import torch
 
 from vllm_omni.entrypoints.omni import Omni

examples/offline_inference/text_to_image/text_to_image.py

@@ -64,8 +64,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--guidance-scale",
         type=float,
-        default=1.0,
-        help="Classifier-free guidance scale.",
+        default=4.0,
+        help="Classifier-free guidance scale. HunyuanImage3 recommends 4.0-5.0.",
     )
     parser.add_argument("--height", type=int, default=1024, help="Height of generated image.")
     parser.add_argument("--width", type=int, default=1024, help="Width of generated image.")
@@ -242,6 +242,19 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable logging of diffusion pipeline stats.",
     )
+    parser.add_argument(
+        "--use-system-prompt",
+        type=str,
+        default=None,
+        choices=["None", "dynamic", "en_vanilla", "en_recaption", "en_think_recaption", "en_unified", "custom"],
+        help="System prompt preset for generation. Recommended: en_unified.",
+    )
+    parser.add_argument(
+        "--system-prompt",
+        type=str,
+        default=None,
+        help=("Custom system prompt. Used when --use-system-prompt is custom. "),
+    )
     return parser.parse_args()
 
 
@@ -382,13 +395,13 @@ def main():
         )
 
     generation_start = time.perf_counter()
-
     extra_args = {
         "timesteps_shift": args.timesteps_shift,
         "cfg_schedule": args.cfg_schedule,
         "use_norm": args.use_norm,
+        "use_system_prompt": args.use_system_prompt,
+        "system_prompt": args.system_prompt,
     }
-
     if lora_request:
         extra_args["lora_request"] = lora_request
         extra_args["lora_scale"] = args.lora_scale

examples/offline_inference/text_to_video/text_to_video.py

@@ -56,8 +56,13 @@ def parse_args() -> argparse.Namespace:
         "Examples: Wan-AI/Wan2.2-T2V-A14B-Diffusers, "
         "hunyuanvideo-community/HunyuanVideo-1.5-480p_t2v",
     )
+    parser.add_argument(
+        "--model-class-name",
+        default=None,
+        help="Override model class name (e.g., LTX2TwoStagesVideoPipeline).",
+    )
     parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
-    parser.add_argument("--negative-prompt", default="", help="Negative prompt (Wan2.2 only).")
+    parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
     parser.add_argument("--seed", type=int, default=42, help="Random seed.")
     parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
     parser.add_argument(
@@ -185,6 +190,7 @@ def parse_args() -> argparse.Namespace:
 
 def main():
     args = parse_args()
+    model_class_name = args.model_class_name
 
     preset = _detect_preset(args.model)
     for key, default_val in preset.items():
@@ -229,6 +235,7 @@ def main():
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
+        model_class_name=model_class_name,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,

examples/offline_inference/vace/vace_video_generation.md

@@ -0,0 +1,88 @@
+# VACE Video Generation
+
+[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single model.
+
+| Model | Architecture | Model Weights (bf16) | HuggingFace |
+|-------|-------------|----------------------|-------------|
+| Wan2.1-VACE (1.3B) | Wan2.1 | ~10 GB | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
+| Wan2.1-VACE (14B) | Wan2.1 | ~38 GB | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
+
+## Text-to-Video (T2V)
+
+```bash
+python vace_video_generation.py \
+  --mode t2v \
+  --prompt "A sleek robot stands in a vast warehouse filled with boxes" \
+  --height 480 --width 832 --num-frames 81 \
+  --num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
+  --output t2v_output.mp4
+```
+
+## Image-to-Video (I2V)
+
+First frame is kept, remaining frames are generated:
+
+```bash
+python vace_video_generation.py \
+  --mode i2v \
+  --image astronaut.jpg \
+  --prompt "An astronaut emerging from a cracked egg on the moon" \
+  --height 480 --width 832 --num-frames 81 \
+  --output i2v_output.mp4
+```
+
+## First-Last-Frame Interpolation (FLF2V)
+
+```bash
+python vace_video_generation.py \
+  --mode flf2v \
+  --image first_frame.jpg --last-image last_frame.jpg \
+  --prompt "A bird takes off from a branch and lands on another" \
+  --height 512 --width 512 --num-frames 81 \
+  --output flf2v_output.mp4
+```
+
+## Inpainting
+
+Center vertical stripe is masked and regenerated:
+
+```bash
+python vace_video_generation.py \
+  --mode inpaint \
+  --image scene.jpg \
+  --prompt "Shrek walks out of a building" \
+  --height 480 --width 832 --num-frames 81 \
+  --output inpaint_output.mp4
+```
+
+## Reference Image-guided (R2V)
+
+```bash
+python vace_video_generation.py \
+  --mode r2v \
+  --image reference.jpg \
+  --prompt "Camera slowly zooms out from the character" \
+  --height 480 --width 832 --num-frames 81 \
+  --output r2v_output.mp4
+```
+
+## Key Arguments
+
+- `--mode`: VACE task mode (`t2v`, `i2v`, `flf2v`, `inpaint`, `r2v`).
+- `--model`: Model ID (default: `Wan-AI/Wan2.1-VACE-1.3B-diffusers`).
+- `--image`: Input image for I2V, inpainting, and R2V modes.
+- `--last-image`: Last frame image for FLF2V mode.
+- `--prompt`: Text description of desired video.
+- `--height/--width`: Output resolution (default 480x832). Dimensions should be multiples of 16.
+- `--num-frames`: Number of frames (default 81).
+- `--guidance-scale`: CFG scale (default 5.0).
+- `--flow-shift`: Scheduler flow shift (default 5.0).
+- `--num-inference-steps`: Number of denoising steps (default 30).
+- `--fps`: Frames per second for the saved MP4 (default 16).
+- `--output`: Path to save the generated video.
+- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
+- `--ulysses-degree`: Ulysses sequence parallelism degree for multi-GPU.
+- `--cfg-parallel-size`: CFG parallel size for multi-GPU.
+- `--tensor-parallel-size`: Tensor parallel size.
+
+> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options.

examples/offline_inference/vace/vace_video_generation.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""VACE video generation example.
+
+VACE (Video All-in-one Creation Engine) supports multiple video tasks:
+  - T2V:        Text-to-Video
+  - I2V:        Image-to-Video (first frame conditioning)
+  - V2LF:       Video-to-Last-Frame
+  - FLF2V:      First-Last-Frame interpolation
+  - Inpainting:  Masked region generation
+  - R2V:        Reference image-guided generation
+
+Usage examples:
+  # T2V (text-to-video)
+  python vace_video_generation.py --mode t2v --prompt "A robot in a warehouse"
+
+  # I2V (image-to-video, first frame kept)
+  python vace_video_generation.py --mode i2v --image input.jpg --prompt "..."
+
+  # FLF2V (first-last frame interpolation)
+  python vace_video_generation.py --mode flf2v --image first.jpg --last-image last.jpg
+
+  # R2V (reference image guided)
+  python vace_video_generation.py --mode r2v --image ref.jpg --prompt "..."
+"""
+
+import argparse
+import time
+from pathlib import Path
+
+import numpy as np
+import PIL.Image
+import torch
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="VACE video generation.")
+    parser.add_argument(
+        "--model",
+        default="Wan-AI/Wan2.1-VACE-14B-diffusers",
+        help="VACE model ID or local path.",
+    )
+    parser.add_argument(
+        "--mode",
+        default="t2v",
+        choices=["t2v", "i2v", "v2lf", "flf2v", "inpaint", "r2v"],
+        help="Generation mode.",
+    )
+    parser.add_argument("--prompt", default="A cat walking in a garden", help="Text prompt.")
+    parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
+    parser.add_argument("--image", type=str, default=None, help="Input image path (for I2V, R2V, FLF2V, inpaint).")
+    parser.add_argument("--last-image", type=str, default=None, help="Last frame image path (for FLF2V).")
+    parser.add_argument("--video-dir", type=str, default=None, help="Directory of video frames (for inpaint).")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
+    parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG guidance scale.")
+    parser.add_argument("--height", type=int, default=480, help="Video height.")
+    parser.add_argument("--width", type=int, default=832, help="Video width.")
+    parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
+    parser.add_argument("--num-inference-steps", type=int, default=30, help="Sampling steps.")
+    parser.add_argument("--flow-shift", type=float, default=5.0, help="Scheduler flow_shift.")
+    parser.add_argument("--output", type=str, default="vace_output.mp4", help="Output video path.")
+    parser.add_argument("--fps", type=int, default=16, help="Output video FPS.")
+    parser.add_argument("--vae-use-tiling", action="store_true", default=True, help="Enable VAE tiling.")
+    parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
+    parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
+    parser.add_argument("--ring-degree", type=int, default=1, help="Ring attention degree.")
+    parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
+    return parser.parse_args()
+
+
+def build_prompts(args):
+    """Build prompt dict with multi_modal_data based on mode."""
+    h, w, nf = args.height, args.width, args.num_frames
+
+    gray = PIL.Image.new("RGB", (w, h), (128, 128, 128))
+    mask_black = PIL.Image.new("L", (w, h), 0)
+    mask_white = PIL.Image.new("L", (w, h), 255)
+
+    prompt_data = {
+        "prompt": args.prompt,
+        "negative_prompt": args.negative_prompt,
+    }
+
+    if args.mode == "t2v":
+        return prompt_data
+
+    if args.mode == "r2v":
+        assert args.image, "--image required for R2V mode"
+        ref_img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
+        prompt_data["multi_modal_data"] = {"reference_images": [ref_img]}
+        return prompt_data
+
+    if args.mode == "i2v":
+        assert args.image, "--image required for I2V mode"
+        img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
+        prompt_data["multi_modal_data"] = {
+            "video": [img] + [gray] * (nf - 1),
+            "mask": [mask_black] + [mask_white] * (nf - 1),
+        }
+        return prompt_data
+
+    if args.mode == "v2lf":
+        assert args.image, "--image required for V2LF mode"
+        img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
+        prompt_data["multi_modal_data"] = {
+            "video": [gray] * (nf - 1) + [img],
+            "mask": [mask_white] * (nf - 1) + [mask_black],
+        }
+        return prompt_data
+
+    if args.mode == "flf2v":
+        assert args.image and args.last_image, "--image and --last-image required for FLF2V"
+        first = PIL.Image.open(args.image).convert("RGB").resize((w, h))
+        last = PIL.Image.open(args.last_image).convert("RGB").resize((w, h))
+        prompt_data["multi_modal_data"] = {
+            "video": [first] + [gray] * (nf - 2) + [last],
+            "mask": [mask_black] + [mask_white] * (nf - 2) + [mask_black],
+        }
+        return prompt_data
+
+    if args.mode == "inpaint":
+        assert args.image, "--image required for inpaint mode"
+        img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
+        d = 80
+        frames, masks = [], []
+        for _ in range(nf):
+            base = np.array(img).copy()
+            mask = PIL.Image.new("L", (w, h), 0)
+            stripe = PIL.Image.new("L", (2 * d, h), 255)
+            mask.paste(stripe, (w // 2 - d, 0))
+            base[np.array(mask) > 128] = 128
+            frames.append(PIL.Image.fromarray(base))
+            masks.append(mask)
+        prompt_data["multi_modal_data"] = {"video": frames, "mask": masks}
+        return prompt_data
+
+    raise ValueError(f"Unknown mode: {args.mode}")
+
+
+def main():
+    args = parse_args()
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+    )
+
+    omni = Omni(
+        model=args.model,
+        vae_use_tiling=args.vae_use_tiling,
+        flow_shift=args.flow_shift,
+        enforce_eager=args.enforce_eager,
+        parallel_config=parallel_config,
+    )
+
+    prompt_data = build_prompts(args)
+
+    print(f"\n{'=' * 60}")
+    print(f"VACE {args.mode.upper()} Generation")
+    print(f"  Model: {args.model}")
+    print(f"  Size: {args.width}x{args.height}, {args.num_frames} frames, {args.num_inference_steps} steps")
+    print(f"{'=' * 60}\n")
+
+    start = time.perf_counter()
+    outputs = omni.generate(
+        prompt_data,
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            num_frames=args.num_frames,
+            num_inference_steps=args.num_inference_steps,
+            guidance_scale=args.guidance_scale,
+            generator=generator,
+        ),
+    )
+    elapsed = time.perf_counter() - start
+
+    video = outputs[0].images
+    if isinstance(video, list):
+        video = video[0]
+    if isinstance(video, torch.Tensor):
+        video = video.cpu().numpy()
+    if video.ndim == 5:
+        video = video[0]
+    print(f"Output shape: {video.shape}, Time: {elapsed:.1f}s")
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    from diffusers.utils import export_to_video
+
+    if np.issubdtype(video.dtype, np.integer):
+        video = video.astype(np.float32) / 255.0
+    export_to_video(list(video), str(output_path), fps=args.fps)
+    print(f"Saved to {output_path}")
+
+    omni.close()
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/x_to_video_audio/download_dreamid_omni.py

@@ -82,7 +82,6 @@ def main(output_dir: str):
 
     data = {
         "_class_name": "DreamIDOmniPipeline",
-        "fusion": "DreamID-Omni/dreamid_omni.safetensors",
     }
 
     with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
@@ -90,6 +89,12 @@ def main(output_dir: str):
 
     print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}")
 
+    transformer_dir = os.path.join(output_dir, "transformer")
+    os.makedirs(transformer_dir, exist_ok=True)
+    with open(os.path.join(transformer_dir, "config.json"), "w", encoding="utf-8") as f:
+        json.dump({"fusion": "DreamID-Omni/dreamid_omni.safetensors"}, f)
+    print(f"transformer/config.json created at {os.path.join(transformer_dir, 'config.json')}")
+
     # now we download the dependency code
     download_dependency()
 

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -24,7 +24,9 @@ dreamid_omni/
 │   ├── models_t5_umt5-xxl-enc-bf16.pth
 │   ├── Wan2.2_VAE.pth
 │
-├── model_index.json # create by download_dreamid_omni.py
+├── model_index.json
+└── transformer/
+    └── config.json   # create by download_dreamid_omni.py
 ```
 
 ### Run the Inference