Compare commits

..
Author SHA1 Message Date
Roger Wang 171826dd8b release
Signed-off-by: Roger Wang <hey@rogerw.io>
2026-03-01 05:16:03 +00:00
825 changed files with 18645 additions and 119559 deletions

No files matched your search

+1 -9
View File
@@ -90,18 +90,10 @@ upload_pipeline() {
FAIL_FAST=$(fail_fast)
cd .buildkite
# Select test definition file: merge suite for main, ready suite for PRs
if [[ $BUILDKITE_BRANCH == "main" ]]; then
TEST_YAML="test-amd-merge.yml"
else
TEST_YAML="test-amd-ready.yaml"
fi
(
set -x
# Output pipeline.yaml with all blank lines removed
minijinja-cli test-template.j2 "$TEST_YAML" \
minijinja-cli test-template.j2 test-amd.yaml \
-D branch="$BUILDKITE_BRANCH" \
-D list_file_diff="$LIST_FILE_DIFF" \
-D run_all="$RUN_ALL" \
-51
View File
@@ -1,51 +0,0 @@
#!/bin/bash
# vllm-omni Intel bootstrap
# Uses static pipeline-intel.yaml for Intel XPU tests
# Based on: bootstrap-amd-omni.sh
set -euo pipefail
if [[ -z "${DOCS_ONLY_DISABLE:-}" ]]; then
DOCS_ONLY_DISABLE=0
fi
upload_pipeline() {
echo "--- 🛠 Preparing Intel pipeline"
buildkite-agent pipeline upload .buildkite/pipeline-intel.yaml
}
get_diff() {
$(git add .)
echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
}
get_diff_main() {
$(git add .)
echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
}
file_diff=$(get_diff)
if [[ $BUILDKITE_BRANCH == "main" ]]; then
file_diff=$(get_diff_main)
fi
# ----------------------------------------------------------------------
# Early exit: skip pipeline if all changed files are under docs/
# ----------------------------------------------------------------------
if [[ "${DOCS_ONLY_DISABLE}" != "1" ]] && [[ -n "${file_diff:-}" ]]; then
docs_only=1
while IFS= read -r f; do
[[ -z "$f" ]] && continue
if [[ "$f" != docs/* ]]; then
docs_only=0
break
fi
done < <(printf '%s\n' "$file_diff" | tr ' ' '\n' | tr -d '\r')
if [[ "$docs_only" -eq 1 ]]; then
buildkite-agent annotate ":memo: CI skipped — docs only" --style "info"
exit 0
fi
fi
upload_pipeline
-20
View File
@@ -1,20 +0,0 @@
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"
-16
View File
@@ -1,16 +0,0 @@
# .buildkite/pipeline-intel.yaml
steps:
- group: "INTEL Tests"
steps:
- label: "Intel: XPU Test"
command: "bash .buildkite/scripts/hardware_ci/run-xpu-test.sh"
agents:
queue: "intel-gpu-omni"
env:
DOCKER_BUILDKIT: "1"
# Buildkite will automatically replace this with the actual commit hash
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
VLLM_VERSION: "v0.18.0"
priority: 100
timeout_in_minutes: 60
soft_fail: true
+4 -4
View File
@@ -13,7 +13,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: build.branch != "main"
commands:
- buildkite-agent pipeline upload .buildkite/test-ready.yml
agents:
@@ -23,17 +23,17 @@ steps:
- label: "Upload Merge Pipeline"
depends_on: image-build
key: upload-merge-pipeline
if: build.branch == "main" && build.env("NIGHTLY") != "1"
if: build.branch == "main"
commands:
- buildkite-agent pipeline upload .buildkite/test-merge.yml
agents:
queue: "cpu_queue_premerge"
# L4 Test — main+NIGHTLY=1 (scheduled), or PR with label nightly-test (e.g. add label then Rebuild)
# L4 Test
- label: "Upload Nightly Pipeline"
depends_on: image-build
key: upload-nightly-pipeline
if: '(build.branch == "main" && build.env("NIGHTLY") == "1") || (build.branch != "main" && build.pull_request.labels includes "nightly-test")'
if: build.env("NIGHTLY") == "1"
commands:
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
agents:
-38
View File
@@ -1,38 +0,0 @@
steps:
- label: ":docker: Build image"
key: image-build
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 ."
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
agents:
queue: "cpu_queue_premerge"
# L2 Test
- label: "Upload Ready Pipeline"
depends_on: image-build
key: upload-ready-pipeline
commands:
- buildkite-agent pipeline upload .buildkite/test-ready.yml
agents:
queue: "cpu_queue_premerge"
# L3 Test
- label: "Upload Merge Pipeline"
depends_on: image-build
key: upload-merge-pipeline
commands:
- buildkite-agent pipeline upload .buildkite/test-merge.yml
agents:
queue: "cpu_queue_premerge"
# L4 Test
- label: "Upload Nightly Pipeline"
depends_on: image-build
key: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
agents:
queue: "cpu_queue_premerge"
@@ -1,87 +0,0 @@
#!/usr/bin/env bash
set -ex
# Generate and upload wheel indices for all vllm-omni wheels in the commit directory.
# This script should run once after all wheels have been built and uploaded.
# All paths are under the omni/ prefix in the vllm-wheels S3 bucket.
# ======== setup ========
BUCKET="vllm-wheels"
INDICES_OUTPUT_DIR="indices"
PYTHON="${PYTHON_PROG:-python3}"
SUBPATH="omni/$BUILDKITE_COMMIT"
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
# detect if python3.12+ is available
has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
if [[ "$has_new_python" -eq 0 ]]; then
# use new python from docker
docker pull python:3-slim
PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3"
fi
echo "Using python interpreter: $PYTHON"
echo "Python version: $($PYTHON --version)"
# ======== generate and upload indices ========
# list all wheels in the commit directory
echo "Existing wheels on S3:"
aws s3 ls "$S3_COMMIT_PREFIX" || echo "(no objects found)"
obj_json="objects.json"
aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$SUBPATH/" --delimiter / --output json > "$obj_json"
mkdir -p "$INDICES_OUTPUT_DIR"
# HACK: we do not need regex module here, but it is required by pre-commit hook
# To avoid any external dependency, we simply replace it back to the stdlib re module
sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
# Generate indices -- the version is just the commit hash (not omni/{commit})
# because relative paths are computed between the index and wheel directories,
# both of which live under the omni/ prefix in S3.
$PYTHON .buildkite/scripts/generate-nightly-index.py \
--version "$BUILDKITE_COMMIT" \
--current-objects "$obj_json" \
--output-dir "$INDICES_OUTPUT_DIR" \
--comment "commit $BUILDKITE_COMMIT"
# copy indices to /omni/{commit}/ unconditionally
echo "Uploading indices to $S3_COMMIT_PREFIX"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
# copy to /omni/nightly/ when NIGHTLY=1
if [[ "${NIGHTLY:-}" == "1" ]]; then
echo "Uploading indices to overwrite /omni/nightly/"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/nightly/"
fi
# detect version from any wheel in the commit directory
first_wheel_key=$($PYTHON -c "import json; obj=json.load(open('$obj_json')); print(next((c['Key'] for c in obj.get('Contents', []) if c['Key'].endswith('.whl')), ''))")
if [[ -z "$first_wheel_key" ]]; then
echo "Error: No wheels found in $S3_COMMIT_PREFIX"
exit 1
fi
first_wheel=$(basename "$first_wheel_key")
aws s3 cp "s3://$BUCKET/${first_wheel_key}" "/tmp/${first_wheel}"
version=$(unzip -p "/tmp/${first_wheel}" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
rm -f "/tmp/${first_wheel}"
echo "Version in wheel: $version"
pure_version="${version%%+*}"
echo "Pure version (without variant): $pure_version"
# re-generate and copy to /omni/{version}/ only if it does not have "dev" in the version
if [[ "$version" != *"dev"* ]]; then
echo "Re-generating indices for /omni/$pure_version/"
rm -rf "${INDICES_OUTPUT_DIR:?}"
mkdir -p "$INDICES_OUTPUT_DIR"
# wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
$PYTHON .buildkite/scripts/generate-nightly-index.py \
--version "$pure_version" \
--wheel-dir "$BUILDKITE_COMMIT" \
--current-objects "$obj_json" \
--output-dir "$INDICES_OUTPUT_DIR" \
--comment "version $pure_version"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/$pure_version/"
fi
@@ -1,193 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import json
import sys
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from typing import Any
from urllib.parse import quote
import regex as re
def normalize_package_name(name: str) -> str:
"""Normalize package name per PEP 503."""
return re.sub(r"[-_.]+", "-", name).lower()
if not sys.version_info >= (3, 12):
raise RuntimeError("This script requires Python 3.12 or higher.")
INDEX_HTML_TEMPLATE = """<!DOCTYPE html>
<html>
<!-- {comment} -->
<meta name="pypi:repository-version" content="1.0">
<body>
{items}
</body>
</html>
"""
@dataclass
class WheelFileInfo:
package_name: str
version: str
build_tag: str | None
python_tag: str
abi_tag: str
platform_tag: str
filename: str
def parse_from_filename(file: str) -> WheelFileInfo:
"""
Parse wheel filename per PEP 427:
{package_name}-{version}(-{build_tag})?-{python_tag}-{abi_tag}-{platform_tag}.whl
"""
wheel_file_re = re.compile(
r"^(?P<package_name>.+)-(?P<version>[^-]+?)(-(?P<build_tag>[^-]+))?-(?P<python_tag>[^-]+)-(?P<abi_tag>[^-]+)-(?P<platform_tag>[^-]+)\.whl$"
)
match = wheel_file_re.match(file)
if not match:
raise ValueError(f"Invalid wheel file name: {file}")
return WheelFileInfo(
package_name=match.group("package_name"),
version=match.group("version"),
build_tag=match.group("build_tag"),
python_tag=match.group("python_tag"),
abi_tag=match.group("abi_tag"),
platform_tag=match.group("platform_tag"),
filename=file,
)
def generate_project_list(package_names: list[str], comment: str = "") -> str:
"""Generate top-level PEP 503 project list HTML."""
href_tags = []
for name in sorted(package_names):
href_tags.append(f' <a href="{name}/">{name}/</a><br/>')
return INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
def generate_package_index(
wheel_files: list[WheelFileInfo],
wheel_base_dir: Path,
index_base_dir: Path,
comment: str = "",
) -> tuple[str, str]:
"""Generate package index HTML and metadata JSON linking to wheel files."""
href_tags = []
metadata = []
for file in sorted(wheel_files, key=lambda x: x.filename):
relative_path = wheel_base_dir.relative_to(index_base_dir, walk_up=True) / file.filename
# handle '+' in URL; avoid double-encoding '/' and '%2B' (AWS S3 behavior)
file_path_quoted = quote(relative_path.as_posix(), safe=":%/")
href_tags.append(f' <a href="{file_path_quoted}">{file.filename}</a><br/>')
file_meta = asdict(file)
file_meta["path"] = file_path_quoted
metadata.append(file_meta)
index_str = INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
metadata_str = json.dumps(metadata, indent=2)
return index_str, metadata_str
def generate_index(
whl_files: list[str],
wheel_base_dir: Path,
index_base_dir: Path,
comment: str = "",
):
"""
Generate PEP 503 index for all wheel files.
Output structure:
index_base_dir/
index.html # project list linking to vllm-omni/
vllm-omni/
index.html # package index linking to wheel files
metadata.json # machine-readable metadata
"""
parsed_files = [parse_from_filename(f) for f in whl_files]
if not parsed_files:
print("No wheel files found, skipping index generation.")
return
comment_str = f" ({comment})" if comment else ""
comment_tmpl = f"Generated on {datetime.now().isoformat()}{comment_str}"
# Group by normalized package name
packages: dict[str, list[WheelFileInfo]] = {}
for file in parsed_files:
name = normalize_package_name(file.package_name)
packages.setdefault(name, []).append(file)
print(f"Found packages: {list(packages.keys())}")
# Generate per-package index
for package, files in packages.items():
package_dir = index_base_dir / package
package_dir.mkdir(parents=True, exist_ok=True)
index_str, metadata_str = generate_package_index(files, wheel_base_dir, package_dir, comment)
with open(package_dir / "index.html", "w") as f:
f.write(index_str)
with open(package_dir / "metadata.json", "w") as f:
f.write(metadata_str)
# Generate top-level project list
project_list_str = generate_project_list(sorted(packages.keys()), comment_tmpl)
with open(index_base_dir / "index.html", "w") as f:
f.write(project_list_str)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Generate PEP 503 wheel index from S3 object listing.")
parser.add_argument("--version", type=str, required=True, help="Version string (e.g., commit hash)")
parser.add_argument("--current-objects", type=str, required=True, help="Path to JSON from S3 list-objects-v2")
parser.add_argument("--output-dir", type=str, required=True, help="Directory to write index files")
parser.add_argument("--wheel-dir", type=str, default=None, help="Wheel directory (defaults to --version)")
parser.add_argument("--comment", type=str, default="", help="Comment for generated HTML")
args = parser.parse_args()
version = args.version
if "\\" in version or "/" in version:
raise ValueError("Version string must not contain slashes or backslashes.")
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
with open(args.current_objects) as f:
current_objects: dict[str, list[dict[str, Any]]] = json.load(f)
wheel_files = [
item["Key"].split("/")[-1] for item in current_objects.get("Contents", []) if item["Key"].endswith(".whl")
]
print(f"Found {len(wheel_files)} wheel files for version {version}: {wheel_files}")
# For release versions, filter to only matching non-dev wheels
PY_VERSION_RE = re.compile(r"^\d+\.\d+\.\d+([a-zA-Z0-9.+-]*)?$")
if PY_VERSION_RE.match(version):
wheel_files = [f for f in wheel_files if version in f and "dev" not in f]
print(f"Non-nightly version detected, wheel files used: {wheel_files}")
else:
print("Nightly version detected, keeping all wheel files.")
wheel_dir = (args.wheel_dir or version).strip().rstrip("/")
wheel_base_dir = Path(output_dir).parent / wheel_dir
index_base_dir = Path(output_dir)
generate_index(
whl_files=wheel_files,
wheel_base_dir=wheel_base_dir,
index_base_dir=index_base_dir,
comment=args.comment.strip(),
)
print(f"Successfully generated index in {output_dir}")
+14 -22
View File
@@ -64,22 +64,15 @@ while true; do
done
echo "--- Pulling container"
## 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"
image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
# 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
# Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
sudo apt-get update && sudo apt-get install -y awscli
# Use safe docker login helper to prevent race conditions
source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
safe_docker_login_ecr_public
# Pull the container from ECR Public Gallery
docker pull "${image_name}"
@@ -94,11 +87,7 @@ HF_CACHE="$(realpath ~)/huggingface"
mkdir -p "${HF_CACHE}"
HF_MOUNT="/root/.cache/huggingface"
if [[ -n "${TEST_COMMAND:-}" ]]; then
commands="$TEST_COMMAND"
else
commands="$@"
fi
commands=$@
echo "Commands:$commands"
PARALLEL_JOB_COUNT=8
@@ -113,7 +102,6 @@ if [[ -z "$render_gid" ]]; then
fi
# check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs.
# TODO: @tjtanaa reenable to run VLLM_ROCM_USE_AITER=1 when AITER is shipped with prebuilt kernels.
if [[ $commands == *"--shard-id="* ]]; then
# assign job count as the number of shards used
commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g')
@@ -130,9 +118,11 @@ if [[ $commands == *"--shard-id="* ]]; then
--rm \
-e MIOPEN_DEBUG_CONV_DIRECT=0 \
-e MIOPEN_DEBUG_CONV_GEMM=0 \
-e VLLM_ROCM_USE_AITER=0 \
-e VLLM_ROCM_USE_AITER=1 \
-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}" \
@@ -163,8 +153,10 @@ else
--rm \
-e MIOPEN_DEBUG_CONV_DIRECT=0 \
-e MIOPEN_DEBUG_CONV_GEMM=0 \
-e VLLM_ROCM_USE_AITER=0 \
-e VLLM_ROCM_USE_AITER=1 \
-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}" \
@@ -1,56 +0,0 @@
#!/bin/bash
# This script build the XPU docker image and run the offline inference inside the container.
set -ex
omni_source_dir=$(git rev-parse --show-toplevel)
base_image_name="xpu/vllm-omni-ci-base:${VLLM_VERSION:?VLLM_VERSION must be set}"
image_name="xpu/vllm-omni-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set}"
container_name="xpu_${BUILDKITE_COMMIT}_$(
tr -dc A-Za-z0-9 </dev/urandom | head -c 10
echo
)"
cd "${omni_source_dir}"
if [ -z "$(docker images -q "${base_image_name}")" ]; then
docker build --target vllm-base -t "${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -f docker/Dockerfile.xpu .
fi
# Try building the docker image
docker build --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -t "${image_name}" -f docker/Dockerfile.xpu .
# Setup cleanup
remove_docker_container() {
docker rm -f "${container_name}" || true
docker image rm -f "${image_name}" || true
docker system prune -f || true
}
trap remove_docker_container EXIT
HF_CACHE="${HF_CACHE:-$(realpath ~)/.cache/huggingface}"
mkdir -p "${HF_CACHE}"
HF_MOUNT="/root/.cache/huggingface"
time timeout -k 30 30m docker run \
--device /dev/dri:/dev/dri \
--net=host \
--ipc=host \
-v /dev/dri/by-path:/dev/dri/by-path \
-v "${HF_CACHE}:${HF_MOUNT}" \
--security-opt seccomp=unconfined \
--entrypoint="" \
-e VLLM_LOGGING_LEVEL \
-e VLLM_OMNI_LOGGING_LEVEL \
-e HF_TOKEN \
-e ZE_AFFINITY_MASK \
--name "${container_name}" \
"${image_name}" \
bash -c '
set -e
echo $ZE_AFFINITY_MASK
pip install tblib==3.1.0
cd /workspace/vllm-omni
pytest -v -s -m "core_model and xpu and B60"
pytest -v -s -m "advanced_model and xpu and B60"
'
@@ -1,33 +0,0 @@
#!/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."
-215
View File
@@ -1,215 +0,0 @@
steps:
- label: "Simple Unit Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
- label: "Diffusion Model Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
- label: "Diffusion Images API LoRA E2E"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
- label: "Diffusion Model CPU offloading Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
# - label: "Audio Generation Model Test"
# agent_pool: mi325_1
# depends_on: amd-build
# mirror_hardwares: [amdproduction]
# grade: Blocking
# commands:
# - export GPU_ARCHS=gfx942
# - export VLLM_LOGGING_LEVEL=DEBUG
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
- label: "Diffusion Cache Backend Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
- label: "Diffusion Sequence Parallelism Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
# merge-only tests
- label: "Diffusion Tensor Parallelism Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
- label: "Diffusion GPU Worker Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
- label: "Benchmark & Engine Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/benchmarks/test_serve_cli.py
EXIT1=\$?
pytest -s -v tests/engine/test_async_omni_engine_abort.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
- label: "Omni Model Test Qwen2-5-Omni"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
- label: "Omni Model Test Qwen3-Omni"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
- label: "Qwen3-TTS CustomVoice E2E Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
'
- label: "Qwen3-TTS Base E2E Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_base.py
'
- label: "Diffusion Image Edit Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
# split Bagel Model Test with H100 (Real Weights) into three tests
- label: "Bagel Text2Img Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ROCM_USE_AITER_RMSNORM=0
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
- label: "Bagel Img2Img Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ROCM_USE_AITER_RMSNORM=0
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
- label: "Bagel Online Serving Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- export VLLM_IMAGE_FETCH_TIMEOUT=60
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ROCM_USE_AITER_RMSNORM=0
- timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
-179
View File
@@ -1,179 +0,0 @@
steps:
- label: "Simple Unit Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
- label: "Diffusion Model Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
- label: "Diffusion Model CPU offloading Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
# - label: "Audio Generation Model Test"
# agent_pool: mi325_1
# depends_on: amd-build
# mirror_hardwares: [amdproduction]
# grade: Blocking
# commands:
# - export GPU_ARCHS=gfx942
# - export VLLM_LOGGING_LEVEL=DEBUG
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
- label: "Diffusion Cache Backend Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
- label: "Diffusion Sequence Parallelism Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
- label: "Diffusion GPU Worker Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
- label: "Benchmark & Engine Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
timeout 30m bash -c '
set +e
pytest -s -v tests/benchmarks/test_serve_cli.py
EXIT1=\$?
pytest -s -v tests/engine/test_async_omni_engine_abort.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
- label: "Omni Model Test Qwen2-5-Omni"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 17m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
- label: "Omni Model Test Qwen3-Omni"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
- timeout 10m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
- label: "Qwen3-TTS E2E Test"
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
- label: "Diffusion Image Edit Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
- label: "Bagel Text2Img Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ROCM_USE_AITER_RMSNORM=0
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model" -k "rocm"
- label: "Bagel Img2Img Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ROCM_USE_AITER_RMSNORM=0
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model" -k "rocm"
- label: "Bagel Online Serving Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- export VLLM_IMAGE_FETCH_TIMEOUT=60
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ROCM_USE_AITER_RMSNORM=0
- timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model" -k "rocm"
+2 -3
View File
@@ -1,7 +1,7 @@
steps:
- label: "Diffusion Model Test"
timeout_in_minutes: 30
timeout_in_minutes: 20
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
@@ -11,7 +11,7 @@ steps:
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
- label: "Diffusion Images API LoRA E2E"
timeout_in_minutes: 30
timeout_in_minutes: 20
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
@@ -101,7 +101,6 @@ steps:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- export VLLM_ROCM_USE_AITER=0
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
+44 -161
View File
@@ -12,12 +12,11 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Model Test"
timeout_in_minutes: 30
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
commands:
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
@@ -30,12 +29,11 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Images API LoRA E2E"
timeout_in_minutes: 30
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
commands:
- pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
@@ -48,7 +46,6 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -67,7 +64,6 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -85,7 +81,6 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -103,7 +98,6 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -112,7 +106,6 @@ steps:
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
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
@@ -123,7 +116,6 @@ steps:
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -142,7 +134,6 @@ steps:
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -161,17 +152,21 @@ steps:
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Engine Test"
- label: "Benchmark & Engine Test"
depends_on: upload-merge-pipeline
commands:
- |
timeout 15m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
set +e
pytest -s -v tests/benchmarks/test_serve_cli.py
EXIT1=$$?
pytest -s -v tests/engine/test_async_omni_engine_abort.py
EXIT2=$$?
exit $$((EXIT1 | EXIT2))
'
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
@@ -183,18 +178,16 @@ steps:
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test"
timeout_in_minutes: 30
timeout_in_minutes: 15
depends_on: upload-merge-pipeline
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
- pytest -s -v 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:
@@ -202,53 +195,6 @@ steps:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Qwen3-TTS CustomVoice E2E Test"
depends_on: upload-merge-pipeline
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
'
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 -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_base.py
'
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:
@@ -259,9 +205,9 @@ steps:
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"
agents:
queue: "mithril-h100-pool"
plugins:
@@ -280,11 +226,6 @@ 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:
@@ -320,11 +261,6 @@ 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:
@@ -336,90 +272,37 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Bagel Model Test with H100 (Real Weights)"
timeout_in_minutes: 60
depends_on: upload-merge-pipeline
commands:
- |
timeout 55m bash -c '
set -e
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_TEST_CLEAN_GPU_MEMORY=1
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory"
pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -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"
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Voxtral-TTS E2E Test"
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -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"
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
# - label: "Bagel Text2Img Model Test with H100"
# timeout_in_minutes: 30
# depends_on: upload-merge-pipeline
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# resources:
# limits:
# nvidia.com/gpu: 1
# volumeMounts:
# - name: devshm
# mountPath: /dev/shm
# - name: hf-cache
# mountPath: /root/.cache/huggingface
# env:
# - name: HF_HOME
# value: /root/.cache/huggingface
# nodeSelector:
# node.kubernetes.io/instance-type: gpu-h100-sxm
# volumes:
# - name: devshm
# emptyDir:
# medium: Memory
# - name: hf-cache
# hostPath:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
+53 -421
View File
@@ -1,8 +1,8 @@
steps:
- label: ":full_moon: Omni Model Test with H100"
- label: "Omni Model Test with H100"
timeout_in_minutes: 90
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
@@ -11,11 +11,7 @@ steps:
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))
exit $$((EXIT1 | EXIT2))
agents:
queue: "mithril-h100-pool"
plugins:
@@ -34,11 +30,6 @@ 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:
@@ -50,417 +41,58 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- 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
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/examples/ -m "advanced_model and L4 and omni" --run-level "advanced_model"
- pytest -s -v 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:
- 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: Diffusion Model Test with H100"
timeout_in_minutes: 120
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion" -m "advanced_model and diffusion and H100" --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 Model Wan22 completed 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"
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 Model Test with L4"
- label: "Omni Model Test"
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model"
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
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: Documentation 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"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/examples/online_serving/test_text_to_image.py tests/examples/offline_inference/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
- 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: Omni Model Perf Test & Test Case Statistics"
key: nightly-omni-performance
timeout_in_minutes: 180
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export BENCHMARK_DIR=tests/dfx/perf/results
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
- buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
- python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: GEBench Accuracy Test with H100"
key: nightly-gebench-accuracy
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: GEdit-Bench Accuracy Test with H100"
key: nightly-gedit-bench-accuracy
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
value: "120"
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Wan22 I2V Accuracy Test with H100"
key: nightly-wan22-i2v-accuracy
timeout_in_minutes: 180
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --run-level advanced_model
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:c392ce21e9cf9ea65c52b866447793db10e0261c
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: Qwen-Image Diffusion Perf Test with H100"
key: nightly-qwen-image-performance
timeout_in_minutes: 300
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export CACHE_DIT_VERSION=1.3.0
- pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
- buildkite-agent artifact upload "tests/dfx/perf/results/benchmark_results_*.json"
- buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 4
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
# No need to run this step for PRs with label nightly-test
- label: ":email: Nightly Collection & Email"
key: nightly-perf-distribution
depends_on:
- nightly-omni-performance
- nightly-qwen-image-performance
if: build.env("NIGHTLY") == "1"
commands:
- pip install openpyxl
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
- export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-qwen-image-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-omni-performance
- python tools/nightly/generate_nightly_perf_excel.py
- python tools/nightly/generate_nightly_perf_html.py
- python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html"
- buildkite-agent artifact upload "tests/dfx/perf/results/*.xlsx"
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
- 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"
agents:
queue: "cpu_queue_premerge"
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
shm-size: "8gb"
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Perf Test"
timeout_in_minutes: 120
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/perf/scripts/run_benchmark.py
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 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
+93 -315
View File
@@ -5,24 +5,6 @@ steps:
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
agents:
queue: "gpu_1_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Voxtral TTS CUDA Unit Test"
timeout_in_minutes: 10
depends_on: upload-ready-pipeline
commands:
- "timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py"
agents:
queue: "gpu_1_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -36,41 +18,7 @@ steps:
- label: "Diffusion Model Test"
depends_on: upload-ready-pipeline
commands:
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
agents:
queue: "gpu_1_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Batching Test"
depends_on: upload-ready-pipeline
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Custom Pipeline Test"
depends_on: upload-ready-pipeline
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
agents:
queue: "gpu_1_queue"
plugins:
@@ -86,7 +34,15 @@ steps:
- label: "Diffusion Model CPU offloading Test"
depends_on: upload-ready-pipeline
commands:
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
EXIT1=$$?
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
EXIT2=$$?
exit $$((EXIT1 | EXIT2))
'
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
@@ -96,7 +52,6 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -113,7 +68,6 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -130,7 +84,6 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -148,7 +101,6 @@ steps:
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -166,18 +118,22 @@ steps:
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Engine Test"
- label: "Benchmark & Engine Test"
depends_on: upload-ready-pipeline
commands:
- |
timeout 15m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
set +e
pytest -s -v tests/benchmarks/test_serve_cli.py
EXIT1=$$?
pytest -s -v tests/engine/test_async_omni_engine_abort.py
EXIT2=$$?
exit $$((EXIT1 | EXIT2))
'
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
@@ -189,7 +145,6 @@ steps:
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -201,7 +156,7 @@ steps:
timeout 17m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "core_model" --run-level "core_model"
pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
'
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
@@ -212,89 +167,47 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test with H100"
depends_on: upload-ready-pipeline
commands:
- |
timeout 20m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "MiMo-Audio E2E Test with H100"
depends_on: upload-ready-pipeline
commands:
- |
timeout 30m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 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: "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: "Qwen3-TTS E2E Test"
depends_on: upload-ready-pipeline
@@ -303,11 +216,10 @@ steps:
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 "core_model" --run-level "core_model"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
'
agents:
queue: "gpu_1_queue"
queue: "gpu_4_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -316,50 +228,9 @@ steps:
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Voxtral-TTS E2E Test"
timeout_in_minutes: 20
depends_on: upload-ready-pipeline
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
# - label: "Diffusion Image Edit Test with H100 (1 GPU)"
# depends_on: upload-ready-pipeline
# commands:
@@ -397,132 +268,39 @@ steps:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
- label: "Bagel Text2Img Model Test with H100"
depends_on: upload-ready-pipeline
commands:
- |
timeout 30m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_TEST_CLEAN_GPU_MEMORY=1
pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model"
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Bagel Img2Img Model Test with H100"
depends_on: upload-ready-pipeline
commands:
- |
timeout 30m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_TEST_CLEAN_GPU_MEMORY=1
pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model"
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Bagel Online Serving Test with H100"
depends_on: upload-ready-pipeline
commands:
- |
timeout 40m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_TEST_CLEAN_GPU_MEMORY=1
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model"
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
# - label: "Bagel Text2Img Model Test with H100"
# depends_on: upload-ready-pipeline
# commands:
# - |
# timeout 30m bash -c '
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
# pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
# '
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# resources:
# limits:
# nvidia.com/gpu: 1
# volumeMounts:
# - name: devshm
# mountPath: /dev/shm
# - name: hf-cache
# mountPath: /root/.cache/huggingface
# env:
# - name: HF_HOME
# value: /root/.cache/huggingface
# nodeSelector:
# node.kubernetes.io/instance-type: gpu-h100-sxm
# volumes:
# - name: devshm
# emptyDir:
# medium: Memory
# - name: hf-cache
# hostPath:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
+6 -12
View File
@@ -3,7 +3,7 @@
Last synced: 2025-12-15
Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests
#}
{% set docker_image_amd = "rocm/vllm-ci:$BUILDKITE_COMMIT-rocm-omni" %}
{% set docker_image_amd = "public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT-rocm-omni" %}
{% set default_working_dir = "/app/vllm-omni" %}
- group: "AMD Tests"
@@ -13,6 +13,7 @@
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"
@@ -21,15 +22,13 @@
retry:
automatic:
- exit_status: -1 # Agent was lost
limit: 2
limit: 1
- exit_status: -10 # Agent was lost
limit: 2
- exit_status: 128 # Git connectivity issues
limit: 2
limit: 1
- exit_status: 1 # Machine occasionally fail
limit: 1
agents:
queue: amd-cpu
queue: cpu_queue_premerge
{% for step in steps %}
{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
@@ -41,14 +40,9 @@
{% else %}
queue: amd_mi325_1
{% endif %}
{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
{% set indented_cmd = cmd_body | replace("\n", "\n ") %}
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh "(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }} ; {{ step.command or (step.commands | join(" && ")) | safe }}"
env:
DOCKER_BUILDKIT: "1"
TEST_COMMAND: |-
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
{{ indented_cmd | safe }}
priority: 100
{% if step.grade and step.grade == "Blocking" %}
soft_fail: false
+19 -12
View File
@@ -74,21 +74,28 @@ body:
If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. For example:
```python
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm import SamplingParams
from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config
omni = Omni(
model="Qwen/Qwen-Image",
stage_configs_path="/path/to/stage_configs.yaml",
# Create stage configurations
ar_config = create_ar_stage_config(
stage_id=0,
model_path="Qwen/Qwen3-0.6B",
input_modalities=["text"],
output_modalities=["text"]
)
prompts = [{"prompt": "A scenic watercolor painting of a lighthouse at sunset"}]
sampling_params_list = [
SamplingParams(max_tokens=1),
OmniDiffusionSamplingParams(num_outputs_per_prompt=1),
]
outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
dit_config = create_dit_stage_config(
stage_id=1,
model_path="stabilityai/stable-diffusion-2-1",
input_modalities=["text"],
output_modalities=["image"]
)
# Initialize OmniLLM
omni = OmniLLM([ar_config, dit_config])
# Generate
outputs = omni.generate(prompt="A scenic watercolor painting of a lighthouse at sunset")
```
If the code is too long (hopefully, it isn't), feel free to put it in a public gist and link it in the issue: https://gist.github.com.
+3 -3
View File
@@ -39,10 +39,10 @@ jobs:
python-version: ["3.11", "3.12"]
steps:
- uses: actions/checkout@v6
- uses: actions/checkout@v4
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v6
uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
@@ -57,7 +57,7 @@ jobs:
bash scripts/build_wheel.sh --python python
- name: Upload Wheel Artifact
uses: actions/upload-artifact@v7
uses: actions/upload-artifact@v4
with:
name: vllm-omni-wheel-py${{ matrix.python-version }}
path: dist/*.whl
+3 -3
View File
@@ -24,12 +24,12 @@ jobs:
github.event.comment.author_association == 'OWNER')
steps:
- name: Checkout repository
uses: actions/checkout@v6.0.2
uses: actions/checkout@v4.2.2
with:
fetch-depth: 0
- name: Set up Python
uses: actions/setup-python@v6.2.0
uses: actions/setup-python@v5.3.0
with:
python-version: '3.11'
cache: 'pip'
@@ -54,7 +54,7 @@ jobs:
- name: Upload review logs
if: always()
uses: actions/upload-artifact@v7.0.0
uses: actions/upload-artifact@v4.4.3
with:
name: pr-review-logs-${{ github.event.issue.number || github.event.pull_request.number }}
path: /tmp/pr_review_*.log
+2 -2
View File
@@ -18,8 +18,8 @@ jobs:
pre-commit:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- uses: actions/setup-python@v6
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: "3.12"
cache: 'pip'
-1
View File
@@ -8,7 +8,6 @@ repos:
# list of supported hooks: https://pre-commit.com/hooks.html
- id: check-yaml
args: ["--unsafe"]
exclude: "examples/online_serving/chart-helm/templates/"
- id: debug-statements
- id: end-of-file-fixer
- id: mixed-line-ending
+3 -5
View File
@@ -1,6 +1,6 @@
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
<img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
</picture>
</p>
@@ -9,16 +9,14 @@ Easy, fast, and cheap omni-modality model serving for everyone
</h3>
<p align="center">
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
</p>
---
*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.
- [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
+1 -16
View File
@@ -47,7 +47,6 @@ If no, check your shell running the ComfyUI process. There may be some error mes
This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
- **Generate Image** for text-to-image and image-to-image tasks
- **Generate Video** for text-to-video and image-to-video tasks
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
- **TTS** and **TTS Voice Clone** for TTS tasks
@@ -60,7 +59,7 @@ To build a simple workflow yourself,
- Drag a generation node onto the canvas.
- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
- For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
- For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
@@ -105,20 +104,6 @@ You can configure per-stage sampling parameters for multi-stage models.
> [!TIP]
> The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
### Text-to-video and image-to-video generation (e.g., Wan)
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video Generation**)
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg">
<img alt="vLLM-Omni video generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg" width=55%>
</picture>
</p>
> [!TIP]
> The node automatically choose text-to-video or image-to-video API endpoints depending on whether you connect an image input or not.
### TTS (e.g., Qwen TTS series)
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni TTS**)
-9
View File
@@ -14,21 +14,17 @@ from .comfyui_vllm_omni.nodes import (
VLLMOmniARSampling,
VLLMOmniDiffusionSampling,
VLLMOmniGenerateImage,
VLLMOmniGenerateVideo,
VLLMOmniQwenTTSParams,
VLLMOmniRemoteLoRA,
VLLMOmniSamplingParamsList,
VLLMOmniTTS,
VLLMOmniUnderstanding,
VLLMOmniVoiceClone,
VLLMOmniWanParams,
)
# A dictionary that contains all nodes you want to export with their names
NODE_CLASS_MAPPINGS = {
# === Generation ===
"VLLMOmniGenerateImage": VLLMOmniGenerateImage,
"VLLMOmniGenerateVideo": VLLMOmniGenerateVideo,
"VLLMOmniUnderstanding": VLLMOmniUnderstanding,
"VLLMOmniTTS": VLLMOmniTTS,
"VLLMOmniVoiceClone": VLLMOmniVoiceClone,
@@ -36,16 +32,13 @@ NODE_CLASS_MAPPINGS = {
"VLLMOmniARSampling": VLLMOmniARSampling,
"VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
"VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList,
"VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA,
"VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
"VLLMOmniWanParams": VLLMOmniWanParams,
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
# === Generation ===
"VLLMOmniGenerateImage": "Generate Image",
"VLLMOmniGenerateVideo": "Generate Video",
"VLLMOmniUnderstanding": "Multimodality Understanding",
"VLLMOmniTTS": "TTS (Text to Speech)",
"VLLMOmniVoiceClone": "TTS Voice Cloning",
@@ -53,9 +46,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"VLLMOmniARSampling": "AR Sampling Params",
"VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
"VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List",
"VLLMOmniRemoteLoRA": "LoRA",
"VLLMOmniQwenTTSParams": "Qwen TTS Params",
"VLLMOmniWanParams": "Wan Video Params",
}
WEB_DIRECTORY = "./web"
@@ -11,7 +11,6 @@ from .utils.types import (
AutoregressionSamplingParams,
DiffusionSamplingParams,
QwenTTSModelSpecificParams,
WanModelSpecificParams,
)
from .utils.validators import (
add_sampling_parameters_to_stage,
@@ -57,7 +56,6 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
# "video": ("VIDEO",),
# "audio": ("AUDIO",),
"sampling_params": ("SAMPLING_PARAMS",),
"lora": ("REMOTE_LORA",),
},
}
@@ -78,11 +76,9 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
audio: AudioInput | None = None, # Hidden & unused
video: VideoInput | None = None, # Hidden & unused
sampling_params: dict | list[dict] | None = None,
lora: dict | None = None,
**kwargs,
):
if kwargs:
logger.info("Uncaught kwargs: %s", kwargs)
logger.info("Uncaught kwargs: %s", kwargs)
logger.debug("Got sampling params: %s", sampling_params)
validate_model_and_sampling_params_types(model, sampling_params)
if image is None and mask is not None:
@@ -109,7 +105,6 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
height=height,
negative_prompt=negative_prompt,
sampling_params=sampling_params,
lora=lora,
)
return (output,)
elif image is not None and audio is None and video is None:
@@ -124,7 +119,6 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
negative_prompt=negative_prompt,
mask=mask,
sampling_params=sampling_params,
lora=lora,
)
return (output,)
@@ -142,91 +136,11 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
audio=audio,
video=video,
sampling_params=sampling_params,
lora=lora,
)
return (output,)
class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"url": ("STRING", {"default": "http://localhost:8000/v1"}),
"model": ("STRING", {"default": "Wan-AI/Wan2.2-T2V-A14B-Diffusers"}),
"prompt": ("STRING", {"multiline": True}),
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
"width": ("INT", {"default": 832, "min": 1}),
"height": ("INT", {"default": 480, "min": 1}),
"fps": ("INT", {"default": 16, "min": 1}),
"num_frames": ("INT", {"default": 41, "min": 1}),
},
"optional": {
"image": ("IMAGE",),
"sampling_params": ("SAMPLING_PARAMS",),
"lora": ("REMOTE_LORA",),
"model_params": ("VIDEO_PARAMS",),
},
}
RETURN_TYPES = ("VIDEO",)
RETURN_NAMES = ("video",)
FUNCTION = "generate"
async def generate(
self,
url: str,
model: str,
prompt: str,
width: int,
height: int,
fps: int,
num_frames: int,
negative_prompt: str | None = None,
image: torch.Tensor | None = None,
sampling_params: dict | list[dict] | None = None,
model_params: dict | None = None,
lora: dict | None = None,
**kwargs,
):
if kwargs:
logger.info("Uncaught kwargs: %s", kwargs)
logger.debug("Got sampling params: %s", sampling_params)
logger.debug("Got model params: %s", model_params)
validate_model_and_sampling_params_types(model, sampling_params)
# Currently, all video generation models are single-stage diffusion models
if isinstance(sampling_params, list):
if len(sampling_params) != 1:
raise ValueError(
"Video generation expects a single sampling params group. "
"Please provide one Diffusion sampling node."
)
sampling_params = sampling_params[0]
if sampling_params is not None:
sampling_params.pop("type", None) # internal fields
if model_params is not None:
model_params.pop("type", None) # internal fields
client = VLLMOmniClient(url)
output = await client.generate_video(
model=model,
prompt=prompt,
image=image, # image present => i2v, absent => t2v
width=width,
height=height,
num_frames=num_frames,
fps=fps,
negative_prompt=negative_prompt,
sampling_params=sampling_params,
lora=lora,
model_params=model_params,
)
return (output,)
class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
@classmethod
def INPUT_TYPES(cls):
@@ -274,8 +188,7 @@ class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
use_audio_in_video: bool = True,
**kwargs,
) -> tuple[str, AudioInput]:
if kwargs:
logger.info("Uncaught kwargs: %s", kwargs)
logger.info("Uncaught kwargs: %s", kwargs)
logger.debug("Got sampling params: %s", sampling_params)
validate_model_and_sampling_params_types(model, sampling_params)
@@ -634,52 +547,6 @@ class VLLMOmniSamplingParamsList:
return (params,)
class VLLMOmniRemoteLoRA:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"local_path": ("STRING", {"default": ""}),
"name": ("STRING", {"default": ""}),
"scale": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1},
),
"int_id": (
"INT",
{
"default": 0,
"min": 0,
"step": 1,
"tooltip": "0 means it is not set and the server can derive it.",
},
),
}
}
RETURN_TYPES = ("REMOTE_LORA",)
RETURN_NAMES = ("lora",)
FUNCTION = "get_lora"
CATEGORY = "vLLM-Omni"
@classmethod
def VALIDATE_INPUTS(cls, local_path, name) -> str | Literal[True]:
if not local_path.strip() or not name.strip():
return "Both local_path and name must be provided."
return True
def get_lora(self, local_path: str, name: str, scale: float, int_id: int):
local_path = local_path.strip()
name = name.strip()
lora = {
"local_path": local_path or None,
"name": name or None,
"scale": float(scale),
"int_id": int(int_id) if int_id > 0 else None,
}
return (lora,)
class VLLMOmniQwenTTSParams:
@classmethod
def INPUT_TYPES(cls):
@@ -705,32 +572,3 @@ class VLLMOmniQwenTTSParams:
def get_params(self, **kwargs):
return (QwenTTSModelSpecificParams(kwargs),)
class VLLMOmniWanParams:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"guidance_scale_2": (
"FLOAT",
{"default": 4.0, "min": 0.0, "max": 20.0, "step": 0.1},
),
"boundary_ratio": (
"FLOAT",
{"default": 0.875, "min": 0.0, "max": 1.0, "step": 0.001},
),
"flow_shift": (
"FLOAT",
{"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1},
),
}
}
RETURN_TYPES = ("VIDEO_PARAMS",)
RETURN_NAMES = ("Wan video params",)
FUNCTION = "get_params"
CATEGORY = "vLLM-Omni/Video Params"
def get_params(self, **kwargs):
return (WanModelSpecificParams(kwargs),)
@@ -6,8 +6,6 @@ The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@do
Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
"""
import asyncio
import json
from typing import Any
import aiohttp
@@ -20,7 +18,6 @@ from .format import (
base64_to_audio,
base64_to_image_tensor,
bytes_to_audio,
bytes_to_video,
image_tensor_to_base64,
image_tensor_to_png_bytes,
video_to_base64,
@@ -32,43 +29,10 @@ from .types import AudioFormat
logger = get_logger(__name__)
async def url_json(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> dict[str, Any]:
try:
async with getattr(session, verb)(url, **kwargs) as response:
if not response.ok:
error_text = await response.text()
raise (ValueError if response.status < 500 else RuntimeError)(
f"vLLM-Omni API returned status {response.status}: {error_text}"
)
try:
return await response.json()
except aiohttp.ContentTypeError as e:
raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
except aiohttp.ClientError as e:
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
async def url_bytes(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> bytes:
try:
async with getattr(session, verb)(url, **kwargs) as response:
if not response.ok:
error_text = await response.text()
raise (ValueError if response.status < 500 else RuntimeError)(
f"vLLM-Omni API returned status {response.status}: {error_text}"
)
return await response.read()
except aiohttp.ClientError as e:
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
class VLLMOmniClient:
def __init__(
self, base_url: str, timeout: float | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
):
def __init__(self, base_url: str, timeout: float = 300.0):
self.base_url = base_url
self.timeout = aiohttp.ClientTimeout(total=timeout)
self.poll_interval = poll_interval
self.max_poll_duration = max_poll_duration
async def generate_image(
self,
@@ -79,13 +43,12 @@ class VLLMOmniClient:
height: int,
negative_prompt: str | None = None,
sampling_params: dict | None = None,
lora: dict | None = None,
) -> torch.Tensor:
"""Run text-to-image generation via DALLE API"""
await self._check_model_exist(model)
size = f"{width}x{height}"
payload: dict[str, Any] = {
payload = {
"model": model,
"prompt": prompt,
"size": size,
@@ -95,8 +58,6 @@ class VLLMOmniClient:
payload["negative_prompt"] = negative_prompt
if sampling_params is not None:
payload.update(sampling_params)
if lora is not None:
payload["lora"] = lora
logger.debug("img gen payload: %s", payload)
url = self.base_url + "/images/generations"
@@ -149,7 +110,6 @@ class VLLMOmniClient:
negative_prompt: str | None = None,
mask: torch.Tensor | None = None,
sampling_params: dict | None = None,
lora: dict | None = None,
) -> torch.Tensor:
"""Run image editing via DALLE API"""
await self._check_model_exist(model)
@@ -171,8 +131,6 @@ class VLLMOmniClient:
if sampling_params is not None:
for k, v in sampling_params.items():
form.add_field(k, str(v))
if lora is not None:
form.add_field("lora", json.dumps(lora, ensure_ascii=False))
if mask is not None:
mask_filename = "mask.png"
form.add_field(
@@ -225,7 +183,6 @@ class VLLMOmniClient:
audio: AudioInput | None = None,
video: VideoInput | None = None,
sampling_params: dict | list[dict] | None = None,
lora: dict | None = None,
) -> torch.Tensor:
payload = VLLMOmniClient._prepare_chat_completion_messages(
model=model,
@@ -236,8 +193,6 @@ class VLLMOmniClient:
video=video,
sampling_params=sampling_params,
modalities=["image"],
# === below are additional `extra_body` fields, handled by **kwargs ===
lora=lora,
)
choices = await self._generate_base_chat_completion(model, payload)
@@ -251,87 +206,6 @@ class VLLMOmniClient:
return torch.stack(image_tensors, dim=0)
async def generate_video(
self,
*,
model: str,
prompt: str,
width: int,
height: int,
num_frames: int,
fps: int,
negative_prompt: str | None = None,
image: torch.Tensor | None = None,
sampling_params: dict | None = None,
model_params: dict | None = None,
lora: dict | None = None,
**extra_body,
) -> VideoInput:
form = aiohttp.FormData()
form.add_field("model", model)
form.add_field("prompt", prompt)
form.add_field("width", str(width))
form.add_field("height", str(height))
form.add_field("num_frames", str(num_frames))
form.add_field("fps", str(fps))
if negative_prompt:
form.add_field("negative_prompt", negative_prompt)
if sampling_params is not None:
for k, v in sampling_params.items():
form.add_field(k, str(v))
if model_params is not None:
for k, v in model_params.items():
form.add_field(k, str(v))
if lora is not None:
form.add_field("lora", json.dumps(lora, ensure_ascii=False))
if extra_body:
form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False))
if image is not None:
image_filename = "image.png" # Required for multipart form
form.add_field(
"input_reference",
image_tensor_to_png_bytes(image, image_filename),
filename=image_filename,
content_type="image/png",
)
async with aiohttp.ClientSession(timeout=self.timeout) as session:
# Start the video generation job
url = f"{self.base_url}/videos"
data = await url_json(session, url, "post", data=form)
if (job_id := data.get("id", None)) is None:
raise RuntimeError("API response missing job 'id' field - expected OpenAI compliant format")
if (job_status := data.get("status", None)) is None:
raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
# Poll for video generation job completion
deadline = asyncio.get_running_loop().time() + self.max_poll_duration
url = f"{self.base_url}/videos/{job_id}"
while job_status not in {"completed", "failed"}:
await asyncio.sleep(self.poll_interval)
data = await url_json(session, url)
if (job_status := data.get("status", None)) is None:
raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
if asyncio.get_running_loop().time() >= deadline:
raise RuntimeError(f"Timed out waiting for video job {job_id} to complete")
if job_status == "failed":
raise RuntimeError(f"Video job failed: {data}")
# Retrieve completed content
video_bytes = await url_bytes(session, f"{url}/content")
# Decode video and make a best effort at cleaning up server resources
try:
return bytes_to_video(video_bytes)
finally:
try:
await url_json(session, url, "delete")
except Exception as exc:
logger.warning("Failed to clean up video job %s: %s", job_id, exc)
async def generate_understanding_chat_completion(
self,
*,
@@ -6,17 +6,12 @@ Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed und
import base64
import mimetypes
from fractions import Fraction
from io import BytesIO
import av
import numpy as np
import torch
from av.audio.frame import AudioFrame
from av.audio.resampler import AudioResampler
from av.video.frame import VideoFrame
from comfy_api.input import AudioInput, VideoInput
from comfy_api.latest import InputImpl, Types
from comfy_extras import nodes_audio
from PIL import Image
@@ -143,75 +138,6 @@ def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
return f"data:{mime_type};base64,{base64_str}"
def bytes_to_video(video_bytes: bytes) -> VideoInput:
video_buffer = BytesIO(video_bytes)
try:
with av.open(video_buffer, mode="r") as container:
video_stream = next((s for s in container.streams if s.type == "video"), None)
if video_stream is None:
raise ValueError("No video stream found in decoded payload.")
frames = []
for frame in container.decode(video_stream):
if not isinstance(frame, VideoFrame):
continue
image = frame.to_ndarray(format="rgb24")
frames.append(torch.from_numpy(image).float() / 255.0)
if len(frames) == 0:
raise ValueError("No video frames found in decoded payload.")
images = torch.stack(frames, dim=0)
frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
audio: AudioInput | None = None
if len(container.streams.audio):
audio_stream = container.streams.audio[-1]
audio_frames = []
resampler = AudioResampler(format="fltp")
for frame in container.decode(audio_stream):
if not isinstance(frame, AudioFrame):
continue
resampled = resampler.resample(frame)
if not isinstance(resampled, list):
resampled = [resampled]
for audio_frame in resampled:
if audio_frame is not None:
audio_frames.append(audio_frame.to_ndarray())
if len(audio_frames) > 0:
audio_data = np.concatenate(audio_frames, axis=1)
sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1
audio = {
"waveform": torch.from_numpy(audio_data).unsqueeze(0),
"sample_rate": sample_rate,
}
components = Types.VideoComponents(
images=images,
frame_rate=frame_rate,
audio=audio,
metadata=container.metadata if container.metadata else None,
)
except Exception as e:
raise RuntimeError(f"Failed to decode video: {e}")
return InputImpl.VideoFromComponents(components)
def base64_to_video(base64_str: str) -> VideoInput:
if base64_str.startswith("data:video"):
_, base64_str = base64_str.split(",", 1)
try:
video_bytes = base64.b64decode(base64_str)
except Exception as e:
raise ValueError(f"Invalid base64 string: {e}")
return bytes_to_video(video_bytes)
def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO:
waveform = audio["waveform"][0] # Shape: (C, T)
sample_rate = audio["sample_rate"]
@@ -23,10 +23,6 @@ class QwenTTSModelSpecificParams(dict):
pass
class WanModelSpecificParams(dict):
pass
class ModelMode(Enum):
IMAGE_GENERATION = auto()
VIDEO_GENERATION = auto()
@@ -16,7 +16,7 @@ def validate_model_and_sampling_params_types(
# Skip if no spec or no sampling params
pipeline_spec, _ = lookup_model_spec(model_name)
if pipeline_spec is None:
logger.info(f"skipping sampling params check because spec for {model_name} is not found")
logger.info("skipping sampling params check because spec is not found")
return
if sampling_param_list is None:
return
@@ -70,17 +70,13 @@ def add_sampling_parameters_to_stage(
stages = pipeline_spec["stages"]
if isinstance(sampling_param_list, dict):
sampling_param_list = sampling_param_list.__class__(sampling_param_list)
sampling_param_list.update(params_to_add)
elif sampling_param_list is None:
sampling_param_list = params_to_add.copy()
else:
for i, stage in enumerate(stages):
if stage == stage_type:
stage_param = sampling_param_list[i]
stage_param = stage_param.__class__(stage_param)
stage_param.update(params_to_add)
sampling_param_list[i] = stage_param
sampling_param_list[i].update(params_to_add)
return sampling_param_list
Binary file not shown.

Before

Width:  |  Height:  |  Size: 129 KiB

@@ -1,513 +0,0 @@
{
"id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21",
"revision": 0,
"last_node_id": 14,
"last_link_id": 28,
"nodes": [
{
"id": 10,
"type": "MarkdownNote",
"pos": [
227.99306819575278,
-231.24143306069843
],
"size": [
240.3326478922474,
136.3505820791881
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [],
"title": "Note: Task and Input",
"properties": {},
"widgets_values": [
"vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload depending on whether an input image is present or not."
],
"color": "#432",
"bgcolor": "#000"
},
{
"id": 4,
"type": "LoadImage",
"pos": [
496.31859627609606,
-229.71277089860084
],
"size": [
270,
314
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"type": "COMBO",
"widget": {
"name": "image"
},
"link": null
},
{
"localized_name": "choose file to upload",
"name": "upload",
"type": "IMAGEUPLOAD",
"widget": {
"name": "upload"
},
"link": null
}
],
"outputs": [
{
"localized_name": "IMAGE",
"name": "IMAGE",
"type": "IMAGE",
"links": [
25
]
},
{
"localized_name": "MASK",
"name": "MASK",
"type": "MASK",
"links": null
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"cute-cat.jpg",
"image"
]
},
{
"id": 8,
"type": "VLLMOmniDiffusionSampling",
"pos": [
465.5140218220927,
148.80072646828987
],
"size": [
284.205078125,
226
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"localized_name": "n",
"name": "n",
"type": "INT",
"widget": {
"name": "n"
},
"link": null
},
{
"localized_name": "num_inference_steps",
"name": "num_inference_steps",
"type": "INT",
"widget": {
"name": "num_inference_steps"
},
"link": null
},
{
"localized_name": "guidance_scale",
"name": "guidance_scale",
"type": "FLOAT",
"widget": {
"name": "guidance_scale"
},
"link": null
},
{
"localized_name": "true_cfg_scale",
"name": "true_cfg_scale",
"type": "FLOAT",
"widget": {
"name": "true_cfg_scale"
},
"link": null
},
{
"localized_name": "vae_use_slicing",
"name": "vae_use_slicing",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_slicing"
},
"link": null
},
{
"localized_name": "vae_use_tiling",
"name": "vae_use_tiling",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_tiling"
},
"link": null
},
{
"localized_name": "seed",
"name": "seed",
"type": "INT",
"widget": {
"name": "seed"
},
"link": null
}
],
"outputs": [
{
"localized_name": "diffusion sampling params",
"name": "diffusion sampling params",
"type": "SAMPLING_PARAMS",
"links": [
27
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniDiffusionSampling"
},
"widgets_values": [
1,
50,
1,
1,
false,
false,
1992,
"randomize"
]
},
{
"id": 14,
"type": "VLLMOmniWanParams",
"pos": [
518.6005521845747,
428.4455663719664
],
"size": [
270,
106
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"localized_name": "guidance_scale_2",
"name": "guidance_scale_2",
"type": "FLOAT",
"widget": {
"name": "guidance_scale_2"
},
"link": null
},
{
"localized_name": "boundary_ratio",
"name": "boundary_ratio",
"type": "FLOAT",
"widget": {
"name": "boundary_ratio"
},
"link": null
},
{
"localized_name": "flow_shift",
"name": "flow_shift",
"type": "FLOAT",
"widget": {
"name": "flow_shift"
},
"link": null
}
],
"outputs": [
{
"localized_name": "Wan video params",
"name": "Wan video params",
"type": "VIDEO_PARAMS",
"links": [
28
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniWanParams"
},
"widgets_values": [
4,
0.875,
5
]
},
{
"id": 12,
"type": "VLLMOmniGenerateVideo",
"pos": [
827.2566336087867,
-73.77449831827545
],
"size": [
400,
346
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"shape": 7,
"type": "IMAGE",
"link": 25
},
{
"localized_name": "sampling_params",
"name": "sampling_params",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 27
},
{
"localized_name": "lora",
"name": "lora",
"shape": 7,
"type": "REMOTE_LORA",
"link": null
},
{
"localized_name": "model_params",
"name": "model_params",
"shape": 7,
"type": "VIDEO_PARAMS",
"link": 28
},
{
"localized_name": "url",
"name": "url",
"type": "STRING",
"widget": {
"name": "url"
},
"link": null
},
{
"localized_name": "model",
"name": "model",
"type": "STRING",
"widget": {
"name": "model"
},
"link": null
},
{
"localized_name": "prompt",
"name": "prompt",
"type": "STRING",
"widget": {
"name": "prompt"
},
"link": null
},
{
"localized_name": "negative_prompt",
"name": "negative_prompt",
"type": "STRING",
"widget": {
"name": "negative_prompt"
},
"link": null
},
{
"localized_name": "width",
"name": "width",
"type": "INT",
"widget": {
"name": "width"
},
"link": null
},
{
"localized_name": "height",
"name": "height",
"type": "INT",
"widget": {
"name": "height"
},
"link": null
},
{
"localized_name": "fps",
"name": "fps",
"type": "INT",
"widget": {
"name": "fps"
},
"link": null
},
{
"localized_name": "num_frames",
"name": "num_frames",
"type": "INT",
"widget": {
"name": "num_frames"
},
"link": null
}
],
"outputs": [
{
"localized_name": "video",
"name": "video",
"type": "VIDEO",
"links": [
26
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniGenerateVideo"
},
"widgets_values": [
"http://localhost:8000/v1",
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
"Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.",
"Cartoonish",
900,
1200,
12,
41
]
},
{
"id": 13,
"type": "SaveVideo",
"pos": [
1257.5441369501657,
-44.12956394419863
],
"size": [
374,
372.34896826553796
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"localized_name": "video",
"name": "video",
"type": "VIDEO",
"link": 26
},
{
"localized_name": "filename_prefix",
"name": "filename_prefix",
"type": "STRING",
"widget": {
"name": "filename_prefix"
},
"link": null
},
{
"localized_name": "format",
"name": "format",
"type": "COMBO",
"widget": {
"name": "format"
},
"link": null
},
{
"localized_name": "codec",
"name": "codec",
"type": "COMBO",
"widget": {
"name": "codec"
},
"link": null
}
],
"outputs": [],
"properties": {},
"widgets_values": [
"video/ComfyUI",
"auto",
"auto"
]
}
],
"links": [
[
25,
4,
0,
12,
0,
"IMAGE"
],
[
26,
12,
0,
13,
0,
"VIDEO"
],
[
27,
8,
0,
12,
1,
"SAMPLING_PARAMS"
],
[
28,
14,
0,
12,
3,
"VIDEO_PARAMS"
]
],
"groups": [
{
"id": 1,
"title": "Input",
"bounding": [
213.1706010087147,
-313.20095750667554,
560.8246471437027,
407.4246532472182
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
}
],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 1.1469075819486894,
"offset": [
101.7632950847089,
492.29121889347806
]
}
},
"version": 0.4
}
-2
View File
@@ -15,8 +15,6 @@ This document explains the benchmark architecture across all benchmark assets in
## Directory layout
- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
- `benchmarks/accuracy/` — accuracy benchmark integrations that adapt external
benchmark suites to vLLM-Omni serving and evaluation flows.
- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
## Reference workflows
-1
View File
@@ -1 +0,0 @@
"""Benchmark helpers and runnable benchmark entrypoints."""
-27
View File
@@ -1,27 +0,0 @@
# Accuracy Benchmarks
This directory hosts accuracy benchmark integrations that run entirely through a
local `vllm-omni serve` deployment.
Current integrations:
- `text_to_image/`: GEBench generation + local judge scoring flow.
- `image_to_image/`: GEdit-Bench generation + local VIEScore-style scoring flow.
Design notes:
- Generation is executed through the OpenAI-compatible endpoints exposed by
`vllm-omni serve`.
- Evaluation is also executed through a local OpenAI-compatible judge model
served by `vllm-omni`.
- Both generation and judge requests accept either `http://host:port` or
`http://host:port/v1`.
- Output directory layout intentionally stays close to the upstream repos.
Test guidance:
- Local static/self-checks live in `tests/benchmarks/test_accuracy_bench_utils.py`.
- End-to-end generation/evaluation should be validated in a remote GPU
environment. In the current repo marker system there is `L4` but no `L5`
marker, so benchmark smoke tests should be wired as `advanced_model +
benchmark + L4` when GPU capacity is available.
-1
View File
@@ -1 +0,0 @@
"""Accuracy benchmark integrations for vLLM-Omni."""
-199
View File
@@ -1,199 +0,0 @@
from __future__ import annotations
import base64
import io
import json
from pathlib import Path
from typing import Any
import requests
from PIL import Image
def ensure_dir(path: Path) -> Path:
path.mkdir(parents=True, exist_ok=True)
return path
def load_json(path: Path) -> dict[str, Any]:
with path.open("r", encoding="utf-8") as handle:
return json.load(handle)
def write_json(path: Path, payload: dict[str, Any]) -> None:
ensure_dir(path.parent)
with path.open("w", encoding="utf-8") as handle:
json.dump(payload, handle, indent=2, ensure_ascii=False)
def save_image(path: Path, image: Image.Image) -> None:
ensure_dir(path.parent)
image.save(path)
def find_first_image(folder: Path, stem: str | None = None) -> Path | None:
patterns = [f"{stem}.*"] if stem else ["*.png", "*.jpg", "*.jpeg", "*.webp"]
for pattern in patterns:
for candidate in sorted(folder.glob(pattern)):
if candidate.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"}:
return candidate
return None
def extract_json_object(raw_text: str) -> dict[str, Any]:
raw_text = raw_text.strip()
delimiter = "||V^=^V||"
if raw_text.count(delimiter) >= 2:
start = raw_text.find(delimiter) + len(delimiter)
end = raw_text.rfind(delimiter)
raw_text = raw_text[start:end].strip()
start = raw_text.find("{")
end = raw_text.rfind("}")
if start == -1 or end == -1 or end < start:
raise ValueError(f"Could not find JSON object in: {raw_text[:200]}")
return json.loads(raw_text[start : end + 1])
def build_openai_url(base_url: str, api_path: str) -> str:
base = base_url.rstrip("/")
normalized_path = api_path if api_path.startswith("/") else f"/{api_path}"
if base.endswith(normalized_path):
return base
if base.endswith("/v1"):
return f"{base}{normalized_path}"
return f"{base}/v1{normalized_path}"
def pil_to_base64(image: Image.Image, image_format: str = "PNG") -> str:
buffer = io.BytesIO()
image.save(buffer, format=image_format)
return base64.b64encode(buffer.getvalue()).decode("utf-8")
def pil_to_data_url(image: Image.Image, image_format: str = "PNG") -> str:
return f"data:image/{image_format.lower()};base64,{pil_to_base64(image, image_format=image_format)}"
def decode_base64_image(encoded: str) -> Image.Image:
image = Image.open(io.BytesIO(base64.b64decode(encoded)))
image.load()
return image.convert("RGB")
def pil_to_png_bytes(image: Image.Image) -> bytes:
buffer = io.BytesIO()
image.save(buffer, format="PNG")
return buffer.getvalue()
class VllmOmniImageClient:
"""Thin OpenAI-compatible image client for vLLM-Omni serving."""
def __init__(self, base_url: str, api_key: str = "EMPTY", timeout: int = 600):
self.base_url = base_url.rstrip("/")
self.api_key = api_key
self.timeout = timeout
@property
def _headers(self) -> dict[str, str]:
return {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
}
def generate_text_to_image(
self,
*,
model: str,
prompt: str,
width: int,
height: int,
num_inference_steps: int = 20,
guidance_scale: float | None = None,
seed: int | None = None,
output_compression: int | None = None,
) -> Image.Image:
payload: dict[str, Any] = {
"model": model,
"prompt": prompt,
"n": 1,
"size": f"{width}x{height}",
"response_format": "b64_json",
"num_inference_steps": num_inference_steps,
}
if guidance_scale is not None:
payload["guidance_scale"] = guidance_scale
if seed is not None:
payload["seed"] = seed
if output_compression is not None:
payload["output_compression"] = output_compression
response = requests.post(
build_openai_url(self.base_url, "/images/generations"),
json=payload,
headers=self._headers,
timeout=self.timeout,
)
response.raise_for_status()
return decode_base64_image(response.json()["data"][0]["b64_json"])
def generate_image_edit(
self,
*,
model: str,
prompt: str,
images: Image.Image | list[Image.Image],
width: int,
height: int,
num_inference_steps: int = 20,
guidance_scale: float | None = None,
seed: int | None = None,
negative_prompt: str | None = None,
output_compression: int | None = None,
) -> Image.Image:
if not isinstance(images, list):
images = [images]
data: dict[str, Any] = {
"model": model,
"prompt": prompt,
"n": 1,
"size": f"{width}x{height}",
"response_format": "b64_json",
"num_inference_steps": str(num_inference_steps),
}
if guidance_scale is not None:
data["guidance_scale"] = str(guidance_scale)
if seed is not None:
data["seed"] = str(seed)
if negative_prompt:
data["negative_prompt"] = negative_prompt
if output_compression is not None:
data["output_compression"] = str(output_compression)
files = [
(
"image[]" if len(images) > 1 else "image",
(f"image_{index}.png", pil_to_png_bytes(image), "image/png"),
)
for index, image in enumerate(images)
]
edit_paths = ["/images/edits", "/images/edit"]
last_response: requests.Response | None = None
for api_path in edit_paths:
response = requests.post(
build_openai_url(self.base_url, api_path),
data=data,
files=files,
headers={"Authorization": f"Bearer {self.api_key}"},
timeout=self.timeout,
)
last_response = response
if response.status_code != 404:
response.raise_for_status()
return decode_base64_image(response.json()["data"][0]["b64_json"])
assert last_response is not None
last_response.raise_for_status()
raise ValueError("No image payload returned from image edit endpoint")
@@ -1,103 +0,0 @@
# GEdit-Bench on vLLM-Omni
This integration adapts the upstream `stepfun-ai/Step1X-Edit/GEdit-Bench`
evaluation flow into `vllm-omni/benchmarks/accuracy/image_to_image`.
Upstream mapping:
- `run_gedit_score.py` -> `run_gedit_bench.py evaluate`
- `calculate_statistics.py` -> `run_gedit_bench.py summarize`
- upstream output layout under `results/<method>/fullset/...` is preserved
What changed:
- The upstream repo mainly ships evaluation scripts. This integration adds a
generation runner that uses the local `vllm-omni` OpenAI-compatible
`/v1/images/edits` endpoint to produce benchmark outputs in the expected
directory structure.
- The evaluator keeps the same VIEScore-style decomposition:
- `semantics_score`
- `quality_score`
- `overall_score = sqrt(semantics_score * quality_score)`
- Judge calls are routed to a local OpenAI-compatible model served by
`vllm-omni`, not a remote provider.
Dataset:
- Default `--dataset-ref` is `stepfun-ai/GEdit-Bench`
- You can also pass a local dataset directory previously saved with
Hugging Face `datasets`
Example usage:
```bash
python benchmarks/accuracy/image_to_image/run_gedit_bench.py generate \
--output-root benchmarks/accuracy/image_to_image/results \
--base-url http://127.0.0.1:8000 \
--model Qwen/Qwen-Image-Edit \
--model-name qwen_image_edit \
--dataset-ref stepfun-ai/GEdit-Bench \
--task-type all \
--instruction-language en
```
```bash
python benchmarks/accuracy/image_to_image/run_gedit_bench.py evaluate \
--output-root benchmarks/accuracy/image_to_image/results \
--model-name qwen_image_edit \
--save-dir benchmarks/accuracy/image_to_image/scores \
--dataset-ref stepfun-ai/GEdit-Bench \
--judge-base-url http://127.0.0.1:8000 \
--judge-model Qwen/Qwen2.5-VL-7B-Instruct \
--judge-api-key EMPTY
```
```bash
python benchmarks/accuracy/image_to_image/run_gedit_bench.py summarize \
--csv-path benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_vie_score.csv \
--language en
```
Example summary output:
```json
{
"language": "all",
"languages": {
"en": {
"overall": {"count": 110, "Q_SC": 6.58, "Q_PQ": 5.89, "Q_O": 5.86},
"intersection": {"count": 78, "Q_SC": 6.50, "Q_PQ": 5.66, "Q_O": 5.65}
},
"cn": {
"overall": {"count": 110, "Q_SC": 6.90, "Q_PQ": 5.78, "Q_O": 6.11},
"intersection": {"count": 63, "Q_SC": 7.22, "Q_PQ": 5.59, "Q_O": 6.28}
}
}
}
```
Example generated images to inspect:
- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/background_change/en/<sample>.png`
- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/text_change/en/<sample>.png`
- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/subject-replace/cn/<sample>.png`
Example score artifacts to inspect together with the images:
- `benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_vie_score.csv`
- `benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_summary.json`
What to expect:
- `Q_SC` measures instruction following and content preservation.
- `Q_PQ` measures image naturalness and artifact quality.
- `Q_O` is the combined overall score; higher is better.
- `overall.count` is the number of evaluated samples for that language, while `intersection.count` is the subset with `intersection_exist == True`.
Notes:
- This flow requires the optional Hugging Face `datasets` package.
- `generate` writes `generation_manifest.json` with local output coverage.
- The current repo marker set exposes `L4` but not `L5`, so if you promote an
end-to-end smoke test into CI, use the existing `advanced_model`, `benchmark`,
and `L4` markers or introduce a new repo-wide marker explicitly first.
@@ -1 +0,0 @@
"""GEdit-Bench integration for vLLM-Omni."""
@@ -1,718 +0,0 @@
from __future__ import annotations
import argparse
import csv
import json
import logging
import math
import statistics
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from io import BytesIO
from pathlib import Path
from typing import Any
import requests
from PIL import Image
from tqdm.auto import tqdm
from benchmarks.accuracy.common import (
VllmOmniImageClient,
build_openai_url,
extract_json_object,
write_json,
)
GROUPS = [
"background_change",
"color_alter",
"material_alter",
"motion_change",
"ps_human",
"style_change",
"subject-add",
"subject-remove",
"subject-replace",
"text_change",
"tone_transfer",
]
DEFAULT_SAMPLES_PER_GROUP = 10
logger = logging.getLogger(__name__)
def infer_model_name(model: str) -> str:
normalized = model.rstrip("/\\")
name = Path(normalized).name
return name or normalized
def resolve_model_name(*, model_name: str | None, model: str | None = None, output_root: Path | None = None) -> str:
if model_name:
return model_name
if model:
return infer_model_name(model)
if output_root is not None:
candidates = sorted(path.name for path in output_root.iterdir() if path.is_dir())
if len(candidates) == 1:
return candidates[0]
if not candidates:
raise ValueError(f"Could not infer model-name from empty output root: {output_root}")
raise ValueError(
f"Could not infer model-name from output root {output_root}; multiple candidates found: {candidates}"
)
raise ValueError("model-name is required when it cannot be inferred from model or output-root")
def parse_score_payload(raw_text: str) -> dict[str, Any]:
try:
parsed = extract_json_object(raw_text)
except Exception:
stripped = raw_text.strip()
if stripped.startswith("[") and stripped.endswith("]"):
parsed = {"score": json.loads(stripped), "reasoning": ""}
else:
raise
score = parsed.get("score", [])
if not isinstance(score, list):
score = [score]
parsed["score"] = [int(value) for value in score]
return parsed
def _utc_timestamp() -> str:
return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
by_task: dict[str, list[dict[str, Any]]] = defaultdict(list)
by_language: dict[str, list[dict[str, Any]]] = defaultdict(list)
for record in records:
by_task[record["task_type"]].append(record)
by_language[record["instruction_language"]].append(record)
return {
"count": len(records),
"by_task": {
task_type: {
"count": len(rows),
"samples": sorted(row["key"] for row in rows),
}
for task_type, rows in sorted(by_task.items())
},
"by_language": {
language: {
"count": len(rows),
"samples": sorted(row["key"] for row in rows),
}
for language, rows in sorted(by_language.items())
},
}
def select_balanced_gedit_rows(
rows: list[dict[str, Any]],
*,
task_type: str = "all",
instruction_language: str = "all",
samples_per_group: int | None,
) -> list[dict[str, Any]]:
filtered_rows = []
for row in rows:
if task_type != "all" and row["task_type"] != task_type:
continue
if instruction_language != "all" and row["instruction_language"] != instruction_language:
continue
filtered_rows.append(row)
if samples_per_group is None:
return filtered_rows
def _select_rows_for_group(group_rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
if instruction_language != "all":
return group_rows[:samples_per_group]
rows_by_language: dict[str, list[dict[str, Any]]] = defaultdict(list)
for row in group_rows:
rows_by_language[str(row.get("instruction_language", "")).strip()].append(row)
ordered_languages = ["en", "cn"]
per_language_quota = samples_per_group // max(len(ordered_languages), 1)
remainder = samples_per_group % max(len(ordered_languages), 1)
selected_rows: list[dict[str, Any]] = []
leftovers: list[dict[str, Any]] = []
for index, language in enumerate(ordered_languages):
language_rows = rows_by_language.get(language, [])
quota = per_language_quota + (1 if index < remainder else 0)
selected_rows.extend(language_rows[:quota])
leftovers.extend(language_rows[quota:])
if len(selected_rows) < samples_per_group:
selected_rows.extend(leftovers[: samples_per_group - len(selected_rows)])
return selected_rows
if task_type != "all":
return _select_rows_for_group(filtered_rows)
grouped_rows: dict[str, list[dict[str, Any]]] = defaultdict(list)
for row in filtered_rows:
grouped_rows[row["task_type"]].append(row)
selected: list[dict[str, Any]] = []
for group in GROUPS:
selected.extend(_select_rows_for_group(grouped_rows.get(group, [])))
return selected
def summarize_gedit_rows(rows: list[dict[str, Any]], language: str = "all") -> dict[str, Any]:
return summarize_gedit_rows_with_backbone(rows, language=language)
def _mean_or_none(values: list[float]) -> float | None:
return statistics.fmean(values) if values else None
def _to_q_metrics(section: dict[str, float | None]) -> dict[str, float | None]:
return {
"count": section.get("count"),
"Q_SC": section.get("avg_semantics"),
"Q_PQ": section.get("avg_quality"),
"Q_O": section.get("avg_overall"),
}
def _summarize_gedit_rows_single_language(
rows: list[dict[str, Any]],
*,
language: str,
) -> dict[str, Any]:
filtered_rows = [row for row in rows if str(row.get("instruction_language", "")).strip() == language]
def _to_bool(value: Any) -> bool:
if isinstance(value, bool):
return value
return str(value).strip().lower() in {"1", "true", "yes"}
per_group: dict[str, dict[str, float | None]] = {}
intersection_group: dict[str, dict[str, float | None]] = {}
for group in GROUPS:
group_rows = [row for row in filtered_rows if row.get("task_type") == group]
semantics = [float(row["semantics_score"]) for row in group_rows]
quality = [float(row["quality_score"]) for row in group_rows]
overall = [math.sqrt(float(row["semantics_score"]) * float(row["quality_score"])) for row in group_rows]
per_group[group] = {
"count": len(group_rows),
"avg_semantics": _mean_or_none(semantics),
"avg_quality": _mean_or_none(quality),
"avg_overall": _mean_or_none(overall),
}
intersection_rows = [row for row in group_rows if _to_bool(row.get("intersection_exist", False))]
intersection_semantics = [float(row["semantics_score"]) for row in intersection_rows]
intersection_quality = [float(row["quality_score"]) for row in intersection_rows]
intersection_overall = [
math.sqrt(float(row["semantics_score"]) * float(row["quality_score"])) for row in intersection_rows
]
intersection_group[group] = {
"count": len(intersection_rows),
"avg_semantics": _mean_or_none(intersection_semantics),
"avg_quality": _mean_or_none(intersection_quality),
"avg_overall": _mean_or_none(intersection_overall),
}
overall_section = {
"count": len(filtered_rows),
"avg_semantics": _mean_or_none(
[score["avg_semantics"] for score in per_group.values() if score["avg_semantics"] is not None]
),
"avg_quality": _mean_or_none(
[score["avg_quality"] for score in per_group.values() if score["avg_quality"] is not None]
),
"avg_overall": _mean_or_none(
[score["avg_overall"] for score in per_group.values() if score["avg_overall"] is not None]
),
}
intersection_section = {
"count": sum(score["count"] for score in intersection_group.values()),
"avg_semantics": _mean_or_none(
[score["avg_semantics"] for score in intersection_group.values() if score["avg_semantics"] is not None]
),
"avg_quality": _mean_or_none(
[score["avg_quality"] for score in intersection_group.values() if score["avg_quality"] is not None]
),
"avg_overall": _mean_or_none(
[score["avg_overall"] for score in intersection_group.values() if score["avg_overall"] is not None]
),
}
return {
"language": language,
"by_group": {group: _to_q_metrics(section) for group, section in per_group.items()},
"overall": _to_q_metrics(overall_section),
"intersection": _to_q_metrics(intersection_section),
}
def summarize_gedit_rows_with_backbone(
rows: list[dict[str, Any]],
*,
language: str = "all",
) -> dict[str, Any]:
if language == "all":
return {
"language": "all",
"languages": {
single_language: _summarize_gedit_rows_single_language(
rows,
language=single_language,
)
for single_language in ["en", "cn"]
},
}
return _summarize_gedit_rows_single_language(rows, language=language)
def _require_datasets():
try:
from datasets import load_dataset, load_from_disk
except ImportError as exc:
raise ImportError("GEdit-Bench requires the optional `datasets` package.") from exc
return load_dataset, load_from_disk
def _load_gedit_dataset(dataset_ref: str):
load_dataset, load_from_disk = _require_datasets()
dataset_path = Path(dataset_ref)
if dataset_path.exists():
if (dataset_path / "state.json").exists() and (dataset_path / "dataset_info.json").exists():
return load_from_disk(str(dataset_path))
return load_dataset(str(dataset_path))
return load_dataset(dataset_ref)
def _resolve_gedit_split(dataset_obj: Any) -> Any:
if isinstance(dataset_obj, dict):
if "train" in dataset_obj:
return dataset_obj["train"]
return dataset_obj
try:
return dataset_obj["train"]
except Exception:
return dataset_obj
def _to_pil_image(value: Any) -> Image.Image:
if isinstance(value, Image.Image):
return value.convert("RGB")
if isinstance(value, dict) and "bytes" in value:
image = Image.open(BytesIO(value["bytes"]))
image.load()
return image.convert("RGB")
raise TypeError(f"Unsupported image payload type: {type(value)!r}")
class LocalVIEScorer:
def __init__(self, *, base_url: str, api_key: str, model: str, timeout: int = 600):
self.base_url = base_url.rstrip("/")
self.api_key = api_key
self.model = model
self.timeout = timeout
self.sc_prompt = (
"You are evaluating image editing quality.\n"
"Two images are provided: the source image and the edited image.\n"
'Return JSON only in the format {"score": [edit_success, content_preservation], '
'"reasoning": "..."}.\n'
"Each score must be an integer from 0 to 10.\n"
"Editing instruction: <instruction>"
)
self.pq_prompt = (
"You are evaluating image quality.\n"
'Return JSON only in the format {"score": [naturalness, artifact_free], "reasoning": "..."}.\n'
"Each score must be an integer from 0 to 10."
)
def _request(self, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
from benchmarks.accuracy.common import pil_to_data_url
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
for image in images:
content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
response = requests.post(
build_openai_url(self.base_url, "/chat/completions"),
json={
"model": self.model,
"messages": [{"role": "user", "content": content}],
"temperature": 0,
},
headers={
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
},
timeout=self.timeout,
)
response.raise_for_status()
message_content = response.json()["choices"][0]["message"]["content"]
if isinstance(message_content, list):
text = "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
else:
text = str(message_content)
return parse_score_payload(text)
def evaluate(self, source_image: Image.Image, edited_image: Image.Image, instruction: str) -> dict[str, float]:
sc_payload = self._request(self.sc_prompt.replace("<instruction>", instruction), [source_image, edited_image])
pq_payload = self._request(self.pq_prompt, [edited_image])
semantics = float(min(sc_payload["score"])) if sc_payload["score"] else 0.0
quality = float(min(pq_payload["score"])) if pq_payload["score"] else 0.0
overall = math.sqrt(semantics * quality)
return {
"semantics_score": semantics,
"quality_score": quality,
"overall_score": overall,
}
class GEditBenchRunner:
def __init__(
self,
*,
dataset_ref: str,
output_root: Path,
base_url: str,
model: str,
api_key: str = "EMPTY",
width: int = 512,
height: int = 512,
num_inference_steps: int = 20,
guidance_scale: float | None = None,
seed: int | None = 42,
):
self.dataset_ref = dataset_ref
self.output_root = output_root
self.model = model
self.width = width
self.height = height
self.num_inference_steps = num_inference_steps
self.guidance_scale = guidance_scale
self.seed = seed
self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
def generate(
self,
*,
model_name: str,
task_type: str = "all",
instruction_language: str = "all",
workers: int = 1,
max_samples: int | None = None,
samples_per_group: int | None = None,
) -> list[dict[str, Any]]:
dataset = _resolve_gedit_split(_load_gedit_dataset(self.dataset_ref))
rows = select_balanced_gedit_rows(
list(dataset),
task_type=task_type,
instruction_language=instruction_language,
samples_per_group=samples_per_group,
)
if max_samples is not None:
rows = rows[:max_samples]
outputs: list[dict[str, Any]] = []
total = len(rows)
if workers <= 1:
with tqdm(total=total, desc="GEdit generate", unit="sample") as progress:
for item in rows:
result = self._safe_generate_one(model_name, item)
if result:
outputs.append(result)
progress.update(1)
return outputs
with tqdm(total=total, desc="GEdit generate", unit="sample") as progress:
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = [executor.submit(self._safe_generate_one, model_name, item) for item in rows]
for future in as_completed(futures):
result = future.result()
if result:
outputs.append(result)
progress.update(1)
return outputs
def _safe_generate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
try:
return self._generate_one(model_name, item)
except Exception:
logger.exception("Failed to generate GEdit-Bench sample %s", item.get("key", "<unknown>"))
return None
def _generate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
output_path = (
self.output_root
/ model_name
/ "fullset"
/ item["task_type"]
/ item["instruction_language"]
/ f"{item['key']}.png"
)
output_path.parent.mkdir(parents=True, exist_ok=True)
if output_path.exists():
return {
"task_type": item["task_type"],
"instruction_language": item["instruction_language"],
"key": item["key"],
"output_path": str(output_path),
}
source_image = _to_pil_image(item["input_image_raw"])
edited_image = self.client.generate_image_edit(
model=self.model,
prompt=item["instruction"],
images=source_image,
width=self.width,
height=self.height,
num_inference_steps=self.num_inference_steps,
guidance_scale=self.guidance_scale,
seed=self.seed,
)
edited_image.save(output_path)
return {
"task_type": item["task_type"],
"instruction_language": item["instruction_language"],
"key": item["key"],
"output_path": str(output_path),
}
class GEditBenchEvaluator:
def __init__(self, *, dataset_ref: str, output_root: Path, scorer: LocalVIEScorer):
self.dataset_ref = dataset_ref
self.output_root = output_root
self.scorer = scorer
def evaluate(
self,
*,
model_name: str,
save_dir: Path,
task_type: str = "all",
instruction_language: str = "all",
workers: int = 1,
max_samples: int | None = None,
samples_per_group: int | None = None,
) -> dict[str, Any]:
dataset = _resolve_gedit_split(_load_gedit_dataset(self.dataset_ref))
rows = select_balanced_gedit_rows(
list(dataset),
task_type=task_type,
instruction_language=instruction_language,
samples_per_group=samples_per_group,
)
if max_samples is not None:
rows = rows[:max_samples]
results: list[dict[str, Any]] = []
total = len(rows)
if workers <= 1:
with tqdm(total=total, desc="GEdit evaluate", unit="sample") as progress:
for item in rows:
result = self._safe_evaluate_one(model_name, item)
if result:
results.append(result)
progress.update(1)
else:
with tqdm(total=total, desc="GEdit evaluate", unit="sample") as progress:
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = [executor.submit(self._safe_evaluate_one, model_name, item) for item in rows]
for future in as_completed(futures):
result = future.result()
if result:
results.append(result)
progress.update(1)
save_dir.mkdir(parents=True, exist_ok=True)
base_name = f"{model_name}_{task_type}_{instruction_language}"
timestamp = _utc_timestamp()
csv_path = save_dir / f"{base_name}_vie_score.csv"
timestamped_csv_path = save_dir / f"{base_name}_vie_score_{timestamp}.csv"
def _write_csv(path: Path) -> None:
with path.open("w", encoding="utf-8", newline="") as handle:
writer = csv.DictWriter(
handle,
fieldnames=[
"key",
"task_type",
"edited_image",
"instruction",
"semantics_score",
"quality_score",
"overall_score",
"intersection_exist",
"instruction_language",
],
)
writer.writeheader()
for row in results:
writer.writerow(row)
_write_csv(csv_path)
_write_csv(timestamped_csv_path)
summary = summarize_gedit_rows_with_backbone(
results,
language=instruction_language,
)
summary_path = save_dir / f"{base_name}_summary.json"
timestamped_summary_path = save_dir / f"{base_name}_summary_{timestamp}.json"
write_json(summary_path, summary)
write_json(timestamped_summary_path, summary)
return {
"results": results,
"summary": summary,
"csv_path": str(csv_path),
"summary_path": str(summary_path),
"timestamped_csv_path": str(timestamped_csv_path),
"timestamped_summary_path": str(timestamped_summary_path),
}
def _safe_evaluate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
try:
return self._evaluate_one(model_name, item)
except Exception:
logger.exception("Failed to evaluate GEdit-Bench sample %s", item.get("key", "<unknown>"))
return None
def _evaluate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
edited_image_path = (
self.output_root
/ model_name
/ "fullset"
/ item["task_type"]
/ item["instruction_language"]
/ f"{item['key']}.png"
)
if not edited_image_path.exists():
return None
source_image = _to_pil_image(item["input_image_raw"])
edited_image = Image.open(edited_image_path).convert("RGB")
scores = self.scorer.evaluate(source_image, edited_image, item["instruction"])
return {
"key": item["key"],
"task_type": item["task_type"],
"edited_image": str(edited_image_path),
"instruction": item["instruction"],
"semantics_score": scores["semantics_score"],
"quality_score": scores["quality_score"],
"overall_score": scores["overall_score"],
"intersection_exist": item.get("Intersection_exist", False),
"instruction_language": item["instruction_language"],
}
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Run the GEdit-Bench integration against a local vLLM-Omni server.")
subparsers = parser.add_subparsers(dest="command", required=True)
generate = subparsers.add_parser("generate")
generate.add_argument("--dataset-ref", type=str, default="stepfun-ai/GEdit-Bench")
generate.add_argument("--output-root", type=Path, required=True)
generate.add_argument("--base-url", type=str, required=True)
generate.add_argument("--model", type=str, required=True)
generate.add_argument("--model-name", type=str, default=None)
generate.add_argument("--api-key", type=str, default="EMPTY")
generate.add_argument("--task-type", choices=["all", *GROUPS], default="all")
generate.add_argument("--instruction-language", choices=["all", "en", "cn"], default="all")
generate.add_argument("--width", type=int, default=512)
generate.add_argument("--height", type=int, default=512)
generate.add_argument("--num-inference-steps", type=int, default=20)
generate.add_argument("--guidance-scale", type=float, default=None)
generate.add_argument("--seed", type=int, default=42)
generate.add_argument("--workers", type=int, default=1)
generate.add_argument("--max-samples", type=int, default=None)
generate.add_argument("--samples-per-group", type=int, default=None)
evaluate = subparsers.add_parser("evaluate")
evaluate.add_argument("--dataset-ref", type=str, default="stepfun-ai/GEdit-Bench")
evaluate.add_argument("--output-root", type=Path, required=True)
evaluate.add_argument("--model-name", type=str, default=None)
evaluate.add_argument("--save-dir", type=Path, required=True)
evaluate.add_argument("--task-type", choices=["all", *GROUPS], default="all")
evaluate.add_argument("--instruction-language", choices=["all", "en", "cn"], default="all")
evaluate.add_argument("--judge-base-url", type=str, required=True)
evaluate.add_argument("--judge-model", type=str, required=True)
evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
evaluate.add_argument("--workers", type=int, default=1)
evaluate.add_argument("--max-samples", type=int, default=None)
evaluate.add_argument("--samples-per-group", type=int, default=None)
summarize = subparsers.add_parser("summarize")
summarize.add_argument("--csv-path", type=Path, required=True)
summarize.add_argument("--language", choices=["all", "en", "cn"], default="all")
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
if args.command == "generate":
model_name = resolve_model_name(model_name=args.model_name, model=args.model)
runner = GEditBenchRunner(
dataset_ref=args.dataset_ref,
output_root=args.output_root,
base_url=args.base_url,
model=args.model,
api_key=args.api_key,
width=args.width,
height=args.height,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
)
records = runner.generate(
model_name=model_name,
task_type=args.task_type,
instruction_language=args.instruction_language,
workers=args.workers,
max_samples=args.max_samples,
samples_per_group=args.samples_per_group,
)
payload = {"records": records, "summary": summarize_generated_records(records)}
write_json(args.output_root / model_name / "generation_manifest.json", payload)
return 0
if args.command == "evaluate":
model_name = resolve_model_name(model_name=args.model_name, output_root=args.output_root)
scorer = LocalVIEScorer(
base_url=args.judge_base_url,
api_key=args.judge_api_key,
model=args.judge_model,
)
evaluator = GEditBenchEvaluator(dataset_ref=args.dataset_ref, output_root=args.output_root, scorer=scorer)
evaluator.evaluate(
model_name=model_name,
save_dir=args.save_dir,
task_type=args.task_type,
instruction_language=args.instruction_language,
workers=args.workers,
max_samples=args.max_samples,
samples_per_group=args.samples_per_group,
)
return 0
if args.command == "summarize":
with args.csv_path.open("r", encoding="utf-8", newline="") as handle:
rows = list(csv.DictReader(handle))
summary = summarize_gedit_rows_with_backbone(rows, language=args.language)
print(json.dumps(summary, indent=2, ensure_ascii=False))
return 0
parser.error(f"Unknown command: {args.command}")
return 1
@@ -1,13 +0,0 @@
# ruff: noqa: E402, I001
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[3]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from benchmarks.accuracy.image_to_image.gedit_bench import main
if __name__ == "__main__":
raise SystemExit(main())
-103
View File
@@ -1,103 +0,0 @@
# GEBench on vLLM-Omni
This integration adapts the upstream `stepfun-ai/GEBench` scripts into
`vllm-omni/benchmarks/accuracy/text_to_image`.
Upstream mapping:
- `scripts/generate.py` -> `run_gebench.py generate`
- `scripts/evaluate.py` -> `run_gebench.py evaluate`
- upstream prompt / judge logic -> `gbench.py`
What changed:
- Generation calls the local OpenAI-compatible `vllm-omni` endpoints:
- `/v1/images/generations` for text-only frame generation
- `/v1/images/edits` for image-conditioned GUI transition generation
- Evaluation still keeps the GEBench scoring dimensions:
- `goal`
- `logic`
- `cons`
- `ui`
- `qual`
- Judge calls are also routed to a local OpenAI-compatible model served by
`vllm-omni` instead of a remote service.
Expected dataset layout:
- Clone the benchmark dataset from Hugging Face into a local directory:
```bash
git clone https://huggingface.co/datasets/stepfun-ai/GEBench /path/to/GEBench
```
Example usage:
```bash
python benchmarks/accuracy/text_to_image/run_gebench.py generate \
--dataset-root /path/to/GEBench \
--output-root benchmarks/accuracy/text_to_image/outputs \
--base-url http://127.0.0.1:8000 \
--model Tongyi-MAI/Z-Image-Turbo \
--data-type type3
```
```bash
python benchmarks/accuracy/text_to_image/run_gebench.py evaluate \
--dataset-root /path/to/GEBench \
--output-root benchmarks/accuracy/text_to_image/outputs \
--data-type type3 \
--judge-base-url http://127.0.0.1:8000 \
--judge-model Qwen/Qwen2.5-VL-7B-Instruct \
--judge-api-key EMPTY
```
```bash
python benchmarks/accuracy/text_to_image/run_gebench.py summarize \
--output-root benchmarks/accuracy/text_to_image/outputs
```
Example summary output:
```json
{
"generation": {
"count": 20,
"by_type": {
"type3": {"count": 10},
"type4": {"count": 10}
}
},
"evaluation": {
"count": 20,
"overall_mean": 0.52,
"by_type": {
"type3": {"count": 10, "overall_mean": 0.50, "overall_mean_100": 50.0},
"type4": {"count": 10, "overall_mean": 0.54, "overall_mean_100": 54.0}
}
}
}
```
Example generated images to inspect:
- `benchmarks/accuracy/text_to_image/outputs/03_trajectory_text_fictionalapp/<lang_device>/<sample>/frame0.png`
- `benchmarks/accuracy/text_to_image/outputs/03_trajectory_text_fictionalapp/<lang_device>/<sample>/frame5.png`
- `benchmarks/accuracy/text_to_image/outputs/04_trajectory_text_realapp/<lang_device>/<sample>/frame0.png`
- `benchmarks/accuracy/text_to_image/outputs/04_trajectory_text_realapp/<lang_device>/<sample>/frame5.png`
What to expect:
- `overall_mean` is normalized to `0.0 ~ 1.0`; higher is better.
- `frame0.png` is the initial GUI frame; `frame5.png` is the final trajectory frame most often used for quick inspection.
- For full debugging, inspect the whole `frame0.png` ... `frame5.png` sequence for one sample directory.
Notes:
- GEBench upstream leaves type3/type4 generation unfinished. This integration
fills that gap with a trajectory runner that generates `frame0.png` followed
by `frame1.png` ... `frame5.png`.
- Type1/2/5 require an image-edit capable model exposed through
`vllm-omni serve`.
- `summarize` will report both generated coverage and any existing evaluation
summary files.
@@ -1 +0,0 @@
"""GEBench integration for vLLM-Omni."""
-927
View File
@@ -1,927 +0,0 @@
from __future__ import annotations
import argparse
import json
import statistics
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
import requests
from PIL import Image
from benchmarks.accuracy.common import (
VllmOmniImageClient,
build_openai_url,
ensure_dir,
extract_json_object,
find_first_image,
load_json,
pil_to_data_url,
save_image,
write_json,
)
TYPE_TO_FOLDER = {
"type1": "01_single_step",
"type2": "02_multi_step",
"type3": "03_trajectory_text_fictionalapp",
"type4": "04_trajectory_text_realapp",
"type5": "05_grounding_data",
}
SCORE_KEYS = ("goal", "logic", "cons", "ui", "qual")
DEFAULT_SAMPLES_PER_TYPE = 10
def _utc_timestamp() -> str:
return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
def _write_json_with_timestamp(path: Path, payload: dict[str, Any]) -> Path:
write_json(path, payload)
timestamped_path = path.with_name(f"{path.stem}_{_utc_timestamp()}{path.suffix}")
write_json(timestamped_path, payload)
return timestamped_path
@dataclass(frozen=True)
class GEBenchSampleSpec:
sample_path: Path
metadata: dict[str, Any]
sample_name: str
lang_device: str
def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
for record in records:
by_type[record["data_type"]].append(record)
return {
"count": len(records),
"by_type": {
data_type: {
"count": len(rows),
}
for data_type, rows in sorted(by_type.items())
},
}
def summarize_gebench_results(results: list[dict[str, Any]]) -> dict[str, Any]:
by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
for result in results:
by_type[result["data_type"]].append(result)
summary: dict[str, Any] = {
"count": len(results),
"overall_mean": statistics.fmean(r["overall"] for r in results) if results else 0.0,
"by_type": {},
}
for data_type, rows in by_type.items():
score_means: dict[str, float] = {}
all_score_keys = {key for row in rows for key in row.get("scores", {}).keys()}
for score_key in all_score_keys:
values = [row["scores"][score_key] for row in rows if score_key in row.get("scores", {})]
score_means[score_key] = statistics.fmean(values) if values else 0.0
overall_mean = statistics.fmean(row["overall"] for row in rows)
summary["by_type"][data_type] = {
"count": len(rows),
"overall_mean": overall_mean,
"overall_mean_100": overall_mean * 100.0,
"score_means": score_means,
}
return summary
def select_balanced_gebench_samples(
sample_paths_by_type: dict[str, list[Any]],
*,
samples_per_type: int | None,
) -> dict[str, list[Any]]:
if samples_per_type is None:
return {data_type: list(paths) for data_type, paths in sample_paths_by_type.items()}
return {data_type: list(paths)[:samples_per_type] for data_type, paths in sample_paths_by_type.items()}
def collect_gebench_generation_summary(output_root: Path) -> dict[str, Any]:
records: list[dict[str, Any]] = []
for data_type, folder_name in TYPE_TO_FOLDER.items():
type_root = output_root / folder_name
if not type_root.exists():
continue
for lang_dir in sorted(path for path in type_root.iterdir() if path.is_dir()):
for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir()):
expected = sample_dir / "frame5.png" if data_type in {"type2", "type3", "type4"} else None
if expected is None:
expected = find_first_image(sample_dir)
elif not expected.exists():
expected = None
if expected is None:
continue
records.append(
{
"data_type": data_type,
"sample_name": f"{lang_dir.name}/{sample_dir.name}",
"output_path": str(expected),
}
)
return summarize_generated_records(records)
def _normalize_score_key(key: str) -> str:
mapping = {
"goal": "goal",
"logic": "logic",
"cons": "cons",
"consistency": "cons",
"ui": "ui",
"qual": "qual",
"quality": "qual",
}
return mapping.get(key.lower(), key.lower())
def _normalize_scores(raw_scores: dict[str, Any]) -> dict[str, int]:
scores: dict[str, int] = {}
for key, value in raw_scores.items():
normalized = _normalize_score_key(key)
if normalized not in SCORE_KEYS:
continue
scalar = value.get("s", 0) if isinstance(value, dict) else value
try:
scores[normalized] = int(scalar)
except (TypeError, ValueError):
scores[normalized] = 0
for key in SCORE_KEYS:
scores.setdefault(key, 0)
return scores
def _compute_overall(scores: dict[str, int]) -> float:
return sum(scores.values()) / (len(SCORE_KEYS) * 5.0)
def _iter_sample_paths(dataset_root: Path, data_type: str) -> list[Path]:
data_dir = dataset_root / TYPE_TO_FOLDER[data_type]
if not data_dir.exists():
raise FileNotFoundError(f"Dataset path not found: {data_dir}")
samples: list[Path] = []
for lang_dir in sorted(path for path in data_dir.iterdir() if path.is_dir()):
for child in sorted(lang_dir.iterdir()):
if child.is_dir():
samples.append(child)
elif child.suffix.lower() == ".json":
samples.append(child)
return samples
def _sample_name_from_metadata(metadata: dict[str, Any], sample_path: Path, item_index: int | None = None) -> str:
sample_id = metadata.get("id") or metadata.get("sample_id") or metadata.get("name")
if sample_id:
return str(sample_id)
if item_index is not None:
return f"{sample_path.stem}_{item_index:04d}"
return sample_path.stem if sample_path.is_file() else sample_path.name
def _expand_sample_path(sample_path: Path) -> list[GEBenchSampleSpec]:
if sample_path.is_dir():
metadata = _load_metadata(sample_path)
return [
GEBenchSampleSpec(
sample_path=sample_path,
metadata=metadata,
sample_name=_sample_name(sample_path),
lang_device=_lang_device(sample_path, metadata),
)
]
payload = load_json(sample_path)
if isinstance(payload, dict):
return [
GEBenchSampleSpec(
sample_path=sample_path,
metadata=payload,
sample_name=_sample_name_from_metadata(payload, sample_path),
lang_device=_lang_device(sample_path, payload),
)
]
if isinstance(payload, list):
specs: list[GEBenchSampleSpec] = []
for item_index, item in enumerate(payload):
if not isinstance(item, dict):
continue
specs.append(
GEBenchSampleSpec(
sample_path=sample_path,
metadata=item,
sample_name=_sample_name_from_metadata(item, sample_path, item_index),
lang_device=_lang_device(sample_path, item),
)
)
return specs
raise TypeError(f"Unsupported metadata payload type for sample {sample_path}: {type(payload)!r}")
def _iter_sample_specs(dataset_root: Path, data_type: str) -> list[GEBenchSampleSpec]:
sample_specs: list[GEBenchSampleSpec] = []
for sample_path in _iter_sample_paths(dataset_root, data_type):
sample_specs.extend(_expand_sample_path(sample_path))
return sample_specs
def _load_metadata(sample_path: Path) -> dict[str, Any]:
if sample_path.is_file():
return load_json(sample_path)
for candidate in ("meta_data.json", "metadata.json"):
meta_path = sample_path / candidate
if meta_path.exists():
return load_json(meta_path)
raise FileNotFoundError(f"Metadata not found for sample: {sample_path}")
def _sample_name(sample_path: Path) -> str:
return sample_path.stem if sample_path.is_file() else sample_path.name
def _lang_device(sample_path: Path, metadata: dict[str, Any]) -> str:
return str(metadata.get("lang_device") or sample_path.parent.name)
def _resolve_referenced_image(
*,
metadata: dict[str, Any],
sample_path: Path,
dataset_root: Path,
data_type: str,
) -> Image.Image | None:
for key in ("image", "input_image", "initial_image", "reference_image"):
image_ref = metadata.get(key)
if not image_ref:
continue
candidate = dataset_root / TYPE_TO_FOLDER[data_type] / str(image_ref)
if candidate.exists():
image = Image.open(candidate)
image.load()
return image.convert("RGB")
if sample_path.is_dir():
local_image = find_first_image(sample_path)
if local_image:
image = Image.open(local_image)
image.load()
return image.convert("RGB")
return None
def _trajectory_steps(metadata: dict[str, Any]) -> list[dict[str, Any]]:
for key in ("trajectory", "steps", "frames"):
value = metadata.get(key)
if isinstance(value, list):
return [step for step in value if isinstance(step, dict)]
extracted: list[dict[str, Any]] = []
for index in range(1, 6):
value = metadata.get(f"step{index}") or metadata.get(str(index))
if isinstance(value, dict):
extracted.append(value)
return extracted
def _text_or_default(value: Any, default: str = "") -> str:
return str(value).strip() if value is not None else default
def _type1_prompt(metadata: dict[str, Any]) -> str:
caption = _text_or_default(metadata.get("caption") or metadata.get("instruction"), "Transform the reference GUI.")
return (
"Using the reference GUI screenshot, generate the next GUI state after the requested interaction.\n\n"
f"Requested change:\n{caption}\n\n"
"Requirements:\n"
"- Preserve layout, visual identity, and unrelated regions.\n"
"- Only apply the requested state change.\n"
"- Keep all text and controls readable.\n"
)
def _type2_prompt(goal: str, step_num: int) -> str:
return (
"Generate the next GUI state for a multi-step task.\n\n"
f"Overall goal: {goal}\n"
f"Current progress step: {step_num}/5\n\n"
"Requirements:\n"
"- The change should be incremental and plausible.\n"
"- Preserve layout and visual identity.\n"
"- Make text/buttons readable.\n"
)
def _type34_initial_prompt(metadata: dict[str, Any], first_step: dict[str, Any]) -> str:
app_name = _text_or_default(metadata.get("app_name"), "App")
final_goal = _text_or_default(metadata.get("final_goal") or metadata.get("instruction"), "Complete the task.")
visual_description = _text_or_default(
metadata.get("visual_description") or first_step.get("visual_description") or first_step.get("description"),
"A clean product-quality app home screen.",
)
return (
"Generate the first GUI frame for a task trajectory.\n\n"
f"App name: {app_name}\n"
f"Final goal: {final_goal}\n"
f"Visual description:\n{visual_description}\n\n"
"Requirements:\n"
"- Generate a production-looking UI screenshot only.\n"
"- Keep the layout coherent and readable.\n"
)
def _type34_next_prompt(step_num: int, step_info: dict[str, Any]) -> str:
action = _text_or_default(step_info.get("action") or step_info.get("instruction"), "Continue the task.")
visual_description = _text_or_default(
step_info.get("visual_description") or step_info.get("description"),
"Reflect the expected next GUI state.",
)
return (
"Using the previous frame as reference, generate the next GUI frame.\n\n"
f"Step {step_num} action: {action}\n"
f"Expected visual state:\n{visual_description}\n\n"
"Requirements:\n"
"- Only change UI regions affected by this action.\n"
"- Preserve persistent bars, layout, and style.\n"
"- Keep text and icons readable.\n"
)
def _type5_prompt(metadata: dict[str, Any]) -> str:
grounding = metadata.get("grounding") or {}
explanation = _text_or_default(
metadata.get("grounding_explanation") or grounding.get("effect") or grounding.get("description"),
"Predict the immediate GUI reaction to the indicated target.",
)
return (
"Using the reference GUI screenshot, predict the immediate GUI state after the grounded interaction.\n\n"
f"Expected effect: {explanation}\n"
f"Grounding metadata: {json.dumps(grounding, ensure_ascii=False)}\n\n"
"Requirements:\n"
"- Apply only the interaction-triggered change.\n"
"- Preserve unrelated regions.\n"
"- Keep the UI realistic and readable.\n"
)
def _make_storyboard_image(
frames: list[Image.Image],
*,
columns: int = 3,
background_color: tuple[int, int, int] = (255, 255, 255),
) -> Image.Image:
if not frames:
raise ValueError("Expected at least one frame to build a storyboard image.")
normalized = [frame.convert("RGB") for frame in frames]
frame_width = max(frame.width for frame in normalized)
frame_height = max(frame.height for frame in normalized)
rows = (len(normalized) + columns - 1) // columns
storyboard = Image.new("RGB", (frame_width * columns, frame_height * rows), color=background_color)
for index, frame in enumerate(normalized):
x_offset = (index % columns) * frame_width
y_offset = (index // columns) * frame_height
if frame.size != (frame_width, frame_height):
frame = frame.resize((frame_width, frame_height))
storyboard.paste(frame, (x_offset, y_offset))
return storyboard
def _trajectory_judge_payload(frames: list[Image.Image]) -> tuple[str, list[Image.Image]]:
storyboard = _make_storyboard_image(frames, columns=3)
prompt_suffix = (
"The attached image is a storyboard containing six frames arranged left-to-right, "
"top-to-bottom as frame0, frame1, frame2, frame3, frame4, frame5."
)
return prompt_suffix, [storyboard]
class LocalJudgeClient:
def __init__(self, base_url: str, api_key: str, model: str, timeout: int = 600):
self.base_url = base_url.rstrip("/")
self.api_key = api_key
self.model = model
self.timeout = timeout
def _build_scoring_prompt(self, task_prompt: str) -> str:
return (
"You are an expert evaluator for GUI image editing and GUI trajectory generation.\n"
"Evaluate whether the generated image(s) satisfy the task.\n\n"
"Score these five dimensions from 0 to 5:\n"
"- goal: whether the user goal is completed correctly\n"
"- logic: whether the transition/state change is logically correct\n"
"- cons: whether unrelated regions remain consistent\n"
"- ui: whether the UI layout/components remain realistic and coherent\n"
"- qual: whether the images are visually clear and artifact-free\n\n"
"Return JSON only. Do not add any prose outside JSON.\n"
"Use exactly this schema:\n"
"{\n"
' "goal": 0,\n'
' "logic": 0,\n'
' "cons": 0,\n'
' "ui": 0,\n'
' "qual": 0,\n'
' "reasoning": "short explanation"\n'
"}\n\n"
"Scoring task:\n"
f"{task_prompt}"
)
def _request_text(self, prompt: str, images: list[Image.Image]) -> str:
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
for image in images:
content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
response = requests.post(
build_openai_url(self.base_url, "/chat/completions"),
json={
"model": self.model,
"messages": [{"role": "user", "content": content}],
"temperature": 0,
},
headers={
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
},
timeout=self.timeout,
)
response.raise_for_status()
message_content = response.json()["choices"][0]["message"]["content"]
if isinstance(message_content, list):
return "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
return str(message_content)
def evaluate(self, *, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
primary_prompt = self._build_scoring_prompt(prompt)
text = self._request_text(primary_prompt, images)
try:
return extract_json_object(text)
except ValueError:
retry_prompt = (
self._build_scoring_prompt(prompt) + "\n\nYour previous response was not valid JSON. "
"Return only the JSON object with integer scores."
)
retry_text = self._request_text(retry_prompt, images)
try:
return extract_json_object(retry_text)
except ValueError:
return {
"goal": 0,
"logic": 0,
"cons": 0,
"ui": 0,
"qual": 0,
"reasoning": retry_text.strip() or text.strip() or "Judge response was not valid JSON.",
}
class GEBenchRunner:
def __init__(
self,
*,
dataset_root: Path,
output_root: Path,
base_url: str,
model: str,
api_key: str = "EMPTY",
width: int = 768,
height: int = 576,
num_inference_steps: int = 8,
output_compression: int | None = 98,
guidance_scale: float | None = None,
seed: int | None = 42,
):
self.dataset_root = dataset_root
self.output_root = output_root
self.model = model
self.width = width
self.height = height
self.num_inference_steps = num_inference_steps
self.output_compression = output_compression
self.guidance_scale = guidance_scale
self.seed = seed
self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
def generate(
self,
*,
data_type: str,
workers: int = 1,
max_samples: int | None = None,
samples_per_type: int | None = None,
) -> list[dict[str, Any]]:
sample_specs = select_balanced_gebench_samples(
{data_type: _iter_sample_specs(self.dataset_root, data_type)},
samples_per_type=samples_per_type,
)[data_type]
if max_samples is not None:
sample_specs = sample_specs[:max_samples]
results: list[dict[str, Any]] = []
if workers <= 1:
for sample_spec in sample_specs:
result = self._generate_one(data_type, sample_spec)
if result:
results.append(result)
return results
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = [executor.submit(self._generate_one, data_type, sample_spec) for sample_spec in sample_specs]
for future in as_completed(futures):
result = future.result()
if result:
results.append(result)
return results
def _generate_one(self, data_type: str, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
sample_path = sample_spec.sample_path
metadata = sample_spec.metadata
lang_device = sample_spec.lang_device
sample_name = sample_spec.sample_name
output_dir = ensure_dir(self.output_root / TYPE_TO_FOLDER[data_type] / lang_device / sample_name)
if data_type == "type1":
output_path = output_dir / "generated.png"
if output_path.exists():
return {
"data_type": data_type,
"sample_name": f"{lang_device}/{sample_name}",
"output_path": str(output_path),
}
source = _resolve_referenced_image(
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
)
if source is None:
return None
generated = self.client.generate_image_edit(
model=self.model,
prompt=_type1_prompt(metadata),
images=source,
width=self.width,
height=self.height,
num_inference_steps=self.num_inference_steps,
output_compression=self.output_compression,
guidance_scale=self.guidance_scale,
seed=self.seed,
)
save_image(output_path, generated)
return {
"data_type": data_type,
"sample_name": f"{lang_device}/{sample_name}",
"output_path": str(output_path),
}
if data_type == "type2":
goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
source = _resolve_referenced_image(
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
)
if source is None:
return None
frame0_path = output_dir / "frame0.png"
if not frame0_path.exists():
save_image(frame0_path, source)
previous = source
for step_num in range(1, 6):
frame_path = output_dir / f"frame{step_num}.png"
if frame_path.exists():
previous = Image.open(frame_path).convert("RGB")
continue
generated = self.client.generate_image_edit(
model=self.model,
prompt=_type2_prompt(goal, step_num),
images=previous,
width=self.width,
height=self.height,
num_inference_steps=self.num_inference_steps,
output_compression=self.output_compression,
guidance_scale=self.guidance_scale,
seed=self.seed,
)
save_image(frame_path, generated)
previous = generated
output_path = output_dir / "frame5.png"
return {
"data_type": data_type,
"sample_name": f"{lang_device}/{sample_name}",
"output_path": str(output_path),
}
if data_type in {"type3", "type4"}:
steps = _trajectory_steps(metadata)
frame0_path = output_dir / "frame0.png"
if frame0_path.exists():
previous = Image.open(frame0_path).convert("RGB")
else:
previous = self.client.generate_text_to_image(
model=self.model,
prompt=_type34_initial_prompt(metadata, steps[0] if steps else {}),
width=self.width,
height=self.height,
num_inference_steps=self.num_inference_steps,
output_compression=self.output_compression,
guidance_scale=self.guidance_scale,
seed=self.seed,
)
save_image(frame0_path, previous)
for step_num in range(1, 6):
frame_path = output_dir / f"frame{step_num}.png"
if frame_path.exists():
previous = Image.open(frame_path).convert("RGB")
continue
step_info = steps[step_num - 1] if step_num - 1 < len(steps) else {}
generated = self.client.generate_image_edit(
model=self.model,
prompt=_type34_next_prompt(step_num, step_info),
images=previous,
width=self.width,
height=self.height,
num_inference_steps=self.num_inference_steps,
output_compression=self.output_compression,
guidance_scale=self.guidance_scale,
seed=self.seed,
)
save_image(frame_path, generated)
previous = generated
output_path = output_dir / "frame5.png"
return {
"data_type": data_type,
"sample_name": f"{lang_device}/{sample_name}",
"output_path": str(output_path),
}
if data_type == "type5":
output_path = output_dir / "generated.png"
if output_path.exists():
return {
"data_type": data_type,
"sample_name": f"{lang_device}/{sample_name}",
"output_path": str(output_path),
}
source = _resolve_referenced_image(
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
)
if source is None:
return None
generated = self.client.generate_image_edit(
model=self.model,
prompt=_type5_prompt(metadata),
images=source,
width=self.width,
height=self.height,
num_inference_steps=self.num_inference_steps,
output_compression=self.output_compression,
guidance_scale=self.guidance_scale,
seed=self.seed,
)
save_image(output_path, generated)
return {
"data_type": data_type,
"sample_name": f"{lang_device}/{sample_name}",
"output_path": str(output_path),
}
raise ValueError(f"Unsupported data type: {data_type}")
class GEBenchEvaluator:
def __init__(self, *, dataset_root: Path, output_root: Path, judge: LocalJudgeClient):
self.dataset_root = dataset_root
self.output_root = output_root
self.judge = judge
def evaluate(
self,
*,
data_type: str,
workers: int = 1,
max_samples: int | None = None,
samples_per_type: int | None = None,
) -> dict[str, Any]:
output_type_dir = self.output_root / TYPE_TO_FOLDER[data_type]
sample_specs_by_name = {
(spec.lang_device, spec.sample_name): spec for spec in _iter_sample_specs(self.dataset_root, data_type)
}
if not output_type_dir.exists():
payload = {"data_type": data_type, "results": [], "summary": summarize_gebench_results([])}
write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
return payload
sample_dirs = [
sample_dir
for lang_dir in sorted(path for path in output_type_dir.iterdir() if path.is_dir())
for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir())
if (lang_dir.name, sample_dir.name) in sample_specs_by_name
]
sample_dirs = select_balanced_gebench_samples(
{data_type: sample_dirs},
samples_per_type=samples_per_type,
)[data_type]
if max_samples is not None:
sample_dirs = sample_dirs[:max_samples]
results: list[dict[str, Any]] = []
if workers <= 1:
for sample_dir in sample_dirs:
result = self._evaluate_one(
data_type,
sample_dir,
sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
)
if result:
results.append(result)
else:
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = [
executor.submit(
self._evaluate_one,
data_type,
sample_dir,
sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
)
for sample_dir in sample_dirs
]
for future in as_completed(futures):
result = future.result()
if result:
results.append(result)
payload = {"data_type": data_type, "results": results, "summary": summarize_gebench_results(results)}
write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
return payload
def _evaluate_one(self, data_type: str, sample_dir: Path, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
lang_device = sample_dir.parent.name
sample_name = sample_dir.name
dataset_sample = sample_spec.sample_path
metadata = sample_spec.metadata
if data_type == "type1":
source = _resolve_referenced_image(
metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
)
generated_path = find_first_image(sample_dir)
if source is None or generated_path is None:
return None
generated = Image.open(generated_path).convert("RGB")
raw_scores = self.judge.evaluate(prompt=_type1_prompt(metadata), images=[source, generated])
elif data_type == "type2":
frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
prompt_suffix, judge_images = _trajectory_judge_payload(frames)
raw_scores = self.judge.evaluate(
prompt=f"Evaluate a six-frame GUI trajectory.\nTask: {goal}\n{prompt_suffix}",
images=judge_images,
)
elif data_type in {"type3", "type4"}:
frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
instruction = _text_or_default(metadata.get("instruction") or metadata.get("caption"), "Complete the task.")
prompt_suffix, judge_images = _trajectory_judge_payload(frames)
raw_scores = self.judge.evaluate(
prompt=f"Evaluate a six-frame GUI trajectory.\nInstruction: {instruction}\n{prompt_suffix}",
images=judge_images,
)
elif data_type == "type5":
source = _resolve_referenced_image(
metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
)
generated_path = find_first_image(sample_dir)
if source is None or generated_path is None:
return None
generated = Image.open(generated_path).convert("RGB")
raw_scores = self.judge.evaluate(prompt=_type5_prompt(metadata), images=[source, generated])
else:
raise ValueError(f"Unsupported data type: {data_type}")
scores = _normalize_scores(raw_scores)
return {
"sample_name": f"{lang_device}/{sample_name}",
"data_type": data_type,
"scores": scores,
"overall": _compute_overall(scores),
"raw_scores": raw_scores,
}
def _data_types_arg(value: str) -> list[str]:
return list(TYPE_TO_FOLDER.keys()) if value == "all" else [value]
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Run local GEBench generation and scoring against vLLM-Omni.")
subparsers = parser.add_subparsers(dest="command", required=True)
generate = subparsers.add_parser("generate")
generate.add_argument("--dataset-root", type=Path, required=True)
generate.add_argument("--output-root", type=Path, required=True)
generate.add_argument("--base-url", type=str, required=True)
generate.add_argument("--model", type=str, required=True)
generate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
generate.add_argument("--api-key", type=str, default="EMPTY")
generate.add_argument("--width", type=int, default=768)
generate.add_argument("--height", type=int, default=576)
generate.add_argument("--num-inference-steps", type=int, default=8)
generate.add_argument("--output-compression", type=int, default=98)
generate.add_argument("--guidance-scale", type=float, default=None)
generate.add_argument("--seed", type=int, default=42)
generate.add_argument("--workers", type=int, default=1)
generate.add_argument("--max-samples", type=int, default=None)
generate.add_argument("--samples-per-type", type=int, default=None)
evaluate = subparsers.add_parser("evaluate")
evaluate.add_argument("--dataset-root", type=Path, required=True)
evaluate.add_argument("--output-root", type=Path, required=True)
evaluate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
evaluate.add_argument("--judge-base-url", type=str, required=True)
evaluate.add_argument("--judge-model", type=str, required=True)
evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
evaluate.add_argument("--workers", type=int, default=1)
evaluate.add_argument("--max-samples", type=int, default=None)
evaluate.add_argument("--samples-per-type", type=int, default=None)
summarize = subparsers.add_parser("summarize")
summarize.add_argument("--output-root", type=Path, required=True)
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
if args.command == "generate":
runner = GEBenchRunner(
dataset_root=args.dataset_root,
output_root=args.output_root,
base_url=args.base_url,
model=args.model,
api_key=args.api_key,
width=args.width,
height=args.height,
num_inference_steps=args.num_inference_steps,
output_compression=args.output_compression,
guidance_scale=args.guidance_scale,
seed=args.seed,
)
records: list[dict[str, Any]] = []
for data_type in _data_types_arg(args.data_type):
records.extend(
runner.generate(
data_type=data_type,
workers=args.workers,
max_samples=args.max_samples,
samples_per_type=args.samples_per_type,
)
)
payload = {"records": records, "summary": summarize_generated_records(records)}
write_json(args.output_root / "generation_manifest.json", payload)
return 0
if args.command == "evaluate":
judge = LocalJudgeClient(
base_url=args.judge_base_url,
api_key=args.judge_api_key,
model=args.judge_model,
)
evaluator = GEBenchEvaluator(dataset_root=args.dataset_root, output_root=args.output_root, judge=judge)
combined_results: list[dict[str, Any]] = []
for data_type in _data_types_arg(args.data_type):
payload = evaluator.evaluate(
data_type=data_type,
workers=args.workers,
max_samples=args.max_samples,
samples_per_type=args.samples_per_type,
)
combined_results.extend(payload["results"])
_write_json_with_timestamp(
args.output_root / "evaluations" / "summary.json",
{"summary": summarize_gebench_results(combined_results)},
)
return 0
if args.command == "summarize":
generation_summary = collect_gebench_generation_summary(args.output_root)
evaluation_dir = args.output_root / "evaluations"
result_records: list[dict[str, Any]] = []
if evaluation_dir.exists():
for file_path in sorted(evaluation_dir.glob("type*.json")):
payload = load_json(file_path)
result_records.extend(payload.get("results", []))
payload: dict[str, Any] = {"generation": generation_summary}
if result_records:
payload["evaluation"] = summarize_gebench_results(result_records)
_write_json_with_timestamp(args.output_root / "summary.json", payload)
print(json.dumps(payload, indent=2, ensure_ascii=False))
return 0
parser.error(f"Unknown command: {args.command}")
return 1
@@ -1,13 +0,0 @@
# ruff: noqa: E402, I001
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[3]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from benchmarks.accuracy.text_to_image.gbench import main
if __name__ == "__main__":
raise SystemExit(main())
+1 -22
View File
@@ -43,27 +43,6 @@ The benchmark supports three dataset modes via `--dataset`:
### VBench dataset
`vbench` only provides prompt data (and image paths for i2v/i2i); it does not carry
per-request generation fields. In this mode, all requests share CLI values:
`--width --height --num-frames --fps --num-inference-steps`
(pass `--width` and `--height` together).
Example (`t2v`):
```bash
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--base-url http://localhost:8099 \
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--task t2v \
--dataset vbench \
--num-prompts 50 \
--width 640 --height 480 \
--num-frames 81 --fps 16 \
--num-inference-steps 40
```
Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths).
If you use i2v/i2i bench datasets and need auto-download support, you may need:
```bash
@@ -111,7 +90,7 @@ Common optional flags:
Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`.
- For `vbench` / `random`: these CLI flags act as global defaults for all generated requests.
- For `trace`: requests can carry their own fields (e.g., `width/height/num_frames/num_inference_steps`), with overrides/fallbacks as below.
- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`).
Precedence rules for `trace` (i.e., what actually gets sent):
+4 -151
View File
@@ -1,4 +1,3 @@
import asyncio
import base64
import mimetypes
import os
@@ -36,7 +35,6 @@ class RequestFuncOutput:
error: str = ""
start_time: float = 0.0
response_body: dict[str, Any] = field(default_factory=dict)
stage_durations: dict[str, float] = field(default_factory=dict)
peak_memory_mb: float = 0.0
slo_achieved: bool | None = None
@@ -57,7 +55,6 @@ async def async_request_chat_completions(
input: RequestFuncInput,
session: aiohttp.ClientSession,
pbar: tqdm | None = None,
enable_diffusion_pipeline_profiler: bool = False,
) -> RequestFuncOutput:
output = RequestFuncOutput()
output.start_time = time.perf_counter()
@@ -109,19 +106,8 @@ async def async_request_chat_completions(
resp_json = await response.json()
output.response_body = resp_json
output.success = True
try:
choices = resp_json.get("choices", [])
if choices and isinstance(choices, list):
msg = choices[0].get("message", {})
if isinstance(msg, dict):
content = msg.get("content", [])
if content and isinstance(content, list) and len(content) > 0:
first_item = content[0]
if isinstance(first_item, dict):
output.stage_durations = first_item.get("stage_durations") or {}
output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
except (IndexError, TypeError, AttributeError):
pass
if "peak_memory_mb" in resp_json:
output.peak_memory_mb = resp_json["peak_memory_mb"]
else:
output.error = f"HTTP {response.status}: {await response.text()}"
output.success = False
@@ -206,140 +192,7 @@ async def async_request_openai_images(
return output
async def async_request_v1_videos(
input: RequestFuncInput,
session: aiohttp.ClientSession,
pbar: tqdm | None = None,
) -> RequestFuncOutput:
output = RequestFuncOutput()
output.start_time = time.perf_counter()
files = dict(input.extra_body)
if input.prompt:
files.setdefault("prompt", input.prompt)
if input.width and input.height:
files.setdefault("height", input.height)
files.setdefault("width", input.width)
if input.num_frames:
files.setdefault("num_frames", input.num_frames)
if input.num_inference_steps:
files.setdefault("num_inference_steps", input.num_inference_steps)
if input.seed is not None:
files.setdefault("seed", input.seed)
if input.fps:
files.setdefault("fps", input.fps)
form = aiohttp.FormData()
for k, v in files.items():
form.add_field(k, str(v))
image_file = None
if input.image_paths and len(input.image_paths) > 0:
image_path = input.image_paths[0]
image_file = open(image_path, "rb")
form.add_field(
"input_reference",
image_file,
filename=os.path.basename(image_path),
content_type="application/octet-stream",
)
job_id = None
job_status = None
poll_json = {}
resp_json = {}
try:
# invoke a post request (POST /v1/videos)
async with session.post(input.api_url, data=form) as response:
if response.status == 200:
resp_json = await response.json()
job_id = resp_json.get("id")
job_status = resp_json.get("status")
if not job_id or not job_status:
output.error = "API response missing job 'id' or 'status' field."
output.success = False
return output
else:
output.error = f"HTTP {response.status}: {await response.text()}"
output.success = False
return output
# invoke a poll request (GET /v1/videos/{video_id})
poll_interval = 2.0 # Unit(s)
timeout_seconds = 600.0
deadline = time.perf_counter() + timeout_seconds
job_url = f"{input.api_url}/{job_id}"
while job_status not in {"completed", "failed"}:
await asyncio.sleep(poll_interval)
async with session.get(job_url) as poll_response:
if poll_response.status != 200:
output.error = f"Polling failed HTTP {poll_response.status}: {await poll_response.text()}"
output.success = False
return output
poll_json = await poll_response.json()
job_status = poll_json.get("status")
if time.perf_counter() >= deadline:
output.error = f"Timed out waiting for video job {job_id} to complete."
output.success = False
return output
if job_status == "failed":
output.error = f"Video job failed: {poll_json}"
output.success = False
return output
# invoke a get request (GET /v1/videos/{video_id}/content)
content_url = f"{job_url}/content"
async with session.get(content_url) as content_response:
if content_response.status != 200:
output.error = (
f"Content retrieval failed HTTP {content_response.status}: {await content_response.text()}"
)
output.success = False
return output
video_bytes = await content_response.read()
output.response_body = video_bytes
output.success = True
if "peak_memory_mb" in poll_json:
output.peak_memory_mb = poll_json["peak_memory_mb"]
elif "peak_memory_mb" in resp_json:
output.peak_memory_mb = resp_json["peak_memory_mb"]
except Exception as e:
output.error = str(e)
output.success = False
finally:
if image_file is not None:
image_file.close()
if job_id is not None:
try:
async with session.delete(f"{input.api_url}/{job_id}") as _:
pass
except Exception as e:
print(f"Failed to clean up video job {job_id}: {e}")
output.latency = time.perf_counter() - output.start_time
if output.success and input.slo_ms is not None:
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
if pbar:
pbar.update(1)
return output
backends_function_mapping = {
"2i": {
"vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
"openai": (async_request_openai_images, "/v1/images/generations"),
},
"2v": {
"v1/videos": (async_request_v1_videos, "/v1/videos"),
},
"vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
"openai": (async_request_openai_images, "/v1/images/generations"),
}
@@ -1,4 +1,4 @@
# adapted from fastvideo
# adapted from sglang and fastvideo
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -9,7 +9,6 @@ If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
Supports multiple backends:
- vllm-omni: Uses /v1/chat/completions endpoint (default)
- openai: Uses /v1/images/generations endpoint
- v1/videos: Use /v1/videos endpoint
Usage:
# Video (vllm-omni backend)
@@ -29,16 +28,6 @@ Usage:
--backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
--height 1024 --width 1024
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni --dataset random --task t2i --num-prompts 1 \
--max-concurrency 1 --enable-negative-prompt \
--random-request-config '[
{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
{"width":768,"height":768,"num_inference_steps":20,"weight":0.25},
{"width":1024,"height":1024,"num_inference_steps":25,"weight":0.45},
{"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
]'
i2i:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni --dataset vbench --task i2i --num-prompts 10
@@ -49,16 +38,6 @@ Usage:
--backend openai --dataset vbench --task t2i --num-prompts 10 \
--height 1024 --width 1024 --port 3000
# Video (v1/vedeos)
t2v:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend v1/videos --dataset random --task t2v --num-prompts 1 \
--max-concurrency 1 --enable-negative-prompt \
--random-request-config '[
{"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1}
]'
"""
import argparse
@@ -66,10 +45,7 @@ import ast
import asyncio
import glob
import json
import logging
import os
import random
import tempfile
import time
import uuid
from abc import ABC, abstractmethod
@@ -81,11 +57,8 @@ import aiohttp
import numpy as np
import requests
from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
from PIL import Image
from tqdm.asyncio import tqdm
logger = logging.getLogger(__name__)
class BaseDataset(ABC):
def __init__(self, args, api_url: str, model: str):
@@ -289,9 +262,6 @@ class VBenchDataset(BaseDataset):
def _resize_data(self, data: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Resize data to match num_prompts."""
if not data:
raise ValueError("No benchmark data available. Install Pillow or provide --dataset-path.")
if not self.args.num_prompts:
return data
@@ -515,14 +485,6 @@ class TraceDataset(BaseDataset):
timestamp = self._coerce_float(row.get("timestamp"))
slo_ms = self._coerce_float(row.get("slo_ms"))
image_paths = row.get("image_paths")
if not image_paths:
single = row.get("image_path")
image_paths = [single] if single else None
if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i"]:
raise ValueError(
f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row."
)
override_w = self.args.width
override_h = self.args.height
@@ -554,66 +516,26 @@ class TraceDataset(BaseDataset):
class RandomDataset(BaseDataset):
def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool = False):
super().__init__(args, api_url, model)
def __init__(self, args, api_url: str, model: str):
self.args = args
self.api_url = api_url
self.model = model
self.num_prompts = args.num_prompts
self.enable_negative_prompt = enable_negative_prompt
self.random_request_config = getattr(args, "random_request_config", None)
if self.random_request_config:
self.random_request_config = json.loads(self.random_request_config)
self._weights = [p["weight"] for p in self.random_request_config]
self.random_request_config = [
{k: v for k, v in p.items() if k != "weight"} for p in self.random_request_config
]
seed = getattr(args, "random_request_seed", 42)
self._rng = random.Random(seed)
self._sampled_requests = self._rng.choices(
self.random_request_config,
weights=self._weights,
k=self.num_prompts,
)
else:
self._sampled_requests = None
# Random image generate
if self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
img = Image.new("RGB", (512, 512), (255, 255, 255))
image_path = os.path.join(tempfile.gettempdir(), "diffusion_benchmark_random_image.png")
self._random_image_path = [image_path]
img.save(image_path)
else:
self._random_image_path = None
def __len__(self) -> int:
return self.num_prompts
def __getitem__(self, idx: int) -> RequestFuncInput:
extra_body = {}
if self.enable_negative_prompt:
extra_body["negative_prompt"] = f"Negative prompt {idx} for benchmarking diffusion models"
params = {
"width": self.args.width,
"height": self.args.height,
"num_frames": self.args.num_frames,
"num_inference_steps": self.args.num_inference_steps,
"fps": self.args.fps,
}
if self._sampled_requests:
profile = self._sampled_requests[idx]
params.update(profile)
return RequestFuncInput(
prompt=f"Random prompt {idx} for benchmarking diffusion models",
api_url=self.api_url,
model=self.model,
width=self.args.width,
height=self.args.height,
num_frames=self.args.num_frames,
num_inference_steps=self.args.num_inference_steps,
seed=self.args.seed,
extra_body=extra_body,
image_paths=self._random_image_path,
**params,
fps=self.args.fps,
)
def get_requests(self) -> list[RequestFuncInput]:
@@ -725,19 +647,19 @@ async def iter_requests(
requests_list: list[RequestFuncInput],
request_rate: float,
) -> AsyncGenerator[RequestFuncInput, None]:
"""Yield requests using a Poisson process if request_rate is set.
"""Yield requests using a fixed interval if request_rate is set.
- If request_rate is inf, all requests are yielded immediately (no sleep).
- Otherwise, inter-arrival times follow an exponential distribution.
- Otherwise, requests are emitted at a fixed cadence of 1 / request_rate seconds.
"""
if request_rate != float("inf"):
if request_rate <= 0:
raise ValueError(f"request_rate must be positive or inf, got {request_rate}.")
interval_s = 1.0 / float(request_rate)
for i, req in enumerate(requests_list):
if request_rate != float("inf") and i > 0:
interval_s = random.expovariate(request_rate)
await asyncio.sleep(interval_s)
yield req
@@ -756,15 +678,6 @@ def calculate_metrics(
latencies = [o.latency for o in success_outputs]
peak_memories = [o.peak_memory_mb for o in success_outputs if o.peak_memory_mb > 0]
# Aggregate per-stage durations across all successful requests that reported them.
stage_duration_lists: dict[str, list[float]] = {}
for o in success_outputs:
for stage, duration in (o.stage_durations or {}).items():
stage_duration_lists.setdefault(stage, []).append(duration)
stage_durations_mean = {s: float(np.mean(v)) for s, v in stage_duration_lists.items()}
stage_durations_p50 = {s: float(np.percentile(v, 50)) for s, v in stage_duration_lists.items()}
stage_durations_p99 = {s: float(np.percentile(v, 99)) for s, v in stage_duration_lists.items()}
metrics = {
"duration": total_duration,
"completed_requests": num_success,
@@ -773,14 +686,10 @@ def calculate_metrics(
"latency_mean": np.mean(latencies) if latencies else 0,
"latency_median": np.median(latencies) if latencies else 0,
"latency_p99": np.percentile(latencies, 99) if latencies else 0,
"latency_p95": np.percentile(latencies, 95) if latencies else 0,
"latency_p50": np.percentile(latencies, 50) if latencies else 0,
"peak_memory_mb_max": max(peak_memories) if peak_memories else 0,
"peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0,
"peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0,
"stage_durations_mean": stage_durations_mean,
"stage_durations_p50": stage_durations_p50,
"stage_durations_p99": stage_durations_p99,
}
if slo_enabled:
@@ -833,32 +742,8 @@ async def benchmark(args):
if args.base_url is None:
args.base_url = f"http://{args.host}:{args.port}"
VIDEO_TASKS = {"t2v", "i2v", "ti2v"}
IMAGE_TASKS = {"t2i", "i2i", "ti2i"}
if args.task in VIDEO_TASKS:
task_type = "2v"
elif args.task in IMAGE_TASKS:
task_type = "2i"
else:
raise ValueError(
f"Unsupported task: '{args.task}'. "
f"Valid video tasks: {sorted(VIDEO_TASKS)}, "
f"Valid image tasks: {sorted(IMAGE_TASKS)}"
)
valid_backends = sorted(backends_function_mapping[task_type].keys())
if args.backend not in valid_backends:
logger.error(
f"Invalid backend '{args.backend}' for task '{args.task}' (task type: '{task_type}').\n"
f"Valid backends for this task type: {valid_backends}\n"
f"Example usage: --task {args.task} --backend {valid_backends[0]}"
)
raise ValueError("Backend validation failed. See log above for valid options.")
# Setup API URL and request function based on backend
request_func, api_url = backends_function_mapping[task_type][args.backend]
request_func, api_url = backends_function_mapping[args.backend]
api_url = f"{args.base_url}{api_url}"
if args.dataset == "vbench":
@@ -866,7 +751,7 @@ async def benchmark(args):
elif args.dataset == "trace":
dataset = TraceDataset(args, api_url, args.model)
elif args.dataset == "random":
dataset = RandomDataset(args, api_url, args.model, args.enable_negative_prompt)
dataset = RandomDataset(args, api_url, args.model)
else:
raise ValueError(f"Unknown dataset: {args.dataset}")
@@ -904,8 +789,6 @@ async def benchmark(args):
warm_req,
num_inference_steps=args.warmup_num_inference_steps,
)
if args.task == "t2v":
warm_req = replace(warm_req, num_frames=1)
warm_out = await limited_request_func(warm_req, session, None)
warmup_pairs.append((warm_req, warm_out))
@@ -964,7 +847,6 @@ async def benchmark(args):
print("{:<40} {:<15.4f}".format("Latency Mean (s):", metrics["latency_mean"]))
print("{:<40} {:<15.4f}".format("Latency Median (s):", metrics["latency_median"]))
print("{:<40} {:<15.4f}".format("Latency P99 (s):", metrics["latency_p99"]))
print("{:<40} {:<15.4f}".format("Latency P95 (s):", metrics["latency_p95"]))
if args.slo:
print(f"{'-' * 50}")
@@ -978,12 +860,6 @@ async def benchmark(args):
print("{:<40} {:<15.2f}".format("Peak Memory Mean (MB):", metrics["peak_memory_mb_mean"]))
print("{:<40} {:<15.2f}".format("Peak Memory Median (MB):", metrics["peak_memory_mb_median"]))
if metrics["stage_durations_mean"]:
print(f"{'-' * 50}")
print("Stage Durations Mean (s):")
for stage, val in metrics["stage_durations_mean"].items():
print("{:<40} {:<15.4f}".format(f" {stage}:", val))
print("\n" + "=" * 60)
if args.output_file:
@@ -1007,7 +883,7 @@ if __name__ == "__main__":
"--backend",
type=str,
default="vllm-omni",
choices=["vllm-omni", "openai", "v1/videos"],
choices=["vllm-omni", "openai"],
help="Backend to target the benchmark to.",
)
parser.add_argument(
@@ -1096,25 +972,6 @@ if __name__ == "__main__":
help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).",
)
parser.add_argument("--disable-tqdm", action="store_true", help="Disable progress bar.")
parser.add_argument(
"--enable-negative-prompt",
action="store_true",
default=False,
help="Generate negative prompts when using the random dataset.",
)
parser.add_argument(
"--random-request-config",
type=str,
default=None,
help=(
"JSON string defining random request profiles. "
"Each profile may contain: width, height, num_inference_steps, etc. "
"The 'weight' field controls sampling probability (relative weight). "
"Example: "
'[{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},'
'{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
),
)
args = parser.parse_args()
@@ -1,169 +0,0 @@
# Qwen-Image Serving Performance Dashboard
This document describes how to deploy and benchmark **Qwen-Image** using vLLM-Omni. It includes service startup configuration, acceleration-related options, benchmark methodology, dataset settings, and performance results.
---
# 1. Overview
Qwen-Image is a multimodal text-to-image generation model served through the vLLM-Omni infrastructure.
This document covers:
* Service launch configuration (including acceleration options)
* Benchmark scripts and usage
* Dataset and workload settings
* Performance measurement results
* Reproducibility guidelines
---
# 2. Test Environment
| Component | Specification |
|------------|----------------|
| GPU | NVIDIA A100-SXM4-80GB |
| Diffusion Attention Backend | FlashAttention |
# 3. Service Launch Configuration
## 3.1 Basic Serving Command
```bash
vllm serve Qwen/Qwen-Image --omni \
--port 8091
```
## 3.2 Key Parameters
| Parameter | Description |
| --------------------- | ------------------------ |
| `--cfg-parallel-size` | CFG parallelism degree |
| `--ulysses-degree` | Ulysses parallel degree |
| `--vae-patch-parallel-size` | VAE parallel degree |
| `--tensor-parallel-size` | Tensor parallelism degree |
Record these parameters when reporting performance results.
---
# 4. Benchmark Script
## 4.1 Benchmark Entry
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni \
--dataset <DATASET_NAME> \
--task t2i \
--num-prompts <N> \
--max-concurrency <C> \
--enable-negative-prompt \
--random-request-config <CFG>
```
## 4.2 Key Benchmark Arguments
| Parameter | Description |
| ---------------------- | --------------------------------- |
| `--backend` | Serving backend (use `vllm-omni`) |
| `--dataset` | Dataset name (`random` or custom) |
| `--task` | Task type (e.g., `t2i`) |
| `--num-prompts` | Total number of requests |
| `--max-concurrency` | Client-side concurrency |
| `--random-request-config`| JSON string defining random request |
---
# 5. Dataset & Workload Settings
## 5.1 Recommended Evaluation Configurations
### Dataset A ( 512 Resolution)
* Dataset: `random`
* Task: t2i
* Concurrency: 1
* Mix Resolution
```
[
{"width":512,"height":512,"num_inference_steps":20,"weight":1}
]
```
### Dataset B (1536 Resolution)
* Dataset: `random`
* Task: t2i
* Concurrency: 1
* Mix Resolution
```
[
{"width":1536,"height":1536,"num_inference_steps":35,"weight":1}
]
```
### Dataset C (Mix Resolution)
* Dataset: `random`
* Task: t2i
* Concurrency: 1
* Mix Resolution
```
[
{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
{"width":768,"height":768,"num_inference_steps":20,"weight":0.25},
{"width":1024,"height":1024,"num_inference_steps":25,"weight":0.45},
{"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
]
```
---
## 5.2 Example Benchmark Command
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni \
--dataset random \
--task t2i \
--num-prompts 1 \
--max-concurrency 1 \
--enable-negative-prompt \
--random-request-config '[
{"width":512,"height":512,"num_inference_steps":20,"weight":1}
]'
```
---
# 6. Performance Metrics
The following metrics are collected during benchmarking:
| Metric | Description | Unit |
| ------------------ | ----------------------------- | ------- |
| Mean Latency | Mean of latency | seconds |
| P99 Latency | P99 of latency | seconds |
---
# 7. Performance Results
| Dataset Configuration | Max Concur. | CFG | Usp | Tp | VAE Parallel | Mean Latency (s) | P99 Latency (s) |
|-----------------------|-----|-----|-----|----|--------------|------------------|------------------|
| Dataset A | 1 | 2 | 2 | Off | Off | 2.2087 | 2.2087 |
| Dataset B | 1 | 2 | 2 | Off | Off | 19.6739 | 19.6739 |
| Dataset C | 1 | 2 | 2 | Off | Off | 5.67259 | 18.6234 |
---
# 8. Reproducibility Checklist
To ensure consistent and comparable benchmark results:
* Record GPU type
* Record parallel configuration
* Record benchmark parameters (resolution, concurrency, number of prompts)
* Ensure no background workload on GPUs during testing
---
This document serves as the official Qwen-Image serving performance reference under vLLM-Omni.
@@ -1,170 +0,0 @@
# Wan2.2 Serving Performance Dashboard
This document describes how to deploy and benchmark **Wan-AI/Wan2.2-T2V-A14B-Diffusers** using vLLM-Omni. It includes service startup configuration, acceleration-related options, benchmark methodology, dataset settings, and performance results.
---
# 1. Overview
Wan-AI/Wan2.2-T2V-A14B-Diffusers is a multimodal text-to-video generation model served through the vLLM-Omni infrastructure.
This document covers:
* Service launch configuration (including acceleration options)
* Benchmark scripts and usage
* Dataset and workload settings
* Performance measurement results
* Reproducibility guidelines
---
# 2. Test Environment
| Component | Specification |
|------------|----------------|
| GPU | NVIDIA A100-SXM4-80GB |
| Diffusion Attention Backend | FlashAttention |
# 3. Service Launch Configuration
## 3.1 Basic Serving Command
```bash
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni \
--port 8091
```
## 3.2 Key Parameters
| Parameter | Description |
| --------------------- | ------------------------ |
| `--cfg-parallel-size` | CFG parallelism degree |
| `--ulysses-degree` | Ulysses parallel degree |
| `--vae-patch-parallel-size` | VAE parallel degree |
| `--tensor-parallel-size` | Tensor parallelism degree |
| `--use-hsdp` | Enable Hybrid Sharded Data Parallel to shard model weights across GPUs |
Record these parameters when reporting performance results.
---
# 4. Benchmark Script
## 4.1 Benchmark Entry
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend v1/videos \
--dataset <DATASET_NAME> \
--task t2v \
--num-prompts <N> \
--max-concurrency <C> \
--enable-negative-prompt \
--random-request-config <CFG>
```
## 4.2 Key Benchmark Arguments
| Parameter | Description |
| ---------------------- | --------------------------------- |
| `--backend` | Serving backend (use `v1/videos`) |
| `--dataset` | Dataset name (`random` or custom) |
| `--task` | Task type (e.g., `t2v`) |
| `--num-prompts` | Total number of requests |
| `--max-concurrency` | Client-side concurrency |
| `--random-request-config`| JSON string defining random request |
---
# 5. Dataset & Workload Settings
## 5.1 Recommended Evaluation Configurations
### Dataset A (480p)
* Dataset: `random`
* Task: t2v
* Concurrency: 1
* Mix Resolution
```
[
{"width":854,"height":480,"num_inference_steps":3,"num_frames":80,"fps":16,"weight":1}
]
```
### Dataset B (720p)
* Dataset: `random`
* Task: t2v
* Concurrency: 1
* Mix Resolution
```
[
{"width":1280,"height":720,"num_inference_steps":6,"num_frames":80,"fps":16,"weight":1}
]
```
### Dataset C (Mix Resolution)
* Dataset: `random`
* Task: t2v
* Concurrency: 1
* Mix Resolution
```
[
{"width":854,"height":480,"num_inference_steps":3,"num_frames":80,"fps":16,"weight":0.15},
{"width":854,"height":480,"num_inference_steps":4,"num_frames":120,"fps":24,"weight":0.25},
{"width":1280,"height":720,"num_inference_steps":6,"num_frames":80,"fps":16,"weight":0.6}
]
```
---
## 5.2 Example Benchmark Command
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend v1/videos \
--dataset random \
--task t2v \
--num-prompts 1 \
--max-concurrency 1 \
--enable-negative-prompt \
--random-request-config '[
{"width":854,"height":480,"num_inference_steps":18,"num_frames": 33,"fps":16",weight":1}
]'
```
---
# 6. Performance Metrics
The following metrics are collected during benchmarking:
| Metric | Description | Unit |
| ------------------ | ----------------------------- | ------- |
| Mean Latency | Mean of latency | seconds |
| P99 Latency | P99 of latency | seconds |
---
# 7. Performance Results
| Dataset Configuration | Max Concur. | CFG | Usp | Tp | Hsdp | VAE Parallel | Mean Latency (s) | P99 Latency (s) |
|-----------------------|-----|-----|-----|-----|----|--------------|------------------|------------------|
| Dataset A | 1 | 2 | 2 | 1 | On | 1 | 24.6766 | 24.6766 |
| Dataset A | 1 | 2 | 2 | 1 | On | 4 | 21.6810 | 21.6810 |
| Dataset B | 1 | 2 | 2 | 1 | On | 1 | 124.6639 | 124.6639 |
| Dataset B | 1 | 2 | 2 | 1 | On | 4 | 117.44 | 117.44 |
| Dataset C | 1 | 2 | 2 | 1 | On | 1 | 79.2175 | 124.2565 |
| Dataset C | 1 | 2 | 2 | 1 | On | 4 | 74.4977 | 117.710 |
---
# 8. Reproducibility Checklist
To ensure consistent and comparable benchmark results:
* Record GPU type
* Record parallel configuration
* Record benchmark parameters (resolution, concurrency, number of prompts)
* Ensure no background workload on GPUs during testing
---
This document serves as the official Wan2.2 serving performance reference under vLLM-Omni.
@@ -1,460 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Benchmark quantization quality loss for diffusion models (image & video).
Generates outputs with BF16 (baseline) and a quantized config using the same
seed, then computes LPIPS perceptual distance between them. Results are printed
as a Markdown table ready to paste into a PR description.
Requirements:
pip install lpips Pillow numpy
Image example (text-to-image):
python benchmarks/diffusion/quantization_quality.py \
--model Tongyi-MAI/Z-Image-Turbo \
--task t2i \
--quantization fp8 \
--prompts \
"an aerial view of a coral reef with crystal clear turquoise water" \
"a campfire in a dark forest with sparks rising into a starry sky" \
"a gourmet dessert plate with chocolate mousse and gold leaf" \
--height 1024 --width 1024 \
--num-inference-steps 50 --seed 42
Video example (text-to-video):
python benchmarks/diffusion/quantization_quality.py \
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--task t2v \
--quantization fp8 \
--prompts \
"A serene lakeside sunrise with mist over the water" \
"A cat walking across a wooden bridge in autumn" \
--height 720 --width 1280 \
--num-frames 81 --num-inference-steps 40 --seed 42
Multiple quantization methods:
python benchmarks/diffusion/quantization_quality.py \
--model Tongyi-MAI/Z-Image-Turbo \
--task t2i \
--quantization fp8 int8 bitsandbytes \
--prompts "a cup of coffee on the table" \
--height 1024 --width 1024 \
--num-inference-steps 50 --seed 42
Output directory structure (--output-dir, default: ./quant_bench_output):
quant_bench_output/
baseline/ # BF16 outputs
<method>/ # Quantized outputs per method
results.md # Markdown table
"""
import argparse
import gc
import time
from pathlib import Path
import numpy as np
import torch
def compute_lpips_images(
baseline_images: list,
quantized_images: list,
net: str = "alex",
) -> list[float]:
"""Compute LPIPS between paired lists of PIL images."""
import lpips
from torchvision import transforms
loss_fn = lpips.LPIPS(net=net).eval()
if torch.cuda.is_available():
loss_fn = loss_fn.cuda()
transform = transforms.Compose(
[
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
]
)
scores = []
for img_bl, img_qt in zip(baseline_images, quantized_images):
t_bl = transform(img_bl.convert("RGB")).unsqueeze(0)
t_qt = transform(img_qt.convert("RGB")).unsqueeze(0)
if torch.cuda.is_available():
t_bl, t_qt = t_bl.cuda(), t_qt.cuda()
with torch.no_grad():
score = loss_fn(t_bl, t_qt).item()
scores.append(score)
return scores
def compute_lpips_video(
baseline_frames: np.ndarray,
quantized_frames: np.ndarray,
net: str = "alex",
) -> float:
"""Compute mean per-frame LPIPS for a video pair.
Args:
baseline_frames: (F, H, W, C) float array in [0, 1].
quantized_frames: same shape.
Returns:
Mean LPIPS across all frames.
"""
import lpips
loss_fn = lpips.LPIPS(net=net).eval()
if torch.cuda.is_available():
loss_fn = loss_fn.cuda()
num_frames = min(len(baseline_frames), len(quantized_frames))
scores = []
for i in range(num_frames):
# Convert (H, W, C) float [0,1] -> (1, C, H, W) float [-1, 1]
f_bl = torch.from_numpy(baseline_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
f_qt = torch.from_numpy(quantized_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
if torch.cuda.is_available():
f_bl, f_qt = f_bl.cuda(), f_qt.cuda()
with torch.no_grad():
score = loss_fn(f_bl, f_qt).item()
scores.append(score)
return float(np.mean(scores))
def _build_omni_kwargs(args, quantization=None):
"""Build kwargs dict for Omni() constructor."""
from vllm_omni.diffusion.data import DiffusionParallelConfig
parallel_config = DiffusionParallelConfig(
ulysses_degree=args.ulysses_degree,
ring_degree=args.ring_degree,
tensor_parallel_size=args.tensor_parallel_size,
)
kwargs = {
"model": args.model,
"parallel_config": parallel_config,
"enforce_eager": args.enforce_eager,
}
if quantization:
kwargs["quantization_config"] = quantization
return kwargs
def _generate_image(omni, args, prompt, seed):
"""Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
torch.cuda.reset_peak_memory_stats()
start = time.perf_counter()
outputs = omni.generate(
{"prompt": prompt},
OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
generator=generator,
num_inference_steps=args.num_inference_steps,
),
)
elapsed = time.perf_counter() - start
peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
first = outputs[0]
req_out = first.request_output[0] if hasattr(first, "request_output") else first
img = req_out.images[0]
return img, elapsed, peak_mem
def _generate_video(omni, args, prompt, seed):
"""Generate a video and return (np.ndarray [F,H,W,C], time_seconds, memory_gib)."""
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
torch.cuda.reset_peak_memory_stats()
start = time.perf_counter()
outputs = omni.generate(
{"prompt": prompt, "negative_prompt": ""},
OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
generator=generator,
guidance_scale=args.guidance_scale,
num_inference_steps=args.num_inference_steps,
num_frames=args.num_frames,
),
)
elapsed = time.perf_counter() - start
peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
first = outputs[0]
if hasattr(first, "request_output") and isinstance(first.request_output, list):
inner = first.request_output[0]
if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
frames = inner.images[0] if inner.images else None
else:
frames = inner
elif hasattr(first, "images") and first.images:
frames = first.images
else:
raise ValueError("Could not extract video frames from output.")
if isinstance(frames, torch.Tensor):
video = frames.detach().cpu()
if video.dim() == 5:
video = video[0].permute(1, 2, 3, 0) if video.shape[1] in (3, 4) else video[0]
elif video.dim() == 4 and video.shape[0] in (3, 4):
video = video.permute(1, 2, 3, 0)
if video.is_floating_point():
video = video.clamp(-1, 1) * 0.5 + 0.5
frames_array = video.float().numpy()
else:
frames_array = np.asarray(frames)
if frames_array.ndim == 5:
frames_array = frames_array[0]
return frames_array, elapsed, peak_mem
def _unload_omni(omni):
"""Delete Omni instance and free GPU memory."""
del omni
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.synchronize()
def run_benchmark(args):
from vllm_omni.entrypoints.omni import Omni
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
is_video = args.task == "t2v"
prompts = args.prompts
seed = args.seed
# Determine configs to benchmark
configs = [] # list of (label, quantization_method)
for method in args.quantization:
configs.append((method, method))
# --- Baseline run ---
print("\n" + "=" * 60)
print("Running BF16 baseline...")
print("=" * 60)
bl_kwargs = _build_omni_kwargs(args, quantization=None)
omni_bl = Omni(**bl_kwargs)
baseline_outputs = {} # prompt -> (output, time, mem)
for prompt in prompts:
print(f" Generating: {prompt[:60]}...")
if is_video:
out, t, mem = _generate_video(omni_bl, args, prompt, seed)
else:
out, t, mem = _generate_image(omni_bl, args, prompt, seed)
baseline_outputs[prompt] = (out, t, mem)
bl_avg_time = np.mean([v[1] for v in baseline_outputs.values()])
bl_mem = baseline_outputs[prompts[0]][2] # use first prompt's memory
_unload_omni(omni_bl)
# Save baseline outputs
bl_dir = output_dir / "baseline"
bl_dir.mkdir(parents=True, exist_ok=True)
for i, prompt in enumerate(prompts):
out = baseline_outputs[prompt][0]
if is_video:
try:
from diffusers.utils import export_to_video
frames_list = list(out) if isinstance(out, np.ndarray) and out.ndim == 4 else out
export_to_video(frames_list, str(bl_dir / f"prompt_{i}.mp4"), fps=args.fps)
except ImportError:
np.save(bl_dir / f"prompt_{i}.npy", out)
else:
out.save(bl_dir / f"prompt_{i}.png")
# --- Quantized runs ---
all_results = [] # list of dicts
for config_label, quant_method in configs:
print(f"\n{'=' * 60}")
print(f"Running: {config_label}...")
print("=" * 60)
qt_kwargs = _build_omni_kwargs(args, quantization=quant_method)
omni_qt = Omni(**qt_kwargs)
qt_outputs = {}
for prompt in prompts:
print(f" Generating: {prompt[:60]}...")
if is_video:
out, t, mem = _generate_video(omni_qt, args, prompt, seed)
else:
out, t, mem = _generate_image(omni_qt, args, prompt, seed)
qt_outputs[prompt] = (out, t, mem)
qt_avg_time = np.mean([v[1] for v in qt_outputs.values()])
qt_mem = qt_outputs[prompts[0]][2]
_unload_omni(omni_qt)
# Save quantized outputs
qt_dir = output_dir / config_label.replace(" ", "_")
qt_dir.mkdir(parents=True, exist_ok=True)
# Compute LPIPS per prompt
per_prompt = []
for i, prompt in enumerate(prompts):
bl_out = baseline_outputs[prompt][0]
qt_out = qt_outputs[prompt][0]
if is_video:
lpips_score = compute_lpips_video(bl_out, qt_out, net=args.lpips_net)
try:
from diffusers.utils import export_to_video
frames_list = list(qt_out) if isinstance(qt_out, np.ndarray) and qt_out.ndim == 4 else qt_out
export_to_video(frames_list, str(qt_dir / f"prompt_{i}.mp4"), fps=args.fps)
except ImportError:
np.save(qt_dir / f"prompt_{i}.npy", qt_out)
else:
lpips_score = compute_lpips_images([bl_out], [qt_out], net=args.lpips_net)[0]
qt_out.save(qt_dir / f"prompt_{i}.png")
per_prompt.append({"prompt": prompt, "lpips": lpips_score})
mean_lpips = np.mean([p["lpips"] for p in per_prompt])
speedup = bl_avg_time / qt_avg_time if qt_avg_time > 0 else float("inf")
mem_reduction = (bl_mem - qt_mem) / bl_mem * 100
all_results.append(
{
"config": config_label,
"avg_time": qt_avg_time,
"speedup": speedup,
"memory_gib": qt_mem,
"mem_reduction_pct": mem_reduction,
"mean_lpips": mean_lpips,
"per_prompt": per_prompt,
}
)
# --- Print results ---
print("\n\n")
print("=" * 80)
print("RESULTS")
print("=" * 80)
# Summary table
lines = []
lines.append(f"## Quantization Quality Benchmark — {args.model.split('/')[-1]}")
lines.append(
f"Setup: {args.height}x{args.width}, {args.num_inference_steps} steps, "
f"seed={args.seed}, LPIPS ({args.lpips_net})"
)
if is_video:
lines.append(f"Video: {args.num_frames} frames")
lines.append("")
lines.append("### Summary")
lines.append("")
lines.append("| Config | Avg Time | Speedup | Memory (GiB) | Mem Reduction | Mean LPIPS |")
lines.append("|--------|----------|---------|--------------|---------------|------------|")
lines.append(f"| BF16 baseline | {bl_avg_time:.2f}s | 1.00x | {bl_mem:.2f} | — | (ref) |")
for r in all_results:
lines.append(
f"| {r['config']} | {r['avg_time']:.2f}s | {r['speedup']:.2f}x "
f"| {r['memory_gib']:.2f} | {r['mem_reduction_pct']:.0f}% "
f"| {r['mean_lpips']:.4f} |"
)
lines.append("")
lines.append("> LPIPS < 0.01 = imperceptible, > 0.1 = clearly noticeable.")
lines.append("")
# Per-prompt table
if len(prompts) > 1:
lines.append("### Per-Prompt LPIPS")
lines.append("")
header = "| Prompt |"
sep = "|--------|"
for r in all_results:
header += f" {r['config']} |"
sep += "--------|"
lines.append(header)
lines.append(sep)
for i, prompt in enumerate(prompts):
short = prompt[:50] + "..." if len(prompt) > 50 else prompt
row = f"| {short} |"
for r in all_results:
row += f" {r['per_prompt'][i]['lpips']:.4f} |"
lines.append(row)
lines.append("")
md = "\n".join(lines)
print(md)
# Save markdown
results_path = output_dir / "results.md"
results_path.write_text(md, encoding="utf-8")
print(f"\nResults saved to {results_path}")
print(f"Baseline outputs in {bl_dir}")
for r in all_results:
qt_dir = output_dir / r["config"].replace(" ", "_")
print(f"Quantized outputs in {qt_dir}")
def parse_args():
parser = argparse.ArgumentParser(
description="Benchmark quantization quality loss for diffusion models.",
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument("--model", required=True, help="Model name or local path.")
parser.add_argument(
"--task",
default="t2i",
choices=["t2i", "t2v"],
help="Task type: t2i (text-to-image) or t2v (text-to-video).",
)
parser.add_argument(
"--quantization",
nargs="+",
required=True,
help="One or more quantization methods to benchmark (e.g. fp8 int8 bitsandbytes).",
)
parser.add_argument(
"--prompts",
nargs="+",
default=["a cup of coffee on the table"],
help="One or more prompts to generate.",
)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--height", type=int, default=1024)
parser.add_argument("--width", type=int, default=1024)
parser.add_argument("--num-inference-steps", type=int, default=50)
parser.add_argument("--num-frames", type=int, default=81, help="Number of video frames (t2v only).")
parser.add_argument("--fps", type=int, default=24, help="Video FPS for saving (t2v only).")
parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (used for video).")
parser.add_argument("--output-dir", type=str, default="./quant_bench_output", help="Directory to save outputs.")
parser.add_argument(
"--lpips-net",
type=str,
default="alex",
choices=["alex", "vgg", "squeeze"],
help="LPIPS backbone network.",
)
parser.add_argument("--ulysses-degree", type=int, default=1)
parser.add_argument("--ring-degree", type=int, default=1)
parser.add_argument("--tensor-parallel-size", type=int, default=1)
parser.add_argument("--enforce-eager", action="store_true")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
run_benchmark(args)
+3 -3
View File
@@ -56,9 +56,9 @@ What it does:
- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
- Uses `benchmarks/build_dataset/top100.txt` and writes to:
- Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
- `omni_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
- `omni_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
- `omni_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
- `omni_llm_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
- `omni_llm_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
- `omni_llm_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
- Outputs: `benchmarks/qwen3-omni/vllm_omni/outputs/` — ~100 text and `.wav` files.
Key checks:
@@ -26,12 +26,12 @@ else
--log-stats \
--log-dir $log_dir
echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
echo " - omni_pipeline_text run dir/base name"
echo " - omni_pipeline_text.orchestrator.stats.jsonl orchestrator-stage latency stats"
echo " - omni_pipeline_text.overall.stats.jsonl overall latency/TPS stats"
echo " - omni_pipeline_text.stage0.log per-stage detailed logs"
echo " - omni_pipeline_text.stage1.log"
echo " - omni_pipeline_text.stage2.log"
echo " - omni_llm_pipeline_text run dir/base name"
echo " - omni_llm_pipeline_text.orchestrator.stats.jsonl orchestrator-stage latency stats"
echo " - omni_llm_pipeline_text.overall.stats.jsonl overall latency/TPS stats"
echo " - omni_llm_pipeline_text.stage0.log per-stage detailed logs"
echo " - omni_llm_pipeline_text.stage1.log"
echo " - omni_llm_pipeline_text.stage2.log"
echo "Key checks: overall.stats.jsonl for end-to-end latency/TPS; orchestrator.stats.jsonl for stable per-stage latency; stage*.log for errors or long tails."
echo " - outputs/ Generated txt and wav files, there should be 100 text and wav files generated respectively"
fi
-100
View File
@@ -1,100 +0,0 @@
# Qwen3-TTS Benchmark
Benchmarks for Qwen3-TTS text-to-speech models, comparing vLLM-Omni streaming serving against HuggingFace Transformers offline inference.
## Prerequisites
```bash
pip install matplotlib aiohttp soundfile numpy tqdm
pip install qwen_tts # for HF baseline
```
## Quick Start
Run the full benchmark (vllm-omni + HF baseline) with a single command:
```bash
cd benchmarks/qwen3-tts
bash run_benchmark.sh
```
Results (JSON + PNG plots) are saved to `results/`.
### Common options
```bash
# Only vllm-omni (skip HF baseline)
bash run_benchmark.sh --async-only
# Only HF baseline
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 bs16 config for higher throughput
STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs16.yaml bash run_benchmark.sh --async-only
# Custom GPU, prompt count, concurrency levels
GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
```
## Manual Steps
### 1) Start the vLLM-Omni server
```bash
CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
"Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
--omni --host 127.0.0.1 --port 8000 \
--stage-configs-path benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml \
--trust-remote-code
```
### 2) Run online serving benchmark
```bash
python benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py \
--port 8000 \
--num-prompts 50 \
--max-concurrency 1 4 10 \
--config-name "async_chunk" \
--result-dir results/
```
### 3) Run HuggingFace baseline
```bash
python benchmarks/qwen3-tts/transformers/bench_tts_hf.py \
--model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
--num-prompts 50 \
--gpu-device 0 \
--result-dir results/
```
### 4) Generate comparison plots
```bash
python benchmarks/qwen3-tts/plot_results.py \
--results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
--labels "vllm-omni" "hf_transformers" \
--output results/comparison.png
```
## Stage Configs
| Config | max_num_seqs | Description |
|--------|:------------:|-------------|
| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lowest latency) |
| `vllm_omni/configs/qwen3_tts_bs16.yaml` | 16 | High-throughput concurrent processing |
All configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
The model is specified via the CLI `--model` flag (or `MODEL` env var), so the same configs work for both the 0.6B and 1.7B model variants.
## Metrics
- **TTFP (Time to First Audio Packet)**: Time from request to first audio chunk (streaming latency)
- **E2E (End-to-End Latency)**: Total time from request to complete audio response
- **RTF (Real-Time Factor)**: E2E latency / audio duration. RTF < 1.0 means faster-than-real-time synthesis
- **Throughput**: Total audio seconds generated per wall-clock second
-254
View File
@@ -1,254 +0,0 @@
"""Plot Qwen3-TTS benchmark results.
Generates comparison bar charts similar to the async_chunk design doc:
- TTFP (Time-to-First-Packet) across concurrency levels
- E2E latency across concurrency levels
- RTF (Real-Time Factor) across concurrency levels
Usage:
# Compare two configs (async_chunk vs no_async_chunk):
python plot_results.py \
--results results/bench_async_chunk_*.json results/bench_no_async_chunk_*.json \
--labels "async_chunk" "no_async_chunk" \
--output results/qwen3_tts_benchmark.png
# Single config:
python plot_results.py \
--results results/bench_async_chunk_*.json \
--labels "async_chunk" \
--output results/qwen3_tts_benchmark.png
"""
import argparse
import json
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
def load_results(result_files: list[str]) -> list[list[dict]]:
"""Load benchmark results from JSON files."""
all_results = []
for f in result_files:
with open(f) as fh:
data = json.load(fh)
all_results.append(data)
return all_results
def plot_comparison(
all_results: list[list[dict]],
labels: list[str],
output_path: str,
title_prefix: str = "Qwen3-TTS",
):
"""Generate comparison bar charts."""
n_configs = len(all_results)
# Collect concurrency levels present in ALL configs (skip missing data)
all_concurrencies = [set(r["concurrency"] for r in results) for results in all_results]
concurrencies = sorted(set.union(*all_concurrencies))
# Build data arrays, using None for missing concurrency levels
ttfp_data = {label: [] for label in labels}
e2e_data = {label: [] for label in labels}
rtf_data = {label: [] for label in labels}
throughput_data = {label: [] for label in labels}
for results, label in zip(all_results, labels):
conc_map = {r["concurrency"]: r for r in results}
for c in concurrencies:
r = conc_map.get(c)
ttfp_data[label].append(r["mean_ttfp_ms"] if r else None)
e2e_data[label].append(r["mean_e2e_ms"] if r else None)
rtf_data[label].append(r["mean_rtf"] if r else None)
throughput_data[label].append(r["audio_throughput"] if r else None)
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
fig.suptitle(f"{title_prefix} Performance Benchmark", fontsize=16, fontweight="bold")
x = np.arange(len(concurrencies))
width = 0.35 if n_configs == 2 else 0.5
if n_configs > 1:
offsets = np.linspace(-width / 2 * (n_configs - 1), width / 2 * (n_configs - 1), n_configs)
else:
offsets = [0]
colors = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107"]
def plot_metric(ax, data_dict, ylabel, title, fmt=".1f"):
bars = []
for i, (label, values) in enumerate(data_dict.items()):
# Replace None with 0 for plotting, but track which are missing
plot_values = [v if v is not None else 0 for v in values]
color = colors[i % len(colors)]
bar = ax.bar(x + offsets[i], plot_values, width, label=label, color=color, alpha=0.85)
bars.append(bar)
# Add value labels on bars (skip None/missing data)
max_val = max((v for v in values if v is not None), default=1)
for rect, val in zip(bar, values):
if val is not None and val > 0:
ax.text(
rect.get_x() + rect.get_width() / 2,
rect.get_height() + max_val * 0.02,
f"{val:{fmt}}",
ha="center",
va="bottom",
fontsize=9,
fontweight="bold",
)
ax.set_xlabel("Concurrency", fontsize=12)
ax.set_ylabel(ylabel, fontsize=12)
ax.set_title(title, fontsize=13, fontweight="bold")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.legend(fontsize=10)
ax.grid(axis="y", alpha=0.3)
ax.set_axisbelow(True)
plot_metric(axes[0, 0], ttfp_data, "TTFP (ms)", "Time to First Audio Packet (TTFP)")
plot_metric(axes[0, 1], e2e_data, "E2E Latency (ms)", "End-to-End Latency (E2E)")
plot_metric(axes[1, 0], rtf_data, "RTF", "Real-Time Factor (RTF)", fmt=".3f")
plot_metric(axes[1, 1], throughput_data, "Audio-sec / Wall-sec", "Audio Throughput", fmt=".2f")
plt.tight_layout()
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
def plot_single_summary(results: list[dict], label: str, output_path: str):
"""Generate a single-config summary with percentile breakdown."""
concurrencies = [r["concurrency"] for r in results]
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
fig.suptitle(f"Qwen3-TTS Benchmark - {label}", fontsize=15, fontweight="bold")
# TTFP breakdown
ax = axes[0]
means = [r["mean_ttfp_ms"] for r in results]
medians = [r["median_ttfp_ms"] for r in results]
p90s = [r["p90_ttfp_ms"] for r in results]
p99s = [r["p99_ttfp_ms"] for r in results]
x = np.arange(len(concurrencies))
w = 0.2
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_xlabel("Concurrency")
ax.set_ylabel("TTFP (ms)")
ax.set_title("Time to First Audio Packet")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3)
# E2E breakdown
ax = axes[1]
means = [r["mean_e2e_ms"] for r in results]
medians = [r["median_e2e_ms"] for r in results]
p90s = [r["p90_e2e_ms"] for r in results]
p99s = [r["p99_e2e_ms"] for r in results]
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_xlabel("Concurrency")
ax.set_ylabel("E2E Latency (ms)")
ax.set_title("End-to-End Latency")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3)
# RTF
ax = axes[2]
means = [r["mean_rtf"] for r in results]
medians = [r["median_rtf"] for r in results]
ax.bar(x - 0.15, means, 0.3, label="mean", color="#2196F3")
ax.bar(x + 0.15, medians, 0.3, label="median", color="#4CAF50")
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_xlabel("Concurrency")
ax.set_ylabel("RTF")
ax.set_title("Real-Time Factor")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3)
plt.tight_layout()
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
def print_comparison_table(all_results: list[list[dict]], labels: list[str]):
"""Print a markdown-formatted comparison table."""
concurrencies = sorted(set(r["concurrency"] for r in all_results[0]))
print("\n## Benchmark Results\n")
header = "| Metric | Concurrency |"
sep = "| --- | --- |"
for label in labels:
header += f" {label} |"
sep += " --- |"
print(header)
print(sep)
for metric, key, fmt in [
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
("E2E (ms)", "mean_e2e_ms", ".1f"),
("RTF", "mean_rtf", ".3f"),
("Throughput (audio-s/s)", "audio_throughput", ".2f"),
]:
for c in concurrencies:
row = f"| {metric} | {c} |"
for results in all_results:
conc_map = {r["concurrency"]: r for r in results}
val = conc_map.get(c, {}).get(key, 0)
row += f" {val:{fmt}} |"
print(row)
# Improvement calculation (only if 2 configs)
if len(all_results) == 2:
print(f"\n## Improvement ({labels[0]} vs {labels[1]})\n")
print("| Metric | Concurrency | Improvement |")
print("| --- | --- | --- |")
for metric, key in [("TTFP", "mean_ttfp_ms"), ("E2E", "mean_e2e_ms"), ("RTF", "mean_rtf")]:
for c in concurrencies:
m0 = {r["concurrency"]: r for r in all_results[0]}
m1 = {r["concurrency"]: r for r in all_results[1]}
v0 = m0.get(c, {}).get(key, 0)
v1 = m1.get(c, {}).get(key, 0)
if v1 > 0:
pct = (v1 - v0) / v1 * 100
print(f"| {metric} | {c} | {pct:+.1f}% |")
def parse_args():
parser = argparse.ArgumentParser(description="Plot Qwen3-TTS benchmark results")
parser.add_argument(
"--results", type=str, nargs="+", required=True, help="Path(s) to result JSON files (one per config)"
)
parser.add_argument(
"--labels", type=str, nargs="+", required=True, help="Labels for each config (must match --results count)"
)
parser.add_argument("--output", type=str, default="results/qwen3_tts_benchmark.png", help="Output image path")
parser.add_argument("--title", type=str, default="Qwen3-TTS", help="Title prefix for the plot")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
assert len(args.results) == len(args.labels), "--results and --labels must have the same count"
all_results = load_results(args.results)
print_comparison_table(all_results, args.labels)
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
if len(all_results) == 1:
plot_single_summary(all_results[0], args.labels[0], args.output)
else:
plot_comparison(all_results, args.labels, args.output, title_prefix=args.title)
-3
View File
@@ -1,3 +0,0 @@
# Benchmark results are machine-specific - do not commit
*
!.gitignore
-272
View File
@@ -1,272 +0,0 @@
#!/bin/bash
# Qwen3-TTS Benchmark Runner
#
# Compares vllm-omni streaming serving vs HuggingFace transformers offline inference.
# Produces JSON results and comparison plots.
#
# Usage:
# # Full comparison (vllm-omni + HF):
# bash run_benchmark.sh
#
# # Only vllm-omni async_chunk config:
# bash run_benchmark.sh --async-only
#
# # Only HuggingFace baseline:
# bash run_benchmark.sh --hf-only
#
# # vllm-omni only (skip HF):
# bash run_benchmark.sh --skip-hf
#
# # Custom settings:
# GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
#
# # Use 1.7B model:
# MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice 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
#
# Environment variables:
# GPU_DEVICE - GPU index to use (default: 0)
# NUM_PROMPTS - Number of prompts per concurrency level (default: 50)
# CONCURRENCY - Space-separated concurrency levels (default: "1 4 10")
# MODEL - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
# PORT - Server port (default: 8000)
# 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)
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
# Defaults
GPU_DEVICE="${GPU_DEVICE:-0}"
NUM_PROMPTS="${NUM_PROMPTS:-50}"
CONCURRENCY="${CONCURRENCY:-1 4 10}"
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
PORT="${PORT:-8000}"
GPU_MEM_TALKER="${GPU_MEM_TALKER:-0.3}"
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-0.2}"
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)"
# Parse args
RUN_ASYNC=true
RUN_HF=true
for arg in "$@"; do
case "$arg" in
--async-only) RUN_HF=false ;;
--hf-only) RUN_ASYNC=false ;;
--skip-hf) RUN_HF=false ;;
esac
done
mkdir -p "${RESULT_DIR}"
echo "============================================================"
echo " Qwen3-TTS Benchmark"
echo "============================================================"
echo " GPU: ${GPU_DEVICE}"
echo " Model: ${MODEL}"
echo " Prompts: ${NUM_PROMPTS}"
echo " Concurrency: ${CONCURRENCY}"
echo " Port: ${PORT}"
echo " Stage config: ${STAGE_CONFIG}"
echo " Results: ${RESULT_DIR}"
echo "============================================================"
# Prepare stage config with correct GPU device and memory settings
prepare_config() {
local config_template="$1"
local config_name="$2"
local output_path="${RESULT_DIR}/${config_name}_stage_config.yaml"
# Use sed to patch GPU device and memory utilization
sed \
-e "s/devices: \"0\"/devices: \"${GPU_DEVICE}\"/g" \
-e "s/gpu_memory_utilization: 0.3/gpu_memory_utilization: ${GPU_MEM_TALKER}/g" \
-e "s/gpu_memory_utilization: 0.2/gpu_memory_utilization: ${GPU_MEM_CODE2WAV}/g" \
"${config_template}" > "${output_path}"
echo "${output_path}"
}
# Start server and wait for it to be ready
start_server() {
local stage_config="$1"
local config_name="$2"
local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log"
echo ""
echo "Starting server with config: ${config_name}"
echo " Stage config: ${stage_config}"
echo " Log file: ${log_file}"
VLLM_WORKER_MULTIPROC_METHOD=spawn \
CUDA_VISIBLE_DEVICES="${GPU_DEVICE}" \
python -m vllm_omni.entrypoints.cli.main serve "${MODEL}" \
--omni \
--host 127.0.0.1 \
--port "${PORT}" \
--stage-configs-path "${stage_config}" \
--stage-init-timeout 120 \
--trust-remote-code \
--disable-log-stats \
> "${log_file}" 2>&1 &
SERVER_PID=$!
echo " Server PID: ${SERVER_PID}"
# Wait for server to be ready
echo " Waiting for server to be ready..."
local max_wait=300
local waited=0
while [ ${waited} -lt ${max_wait} ]; do
if curl -sf "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then
echo " Server is ready! (waited ${waited}s)"
return 0
fi
# Check if process is still alive
if ! kill -0 ${SERVER_PID} 2>/dev/null; then
echo " ERROR: Server process died. Check log: ${log_file}"
tail -20 "${log_file}"
return 1
fi
sleep 2
waited=$((waited + 2))
done
echo " ERROR: Server did not start within ${max_wait}s. Check log: ${log_file}"
kill ${SERVER_PID} 2>/dev/null || true
return 1
}
# Stop the server
stop_server() {
if [ -n "${SERVER_PID:-}" ]; then
echo " Stopping server (PID: ${SERVER_PID})..."
kill ${SERVER_PID} 2>/dev/null || true
wait ${SERVER_PID} 2>/dev/null || true
# Kill any remaining child processes on the port
local pids
pids=$(lsof -ti:${PORT} 2>/dev/null || true)
if [ -n "${pids}" ]; then
echo " Cleaning up remaining processes on port ${PORT}..."
echo "${pids}" | xargs kill -9 2>/dev/null || true
fi
echo " Server stopped."
SERVER_PID=""
fi
}
# Cleanup on exit
trap 'stop_server' EXIT
# Run benchmark for a given config
run_bench() {
local config_name="$1"
local config_template="$2"
echo ""
echo "============================================================"
echo " Benchmarking: ${config_name}"
echo "============================================================"
local stage_config
stage_config=$(prepare_config "${config_template}" "${config_name}")
start_server "${stage_config}" "${config_name}"
# Convert concurrency string to args
local conc_args=""
for c in ${CONCURRENCY}; do
conc_args="${conc_args} ${c}"
done
cd "${PROJECT_ROOT}"
python "${SCRIPT_DIR}/vllm_omni/bench_tts_serve.py" \
--host 127.0.0.1 \
--port "${PORT}" \
--num-prompts "${NUM_PROMPTS}" \
--max-concurrency ${conc_args} \
--num-warmups "${NUM_WARMUPS}" \
--config-name "${config_name}" \
--result-dir "${RESULT_DIR}"
stop_server
# Allow GPU memory to settle
sleep 5
}
# Run vllm-omni benchmark
if [ "${RUN_ASYNC}" = true ]; then
run_bench "async_chunk" "${SCRIPT_DIR}/${STAGE_CONFIG}"
fi
# Run HuggingFace baseline benchmark
if [ "${RUN_HF}" = true ]; then
echo ""
echo "============================================================"
echo " Benchmarking: HuggingFace transformers (offline)"
echo "============================================================"
cd "${PROJECT_ROOT}"
python "${SCRIPT_DIR}/transformers/bench_tts_hf.py" \
--model "${MODEL}" \
--num-prompts "${NUM_PROMPTS}" \
--num-warmups "${NUM_WARMUPS}" \
--gpu-device "${GPU_DEVICE}" \
--config-name "hf_transformers" \
--result-dir "${RESULT_DIR}"
# Allow GPU memory to settle
sleep 5
fi
# Plot results
echo ""
echo "============================================================"
echo " Generating plots..."
echo "============================================================"
RESULT_FILES=""
LABELS=""
if [ "${RUN_ASYNC}" = true ]; then
ASYNC_FILE=$(ls -t "${RESULT_DIR}"/bench_async_chunk_*.json 2>/dev/null | head -1)
if [ -n "${ASYNC_FILE}" ]; then
RESULT_FILES="${ASYNC_FILE}"
LABELS="async_chunk"
fi
fi
if [ "${RUN_HF}" = true ]; then
HF_FILE=$(ls -t "${RESULT_DIR}"/bench_hf_transformers_*.json 2>/dev/null | head -1)
if [ -n "${HF_FILE}" ]; then
if [ -n "${RESULT_FILES}" ]; then
RESULT_FILES="${RESULT_FILES} ${HF_FILE}"
LABELS="${LABELS} hf_transformers"
else
RESULT_FILES="${HF_FILE}"
LABELS="hf_transformers"
fi
fi
fi
if [ -n "${RESULT_FILES}" ]; then
python "${SCRIPT_DIR}/plot_results.py" \
--results ${RESULT_FILES} \
--labels ${LABELS} \
--output "${RESULT_DIR}/qwen3_tts_benchmark_${TIMESTAMP}.png"
fi
echo ""
echo "============================================================"
echo " Benchmark complete!"
echo " Results: ${RESULT_DIR}"
echo "============================================================"
@@ -1,281 +0,0 @@
"""Benchmark Qwen3-TTS using HuggingFace transformers (qwen_tts library).
Measures E2E latency, RTF, and audio duration for offline (non-serving) inference.
Results are saved in the same JSON format as bench_tts_serve.py for unified plotting.
Usage:
python bench_tts_hf.py \
--model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
--num-prompts 50 \
--num-warmups 3 \
--gpu-device 0 \
--result-dir results/
"""
import argparse
import json
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
import numpy as np
import soundfile as sf
import torch
PROMPTS = [
"Hello, welcome to the voice synthesis benchmark test.",
"She said she would be here by noon, but nobody showed up.",
"The quick brown fox jumps over the lazy dog near the riverbank.",
"I can't believe how beautiful the sunset looks from up here on the mountain.",
"Please remember to bring your identification documents to the appointment tomorrow morning.",
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
"The train leaves at half past seven, so we need to arrive at the station before then.",
"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.",
]
@dataclass
class BenchmarkResult:
config_name: str = ""
concurrency: int = 1 # always 1 for offline
num_prompts: int = 0
completed: int = 0
failed: int = 0
duration_s: float = 0.0
# TTFP stats - not applicable for HF offline, set to E2E for compatibility
mean_ttfp_ms: float = 0.0
median_ttfp_ms: float = 0.0
std_ttfp_ms: float = 0.0
p90_ttfp_ms: float = 0.0
p95_ttfp_ms: float = 0.0
p99_ttfp_ms: float = 0.0
# E2E stats (ms)
mean_e2e_ms: float = 0.0
median_e2e_ms: float = 0.0
std_e2e_ms: float = 0.0
p90_e2e_ms: float = 0.0
p95_e2e_ms: float = 0.0
p99_e2e_ms: float = 0.0
# RTF stats
mean_rtf: float = 0.0
median_rtf: float = 0.0
std_rtf: float = 0.0
p99_rtf: float = 0.0
# Audio stats
mean_audio_duration_s: float = 0.0
total_audio_duration_s: float = 0.0
audio_throughput: float = 0.0
request_throughput: float = 0.0
# Per-request details
per_request: list = field(default_factory=list)
def run_benchmark(args):
from qwen_tts import Qwen3TTSModel
device = f"cuda:{args.gpu_device}"
print(f"Loading model: {args.model} on {device}")
model = Qwen3TTSModel.from_pretrained(
args.model,
device_map=device,
dtype=torch.bfloat16,
)
print("Model loaded.")
# Build prompt list
prompts = [PROMPTS[i % len(PROMPTS)] for i in range(args.num_prompts)]
# Warmup
if args.num_warmups > 0:
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,
)
# Sync GPU
torch.cuda.synchronize(device)
print("Warmup done.")
# Benchmark
print(f"Running {args.num_prompts} requests sequentially...")
e2e_times = []
rtfs = []
audio_durations = []
per_request = []
failed = 0
audio_dir = None
if args.save_audio:
audio_dir = Path(args.result_dir) / "audio_hf"
audio_dir.mkdir(parents=True, exist_ok=True)
total_start = time.perf_counter()
for i, prompt in enumerate(prompts):
try:
torch.cuda.synchronize(device)
st = time.perf_counter()
wavs, sr = model.generate_custom_voice(
text=prompt,
language=args.language,
speaker=args.voice,
)
torch.cuda.synchronize(device)
elapsed = time.perf_counter() - st
# Compute audio duration
audio_samples = wavs[0]
if isinstance(audio_samples, torch.Tensor):
audio_samples = audio_samples.cpu().numpy()
audio_dur = len(audio_samples) / sr
rtf = elapsed / audio_dur if audio_dur > 0 else 0.0
e2e_times.append(elapsed)
rtfs.append(rtf)
audio_durations.append(audio_dur)
per_request.append(
{
"e2e_ms": elapsed * 1000,
"ttfp_ms": elapsed * 1000, # no streaming, TTFP = E2E
"rtf": rtf,
"audio_duration_s": audio_dur,
"prompt": prompt,
}
)
if audio_dir:
sf.write(str(audio_dir / f"output_{i:04d}.wav"), audio_samples, sr)
if (i + 1) % 10 == 0 or i == 0:
print(
f" [{i + 1}/{args.num_prompts}] e2e={elapsed * 1000:.0f}ms rtf={rtf:.3f} audio={audio_dur:.2f}s"
)
except Exception as e:
print(f" [{i + 1}/{args.num_prompts}] FAILED: {e}")
failed += 1
total_duration = time.perf_counter() - total_start
completed = len(e2e_times)
# Compute stats
result = BenchmarkResult(
config_name=args.config_name,
concurrency=1,
num_prompts=args.num_prompts,
completed=completed,
failed=failed,
duration_s=total_duration,
)
if e2e_times:
e2e_ms = [t * 1000 for t in e2e_times]
result.mean_e2e_ms = float(np.mean(e2e_ms))
result.median_e2e_ms = float(np.median(e2e_ms))
result.std_e2e_ms = float(np.std(e2e_ms))
result.p90_e2e_ms = float(np.percentile(e2e_ms, 90))
result.p95_e2e_ms = float(np.percentile(e2e_ms, 95))
result.p99_e2e_ms = float(np.percentile(e2e_ms, 99))
# For HF offline, TTFP = E2E (no streaming)
result.mean_ttfp_ms = result.mean_e2e_ms
result.median_ttfp_ms = result.median_e2e_ms
result.std_ttfp_ms = result.std_e2e_ms
result.p90_ttfp_ms = result.p90_e2e_ms
result.p95_ttfp_ms = result.p95_e2e_ms
result.p99_ttfp_ms = result.p99_e2e_ms
result.mean_rtf = float(np.mean(rtfs))
result.median_rtf = float(np.median(rtfs))
result.std_rtf = float(np.std(rtfs))
result.p99_rtf = float(np.percentile(rtfs, 99))
result.mean_audio_duration_s = float(np.mean(audio_durations))
result.total_audio_duration_s = float(np.sum(audio_durations))
result.audio_throughput = result.total_audio_duration_s / total_duration
result.request_throughput = completed / total_duration
result.per_request = per_request
# Print summary in standardized performance template
W = 50
print("")
print(f"{'=' * W}")
print(f"{'Serving Benchmark Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Successful requests:':<40}{completed:<10}")
print(f"{'Failed requests:':<40}{failed:<10}")
print(f"{'Maximum request concurrency:':<40}{1:<10}")
print(f"{'Benchmark duration (s):':<40}{total_duration:<10.2f}")
print(f"{'Request throughput (req/s):':<40}{result.request_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'End-to-end Latency':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean E2EL (ms):':<40}{result.mean_e2e_ms:<10.2f}")
print(f"{'Median E2EL (ms):':<40}{result.median_e2e_ms:<10.2f}")
print(f"{'P99 E2EL (ms):':<40}{result.p99_e2e_ms:<10.2f}")
print(f"{'=' * W}")
print(f"{'Audio Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Total audio duration generated (s):':<40}{result.total_audio_duration_s:<10.2f}")
print(f"{'Audio throughput (audio duration/s):':<40}{result.audio_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'Time to First Packet':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_TTFP (ms):':<40}{result.mean_ttfp_ms:<10.2f}")
print(f"{'Median AUDIO_TTFP (ms):':<40}{result.median_ttfp_ms:<10.2f}")
print(f"{'P99 AUDIO_TTFP (ms):':<40}{result.p99_ttfp_ms:<10.2f}")
print(f"{'-' * W}")
print(f"{'Real Time Factor':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_RTF:':<40}{result.mean_rtf:<10.3f}")
print(f"{'Median AUDIO_RTF:':<40}{result.median_rtf:<10.3f}")
print(f"{'P99 AUDIO_RTF:':<40}{result.p99_rtf:<10.3f}")
print(f"{'=' * W}")
print("")
# Save results (as a list with single concurrency=1 entry, matching serve format)
result_dir = Path(args.result_dir)
result_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
with open(result_file, "w") as f:
json.dump([asdict(result)], f, indent=2)
print(f"Results saved to {result_file}")
return result
def parse_args():
parser = argparse.ArgumentParser(description="Qwen3-TTS HuggingFace Benchmark")
parser.add_argument(
"--model", type=str, default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", help="HuggingFace model name or path"
)
parser.add_argument("--num-prompts", type=int, default=50)
parser.add_argument("--num-warmups", type=int, default=3)
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(
"--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
)
parser.add_argument("--result-dir", type=str, default="results")
parser.add_argument("--save-audio", action="store_true", help="Save generated audio files")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
run_benchmark(args)
@@ -1,301 +0,0 @@
"""Benchmark comparing async_chunk on vs off for Qwen3-TTS.
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across
concurrency levels for both async_chunk modes. Saves results as JSON.
Usage:
# Run against a server already serving with a given config:
python bench_async_chunk.py \
--host 127.0.0.1 --port 8000 \
--config-name async_chunk_on \
--num-prompts 50 \
--max-concurrency 1 10 \
--result-dir results/
"""
import argparse
import asyncio
import json
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
import aiohttp
import numpy as np
from tqdm.asyncio import tqdm
PROMPTS = [
"Hello, welcome to the voice synthesis benchmark test.",
"She said she would be here by noon, but nobody showed up.",
"The quick brown fox jumps over the lazy dog near the riverbank.",
"I can't believe how beautiful the sunset looks from up here on the mountain.",
"Please remember to bring your identification documents to the appointment tomorrow morning.",
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
"The train leaves at half past seven, so we need to arrive at the station before then.",
"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.",
]
@dataclass
class RequestResult:
success: bool = False
ttfp: float = 0.0
e2e: float = 0.0
audio_bytes: int = 0
audio_duration: float = 0.0
rtf: float = 0.0
prompt: str = ""
error: str = ""
@dataclass
class BenchmarkResult:
config_name: str = ""
concurrency: int = 0
num_prompts: int = 0
completed: int = 0
failed: int = 0
duration_s: float = 0.0
mean_ttfp_ms: float = 0.0
median_ttfp_ms: float = 0.0
std_ttfp_ms: float = 0.0
p90_ttfp_ms: float = 0.0
p95_ttfp_ms: float = 0.0
p99_ttfp_ms: float = 0.0
mean_e2e_ms: float = 0.0
median_e2e_ms: float = 0.0
std_e2e_ms: float = 0.0
p90_e2e_ms: float = 0.0
p95_e2e_ms: float = 0.0
p99_e2e_ms: float = 0.0
mean_rtf: float = 0.0
median_rtf: float = 0.0
std_rtf: float = 0.0
mean_audio_duration_s: float = 0.0
total_audio_duration_s: float = 0.0
audio_throughput: float = 0.0
request_throughput: float = 0.0
per_request: list = field(default_factory=list)
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
return num_bytes / sample_width / sample_rate
async def send_tts_request(
session: aiohttp.ClientSession,
api_url: str,
prompt: str,
voice: str = "vivian",
language: str = "English",
stream: bool = True,
pbar: tqdm | None = None,
) -> RequestResult:
payload = {
"input": prompt,
"voice": voice,
"language": language,
"stream": stream,
"response_format": "pcm",
}
result = RequestResult(prompt=prompt)
st = time.perf_counter()
try:
async with session.post(api_url, json=payload) as response:
if response.status != 200:
result.error = f"HTTP {response.status}: {await response.text()}"
return result
first_chunk = True
total_bytes = 0
async for chunk in response.content.iter_any():
if first_chunk and len(chunk) > 0:
result.ttfp = time.perf_counter() - st
first_chunk = False
total_bytes += len(chunk)
result.e2e = time.perf_counter() - st
result.audio_bytes = total_bytes
result.audio_duration = pcm_bytes_to_duration(total_bytes)
if result.audio_duration > 0:
result.rtf = result.e2e / result.audio_duration
result.success = True
except Exception as e:
result.error = str(e)
result.e2e = time.perf_counter() - st
if pbar:
pbar.update(1)
return result
async def run_benchmark(
host: str,
port: int,
num_prompts: int,
max_concurrency: int,
num_warmups: int = 3,
voice: str = "vivian",
language: str = "English",
stream: bool = True,
) -> BenchmarkResult:
api_url = f"http://{host}:{port}/v1/audio/speech"
connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600))
if num_warmups > 0:
print(f" Warming up with {num_warmups} requests...")
warmup_tasks = [
send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream)
for i in range(num_warmups)
]
await asyncio.gather(*warmup_tasks)
print(" Warmup done.")
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
semaphore = asyncio.Semaphore(max_concurrency)
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
async def limited_request(prompt):
async with semaphore:
return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar)
start_time = time.perf_counter()
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
results: list[RequestResult] = await asyncio.gather(*tasks)
duration = time.perf_counter() - start_time
pbar.close()
await session.close()
successful = [r for r in results if r.success]
failed = [r for r in results if not r.success]
bench = BenchmarkResult(
concurrency=max_concurrency,
num_prompts=num_prompts,
completed=len(successful),
failed=len(failed),
duration_s=duration,
)
if successful:
ttfps = [r.ttfp * 1000 for r in successful]
e2es = [r.e2e * 1000 for r in successful]
rtfs = [r.rtf for r in successful]
audio_durs = [r.audio_duration for r in successful]
bench.mean_ttfp_ms = float(np.mean(ttfps))
bench.median_ttfp_ms = float(np.median(ttfps))
bench.std_ttfp_ms = float(np.std(ttfps))
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
bench.mean_e2e_ms = float(np.mean(e2es))
bench.median_e2e_ms = float(np.median(e2es))
bench.std_e2e_ms = float(np.std(e2es))
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
bench.mean_rtf = float(np.mean(rtfs))
bench.median_rtf = float(np.median(rtfs))
bench.std_rtf = float(np.std(rtfs))
bench.mean_audio_duration_s = float(np.mean(audio_durs))
bench.total_audio_duration_s = float(np.sum(audio_durs))
bench.audio_throughput = bench.total_audio_duration_s / duration
bench.request_throughput = len(successful) / duration
bench.per_request = [
{
"ttfp_ms": r.ttfp * 1000,
"e2e_ms": r.e2e * 1000,
"rtf": r.rtf,
"audio_duration_s": r.audio_duration,
"prompt": r.prompt,
}
for r in successful
]
print(f"\n{'=' * 60}")
print(f" Concurrency: {max_concurrency} | Completed: {bench.completed} | Failed: {bench.failed}")
print(f" Duration: {duration:.2f}s | Throughput: {bench.request_throughput:.2f} req/s")
print(
f" TTFP (ms): mean={bench.mean_ttfp_ms:.1f} median={bench.median_ttfp_ms:.1f}"
f" p90={bench.p90_ttfp_ms:.1f} p99={bench.p99_ttfp_ms:.1f}"
)
print(
f" E2E (ms): mean={bench.mean_e2e_ms:.1f} median={bench.median_e2e_ms:.1f}"
f" p90={bench.p90_e2e_ms:.1f} p99={bench.p99_e2e_ms:.1f}"
)
print(f" RTF: mean={bench.mean_rtf:.3f} median={bench.median_rtf:.3f}")
print(f" Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec")
print(f"{'=' * 60}\n")
if failed:
for r in failed[:3]:
print(f" [ERROR] {r.error[:200]}")
return bench
async def main(args):
all_results = []
for concurrency in args.max_concurrency:
result = await run_benchmark(
host=args.host,
port=args.port,
num_prompts=args.num_prompts,
max_concurrency=concurrency,
num_warmups=args.num_warmups,
voice=args.voice,
language=args.language,
stream=args.stream,
)
result.config_name = args.config_name
all_results.append(asdict(result))
result_dir = Path(args.result_dir)
result_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
with open(result_file, "w") as f:
json.dump(all_results, f, indent=2)
print(f"Results saved to {result_file}")
def parse_args():
parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client")
parser.add_argument("--host", type=str, default="127.0.0.1")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--num-prompts", type=int, default=50)
parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10])
parser.add_argument("--num-warmups", type=int, default=3)
parser.add_argument("--voice", type=str, default="vivian")
parser.add_argument("--language", type=str, default="English")
parser.add_argument("--stream", action="store_true", default=True)
parser.add_argument("--no-stream", dest="stream", action="store_false")
parser.add_argument("--config-name", type=str, default="async_chunk_on")
parser.add_argument("--result-dir", type=str, default="results")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
asyncio.run(main(args))
@@ -1,348 +0,0 @@
"""Benchmark client for Qwen3-TTS via /v1/audio/speech endpoint.
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF (Real-Time Factor)
across configurable concurrency levels. Saves results as JSON for plotting.
Usage:
python bench_tts_serve.py \
--host 127.0.0.1 --port 8000 \
--num-prompts 50 \
--max-concurrency 1 4 10 \
--result-dir results/
"""
import argparse
import asyncio
import json
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
import aiohttp
import numpy as np
from tqdm.asyncio import tqdm
PROMPTS = [
"Hello, welcome to the voice synthesis benchmark test.",
"She said she would be here by noon, but nobody showed up.",
"The quick brown fox jumps over the lazy dog near the riverbank.",
"I can't believe how beautiful the sunset looks from up here on the mountain.",
"Please remember to bring your identification documents to the appointment tomorrow morning.",
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
"The train leaves at half past seven, so we need to arrive at the station before then.",
"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.",
]
@dataclass
class RequestResult:
success: bool = False
ttfp: float = 0.0 # Time to first audio packet (seconds)
e2e: float = 0.0 # End-to-end latency (seconds)
audio_bytes: int = 0 # Total audio bytes received
audio_duration: float = 0.0 # Audio duration in seconds (estimated from PCM)
rtf: float = 0.0 # Real-time factor = e2e / audio_duration
prompt: str = ""
error: str = ""
@dataclass
class BenchmarkResult:
config_name: str = ""
concurrency: int = 0
num_prompts: int = 0
completed: int = 0
failed: int = 0
duration_s: float = 0.0
# TTFP stats (ms)
mean_ttfp_ms: float = 0.0
median_ttfp_ms: float = 0.0
std_ttfp_ms: float = 0.0
p90_ttfp_ms: float = 0.0
p95_ttfp_ms: float = 0.0
p99_ttfp_ms: float = 0.0
# E2E stats (ms)
mean_e2e_ms: float = 0.0
median_e2e_ms: float = 0.0
std_e2e_ms: float = 0.0
p90_e2e_ms: float = 0.0
p95_e2e_ms: float = 0.0
p99_e2e_ms: float = 0.0
# RTF stats
mean_rtf: float = 0.0
median_rtf: float = 0.0
std_rtf: float = 0.0
p99_rtf: float = 0.0
# Audio stats
mean_audio_duration_s: float = 0.0
total_audio_duration_s: float = 0.0
audio_throughput: float = 0.0 # audio_duration / wall_time
request_throughput: float = 0.0 # requests / second
# Per-request details
per_request: list = field(default_factory=list)
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
"""Convert raw PCM byte count to duration in seconds."""
num_samples = num_bytes / sample_width
return num_samples / sample_rate
async def send_tts_request(
session: aiohttp.ClientSession,
api_url: str,
prompt: str,
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",
}
result = RequestResult(prompt=prompt)
st = time.perf_counter()
try:
async with session.post(api_url, json=payload) as response:
if response.status != 200:
result.error = f"HTTP {response.status}: {await response.text()}"
result.success = False
return result
first_chunk = True
total_bytes = 0
async for chunk in response.content.iter_any():
if first_chunk and len(chunk) > 0:
result.ttfp = time.perf_counter() - st
first_chunk = False
total_bytes += len(chunk)
result.e2e = time.perf_counter() - st
result.audio_bytes = total_bytes
result.audio_duration = pcm_bytes_to_duration(total_bytes)
if result.audio_duration > 0:
result.rtf = result.e2e / result.audio_duration
result.success = True
except Exception as e:
result.error = str(e)
result.success = False
result.e2e = time.perf_counter() - st
if pbar:
pbar.update(1)
return result
async def run_benchmark(
host: str,
port: int,
num_prompts: int,
max_concurrency: int,
num_warmups: int = 3,
voice: str = "vivian",
language: str = "English",
) -> BenchmarkResult:
"""Run benchmark at a given concurrency level."""
api_url = f"http://{host}:{port}/v1/audio/speech"
connector = aiohttp.TCPConnector(
limit=max_concurrency,
limit_per_host=max_concurrency,
keepalive_timeout=60,
)
session = aiohttp.ClientSession(
connector=connector,
timeout=aiohttp.ClientTimeout(total=600),
)
# Warmup
if num_warmups > 0:
print(f" Warming up with {num_warmups} requests...")
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))
await asyncio.gather(*warmup_tasks)
print(" Warmup done.")
# Build request list
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
# Run benchmark
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
semaphore = asyncio.Semaphore(max_concurrency)
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
async def limited_request(prompt):
async with semaphore:
return await send_tts_request(session, api_url, prompt, voice, language, pbar)
start_time = time.perf_counter()
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
results: list[RequestResult] = await asyncio.gather(*tasks)
duration = time.perf_counter() - start_time
pbar.close()
await session.close()
# Compute stats
successful = [r for r in results if r.success]
failed = [r for r in results if not r.success]
bench = BenchmarkResult(
concurrency=max_concurrency,
num_prompts=num_prompts,
completed=len(successful),
failed=len(failed),
duration_s=duration,
)
if successful:
ttfps = [r.ttfp * 1000 for r in successful] # convert to ms
e2es = [r.e2e * 1000 for r in successful]
rtfs = [r.rtf for r in successful]
audio_durs = [r.audio_duration for r in successful]
bench.mean_ttfp_ms = float(np.mean(ttfps))
bench.median_ttfp_ms = float(np.median(ttfps))
bench.std_ttfp_ms = float(np.std(ttfps))
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
bench.mean_e2e_ms = float(np.mean(e2es))
bench.median_e2e_ms = float(np.median(e2es))
bench.std_e2e_ms = float(np.std(e2es))
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
bench.mean_rtf = float(np.mean(rtfs))
bench.median_rtf = float(np.median(rtfs))
bench.std_rtf = float(np.std(rtfs))
bench.p99_rtf = float(np.percentile(rtfs, 99))
bench.mean_audio_duration_s = float(np.mean(audio_durs))
bench.total_audio_duration_s = float(np.sum(audio_durs))
bench.audio_throughput = bench.total_audio_duration_s / duration
bench.request_throughput = len(successful) / duration
bench.per_request = [
{
"ttfp_ms": r.ttfp * 1000,
"e2e_ms": r.e2e * 1000,
"rtf": r.rtf,
"audio_duration_s": r.audio_duration,
"prompt": r.prompt,
}
for r in successful
]
# Print summary in standardized performance template
W = 50
print("")
print(f"{'=' * W}")
print(f"{'Serving Benchmark Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Successful requests:':<40}{bench.completed:<10}")
print(f"{'Failed requests:':<40}{bench.failed:<10}")
print(f"{'Maximum request concurrency:':<40}{max_concurrency:<10}")
print(f"{'Benchmark duration (s):':<40}{duration:<10.2f}")
print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'End-to-end Latency':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
print(f"{'=' * W}")
print(f"{'Audio Result':^{W}}")
print(f"{'=' * W}")
print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
print(f"{'-' * W}")
print(f"{'Time to First Packet':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
print(f"{'-' * W}")
print(f"{'Real Time Factor':^{W}}")
print(f"{'-' * W}")
print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
print(f"{'=' * W}")
print("")
if failed:
for r in failed[:3]:
print(f" [ERROR] {r.error[:200]}")
return bench
async def main(args):
all_results = []
for concurrency in args.max_concurrency:
result = await run_benchmark(
host=args.host,
port=args.port,
num_prompts=args.num_prompts,
max_concurrency=concurrency,
num_warmups=args.num_warmups,
voice=args.voice,
language=args.language,
)
result.config_name = args.config_name
all_results.append(asdict(result))
# Save results
result_dir = Path(args.result_dir)
result_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
with open(result_file, "w") as f:
json.dump(all_results, f, indent=2)
print(f"Results saved to {result_file}")
return all_results
def parse_args():
parser = argparse.ArgumentParser(description="Qwen3-TTS Benchmark Client")
parser.add_argument("--host", type=str, default="127.0.0.1")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--num-prompts", type=int, default=50, help="Number of prompts per concurrency level")
parser.add_argument( # noqa: E501
"--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("--voice", type=str, default="vivian")
parser.add_argument("--language", type=str, default="English")
parser.add_argument(
"--config-name", type=str, default="async_chunk", help="Label for this config (used in filenames)"
)
parser.add_argument("--result-dir", type=str, default="results")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
asyncio.run(main(args))
@@ -1,93 +0,0 @@
# Qwen3-TTS batch_size=1 config (streaming with async_chunk)
# 2-stage pipeline: Talker -> Code2Wav
async_chunk: true
stage_args:
- stage_id: 0
stage_type: llm
is_comprehension: true
runtime:
devices: "0"
engine_args:
max_num_seqs: 1
model_stage: qwen3_tts
model_arch: Qwen3TTSTalkerForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
enforce_eager: false
trust_remote_code: true
async_scheduling: true
enable_prefix_caching: false
engine_output_type: latent
gpu_memory_utilization: 0.3
distributed_executor_backend: "mp"
max_num_batched_tokens: 512
max_model_len: 4096
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
output_connectors:
to_stage_1: connector_of_shared_memory
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 4096
seed: 42
detokenize: false
repetition_penalty: 1.05
stop_token_ids: [2150]
- stage_id: 1
stage_type: llm
runtime:
devices: "0"
engine_args:
max_num_seqs: 1
model_stage: code2wav
model_arch: Qwen3TTSCode2Wav
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
async_scheduling: true
enable_prefix_caching: false
engine_output_type: audio
gpu_memory_utilization: 0.3
distributed_executor_backend: "mp"
max_num_batched_tokens: 8192
max_model_len: 32768
engine_input_source: [0]
final_output: true
final_output_type: audio
input_connectors:
from_stage_0: connector_of_shared_memory
tts_args:
max_instructions_length: 500
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 65536
seed: 42
detokenize: true
repetition_penalty: 1.0
runtime:
enabled: true
defaults:
window_size: -1
max_inflight: 1
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
connector_get_max_wait: 300
codec_chunk_frames: 25
codec_left_context_frames: 25
edges:
- from: 0
to: 1
window_size: -1
@@ -1,94 +0,0 @@
# Qwen3-TTS max_num_seqs=16 config (streaming with async_chunk)
# High-throughput concurrent request processing
# 2-stage pipeline: Talker -> Code2Wav
async_chunk: true
stage_args:
- stage_id: 0
stage_type: llm
is_comprehension: true
runtime:
devices: "0"
engine_args:
max_num_seqs: 16
model_stage: qwen3_tts
model_arch: Qwen3TTSTalkerForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
enforce_eager: false
trust_remote_code: true
async_scheduling: true
enable_prefix_caching: false
engine_output_type: latent
gpu_memory_utilization: 0.3
distributed_executor_backend: "mp"
max_num_batched_tokens: 4096
max_model_len: 4096
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
output_connectors:
to_stage_1: connector_of_shared_memory
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 4096
seed: 42
detokenize: false
repetition_penalty: 1.05
stop_token_ids: [2150]
- stage_id: 1
stage_type: llm
runtime:
devices: "0"
engine_args:
max_num_seqs: 16
model_stage: code2wav
model_arch: Qwen3TTSCode2Wav
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
async_scheduling: true
enable_prefix_caching: false
engine_output_type: audio
gpu_memory_utilization: 0.2
distributed_executor_backend: "mp"
max_num_batched_tokens: 16384
max_model_len: 32768
engine_input_source: [0]
final_output: true
final_output_type: audio
input_connectors:
from_stage_0: connector_of_shared_memory
tts_args:
max_instructions_length: 500
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 65536
seed: 42
detokenize: true
repetition_penalty: 1.0
runtime:
enabled: true
defaults:
window_size: -1
max_inflight: 16
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
connector_get_max_wait: 300
codec_chunk_frames: 25
codec_left_context_frames: 25
edges:
- from: 0
to: 1
window_size: -1
@@ -1,94 +0,0 @@
# Qwen3-TTS batch_size=4 config (streaming with async_chunk)
# Enables concurrent request processing
# 2-stage pipeline: Talker -> Code2Wav
async_chunk: true
stage_args:
- stage_id: 0
stage_type: llm
is_comprehension: true
runtime:
devices: "0"
engine_args:
max_num_seqs: 4
model_stage: qwen3_tts
model_arch: Qwen3TTSTalkerForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
enforce_eager: false
trust_remote_code: true
async_scheduling: true
enable_prefix_caching: false
engine_output_type: latent
gpu_memory_utilization: 0.3
distributed_executor_backend: "mp"
max_num_batched_tokens: 512
max_model_len: 4096
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
output_connectors:
to_stage_1: connector_of_shared_memory
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 4096
seed: 42
detokenize: false
repetition_penalty: 1.05
stop_token_ids: [2150]
- stage_id: 1
stage_type: llm
runtime:
devices: "0"
engine_args:
max_num_seqs: 4
model_stage: code2wav
model_arch: Qwen3TTSCode2Wav
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
async_scheduling: true
enable_prefix_caching: false
engine_output_type: audio
gpu_memory_utilization: 0.2
distributed_executor_backend: "mp"
max_num_batched_tokens: 8192
max_model_len: 32768
engine_input_source: [0]
final_output: true
final_output_type: audio
input_connectors:
from_stage_0: connector_of_shared_memory
tts_args:
max_instructions_length: 500
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 65536
seed: 42
detokenize: true
repetition_penalty: 1.0
runtime:
enabled: true
defaults:
window_size: -1
max_inflight: 4
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
codec_streaming: true
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 3000
connector_get_max_wait: 300
codec_chunk_frames: 25
codec_left_context_frames: 25
edges:
- from: 0
to: 1
window_size: -1
@@ -1,249 +0,0 @@
"""Plot TTFP comparison: async_chunk off vs on.
Generates a bar chart with improvement arrows, matching the Qwen3-Omni
async_chunk benchmark figure style.
Usage:
python plot_async_chunk.py \
--off results/bench_async_chunk_off_*.json \
--on results/bench_async_chunk_on_*.json \
--output results/qwen3_tts_async_chunk_ttfp.png
# Also supports E2E and RTF metrics:
python plot_async_chunk.py \
--off results/bench_async_chunk_off_*.json \
--on results/bench_async_chunk_on_*.json \
--metric e2e \
--output results/qwen3_tts_async_chunk_e2e.png
"""
import argparse
import json
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
METRIC_CONFIG = {
"ttfp": {
"key": "mean_ttfp_ms",
"ylabel": "TTFP (s)",
"title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency",
"to_seconds": True,
},
"e2e": {
"key": "mean_e2e_ms",
"ylabel": "E2E (s)",
"title": "E2E Latency - Qwen3-TTS, by concurrency",
"to_seconds": True,
},
"rtf": {
"key": "mean_rtf",
"ylabel": "RTF",
"title": "Real-Time Factor - Qwen3-TTS, by concurrency",
"to_seconds": False,
},
}
def load_results(path: str) -> list[dict]:
with open(path) as f:
return json.load(f)
def plot_ttfp_comparison(
off_results: list[dict],
on_results: list[dict],
metric: str,
output_path: str,
title_override: str | None = None,
):
cfg = METRIC_CONFIG[metric]
key = cfg["key"]
to_seconds = cfg["to_seconds"]
off_map = {r["concurrency"]: r for r in off_results}
on_map = {r["concurrency"]: r for r in on_results}
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
off_vals = []
on_vals = []
for c in concurrencies:
v_off = off_map[c][key]
v_on = on_map[c][key]
if to_seconds:
v_off /= 1000.0
v_on /= 1000.0
off_vals.append(v_off)
on_vals.append(v_on)
fig, ax = plt.subplots(figsize=(8, 6))
x = np.arange(len(concurrencies))
width = 0.3
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none")
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
# Draw improvement arrows and labels
for i in range(len(concurrencies)):
v_off = off_vals[i]
v_on = on_vals[i]
if v_on > 0:
improvement = v_off / v_on
else:
improvement = float("inf")
# Arrow from top of off-bar to top of on-bar
arrow_start_x = x[i] - width / 2
arrow_start_y = v_off * 0.95
arrow_end_x = x[i] + width / 2
arrow_end_y = v_on * 1.05
ax.annotate(
"",
xy=(arrow_end_x, arrow_end_y),
xytext=(arrow_start_x, arrow_start_y),
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
)
# Improvement label
label_x = (arrow_start_x + arrow_end_x) / 2
label_y = arrow_start_y + (v_off - v_on) * 0.15
ax.text(
label_x,
label_y,
f"{improvement:.1f}x improvement",
ha="center",
va="bottom",
fontsize=10,
color="red",
fontweight="bold",
)
title = title_override or cfg["title"]
ax.set_title(title, fontsize=13, fontweight="bold")
ax.set_ylabel(cfg["ylabel"], fontsize=12)
ax.set_xlabel("Max concurrency", fontsize=12)
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
ax.set_yscale("log")
ax.legend(loc="upper left", fontsize=11)
ax.grid(axis="y", alpha=0.3, linestyle="--")
ax.set_axisbelow(True)
plt.tight_layout()
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str):
"""Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons."""
off_map = {r["concurrency"]: r for r in off_results}
on_map = {r["concurrency"]: r for r in on_results}
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold")
for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]):
cfg = METRIC_CONFIG[metric]
key = cfg["key"]
to_seconds = cfg["to_seconds"]
off_vals = []
on_vals = []
for c in concurrencies:
v_off = off_map[c][key]
v_on = on_map[c][key]
if to_seconds:
v_off /= 1000.0
v_on /= 1000.0
off_vals.append(v_off)
on_vals.append(v_on)
x = np.arange(len(concurrencies))
width = 0.3
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB")
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
for i in range(len(concurrencies)):
if on_vals[i] > 0:
improvement = off_vals[i] / on_vals[i]
ax.annotate(
"",
xy=(x[i] + width / 2, on_vals[i] * 1.05),
xytext=(x[i] - width / 2, off_vals[i] * 0.95),
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
)
label_y = off_vals[i] * 0.85
ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold")
ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold")
ax.set_ylabel(cfg["ylabel"], fontsize=11)
ax.set_xlabel("Max concurrency", fontsize=11)
ax.set_xticks(x)
ax.set_xticklabels([str(c) for c in concurrencies])
if metric != "rtf":
ax.set_yscale("log")
ax.legend(fontsize=9)
ax.grid(axis="y", alpha=0.3, linestyle="--")
ax.set_axisbelow(True)
plt.tight_layout()
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
plt.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"Plot saved to {output_path}")
plt.close()
def print_table(off_results: list[dict], on_results: list[dict]):
off_map = {r["concurrency"]: r for r in off_results}
on_map = {r["concurrency"]: r for r in on_results}
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
print("\n## Benchmark Results: async_chunk off vs on\n")
print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |")
print("| --- | --- | --- | --- | --- |")
for name, key, fmt in [
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
("E2E (ms)", "mean_e2e_ms", ".1f"),
("RTF", "mean_rtf", ".3f"),
("Throughput", "audio_throughput", ".2f"),
]:
for c in concurrencies:
v_off = off_map[c].get(key, 0)
v_on = on_map[c].get(key, 0)
if v_on > 0 and key != "audio_throughput":
ratio = f"{v_off / v_on:.1f}x"
elif v_off > 0 and key == "audio_throughput":
ratio = f"{v_on / v_off:.1f}x"
else:
ratio = "N/A"
print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |")
def parse_args():
parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS")
parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off")
parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on")
parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"])
parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png")
parser.add_argument("--title", type=str, default=None, help="Custom title override")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
off_results = load_results(args.off)
on_results = load_results(args.on)
print_table(off_results, on_results)
if args.metric == "all":
plot_all_metrics(off_results, on_results, args.output)
else:
plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title)
@@ -1,167 +0,0 @@
#!/bin/bash
# Qwen3-TTS async_chunk on vs off Benchmark
#
# Starts two servers (async_chunk on and off), benchmarks both,
# and generates comparison plots.
#
# Usage:
# bash run_async_chunk_benchmark.sh
#
# Environment variables:
# GPU_DEVICE - GPU index (default: 0)
# MODEL - Model path (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
# NUM_PROMPTS - Prompts per concurrency level (default: 50)
# CONCURRENCY - Space-separated concurrency levels (default: "1 10")
# PORT_ON - Port for async_chunk on server (default: 8000)
# PORT_OFF - Port for async_chunk off server (default: 8001)
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
cd "$PROJECT_ROOT"
GPU_DEVICE="${GPU_DEVICE:-0}"
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
NUM_PROMPTS="${NUM_PROMPTS:-50}"
CONCURRENCY="${CONCURRENCY:-1 10}"
NUM_WARMUPS="${NUM_WARMUPS:-3}"
PORT_ON="${PORT_ON:-8000}"
PORT_OFF="${PORT_OFF:-8001}"
RESULT_DIR="${SCRIPT_DIR}/results"
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
STAGE_CONFIG_ON="vllm_omni/model_executor/stage_configs/qwen3_tts.yaml"
STAGE_CONFIG_OFF="vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml"
mkdir -p "${RESULT_DIR}"
echo "============================================================"
echo " Qwen3-TTS async_chunk Benchmark"
echo "============================================================"
echo " GPU: ${GPU_DEVICE}"
echo " Model: ${MODEL}"
echo " Prompts: ${NUM_PROMPTS}"
echo " Concurrency: ${CONCURRENCY}"
echo " Port (on/off): ${PORT_ON} / ${PORT_OFF}"
echo " Results: ${RESULT_DIR}"
echo "============================================================"
cleanup() {
echo "Cleaning up servers..."
kill "$PID_ON" 2>/dev/null || true
kill "$PID_OFF" 2>/dev/null || true
wait "$PID_ON" 2>/dev/null || true
wait "$PID_OFF" 2>/dev/null || true
}
trap cleanup EXIT
wait_for_server() {
local port=$1
local name=$2
local max_wait=300
local elapsed=0
echo "Waiting for ${name} server on port ${port}..."
while ! curl -s "http://localhost:${port}/health" >/dev/null 2>&1; do
sleep 5
elapsed=$((elapsed + 5))
if [ $elapsed -ge $max_wait ]; then
echo "ERROR: ${name} server failed to start within ${max_wait}s"
exit 1
fi
done
echo "${name} server ready (${elapsed}s)"
}
# ---- Phase 1: Start async_chunk ON server ----
echo ""
echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..."
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
--stage-configs-path "${STAGE_CONFIG_ON}" \
--host 0.0.0.0 --port "${PORT_ON}" \
--trust-remote-code --enforce-eager --omni \
> "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 &
PID_ON=$!
wait_for_server "${PORT_ON}" "async_chunk_on"
echo "[Phase 1] Benchmarking async_chunk ON..."
# shellcheck disable=SC2086
python "${SCRIPT_DIR}/bench_async_chunk.py" \
--host 127.0.0.1 --port "${PORT_ON}" \
--config-name "async_chunk_on" \
--num-prompts "${NUM_PROMPTS}" \
--max-concurrency ${CONCURRENCY} \
--num-warmups "${NUM_WARMUPS}" \
--result-dir "${RESULT_DIR}"
echo "[Phase 1] Stopping async_chunk ON server..."
kill "$PID_ON" 2>/dev/null || true
wait "$PID_ON" 2>/dev/null || true
sleep 5
# ---- Phase 2: Start async_chunk OFF server ----
echo ""
echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..."
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
--stage-configs-path "${STAGE_CONFIG_OFF}" \
--host 0.0.0.0 --port "${PORT_OFF}" \
--trust-remote-code --enforce-eager --omni \
> "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 &
PID_OFF=$!
wait_for_server "${PORT_OFF}" "async_chunk_off"
echo "[Phase 2] Benchmarking async_chunk OFF (non-streaming)..."
# shellcheck disable=SC2086
python "${SCRIPT_DIR}/bench_async_chunk.py" \
--host 127.0.0.1 --port "${PORT_OFF}" \
--config-name "async_chunk_off" \
--num-prompts "${NUM_PROMPTS}" \
--max-concurrency ${CONCURRENCY} \
--num-warmups "${NUM_WARMUPS}" \
--no-stream \
--result-dir "${RESULT_DIR}"
echo "[Phase 2] Stopping async_chunk OFF server..."
kill "$PID_OFF" 2>/dev/null || true
wait "$PID_OFF" 2>/dev/null || true
# ---- Phase 3: Plot results ----
echo ""
echo "[Phase 3] Generating plots..."
# Find the latest result files
RESULT_ON=$(ls -t "${RESULT_DIR}"/bench_async_chunk_on_*.json 2>/dev/null | head -1)
RESULT_OFF=$(ls -t "${RESULT_DIR}"/bench_async_chunk_off_*.json 2>/dev/null | head -1)
if [ -z "$RESULT_ON" ] || [ -z "$RESULT_OFF" ]; then
echo "ERROR: Could not find result files. Check logs in ${RESULT_DIR}/"
exit 1
fi
echo " ON results: ${RESULT_ON}"
echo " OFF results: ${RESULT_OFF}"
# TTFP comparison (main figure)
python "${SCRIPT_DIR}/plot_async_chunk.py" \
--off "${RESULT_OFF}" \
--on "${RESULT_ON}" \
--metric ttfp \
--output "${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
# All metrics comparison
python "${SCRIPT_DIR}/plot_async_chunk.py" \
--off "${RESULT_OFF}" \
--on "${RESULT_ON}" \
--metric all \
--output "${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
echo ""
echo "============================================================"
echo " Benchmark complete!"
echo " Results: ${RESULT_DIR}/"
echo " Plots:"
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
echo "============================================================"
+3 -2
View File
@@ -1,13 +1,14 @@
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
ARG VLLM_BASE_TAG=v0.18.0
ARG VLLM_BASE_TAG=v0.16.0
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
ARG APP_DIR=/workspace/vllm-omni
WORKDIR ${APP_DIR}
COPY . .
# Install system dependencies
RUN apt-get update && \
apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
+5 -18
View File
@@ -1,30 +1,17 @@
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.17.0rc1
ARG VLLM_ASCEND_TAG=v0.14.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-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
RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
RUN pip install -v -e .
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
COPY . .
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
source /usr/local/Ascend/nnal/atb/set_env.sh && \
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
RUN pip install -v -e .
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+5 -18
View File
@@ -1,30 +1,17 @@
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
ARG VLLM_ASCEND_TAG=v0.14.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-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
RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
RUN pip install -v -e .
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
COPY . .
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
source /usr/local/Ascend/nnal/atb/set_env.sh && \
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
RUN pip install -v -e .
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+3 -3
View File
@@ -1,4 +1,4 @@
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.18.0
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.16.0
FROM ${BASE_IMAGE} AS base
# Declare a variable to know if we want to use the nightly build or the stable build.
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
# we should swap over to use stable release ASAP.
# We should avoid relying on custom commits.
ARG USE_NIGHTLY_BUILD=0
ARG VLLM_VERSION_OR_COMMIT_HASH=89138b21cc246ae944c741d5c399c148e2b770ab
ARG VLLM_VERSION_OR_COMMIT_HASH=2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
ARG ARG_PYTORCH_ROCM_ARCH
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
@@ -19,7 +19,7 @@ WORKDIR ${COMMON_WORKDIR}
# Step 1: Setup - Install system dependencies
RUN apt-get update && \
apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
+4 -17
View File
@@ -9,12 +9,12 @@ ARG PYTHON_VERSION=3.12
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
add-apt-repository -y ppa:kobuk-team/intel-graphics
RUN apt clean && apt-get update -y && \
apt-get install -y --no-install-recommends --fix-missing \
curl \
espeak-ng \
ffmpeg \
git \
libsndfile1 \
@@ -31,22 +31,9 @@ RUN apt clean && apt-get update -y && \
python3-pip
RUN apt update && apt upgrade -y && \
apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing intel-ocloc && \
apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
# Install UMD with fixed version
RUN mkdir neo && \
cd neo && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
dpkg -i *.deb && \
cd .. && \
rm -rf neo
ENV PATH="/root/.local/bin:$PATH"
ENV VIRTUAL_ENV="/opt/venv"
ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
@@ -76,7 +63,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.16.0
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
WORKDIR /workspace/vllm
+5 -33
View File
@@ -6,7 +6,6 @@ nav:
- getting_started/installation/*
- Serving:
- OpenAI-Compatible API:
- Diffusion Chat API: serving/diffusion_chat_api.md
- Image Generation: serving/image_generation_api.md
- Image Edit: serving/image_edit_api.md
- Text to Speech: serving/speech_api.md
@@ -14,14 +13,9 @@ nav:
- examples/README.md
- Offline Inference:
- BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
- CosyVoice3: user_guide/examples/offline_inference/cosyvoice3.md
- Fish Speech S2 Pro: user_guide/examples/offline_inference/fish_speech.md
- GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
- Helios Video Generation: user_guide/examples/offline_inference/helios.md
- HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
- Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
- MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
- MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
- Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
- Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
@@ -29,12 +23,8 @@ nav:
- Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
- Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
- Voxtral TTS Offline Inference: user_guide/examples/offline_inference/voxtral_tts.md
- X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
- Online Serving:
- BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
- vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
- Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md
- GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md
@@ -54,28 +44,18 @@ nav:
- Features:
- Sleep Mode: features/sleep_mode.md
- Diffusion Features:
- 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
- Acceleration Overview: user_guide/diffusion_acceleration.md
- TeaCache: user_guide/diffusion/teacache.md
- Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
- Quantization:
- Overview: user_guide/diffusion/quantization/overview.md
- FP8: user_guide/diffusion/quantization/fp8.md
- Int8: user_guide/diffusion/quantization/int8.md
- GGUF: user_guide/diffusion/quantization/gguf.md
- Parallelism:
- 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
- Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.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:
@@ -86,9 +66,7 @@ nav:
- Model Implementation:
- contributing/model/README.md
- contributing/model/adding_omni_model.md
- contributing/model/adding_tts_model.md
- contributing/model/adding_diffusion_model.md
- contributing/model/adding_quantization_model.md
- CI: contributing/ci
- Design Documents:
- design/index.md
@@ -100,25 +78,19 @@ nav:
- design/feature/cfg_parallel.md
- design/feature/sequence_parallel.md
- design/feature/tensor_parallel.md
- design/feature/vae_parallel.md
- design/feature/hsdp.md
- design/feature/cache_dit.md
- design/feature/teacache.md
- design/feature/async_chunk_design.md
- design/feature/vae_parallel.md
- design/feature/diffusion_step_execution.md
- Module Design:
- design/module/ar_module.md
- design/module/dit_module.md
- design/module/entrypoint_module.md
- design/module/async_omni_architecture.md
- Docs Guide: contributing/DOCS_GUIDE.md
- API Reference:
- api/README.md
- api/vllm_omni
- CLI Reference: cli
- Community:
- Governance: community/governance.md
- community/*
- Slack: https://slack.vllm.ai
- Blog: https://blog.vllm.ai
+1 -1
View File
@@ -60,5 +60,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
For more information, checkout the following:
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1)
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
+9 -42
View File
@@ -6,21 +6,24 @@ 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.async_omni_llm.AsyncOmniLLM][]
- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
- [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
- [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
- [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
- [vllm_omni.entrypoints.omni.Omni][]
- [vllm_omni.entrypoints.omni_base.OmniBase][]
- [vllm_omni.entrypoints.pd_utils.PDDisaggregationMixin][]
- [vllm_omni.entrypoints.omni.OmniBase][]
- [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]
- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
- [vllm_omni.entrypoints.omni_stage.OmniStage][]
- [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
- [vllm_omni.entrypoints.zmq_utils.ZmqQueue][]
## Inputs
Input data structures for multi-modal inputs.
- [vllm_omni.inputs.data.OmniCustomPrompt][]
- [vllm_omni.inputs.data.OmniDiffusionSamplingParams][]
- [vllm_omni.inputs.data.OmniEmbedsPrompt][]
- [vllm_omni.inputs.data.OmniTextPrompt][]
@@ -45,19 +48,11 @@ Engine classes for offline and online inference.
- [vllm_omni.engine.OmniEngineCoreOutputs][]
- [vllm_omni.engine.OmniEngineCoreRequest][]
- [vllm_omni.engine.PromptEmbedsPayload][]
- [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
- [vllm_omni.engine.arg_utils.OmniEngineArgs][]
- [vllm_omni.engine.async_omni_engine.AsyncOmniEngine][]
- [vllm_omni.engine.mm_outputs.MultimodalCompletionOutput][]
- [vllm_omni.engine.mm_outputs.MultimodalPayload][]
- [vllm_omni.engine.orchestrator.Orchestrator][]
- [vllm_omni.engine.orchestrator.OrchestratorRequestState][]
- [vllm_omni.engine.output_modality.OutputModality][]
- [vllm_omni.engine.output_modality.TensorAccumulationStrategy][]
- [vllm_omni.engine.input_processor.OmniInputProcessor][]
- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
- [vllm_omni.engine.output_processor.OmniRequestState][]
- [vllm_omni.engine.stage_engine_core_client.StageEngineCoreClient][]
- [vllm_omni.engine.stage_init_utils.StageMetadata][]
- [vllm_omni.engine.stage_init_utils.StartedLlmStage][]
## Core
@@ -69,19 +64,6 @@ Core scheduling and caching components.
- [vllm_omni.core.sched.output.OmniCachedRequestData][]
- [vllm_omni.core.sched.output.OmniNewRequestData][]
- [vllm_omni.core.sched.output.OmniSchedulerOutput][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.BASECFM][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.CausalConditionalCFM][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.CausalMaskedDiffWithDiT][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.ConditionalCFM][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConv1d][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConv1dUpsample][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConvRNNF0Predictor][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalHiFTGenerator][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.HiFTGenerator][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SineGen][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SineGen2][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.Snake][]
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SourceModuleHnNSF][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.EuclideanCodebook][]
@@ -93,17 +75,9 @@ Core scheduling and caching components.
Configuration classes.
- [vllm_omni.config.model.OmniModelConfig][]
- [vllm_omni.config.stage_config.ModelPipeline][]
- [vllm_omni.config.stage_config.StageConfig][]
- [vllm_omni.config.stage_config.StageConfigFactory][]
- [vllm_omni.config.stage_config.StageType][]
- [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
- [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
- [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
- [vllm_omni.model_executor.models.cosyvoice3.config.CosyVoice3Config][]
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechConfig][]
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechFastARConfig][]
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechSlowARConfig][]
- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioConfig][]
- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioTokenizerConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSConfig][]
@@ -117,10 +91,6 @@ Configuration classes.
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
## Workers
@@ -131,8 +101,6 @@ Worker classes and model runners for distributed inference.
- [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
- [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
- [vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase][]
- [vllm_omni.diffusion.worker.utils.DiffusionRequestState][]
- [vllm_omni.diffusion.worker.utils.RunnerOutput][]
- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.ExecuteModelState][]
- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.NPUARModelRunner][]
- [vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker][]
@@ -150,6 +118,5 @@ Worker classes and model runners for distributed inference.
- [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
- [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
- [vllm_omni.worker.gpu_memory_utils.parse_cuda_visible_devices][]
- [vllm_omni.worker.gpu_model_runner.CUDAGraphWrapper][]
- [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
- [vllm_omni.worker.mixins.OmniWorkerMixin][]
Binary file not shown.

Before

Width:  |  Height:  |  Size: 98 KiB

After

Width:  |  Height:  |  Size: 127 KiB

-54
View File
@@ -1,54 +0,0 @@
# Governance
vLLM-Omni's governance is inspired by the [vLLM governance process](https://docs.vllm.ai/en/latest/governance/process/). We share the same commitment to open source community and meritocratic norms.
## Values
vLLM-Omni aims to be the fastest and easiest-to-use omni-modality inference and serving engine. Our values are aligned with [vLLM's values](https://docs.vllm.ai/en/latest/governance/process/#values):
### Design Values
1. **Top performance**: System performance is our top priority. We monitor overheads, optimize kernels, and publish benchmarks. We never leave performance on the table.
2. **Ease of use**: vLLM-Omni must be simple to install, configure, and operate. We provide clear documentation, fast startup, clean logs, helpful error messages, and monitoring guides. Many users fork our code or study it deeply, so we keep it readable and modular.
3. **Wide coverage**: vLLM-Omni supports frontier models and high-performance accelerators. We make it easy to add new models and hardware. vLLM-Omni + PyTorch form a simple interface that avoids complexity.
4. **Production ready**: vLLM-Omni runs 24/7 in production. It must be easy to operate and monitor for health issues.
5. **Extensibility**: vLLM-Omni serves as fundamental omni-modality infrastructure. Our codebase cannot cover every use case, so we design for easy forking and customization.
### Collaboration Values
1. **Tightly Knit and Fast-Moving**: Our maintainer team is aligned on vision, philosophy, and roadmap. We work closely to unblock each other and move quickly.
2. **Individual Merit**: No one buys their way into governance. Committer status belongs to individuals, not companies. We reward contribution, maintenance, and project stewardship.
## Project Maintainers
### Lead Maintainers
Lead maintainers are responsible for the overall direction and strategy of the project:
- [@Gaohan123](https://github.com/Gaohan123)
- [@hsliuustc0106](https://github.com/hsliuustc0106)
- [@ywang96](https://github.com/ywang96)
### Active Committers
Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
- [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism
- [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
## Meetings
Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
## Committer Nomination Process
Every month, any active committer can nominate new committer(s) to the project. Up to **two new committers** will be admitted per month based on the quality and impact of their contributions.
+4 -6
View File
@@ -2,7 +2,7 @@
This section lists the most common options for running vLLM-Omni.
For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.16.0/configuration/index.html)
For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html)
Currently, the main options are maintained by stage configs for each model.
@@ -14,10 +14,8 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
- **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
## Multi-Stage Recipes
- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
## Optimization Features
- **[Diffusion Features Overview](../user_guide/diffusion_features.md)** - Complete overview of all diffusion model features and supported models
- **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
- **[Cache-DiT Configuration](../user_guide/diffusion/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
- **[Parallelism Configuration](../user_guide/diffusion/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models
+1 -1
View File
@@ -158,7 +158,7 @@ The stage initialized but ran out of memory during processing.
**Solutions:**
1. Reduce `max_num_batched_tokens`
2. Reduce `max_num_seqs` in engine_args
2. Reduce `max_batch_size` in runtime config
3. Lower `gpu_memory_utilization` slightly
4. Enable quantization if supported
-177
View File
@@ -1,177 +0,0 @@
# Prefill-Decode (PD) Disaggregation
PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
stages so prompt processing and token generation can run on different workers.
This is documented as a stage-config recipe instead of a bundled YAML because the
deployment-specific values usually change per environment:
- GPU placement
- `tensor_parallel_size`
- connector backend and connector ports
- connector IPs or bootstrap addresses
Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml)
and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
changes below.
## Requirements
- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
- A KV connector supported by vLLM, such as `MooncakeConnector`
- Matching `tensor_parallel_size` on the prefill and decode thinker stages
## 1. Split the thinker into prefill and decode stages
Replace the original thinker stage with two stages:
```yaml
stage_args:
- stage_id: 0
stage_type: llm
is_prefill_only: true
runtime:
devices: "0"
engine_args:
max_num_seqs: 16
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.9
enforce_eager: true
trust_remote_code: true
engine_output_type: latent
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
hf_config_name: thinker_config
tensor_parallel_size: 1
kv_transfer_config:
kv_connector: "MooncakeConnector"
kv_role: "kv_producer"
kv_rank: 0
kv_parallel_size: 2
kv_connector_extra_config:
mooncake_bootstrap_port: 25201
final_output: false
is_comprehension: true
default_sampling_params:
temperature: 0.4
top_p: 0.9
top_k: 1
max_tokens: 2048
seed: 42
detokenize: True
repetition_penalty: 1.05
- stage_id: 1
stage_type: llm
is_decode_only: true
runtime:
devices: "1"
engine_args:
max_num_seqs: 64
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.9
enforce_eager: true
trust_remote_code: true
engine_output_type: latent
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
hf_config_name: thinker_config
tensor_parallel_size: 1
kv_transfer_config:
kv_connector: "MooncakeConnector"
kv_role: "kv_consumer"
kv_rank: 1
kv_parallel_size: 2
kv_connector_extra_config:
mooncake_bootstrap_port: 25202
engine_input_source: [0]
final_output: true
final_output_type: text
is_comprehension: true
default_sampling_params:
temperature: 0.4
top_p: 0.9
top_k: 1
max_tokens: 2048
seed: 42
detokenize: True
repetition_penalty: 1.05
```
Notes:
- `is_prefill_only: true` marks the thinker stage that only saves KV.
- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
- `kv_transfer_config` is required on both stages.
- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
prefill side only processes the prompt and exports KV.
## 2. Shift the downstream stages by one index
After inserting the extra thinker stage, renumber the remaining stages:
```yaml
- stage_id: 2
runtime:
devices: "2"
engine_input_source: [1]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
- stage_id: 3
runtime:
devices: "2"
engine_args:
max_num_seqs: 1
engine_input_source: [2]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
```
Compared with the default Qwen3-Omni config:
- the talker becomes stage `2` instead of stage `1`
- the code2wav stage becomes stage `3` instead of stage `2`
- the talker now reads from decode stage `1`
## 3. Add runtime edges for the four-stage pipeline
```yaml
runtime:
enabled: true
defaults:
window_size: -1
max_inflight: 1
edges:
- from: 0
to: 1
window_size: -1
- from: 1
to: 2
window_size: -1
- from: 2
to: 3
window_size: -1
```
## 4. Launch with your custom config
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--stage-configs-path /path/to/qwen3_omni_pd.yaml
```
## Operational Notes
- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
pair. Keep prefill and decode at the same `tensor_parallel_size`.
- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
- Choose connector ports and addresses that match your deployment. The values
shown above are examples only.
+7 -15
View File
@@ -18,7 +18,7 @@ If users want to modify some part of it. The custom stage_configs file can be in
For offline (Assume necessary dependencies have ben imported):
```python
model_name = "Qwen/Qwen2.5-Omni-7B"
omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
omni_llm = OmniLLM(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
```
For online serving:
@@ -30,15 +30,15 @@ vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to
Below is a specific example of stage_configs.yaml in Qwen2.5-omni.
```python
# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
# stage config for running qwen2.5-omni with architecture of OmniLLM.
stage_args:
- stage_id: 0 # mark the unique id for each stage
runtime: # The disaggregated configuration
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1 # the batch_size for offline inference
engine_args: # Engine arguments for a certain engine
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
worker_type: ar # The specific worker used
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used
@@ -62,9 +62,9 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 3
engine_args:
model_stage: talker
max_num_seqs: 3
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -88,9 +88,9 @@ stage_args:
runtime:
process: true
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
@@ -135,14 +135,6 @@ Each stage in the `stage_args` list contains the following configuration options
A unique identifier for each stage in the multi-stage pipeline. Stages are numbered sequentially starting from 0, and this ID is used to reference stages in inter-stage dependencies (e.g., `engine_input_source`).
### `prompt_expand_func` (Optional)
A custom Python function hook for the LLM stage (Stage 0) that expands a single incoming prompt object into multiple prompts. This is primarily used for multi-modal Classifier-Free Guidance (CFG), where it generates the necessary companion requests (like a negative text prompt) and tags them with internal roles (e.g., `cfg_text`). This ensures the upstream LLM generates the needed contextual hidden states for both the conditional and unconditional generations simultaneously.
### `cfg_kv_collect_func` (Optional)
A custom Python function hook for downstream diffusion stages (Stage 1+) to collect, map, and process the KV caches transferred from the companion requests fired by `prompt_expand_func`. It aggregates the hidden condition states cleanly (e.g., binding them as `cfg_text_past_key_values` and `cfg_text_kv_metadata`), allowing the diffusion runtime to perform CFG smoothly without redundantly evaluating text paths on the DiT workers.
### `runtime`
Configuration for disaggregated execution of the stage, controlling how the stage is deployed and executed.
@@ -159,9 +151,9 @@ Visible devices for this stage, specified as a string. This controls which GPU d
Default: `"0"`
#### `engine_args.max_num_seqs`
#### `runtime.max_batch_size`
The maximum number of sequences for concurrent processing in this stage. For LLM stages, this controls the vLLM scheduler's maximum concurrent sequences. For all stage types, this also controls how many tasks can be batched together in the task processing loop.
The maximum batch size for offline inference in this stage. This limits how many sequences can be processed together in a single batch during offline inference operations.
Default: `1`
@@ -1,12 +1,12 @@
# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
# stage config for running qwen2.5-omni with architecture of OmniLLM.
stage_args:
- stage_id: 0
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -30,9 +30,9 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -56,9 +56,9 @@ stage_args:
runtime:
process: true
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+9 -5
View File
@@ -71,7 +71,14 @@ For additional features and advanced configurations, refer to the:
### Testing
vLLM-Omni uses `pytest` to test the codebase.
Please refer to the [test instructions](./ci/test_guide.md) for detailed testing information.
```bash
# Run all tests
pytest tests/
# Run tests for a single test file with detailed output
pytest -s -v tests/test_omni_llm.py
```
!!! warning
Currently, not all unit tests pass when run on CPU platforms. If you don't have access to a GPU platform to run unit tests locally, rely on the continuous integration system to run the tests for now.
@@ -107,7 +114,7 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
- `[CI/Build]` for build or continuous integration improvements.
- `[Doc]` for documentation fixes and improvements.
- `[Model]` for adding a new model or improving an existing model. Model name should appear in the title.
- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `Omni`/`AsyncOmni`, etc.)
- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `OmniLLM` class, etc.)
- `[Kernel]` for changes affecting CUDA kernels or other compute kernels.
- `[Core]` for changes in the core vLLM-Omni logic (e.g., `OmniProcessor`, `OmniARScheduler`, etc.)
- `[Hardware][Vendor]` for hardware-specific changes. Vendor name should appear in the prefix, such as [Ascend] for Ascend NPUs.
@@ -116,9 +123,6 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
!!! note
If the PR spans more than one category, please include all relevant prefixes.
### Local Test
Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
### Code Quality
The PR needs to meet the following code quality standards:
-6
View File
@@ -1,6 +0,0 @@
nav:
- CI_5levels.md
- failures.md
- test_guide.md
- test_markers.md
- test_style.md
+91 -80
View File
@@ -67,8 +67,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
<td>Online & Offline (multiple deployment scenarios):<br>real model, normal inference function, normal accuracy</td>
<td>&lt;30min</td>
<td>
/tests/e2e/online_serving/test_{model_name}.py<br>
/tests/e2e/offline_inference/test_{model_name}.py
/tests/e2e/online_serving/test_{model_name}_expansion.py<br>
/tests/e2e/offline_inference/test_{model_name}_expansion.py
</td>
<td>
<a href="#chapter-2-l3-level-testing-core-integration-performance-and-accuracy-verification">Chapter 2</a><br>
@@ -86,7 +86,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
/tests/e2e/online_serving/test_{model_name}_expansion.py<br>
/tests/e2e/offline_inference/test_{model_name}_expansion.py<br>
<strong>Performance:</strong><br>
/tests/dfx/perf/tests/test.json<br>
/tests/e2e/perf/nightly.json<br>
<strong>Doc Test:</strong><br>
tests/example/online_serving/test_{model_name}.py<br>
tests/example/offline_inference/test_{model_name}.py
@@ -104,7 +104,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
<td> Depends on reality </td>
<td>
<strong>Stability:</strong><br>
/tests/dfx/stability/tests/test.json<br>
tests/e2e/stability/weekly.json<br>
<strong>Reliability:</strong><br>
tests/e2e/reliability/test_{model_name}.py
</td>
@@ -168,6 +168,11 @@ vllm_omni/ tests/
│ └── arg_utils.py │ └── test_arg_utils.py ⬜
│
├── entrypoints/ → ├── entrypoints/
│ ├── omni.py │ ├── test_omni.py ⬜ (E2E covered by e2e/offline, e2e/online)
│ ├── omni_llm.py │ ├── test_omni_llm.py ✅
│ ├── omni_stage.py │ ├── test_omni_stage.py ⬜ (partial in test_omni_stage_diffusion_config.py)
│ ├── omni_diffusion.py │ ├── test_omni_diffusion.py ✅
│ ├── async_omni.py │ ├── test_async_omni.py ✅ actually in e2e/online_serving/test_async_omni.py
│ ├── async_omni_diffusion.py │ ├── test_async_omni_diffusion_config.py ✅
│ ├── stage_utils.py │ ├── test_stage_utils.py ✅
│ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve)
@@ -225,11 +230,9 @@ vllm_omni/ tests/
│
└── e2e (tests side only) → ├── e2e/
├── online_serving/ ✅ non-empty
│ ├── test_qwen2_5_omni.py
│ ├── test_async_omni.py
│ ├── test_qwen3_omni.py
│ ├── test_qwen3_omni_expansion.py
│ ├── test_mimo_audio.py
│ ├── test_image_gen_edit.py
│ ├── test_images_generations_lora.py
│ └── stage_configs/
@@ -505,9 +508,7 @@ L3 level testing executes after code is merged into the main branch. Its core pu
**2.4.5 Request Execution**
```python
openai_client.send_omni_request(request_config, request_num=1) # for omni-understanding models
# or
openai_client.send_diffusion_request(request_config, request_num=1) # for diffusion models
openai_client.send_request(request_config, request_num=1)
```
**Explanation**:
@@ -516,7 +517,7 @@ L3 level testing executes after code is merged into the main branch. Its core pu
**Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n.
**Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
**Implicit Validation**: The send_request method internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
- ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
@@ -531,13 +532,13 @@ L4 level testing is a comprehensive quality audit before a version release. It e
### 3.2 Testing Content and Scope
- ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling.
- ***Performance Testing***: Uses the `tests/dfx/perf/tests/test.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
- ***Performance Testing***: Uses the `/tests/e2e/perf/nightly.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
- ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description.
### 3.3 Test Directory and Execution Files
- ***Functional Testing***: Same directories as L3.
- ***Performance Test Configuration***: `tests/dfx/perf/tests/test.json`
- ***Performance Test Configuration***: `/tests/e2e/perf/nightly.json`
- ***Documentation Example Tests***:
- - `tests/example/online_serving/test_{model_name}.py`
- `tests/example/offline_inference/test_{model_name}.py`
@@ -547,52 +548,13 @@ L4 level testing is a comprehensive quality audit before a version release. It e
- ***Trigger Timing***: **`Nightly`**, automatically executed every night.
- ***Execution Environment***: ***GPU*** server clusters to meet the resource demands of performance testing.
- ***Script Example***:
??? example "Test Examples: Documentation Example Tests"
--8<-- "docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md"
??? example "Test Examples: Performance Tests"
--8<-- "docs/contributing/ci/test_examples/l4_performance_tests.inc.md"
??? example "Test Examples: Functionality Tests"
--8<-- "docs/contributing/ci/test_examples/l4_functionality_tests.inc.md"
- ***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
### 4.1 Testing Purpose
L5 level testing focuses on the performance of model services under ***long-running*** and ***abnormal fault*** scenarios. It aims to uncover deep-seated issues that only manifest under sustained pressure or extreme conditions, such as memory leaks, resource contention, gradual performance degradation, and lack of fault tolerance mechanisms. This is the final, yet crucial, line of defense for ensuring service high availability and production environment robustness.
### 4.2 Testing Content and Scope
- ***Long-term Stability (Stability) Testing***: Uses the `tests/dfx/stability/tests/test.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
- ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
### 4.3 Test Directory and Execution Files
- ***Stability Test Configuration***: `tests/dfx/stability/tests/test.json`
- ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
### 4.4 Execution Method and Example
- ***Trigger Timing***: **`Weekly`** (weekly) or **`Days before Release`** (several days before a major release). Due to long execution times, the frequency is lower.
- ***Execution Environment***: ***GPU*** servers, requiring a stable and exclusive testing environment.
- ***Script Example***:
<details>
<summary> Test Examples</summary>
When you want to add L4-level performance test cases, you can refer to the following format for case addition in tests/perf/tests/test.json:
When you want to add L5-level stability test cases, you can refer to the following format for case addition in `tests/dfx/stability/tests/test.json`:
```json
```JSON
{
"test_name": "test_qwen3_omni_stability",
"test_name": "test_qwen3_omni",
"server_params": {
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"stage_config_name": "qwen3_omni.yaml"
@@ -600,15 +562,17 @@ When you want to add L5-level stability test cases, you can refer to the followi
"benchmark_params": [
{
"dataset_name": "random",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"duration_sec": 43200,
"request_rate": 0.5,
"num_prompts_per_batch": 20,
"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"
"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
}
}
]
}
@@ -618,20 +582,20 @@ When you want to add L5-level stability test cases, you can refer to the followi
***Overview***
| Field | Required | Description |
| ---------------- | -------- | --------------------------------------------------------------------------- |
| test_name | Yes | Unique identifier for the stability test case |
| server_params | Yes | Server-side configuration parameters (model, stage configuration, etc.) |
| benchmark_params | Yes | Stability benchmark running parameters (supports multiple configurations) |
| 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 |
| 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)
@@ -643,23 +607,70 @@ Supports incremental modifications based on the basic configuration:
| delete | Delete specified configuration items |
***Example***:
You can refer to Test Examples in Chapter 3.4
```
"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
For stability testing, the key parameters are:
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:
- **duration_sec**: Total duration (in seconds) during which benchmark traffic is sent. The stability benchmark will keep sending batches until this duration is reached.
- **request_rate** / **max_concurrency**: Exactly one of them must be specified. They control how the traffic is generated for each batch:
- `request_rate`: Number of requests per second. The benchmark will send `num_prompts_per_batch` requests at the given rate.
- `max_concurrency`: Maximum number of concurrent requests. When this is used, `request_rate` is set to `inf` internally.
- **num_prompts_per_batch**: Number of prompts sent in each batch. Multiple batches will be executed sequentially within `duration_sec`.
All other optional parameters follow the same rules as the in Chapter 3.4.
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>
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_{model_name}.py`
- - ***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
### 4.1 Testing Purpose
L5 level testing focuses on the performance of model services under ***long-running*** and ***abnormal fault*** scenarios. It aims to uncover deep-seated issues that only manifest under sustained pressure or extreme conditions, such as memory leaks, resource contention, gradual performance degradation, and lack of fault tolerance mechanisms. This is the final, yet crucial, line of defense for ensuring service high availability and production environment robustness.
### 4.2 Testing Content and Scope
- ***Long-term Stability (Stability) Testing***: Uses the `tests/e2e/stability/weekly.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
- ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
### 4.3 Test Directory and Execution Files
- ***Stability Test Configuration***: `tests/e2e/stability/weekly.json`
- ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
### 4.4 Execution Method and Example
- ***Trigger Timing***: **`Weekly`** (weekly) or **`Days before Release`** (several days before a major release). Due to long execution times, the frequency is lower.
- ***Execution Environment***: ***GPU*** servers, requiring a stable and exclusive testing environment.
- ***Script Example***:
<details>
<summary> Test Examples</summary>
```python
# WIP
```
</details>
- - ***Stability***: (Execution would be driven by the configuration in `weekly.json`).
- ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py`
## Summary
+4 -100
View File
@@ -1,106 +1,10 @@
# CI Failures
## Overview of CI Checks
What should I do when a CI job fails on my PR, but I don't think my PR caused the failure?
When you open a PR against vLLM-Omni, several CI checks run automatically:
| 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). |
## Common Case 1:
## Step 1: Identify the Failing Check
🚨 Error: The command was interrupted by a signal: signal: terminated
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 |
Reason: The test is terminated due to exceed the time limits. Sometimes rebuild the test will pass.
@@ -1,49 +0,0 @@
**Preferred Test Strategy**
Use one of the following patterns depending on page type:
- **Dynamic code-block extraction (preferred for offline docs)**
- Extract Python/Bash code blocks from markdown AST analyzer, then execute them directly in tests.
- Benefit: test logic stays automatically aligned with docs.
- Basic idea: Use `ReadmeSnippet.extract_readme_snippets` to extract a list of code blocks as a global variable in file,
use this list as `pytest.mark.parametrize` parameters, and pass each snippet item to `example_runner.run` inside the parametrized test.
Additionally pass an `output_subfolder` argument for the 2nd-level output folder explained in **Output Directory Structure** below.
If any extra environment variable is need for a test (e.g., the example script reads it), `example_runner.run` also accepts a 3rd `env` parameter.
- See [tests/examples/offline_inference/test_text_to_image.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/examples/offline_inference/test_text_to_image.py) for reference implementation.
- **Explicit copied scripts (used by online docs for now until further update)**
- For online serving pages, it is acceptable to copy code from docs into dedicated test functions, because only client-side, request-sending scripts are tested.
- Benefit: dynamic extraction is overly complex: need to tell server-launch and client-request scripts.
- Requirement: copied test code must be kept in sync with doc updates.
**Test Case Naming Convention**
- Dynamic code extraction (auto-generated internally):
- `test_{single_function_name_matching_file_name}[h2_heading_00X]`
- Example: `test_text_to_image[basic_usage_001]`
- Explicit copied scripts:
- `test_{h2_heading_00X}[{dummy_param_id_for_omni_server}]`
- Example: `test_api_calls_001[omni_server0]`
**Runtime Configuration**
In the example code tests, do **not** reduce `num_inference_steps` just to speed up the tests unless there is a strong CI reliability reason to do otherwise.
**Skipping Rules**
You may skip examples falling in the following categories using `pytest.mark.skip` or `pytest.skip`:
- Gradio UI scripts
- Scenarios that significantly overlap with existing tests and add little new coverage.
**Output Directory Structure**
Use a three-layer output structure to store output artifacts:
1. Root output directory
- Auto-detected from `OUTPUT_DIR` env var or auto-generated under `/tmp`.
2. Doc-page directory
- Define and use a clear page-level folder name in each `test_*.py` yourself (abbreviations are acceptable, e.g., `example_offline_t2i`).
3. Test-case directory
- Must match the case identifier (e.g., `basic_usage_001`).
- Auto-generated for dynamic extracted tests.
@@ -1,46 +0,0 @@
**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).
@@ -1,89 +0,0 @@
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).
-75
View File
@@ -1,75 +0,0 @@
# Test Guide
## Setting Up the Test Environment
### Creating a Container
vLLM-Omni provides an official Docker image for deployment. These images are built upon vLLM Docker images and are available on [Docker Hub](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which vLLM release it is based on.
For a local test environment, you can follow the steps below to create a container:
## Installing Dependencies
### vLLM & vLLM-Omni
vLLM-Omni is built based on vLLM. You can follow [install guide](../../getting_started/installation/README.md) to build your local environment.
### Test Case Dependencies
When running test cases, you may need to install the following dependencies:
```bash
uv pip install ".[dev]"
apt-get install -y ffmpeg
```
## Running Tests
Our test scripts use the pytest framework. First, please use `git clone https://github.com/vllm-project/vllm-omni.git` to download the vllm-omni source code. Then, in the root directory of vllm-omni, you can run the following commands in your local test environment to execute the corresponding test cases.
=== "L1 level"
```bash
cd tests
pytest -s -v -m "core_model and cpu"
```
The latest test command is available in the "Simple Unit Test" step of this [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
=== "L2 level"
```bash
cd tests
pytest -s -v -m "core_model and not cpu" --run-level=core_model
```
If you only want to run a specific test case, you can use:
```bash
pytest -s -v test_xxxx.py --run-level=core_model
```
If you only want to run specific test cases on a particular platform, you can use:
```bash
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
```
The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
=== "L3 level & L4 level"
```bash
cd tests
pytest -s -v -m "advanced_model" --run-level=advanced_model
```
If you only want to run L3 test case, you can use:
```bash
pytest -s -v e2e/ --ignore-glob='*expansion.py' -m "advanced_model" --run-level=advanced_model
```
If you only want to run a specific test case, you can use:
```bash
pytest -s -v test_xxxx.py --run-level=advanced_model
```
If you only want to run specific test cases on a particular platform, you can use:
```bash
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
```
Note: To run performance tests, use:
```bash
pytest -s -v perf/scripts/run_benchmark.py
```
The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
The latest L4 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml).
You can find more information about markers in the documentation: [marker doc](./tests_markers.md)
## Adding New Test Cases
Please refer to the [L5 Layering Specification document](./CI_5levels.md).
+9 -30
View File
@@ -17,7 +17,6 @@ Defined in `pyproject.toml`:
| `gpu` | Tests that run on GPU * |
| `cuda` | Tests that run on CUDA * |
| `rocm` | Tests that run on AMD/ROCm * |
| `xpu` | Tests that run on Intel XPU * |
| `npu` | Tests that run on NPU/Ascend * |
| `H100` | Tests that require H100 GPU * |
| `L4` | Tests that require L4 GPU * |
@@ -33,7 +32,7 @@ Defined in `pyproject.toml`:
| `slow` | Slow tests (may skip in quick CI) |
| `benchmark` | Benchmark tests |
\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility).
\* Means those markers are auto-added, and they will be added by the `@hardware_test` decorator.
### Example usage for markers
@@ -50,13 +49,12 @@ from tests.utils import hardware_test
def test_video_to_audio()
...
```
### Decorator: `@hardware_test`
This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions:
1. **Applies platform and resource markers**
Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`).
Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `A2`, `A3`).
```
@pytest.mark.cuda
@pytest.mark.L4
@@ -81,14 +79,14 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
- Single call for multiple platforms:
```python
@hardware_test(
res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
num_cards={"cuda": 2, "rocm": 2, "xpu": 2, "npu": 2},
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
num_cards={"cuda": 2, "rocm": 2, "npu": 2},
)
```
or
```python
@hardware_test(
res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
num_cards=2,
)
```
@@ -100,24 +98,6 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
- ROCm only: `pytest -m "distributed_rocm and MI325"`
- NPU only: `pytest -m "distributed_npu"`
### Function: `hardware_marks`
`hardware_marks` returns a list of pytest mark objects with the same signature as `@hardware_test`. Use it when you need more flexibility, such as attaching hardware marks to individual `pytest.param` entries rather than an entire test function.
```python
from tests.utils import hardware_marks
MULTI_CARD_MARKS = hardware_marks(
res={"cuda": "H100", "rocm": "MI325", "npu": "A2"}, num_cards=2
)
@pytest.mark.parametrize("omni_server", [
pytest.param(OmniServerParams(...), id="case_001", marks=MULTI_CARD_MARKS),
], indirect=True)
def test_feature(omni_server):
...
```
## Add Support for a New Platform
If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps:
@@ -152,13 +132,12 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
# Optionally: add skipif_tpu when implemented
return [test_platform, test_resource, test_distributed]
```
3. **Update `hardware_marks` to recognize your new platform**:
In the relevant place (see the `hardware_marks` implementation), add:
3. **Update `hardware_test` to recognize your new platform**:
In the relevant place (see the `hardware_test` implementation), add:
```python
if platform == "tpu":
marks = tpu_marks(res=resource, num_cards=cards)
```
(`hardware_test` calls `hardware_marks` internally, so both will pick up the change.)
4. **(Recommended) Add a test using your new markers**:
```python
@hardware_test(
@@ -172,7 +151,7 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
**Summary**:
- Add pytest markers for your new platform/resources
- Implement a marker function (`xxx_marks`)
- Plug into `hardware_marks`
- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation!
- Plug into `hardware_test`
- You're done: tests decorated with `@hardware_test` using your platform now automatically get the correct markers, distribution, and isolation!
See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).
+8 -10
View File
@@ -24,6 +24,8 @@ End-to-end tests verify the complete functionality of a system or component. For
- **`tests/e2e/online_serving/`**: Tests for online serving scenarios (e.g., API server tests)
**Example:** The test file for `vllm_omni/entrypoints/omni_llm.py` should be located at `tests/entrypoints/test_omni_llm.py`.
## Test Directory Structure
The ideal directory structure mirrors the source code organization. Legend: `✅` = test exists, `⬜` = suggested to add.
@@ -73,6 +75,11 @@ vllm_omni/ tests/
│ └── arg_utils.py │ └── test_arg_utils.py ⬜
│
├── entrypoints/ → ├── entrypoints/
│ ├── omni.py │ ├── test_omni.py ⬜ (E2E covered by e2e/offline, e2e/online)
│ ├── omni_llm.py │ ├── test_omni_llm.py ✅
│ ├── omni_stage.py │ ├── test_omni_stage.py ⬜ (partial in test_omni_stage_diffusion_config.py)
│ ├── omni_diffusion.py │ ├── test_omni_diffusion.py ✅
│ ├── async_omni.py │ ├── test_async_omni.py ✅ actually in e2e/online_serving/test_async_omni.py
│ ├── async_omni_diffusion.py │ ├── test_async_omni_diffusion_config.py ✅
│ ├── stage_utils.py │ ├── test_stage_utils.py ✅
│ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve)
@@ -130,11 +137,9 @@ vllm_omni/ tests/
│
└── e2e (tests side only) → ├── e2e/
├── online_serving/ ✅ non-empty
│ ├── test_qwen2_5_omni.py
│ ├── test_async_omni.py
│ ├── test_qwen3_omni.py
│ ├── test_qwen3_omni_expansion.py
│ ├── test_mimo_audio.py
│ ├── test_image_gen_edit.py
│ ├── test_images_generations_lora.py
│ └── stage_configs/
@@ -153,25 +158,18 @@ vllm_omni/ tests/
├── test_diffusion_layerwise_offload.py
├── test_diffusion_lora.py
├── test_sequence_parallel.py
├── test_qwen_image_edit_expansion.py
└── stage_configs/
├── qwen2_5_omni_ci.yaml
├── qwen3_omni_ci.yaml
├── bagel_*.yaml
└── npu/, rocm/, etc.
examples/ tests
│ └── examples
├── online_serving/ → ├── online_serving/
│ └── {doc_page_title}/README.md │ └── test_{doc_page_title}.py ⬜
└── offline_inference/ → └── offline_inference/
└── {doc_page_title}/README.md └── test_{doc_page_title}.py ⬜
```
### Naming Conventions
- **Unit Tests**: Use `test_<module_name>.py` format. Example: `stage_utils.py` → `test_stage_utils.py`
- **Unit Tests**: Use `test_<module_name>.py` format. Example: `omni_llm.py` → `test_omni_llm.py`
- **E2E Tests**: Place in `tests/e2e/offline_inference/` or `tests/e2e/online_serving/` with descriptive names. Example: `tests/e2e/offline_inference/test_qwen3_omni.py`, `tests/e2e/offline_inference/test_diffusion_model.py`
@@ -651,9 +651,8 @@ For a fair comparison, keep the same **prompt**, **seed**, **resolution**, **num
#### 5.2 Functionality Check in CI
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 [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.
To ensure project maintainability and sustainable development, we encourage contributors to submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
---
@@ -721,22 +720,6 @@ See detailed guide: [How to add Sequence Parallel support](../../design/feature/
omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2)
```
### Step Execution
See detailed design guide: [How to add step execution support](../../design/feature/diffusion_step_execution.md)
Use this only when your pipeline can be split into stable request-scoped and
step-scoped phases. The reference implementation is
`QwenImagePipeline`, which maps its request-level `forward()` into:
1. `prepare_encode()` for prompt encoding, latent init, timestep prep, and per-request scheduler setup.
2. `denoise_step()` for one transformer/noise prediction.
3. `step_scheduler()` for one scheduler update and `step_index` advance.
4. `post_decode()` for the final VAE decode.
Do not enable `step_execution=True` until those four methods are implemented
and validated against the request-level path.
### Cache Acceleration
#### TeaCache
@@ -813,107 +796,6 @@ class WanTransformer3DModel(nn.Module):
---
### Diffusion Timing (Performance Profiling)
When adapting a new diffusion model, it is often useful to analyze the latency of key components such as text encoding, diffusion denoising, and VAE decoding.
vLLM-Omni provides a timing utility via `DiffusionPipelineProfilerMixin` to help developers quickly identify performance bottlenecks.
!!! info
`DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md#3-profiling-diffusion-models). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
This tool automatically measures the execution time of selected pipeline modules and prints the results in the logs.
**Enabling Diffusion Timing**
Enable timing by setting:
```
vllm serve Qwen/Qwen-Image --omni --port 8091 --enable-diffusion-pipeline-profiler
```
You can optionally specify which modules to profile:
```
class YourPipeline(xxx, DiffusionPipelineProfilerMixin):
def __init__(self, xxx):
...
self.setup_diffusion_pipeline_profiler(profiler_targets=["diffuse"], enable_diffusion_pipeline_profiler)
```
If not specified, the default targets are used:
```
["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"]
```
**Adding DiffusionPipelineProfilerMixin to a Pipeline**
To enable timing support in your pipeline, inherit from DiffusionPipelineProfilerMixin.
```python
from vllm_omni.diffusion.utils.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
# Optional: Specify custom timing targets
_PROFILER_TARGETS = ["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"]
def __init__(
self,
*,
od_config: OmniDiffusionConfig,
prefix: str = "",
):
super().__init__()
self.od_config = od_config
self.parallel_config = od_config.parallel_config
# initialize pipeline components
...
# initialize timing profiler
self.setup_diffusion_pipeline_profiler(enable_diffusion_pipeline_profiler)
```
The mixin dynamically wraps selected methods and records their execution time during inference.
If you need to fetch the execution time of different modules, you will need to pass `self.stage_durations` to `DiffusionOutput`, as shown below:
```diff
- return DiffusionOutput(output=img)
+ return DiffusionOutput(
output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
)
```
**Pipeline Design for Timing**
The current diffusion timing utility is function-based, meaning it measures the execution time of individual methods.
When implementing a new pipeline, avoid putting all logic inside a single function (e.g., forward). Instead, structure the pipeline in a modular way by separating key stages into independent methods, such as the diffusion loop.
For example:
```
def forward(self, req: OmniDiffusionRequest):
prompt_embeds = self.encode_prompt(req)
latents = self.diffuse(prompt_embeds, req)
images = self.vae.decode(latents)
return DiffusionOutput(output=images)
```
This allows the timing utility to measure each stage (e.g., encode_prompt, diffuse, vae.decode) separately and helps identify performance bottlenecks more easily.
**Default Profiled Modules**
By default, the following pipeline modules are timed:
```
vae.encode
vae.decode
diffuse
text_encoder.forward
tokenizer.forward
```
**Example Output**
When enabled, timing logs appear like this:
```
[DiffusionTiming] text_encoder.forward took 0.018s
[DiffusionTiming] diffuse took 2.412s
[DiffusionTiming] vae.decode took 0.063s
```
These measurements help identify bottlenecks during model adaptation and optimization
## Troubleshooting
@@ -965,7 +847,7 @@ hidden_states = hidden_states.reshape(batch_size, seq_len, -1)
1. **Reduce batch size:**
```python
omni.generate(prompts=[...], max_num_seqs=2)
omni.generate(prompts=[...], max_batch_size=2)
```
2. **Use smaller image size:**
+35 -27
View File
@@ -330,46 +330,54 @@ Stage transitions are the mechanism by which outputs from one stage are converte
### Where Stage Transitions Are Called
Stage transitions happen automatically in the runtime orchestrator. Here's the detailed flow:
Stage transitions happen automatically in the orchestrator (`OmniLLM` class) during the generation loop. Here's the detailed flow:
1. **Location**: `vllm_omni/engine/orchestrator.py` in `_forward_to_next_stage()`
1. **Location**: `vllm_omni/entrypoints/omni_llm.py` in the `_run_generation()` method
2. **Trigger**: When a stage completes processing and produces outputs
3. **Execution Flow**:
```python
# In orchestrator.py
next_stage_id = stage_id + 1
next_client = self.stage_clients[next_stage_id]
params = req_state.sampling_params_list[next_stage_id]
# In omni_llm.py, _run_generation() method (around line 345-460)
# Save current stage outputs so stage_input_processors can consume them.
self.stage_clients[stage_id].set_engine_outputs([output])
# Main orchestrator loop polls each stage for completed requests
for stage_id, stage in enumerate(self.stage_list):
result = stage.try_collect() # Get completed request
if result is None:
continue
# THIS IS WHERE STAGE TRANSITION HAPPENS
next_inputs = next_client.process_engine_inputs(
stage_list=self.stage_clients,
prompt=req_state.prompt,
)
# Store outputs from this stage
engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
stage.set_engine_outputs(engine_outputs)
# Build and submit request(s) to the next stage.
for next_input in next_inputs:
request = build_engine_core_request_from_tokens(
request_id=req_id,
prompt=next_input,
params=params,
model_config=self.stage_vllm_configs[next_stage_id].model_config,
)
await next_client.add_request_async(request)
# Check if there's a next stage to forward to
next_stage_id = stage_id + 1
if next_stage_id < len(self.stage_list):
next_stage: OmniStage = self.stage_list[next_stage_id]
# THIS IS WHERE STAGE TRANSITION HAPPENS
next_inputs = next_stage.process_engine_inputs(
self.stage_list,
[request_id_to_prompt[req_id]]
)
# Submit to next stage
task = {
"type": OmniStageTaskType.GENERATE,
"request_id": req_id,
"engine_inputs": next_inputs[0],
"sampling_params": sampling_params_list[next_stage_id],
}
next_stage.submit(task)
```
### How Stage Transitions Work
The stage transition process follows these steps:
1. **Stage Completion**: When a stage finishes processing a request, the orchestrator stores outputs via `stage_client.set_engine_outputs(...)`
1. **Stage Completion**: When a stage finishes processing a request, it stores outputs via `stage.set_engine_outputs(engine_outputs)`
2. **Transition Detection**: The orchestrator checks if there's a next stage and calls `process_engine_inputs()` on it
3. **Input Processing**: The stage input processor configured in stage YAML (under `vllm_omni/model_executor/stage_input_processors/`) handles the transition:
3. **Input Processing**: The `process_engine_inputs()` method in `OmniStage` (`omni_stage.py`) handles the transition:
```python
def process_engine_inputs(
self, stage_list: list[Any], prompt: OmniTokensPrompt | TextPrompt = None
@@ -452,7 +460,7 @@ def forward(self, ...):
These keys are then accessible in your stage transition function:
```python
# In stage_input_processors/qwen3_omni.py
thinker_prefill_embeddings = output.multimodal_output["0"] # Access by key
thinker_embeddings = output.multimodal_output["0"] # Access by key
thinker_hidden_states = output.multimodal_output["24"]
```
@@ -505,11 +513,11 @@ def thinker2talker(
for thinker_output in thinker_outputs:
output = thinker_output.outputs[0]
# Extract thinker embeddings and hidden states
thinker_prefill_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
thinker_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda()
info = {
"thinker_prefill_embeddings": thinker_prefill_embeddings,
"thinker_embeddings": thinker_embeddings,
"thinker_hidden_states": thinker_hidden_states,
"thinker_sequences": thinker_output.prompt_token_ids + output.token_ids,
"thinker_input_ids": thinker_output.prompt_token_ids,
@@ -1,478 +0,0 @@
# Support FP8 Quantization
This section describes how to add FP8 quantization support to a diffusion model. We use the Qwen-Image pipeline as the reference implementation.
---
## Table of Contents
- [Overview](#overview)
- [Step-by-Step Implementation](#step-by-step-implementation)
- [Testing](#testing)
- [Troubleshooting](#troubleshooting)
- [Reference Implementations](#reference-implementations)
- [Summary](#summary)
---
## Overview
### What is FP8 Quantization?
FP8 quantization reduces model memory footprint by storing weights in 8-bit floating point (`float8_e4m3fn`) instead of 16-bit (`bfloat16`/`float16`). This saves ~50% of weight memory.
A diffusion pipeline typically has three quantizable components:
| Component | Layer Types | Quantization Approach |
|-----------|-----------|----------------------|
| **DiT (transformer)** | `nn.Linear` | Replace with vLLM quantized linear layers |
| **Text encoder** | `nn.Linear` | FP8 weight storage via hooks |
| **VAE** | `nn.Conv2d`, `nn.Conv3d` | FP8 weight storage via hooks |
### Architecture
The quantization system has two mechanisms:
**1. vLLM Quantized Linear Layers (for DiT)**
vLLM provides drop-in replacements for `nn.Linear` that perform FP8 W8A8 computation on supported hardware. These are configured via `quant_config` passed to layer constructors.
Key classes:
| Class | Location | Purpose |
|-------|----------|---------|
| `build_quant_config` | `vllm_omni/quantization/__init__.py` | Build quantization config from string/dict/per-component spec |
| `ComponentQuantizationConfig` | `vllm_omni/quantization/component_config.py` | Per-component routing by layer prefix |
**2. FP8 Weight Storage via Hooks (for text encoder and VAE)**
For models loaded via `from_pretrained()` (outside vLLM's weight pipeline), we use forward hooks to store weights in FP8 and dequantize on-the-fly:
```
At rest: weight.data = fp8_tensor (half memory)
Pre-hook: weight.data = fp8_weight.to(bf16) * scale
Forward: Normal computation in BF16
Post-hook: weight.data = fp8_weight (back to FP8)
```
Key function:
| Function | Location | Purpose |
|----------|----------|---------|
| `apply_fp8_weight_storage` | `vllm_omni/diffusion/models/utils.py` | Apply FP8 hooks to Linear/Conv layers (introduced in [#1412](https://github.com/vllm-project/vllm-omni/pull/1412)) |
---
## Step-by-Step Implementation
### Step 1: Add `quant_config` to the Transformer
Pass `quant_config` to vLLM parallel linear layers in the transformer. This enables FP8 W8A8 for DiT linear layers.
#### 1.1: Accept `quant_config` in Constructor
```python
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
class YourTransformer2DModel(nn.Module):
def __init__(
self,
*,
od_config: OmniDiffusionConfig,
quant_config: QuantizationConfig | None = None,
# ... other params
):
self.quant_config = quant_config
# Pass to sub-modules
self.blocks = nn.ModuleList([
YourTransformerBlock(quant_config=quant_config, ...)
for _ in range(num_layers)
])
```
#### 1.2: Pass `quant_config` to vLLM Linear Layers
Replace `nn.Linear` with vLLM parallel linear layers that accept `quant_config`:
```python
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
QKVParallelLinear,
RowParallelLinear,
)
class YourAttentionBlock(nn.Module):
def __init__(self, hidden_size, num_heads, quant_config=None, prefix=""):
# QKV projection
self.to_qkv = QKVParallelLinear(
hidden_size=hidden_size,
head_size=hidden_size // num_heads,
total_num_heads=num_heads,
quant_config=quant_config,
prefix=f"{prefix}.to_qkv",
)
# Output projection
self.to_out = RowParallelLinear(
input_size=hidden_size,
output_size=hidden_size,
quant_config=quant_config,
prefix=f"{prefix}.to_out",
)
class YourFeedForward(nn.Module):
def __init__(self, hidden_size, intermediate_size, quant_config=None, prefix=""):
self.gate_up_proj = ColumnParallelLinear(
input_size=hidden_size,
output_size=intermediate_size * 2,
quant_config=quant_config,
prefix=f"{prefix}.gate_up_proj",
)
self.down_proj = RowParallelLinear(
input_size=intermediate_size,
output_size=hidden_size,
quant_config=quant_config,
prefix=f"{prefix}.down_proj",
)
```
!!! important
Always pass the `prefix` parameter. It is used by `ignored_layers` matching to skip sensitive layers.
#### 1.3: Identify Sensitive Layers
Some layers degrade output quality when quantized. Common sensitive layers:
| Layer | Why Sensitive | Typical Models | Quant Methods Affected |
|-------|-------------|---------------|----------------------|
| `img_mlp` | Processes denoising latents with shifting dynamic range | Qwen-Image | FP8, Int8 |
| `feed_forward` | FFN layers in DiT blocks, large dynamic range | Z-Image | Int8 |
| `proj_out` | Final output projection, small quantization errors amplified | Various | FP8, Int8 |
| `lm_head` | Output vocabulary projection, precision-critical | Omni (thinker) | NVFP4 |
| `mlp.gate` | MoE router gates, routing decisions are precision-critical | MoE models | NVFP4 |
**Quantization quality benchmarks** (LPIPS: lower = better, <0.01 imperceptible, >0.1 noticeable):
| Model | Method | All Layers | With `ignored_layers` | Recommendation |
|-------|--------|-----------|----------------------|----------------|
| Qwen-Image-2512 | Int8 | LPIPS 0.0197 | 0.0027 (skip `img_mlp`) | Skip `img_mlp` |
| Z-Image-Turbo | Int8 | LPIPS 0.1597 | 0.0290 (skip `feed_forward`) | Skip `feed_forward` |
| Z-Image-Turbo | FP8 | LPIPS ~0.005 | — | All layers OK |
To identify sensitive layers for a new model:
1. Run inference with all layers quantized
2. Compare output quality with BF16 baseline using LPIPS (see `benchmarks/diffusion/quantization_quality.py`)
3. Selectively disable quantization per layer (via `ignored_layers`) until quality matches
4. Document the recommended `ignored_layers` in the model's supported models table
### Step 2: Wire Up Quantization in the Pipeline
#### 2.1: Create `quant_config` and Pass to Transformer
In the pipeline's `__init__`, extract the vLLM quant config and pass it to the transformer:
```python
from vllm_omni.quantization import build_quant_config
class YourPipeline(nn.Module):
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__()
# ... load scheduler, text_encoder, vae ...
# Build quant config from OmniDiffusionConfig
quant_config = build_quant_config(od_config.quantization_config)
# Pass to transformer
self.transformer = YourTransformer2DModel(
od_config=od_config,
quant_config=quant_config,
# ... other kwargs
)
```
#### 2.2: Apply FP8 Weight Storage to Text Encoder and VAE
For models loaded via `from_pretrained()`, apply hook-based FP8 weight storage.
!!! note
The `apply_fp8_weight_storage` function is provided by `vllm_omni/diffusion/models/utils.py`, introduced in [PR #1412](https://github.com/vllm-project/vllm-omni/pull/1412).
```python
from vllm_omni.diffusion.models.utils import apply_fp8_weight_storage
class YourPipeline(nn.Module):
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__()
self.text_encoder = SomeTextEncoder.from_pretrained(...)
self.vae = SomeVAE.from_pretrained(...).to(self.device)
# Apply FP8 weight storage to text encoder and VAE
if (
od_config.quantization_config is not None
and getattr(od_config.quantization_config, "quant_method", None) == "fp8"
):
apply_fp8_weight_storage(self.vae)
apply_fp8_weight_storage(self.text_encoder)
# ... rest of init
```
`apply_fp8_weight_storage` walks all `nn.Linear`, `nn.Conv2d`, and `nn.Conv3d` modules in the model and applies FP8 hooks to each.
#### 2.3: Fix `load_weights` to Include Pre-loaded Components
Since text encoders and VAEs are loaded via `from_pretrained()` (not through vLLM's weight pipeline), their parameters won't appear in the `loaded_weights` set. Mark them explicitly:
```python
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self)
loaded_weights = loader.load_weights(weights)
# Mark pre-loaded component weights as loaded
loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
return loaded_weights
```
### Step 3: Add to Supported Models Documentation
Update the FP8 supported models table in `docs/user_guide/diffusion/quantization/fp8.md`:
```markdown
| Your-Model | `org/Your-Model` | ✅ | ✅ | ✅ | `sensitive_layer` |
```
---
## Testing
### Quick Verification
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# Test with FP8 enabled
omni = Omni(model="your-model", quantization="fp8")
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
### Quality Gate Test (LPIPS)
We provide a pytest-based quality gate at `tests/diffusion/quantization/test_quantization_quality.py`.
It generates outputs with both BF16 and your quantized method using the same seed, computes
[LPIPS](https://github.com/richzhang/PerceptualSimilarity) perceptual distance, and **fails if it
exceeds a threshold**. This is the recommended way to validate that a quantization method does not
unacceptably degrade output quality.
**Step 1: Add your test config** — edit `QUALITY_CONFIGS` in the test file:
```python
QualityTestConfig(
id="int8_z_image", # pytest ID
model="Tongyi-MAI/Z-Image-Turbo",
quantization="int8", # your quantization method
task="t2i", # "t2i" or "t2v"
prompt="a cup of coffee on a wooden table, morning light",
max_lpips=0.15, # fail threshold (see guide below)
num_inference_steps=20,
),
```
**Step 2: Run the test:**
```bash
# Install dependency
pip install lpips
# Run all quality tests
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m ""
# Run only your method
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "int8"
# Run a specific model
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image"
```
**Step 3: Read the output:**
```
Quantization Quality: int8_z_image
Model: Tongyi-MAI/Z-Image-Turbo
Method: int8
LPIPS: 0.0290 (threshold: 0.15)
BF16 memory: 19.15 GiB
Quant memory: 14.80 GiB (23% reduction)
Result: PASS
```
**LPIPS threshold guide:**
| LPIPS Range | Meaning | Recommendation |
|-------------|---------|----------------|
| < 0.01 | Imperceptible | Ideal — no visible difference |
| 0.01 – 0.05 | Minor differences | Good — acceptable for most use cases |
| 0.05 – 0.15 | Noticeable on close inspection | Acceptable with `ignored_layers` tuning |
| 0.15 – 0.30 | Clearly visible | Investigate — likely needs `ignored_layers` |
| > 0.30 | Significant degradation | Quantization method may not suit this model |
!!! tip
If your test fails, don't raise the threshold — instead, identify sensitive layers and add
them to `ignored_layers` (see next section). Then re-run the test with the lower LPIPS.
### Quality Benchmark (Detailed)
For a full multi-prompt benchmark with Markdown output (suitable for PR descriptions), use the
standalone benchmark script:
```bash
python benchmarks/diffusion/quantization_quality.py \
--model Tongyi-MAI/Z-Image-Turbo \
--quantization fp8 \
--prompts \
"a cup of coffee on the table" \
"an aerial view of a coral reef" \
--height 1024 --width 1024 \
--num-inference-steps 50 --seed 42
```
This generates BF16 and quantized outputs, computes per-prompt LPIPS, and prints a Markdown
summary table ready to paste into your PR.
### Identifying Sensitive Layers with `ignored_layers`
Not all layers tolerate quantization equally. **`ignored_layers`** lets you keep specific layers
in BF16 while quantizing everything else. This is critical for achieving good quality with
aggressive quantization methods (Int8, NVFP4).
**Why some layers are sensitive:**
| Layer Type | Why Sensitive | Typical Impact |
|------------|--------------|----------------|
| `img_mlp` | Processes denoising latents with shifting dynamic range across timesteps. No built-in normalization to absorb quantization error. | Color shifts, blurriness |
| `feed_forward` | FFN layers with large dynamic range in DiT blocks. | Artifacts, loss of detail |
| `proj_out` | Final output projection — small errors get amplified since there are no subsequent layers to correct them. | Overall quality loss |
| `lm_head` | Vocabulary projection in LLMs — precision-critical for token selection. | Garbage text output |
| `mlp.gate` | MoE router gates — routing decisions are binary and precision-critical. | Wrong expert selection, quality collapse |
**How to find sensitive layers for your model:**
1. Run the quality gate test with all layers quantized:
```bash
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "your_model"
```
2. If LPIPS exceeds your threshold, try skipping common sensitive layers one at a time:
```python
QualityTestConfig(
id="fp8_your_model_skip_mlp",
model="your-org/Your-Model",
quantization="fp8",
task="t2i",
prompt="a cup of coffee on a wooden table",
max_lpips=0.05,
),
```
And in the test, modify the Omni call:
```python
omni_qt = Omni(
model=config.model,
quantization_config={
"method": config.quantization,
"ignored_layers": ["img_mlp"],
},
)
```
3. Compare LPIPS with and without `ignored_layers` to isolate the problematic layers.
4. Document the recommended `ignored_layers` in the supported models table.
**Real-world examples:**
| Model | Method | All Layers (LPIPS) | With `ignored_layers` | Which Layers to Skip |
|-------|--------|-------------------|----------------------|---------------------|
| Qwen-Image-2512 | Int8 | 0.0197 | 0.0027 | `img_mlp` |
| Z-Image-Turbo | Int8 | 0.1597 | 0.0290 | `feed_forward` |
| Z-Image-Turbo | FP8 | ~0.005 | — | None needed |
### What to Verify (Checklist)
1. **Model loads without errors** — Check logs for quantization initialization messages
2. **Quality gate passes** — LPIPS within threshold vs BF16 baseline
3. **Memory reduction** — Expect ~30-50% savings on quantized components
4. **Sensitive layers identified** — Document which layers (if any) need `ignored_layers`
5. **`ignored_layers` pattern works** — Verify that the layer name prefix matches your transformer's module names
---
## Troubleshooting
### Issue: Output Quality Degradation
**Symptoms:** Generated images have visible artifacts, color shifts, or blurriness compared to BF16 baseline.
**Causes & Solutions:**
- **Sensitive layers being quantized:**
**Solution:** Identify and skip sensitive layers:
```python
omni = Omni(
model="your-model",
quantization_config={
"method": "fp8",
"ignored_layers": ["img_mlp"],
},
)
```
- **All layers sensitive (rare):**
**Solution:** The model may not be suitable for FP8 quantization. Consider weight-only quantization or other methods.
### Issue: Missing Weights Warning
**Symptoms:** Warnings about unloaded weights for `vae.*` or `text_encoder.*` parameters.
**Solution:** Ensure `load_weights` marks pre-loaded component weights. See [Step 2.3](#23-fix-load_weights-to-include-pre-loaded-components).
### Issue: CUDA Out of Memory
**Symptoms:** OOM errors during model loading despite FP8 being enabled.
**Causes & Solutions:**
- **Text encoder/VAE loaded in BF16 before FP8 conversion:**
FP8 weight storage is applied after `from_pretrained()`, so peak load memory is still BF16. This is a one-time cost during initialization.
---
## Reference Implementations
| Model | Pipeline File | Transformer File | Notes |
|-------|-------------|-----------------|-------|
| **Z-Image** | `models/z_image/pipeline_z_image.py` | `models/z_image/z_image_transformer.py` | DiT + text encoder FP8 |
| **Qwen-Image** | `models/qwen_image/pipeline_qwen_image.py` | `models/qwen_image/qwen_image_transformer.py` | DiT + text encoder + VAE FP8 |
All files are under `vllm_omni/diffusion/`.
---
## Summary
Adding FP8 quantization support to a new model:
1. ✅ **Add `quant_config` to transformer** — Accept and pass to vLLM parallel linear layers with proper `prefix`
2. ✅ **Identify sensitive layers** — Profile which layers need `ignored_layers` for quality
3. ✅ **Apply FP8 weight storage** — Call `apply_fp8_weight_storage()` on text encoder and VAE in pipeline `__init__`
4. ✅ **Fix `load_weights`** — Mark pre-loaded component weights as loaded
5. ✅ **Update docs** — Add model to the supported models table in `docs/user_guide/diffusion/quantization/fp8.md`
6. ✅ **Test** — Verify quality, memory savings, and sensitive layer behavior
-584
View File
@@ -1,584 +0,0 @@
# Adding a TTS Model
This guide walks through adding a new TTS model to vLLM-Omni, using **Qwen3-TTS**
as a reference. Qwen3-TTS demonstrates the standard two-stage TTS pipeline and the
key optimizations all TTS models in this repo should follow.
## Table of Contents
1. [Overview](#overview)
2. [Directory Structure](#directory-structure)
3. [Step-by-Step Implementation](#step-by-step-implementation)
4. [Key Components](#key-components)
5. [Model Registration](#model-registration)
6. [Stage Configuration](#stage-configuration)
7. [Stage Input Processors](#stage-input-processors)
8. [Testing](#testing)
9. [Summary](#summary)
## Overview
vLLM-Omni supports TTS models as multi-stage pipelines where each stage runs independently
and can be placed on different devices. Qwen3-TTS has two stages:
| Stage | Name | Input | Output |
|-------|------|-------|--------|
| 0 | Code Predictor (AR) | Text tokens | Discrete RVQ codec codes |
| 1 | Code2Wav (Decoder) | RVQ codec codes | Audio waveform |
Each stage is a separate model class configured independently via YAML. The two stages
are connected by the `async_chunk` framework, which enables inter-stage streaming for
low first-packet latency (see [Async Chunk Design](../../design/feature/async_chunk_design.md)).
### Without async_chunk (batch mode)
Stage 0 runs to completion before Stage 1 starts, resulting in long first-packet latency:
```mermaid
flowchart TB
subgraph stage0["Stage 0: AR Code Predictor"]
direction LR
P[Prefill] --> D1[Decode 1]
D1 --> D2[Decode 2]
D2 --> Dots1["..."]
Dots1 --> DN[Decode N]
end
subgraph stage1["Stage 1: Code2Wav"]
direction LR
DEC[Decode all codes at once]
end
stage0 -- "all N codes" --> stage1
stage1 --> FPL["First Packet Latency = Stage 0 + Stage 1"]
style stage0 fill:#dae8fc,stroke:#6c8ebf
style stage1 fill:#f8d7c8,stroke:#d4856a
style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px
```
### With async_chunk (streaming mode)
Stage 0 sends codec codes to Stage 1 every `chunk_size=25` tokens. Stage 1 begins decoding
immediately, reducing first-packet latency from the full AR time to just the first chunk:
```mermaid
flowchart TB
subgraph stage0["Stage 0: Code Predictor (AR)"]
direction LR
P[Prefill] --> D1["Decode 1-25"]
D1 --> D2["Decode 26-50"]
D2 --> Dots1["..."]
Dots1 --> DN["Decode N"]
end
subgraph stage1["Stage 1: Code2Wav"]
direction LR
C1["Chunk 1\n(25 frames)"] --> C2["Chunk 2\n(context + 25)"]
C2 --> Dots2["..."]
Dots2 --> CN["Final chunk"]
end
D1 -. "chunk 1 (25 codes)" .-> C1
D2 -. "chunk 2 (context + 25)" .-> C2
DN -. "final" .-> CN
stage0 --> FPL["⏱ First Packet Latency = Prefill + 25 decode steps only"]
style stage0 fill:#dae8fc,stroke:#6c8ebf
style stage1 fill:#e8d4f8,stroke:#8a6cad
style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px
```
Key parameters: `chunk_size=25`, `left_context_size=25` (validated defaults from Qwen3-TTS
and Qwen3-Omni).
## Directory Structure
When adding a new TTS model, create the following structure:
```
vllm_omni/model_executor/models/
your_model_name/
__init__.py
your_model.py # Unified class (stage dispatch)
your_model_ar_stage.py # Stage 0: AR stage
your_model_decoder.py # Stage 1: audio decoder
vllm_omni/model_executor/stage_input_processors/
your_model_name.py # Stage 0 -> Stage 1 transition
vllm_omni/model_executor/stage_configs/
your_model_name.yaml # Batch mode config
your_model_name_async_chunk.yaml # Streaming mode config
```
**Qwen3-TTS reference files:**
| File | Purpose |
|------|---------|
| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | Stage 0 - optimized AR |
| `models/qwen3_tts/qwen3_tts_code2wav.py` | Stage 1 - decoder |
| `stage_configs/qwen3_tts.yaml` | Stage config (async_chunk enabled) |
| `stage_configs/qwen3_tts_batch.yaml` | Batch mode config |
| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
## Step-by-Step Implementation
### Step 1: Implement Stage 0 - AR Stage
Stage 0 is the autoregressive stage that generates intermediate audio representations.
**It must use vLLM's native decoder layers with fused ops and PagedAttention** for the LLM
backbone - this is the primary source of speedup over HuggingFace inference.
#### 1.1 Use vLLM Decoder Layers Directly
Build your transformer layers from the corresponding vLLM decoder layer class (e.g.
`Qwen3DecoderLayer` for Qwen3-based backbones, or the equivalent for LLaMA, Qwen2, etc.).
Do not wrap the HuggingFace model directly - that bypasses PagedAttention and fused kernels.
```python
# your_model_ar_stage.py
from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer
class YourTTSARStage(nn.Module):
def __init__(self, config, vllm_config, prefix):
self.layers = nn.ModuleList([
Qwen3DecoderLayer(
config, vllm_config=vllm_config, prefix=f"{prefix}.layers.{i}"
)
for i in range(config.num_hidden_layers)
])
self.lm_head = ParallelLMHead(config.codec_size, config.hidden_size)
```
See `qwen3_tts_code_predictor_vllm.py` for the full implementation.
#### 1.2 Forward Pass
Implement `forward()` to return an `OmniOutput` with intermediate data for Stage 1:
```python
def forward(self, input_ids, positions, intermediate_tensors=None,
inputs_embeds=None, **kwargs) -> OmniOutput:
hidden_states = self.run_layers(input_ids, positions, intermediate_tensors, inputs_embeds)
logits = self.lm_head(hidden_states)
return OmniOutput(
text_hidden_states=hidden_states,
multimodal_outputs={
"audio_codes": self.extract_codes(logits),
},
)
```
The keys in `multimodal_outputs` are what your stage input processor will read to build
Stage 1 inputs.
#### 1.3 Weight Loading with Fused QKV
When using vLLM's fused `QKVParallelLinear`, pack the HF `q_proj`/`k_proj`/`v_proj` weights
into `qkv_proj` using `stacked_params_mapping`. See the `load_weights()` method in
`qwen3_tts_code_predictor_vllm.py` for the standard pattern - it can be reused as-is
for any Qwen-family backbone.
#### 1.4 Custom Stop Condition (if needed)
Some TTS models use a learned stop head rather than an EOS token. If your model does this,
implement it inside `sample()`:
```python
def sample(self, logits, sampling_metadata) -> SamplerOutput | None:
output = self.sampler(logits, sampling_metadata)
if self._stop_head_fired():
output = mark_as_finished(output)
return output
```
### Step 2: Implement Stage 1 - Decoder
Stage 1 decodes Stage 0 output into audio. It runs outside the scheduler (no PagedAttention
needed). Implement `chunked_decode_streaming()` to support async_chunk streaming:
```python
# your_model_decoder.py
class YourTTSDecoder(nn.Module):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
# Initialize your audio decoder (SpeechTokenizer, HiFiGAN, etc.)
def forward(self, codes: torch.Tensor, **kwargs) -> torch.Tensor:
return self.decoder(codes)
def chunked_decode_streaming(self, codes, chunk_size=25,
left_context_size=25) -> torch.Tensor:
"""Decode with a sliding context window for smooth chunk boundaries."""
end_index = codes.shape[-1]
context_size = 0 if end_index <= chunk_size else left_context_size
wav_chunk = self(codes)
# Trim left context to avoid duplicate audio
return wav_chunk[..., context_size * self.total_upsample:]
```
### Step 3: Implement the Unified Model Class
The unified class dispatches to the correct stage based on `model_stage` in the config:
```python
# your_model.py
class YourTTSModelForConditionalGeneration(nn.Module, SupportsPP):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
self.model_stage = vllm_config.model_config.model_stage
if self.model_stage == "ar_stage":
ar_vllm_config = vllm_config.with_hf_config(
vllm_config.model_config.hf_config.ar_config,
architectures=["YourTTSARStageForConditionalGeneration"],
)
self.ar_stage = init_vllm_registered_model(
vllm_config=ar_vllm_config,
prefix=maybe_prefix(prefix, "ar"),
hf_config=ar_vllm_config.model_config.hf_config,
architectures=["YourTTSARStageForConditionalGeneration"],
)
self.model = self.ar_stage
elif self.model_stage == "decoder":
self.decoder = YourTTSDecoder(vllm_config=vllm_config, prefix=prefix)
self.model = self.decoder
```
### Step 4: Create `__init__.py`
```python
# vllm_omni/model_executor/models/your_model_name/__init__.py
from .your_model import YourTTSModelForConditionalGeneration
__all__ = ["YourTTSModelForConditionalGeneration"]
```
## Key Components
### Model Interfaces
Your unified model class should implement the appropriate interfaces:
- **`SupportsPP`**: Required for pipeline parallelism support (all models should implement this)
- **`SupportsMultiModal`**: Only if your model accepts multimodal inputs (e.g. reference audio for voice cloning)
### Output Format
Use `OmniOutput` so the orchestrator can route intermediate data between stages:
```python
from vllm_omni.model_executor.models.output_templates import OmniOutput
return OmniOutput(
text_hidden_states=hidden_states,
multimodal_outputs={
"audio_codes": codec_codes,
},
)
```
### Weight Loading from a Single Checkpoint
If both stages load from one checkpoint, separate them by prefix in the unified class:
```python
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
ar_weights, decoder_weights = [], []
for name, tensor in weights:
if name.startswith("decoder."):
decoder_weights.append((name, tensor))
else:
ar_weights.append((name, tensor))
if self.model_stage == "ar_stage":
return self.ar_stage.load_weights(ar_weights)
elif self.model_stage == "decoder":
return self.decoder.load_weights(decoder_weights)
```
## Model Registration
Register all stage classes in `vllm_omni/model_executor/models/registry.py`:
```python
_OMNI_MODELS = {
# (package_name, module_name, class_name)
"YourTTSModelForConditionalGeneration": (
"your_model_name", "your_model",
"YourTTSModelForConditionalGeneration",
),
"YourTTSARStageForConditionalGeneration": (
"your_model_name", "your_model_ar_stage",
"YourTTSARStageForConditionalGeneration",
),
"YourTTSDecoder": (
"your_model_name", "your_model_decoder",
"YourTTSDecoder",
),
}
```
The registry uses lazy loading - model classes are only imported when needed.
## Stage Configuration
Each stage has a `worker_type` that determines how it is scheduled:
- `worker_type: ar` - autoregressive stage, uses `OmniARScheduler` with PagedAttention
- `worker_type: generation` - non-AR stage (e.g. decoder), uses `OmniGenerationScheduler`
Key configuration fields:
| Field | Description |
|-------|-------------|
| `model_stage` | Which stage to initialize (`ar_stage`, `decoder`, etc.) |
| `model_arch` | Architecture name, must match `registry.py` |
| `engine_input_source` | List of upstream stage IDs that provide input (e.g. `[0]`) |
| `engine_output_type` | Output type: `latent` for intermediate, `audio` for final |
| `custom_process_next_stage_input_func` | Async chunk processor function path (streaming only) |
| `final_output` | Whether this stage produces the final user-facing output |
| `final_output_type` | Type of final output (`audio`, `text`, etc.) |
### Batch mode
```yaml
# stage_configs/your_model_name.yaml
stage_args:
- stage_id: 0
stage_type: llm
runtime:
devices: "0"
engine_args:
model_stage: ar_stage
max_num_seqs: 64
model_arch: YourTTSModelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
engine_output_type: latent
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 2048
- stage_id: 1
stage_type: llm
runtime:
devices: "0"
engine_args:
model_stage: decoder
model_arch: YourTTSModelForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
engine_output_type: audio
engine_input_source: [0]
final_output: true
final_output_type: audio
```
### Streaming mode (async_chunk)
Add `async_chunk: true` at the top level and specify `custom_process_next_stage_input_func`
on Stage 0 to define how intermediate outputs are chunked and forwarded:
```yaml
# stage_configs/your_model_name_async_chunk.yaml
async_chunk: true
stage_args:
- stage_id: 0
stage_type: llm
runtime:
devices: "0"
engine_args:
model_stage: ar_stage
max_num_seqs: 64
model_arch: YourTTSModelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
engine_output_type: latent
custom_process_next_stage_input_func: >
vllm_omni.model_executor.stage_input_processors.your_model_name.ar2decoder_async_chunk
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 2048
- stage_id: 1
stage_type: llm
runtime:
devices: "0"
engine_args:
model_stage: decoder
model_arch: YourTTSModelForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
engine_output_type: audio
engine_input_source: [0]
final_output: true
final_output_type: audio
```
## Stage Input Processors
Stage input processors convert Stage 0 outputs into Stage 1 inputs. Create yours in
`vllm_omni/model_executor/stage_input_processors/your_model_name.py`.
See `stage_input_processors/qwen3_tts.py` for the full reference implementation.
### Data structures
Understanding what's available in stage outputs:
- `stage_list[source_id].engine_outputs` - list of `EngineCoreOutput` objects
- Each `EngineCoreOutput` has `outputs` - list of `RequestOutput` objects
- Each `RequestOutput` has:
- `token_ids` - generated token IDs
- `multimodal_output` - dict with keys matching your model's `OmniOutput.multimodal_outputs`
- `prompt_token_ids` - original prompt token IDs
### Batch mode (non-streaming)
Collects all Stage 0 outputs and forwards them to Stage 1 in one shot:
```python
def ar2decoder(
stage_list: list[Any],
engine_input_source: list[int],
prompt: OmniTokensPrompt | TextPrompt | None = None,
requires_multimodal_data: bool = False,
) -> list[OmniTokensPrompt]:
source_id = engine_input_source[0]
decoder_inputs = []
for output in stage_list[source_id].engine_outputs:
result = output.outputs[0]
codes = result.multimodal_output["audio_codes"].cpu()
decoder_inputs.append(
OmniTokensPrompt(prompt_token_ids=codes.reshape(-1).tolist())
)
return decoder_inputs
```
### Streaming mode (async_chunk)
Buffers Stage 0 outputs and forwards a chunk to Stage 1 once `chunk_size` frames
have accumulated. The function signature follows the `OmniChunkTransferAdapter` protocol:
```python
def ar2decoder_async_chunk(
transfer_manager: Any,
pooling_output: dict[str, Any] | None,
request: Any,
is_finished: bool = False,
) -> dict[str, Any] | None:
"""Forward chunks of AR output to the decoder stage."""
request_id = request.external_req_id
finished = bool(is_finished or request.is_finished())
# Extract and buffer the latest frame
if isinstance(pooling_output, dict):
frame = extract_frame(pooling_output)
if frame is not None:
transfer_manager.code_prompt_token_ids[request_id].append(
frame.cpu().tolist()
)
elif not finished:
return None
# Read chunk config from connector
chunk_size = 25
left_context_size = 25
length = len(transfer_manager.code_prompt_token_ids[request_id])
if length <= 0:
if finished:
return {"codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
return None
# Wait until a full chunk is ready (or request is finished)
chunk_length = length % chunk_size
if chunk_length != 0 and not finished:
return None
# Build context window: left_context + chunk
context_length = chunk_length if chunk_length != 0 else chunk_size
end_index = min(length, left_context_size + context_length)
window = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
return {
"codes": torch.tensor(window).transpose(0, 1).reshape(-1).tolist(),
"left_context_size": max(0, int(end_index - context_length)),
"finished": torch.tensor(finished, dtype=torch.bool),
}
```
Key points:
- `transfer_manager` is the `OmniChunkTransferAdapter` that owns the chunk lifecycle
- Each call appends one AR decode step's output; a chunk is emitted every `chunk_size` steps
- The final (possibly partial) chunk is flushed when `is_finished` is true
- `left_context_size` frames of overlap are included for smooth audio boundaries
## Testing
For general testing conventions, see [tests_style.md](../ci/tests_style.md).
Recommended test cases for a new TTS model:
1. **Single request** - verify waveform output shape and sample rate
2. **Batched requests** - verify each request in the batch finishes independently
3. **async_chunk streaming** - verify audio chunks arrive incrementally and decode correctly
4. **Speaker conditioning** (if applicable) - verify different speaker inputs produce different outputs
Reference test: `tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py`
## Adding a Model Recipe
After implementing and testing your model, add a model recipe to the
[vllm-project/recipes](https://github.com/vllm-project/recipes) repository so users can
get started quickly. See [Adding an Omni-Modality Model](./adding_omni_model.md#adding-a-model-recipe)
for the expected format.
## Summary
Adding a TTS model to vLLM-Omni involves:
1. **Create model directory** with AR stage, decoder stage, and unified class
2. **AR stage** - use vLLM's native decoder layers with fused QKV; do not wrap HF directly
3. **Decoder stage** - thin wrapper around your audio decoder; implement `chunked_decode_streaming()`
4. **Unified class** - dispatches on `model_stage`; same structure as `Qwen3TTSModelForGeneration`
5. **Register** all stage classes in `registry.py`
6. **YAML configs** - provide both batch and `async_chunk` variants
7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25
8. **Tests** - cover single request, batching, and async_chunk streaming
9. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
### Qwen3-TTS Reference Files
| File | Purpose |
|------|---------|
| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | AR stage with vLLM fused ops |
| `models/qwen3_tts/qwen3_tts_code2wav.py` | Decoder stage with `chunked_decode_streaming()` |
| `stage_configs/qwen3_tts.yaml` | Stage configuration |
| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
For more information, see:
- [Architecture Overview](../../design/architecture_overview.md)
- [Async Chunk Design](../../design/feature/async_chunk_design.md)
- [Stage Configuration Guide](../../configuration/stage_configs.md)
+31 -101
View File
@@ -4,57 +4,49 @@
vLLM-Omni uses the PyTorch Profiler to analyze performance across both **multi-stage omni-modality models** and **diffusion models**.
### 1. Configure Profiling in the Stage YAML
### 1. Set the Output Directory
Before running any script, set this environment variable. The system detects this and automatically saves traces here.
Enable profiling by adding `profiler_config` under `engine_args` for the stage(s) you want to profile in your stage config YAML:
```yaml
stage_args:
- stage_id: 0
stage_type: llm
engine_args:
# ... other engine args ...
profiler_config:
profiler: torch
torch_profiler_dir: ./perf
```bash
export VLLM_TORCH_PROFILER_DIR=./profiles
```
| Field | Description |
|---|---|
| `profiler` | Profiler backend to use. Currently supports `torch`. |
| `torch_profiler_dir` | Directory where trace files are saved. Created automatically if it doesn't exist. |
> **Tip:** Only enable `profiler_config` on stages you actually need to profile. Stages without it will not start a profiler, keeping overhead minimal.
### 2. Profiling Omni-Modality Models
It is best to limit profiling to one iteration to keep trace files manageable.
```bash
export VLLM_PROFILER_MAX_ITERS=1
```
**Selective Stage Profiling**
It is highly recommended to profile specific stages to prevent producing overly large trace files:
The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
```python
# Profile all stages
omni_llm.start_profile()
# Only profile Stage 1
omni_llm.start_profile(stages=[1])
```
```python
# Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
omni_llm.start_profile(stages=[0, 2])
```
> **Important:** Always pass the same `stages` list to both `start_profile()` and `stop_profile()`. If you omit `stages` from `stop_profile()`, it defaults to stopping all stages — including ones that were never started — which will produce errors.
**Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
```python
profiler_stages = [0] # Only profile the stages you need
from vllm_omni import omni_llm
# 1. Start profiling
omni.start_profile(stages=profiler_stages)
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
# 1. Start profiling if enabled
if profiler_enabled:
omni_llm.start_profile(stages=[0])
# Initialize generator
omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
total_requests = len(prompts)
processed_count = 0
@@ -72,8 +64,7 @@ for stage_outputs in omni_generator:
print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
# Stop the profiler while workers are still active
# Pass the same stages list used in start_profile()
omni_llm.stop_profile(stages=profiler_stages)
omni_llm.stop_profile()
# Wait for traces to flush to disk
print("[Info] Waiting 30s for workers to write trace files to disk...")
@@ -84,38 +75,24 @@ omni_llm.close()
```
**CLI Usage** (using `end2end.py`):
```bash
# Profile only Stage 0 (Thinker)
python end2end.py --output-wav output_audio \
--query-type text --enable-profiler --profiler-stages 0
# Profile Stage 0 and Stage 2
python end2end.py --output-wav output_audio \
--query-type text --enable-profiler --profiler-stages 0 2
# Profile all stages (omit --profiler-stages)
python end2end.py --output-wav output_audio \
--query-type text --enable-profiler
```
**Examples**:
1. **Qwen2.5-Omni**: [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
2. **Qwen3-Omni**: [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
### 3. Profiling diffusion models
Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding. Standalone diffusion scripts use `--profiler-dir` to enable profiling.
Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding.
**CLI Usage:**
```bash
```python
python image_to_video.py \
--model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
--image qwen-bear.png \
--prompt "A cat playing with yarn, smooth motion" \
--profiler-dir \
\
# Minimize Spatial Dimensions (Optional but helpful):
# Drastically reduces memory usage so the profiler doesn't
@@ -145,9 +122,8 @@ python image_to_video.py \
--flow-shift 12.0 \
--fps 16 \
--output i2v_output.mp4
```
> **Note:** For diffusion stages within a multi-stage omni pipeline, use `profiler_config` in the stage YAML instead (see Section 1).
```
**Examples**:
@@ -155,62 +131,16 @@ python image_to_video.py \
2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**: [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
### 4. Profiling Online Serving
### 4. Analyzing Omni Traces
When `profiler_config` is set in the stage YAML, the server automatically exposes `/start_profile` and `/stop_profile` HTTP endpoints.
**1. Start the server** with a stage YAML that has `profiler_config` enabled:
```bash
vllm serve Qwen/Qwen2.5-Omni-7B \
--omni \
--stage-configs-path qwen2_5_omni.yaml \
--port 8091
```
Or for one stage diffusion models:
```bash
vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --port 8091 --profiler-config '{"profiler": "torch", "torch_profiler_dir": "./vllm_profile"}'
```
**2. Start profiling** by sending a POST request:
```bash
# Profile all stages that have profiler_config set
curl -X POST http://localhost:8091/start_profile
# Profile specific stages only
curl -X POST http://localhost:8091/start_profile \
-H "Content-Type: application/json" \
-d '{"stages": [0]}'
```
**3. Send your inference requests** as normal while the profiler is running.
**4. Stop profiling** and collect traces:
```bash
# Stop all stages
curl -X POST http://localhost:8091/stop_profile
# Stop specific stages (must match the stages you started)
curl -X POST http://localhost:8091/stop_profile \
-H "Content-Type: application/json" \
-d '{"stages": [0]}'
```
Trace files are written to the `torch_profiler_dir` specified in your stage YAML.
> **Important:** Always stop the same stages you started. Stopping a stage that was never started will produce errors.
### 5. Analyzing Traces
Output files are saved to the `torch_profiler_dir` specified in your stage YAML config.
Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
**Output**
**Chrome Trace** (`.json.gz`): Visual timeline of kernels and stages. Open in Perfetto UI.
**Chrome Trace** (```.json.gz```): Visual timeline of kernels and stages. Open in Perfetto UI.
**Viewing Tools:**
- [Perfetto](https://ui.perfetto.dev/) (recommended)
- `chrome://tracing` (Chrome only)
- [Perfetto](https://ui.perfetto.dev/)(recommended)
- ```chrome://tracing```(Chrome only)
**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation: [vLLM Profiling Guide](https://docs.vllm.ai/en/stable/contributing/profiling/)
+5 -12
View File
@@ -67,12 +67,12 @@ According to analysis for current popular open-source models, most of them have
| Component | Description |
| ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
| **OmniRouter** | provide an intelligent router for Omni-modality requests dispatch |
| **EntryPoints** | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni), while `AsyncOmniEngine` and `Orchestrator` coordinate multi-stage AR/DiT execution |
| **EntryPoints** | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni) and provide the OmniStage abstraction for different AR/DiT stages |
| **AR** | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management |
| **Diffusion** | natively implemented and optimized using acceleration components |
| **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages |
Disaggregated stages are managed through stage configuration. In Qwen3-Omni, Thinker/Talker/Code2wav are declared as separate configured stages, and runtime routing is handled by `Orchestrator` over `StageEngineCoreClient` / `StageDiffusionClient`.
Disaggregated stages are managed through configuration, such as in the Qwen3-Omni example, where stages like Thinker, Talker, and Code2wav are defined as separate OmniStage instances with specific resources and input/output type.
## Main features
@@ -92,13 +92,6 @@ The framework achieves high performance through several optimization techniques:
* **Quantization:** Supports various quantization implementations including FP8 and AWQ.
* **FusedOps:** Allows for custom and third-party integration.
### Classifier-Free Guidance (CFG) Companion Flow
vLLM-Omni natively models Classifier-Free Guidance (CFG) across disaggregated multi-stage setups via a "companion request" paradigm, eliminating redundant textual/multimodal context computation boundaries:
1. **Prompt Expansion:** In the initial autoregressive (AR) stage, a customized `prompt_expand_func` hook intercepts incoming generation prompts and pairs them directly with negative companion prompts (e.g., a default negative prompt) on the fly, tagging the secondary prompt with a specific internal role (`cfg_text`).
2. **Synchronized KV Cache Transfer:** The AR stage evaluates both the primary and companion sequence batches concurrently. The `OmniConnector` captures these specific structural dependencies and reliably passes the positive and negative outcome KV caches seamlessly across stage boundaries via shared memory or network protocols.
3. **KV Cache Collection & Injection:** Upon reaching the downstream Diffusion (DiT) Engine, an assigned `cfg_kv_collect_func` automatically intercepts the mapped companion caches (`cfg_text_past_key_values`). These auxiliary dependencies are natively gathered and seamlessly bound to the primary generation sequence variables, enabling the DiT Engine to cleanly implement cross-attention CFG guidance over accurate conditioning and unconditioning structures in parallel.
### Flexibility and Usability
vLLM-Omni is designed to be flexible and straightforward for users:
@@ -127,10 +120,10 @@ Taking **Qwen3-Omni** as an example:
The **Omni** class provides a Python interface for offline batched inference. Users initialize the Omni class with a Hugging Face model name and use the generate method, passing inputs that include both text prompts and multi-modal data:
```
# Create an omni runtime with HF model name.
# Create an omni_lm with HF model name.
from vllm_omni.entrypoints.omni import Omni
omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
omni_lm = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
# Example prompts.
om_inputs = {"prompt": prompt,
@@ -140,7 +133,7 @@ om_inputs = {"prompt": prompt,
}}
# Generate texts and audio from the multi-modality inputs.
outputs = omni.generate(om_inputs, sampling_params_list)
outputs = omni_lm.generate(om_inputs, sampling_params_list)
```
## Online Serving
+3 -10
View File
@@ -19,7 +19,7 @@ The `async_chunk` feature enables asynchronous, chunked processing of data acros
For qwen3-omni:
- **Thinker → Talker**: Per decode step (typically chunk_size=1)
- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. During the initial phase, a dynamic initial chunk size (IC) is automatically selected based on server load to reduce TTFA. Use the per-request `initial_codec_chunk_frames` API field to override.
- **Talker → Code2Wav**: Accumulated to code2wav chunk_size(default=25, currently only support default, will support chunk_size soon) before sending
- **Code2Wav**: Streaming decode with code2wav chunk_size
With `async_chunk`:
@@ -67,13 +67,6 @@ Enabling **async_chunk** (False→True) sharply reduces time-to-first-audio (TTF
</picture>
</p>
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png">
<img alt="E2E Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png" width=100%>
</picture>
</p>
## Architecture
### Data Flow
@@ -171,7 +164,7 @@ stage_args:
- `custom_process_next_stage_input_func: str`: Path to custom chunk processing function; receives `(transfer_manager, pooling_output, request)`. For qwen3-omni: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
- `stage_connector_config: dict`: Connector configuration
- `worker_type: str`: Model type, e.g. `"ar"` or `"generation"` (used by OmniChunkTransferAdapter for mode-specific payload handling)
- `max_num_seqs: int`: Maximum number of sequences for concurrent processing in the stage
- `max_batch_size: int`: Maximum batch size for the stage
### Connector Configuration
@@ -195,9 +188,9 @@ stage_args:
- stage_id: 2 # code2wav stage
runtime:
devices: "1"
max_batch_size: 64 # Enables batched audio generation
engine_args:
model_stage: code2wav
max_num_seqs: 64 # Enables batched audio generation
```
## Related Files
+5 -5
View File
@@ -113,13 +113,13 @@ def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], No
Some models require custom handling:
- **Single or dual-transformer:** Models that may use one or two transformers (e.g., Wan2.2)
- **Dual-transformer:** Models with separate high-noise and low-noise transformers (e.g., Wan2.2)
- **Multi-block-list:** Models with multiple block lists in one transformer (e.g., LongCatImage with `transformer_blocks` + `single_transformer_blocks`)
- **Special forward patterns:** Models with non-standard block execution patterns
### Example 1: Single or Dual-Transformer Model (Wan2.2)
### Example 1: Dual-Transformer Model (Wan2.2)
Wan2.2 can use either a single transformer or two transformers (one for high-noise steps and one for low-noise steps). The implementation automatically detects the mode based on the presence of `transformer_2`.
Wan2.2 uses two transformers: one for high-noise steps and one for low-noise steps.
**Key difference:** Use `BlockAdapter` to wrap multiple transformers with separate configurations.
@@ -256,7 +256,7 @@ cache_config={
}
```
Check the [user guide for cache_dit](../../user_guide/diffusion/cache_acceleration/cache_dit.md) for more adjustable parameters.
Check the [user guide for cache_dit](../../user_guide/diffusion/cache_dit_acceleration.md) for more adjustable parameters.
---
@@ -267,7 +267,7 @@ Complete examples in the codebase:
| Model | Path | Pattern | Notes |
|-------|------|---------|-------|
| **Standard DiT** | [`cache_dit_backend.py::enable_cache_for_dit`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Default enabler | Single transformer, automatic |
| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Single or dual-transformer | Auto-detects mode based on transformer_2 presence |
| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Dual-transformer | Separate high/low noise transformers |
| **LongCat** | [`cache_dit_backend.py::enable_cache_for_longcat_image`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_longcat_image) | Multi-block-list | Two block lists in one transformer |
| **BAGEL** | [`cache_dit_backend.py::enable_cache_for_bagel`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_bagel) | Omni model | Complex architecture |
+10 -67
View File
@@ -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 | All ranks step locally (no broadcast needed) |
| [`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 |
| [`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()`
- All ranks compute CFG combine locally (deterministic, identical results)
- Combined on rank 0 using CFG formula
- **Sequential mode** (when `cfg_world_size == 1`):
- Single rank computes both positive and negative predictions
@@ -55,7 +55,14 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
`scheduler_step_maybe_with_cfg()` ensures consistent latent states across all ranks:
- 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.
- **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
---
@@ -170,70 +177,6 @@ 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
Loaded 100 of 825 files, more files were not shown because too many files have changed in this diff. Show more