Compare commits

..
Author SHA1 Message Date
khluuandClaude Opus 4.6 3e280145cb Show all S3 wheels vs filtered wheels for verification
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-05-17 03:18:49 -07:00
khluuandClaude Opus 4.6 f5647e7e42 Allow RC wheels (e.g. v0.21.0rc1) in PyPI publish step
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-05-17 03:17:42 -07:00
khluuandClaude Opus 4.6 f825e007f5 Comment out twine upload; dry-run to verify wheel list first
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-05-17 03:12:12 -07:00
khluuandClaude Opus 4.6 c5f1dc4f44 Add PyPI publish step to release pipeline
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-05-17 03:07:29 -07:00
khluuandClaude Opus 4.7 65ae600835 Drop nightly Docker image flow; image builds are release-only
Wheels still build on every commit (existing nightly wheel flow), but
Docker images are now version-tagged releases only. Remove:
  - the NIGHTLY conditional on the build/publish block steps
  - the nightly DockerHub publish step
  - push-nightly-builds.sh and cleanup-nightly-builds.sh

The image flow now: manual unblock -> build per-arch images -> manual
unblock with release-version -> publish-release-images.sh tags
vllm/vllm-omni:vX.Y.Z + latest on DockerHub.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-07 18:03:55 -07:00
khluuandClaude Opus 4.7 9ce9e7dd48 Drop ECR multi-arch manifest step; only publish manifest to DockerHub
The DockerHub manifest is the only one consumers actually pull. The ECR
multi-arch manifest was unused, so remove the step and have the
nightly/release publish steps depend on the per-arch build steps
directly.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-07 17:43:27 -07:00
khluuandClaude Opus 4.7 4f12f0c733 Refactor: unify release pipeline with NIGHTLY=1 option (mirror vllm)
- Rename existing CI .buildkite/release-pipeline.yaml -> rebase-pipeline.yaml
  (it builds the CI image and uploads premerge child pipelines, not a release)
- Promote .buildkite/nightly-release-pipeline.yaml -> release-pipeline.yaml,
  refactored to mirror vllm/vllm's release-pipeline.yaml structure:
    - release-version input field
    - block step gating release Docker image builds (NIGHTLY != "1")
    - per-arch x86_64 / aarch64 docker builds pushing to ECR staging
    - multi-arch manifest in ECR
    - nightly DockerHub publish (NIGHTLY == "1") via push-nightly-builds.sh
    - block-gated release DockerHub publish (NIGHTLY != "1") via
      publish-release-images.sh, tagging vX.Y.Z + latest
- Add scripts adapted from vllm:
    push-nightly-builds.sh, cleanup-nightly-builds.sh,
    publish-release-images.sh — all targeting Docker Hub repo vllm/vllm-omni

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-07 17:19:55 -07:00
khluuandClaude Opus 4.7 afbb8af51a Tag and push DockerHub images directly from build steps
Tag each per-arch image with both the ECR commit-pinned tag and the
DockerHub tags (vllm/vllm-omni:$COMMIT-$arch and :nightly-$arch) and
push them to DockerHub from the build step itself. The publish step
now just creates and pushes the multi-arch manifests inline, so the
push-nightly-builds.sh helper is no longer needed.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-07 17:08:27 -07:00
khluuandClaude Opus 4.7 49afe9d31e [CI/Build] Add x86_64 and aarch64 docker image builds to release pipeline
Add release Docker image build steps to the nightly release pipeline
that build per-arch images from docker/Dockerfile.cuda, push them to
ECR, create a multi-arch manifest, and publish it to Docker Hub under
the nightly tag at vllm/vllm-omni.

Mirrors the pattern used by vllm-project/vllm's release pipeline.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-07 16:14:14 -07:00
838 changed files with 18866 additions and 92153 deletions

No files matched your search

+24 -105
View File
@@ -56,88 +56,6 @@ check_run_all_label() {
fi
}
is_docs_only_change() {
local file_path
local has_any=0
while IFS= read -r file_path; do
[[ -z "${file_path}" ]] && continue
has_any=1
if [[ "${file_path}" == docs/* ]]; then
continue
fi
if [[ "${file_path}" == *.md ]]; then
continue
fi
if [[ "${file_path}" == "mkdocs.yml" ]]; then
continue
fi
return 1
done
[[ "${has_any}" -eq 1 ]]
}
resolve_skip_ci() {
local is_pr_build=0
local files
local base_branch base_ref
if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" && -n "${BUILDKITE_PULL_REQUEST:-}" ]]; then
is_pr_build=1
fi
if [[ "${is_pr_build}" -eq 1 ]]; then
base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
if ! git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
echo "resolve_skip_ci: origin/${base_branch} not found locally; trying fetch" >&2
git fetch --depth=200 origin "${base_branch}" >/dev/null 2>&1 || true
fi
base_ref=""
if git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
base_ref="origin/${base_branch}"
elif git rev-parse --verify "${base_branch}" >/dev/null 2>&1; then
base_ref="${base_branch}"
else
echo "resolve_skip_ci: cannot resolve PR base ${base_branch}; skip-ci=0" >&2
echo -n 0
return 0
fi
if ! files="$(git diff --name-only "${base_ref}...${BUILDKITE_COMMIT}" 2>/dev/null)"; then
echo "resolve_skip_ci: failed to compute PR changed files; skip-ci=0" >&2
echo -n 0
return 0
fi
elif [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then
if ! git rev-parse --verify "${BUILDKITE_COMMIT}^" >/dev/null 2>&1; then
echo "resolve_skip_ci: commit has no parent on main; skip-ci=0" >&2
echo -n 0
return 0
fi
if ! files="$(git diff --name-only "${BUILDKITE_COMMIT}^..${BUILDKITE_COMMIT}" 2>/dev/null)"; then
echo "resolve_skip_ci: failed to compute main changed files; skip-ci=0" >&2
echo -n 0
return 0
fi
else
echo "resolve_skip_ci: not PR/main build; skip-ci=0" >&2
echo -n 0
return 0
fi
if is_docs_only_change <<< "${files}"; then
echo "resolve_skip_ci: docs-only change detected; skip-ci=1" >&2
echo -n 1
return 0
fi
echo "resolve_skip_ci: non-doc changes detected; skip-ci=0" >&2
echo -n 0
}
if [[ -z "${COV_ENABLED:-}" ]]; then
COV_ENABLED=0
fi
@@ -173,24 +91,8 @@ upload_pipeline() {
cd .buildkite
# Select test definition file: merge suite for main, ready suite for PRs.
# For debugging, DEBUG_TEST_YAML can override the selection — accepts
# "merge" or "ready" (case-insensitive).
if [[ -n "${DEBUG_TEST_YAML:-}" ]]; then
case "${DEBUG_TEST_YAML,,}" in
merge)
TEST_YAML="test-amd-merge.yml"
;;
ready)
TEST_YAML="test-amd-ready.yaml"
;;
*)
echo "ERROR: DEBUG_TEST_YAML must be 'merge' or 'ready', got '$DEBUG_TEST_YAML'" >&2
exit 1
;;
esac
echo "DEBUG_TEST_YAML override: using $TEST_YAML"
elif [[ $BUILDKITE_BRANCH == "main" ]]; then
# 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"
@@ -238,13 +140,30 @@ fi
# Early exit start: skip pipeline if conditions are met
# ----------------------------------------------------------------------
# Match CUDA Buildkite skip-ci behavior for docs/markdown-only changes.
# skip pipeline if all changed files are under docs/
if [[ "${DOCS_ONLY_DISABLE}" != "1" ]]; then
SKIP_CI="$(resolve_skip_ci)"
if [[ "${SKIP_CI}" == "1" ]]; then
echo "[docs-only] Docs/markdown-only changes detected. Exiting before pipeline upload."
exit 0
if [[ -n "${file_diff:-}" ]]; then
docs_only=1
# Robust iteration over newline-separated file_diff
while IFS= read -r f; do
[[ -z "$f" ]] && continue
# **Policy:** only skip if *every* path starts with docs/
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 changes detected
\`\`\`
${file_diff}
\`\`\`" --style "info" || true
echo "[docs-only] All changes are under docs/. Exiting before pipeline upload."
exit 0
fi
fi
fi
# ----------------------------------------------------------------------
+1 -1
View File
@@ -10,7 +10,7 @@ steps:
DOCKER_BUILDKIT: "1"
# Buildkite will automatically replace this with the actual commit hash
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
VLLM_VERSION: "v0.21.0"
VLLM_VERSION: "v0.20.0"
priority: 100
timeout_in_minutes: 60
soft_fail: false
-110
View File
@@ -1,110 +0,0 @@
# steps:
# - label: ":buildkit: Build and Push NPU Test Image - A3"
# key: image-build-a3
# agents:
# queue: "ascend-a3"
# resource_class: "npu-2"
# plugins:
# - kubernetes:
# metadata:
# annotations:
# vault.hashicorp.com/agent-init-first: "true"
# vault.hashicorp.com/agent-inject: "true"
# vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
# vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
# vault.hashicorp.com/agent-inject-perms-config.json: "0400"
# vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
# vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
# vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
# vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
# vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
# vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
# vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
# vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
# vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
# vault.hashicorp.com/agent-pre-populate-only: "true"
# vault.hashicorp.com/agent-run-as-group: "1000"
# vault.hashicorp.com/agent-run-as-user: "1000"
# vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
# vault.hashicorp.com/role: ascend-gha-runners
# vault.hashicorp.com/secret-volume-path: /home/user/.docker/
# vault.hashicorp.com/tls-skip-verify: "true"
# podSpecPatch:
# volumes:
# - name: token-vol
# projected:
# defaultMode: 420
# sources:
# - serviceAccountToken:
# audience: api
# expirationSeconds: 600
# path: token
# env:
# VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
# IMAGE_NAME: "vllm-omni-ci-npu"
# IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
# BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
# command: |
# set -ex
# echo "--- Building and pushing NPU Test Image"
# echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
# echo "buildkitd address: $${BUILDKITD_ADDR}"
# git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf "https://github.com/"
# if ! command -v buildctl &> /dev/null; then
# echo "Installing buildctl..."
# mkdir -p /tmp/buildkit
# BUILDKIT_VERSION="v0.29.0"
# wget -q "https://gh-proxy.test.osinfra.cn/https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
# tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
# cp /tmp/buildkit/bin/buildctl /usr/local/bin/
# fi
# export DOCKER_CONFIG=/home/user/.docker
# buildctl \
# --addr="$${BUILDKITD_ADDR}" \
# --tlscacert=/home/user/.docker/ca.pem \
# --tlscert=/home/user/.docker/cert.pem \
# --tlskey=/home/user/.docker/key.pem \
# build \
# --frontend dockerfile.v0 \
# --local context=. \
# --local dockerfile=./docker \
# --opt filename=Dockerfile.npu.ci.a3 \
# --secret id=dockerconfig,src=/home/user/.docker/config.json \
# --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
# --progress=plain
# echo "--- Image pushed successfully"
# echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
# # - label: "Upload Ready Pipeline - A3"
# # depends_on: image-build-a3
# # key: upload-ready-pipeline-a3
# # if: build.branch != "main" && build.pull_request.labels includes "npu-test"
# # agents:
# # queue: "ascend-a3"
# # resource_class: "npu-2"
# # command: |
# # buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml
# # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
# - 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" ||
# build.pull_request.labels includes "omni-test" ||
# build.pull_request.labels includes "tts-test" ||
# build.pull_request.labels includes "diffusion-x2iat-test" ||
# build.pull_request.labels includes "diffusion-x2v-test"
# ))
# commands:
# - buildkite-agent pipeline upload .buildkite/test-npu-nightly.yaml
# agents:
# queue: "ascend-a3"
# resource_class: "npu-2"
-112
View File
@@ -1,112 +0,0 @@
steps:
- label: ":buildkit: Build and Push NPU Test Image - A2"
key: image-build
agents:
queue: "ascend-a2b3"
resource_class: "npu-1"
plugins:
- kubernetes:
metadata:
annotations:
vault.hashicorp.com/agent-init-first: "true"
vault.hashicorp.com/agent-inject: "true"
vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
vault.hashicorp.com/agent-inject-perms-config.json: "0400"
vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
vault.hashicorp.com/agent-pre-populate-only: "true"
vault.hashicorp.com/agent-run-as-group: "1000"
vault.hashicorp.com/agent-run-as-user: "1000"
vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
vault.hashicorp.com/role: ascend-gha-runners
vault.hashicorp.com/secret-volume-path: /home/user/.docker/
vault.hashicorp.com/tls-skip-verify: "true"
podSpecPatch:
volumes:
- name: token-vol
projected:
defaultMode: 420
sources:
- serviceAccountToken:
audience: api
expirationSeconds: 600
path: token
env:
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
IMAGE_NAME: "vllm-omni-ci-npu"
IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
command: |
set -ex
echo "--- Building and pushing NPU Test Image"
echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
echo "buildkitd address: $${BUILDKITD_ADDR}"
if ! command -v buildctl &> /dev/null; then
echo "Installing buildctl..."
mkdir -p /tmp/buildkit
BUILDKIT_VERSION="v0.29.0"
wget -q "https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
cp /tmp/buildkit/bin/buildctl /usr/local/bin/
fi
export DOCKER_CONFIG=/home/user/.docker
buildctl \
--addr="$${BUILDKITD_ADDR}" \
--tlscacert=/home/user/.docker/ca.pem \
--tlscert=/home/user/.docker/cert.pem \
--tlskey=/home/user/.docker/key.pem \
build \
--frontend dockerfile.v0 \
--local context=. \
--local dockerfile=./docker \
--opt filename=Dockerfile.npu.ci \
--secret id=dockerconfig,src=/home/user/.docker/config.json \
--output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
--progress=plain
echo "--- Image pushed successfully"
echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
- label: "Upload Ready Pipeline"
depends_on: image-build
key: upload-ready-pipeline
if: >-
build.branch != "main" && (
build.pull_request.labels includes "npu-test" ||
build.pull_request.labels includes "ready"
)
agents:
queue: "ascend-a2b3"
resource_class: "npu-1"
command: |
buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml
# L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
- 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" ||
build.pull_request.labels includes "omni-test" ||
build.pull_request.labels includes "tts-test" ||
build.pull_request.labels includes "diffusion-x2iat-test" ||
build.pull_request.labels includes "diffusion-x2v-test"
))
commands:
- buildkite-agent pipeline upload .buildkite/test-npu-nightly.yaml
agents:
queue: "ascend-a2b3"
resource_class: "npu-1"
@@ -52,7 +52,5 @@ time timeout -k 30 30m docker run \
pip install tblib==3.1.0
cd /workspace/vllm-omni
pytest -v -s -m "core_model and xpu and B60"
pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py
export VLLM_XPU_USE_SAMPLER_KERNEL=0 # NOTE: Remove this after vLLM v0.21.1 is merged. Fixes Qwen2-5 omni-expansion tests.
pytest -v -s -m "advanced_model and xpu and B60"
'
@@ -14,30 +14,30 @@ if [[ -z "$RELEASE_VERSION" ]]; then
fi
echo "Release version: $RELEASE_VERSION"
if [[ -z "$PYPI_TOKEN" ]]; then
echo "[FATAL] PYPI_TOKEN is not set."
exit 1
else
export TWINE_USERNAME="__token__"
export TWINE_PASSWORD="$PYPI_TOKEN"
fi
# if [[ -z "$PYPI_TOKEN" ]]; then
# echo "[FATAL] PYPI_TOKEN is not set."
# exit 1
# else
# export TWINE_USERNAME="__token__"
# export TWINE_PASSWORD="$PYPI_TOKEN"
# fi
set -x
if ! command -v uv &> /dev/null; then
curl -LsSf https://astral.sh/uv/install.sh | sh
export PATH="$HOME/.local/bin:$PATH"
fi
uv venv --python 3.12 /tmp/vllm-omni-release-env
source /tmp/vllm-omni-release-env/bin/activate
uv pip install twine
ALL_DIR=/tmp/vllm-omni-release-all
DIST_DIR=/tmp/vllm-omni-release-dist
mkdir -p "$DIST_DIR"
mkdir -p "$ALL_DIR" "$DIST_DIR"
echo "Downloading wheels from S3:"
aws s3 ls "$S3_COMMIT_PREFIX"
echo "========================================"
echo "All wheels in S3 (no filters):"
echo "========================================"
aws s3 cp --recursive "$S3_COMMIT_PREFIX" "$ALL_DIR"
ls -la "$ALL_DIR"
echo ""
echo "========================================"
echo "Filtered wheels (include: vllm_omni-${RELEASE_VERSION}*.whl, exclude: *dev*):"
echo "========================================"
aws s3 cp --recursive --exclude "*" --include "vllm_omni-${RELEASE_VERSION}*.whl" --exclude "*dev*" "$S3_COMMIT_PREFIX" "$DIST_DIR"
ls -la "$DIST_DIR"
@@ -47,6 +47,12 @@ if [[ -z "$PYPI_WHEEL_FILES" ]]; then
exit 1
fi
python3 -m twine check $PYPI_WHEEL_FILES
python3 -m twine upload --non-interactive --verbose $PYPI_WHEEL_FILES
echo "Wheels uploaded to PyPI"
echo ""
echo "========================================"
echo "Wheels that would be uploaded to PyPI:"
echo "========================================"
echo "$PYPI_WHEEL_FILES"
# python3 -m twine check $PYPI_WHEEL_FILES
# python3 -m twine upload --non-interactive --verbose $PYPI_WHEEL_FILES
# echo "Wheels uploaded to PyPI"
@@ -1,18 +1,10 @@
#!/usr/bin/env bash
# Evaluate docs-only skip-ci and upload continuation steps from the given pipeline YAML
# Evaluate docs-only skip-ci and upload continuation steps from the same `.buildkite/pipeline.yml`
# (YAML document after the first `---`). Buildkite `if` is evaluated at upload time.
#
# Usage: upload_pipeline_with_skip_ci.sh [pipeline_yaml]
# pipeline_yaml: path relative to repo root or absolute (default: .buildkite/pipeline.yml)
set -euo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
PIPELINE_ARG="${1:-.buildkite/pipeline.yml}"
if [[ "${PIPELINE_ARG}" = /* ]]; then
PIPELINE_YML="${PIPELINE_ARG}"
else
PIPELINE_YML="${ROOT}/${PIPELINE_ARG}"
fi
PIPELINE_YML="${ROOT}/.buildkite/pipeline.yml"
# Prints a single digit to stdout: 1 = skip image CI, 0 = run. Logs go to stderr.
is_docs_only_change() {
@@ -29,7 +21,7 @@ is_docs_only_change() {
if [[ "${file_path}" == *.md ]]; then
continue
fi
if [[ "${file_path}" == "mkdocs.yml" ]]; then
if [[ "${file_path}" == "mkdocs.yaml" ]]; then
continue
fi
return 1
@@ -112,13 +104,12 @@ import pathlib
path = pathlib.Path(os.environ["PIPELINE_YML"])
text = path.read_text(encoding="utf-8")
sep = "\n---\n"
# Two supported layouts:
# - multi-doc: doc 1 = bootstrap (loaded by Buildkite), doc 2 (after `---`) = uploaded steps.
# - single-doc: whole file is the uploaded steps; caller (UI init step) is the bootstrap.
if sep in text:
_, continuation = text.split(sep, 1)
else:
continuation = text
if sep not in text:
raise SystemExit(
"upload_pipeline_with_skip_ci: .buildkite/pipeline.yml must contain a '\\n---\\n' separator "
"(document 1 = bootstrap, document 2 = uploaded steps)"
)
_, continuation = text.split(sep, 1)
skip = os.environ.get("SKIP_CI") == "1"
# When docs-only skip-ci: skip default CI image, but still build for L4 nightly (PR label nightly-test or
@@ -135,7 +126,7 @@ if skip:
else:
rep = "'true'"
ready_rep = "'build.branch != \"main\" && build.pull_request.labels includes \"ready\"'"
merge_rep = "'(build.branch == \"main\" && build.env(\"NIGHTLY\") != \"1\" && build.env(\"WEEKLY\") != \"1\") || (build.branch != \"main\" && build.pull_request.labels includes \"merge-test\")'"
merge_rep = "'(build.branch == \"main\" && build.env(\"NIGHTLY\") != \"1\") || (build.branch != \"main\" && build.pull_request.labels includes \"merge-test\")'"
rendered = (
continuation
.replace("__IMAGE_BUILD_IF__", rep)
+33 -29
View File
@@ -9,7 +9,16 @@ steps:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- "timeout 20m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --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
@@ -17,7 +26,7 @@ steps:
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
- timeout 20m pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
- label: "Diffusion Model CPU offloading Test"
agent_pool: mi325_1
@@ -30,9 +39,9 @@ steps:
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
@@ -67,7 +76,7 @@ steps:
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
- timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
# merge-only tests
@@ -77,15 +86,7 @@ steps:
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
- label: "Diffusion Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
- pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
- label: "Diffusion GPU Worker Test"
agent_pool: mi325_2
@@ -123,6 +124,7 @@ steps:
grade: Blocking
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 tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
- label: "Qwen3-TTS CustomVoice E2E Test"
@@ -153,6 +155,17 @@ steps:
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
'
- 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
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
# Need to debug before reneable numerical changes across large PRs
# # split Bagel Model Test with H100 (Real Weights) into three tests
@@ -163,10 +176,11 @@ steps:
# 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.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
# - 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 (2/3)"
# agent_pool: mi325_1
@@ -175,10 +189,11 @@ steps:
# 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.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
# - 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 (3/3)"
# agent_pool: mi325_1
@@ -187,23 +202,12 @@ steps:
# 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.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
- label: "Diffusion · Wan22 Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
"
# - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
- label: "Voxtral-TTS E2E Test"
agent_pool: mi325_1
+43 -30
View File
@@ -9,7 +9,7 @@ steps:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- "timeout 20m pytest -vvvv -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
- label: "Voxtral TTS CUDA Unit Test"
agent_pool: mi325_1
@@ -19,13 +19,21 @@ steps:
commands:
- timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
- label: "Diffusion Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
- label: "Diffusion Batching Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/diffusion/batching/test_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
- label: "Custom Pipeline Test"
agent_pool: mi325_1
@@ -46,9 +54,9 @@ steps:
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
@@ -65,6 +73,17 @@ steps:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
- label: "AudioX Online 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_audiox_online.py
- label: "Diffusion Cache Backend Test"
agent_pool: mi325_1
depends_on: amd-build
@@ -84,7 +103,7 @@ steps:
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
- 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
@@ -94,27 +113,6 @@ steps:
commands:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
- label: "Diffusion · Model Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
- label: "Diffusion · Wan22 Test"
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 40m bash -c "
export VLLM_TEST_CLEAN_GPU_MEMORY=1
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
"
- label: "Engine Test"
agent_pool: mi325_1
depends_on: amd-build
@@ -135,6 +133,7 @@ steps:
# 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 "core_model" --run-level "core_model"
@@ -176,6 +175,17 @@ steps:
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_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
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
# Need to debug before reneable numerical changes across large PRs
# - label: "Bagel Text2Img Model Test"
@@ -185,10 +195,11 @@ steps:
# 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.py -m "core_model" --run-level "core_model" -k "rocm"
# - 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
@@ -197,10 +208,11 @@ steps:
# 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.py -m "core_model" --run-level "core_model" -k "rocm"
# - 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
@@ -209,11 +221,12 @@ steps:
# 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.py -m "core_model" --run-level "core_model" -k "rocm"
# - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model" -k "rocm"
- label: "CosyVoice3-TTS E2E Test"
agent_pool: mi325_1
@@ -223,5 +236,5 @@ steps:
commands:
- |
timeout 20m bash -c '
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model"
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts.py -m "core_model" --run-level "core_model"
'
+22 -7
View File
@@ -8,7 +8,7 @@ steps:
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- pytest -s -v tests/e2e/offline_inference/test_z_image.py
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
- label: "Diffusion Images API LoRA E2E"
timeout_in_minutes: 30
@@ -20,7 +20,7 @@ steps:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
- pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
- label: "Diffusion Model CPU offloading Test"
timeout_in_minutes: 20
@@ -32,8 +32,8 @@ steps:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
- pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
- pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
- label: "Diffusion Cache Backend Test"
timeout_in_minutes: 15
@@ -57,7 +57,7 @@ steps:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
- label: "Diffusion Tensor Parallelism Test"
timeout_in_minutes: 20
@@ -67,7 +67,7 @@ steps:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
- pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
- label: "Diffusion GPU Worker Test"
timeout_in_minutes: 20
@@ -100,11 +100,25 @@ steps:
commands:
- 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
- label: "Diffusion Image Edit Test"
timeout_in_minutes: 15
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
- label: "Omni Sleep Mode Test"
timeout_in_minutes: 40
agent_pool: mi325_2
@@ -115,4 +129,5 @@ steps:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/entrypoints/test_omni_sleep_mode.py -m "advanced_model and omni and MI325" --run-level "advanced_model"
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/e2e/offline_inference/test_omni_sleep_mode.py -m "advanced_model and omni and MI325" --run-level "advanced_model"
+431 -752
View File
File diff suppressed because it is too large. Load diff
+154 -530
View File
@@ -1,6 +1,5 @@
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_USE_FLASHINFER_MOE_FP16: 0
HF_HUB_DOWNLOAD_TIMEOUT: 300
HF_HUB_ETAG_TIMEOUT: 60
@@ -14,7 +13,7 @@ steps:
- label: ":full_moon: Omni · Function Test with H100"
timeout_in_minutes: 90
commands:
- pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py
- pytest -s -v tests/e2e/ -m "full_model and H100 and omni" --run-level "full_model" --ignore=tests/e2e/accuracy
agents:
queue: "mithril-h100-pool"
plugins:
@@ -170,14 +169,14 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Omni · Perf Test · No Async Chunk"
key: nightly-omni-performance-no-async-chunk
- label: ":full_moon: Omni · Perf Test"
key: nightly-omni-performance
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
@@ -215,146 +214,11 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Omni · Perf Test · Async Chunk"
key: nightly-omni-performance-async-chunk
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
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 · Perf Test · vLLM Text"
key: nightly-omni-performance-vllm-text
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
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 · Perf Test · Multi-Replica"
key: nightly-omni-performance-multi-replicas
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
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: 3
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 · Multi-Replica Startup Test with 4x H100"
timeout_in_minutes: 45
commands:
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py -m "full_model" --run-level "core_model"
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py -m "core_model" --run-level "core_model"
agents:
queue: "mithril-h100-pool"
plugins:
@@ -398,7 +262,7 @@ steps:
timeout_in_minutes: 90
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
- pytest -s -v tests/e2e/ -m "full_model and L4 and omni" -k "not test_qwen2_5" --run-level "full_model" --ignore=tests/e2e/accuracy
agents:
queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
@@ -434,7 +298,7 @@ steps:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
@@ -468,65 +332,13 @@ steps:
build.pull_request.labels includes "nightly-test" ||
build.pull_request.labels includes "diffusion-x2iat-test"
steps:
# Use explicit file shards so single-card cases do not reserve 2+ GPUs,
# and Bagel feature coverage is not duplicated via a broad tests/e2e sweep.
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU"
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100"
timeout_in_minutes: 120
commands:
- >-
pytest -sv
tests/e2e/online_serving/test_qwen_image_expansion.py
tests/e2e/online_serving/test_qwen_image_edit_expansion.py
tests/e2e/online_serving/test_qwen_image_layered_expansion.py
tests/e2e/online_serving/test_longcat_image_expansion.py
tests/e2e/online_serving/test_longcat_image_edit_expansion.py
tests/e2e/online_serving/test_flux_2_dev_expansion.py
tests/e2e/online_serving/test_bagel_expansion.py
-m "full_model and diffusion and H100 and not distributed_cuda"
--run-level "full_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: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Qwen-Image"
timeout_in_minutes: 120
commands:
- >-
pytest -sv
tests/e2e/online_serving/test_qwen_image_expansion.py
tests/e2e/online_serving/test_qwen_image_edit_expansion.py
-m "full_model and diffusion and H100 and distributed_cuda"
--run-level "full_model"
- pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy
# revoke after bagel optimization finished
# Keep Bagel expansion and multi-replica tests in their dedicated H100 jobs below.
# - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_bagel_multi_replicas.py
agents:
queue: "mithril-h100-pool"
plugins:
@@ -561,60 +373,10 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Layered/LongCat/FLUX"
timeout_in_minutes: 90
commands:
- >-
pytest -sv
tests/e2e/online_serving/test_qwen_image_layered_expansion.py
tests/e2e/online_serving/test_longcat_image_expansion.py
tests/e2e/online_serving/test_longcat_image_edit_expansion.py
tests/e2e/online_serving/test_flux_2_dev_expansion.py
tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
-m "full_model and diffusion and H100 and distributed_cuda"
--run-level "full_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 X2I(&A&T) · Function Test with H100 · Multi-GPU Bagel"
- label: ":full_moon: Diffusion X2I(&A&T) · Bagel Function Test with H100"
timeout_in_minutes: 120
commands:
- >-
pytest -sv
tests/e2e/online_serving/test_bagel_expansion.py
-m "full_model and diffusion and H100 and distributed_cuda"
--run-level "full_model"
- pytest -sv tests/e2e/online_serving/test_bagel_expansion.py -m "full_model and diffusion and H100" --run-level "full_model"
agents:
queue: "mithril-h100-pool"
plugins:
@@ -649,6 +411,44 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
# revoke after bagel optimization finished
# - label: ":full_moon: Diffusion X2I(&A&T) · BAGEL Multi-Replica with H100"
# timeout_in_minutes: 120
# commands:
# - pytest -s -v tests/e2e/online_serving/test_bagel_multi_replicas.py -m "full_model and diffusion and H100" --run-level "full_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: 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
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
timeout_in_minutes: 60
@@ -671,6 +471,7 @@ steps:
- label: ":full_moon: Diffusion X2I(&A&T) · Doc Test"
timeout_in_minutes: 60
commands:
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/examples/*/test_text_to_image.py -m "full_model and example and H100" --run-level "full_model"
agents:
queue: "mithril-h100-pool"
@@ -706,6 +507,87 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion X2I(&A&T) · GEBench Accuracy Test"
timeout_in_minutes: 60
commands:
- pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level full_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: Diffusion X2I(&A&T) · GEdit-Bench Accuracy Test"
timeout_in_minutes: 60
commands:
- pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level full_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: Diffusion X2I(&A&T) · Accuracy Test"
timeout_in_minutes: 180
commands:
@@ -747,164 +629,28 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image"
key: nightly-diffusion-x2iat-performance-qwen-image
timeout_in_minutes: 180
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test"
key: nightly-diffusion-x2iat-performance
timeout_in_minutes: 360
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
# [HACK]: run upload in the same command block as pytest.
# Because `exit` aborts the entire commands list.
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
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
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit"
key: nightly-diffusion-x2iat-performance-qwen-image-edit
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- |
set +e
EXIT1=$$?
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
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
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
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
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Layered"
key: nightly-diffusion-x2iat-performance-qwen-image-layered
timeout_in_minutes: 120
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- |
set +e
EXIT2=$$?
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json
EXIT3=$$?
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
EXIT=$$?
EXIT4=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
exit $$((EXIT1 | EXIT2 | EXIT3 | EXIT4))
agents:
queue: "mithril-h100-pool"
plugins:
@@ -914,55 +660,7 @@ steps:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 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: Diffusion X2I(&A&T) · Perf Test · BAGEL"
key: nightly-diffusion-x2iat-performance-bagel
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
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
nvidia.com/gpu: 4
volumeMounts:
- name: devshm
mountPath: /dev/shm
@@ -1053,7 +751,7 @@ steps:
- label: ":full_moon: Diffusion X2V · Function Test"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py --run-level "full_model"
agents:
queue: "mithril-h100-pool"
plugins:
@@ -1092,9 +790,7 @@ steps:
timeout_in_minutes: 180
commands:
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m advanced_model --run-level advanced_model
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
agents:
queue: "mithril-h100-pool"
@@ -1177,64 +873,6 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- group: ":card_index_dividers: Diffusion · Quantization Test"
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "diffusion-x2iat-test"
steps:
- label: ":full_moon: Quantization Test with H100"
commands:
- timeout 60m pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_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: ":full_moon: Quantization Test with L4"
timeout_in_minutes: 60
commands:
- pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_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: ":bar_chart: Testcase Statistics"
key: nightly-testcase-statistics
timeout_in_minutes: 120
@@ -1252,7 +890,7 @@ steps:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
@@ -1281,33 +919,19 @@ steps:
- label: ":email: Nightly Collection & Email"
key: nightly-perf-distribution
depends_on:
- nightly-omni-performance-no-async-chunk
- nightly-omni-performance-async-chunk
- nightly-omni-performance-vllm-text
- nightly-omni-performance-multi-replicas
- nightly-omni-performance
- nightly-tts-performance
- nightly-diffusion-x2iat-performance-qwen-image
- nightly-diffusion-x2iat-performance-qwen-image-edit
- nightly-diffusion-x2iat-performance-qwen-image-edit-2511
- nightly-diffusion-x2iat-performance-qwen-image-layered
- nightly-diffusion-x2iat-performance-bagel
- nightly-diffusion-x2iat-performance
- nightly-diffusion-x2v-performance
- nightly-testcase-statistics
if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
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-tts-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
- python tools/nightly/generate_nightly_perf_excel.py
-24
View File
@@ -1,24 +0,0 @@
steps:
- group: ":card_index_dividers: Diffusion X2V Model Test"
key: nightly-diffusion-x2v-group
depends_on: upload-nightly-pipeline
if: >-
build.env("NIGHTLY") == "1" ||
build.pull_request.labels includes "nightly-test" ||
build.pull_request.labels includes "diffusion-x2v-test"
steps:
- label: ":full_moon: Diffusion X2V · Function Test"
agents:
queue: "ascend-a2b3"
resource_class: "npu-4"
image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}"
plugins:
- kubernetes:
podSpecPatch:
imagePullSecrets:
- name: swr-secret
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
HF_TOKEN: "${HF_TOKEN}"
commands:
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model"
-21
View File
@@ -1,21 +0,0 @@
steps:
- group: ":card_index_dividers: TTS Model Test"
key: ready-tts-group
depends_on: upload-ready-pipeline
steps:
- label: ":sound: TTS · Qwen3-TTS CustomVoice Test"
agents:
queue: "ascend-a2b3"
resource_class: "npu-1"
image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}"
plugins:
- kubernetes:
podSpecPatch:
imagePullSecrets:
- name: swr-secret
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
HF_TOKEN: "${HF_TOKEN}"
commands:
- pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model' --run-level 'core_model'
+546 -449
View File
File diff suppressed because it is too large. Load diff
-63
View File
@@ -1,6 +1,5 @@
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_USE_FLASHINFER_MOE_FP16: 0
HF_HUB_DOWNLOAD_TIMEOUT: 300
HF_HUB_ETAG_TIMEOUT: 60
@@ -84,65 +83,3 @@ steps:
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
# Invalid HTTP / WS payloads vs live Omni servers (Issue #3649 tracks known skips).
- group: "Reliability Test - Invalid parameters Test"
depends_on: upload-weekly-pipeline
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
steps:
- label: "Invalid parameters Test · H100"
timeout_in_minutes: 120
commands:
- pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100"
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
# Stable Audio Open, Qwen3-TTS speech/batch/voices paths marked ``L4``.
- label: "Invalid parameters Test · L4"
timeout_in_minutes: 120
commands:
- pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4"
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"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
+4 -4
View File
@@ -515,15 +515,15 @@ After adding parallelism support, update:
### Step 11: Add CPU Offload Support
Implement `SupportsComponentDiscovery` on your pipeline class to enable
Implement `SupportsModuleOffload` on your pipeline class to enable
`--enable-cpu-offload` and `--enable-layerwise-offload`. The protocol
declares which submodules the offloader should manage:
```python
from typing import ClassVar
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
from vllm_omni.diffusion.models.interface import SupportsModuleOffload
class YourPipeline(nn.Module, SupportsComponentDiscovery):
class YourPipeline(nn.Module, SupportsModuleOffload):
_dit_modules: ClassVar[list[str]] = ["transformer"]
_encoder_modules: ClassVar[list[str]] = ["text_encoder"]
_vae_modules: ClassVar[list[str]] = ["vae"]
@@ -540,7 +540,7 @@ class YourPipeline(nn.Module, SupportsComponentDiscovery):
All attribute names support dotted paths for nested submodules
(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
Pipelines without `SupportsComponentDiscovery` fall back to scanning
Pipelines without `SupportsModuleOffload` fall back to scanning
well-known attribute names (`transformer`, `text_encoder`, `vae`,
etc.), which fails for non-standard names.
@@ -92,7 +92,6 @@ class TPSelfAttention(nn.Module):
softmax_scale=1.0 / (self.head_dim ** 0.5),
causal=False,
num_kv_heads=self.to_qkv.num_kv_heads, # Local KV heads per GPU
role="self",
)
def forward(self, x):
@@ -61,7 +61,6 @@ class SelfAttentionBlock(nn.Module):
softmax_scale=1.0 / (self.head_dim ** 0.5),
causal=False,
num_kv_heads=num_heads,
role="self",
)
def forward(self, x, attn_mask=None):
@@ -76,24 +75,6 @@ class SelfAttentionBlock(nn.Module):
return self.to_out(out)
```
**Declaring the attention role.** Always pass a `role` string to `Attention(...)`. It is the lookup key users target with `--diffusion-attention-config.per_role.<role>.backend` to swap kernels per site without touching model code.
| Convention | When to use |
|---|---|
| `role="self"` | Q/K/V come from the same hidden state |
| `role="cross"` | K/V come from a separate `encoder_hidden_states`; pair with `skip_sequence_parallel=True` if K/V is replicated across SP ranks |
For multi-modal sites that don't fit `self` / `cross`, use a dot-namespaced role and provide `role_category` so user config can fall back to the generic category:
```python
self.audio_to_video_attn = Attention(
num_heads=num_heads, head_size=self.head_dim,
softmax_scale=1.0 / (self.head_dim ** 0.5), causal=False,
role="mymodel.audio_to_video",
role_category="cross",
)
```
#### Fused QKV with TP (Advanced)
For tensor parallelism, use vLLM's parallel linear layers:
@@ -123,7 +104,6 @@ class TPSelfAttention(nn.Module):
softmax_scale=1.0 / (self.head_dim ** 0.5),
causal=False,
num_kv_heads=num_heads,
role="self",
)
```
@@ -68,7 +68,7 @@ out = out.reshape(batch, seq_len, -1)
### Different Output vs Diffusers Reference
**Common causes**:
1. **Attention backend difference**: FlashAttention vs SDPA may produce slightly different results. Force SDPA via `--diffusion-attention-backend TORCH_SDPA` (or `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA`) to match diffusers; for finer-grained per-role overrides see `--diffusion-attention-config.per_role.<role>.backend`
1. **Attention backend difference**: FlashAttention vs SDPA may produce slightly different results. Set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` to match diffusers
2. **Float precision**: vLLM-Omni may use bfloat16 where diffusers uses float32 for some operations
3. **Missing normalization**: Check all LayerNorm/RMSNorm are preserved
4. **Scheduler rounding**: Some schedulers have numerical sensitivity
-510
View File
@@ -1,510 +0,0 @@
---
name: diffusion-perf-opt
description: Diagnose and optimize vLLM Omni diffusion workloads, especially Wan/Qwen/Flux-style image and video generation. Use when Codex is asked to analyze profiling traces, choose parallel strategies, inspect torch profiler trace.json or trace.json.gz timelines, estimate optimization ROI, investigate GPU idle/free bubbles, compare USP/CFG/HSDP/VAE parallelism, or design operator/host/quantization optimizations for vLLM Omni.
---
# vLLM Omni Performance Optimization
Use this skill to run a disciplined optimization loop for vLLM Omni diffusion workloads. Keep two ideas separate: real performance baselines are collected with low overhead, while torch profiler traces are diagnostic artifacts and may distort latency.
## First Questions
Before proposing changes, ask for the optimization scene if it is not already known:
- GPU model, card count, topology, and whether NVLink is present.
- Model and pipeline, for example Wan2.2 I2V A14B.
- User workload: resolution, frames, steps, batch/concurrency, CFG scales, prompt/image inputs.
- Three runnable commands for each model/workload, or scripts that generate them:
- **Server startup command**: exact `vllm serve` command, environment variables, model path, port, parallelism flags, profiler flags, and precision/compile settings.
- **User request command**: exact single-request client command (for example curl or Python client) with fixed prompt/media/seed/size/steps. Use this to validate correctness and collect per-request stage timings.
- **User benchmark command**: exact repeatable benchmark command or script with warmup count, measured iteration count, concurrency/batch policy, output directory, and summary format.
- Current enabled strategies: USP/SP, CFG parallel, HSDP/FSDP, VAE patch parallel, torch.compile, profiler options.
- Optimization target: latency, throughput, memory, cost, or quality-preserving speed.
- Precision/quality tolerance: bf16/fp8/quantization/sparsity/approximate attention allowed or not.
- Quality validation method: exact-output comparison when feasible, or metrics
such as SSIM, PSNR, LPIPS, cosine similarity, MAE, MSE, temporal flicker
checks, baseline self-run variance, and reviewable output artifacts.
## Workflow
1. **Freeze the measurement protocol and commands.**
- Before analyzing parallel strategies or traces, establish the three commands for the target model/workload: server startup, single user request, and user benchmark.
- Prefer checked-in scripts or one-off shell scripts over free-form commands in chat. The scripts should make workload variables explicit, including model path, port/base URL, prompt/media inputs, resolution, frames, steps, seed, warmup, measured iterations, concurrency, and output path.
- Measurement reports must preserve the concrete process, not just the final numbers. For each tested configuration, write down the server startup command, the user request command, the user benchmark command or polling loop, the final server response metrics, and the client-side timing/HTTP result.
- Present all duration and latency values in **milliseconds (ms)** in tables and summaries. If an API returns seconds, convert to ms before reporting. Keep raw units only inside quoted raw logs or source snippets.
- Keep profiler modes separate in the command set:
- Baseline/benchmark commands should avoid torch profiler and stack collection.
- Diagnostic commands may enable torch profiler after a baseline identifies a bottleneck.
- If any of the three commands is missing, ask for it or create a proposed script before proceeding.
- Make clear which command is authoritative for correctness validation and which command is authoritative for performance numbers.
2. **Collect a real baseline.**
- Disable PyTorch profiler and stack collection.
- Disable or fix torch.compile state so A/B is fair.
- Avoid `--enforce-eager` for production-speed baselines unless eager is the target.
- Prefer `--log-stats` and `--enable-diffusion-pipeline-profiler` for low-overhead stage timing. PR 3069 has the relevant metrics/log-stats changes; if local code does not include them, fetch or cherry-pick the minimal metrics changes rather than merging unrelated PR drift.
- Run warmup requests and exclude first-request lazy init.
3. **Model the parallel strategy before testing.**
- Estimate compute and communication for self-attention, cross-attention, FFN, CFG, VAE encode/decode, and HSDP.
- Select a small candidate matrix rather than exhaustively testing everything.
- Typical 2-card candidates: `USP=2`, `CFG=2`, `USP=1/HSDP on-off`, VAE parallel on-off if memory allows.
- Typical 4-card candidates: `USP=4`, `CFG=2 x USP=2`, `USP=2 x HSDP`, VAE parallel on-off.
- Typical 8-card candidates for official CFG-enabled video diffusion models:
- Primary candidate: `CFG=2 x USP=4` with VAE patch parallel across all 8 ranks.
- Compare against: `CFG=1 x USP=8` to test whether larger sequence parallel groups beat CFG branch parallelism.
- Isolate VAE parallelism: keep the DiT strategy fixed and compare VAE patch parallel on/off or different VAE patch sizes.
- If long-video `diffuse` remains dominant and Ulysses all-to-all is suspected, test a hybrid sequence strategy such as `CFG=2 x USP=2 x Ring=2`.
- Test HSDP off only if the model fits without it; treat HSDP primarily as a memory strategy until A/B proves otherwise.
- Prefer USP/SP for long video token sequences; prefer CFG parallel when CFG doubles transformer forwards and sequence length is modest.
- Convert workload shape to latent/patch token counts before choosing candidates. For video models, record frames, latent frames, latent HxW, patch size, approximate token count, and which stages should scale with the token count.
4. **Search for the best parallel configuration.**
- Start with a small matrix that answers one question per comparison:
- `CFG parallel vs larger USP`: compare `CFG=2 x USP=world/2` against `CFG=1 x USP=world` for CFG-enabled workloads.
- `VAE patch parallel value`: compare the best DiT strategy with VAE patch parallel enabled and disabled.
- `HSDP cost`: compare HSDP on/off only when both configurations fit memory.
- `Ulysses vs Ring`: test a Ring/Ulysses hybrid only after long-sequence `diffuse` is confirmed dominant.
- For each configuration, create or record a stable config id, for example `A_cfg2_usp4_vaepp8_hsdp_tiling`.
- For each config id, capture:
- exact server command and environment variables,
- observed distributed setup from logs, such as SP groups, CFG groups, HSDP shard/replicate sizes, VAE patch size,
- exact request command for every scenario,
- final server response metrics and client-side elapsed time,
- output artifact paths and any failed/empty responses.
- Use one warmup request per scenario, then at least three measured repeats for the shortlisted config. For exploratory matrix pruning, one measured request is acceptable only if the margin is large; label it as one-shot.
- Compare configurations by stage, not only end-to-end latency. A configuration can improve `diffuse` while hurting `vae.decode`; record both effects.
- Select the best config only after checking the target metric, dominant stages, memory headroom, and output correctness.
5. **Run targeted A/B tests.**
- Change one variable per test.
- Keep model, input, seed, request parameters, GPU placement, and warmup policy fixed.
- Record latency, stage timings, memory, output quality, and logs.
- Report comparison tables in ms. Include at least end-to-end client time, server `inference_time`, server stage generation time if available, `vae.encode`, `diffuse`, `vae.decode`, and peak memory.
6. **Enforce the quality and precision gate for every optimization.**
- Do not mark an optimization ready until quality is checked on the same
workload: model, prompt/media, seed, shape, steps, scheduler, dtype,
backend, parallelism, and output encoding.
- For math-preserving changes, require exact/near-exact agreement or show
that differences are within baseline self-run variance.
- For precision, quantization, approximate attention, sparse/custom kernels,
or backend changes, include reviewable artifacts and metrics such as SSIM,
PSNR, LPIPS, cosine similarity, MAE, MSE, and temporal flicker checks.
- If quality tolerance is not stated, default to quality-preserving behavior.
Failed or inconclusive quality validation blocks a ready-to-merge claim.
7. **Collect diagnostic trace only after narrowing hypotheses.**
- Use torch profiler for a small number of requests.
- Run two separate diagnostic traces instead of mixing concerns:
- **Operator/shape trace**: enable `torch_profiler_record_shapes=True` and keep stack collection disabled. Use this to rank CUDA kernels, NCCL collectives, attention/MLP/norm/RoPE work, and shape-specific hot operators.
- **Host-stack trace**: enable `torch_profiler_with_stack=True` and normally keep shape collection disabled. Use this to map CPU/Python host gaps, synchronization points, scheduler paths, and request handling overhead.
- Keep both trace commands and reports separate from baseline/benchmark commands. Torch profiler latency is diagnostic only and must not be used as the final latency claim.
- Prefer profiling only the narrowed dominant scenario, for example the highest-resolution/video-length case where `diffuse` dominates.
- If profiler endpoints are available, run one warmup request first, then call `/start_profile`, run one profiled request, and call `/stop_profile`. This keeps model initialization and warmup out of the diagnostic trace.
- Start by analyzing rank 0 only. Expand to more ranks only if rank 0 suggests imbalance, unclear GPU idle/free bubbles, high NCCL wait, server timing mismatch, CFG branch imbalance, or USP group stragglers.
- Diagnostic reports must be written to disk and preserve: server command, profiler config, warmup/request/polling commands, trace artifact paths, rank analyzed, analyzer output or summary, and the decision about whether additional ranks are necessary.
- Analyze both rank-level balance and device-level free bubbles when additional ranks are opened.
8. **Analyze host, communication, and operators.**
- Find GPU idle/free intervals and map each large gap to the enclosing CPU/Python code.
- Separate real GPU idle from profiler overhead such as CUPTI `Command Buffer Full`.
- Compare NCCL kernel time to user annotations; annotations can overcount nested intervals.
- Rank operator work by total CUDA time and by repeated small-kernel launch count.
9. **Produce an optimization plan.**
- Classify candidates as P0/P1/P2.
- For each candidate, state necessity, expected benefit, implementation path, validation plan, and quality risk.
- Include a concrete quality/precision validation gate for each candidate.
- Do not implement high-risk operator rewrites before proving the operator is a bottleneck for the target shapes.
- End the plan with a user-facing candidate selection table. The assistant
should not automatically choose a risky optimization just because it is
technically possible. Present the options clearly and let the user decide
which item is worth implementing based on latency target, engineering
budget, memory headroom, and quality tolerance.
- Organize the plan by optimization layer, not by a flat list of ideas:
- host/runtime optimization,
- measurement/benchmark reliability,
- parallelism and communication,
- VAE encode/decode and media pre/post processing,
- operator fusion and layout cleanup,
- attention main-path optimization,
- algorithmic, precision, or approximation changes.
- For each layer, tie every candidate back to evidence from baseline
metrics, diagnostic trace, source code, or output quality requirements.
Do not list generic optimizations without a trace or workload reason.
## Priority Rules
- **P0:** low risk, likely useful, or required for trustworthy measurement. Examples: real baseline, warmup, targeted parallel A/B, disabling avoidable `empty_cache`, scheduler coefficient caching.
- **P1:** meaningful code changes with contained risk. Examples: cross-attention KV caching, VAE gather/broadcast reduction, AdaLayerNorm/RMSNorm/RoPE fusion after trace evidence.
- **P2:** high implementation or quality risk. Examples: FA to LA replacement, custom Triton/CUDA fused kernels, FP8/quantization, sparsity/Rainfusion-style acceleration.
Every implemented optimization needs A/B validation and a passing
quality/precision gate, including math-preserving changes such as padding trim,
layout cleanup, cache reuse, or host/runtime cleanup. Use objective metrics
when available, keep reviewable artifacts, and compare against baseline
self-run variance for precision, quantization, approximate kernels, or backend
changes. If quality validation fails or is inconclusive, do not present the
optimization as ready to merge.
## Optimization Layers
After baseline, parallel-search, and diagnostic traces, summarize optimization
opportunities by layer. This is the core of the performance analysis: the goal
is to connect evidence to a scoped implementation and a validation plan.
### Host and Runtime Optimization
Purpose: remove CPU/Python stalls, synchronization points, allocator overhead,
and request-path overhead that leave GPU lanes empty.
Evidence to look for:
- High `idle_pct` or large `GAP` blocks in `trace_analyzer.py`.
- Host-stack trace lines such as `torch.cuda.empty_cache`,
`cudaStreamSynchronize`, `cudaDeviceSynchronize`, Python locks, scheduler
waits, image/video preprocessing, or repeated small allocation paths.
- Difference between client wall-clock time and server `inference_time_s`.
Typical candidates:
- Make avoidable `torch.cuda.empty_cache()` optional or guard it by memory
headroom.
- Cache scheduler coefficients, timesteps, masks, or other tiny repeated CPU
computations when the request shape/steps are fixed.
- Remove avoidable host-device synchronizations and blocking logging/stat calls.
- Move expensive preprocessing out of the critical path or cache fixed prompt,
image, and transform work for benchmark scenarios.
- Ensure benchmark scripts record client-side elapsed time, HTTP status, output
path, and server response metrics.
Priority guidance:
- Usually P0 when the change is measurement reliability or an obvious removable
synchronization.
- Usually P1 when it changes scheduling, memory lifetime, or request execution
order.
Validation:
- Re-run non-profiler baseline with same workload and seed.
- Confirm peak memory headroom if disabling cache cleanup.
- Confirm generated output exists and quality/seed behavior is unchanged.
### Parallelism and Communication Optimization
Purpose: choose the right decomposition for CFG branches, sequence tokens, model
weights, VAE tiles, and rank topology.
Evidence to look for:
- Baseline A/B across `CFG`, `USP/SP`, `Ring`, `HSDP/FSDP`, and VAE patch
parallelism.
- Stage timing shifts: `diffuse`, `vae.encode`, `vae.decode`, and server
end-to-end.
- NCCL kernel time from trace, not only `user_annotation` time.
- Rank imbalance across SP group ranks or CFG branch ranks.
- Memory headroom and OOM risk.
Typical candidates:
- `CFG=2 x USP=world/2` versus `CFG=1 x USP=world` for CFG-enabled models.
- VAE patch parallel on/off or patch size tuning.
- HSDP/FSDP on/off only if both configurations fit memory.
- Ulysses versus Ulysses+Ring only after long-sequence `diffuse` is confirmed
dominant and all-to-all is suspected.
- Rank mapping/topology changes if all-rank traces show stragglers or NCCL wait.
- Buffer reuse or preallocation for FSDP/HSDP all-gather paths.
Priority guidance:
- P0/P1 for configuration-only changes with strong measured wins.
- P1 for buffer reuse or rank mapping changes.
- P2 for invasive distributed algorithm changes.
Validation:
- Measure by stage and memory, not only end-to-end.
- Use one-variable A/B with identical prompt/media/seed/shape/steps.
- When communication is suspected, compare rank0-3 in one USP group and rank0
versus rank4 across CFG branches for `CFG=2 x USP=4`.
### VAE and Media Pipeline Optimization
Purpose: reduce encode/decode, tiling, split/gather, and media conversion time.
Evidence to look for:
- Large `vae.encode` or `vae.decode` in low-overhead stage timings.
- Host-stack gaps in VAE tile split, gather, merge, broadcast, or image/video
transforms.
- VAE kernels or cuDNN convolution in operator trace.
- Whether every rank needs the final decoded tensor.
Typical candidates:
- Keep VAE patch parallel enabled when it has clear measured benefit.
- Reduce VAE gather/broadcast to only ranks that need the final media output.
- Reuse tile metadata, split buffers, or gather buffers.
- Evaluate bf16/autocast behavior for VAE only with visual quality checks.
- Avoid redundant image conversion, resize, or tensor construction in repeated
benchmark runs.
Priority guidance:
- P0/P1 if VAE is a large share of the target workload or if a host gap is
obvious and low risk.
- Lower priority when `diffuse` dominates and VAE is already patch-parallelized.
Validation:
- Compare `vae.encode`, `vae.decode`, server end-to-end, and peak memory.
- Check output video integrity, artifacts, flicker, and seed stability.
### Operator Fusion and Layout Cleanup
Purpose: reduce high-frequency small kernels, memory bandwidth pressure, layout
conversions, and launch overhead in transformer and VAE blocks.
Evidence to look for:
- Top operator tables showing many `aten::copy_`, `aten::cat`,
`split_with_sizes_copy`, `aten::add`, `aten::mul`, `aten::div`, norm,
activation, RoPE, or reshape/layout kernels.
- `ops_rankN.xlsx` `by_shape` sheet showing repeated small shapes inside the
same block path.
- Trace lanes showing many short kernels between larger GEMM/attention kernels.
- Source code patterns with repeated elementwise chains or layout conversions.
Typical fusion targets:
- AdaLayerNorm / RMSNorm / LayerNorm plus scale/shift fusion.
- RoPE fusion with Q/K layout preparation when shapes are stable.
- Residual add, scale, gate, and elementwise chains.
- MLP gate/up/down cleanup, such as fusing activation and multiply around
GEMM outputs when feasible.
- QKV projection and reshape/split/cat path cleanup.
- Attention pre/post layout cleanup to avoid unnecessary copies, cats, and
splits around sequence parallel all-to-all.
Priority guidance:
- P1 when implemented with existing PyTorch/Triton/local helper patterns and
validated against exact outputs or tolerances.
- P2 when it requires custom CUDA/Triton kernels, changes numerics, or touches
attention math directly.
Validation:
- First prove the operator family is material for the target shape.
- Use non-profiler A/B for latency and stage timing.
- Use quality regression checks for generated video stability.
- Check compile behavior and graph breaks if using `torch.compile`.
### Attention Main-Path Optimization
Purpose: address the dominant self-attention cost when FlashAttention or other
attention kernels dominate CUDA time.
Evidence to look for:
- Operator trace where FlashAttention/SDPA kernels dominate total CUDA time.
- Attention shape from `ops_rankN.xlsx` `by_shape`, model code, or trace
metadata.
- Whether attention cost scales with latent frames, latent H/W, patch size, or
CFG duplication.
- Layout/copy/all-to-all work around attention.
Typical candidates:
- Verify the attention backend and shape are on the intended fast kernel path.
- Compare supported attention backends only with identical workload and quality
settings.
- Reduce attention input size by safe model/config choices when allowed:
latent resolution, frame count, patching, boundary ratio, or windowing.
- Remove avoidable layout conversions before/after attention.
- Reuse condition-side KV or other static inputs if the model structure allows.
- Consider custom kernels, sparse/window/linear attention, or approximation only
after quality risk is accepted.
Priority guidance:
- P1 for backend/config/layout changes with preserved math.
- P2 for approximate attention, sparsity, custom kernels, or any change that can
alter quality/temporal consistency.
Validation:
- Always include output quality and seed behavior checks.
- Compare `diffuse`, server end-to-end, peak memory, and attention kernel time
in diagnostic traces if needed.
### Algorithmic, Precision, and Approximation Optimization
Purpose: reduce mathematical work or precision cost beyond local code cleanup.
Evidence to look for:
- A single operator family dominates even after low-risk runtime, parallel, and
fusion work.
- Memory bandwidth or compute utilization suggests precision or quantization
could matter.
- The user explicitly allows quality-preserving or approximate methods.
Typical candidates:
- FP8/quantization for transformer or selected projections.
- Sparsity or Rainfusion-style acceleration.
- Reduced steps, scheduler changes, distillation, or caching across frames.
- Approximate attention or linear attention.
Priority guidance:
- Usually P2 because quality, numerics, and implementation risk are high.
Validation:
- Requires strict A/B, visual quality review, temporal flicker checks, seed
stability, and possibly human evaluation.
### Interpolation, Super-Resolution, and E2E Pipeline Optimization
Purpose: optimize the whole user-visible video product, not only the base
diffusion invocation. Some deployments trade base-model latency against
post-processing, interpolation, or super-resolution stages.
Evidence to look for:
- E2E latency breakdown across base generation, interpolation, super-resolution,
encoding, storage, and response streaming.
- Fast/slow GPU or fast/slow stage analysis across multiple cards and pipeline
stages.
- User quality target: resolution, FPS, temporal smoothness, and acceptable
post-processing artifacts.
Typical candidates:
- Add or optimize a frame interpolation stage when it reduces required base
model frames for the same perceived FPS.
- Add or optimize a super-resolution model when generating lower base
resolution plus SR is faster for the target quality.
- Analyze E2E2 pipeline behavior: client request, service scheduling, diffusion,
VAE/media, post-process, file write, and response.
- Identify fast/slow cards or stages and rebalance pipeline placement.
Priority guidance:
- P1 when using proven interpolation/SR components without changing diffusion
math.
- P2 when quality risk is high or the pipeline adds significant operational
complexity.
Validation:
- Measure E2E wall-clock, per-stage server timings, output FPS/resolution,
artifacts, flicker, and user-visible quality.
### Optimization Candidate Library
Use this table as a compact menu, not as automatic recommendations. Pick items
only when baseline metrics, trace evidence, source inspection, or quality
tolerance supports them.
| Layer | Candidate | Evidence | Priority | Validation focus |
|---|---|---|---|---|
| Measurement | Freeze server/request/benchmark commands | Missing or drifting commands | P0 | Repeatable non-profiler A/B |
| Measurement | Separate baseline and diagnostic profiler runs | Profiler used for latency claims | P0 | Low-overhead stage timings |
| Host/runtime | Guard or remove avoidable `empty_cache` | Host-stack gaps or sync stalls | P0/P1 | Latency, peak memory, OOM safety |
| Host/runtime | Cache scheduler coefficients/timesteps | Repeated tiny CPU/GPU work | P0/P1 | Same seed/output, stage timing |
| Host/runtime | Reduce framework scheduling overhead | Client time exceeds server time | P1 | E2E latency, throughput |
| Parallel | `CFG=2 x USP=world/2` vs `USP=world` | CFG doubles forward work | P0/P1 | `diffuse`, NCCL, memory |
| Parallel | Tune VAE patch parallelism | VAE encode/decode is material | P0/P1 | VAE time, output correctness |
| Parallel | HSDP on/off or buffer reuse | HSDP affects memory/all-gather | P1 | Memory, latency, OOM risk |
| Parallel | Ulysses vs Ulysses+Ring | Long sequence all-to-all suspected | P1/P2 | Rank balance, NCCL kernels |
| Cross-attn | Disable SP for short condition tokens | Cross-attn comm exceeds compute | P1 | `diffuse`, correctness |
| VAE/media | Reduce VAE gather/broadcast | Rank traces show VAE wait | P1 | Rank balance, output file |
| VAE/media | Reuse tile metadata/buffers | Tile split/merge host gaps | P1 | `vae.encode/decode`, memory |
| VAE/media | VAE bf16/autocast | VAE float kernels are slow | P1/P2 | Artifacts, flicker, seed stability |
| Operator fusion | AdaLayerNorm/LayerNorm fusion | Norm plus scale/shift kernels | P1 | Numeric tolerance, latency |
| Operator fusion | RMSNorm fusion | Many small RMSNorm kernels | P1 | Numeric tolerance, latency |
| Operator fusion | RoPE cache/fuse/layout cleanup | RoPE copy/reshape kernels | P1/P2 | Kernel count, correctness |
| Operator layout | QKV or attention layout cleanup | Copy/cat/split around attention | P1 | Copy kernels, compile behavior |
| Attention | Verify backend fast path | FA/SDPA dominates trace | P1 | `diffuse`, attention kernels |
| Attention | FA to LA or selected-head LA | Attention remains dominant | P2 | Quality, temporal stability |
| Precision | Transformer FP8/quantization | Compute/bandwidth bound and allowed | P2 | Quality, speed, stability |
| Sparsity | Rainfusion-style acceleration | DiT compute remains dominant | P2 | Prompt diversity, quality |
| Pipeline | Frame interpolation | Fewer base frames can meet FPS | P1/P2 | E2E latency, motion artifacts |
| Pipeline | Super-resolution | Lower base res plus SR may win | P1/P2 | Detail quality, artifacts |
| E2E | Fast/slow-card analysis | Multi-card stragglers | P0/P1 | Per-rank/stage wall-clock |
### Optimization Plan Template
Use this table shape when reporting the next work items:
| Priority | Layer | Candidate | Evidence | Expected benefit | Implementation path | Validation | Quality risk |
|---|---|---|---|---|---|---|---|
| P0 | Host/runtime | Guard `empty_cache` | Host-stack gap points to `torch.cuda.empty_cache` | Small latency reduction, less idle | Add config/env guard | Non-profiler A/B, memory check | Low |
| P1 | Operator fusion | RMSNorm/AdaLayerNorm fusion | High-frequency norm/elementwise kernels | Lower launch/bandwidth overhead | Use existing fusion helper or targeted Triton | A/B + output check | Medium |
| P1/P2 | Attention | Attention layout/backend investigation | FA kernel dominates CUDA time | Potentially large | Inspect shapes/backend and remove layout copies | A/B + trace + quality | Medium/high |
Then present a short selection prompt using the same rows:
```text
Which candidate should we implement next?
1. P0 Host/runtime: guard empty_cache
- Expected benefit: small but low-risk latency reduction.
- Risk: possible memory increase/OOM if memory headroom is insufficient.
2. P1 Operator fusion: inspect by_shape and implement first norm/RoPE/layout fusion
- Expected benefit: medium if high-frequency small kernels are confirmed.
- Risk: numerical/compile/quality validation needed.
3. P1/P2 Attention: FA/LA/backend/layout investigation
- Expected benefit: potentially large.
- Risk: high quality and implementation risk.
```
If the user has not chosen an item, default to explaining tradeoffs and asking
which candidate to execute. Only proceed autonomously on low-risk P0 measurement
or instrumentation fixes.
## Analysis Helpers
PyTorch profiler traces are Chrome/Perfetto-compatible JSON files, usually
`trace_rankN.json` or `trace_rankN.json.gz`. They normally contain a top-level
`traceEvents` list, though some exporters emit the raw event list directly.
Use the checked-in analyzer from the repository root:
```bash
python3 .claude/skills/diffusion-perf-opt/scripts/trace_analyzer.py \
vllm_profile/.../trace_rank0.json.gz \
--min-gap-ms 5 \
--topn 20
```
For rank imbalance or communication questions, pass all relevant rank traces in
one command. For host gaps, lower `--min-gap-ms` to `1` and use a host-stack
trace. Read `gpu_span_s`, `busy_union_s`, `idle_union_s`, `idle_pct`, `GAP`
blocks, `Top GPU/operator events by total duration`, and `Top NCCL-like events
by category`. Treat `cat=user_annotation` NCCL ranges as enclosing annotations;
prefer `cat=kernel` or `cat=gpu_user_annotation` for real device work.
The analyzer summarizes timing only. It does not parse tensor shapes, attribute
overlap to individual streams, prove quality, or provide final latency claims.
Use `ops_rankN.xlsx` or PyTorch key averages for shape analysis, and re-test any
optimization with non-profiler baseline commands.
Read `references/optimization-playbook.md` when drafting the optimization table or comparing candidate techniques.
## vLLM Omni Heuristics
- Cross-attention usually should not use USP/SP when text/image condition token count is much smaller than latent video tokens. Confirm via trace; in Wan2.2 I2V, self-attention dominates cross-attention.
- VAE bf16/autocast is often worthwhile but requires visual quality checks.
- VAE patch parallel can help decode/encode but may add gather/merge/broadcast overhead. Check whether all ranks need the final decoded tensor.
- HSDP/FSDP is primarily a memory strategy. If the model fits without it, run an on/off latency comparison.
- Scheduler work can create small host/device gaps; cache tiny solve coefficients when timesteps/order are known.
- `torch.cuda.empty_cache()` can prevent OOM but creates synchronization/idle. Make it optional if memory headroom is sufficient.
- `Command Buffer Full` in profiler output is profiler overhead, not a model optimization target.
@@ -1,4 +0,0 @@
interface:
display_name: "vLLM Omni Perf Opt"
short_description: "Profile and optimize vLLM Omni diffusion performance."
default_prompt: "Analyze and optimize a vLLM Omni diffusion workload using baselines, traces, parallel strategy modeling, and prioritized optimization plans."
@@ -1,82 +0,0 @@
# Optimization Playbook
Use this reference when turning trace evidence into prioritized vLLM Omni optimizations.
## Parallel Strategy
Start from workload shape:
- Long video tokens, high resolution, many frames: self-attention and FFN dominate. USP/SP is usually valuable.
- CFG enabled with positive and negative prompts: transformer forward count doubles. CFG parallel may beat USP on small/moderate sequence lengths.
- Cross-attention: skip SP/USP when condition tokens are much smaller than latent tokens; extra all-to-all can exceed compute saved.
- HSDP/FSDP: use to fit memory first. If memory allows, compare without HSDP because all-gathers and lazy init add overhead.
- VAE parallel: useful when encode/decode is material. Check gather/merge/broadcast and whether every rank needs final decoded output.
Candidate matrices:
- 2 cards: `USP=2`, `CFG=2`, HSDP on/off, VAE PP on/off.
- 4 cards: `USP=4`, `CFG=2 x USP=2`, `USP=2 x HSDP`, VAE PP on/off.
- 8 cards: add pipeline/stage parallel only after measuring orchestration and inter-stage transfer overhead.
## Host And Free Bubble Analysis
Look for device intervals where no kernel, memcpy, or memset is active.
Common causes:
- First request FSDP/HSDP lazy init.
- Python preprocessing: image load, resize, `to_tensor`, tokenizer, request marshaling.
- `torch.cuda.empty_cache()` or explicit synchronizations.
- Scheduler small linear algebra such as tiny `torch.linalg.solve`.
- VAE tile split/merge Python loops.
- Profiler overhead such as CUPTI command buffer flush/full.
Also inspect trace-external time:
- API request parsing and multipart upload.
- Engine/orchestrator dispatch.
- Worker RPC and queue wait.
- Output serialization, video write, base64, HTTP response.
## Operator Optimization Candidates
Evaluate necessity before implementation:
- **FA to LA / LA for head:** high potential only if attention is the main CUDA hotspot and target shapes fit LA well.
- **Cross-attention BSND/no transpose:** do only if transpose/copy around cross-attn is measurable.
- **AdaLayerNorm fusion:** fuse norm, scale/shift, gate, residual, and dtype conversion if repeated elementwise kernels are significant.
- **RMSNorm fusion:** useful for Q/K norms and VAE norms when small-kernel launch count is high.
- **FP32 LayerNorm to bf16:** quality-sensitive; run visual regression.
- **LA preprocess:** integrate only if preprocess overhead is cached or amortized.
- **RoPE move/cache/fuse:** useful when RoPE small kernels are visible around every self-attn block.
- **VAE RMSNorm replacement:** only after VAE internal trace shows norm as a real component of decode/encode.
- **Scheduler fusion/cache:** cache coefficient solves and fuse CFG combine with scheduler arithmetic when possible.
Before writing a custom kernel, check:
- Local vLLM Omni implementations.
- Upstream vLLM kernels and attention backends.
- SGLang, FlashInfer, FlashAttention, and Triton examples.
- Shape constraints and layout expectations.
## Quantization And Sparsity
Treat these as high-risk optimizations:
- Rainfusion/sparse diffusion can be high reward but needs model-specific quality validation.
- VAE decode quantization can reduce decode time but usually has smaller end-to-end upside than transformer attention/FFN work.
- FP8 or approximate kernels require seed-stable visual regression, temporal flicker checks, and prompt diversity.
## Output Format For Recommendations
Return a table with:
- Priority: P0/P1/P2.
- Optimization.
- Evidence from trace or code.
- Expected benefit.
- Risk/quality impact.
- Implementation path.
- Validation command or A/B test.
State clearly when a recommendation is a hypothesis requiring new trace with shapes.
@@ -1,144 +0,0 @@
#!/usr/bin/env python3
"""Summarize torch profiler Chrome traces for vLLM Omni diffusion optimization."""
from __future__ import annotations
import argparse
import collections
import gzip
import json
from pathlib import Path
from typing import Any
GPU_CATS = {"kernel", "gpu_memcpy", "gpu_memset"}
CPU_CATS = {"python_function", "user_annotation", "cpu_op", "cuda_runtime", "cuda_driver"}
def open_trace(path: Path):
if path.suffix == ".gz":
return gzip.open(path, "rt")
return path.open("rt")
def event_name(event: dict[str, Any]) -> str:
return str(event.get("name", ""))
def summarize_trace(path: Path, min_gap_us: float, topn: int) -> None:
with open_trace(path) as f:
data = json.load(f)
events = data if isinstance(data, list) else data.get("traceEvents", [])
gpu: list[tuple[float, float, float, str, str, Any, Any]] = []
cpu: list[tuple[float, float, float, str, str, Any, Any]] = []
by_gpu_name: dict[str, list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
by_name: dict[str, list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
nccl: dict[tuple[str, str], list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
for event in events:
dur = event.get("dur")
ts = event.get("ts")
if dur is None or ts is None or dur <= 0:
continue
cat = str(event.get("cat", ""))
name = event_name(event)
row = (float(ts), float(ts + dur), float(dur), name, cat, event.get("pid"), event.get("tid"))
stat = by_name[name]
stat[0] += 1
stat[1] += dur
stat[2] = max(stat[2], dur)
if cat in GPU_CATS:
gpu.append(row)
gstat = by_gpu_name[name]
gstat[0] += 1
gstat[1] += dur
gstat[2] = max(gstat[2], dur)
elif cat in CPU_CATS:
cpu.append(row)
if "nccl" in name.lower():
nstat = nccl[(cat, name)]
nstat[0] += 1
nstat[1] += dur
nstat[2] = max(nstat[2], dur)
print(f"\n== {path}")
print(f"events={len(events)} gpu_events={len(gpu)} cpu_events={len(cpu)}")
if not gpu:
print("No GPU events found.")
return
gpu.sort()
merged: list[list[Any]] = []
for start, end, dur, name, cat, pid, tid in gpu:
if not merged or start > merged[-1][1]:
merged.append([start, end, [(start, end, dur, name, cat, pid, tid)]])
else:
merged[-1][1] = max(merged[-1][1], end)
merged[-1][2].append((start, end, dur, name, cat, pid, tid))
span = max(end for _, end, *_ in gpu) - min(start for start, *_ in gpu)
busy = sum(end - start for start, end, _ in merged)
idle = span - busy
print(
f"gpu_span_s={span / 1e6:.3f} "
f"busy_union_s={busy / 1e6:.3f} "
f"idle_union_s={idle / 1e6:.3f} "
f"idle_pct={idle / span * 100:.2f}"
)
interesting_cpu = [
row
for row in cpu
if row[2] >= 1000
and (
row[4] in {"python_function", "user_annotation"}
or "cudaStreamSynchronize" in row[3]
or "cudaDeviceSynchronize" in row[3]
or "cudaLaunch" in row[3]
or "cudaMemcpy" in row[3]
)
]
gaps = []
for idx in range(1, len(merged)):
gap_start = merged[idx - 1][1]
gap_end = merged[idx][0]
gap_dur = gap_end - gap_start
if gap_dur >= min_gap_us:
prev_event = max(merged[idx - 1][2], key=lambda x: x[1])
next_event = min(merged[idx][2], key=lambda x: x[0])
mid = (gap_start + gap_end) / 2
containers = [row for row in interesting_cpu if row[0] <= mid <= row[1]]
containers = sorted(containers, key=lambda x: x[2])[:8]
gaps.append((gap_dur, gap_start, gap_end, prev_event, next_event, containers))
print(f"gaps_ge_{min_gap_us / 1000:.3f}ms count={len(gaps)} sum_s={sum(g[0] for g in gaps) / 1e6:.3f}")
for gap_dur, gap_start, gap_end, prev_event, next_event, containers in sorted(gaps, reverse=True)[:topn]:
print(f"\nGAP {gap_dur / 1000:.3f} ms ts={gap_start:.0f}->{gap_end:.0f}")
print(f" prev {prev_event[4]} {prev_event[2] / 1000:.3f} ms {prev_event[3][:160]}")
print(f" next {next_event[4]} {next_event[2] / 1000:.3f} ms {next_event[3][:160]}")
for row in containers:
print(f" in {row[4]} {row[2] / 1000:.3f} ms {row[3][:180]}")
print("\nTop GPU/operator events by total duration:")
for name, (count, total, max_dur) in sorted(by_gpu_name.items(), key=lambda kv: kv[1][1], reverse=True)[:topn]:
print(f" {int(count):8d} total={total / 1e6:9.3f}s max={max_dur / 1000:9.3f}ms {name[:180]}")
print("\nTop NCCL-like events by category:")
for (cat, name), (count, total, max_dur) in sorted(nccl.items(), key=lambda kv: kv[1][1], reverse=True)[:topn]:
print(f" {int(count):8d} total={total / 1e6:9.3f}s max={max_dur / 1000:9.3f}ms cat={cat} {name[:160]}")
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("traces", nargs="+", type=Path, help="trace.json or trace.json.gz files")
parser.add_argument("--min-gap-ms", type=float, default=5.0, help="minimum GPU idle gap to print")
parser.add_argument("--topn", type=int, default=20, help="number of gaps/hotspots to print")
args = parser.parse_args()
for trace in args.traces:
summarize_trace(trace, args.min_gap_ms * 1000.0, args.topn)
if __name__ == "__main__":
main()
-3
View File
@@ -18,9 +18,6 @@ include:
- `add-diffusion-model`: guides integration of a new diffusion model into
`vllm-omni`
- `diffusion-perf-opt`: guides diffusion model performance optimization,
including profiling traces, parallel strategies, stage timing analysis, and
benchmark-driven tuning
- `add-omni-model`: covers addition of new omni-modality model support
- `add-tts-model`: covers integration of new TTS models and related serving
workflows
+10 -59
View File
@@ -3,64 +3,15 @@
# More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
# Default owners (fallback reviewers for all files)
* @Gaohan123 @hsliuustc0106 @ywang96
# later, we can add more reviewers here
* @hsliuustc0106
# Repository operations
/.github/ @hsliuustc0106
# Assign reviewers by directory
/.github/ @owner1
/tests/ @owner1
/docs/ @owner1
# Documentation and Recipes
/docs/ @hsliuustc0106 @Gaohan123 @david6666666
/recipes @hsliuustc0106 @Gaohan123 @david6666666
# Skills
/.claude/ @hsliuustc0106
# Tests
/tests/ @yenuo26
# CI
/.buildkite/ @yenuo26 @congw729
# Engine and model_executor
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan
/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan
/vllm_omni/model_executor/ @tzhouam @gcanlin
/vllm_omni/worker/ @tzhouam @gcanlin
# Diffusion
/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666
/vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
# Diffusion CPU offloading
/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
# Omni and TTS support
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter
# Benchmarks
/benchmarks/accuracy/ @david6666666
/benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa
/benchmarks/distributed/ @princepride
/benchmarks/glm_image/ @princepride @JaredforReal
/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride
# Usability and Observability
/vllm_omni/deploy/ @tzhouam @lishunyang12
/vllm_omni/metrics/ @lishunyang12
/vllm_omni/profiler/ @gcanlin
# ComfyUI
/apps/ @fhfuih
# Large Scale Serving
/vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
# Quantization and configuration
/vllm_omni/config/ @lishunyang12
/vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
# Hardware plugins and accelerators
/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
# Example: reviewers for specific directories
# /examples/offline_inference/ @offline-inference-team
# /vllm_omni/core/ @core-team
# /vllm_omni/model_executor/ @model-team
-1
View File
@@ -239,7 +239,6 @@ data/
datasets/
*.csv
*.json
!tests/dfx/perf/tests/*.json
!apps/ComfyUI-vLLM-Omni/example_workflows/*.json
*.jsonl
*.parquet
-5
View File
@@ -10,11 +10,6 @@ build:
jobs:
post_checkout:
- git fetch --unshallow || true
pre_install:
# Pre-install CPU torch so the subsequent `pip install .[docs]` does not
# pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
# 1 GB memory cap and time out the mkdocs build.
- python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
mkdocs:
configuration: mkdocs.yml
+2 -2
View File
@@ -9,19 +9,19 @@ 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://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</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/111-L8zF7A1j_YI_cR8JsblofdScdRr2f/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><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/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
</p>
---
*Latest News* 🔥
- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
- [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).
- [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
---
+35 -51
View File
@@ -1,60 +1,44 @@
# Benchmarks
# Benchmarks Overview and Architecture
This directory contains benchmark suites for evaluating different model families and infrastructure components in vLLM-Omni. Each subfolder targets a different benchmark family with its own scripts, configs, and metrics. See the per-subfolder READMEs for detailed usage.
This document explains the benchmark architecture across all benchmark assets in this repo. It describes what we measure, and where to find or plug in new scenarios. Per-task details remain in subfolder READMEs (e.g., `benchmarks/<model>/README.md`).
## Benchmark families
## Scope and goals
- Establish repeatable latency/throughput measurements for multimodal LLM pipelines.
- Provide both HF Transformers (offline) and vLLM-Omni (multi-stage/pipeline) baselines.
- Make it easy to plug in new datasets and models with minimal changes to the runner scripts.
### [TTS](tts/README.md) — Text-to-Speech
## Dataset and inputs
- Default example: SeedTTS top-100 prompts (`benchmarks/build_dataset/top100.txt`) via `benchmarks/build_dataset/`.
- Extensible: drop in new prompt files or modality-aligned payloads; keep the expected format for the consuming scripts (e.g., one prompt per line).
- If you add a new dataset, document it under `benchmarks/<model>/README.md` and point scripts to your data path.
Model-agnostic serving benchmarks for TTS models, including Qwen3-TTS and VoxCPM2.
## 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.
- **Layout**: `tts/bench_tts.py` (serving benchmark driver), `tts/model_configs.yaml` (model registry), `tts/plot_results.py` (result plotting)
- **Dataset**: Seed-TTS full or text-only datasets, plus bundled smoke/design prompts under `build_dataset/`
- **Key metrics**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor), throughput (audio seconds / wall-clock second)
## Reference workflows
- **HF Transformers (offline, single process)**
Script (example): `benchmarks/<model>/transformers/eval_qwen3_moe_omni_transformers.sh`
Outputs: `benchmark_results/perf_stats.json`, `benchmark_results/results.json`, `benchmark_results/audio/` (if audio is produced).
### [Diffusion](diffusion/README.md) — Image and Video Generation
- **vLLM-Omni end-to-end pipeline**
Script (example): `benchmarks/<model>/vllm_omni/eval_qwen3_moe_omni.sh`
Outputs: `vllm_omni/logs/*.stats.jsonl` (per-stage/overall latency & TPS), `vllm_omni/logs/stage*.log`, `vllm_omni/outputs/` (text/audio artifacts).
Online-serving benchmark for diffusion image/video models, sending requests to the configured vLLM serving endpoint (`/v1/chat/completions`, `/v1/images/generations`, or `/v1/videos`, depending on backend/task).
- **Adding a new task/model**
1) Create `benchmarks/<model>/transformers/` and/or `benchmarks/<model>/vllm_omni/` with scripts referencing your model and dataset.
2) Add a task README describing dataset, configs, and expected outputs.
3) Keep the output/log structure similar for easy comparison (perf_stats/results/audio or text outputs; stats.jsonl/logs for pipeline).
- **Tasks**: text-to-image, text-to-video, image-to-image, image-to-video, text+image-to-image, text+image-to-video
- **Datasets**: `vbench`, `trace`, `random`
- **Key metrics**: request throughput, latency percentiles, optional SLO attainment
## Metrics to watch
- **Throughput**: `overall_tps`, `*_tps_avg` per stage.
- **Latency distribution**: look for long tails in `*.stats.jsonl`.
- **Quality/completeness**: missing outputs or errors in stage logs indicate pipeline failures or misconfigurations.
### [GLM-Image](glm_image/README.md) — Text-to-Image and Image-to-Image
Benchmarks for GLM-Image performance across HuggingFace baseline, vLLM-Omni offline inference, and vLLM-Omni online serving.
- **Layout**: `glm_image/huggingface/` (HF baseline), `glm_image/vllm-omni/` (offline inference), `glm_image/benchmark_glm_image.py` (online serving)
- **Tasks**: text-to-image and image-to-image
- **Key metrics**: request/image throughput, latency percentiles, optional per-stage pipeline timings
### [Distributed](distributed/omni_connectors/README.md) — RDMA Connector Testing
RDMA environment setup and transfer tests for `MooncakeTransferEngineConnector`, including pytest-based single-node checks and manual cross-node benchmarks.
- **Transfer modes**: `copy`, `zerocopy`, `gpu` (GPUDirect)
- **Supports**: single-node pytest suites and manual multi-node/cross-node transfer testing
### [Accuracy](accuracy/README.md) — Image Generation and Editing Quality
Accuracy benchmarks for image generation/editing models, adapting external suites to vLLM-Omni serving and local judge-evaluation flows.
- **Layout**: `accuracy/text_to_image/` (GEBench), `accuracy/image_to_image/` (GEdit-Bench)
- **Method**: generation and judge scoring both run through local `vllm-omni serve` endpoints
### Common serving metrics framework
`vllm_omni/benchmarks/` extends `vllm bench serve --omni` with Omni-specific datasets, backends, and multimodal metrics. Key metrics include:
- **Text output**: TTFT (time to first token), TPOT (time per output token), ITL (inter-token latency)
- **Audio output**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor)
- **Throughput**: request throughput, output token throughput, total token throughput, audio throughput
See `vllm_omni/benchmarks/serve.py` for the `vllm bench serve --omni` runner wrapper and `vllm_omni/benchmarks/metrics/` for Omni metric definitions.
## Adding a new benchmark
1. Create a subfolder under `benchmarks/<name>/` with scripts, configs if needed, and a `README.md`.
2. If comparing against another runtime, use clear backend subfolders where applicable, such as `huggingface/` and `vllm-omni/`, or follow the shared TTS serving benchmark pattern in `tts/`.
3. Add reusable dataset or prompt-building utilities to `build_dataset/` if applicable.
4. Update this README with a link to the new benchmark family.
## Troubleshooting
- Verify GPU/driver/FlashAttention2 requirements for your chosen model/config.
- Ensure network access for dataset/model downloads (Google Drive, Hugging Face, etc.).
- If outputs are missing or slow, inspect per-stage logs and `*.stats.jsonl` for errors, stragglers, or contention.
@@ -18,13 +18,3 @@ vd017|||The quarterly results exceed expectations across all major metrics.|A sh
vd018|||Chapter one. The morning sun filtered gently through the forest canopy.|A smooth, rich male audiobook narrator voice, expressive and engaging.
vd019|||To be or not to be, that is the question.|A theatrical female voice, dramatic and expressive, stage projection quality.
vd020|||And that is all for tonight. Stay well out there, everyone.|A warm, velvety male late-night radio DJ voice, smooth and intimate.
vd021|||Please fasten your seatbelts; we will be experiencing some turbulence shortly.|A calm, professional female flight attendant voice, reassuring and clear.
vd022|||The fossil record suggests these creatures lived more than sixty million years ago.|A measured male documentary narrator voice, scientific and curious.
vd023|||Hey kiddo, did you finish your homework before dinner today?|A gentle, caring father voice, warm and slightly tired after a long day.
vd024|||Live from courtside, the home team has just taken a commanding lead.|A high-energy male sports broadcaster, urgent and engaged.
vd025|||Please remain on the line while I transfer you to the next available agent.|A polite, neutral female call-center voice, professional and patient.
vd026|||I have a dream that one day this nation will rise up and live out the true meaning of its creed.|A powerful, resonant male orator voice, slow and impassioned.
vd027|||Mama always said life was like a box of chocolates; you never know what you're gonna get.|A folksy male voice with a Southern American accent, warm and reflective.
vd028|||Ahoy there, matey! Hoist the colors and prepare to set sail.|A theatrical pirate male voice, gruff and boisterous.
vd029|||Ladies and gentlemen, please welcome to the stage our keynote speaker.|A polished female event MC voice, gracious and articulate.
vd030|||Once you've completed the form, click submit and we'll review it within two business days.|A friendly female product-onboarding voice, encouraging and clear.
@@ -18,13 +18,3 @@ smoke017|||In the distance, the mountains were shrouded in thick morning fog.
smoke018|||Our company reported record revenue for the fourth quarter of the fiscal year.
smoke019|||She explained the new policy in detail during the staff meeting this morning.
smoke020|||The children laughed and played in the garden until the sun began to set.
smoke021|||A long line of cars stretched along the coastal highway after the lighthouse opened.
smoke022|||The chef recommended the seared scallops with a lemon butter sauce as the special.
smoke023|||Our quarterly engineering review begins promptly at two in the afternoon on Thursday.
smoke024|||Despite the heavy rain, the marathon runners kept a steady pace through downtown.
smoke025|||She finished the crossword puzzle in under fifteen minutes while sipping her tea.
smoke026|||The new release adds streaming output, faster startup, and better error reporting.
smoke027|||His grandfather taught him how to whittle small wooden birds from cedar offcuts.
smoke028|||Travelers were advised to check the timetable since trains were running on a holiday schedule.
smoke029|||The startup launched a public beta after eighteen months of closed development.
smoke030|||At night the desert sky filled with stars so bright the Milky Way was clearly visible.
+33 -3
View File
@@ -30,8 +30,10 @@ python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
**Notes**
- By default, image tasks talk to `http://<host>:<port>/v1/chat/completions`; video tasks talk to `/v1/videos`.
- The benchmark talks to `http://<host>:<port>/v1/chat/completions`.
- If you run the server on another host or port, pass `--base-url` accordingly.
- To compare diffusion multi-replica serving, restart the server with different
`runtime.num_replicas` values and keep the benchmark parameters unchanged.
## 2. Supported Datasets
@@ -93,11 +95,39 @@ You can point to your own trace using `--dataset-path`.
## 3. Benchmark Parameters
### Multi-replica sweep
For replica scaling, use the same model, prompt count, image/video size, and
step count for every run. Change only the serving stage config:
```yaml
runtime:
devices: "0,1,2,3"
num_replicas: 4
```
Then run a fixed benchmark and write one JSON file per replica count:
```bash
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--base-url http://localhost:8099 \
--model Qwen/Qwen-Image \
--task t2i \
--dataset random \
--num-prompts 32 \
--max-concurrency 8 \
--request-rate inf \
--num-inference-steps 10 \
--width 512 --height 512 \
--output-file diffusion_replica_4.json
```
Repeat for `num_replicas=1`, `2`, `3`, and `4`, using matching `runtime.devices`.
### Basic flags
- `--base-url`: Server address; `--endpoint` selects the path appended to this base URL.
- `--base-url`: Server address (the script calls `.../v1/chat/completions`).
- `--model`: The OpenAI-compatible `model` field.
- `--endpoint`: API endpoint path. Leading `/` is optional, e.g. `/v1/videos` or `v1/videos`.
- `--task`: Task type (e.g., `t2i`, `t2v`, `i2i`, `i2v`).
- `--dataset`: Dataset mode (`vbench` / `trace` / `random`).
- `--num-prompts`: Number of requests to send.
+4 -123
View File
@@ -10,8 +10,6 @@ from typing import Any
import aiohttp
from tqdm import tqdm
DEFAULT_EDITS_BOT_TASK = "think"
@dataclass
class RequestFuncInput:
@@ -29,7 +27,6 @@ class RequestFuncInput:
extra_body: dict[str, Any] = field(default_factory=dict)
image_paths: list[str] | None = None
request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
default_bot_task: str | None = DEFAULT_EDITS_BOT_TASK
@dataclass
@@ -56,95 +53,6 @@ def _encode_image_as_data_url(path: str) -> str:
return f"data:{mime};base64,{encoded}"
async def async_request_image_edits(
input: RequestFuncInput,
session: aiohttp.ClientSession,
pbar: tqdm | None = None,
enable_diffusion_pipeline_profiler: bool = False,
) -> RequestFuncOutput:
"""POST /v1/images/edits (multipart)."""
del enable_diffusion_pipeline_profiler
output = RequestFuncOutput()
output.start_time = time.perf_counter()
extra_body = dict(input.extra_body)
width = input.width or extra_body.get("width") or 1024
height = input.height or extra_body.get("height") or 1024
edits_url = input.api_url
form = aiohttp.FormData()
form.add_field("model", input.model)
form.add_field("prompt", input.prompt)
form.add_field("size", f"{width}x{height}")
form.add_field("response_format", "b64_json")
if input.num_inference_steps is not None:
form.add_field("num_inference_steps", str(input.num_inference_steps))
elif extra_body.get("num_inference_steps") is not None:
form.add_field("num_inference_steps", str(extra_body["num_inference_steps"]))
if input.seed is not None:
form.add_field("seed", str(input.seed))
elif extra_body.get("seed") is not None:
form.add_field("seed", str(extra_body["seed"]))
if extra_body.get("guidance_scale") is not None:
form.add_field("guidance_scale", str(extra_body["guidance_scale"]))
if extra_body.get("negative_prompt") is not None:
form.add_field("negative_prompt", str(extra_body["negative_prompt"]))
if extra_body.get("true_cfg_scale") is not None:
form.add_field("true_cfg_scale", str(extra_body["true_cfg_scale"]))
if extra_body.get("sys_type") is not None:
form.add_field("sys_type", str(extra_body["sys_type"]))
if extra_body.get("system_prompt") is not None:
form.add_field("system_prompt", str(extra_body["system_prompt"]))
bot_task = extra_body.get("bot_task")
if bot_task is None and input.default_bot_task is not None:
bot_task = input.default_bot_task
if bot_task is not None:
form.add_field("bot_task", str(bot_task))
assert input.image_paths is not None
for img_path in input.image_paths:
if not os.path.exists(img_path):
output.error = f"Image file not found: {img_path}"
output.success = False
if pbar:
pbar.update(1)
return output
with open(img_path, "rb") as img_f:
image_bytes = img_f.read()
form.add_field(
"image",
image_bytes,
filename=os.path.basename(img_path),
content_type=_guess_mime_type(img_path),
)
try:
async with session.post(edits_url, data=form) as response:
if response.status == 200:
resp_json = await response.json()
output.response_body = resp_json
output.success = True
else:
output.error = f"HTTP {response.status}: {await response.text()}"
output.success = False
except Exception as e:
output.error = str(e)
output.success = False
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
async def async_request_chat_completions(
input: RequestFuncInput,
session: aiohttp.ClientSession,
@@ -243,7 +151,7 @@ async def async_request_chat_completions(
return output
async def async_request_openai_image_generations(
async def async_request_openai_images(
input: RequestFuncInput,
session: aiohttp.ClientSession,
pbar: tqdm | None = None,
@@ -440,39 +348,12 @@ async def async_request_v1_videos(
return output
LEGACY_BACKEND_ENDPOINT_ALIASES = {
"vllm-omni": "/v1/chat/completions",
"openai": "/v1/images/generations",
}
def normalize_endpoint(value: str) -> str:
endpoint = str(value).strip()
if not endpoint:
raise ValueError("endpoint must not be empty.")
endpoint = LEGACY_BACKEND_ENDPOINT_ALIASES.get(
endpoint,
LEGACY_BACKEND_ENDPOINT_ALIASES.get(endpoint.lstrip("/"), endpoint),
)
if not endpoint.startswith("/"):
endpoint = f"/{endpoint}"
return endpoint
def endpoint_filename_token(value: str) -> str:
token = normalize_endpoint(value).lstrip("/")
for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"):
token = token.replace(bad, "_")
return token or "endpoint"
backends_function_mapping = {
"2i": {
"/v1/chat/completions": (async_request_chat_completions, "/v1/chat/completions"),
"/v1/images/generations": (async_request_openai_image_generations, "/v1/images/generations"),
"/v1/images/edits": (async_request_image_edits, "/v1/images/edits"),
"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"),
"v1/videos": (async_request_v1_videos, "/v1/videos"),
},
}
@@ -1,328 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Standalone diagnostic for the diffusion attention backends added in #3079.
Exercises the same synthetic attention shape through each backend and each
torch SDPA sub-kernel, with and without an attn_mask, so we can pinpoint why
CUDNN_ATTN underperforms on some SKUs. Two typical causes:
(a) cuDNN lacks a tuned kernel for this (SM, head_dim, seq) combination
and silently falls back to MATH.
(b) attn_mask is non-None and the cuDNN SDPA dispatch rejects it,
walking CUDNN -> FLASH -> MATH.
Run:
python benchmarks/diffusion/bench_attention_backends.py --preset hv15
python benchmarks/diffusion/bench_attention_backends.py --preset wan22
python benchmarks/diffusion/bench_attention_backends.py \
--batch 1 --heads 24 --seq 14336 --head-dim 128
Optional extras:
pip install --pre flash-attn-4 # FA4 is currently pre-release only
pip install -U flashinfer # latest FlashInfer (0.6.9)
The table at the end is the data we want on the PR — surface the row where a
backend is >1.5x the SDPA baseline, that's the one to gate off the auto-route.
"""
from __future__ import annotations
import argparse
import time
import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
# Presets mirror the DiT hot-path attention call (single-stream block) at the
# shapes we run in the PR's validation configs.
_PRESETS = {
# HunyuanVideo-1.5 480p/33f: (30 * 52 * 9) latent + ~256 text tokens.
"hv15": {"batch": 1, "heads": 24, "seq": 14336, "head_dim": 128},
# Wan 2.2 480p/33f rough estimate; adjust when we measure real shapes.
"wan22": {"batch": 1, "heads": 40, "seq": 16384, "head_dim": 128},
# Smaller image-gen shape for quick smoke tests.
"flux": {"batch": 1, "heads": 24, "seq": 4096, "head_dim": 128},
}
_SDPA_BACKENDS = [
("CUDNN_ATTENTION", [SDPBackend.CUDNN_ATTENTION]),
("FLASH_ATTENTION", [SDPBackend.FLASH_ATTENTION]),
("EFFICIENT_ATTENTION", [SDPBackend.EFFICIENT_ATTENTION]),
("MATH", [SDPBackend.MATH]),
# The PR's CUDNN_ATTN impl uses this priority chain; keep it here to show
# what torch actually picks when multiple are allowed.
("CUDNN_ATTN_CHAIN", [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]),
]
def _env_report() -> None:
if not torch.cuda.is_available():
raise SystemExit("CUDA not available; this diagnostic needs a GPU.")
cc = torch.cuda.get_device_capability()
print("=" * 72)
print(f"GPU : {torch.cuda.get_device_name()}")
print(f"SM capability : sm_{cc[0]}{cc[1]}")
print(f"torch : {torch.__version__}")
print(f"cuDNN version : {torch.backends.cudnn.version()}")
print("flashinfer : ", end="")
try:
import flashinfer
print(getattr(flashinfer, "__version__", "present"))
except Exception as e:
print(f"not installed ({type(e).__name__})")
print("=" * 72)
def _make_qkv(batch: int, heads: int, seq: int, head_dim: int, device: str, dtype: torch.dtype):
# Layout: (B, S, H, D) — what the PR's backends accept as input.
q = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
k = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
v = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
return q, k, v
def _make_mask(batch: int, seq: int, device: str, dtype: torch.dtype, pad_tokens: int = 128) -> torch.Tensor:
# Mirrors a text-encoder padding mask: last `pad_tokens` positions masked.
mask = torch.zeros(batch, 1, seq, seq, device=device, dtype=dtype)
if pad_tokens > 0:
mask[..., -pad_tokens:] = float("-inf")
return mask
def _time_call(fn, *args, warmup: int = 3, iters: int = 10) -> tuple[float, str]:
"""Return (median ms, error string). Skips on any exception so that a
backend that rejects our args (wrong dtype, missing JIT module, unsupported
kwarg value) doesn't abort the whole sweep."""
try:
for _ in range(warmup):
fn(*args)
torch.accelerator.synchronize()
times = []
for _ in range(iters):
torch.accelerator.synchronize()
t0 = time.perf_counter()
fn(*args)
torch.accelerator.synchronize()
times.append((time.perf_counter() - t0) * 1000.0)
times.sort()
return times[len(times) // 2], ""
except Exception as e: # noqa: BLE001 — probe script, keep the whole table going
msg = str(e).split("\n", 1)[0][:60]
return float("nan"), f"{type(e).__name__}: {msg}" if msg else type(e).__name__
def _run_sdpa_variants(q, k, v, attn_mask, scale: float) -> list[tuple[str, float, str]]:
# q/k/v come in (B, S, H, D); F.sdpa wants (B, H, S, D).
q_t = q.permute(0, 2, 1, 3)
k_t = k.permute(0, 2, 1, 3)
v_t = v.permute(0, 2, 1, 3)
rows: list[tuple[str, float, str]] = []
for name, backends in _SDPA_BACKENDS:
def _call(mask=attn_mask):
with sdpa_kernel(backends):
return torch.nn.functional.scaled_dot_product_attention(
q_t, k_t, v_t, attn_mask=mask, dropout_p=0.0, is_causal=False, scale=scale
)
ms, err = _time_call(_call)
rows.append((name, ms, err))
return rows
def _run_flashinfer(q, k, v, scale: float, backend: str | None = None, attn_mask=None) -> tuple[float, str]:
"""Call FlashInfer's dense single-prefill.
``backend`` hints at cutlass/fa3/trtllm if the installed version exposes
that kwarg. ``attn_mask`` is the same additive 4D mask we pass to SDPA —
we convert it to the 2D boolean ``custom_mask`` form FlashInfer accepts.
Per FlashInfer docs, ``custom_mask`` only applies when ``causal=False``.
"""
try:
import inspect
from flashinfer.prefill import single_prefill_with_kv_cache
except Exception as e:
return float("nan"), f"import-{type(e).__name__}"
kwargs: dict = {"sm_scale": scale, "causal": False, "return_lse": False}
if backend is not None:
sig = inspect.signature(single_prefill_with_kv_cache)
if "backend" not in sig.parameters:
return float("nan"), "no-backend-kwarg"
kwargs["backend"] = backend
if attn_mask is not None:
# Collapse (B, 1, S, S) additive float mask to (S, S) boolean.
# FlashInfer expects True = keep, False = masked out.
mask_2d = attn_mask[0, 0]
kwargs["custom_mask"] = mask_2d != float("-inf")
def _call():
out = single_prefill_with_kv_cache(q[0], k[0], v[0], **kwargs)
return out.unsqueeze(0)
return _time_call(_call)
def _run_fa4(q, k, v, scale: float) -> tuple[float, str]:
"""Call FlashAttention-4 directly (``pip install flash-attn-4``). FA4
ships a Blackwell-native kernel via CuTe-DSL; on sm_120 it should beat
cuDNN by ~20%. API lives under ``flash_attn.cute``; the older
``flash_attn.flash_attn_func`` path is FA2/FA3 only."""
try:
from flash_attn.cute import flash_attn_func
except Exception as e:
return float("nan"), f"import-{type(e).__name__}"
def _call():
# FA4 accepts (B, S, H, D) directly — same layout the PR's backends use.
return flash_attn_func(q, k, v, softmax_scale=scale, causal=False)
return _time_call(_call)
def _run_flashinfer_cudnn_batch(q, k, v, scale: float) -> tuple[float, str]:
"""Call FlashInfer's direct cuDNN wrapper. Bypasses PyTorch SDPA dispatch
(which has a few hundred ns of overhead per call) and talks to cuDNN FMHA
straight. Useful as a ceiling for 'pure cuDNN, no SDPA' on Blackwell."""
try:
from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
except Exception as e:
return float("nan"), f"import-{type(e).__name__}"
b, s, h, d = q.shape
qo_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
kv_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
def _call():
return cudnn_batch_prefill_with_kv_cache(
q.reshape(b * s, h, d),
k.reshape(b * s, h, d),
v.reshape(b * s, h, d),
qo_indptr=qo_indptr,
kv_indptr=kv_indptr,
max_qo_len=s,
max_kv_len=s,
sm_scale=scale,
causal=False,
)
return _time_call(_call)
def _print_table(title: str, rows: list[tuple[str, float, str]], baseline_name: str | None = None) -> None:
print(f"\n{title}")
print("-" * 88)
header = f"{'backend':<24} {'median (ms)':>14} {'vs baseline':>12} status"
print(header)
print("-" * 88)
baseline_ms = None
if baseline_name is not None:
for name, ms, _ in rows:
if name == baseline_name and ms == ms: # not NaN
baseline_ms = ms
break
for name, ms, err in rows:
ms_str = f"{ms:>14.3f}" if ms == ms else f"{'n/a':>14}"
if baseline_ms is not None and ms == ms:
ratio = baseline_ms / ms
ratio_str = f"{ratio:>11.2f}x"
else:
ratio_str = f"{'—':>12}"
status = "ok" if not err else f"FAILED ({err})"
print(f"{name:<24} {ms_str} {ratio_str} {status}")
def _pick_winner(rows: list[tuple[str, float, str]]) -> tuple[str, float] | None:
"""Return (backend, ms) of the fastest non-failing row, or None."""
ok_rows = [(n, m) for n, m, err in rows if not err and m == m]
if not ok_rows:
return None
return min(ok_rows, key=lambda x: x[1])
def _bench_one_shape(shape: dict, dtype: torch.dtype, device: str) -> tuple[list, list]:
torch.manual_seed(0)
q, k, v = _make_qkv(**shape, device=device, dtype=dtype)
scale = 1.0 / (shape["head_dim"] ** 0.5)
rows_nomask = _run_sdpa_variants(q, k, v, attn_mask=None, scale=scale)
rows_nomask.append(("FLASHINFER (default)", *_run_flashinfer(q, k, v, scale)))
# `trtllm-gen` has no sm_120 cubins (NVIDIA/TensorRT-LLM#11799) — skipped.
for fi_backend in ("fa2", "fa3", "cutlass", "auto"):
rows_nomask.append((f"FLASHINFER ({fi_backend})", *_run_flashinfer(q, k, v, scale, backend=fi_backend)))
rows_nomask.append(("FLASHINFER (cudnn-batch)", *_run_flashinfer_cudnn_batch(q, k, v, scale)))
rows_nomask.append(("FA4 (direct)", *_run_fa4(q, k, v, scale)))
mask = _make_mask(shape["batch"], shape["seq"], device=device, dtype=dtype, pad_tokens=256)
rows_mask = _run_sdpa_variants(q, k, v, attn_mask=mask, scale=scale)
rows_mask.append(("FLASHINFER (dense)", *_run_flashinfer(q, k, v, scale, attn_mask=mask)))
rows_mask.append(("FA4 (direct)", float("nan"), "mask-not-supported"))
return rows_nomask, rows_mask
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--preset", choices=list(_PRESETS.keys()), default="hv15")
parser.add_argument("--sweep", action="store_true", help="Run all presets and print a ranking")
parser.add_argument("--batch", type=int, default=None)
parser.add_argument("--heads", type=int, default=None)
parser.add_argument("--seq", type=int, default=None)
parser.add_argument("--head-dim", type=int, default=None)
parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16")
parser.add_argument("--device", default="cuda:0")
args = parser.parse_args()
dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
_env_report()
print(f"dtype : {args.dtype}")
presets = list(_PRESETS.keys()) if args.sweep else [args.preset]
summary: dict[str, tuple[str, float] | None] = {}
for preset_name in presets:
shape = dict(_PRESETS[preset_name])
if not args.sweep:
for k in ("batch", "heads", "seq", "head_dim"):
v = getattr(args, k if k != "head_dim" else "head_dim")
if v is not None:
shape[k] = v
print("\n" + "=" * 88)
print(
f"Preset: {preset_name} | batch={shape['batch']} heads={shape['heads']} "
f"seq={shape['seq']} head_dim={shape['head_dim']}"
)
print("=" * 88)
rows_nomask, rows_mask = _bench_one_shape(shape, dtype, args.device)
_print_table(f"[{preset_name}] No attention mask", rows_nomask, baseline_name="CUDNN_ATTENTION")
_print_table(
f"[{preset_name}] With attention mask (pad 256 tokens)", rows_mask, baseline_name="CUDNN_ATTENTION"
)
summary[preset_name] = _pick_winner(rows_nomask)
if args.sweep or len(presets) > 1:
print("\n" + "=" * 88)
print("Winners per preset (no-mask path)")
print("=" * 88)
for preset_name, winner in summary.items():
if winner is None:
print(f" {preset_name:<10} — no successful backend")
else:
name, ms = winner
print(f" {preset_name:<10} {name:<24} {ms:>8.3f} ms")
print("\nNotes:")
print(" * Ratios are relative to CUDNN_ATTENTION. >1.0x means faster than cuDNN.")
print(" * Mask-path winner inherits CUDNN_ATTN's fallback in the PR's backends.")
print(" * `trtllm-gen` and FA4 4.0.0b10 are known-broken on sm_120 as of Apr 2026.")
print(" * For e2e timings, run text_to_video.py / text_to_image.py while varying DIFFUSION_ATTENTION_BACKEND.")
if __name__ == "__main__":
main()
@@ -6,34 +6,31 @@
Benchmark online serving for diffusion models (Image/Video Generation).
If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
Supports multiple endpoints:
- /v1/chat/completions: OpenAI chat-compatible image requests (e.g. t2i, Qwen i2i)
- /v1/images/edits: OpenAI image edit / IT2I (multipart; e.g. Hunyuan --bot-task think)
- /v1/images/generations: OpenAI image generation requests
- /v1/videos: Async video jobs
Legacy --backend vllm-omni and openai are aliases for chat/completions and images/generations.
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 (/v1/videos endpoint)
# Video (v1/videos backend)
t2v:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--endpoint /v1/videos --dataset vbench --task t2v --num-prompts 10 \
--backend v1/videos --dataset vbench --task t2v --num-prompts 10 \
--height 480 --width 640 --fps 16 --num-frames 80
i2v:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--endpoint /v1/videos --dataset vbench --task i2v --num-prompts 10
--backend v1/videos --dataset vbench --task i2v --num-prompts 10
# Image (/v1/chat/completions endpoint)
# Image (vllm-omni backend)
t2i:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--endpoint /v1/chat/completions --dataset vbench --task t2i --num-prompts 10 \
--backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
--height 1024 --width 1024
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--endpoint /v1/chat/completions --dataset random --task t2i --num-prompts 1 \
--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},
@@ -42,32 +39,20 @@ Usage:
{"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
]'
ti2i (Hunyuan / OpenAI image edit API):
i2i:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--endpoint /v1/images/edits --dataset random --task ti2i --num-prompts 10 \
--bot-task think
--backend vllm-omni --dataset vbench --task i2i --num-prompts 10
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--endpoint /v1/images/edits \
--dataset custom \
--dataset-path custom_requests.jsonl \
--task ti2i \
--bot-task think
i2i (chat-based models such as Qwen-Image-Edit):
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--endpoint /v1/chat/completions --dataset vbench --task i2i --num-prompts 10
# Image (/v1/images/generations endpoint)
# Image (openai backend)
t2i:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--endpoint /v1/images/generations --dataset vbench --task t2i --num-prompts 10 \
--backend openai --dataset vbench --task t2i --num-prompts 10 \
--height 1024 --width 1024 --port 3000
# Video (v1/videos)
t2v:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--endpoint /v1/videos --dataset random --task t2v --num-prompts 1 \
--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}
@@ -95,12 +80,7 @@ from typing import Any
import aiohttp
import numpy as np
import requests
from backends import (
RequestFuncInput,
RequestFuncOutput,
backends_function_mapping,
normalize_endpoint,
)
from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
from PIL import Image
from tqdm.asyncio import tqdm
@@ -147,7 +127,7 @@ class VBenchDataset(BaseDataset):
def _load_data(self) -> list[dict[str, Any]]:
if self.args.task == "t2v":
return self._load_t2v_prompts()
elif self.args.task in ["i2v", "ti2v", "ti2i", "i2i", "it2i"]:
elif self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
return self._load_i2v_data()
else:
return self._load_t2v_prompts()
@@ -539,7 +519,7 @@ class TraceDataset(BaseDataset):
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", "it2i"]:
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."
)
@@ -573,135 +553,6 @@ class TraceDataset(BaseDataset):
return [self[i] for i in range(len(self))]
class CustomDataset(BaseDataset):
"""
Custom dataset that loads requests from a JSONL file.
Each line in the JSONL file should be a JSON object with the following fields:
- prompt (required): The text prompt for the request
- width (optional): Image/video width
- height (optional): Image height
- num_inference_steps (optional): Number of diffusion steps
- seed (optional): Random seed
- image_paths (optional): List of input image paths for i2i/i2v/ti2i tasks
- image_urls (optional): List of input image URLs (alternative to image_paths)
Example JSONL for ti2i:
{"prompt": "Add sunset lighting", "width": 1024, "height": 1024, "image_urls": ["https://example.com/image.jpg"]}
"""
def __init__(self, args, api_url: str, model: str):
super().__init__(args, api_url, model)
self.dataset_path = args.dataset_path
if not self.dataset_path:
raise ValueError("--dataset-path must be provided when using 'custom' dataset")
self.load_data()
def load_data(self) -> None:
"""Load data from JSONL file."""
import pandas as pd
if not self.dataset_path.endswith(".jsonl"):
raise ValueError("Custom dataset must be a JSONL file")
try:
df = pd.read_json(path_or_buf=self.dataset_path, lines=True)
except Exception as e:
raise ValueError(f"Failed to load JSONL file {self.dataset_path}: {e}")
if "prompt" not in df.columns:
raise ValueError("JSONL file must contain a 'prompt' column")
self.data = df.to_dict("records")
print(f"Loaded {len(self.data)} requests from {self.dataset_path}")
def _resolve_image_paths(self, item: dict) -> list[str] | None:
# Handle local file paths
if "image_paths" in item and item["image_paths"]:
paths = item["image_paths"]
if isinstance(paths, str):
paths = [paths]
# Verify all paths exist
valid_paths = []
for path in paths:
if os.path.exists(path):
valid_paths.append(path)
else:
raise ValueError(f"Image file not found: {path}")
return valid_paths if valid_paths else None
# Handle URLs - download to temp files
if "image_urls" in item and item["image_urls"]:
urls = item["image_urls"]
if isinstance(urls, str):
urls = [urls]
downloaded_paths = []
for url in urls:
try:
response = requests.get(url, timeout=30)
response.raise_for_status()
# Create temp file with appropriate extension
suffix = os.path.splitext(url)[1] or ".png"
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
tmp.write(response.content)
downloaded_paths.append(tmp.name)
except Exception as e:
raise ValueError(f"Failed to download image from {url}: {e}")
return downloaded_paths if downloaded_paths else None
return None
def __len__(self) -> int:
return len(self.data)
def __getitem__(self, idx: int) -> RequestFuncInput:
item = self.data[idx]
# Extract fields with fallback to args defaults
prompt = item["prompt"]
width = item.get("width", self.args.width)
height = item.get("height", self.args.height)
num_inference_steps = item.get("num_inference_steps", self.args.num_inference_steps)
seed = item.get("seed", self.args.seed)
reserved_keys = {"prompt", "width", "height", "num_inference_steps", "seed", "image_paths", "image_urls"}
extra_body = {k: v for k, v in item.items() if k not in reserved_keys}
# Handle image paths/URLs
image_paths = self._resolve_image_paths(item)
return RequestFuncInput(
prompt=prompt,
api_url=self.api_url,
model=self.model,
seed=seed,
image_paths=image_paths,
extra_body=extra_body,
width=width,
height=height,
num_inference_steps=num_inference_steps,
)
def get_requests(self) -> list[RequestFuncInput]:
"""Get all requests, cycling through data if num_prompts > dataset size."""
num_requests = getattr(self.args, "num_prompts", len(self.data))
if num_requests <= len(self.data):
return [self[i] for i in range(num_requests)]
# Cycle through the dataset to reach num_prompts
requests = []
for i in range(num_requests):
idx = i % len(self.data)
requests.append(self[idx])
return requests
class RandomDataset(BaseDataset):
def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool = False):
super().__init__(args, api_url, model)
@@ -729,7 +580,7 @@ class RandomDataset(BaseDataset):
self._sampled_requests = None
# Random image generate
if self.args.task in ["i2v", "ti2v", "ti2i", "i2i", "it2i"]:
if self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
self._random_image_path = self._generate_random_image_paths()
else:
self._random_image_path = None
@@ -1032,23 +883,13 @@ def wait_for_service(base_url: str, timeout: int = 120) -> None:
time.sleep(1)
def _default_endpoint_for_task(task: str) -> str:
if task in {"t2v", "i2v", "ti2v"}:
return "/v1/videos"
if task in {"i2i", "ti2i", "it2i"}:
return "/v1/images/edits"
if task == "t2i":
return "/v1/chat/completions"
raise ValueError(f"Unsupported task for endpoint resolution: {task}")
async def benchmark(args):
# Construct base_url if not provided
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", "it2i"}
IMAGE_TASKS = {"t2i", "i2i", "ti2i"}
if args.task in VIDEO_TASKS:
task_type = "2v"
@@ -1061,23 +902,18 @@ async def benchmark(args):
f"Valid image tasks: {sorted(IMAGE_TASKS)}"
)
raw_endpoint = args.endpoint if args.endpoint is not None else args.backend
if raw_endpoint is None:
raw_endpoint = _default_endpoint_for_task(args.task)
args.endpoint = normalize_endpoint(raw_endpoint)
valid_backends = sorted(backends_function_mapping[task_type].keys())
valid_endpoints = sorted(backends_function_mapping[task_type].keys())
if args.endpoint not in valid_endpoints:
if args.backend not in valid_backends:
logger.error(
f"Invalid endpoint '{args.endpoint}' for task '{args.task}' (task type: '{task_type}').\n"
f"Valid endpoints for this task type: {valid_endpoints}\n"
f"Example usage: --task {args.task} --endpoint {valid_endpoints[0]}"
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("Endpoint validation failed. See log above for valid options.")
raise ValueError("Backend validation failed. See log above for valid options.")
# Setup API URL and request function based on endpoint.
request_func, api_url = backends_function_mapping[task_type][args.endpoint]
# Setup API URL and request function based on backend
request_func, api_url = backends_function_mapping[task_type][args.backend]
api_url = f"{args.base_url}{api_url}"
if args.dataset == "vbench":
@@ -1086,8 +922,6 @@ async def benchmark(args):
dataset = TraceDataset(args, api_url, args.model)
elif args.dataset == "random":
dataset = RandomDataset(args, api_url, args.model, args.enable_negative_prompt)
elif args.dataset == "custom":
dataset = CustomDataset(args, api_url, args.model)
else:
raise ValueError(f"Unknown dataset: {args.dataset}")
@@ -1095,10 +929,6 @@ async def benchmark(args):
requests_list = dataset.get_requests()
print(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.")
if args.endpoint == "/v1/images/edits":
for req in requests_list:
req.default_bot_task = args.bot_task
# Limit concurrency
if args.max_concurrency is not None:
semaphore = asyncio.Semaphore(args.max_concurrency)
@@ -1146,17 +976,15 @@ async def benchmark(args):
metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo)
# Add configuration info to metrics for JSON output
metrics["endpoint"] = args.endpoint
metrics["backend"] = args.backend
metrics["model"] = args.model
metrics["dataset"] = args.dataset
metrics["task"] = args.task
if args.endpoint == "/v1/images/edits":
metrics["bot_task"] = args.bot_task
print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))
print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="="))
# Section 1: Configuration
print("{:<40} {:<15}".format("Endpoint:", args.endpoint))
print("{:<40} {:<15}".format("Backend:", args.backend))
print("{:<40} {:<15}".format("Model:", args.model))
print("{:<40} {:<15}".format("Dataset:", args.dataset))
print("{:<40} {:<15}".format("Task:", args.task))
@@ -1219,30 +1047,25 @@ if __name__ == "__main__":
parser.add_argument("--host", type=str, default="localhost", help="Server host.")
parser.add_argument("--port", type=int, default=8091, help="Server port.")
parser.add_argument("--model", type=str, default="default", help="Model name.")
parser.add_argument(
"--endpoint",
type=str,
default=None,
help=("Endpoint path to target. Leading '/' is optional, e.g. /v1/videos or v1/videos."),
)
parser.add_argument(
"--backend",
type=str,
default=None,
help=argparse.SUPPRESS,
default="vllm-omni",
choices=["vllm-omni", "openai", "v1/videos"],
help="Backend to target the benchmark to.",
)
parser.add_argument(
"--dataset",
type=str,
default="vbench",
choices=["vbench", "trace", "random", "custom"],
choices=["vbench", "trace", "random"],
help="Dataset to use.",
)
parser.add_argument(
"--task",
type=str,
default="t2v",
choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "it2i", "t2i"],
choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "t2i"],
help="Task type.",
)
parser.add_argument(
@@ -1349,15 +1172,9 @@ if __name__ == "__main__":
default=1,
help=(
"Number of synthetic input images to attach for image-conditioned tasks "
"(i2v, ti2v, ti2i, i2i, it2i) when using random dataset."
"(i2v, ti2v, ti2i, i2i) when using random dataset."
),
)
parser.add_argument(
"--bot-task",
type=str,
default="think",
help=("bot_task form field for --endpoint /v1/images/edits (think, recaption, think_recaption, vanilla)."),
)
args = parser.parse_args()
asyncio.run(benchmark(args))
@@ -58,7 +58,7 @@ Record these parameters when reporting performance results.
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--endpoint /v1/chat/completions \
--backend vllm-omni \
--dataset <DATASET_NAME> \
--task t2i \
--num-prompts <N> \
@@ -71,7 +71,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
| Parameter | Description |
| ---------------------- | --------------------------------- |
| `--endpoint` | API endpoint (use `/v1/chat/completions`; leading `/` is optional) |
| `--backend` | Serving backend (use `vllm-omni`) |
| `--dataset` | Dataset name (`random` or custom) |
| `--task` | Task type (e.g., `t2i`) |
| `--num-prompts` | Total number of requests |
@@ -129,7 +129,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--endpoint /v1/chat/completions \
--backend vllm-omni \
--dataset random \
--task t2i \
--num-prompts 1 \
@@ -53,7 +53,7 @@ Record these parameters when reporting performance results.
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--endpoint /v1/videos \
--backend v1/videos \
--dataset <DATASET_NAME> \
--task t2v \
--num-prompts <N> \
@@ -66,7 +66,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
| Parameter | Description |
| ---------------------- | --------------------------------- |
| `--endpoint` | API endpoint (use `/v1/videos`; leading `/` is optional) |
| `--backend` | Serving backend (use `v1/videos`) |
| `--dataset` | Dataset name (`random` or custom) |
| `--task` | Task type (e.g., `t2v`) |
| `--num-prompts` | Total number of requests |
@@ -120,7 +120,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--endpoint /v1/videos \
--backend v1/videos \
--dataset random \
--task t2v \
--num-prompts 1 \
+8 -35
View File
@@ -34,16 +34,6 @@ Video example (text-to-video):
--height 720 --width 1280 \
--num-frames 81 --num-inference-steps 40 --seed 42
LTX-2 example (text-to-video; audio output is generated but not scored by LPIPS):
python benchmarks/diffusion/quantization_quality.py \
--model Lightricks/LTX-Video-2 \
--task t2v \
--quantization fp8 int8 \
--prompts \
"A serene lakeside sunrise with mist over the water" \
--height 704 --width 1216 \
--num-frames 121 --num-inference-steps 40 --seed 42
Multiple quantization methods:
python benchmarks/diffusion/quantization_quality.py \
--model Tongyi-MAI/Z-Image-Turbo \
@@ -158,7 +148,6 @@ def _build_omni_kwargs(args, quantization=None):
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.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
@@ -176,9 +165,8 @@ def _generate_image(omni, args, prompt, seed):
elapsed = time.perf_counter() - start
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
req_out = OmniRequestOutput.unwrap_result(outputs)
if not req_out.images:
raise ValueError("Could not extract image output from result.")
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
@@ -214,19 +202,10 @@ def _generate_video(omni, args, prompt, seed):
else:
frames = inner
elif hasattr(first, "images") and first.images:
frames = first.images[0]
frames = first.images
else:
raise ValueError("Could not extract video frames from output.")
# LTX-2 (and similar audio+video models) may surface a dict or (video, audio) tuple
if isinstance(frames, dict):
frames = frames.get("video") or frames.get("frames")
elif isinstance(frames, tuple) and len(frames) == 2:
frames = frames[0]
if frames is None:
raise ValueError("Could not extract video frames from output.")
if isinstance(frames, torch.Tensor):
video = frames.detach().cpu()
if video.dim() == 5:
@@ -244,11 +223,9 @@ def _generate_video(omni, args, prompt, seed):
return frames_array, elapsed, peak_mem
def _free_gpu_memory():
"""Force GC and release cached GPU memory.
Must be called AFTER the caller has dropped (i.e., via `del`)
every reference to the Omni instance"""
def _unload_omni(omni):
"""Delete Omni instance and free GPU memory."""
del omni
gc.collect()
if torch.cuda.is_available():
torch.accelerator.empty_cache()
@@ -288,9 +265,7 @@ def run_benchmark(args):
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
omni_bl.shutdown()
del omni_bl
_free_gpu_memory()
_unload_omni(omni_bl)
# Save baseline outputs
bl_dir = output_dir / "baseline"
@@ -330,9 +305,7 @@ def run_benchmark(args):
qt_avg_time = np.mean([v[1] for v in qt_outputs.values()])
qt_mem = qt_outputs[prompts[0]][2]
omni_qt.shutdown()
del omni_qt
_free_gpu_memory()
_unload_omni(omni_qt)
# Save quantized outputs
qt_dir = output_dir / config_label.replace(" ", "_")
+4 -21
View File
@@ -40,7 +40,7 @@ vllm bench serve --omni \
--num-prompts 20 --num-warmups 2 \
--extra-body '{"task_type":"Base"}' \
--max-concurrency 1 --request-rate inf \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
--save-result --result-dir ./results
```
@@ -58,7 +58,7 @@ vllm bench serve --omni \
--num-prompts 20 --num-warmups 2 \
--extra-body '{"voice":"Vivian","language":"English","task_type":"CustomVoice"}' \
--max-concurrency 1 --request-rate inf \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
--save-result --result-dir ./results
```
@@ -76,28 +76,10 @@ vllm bench serve --omni \
--num-prompts 20 --num-warmups 2 \
--extra-body '{"task_type":"VoiceDesign","language":"English"}' \
--max-concurrency 1 --request-rate inf \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
--save-result --result-dir ./results
```
#### Streaming continuity (`audio_underrun`)
`audio_underrun` (seconds) is the per-request worst-case buffer deficit
under a realtime-rate player simulation. The audio-speech backend records
chunk arrival times during the SSE stream and surfaces:
- `Mean / Median / P{50,99} AUDIO_UNDERRUN (s)` per the standard percentile
output - any positive value means at least one inter-chunk gap was longer
than the chunk's audio duration.
- `Streaming continuity OK rate` - fraction of requests whose worst-case
underrun stayed under the threshold (default 100 ms, the commonly cited
"audible gap" budget). Override the threshold with
`VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S=<float>`.
This captures the failure mode where `RTF_p50 < 1` (server keeps up in
aggregate) but per-stream chunk arrival is bursty enough that listeners
still hear gaps - common at high concurrency on streaming TTS pipelines.
#### Add WER / SIM / UTMOS to any of the above
Append `--seed-tts-wer-eval` (and optionally `SEED_TTS_EVAL_DEVICE=cuda:0`
@@ -233,6 +215,7 @@ sentinel for regressions in this area.
benchmarks/tts/
├── README.md (this file)
├── bench_tts.py CLI — serve-mode benchmark driver
├── bench_voxcpm_offline.py CLI — offline VoxCPM benchmark (sync + streaming)
├── plot_results.py Generate per-task / per-concurrency curves
└── model_configs.yaml Model registry (supported tasks + extra body)
```
+1 -1
View File
@@ -118,7 +118,7 @@ def build_bench_args(
"--num-warmups",
"2",
"--percentile-metrics",
"ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun",
"ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
]
if resolved_dataset_path:
+922
View File
@@ -0,0 +1,922 @@
"""Offline VoxCPM benchmark for vLLM Omni.
Supports both:
- sync one-shot (Omni.generate)
- streaming (AsyncOmni.generate with async_chunk config)
- text-only synthesis
- voice cloning
- text/clone batch inputs from txt or jsonl
Usage::
# Sync (default voice)
python benchmarks/tts/bench_voxcpm_offline.py \\
--model /path/to/VoxCPM \\
--text "Hello world" \\
--output-dir results/audio/
# Streaming (async_chunk)
python benchmarks/tts/bench_voxcpm_offline.py \\
--model /path/to/VoxCPM \\
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \\
--txt-prompts prompts.txt \\
--output-dir results/audio/
# Voice cloning batch via JSONL
python benchmarks/tts/bench_voxcpm_offline.py \\
--model /path/to/VoxCPM \\
--jsonl-prompts prompts.jsonl \\
--output-dir results/audio/
"""
from __future__ import annotations
import asyncio
import json
import logging
import os
import tempfile
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import torch
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import AsyncOmni, Omni
def _find_repo_root(start: Path) -> Path:
"""Walk up from ``start`` until a repo marker is found.
Falls back to ``parents[2]`` for backwards compatibility if no marker hits
(which can only happen in unusual checkouts — the tree should always have
pyproject.toml + vllm_omni/ at the top level).
"""
for candidate in [start, *start.parents]:
if (candidate / "pyproject.toml").is_file() and (candidate / "vllm_omni").is_dir():
return candidate
return start.parents[2]
REPO_ROOT = _find_repo_root(Path(__file__).resolve())
DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml"
DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml"
logger = logging.getLogger(__name__)
@dataclass(frozen=True, slots=True)
class PromptSpec:
text: str
label: str
ref_audio: str | None = None
ref_text: str | None = None
def _require_soundfile():
try:
import soundfile as sf # type: ignore
except ModuleNotFoundError as exc:
raise RuntimeError(
"soundfile is required to write VoxCPM benchmark WAV outputs. Install it with: pip install soundfile"
) from exc
return sf
def _build_prompt(
args,
*,
text: str,
ref_audio: str | None = None,
ref_text: str | None = None,
global_request_id: str | None = None,
) -> dict[str, Any]:
additional_information: dict[str, list[Any]] = {
"text": [text],
"cfg_value": [args.cfg_value],
"inference_timesteps": [args.inference_timesteps],
"min_len": [args.min_len],
"max_new_tokens": [args.max_new_tokens],
}
if args.streaming_prefix_len is not None:
additional_information["streaming_prefix_len"] = [args.streaming_prefix_len]
if ref_audio:
additional_information["ref_audio"] = [ref_audio]
if ref_text:
additional_information["ref_text"] = [ref_text]
if global_request_id is not None:
additional_information["global_request_id"] = [global_request_id]
return {
"prompt_token_ids": [1],
"additional_information": additional_information,
}
def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor:
audio = mm.get("audio", mm.get("model_outputs"))
if audio is None:
raise ValueError("No audio output found in multimodal output.")
if isinstance(audio, list):
parts = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio]
audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0)
if not isinstance(audio, torch.Tensor):
audio = torch.as_tensor(audio)
return audio.float().cpu().reshape(-1)
def _extract_sample_rate(mm: dict[str, Any]) -> int:
sr_raw = mm.get("sr", 24000)
if isinstance(sr_raw, list) and sr_raw:
sr_raw = sr_raw[-1]
if hasattr(sr_raw, "item"):
return int(sr_raw.item())
return int(sr_raw)
def _emit_offline_metrics(
*,
request_id: str,
elapsed_s: float,
first_audio_elapsed: float | None,
audio_duration_s: float,
) -> None:
metrics = {
"request_id": request_id,
"ttfp_ms": round(first_audio_elapsed * 1000.0, 3) if first_audio_elapsed is not None else None,
"audio_duration_s": round(audio_duration_s, 6),
"rtf": round(elapsed_s / audio_duration_s, 6) if audio_duration_s > 0 else None,
}
print(f"[OfflineMetrics] {metrics}")
def _write_audio_tensor(output_path: Path, audio_tensor: Any, sample_rate: int) -> None:
sf = _require_soundfile()
if isinstance(audio_tensor, torch.Tensor):
audio_np = audio_tensor.float().cpu().clamp(-1.0, 1.0).numpy()
else:
audio_np = torch.as_tensor(audio_tensor).float().cpu().clamp(-1.0, 1.0).numpy()
sf.write(
output_path,
audio_np,
sample_rate,
format="WAV",
subtype="PCM_16",
)
def _save_wav(mm: dict[str, Any], output_dir: Path, request_id: str) -> Path:
output_dir.mkdir(parents=True, exist_ok=True)
output_path = output_dir / f"output_{request_id}.wav"
_write_audio_tensor(output_path, _extract_audio_tensor(mm), _extract_sample_rate(mm))
return output_path
def _iter_request_multimodal_outputs(request_output: Any):
outputs = getattr(request_output, "outputs", None)
if outputs:
for output in outputs:
mm = getattr(output, "multimodal_output", None)
if isinstance(mm, dict):
yield mm
mm = getattr(request_output, "multimodal_output", None)
if isinstance(mm, dict):
yield mm
def _read_non_empty_lines(path: str) -> list[str]:
with open(path, encoding="utf-8") as f:
return [line.strip() for line in f if line.strip()]
def _load_prompt_specs(args) -> list[PromptSpec]:
specs: list[PromptSpec] = []
if args.txt_prompts is not None:
texts = _read_non_empty_lines(args.txt_prompts)
if not texts:
raise ValueError(f"No prompts found in {args.txt_prompts}")
for idx, text in enumerate(texts, start=1):
specs.append(
PromptSpec(
text=text,
label=f"item{idx:03d}",
ref_audio=args.ref_audio,
ref_text=args.ref_text,
)
)
return specs
if args.jsonl_prompts is not None:
with open(args.jsonl_prompts, encoding="utf-8") as f:
for line_no, raw_line in enumerate(f, start=1):
line = raw_line.strip()
if not line:
continue
try:
item = json.loads(line)
except json.JSONDecodeError as exc:
raise ValueError(f"{args.jsonl_prompts}:{line_no} is not valid JSON: {exc}") from exc
if not isinstance(item, dict):
raise ValueError(f"{args.jsonl_prompts}:{line_no} must be a JSON object")
text = item.get("text")
if not isinstance(text, str) or not text.strip():
raise ValueError(f"{args.jsonl_prompts}:{line_no} requires non-empty string field 'text'")
ref_audio = item.get("ref_audio", args.ref_audio)
ref_text = item.get("ref_text", args.ref_text)
if (ref_audio is None) != (ref_text is None):
raise ValueError(
f"{args.jsonl_prompts}:{line_no} must provide both 'ref_audio' and 'ref_text' together"
)
specs.append(
PromptSpec(
text=text.strip(),
label=f"item{len(specs) + 1:03d}",
ref_audio=ref_audio,
ref_text=ref_text,
)
)
if not specs:
raise ValueError(f"No prompts found in {args.jsonl_prompts}")
return specs
specs.append(
PromptSpec(
text=args.text,
label="item001",
ref_audio=args.ref_audio,
ref_text=args.ref_text,
)
)
return specs
def _build_prompt_for_spec(args, spec: PromptSpec, *, global_request_id: str | None = None) -> dict[str, Any]:
return _build_prompt(
args,
text=spec.text,
ref_audio=spec.ref_audio,
ref_text=spec.ref_text,
global_request_id=global_request_id,
)
def _count_voice_clone_prompts(prompt_specs: list[PromptSpec]) -> int:
return sum(1 for spec in prompt_specs if spec.ref_audio is not None)
def _get_warmup_specs(prompt_specs: list[PromptSpec]) -> list[PromptSpec]:
return prompt_specs[:1]
def _extract_stream_finished(stage_output: Any) -> bool:
request_output = getattr(stage_output, "request_output", None)
request_finished = getattr(request_output, "finished", None)
if request_finished is not None:
return bool(request_finished)
return bool(getattr(stage_output, "finished", False))
def _build_profiled_stage_config(
stage_configs_path: str,
profiler_dir: str,
) -> str:
stage_config_path = Path(stage_configs_path)
yaml_text = stage_config_path.read_text(encoding="utf-8")
injected_lines: list[str] = []
injected_count = 0
for line in yaml_text.splitlines():
injected_lines.append(line)
if line.strip() != "engine_args:":
continue
indent = line[: len(line) - len(line.lstrip())]
child_indent = indent + " "
grandchild_indent = child_indent + " "
injected_lines.extend(
[
f"{child_indent}profiler_config:",
f'{grandchild_indent}profiler: "torch"',
f'{grandchild_indent}torch_profiler_dir: "{profiler_dir}"',
f"{grandchild_indent}torch_profiler_with_stack: true",
]
)
injected_count += 1
if injected_count == 0:
raise ValueError(f"No engine_args block found in stage config: {stage_configs_path}")
tmp = tempfile.NamedTemporaryFile(
mode="w",
encoding="utf-8",
delete=False,
suffix=".yaml",
prefix=f"{stage_config_path.stem}_profile_",
)
tmp.write("\n".join(injected_lines) + "\n")
tmp.close()
return tmp.name
def parse_args():
parser = FlexibleArgumentParser(
description="Offline split-stage VoxCPM inference with vLLM Omni (auto sync/streaming by stage config)"
)
parser.add_argument(
"--model",
type=str,
default=os.environ.get("VOXCPM_MODEL"),
help="Local VoxCPM model directory. Defaults to $VOXCPM_MODEL.",
)
parser.add_argument(
"--text",
type=str,
default="This is a split-stage VoxCPM synthesis example running on vLLM Omni.",
help="Text to synthesize. Ignored when --txt-prompts or --jsonl-prompts is used.",
)
parser.add_argument(
"--txt-prompts",
type=str,
default=None,
help="Path to a .txt file with one synthesis text per line.",
)
parser.add_argument(
"--jsonl-prompts",
type=str,
default=None,
help=(
"Path to a .jsonl file. Each line must contain at least {'text': ...}; "
"clone rows can also set ref_audio/ref_text, and ref_text must be the "
"real transcript of ref_audio."
),
)
parser.add_argument(
"--ref-audio",
type=str,
default=None,
help=(
"Optional reference audio path for voice cloning. With --txt-prompts, "
"the same reference is applied to every line."
),
)
parser.add_argument(
"--ref-text",
type=str,
default=None,
help=(
"Real transcript of the reference audio. Placeholder text or mismatched "
"text will usually produce noisy/electronic clone audio."
),
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=str(DEFAULT_STAGE_SYNC),
help="Stage config YAML path. Routing is selected only from this path.",
)
parser.add_argument(
"--cfg-value",
type=float,
default=2.0,
help="Classifier-free guidance value for VoxCPM.",
)
parser.add_argument(
"--inference-timesteps",
type=int,
default=10,
help="Number of inference timesteps.",
)
parser.add_argument(
"--min-len",
type=int,
default=2,
help="Minimum generated token length.",
)
parser.add_argument(
"--max-new-tokens",
type=int,
default=4096,
help="Maximum generated token length.",
)
parser.add_argument(
"--streaming-prefix-len",
type=int,
default=None,
help="VoxCPM streaming window (optional, streaming mode only).",
)
parser.add_argument(
"--output-dir",
type=str,
default=None,
help="Directory for output WAV files.",
)
parser.add_argument(
"--stage-init-timeout",
type=int,
default=600,
help="Stage initialization timeout in seconds.",
)
parser.add_argument(
"--log-stats",
dest="log_stats",
action="store_true",
help="Enable vLLM Omni stats logging.",
)
parser.add_argument(
"--no-log-stats",
dest="log_stats",
action="store_false",
help="Disable vLLM Omni stats logging.",
)
parser.set_defaults(log_stats=True)
parser.add_argument(
"--num-runs",
type=int,
default=1,
help="Number of full inference runs (same prompt each time). Default 1.",
)
parser.add_argument(
"--warmup-runs",
type=int,
default=0,
help=(
"Optional number of warmup passes before measured runs. Warmup uses only "
"the first prompt and does not save outputs."
),
)
parser.add_argument(
"--enable-profiler",
action="store_true",
help=(
"Enable torch profiler for the configured stages. A temporary profiled "
"stage config is generated automatically."
),
)
parser.add_argument(
"--profiler-dir",
type=str,
default=None,
help="Directory for profiler traces. Defaults to <output-dir>/profiler when profiling is enabled.",
)
parser.add_argument(
"--profiler-stages",
type=int,
nargs="*",
default=None,
help="Optional stage ids to profile. Defaults to all stages that have profiler_config.",
)
parser.add_argument(
"--profiler-wait-seconds",
type=float,
default=30.0,
help="Seconds to wait after stop_profile for trace files to flush.",
)
args = parser.parse_args()
if not args.model:
parser.error("--model is required unless $VOXCPM_MODEL is set")
if args.txt_prompts is not None and args.jsonl_prompts is not None:
parser.error("--txt-prompts and --jsonl-prompts are mutually exclusive")
if (args.ref_audio is None) != (args.ref_text is None):
parser.error("--ref-audio and --ref-text must be provided together")
if args.num_runs < 1:
parser.error("--num-runs must be >= 1")
if args.warmup_runs < 0:
parser.error("--warmup-runs must be >= 0")
if args.output_dir is None:
args.output_dir = (
"output_audio_streaming" if _is_streaming_stage_config(args.stage_configs_path) else "output_audio"
)
if args.enable_profiler and args.profiler_dir is None:
args.profiler_dir = str(Path(args.output_dir) / "profiler")
try:
args.prompt_specs = _load_prompt_specs(args)
except ValueError as exc:
parser.error(str(exc))
return args
def _is_streaming_stage_config(stage_configs_path: str) -> bool:
cfg_name = Path(stage_configs_path).name.lower()
return "async_chunk" in cfg_name
async def _collect_streaming_audio(
omni: AsyncOmni,
args: Any,
spec: PromptSpec,
request_id: str,
*,
phase_label: str,
prompt_index: int,
prompt_count: int,
print_prompt: bool = False,
) -> tuple[torch.Tensor, int, float, float | None]:
prompt = _build_prompt_for_spec(args, spec, global_request_id=request_id)
delta_chunks: list[torch.Tensor] = []
sample_rate = 24000
chunk_i = 0
prev_total_samples = 0
t_start = time.perf_counter()
first_audio_elapsed: float | None = None
if print_prompt:
print(f"---prompt---:{prompt}")
async for stage_output in omni.generate(prompt, request_id=request_id):
mm = getattr(stage_output, "multimodal_output", None)
if not isinstance(mm, dict):
ro = getattr(stage_output, "request_output", None)
if ro is None:
continue
mm = getattr(ro, "multimodal_output", None)
if not isinstance(mm, dict) and getattr(ro, "outputs", None):
seq = ro.outputs[0]
mm = getattr(seq, "multimodal_output", None)
if not isinstance(mm, dict):
continue
sample_rate = _extract_sample_rate(mm)
try:
w = _extract_audio_tensor(mm)
n = int(w.numel())
if n == 0:
continue
finished = _extract_stream_finished(stage_output)
if n > prev_total_samples:
delta = w.reshape(-1)[prev_total_samples:]
prev_total_samples = n
elif finished and n == prev_total_samples:
delta = w.reshape(-1)[:0]
else:
delta = w.reshape(-1)
prev_total_samples += int(delta.numel())
if int(delta.numel()) > 0:
delta_chunks.append(delta)
if first_audio_elapsed is None and int(delta.numel()) > 0:
first_audio_elapsed = time.perf_counter() - t_start
logger.info(
"%s prompt=%d/%d chunk=%d delta_samples=%d buf_len=%d finished=%s",
phase_label,
prompt_index + 1,
prompt_count,
chunk_i,
int(delta.numel()),
n,
finished,
)
chunk_i += 1
except ValueError:
if not _extract_stream_finished(stage_output):
logger.debug("skip non-audio partial output chunk=%d", chunk_i)
if not delta_chunks:
raise RuntimeError("No audio chunks received; check stage config and logs.")
audio_cat = torch.cat([c.reshape(-1) for c in delta_chunks], dim=0)
elapsed = time.perf_counter() - t_start
return audio_cat, sample_rate, elapsed, first_audio_elapsed
async def _abort_streaming_residual_work(
omni: AsyncOmni,
request_id: str,
*,
settle_seconds: float = 0.1,
) -> None:
"""Stop any late stage-0 work once the final audio has been collected."""
await omni.engine.abort_async([request_id])
if settle_seconds > 0:
await asyncio.sleep(settle_seconds)
async def _run_streaming_single(
omni: AsyncOmni,
args: Any,
spec: PromptSpec,
output_dir: Path,
request_id: str,
*,
run_index: int,
num_runs: int,
prompt_index: int,
prompt_count: int,
) -> Path:
audio_cat, sample_rate, elapsed, first_audio_elapsed = await _collect_streaming_audio(
omni,
args,
spec,
request_id,
phase_label=f"run={run_index + 1}/{num_runs}",
prompt_index=prompt_index,
prompt_count=prompt_count,
print_prompt=(run_index == 0 and prompt_index == 0),
)
await _abort_streaming_residual_work(omni, request_id)
output_path = output_dir / f"output_run{run_index + 1}_{spec.label}.wav"
_write_audio_tensor(output_path, audio_cat, sample_rate)
audio_duration_s = float(audio_cat.numel()) / float(sample_rate) if sample_rate > 0 else 0.0
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
rtf_text = f", rtf={elapsed / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
print(
f"Saved (streaming) run {run_index + 1}/{num_runs}, "
f"prompt {prompt_index + 1}/{prompt_count}: {output_path} ({elapsed:.2f}s{ttfp_text}{rtf_text})"
)
_emit_offline_metrics(
request_id=request_id,
elapsed_s=elapsed,
first_audio_elapsed=first_audio_elapsed,
audio_duration_s=audio_duration_s,
)
return output_path
async def _run_streaming_warmup(args, omni: AsyncOmni) -> None:
if args.warmup_runs == 0:
return
warmup_specs = _get_warmup_specs(args.prompt_specs)
print(
f"Warmup: {args.warmup_runs} run(s) using the first prompt "
f"({len(warmup_specs)} prompt(s)); outputs will be discarded."
)
for warmup_index in range(args.warmup_runs):
t_warmup = time.perf_counter()
tasks = []
request_ids: list[str] = []
for prompt_index, spec in enumerate(warmup_specs):
request_id = f"warmup_stream_{warmup_index + 1}_{spec.label}_{uuid.uuid4().hex[:8]}"
request_ids.append(request_id)
tasks.append(
_collect_streaming_audio(
omni,
args,
spec,
request_id,
phase_label=f"warmup={warmup_index + 1}/{args.warmup_runs}",
prompt_index=prompt_index,
prompt_count=len(warmup_specs),
)
)
results = await asyncio.gather(*tasks)
for request_id in request_ids:
await _abort_streaming_residual_work(omni, request_id)
total_samples = sum(int(audio.numel()) for audio, _, _, _ in results)
warmup_ttfps = [ttfp for _, _, _, ttfp in results if ttfp is not None]
ttfp_text = f", ttfp={min(warmup_ttfps):.2f}s" if warmup_ttfps else ""
print(
f"Warmup (streaming) {warmup_index + 1}/{args.warmup_runs} finished: "
f"{len(results)} prompt(s), {total_samples} sample(s) "
f"({time.perf_counter() - t_warmup:.2f}s{ttfp_text})"
)
async def _run_streaming(args) -> list[Path]:
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
omni = AsyncOmni(
model=args.model,
stage_configs_path=args.stage_configs_path,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
await _run_streaming_warmup(args, omni)
profiler_started = False
if args.enable_profiler:
profile_prefix = f"voxcpm_streaming_{int(time.time())}"
stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured"
print(f"Starting profiler (streaming): stages={stages_text}, dir={args.profiler_dir}")
await omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages)
profiler_started = True
t_total = time.perf_counter()
total_elapsed = 0.0
paths: list[Path] = []
prompt_specs: list[PromptSpec] = args.prompt_specs
try:
for run in range(args.num_runs):
for prompt_index, spec in enumerate(prompt_specs):
request_id = f"stream_{run + 1}_{spec.label}_{uuid.uuid4().hex[:8]}"
paths.append(
await _run_streaming_single(
omni,
args,
spec,
output_dir,
request_id,
run_index=run,
num_runs=args.num_runs,
prompt_index=prompt_index,
prompt_count=len(prompt_specs),
)
)
total_elapsed = time.perf_counter() - t_total
finally:
if profiler_started:
print("Stopping profiler (streaming)...")
await omni.stop_profile(stages=args.profiler_stages)
if args.profiler_wait_seconds > 0:
print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...")
await asyncio.sleep(args.profiler_wait_seconds)
print(
f"All streaming runs finished: {args.num_runs} run(s), "
f"{len(prompt_specs)} prompt(s), {len(paths)} file(s) in {total_elapsed:.2f}s total"
)
return paths
def _run_sync(args) -> list[Path]:
output_dir = Path(args.output_dir)
omni = Omni(
model=args.model,
stage_configs_path=args.stage_configs_path,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
def _run_sync_single(
spec: PromptSpec,
*,
request_prefix: str,
save_outputs: bool,
run_index: int | None = None,
) -> tuple[list[Path], int, float | None, float, float, str]:
global_request_id = f"{request_prefix}_{spec.label}"
prompt = _build_prompt_for_spec(args, spec, global_request_id=global_request_id)
if save_outputs and run_index == 0 and spec.label == "item001":
print(f"---prompt---:{prompt}")
saved_paths: list[Path] = []
output_count = 0
first_audio_elapsed: float | None = None
total_audio_duration_s = 0.0
metrics_request_id = global_request_id
t_start = time.perf_counter()
for stage_outputs in omni.generate(prompt):
request_output = stage_outputs.request_output
if request_output is None:
continue
request_output_id = getattr(request_output, "request_id", None)
if isinstance(request_output_id, str) and request_output_id:
metrics_request_id = request_output_id
for j, mm in enumerate(_iter_request_multimodal_outputs(request_output)):
output_count += 1
if first_audio_elapsed is None:
try:
audio_tensor = _extract_audio_tensor(mm)
if int(audio_tensor.numel()) > 0:
first_audio_elapsed = time.perf_counter() - t_start
total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm))
except ValueError:
pass
else:
try:
audio_tensor = _extract_audio_tensor(mm)
total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm))
except ValueError:
pass
if not save_outputs:
continue
save_stem = f"run{run_index + 1}_{spec.label}" if j == 0 else f"run{run_index + 1}_{spec.label}_{j}"
saved_paths.append(_save_wav(mm, output_dir, save_stem))
if output_count == 0:
raise RuntimeError("No output from Omni.generate")
elapsed_s = time.perf_counter() - t_start
return saved_paths, output_count, first_audio_elapsed, elapsed_s, total_audio_duration_s, metrics_request_id
if args.warmup_runs:
warmup_specs = _get_warmup_specs(args.prompt_specs)
print(
f"Warmup: {args.warmup_runs} run(s) using the first prompt "
f"({len(warmup_specs)} prompt(s)); outputs will be discarded."
)
for warmup_index in range(args.warmup_runs):
t_warmup = time.perf_counter()
_, output_count, first_audio_elapsed, elapsed_s, audio_duration_s, _ = _run_sync_single(
warmup_specs[0],
request_prefix=f"warmup_sync{warmup_index + 1}",
save_outputs=False,
)
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
print(
f"Warmup (sync) {warmup_index + 1}/{args.warmup_runs} finished: "
f"{output_count} output(s) ({time.perf_counter() - t_warmup:.2f}s{ttfp_text}{rtf_text})"
)
profiler_started = False
if args.enable_profiler:
profile_prefix = f"voxcpm_sync_{int(time.time())}"
stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured"
print(f"Starting profiler (sync): stages={stages_text}, dir={args.profiler_dir}")
omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages)
profiler_started = True
t_total = time.perf_counter()
total_elapsed = 0.0
saved_paths: list[Path] = []
prompt_specs: list[PromptSpec] = args.prompt_specs
try:
for run in range(args.num_runs):
t_run = time.perf_counter()
run_paths: list[Path] = []
for prompt_index, spec in enumerate(prompt_specs):
prompt_paths, _, first_audio_elapsed, elapsed_s, audio_duration_s, metrics_request_id = (
_run_sync_single(
spec,
request_prefix=f"sync_run{run + 1}_{prompt_index + 1:03d}",
save_outputs=True,
run_index=run,
)
)
run_paths.extend(prompt_paths)
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
print(
f"Saved (sync) run {run + 1}/{args.num_runs}, "
f"prompt {prompt_index + 1}/{len(prompt_specs)}: {len(prompt_paths)} file(s){ttfp_text}{rtf_text}"
)
_emit_offline_metrics(
request_id=metrics_request_id,
elapsed_s=elapsed_s,
first_audio_elapsed=first_audio_elapsed,
audio_duration_s=audio_duration_s,
)
saved_paths.extend(run_paths)
print(
f"Run {run + 1}/{args.num_runs} finished: {len(run_paths)} file(s) ({time.perf_counter() - t_run:.2f}s)"
)
for path in run_paths:
print(f" {path}")
total_elapsed = time.perf_counter() - t_total
finally:
if profiler_started:
print("Stopping profiler (sync)...")
omni.stop_profile(stages=args.profiler_stages)
if args.profiler_wait_seconds > 0:
print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...")
time.sleep(args.profiler_wait_seconds)
print(
f"All sync runs finished: {args.num_runs} run(s), "
f"{len(prompt_specs)} prompt(s), {len(saved_paths)} file(s) in {total_elapsed:.2f}s total"
)
return saved_paths
def main(args) -> int:
logging.basicConfig(level=logging.INFO)
profiled_stage_config_path: str | None = None
original_stage_config_path = args.stage_configs_path
if args.enable_profiler:
Path(args.profiler_dir).mkdir(parents=True, exist_ok=True)
profiled_stage_config_path = _build_profiled_stage_config(
args.stage_configs_path,
str(Path(args.profiler_dir).resolve()),
)
args.stage_configs_path = profiled_stage_config_path
is_streaming = _is_streaming_stage_config(args.stage_configs_path)
voice_clone_count = _count_voice_clone_prompts(args.prompt_specs)
print(f"Model: {args.model}")
print(f"Stage config: {original_stage_config_path}")
print(f"Route: {'streaming' if is_streaming else 'sync'} (from stage-configs-path)")
print(f"Prompt count: {len(args.prompt_specs)}")
print("Batch mode: sequential (aligned with native VoxCPM)")
print(f"Warmup runs: {args.warmup_runs}")
print(f"Voice cloning prompts: {voice_clone_count}/{len(args.prompt_specs)}")
if args.enable_profiler:
print(f"Profiler: enabled (dir={args.profiler_dir}, stages={args.profiler_stages or 'all-configured'})")
print(f"Profiled stage config: {args.stage_configs_path}")
if voice_clone_count:
print("Voice cloning note: --ref-text/ref_text must match the spoken content of the reference audio.")
print(f"Num runs: {args.num_runs}")
try:
if is_streaming:
asyncio.run(_run_streaming(args))
else:
_run_sync(args)
finally:
if profiled_stage_config_path is not None and os.path.exists(profiled_stage_config_path):
os.unlink(profiled_stage_config_path)
return 0
if __name__ == "__main__":
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
raise SystemExit(main(parse_args()))
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
ARG VLLM_BASE_TAG=v0.21.0
ARG VLLM_BASE_TAG=v0.20.0
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
ARG APP_DIR=/workspace/vllm-omni
WORKDIR ${APP_DIR}
+1 -1
View File
@@ -1,4 +1,4 @@
ARG BASE_IMAGE=vllm/vllm-openai:v0.21.0
ARG BASE_IMAGE=vllm/vllm-openai:v0.20.0
FROM ${BASE_IMAGE}
ARG COMMON_WORKDIR=/app
-48
View File
@@ -1,48 +0,0 @@
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
ARG VLLM_ASCEND_TAG=nightly-main
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
COPY . .
RUN apt-get update && \
apt-get install -y --no-install-recommends \
ffmpeg \
libgl1 \
libglib2.0-0 \
libgomp1 \
libsm6 \
libxext6 \
libxrender1 \
libxcb1 \
espeak-ng && \
rm -rf /var/cache/apt/* && \
rm -rf /var/lib/apt/lists/*
RUN 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 -U pip && \
python3 -m pip install -v \
"pytest>=7.0.0" \
"pytest-asyncio>=0.21.0" \
"pytest-cov>=4.0.0" \
"pytest-mock>=3.10.0" \
"openpyxl>=3.0.0" \
"psutil>=7.2.0" \
"soundfile>=0.13.1" \
"pyttsx3>=2.99" \
"opencv-python>=4.12.0.88" \
"imageio[ffmpeg]>=0.6.0" \
--index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
--default-timeout=60 && \
python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \
--index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
--default-timeout=60
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
ENTRYPOINT []
-48
View File
@@ -1,48 +0,0 @@
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
ARG VLLM_ASCEND_TAG=nightly-main-a3
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
COPY . .
RUN apt-get update && \
apt-get install -y --no-install-recommends \
ffmpeg \
libgl1 \
libglib2.0-0 \
libgomp1 \
libsm6 \
libxext6 \
libxrender1 \
libxcb1 \
espeak-ng && \
rm -rf /var/cache/apt/* && \
rm -rf /var/lib/apt/lists/*
RUN 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 -U pip && \
python3 -m pip install -v \
"pytest>=7.0.0" \
"pytest-asyncio>=0.21.0" \
"pytest-cov>=4.0.0" \
"pytest-mock>=3.10.0" \
"openpyxl>=3.0.0" \
"psutil>=7.2.0" \
"soundfile>=0.13.1" \
"pyttsx3>=2.99" \
"opencv-python>=4.12.0.88" \
"imageio[ffmpeg]>=0.6.0" \
--index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
--default-timeout=60 && \
python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \
--index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
--default-timeout=60
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
ENTRYPOINT []
+1 -1
View File
@@ -1,4 +1,4 @@
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.20.0
FROM ${BASE_IMAGE} AS base
# Declare a variable to know if we want to use the nightly build or the stable build.
+1 -6
View File
@@ -88,7 +88,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.21.0
ARG VLLM_VERSION=v0.20.0
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
WORKDIR /workspace/vllm
@@ -133,11 +133,6 @@ ENTRYPOINT []
FROM ${VLLM_BASE} AS vllm-omni
RUN apt-get update && \
apt-get install -y espeak-ng ffmpeg git jq && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
WORKDIR /workspace/vllm-omni
COPY . .
+1 -6
View File
@@ -6,11 +6,10 @@ nav:
- getting_started/installation/*
- Serving:
- OpenAI-Compatible API:
- Chat Completions API: serving/chat_completions_api.md
- 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
- Videos: serving/videos_api.md
- Audio Generation: serving/audio_generate_api.md
- Streaming Video Input: serving/video_stream_api.md
- Examples:
@@ -77,12 +76,8 @@ nav:
- Quantization:
- Overview: user_guide/quantization/overview.md
- Online Quantization: user_guide/quantization/online.md
- Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
- FP8 W8A8: user_guide/quantization/fp8.md
- Int8 W8A8: user_guide/quantization/int8.md
- MXFP8 W8A8: user_guide/quantization/mxfp8.md
- MXFP4 W4A4: user_guide/quantization/mxfp4.md
- ModelOpt: user_guide/quantization/modelopt.md
- GGUF: user_guide/quantization/gguf.md
- AutoRound: user_guide/quantization/autoround.md
- msModelSlim: user_guide/quantization/msmodelslim.md
+1 -3
View File
@@ -99,6 +99,7 @@ Configuration classes.
- [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][]
@@ -115,13 +116,10 @@ Configuration classes.
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
- [vllm_omni.transformers_utils.configs.cosyvoice3.CosyVoice3Config][]
- [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][]
- [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][]
- [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][]
## Workers
Binary file not shown.

Before

Width:  |  Height:  |  Size: 402 KiB

After

Width:  |  Height:  |  Size: 182 KiB

+3 -26
View File
@@ -39,34 +39,11 @@ Committers have write access and merge rights. They typically have deep expertis
- [@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
- [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs
- [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache
- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
- [@wtomin](https://github.com/wtomin): Diffusion models, parallel, and docs
- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism
- [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
- [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
- [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
## Diffusion Workload Division
This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review.
### Runtime-related
- Diffusion models: @RuixiangMa, @wtomin
### Optimization
- Parallel: @RuixiangMa, @wtomin
- Attention backend: @ZJY0516
- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
- Offload: @yuanheng-zhao
- Quantization: @lishunyang12, @david6666666
- Kernel fusion / communication-computation: @ZJY0516
### Docs & Test
- Diffusion docs: @RuixiangMa, @wtomin
- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
## Meetings
+4 -4
View File
@@ -152,8 +152,6 @@ vllm_omni/ tests/
│ │ │ │ └── test_lora_manager.py ✅
│ ├── models/ │ ├── models/
│ │ ├── qwen_image/ │ │ ├── qwen_image/ (e2e coverage)
│ │ ├── ovis_image/ │ │ ├── ovis_image/
│ │ │ └── ... │ │ │ └── test_ovis_image.py ✅
│ │ ├── z_image/ │ │ └── z_image/
│ │ └── ... │ │ └── test_zimage_tp_constraints.py ✅
│ └── worker/ │ └── worker/
@@ -234,13 +232,15 @@ vllm_omni/ tests/
│ ├── test_qwen3_omni.py
│ ├── test_qwen3_omni_expansion.py
│ ├── test_mimo_audio.py
│ ├── test_image_gen_edit.py
│ └── test_images_generations_lora.py
└── offline_inference/ ✅
├── test_qwen2_5_omni.py
├── test_qwen3_omni.py
├── test_bagel_text2img.py
├── test_z_image.py
├── test_wan22.py
├── test_t2i_model.py
├── test_t2v_model.py
├── test_ovis_image.py
├── test_zimage_tensor_parallel.py
├── test_cache_dit.py
├── test_teacache.py
+3 -9
View File
@@ -80,7 +80,7 @@ Our test scripts use the pytest framework. First, please use `git clone https://
```bash
cd tests
pytest -s -v -m "full_model and (omni or tts) and H100" --run-level=full_model
pytest -s -v -m "full_model and omni and H100" --run-level=full_model
```
=== "L5 level"
@@ -91,16 +91,10 @@ Our test scripts use the pytest framework. First, please use `git clone https://
cd tests
# Stability: Qwen3-Omni
pytest -s -v dfx/stability/scripts/test_stability_qwen3_omni.py -m slow
pytest -s -v dfx/stability/scripts/test_stability_qwen3_omni.py
# Stability: Wan2.2 (v1/videos diffusion benchmark loop)
pytest -s -v dfx/stability/scripts/test_stability_wan22.py -m slow
# Reliability: Qwen3-Omni
pytest -s -v dfx/reliability/test_reliability_qwen3_omni.py -m slow
# Reliability: Wan2.2
pytest -s -v dfx/reliability/test_reliability_wan22.py -m slow
pytest -s -v dfx/stability/scripts/test_stability_wan22.py
```
+1 -2
View File
@@ -11,8 +11,7 @@ Defined in `pyproject.toml`:
| `advanced_model` | L3 tests (run on each merge to main) |
| `full_model` | L4 tests (run nightly) |
| `diffusion` | Diffusion model tests |
| `omni` | Omni multimodal model tests |
| `tts` | TTS model tests |
| `omni` | Omni model tests |
| `cache` | Cache backend tests |
| `parallel` | Parallelism/distributed tests |
| `cpu` | Tests that run on CPU |
+5 -4
View File
@@ -54,8 +54,6 @@ vllm_omni/ tests/
│ │ │ │ └── test_lora_manager.py ✅
│ ├── models/ │ ├── models/
│ │ ├── qwen_image/ │ │ ├── qwen_image/ (e2e coverage)
│ │ ├── ovis_image/ │ │ ├── ovis_image/
│ │ │ └── ... │ │ │ └── test_ovis_image.py ✅
│ │ ├── z_image/ │ │ └── z_image/
│ │ └── ... │ │ └── test_zimage_tp_constraints.py ✅
│ └── worker/ │ └── worker/
@@ -136,13 +134,15 @@ vllm_omni/ tests/
│ ├── test_qwen3_omni.py
│ ├── test_qwen3_omni_expansion.py
│ ├── test_mimo_audio.py
│ ├── test_image_gen_edit.py
│ └── test_images_generations_lora.py
└── offline_inference/ ✅
├── test_qwen2_5_omni.py
├── test_qwen3_omni.py
├── test_bagel_text2img.py
├── test_z_image.py
├── test_wan22.py
├── test_t2i_model.py
├── test_t2v_model.py
├── test_ovis_image.py
├── test_zimage_tensor_parallel.py
├── test_cache_dit.py
├── test_teacache.py
@@ -213,6 +213,7 @@ E2E Online tests for Qwen3-Omni model with mix input and audio+text output.
import os
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
import threading
from pathlib import Path
@@ -115,20 +115,17 @@ class YourAttentionBlock(nn.Module):
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
class YourSelfAttentionBlock(nn.Module):
class YourAttentionBlock(nn.Module):
def __init__(self, ...):
super().__init__()
# Initialize vLLM-Omni's Attention layer.
# `role` lets users target this site with --diffusion-attention-config
# (e.g. --diffusion-attention-config.per_role.self.backend SAGE_ATTN).
# Initialize vLLM-Omni's Attention layer
self.attn = Attention(
num_heads=self.num_heads,
head_size=self.head_dim,
softmax_scale=1.0 / (self.head_dim ** 0.5),
causal=False, # Diffusion models typically use bidirectional attention
num_kv_heads=self.num_kv_heads,
role="self",
)
def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, ...):
@@ -136,6 +133,7 @@ class YourSelfAttentionBlock(nn.Module):
# Create attention metadata
attn_metadata = AttentionMetadata(attn_mask=attention_mask)
hidden_states = self.attn(query, key, value, attn_metadata=attn_metadata)
```
**Key Points:**
@@ -143,35 +141,8 @@ class YourSelfAttentionBlock(nn.Module):
- **Attention layer initialization:** Done in `__init__`, not per-forward
- **Tensor shapes:** vLLM-Omni `Attention` expects QKV to have `[B, seq, num_heads, head_dim]` shape
- **AttentionMetadata:** Wraps attention mask and other metadata
- **Role:** Tag every `Attention` site with a `role` string so users can configure backends per role (see below)
**Declaring attention roles**
The `role` argument is a free-form string that identifies this attention site. Users can match it from `--diffusion-attention-config.per_role.<role>.*` to swap backends without touching model code. Two conventions cover the common cases:
| Convention | When to use | Example |
|---|---|---|
| `"self"` | Q/K/V come from the same hidden state | DiT self-attention block |
| `"cross"` | K/V come from a separate `encoder_hidden_states` | Text-conditioned cross-attention |
For multi-modal or unusual sites, use a dot-namespaced role and pair it with `role_category` so it can fall back to the generic config when nothing model-specific is set:
```python
# A model-specific cross-attention site that user config can target
# either as 'mymodel.audio_to_video' (exact) or as 'cross' (category fallback).
self.audio_to_video_attn = Attention(
num_heads=self.num_heads,
head_size=self.head_dim,
softmax_scale=1.0 / (self.head_dim ** 0.5),
causal=False,
role="mymodel.audio_to_video",
role_category="cross",
)
```
For cross-attention sites whose K/V are replicated across ranks (e.g. text encoder output), pass `skip_sequence_parallel=True` to opt this layer out of sequence-parallel sharding.
**Attention backends:** When the user does not configure a backend, vLLM-Omni asks the current platform for its default (typically `FLASH_ATTN` on CUDA when available). Users override the default via `--diffusion-attention-backend`, the `DIFFUSION_ATTENTION_BACKEND` env var, or finer-grained `--diffusion-attention-config.per_role.*` flags. See [Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md) for the full configuration surface.
**Attention backends:** vLLM-Omni automatically selects the attention backend given the environmental variable `DIFFUSION_ATTENTION_BACKEND`. The default attention backend is `FLASH_ATTN` for diffusion models.
#### 1.3: Replace Imports and Utilities
@@ -83,10 +83,7 @@ There are two important details:
Today that compatibility rule is still shape-sensitive. `height`, `width`,
`num_frames`, and CFG-related fields remain part of the key, so different
resolutions or incompatible guidance settings do **not** co-batch yet. The
key also covers LoRA identity (`lora_int_id`, `lora_scale`), so requests
targeting different adapters or scales run in separate batches and the
worker can activate exactly one adapter per step.
resolutions or incompatible guidance settings do **not** co-batch yet.
The current batching unit is one `OmniDiffusionRequest`. Requests with
multiple prompts do not participate in batching today.
@@ -9,7 +9,6 @@ Backend-specific setup lives in separate docs:
- [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
- [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
- [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md)
## Overview
@@ -24,7 +23,6 @@ Current connectors operate in D2H2D (device to host to device) mode.
| Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
| Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
| Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. |
## Core API
@@ -1,211 +0,0 @@
# YuanrongTransferEngineConnector
## When to Use
Use `YuanrongTransferEngineConnector` for high-performance peer-to-peer KV cache
transfer on Ascend NPU. It is different from `YuanrongConnector`:
- `YuanrongConnector` uses Yuanrong Datasystem as a distributed KV store and
requires Datasystem workers plus etcd.
- `YuanrongTransferEngineConnector` uses Yuanrong TransferEngine directly. It
registers local memory in each stage worker and lets the receiver pull data
from the sender with TransferEngine.
The implementation currently lives under
`vllm_omni/platforms/npu/omni_connectors/` because the supported backend is
Ascend NPU only. The connector is still registered with the generic
OmniConnector factory under the name `YuanrongTransferEngineConnector`.
For Ascend P2P transfer, the connector should use NPU memory as its transfer
pool. The Mooncake-style CPU pool is not supported for this connector.
## Prerequisites
Install Yuanrong Datasystem/TransferEngine Python bindings in the runtime
environment. The packaged install is:
```bash
pip install openyuanrong-datasystem
```
To build from source, follow the upstream project instructions:
- Source repository: [openeuler/yuanrong-datasystem](https://atomgit.com/openeuler/yuanrong-datasystem)
- Installation guide: [Yuanrong Datasystem Linux installation](https://pages.openeuler.openatom.cn/openyuanrong-datasystem/docs/zh-cn/latest/installation/installation_linux.html)
The Ascend driver, CANN runtime, HCCL/HCCP, and NPU device network must also be
configured on every machine participating in the transfer.
## Check NPU Device IPv4
Yuanrong TransferEngine with `protocol: "ascend"` uses the NPU device network,
not only the host management IP. Each participating NPU must have an IPv4 HCCN
device IP.
Check the device IP, link, and health state:
```bash
for i in {0..7}; do
echo "===== NPU $i ====="
hccn_tool -i "$i" -ip -g
hccn_tool -i "$i" -link -g
hccn_tool -i "$i" -net_health -g
done
```
Expected output for each NPU should include an IPv4 address and healthy link:
```text
ipaddr:10.86.1.80
netmask:255.255.255.0
link status: UP
net health status: Success
```
If the output says `Get ipconf failed, because no ip was preset there!`, the NPU
device IP is not configured and Ascend P2P will fail. If the address is IPv6,
configure IPv4 for this connector path or validate that the underlying Yuanrong
TransferEngine build supports the IPv6 HCCN path before using it.
## Check Device-to-Device Connectivity
After collecting the target-stage NPU IPv4 addresses, ping them from each
source-stage NPU:
```bash
# Example: AR uses NPU 0..3 and DiT uses NPU 4..7.
# Replace the destination IPs with the values reported by hccn_tool -i 4..7 -ip -g.
for src in 0 1 2 3; do
for dst_ip in 10.86.1.78 10.86.1.79 10.86.1.80 10.86.1.81; do
echo "===== NPU $src -> $dst_ip ====="
hccn_tool -i "$src" -ping -g address "$dst_ip" || true
done
done
```
The expected result is `0.00% packet loss` for every required source/target pair.
This verifies basic HCCN IPv4 connectivity. TransferEngine can still fail later
if RDMA/QP/HCCL setup is broken, but failed ping means the device network must be
fixed first.
## YAML Configuration
For Ascend P2P, configure the connector with `protocol: "ascend"` and an NPU
memory pool:
```yaml
runtime:
enabled: true
defaults:
window_size: -1
max_inflight: 1
connectors:
yuanrong_te_connector:
name: YuanrongTransferEngineConnector
extra:
host: "auto"
zmq_port: 50051
rpc_port: "auto"
protocol: "ascend"
device_name: "auto"
memory_pool_size: 1073741824
memory_pool_device: "npu"
edges:
- from: 0
to: 1
window_size: -1
```
Important fields:
| Parameter | Recommended Value | Notes |
|---|---|---|
| `host` | `"auto"` | Advertises a routable local host IP for ZMQ metadata exchange. |
| `zmq_port` | `50051` | Base ZMQ port. The runtime applies stage/rank offsets. |
| `rpc_port` | `"auto"` | Lets TransferEngine choose an available RPC port. |
| `protocol` | `"ascend"` | Uses the Ascend TransferEngine backend. |
| `device_name` | `"auto"` | Each TP worker resolves to its local logical NPU, for example `npu:0`, `npu:1`, etc. |
| `memory_pool_device` | `"npu"` | Required for the Ascend fast path. This registers NPU memory with TransferEngine. |
| `memory_pool_size` | `1073741824` | 1 GiB per worker is a practical starting point for KV-cache transfer validation. Increase if the transfer pool is exhausted. |
`memory_pool_device: "cpu"` is rejected for
`YuanrongTransferEngineConnector` with `protocol: "ascend"`. CPU pool registration
can lead to Ascend P2P failures such as HCCL RA init errors, QP timeouts, or
FFTS/SDMA runtime errors. `memory_pool_device: "cpu"` is common for
`MooncakeTransferEngineConnector`, but Yuanrong TE on Ascend only supports
`"npu"`.
## 4-Card AR-to-DiT TP Example
For AR 4-way TP sending KV cache to DiT 4-way TP, set the stage TP topology in
both stage `omni_kv_config` blocks:
```yaml
stage_args:
- stage_id: 0
runtime:
devices: "0,1,2,3"
engine_args:
tensor_parallel_size: 4
omni_kv_config:
need_send_cache: true
rank_mapping:
from_tp: 4
to_tp: 4
output_connectors:
to_stage_1: yuanrong_te_connector
- stage_id: 1
runtime:
devices: "4,5,6,7"
engine_args:
parallel_config:
tensor_parallel_size: 4
omni_kv_config:
need_recv_cache: true
rank_mapping:
from_tp: 4
to_tp: 4
input_connectors:
from_stage_0: yuanrong_te_connector
```
In this topology each DiT rank receives the corresponding AR rank shard:
```text
DiT rank 0 <- AR rank 0
DiT rank 1 <- AR rank 1
DiT rank 2 <- AR rank 2
DiT rank 3 <- AR rank 3
```
## Validate a Successful Run
Search the inference log for Yuanrong TE receive-side messages:
```bash
grep -E "\[YR GET\]|Successfully received KV cache|transfer_engine get failed" inference.log
```
A successful Ascend fast-path transfer contains lines like:
```text
[YR GET] ... fast_path, zero-copy
Successfully received KV cache ... across 1 key(s)
Applied CFG KV caches
[KV Reuse] ...
```
Sender-side `KV transfer OK` only means that the sender registered the payload
and metadata. The actual P2P data transfer happens on the receiver side during
`[YR GET]`.
## Troubleshooting
| Symptom | Likely Cause | Action |
|---|---|---|
| `IPv4 device IP not found` | NPU HCCN IPv4 is missing. | Run `hccn_tool -i <id> -ip -g` and configure IPv4 for every participating NPU. |
| `P2PCommInitRootInfo failed`, `ra init failed`, `RA_QP_STATUS_TIMEOUT` | Ascend P2P/RDMA/QP setup failed. | First verify device IPv4, link health, and device-to-device ping. Then test a minimal Yuanrong TE pair outside vLLM. |
| `fftsplus sdma error`, `context is abort` | Runtime SDMA/FFTS failure, often after an invalid or failed P2P transfer. | Use `memory_pool_device: "npu"` for Ascend TE, reduce `memory_pool_size` if OOM occurs, and isolate the failing NPU pair with a Yuanrong TE demo. |
| Sender logs `KV transfer OK` but receiver never logs `[YR GET]` | Sender `put()` succeeded, receiver P2P read failed or timed out. | Check receiver logs and search for `transfer_engine get failed`. |
| Pool allocation fails or OOM | Transfer pool is too small or too large for available NPU memory. | Start with 1 GiB (`1073741824`), reduce to 512 MiB for validation, or increase for larger KV payloads. |
-356
View File
@@ -1,356 +0,0 @@
# Pipeline Parallel
This section describes how to add Pipeline Parallelism (PP) to a diffusion pipeline. We use the Wan2.2 text-to-video and
image-to-video pipelines as the reference implementations.
## Implementation Checklist
Adding Pipeline Parallel support requires:
1. ✅ **Inherit the mixin** - add `PipelineParallelMixin` before `CFGParallelMixin` in the pipeline class
2. ✅ **Make stage forward resumable** - support `intermediate_tensors` in the transformer forward path
3. ✅ **Return the right object type** - `IntermediateTensors` on non-last PP ranks, final model output on the last rank
4. ✅ **Use the standard denoising loop** - call `predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()`
5. ✅ **Test parity** - compare PP results against the single-GPU baseline
---
## Table of Contents
- [Overview](#overview)
- [Step-by-Step Implementation](#step-by-step-implementation)
- [Testing](#testing)
- [Troubleshooting](#troubleshooting)
- [Reference Implementations](#reference-implementations)
---
## Overview
### What is Pipeline Parallelism?
Pipeline Parallelism splits the denoising transformer into multiple sequential stages and places each stage on a
different rank. Instead of every rank holding the full DiT, each PP rank owns only a slice of the layers.
For each denoising step:
1. Rank 0 starts the forward pass using the current latents.
2. Each intermediate rank receives hidden states from the previous rank, runs its local layer slice, and forwards the
intermediate tensors downstream.
3. The last PP rank produces the final noise prediction.
4. The last PP rank applies the scheduler step and sends the updated latents back to rank 0 for the next timestep.
This reduces per-rank model memory and enables larger diffusion transformers to run across multiple GPUs. It can also be
combined with CFG-Parallel, where each PP pipeline carries one CFG branch.
### Architecture
vLLM-Omni provides `PipelineParallelMixin` to encapsulate the PP communication pattern for diffusion pipelines.
| Method | Purpose | Automatic Behavior |
|-----------------------------------|-------------------------------|---------------------------------------------------------------------------------|
| `diffuse()` | Denoising loop boundary | Wrapped by `PipelineParallelMixin` to flush pending async PP sends on exit |
| `predict_noise_maybe_with_cfg()` | Predict noise with PP support | Runs partial forwards on non-last PP ranks, combines with CFG logic when needed |
| `scheduler_step_maybe_with_cfg()` | Step scheduler with PP sync | Runs scheduler on the last PP rank, returns updated latents to rank 0 |
| `_sync_pp_send()` | Flush pending async sends | Waits for outstanding `isend` handles before later collectives or decode |
`PipelineParallelMixin` is intentionally a pipeline-level abstraction. Your model-specific `predict_noise()` still
defines how a local stage executes.
### How It Works
`predict_noise_maybe_with_cfg()` automatically switches between these modes:
- **PP disabled** (`pipeline_parallel_size == 1`):
- Falls back to `CFGParallelMixin.predict_noise_maybe_with_cfg()`
- **PP only** (`pipeline_parallel_size > 1`, `cfg_parallel_size == 1`):
- Rank 0 starts with the input latents
- Middle ranks receive `intermediate_tensors`, run their local layer range, and asynchronously send downstream
- The last rank returns the final noise prediction
- When CFG is enabled in this mode (sequential CFG), both the positive and negative branches run through the same
PP pipeline, doubling the communication volume per denoising step. Prefer PP + CFG-Parallel when
`cfg_parallel_size > 1` is available to avoid this overhead.
- **PP + CFG-Parallel** (`pipeline_parallel_size > 1`, `cfg_parallel_size > 1`):
- Each PP pipeline carries one CFG branch
- The last PP rank all-gathers across the CFG group
- CFG combination happens on every last PP rank in the CFG group, matching the non-PP CFG-parallel behavior
`scheduler_step_maybe_with_cfg()` keeps the denoising loop consistent:
- **PP disabled**:
- Falls back to `scheduler_step_maybe_with_cfg()`
- **PP enabled**:
- Only the last PP rank has `noise_pred` and runs the scheduler step
- The resulting latents are sent back to rank 0
- Rank 0 receives an `AsyncLatents` wrapper that resolves only when the tensor is next consumed
This asynchronous design avoids unnecessary blocking between denoising steps.
When a pipeline class is defined, `PipelineParallelMixin` wraps its `diffuse()` method and calls `_sync_pp_send()` in a
`finally` block after `diffuse()` returns or raises. This keeps model code free of explicit PP cleanup while still
ensuring that the final non-blocking PP send completes before decode or later collectives.
---
## Step-by-Step Implementation
### Step 1: Inherit `PipelineParallelMixin` and `CFGParallelMixin`
`PipelineParallelMixin` **requires** `CFGParallelMixin` and must appear before it in the class MRO. This is enforced at
class definition time via `__init_subclass__`: defining a pipeline that inherits `PipelineParallelMixin` without
`CFGParallelMixin`, or after `CFGParallelMixin`, raises a `TypeError` immediately on import.
`PipelineParallelMixin` delegates noise prediction, CFG combination, and scheduler stepping to `CFGParallelMixin`, which
supplies `predict_noise()`, `predict_noise_maybe_with_cfg()`, `scheduler_step_maybe_with_cfg()`, and
`combine_cfg_noise()`.
**Example:**
```python
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.pipeline_parallel import PipelineParallelMixin
import torch.nn as nn
class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin):
...
```
The order matters: `PipelineParallelMixin` must be listed before `CFGParallelMixin` so calls to
`predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()` resolve to the PP-aware wrappers, while their
`super()` calls delegate to the CFG implementation when PP is disabled or after the last PP stage.
### Step 2: Make model forward and `predict_noise()` PP-aware
The PP mixin injects `intermediate_tensors` into the normal `predict_noise()` call. Your model forward path must support
two inputs:
- Normal input from rank 0, typically passed as `hidden_states` or `x`
- `intermediate_tensors` from upstream PP ranks
The standard model forward pattern is:
1. If `intermediate_tensors` is present, read the local hidden state from it.
2. Run only this rank's layer slice.
3. Return `IntermediateTensors(...)` on non-last PP ranks.
4. Return the final model output on the last PP rank. `CFGParallelMixin.predict_noise()` already follows this contract
for common pipelines.
**Minimal example:**
```python
from vllm.sequence import IntermediateTensors
from vllm_omni.diffusion.distributed.parallel_state import get_pp_group
def forward(self, hidden_states=None, intermediate_tensors=None, **kwargs):
if intermediate_tensors is not None:
hidden_states = intermediate_tensors["hidden_states"]
for i in range(self.start_layer, self.end_layer):
hidden_states = self.layers[i](hidden_states)
pp_group = get_pp_group()
if not pp_group.is_last_rank:
return IntermediateTensors({"hidden_states": hidden_states})
return (hidden_states,)
```
### Step 3: Partition the transformer layers
The local module on each PP rank must expose only that rank's layer slice. In the reference tests and vLLM model
implementations, this is typically done with vLLM utilities such as `make_layers(...)`, which fill missing ranges with
`PPMissingLayer`. Besides partitioning the layers themselves, model authors should also wire
`make_empty_intermediate_tensors_factory(...)` for intermediate tensor allocation and `is_pp_missing_parameter(...)` for
PP-aware weight loading.
To prepare a transformer for PP, implement the following pieces in order.
#### 3.1 Split the transformer layers across PP ranks
Each PP rank should own only its local layer range, typically exposed as `[start_layer, end_layer)`. In practice this is
usually done with `make_layers(...)`, which constructs local layers and fills missing ranges with `PPMissingLayer`.
The goal is:
- every PP rank knows its `[start_layer, end_layer)` range
- non-local layers do not execute on this rank
- the forward can resume from incoming `intermediate_tensors`
By default, layers are distributed evenly across PP ranks using `get_pp_indices()`. When the number of layers is not
divisible by the PP size, remaining layers are assigned to middle partitions to balance compute and memory. You can
override this with the `VLLM_PP_LAYER_PARTITION` environment variable to specify an exact per-rank layer count:
```bash
# Example: 40 layers across 4 PP ranks, assigning 8 / 12 / 12 / 8 layers
export VLLM_PP_LAYER_PARTITION=8,12,12,8
```
The value must be a comma-separated list of integers whose length equals `pipeline_parallel_size` and whose sum equals
the total number of transformer layers.
#### 3.2 Expose `make_empty_intermediate_tensors`
Your transformer module should expose `self.make_empty_intermediate_tensors`, typically created with
`make_empty_intermediate_tensors_factory(...)`.
This is important because PP ranks need a consistent way to allocate placeholder `IntermediateTensors` with the expected
keys and hidden dimension.
**Example:**
```python
from vllm.model_executor.models.utils import make_empty_intermediate_tensors_factory
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
["hidden_states"],
inner_dim,
)
```
For Wan2.2, the intermediate payload between PP stages is the token sequence stored under `"hidden_states"`, so the
factory is created with that key and the transformer hidden size.
#### 3.3 Return `IntermediateTensors` on non-last PP ranks
Your model `forward()` or custom `predict_noise()` implementation should consume `intermediate_tensors` on non-first
ranks and return `IntermediateTensors(...)` on non-last ranks.
This lets each PP stage resume from the upstream hidden states and pass the local result to the next stage.
#### 3.4 Skip non-local weights during `load_weights()`
When a model is PP-partitioned, many parameters in the checkpoint belong to layers that do not exist on the current
rank. `load_weights()` must therefore skip parameters for missing PP stages using `is_pp_missing_parameter(...)`.
If this is not done, weight loading will either fail or incorrectly try to load tensors into `PPMissingLayer`
placeholders.
The Wan2.2 transformer is the best reference here: it uses `is_pp_missing_parameter(...)` before loading both remapped
and fused parameters.
If your model has multiple transformer variants, PP still works as long as each selected transformer obeys the same
contract.
### Step 4: Use the standard denoising contract
vLLM-Omni diffusion pipelines already route denoising through `diffuse()`, `predict_noise_maybe_with_cfg()`, and
`scheduler_step_maybe_with_cfg()`. `PipelineParallelMixin` overrides those standard helpers when
`pipeline_parallel_size > 1`, so model integrations should not add separate PP-specific helper names or manual post-loop
synchronization.
In PP mode:
- Non-last PP ranks return `None` from `predict_noise_maybe_with_cfg()` after sending `IntermediateTensors` downstream.
- The last PP rank returns the final prediction, and with CFG-Parallel every last PP rank in the CFG group receives the
combined result.
- Rank 0 receives `AsyncLatents` from `scheduler_step_maybe_with_cfg()`, which resolves only when the latents are
consumed.
- Pending non-blocking sends are flushed automatically when `diffuse()` exits.
---
## Testing
Use an offline inference script with `pipeline_parallel_size > 1`.
```bash
python examples/offline_inference/text_to_video/text_to_video.py \
--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--width=1280 \
--height=704 \
--guidance-scale=5.0 \
--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
--output=t2v_5B_pp2.mp4 \
--pipeline-parallel-size=2
```
For PP + CFG-Parallel together:
```bash
python examples/offline_inference/text_to_video/text_to_video.py \
--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--width=1280 \
--height=704 \
--guidance-scale=5.0 \
--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
--output=t2v_5B_pp2_cfg2.mp4 \
--pipeline-parallel-size=2 \
--cfg-parallel-size=2
```
**Verify:**
1. The run completes without hangs at the PP stage boundary
2. Output quality matches the non-PP baseline within normal numerical variance
3. Peak memory per GPU drops compared with the single-rank model
4. There are no outstanding communication errors before decode
---
## Troubleshooting
### Issue: `TypeError` on import — `CFGParallelMixin` missing
**Symptoms:** Importing a pipeline that inherits `PipelineParallelMixin` raises:
```
TypeError: YourPipeline inherits PipelineParallelMixin but not CFGParallelMixin.
```
or:
```
TypeError: YourPipeline must inherit PipelineParallelMixin before CFGParallelMixin ...
```
**Cause:** `PipelineParallelMixin` enforces via `__init_subclass__` that the subclass also inherits `CFGParallelMixin`
and lists `PipelineParallelMixin` first in the MRO.
**Solution:** Add `CFGParallelMixin` to the pipeline's base classes after `PipelineParallelMixin`:
```python
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.pipeline_parallel import PipelineParallelMixin
class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin):
...
```
### Issue: Non-last PP ranks crash when calling `predict_noise`
**Symptoms:** Shape or missing-input errors on ranks other than the first or last PP rank.
**Cause:** The model forward path assumes direct input tensors and ignores `intermediate_tensors`.
**Solution:** Update the transformer `forward()` or custom `predict_noise()` path to load hidden states from
`intermediate_tensors` when present.
### Issue: PP outputs differ from single-GPU baseline
**Symptoms:** The PP run finishes but produces numerically inconsistent outputs.
**Causes & Solutions:**
- **Local layer partitioning is wrong**
- Verify each rank runs only its `[start_layer, end_layer)` slice
- **Non-last ranks return a plain tensor instead of `IntermediateTensors`**
- Return `IntermediateTensors({...})` until the last PP stage
- **CFG branch wiring is incorrect**
- When CFG is enabled, confirm positive and negative kwargs are passed exactly as in the non-PP path
---
## Reference Implementations
Complete examples in the codebase:
| Component | Path | Notes |
|-------------------------|------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------|
| `PipelineParallelMixin` | `vllm_omni/diffusion/distributed/pipeline_parallel.py` | Core PP communication and scheduler helpers |
| `CFGParallelMixin` | `vllm_omni/diffusion/distributed/cfg_parallel.py` | Default `predict_noise()` tuple normalization and CFG helper fallback |
| Wan2.2 transformer | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Reference for layer partitioning, `IntermediateTensors`, `make_empty_intermediate_tensors`, and PP-aware weight loading |
| Wan2.2 T2V pipeline | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py` | Reference PP + CFG integration for text-to-video |
| Wan2.2 I2V pipeline | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py` | Reference PP + CFG integration for image-to-video |
| PP tests | `tests/diffusion/distributed/test_pipeline_parallel.py` | Baseline parity and async communication tests |
+59 -55
View File
@@ -169,13 +169,15 @@ class DiffusionRequestStatus(enum.IntEnum):
@dataclass
class DiffusionRequestState:
request_id: str
sched_req_id: str
req: OmniDiffusionRequest
status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
```
**Design Features**:
- **Scheduler-owned ID**: Each `OmniDiffusionRequest` is tracked by an internal `sched_req_id`, separated from public `request_id` values.
- **Explicit lifecycle**: Requests move through waiting, running, optional preemption, and terminal states.
- **Centralized error handling**: Completion, abort, and error states are all normalized in the scheduler layer.
@@ -186,6 +188,7 @@ class DiffusionRequestState:
class _BaseScheduler(SchedulerInterface):
def __init__(self) -> None:
self._request_states = {}
self._request_id_to_sched_req_id = {}
self._waiting = deque()
self._running = []
self._finished_req_ids = set()
@@ -196,8 +199,7 @@ class _BaseScheduler(SchedulerInterface):
- **Common state storage**: Shared request maps and waiting/running sets live in the base class.
- **Shared cleanup logic**: Duplicate-request checks, finish handling, and state removal are centralized instead of
duplicated in each policy.
- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
- **Current constraint boundary**: `_BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`, but request-mode diffusion is still clamped back to `1` by the engine. The step-wise path can keep this above `1` for compatible-request batching.
@@ -217,12 +219,12 @@ class RequestScheduler(_BaseScheduler):
- **Single-request admission**: `RequestScheduler` still admits one active request at a time because request-mode execution completes a whole request per dispatch.
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished request ids.
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished scheduler ids.
#### 2.5 Engine-Driven Execution Loop
```python
request_id = scheduler.add_request(request)
sched_req_id = scheduler.add_request(request)
while True:
sched_output = scheduler.schedule()
output = executor.add_req(req)
@@ -453,7 +455,7 @@ To learn more about the diffusion pipeline and how to add a new diffusion pipeli
#### Architecture
The attention system uses a **role-aware backend selector pattern**. Each `Attention` site declares a semantic `role` (`"self"`, `"cross"`, or a model-specific string). The selector consults the user's `AttentionConfig` to pick a backend per role and falls back to the platform default when nothing is configured.
The attention system uses a **backend selector pattern** that automatically chooses the optimal attention implementation based on hardware and model configuration.
#### Backend Selection
@@ -461,90 +463,92 @@ The attention system uses a **role-aware backend selector pattern**. Each `Atten
```python
class Attention(nn.Module):
def __init__(self, num_heads, head_size, causal, softmax_scale, *,
role="self", role_category=None, ...):
# Resolve backend for this role from the active OmniDiffusionConfig
config = get_current_diffusion_config_or_none()
attention_config = config.attention_config if config is not None else None
attn_backend_cls, spec = get_attn_backend_for_role(
role=role,
head_size=head_size,
attention_config=attention_config,
role_category=role_category,
)
self.attn_impl_cls = attn_backend_cls.get_impl_cls()
self.attention = self.attn_impl_cls(..., backend_kwargs=spec.extra if spec else None)
def __init__(self, num_heads, head_size, causal, softmax_scale, ...):
# Auto-select backend
self.attn_backend = get_attn_backend(-1)
self.attn_impl_cls = self.attn_backend.get_impl_cls()
self.attention = self.attn_impl_cls(...)
```
**Available Backends**:
- **FlashAttention** (`FLASH_ATTN`): Memory-efficient kernel that dispatches to FA2/FA3 on CUDA, AITER on ROCm, and `mindiesd` on Ascend NPU.
- **FlashAttention**: Optimized CUDA kernel (FA2/FA3) - memory efficient via tiling
- **SDPA** (`TORCH_SDPA`): PyTorch's scaled dot-product attention — cross-platform fallback.
- **SDPA**: PyTorch's scaled dot-product attention - default, cross-platform
- **SageAttention** (`SAGE_ATTN`): Quantized attention implementation from the SageAttention library.
- **SageAttention**: Sparse attention implementation from SageAttention library
- **AscendAttention**: NPU-optimized attention for Ascend hardware
These backends provide the **kernel implementations** for attention computation. For attention-level sequence parallelism strategies (Ring Attention, Ulysses), see [Parallel Attention](#52-parallel-attention).
#### Backend Selection Mechanism
Selection is driven by `AttentionConfig` on `OmniDiffusionConfig`, which carries a global `default` plus a `per_role` map of `AttentionSpec` entries. A role string is resolved in this order:
```python
def get_attn_backend_for_role(role, head_size, attention_config=None, role_category=None):
# 1. attention_config.per_role[role] — exact match
# 2. attention_config.per_role[role_category] — category fallback
# 3. attention_config.default — global default
# 4. Platform default — hardware-specific
if attention_config is not None:
spec, source = attention_config.resolve_with_source(
role=role, role_category=role_category,
)
if spec is not None:
return load_backend(spec.backend), spec
return current_omni_platform.get_diffusion_attn_backend_cls(...), None
def get_attn_backend(head_size: int) -> type[AttentionBackend]:
# Check environment variable
backend_name = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
if backend_name:
return load_backend(backend_name.upper())
# Default to SDPA
return SDPABackend
```
**Selection Priority**:
1. **Per-role override** (`--diffusion-attention-config.per_role.<role>.backend`) — finest control; matched on the layer's exact `role` string.
1. **Environment Variable**: `DIFFUSION_ATTENTION_BACKEND` for manual override
2. **Role category fallback** (`--diffusion-attention-config.per_role.<category>.backend`) — used when an exact match is missing and the layer declared `role_category` (e.g. a `"mymodel.audio_to_video"` site can fall back to `"cross"`).
- Valid values: `FLASH_ATTN`, `TORCH_SDPA`, `SAGE_ATTN`, `ASCEND`
3. **Global default** (`--diffusion-attention-backend FLASH_ATTN`, or `--diffusion-attention-config.default.backend`, or `DIFFUSION_ATTENTION_BACKEND` env var).
- Example: `export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN`
4. **Platform default** — `current_omni_platform.get_diffusion_attn_backend_cls(...)` picks the best-available kernel for the hardware.
2. **Automatic Fallback**: Falls back to SDPA if selected backend unavailable
`AttentionSpec.extra` is forwarded to the backend constructor as `backend_kwargs`, so backend-specific parameters (e.g. SparseBlock `block_size`) can be set without changing model code.
For the user-facing CLI surface, see [Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). For the role declaration contract on the model side, see [Adding a Diffusion Model](../../contributing/model/adding_diffusion_model.md).
3. **Hardware Detection**: Can select based on device type (NPU, CUDA, etc.)
**Backend Availability**:
- **SDPA**: Always available (PyTorch built-in)
- **FlashAttention**: Requires `flash-attn` on CUDA, `aiter` on ROCm, or `mindiesd` on Ascend NPU
- **FlashAttention**: Requires `flash-attn` package installed
- **SageAttention**: Requires `sage-attention` package (from THU-ML GitHub)
- **AscendAttention**: Only available on Ascend NPU hardware
#### Attention Backend Registry
**Location**: `vllm_omni/diffusion/attention/selector.py` + `vllm_omni/platforms/<device>/platform.py`
**Location**: `vllm_omni/diffusion/attention/selector.py`
Backend resolution is delegated to the active platform: `current_omni_platform.get_diffusion_attn_backend_cls(selected_backend, head_size)` returns a fully-qualified class path that `_load_backend_cls` imports lazily, with the result cached by `(backend_name, head_size)`.
The attention system uses a **registry pattern** to manage and dynamically load attention backends. This allows for easy extension and runtime selection of backends.
**Registry Structure**:
```python
@cache
def _cached_get_backend_cls(backend_name: str | None, head_size: int) -> type[AttentionBackend]:
backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls(
selected_backend=backend_name,
head_size=head_size,
)
return _load_backend_cls(backend_cls_path)
# Registry mapping backend names to their module paths and class names
_BACKEND_CONFIG = {
"FLASH_ATTN": {
"module": "vllm_omni.diffusion.attention.backends.flash_attn",
"class": "FlashAttentionBackend",
},
"TORCH_SDPA": {
"module": "vllm_omni.diffusion.attention.backends.sdpa",
"class": "SDPABackend",
},
"SAGE_ATTN": {
"module": "vllm_omni.diffusion.attention.backends.sage_attn",
"class": "SageAttentionBackend",
},
"ASCEND": {
"module": "vllm_omni.diffusion.attention.backends.ascend_attn",
"class": "AscendAttentionBackend",
},
}
```
Each platform (`cuda`, `rocm`, `xpu`, `musa`, `npu`) maps backend names like `"FLASH_ATTN"`, `"TORCH_SDPA"`, `"SAGE_ATTN"` to the right backend class path for that hardware, including head-size compatibility checks. Passing `selected_backend=None` lets the platform pick its own default.
#### Attention Backend Integration
The `Attention` layer integrates backends through a unified interface. Here's how **FlashAttentionBackend** is integrated as an example:
+1 -1
View File
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
### Pre-built wheels
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, and 0.21.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, and 0.18.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
=== "NVIDIA CUDA"
@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
uv pip install vllm==0.21.0 --torch-backend=auto
uv pip install vllm --torch-backend=auto
```
#### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library
```bash
uv pip install vllm==0.21.0 --torch-backend=auto
uv pip install vllm==0.20.0 --torch-backend=auto
```
The 0.21.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
The 0.20.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.21.0
git checkout v0.20.0
```
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
```bash
#For CUDA 13.0 (the default for v0.21.0; the wheel filename has no `+cu130` suffix)
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.21.0/vllm-0.21.0-cp38-abi3-manylinux_2_35_x86_64.whl
#For CUDA 13.0 (the default for v0.20.0; the wheel filename has no `+cu130` suffix)
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.20.0/vllm-0.20.0-cp38-abi3-manylinux_2_35_x86_64.whl
```
Install vllm with command below (If you have no existing PyTorch).
```bash
@@ -102,11 +102,11 @@ docker run --runtime nvidia --gpus 2 \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
--ipc=host \
vllm/vllm-omni:v0.21.0 \
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
vllm/vllm-omni:v0.18.0 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
```
!!! tip
The CUDA image does not define a default entrypoint, so include `vllm serve ... --omni` after the image name.
You can use this docker image to serve models the same way you would with in vLLM! To do so, make sure you overwrite the default entrypoint (`vllm serve --omni`) which works only for models supported in the vLLM-Omni project.
# --8<-- [end:pre-built-images]
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
```
#### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library
```bash
uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
```
The pre-built 0.21.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
The pre-built 0.20.0 vLLM wheel targets ROCm 7.2.1. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.21.0
git checkout v0.20.0
python3 -m pip install -r requirements/rocm.txt
python3 setup.py develop
```
@@ -130,7 +130,7 @@ docker run --rm \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
vllm/vllm-omni-rocm:v0.21.0 \
vllm/vllm-omni-rocm:v0.20.0 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
@@ -149,7 +149,7 @@ docker run --rm -it \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--entrypoint bash \
vllm/vllm-omni-rocm:v0.21.0
vllm/vllm-omni-rocm:v0.20.0
```
# --8<-- [end:pre-built-images]
+3 -3
View File
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
source .venv/bin/activate
# On CUDA
uv pip install vllm==0.21.0 --torch-backend=auto
uv pip install vllm==0.20.0 --torch-backend=auto
# On ROCm
uv pip install vllm==0.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721
uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
@@ -35,7 +35,7 @@ For additional installation methods — please see the [installation guide](inst
!!! note
It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.21.0` and vLLM Omni `0.21.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.20.0` and vLLM Omni `0.20.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
## Offline Inference
+7 -16
View File
@@ -16,7 +16,6 @@ ROOT_DIR_RELATIVE = "../../../../.."
EXAMPLE_DIR = ROOT_DIR / "examples"
EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
NAV_FILE = ROOT_DIR / "docs/.nav.yml"
MAX_INLINE_MATERIAL_SIZE = 8 * 1024
def fix_case(text: str) -> str:
@@ -192,21 +191,18 @@ class Example:
return re.sub(link_pattern, replace_link, content)
def github_url(self, path: Path) -> str:
url = "https://github.com/vllm-project/vllm-omni/"
url += "tree/main" if path.is_dir() else "blob/main"
return f"{url}/{path.relative_to(ROOT_DIR).as_posix()}"
def generate(self) -> str:
content = f"# {self.title}\n\n"
content += f"Source <{self.github_url(self.path)}>.\n\n"
url = "https://github.com/vllm-project/vllm-omni/"
url += "tree/main" if self.path.is_dir() else "blob/main"
content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"
# Use long code fence to avoid issues with
# included files containing code fences too
code_fence = "``````"
if self.is_code:
main_file_rel = self.main_file.relative_to(ROOT_DIR).as_posix()
main_file_rel = self.main_file.relative_to(ROOT_DIR)
content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel}"\n{code_fence}\n'
else:
with open(self.main_file, encoding="utf-8") as f:
@@ -220,15 +216,10 @@ class Example:
content += "## Example materials\n\n"
for file in sorted(self.other_files):
content += f'??? abstract "{file.relative_to(self.path).as_posix()}"\n'
if file.stat().st_size > MAX_INLINE_MATERIAL_SIZE:
content += (
f" Large file omitted from the rendered docs. View it on GitHub: <{self.github_url(file)}>.\n\n"
)
continue
content += f'??? abstract "{file.relative_to(self.path)}"\n'
if file.suffix != ".md":
content += f" {code_fence}{file.suffix[1:]}\n"
content += f' --8<-- "{file.relative_to(ROOT_DIR).as_posix()}"\n'
content += f' --8<-- "{file.relative_to(ROOT_DIR)}"\n'
if file.suffix != ".md":
content += f" {code_fence}\n"
@@ -300,7 +291,7 @@ def update_nav_file(examples: list[Example]):
for example in category_examples:
doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
rel_path = doc_path.relative_to(ROOT_DIR / "docs")
category_items.append({example.title: rel_path.as_posix()})
category_items.append({example.title: str(rel_path)})
if category_items:
# Format category name (e.g., "offline_inference" -> "Offline Inference")
-29
View File
@@ -1,29 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Reduce browser search-index memory for generated API reference pages."""
from material.plugins.search.plugin import SearchIndex
_add_entry_from_context = SearchIndex.add_entry_from_context
def add_entry_from_context(self: SearchIndex, page) -> None:
search = page.meta.get("search") or {}
if search.get("exclude"):
return
if str(page.url).startswith("api/"):
entry = {
"location": page.url,
"title": str(page.meta.get("title", page.title)),
"text": "",
}
if "boost" in search:
entry["boost"] = search["boost"]
self.entries.append(entry)
return
_add_entry_from_context(self, page)
SearchIndex.add_entry_from_context = add_entry_from_context
+18 -42
View File
@@ -1,44 +1,20 @@
// Enables MathJax rendering only on pages that contain math.
// Enables MathJax rendering
window.MathJax = {
tex: {
inlineMath: [["\\(", "\\)"]],
displayMath: [["\\[", "\\]"]],
processEscapes: true,
processEnvironments: true,
},
options: {
ignoreHtmlClass: ".*|",
processHtmlClass: "arithmatex",
},
};
tex: {
inlineMath: [["\\(", "\\)"]],
displayMath: [["\\[", "\\]"]],
processEscapes: true,
processEnvironments: true
},
options: {
ignoreHtmlClass: ".*|",
processHtmlClass: "arithmatex"
}
};
let mathJaxLoading;
function loadMathJax() {
if (window.MathJax.typesetPromise) {
return Promise.resolve();
}
if (!mathJaxLoading) {
mathJaxLoading = new Promise((resolve, reject) => {
const script = document.createElement("script");
script.src = "https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js";
script.async = true;
script.onload = resolve;
script.onerror = reject;
document.head.appendChild(script);
});
}
return mathJaxLoading;
}
document$.subscribe(() => {
if (!document.querySelector(".arithmatex")) {
return;
}
loadMathJax().then(() => {
MathJax.startup.output.clearCache();
MathJax.typesetClear();
MathJax.texReset();
MathJax.typesetPromise();
});
});
document$.subscribe(() => {
MathJax.startup.output.clearCache()
MathJax.typesetClear()
MathJax.texReset()
MathJax.typesetPromise()
})
-2
View File
@@ -1,2 +0,0 @@
// Opt out of bfcache so large generated API pages are released after navigation.
window.addEventListener("unload", () => {});
+18 -39
View File
@@ -1,42 +1,21 @@
// Load Mermaid only on pages that contain diagrams.
let mermaidLoading;
// Initialize Mermaid for diagram rendering
mermaid.initialize({
startOnLoad: false,
theme: 'default',
securityLevel: 'loose',
flowchart: {
useMaxWidth: true,
htmlLabels: true
}
});
function loadMermaid() {
if (window.mermaid) {
return Promise.resolve();
// Render Mermaid diagrams when page content is ready
document$.subscribe(() => {
const mermaidElements = document.querySelectorAll('.mermaid');
if (mermaidElements.length > 0) {
mermaid.run({
querySelector: '.mermaid',
nodes: mermaidElements
});
}
if (!mermaidLoading) {
mermaidLoading = new Promise((resolve, reject) => {
const script = document.createElement("script");
script.src = "https://unpkg.com/mermaid@10/dist/mermaid.min.js";
script.async = true;
script.onload = () => {
mermaid.initialize({
startOnLoad: false,
theme: "default",
securityLevel: "loose",
flowchart: {
useMaxWidth: true,
htmlLabels: true,
},
});
resolve();
};
script.onerror = reject;
document.head.appendChild(script);
});
}
return mermaidLoading;
}
document$.subscribe(() => {
if (!document.querySelector(".mermaid")) {
return;
}
loadMermaid().then(() => {
mermaid.run({
querySelector: ".mermaid",
nodes: document.querySelectorAll(".mermaid"),
});
});
});
-6
View File
@@ -19,7 +19,6 @@ th {
|--------------|--------|-------------------|------------|---------|-----|-----------|
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | |
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
@@ -59,20 +58,15 @@ th {
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | |
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
@@ -1,4 +1,4 @@
# Chat Completions API
# Diffusion Chat Completions API
vLLM-Omni supports generating and editing images via the `/v1/chat/completions`
endpoint using diffusion models. This page explains how to pass generation
+5 -70
View File
@@ -13,7 +13,7 @@ For example...
```bash
# Qwen-Image
vllm serve Qwen/Qwen-Image-Edit-2511 --omni --port 8000
```
### Generate Images
@@ -53,7 +53,7 @@ result = client.images.edit(
output_format='jpeg',
# url格式
extra_body={
"url": [input_image_url,input_image_url],
"url": [input_image_url1,input_image_url],
"num_inference_steps": 50,
"guidance_scale": 1,
"seed": 777,
@@ -92,8 +92,7 @@ Content-Type: multipart/form-data
| `user` | string | null | User identifier for tracking |
| `output_format` | string | "png" | The format in which the generated images are returned. Must be one of "png", "jpg", "jpeg", "webp". |
| `output_compression` | integer | 100 | The compression level (0-100%) for the generated images. |
| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s). |
| `stream` | boolean | false | Return Server-Sent Events instead of a single JSON response. Currently supported only by multi-stage image edit pipelines that produce AR text before the final image, such as HunyuanImage3 IT2I AR+DiT. |
| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s).
#### vllm-omni Extension Parameters
@@ -110,8 +109,6 @@ Content-Type: multipart/form-data
### Response Format
When `stream=false` or omitted, the endpoint returns the standard image edit response:
```json
{
"created": 1701234567,
@@ -123,56 +120,10 @@ When `stream=false` or omitted, the endpoint returns the standard image edit res
}
],
"output_format": null,
"size": null
"size": null,
}
```
### Streaming Response Format
Set `stream=true` when you want to receive the HunyuanImage3 IT2I AR
recaption text before the final edited image is ready. Streaming is only
available for multi-stage image edit pipelines where stage 0 produces AR text
and a later diffusion stage produces the image. Single-stage diffusion image
edit pipelines reject `stream=true` with HTTP 400.
The response uses Server-Sent Events with `Content-Type: text/event-stream`.
Each event is sent as one `data:` line. The stream order is:
1. One or more AR text delta chunks.
2. One final image chunk.
3. `data: [DONE]`.
AR text chunks expose only the generated text delta and the AR completion
`index`. They do not expose token ids, ratio tokens, KV cache data, or internal
prompt token ids.
Example AR delta event:
```text
data: {"object":"image.edit.chunk","type":"ar_delta","delta":"A close-up product photo...","index":0,"created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct"}
```
Example final image event:
```text
data: {"object":"image.edit.chunk","type":"image","data":[{"b64_json":"<base64-encoded PNG>","url":null,"revised_prompt":null}],"output_format":"png","size":"1024x1024","created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct"}
```
Terminal event:
```text
data: [DONE]
```
If the engine fails after the SSE response has started, the stream emits an
error chunk followed by `[DONE]`:
```text
data: {"object":"error","created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct","error":{"message":"<error message>","type":"server_error","code":500}}
data: [DONE]
```
## Examples
### Multiple Images input
@@ -183,28 +134,12 @@ curl -s -D >(grep -i x-request-id >&2) \
-X POST "http://localhost:8000/v1/images/edits" \
-F "model=xxx" \
-F "image=@xx.png" \
-F "image=@xx.png" \
-F "image=@xx.png"
-F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
-F "size=1024x1024" \
-F "output_format=png"
```
### Streaming HunyuanImage3 IT2I AR Text
```bash
curl -N -X POST "http://localhost:8000/v1/images/edits" \
-F "model=tencent/HunyuanImage-3.0-Instruct" \
-F "image=@./input.png" \
-F "prompt=Turn the product into a clean studio advertisement." \
-F "size=1024x1024" \
-F "output_format=png" \
-F "stream=true"
```
Use the `ar_delta` events for progressive display of the AR-generated
recaption. Decode the `b64_json` field from the final `image` event to get the
edited image.
## Parameter Handling
+4 -13
View File
@@ -5,9 +5,6 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
- **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
- **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
- **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
- **CosyVoice3** (`FunAudioLLM/Fun-CosyVoice3-0.5B-2512`) -- 2-stage talker + flow-matching code2wav. Voice cloning via `ref_audio` + `ref_text` (no presets). Output: 24 kHz.
See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, and MOSS-TTS-Nano.
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
@@ -29,10 +26,6 @@ vllm serve fishaudio/s2-pro --omni --port 8091
# Voxtral TTS
vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
# CosyVoice3 (voice cloning only — supply ref_audio + ref_text per request)
vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 \
--omni --port 8091 --trust-remote-code
```
### Generate Speech
@@ -538,18 +531,16 @@ for result in response.json()["results"]:
All items are fanned out to `generate()` concurrently. The engine's stage worker automatically batches them up to the configured `max_batch_size` and queues the rest — no client-side throttling needed.
For best throughput, set both stages' `max_num_seqs` above 1 via `--stage-overrides`. On the current Qwen3-TTS CustomVoice benchmark, stage 1 performed best at `max_num_seqs: 10`:
For best throughput, set both stages' `max_num_seqs` to ≥4 via `--stage-overrides`:
```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--omni --port 8091 --trust-remote-code --enforce-eager \
--stage-overrides '{"0":{"max_num_seqs":10,"gpu_memory_utilization":0.2},
"1":{"max_num_seqs":10,"gpu_memory_utilization":0.2}}'
--stage-overrides '{"0":{"max_num_seqs":4,"gpu_memory_utilization":0.2},
"1":{"max_num_seqs":4,"gpu_memory_utilization":0.2}}'
```
The bundled `qwen3_tts.yaml` uses a multi-request default and lets stage 1 batch chunks across in-flight requests. For latency-sensitive deployments, avoid forcing stage 1 back to `max_num_seqs: 1`; benchmark before reducing it below `10`.
The bundled config also sets `initial_codec_chunk_frames: 1`. This emits only the first audio chunk early for lower TTFA, then returns to the normal `codec_chunk_frames` window so Code2Wav does not repeatedly decode tiny overlapping chunks.
The bundled `qwen3_tts.yaml` uses `max_num_seqs: 1` (single request) on both stages. Bumping to 4 yields roughly 4× throughput on the talker and lets stage 1 batch chunks across in-flight requests.
## Supported Models
-152
View File
@@ -1,152 +0,0 @@
# Videos API
vLLM-Omni provides an OpenAI-compatible video generation API for diffusion
video models. The API supports asynchronous video jobs through `/v1/videos` and
a synchronous benchmark-oriented endpoint through `/v1/videos/sync`.
Each server instance runs a single model specified at startup with
`vllm serve <model> --omni`.
## Quick Start
### Start the Server
```bash
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091
```
### Create a Video Job
```bash
create_response=$(curl -s http://localhost:8091/v1/videos \
-F "prompt=A cinematic tracking shot of a mountain lake at sunrise" \
-F "width=1280" \
-F "height=720" \
-F "num_frames=80" \
-F "fps=16" \
-F "num_inference_steps=40")
video_id=$(echo "${create_response}" | jq -r '.id')
```
### Poll and Download
```bash
curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
```
## API Reference
### Endpoints
| Endpoint | Method | Description |
|----------|--------|-------------|
| `/v1/videos` | `POST` | Create an asynchronous video generation job |
| `/v1/videos/sync` | `POST` | Generate a video synchronously and return raw video bytes |
| `/v1/videos/{video_id}` | `GET` | Retrieve job status and metadata |
| `/v1/videos` | `GET` | List stored video jobs |
| `/v1/videos/{video_id}/content` | `GET` | Download generated video content |
| `/v1/videos/{video_id}` | `DELETE` | Delete a video job and stored output |
### Request Parameters
`POST /v1/videos` and `POST /v1/videos/sync` accept `multipart/form-data`.
#### OpenAI-style fields
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `prompt` | string | **required** | Text prompt for video generation |
| `model` | string | server's model | Optional model name |
| `seconds` | string | null | Requested clip duration in seconds |
| `size` | string | null | Requested output size in `WIDTHxHEIGHT` format |
| `user` | string | null | Optional user identifier |
#### vLLM-Omni extension fields
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input_reference` | file | null | Uploaded reference image for image-to-video requests |
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` |
| `width` | integer | model default | Output video width |
| `height` | integer | model default | Output video height |
| `num_frames` | integer | model default | Number of generated frames |
| `fps` | integer | model default | Output frames per second |
| `num_inference_steps` | integer | model default | Number of diffusion steps |
| `guidance_scale` | number | null | CFG guidance scale for the low-noise stage |
| `guidance_scale_2` | number | null | CFG guidance scale for the high-noise stage |
| `boundary_ratio` | number | null | Boundary split ratio for multi-stage denoising |
| `flow_shift` | number | null | Scheduler flow-shift value |
| `true_cfg_scale` | number | null | True CFG scale when supported by the model |
| `seed` | integer | null | Random seed for reproducibility |
| `negative_prompt` | string | null | Text describing what to avoid in the generated video |
| `enable_frame_interpolation` | boolean | null | Enable post-generation frame interpolation |
| `frame_interpolation_exp` | integer | null | Interpolation exponent; `1=2x`, `2=4x`, and so on |
| `frame_interpolation_scale` | number | null | RIFE inference scale |
| `frame_interpolation_model_path` | string | null | Local path or Hugging Face repo for the interpolation model |
| `lora` | string | null | JSON-encoded LoRA configuration object |
| `extra_params` | string | null | JSON-encoded object for additional model-specific parameters |
### Create Response
`POST /v1/videos` returns a job record:
```json
{
"id": "video-123",
"status": "queued",
"created_at": 1701234567
}
```
The final content is available from `/v1/videos/{video_id}/content` after the
job status becomes `completed`.
### Synchronous Response
`POST /v1/videos/sync` blocks until generation finishes and returns raw video
bytes. It is useful for benchmarks and simple scripts that do not need job
storage or polling.
## Examples
### Image-to-Video
```bash
curl -s http://localhost:8091/v1/videos \
-F "prompt=animate this image with subtle camera movement" \
-F "input_reference=@input.png" \
-F "width=1280" \
-F "height=720" \
-F "num_frames=80" \
-F "fps=16"
```
### Synchronous Generation
```bash
curl -X POST http://localhost:8091/v1/videos/sync \
-F "prompt=A small robot walking through a neon city" \
-F "width=854" \
-F "height=480" \
-F "num_frames=80" \
-F "fps=16" \
-o output.mp4
```
## Storage
Set `VLLM_OMNI_STORAGE_PATH` to control where asynchronous video outputs are
stored:
```bash
export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
```
## Model-Specific Examples
For complete text-to-video and image-to-video walkthroughs, see:
- [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md)
- [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md)
+29 -206
View File
@@ -1,170 +1,42 @@
# Diffusion Attention Backends
This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, the per-platform defaults, and how to use SageAttention.
This document describes the diffusion attention backends available in vLLM-Omni, how to select them, and how to use SageAttention.
## Overview
Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured.
Diffusion attention backend selection is controlled by the `DIFFUSION_ATTENTION_BACKEND` environment variable and resolved in `vllm_omni.diffusion.attention.selector`.
This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector.
This backend is used by diffusion attention layers such as the DiT attention in video and image generation models.
The full set of backends and their platform defaults is in the **Backend Options** and **Platform Defaults** sections below. If no attention backend is configured, vLLM-Omni asks the current platform to choose the default.
On CUDA, the practical choices today are:
- `FLASH_ATTN`: FlashAttention backend. This is the default on supported CUDA systems when FlashAttention is installed.
- `TORCH_SDPA`: PyTorch `scaled_dot_product_attention`.
- `SAGE_ATTN`: SageAttention backend, if `sageattention` is installed.
If `DIFFUSION_ATTENTION_BACKEND` is unset, vLLM-Omni asks the current platform to choose the default backend. On CUDA, that normally means `FLASH_ATTN` when available, otherwise `TORCH_SDPA`.
## Backend Options
| Value | Notes |
|---|---|
| `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. |
| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
| `FLASH_ATTN` | Default on CUDA when FlashAttention is available. Good default for most diffusion workloads. |
| `TORCH_SDPA` | Most conservative fallback. Useful for debugging or compatibility. |
| `SAGE_ATTN` | Requires `sageattention`. Can improve performance on some workloads, but output quality must be validated model-by-model. |
## Configuration
## Selection Priority
Diffusion attention backends can be configured three ways, in priority order:
Diffusion attention backend selection follows this order:
1. **`--diffusion-attention-config`** — structured per-role config (highest priority).
2. **`--diffusion-attention-backend` / `DIFFUSION_ATTENTION_BACKEND` env var** — global shorthand that sets the default backend.
3. **Platform default** — used when nothing is configured.
1. `DIFFUSION_ATTENTION_BACKEND`
2. Platform default
`--diffusion-attention-backend` is shorthand for `--diffusion-attention-config.default.backend`. It may be combined with `--diffusion-attention-config.per_role.*` overrides, but is mutually exclusive with `--diffusion-attention-config.default.backend`.
### Global default
Set the default backend for every diffusion attention layer:
Example:
```bash
# CLI flag
vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN
# Environment variable (also recognized for backwards compatibility)
export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN
```
### Per-role configuration
Roles are free-form strings declared by each diffusion model. The two common categories are `"self"` and `"cross"`; model-specific roles (e.g. `"ltx2.audio_to_video"`) may also be declared. A role string is matched in this order:
1. Exact `per_role[role]` match
2. `per_role[role_category]` fallback (e.g. `"ltx2.audio_to_video"` → `"cross"`)
3. `default`
4. Platform default
Use vLLM-style dotted flags or one JSON blob:
```bash
# Dotted flags
vllm-omni serve <model> \
--diffusion-attention-config.default.backend FLASH_ATTN \
--diffusion-attention-config.per_role.cross.backend TORCH_SDPA
# JSON
vllm-omni serve <model> \
--diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
```
Backends may also accept backend-specific parameters via `extra`:
```bash
--diffusion-attention-config.per_role.self.backend SPARSE_BLOCK \
--diffusion-attention-config.per_role.self.extra.block_size 128
```
### Programmatic API
When constructing `OmniDiffusionConfig` directly:
```python
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
config = OmniDiffusionConfig(
attention_config=AttentionConfig(
default=AttentionSpec(backend="FLASH_ATTN"),
per_role={
"cross": AttentionSpec(backend="TORCH_SDPA"),
},
),
...,
)
```
A plain dict is also accepted and normalized to `AttentionConfig`.
## Platform Defaults
### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
Auto-route preference, in order:
1. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
2. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
3. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
4. `TORCH_SDPA` — last resort
The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
**Why CUDNN_ATTN by default**: on mask-heavy diffusion models (HunyuanVideo-1.5, Qwen-Image), cuDNN's pinned FMHA kernel sidesteps a PyTorch SDPA dispatch quirk where the unpinned dispatcher picks `EFFICIENT_ATTENTION` (~25 ms) for masked calls instead of cuDNN (~11 ms). The pin gives 2× e2e on HV-1.5 with no regression on lighter models.
### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86)
Auto-route preference:
1. `FLASH_ATTN` — when `flash-attn` is installed
2. `TORCH_SDPA` — fallback
`CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware.
## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
| Model | Shape | TORCH_SDPA | CUDNN_ATTN | FLASHINFER_ATTN |
|---|---|---|---|---|
| HunyuanVideo-1.5 (T2V) | 480p / 33f / 50 steps | 147.05 s | **73.02 s** | 127.84 s |
| Wan 2.2 14B (T2V) | 480p / 33f / 40 steps | 117.75 s | 117.17 s | **115.07 s** |
| Qwen-Image (T2I) | 1024² / 50 steps | 17.41 s | **15.14 s** | 16.02 s |
| FLUX.2-dev (T2I) | 1024² / 50 steps, TP=2 | 53.62 s | **53.30 s** | 54.94 s |
Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
## Known Limitations
### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
```bash
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
--model Lightricks/LTX-2 ...
```
### FA4 not yet integrated
FlashAttention-4 (released March 2026) targets Blackwell natively and reportedly beats cuDNN by ~20% on B200. As of this writing the `flash-attn-4 4.0.0b10` wheel crashes with `AttributeError: 'NoneType' object has no attribute '_trait'` during JIT on sm_120. Not yet wired into vLLM-Omni; revisit when stable lands.
## Choosing a Backend Manually
### When to override the default
- **Quality validation**: compare a new backend against `TORCH_SDPA` as the reference, since SDPA's default dispatcher is the most extensively tested.
- **Lossy speedup hunting**: try `SAGE_ATTN` (INT8 quantized) on diffusion outputs — typically indistinguishable visually but always validate.
- **Workaround for known issues**: see Known Limitations above.
### Verifying which backend is in use
The startup log prints one of:
```
Using diffusion attention backend 'CUDNN_ATTN' # explicit override
Defaulting to diffusion attention backend CUDNN_ATTN ... # auto-route
Defaulting to diffusion attention backend SDPA # nothing else available
```
If you don't see one of these, the model didn't reach diffusion stage init — check earlier logs for failures.
## SageAttention Installation
vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
@@ -185,33 +57,11 @@ Quick check:
python -c "import sageattention; print(sageattention.__file__)"
```
## Usage Examples
## Usage
### Default (auto-route)
### Enable SageAttention
```bash
python examples/offline_inference/text_to_video/text_to_video.py \
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
--prompt "A dog running across a field of golden wheat." \
--height 480 --width 832 --num-frames 33 \
--num-inference-steps 50 --seed 42 --guidance-scale 6.0 \
--output hv15.mp4
```
On Blackwell this picks `CUDNN_ATTN` automatically. Check the log for the `Defaulting to ...` line.
### Explicit backend selection
```bash
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
--model Lightricks/LTX-2 \
--prompt "A dog running across a field of golden wheat." \
--height 480 --width 832 --num-frames 33 \
--num-inference-steps 40 --seed 42 --guidance-scale 4.0 \
--output ltx2.mp4
```
### SageAttention (lossy)
Example: HunyuanVideo-1.5 text-to-video
```bash
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
@@ -220,7 +70,7 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
--height 480 --width 832 --num-frames 33 \
--num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
--tensor-parallel-size 2 \
--output hv15_sage.mp4
--output ../tmp/hv15_modelopt_sage.mp4
```
Example: Wan2.2 TI2V 5B
@@ -235,21 +85,7 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
--output outputs/wan22_sage.mp4
```
### Mixed backends across roles
Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention:
```bash
python examples/offline_inference/text_to_video/text_to_video.py \
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--prompt "A dog running across a field of golden wheat." \
--diffusion-attention-config.per_role.self.backend FLASH_ATTN \
--diffusion-attention-config.per_role.cross.backend TORCH_SDPA \
--tensor-parallel-size 2 \
--output outputs/wan22_mixed.mp4
```
### Compare against FlashAttention
### Compare Against FlashAttention
Unset the backend override, or explicitly use `FLASH_ATTN`:
@@ -265,13 +101,13 @@ python examples/offline_inference/text_to_video/text_to_video.py \
## Validation Guidance
Don't assume a faster attention backend is numerically interchangeable with `TORCH_SDPA`.
Do not assume that a faster attention backend is numerically interchangeable with `FLASH_ATTN`.
Always compare:
- End-to-end runtime
- Diffusion-stage runtime (`add_req_and_wait` line in DiffusionEngine.step breakdown)
- Output quality against a known-good baseline (CLIP similarity, frame-level diff, or visual review)
- DiT / diffusion stage runtime
- Output quality against a known-good baseline
At minimum, keep the same:
@@ -279,19 +115,6 @@ At minimum, keep the same:
- prompt
- seed
- resolution
- frame count / step count
- parallel config (TP / CFG-parallel / Ulysses degrees)
## Reproducing the Benchmark Table
The end-to-end numbers above were collected by running `text_to_video.py` /
`text_to_image.py` with the same prompt and seed while varying
`DIFFUSION_ATTENTION_BACKEND`. For a quick kernel-level comparison of the
backends without loading a model:
```bash
python benchmarks/diffusion/bench_attention_backends.py --preset hv15
```
It runs all three BF16 backends on representative DiT attention shapes and
prints a ranking table at the end.
- frame count
- inference steps
- parallel config
@@ -38,15 +38,14 @@ vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload
### To Support a Model
Implement the `SupportsComponentDiscovery` protocol to declare which
submodules serve as pipeline components (used by offloading, HSDP
sharding, and other framework features):
Implement the `SupportsModuleOffload` protocol to declare which
submodules participate in offloading:
```python
from typing import ClassVar
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
from vllm_omni.diffusion.models.interface import SupportsModuleOffload
class MyPipeline(nn.Module, SupportsComponentDiscovery):
class MyPipeline(nn.Module, SupportsModuleOffload):
_dit_modules: ClassVar[list[str]] = ["transformer"]
_encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
_vae_modules: ClassVar[list[str]] = ["vae"]
@@ -159,7 +158,7 @@ class Flux2Transformer2DModel(nn.Module):
The offloader discovers pipeline components in two ways:
1. **Protocol-based** (preferred): If the pipeline implements
`SupportsComponentDiscovery`, its `_dit_modules`, `_encoder_modules`,
`SupportsModuleOffload`, its `_dit_modules`, `_encoder_modules`,
`_vae_modules`, and `_resident_modules` class variables are used
directly. All attribute names support dotted paths (e.g.
`"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules.
@@ -28,6 +28,7 @@ Frame interpolation is currently supported for:
- `WanPipeline` (Wan2.2 text-to-video)
- `WanImageToVideoPipeline`
- `Wan22TI2VPipeline`
## Request Parameters
@@ -0,0 +1,80 @@
# Diffusion Multi-Replica Serving
Diffusion stages can be replicated in online serving by setting `runtime.num_replicas` on a `stage_type: diffusion` stage. Each replica owns an independent diffusion engine process and receives requests through round-robin routing with request affinity.
## When to Use
Use multi-replica serving when one diffusion process does not saturate the available GPUs, or when you want higher request throughput for many independent image/video generation requests. It improves concurrency throughput; it does not make a single request faster.
## Configuration
A replica consumes the devices required by that diffusion stage. For single-device diffusion pipelines, set `runtime.devices` to one device per replica:
```yaml
stages:
- stage_id: 0
stage_type: diffusion
runtime:
devices: "0,1,2,3"
num_replicas: 4
engine_args:
model: Qwen/Qwen-Image
parallel_config:
world_size: 1
```
For multi-device diffusion pipelines, `parallel_config.world_size` is treated as the number of devices per replica. For example, four 2-GPU replicas need eight devices:
```yaml
stages:
- stage_id: 0
stage_type: diffusion
runtime:
devices: "0,1,2,3,4,5,6,7"
num_replicas: 4
engine_args:
model: Wan-AI/Wan2.2-T2V-A14B-Diffusers
parallel_config:
world_size: 2
```
The runtime splits the device list into contiguous chunks per replica. With the second example, replica 0 uses `0,1`, replica 1 uses `2,3`, replica 2 uses `4,5`, and replica 3 uses `6,7`.
## Serving
Start the server with the stage config:
```bash
vllm serve Qwen/Qwen-Image \
--omni \
--port 8099 \
--stage-configs-path /path/to/diffusion_multi_replica.yaml
```
Then send normal OpenAI-compatible diffusion chat requests. The orchestrator routes new requests across replicas and keeps updates/abort calls on the same replica for each request.
## Benchmarking
Benchmark each replica count by changing `runtime.num_replicas` and `runtime.devices`, restarting the server, then running:
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--base-url http://localhost:8099 \
--model Qwen/Qwen-Image \
--task t2i \
--dataset random \
--num-prompts 8 \
--request-rate inf \
--num-inference-steps 2 \
--width 256 --height 256 \
--output-file diffusion_replica_4.json
```
For a quick scaling sweep, test `num_replicas` values `1`, `2`, `3`, and `4` with the same prompt count and generation parameters. Compare throughput and end-to-end latency from the JSON outputs.
## Notes and Limitations
- Multi-replica serving is currently supported for diffusion stages. Single-stage mode still rejects non-diffusion replicas.
- `runtime.devices` must provide exactly `num_replicas * devices_per_replica` devices.
- Request-level latency may stay similar or slightly increase because each request still runs on one replica.
- Throughput improves only when enough concurrent requests are issued and each replica has enough GPU memory.
@@ -4,14 +4,13 @@ This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion
## Supported Methods
| Method | Description |
|----------------------------------------------------|---------------------------------------------------------------------------------------------------------------------|
| **[Tensor Parallelism](tensor_parallel.md)** | Shards DiT weights across GPUs to reduce per-GPU memory |
| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation |
| **[Pipeline Parallelism](pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages to reduce per-GPU model memory |
| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs |
| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) |
| Method | Description |
|--------|-------------|
| **[Tensor Parallelism](tensor_parallel.md)** | Shards DiT weights across GPUs to reduce per-GPU memory |
| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation |
| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs |
| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) |
See [Supported Models](../../diffusion_features.md#supported-models) for per-model compatibility.
@@ -1,205 +0,0 @@
# Pipeline Parallelism Guide
## Table of Content
- [Overview](#overview)
- [Quick Start](#quick-start)
- [Example Script](#example-script)
- [Configuration Parameters](#configuration-parameters)
- [Best Practices](#best-practices)
- [Troubleshooting](#troubleshooting)
- [Summary](#summary)
---
## Overview
Pipeline Parallelism splits the denoising transformer block-wise into sequential stages across GPUs. Each rank owns only
part of the transformer, which reduces per-GPU model memory and enables larger diffusion models to run across multiple
devices.
It can also be combined with other distributed methods such as CFG-Parallel, Tensor Parallelism, and Sequence
Parallelism.
See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
---
## Quick Start
### Basic Usage
Simplest working example:
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
parallel_config=DiffusionParallelConfig(
pipeline_parallel_size=2,
),
)
outputs = omni.generate(
{"prompt": "A cinematic drone shot over snowy mountains"},
OmniDiffusionSamplingParams(
num_inference_steps=40,
num_frames=81,
height=704,
width=1280,
),
)
```
---
## Example Script
### Offline Inference
Use python scripts under:
- `examples/offline_inference/text_to_video/text_to_video.py`
- `examples/offline_inference/image_to_video/image_to_video.py`
Text-to-video example:
```bash
python examples/offline_inference/text_to_video/text_to_video.py \
--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--width=1280 \
--height=704 \
--guidance-scale=5.0 \
--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
--output=t2v_5B_pp2.mp4 \
--pipeline-parallel-size=2
```
Pipeline Parallelism can also be combined with CFG-Parallel:
```bash
python examples/offline_inference/text_to_video/text_to_video.py \
--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--width=1280 \
--height=704 \
--guidance-scale=5.0 \
--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
--output=t2v_5B_pp2_cfg2.mp4 \
--pipeline-parallel-size=2 \
--cfg-parallel-size=2
```
### Online Serving
Enable Pipeline Parallelism in online serving:
```bash
# Default PP configuration
vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni --port 8091 --pipeline-parallel-size 2
# PP + CFG-Parallel
vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni --port 8091 \
--pipeline-parallel-size 2 \
--cfg-parallel-size 2
```
---
## Configuration Parameters
In `DiffusionParallelConfig`
| Parameter | Type | Default | Description |
|--------------------------|------|---------|------------------------------------------------------------------------------------------------------------------|
| `pipeline_parallel_size` | int | 1 | Number of pipeline-parallel stages. Set to a value greater than 1 to split the denoising transformer across GPUs |
> [!NOTE]
> Total GPU count is the product of all enabled distributed dimensions, for example
`pipeline_parallel_size * cfg_parallel_size * tensor_parallel_size * ulysses_degree * ring_degree`.
### Manual Layer Partitioning
By default, transformer layers are distributed evenly across PP ranks. You can override this with the
`VLLM_PP_LAYER_PARTITION` environment variable to assign a specific number of layers to each rank:
```bash
# Example: 40 layers across 4 PP ranks, assigning 8 / 12 / 12 / 8 layers
export VLLM_PP_LAYER_PARTITION=8,12,12,8
```
The value must be a comma-separated list of integers whose length equals `pipeline_parallel_size` and whose sum equals
the total number of transformer layers. This is useful when you want to balance memory or compute asymmetrically across
ranks.
---
## Best Practices
### When to Use
**Good for:**
- Large diffusion transformers that do not fit comfortably on one GPU
- Multi-GPU setups where reducing per-GPU model memory is more important than minimizing communication
- Combining with CFG-Parallel or other distributed methods on supported models
**Not for:**
- Single GPU setups
- Models that do not support Pipeline Parallelism
(check [supported models](../../diffusion_features.md#supported-models))
- Very small models where inter-stage communication overhead may outweigh the benefit
### Expected Behavior
Pipeline Parallelism primarily reduces per-GPU memory usage by splitting the transformer into stage-local blocks.
Depending on the model, topology, and resolution, it may also help execution fit into available hardware, but it is not
primarily a latency optimization.
---
## Troubleshooting
### Common Issue 1: No benefit from Pipeline Parallelism
**Symptoms**: PP is enabled, but latency does not improve or becomes slightly worse.
**Solutions**:
1. **Check your goal:**
```python
# PP is mainly for memory scaling, not guaranteed latency speedup
parallel_config = DiffusionParallelConfig(pipeline_parallel_size=2)
```
2. **Check model support:**
- Verify your model in [supported models](../../diffusion_features.md#supported-models)
- PP is currently validated only on selected pipelines
3. **Combine with other methods when appropriate:**
- PP can be combined with CFG-Parallel, Tensor Parallelism, or Sequence Parallelism on supported models
### Common Issue 2: PP pipeline fails at import
**Symptoms**: Importing a custom pipeline raises a `TypeError` about `CFGParallelMixin` or mixin order.
**Solutions**:
1. Inherit both mixins.
2. Put `PipelineParallelMixin` before `CFGParallelMixin` in the class base list.
3. Use `class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin): ...` as the reference pattern.
---
## Summary
1. ✅ **Enable Pipeline Parallelism** - Set `pipeline_parallel_size > 1` in `DiffusionParallelConfig`
2. ✅ **Use Supported Models** - Verify your model supports PP in
[supported models](../../diffusion_features.md#supported-models)
3. ✅ **Combine When Needed** - PP can be combined with CFG-Parallel and other distributed methods on supported pipelines
4. ✅ **Scale for Memory** - Use PP primarily to reduce per-GPU model memory and fit larger transformers
@@ -60,9 +60,6 @@ batching only becomes relevant when `--max-num-seqs > 1`.
- `cache_backend` is not supported together with step execution.
- Unsupported pipelines fail early during model loading.
- Request-mode extras such as KV transfer are not wired into step mode yet.
- LoRA is supported in step mode, but each batch must use a single adapter:
requests with different `lora_request` or `lora_scale` are scheduled into
separate batches.
## When To Use It
+50 -52
View File
@@ -15,7 +15,7 @@ vLLM-Omni supports various advanced features for diffusion models:
- Acceleration: **cache methods**, **parallelism methods**, **startup optimizations**
- Memory optimization: **cpu offloading**, **quantization**
- Extensions: **LoRA inference**, **frame interpolation**
- Execution modes: **step execution**
- Execution modes: **step execution**, **multi-replica serving**
## Supported Features
@@ -35,15 +35,15 @@ Cache methods trade minimal quality for significant speedup. Quality loss is typ
Parallelism methods distribute computation across GPUs without quality loss (mathematically equivalent to single-GPU).
| Method | Description | Best For |
|------------------------------------------------------------------------|--------------------------------------------------------------------------|-------------------------------------------------------------------|
| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via all-to-all communication | High-resolution images (>1536px) or long videos with 2-8 GPUs |
| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via ring-based communication | Videos, very long sequences, memory-constrained, with 2-8 GPUs |
| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)** | Splits CFG positive/negative branches across devices | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs |
| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs |
| **[Pipeline Parallelism](diffusion/parallelism/pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages | Large diffusion transformers that need lower per-GPU model memory |
| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
| Method | Description | Best For |
|--------|-------------|----------|
| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via all-to-all communication | High-resolution images (>1536px) or long videos with 2-8 GPUs |
| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via ring-based communication | Videos, very long sequences, memory-constrained, with 2-8 GPUs |
| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)** | Splits CFG positive/negative branches across devices | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs |
| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs |
| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
| **[Multi-Replica Serving](diffusion/multi_replica_serving.md)** | Runs multiple independent diffusion stage replicas and routes requests across them | Online serving throughput with many concurrent requests |
#### Startup Optimization
@@ -106,33 +106,31 @@ The following tables show which models support each feature:
### ImageGen
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|--------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.2-dev** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **Hidream-I1-Full** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **SenseNova-U1** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **OmniGen2** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
> Notes:
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
@@ -140,15 +138,15 @@ The following tables show which models support each feature:
### VideoGen
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
**Frame Interpolation Support**
@@ -157,9 +155,9 @@ The following tables show which models support each feature:
### AudioGen
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|-----------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ |
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ |
## Feature Compatibility
@@ -185,7 +183,7 @@ The following tables show which models support each feature:
| **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | | |
| **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | | |
| **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | |
| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ✅ | |
| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ❌ | |
!!! info
@@ -194,7 +192,7 @@ The following tables show which models support each feature:
3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
4. CPU Offloading (Layerwise) supports single-card for now.
5. Using FP8-Quant as an example of qunatization methods.
6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT). LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT) or LoRA.
## Multi-Thread Weight Loading
@@ -259,7 +257,7 @@ Measured on NVIDIA H800:
**Parallelism Methods:**
- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, Pipeline Parallelism, HSDP, and Expert Parallelism
- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, HSDP, and Expert Parallelism
**Memory Optimization:**
@@ -17,13 +17,12 @@ list of supported architectures across all modalities, see
| Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
|---|---|---|---|---|---|---|
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | — | continuation (`--ref-audio` + `--ref-text`) | 48 kHz |
| VoxCPM | local model dir | split | ✓ | ✓ (`voxcpm_async_chunk.yaml`) | — | 24 kHz |
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ (`async_chunk: true` default) | — | 22.05 kHz |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ (`--streaming`) | — | 44.1 kHz |
| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design (`--instruct`), language (`--lang`) | 24 kHz |
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants (`--query-type`) | 24 kHz |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets (`--voice`) | 24 kHz |
| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` (+ `k2-fsa/OmniVoice/audio_tokenizer/` codec) | 2 (talker + code2wav, DualFFN) | ✓ (`--ref-audio` + `--ref-text`) | — | — | 24 kHz |
## Common Quick Start
@@ -73,6 +72,61 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
---
## VoxCPM
Split-stage TTS. The hub example covers single TTS, single voice cloning, and streaming. Use `benchmarks/voxcpm/` for warmup, batch JSONL prompts, profiler injection, and offline TTFP / RTF measurement.
### Prerequisites
```bash
pip install voxcpm soundfile
# or use a local source tree:
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
```
If the native VoxCPM `config.json` does not contain HF metadata such as `model_type`, prepare a persistent HF-compatible config directory and point the stage configs to it via `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH`:
```bash
export VOXCPM_MODEL=/path/to/voxcpm-model
export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
```
### Quick start
```bash
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
--model "$VOXCPM_MODEL" \
--text "This is a split-stage VoxCPM synthesis example running on vLLM Omni."
```
### Voice cloning
```bash
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
--model "$VOXCPM_MODEL" \
--text "This sentence is synthesized with a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "The exact transcript spoken in reference.wav."
```
### Streaming
Pass the async-chunk stage config:
```bash
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
--model "$VOXCPM_MODEL" \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
--text "This is a split-stage VoxCPM streaming example running on vLLM Omni."
```
### Notes
- `voxcpm.yaml` is the default non-streaming stage config; `voxcpm_async_chunk.yaml` enables streaming.
- Streaming is currently single-request oriented.
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades quality.
- For online serving, see [`examples/online_serving/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/voxcpm/README.md). For benchmark reporting, see [`benchmarks/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/benchmarks/voxcpm/README.md).
---
## CosyVoice3
2-stage TTS pipeline (`talker` + `code2wav`) at 22.05 kHz.
@@ -160,35 +214,6 @@ Streaming requires `async_chunk: true` in the stage config.
---
## GLM-TTS
2-stage TTS pipeline (AR + DiT flow-matching) at 24 kHz. Every request requires reference audio and its transcript for zero-shot voice cloning.
### Quick start
```bash
python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
--model zai-org/GLM-TTS \
--text "你好,这是语音合成测试。" \
--ref-audio /path/to/reference.wav \
--ref-text "这是参考音频的文本内容。" \
--output-dir ./output
```
### Architecture
```
Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz)
(Llama-based) (32k vocab) (Flow Matching)
```
### Notes
- `--ref-audio` and `--ref-text` are **required** together; GLM-TTS does not support text-only synthesis.
- Reference audio should be 3-10 seconds.
- First run may be slow due to lazy loading of WhisperVQ tokenizer and CampPlus ONNX speaker embedder.
- Default sampling: temperature=1.0, top_k=25, top_p=0.8 (RAS method).
- The `--model` path should point to the repository root (not `llm/` subdirectory).
---
## OmniVoice
Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design).
@@ -334,46 +359,6 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
- `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
- Run `--help` for the full argument surface.
---
## higgs-audio v2
2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) over the shared-memory connector. Stage 1 builds on the HiggsAudio decoder kernel at `vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py`.
### Prerequisites
Voice clone uses HF's `HiggsAudioV2TokenizerModel`, instantiated from `k2-fsa/OmniVoice/audio_tokenizer/` — the boson-ai standalone tokenizer Hub repo's `model.safetensors` is actually the 3B talker LM, so we point HF at k2's repackaged codec weights instead. Only the `audio_tokenizer/` subdir (~806 MB) is downloaded.
```bash
pip install -U "transformers>=5.3.0"
```
### Quick start (plain TTS)
```bash
python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
--texts "Hello world." "The quick brown fox jumps over the lazy dog." \
--output-dir results/wavs
```
### Voice cloning
Pass both `--ref-audio` and `--ref-text` together:
```bash
python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
--texts "Hello, this is a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "Exact transcript spoken in reference.wav." \
--output-dir results/wavs
```
### Notes
- Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
- Out of scope (rejected with 4xx by the request validator): multi-speaker `[SPEAKERn]` dialogue, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
## Example materials
??? abstract "cosyvoice3/end2end.py"
@@ -384,13 +369,6 @@ python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
``````py
--8<-- "examples/offline_inference/text_to_speech/fish_speech/end2end.py"
``````
??? abstract "higgs_audio_v2/end2end.py"
``````py
--8<-- "examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py"
??? abstract "glm_tts/end2end.py"
``````py
--8<-- "examples/offline_inference/text_to_speech/glm_tts/end2end.py"
``````
??? abstract "omnivoice/end2end.py"
``````py
--8<-- "examples/offline_inference/text_to_speech/omnivoice/end2end.py"
@@ -403,6 +381,10 @@ python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
``````py
--8<-- "examples/offline_inference/text_to_speech/qwen3_tts/end2end.py"
``````
??? abstract "voxcpm/end2end.py"
``````py
--8<-- "examples/offline_inference/text_to_speech/voxcpm/end2end.py"
``````
??? abstract "voxcpm2/end2end.py"
``````py
--8<-- "examples/offline_inference/text_to_speech/voxcpm2/end2end.py"
@@ -101,7 +101,7 @@ guides.
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
[Chat Completions API guide](../../../serving/chat_completions_api.md)).
[Diffusion Chat API guide](../../../serving/diffusion_chat_api.md)).
When using the dedicated [`/v1/images/generations`](../../../serving/image_generation_api.md)
or [`/v1/images/edits`](../../../serving/image_edit_api.md) endpoints, pass
the supported generation controls as top-level fields directly. For image
@@ -208,7 +208,7 @@ Wrap generation parameters inside `extra_body` in the request JSON:
```
For details on how generation parameters are handled across different clients, see the
[Chat Completions API guide](../../../../serving/chat_completions_api.md).
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md).
### Layered Image Generation (Qwen-Image-Layered)
@@ -361,7 +361,7 @@ Provide multiple images in `content` (order matters):
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
[Chat Completions API guide](../../../../serving/chat_completions_api.md)).
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
When using the dedicated [`/v1/images/edits`](../../../../serving/image_edit_api.md)
endpoint, pass the supported generation controls as top-level form fields
directly. For image dimensions and count, use `size` and `n` rather than
@@ -216,7 +216,7 @@ Wrap generation parameters inside `extra_body` in the request JSON:
```
For details on how generation parameters are handled across different clients, see the
[Chat Completions API guide](../../../../serving/chat_completions_api.md).
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md).
### Multimodal Input (Text + Structured Content)
@@ -237,7 +237,7 @@ Wrap generation parameters inside `extra_body` in the request JSON:
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
[Chat Completions API guide](../../../../serving/chat_completions_api.md)).
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
When using the dedicated [`/v1/images/generations`](../../../../serving/image_generation_api.md)
endpoint, pass the supported generation controls as top-level JSON fields
directly. For image dimensions and count, use `size` and `n` rather than
@@ -17,16 +17,14 @@ For the full list of supported architectures across all modalities, see
| Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
|---|---|---|---|---|---|
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — |
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
| OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
| VoxCPM | local model dir | ✓ | ✓ (PCM stream) | — | — |
| VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
For offline inference of any of these models, see the [offline TTS hub](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md).
CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3) instead.
## Common Quick Start
@@ -98,54 +96,6 @@ For full request-shape documentation (all parameters, response formats, error co
---
## CosyVoice3
2-stage TTS (`talker` + flow-matching `code2wav`) at 24 kHz. Voice cloning only — every request needs `ref_audio` + `ref_text`; there are no built-in voice presets.
### Prerequisites
```bash
huggingface-cli download FunAudioLLM/Fun-CosyVoice3-0.5B-2512
```
If your downloaded checkpoint lacks `config.json`, add one with `{"model_type": "cosyvoice3", "architectures": ["CosyVoice3Model"]}` (the loader reads `model_type` to select the class).
### Launch
```bash
vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code
# or:
./cosyvoice3/run_server.sh
```
Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path.
### CLI client
The client defaults to the official upstream zero-shot prompt, so it runs without extra flags:
```bash
cd examples/online_serving/text_to_speech/cosyvoice3
python speech_client.py --text "收到好友从远方寄来的生日礼物。"
```
Pass your own reference clip and transcript for a different voice:
```bash
python speech_client.py --text "Hello, this is a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "Transcript of the reference audio."
```
Stream PCM instead of WAV:
```bash
python speech_client.py --text "Hello world" --stream --output output.pcm
```
The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-text`, `--response-format`, `--stream`, `--output`.
### Notes
- Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
- For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3).
---
## Fish Speech S2 Pro
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
@@ -194,84 +144,6 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server
---
## higgs-audio v2
2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) that streams chunks back to the client.
### Prerequisites
Voice clone uses HF's `HiggsAudioV2TokenizerModel` loaded from `k2-fsa/OmniVoice/audio_tokenizer/` (~806 MB subdir; the boson-ai standalone tokenizer Hub repo's `model.safetensors` is the 3B talker LM, not the codec):
```bash
pip install -U "transformers>=5.3.0"
```
### Launch
```bash
GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
```
Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
### Sending requests
```bash
# Plain TTS
python higgs_audio_v2/batch_speech_client.py \
--base-url http://localhost:8094 \
--output-dir /tmp/higgs_out \
--prompts "Hello world." "The quick brown fox jumps over the lazy dog."
# Voice cloning — pass a reference clip and its transcript together
python higgs_audio_v2/batch_speech_client.py \
--base-url http://localhost:8094 \
--output-dir /tmp/higgs_clone \
--ref-audio /path/to/reference.wav \
--ref-text "Exact transcript spoken in reference.wav." \
--prompts "Hello, this is a cloned voice."
```
### Notes
- Output: 24 kHz mono.
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
- Out of scope (rejected with explicit 4xx): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
## GLM-TTS
2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.
### Launch
```bash
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
# or:
bash examples/online_serving/text_to_speech/glm_tts/run_server.sh /path/to/GLM-TTS
```
### Sending requests
```bash
# Voice cloning (required)
python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
--text "你好,这是语音克隆测试。" \
--ref-audio file:///path/to/ref.wav \
--ref-text "这是参考音频的文本内容。"
# Custom format
python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
--text "Hello, this is a voice cloning test." \
--ref-audio file:///path/to/ref.wav \
--ref-text "Transcript of the reference audio." \
--response-format mp3 -o output.mp3
```
### Gradio demo
```bash
bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
```
### Notes
- Output: 24 kHz mono WAV via HiFT vocoder.
- `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds.
- Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer.
---
## OmniVoice
Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline.
@@ -324,10 +196,10 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
./qwen3_tts/run_server.sh Base
```
### Executor backend
Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override.
### Choosing an executor backend (uniproc vs mp)
Stage configs ship with the chunked-streaming default. To use the uniproc executor (lower IPC overhead for the Base cloning task), pass `--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml`. See [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) and [#2604](https://github.com/vllm-project/vllm-omni/pull/2604) for the full investigation.
To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.
To opt out of chunked streaming, pass `--no-async-chunk` instead — the pipeline auto-dispatches to the end-to-end codec processor.
### Sending requests
```bash
@@ -409,11 +281,71 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
### Notes
- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
---
## VoxCPM
Split-stage TTS at 24 kHz.
### Prerequisites
```bash
pip install voxcpm
# or use a local source tree:
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
```
If the native VoxCPM `config.json` lacks HF `model_type`, set up an HF-compatible config dir:
```bash
export VOXCPM_MODEL=/path/to/voxcpm-model
export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
```
### Launch
```bash
export VOXCPM_MODEL=/path/to/voxcpm-model
./voxcpm/run_server.sh # async-chunk streaming (default)
./voxcpm/run_server.sh sync # non-streaming
```
Or directly:
```bash
vllm serve "$VOXCPM_MODEL" \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
--trust-remote-code --enforce-eager --omni --port 8091
```
### Sending requests
```bash
# Basic TTS
python voxcpm/openai_speech_client.py \
--model "$VOXCPM_MODEL" \
--text "This is a VoxCPM online text-to-speech example."
# Voice cloning
python voxcpm/openai_speech_client.py \
--model "$VOXCPM_MODEL" \
--text "This sentence is synthesized with a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "The exact transcript spoken in reference.wav."
# Streaming PCM
python voxcpm/openai_speech_client.py \
--model "$VOXCPM_MODEL" \
--text "This is a streaming VoxCPM request." \
--stream --output voxcpm_stream.pcm
```
### Notes
- `voxcpm.yaml` for one-shot decode; `voxcpm_async_chunk.yaml` for single-request streaming. Do not use the async-chunk config for concurrent requests or `/v1/audio/speech/batch`.
- Generic TTS fields not supported by VoxCPM: `voice`, `instructions`, `language`, `speaker_embedding`, `x_vector_only_mode`.
- For benchmark measurement, see [`benchmarks/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/benchmarks/voxcpm/README.md).
---
## VoxCPM2
@@ -475,14 +407,6 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
## Example materials
??? abstract "cosyvoice3/run_server.sh"
--8<-- "examples/online_serving/text_to_speech/cosyvoice3/run_server.sh"
??? abstract "cosyvoice3/speech_client.py"
--8<-- "examples/online_serving/text_to_speech/cosyvoice3/speech_client.py"
??? abstract "fish_speech/gradio_demo.py"
``````py
--8<-- "examples/online_serving/text_to_speech/fish_speech/gradio_demo.py"
@@ -499,29 +423,6 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
``````py
--8<-- "examples/online_serving/text_to_speech/fish_speech/speech_client.py"
``````
??? abstract "higgs_audio_v2/batch_speech_client.py"
``````py
--8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py"
``````
??? abstract "higgs_audio_v2/run_server.sh"
``````sh
--8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh"
??? abstract "glm_tts/gradio_demo.py"
``````py
--8<-- "examples/online_serving/text_to_speech/glm_tts/gradio_demo.py"
``````
??? abstract "glm_tts/openai_speech_client.py"
``````py
--8<-- "examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py"
``````
??? abstract "glm_tts/run_gradio_demo.sh"
``````sh
--8<-- "examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh"
``````
??? abstract "glm_tts/run_server.sh"
``````sh
--8<-- "examples/online_serving/text_to_speech/glm_tts/run_server.sh"
``````
??? abstract "omnivoice/run_server.sh"
``````sh
--8<-- "examples/online_serving/text_to_speech/omnivoice/run_server.sh"
@@ -562,6 +463,14 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
``````py
--8<-- "examples/online_serving/text_to_speech/qwen3_tts/tts_common.py"
``````
??? abstract "voxcpm/openai_speech_client.py"
``````py
--8<-- "examples/online_serving/text_to_speech/voxcpm/openai_speech_client.py"
``````
??? abstract "voxcpm/run_server.sh"
``````sh
--8<-- "examples/online_serving/text_to_speech/voxcpm/run_server.sh"
``````
??? abstract "voxcpm2/gradio_demo.py"
``````py
--8<-- "examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py"
+2 -4
View File
@@ -32,9 +32,7 @@ guide. AutoRound is Intel-supported.
|-------|------------|-------|--------|---------|
| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
| Qwen-Image | Not listed | Diffusion transformer | W4A16 | Not validated |
| Wan2.2-I2V | `Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
| Wan2.2-T2V | `Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
| Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
| Wan2.2 | Not listed | Diffusion transformer | W4A16 | Not validated |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
@@ -52,8 +50,8 @@ vLLM-Omni's runtime module names.
| Model | Scope | Status | Notes |
|-------|-------|--------|-------|
| GLM-Image | Diffusion transformer | ✅ | `Intel/GLM-Image-int4-AutoRound` |
| BAGEL | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
| GLM-Image | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
## Configuration
+6 -6
View File
@@ -2,11 +2,11 @@
## Overview
FP8 quantization converts BF16/FP16 weights to FP8 at model load time. Online
activation scaling is the default and does not require calibration. Static
activation scaling is supported when calibrated scale information is available.
For ModelOpt-produced pre-quantized checkpoints, see
[ModelOpt Quantization](modelopt.md).
FP8 quantization converts BF16/FP16 weights to FP8 at model load time, or loads
a checkpoint whose target stage already declares an FP8 quantization config.
Online activation scaling is the default and does not require calibration.
Static activation scaling is supported when calibrated scale information is
available.
Some architectures can quantize all linear layers. Others have
quality-sensitive layers that should stay in BF16 through `ignored_layers`.
@@ -46,7 +46,7 @@ guide. FP8 on Ampere may use a weight-only path where available.
| Model | Scope | Format | Status |
|-------|-------|--------|--------|
| Qwen3-Omni | Thinker language-model stage | [ModelOpt](modelopt.md) `quant_algo=FP8` | Tested for thinker memory reduction |
| Qwen3-Omni | Thinker language-model stage | ModelOpt `quant_algo=FP8` | Tested for thinker memory reduction |
| Qwen3-TTS | TTS language-model stage | Checkpoint config | Not validated |
Audio encoder, vision encoder, talker, and code2wav stay in BF16 unless a
-191
View File
@@ -1,191 +0,0 @@
# ModelOpt Quantization
## Overview
ModelOpt quantization loads checkpoints produced by NVIDIA ModelOpt. The
quantized weights and scale tensors are generated before serving, so inference
does not run online calibration or convert a BF16 checkpoint at startup.
vLLM-Omni validates ModelOpt FP8, ModelOpt NVFP4, and ModelOpt mixed
FP8/NVFP4 checkpoint loading for diffusion transformer stages. The loader
auto-detects supported ModelOpt checkpoint configs and keeps non-transformer
components, such as the tokenizer, scheduler, text encoder, vision/audio
encoder, and VAE, on the base checkpoint unless a model-specific guide says
otherwise.
!!! note
ModelOpt checkpoints are pre-quantized checkpoints. Do not pass
`--quantization fp8` for these checkpoints. The checkpoint
`quantization_config` selects the ModelOpt path.
!!! note
`--force-cutlass-fp8`, `--linear-backend cutlass`, and
`--moe-backend cutlass` are runtime backend selections for checkpoints that
already carry supported ModelOpt quantized weights and scales. They do not
quantize BF16 checkpoints at startup.
## Supported ModelOpt Checkpoint Formats
vLLM-Omni treats ModelOpt checkpoints as pre-quantized checkpoints. The
checkpoint config must identify ModelOpt as the quantization method or producer,
and the quantization algorithm must be one of the validated algorithms below.
| Checkpoint field | Supported value |
|------------------|-----------------|
| `method` / `quant_method` | `modelopt`, `modelopt_fp4`, `modelopt_mixed` |
| `producer.name` | `modelopt` |
| `quant_algo` | `FP8`, `FP8_PER_CHANNEL_PER_TOKEN`, `NVFP4`, `MIXED_PRECISION` |
| `quant_algo` | Runtime method | Typical use |
|--------------|----------------|-------------|
| `FP8`, `FP8_PER_CHANNEL_PER_TOKEN` | `modelopt` | FP8 diffusion transformer checkpoints |
| `NVFP4` | `modelopt_fp4` | NVFP4 diffusion transformer checkpoints |
| `MIXED_PRECISION` | `modelopt_mixed` | Mixed FP8/NVFP4 checkpoints with a ModelOpt per-layer policy |
For multi-component diffusion or omni models, only the checkpoint component
that contains ModelOpt quantized weights should use the ModelOpt quantization
method. Encoders, decoders, tokenizers, schedulers, and other BF16 components
stay unquantized unless the model-specific recipe validates otherwise.
## Hardware Support
| Device | Support |
|--------|---------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
| NVIDIA Ampere GPU (SM 80+) | ⭕ |
| AMD ROCm | ⭕ |
| Intel XPU | ⭕ |
| Ascend NPU | ❌ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
The optional CUTLASS FP8 runtime override requires CUDA SM89+. ModelOpt NVFP4
and mixed FP8/NVFP4 diffusion checkpoints are currently validated on Blackwell
CUDA systems in the recipes below; other CUDA generations require separate
backend and quality validation.
## Model Type Support
### Diffusion Model
| Model | HF checkpoint | Scope | Status |
|-------|---------------|-------|--------|
| Qwen-Image 2512 | `feizhai123/qwen-image-2512-modelopt-fp8-dynamic-all` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
| Qwen-Image 2512 | `feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy` | Diffusion transformer | Validated for ModelOpt mixed FP8/NVFP4 checkpoints |
| Z-Image | `feizhai123/z-image-modelopt-fp8-conservative` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
| FLUX.2-dev | `feizhai123/flux2-dev-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
| FLUX.2-klein 4B | `feizhai123/flux2-klein-4b-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
| HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-fp8` | MoE diffusion transformer | Validated for ModelOpt FP8 checkpoints |
| HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8` | MoE diffusion transformer | Validated for ModelOpt mixed FP8/NVFP4 checkpoints |
| Wan2.2 | Not available | Diffusion transformer | Not validated |
For full serving commands and benchmark context, see
[`recipes/Qwen/Qwen-Image.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image.md)
and
[`recipes/Tencent/HunyuanImage-3.0-Instruct.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md).
### Multi-Stage Omni/TTS Model
| Model | Scope | Status |
|-------|-------|--------|
| Qwen3-Omni | Thinker language-model stage | ModelOpt FP8 checkpoint path |
| Qwen3-TTS | TTS language-model stage | Not validated |
Audio encoder, vision encoder, talker, and code2wav stages stay in BF16 unless
a model-specific guide documents otherwise.
### Multi-Stage Diffusion Model
ModelOpt checkpoints must be routed to the stage whose checkpoint contains the
ModelOpt `quantization_config`. BAGEL and GLM-Image are not listed as validated
ModelOpt targets yet.
## Configuration
For pre-quantized ModelOpt checkpoints, no `--quantization fp8` flag is needed.
The checkpoint config selects the ModelOpt path.
Online serving:
```bash
vllm serve <modelopt-checkpoint> \
--omni \
--tensor-parallel-size <N> \
--linear-backend cutlass \
--force-cutlass-fp8
```
For mixed FP8/NVFP4 MoE checkpoints, also select the validated MoE backend:
```bash
vllm serve <modelopt-mixed-moe-checkpoint> \
--omni \
--tensor-parallel-size <N> \
--enable-expert-parallel \
--linear-backend cutlass \
--moe-backend cutlass \
--force-cutlass-fp8
```
Offline inference:
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model <modelopt-checkpoint> \
--tensor-parallel-size <N> \
--prompt "a red ceramic teapot on a wooden table" \
--height 1024 \
--width 1024 \
--num-inference-steps 20 \
--seed 42 \
--output outputs/modelopt.png
```
Python API:
```python
from vllm_omni import Omni
omni = Omni(
model="<modelopt-checkpoint>",
tensor_parallel_size=2,
force_cutlass_fp8=True,
)
```
## Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `force_cutlass_fp8` / `--force-cutlass-fp8` | bool | `False` | Force CUTLASS FP8 linear kernels for supported ModelOpt FP8 diffusion stages on CUDA SM89+ |
| `--linear-backend cutlass` | str | auto | Select the validated CUTLASS linear backend for supported ModelOpt NVFP4 or mixed FP8/NVFP4 diffusion stages |
| `--moe-backend cutlass` | str | auto | Select the validated CUTLASS MoE backend for supported ModelOpt mixed MoE checkpoints |
## Validation and Notes
1. Compare the ModelOpt checkpoint against the BF16 baseline with the same
prompt, resolution, seed, and inference steps.
2. Use `tests/diffusion/quantization/test_quantization_quality.py` with
`VLLM_OMNI_QUALITY_CONFIGS` to validate local baseline and quantized model
paths.
3. For HunyuanImage-3.0 quantized DiT checkpoints, the opt-in accuracy check is:
```bash
CUDA_VISIBLE_DEVICES=2,3 \
HUNYUAN_IMAGE3_RUN_QUANT_ACCURACY=1 \
HUNYUAN_IMAGE3_QUANT_DEVICES=0,1 \
HUNYUAN_IMAGE3_QUANT_TP=2 \
HUNYUAN_IMAGE3_BF16_MODEL=/path/to/hunyuan-image3-bf16 \
HUNYUAN_IMAGE3_FP8_MODEL=/path/to/hunyuan-image3-modelopt-fp8 \
HUNYUAN_IMAGE3_NVFP4_MODEL=/path/to/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8 \
PYTHONPATH=/path/to/vllm-omni:${PYTHONPATH:-} \
python -m pytest -s -v \
tests/e2e/accuracy/test_hunyuan_image3.py \
-k quantized_dit_matches_bf16_accuracy
```
4. Report CLIP score deltas, SSIM, PSNR, throughput, latency, and peak memory
when adding a new validated ModelOpt diffusion checkpoint.
5. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this
ModelOpt path only when the checkpoint already contains ModelOpt quantized
weights and scales.
-479
View File
@@ -1,479 +0,0 @@
# W4A4 MXFP4 Quantization
## Overview
W4A4 MXFP4 (Microscaling FP4) quantizes both weights and activations to FP4
(`float4_e2m1fn_x2`, packed 2 values per byte) using the OCP MX format: groups
of 32 K-dimension elements share a single `float8_e8m0fnu` exponent scale.
vLLM-Omni provides two quantization methods with different scale structures:
| Method | Scale structure | Mode | Use case |
|--------|----------------|------|----------|
| `mxfp4` | Single-scale (per-32 fine only) | Online only | Quick accuracy baseline; no checkpoint prep needed |
| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | Production; better accuracy; offline recommended |
!!! tip "Recommended: `mxfp4_dualscale` offline"
For production deployments, use the `mxfp4_dualscale` offline mode with a
pre-quantized checkpoint produced by msModelSlim. Offline checkpoints load
calibrated `mul_scale` tensors from disk, providing measurably better accuracy
than any online method. The one-time preprocessing cost amortises across all
subsequent inference runs.
Use `mxfp4` online only for quick experimentation where preprocessing time
is not acceptable and accuracy loss is tolerable.
!!! warning "Online single-scale ≠ Offline dual-scale"
`mxfp4_dualscale` offline mode uses `NPUMxfp4DualScaleLinearMethod`:
fine scale (per-32 K), coarse scale (per-512 K), and per-input-channel
`mul_scale` from calibration — all loaded from the checkpoint.
`mxfp4_dualscale` online mode uses `NPUMxfp4DualScaleOnlineLinearMethod`:
dual-level scales computed on the fly from BF16 weights; no calibration
`mul_scale` is available. Loading an offline checkpoint with the online
method (or vice versa) will produce incorrect results or shape errors.
## Hardware Support
| Device | Support |
|--------|---------|
| NVIDIA Blackwell GPU (SM 100+) | ⭕ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
| NVIDIA Ampere GPU (SM 80+) | ⭕ |
| AMD ROCm | ⭕ |
| Intel XPU | ⭕ |
| Ascend NPU (Atlas 950 A5) | ✅ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
## Model Type Support
### Diffusion Model (Wan2.2)
| Model | Online | Offline | Notes |
|-------|--------|---------|-------|
| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade (`transformer` + `transformer_2`); both transformers quantized with the same config |
| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade; same scheme as T2V-A14B |
| Wan2.2-TI2V-5B | ❌ | ❌ | Parameter count too small; W4A4 causes unacceptable accuracy loss |
The choice between `mxfp4` and `mxfp4_dualscale` in **online mode** is about
quantization quality, not model compatibility — both work on cascade (A14B) and
single-transformer models alike, the same as `mxfp8` online:
- `mxfp4`: single-scale, lower overhead, simpler compute, online only
- `mxfp4_dualscale`: dual-scale + optional BF16 fallback, better accuracy, online **and** offline
**Offline** checkpoints for A14B are always in `mxfp4_dualscale` format (produced
by the merge script); there is no offline `mxfp4` single-scale format.
!!! note "Per-layer BF16 fallback in offline cascade models"
The A14B offline checkpoint uses `quant_method: mxfp4_dualscale`. Most
linear layers are stored as W4A4 MXFP4 DualScale; precision-sensitive layers
retain their original BF16 weights and are listed in `ignored_layers` inside
each transformer's `config.json`. The two transformers may have different
`ignored_layers` sets — the pipeline reads each transformer's own `config.json`
and rebuilds the config locally when they differ, so routing is always
per-transformer-accurate.
!!! warning "TI2V-5B not supported"
Wan2.2-TI2V-5B is excluded from W4A4 quantization. Its smaller parameter
count makes it significantly more sensitive to 4-bit quantization noise,
resulting in unacceptable accuracy loss. Use [MXFP8](mxfp8.md) for TI2V-5B.
## Configuration
### `mxfp4` — Single-Scale Online Mode
Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to
MXFP4 at load time using `npu_dynamic_mx_quant`. A single block scale
(`float8_e8m0fnu`, one per 32 K elements) is computed on the fly; no
calibration `mul_scale` is available. Applies equally to single-transformer
and cascade (A14B) models — both transformers in a cascade receive the same
quantization config automatically.
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(model="<your-model>", quantization="mxfp4")
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
```bash
# Single-transformer or cascade model — same command
python text_to_video.py --model <your-model> --quantization mxfp4
# Online serving
vllm serve <your-model> --omni --quantization mxfp4
```
### `mxfp4_dualscale` — DualScale Online Mode
Online DualScale mode computes both fine and coarse scales on the fly from BF16
weights using `npu_dynamic_dual_level_mx_quant`. Applies equally to
single-transformer and cascade (A14B) models. Compared to `mxfp4` online,
DualScale provides better quantization accuracy at higher compute cost.
The default configuration keeps the leading 5 transformer blocks in BF16
(`num_bf16_fallback_layers=5`). Accuracy evaluation on Wan2.2-A14B shows this
is sufficient to meet quality requirements and is the recommended setting.
```python
omni = Omni(model="<your-model>", quantization="mxfp4_dualscale")
```
```bash
python text_to_video.py --model <your-model> --quantization mxfp4_dualscale
```
If accuracy debugging identifies additional precision-sensitive layers, they
can be pinned to BF16 via the Python API:
```python
omni = Omni(
model="<your-model>",
quantization_config={
"method": "mxfp4_dualscale",
"ignored_layers": ["blocks.10.attn1.to_q"], # explicit per-layer override
},
)
```
BF16 fallback routing in online mode applies two rules in priority order:
1. **`ignored_layers`** (explicit per-layer override): any layer whose prefix
matches is kept in BF16 regardless of block index.
2. **`num_bf16_fallback_layers`** (coarse leading-block rule): the first N
transformer blocks (`blocks.0` … `blocks.N-1`) fall back to BF16. Defaults
to `5` (recommended). Layers outside `blocks.N.*`
(e.g. `condition_embedder`) are always quantized.
### `mxfp4_dualscale` — DualScale Offline Mode (Recommended)
Offline mode loads a pre-quantized DualScale checkpoint from msModelSlim. A
preprocessing step converts the raw quantized output to the diffusers format
expected by vLLM-Omni and injects the quantization config into each
`transformer/config.json` so that vLLM-Omni auto-detects the offline path
without a `--quantization` flag.
BF16 fallback layers may be interleaved anywhere in the transformer — they are
not restricted to leading blocks. The merge script detects them from
`quant_model_description.json` and writes their prefixes into `ignored_layers`
inside `config.json`. At runtime, each layer's prefix is matched against
`ignored_layers` to decide BF16 vs. MXFP4 DualScale.
#### Checkpoint tensor layout
Each quantized linear layer stores four tensors:
| Tensor | Shape | dtype | Description |
|--------|-------|-------|-------------|
| `weight` | `(N, K)` | float8_e4m3fn | FP4 packed (2 values per byte) |
| `weight_scale` | `(N, K//32)` | uint8 | Fine block scale (`float8_e8m0fnu` bit pattern) |
| `weight_dual_scale` | `(N, K//512, 1)` | float32 | Coarse block scale |
| `mul_scale` | `(K,)` | float32 | Per-input-channel smooth pre-scale (from calibration) |
BF16 fallback layers have no quantization tensors; only the original `weight`
(and optional `bias`) are present, loaded directly from the base checkpoint.
#### Step 1 — Quantize with msModelSlim
```bash
msmodelslim quant \
--model_path /path/to/Wan2.2-T2V-A14B-Diffusers \
--save_path /path/to/wan2_2_t2v_quantized_raw \
--device npu \
--model_type Wan2_2 \
--config_path /path/to/wan2_2_w4a4_mxfp4_dualscale.yaml \
--trust_remote_code True
```
After this step, `--save_path` contains raw quantized safetensors files,
scale files, and a metadata JSON (`quant_model_description*.json`).
For cascade MoE models (T2V-A14B, I2V-A14B), msModelSlim outputs two
subdirectories: `high_noise_model/` (transformer) and `low_noise_model/`
(transformer_2).
#### Step 2 — Preprocess with merge_mxfp4_dualscale_checkpoint.py
The script (`vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py`):
1. Copies the original diffusers model to `--output-path` (VAE, text encoder,
scheduler, etc. are preserved).
2. Remaps tensor names from msModelSlim convention to diffusers convention and
strips `.linear.` / `.div.` wrappers added by the quantization tool.
3. Overlays MXFP4 tensors (weight, fine/coarse scales, `mul_scale`) onto the
BF16 base checkpoint. Non-quantized layers keep their original BF16 weights.
4. Detects all linear layers that remain in BF16 and writes their prefixes into
`ignored_layers` in `config.json`.
5. Injects `quantization_config` so vLLM-Omni auto-detects offline MXFP4
DualScale.
For cascade MoE models, steps 2–5 run separately for each transformer.
```bash
python vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py \
--model-type Wan2.2-T2V-A14B \
--original-model /path/to/Wan2.2-T2V-A14B-Diffusers \
--quant-path /path/to/wan2_2_t2v_quantized_raw \
--output-path /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale
```
| Argument | Description |
|----------|-------------|
| `--model-type` | Model variant: `Wan2.2-T2V-A14B` or `Wan2.2-I2V-A14B` |
| `--original-model` | Root directory of the original BF16 diffusers model |
| `--quant-path` | Root directory of the msModelSlim quantized output |
| `--output-path` | Output directory for the merged model (created by the script) |
The script outputs a complete diffusers model directory at `--output-path`,
with each transformer subfolder containing:
- `diffusion_pytorch_model.safetensors` — MXFP4 weights + scale tensors, with BF16 fallback layers from the base checkpoint
- `config.json` — original transformer config with `quantization_config` injected
- `quant_model_description.json` — quantization metadata (reference only)
The `quantization_config` injected into `config.json` for each transformer:
```json
{
"quant_method": "mxfp4_dualscale",
"is_checkpoint_serialized": true,
"ignored_layers": [
"blocks.0.attn1.to_qkv",
"blocks.0.attn1.to_out",
"proj_out"
]
}
```
`ignored_layers` lists every linear layer that retains its original BF16 weight,
using vllm-omni model parameter names (QKV-fused, FFN underscored, `to_out`
unindexed). The exact entries are determined by the quantization tool (msModelSlim)
and may differ between `transformer` and `transformer_2` in a cascade model.
#### Step 3 — Serve
```bash
python text_to_video.py --model /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale
# Online serving
vllm serve /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale --omni
```
```python
omni = Omni(model="/path/to/Wan2.2-T2V-A14B-MXFP4-DualScale")
```
!!! note
No `--quantization` flag is needed for offline mode. The preprocessing
script injects `quantization_config` into each `transformer/config.json`,
which vLLM-Omni reads automatically to activate the correct offline path.
## Parameters
### `mxfp4` (single-scale, online only)
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `method` | str | — | `"mxfp4"` |
| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16 |
### `mxfp4_dualscale` (dual-scale, online + offline)
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `method` | str | — | `"mxfp4_dualscale"` |
| `is_checkpoint_serialized` | bool | `False` | `True` for offline DualScale checkpoints; auto-set from `config.json` when using the preprocessing script |
| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16. **Works in both modes**: offline — populated by the merge script for interleaved sensitive layers; online — user-supplied for explicit per-layer precision override |
| `num_bf16_fallback_layers` | int | `5` | **Online mode only**: leading N transformer blocks (`blocks.0` … `blocks.N-1`) kept in BF16. Applied after `ignored_layers`; ignored in offline mode. Default of `5` is the evaluated recommended value for Wan2.2-A14B |
#### BF16 fallback priority (online mode)
```
for each linear layer:
if prefix in ignored_layers → BF16 (explicit override, highest priority)
elif block_idx < num_bf16_fallback_layers → BF16 (coarse leading-block rule)
else → MXFP4 DualScale online
```
Layers outside `blocks.N.*` (e.g. `condition_embedder.*`) are always quantized
unless they appear in `ignored_layers`.
## Validation and Notes
1. **Online single-scale (`mxfp4`)** quantizes BF16 weights at load time using
`npu_dynamic_mx_quant` (single-scale). No calibration `mul_scale` is
available — all output partitions receive an identity pre-scale. No offline
checkpoint format exists for this method.
2. **Online dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=False`)**
quantizes BF16 weights using `npu_dynamic_dual_level_mx_quant` (fine + coarse
scales computed on the fly). No calibration `mul_scale`; leading blocks or
explicit `ignored_layers` stay in BF16 for accuracy.
3. **Offline dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=True`)** —
**recommended for production** — loads four tensors per quantized layer: FP4
weight, fine scale (`uint8` reinterpreted as `float8_e8m0fnu`), coarse scale
(`float32`), and per-input-channel `mul_scale` (`float32`). BF16 fallback
layers have no quantization tensors and are routed via `ignored_layers`.
4. **Scale dtype**: fine scales are stored as `uint8` in safetensors (same bit
layout as `float8_e8m0fnu`) and reinterpreted at load time without a lossy
float32 round-trip.
5. **Cascade model config propagation**: in a cascade model (transformer +
transformer_2), vLLM-Omni reads each transformer's own `config.json` and
rebuilds the quant config locally when `ignored_layers` differs between
transformers, ensuring per-layer routing is accurate for each. The first
transformer's config is propagated to `od_config` so the second transformer
can reuse it as a starting point.
6. **Self-attention QKV fusion**: Q, K, V projection weights are fused into a
single `QKVParallelLinear` layer at runtime. `ignored_layers` entries use the
fused name (`attn1.to_qkv`), written automatically by the merge script.
7. W4A4 carries higher quantization noise than W8A8 (16 vs 256 levels). The
DualScale offline method mitigates this with calibrated `mul_scale` smooth
quantization. Use `ignored_layers` and `num_bf16_fallback_layers` to trade
off compression vs. accuracy for precision-sensitive layers.
## Adapting MXFP4 for a New Model
This section is aimed at developers who want to add MXFP4 support to a model
other than Wan2.2. The three integration points are: (1) discovering the correct
runtime layer names, (2) wiring `ignored_layers` into the model, and (3) writing
a merge script for offline checkpoints.
### Step 1 — Discover runtime layer names
`ignored_layers` entries must match the **runtime parameter names** used inside
vllm-omni, which may differ from the names stored in the diffusers checkpoint.
The canonical source of truth is the model's own `named_parameters()`.
```python
from vllm_omni import Omni
# Load the model without quantization to inspect parameter names.
omni = Omni(model="/path/to/your-model") # no --quantization flag
for name, _ in omni.pipeline.transformer.named_parameters():
if "weight" in name and "scale" not in name:
print(name)
```
Compare the printed names against the diffusers checkpoint keys
(`safetensors.safe_open` or `torch.load`) to identify any renames your model
applies. Common patterns that differ in Wan2.2 (and may appear in other
models):
| Diffusers checkpoint name | vllm-omni runtime name | Reason |
|---------------------------|------------------------|--------|
| `attn1.to_q`, `attn1.to_k`, `attn1.to_v` | `attn1.to_qkv` | Self-attention Q/K/V fused into `QKVParallelLinear` |
| `ffn.net.0.proj` | `ffn.net_0.proj` | Dots in sub-module names replaced with underscores |
| `ffn.net.2` | `ffn.net_2` | Same underscore rule |
| `to_out.0` | `to_out` | Sequential index stripped |
If your model has different fusion patterns, inspect `packed_modules_mapping`
on the model class — this dict records how checkpoint keys are mapped to
fused runtime parameters.
!!! warning "Partial QKV fallback is not allowed"
If your model fuses Q, K, V into a single layer, `ignored_layers` must
include **all three or none**. A partial fallback (e.g. `to_q` in BF16 but
`to_k`, `to_v` quantized) cannot be expressed at runtime because they share
one `QKVParallelLinear`. The merge script enforces this and raises an error
if only some of the trio appear as non-quantized.
### Step 2 — Add ignored_layers to the model
#### Online mode
Pass `ignored_layers` directly in the quantization config using the **runtime
names** discovered in Step 1. No code changes to the model are required.
```python
omni = Omni(
model="/path/to/your-model",
quantization={
"method": "mxfp4_dualscale",
"ignored_layers": [
"blocks.0.attn1.to_qkv", # runtime name, not diffusers name
"blocks.0.attn1.to_out",
"blocks.0.ffn.net_0.proj",
],
},
)
```
```bash
# CLI does not support list-typed ignored_layers directly.
# Use the Python API or set ignored_layers in config.json (offline).
python your_script.py --model /path/to/your-model --quantization mxfp4_dualscale
```
The `num_bf16_fallback_layers` coarse rule is an alternative to listing layers
individually: set it to N to keep all linear layers in blocks 0 … N-1 in BF16.
The right value depends on the model's sensitivity; evaluate on a validation
set and pick the smallest N that meets your accuracy target.
#### Offline mode
For offline checkpoints, `ignored_layers` is written into each transformer's
`config.json` by the merge script (see Step 3). No manual editing is needed if
the merge script is correct. The injected block:
```json
{
"quant_method": "mxfp4_dualscale",
"is_checkpoint_serialized": true,
"ignored_layers": [
"blocks.0.attn1.to_qkv",
"blocks.0.attn1.to_out"
]
}
```
To add a layer manually (e.g. to pin an additional layer to BF16 without
re-running the merge script), edit `config.json` inside the transformer
subfolder. Use runtime names, not diffusers checkpoint names.
### Step 3 — Write a merge script for offline mode
The merge script for a new model mirrors
`vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py`. The four
things it must do:
1. **Remap tensor names** from the quantization tool convention to diffusers
convention (strip wrappers like `.linear.`, `.div.`; fix any prefix
differences).
2. **Collect ignored_layers**: after loading, enumerate all `*.weight` keys that
have no corresponding `*.weight_scale` (i.e. layers the tool left in BF16).
Convert diffusers names to vllm-omni runtime names (fuse QKV, rename FFN
sub-modules, etc.). Write the result to `config.json`.
3. **Inject `quantization_config`** into `config.json`:
```python
config["quantization_config"] = {
"quant_method": "mxfp4_dualscale",
"is_checkpoint_serialized": True,
"ignored_layers": ignored_layers, # runtime names
}
```
4. **Save** the merged safetensors and the updated `config.json`.
The key helper to implement is the diffusers-to-runtime name translator
(equivalent to `_diffusers_to_vllm_ignored` in the Wan2.2 merge script).
For each non-quantized diffusers weight key, apply your model's specific
renaming rules and collect the results.
!!! tip "Validate before serving"
After producing the offline checkpoint, load it without a `--quantization`
flag and verify that vLLM-Omni auto-detects the correct method. Check that
the layer count reported in the startup log matches expectations: quantized
layer count + `ignored_layers` count should equal total linear layer count.
Any mismatch indicates a name-mapping bug in the merge script.
-220
View File
@@ -1,220 +0,0 @@
# W8A8 MXFP8 Quantization
## Overview
W8A8 MXFP8 (Microscaling FP8) quantizes both weights and activations to FP8
using the OCP MX format: groups of 32 K-dimension elements share a single
`float8_e8m0fnu` exponent scale. This gives better accuracy than channel-wise
FP8 while keeping the same 8-bit weight footprint.
This method supports three modes:
| Mode | Description |
|------|-------------|
| **Online** | BF16 weights are quantized to MXFP8 at load time — no pre-processing needed |
| **Offline (Native)** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
| **Offline (AutoRound)** | AutoRound MXFP8 checkpoints with `data_type="mx_fp"` — auto-detected from `config.json` |
## Hardware Support
| Device | Online | Offline (Native) | Offline (AutoRound) |
|--------|--------|------------------|---------------------|
| NVIDIA Blackwell GPU (SM 100+) | ⭕ | ⭕ | ⭕ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | ⭕ | ⭕ |
| NVIDIA Ampere GPU (SM 80+) | ⭕ | ⭕ | ⭕ |
| AMD ROCm | ⭕ | ⭕ | ⭕ |
| Intel XPU | ✅ | ❌ | ✅ |
| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ⭕ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
**Note**: Intel XPU only supports AutoRound MXFP8 for offline mode. Use AutoRound quantized checkpoints or online mode for XPU.
## Model Type Support
### Diffusion Model (Wan2.2)
| Model | Mode | Notes |
|-------|------|-------|
| Wan2.2-T2V-A14B | Online + Offline | MoE cascade; quantizes two transformers (`transformer` + `transformer_2`) |
| Wan2.2-I2V-A14B | Online + Offline | MoE cascade; quantizes two transformers (`transformer` + `transformer_2`) |
| Wan2.2-TI2V-5B | Online + Offline | Single transformer |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
| Model | Status | Notes |
|-------|--------|-------|
| Qwen3-Omni | Not validated | — |
| Qwen3-TTS | Not validated | — |
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
| Model | Status | Notes |
|-------|--------|-------|
| BAGEL | Not validated | — |
| GLM-Image | Not validated | — |
## Configuration
### Online Mode
Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to
MXFP8 at load time.
Python API:
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(model="<your-model>", quantization="mxfp8")
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
CLI:
```bash
python text_to_video.py --model <your-model> --quantization mxfp8
# Online serving
vllm serve <your-model> --omni --quantization mxfp8
```
### Offline Mode (Native)
Native offline mode loads a pre-quantized checkpoint from msModelSlim. A preprocessing
step converts the raw quantized output to the diffusers format expected by
vLLM-Omni and injects the quantization config into `transformer/config.json` so
that vLLM-Omni auto-detects the offline path without a `--quantization` flag.
#### Step 1 — Quantize with msModelSlim
```bash
msmodelslim quant \
--model_path /path/to/Wan2.2-TI2V-5B-Diffusers \
--save_path /path/to/wan2_2_ti2v_quantized_raw \
--device npu \
--model_type Wan2_2 \
--config_path /path/to/wan2_2_w8a8f8_mxfp.yaml \
--trust_remote_code True
```
After this step, `--save_path` contains the raw quantized safetensors files and
a metadata JSON (`quant_model_description*.json`).
For cascade MoE models (T2V-A14B, I2V-A14B), msModelSlim outputs two
subdirectories: `high_noise_model/` and `low_noise_model/`.
#### Step 2 — Preprocess with merge_mxfp8_checkpoint.py
The script (`vllm_omni/quantization/tools/merge_mxfp8_checkpoint.py`):
1. Copies the original diffusers model to `--output-path` (VAE, text encoder,
scheduler, etc. are preserved).
2. Remaps tensor names from msModelSlim convention to diffusers convention.
3. Saves the converted weights as `diffusion_pytorch_model.safetensors`.
4. Copies the original `transformer/config.json` and injects
`quantization_config` so that vLLM-Omni auto-detects offline MXFP8.
For cascade MoE models, steps 2–4 run separately for `high_noise_model/` →
`transformer/` and `low_noise_model/` → `transformer_2/`.
```bash
python vllm_omni/quantization/tools/merge_mxfp8_checkpoint.py \
--model-type Wan2.2-TI2V-5B \
--original-model /path/to/Wan2.2-TI2V-5B-Diffusers \
--quant-path /path/to/wan2_2_ti2v_quantized_raw \
--output-path /path/to/Wan2.2-TI2V-5B-MXFP8
```
| Argument | Description |
|----------|-------------|
| `--model-type` | Model variant: `Wan2.2-T2V-A14B`, `Wan2.2-I2V-A14B`, or `Wan2.2-TI2V-5B` |
| `--original-model` | Root directory of the original BF16 diffusers model |
| `--quant-path` | Root directory of the msModelSlim quantized output |
| `--output-path` | Output directory for the merged model (created by the script) |
The script outputs a complete diffusers model directory at `--output-path`,
with each transformer subfolder containing:
- `diffusion_pytorch_model.safetensors` — converted FP8 weights
- `config.json` — original transformer config with `quantization_config` injected
- `quant_model_description.json` — renamed quantization metadata (reference only)
#### Step 3 — Serve
```bash
python text_to_video.py --model /path/to/Wan2.2-TI2V-5B-MXFP8
# Online serving
vllm serve /path/to/Wan2.2-TI2V-5B-MXFP8 --omni
```
Python API:
```python
omni = Omni(model="/path/to/Wan2.2-TI2V-5B-MXFP8")
```
!!! note
No `--quantization` flag is needed for native offline mode. The preprocessing
script injects `quantization_config` into each `transformer/config.json`,
which vLLM-Omni reads automatically to activate the offline MXFP8 method.
### Offline Mode (AutoRound)
AutoRound MXFP8 checkpoints declare `quant_method="auto-round"` with `data_type="mx_fp"` in their `config.json`. These are automatically detected and use the `IncMxfp8OfflineLinearMethod` backend.
To use an AutoRound MXFP8 checkpoint:
```bash
python text_to_video.py --model <autoround-mxfp8-model>
# Online serving
vllm serve <autoround-mxfp8-model> --omni
```
Python API:
```python
omni = Omni(model="<autoround-mxfp8-model>")
```
!!! note
AutoRound MXFP8 checkpoints are auto-detected from `config.json` and do not require a `--quantization` flag. The config must include:
```json
{
"quantization_config": {
"quant_method": "auto-round",
"data_type": "mx_fp",
...
}
}
```
## Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `method` | str | — | Must be `"mxfp8"` |
| `is_checkpoint_mxfp8_serialized` | bool | `False` | `True` for offline pre-quantized checkpoints; auto-set from `config.json` when using the preprocessing script |
| `ignored_layers` | list[str] | `[]` | Layer name substrings to keep in BF16 (e.g. `"to_out"` matches `blocks.0.attn1.to_out.0`) |
## Validation and Notes
1. Online mode quantizes BF16 weights at load time using
`npu_dynamic_mx_quant`. This adds a one-time overhead on the first load
but requires no checkpoint preparation.
2. Offline mode loads FP8 weights directly from the checkpoint. Scales are
stored as `uint8` bytes in safetensors (same bit layout as
`float8_e8m0fnu`) and are reinterpreted at load time without a dtype
conversion.
3. If the offline checkpoint was produced with the old `merge_mxfp8_checkpoint.py`
interface (arguments `--quant-dir`, `--orig-dir`, `--meta-json`,
`--output-dir`), regenerate it with the current script. The old script
wrote a separate `quantization_config.json` that is not read by vLLM-Omni;
the current script injects the config directly into `transformer/config.json`.
+11 -26
View File
@@ -9,24 +9,20 @@ separate quantized checkpoint.
This mode is different from pre-quantized checkpoint formats such as GGUF,
AutoRound, msModelSlim, or serialized Int8 checkpoints. Those formats are
prepared before serving and are documented in their method-specific guides.
For MXFP8 and MXFP4, use this page for load-time quantization from BF16
checkpoints, and use the method-specific pages for offline checkpoints produced
by msModelSlim and the merge tools.
## Hardware Support
| Device | FP8 W8A8 | Int8 W8A8 | MXFP8 W8A8 | MXFP4 W4A4 |
|--------|----------|-----------|------------|------------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ⭕ | ⭕ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ |
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ |
| Intel XPU | ⭕ | ⭕ | ✅ | ⭕ |
| Ascend NPU | ❌ | ✅ | ✅ | ✅ |
| Device | FP8 W8A8 | Int8 W8A8 |
|--------|----------|-----------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ |
| AMD ROCm | ⭕ | ⭕ |
| Intel XPU | ⭕ | ⭕ |
| Ascend NPU | ❌ | ✅ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
guide. FP8 on Ampere may use a weight-only path where available. MXFP8 and
MXFP4 are documented for the Ascend NPU path.
guide. FP8 on Ampere may use a weight-only path where available.
## Model Type Support
@@ -36,8 +32,6 @@ MXFP4 are documented for the Ascend NPU path.
|--------|-------|----------------|--------|
| FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
| Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Validated on Ascend NPU and Intel XPU |
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B | Ascend NPU only; TI2V-5B is not supported |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
@@ -59,9 +53,6 @@ from vllm_omni import Omni
omni_fp8 = Omni(model="<your-model>", quantization="fp8")
omni_int8 = Omni(model="<your-model>", quantization="int8")
omni_mxfp8 = Omni(model="<your-model>", quantization="mxfp8")
omni_mxfp4 = Omni(model="<your-model>", quantization="mxfp4")
omni_mxfp4_dualscale = Omni(model="<your-model>", quantization="mxfp4_dualscale")
```
CLI:
@@ -69,9 +60,6 @@ CLI:
```bash
vllm serve <your-model> --omni --quantization fp8
vllm serve <your-model> --omni --quantization int8
vllm serve <your-model> --omni --quantization mxfp8
vllm serve <your-model> --omni --quantization mxfp4
vllm serve <your-model> --omni --quantization mxfp4_dualscale
```
Per-component routing:
@@ -89,11 +77,10 @@ config = build_quant_config({
| Parameter | Methods | Description |
|-----------|---------|-------------|
| `method` | FP8, Int8, MXFP8, MXFP4 | Quantization method: `"fp8"`, `"int8"`, `"mxfp8"`, `"mxfp4"`, or `"mxfp4_dualscale"` |
| `ignored_layers` | FP8, Int8, MXFP8, MXFP4 | Layer name patterns to keep in BF16/FP16 |
| `method` | FP8, Int8 | Quantization method, usually `"fp8"` or `"int8"` |
| `ignored_layers` | FP8, Int8 | Layer name patterns to keep in BF16/FP16 |
| `activation_scheme` | FP8, Int8 | The runtime value `"dynamic"` selects online activation scaling |
| `weight_block_size` | FP8 | Optional block-wise FP8 weight quantization size |
| `num_bf16_fallback_layers` | MXFP4 DualScale | Leading transformer blocks to keep in BF16 for online `mxfp4_dualscale`; defaults to `5` |
## Validation and Notes
@@ -104,5 +91,3 @@ config = build_quant_config({
model as supported.
4. If a model already ships quantized weights, use the matching pre-quantized
method guide instead of online quantization.
5. For Ascend MXFP4 deployments, prefer offline `mxfp4_dualscale` checkpoints
when production quality is more important than avoiding preprocessing.
+13 -117
View File
@@ -1,4 +1,4 @@
# Quantization
# Quantization
vLLM-Omni exposes quantization through the unified `quantization_config`
path. The same configuration entrypoint is used across diffusion-only models,
@@ -9,20 +9,19 @@ type has a different quantization scope.
| Mode | Guide | Description | Methods |
|------|-------|-------------|---------|
| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8, MXFP4 W4A4 |
| Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA |
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale |
| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8 |
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | GGUF, AutoRound, msModelSlim, serialized Int8 |
## Hardware Support
| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | GGUF | AutoRound | msModelSlim |
|--------|----------|-----------|----------|------------|------------|------|-----------|-------------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ |
| Device | FP8 W8A8 | Int8 W8A8 | GGUF | AutoRound | msModelSlim |
|--------|----------|-----------|------|-----------|-------------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ❌ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ❌ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ✅ | ❌ |
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
| Intel XPU | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
guide. FP8 on Ampere may use a weight-only path where available.
@@ -40,9 +39,6 @@ otherwise.
|--------|-------|------|----------------|--------|
| FP8 W8A8 | [FP8](fp8.md) | Online W8A8 or checkpoint FP8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
| Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
| ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints |
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Online W8A8 or offline pre-quantized | Wan2.2-T2V-A14B, I2V-A14B, TI2V-5B | Ascend NPU only; validated for Wan2.2 |
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | `mxfp4`: online single-scale only; `mxfp4_dualscale`: online or offline dual-scale (offline recommended) | Wan2.2-T2V-A14B, I2V-A14B | Ascend NPU only; validated for Wan2.2 A14B cascade models; TI2V-5B not supported; offline `mxfp4_dualscale` uses calibrated `mul_scale` for best accuracy |
| GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists |
| AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven |
| msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path |
@@ -56,10 +52,8 @@ in BF16 unless the model guide explicitly adds support.
| Method | Guide | Scope | Example models | Status |
|--------|-------|-------|----------------|--------|
| ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
| FP8 | [FP8](fp8.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
| Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints |
| msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
@@ -73,9 +67,6 @@ attached to the intended stage rather than applied globally.
|--------|-------|-------|----------------|--------|
| FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
| Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
| ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation |
| MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
| MXFP4 | [MXFP4](mxfp4.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
| GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed |
| AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed |
| msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
@@ -103,7 +94,7 @@ config = build_quant_config({
| Component | Default quantized? | Notes |
|-----------|--------------------|-------|
| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, MXFP8, MXFP4, GGUF, AutoRound, and msModelSlim |
| Diffusion transformer | Yes | Primary target for FP8, Int8, GGUF, AutoRound, and msModelSlim |
| Text encoder | No | Keep BF16 unless a method-specific guide documents support |
| VAE | No | Keep BF16; storage-only paths are method-specific |
| Scheduler/tokenizer | No | Loaded from the base model repository |
@@ -141,98 +132,3 @@ build_quant_config({"method": "gguf", "gguf_model": "/path/to/model.gguf"})
build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
build_quant_config(None)
```
## Output Similarity Comparison Tool
Use `vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity`
to compare a reference diffusion run with a quantized candidate run using the
same prompt, seed, resolution, scheduler settings, and inference steps. The
tool compares final decoded images or video frames, and also reports generation
latency and worker-reported peak memory when available.
This is useful when validating whether online quantization, an offline
pre-quantized checkpoint, or a new `ignored_layers` choice keeps generation
quality close to the BF16 reference.
### Online Quantization Example
```bash
python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \
--task t2i \
--model Qwen/Qwen-Image \
--candidate-quantization fp8 \
--ignored-layers img_mlp \
--prompt "a cup of coffee on the table" \
--height 512 --width 512 \
--num-inference-steps 20 \
--seed 142 \
--output-json /tmp/qwen_image_fp8_similarity/result.json \
--save-output-dir /tmp/qwen_image_fp8_similarity/images \
--enforce-eager
```
### Offline Checkpoint Example
Use `--candidate-model` when the candidate is already quantized or lives at a
different model path:
```bash
python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \
--task t2i \
--reference-model Qwen/Qwen-Image \
--candidate-model /path/to/qwen-image-fp8-checkpoint \
--prompt "a cup of coffee on the table" \
--height 512 --width 512 \
--num-inference-steps 20 \
--seed 142 \
--output-json /tmp/qwen_image_fp8_checkpoint_similarity/result.json
```
If the checkpoint does not include a loadable quantization config, pass one
explicitly:
```bash
--candidate-quantization-config-json '{"method":"fp8"}'
```
### Output Metrics
The output JSON includes `output_metrics`, `reference_generation`, and
`candidate_generation`.
| Metric | Direction | Meaning |
|--------|-----------|---------|
| `cosine_similarity` | Higher is better | Vector direction similarity between output pixels or frames. Useful as a broad sanity check. |
| `mae` | Lower is better | Mean absolute pixel or frame error. For decoded outputs, values are in uint8 pixel units. |
| `mse` / `rmse` | Lower is better | Squared error and its square root. These penalize localized large differences more than `mae`. |
| `max_abs` | Lower is better | Worst single-element absolute error. Treat it as an outlier/debug signal, not as a release gate. |
| `l2` / `relative_l2` | Lower is better | Absolute and reference-normalized L2 distance. `relative_l2` is easier to compare across resolutions. |
| `psnr_db` | Higher is better | Pixel-space signal-to-noise ratio in dB for uint8 images or frames. |
| `avg_generation_time_s` | Lower is better | Average wall-clock generation time across measured runs. |
| `max_peak_memory_mb` | Lower is better | Maximum worker-reported peak device memory across measured runs, when the worker reports it. |
Recommended starting thresholds for same-seed diffusion comparisons:
| Metric | Smoke threshold | Stricter target | Notes |
|--------|-----------------|-----------------|-------|
| `psnr_db` | `>= 20.0` | `>= 25.0` | Good for quick image or frame regression checks. |
| `mae` | `<= 12.0` | `<= 6.0` | Interpreted in decoded uint8 pixel units. |
| `cosine_similarity` | `>= 0.98` | `>= 0.995` | Less sensitive to global scale than L2-style metrics. |
| `relative_l2` | `<= 0.20` | `<= 0.08` | Useful when comparing across prompts or resolutions. |
These thresholds are heuristics. Tune them by model family, task, resolution,
quantization method, and deployment tolerance. For release gating, pair the
numeric report with visual inspection of saved reference and candidate outputs.
The tool intentionally reports separate quality, latency, and memory metrics
instead of a single consolidated similarity score. A single score can hide
important tradeoffs, for example a candidate with good PSNR but a meaningful
memory regression, or a candidate with low average error but localized visual
artifacts. If you need a project-specific pass/fail gate, define it as an
explicit policy over the individual metrics.
Pixel-level metrics do not measure semantic consistency. For higher-cost
evaluation, you can complement this report with a vision-language judge that
describes the reference and candidate outputs and compares those descriptions.
Keep that semantic check separate from this lightweight tool so users can
choose whether the additional model cost and latency are appropriate.
@@ -1,125 +0,0 @@
# Quantized KV Cache
## Overview
In DiT-based image and video generation, Flash Attention can take a large share
of denoising time, especially for high-resolution or long-frame workloads.
vLLM-Omni supports online FP8 quantization for eligible diffusion Flash
Attention (FA) to reduce FA latency while keeping model weights in their
original dtype.
This feature is configured through `diffusion_kv_cache_dtype` on
`OmniDiffusionConfig` (CLI: `--diffusion-kv-cache-dtype`). It is intentionally
**not** the same as vLLM's `--kv-cache-dtype`, which controls autoregressive
language-model KV cache storage and defaults to `"auto"`. Diffusion FA
quantization uses the dedicated diffusion flags so omni serve does not inherit
that default.
In vLLM-Omni diffusion pipelines, this is a runtime FA path: Q/K/V tensors are
dynamically quantized before the attention operator. It does not quantize model
weights and is separate from [FP8 W8A8](fp8.md), [Int8 W8A8](int8.md), or
pre-quantized checkpoint formats.
If `diffusion_kv_cache_dtype` is not set, behavior is unchanged and attention
runs in the native dtype.
## Hardware Support
| Device | FP8 FA |
|--------|--------|
| Ascend NPU | ✅ |
| NVIDIA GPU | ❌ |
| AMD ROCm | ❌ |
| Intel XPU | ❌ |
Legend: `✅` supported, `❌` unsupported.
FP8 FA is currently implemented only for the NPU Flash Attention backend. Other
backends do not support `diffusion_kv_cache_dtype="fp8"` for diffusion attention
and fall back to native dtype execution.
## Model Type Support
### Diffusion Model
| Model | Scope | Status | Notes |
|-------|-------|--------|-------|
| Wan2.2 | Eligible DiT full-attention FA on Ascend NPU | Tested | Compare quality and latency against a BF16 baseline before production use |
| Other diffusion models | Eligible DiT full-attention FA on Ascend NPU | Not tested | You can try `diffusion_kv_cache_dtype="fp8"`; tune `diffusion_kv_cache_skip_steps` and `diffusion_kv_cache_skip_layers` when higher precision is needed |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
Not tested for FP8 FA. Treat any use as experimental unless a model-specific
guide documents support.
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
Not tested. If the diffusion stage uses the same NPU Flash Attention backend,
`diffusion_kv_cache_dtype` may apply in theory; validate quality and latency for
each stage and model.
## Configuration
Offline diffusion example:
```bash
python examples/offline_inference/image_to_video/image_to_video.py \
--model <your-wan2.2-model> \
--prompt "A cat sitting on a surfboard at the beach" \
--height 1280 \
--width 720 \
--num-frames 61 \
--num-inference-steps 4 \
--ulysses-degree 4 \
--vae-patch-parallel-size 4 \
--diffusion-kv-cache-dtype fp8 \
--diffusion-kv-cache-skip-steps "0,1" \
--diffusion-kv-cache-skip-layers "0-2"
```
Online serving:
```bash
vllm serve <your-model> --omni --diffusion-kv-cache-dtype fp8
```
Stage config:
```yaml
stage_args:
- stage_id: 0
stage_type: diffusion
engine_args:
model_stage: dit
diffusion_kv_cache_dtype: "fp8"
diffusion_kv_cache_skip_steps: "0,1"
diffusion_kv_cache_skip_layers: "0-2"
```
Legacy YAML keys `kv_cache_dtype`, `kv_cache_skip_steps`, and
`kv_cache_skip_layers` are still accepted when constructing
`OmniDiffusionConfig` (for example via `from_kwargs`); prefer the `diffusion_*`
names for new configs.
## Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `diffusion_kv_cache_dtype` | str \| None | `None` | Set to `"fp8"` to enable dynamic FP8 FA on supported attention backends |
| `diffusion_kv_cache_skip_steps` | str \| None | `None` | Denoising step selector to keep in native dtype, for example `"0,1,4-6"` |
| `diffusion_kv_cache_skip_layers` | str \| None | `None` | Transformer layer selector to keep in native dtype, for example `"0-2,10"` |
Selectors use comma-separated integers and inclusive ranges. Listed steps or
layers skip FP8 FA; all other eligible full-attention forwards use the FP8 path.
## Validation and Notes
1. Compare generated images or videos against a BF16 baseline with the same
seed, prompt, resolution, frame count, and denoising steps.
2. Use `diffusion_kv_cache_skip_steps` for denoising steps where quality is more
sensitive.
3. Use `diffusion_kv_cache_skip_layers` for transformer layers that show visible quality
regressions.
4. Report both latency and quality results when enabling this option for a new
model. For image or video models, include visual comparison and quantitative
metrics when available, such as PSNR or SSIM.
@@ -195,25 +195,6 @@ stages:
devices: "0,1"
```
#### Hybrid Sharded Data Parallel (HSDP)
For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
1. **Enable HSDP**: Set `use_hsdp: true`.
2. **Set shard size**: Set `hsdp_shard_size` to the number of GPUs used for sharding (for example, `4`).
3. **Set replicate size**: Usually keep `hsdp_replicate_size: 1` unless you want replicated HSDP groups.
4. **Set devices**: Specify the comma-separated GPU IDs used by the diffusion stage (for example, `"0,1,2,3"`).
Example configuration for HSDP across 4 GPUs:
```yaml
- stage_id: 1
devices: "0,1,2,3"
parallel_config:
use_hsdp: true
hsdp_shard_size: 4
hsdp_replicate_size: 1
```
Then pass the custom deploy YAML:
```bash
@@ -1,76 +0,0 @@
# Covo-Audio-Chat (Offline Inference)
## Setup
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
> **Note**
> Covo-Audio code2wav requires `torchdiffeq`. Install it with: `pip install torchdiffeq`
## Run examples
Get into the example folder:
```bash
cd examples/offline_inference/covo_audio
```
### Audio input chat
Using the default audio asset:
```bash
python end2end.py
```
Using a custom audio file:
```bash
python end2end.py --audio-path /path/to/audio.wav
```
Using a local model:
```bash
python end2end.py -m /path/to/Covo-Audio-Chat --output-dir ./my_output
```
### Command-line Arguments
| Argument | Short | Default | Description |
|----------|-------|---------|-------------|
| `--model-name` | `-m` | `tencent/Covo-Audio-Chat` | Model path or HuggingFace model ID |
| `--text` | `-t` | `请回答这段音频里的问题。` | Text prompt / question for the audio |
| `--audio-path` | `-a` | default audio asset | Path to local audio file |
| `--sampling-rate` | | `16000` | Sampling rate for audio loading (Hz) |
| `--output-dir` | | `./output_audio` | Output directory for generated files |
| `--num-prompts` | | `1` | Number of prompts to generate |
| `--stage-configs-path` | | (auto) | Path to stage configs YAML file |
| `--log-stats` | | `false` | Enable detailed statistics logging |
| `--stage-init-timeout` | | `300` | Stage initialization timeout (seconds) |
| `--batch-timeout` | | `5` | Batching timeout (seconds) |
| `--init-timeout` | | `300` | Overall initialization timeout (seconds) |
| `--shm-threshold-bytes` | | `65536` | Shared memory threshold (bytes) |
## Pipeline
Covo-Audio-Chat uses a 2-stage pipeline:
- **Stage 0 (fused_thinker_talker):** The 7B LLM generates interleaved text and audio tokens in a single autoregressive pass.
- **Stage 1 (code2wav):** A BigVGAN-based vocoder converts the extracted audio codes into a 24kHz WAV waveform.
## Output
The script generates two files per request in the output directory:
- `{request_id}.txt` -- prompt and generated text
- `{request_id}.wav` -- generated audio (24kHz WAV)
## FAQ
If you encounter `ModuleNotFoundError: No module named 'librosa'`, install it with:
```bash
pip install librosa
```
## Environment
- GPU: 1x A100 (80 GiB)
- Stage 0 (7B LLM): ~16 GiB VRAM
- Stage 1 (BigVGAN vocoder): ~2 GiB VRAM
@@ -1,222 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This example shows how to use vLLM-Omni for running offline inference
with the correct prompt format on Covo-Audio-Chat.
Usage:
python end2end.py --audio-path /path/to/audio.wav
"""
import os
import soundfile as sf
from vllm.assets.audio import AudioAsset
from vllm.multimodal.media.audio import load_audio
from vllm.sampling_params import SamplingParams
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.model_executor.models.covo_audio.prompt_utils import (
COVO_AUDIO_INPUT_PREFIX,
build_covo_audio_chat_prompt,
)
SEED = 42
def get_audio_query(
question: str | None = None,
audio_path: str | None = None,
sampling_rate: int = 16000,
) -> dict:
if question is None:
question = "请回答这段音频里的问题。"
user_content = COVO_AUDIO_INPUT_PREFIX + question
prompt = build_covo_audio_chat_prompt(user_content)
if audio_path is None:
audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
else:
import numpy as np
audio_signal, sr = load_audio(audio_path, sr=sampling_rate)
audio_data = (audio_signal.astype(np.float32), sr)
return {
"prompt": prompt,
"multi_modal_data": {"audio": audio_data},
"modalities": ["audio"],
}
def main(args):
query_result = get_audio_query(
question=args.text,
audio_path=args.audio_path,
sampling_rate=args.sampling_rate,
)
omni = Omni(
model=args.model_name,
stage_configs_path=args.stage_configs_path,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
batch_timeout=args.batch_timeout,
init_timeout=args.init_timeout,
shm_threshold_bytes=args.shm_threshold_bytes,
)
# Stage 0: fused_thinker_talker
# stop_token_ids=[151645] (<|im_end|>) and ignore_eos=True are required
# so the model generates interleaved text+audio tokens before stopping.
thinker_sampling_params = SamplingParams(
temperature=0.0,
top_p=1.0,
top_k=-1,
max_tokens=2048,
seed=SEED,
detokenize=True,
repetition_penalty=1.05,
stop_token_ids=[151645],
ignore_eos=True,
)
# Stage 1: code2wav (audio codes, not real token IDs — skip detokenize)
code2wav_sampling_params = SamplingParams(
temperature=0.0,
top_p=1.0,
top_k=-1,
max_tokens=2048,
seed=SEED,
detokenize=False,
repetition_penalty=1.1,
)
sampling_params_list = [
thinker_sampling_params,
code2wav_sampling_params,
]
prompts = [query_result for _ in range(args.num_prompts)]
omni_outputs = omni.generate(prompts, sampling_params_list)
output_dir = args.output_dir
os.makedirs(output_dir, exist_ok=True)
for stage_outputs in omni_outputs:
output = stage_outputs.request_output
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
prompt_text = output.prompt
out_txt = os.path.join(output_dir, f"{request_id}.txt")
lines = [
"Prompt:\n",
str(prompt_text) + "\n",
"vllm_text_output:\n",
str(text_output).strip() + "\n",
]
try:
with open(out_txt, "w", encoding="utf-8") as f:
f.writelines(lines)
except Exception as e:
print(f"[Warn] Failed writing text file {out_txt}: {e}")
print(f"Request ID: {request_id}, Text saved to {out_txt}")
elif stage_outputs.final_output_type == "audio":
request_id = output.request_id
audio_tensor = output.outputs[0].multimodal_output.get("audio")
if audio_tensor is None:
continue
output_wav = os.path.join(output_dir, f"{request_id}.wav")
audio_numpy = audio_tensor.float().detach().cpu().numpy()
if audio_numpy.ndim > 1:
audio_numpy = audio_numpy.flatten()
sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
print(f"Request ID: {request_id}, Audio saved to {output_wav}")
omni.close()
def parse_args():
parser = FlexibleArgumentParser(description="Offline inference demo for Covo-Audio-Chat")
parser.add_argument(
"--model-name",
"-m",
type=str,
default="tencent/Covo-Audio-Chat",
help="Model path or HuggingFace model ID.",
)
parser.add_argument(
"--text",
"-t",
type=str,
default=None,
help="Text prompt / question for the audio.",
)
parser.add_argument(
"--audio-path",
"-a",
type=str,
default=None,
help="Path to local audio file. Uses default asset if not provided.",
)
parser.add_argument(
"--sampling-rate",
type=int,
default=16000,
help="Sampling rate for audio loading (default: 16000).",
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="Path to stage configs YAML file.",
)
parser.add_argument(
"--log-stats",
action="store_true",
default=False,
help="Enable writing detailed statistics.",
)
parser.add_argument(
"--stage-init-timeout",
type=int,
default=300,
help="Timeout for initializing a single stage in seconds.",
)
parser.add_argument(
"--batch-timeout",
type=int,
default=5,
help="Timeout for batching in seconds.",
)
parser.add_argument(
"--init-timeout",
type=int,
default=300,
help="Timeout for initializing stages in seconds.",
)
parser.add_argument(
"--shm-threshold-bytes",
type=int,
default=65536,
help="Threshold for using shared memory in bytes.",
)
parser.add_argument(
"--output-dir",
default="./output_audio",
help="Output directory for generated files.",
)
parser.add_argument(
"--num-prompts",
type=int,
default=1,
help="Number of prompts to generate.",
)
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
main(args)
Loaded 100 of 838 files, more files were not shown because too many files have changed in this diff. Show more