Compare commits
262
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5d38b35143 | ||
|
|
4a35392e25 | ||
|
|
4bebbc38b4 | ||
|
|
b58c24f580 | ||
|
|
0d9d57acd9 | ||
|
|
51b4b1131e | ||
|
|
e818dba016 | ||
|
|
754d2e52fc | ||
|
|
9de9d1f7b5 | ||
|
|
efd955674b | ||
|
|
4d3eed152a | ||
|
|
16a84b29d5 | ||
|
|
b9cb57b631 | ||
|
|
2e8e3057bc | ||
|
|
a715abd447 | ||
|
|
b6bdc5997f | ||
|
|
631251a1f8 | ||
|
|
b6f29ee614 | ||
|
|
0ab1d30056 | ||
|
|
56ca7dd612 | ||
|
|
83bbe39d39 | ||
|
|
5313cf6d48 | ||
|
|
c167b9d691 | ||
|
|
dca369d448 | ||
|
|
f4b28f2398 | ||
|
|
aa1184d737 | ||
|
|
d7ea5d5979 | ||
|
|
e40076e872 | ||
|
|
9a60e11e3d | ||
|
|
fe72d078ca | ||
|
|
0d91fbbbb7 | ||
|
|
ce621be29b | ||
|
|
dd5626e607 | ||
|
|
ac5fbed6c0 | ||
|
|
955fcff828 | ||
|
|
4bca522f01 | ||
|
|
bd4ede391b | ||
|
|
6be59f7d19 | ||
|
|
a33e2eb588 | ||
|
|
c9a8556c24 | ||
|
|
2cdffcea6b | ||
|
|
3f27ffbd4d | ||
|
|
3bf4f2850c | ||
|
|
5e263b6929 | ||
|
|
e1088026fa | ||
|
|
ef7f2f1cd2 | ||
|
|
67e0e10c85 | ||
|
|
2c73cf3522 | ||
|
|
857356d5b7 | ||
|
|
11c4c7f0ff | ||
|
|
26d481fc84 | ||
|
|
77480215f5 | ||
|
|
c4a0990044 | ||
|
|
40a07e0d80 | ||
|
|
0e81ef2870 | ||
|
|
ac69cbd27e | ||
|
|
d460673647 | ||
|
|
f6e3dece09 | ||
|
|
de3a291710 | ||
|
|
c481ccee2b | ||
|
|
b4ab37da22 | ||
|
|
c2a624bec4 | ||
|
|
aca4b7d65c | ||
|
|
65bc968465 | ||
|
|
b968373c88 | ||
|
|
039a09a8e1 | ||
|
|
c83cd45069 | ||
|
|
f8624db93a | ||
|
|
07fd6afb4b | ||
|
|
5b61e7f1f1 | ||
|
|
ce8a7dfd2d | ||
|
|
50fd3a3f85 | ||
|
|
32663f21d5 | ||
|
|
ea4cf77f56 | ||
|
|
6f2ad7b403 | ||
|
|
b85833eeec | ||
|
|
eeb7e698c9 | ||
|
|
e6466cf0e6 | ||
|
|
54277a8dd0 | ||
|
|
4a24a517ab | ||
|
|
5d4b16e6e3 | ||
|
|
3c85ca5536 | ||
|
|
c483a23deb | ||
|
|
2856ff7aea | ||
|
|
56ca5d9a8f | ||
|
|
81ab2f98da | ||
|
|
b25ea13cb0 | ||
|
|
687a44e5c8 | ||
|
|
19f8f42822 | ||
|
|
0a8204cb81 | ||
|
|
b8bd75837e | ||
|
|
576afb6f53 | ||
|
|
1e8dc84150 | ||
|
|
28558cc374 | ||
|
|
282e0b6642 | ||
|
|
1e5f288a91 | ||
|
|
e969d2e99f | ||
|
|
6f784cbc50 | ||
|
|
369a47d5a1 | ||
|
|
b076006c35 | ||
|
|
9f81a0a4b0 | ||
|
|
f36d891ed1 | ||
|
|
354511b805 | ||
|
|
005621ba7f | ||
|
|
4e41b7bc23 | ||
|
|
353ac8a402 | ||
|
|
e49fbd8a2d | ||
|
|
2b85474c23 | ||
|
|
5cf3f7947b | ||
|
|
44cde33eb5 | ||
|
|
a77c56725d | ||
|
|
fde38ca54f | ||
|
|
b64ab05d09 | ||
|
|
a0918ce583 | ||
|
|
f19891e4f1 | ||
|
|
d0b531520d | ||
|
|
1a93818e54 | ||
|
|
bb239fa949 | ||
|
|
703e31fc47 | ||
|
|
e1b30061b6 | ||
|
|
c007d40b14 | ||
|
|
c708aecb18 | ||
|
|
c306511512 | ||
|
|
5fc0bfe01e | ||
|
|
9a53703676 | ||
|
|
bb69cbc9f1 | ||
|
|
33586d845c | ||
|
|
3a26799507 | ||
|
|
6bb18af119 | ||
|
|
21a3035c10 | ||
|
|
e50a066548 | ||
|
|
4c06bb09ed | ||
|
|
5dabbb2f18 | ||
|
|
7491b67446 | ||
|
|
d49356e4ca | ||
|
|
6a837849dd | ||
|
|
5e82d7f311 | ||
|
|
368afea2f2 | ||
|
|
e403acda3f | ||
|
|
04bfbbae61 | ||
|
|
6509303a04 | ||
|
|
68a484589d | ||
|
|
5916a78551 | ||
|
|
38d440109b | ||
|
|
cc38bf195c | ||
|
|
01f500a5bf | ||
|
|
9de3806624 | ||
|
|
704b6bb6e6 | ||
|
|
2cd64cdccb | ||
|
|
01ebc0cdf5 | ||
|
|
77e1e5c01c | ||
|
|
e637726f94 | ||
|
|
bede8e02d8 | ||
|
|
da60663f78 | ||
|
|
0309b3bd19 | ||
|
|
aa0587c407 | ||
|
|
16310c8ef3 | ||
|
|
15a1143695 | ||
|
|
7fc528f7a9 | ||
|
|
fa945b2626 | ||
|
|
ac66282746 | ||
|
|
bfb730ed50 | ||
|
|
359e474db4 | ||
|
|
b1ff695029 | ||
|
|
b9789091e9 | ||
|
|
7f2b1e9386 | ||
|
|
dd0fa02547 | ||
|
|
08cf5a6083 | ||
|
|
4c2bea782e | ||
|
|
b50fe1d8b1 | ||
|
|
954d0642a4 | ||
|
|
c6c7fd9b6d | ||
|
|
49c9b89fac | ||
|
|
7a8b428bdb | ||
|
|
9de131e6b0 | ||
|
|
335e2eb9e4 | ||
|
|
1069a144a8 | ||
|
|
569899e8d9 | ||
|
|
c93359bb35 | ||
|
|
0470e96771 | ||
|
|
539fa57006 | ||
|
|
0016006974 | ||
|
|
239abcb480 | ||
|
|
d34093b360 | ||
|
|
5f53dccd15 | ||
|
|
00c11c8b87 | ||
|
|
a150194b11 | ||
|
|
52f17bf292 | ||
|
|
afa2b09684 | ||
|
|
62c86c9d62 | ||
|
|
8e2e7c74fb | ||
|
|
6bcf2c8894 | ||
|
|
9da06690ea | ||
|
|
e03ab82dd2 | ||
|
|
6e35bfbbae | ||
|
|
6b52db9e2b | ||
|
|
2390350533 | ||
|
|
755a2d9d96 | ||
|
|
c174c95ab0 | ||
|
|
e375b12680 | ||
|
|
52c9ce2dac | ||
|
|
9a1270e6cc | ||
|
|
8e9b497387 | ||
|
|
39e3fd3d12 | ||
|
|
d388583e0b | ||
|
|
7b876324c2 | ||
|
|
b8bd6fbbdb | ||
|
|
37e5954acf | ||
|
|
22ffb85063 | ||
|
|
b33d536f63 | ||
|
|
0fe76857ac | ||
|
|
e8d22d4fa8 | ||
|
|
77e8ba0a7b | ||
|
|
c6edb34d04 | ||
|
|
fb4204ebf6 | ||
|
|
e24f86ae2e | ||
|
|
678af19240 | ||
|
|
5632b78b3d | ||
|
|
345504b68f | ||
|
|
d20a6b5b39 | ||
|
|
9ebfbf9fde | ||
|
|
2e067bbc24 | ||
|
|
19980bb37c | ||
|
|
c8efdbd671 | ||
|
|
95a07f7732 | ||
|
|
47edee1a79 | ||
|
|
05c6c9a0b6 | ||
|
|
d14ea37842 | ||
|
|
1a7fb55b5f | ||
|
|
3a32cd6780 | ||
|
|
f327cd1c89 | ||
|
|
d4cbdff9cd | ||
|
|
8384dead56 | ||
|
|
3aa84c4742 | ||
|
|
fd9a0ee27f | ||
|
|
e3b0afbb97 | ||
|
|
ead87aa521 | ||
|
|
d8cc7a0409 | ||
|
|
a664e40716 | ||
|
|
ee15f39ea7 | ||
|
|
9337bec9ad | ||
|
|
9d1392dcfc | ||
|
|
e18cb89d5a | ||
|
|
554233208f | ||
|
|
e1cdd0c23b | ||
|
|
e82c0a5e1d | ||
|
|
8e1bbc9ab5 | ||
|
|
2903f528d2 | ||
|
|
e6cf26efb7 | ||
|
|
7a34298743 | ||
|
|
b60a661b3f | ||
|
|
339a6da28c | ||
|
|
25b451a552 | ||
|
|
0cdec8210f | ||
|
|
3f504b45d2 | ||
|
|
c2578f7131 | ||
|
|
5f029ee4d5 | ||
|
|
b001151e0d | ||
|
|
ad7c966b1d | ||
|
|
52b53367ef | ||
|
|
c1ba86ab5b | ||
|
|
de4e47237a |
No files matched your search
@@ -56,6 +56,88 @@ 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
|
||||
@@ -91,8 +173,24 @@ upload_pipeline() {
|
||||
|
||||
cd .buildkite
|
||||
|
||||
# Select test definition file: merge suite for main, ready suite for PRs
|
||||
if [[ $BUILDKITE_BRANCH == "main" ]]; then
|
||||
# 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
|
||||
TEST_YAML="test-amd-merge.yml"
|
||||
else
|
||||
TEST_YAML="test-amd-ready.yaml"
|
||||
@@ -140,30 +238,13 @@ fi
|
||||
# Early exit start: skip pipeline if conditions are met
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
# skip pipeline if all changed files are under docs/
|
||||
# Match CUDA Buildkite skip-ci behavior for docs/markdown-only changes.
|
||||
if [[ "${DOCS_ONLY_DISABLE}" != "1" ]]; then
|
||||
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
|
||||
SKIP_CI="$(resolve_skip_ci)"
|
||||
if [[ "${SKIP_CI}" == "1" ]]; then
|
||||
echo "[docs-only] Docs/markdown-only changes detected. Exiting before pipeline upload."
|
||||
exit 0
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
@@ -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.19.0"
|
||||
VLLM_VERSION: "v0.20.0"
|
||||
priority: 100
|
||||
timeout_in_minutes: 60
|
||||
soft_fail: true
|
||||
soft_fail: false
|
||||
@@ -0,0 +1,110 @@
|
||||
# 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"
|
||||
@@ -0,0 +1,108 @@
|
||||
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"
|
||||
# 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"
|
||||
+17
-4
@@ -1,8 +1,9 @@
|
||||
# Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci
|
||||
# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs if nightly-test
|
||||
# / main NIGHTLY so upload-nightly is not skipped together with test-ready/test-merge.
|
||||
#
|
||||
# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite).
|
||||
# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs for the nightly
|
||||
# exception in upload_pipeline_with_skip_ci.sh (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
|
||||
# can be skipped together with test-ready/test-merge. Weekly uploads also depend on image-build.
|
||||
|
||||
# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
|
||||
steps:
|
||||
- label: ":github: Resolve skip-ci & upload pipeline"
|
||||
key: upload-ci-pipeline
|
||||
@@ -61,3 +62,15 @@ steps:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L5 Test — main+WEEKLY=1 (scheduled), or PR with weekly-test label
|
||||
- label: "Upload Weekly Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-weekly-pipeline
|
||||
if: >-
|
||||
(build.branch == "main" && build.env("WEEKLY") == "1") ||
|
||||
(build.branch != "main" && build.pull_request.labels includes "weekly-test")
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-weekly.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
@@ -11,15 +11,6 @@ set -o pipefail
|
||||
export PYTHONPATH=".."
|
||||
|
||||
# Print ROCm version
|
||||
echo "--- Confirming Clean Initial State"
|
||||
while true; do
|
||||
sleep 3
|
||||
if grep -q clean /opt/amdgpu/etc/gpu_state; then
|
||||
echo "GPUs state is \"clean\""
|
||||
break
|
||||
fi
|
||||
done
|
||||
|
||||
echo "--- ROCm info"
|
||||
rocminfo
|
||||
|
||||
@@ -51,18 +42,6 @@ cleanup_docker() {
|
||||
# Call the cleanup docker function
|
||||
cleanup_docker
|
||||
|
||||
echo "--- Resetting GPUs"
|
||||
|
||||
echo "reset" > /opt/amdgpu/etc/gpu_state
|
||||
|
||||
while true; do
|
||||
sleep 3
|
||||
if grep -q clean /opt/amdgpu/etc/gpu_state; then
|
||||
echo "GPUs state is \"clean\""
|
||||
break
|
||||
fi
|
||||
done
|
||||
|
||||
echo "--- Pulling container"
|
||||
## Temporary change to use AMD Docker Hub to store the vllm-omni image
|
||||
# to bypass the rate limit issue with ECR Public Gallery.
|
||||
|
||||
@@ -141,5 +141,5 @@ docker run \
|
||||
"${image_name}" \
|
||||
bash -c '
|
||||
set -e
|
||||
VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
|
||||
'
|
||||
@@ -1,10 +1,18 @@
|
||||
#!/usr/bin/env bash
|
||||
# Evaluate docs-only skip-ci and upload continuation steps from the same `.buildkite/pipeline.yml`
|
||||
# Evaluate docs-only skip-ci and upload continuation steps from the given pipeline YAML
|
||||
# (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_YML="${ROOT}/.buildkite/pipeline.yml"
|
||||
PIPELINE_ARG="${1:-.buildkite/pipeline.yml}"
|
||||
if [[ "${PIPELINE_ARG}" = /* ]]; then
|
||||
PIPELINE_YML="${PIPELINE_ARG}"
|
||||
else
|
||||
PIPELINE_YML="${ROOT}/${PIPELINE_ARG}"
|
||||
fi
|
||||
|
||||
# Prints a single digit to stdout: 1 = skip image CI, 0 = run. Logs go to stderr.
|
||||
is_docs_only_change() {
|
||||
@@ -21,7 +29,7 @@ is_docs_only_change() {
|
||||
if [[ "${file_path}" == *.md ]]; then
|
||||
continue
|
||||
fi
|
||||
if [[ "${file_path}" == "mkdocs.yaml" ]]; then
|
||||
if [[ "${file_path}" == "mkdocs.yml" ]]; then
|
||||
continue
|
||||
fi
|
||||
return 1
|
||||
@@ -104,12 +112,13 @@ import pathlib
|
||||
path = pathlib.Path(os.environ["PIPELINE_YML"])
|
||||
text = path.read_text(encoding="utf-8")
|
||||
sep = "\n---\n"
|
||||
if sep not in text:
|
||||
raise SystemExit(
|
||||
"upload_pipeline_with_skip_ci: .buildkite/pipeline.yml must contain a '\\n---\\n' separator "
|
||||
"(document 1 = bootstrap, document 2 = uploaded steps)"
|
||||
)
|
||||
_, continuation = text.split(sep, 1)
|
||||
# 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
|
||||
|
||||
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
|
||||
@@ -126,7 +135,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.branch != \"main\" && build.pull_request.labels includes \"merge-test\")'"
|
||||
merge_rep = "'(build.branch == \"main\" && build.env(\"NIGHTLY\") != \"1\" && build.env(\"WEEKLY\") != \"1\") || (build.branch != \"main\" && build.pull_request.labels includes \"merge-test\")'"
|
||||
rendered = (
|
||||
continuation
|
||||
.replace("__IMAGE_BUILD_IF__", rep)
|
||||
|
||||
@@ -7,16 +7,9 @@ steps:
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
|
||||
# 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"
|
||||
|
||||
- label: "Diffusion Images API LoRA E2E"
|
||||
agent_pool: mi325_1
|
||||
@@ -24,7 +17,7 @@ steps:
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
|
||||
- timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
agent_pool: mi325_1
|
||||
@@ -37,9 +30,9 @@ steps:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
@@ -74,7 +67,7 @@ steps:
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
|
||||
|
||||
# merge-only tests
|
||||
@@ -84,7 +77,15 @@ steps:
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
|
||||
- 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
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
@@ -103,16 +104,17 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
|
||||
# This is from test-nightly.yml.
|
||||
# This model seems to be important, so I am trying to run it here.
|
||||
# Since this is from test-nightly.yml,
|
||||
# I keep this non-blocking.
|
||||
- label: "Omni Model Test Qwen2-5-Omni"
|
||||
agent_pool: mi325_2
|
||||
agent_pool: mi325_4
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
|
||||
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
agent_pool: mi325_2
|
||||
@@ -120,10 +122,8 @@ steps:
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- timeout 30m 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"
|
||||
- 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"
|
||||
agent_pool: mi325_1
|
||||
@@ -153,17 +153,6 @@ 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
|
||||
@@ -174,11 +163,10 @@ 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_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
|
||||
# - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
|
||||
|
||||
# - label: "Bagel Img2Img Model Test (2/3)"
|
||||
# agent_pool: mi325_1
|
||||
@@ -187,11 +175,10 @@ 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_img2img.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
|
||||
# - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
|
||||
|
||||
# - label: "Bagel Online Serving Test (3/3)"
|
||||
# agent_pool: mi325_1
|
||||
@@ -200,12 +187,23 @@ 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_online.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
|
||||
# - 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'
|
||||
"
|
||||
|
||||
- label: "Voxtral-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
|
||||
@@ -7,7 +7,9 @@ steps:
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
# 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"
|
||||
|
||||
- label: "Voxtral TTS CUDA Unit Test"
|
||||
agent_pool: mi325_1
|
||||
@@ -17,21 +19,13 @@ 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/e2e/offline_inference/test_qwen_image_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
|
||||
- timeout 20m pytest -s -v tests/diffusion/batching/test_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
|
||||
|
||||
- label: "Custom Pipeline Test"
|
||||
agent_pool: mi325_1
|
||||
@@ -52,9 +46,9 @@ steps:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
@@ -90,7 +84,7 @@ steps:
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
@@ -100,6 +94,27 @@ 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
|
||||
@@ -112,19 +127,6 @@ steps:
|
||||
pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
'
|
||||
|
||||
|
||||
# NOTE: This test is not running any thing. It is skipped and deselected.
|
||||
# Currently it is = 1 skipped, 1 deselected, 17 warnings in 0.03s ======
|
||||
# - label: "Omni Model Test Qwen2-5-Omni"
|
||||
# agent_pool: mi325_2
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
# commands:
|
||||
# - export VLLM_LOGGING_LEVEL=DEBUG
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py -m "core_model" --run-level "core_model"
|
||||
|
||||
# - label: "Omni Model Test Qwen3-Omni"
|
||||
# agent_pool: mi325_2
|
||||
# depends_on: amd-build
|
||||
@@ -133,7 +135,6 @@ 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"
|
||||
|
||||
@@ -175,17 +176,6 @@ 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"
|
||||
@@ -195,11 +185,10 @@ 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_text2img.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
# - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
|
||||
# - label: "Bagel Img2Img Model Test"
|
||||
# agent_pool: mi325_1
|
||||
@@ -208,11 +197,10 @@ 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_img2img.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
# - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
|
||||
# - label: "Bagel Online Serving Test"
|
||||
# agent_pool: mi325_1
|
||||
@@ -221,12 +209,11 @@ 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_online.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
# - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
|
||||
- label: "CosyVoice3-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
@@ -236,5 +223,5 @@ steps:
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts.py -m "core_model" --run-level "core_model"
|
||||
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model"
|
||||
'
|
||||
@@ -8,7 +8,7 @@ steps:
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
|
||||
- pytest -s -v tests/e2e/offline_inference/test_z_image.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/e2e/online_serving/test_images_generations_lora.py
|
||||
- pytest -s -v tests/diffusion/lora/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/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
- pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
- pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
|
||||
- pytest -s -v tests/diffusion/offloader/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/e2e/offline_inference/test_sequence_parallel.py -m core_model
|
||||
- pytest -s -v tests/diffusion/distributed/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/e2e/offline_inference/test_zimage_parallelism.py
|
||||
- pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
timeout_in_minutes: 20
|
||||
@@ -88,7 +88,7 @@ steps:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
|
||||
- pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
@@ -100,25 +100,11 @@ 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
|
||||
@@ -129,5 +115,4 @@ steps:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- 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"
|
||||
- pytest -s -v tests/entrypoints/test_omni_sleep_mode.py -m "advanced_model and omni and MI325" --run-level "advanced_model"
|
||||
+685
-388
File diff suppressed because it is too large.
Load diff
+680
-35
@@ -10,10 +10,10 @@ steps:
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test"
|
||||
steps:
|
||||
- label: ":full_moon: Omni · Function Test"
|
||||
- label: ":full_moon: Omni · Function Test with H100"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and H100 and omni" --run-level "advanced_model"
|
||||
- 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
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -48,11 +48,30 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Omni · Function Test with L4"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: ":full_moon: Omni · Doc Test with L4"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/examples/ -m "advanced_model and omni and L4" --run-level "advanced_model"
|
||||
- pytest -s -v tests/examples/ -m "full_model and omni and L4" --run-level "full_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
@@ -70,7 +89,7 @@ steps:
|
||||
- label: ":full_moon: Omni · Doc Test with H100"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/examples/ -m "advanced_model and omni and H100" --run-level "advanced_model"
|
||||
- pytest -s -v tests/examples/ -m "full_model and omni and H100" --run-level "full_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -105,14 +124,149 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Omni · Perf Test"
|
||||
key: nightly-omni-performance
|
||||
- label: ":full_moon: Omni · Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export SEED_TTS_WER_EVAL=1
|
||||
- export SEED_TTS_EVAL_DEVICE=cuda:1
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py -m "full_model" --run-level full_model
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/e2e/accuracy/qwen3_omni/results/qwen_omni_acc/*.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 · No Async Chunk"
|
||||
key: nightly-omni-performance-no-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_qwen_omni.json
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_no_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 · 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
|
||||
@@ -151,6 +305,44 @@ steps:
|
||||
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"
|
||||
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
|
||||
|
||||
- group: ":card_index_dividers: TTS Model Test"
|
||||
key: nightly-tts-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
@@ -160,7 +352,7 @@ steps:
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
|
||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
@@ -230,10 +422,65 @@ steps:
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2iat-test"
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100"
|
||||
# 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"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not hunyuan" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
|
||||
- >-
|
||||
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"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -268,10 +515,98 @@ 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
|
||||
-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"
|
||||
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"
|
||||
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: Diffusion X2I(&A&T) · Function Test with L4"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not hunyuan" -m "advanced_model and diffusion and L4" --run-level "advanced_model"
|
||||
- pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
@@ -289,8 +624,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 "advanced_model and example and H100" --run-level "advanced_model"
|
||||
- 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"
|
||||
plugins:
|
||||
@@ -328,7 +662,7 @@ steps:
|
||||
- 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 advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
|
||||
- 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"
|
||||
@@ -367,7 +701,7 @@ steps:
|
||||
- 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 advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
|
||||
- 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:
|
||||
@@ -409,7 +743,8 @@ steps:
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level advanced_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2i_matches_diffusers' -m full_model --run-level full_model
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -446,28 +781,218 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test"
|
||||
key: nightly-diffusion-x2iat-performance
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image
|
||||
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
|
||||
# [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
|
||||
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
|
||||
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
|
||||
EXIT4=$$?
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$((EXIT1 | EXIT2 | EXIT3 | EXIT4))
|
||||
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
|
||||
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-2509"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2509
|
||||
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_2509_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
|
||||
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=$$?
|
||||
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
|
||||
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) · HunyuanImage3 · DiT Perf Test"
|
||||
key: nightly-hunyuan-image3-performance
|
||||
soft_fail: true
|
||||
timeout_in_minutes: 120
|
||||
if: build.env("RUN_HUNYUAN_IMAGE3_PERF") == "1"
|
||||
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_hunyuan_image_tp4_fp8.json
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_sp2.json
|
||||
EXIT2=$$?
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_cfgp2.json
|
||||
EXIT3=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$((EXIT1 | EXIT2 | EXIT3))
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -514,7 +1039,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 -m "advanced_model" --run-level "advanced_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py --run-level "full_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -552,7 +1077,9 @@ steps:
|
||||
- label: ":full_moon: Diffusion X2V · Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m advanced_model --run-level advanced_model
|
||||
- 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/test_ltx2_3_video_similarity.py -m full_model --run-level full_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"
|
||||
plugins:
|
||||
@@ -587,6 +1114,112 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Perf Test"
|
||||
key: nightly-diffusion-x2v-performance
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json
|
||||
EXIT1=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT1
|
||||
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
|
||||
|
||||
- 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"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- timeout 30m 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
|
||||
@@ -604,7 +1237,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
|
||||
@@ -633,9 +1266,15 @@ steps:
|
||||
- label: ":email: Nightly Collection & Email"
|
||||
key: nightly-perf-distribution
|
||||
depends_on:
|
||||
- nightly-omni-performance
|
||||
- nightly-omni-performance-no-async-chunk
|
||||
- nightly-omni-performance-async-chunk
|
||||
- nightly-omni-performance-vllm-text
|
||||
- nightly-tts-performance
|
||||
- nightly-diffusion-x2iat-performance
|
||||
- nightly-diffusion-x2iat-performance-qwen-image
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit-2509
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
- nightly-diffusion-x2v-performance
|
||||
- nightly-testcase-statistics
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
@@ -643,8 +1282,14 @@ steps:
|
||||
- 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
|
||||
- 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-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-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-2509
|
||||
- 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-x2v-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
|
||||
- python tools/nightly/generate_nightly_perf_excel.py
|
||||
- python tools/nightly/generate_nightly_perf_html.py
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
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"
|
||||
+451
-549
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,85 @@
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HF_HUB_DOWNLOAD_TIMEOUT: 300
|
||||
HF_HUB_ETAG_TIMEOUT: 60
|
||||
|
||||
steps:
|
||||
- label: "Reliability Test - qwen3-omni"
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-weekly-pipeline
|
||||
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
|
||||
commands:
|
||||
- pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
|
||||
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: "Reliability Test - wan22"
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-weekly-pipeline
|
||||
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
|
||||
commands:
|
||||
- pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
|
||||
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
|
||||
@@ -282,10 +282,11 @@ For Omni or custom models, create:
|
||||
|
||||
Required updates:
|
||||
1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table
|
||||
2. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
|
||||
3. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
|
||||
4. `examples/offline_inference/xxx/README.md` — offline example docs
|
||||
5. `examples/online_serve/xxx/README.md` — online serve docs
|
||||
2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table)
|
||||
3. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
|
||||
4. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
|
||||
5. `examples/offline_inference/xxx/README.md` — offline example docs
|
||||
6. `examples/online_serve/xxx/README.md` — online serve docs
|
||||
|
||||
### Step 8: Add E2E Tests (Recommended)
|
||||
|
||||
@@ -512,6 +513,46 @@ After adding parallelism support, update:
|
||||
1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table
|
||||
2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch)
|
||||
|
||||
### Step 11: Add CPU Offload Support
|
||||
|
||||
Implement `SupportsComponentDiscovery` 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
|
||||
|
||||
class YourPipeline(nn.Module, SupportsComponentDiscovery):
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["text_encoder"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae"]
|
||||
_resident_modules: ClassVar[list[str]] = [] # optional
|
||||
```
|
||||
|
||||
- `_dit_modules`: denoising submodules (kept on GPU during diffusion loop)
|
||||
- `_encoder_modules`: encoder/vision submodules (offloaded to CPU during diffusion loop)
|
||||
- `_vae_modules`: VAE(s) (handled by both sequential and layerwise backends)
|
||||
- `_resident_modules`: additional modules to pin on GPU during layerwise
|
||||
offloading (e.g. embedders, connectors). Only used by the layerwise
|
||||
backend. Optional — defaults to `[]`.
|
||||
|
||||
All attribute names support dotted paths for nested submodules
|
||||
(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
|
||||
|
||||
Pipelines without `SupportsComponentDiscovery` fall back to scanning
|
||||
well-known attribute names (`transformer`, `text_encoder`, `vae`,
|
||||
etc.), which fails for non-standard names.
|
||||
|
||||
### Step 12: Performance Profiling
|
||||
|
||||
After verifying correctness and implementing parallelism/caching, profile the model's performance to identify bottlenecks and ensure optimal execution.
|
||||
|
||||
See the [Profiling Single-Stage Diffusion](../../../docs/contributing/profiling.md#3-profiling-single-stage-diffusion) guide for detailed instructions on:
|
||||
1. Using the PyTorch profiler (`profiler: "torch"`) to capture detailed CPU/CUDA traces.
|
||||
2. Using Nsight Systems (`nsys`) with `profiler: "cuda"` for low-overhead CUDA traces.
|
||||
3. Controlling profiling via `omni.start_profile()` and `omni.stop_profile()`.
|
||||
|
||||
---
|
||||
|
||||
## Iterative Development Tips
|
||||
|
||||
@@ -92,6 +92,7 @@ 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,6 +61,7 @@ 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):
|
||||
@@ -75,6 +76,24 @@ 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:
|
||||
@@ -104,6 +123,7 @@ 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. Set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` to match diffusers
|
||||
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`
|
||||
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
|
||||
|
||||
@@ -8,10 +8,90 @@ description: "Integrate a new text-to-speech model into vLLM-Omni from HuggingFa
|
||||
## Overview
|
||||
|
||||
```
|
||||
HF Reference -> Stage Separation -> Online Serving -> Async Chunk -> CUDA Graph
|
||||
(Phase 1) (Phase 2) (Phase 3) (Phase 4) (Phase 5)
|
||||
HF Reference -> Stage Separation -> Online Serving -> Async Chunk -> CUDA Graph -> Pre-commit/DCO
|
||||
(Phase 1) (Phase 2) (Phase 3) (Phase 4) (Phase 5) (Phase 6)
|
||||
```
|
||||
|
||||
Three architecture patterns are supported:
|
||||
|
||||
- **Two-stage pipeline** (e.g. Qwen3-TTS, Fish Speech, CosyVoice3): AR
|
||||
code-predictor → audio decoder, connected via async_chunk for low-latency
|
||||
streaming. Use this for maximum performance.
|
||||
- **Single-stage AR via generator** (e.g. MOSS-TTS-Nano): entire model runs
|
||||
inside one AR worker, streaming audio chunks through a per-request
|
||||
`inference_stream()` generator. Use this when the upstream model bundles AR
|
||||
+ codec inseparably. See [references/single-stage-ar.md](references/single-stage-ar.md).
|
||||
- **Single-stage, vLLM-native base LM + side computation** (e.g. VoxCPM2):
|
||||
the base language model runs under vLLM's PagedAttention as a normal AR
|
||||
model; diffusion / VAE / side computations run outside vLLM and are
|
||||
attached via the runner post-processing hook. This is a distinct pattern
|
||||
from the generator approach above — do not confuse the two.
|
||||
|
||||
The single-stage variants skip Phase 4 (async_chunk) but Phase 5 (CUDA graph)
|
||||
is still encouraged for the inner AR loop.
|
||||
|
||||
## Cross-Cutting Invariants
|
||||
|
||||
These rules apply to every TTS model regardless of architecture (AR vs AR+diffusion, single-stage vs two-stage, codec-based vs VAE-based). They surface repeatedly across PRs — check them at the end of every phase.
|
||||
|
||||
### I1. Streaming output contract
|
||||
|
||||
Pick exactly one per-step semantics for `forward()` and document it in the docstring:
|
||||
|
||||
- **Delta**: yield only new audio samples produced this step. Preferred — linear cost, low memory.
|
||||
- **Cumulative**: re-decode from step 0 every call. O(N²); only acceptable if the codec has no streaming decode path.
|
||||
|
||||
If you choose **delta**, verify the full emit→consolidate→consume chain:
|
||||
|
||||
1. `forward()` returns `{"model_outputs": <new_chunk_only>, ...}`
|
||||
2. `_consolidate_multimodal_tensors()` in `vllm_omni/engine/output_processor.py` concatenates the audio key into one tensor at finish. If it skips the key (`continue`), offline consumers receive only the final chunk. See `output_processor.py` for the concrete list of handled modality keys.
|
||||
3. Streaming consumers (SSE, Gradio) receive per-step deltas; offline consumers (`engine.generate()`) receive a single concatenated tensor.
|
||||
|
||||
Cumulative-vs-delta mismatch is the most common silent bug — offline RTF benchmarks pass, but users hear replays or truncation.
|
||||
|
||||
### I2. Multimodal output consumer hygiene
|
||||
|
||||
`outputs[0].outputs[0].multimodal_output[<key>]` can be any of `Tensor`, `list[Tensor]` (pre-consolidation snapshot), `np.ndarray`, or scalar. When writing tests, examples, and benchmarks:
|
||||
|
||||
- **Never** use `dict.get("a") or dict.get("b")` on tensor values — Python evaluates the tensor's boolean, raising `RuntimeError: Boolean value of Tensor with more than one value is ambiguous`. Use explicit `if x is None` chains.
|
||||
- Always defensively handle the list form: `if isinstance(x, list): x = torch.cat([t.reshape(-1) for t in x], dim=0)`.
|
||||
- Assert `shape` / `dtype` / `duration` explicitly; do not rely on truthiness for presence checks.
|
||||
|
||||
### I3. Hot-loop GPU discipline
|
||||
|
||||
Inside any per-step model loop (AR decode, diffusion solver, CFM Euler, vocoder block loop):
|
||||
|
||||
- No `tensor.item()`, `.cpu()`, or `.tolist()` — each triggers a GPU→CPU sync; at 10 steps × 60 frames × 4 ops that is 2400 syncs per request.
|
||||
- Prefer `dst.copy_(src)` over `dst.fill_(src.item())` when writing a scalar tensor into a buffer.
|
||||
- Prefer `torch.compile(Model.forward, fullgraph=False)` on the whole forward over per-submodule compile — fewer dispatch boundaries, larger fusion regions. Measure before choosing granularity.
|
||||
- No Python-side control flow that depends on tensor values; use `torch.where` / masking instead.
|
||||
|
||||
Profile first, optimize second. See the profiling docs / project memory for the trace-analysis workflow.
|
||||
|
||||
### I4. Validation pyramid
|
||||
|
||||
Offline RTF alone is necessary but not sufficient. Every new TTS model must pass all three:
|
||||
|
||||
| Layer | Catches | Tool |
|
||||
|-------|---------|------|
|
||||
| Offline RTF / duration check | Throughput regressions, missing audio, wrong sample rate | `end2end.py`, pytest e2e |
|
||||
| Browser streaming playback | Delta/cumulative bugs, chunk boundary glitches, TTFP regressions | Gradio demo over `/v1/audio/speech?stream=true` |
|
||||
| Concurrent requests | Per-request state leaks, codec window round-robin gaps | `max_num_seqs>1` smoke test with 4+ parallel prompts |
|
||||
|
||||
Declaring a model "done" without all three has shipped regressions more than once.
|
||||
|
||||
### I5. Per-request state is owned by the request, not the model
|
||||
|
||||
If the model caches *anything* across `forward()` calls (streaming generators, codec buffers, sliding-window pads, CUDA graph state), key it by request ID:
|
||||
|
||||
```python
|
||||
self._state: dict[str, YourState] = {} # request_key → state
|
||||
# fetch: request_key = str(info.get("_omni_req_id", "0"))
|
||||
# free on finish: del self._state[request_key]
|
||||
```
|
||||
|
||||
A shared buffer silently corrupts audio across concurrent requests — the symptom is crosstalk or truncation only under load.
|
||||
|
||||
## Phase 1: HuggingFace Reference
|
||||
|
||||
**Goal**: Understand the reference implementation and verify it produces correct audio.
|
||||
@@ -85,12 +165,54 @@ When audio output is wrong, check in this order:
|
||||
6. **Tensor layout**: Codebook-major vs frame-major ordering
|
||||
7. **dtype**: Float32 for codec decoders (autocast can corrupt audio)
|
||||
|
||||
### Streaming Correctness Rules (single-stage and two-stage)
|
||||
|
||||
These bugs appear in almost every new TTS PR. Check all before the first push. See also the cross-cutting invariants I1 (output contract) and I5 (per-request state) above — the rules below are the Phase 2-specific instances of those invariants:
|
||||
|
||||
- **Accumulate codes across AR steps** — each `forward()` appends new codes; do not reset between steps or audio will be truncated (fish speech: `fix: accumulate audio_codes across steps`)
|
||||
- **Emit delta audio, not full waveform** — in streaming mode yield only the new chunk per step, not the re-decoded full waveform from step 0 (fish speech: `fix: emit delta audio not full waveform`)
|
||||
- **All return paths must emit `model_outputs`** — if any early-return branch skips setting `model_outputs`, the serving layer silently drops that step's audio (fish speech: `fix: ensure ALL return paths emit model_outputs`)
|
||||
- **Per-request state isolation** — for batched concurrent requests, key all state by request ID; a shared buffer corrupts audio across requests (fish speech: `fix: per-request vocode + delta emission`)
|
||||
- **Codec tensor device** — move codec codes to the codec decoder's device before calling decode; mismatches cause silent CPU fallback or crashes (fish speech: `fix: use model device for CUDA stream`)
|
||||
- **AR stage `max_num_seqs`** — set to at least 4 in production configs; for single-stage models this is the only stage. For two-stage models, Stage 0 (AR) needs `max_num_seqs ≥ 4` to pipeline concurrent requests; Stage 1 (codec decoder) typically uses `max_num_seqs: 1` intentionally. Default of 1 everywhere causes audio gaps under concurrency because the codec window round-robins across requests (RFC #2568)
|
||||
|
||||
### Optional Dependency Handling
|
||||
|
||||
Patch optional dependencies (`torchaudio` / `torchcodec` / `soundfile`) at
|
||||
the top of `load_weights()`, not at module import. Failures to do so cause
|
||||
cryptic errors only on environments missing the optional package — after
|
||||
the model is already deployed. See
|
||||
[references/optional-deps.md](references/optional-deps.md) for the full
|
||||
pattern, signature constraints, and MOSS-TTS-Nano reference.
|
||||
|
||||
### Single-Stage AR Pattern (alternative to two-stage)
|
||||
|
||||
When the upstream model cannot be cleanly split into an AR stage and a
|
||||
separate decoder, run the full pipeline inside a single AR worker and
|
||||
stream audio through a per-request `inference_stream()` generator keyed by
|
||||
`_omni_req_id`. Stage config must set `worker_type: ar`,
|
||||
`engine_output_type: audio`, `final_output: true`, `is_comprehension: true`,
|
||||
and `async_chunk: false` at the top level. Only extract params from
|
||||
`additional_information` that you actually forward, or pre-commit fails
|
||||
`ruff F841`.
|
||||
|
||||
Full walkthrough with the complete `forward()` / `_create_stream_gen()`
|
||||
skeleton and stage-config fields:
|
||||
[references/single-stage-ar.md](references/single-stage-ar.md). For an
|
||||
in-tree reference, look for any single-stage AR model under
|
||||
`vllm_omni/model_executor/models/` — e.g. the MOSS-TTS-Nano integration when
|
||||
it lands.
|
||||
|
||||
**VoxCPM2 is a different pattern** and should not reuse this skeleton — it
|
||||
runs the base LM under vLLM PagedAttention with external side-computation.
|
||||
See `plan/voxcpm2_native_ar_design.md`.
|
||||
|
||||
### Deliverables
|
||||
|
||||
- Model files in `vllm_omni/model_executor/models/<model_name>/`
|
||||
- Stage config YAML
|
||||
- Working `end2end.py` with correct audio output
|
||||
- README.md in the example directory
|
||||
- Working `end2end.py` at `examples/offline_inference/text_to_speech/<model>/end2end.py`
|
||||
- New section in `examples/offline_inference/text_to_speech/README.md` (table row + per-model section). Do **not** create a top-level `examples/offline_inference/<model>/` dir or a per-model `README.md` inside `text_to_speech/<model>/` — the hub README is the documented surface and the mkdocs `generate_examples` hook only descends one level into `examples/<category>/`.
|
||||
|
||||
## Phase 3: Online Serving
|
||||
|
||||
@@ -98,10 +220,66 @@ When audio output is wrong, check in this order:
|
||||
|
||||
### Steps
|
||||
|
||||
1. **Register in `serving_speech.py`**:
|
||||
- Add model stage name to `_TTS_MODEL_STAGES` set
|
||||
- Add model detection flag (e.g., `_is_fish_speech`)
|
||||
- Implement prompt builder method (e.g., `_build_fish_speech_prompt()`)
|
||||
1. **Register in `serving_speech.py`** — add all 5 points in a **single commit**;
|
||||
partial integration causes hard-to-debug failures. This file is modified by every
|
||||
model PR and is the most common source of rebase conflicts — see conflict note below.
|
||||
|
||||
**Point 1** — stage constant (near the top, alongside the other `_*_TTS_MODEL_STAGES` sets):
|
||||
```python
|
||||
_YOUR_MODEL_TTS_MODEL_STAGES = {"your_stage_key"}
|
||||
```
|
||||
|
||||
**Point 2** — union into `_TTS_MODEL_STAGES`:
|
||||
```python
|
||||
_TTS_MODEL_STAGES: set[str] = (
|
||||
...
|
||||
| _YOUR_MODEL_TTS_MODEL_STAGES
|
||||
)
|
||||
```
|
||||
|
||||
**Point 3** — model type detection in `_detect_tts_model_type()`:
|
||||
```python
|
||||
if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES:
|
||||
return "your_model"
|
||||
```
|
||||
|
||||
**Point 4** — validation dispatch in `_validate_tts_request()`:
|
||||
```python
|
||||
if self._tts_model_type == "your_model":
|
||||
return self._validate_your_model_request(request)
|
||||
```
|
||||
|
||||
**Point 5** — validation + parameter-builder methods:
|
||||
```python
|
||||
def _validate_your_model_request(self, request) -> str | None:
|
||||
if not request.input or not request.input.strip():
|
||||
return "Input text cannot be empty"
|
||||
return None
|
||||
|
||||
def _build_your_model_params(self, request) -> dict:
|
||||
params = {"text": [request.input]}
|
||||
if request.voice is not None:
|
||||
params["voice"] = [request.voice]
|
||||
return params
|
||||
```
|
||||
Wire `_build_your_model_params` into `_create_tts_request()` alongside the other
|
||||
model-specific param builders.
|
||||
|
||||
> **Two dispatch patterns coexist**: Fish Speech uses a `self._is_fish_speech` boolean
|
||||
> instance attribute checked before `elif self._is_tts`, while all newer models
|
||||
> (CosyVoice3, MOSS-TTS-Nano) use the `_tts_model_type` string returned by
|
||||
> `_detect_tts_model_type()`. For new models, always use the `_tts_model_type` string
|
||||
> pattern — do not add new `_is_*` flags.
|
||||
|
||||
> **Unused variable rule**: only extract fields in `_build_your_model_params` that
|
||||
> are actually forwarded to the model. Unused extractions fail `ruff F841`.
|
||||
> For voice-cloning fields (`ref_audio` → `prompt_audio_path`, `ref_text` →
|
||||
> `prompt_text`), add them to the param builder and verify they reach the model call.
|
||||
|
||||
**Rebase conflict note**: when rebasing onto `main` after another model was merged,
|
||||
`serving_speech.py` will conflict. Resolution: always keep *both* the upstream
|
||||
model's additions and your own — never discard either side.
|
||||
|
||||
2. **Handle model-specific parameters**:
|
||||
- Voice cloning: `ref_audio` encoding and prompt injection
|
||||
- `max_new_tokens` override in sampling params
|
||||
@@ -129,10 +307,28 @@ def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list:
|
||||
|
||||
### Deliverables
|
||||
|
||||
- Updated `serving_speech.py` with model-specific prompt builder
|
||||
- Client scripts and server launcher
|
||||
- Gradio demo with streaming and voice cloning UI
|
||||
- Documentation (offline + online serving docs)
|
||||
- Updated `serving_speech.py` with all 5 integration points (single commit)
|
||||
- Client scripts and server launcher under `examples/online_serving/text_to_speech/<model>/`
|
||||
- Gradio demo with streaming and voice cloning UI in the same dir
|
||||
- E2E online serving test (`tests/e2e/online_serving/test_<model>.py`)
|
||||
- Buildkite CI entry in `.buildkite/test-merge.yml`
|
||||
- New section in `examples/online_serving/text_to_speech/README.md` (table row + per-model section). Do **not** create a top-level `examples/online_serving/<model>/` dir or a per-model `README.md` inside `text_to_speech/<model>/`.
|
||||
|
||||
### E2E test pitfalls to avoid
|
||||
|
||||
- **One `OmniServerParams` set per file.** `omni_server` is module-scoped; a second
|
||||
id in the same file forces mid-module teardown/restart and exposes startup
|
||||
races (`APIConnectionError` on the first request post-restart). Split variants
|
||||
into separate files instead.
|
||||
- **No external URL fetches from the server.** CI and some dev hosts can't
|
||||
reach `raw.githubusercontent.com` over TLS. Inline ref audio as
|
||||
`data:audio/wav;base64,...`; the serving layer accepts both URL and data URL.
|
||||
- **Use the harness readiness gate.** The fixture waits for HTTP 200 on
|
||||
`/health`; don't add `time.sleep` in tests. If warmup is incomplete, make
|
||||
`/health` return non-200 until you're actually ready.
|
||||
- **Mark with `@pytest.mark.core_model` + `hardware_test(res={"cuda": "H100"})`**
|
||||
so the test lands in `test-ready.yml` (triggered by the `ready` label) rather
|
||||
than only nightly.
|
||||
|
||||
## Phase 4: Async Chunk (Streaming)
|
||||
|
||||
@@ -199,35 +395,10 @@ Stage 0 (AR) Stage 1 (Decoder)
|
||||
- Fixed batch size (fall back to eager for other sizes)
|
||||
- No dynamic control flow inside the graph
|
||||
|
||||
### Example: Code Predictor CUDA Graph (Qwen3-TTS)
|
||||
|
||||
```python
|
||||
import torch
|
||||
|
||||
class CodePredictorGraph:
|
||||
"""Captures the 16-step code predictor AR loop as a single CUDA graph."""
|
||||
|
||||
def setup_graph(self, device: torch.device, kv_heads: int = 4, head_dim: int = 64):
|
||||
self.num_steps = 16
|
||||
self.kv_cache = torch.zeros(1, kv_heads, self.num_steps, head_dim, device=device)
|
||||
self.positions = torch.arange(self.num_steps, device=device)
|
||||
self.causal_mask = torch.tril(torch.ones(self.num_steps, self.num_steps, device=device))
|
||||
self.input_buf = torch.zeros(1, 1, kv_heads * head_dim, device=device)
|
||||
self.output_buf = torch.zeros(1, self.num_steps, device=device, dtype=torch.long)
|
||||
# Warm up, then: self.graph = torch.cuda.CUDAGraph(); self.graph.capture(...)
|
||||
|
||||
def run_graph(self, initial_input: torch.Tensor) -> torch.Tensor:
|
||||
self.input_buf.copy_(initial_input)
|
||||
self.graph.replay()
|
||||
return self.output_buf.clone()
|
||||
```
|
||||
|
||||
### Performance Expectations
|
||||
|
||||
Based on Qwen3-TTS code predictor experience:
|
||||
- **3-5x speedup** for the graphed component
|
||||
- Only effective for fixed batch sizes (typically batch_size=1)
|
||||
- Falls back to eager mode for unsupported configurations
|
||||
See [references/cuda-graph-example.md](references/cuda-graph-example.md) for
|
||||
a worked skeleton (Qwen3-TTS code predictor, 16-step AR loop), performance
|
||||
expectations (3–5× on the graphed component for fixed batch_size=1), and the
|
||||
graph-safety constraints you must honor inside the captured region.
|
||||
|
||||
### Deliverables
|
||||
|
||||
@@ -235,10 +406,33 @@ Based on Qwen3-TTS code predictor experience:
|
||||
- Benchmark script comparing eager vs graph performance
|
||||
- Documentation of constraints and fallback behavior
|
||||
|
||||
## Phase 6: Pre-commit and DCO
|
||||
|
||||
**Goal**: Every commit passes `pre-commit` lint and carries a DCO
|
||||
`Signed-off-by` line that matches the author email.
|
||||
|
||||
- Install hooks once: `pre-commit install`.
|
||||
- Run `pre-commit run --files <changed-files>` before every push; accept any
|
||||
auto-fixes, stage, re-commit.
|
||||
- Sign every commit with `git commit -s`. DCO checks that author email and
|
||||
`Signed-off-by` email match — `git config user.email` must match your
|
||||
GitHub account email.
|
||||
|
||||
Common pre-commit failures, recovery commands for missing sign-off, and the
|
||||
full `pre-commit run` invocation for a TTS model:
|
||||
[references/precommit-dco.md](references/precommit-dco.md).
|
||||
|
||||
## Integration Checklist
|
||||
|
||||
Use this checklist when integrating a new TTS model:
|
||||
|
||||
### Cross-Cutting Invariants (verify at end of every phase)
|
||||
- [ ] I1: `forward()` docstring states cumulative vs delta; consolidation path audited end-to-end
|
||||
- [ ] I2: Tests / examples / benchmarks never use `dict.get(a) or dict.get(b)` on tensor values; list form handled
|
||||
- [ ] I3: No `.item()` / `.cpu()` / Python branch on tensor values inside per-step loops
|
||||
- [ ] I4: Offline RTF, browser streaming playback, and concurrent-request smoke test all pass
|
||||
- [ ] I5: Any cross-step cache keyed by `_omni_req_id`; entries freed when the request finishes
|
||||
|
||||
### Phase 1: HF Reference
|
||||
- [ ] Reference model runs and produces correct audio
|
||||
- [ ] Architecture documented (stages, codebooks, tokens, sample rate)
|
||||
@@ -248,17 +442,27 @@ Use this checklist when integrating a new TTS model:
|
||||
- [ ] Model registered in `registry.py`
|
||||
- [ ] Config classes created with `model_type` registration
|
||||
- [ ] Stage 0 (AR) implemented and generates correct tokens
|
||||
- [ ] Stage 1 (Decoder) produces correct audio from tokens
|
||||
- [ ] Stage 1 (Decoder) produces correct audio from tokens — dtype float32 for codec decoder
|
||||
- [ ] Stage 1 `max_num_seqs` ≥ 4 in production config (default 1 causes gaps under concurrency)
|
||||
- [ ] Optional dependency fallbacks handled at `load_weights()` time (torchaudio/soundfile/etc.)
|
||||
- [ ] Streaming: codec codes accumulated across AR steps (not reset per step)
|
||||
- [ ] Streaming: delta audio emitted per chunk, not full re-decoded waveform
|
||||
- [ ] Streaming: all `forward()` return paths emit `model_outputs`
|
||||
- [ ] Streaming: per-request state keyed by request ID (not shared across requests)
|
||||
- [ ] Streaming: codec tensors moved to codec decoder device before decode
|
||||
- [ ] Stage config YAML created
|
||||
- [ ] `end2end.py` produces audio matching reference quality
|
||||
- [ ] README.md written
|
||||
|
||||
### Phase 3: Online Serving
|
||||
- [ ] Model added to `serving_speech.py`
|
||||
- [ ] All 5 `serving_speech.py` integration points added in one commit
|
||||
- [ ] Only extract params in `_build_*_params` that are forwarded to the model call (ruff F841)
|
||||
- [ ] Prompt builder handles text input correctly
|
||||
- [ ] Voice cloning works (if supported)
|
||||
- [ ] All response formats work (wav, mp3, flac, pcm)
|
||||
- [ ] Client scripts and server launcher created
|
||||
- [ ] E2E online serving test written (`tests/e2e/online_serving/test_<model>.py`)
|
||||
- [ ] Buildkite CI entry added to `.buildkite/test-merge.yml`
|
||||
- [ ] Gradio demo working
|
||||
- [ ] Documentation added (offline + online docs, nav, supported models)
|
||||
|
||||
@@ -276,9 +480,25 @@ Use this checklist when integrating a new TTS model:
|
||||
- [ ] Benchmark shows meaningful speedup
|
||||
- [ ] Fallback to eager works for unsupported configs
|
||||
|
||||
### Phase 6: Pre-commit and DCO
|
||||
- [ ] `pre-commit run --files <changed>` passes before every push
|
||||
- [ ] Every commit has `Signed-off-by` matching the author email (`git commit -s`)
|
||||
- [ ] `git config user.email` matches the email registered on your GitHub account
|
||||
- [ ] Details and failure-recovery commands: [references/precommit-dco.md](references/precommit-dco.md)
|
||||
|
||||
## References
|
||||
|
||||
- [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) -- supported models and usage
|
||||
- [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) -- complete example of Phases 1-3
|
||||
- [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) -- complete example of all 5 phases
|
||||
In-skill references (details split out of the main body):
|
||||
|
||||
- [references/single-stage-ar.md](references/single-stage-ar.md) — full `forward()` / generator skeleton for the MOSS-TTS-Nano-style pattern
|
||||
- [references/optional-deps.md](references/optional-deps.md) — torchaudio / torchcodec fallback pattern
|
||||
- [references/cuda-graph-example.md](references/cuda-graph-example.md) — Qwen3-TTS code-predictor CUDA graph skeleton
|
||||
- [references/precommit-dco.md](references/precommit-dco.md) — full pre-commit invocation, failure table, DCO recovery
|
||||
|
||||
Project docs and adjacent skills:
|
||||
|
||||
- [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) — supported models and usage
|
||||
- [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) — complete example of Phases 1–3
|
||||
- [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) — complete example of all 5 phases
|
||||
- [Adding a TTS model (developer guide)](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/model/adding_tts_model.md)
|
||||
- `plan/voxcpm2_native_ar_design.md` — VoxCPM2's vLLM-native AR + side-computation pattern (distinct from the generator-based single-stage described above)
|
||||
@@ -0,0 +1,42 @@
|
||||
# CUDA Graph Example: Qwen3-TTS Code Predictor
|
||||
|
||||
Reference sketch for capturing the 16-step code-predictor AR loop as a single
|
||||
CUDA graph. Adapt the shapes, number of steps, and KV-head layout to your
|
||||
model.
|
||||
|
||||
```python
|
||||
import torch
|
||||
|
||||
class CodePredictorGraph:
|
||||
"""Captures the 16-step code predictor AR loop as a single CUDA graph."""
|
||||
|
||||
def setup_graph(self, device: torch.device, kv_heads: int = 4, head_dim: int = 64):
|
||||
self.num_steps = 16
|
||||
self.kv_cache = torch.zeros(1, kv_heads, self.num_steps, head_dim, device=device)
|
||||
self.positions = torch.arange(self.num_steps, device=device)
|
||||
self.causal_mask = torch.tril(torch.ones(self.num_steps, self.num_steps, device=device))
|
||||
self.input_buf = torch.zeros(1, 1, kv_heads * head_dim, device=device)
|
||||
self.output_buf = torch.zeros(1, self.num_steps, device=device, dtype=torch.long)
|
||||
# Warm up, then: self.graph = torch.cuda.CUDAGraph(); self.graph.capture(...)
|
||||
|
||||
def run_graph(self, initial_input: torch.Tensor) -> torch.Tensor:
|
||||
self.input_buf.copy_(initial_input)
|
||||
self.graph.replay()
|
||||
return self.output_buf.clone()
|
||||
```
|
||||
|
||||
## Performance expectations (Qwen3-TTS code predictor)
|
||||
|
||||
- **3–5× speedup** on the graphed component.
|
||||
- Effective only for fixed batch sizes (typically `batch_size=1`).
|
||||
- Fall back to eager for any shape/config that wasn't captured — do not try
|
||||
to recapture per request.
|
||||
|
||||
## Graph-safety constraints
|
||||
|
||||
- `torch.argmax` instead of `torch.multinomial`.
|
||||
- Fixed batch size.
|
||||
- No Python control flow that branches on tensor values inside the captured
|
||||
region (use `torch.where` / masks).
|
||||
- No `.item()`, `.cpu()`, `.tolist()` — each would break the capture or
|
||||
cause a GPU→CPU sync during replay.
|
||||
@@ -0,0 +1,47 @@
|
||||
# Optional Dependency Handling
|
||||
|
||||
Models that rely on `torchaudio`, `torchcodec`, `soundfile`, or other optional
|
||||
packages must handle the missing-package case at import time, not at call
|
||||
time. Failing to do this causes cryptic errors only on environments without
|
||||
the optional package — after the model is already deployed.
|
||||
|
||||
## Pattern (used in MOSS-TTS-Nano)
|
||||
|
||||
```python
|
||||
def _patch_torchaudio_load() -> None:
|
||||
"""Fallback torchaudio.load/save to soundfile if torchcodec is unavailable."""
|
||||
try:
|
||||
import torchcodec # noqa: F401
|
||||
return # torchcodec present, torchaudio works as-is
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
import soundfile as sf
|
||||
import torchaudio
|
||||
|
||||
def _sf_load(path, **kwargs):
|
||||
data, sr = sf.read(str(path), dtype="float32", always_2d=True)
|
||||
return torch.from_numpy(data).T, sr
|
||||
|
||||
torchaudio.load = _sf_load
|
||||
# patch .save similarly if needed
|
||||
```
|
||||
|
||||
## Rules
|
||||
|
||||
- Mirror the full signature of the replaced function. `torchaudio.load`
|
||||
accepts `frame_offset`, `num_frames`, `normalize`, `channels_first`,
|
||||
`format` — missing any of them causes `TypeError` from calling code.
|
||||
- Catch `except Exception`, not just `ImportError`. `import torchaudio`
|
||||
itself can fail with non-`ImportError` errors on broken installs. Log the
|
||||
exception type and message (`logger.warning("torchaudio probe failed: %s: %s",
|
||||
type(exc).__name__, exc)`) before falling back, so unrelated errors are not
|
||||
silently swallowed.
|
||||
- Call the patch function at the top of `load_weights()` before loading any
|
||||
audio assets. Do not call it at module import time.
|
||||
|
||||
## Reference implementation
|
||||
|
||||
Any in-tree model that patches `torchaudio.load` in its `load_weights()` —
|
||||
e.g. MOSS-TTS-Nano's `modeling_moss_tts_nano.py` once that integration
|
||||
lands.
|
||||
@@ -0,0 +1,54 @@
|
||||
# Pre-commit and DCO
|
||||
|
||||
Every commit must pass `pre-commit` lint and carry a `Signed-off-by` line
|
||||
that matches the commit author email.
|
||||
|
||||
## Pre-commit
|
||||
|
||||
Install hooks once:
|
||||
|
||||
```bash
|
||||
pre-commit install
|
||||
```
|
||||
|
||||
Run before every push on the files you changed:
|
||||
|
||||
```bash
|
||||
pre-commit run --files \
|
||||
vllm_omni/model_executor/models/<model_name>/*.py \
|
||||
vllm_omni/entrypoints/openai/serving_speech.py \
|
||||
vllm_omni/model_executor/models/registry.py \
|
||||
tests/e2e/offline_inference/test_<model_name>.py \
|
||||
tests/e2e/online_serving/test_<model_name>.py
|
||||
```
|
||||
|
||||
When pre-commit **modifies files** (ruff format auto-fix), it exits non-zero
|
||||
but the changes are correct — stage the modified files and re-commit.
|
||||
|
||||
| Failure | Root cause | Fix |
|
||||
|---------|-----------|-----|
|
||||
| `ruff F841` | Variable extracted but never forwarded to model call | Remove the extraction or wire it through |
|
||||
| `ruff E402` | Import added below function definitions | Move to top-level import block |
|
||||
| `ruff format` | Line length, spacing, quote style | Accept auto-fix, stage, re-commit |
|
||||
|
||||
## DCO sign-off
|
||||
|
||||
Every commit must carry `Signed-off-by: Your Name <your@email.com>`. Use
|
||||
`-s`:
|
||||
|
||||
```bash
|
||||
git commit -s -m "feat(<model>): add <Model> TTS support"
|
||||
```
|
||||
|
||||
Or set it permanently: `git config format.signOff true`.
|
||||
|
||||
The DCO check verifies that the commit author email matches the
|
||||
`Signed-off-by` line. Confirm `git config user.email` matches your GitHub
|
||||
account email before committing.
|
||||
|
||||
Fix a missing or mismatched sign-off on the latest commit:
|
||||
|
||||
```bash
|
||||
git commit --amend -s --no-edit
|
||||
git push origin <branch> --force-with-lease
|
||||
```
|
||||
@@ -0,0 +1,108 @@
|
||||
# Single-Stage AR Pattern
|
||||
|
||||
When the upstream model cannot be cleanly split into an AR stage and a separate
|
||||
decoder (e.g. MOSS-TTS-Nano, or any model that bundles AR + codec via an
|
||||
`inference_stream()` generator), run the whole pipeline inside a single AR
|
||||
worker that yields audio chunks per request.
|
||||
|
||||
This is distinct from VoxCPM2's pattern, which also runs in a single stage but
|
||||
uses vLLM's native PagedAttention on the base language model with diffusion /
|
||||
VAE side-computation outside vLLM — see
|
||||
`plan/voxcpm2_native_ar_design.md` for that variant.
|
||||
|
||||
## Implementation
|
||||
|
||||
1. **Single model file** — load both AR LM and codec inside
|
||||
`modeling_<model>.py`.
|
||||
2. **Load weights in `load_weights()`**, not `__init__()` — vLLM initializes
|
||||
distributed state before any CUDA allocations.
|
||||
3. **Stream via a per-request generator** stored in `self._stream_gens`:
|
||||
|
||||
```python
|
||||
class YourModelForCausalLM(nn.Module):
|
||||
def __init__(self, *, vllm_config, prefix=""):
|
||||
super().__init__()
|
||||
self._lm = None # populated in load_weights()
|
||||
self._stream_gens: dict = {} # request_key → generator
|
||||
|
||||
def load_weights(self, weights):
|
||||
# Load self._lm here, after vLLM distributed init
|
||||
...
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids,
|
||||
positions,
|
||||
intermediate_tensors=None,
|
||||
inputs_embeds=None,
|
||||
runtime_additional_information: list[dict] | None = None, # one dict per request
|
||||
**kwargs,
|
||||
) -> OmniOutput:
|
||||
infos = runtime_additional_information or [{}]
|
||||
# Skip dummy/profiling calls
|
||||
if not runtime_additional_information or all(i.get("_is_dummy") for i in infos):
|
||||
self._ar_emit_stop_token = True
|
||||
return OmniOutput(...) # return empty outputs
|
||||
|
||||
outputs, last_flags = [], []
|
||||
for info in infos:
|
||||
request_key = str(info.get("_omni_req_id", "0")) # per-request ID from vLLM
|
||||
if request_key not in self._stream_gens:
|
||||
self._stream_gens[request_key] = self._create_stream_gen(info)
|
||||
try:
|
||||
chunk, is_last = next(self._stream_gens[request_key])
|
||||
except StopIteration:
|
||||
chunk, is_last = torch.zeros(0), True
|
||||
if is_last:
|
||||
del self._stream_gens[request_key]
|
||||
outputs.append(chunk)
|
||||
last_flags.append(is_last)
|
||||
|
||||
self._ar_emit_stop_token = all(last_flags)
|
||||
return OmniOutput(multimodal_outputs={"model_outputs": outputs, ...})
|
||||
|
||||
def _create_stream_gen(self, info: dict):
|
||||
"""Yield (waveform_tensor, is_last) tuples from inference_stream()."""
|
||||
for event in self._lm.inference_stream(...):
|
||||
if event["type"] == "audio":
|
||||
yield event["waveform"], False
|
||||
elif event["type"] == "result":
|
||||
# Fallback: some models emit a single "result" event instead of
|
||||
# incremental "audio" events — handle both paths
|
||||
yield event.get("waveform", torch.zeros(0)), True
|
||||
return
|
||||
yield torch.zeros(0), True
|
||||
|
||||
def compute_logits(self, hidden_states, sampling_metadata):
|
||||
# Emit EOS only after the last chunk so the AR scheduler ends the request
|
||||
...
|
||||
```
|
||||
|
||||
## Key points
|
||||
|
||||
- `runtime_additional_information` is the correct parameter name (not
|
||||
`**kwargs`) — it carries one dict per request in the batch.
|
||||
- The request ID is `info.get("_omni_req_id")` — set by vLLM, not by user code.
|
||||
- Handle both `"audio"` (incremental) and `"result"` (final combined) event
|
||||
types from upstream models.
|
||||
|
||||
## Stage config
|
||||
|
||||
Single stage with `worker_type: ar`, `engine_output_type: audio`,
|
||||
`final_output: true`, `is_comprehension: true`, and `async_chunk: false` at
|
||||
the top level. Omitting any of these causes silent misclassification in the
|
||||
serving layer.
|
||||
|
||||
## Lint discipline
|
||||
|
||||
Only extract variables from `additional_information` that you actually
|
||||
forward to the model call — unused extractions trip `ruff F841` in
|
||||
pre-commit.
|
||||
|
||||
## Reference implementation
|
||||
|
||||
Look for any single-stage AR model under
|
||||
`vllm_omni/model_executor/models/` — e.g. `moss_tts_nano/` when its
|
||||
integration lands. If none is in tree yet, follow the skeleton above and
|
||||
cross-check against the `is_comprehension: true` / `async_chunk: false`
|
||||
dispatch in `vllm_omni/entrypoints/openai/serving_speech.py`.
|
||||
@@ -27,6 +27,7 @@
|
||||
*.mov binary
|
||||
*.mp4 binary
|
||||
*.mp3 binary
|
||||
*.wav binary
|
||||
*.flv binary
|
||||
*.fla binary
|
||||
*.swf binary
|
||||
|
||||
+59
-10
@@ -3,15 +3,64 @@
|
||||
# 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)
|
||||
# later, we can add more reviewers here
|
||||
* @hsliuustc0106
|
||||
* @Gaohan123 @hsliuustc0106 @ywang96
|
||||
|
||||
# Assign reviewers by directory
|
||||
/.github/ @owner1
|
||||
/tests/ @owner1
|
||||
/docs/ @owner1
|
||||
# Repository operations
|
||||
/.github/ @hsliuustc0106
|
||||
|
||||
# Example: reviewers for specific directories
|
||||
# /examples/offline_inference/ @offline-inference-team
|
||||
# /vllm_omni/core/ @core-team
|
||||
# /vllm_omni/model_executor/ @model-team
|
||||
# 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
|
||||
@@ -0,0 +1,223 @@
|
||||
pull_request_rules:
|
||||
- name: label-documentation
|
||||
description: Automatically apply documentation label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^[^/]+\.md$
|
||||
- files~=^docs/
|
||||
- files~=^examples/
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- documentation
|
||||
comment:
|
||||
message: "Documentation preview: https://vllm--{{number}}.org.readthedocs.build/en/{{number}}/"
|
||||
|
||||
- name: comment-pre-commit-failure
|
||||
description: Comment on PR when pre-commit check fails
|
||||
conditions:
|
||||
- status-failure=pre-commit
|
||||
- -closed
|
||||
- -draft
|
||||
actions:
|
||||
comment:
|
||||
message: |
|
||||
Hi @{{author}}, the pre-commit checks have failed. Please run:
|
||||
|
||||
```bash
|
||||
uv pip install pre-commit
|
||||
pre-commit install
|
||||
pre-commit run --all-files
|
||||
```
|
||||
|
||||
Then, commit the changes and push to your branch.
|
||||
|
||||
For future commits, `pre-commit` will run automatically on changed files before each commit.
|
||||
|
||||
> [!TIP]
|
||||
> <details>
|
||||
> <summary>Is <code>mypy</code> or <code>markdownlint</code> failing?</summary>
|
||||
> <br/>
|
||||
> <code>mypy</code> and <code>markdownlint</code> are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:
|
||||
>
|
||||
> ```bash
|
||||
> # For mypy (substitute "3.10" with the failing version if needed)
|
||||
> pre-commit run --hook-stage manual mypy-3.10
|
||||
> # For markdownlint
|
||||
> pre-commit run --hook-stage manual markdownlint
|
||||
> ```
|
||||
> </details>
|
||||
|
||||
- name: comment-dco-failure
|
||||
description: Comment on PR when DCO check fails
|
||||
conditions:
|
||||
- status-failure=dco
|
||||
- -closed
|
||||
- -draft
|
||||
actions:
|
||||
comment:
|
||||
message: |
|
||||
Hi @{{author}}, the DCO check has failed. Please click on DCO in the Checks section for instructions on how to resolve this.
|
||||
|
||||
- name: label-ci-build
|
||||
description: Automatically apply ci/build label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^\.github/
|
||||
- files~=\.buildkite/
|
||||
- files~=^cmake/
|
||||
- files=CMakeLists.txt
|
||||
- files~=^docker/Dockerfile
|
||||
- files~=^requirements.*\.txt
|
||||
- files=setup.py
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- ci/build
|
||||
|
||||
- name: label-frontend
|
||||
description: Automatically apply frontend label
|
||||
conditions:
|
||||
- label != stale
|
||||
- files~=^vllm_omni/entrypoints/
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- frontend
|
||||
|
||||
- name: label-diffusion
|
||||
description: Automatically apply diffusion label
|
||||
conditions:
|
||||
- label != stale
|
||||
- files~=^vllm_omni/diffusion/
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- diffusion
|
||||
|
||||
- name: label-new-model
|
||||
description: Automatically apply new-model label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files=vllm_omni/model_executor/models/registry.py
|
||||
- files=vllm_omni/diffusion/registry.py
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- new-model
|
||||
|
||||
- name: label-npu
|
||||
description: Automatically apply npu label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^vllm_omni/platforms/npu/
|
||||
- title~=(?i)NPU
|
||||
- title~=(?i)Ascend
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- npu
|
||||
assign:
|
||||
users:
|
||||
- "gcanlin"
|
||||
|
||||
- name: label-rocm
|
||||
description: Automatically apply rocm label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^vllm_omni/platforms/rocm/
|
||||
- title~=(?i)AMD
|
||||
- title~=(?i)ROCm
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- rocm
|
||||
assign:
|
||||
users:
|
||||
- "tjtanaa"
|
||||
|
||||
- name: label-xpu
|
||||
description: Automatically apply xpu label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^vllm_omni/platforms/xpu/
|
||||
- title~=(?i)Intel
|
||||
- title~=(?i)XPU
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- xpu
|
||||
assign:
|
||||
users:
|
||||
- "yma11"
|
||||
|
||||
- name: label-modelrunner
|
||||
description: Automatically apply modelrunner label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^vllm_omni/worker/
|
||||
- title~=(?i)ModelRunner
|
||||
- title~=(?i)worker
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- ModelRunner
|
||||
assign:
|
||||
users:
|
||||
- "gcanlin" # To check whether need to sync to NPU
|
||||
|
||||
- name: auto-rebase if approved, ready, and 30 commits behind main
|
||||
conditions:
|
||||
- base = main
|
||||
- label=ready
|
||||
- "#approved-reviews-by >= 1"
|
||||
- "#commits-behind >= 30"
|
||||
- -closed
|
||||
- -draft
|
||||
- -conflict
|
||||
actions:
|
||||
rebase: {}
|
||||
|
||||
- name: ping author on conflicts and add 'needs-rebase' label
|
||||
conditions:
|
||||
- label != stale
|
||||
- conflict
|
||||
- -closed
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- needs-rebase
|
||||
comment:
|
||||
message: |
|
||||
This pull request has merge conflicts that must be resolved before it can be
|
||||
merged. Please rebase the PR, @{{author}}.
|
||||
|
||||
https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork
|
||||
|
||||
- name: remove 'needs-rebase' label when conflict is resolved
|
||||
conditions:
|
||||
- -conflict
|
||||
- -closed
|
||||
actions:
|
||||
label:
|
||||
remove:
|
||||
- needs-rebase
|
||||
|
||||
- name: label-bug
|
||||
description: Automatically apply bug label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- title~=(?i)\bbug\b
|
||||
- title~=(?i)\bbugfix\b
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- bug
|
||||
@@ -1,629 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
PR Reviewer using GLM API for vllm-omni project.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, TypedDict
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
# Type definitions for API responses
|
||||
class PRDetails(TypedDict):
|
||||
"""Type definition for GitHub PR details response."""
|
||||
|
||||
title: str
|
||||
body: str
|
||||
number: int
|
||||
state: str
|
||||
user: dict[str, Any]
|
||||
|
||||
|
||||
class GLMMessage(TypedDict):
|
||||
"""Type definition for GLM API message."""
|
||||
|
||||
role: str
|
||||
content: str
|
||||
|
||||
|
||||
class GLMChoice(TypedDict):
|
||||
"""Type definition for GLM API choice."""
|
||||
|
||||
message: GLMMessage
|
||||
finish_reason: str
|
||||
|
||||
|
||||
class GLMResponse(TypedDict):
|
||||
"""Type definition for GLM API response."""
|
||||
|
||||
choices: list[GLMChoice]
|
||||
usage: dict[str, int] | None
|
||||
|
||||
|
||||
class GitHubComment(TypedDict):
|
||||
"""Type definition for GitHub comment."""
|
||||
|
||||
id: int
|
||||
body: str
|
||||
created_at: str
|
||||
user: dict[str, Any]
|
||||
|
||||
|
||||
# Configuration
|
||||
TRIGGER_PHRASE: str = "@vllm-omni-reviewer"
|
||||
DEFAULT_GLM_API_URL: str = "https://open.bigmodel.cn/api/paas/v4/chat/completions" # noqa: E501
|
||||
DEFAULT_GLM_MODEL: str = "glm-5"
|
||||
DEFAULT_COOLDOWN_MINUTES: int = 5
|
||||
DEFAULT_MAX_RETRIES: int = 3
|
||||
DEFAULT_RETRY_DELAY: float = 1.0
|
||||
MAX_DIFF_SIZE: int = 100_000 # Maximum diff size in characters
|
||||
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
"""Configuration for the PR reviewer."""
|
||||
|
||||
glm_api_url: str
|
||||
glm_model: str
|
||||
cooldown_minutes: int
|
||||
max_retries: int
|
||||
retry_delay: float
|
||||
max_diff_size: int
|
||||
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="[PR Reviewer] %(asctime)s - %(levelname)s - %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
)
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_config() -> Config:
|
||||
"""Load configuration from environment variables with defaults."""
|
||||
return Config(
|
||||
glm_api_url=os.getenv("GLM_API_URL", DEFAULT_GLM_API_URL),
|
||||
glm_model=os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL),
|
||||
cooldown_minutes=int(
|
||||
os.getenv(
|
||||
"PR_REVIEWER_COOLDOWN_MINUTES",
|
||||
str(DEFAULT_COOLDOWN_MINUTES),
|
||||
)
|
||||
),
|
||||
max_retries=int(
|
||||
os.getenv(
|
||||
"PR_REVIEWER_MAX_RETRIES",
|
||||
str(DEFAULT_MAX_RETRIES),
|
||||
)
|
||||
),
|
||||
retry_delay=float(os.getenv("PR_REVIEWER_RETRY_DELAY", str(DEFAULT_RETRY_DELAY))),
|
||||
max_diff_size=int(os.getenv("PR_REVIEWER_MAX_DIFF_SIZE", str(MAX_DIFF_SIZE))), # noqa: E501
|
||||
)
|
||||
|
||||
|
||||
def get_env_var(name: str) -> str:
|
||||
"""
|
||||
Get an environment variable or raise an error.
|
||||
|
||||
Args:
|
||||
name: Name of the environment variable.
|
||||
|
||||
Returns:
|
||||
The value of the environment variable.
|
||||
|
||||
Raises:
|
||||
SystemExit: If the environment variable is not set.
|
||||
"""
|
||||
value = os.environ.get(name)
|
||||
if not value:
|
||||
logger.error(f"Environment variable {name} is not set")
|
||||
sys.exit(1)
|
||||
return value
|
||||
|
||||
|
||||
def check_trigger(comment_body: str) -> bool:
|
||||
"""
|
||||
Check if the comment contains the trigger phrase.
|
||||
|
||||
Args:
|
||||
comment_body: The body of the comment to check.
|
||||
|
||||
Returns:
|
||||
True if the trigger phrase is found, False otherwise.
|
||||
"""
|
||||
return TRIGGER_PHRASE in comment_body
|
||||
|
||||
|
||||
def fetch_pr_diff(
|
||||
repo_name: str,
|
||||
pr_number: int,
|
||||
token: str,
|
||||
max_size: int = MAX_DIFF_SIZE,
|
||||
) -> str | None:
|
||||
"""
|
||||
Fetch the diff for a pull request.
|
||||
|
||||
Args:
|
||||
repo_name: The repository name in format "owner/repo".
|
||||
pr_number: The pull request number.
|
||||
token: GitHub authentication token.
|
||||
max_size: Maximum diff size in characters.
|
||||
|
||||
Returns:
|
||||
The diff content as a string, or None if fetching failed.
|
||||
Returns empty string if diff is larger than max_size.
|
||||
"""
|
||||
url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/vnd.github.v3.diff",
|
||||
}
|
||||
|
||||
logger.info(f"Fetching PR diff from {url}")
|
||||
response = requests.get(url, headers=headers, timeout=30)
|
||||
|
||||
if response.status_code == 200:
|
||||
diff: str = response.text
|
||||
if len(diff) > max_size:
|
||||
logger.warning(
|
||||
f"Diff size ({len(diff)} bytes) exceeds maximum "
|
||||
f"({max_size} bytes), truncating to first "
|
||||
f"{max_size} bytes"
|
||||
)
|
||||
return diff[:max_size] + "\n\n... [Diff truncated due to size] ..."
|
||||
logger.info(f"Successfully fetched diff ({len(diff)} bytes)")
|
||||
return diff
|
||||
else:
|
||||
logger.error(f"Failed to fetch PR diff: {response.status_code}")
|
||||
logger.error(f"Response: {response.text}")
|
||||
return None
|
||||
|
||||
|
||||
def fetch_pr_details(
|
||||
repo_name: str,
|
||||
pr_number: int,
|
||||
token: str,
|
||||
) -> PRDetails | None:
|
||||
"""
|
||||
Fetch PR details including title and description.
|
||||
|
||||
Args:
|
||||
repo_name: The repository name in format "owner/repo".
|
||||
pr_number: The pull request number.
|
||||
token: GitHub authentication token.
|
||||
|
||||
Returns:
|
||||
A dictionary containing PR details, or None if fetching failed.
|
||||
"""
|
||||
url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/vnd.github.v3+json",
|
||||
}
|
||||
|
||||
logger.info(f"Fetching PR details from {url}")
|
||||
response = requests.get(url, headers=headers, timeout=30)
|
||||
|
||||
if response.status_code == 200:
|
||||
return response.json()
|
||||
else:
|
||||
logger.error(f"Failed to fetch PR details: {response.status_code}")
|
||||
return None
|
||||
|
||||
|
||||
def build_review_prompt(pr_title: str, pr_description: str, diff: str) -> str:
|
||||
"""
|
||||
Build the prompt for the GLM-4.7 API.
|
||||
|
||||
Args:
|
||||
pr_title: The title of the pull request.
|
||||
pr_description: The description/body of the pull request.
|
||||
diff: The diff content of the pull request.
|
||||
|
||||
Returns:
|
||||
The formatted prompt string for the API.
|
||||
"""
|
||||
return f"""You are an expert code reviewer for the VLLM-Omni project. \
|
||||
Please review the following pull request:
|
||||
|
||||
## Pull Request Details
|
||||
**Title:** {pr_title}
|
||||
|
||||
**Description:**
|
||||
{pr_description if pr_description else "No description provided."}
|
||||
|
||||
## Code Changes (Diff)
|
||||
{diff}
|
||||
|
||||
## Review Guidelines
|
||||
|
||||
Please provide a comprehensive code review with the following sections:
|
||||
|
||||
### 1. Overview
|
||||
- Brief summary of the changes
|
||||
- Overall assessment (positive, neutral, or concerns)
|
||||
|
||||
### 2. Code Quality
|
||||
- Code style and consistency
|
||||
- Potential bugs or edge cases
|
||||
- Performance considerations
|
||||
- Error handling
|
||||
|
||||
### 3. Architecture & Design
|
||||
- Integration with existing codebase
|
||||
- Design patterns and best practices
|
||||
- Potential improvements
|
||||
|
||||
### 4. Security & Safety
|
||||
- Security concerns (if any)
|
||||
- Resource management
|
||||
- Input validation
|
||||
|
||||
### 5. Testing & Documentation
|
||||
- Test coverage considerations
|
||||
- Documentation completeness
|
||||
- Examples and usage clarity
|
||||
|
||||
### 6. Specific Suggestions
|
||||
- Line-by-line specific feedback (use `file:line` format)
|
||||
- Concrete actionable suggestions
|
||||
- Code examples for improvements (if applicable)
|
||||
|
||||
### 7. Approval Status
|
||||
- **LGTM** (Looks Good To Me) if the PR is ready to merge
|
||||
- **LGTM with suggestions** if the PR is good but has minor suggestions
|
||||
- **Changes requested** if significant changes are needed
|
||||
|
||||
## Important Notes
|
||||
- Be constructive and helpful
|
||||
- Focus on objective technical feedback
|
||||
- Acknowledge good practices when you see them
|
||||
- Prioritize critical issues over nitpicks
|
||||
- If the diff is empty or minimal, acknowledge this and provide
|
||||
any relevant context-specific guidance
|
||||
|
||||
Please format your response in Markdown with clear section headers.
|
||||
"""
|
||||
|
||||
|
||||
def validate_glm_response(data: dict[str, Any]) -> str | None:
|
||||
"""
|
||||
Validate and extract content from GLM API response.
|
||||
|
||||
Args:
|
||||
data: The response data from GLM API.
|
||||
|
||||
Returns:
|
||||
The review content string if valid, None otherwise.
|
||||
"""
|
||||
# Check if choices exists and is a non-empty list
|
||||
if "choices" not in data:
|
||||
logger.error("GLM API response missing 'choices' field")
|
||||
logger.error(f"Response structure: {json.dumps(data, indent=2)}")
|
||||
return None
|
||||
|
||||
choices = data["choices"]
|
||||
if not isinstance(choices, list):
|
||||
logger.error(f"GLM API 'choices' is not a list: {type(choices)}")
|
||||
return None
|
||||
|
||||
if len(choices) == 0:
|
||||
logger.error("GLM API 'choices' is an empty list")
|
||||
return None
|
||||
|
||||
# Check if first choice has message
|
||||
try:
|
||||
first_choice = choices[0]
|
||||
if not isinstance(first_choice, dict):
|
||||
logger.error(f"GLM API choice is not a dict: {type(first_choice)}")
|
||||
return None
|
||||
|
||||
if "message" not in first_choice:
|
||||
logger.error("GLM API choice missing 'message' field")
|
||||
logger.error(f"Choice structure: {json.dumps(first_choice, indent=2)}") # noqa: E501
|
||||
return None
|
||||
|
||||
message = first_choice["message"]
|
||||
if not isinstance(message, dict):
|
||||
logger.error(f"GLM API message is not a dict: {type(message)}")
|
||||
return None
|
||||
|
||||
if "content" not in message:
|
||||
logger.error("GLM API message missing 'content' field")
|
||||
logger.error(f"Message structure: {json.dumps(message, indent=2)}")
|
||||
return None
|
||||
|
||||
content = message["content"]
|
||||
if not isinstance(content, str):
|
||||
logger.error(f"GLM API content is not a string: {type(content)}")
|
||||
return None
|
||||
|
||||
return content
|
||||
|
||||
except (KeyError, IndexError, TypeError) as e:
|
||||
logger.error(f"Failed to parse GLM API response: {e}")
|
||||
logger.error(f"Response: {json.dumps(data, indent=2)}")
|
||||
return None
|
||||
|
||||
|
||||
def call_glm_api(prompt: str, api_key: str, config: Config) -> str | None:
|
||||
"""
|
||||
Call the GLM-4.7 API to get code review with retry logic.
|
||||
|
||||
Args:
|
||||
prompt: The prompt to send to the API.
|
||||
api_key: The GLM API key.
|
||||
config: Configuration object.
|
||||
|
||||
Returns:
|
||||
The review content as a string, or None if all retries failed.
|
||||
"""
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"model": config.glm_model,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"temperature": 0.3,
|
||||
"max_tokens": 32000,
|
||||
"top_p": 0.9,
|
||||
}
|
||||
|
||||
last_error: str | None = None
|
||||
|
||||
for attempt in range(config.max_retries):
|
||||
try:
|
||||
logger.info(f"Calling GLM API ({config.glm_model}) - Attempt {attempt + 1}/{config.max_retries}")
|
||||
response = requests.post(
|
||||
config.glm_api_url,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=120,
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
review = validate_glm_response(data)
|
||||
if review:
|
||||
logger.info(f"Successfully received review ({len(review)} chars)") # noqa: E501
|
||||
return review
|
||||
else:
|
||||
last_error = "Failed to validate API response structure"
|
||||
logger.error(last_error)
|
||||
else:
|
||||
last_error = f"GLM API request failed: {response.status_code} - {response.text}"
|
||||
logger.error(last_error)
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
last_error = f"GLM API request timed out (attempt {attempt + 1})"
|
||||
logger.error(last_error)
|
||||
except requests.exceptions.RequestException as e:
|
||||
last_error = f"GLM API request exception: {e}"
|
||||
logger.error(last_error)
|
||||
except json.JSONDecodeError as e:
|
||||
last_error = f"Failed to decode GLM API response as JSON: {e}"
|
||||
logger.error(last_error)
|
||||
|
||||
# Exponential backoff before retry
|
||||
if attempt < config.max_retries - 1:
|
||||
wait_time: float = config.retry_delay * (2**attempt)
|
||||
logger.info(f"Waiting {wait_time}s before retry...") # noqa: E501
|
||||
time.sleep(wait_time)
|
||||
|
||||
logger.error(
|
||||
f"All {config.max_retries} attempts failed. Last error: {last_error}" # noqa: E501
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def check_cooldown( # noqa: E501
|
||||
repo_name: str,
|
||||
pr_number: int,
|
||||
token: str,
|
||||
cooldown_minutes: int,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if the PR is within the cooldown period.
|
||||
|
||||
Args:
|
||||
repo_name: The repository name in format "owner/repo".
|
||||
pr_number: The pull request number.
|
||||
token: GitHub authentication token.
|
||||
cooldown_minutes: Cooldown period in minutes.
|
||||
|
||||
Returns:
|
||||
True if within cooldown period (should skip), False otherwise.
|
||||
"""
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
url: str = (
|
||||
f"https://api.github.com/repos/{repo_name}/issues/"
|
||||
f"{pr_number}/comments" # noqa: E501
|
||||
)
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/vnd.github.v3+json",
|
||||
}
|
||||
|
||||
logger.info(f"Checking cooldown period ({cooldown_minutes} minutes)")
|
||||
response = requests.get(url, headers=headers, timeout=30)
|
||||
|
||||
if response.status_code != 200:
|
||||
logger.warning(f"Failed to check cooldown: {response.status_code}, proceeding with review")
|
||||
return False
|
||||
|
||||
comments: list[dict[str, Any]] = response.json()
|
||||
cutoff_time: datetime = datetime.utcnow() - timedelta(minutes=cooldown_minutes) # noqa: E501
|
||||
|
||||
for comment in reversed(comments):
|
||||
# Check if this is a bot comment
|
||||
body: str = comment.get("body", "")
|
||||
if "VLLM-Omni PR Review" in body or "PR Reviewer Bot" in body:
|
||||
created_at_str: str = comment.get("created_at", "")
|
||||
try:
|
||||
# Parse GitHub timestamp format
|
||||
created_at = datetime.fromisoformat(created_at_str.replace("Z", "+00:00"))
|
||||
created_at = created_at.replace(tzinfo=None)
|
||||
if created_at > cutoff_time:
|
||||
logger.info(f"PR is within cooldown period (last review: {created_at_str})")
|
||||
return True
|
||||
except ValueError:
|
||||
logger.warning(f"Failed to parse comment timestamp: {created_at_str}") # noqa: E501
|
||||
continue
|
||||
|
||||
logger.info("PR is outside cooldown period, proceeding with review")
|
||||
return False
|
||||
|
||||
|
||||
def post_review_comment( # noqa: E501
|
||||
repo_name: str,
|
||||
pr_number: int,
|
||||
token: str,
|
||||
review: str,
|
||||
) -> bool:
|
||||
"""
|
||||
Post the review as a comment on the PR.
|
||||
|
||||
Args:
|
||||
repo_name: The repository name in format "owner/repo".
|
||||
pr_number: The pull request number.
|
||||
token: GitHub authentication token.
|
||||
review: The review content to post.
|
||||
|
||||
Returns:
|
||||
True if posting succeeded, False otherwise.
|
||||
"""
|
||||
url: str = (
|
||||
f"https://api.github.com/repos/{repo_name}/issues/"
|
||||
f"{pr_number}/comments" # noqa: E501
|
||||
)
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/vnd.github.v3+json",
|
||||
}
|
||||
|
||||
# Format the review comment
|
||||
comment_body: str = f"""## 🤖 VLLM-Omni PR Review
|
||||
|
||||
{review}
|
||||
|
||||
---
|
||||
*This review was generated automatically by the VLLM-Omni PR Reviewer Bot
|
||||
using {os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL)}.*
|
||||
"""
|
||||
|
||||
payload: dict[str, str] = {"body": comment_body}
|
||||
|
||||
logger.info(f"Posting review comment to PR #{pr_number}")
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=30)
|
||||
|
||||
if response.status_code == 201:
|
||||
logger.info("Successfully posted review comment")
|
||||
return True
|
||||
else:
|
||||
logger.error(f"Failed to post comment: {response.status_code}")
|
||||
logger.error(f"Response: {response.text}")
|
||||
return False
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""
|
||||
Main entry point for the PR reviewer bot.
|
||||
|
||||
Returns:
|
||||
0 on success, 1 on error.
|
||||
"""
|
||||
logger.info("VLLM-Omni PR Reviewer Bot starting...")
|
||||
|
||||
# Load configuration
|
||||
config: Config = get_config()
|
||||
logger.info(
|
||||
f"Configuration: model={config.glm_model}, "
|
||||
f"cooldown={config.cooldown_minutes}min, "
|
||||
f"max_retries={config.max_retries}"
|
||||
)
|
||||
|
||||
# Get environment variables
|
||||
token: str = get_env_var("GITHUB_TOKEN")
|
||||
api_key: str = get_env_var("GLM_API_KEY")
|
||||
repo_name: str = get_env_var("REPO_NAME")
|
||||
pr_number_str: str = get_env_var("PR_NUMBER")
|
||||
comment_body: str = get_env_var("COMMENT_BODY")
|
||||
|
||||
try:
|
||||
pr_number: int = int(pr_number_str)
|
||||
except ValueError:
|
||||
logger.error(f"Invalid PR number: {pr_number_str}")
|
||||
return 1
|
||||
|
||||
logger.info(f"Repository: {repo_name}")
|
||||
logger.info(f"PR Number: {pr_number}")
|
||||
|
||||
# Check if the comment contains the trigger phrase
|
||||
if not check_trigger(comment_body):
|
||||
logger.info(
|
||||
f"Comment does not contain trigger phrase '{TRIGGER_PHRASE}', exiting" # noqa: E501
|
||||
)
|
||||
return 0
|
||||
|
||||
logger.info("Trigger phrase detected! Starting review process...")
|
||||
|
||||
# Check cooldown period
|
||||
if check_cooldown(repo_name, pr_number, token, config.cooldown_minutes):
|
||||
logger.info("Skipping review due to cooldown period")
|
||||
return 0
|
||||
|
||||
# Fetch PR details
|
||||
logger.info("Step 1/4: Fetching PR details...")
|
||||
pr_details: PRDetails | None = fetch_pr_details(repo_name, pr_number, token) # noqa: E501
|
||||
if not pr_details:
|
||||
logger.error("Failed to fetch PR details")
|
||||
return 1
|
||||
|
||||
pr_title: str = pr_details.get("title", "Unknown")
|
||||
pr_description: str = pr_details.get("body", "")
|
||||
|
||||
logger.info(f"PR Title: {pr_title}")
|
||||
|
||||
# Fetch PR diff
|
||||
logger.info("Step 2/4: Fetching PR diff...")
|
||||
diff: str | None = fetch_pr_diff(repo_name, pr_number, token, config.max_diff_size)
|
||||
if diff is None:
|
||||
logger.error("Failed to fetch PR diff")
|
||||
return 1
|
||||
|
||||
if not diff:
|
||||
logger.warning("Warning: Empty diff - this might be a draft PR or no code changes")
|
||||
|
||||
# Build prompt
|
||||
logger.info("Step 3/4: Building review prompt...")
|
||||
prompt: str = build_review_prompt(pr_title, pr_description, diff)
|
||||
|
||||
# Call GLM API
|
||||
logger.info("Step 4/4: Calling GLM API...")
|
||||
review: str | None = call_glm_api(prompt, api_key, config)
|
||||
if not review:
|
||||
logger.error("Failed to get review from GLM API")
|
||||
return 1
|
||||
|
||||
# Post review comment
|
||||
logger.info("Posting review comment...")
|
||||
if not post_review_comment(repo_name, pr_number, token, review):
|
||||
logger.error("Failed to post review comment")
|
||||
return 1
|
||||
|
||||
logger.info("PR review completed successfully!")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -1,62 +0,0 @@
|
||||
name: VLLM-Omni PR Reviewer
|
||||
|
||||
on:
|
||||
issue_comment:
|
||||
types: [created]
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
issues: write
|
||||
|
||||
jobs:
|
||||
pr-reviewer:
|
||||
name: Review Pull Request
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
# Only run when the comment is from a collaborator/owner/member
|
||||
if: |
|
||||
github.event_name == 'issue_comment' &&
|
||||
github.event.issue.pull_request != null &&
|
||||
contains(github.event.comment.body, '@vllm-omni-reviewer') &&
|
||||
(github.event.comment.author_association == 'MEMBER' ||
|
||||
github.event.comment.author_association == 'COLLABORATOR' ||
|
||||
github.event.comment.author_association == 'OWNER')
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6.0.2
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6.2.0
|
||||
with:
|
||||
python-version: '3.11'
|
||||
cache: 'pip'
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip==24.0
|
||||
pip install requests==2.31.0 pyyaml==6.0.1
|
||||
|
||||
- name: Run PR Reviewer
|
||||
id: reviewer
|
||||
env:
|
||||
GLM_API_KEY: ${{ secrets.GLM_API_KEY }}
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
PR_NUMBER: ${{ github.event.issue.number || github.event.pull_request.number }}
|
||||
COMMENT_BODY: ${{ github.event.comment.body }}
|
||||
REPO_NAME: ${{ github.repository }}
|
||||
PR_REVIEWER_COOLDOWN_MINUTES: 5
|
||||
PR_REVIEWER_MAX_RETRIES: 3
|
||||
run: |
|
||||
python .github/scripts/pr_reviewer.py 2>&1 | tee "/tmp/pr_review_${PR_NUMBER}.log"
|
||||
|
||||
- name: Upload review logs
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v7.0.0
|
||||
with:
|
||||
name: pr-review-logs-${{ github.event.issue.number || github.event.pull_request.number }}
|
||||
path: /tmp/pr_review_*.log
|
||||
retention-days: 7
|
||||
if-no-files-found: ignore
|
||||
@@ -239,6 +239,7 @@ data/
|
||||
datasets/
|
||||
*.csv
|
||||
*.json
|
||||
!tests/dfx/perf/tests/*.json
|
||||
!apps/ComfyUI-vLLM-Omni/example_workflows/*.json
|
||||
*.jsonl
|
||||
*.parquet
|
||||
@@ -264,5 +265,8 @@ tmp_test
|
||||
vllm_omni/_version.py
|
||||
# output files
|
||||
*.wav
|
||||
# Vendored test reference audio (small, intentionally checked in)
|
||||
!tests/assets/**/*.wav
|
||||
.worktrees/
|
||||
# CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
|
||||
tests/.ci_generated/
|
||||
@@ -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://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
|
||||
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/111-L8zF7A1j_YI_cR8JsblofdScdRr2f/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><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.
|
||||
|
||||
---
|
||||
|
||||
@@ -23,5 +23,5 @@ Test guidance:
|
||||
- Local static/self-checks live in `tests/benchmarks/test_accuracy_bench_utils.py`.
|
||||
- End-to-end generation/evaluation should be validated in a remote GPU
|
||||
environment. In the current repo marker system there is `L4` but no `L5`
|
||||
marker, so benchmark smoke tests should be wired as `advanced_model +
|
||||
benchmark + L4` when GPU capacity is available.
|
||||
marker, so benchmark smoke tests should be wired as `full_model +
|
||||
benchmark + L4` for nightly when GPU capacity is available.
|
||||
@@ -99,5 +99,5 @@ Notes:
|
||||
- This flow requires the optional Hugging Face `datasets` package.
|
||||
- `generate` writes `generation_manifest.json` with local output coverage.
|
||||
- The current repo marker set exposes `L4` but not `L5`, so if you promote an
|
||||
end-to-end smoke test into CI, use the existing `advanced_model`, `benchmark`,
|
||||
and `L4` markers or introduce a new repo-wide marker explicitly first.
|
||||
end-to-end smoke test into CI, use the `full_model`, `benchmark`,
|
||||
and `L4` markers for nightly (or `advanced_model` for merge) or introduce a new repo-wide marker explicitly first.
|
||||
@@ -0,0 +1,20 @@
|
||||
vd001|||The quick brown fox jumps over the lazy dog.|A warm, friendly female voice with a slight American Midwest accent, speaking at a moderate pace with natural inflection.
|
||||
vd002|||Welcome to the future of text-to-speech synthesis.|A deep, authoritative male news anchor voice, clear and professional with a measured cadence.
|
||||
vd003|||The sunset painted the sky in brilliant shades of orange and pink.|A gentle elderly female voice, soft and wise, with a slight Southern American accent.
|
||||
vd004|||Scientists have discovered a new species of deep-sea creature.|A young male voice with an Australian accent, curious and enthusiastic.
|
||||
vd005|||Breaking news: a major climate summit opens today in Geneva.|A crisp female newsreader voice, neutral accent, confident and precise.
|
||||
vd006|||In the beginning, there was darkness and silence across the universe.|A rich, dramatic bass male narrator voice, slow and deeply resonant.
|
||||
vd007|||Come closer, I have something important to tell you.|A soft, intimate female voice, slightly whispery, warm and gentle.
|
||||
vd008|||And they're off! The horses race toward the first turn at incredible speed.|An energetic male sports commentator, fast-paced and excited.
|
||||
vd009|||Once upon a time, in a land far away, lived a very clever fox.|A light, playful voice with childlike enthusiasm, bright and clear.
|
||||
vd010|||The ancient manuscript reveals secrets hidden for a thousand years.|A wise, measured elderly male voice, slow and deliberate, British English accent.
|
||||
vd011|||Good evening, ladies and gentlemen, and welcome to our show.|A sophisticated female voice with a slight French accent speaking English, elegant and refined.
|
||||
vd012|||System initialized. Running diagnostics. All systems nominal.|A clear, precise robotic-sounding voice, neutral and monotone with slight synthetic quality.
|
||||
vd013|||I hear what you are saying, and it is completely understandable to feel that way.|A warm, empathetic female therapist voice, calm and reassuring, unhurried pace.
|
||||
vd014|||Attention all units: proceed to grid reference seven-seven-alpha.|A firm, authoritative military male voice, clipped and commanding.
|
||||
vd015|||Oh my goodness, you have to try this amazing new recipe I just found!|An enthusiastic, bubbly female voice, high energy and friendly.
|
||||
vd016|||Dude, the waves were totally amazing out there today. Super happy about it!|A relaxed male voice with a California accent, casual and laid-back.
|
||||
vd017|||The quarterly results exceed expectations across all major metrics.|A sharp, businesslike female voice, confident and efficient, fast-paced delivery.
|
||||
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.
|
||||
@@ -0,0 +1,20 @@
|
||||
smoke001|||The quick brown fox jumps over the lazy dog near the riverbank at sunset.
|
||||
smoke002|||Welcome to the future of text-to-speech synthesis in production systems.
|
||||
smoke003|||Yesterday the team finished rolling out the new authentication flow.
|
||||
smoke004|||She walked carefully across the wet cobblestones, careful not to slip.
|
||||
smoke005|||The conference call is scheduled for nine in the morning, Pacific time.
|
||||
smoke006|||Please remember to save your work before closing the editor.
|
||||
smoke007|||Two plus two equals four, but five hundred and forty three digits is long.
|
||||
smoke008|||I would like a coffee with oat milk and a chocolate croissant please.
|
||||
smoke009|||The library closes at eight on weekdays and six on Saturdays.
|
||||
smoke010|||During the Renaissance, art and science flourished in European cities.
|
||||
smoke011|||He whispered the secret word so quietly that no one else could hear.
|
||||
smoke012|||Our flight departs from gate twenty three at eleven fifteen.
|
||||
smoke013|||The storm knocked out power for six hours, but the backup generator kicked in.
|
||||
smoke014|||Reading a good book on a rainy afternoon is one of life's great pleasures.
|
||||
smoke015|||When the kettle whistled, she poured the hot water over the fresh tea leaves.
|
||||
smoke016|||The algorithm runs in linear time, which is a big improvement over the previous approach.
|
||||
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.
|
||||
@@ -130,9 +130,21 @@ Warmup flags:
|
||||
|
||||
- `--warmup-requests`: Number of warmup requests.
|
||||
- `--warmup-num-inference-steps`: Steps used during warmup.
|
||||
- `--warmup-concurrency`: Maximum concurrent warmup requests. Use this to warm
|
||||
the same batch shape as the measured run instead of warming only batch=`1`.
|
||||
- For `--task t2v`: warmup requests are forced to use `num_frames=1` to make warmup faster and less noisy.
|
||||
|
||||
Traffic / concurrency flags:
|
||||
|
||||
- `--request-rate`: Target request rate (requests/second). If set to `inf`, the script sends all requests immediately.
|
||||
- `--max-concurrency`: Max number of in-flight requests (default: `1`). This can hard-cap the achieved QPS: if it is too small, requests will queue behind the semaphore, and both achieved throughput and observed SLO attainment can be skewed.
|
||||
|
||||
### Batched warmup note
|
||||
|
||||
For batched serving runs, warm the same in-flight shape you plan to measure.
|
||||
For example, a run with `--max-concurrency 8` should usually also use
|
||||
`--warmup-requests 8 --warmup-concurrency 8`; otherwise the first measured
|
||||
batch may still pay compile or CUDA-graph capture cost.
|
||||
|
||||
For a Qwen-Image continuous-batching replay example, see
|
||||
[`performance_dashboard/qwen_image_serving_performance.md`](./performance_dashboard/qwen_image_serving_performance.md).
|
||||
@@ -122,6 +122,18 @@ async def async_request_chat_completions(
|
||||
output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
|
||||
except (IndexError, TypeError, AttributeError):
|
||||
pass
|
||||
|
||||
if (not output.stage_durations or output.peak_memory_mb == 0.0) and isinstance(
|
||||
resp_json.get("metrics"), dict
|
||||
):
|
||||
m = resp_json["metrics"]
|
||||
if not output.stage_durations and isinstance(m.get("stage_durations"), dict):
|
||||
output.stage_durations = m.get("stage_durations") or {}
|
||||
if output.peak_memory_mb == 0.0 and m.get("peak_memory_mb") is not None:
|
||||
try:
|
||||
output.peak_memory_mb = float(m.get("peak_memory_mb") or 0.0)
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
else:
|
||||
output.error = f"HTTP {response.status}: {await response.text()}"
|
||||
output.success = False
|
||||
|
||||
@@ -12,15 +12,15 @@ Supports multiple backends:
|
||||
- v1/videos: Use /v1/videos endpoint
|
||||
|
||||
Usage:
|
||||
# Video (vllm-omni backend)
|
||||
# Video (v1/videos backend)
|
||||
t2v:
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni --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 \
|
||||
--backend vllm-omni --dataset vbench --task i2v --num-prompts 10
|
||||
--backend v1/videos --dataset vbench --task i2v --num-prompts 10
|
||||
|
||||
|
||||
# Image (vllm-omni backend)
|
||||
@@ -49,7 +49,7 @@ Usage:
|
||||
--backend openai --dataset vbench --task t2i --num-prompts 10 \
|
||||
--height 1024 --width 1024 --port 3000
|
||||
|
||||
# Video (v1/vedeos)
|
||||
# Video (v1/videos)
|
||||
t2v:
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend v1/videos --dataset random --task t2v --num-prompts 1 \
|
||||
@@ -751,6 +751,52 @@ async def iter_requests(
|
||||
yield req
|
||||
|
||||
|
||||
def _make_warmup_request(
|
||||
requests_list: list[RequestFuncInput],
|
||||
index: int,
|
||||
args,
|
||||
) -> RequestFuncInput:
|
||||
warm_req = requests_list[index % len(requests_list)]
|
||||
if args.warmup_num_inference_steps is not None:
|
||||
warm_req = replace(
|
||||
warm_req,
|
||||
num_inference_steps=args.warmup_num_inference_steps,
|
||||
)
|
||||
if args.task == "t2v":
|
||||
warm_req = replace(warm_req, num_frames=1)
|
||||
return warm_req
|
||||
|
||||
|
||||
async def _run_warmups(
|
||||
requests_list: list[RequestFuncInput],
|
||||
args,
|
||||
session: aiohttp.ClientSession,
|
||||
request_func,
|
||||
) -> list[tuple[RequestFuncInput, RequestFuncOutput]]:
|
||||
if not args.warmup_requests or not requests_list:
|
||||
return []
|
||||
|
||||
warmup_requests = [_make_warmup_request(requests_list, i, args) for i in range(args.warmup_requests)]
|
||||
warmup_concurrency = min(int(args.warmup_concurrency), len(warmup_requests))
|
||||
warmup_semaphore = asyncio.Semaphore(warmup_concurrency)
|
||||
|
||||
print(
|
||||
f"Running {len(warmup_requests)} warmup request(s) "
|
||||
f"with num_inference_steps={args.warmup_num_inference_steps} "
|
||||
f"and warmup_concurrency={warmup_concurrency}..."
|
||||
)
|
||||
|
||||
async def limited_warmup_request_func(
|
||||
req: RequestFuncInput,
|
||||
) -> RequestFuncOutput:
|
||||
async with warmup_semaphore:
|
||||
return await request_func(req, session, None)
|
||||
|
||||
warmup_tasks = [asyncio.create_task(limited_warmup_request_func(req)) for req in warmup_requests]
|
||||
warmup_outputs = await asyncio.gather(*warmup_tasks)
|
||||
return list(zip(warmup_requests, warmup_outputs))
|
||||
|
||||
|
||||
def calculate_metrics(
|
||||
outputs: list[RequestFuncOutput],
|
||||
total_duration: float,
|
||||
@@ -900,23 +946,12 @@ async def benchmark(args):
|
||||
pbar = tqdm(total=len(requests_list), disable=args.disable_tqdm)
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
warmup_pairs: list[tuple[RequestFuncInput, RequestFuncOutput]] = []
|
||||
if args.warmup_requests and requests_list:
|
||||
print(
|
||||
f"Running {args.warmup_requests} warmup request(s) \
|
||||
with num_inference_steps={args.warmup_num_inference_steps}..."
|
||||
)
|
||||
for i in range(args.warmup_requests):
|
||||
warm_req = requests_list[i % len(requests_list)]
|
||||
if args.warmup_num_inference_steps is not None:
|
||||
warm_req = replace(
|
||||
warm_req,
|
||||
num_inference_steps=args.warmup_num_inference_steps,
|
||||
)
|
||||
if args.task == "t2v":
|
||||
warm_req = replace(warm_req, num_frames=1)
|
||||
warm_out = await limited_request_func(warm_req, session, None)
|
||||
warmup_pairs.append((warm_req, warm_out))
|
||||
warmup_pairs = await _run_warmups(
|
||||
requests_list=requests_list,
|
||||
args=args,
|
||||
session=session,
|
||||
request_func=request_func,
|
||||
)
|
||||
|
||||
if args.slo:
|
||||
# Prefer trace-provided per-request slo_ms. Only populate when missing.
|
||||
@@ -1072,6 +1107,13 @@ if __name__ == "__main__":
|
||||
default=1,
|
||||
help="num_inference_steps used for warmup requests.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup-concurrency",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Maximum number of warmup requests to run concurrently. "
|
||||
"Set this to match the real batch shape when warming up torch.compile or CUDA graphs.",
|
||||
)
|
||||
parser.add_argument("--width", type=int, default=None, help="Image/Video width.")
|
||||
parser.add_argument("--height", type=int, default=None, help="Image/Video height.")
|
||||
parser.add_argument("--num-frames", type=int, default=None, help="Number of frames (for video).")
|
||||
@@ -1135,5 +1177,4 @@ if __name__ == "__main__":
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
asyncio.run(benchmark(args))
|
||||
@@ -33,6 +33,12 @@ vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
To replay step-wise continuous batching, compare
|
||||
`vllm serve Qwen/Qwen-Image --omni --port 8089 --step-execution --max-num-seqs 1`
|
||||
against the same command with `--max-num-seqs 8`. `--step-execution` is the
|
||||
feature gate, and increasing `--max-num-seqs` above `1` lets the scheduler
|
||||
keep more compatible requests active at once.
|
||||
|
||||
## 3.2 Key Parameters
|
||||
|
||||
| Parameter | Description |
|
||||
@@ -70,6 +76,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
| `--task` | Task type (e.g., `t2i`) |
|
||||
| `--num-prompts` | Total number of requests |
|
||||
| `--max-concurrency` | Client-side concurrency |
|
||||
| `--warmup-concurrency` | Warmup concurrency for batch-shape prewarming |
|
||||
| `--random-request-config`| JSON string defining random request |
|
||||
|
||||
---
|
||||
@@ -133,6 +140,12 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
]'
|
||||
```
|
||||
|
||||
For a continuous-batching A/B replay, run the same benchmark command twice
|
||||
against the two serving configs above and keep traffic plus warmup settings
|
||||
identical across both runs. If the measured run uses `--max-concurrency 8`,
|
||||
warm with `--warmup-requests 8 --warmup-concurrency 8` so the first measured
|
||||
batch does not include compile or CUDA-graph capture overhead.
|
||||
|
||||
---
|
||||
|
||||
# 6. Performance Metrics
|
||||
@@ -163,6 +176,7 @@ To ensure consistent and comparable benchmark results:
|
||||
* Record parallel configuration
|
||||
* Record benchmark parameters (resolution, concurrency, number of prompts)
|
||||
* Ensure no background workload on GPUs during testing
|
||||
* When comparing continuous batching, keep warmup settings the same across runs
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -151,7 +151,7 @@ def _generate_image(omni, args, prompt, seed):
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
torch.accelerator.reset_peak_memory_stats()
|
||||
start = time.perf_counter()
|
||||
outputs = omni.generate(
|
||||
{"prompt": prompt},
|
||||
@@ -163,7 +163,7 @@ def _generate_image(omni, args, prompt, seed):
|
||||
),
|
||||
)
|
||||
elapsed = time.perf_counter() - start
|
||||
peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
|
||||
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
|
||||
|
||||
first = outputs[0]
|
||||
req_out = first.request_output[0] if hasattr(first, "request_output") else first
|
||||
@@ -178,7 +178,7 @@ def _generate_video(omni, args, prompt, seed):
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
torch.accelerator.reset_peak_memory_stats()
|
||||
start = time.perf_counter()
|
||||
outputs = omni.generate(
|
||||
{"prompt": prompt, "negative_prompt": ""},
|
||||
@@ -192,7 +192,7 @@ def _generate_video(omni, args, prompt, seed):
|
||||
),
|
||||
)
|
||||
elapsed = time.perf_counter() - start
|
||||
peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
|
||||
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
|
||||
|
||||
first = outputs[0]
|
||||
if hasattr(first, "request_output") and isinstance(first.request_output, list):
|
||||
@@ -228,8 +228,8 @@ def _unload_omni(omni):
|
||||
del omni
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
torch.accelerator.empty_cache()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
|
||||
def run_benchmark(args):
|
||||
|
||||
@@ -614,8 +614,8 @@ Examples:
|
||||
if not torch.cuda.is_available():
|
||||
print("[ERROR] CUDA is not available but GPU mode was requested.")
|
||||
sys.exit(1)
|
||||
if args.gpu_id >= torch.cuda.device_count():
|
||||
print(f"[ERROR] GPU {args.gpu_id} not available. Found {torch.cuda.device_count()} GPUs.")
|
||||
if args.gpu_id >= torch.accelerator.device_count():
|
||||
print(f"[ERROR] GPU {args.gpu_id} not available. Found {torch.accelerator.device_count()} GPUs.")
|
||||
sys.exit(1)
|
||||
print(f"[INFO] Using GPU {args.gpu_id}: {torch.cuda.get_device_name(args.gpu_id)}")
|
||||
|
||||
|
||||
+2
-2
@@ -3,11 +3,11 @@
|
||||
Measures TTFP improvement from DAC-code caching when using uploaded voices.
|
||||
|
||||
Setup:
|
||||
1. Start vllm-omni with Fish Speech S2 Pro (use our feat branch)
|
||||
1. Start vllm-omni with Fish Speech S2 Pro
|
||||
2. Provide a reference audio file for voice cloning
|
||||
|
||||
Usage:
|
||||
python bench_voice_cache.py \
|
||||
python bench_speaker_cache.py \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio." \
|
||||
--num-prompts 20 \
|
||||
@@ -0,0 +1,157 @@
|
||||
# GLM-Image Benchmarks
|
||||
|
||||
Benchmark GLM-Image T2I (text-to-image) and I2I (image-to-image) performance across three backends: HuggingFace baseline, vLLM-Omni offline, and vLLM-Omni online serving.
|
||||
|
||||
## Benchmarks
|
||||
|
||||
| Benchmark | Script | Description |
|
||||
|-----------|--------|-------------|
|
||||
| HuggingFace Baseline | `huggingface/inference.py` | Single-GPU transformers + diffusers pipeline |
|
||||
| vLLM-Omni Offline | `vllm-omni/inference.py` | Offline inference with continuous batching |
|
||||
| vLLM-Omni Online | `benchmark_glm_image.py` | Online serving via `/v1/chat/completions` |
|
||||
|
||||
## HuggingFace Baseline
|
||||
|
||||
Single-request sequential inference using the reference HuggingFace pipeline.
|
||||
|
||||
```bash
|
||||
# T2I
|
||||
CUDA_VISIBLE_DEVICES=0 python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path /path/to/GLM-Image --mode t2i --num-prompts 10
|
||||
|
||||
# I2I
|
||||
CUDA_VISIBLE_DEVICES=0 python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path /path/to/GLM-Image --mode i2i --num-prompts 10
|
||||
```
|
||||
|
||||
### Options
|
||||
|
||||
| Flag | Default | Description |
|
||||
|------|---------|-------------|
|
||||
| `--model-path` | `zai-org/GLM-Image` | Model path |
|
||||
| `--mode` | `t2i` | `t2i` or `i2i` |
|
||||
| `--dataset-path` | `prompt/prompt.json` | Path to prompt.json |
|
||||
| `--num-prompts` | `10` | Number of images to generate |
|
||||
| `--width` / `--height` | `1024` | Output image size |
|
||||
| `--num-inference-steps` | `50` | Diffusion denoising steps |
|
||||
| `--output-dir` | `benchmarks/glm_image/huggingface/outputs` | Output directory |
|
||||
| `--output-file` | - | JSON file for metrics |
|
||||
|
||||
## vLLM-Omni Offline
|
||||
|
||||
Multi-GPU offline inference with pipeline parallelism and continuous batching.
|
||||
|
||||
```bash
|
||||
# T2I
|
||||
CUDA_VISIBLE_DEVICES=0,1 python benchmarks/glm_image/vllm-omni/inference.py \
|
||||
--model-path /path/to/GLM-Image --mode t2i --num-prompts 10
|
||||
|
||||
# I2I
|
||||
CUDA_VISIBLE_DEVICES=0,1 python benchmarks/glm_image/vllm-omni/inference.py \
|
||||
--model-path /path/to/GLM-Image --mode i2i --num-prompts 10
|
||||
```
|
||||
|
||||
### Options
|
||||
|
||||
| Flag | Default | Description |
|
||||
|------|---------|-------------|
|
||||
| `--model-path` | `zai-org/GLM-Image` | Model path |
|
||||
| `--deploy-config` | - | Deploy config YAML |
|
||||
| `--mode` | `t2i` | `t2i` or `i2i` |
|
||||
| `--dataset-path` | `prompt/prompt.json` | Path to prompt.json |
|
||||
| `--num-prompts` | `10` | Number of images to generate |
|
||||
| `--width` / `--height` | `1024` | Output image size |
|
||||
| `--num-inference-steps` | `50` | Diffusion denoising steps |
|
||||
| `--output-dir` | `benchmarks/glm_image/vllm-omni/outputs` | Output directory |
|
||||
| `--output-file` | - | JSON file for metrics |
|
||||
| `--stage-init-timeout` | `600` | Stage initialization timeout (s) |
|
||||
|
||||
### Latency Computation
|
||||
|
||||
In offline mode all requests are submitted simultaneously and processed with continuous batching. The per-request latency is computed by summing the actual per-stage times (with `stage_0_gen_ms` diffed against the previous request to remove accumulated queue/scheduling wait).
|
||||
|
||||
## vLLM-Omni Online Serving
|
||||
|
||||
### Start the server
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1 vllm serve /path/to/GLM-Image \
|
||||
--omni --port 8091 --host 0.0.0.0 \
|
||||
--served-model-name glm-image
|
||||
```
|
||||
|
||||
### Run the benchmark
|
||||
|
||||
```bash
|
||||
# T2I
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode t2i --num-prompts 10 --model glm-image
|
||||
|
||||
# I2I
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode i2i --num-prompts 10 --model glm-image
|
||||
|
||||
# Custom dataset
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode i2i --dataset custom \
|
||||
--dataset-path prompts.json --num-prompts 5
|
||||
```
|
||||
|
||||
### Options
|
||||
|
||||
| Flag | Default | Description |
|
||||
|------|---------|-------------|
|
||||
| `--mode` | `t2i` | `t2i` or `i2i` |
|
||||
| `--dataset` | `prompt` | `prompt`, `random`, or `custom` |
|
||||
| `--dataset-path` | - | JSON file path (required for `custom`) |
|
||||
| `--num-prompts` | `10` | Number of benchmark requests |
|
||||
| `--max-concurrency` | `1` | Max concurrent requests |
|
||||
| `--request-rate` | `inf` | Requests per second (Poisson arrival) |
|
||||
| `--warmup-requests` | `1` | Warmup requests before measurement |
|
||||
| `--width` / `--height` | `1024` | Output image size |
|
||||
| `--num-inference-steps` | `50` | Diffusion denoising steps |
|
||||
| `--seed` | - | Random seed |
|
||||
| `--model` | `default` | Model name (must match `--served-model-name`) |
|
||||
| `--host` | `localhost` | Server host |
|
||||
| `--port` | `8091` | Server port |
|
||||
| `--output-file` | - | JSON output file for metrics |
|
||||
| `--num-input-images` | `1` | Number of input images for random I2I |
|
||||
|
||||
## Dataset
|
||||
|
||||
The default dataset is hosted on [HuggingFace](https://huggingface.co/datasets/JaredforReal/glm-image-bench) (`prompt.json`). It is automatically downloaded and cached to `prompt/prompt.json` on first run. No manual setup needed.
|
||||
|
||||
Each entry contains:
|
||||
|
||||
- `t2i_prompt`: Text prompt for text-to-image generation
|
||||
- `i2i_prompt`: Text prompt for image-to-image editing
|
||||
- `image_url`: Source image URL for I2I (downloaded and cached on first use)
|
||||
|
||||
Custom datasets use the same JSON format and can be provided via `--dataset-path`.
|
||||
|
||||
## Pipeline Timings
|
||||
|
||||
All three benchmarks report per-stage pipeline timings (in milliseconds):
|
||||
|
||||
| Key | Description |
|
||||
|-----|-------------|
|
||||
| `preprocess_ms` | Input preprocessing (tokenization, multimodal encoding) |
|
||||
| `stage_0_gen_ms` | AR (autoregressive) model generation time |
|
||||
| `ar2diffusion_ms` | AR output to diffusion input conversion |
|
||||
| `stage_1_gen_ms` | Diffusion model denoising time |
|
||||
| `queue_wait_ms` | Queue wait time before processing |
|
||||
|
||||
The stages are ordered by execution: `preprocess → stage_0 (AR) → ar2diffusion → stage_1 (Diffusion)`.
|
||||
|
||||
## Sample Results
|
||||
|
||||
Tested on 2x GPU with 10 prompts, 1024x1024, 50 denoising steps:
|
||||
|
||||
| Backend | Mode | Latency Mean (s) | Throughput (img/s) |
|
||||
|---------|------|-------------------|--------------------|
|
||||
| HuggingFace | T2I | 72.6 | 0.014 |
|
||||
| HuggingFace | I2I | 70.9 | 0.014 |
|
||||
| vLLM-Omni Offline | T2I | 35.0 | 0.044 |
|
||||
| vLLM-Omni Offline | I2I | 31.0 | 0.053 |
|
||||
| vLLM-Omni Online | T2I | 38.8 | 0.026 |
|
||||
| vLLM-Omni Online | I2I | 34.7 | 0.029 |
|
||||
Whitespace-only changes.
@@ -0,0 +1,464 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Online serving benchmark for GLM-Image (T2I and I2I modes).
|
||||
|
||||
Sends requests to the /v1/chat/completions endpoint and reports end-to-end
|
||||
latency, throughput, and per-stage durations (when the server is started with
|
||||
--enable-diffusion-pipeline-profiler and/or --enable-ar-profiler).
|
||||
|
||||
Supports three dataset types:
|
||||
- prompt: Use prompt.json (default). T2I uses t2i_prompt, I2I uses i2i_prompt
|
||||
and sends source images from image_url.
|
||||
- random: Generate synthetic prompts (and random images for I2I).
|
||||
- custom: Load from a user-specified JSON file.
|
||||
|
||||
Usage:
|
||||
# T2I with prompt.json (default)
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode t2i --num-prompts 10
|
||||
|
||||
# I2I with prompt.json (downloads source images automatically)
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode i2i --num-prompts 10
|
||||
|
||||
# Random dataset
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode t2i --dataset random --num-prompts 20
|
||||
|
||||
# Custom dataset
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode i2i --dataset custom \
|
||||
--dataset-path my_prompts.json --num-prompts 5
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
import requests as sync_requests
|
||||
from PIL import Image
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
# Import backends from the diffusion benchmark (add parent dirs to path)
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "diffusion"))
|
||||
from backends import RequestFuncOutput
|
||||
|
||||
BENCHMARK_DIR = Path(__file__).resolve().parent
|
||||
DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
|
||||
IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
|
||||
|
||||
DATASET_REPO = "JaredforReal/glm-image-bench"
|
||||
DATASET_FILE = "prompt.json"
|
||||
|
||||
|
||||
def _ensure_prompt_json(dataset_path: str | None) -> str:
|
||||
"""Return path to prompt.json, downloading from HuggingFace if needed."""
|
||||
if dataset_path:
|
||||
return dataset_path
|
||||
local = DEFAULT_PROMPT_JSON
|
||||
if local.exists():
|
||||
return str(local)
|
||||
print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
downloaded = hf_hub_download(
|
||||
repo_id=DATASET_REPO,
|
||||
filename=DATASET_FILE,
|
||||
repo_type="dataset",
|
||||
)
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
import shutil
|
||||
|
||||
shutil.copy2(downloaded, local)
|
||||
print(f"Saved to {local}")
|
||||
except ImportError:
|
||||
url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
|
||||
import urllib.request
|
||||
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
urllib.request.urlretrieve(url, local)
|
||||
print(f"Saved to {local}")
|
||||
return str(local)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class GLMImageRequest:
|
||||
prompt: str
|
||||
image_path: str | None = None # Only for I2I mode
|
||||
|
||||
|
||||
def download_image(url: str) -> str:
|
||||
"""Download an image to cache and return the local path."""
|
||||
IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
fname = url.rsplit("/", 1)[-1]
|
||||
local_path = IMAGE_CACHE_DIR / fname
|
||||
if local_path.exists():
|
||||
return str(local_path)
|
||||
resp = sync_requests.get(url, timeout=30)
|
||||
resp.raise_for_status()
|
||||
local_path.write_bytes(resp.content)
|
||||
return str(local_path)
|
||||
|
||||
|
||||
def encode_image_as_data_url(path: str) -> str:
|
||||
"""Encode a local image file as a base64 data URL."""
|
||||
with open(path, "rb") as f:
|
||||
encoded = base64.b64encode(f.read()).decode("utf-8")
|
||||
ext = Path(path).suffix.lower()
|
||||
mime = {"png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg"}.get(ext, "image/png")
|
||||
return f"data:{mime};base64,{encoded}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Datasets
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PromptDataset:
|
||||
"""Load from prompt.json. T2I uses t2i_prompt, I2I uses i2i_prompt + image_url."""
|
||||
|
||||
def __init__(self, args: argparse.Namespace):
|
||||
path = _ensure_prompt_json(args.dataset_path)
|
||||
with open(path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
|
||||
prompt_key = "t2i_prompt" if args.mode == "t2i" else "i2i_prompt"
|
||||
self.items: list[GLMImageRequest] = []
|
||||
|
||||
for entry in raw:
|
||||
prompt = entry.get(prompt_key, "").strip()
|
||||
if not prompt:
|
||||
continue
|
||||
image_path = None
|
||||
if args.mode == "i2i":
|
||||
url = entry.get("image_url", "")
|
||||
if url:
|
||||
image_path = download_image(url)
|
||||
self.items.append(GLMImageRequest(prompt=prompt, image_path=image_path))
|
||||
|
||||
if args.num_prompts and len(self.items) > args.num_prompts:
|
||||
self.items = self.items[: args.num_prompts]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.items)
|
||||
|
||||
def __getitem__(self, idx: int) -> GLMImageRequest:
|
||||
return self.items[idx]
|
||||
|
||||
def get_requests(self) -> list[GLMImageRequest]:
|
||||
return list(self.items)
|
||||
|
||||
|
||||
class RandomDataset:
|
||||
"""Generate synthetic prompts (and optional random images for I2I)."""
|
||||
|
||||
def __init__(self, args: argparse.Namespace):
|
||||
self.args = args
|
||||
self.num_prompts = args.num_prompts
|
||||
self._random_image_paths: list[str] | None = None
|
||||
if args.mode == "i2i":
|
||||
self._random_image_paths = self._generate_random_images()
|
||||
|
||||
def _generate_random_images(self) -> list[str]:
|
||||
paths: list[str] = []
|
||||
for i in range(self.args.num_input_images):
|
||||
img = Image.new("RGB", (512, 512), (128 + i * 30 % 128, 64, 192))
|
||||
path = os.path.join(tempfile.gettempdir(), f"glm_image_bench_input_{i}.png")
|
||||
img.save(path)
|
||||
paths.append(path)
|
||||
return paths
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.num_prompts
|
||||
|
||||
def __getitem__(self, idx: int) -> GLMImageRequest:
|
||||
image_path = None
|
||||
if self._random_image_paths is not None:
|
||||
image_path = self._random_image_paths[idx % len(self._random_image_paths)]
|
||||
return GLMImageRequest(
|
||||
prompt=f"A beautiful scene with vivid colors and intricate details, prompt {idx}",
|
||||
image_path=image_path,
|
||||
)
|
||||
|
||||
def get_requests(self) -> list[GLMImageRequest]:
|
||||
return [self[i] for i in range(len(self))]
|
||||
|
||||
|
||||
class CustomDataset:
|
||||
"""Load from a user-specified JSON file.
|
||||
|
||||
Expected format:
|
||||
[
|
||||
{"prompt": "A cat sitting on a windowsill"},
|
||||
{"prompt": "Make it look like winter", "image_path": "/path/to/img.png"}
|
||||
]
|
||||
"""
|
||||
|
||||
def __init__(self, args: argparse.Namespace):
|
||||
if not args.dataset_path:
|
||||
raise ValueError("--dataset-path is required for custom dataset")
|
||||
with open(args.dataset_path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
self.items: list[GLMImageRequest] = []
|
||||
for item in raw:
|
||||
self.items.append(
|
||||
GLMImageRequest(
|
||||
prompt=item.get("prompt", ""),
|
||||
image_path=item.get("image_path"),
|
||||
)
|
||||
)
|
||||
if args.num_prompts and len(self.items) > args.num_prompts:
|
||||
self.items = self.items[: args.num_prompts]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.items)
|
||||
|
||||
def __getitem__(self, idx: int) -> GLMImageRequest:
|
||||
return self.items[idx]
|
||||
|
||||
def get_requests(self) -> list[GLMImageRequest]:
|
||||
return list(self.items)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Async request for GLM-Image (chat completions with image support)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def async_glm_image_request(
|
||||
req: GLMImageRequest,
|
||||
api_url: str,
|
||||
model: str,
|
||||
session: aiohttp.ClientSession,
|
||||
pbar: Any,
|
||||
args: argparse.Namespace,
|
||||
) -> RequestFuncOutput:
|
||||
"""Send a single T2I or I2I request via chat completions endpoint."""
|
||||
output = RequestFuncOutput()
|
||||
output.start_time = time.perf_counter()
|
||||
|
||||
# Build messages
|
||||
if req.image_path and args.mode == "i2i":
|
||||
data_url = encode_image_as_data_url(req.image_path)
|
||||
content = [
|
||||
{"type": "text", "text": req.prompt},
|
||||
{"type": "image_url", "image_url": {"url": data_url}},
|
||||
]
|
||||
else:
|
||||
content = req.prompt
|
||||
|
||||
messages = [{"role": "user", "content": content}]
|
||||
|
||||
extra_body: dict[str, Any] = {}
|
||||
if args.height:
|
||||
extra_body["height"] = args.height
|
||||
if args.width:
|
||||
extra_body["width"] = args.width
|
||||
if args.num_inference_steps:
|
||||
extra_body["num_inference_steps"] = args.num_inference_steps
|
||||
if args.seed is not None:
|
||||
extra_body["seed"] = args.seed
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
}
|
||||
if extra_body:
|
||||
payload["extra_body"] = extra_body
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status == 200:
|
||||
resp_json = await response.json()
|
||||
output.response_body = resp_json
|
||||
output.success = True
|
||||
try:
|
||||
choices = resp_json.get("choices", [])
|
||||
if choices and isinstance(choices, list):
|
||||
msg = choices[0].get("message", {})
|
||||
if isinstance(msg, dict):
|
||||
resp_content = msg.get("content", [])
|
||||
if resp_content and isinstance(resp_content, list) and len(resp_content) > 0:
|
||||
first_item = resp_content[0]
|
||||
if isinstance(first_item, dict):
|
||||
output.stage_durations = first_item.get("stage_durations") or {}
|
||||
output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
|
||||
except (IndexError, TypeError, AttributeError):
|
||||
pass
|
||||
else:
|
||||
output.error = f"HTTP {response.status}: {await response.text()}"
|
||||
output.success = False
|
||||
except Exception as e:
|
||||
output.error = str(e)
|
||||
output.success = False
|
||||
|
||||
output.latency = time.perf_counter() - output.start_time
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return output
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def iter_requests(n: int, request_rate: float) -> Any:
|
||||
import random as _random
|
||||
|
||||
for i in range(n):
|
||||
if request_rate != float("inf") and i > 0:
|
||||
await asyncio.sleep(_random.expovariate(request_rate))
|
||||
yield i
|
||||
|
||||
|
||||
def calculate_metrics(outputs: list[RequestFuncOutput], total_duration: float) -> dict[str, Any]:
|
||||
success = [o for o in outputs if o.success]
|
||||
errors = [o for o in outputs if not o.success]
|
||||
latencies = [o.latency for o in success]
|
||||
peak_mems = [o.peak_memory_mb for o in success if o.peak_memory_mb > 0]
|
||||
|
||||
stage_duration_lists: dict[str, list[float]] = {}
|
||||
for o in success:
|
||||
for stage, dur in (o.stage_durations or {}).items():
|
||||
stage_duration_lists.setdefault(stage, []).append(dur)
|
||||
|
||||
return {
|
||||
"duration": total_duration,
|
||||
"completed_requests": len(success),
|
||||
"failed_requests": len(errors),
|
||||
"throughput_qps": len(success) / total_duration if total_duration > 0 else 0,
|
||||
"latency_mean": float(np.mean(latencies)) if latencies else 0,
|
||||
"latency_median": float(np.median(latencies)) if latencies else 0,
|
||||
"latency_p99": float(np.percentile(latencies, 99)) if latencies else 0,
|
||||
"latency_p95": float(np.percentile(latencies, 95)) if latencies else 0,
|
||||
"peak_memory_mb_max": max(peak_mems) if peak_mems else 0,
|
||||
"stage_durations_mean": {s: float(np.mean(v)) for s, v in stage_duration_lists.items()},
|
||||
"stage_durations_p50": {s: float(np.percentile(v, 50)) for s, v in stage_duration_lists.items()},
|
||||
}
|
||||
|
||||
|
||||
async def benchmark(args: argparse.Namespace) -> None:
|
||||
api_url = f"http://{args.host}:{args.port}/v1/chat/completions"
|
||||
|
||||
# Load dataset
|
||||
if args.dataset == "prompt":
|
||||
dataset = PromptDataset(args)
|
||||
elif args.dataset == "random":
|
||||
dataset = RandomDataset(args)
|
||||
elif args.dataset == "custom":
|
||||
dataset = CustomDataset(args)
|
||||
else:
|
||||
raise ValueError(f"Unknown dataset: {args.dataset}")
|
||||
|
||||
glm_requests = dataset.get_requests()
|
||||
print(f"Prepared {len(glm_requests)} requests (mode={args.mode}, dataset={args.dataset})")
|
||||
|
||||
semaphore = asyncio.Semaphore(args.max_concurrency) if args.max_concurrency else None
|
||||
|
||||
async def limited_request(idx: int, req: GLMImageRequest, session: aiohttp.ClientSession, pbar: Any):
|
||||
if semaphore:
|
||||
async with semaphore:
|
||||
return await async_glm_image_request(req, api_url, args.model, session, pbar, args)
|
||||
return await async_glm_image_request(req, api_url, args.model, session, pbar, args)
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# Warmup
|
||||
if args.warmup_requests and glm_requests:
|
||||
print(f"Running {args.warmup_requests} warmup request(s)...")
|
||||
for i in range(args.warmup_requests):
|
||||
await limited_request(i, glm_requests[i % len(glm_requests)], session, None)
|
||||
|
||||
# Main benchmark
|
||||
pbar = tqdm(total=len(glm_requests), disable=args.disable_tqdm)
|
||||
start_time = time.perf_counter()
|
||||
tasks = []
|
||||
async for idx in iter_requests(len(glm_requests), args.request_rate):
|
||||
tasks.append(asyncio.create_task(limited_request(idx, glm_requests[idx], session, pbar)))
|
||||
outputs = await asyncio.gather(*tasks)
|
||||
total_duration = time.perf_counter() - start_time
|
||||
pbar.close()
|
||||
|
||||
# Metrics
|
||||
metrics = calculate_metrics(outputs, total_duration)
|
||||
metrics["mode"] = args.mode
|
||||
metrics["model"] = args.model
|
||||
metrics["dataset"] = args.dataset
|
||||
|
||||
print(f"\n{' GLM-Image Online Benchmark Result ':=^60}")
|
||||
print(f"{'Mode:':<40} {args.mode}")
|
||||
print(f"{'Model:':<40} {args.model}")
|
||||
print(f"{'Dataset:':<40} {args.dataset}")
|
||||
print("-" * 50)
|
||||
print(f"{'Benchmark duration (s):':<40} {metrics['duration']:.2f}")
|
||||
print(f"{'Request rate:':<40} {args.request_rate}")
|
||||
print(f"{'Max concurrency:':<40} {args.max_concurrency}")
|
||||
print(f"{'Successful requests:':<40} {metrics['completed_requests']}/{len(glm_requests)}")
|
||||
print("-" * 50)
|
||||
print(f"{'Throughput (req/s):':<40} {metrics['throughput_qps']:.2f}")
|
||||
print(f"{'Latency Mean (s):':<40} {metrics['latency_mean']:.4f}")
|
||||
print(f"{'Latency Median (s):':<40} {metrics['latency_median']:.4f}")
|
||||
print(f"{'Latency P95 (s):':<40} {metrics['latency_p95']:.4f}")
|
||||
print(f"{'Latency P99 (s):':<40} {metrics['latency_p99']:.4f}")
|
||||
|
||||
if metrics["peak_memory_mb_max"] > 0:
|
||||
print("-" * 50)
|
||||
print(f"{'Peak Memory Max (MB):':<40} {metrics['peak_memory_mb_max']:.2f}")
|
||||
|
||||
if metrics["stage_durations_mean"]:
|
||||
print("-" * 50)
|
||||
print("Stage Durations Mean:")
|
||||
for stage, val in sorted(metrics["stage_durations_mean"].items()):
|
||||
unit = "ms" if stage.endswith("_ms") else "s"
|
||||
print(f" {stage + ':':<38} {val:.4f} ({unit})")
|
||||
|
||||
print("=" * 60)
|
||||
|
||||
if args.output_file:
|
||||
with open(args.output_file, "w") as f:
|
||||
json.dump(metrics, f, indent=2)
|
||||
print(f"Metrics saved to {args.output_file}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Benchmark GLM-Image T2I/I2I online serving.")
|
||||
parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
|
||||
parser.add_argument("--dataset", type=str, default="prompt", choices=["prompt", "random", "custom"])
|
||||
parser.add_argument("--dataset-path", type=str, default=None)
|
||||
parser.add_argument("--num-prompts", type=int, default=10)
|
||||
parser.add_argument("--max-concurrency", type=int, default=1)
|
||||
parser.add_argument("--request-rate", type=float, default=float("inf"))
|
||||
parser.add_argument("--warmup-requests", type=int, default=1)
|
||||
parser.add_argument("--width", type=int, default=1024)
|
||||
parser.add_argument("--height", type=int, default=1024)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=50)
|
||||
parser.add_argument("--seed", type=int, default=None)
|
||||
parser.add_argument("--model", type=str, default="default")
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8091)
|
||||
parser.add_argument("--output-file", type=str, default=None)
|
||||
parser.add_argument("--disable-tqdm", action="store_true")
|
||||
parser.add_argument("--num-input-images", type=int, default=1, help="For random I2I dataset.")
|
||||
args = parser.parse_args()
|
||||
asyncio.run(benchmark(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,291 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
HuggingFace (transformers + diffusers) baseline benchmark for GLM-Image.
|
||||
|
||||
Supports T2I and I2I modes with the prompt.json dataset.
|
||||
Downloads source images for I2I from image_url on first run and caches locally.
|
||||
|
||||
Usage:
|
||||
# T2I mode (text-to-image, no source images needed)
|
||||
python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode t2i --num-prompts 10
|
||||
|
||||
# I2I mode (image-to-image, downloads source images)
|
||||
python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode i2i --num-prompts 10
|
||||
|
||||
# With custom prompt.json
|
||||
python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode i2i --dataset-path prompts.json --num-prompts 5
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
BENCHMARK_DIR = Path(__file__).resolve().parent.parent
|
||||
DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
|
||||
IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
|
||||
|
||||
DATASET_REPO = "JaredforReal/glm-image-bench"
|
||||
DATASET_FILE = "prompt.json"
|
||||
|
||||
|
||||
def _ensure_prompt_json(dataset_path: str | None) -> str:
|
||||
"""Return path to prompt.json, downloading from HuggingFace if needed."""
|
||||
if dataset_path:
|
||||
return dataset_path
|
||||
local = DEFAULT_PROMPT_JSON
|
||||
if local.exists():
|
||||
return str(local)
|
||||
print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
downloaded = hf_hub_download(
|
||||
repo_id=DATASET_REPO,
|
||||
filename=DATASET_FILE,
|
||||
repo_type="dataset",
|
||||
)
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
import shutil
|
||||
|
||||
shutil.copy2(downloaded, local)
|
||||
print(f"Saved to {local}")
|
||||
except ImportError:
|
||||
url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
|
||||
import urllib.request
|
||||
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
urllib.request.urlretrieve(url, local)
|
||||
print(f"Saved to {local}")
|
||||
return str(local)
|
||||
|
||||
|
||||
HEIGHT = 1024
|
||||
WIDTH = 1024
|
||||
SEED = 42
|
||||
NUM_INFERENCE_STEPS = 50
|
||||
GUIDANCE_SCALE = 1.5
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataset
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_dataset(
|
||||
dataset_path: str | None,
|
||||
mode: str,
|
||||
num_prompts: int,
|
||||
) -> list[dict]:
|
||||
"""Load prompts from prompt.json and prepare per-request data."""
|
||||
path = _ensure_prompt_json(dataset_path)
|
||||
with open(path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
|
||||
items = []
|
||||
for entry in raw:
|
||||
if mode == "t2i":
|
||||
prompt_key = "t2i_prompt"
|
||||
else:
|
||||
prompt_key = "i2i_prompt"
|
||||
|
||||
prompt_text = entry.get(prompt_key, "").strip()
|
||||
if not prompt_text:
|
||||
continue
|
||||
|
||||
item = {"prompt": prompt_text}
|
||||
if mode == "i2i":
|
||||
item["image_url"] = entry.get("image_url", "")
|
||||
items.append(item)
|
||||
|
||||
if num_prompts and len(items) > num_prompts:
|
||||
items = items[:num_prompts]
|
||||
return items
|
||||
|
||||
|
||||
def download_image(url: str, cache_dir: Path) -> str:
|
||||
"""Download an image to cache_dir and return the local path."""
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
fname = url.rsplit("/", 1)[-1]
|
||||
local_path = cache_dir / fname
|
||||
if local_path.exists():
|
||||
return str(local_path)
|
||||
print(f" Downloading {url} ...")
|
||||
resp = requests.get(url, timeout=30)
|
||||
resp.raise_for_status()
|
||||
local_path.write_bytes(resp.content)
|
||||
return str(local_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def benchmark(args: argparse.Namespace) -> None:
|
||||
from diffusers.pipelines.glm_image import GlmImagePipeline
|
||||
|
||||
print("=" * 60)
|
||||
print("GLM-Image HuggingFace Baseline Benchmark")
|
||||
print(f"Mode: {args.mode} | Model: {args.model_path}")
|
||||
print(f"Size: {args.height}x{args.width} | Steps: {args.num_inference_steps}")
|
||||
print("=" * 60)
|
||||
|
||||
# Load dataset
|
||||
items = load_dataset(args.dataset_path, args.mode, args.num_prompts)
|
||||
if not items:
|
||||
print("No prompts loaded. Exiting.")
|
||||
return
|
||||
print(f"Loaded {len(items)} prompts for {args.mode} mode")
|
||||
|
||||
# Download I2I source images
|
||||
if args.mode == "i2i":
|
||||
print("Preparing source images...")
|
||||
for item in items:
|
||||
url = item.get("image_url", "")
|
||||
if url:
|
||||
item["image_path"] = download_image(url, IMAGE_CACHE_DIR)
|
||||
else:
|
||||
item["image_path"] = None
|
||||
|
||||
# Load pipeline
|
||||
print(f"\nLoading pipeline from {args.model_path} ...")
|
||||
t0 = time.perf_counter()
|
||||
pipe = GlmImagePipeline.from_pretrained(
|
||||
args.model_path,
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="cuda",
|
||||
)
|
||||
init_time = time.perf_counter() - t0
|
||||
print(f"Pipeline loaded in {init_time:.2f}s")
|
||||
|
||||
# Create output dir
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# Run benchmark
|
||||
generator = torch.Generator(device="cuda").manual_seed(args.seed)
|
||||
latencies = []
|
||||
success = 0
|
||||
failed = 0
|
||||
|
||||
print(f"\nRunning {len(items)} requests sequentially...")
|
||||
print("-" * 60)
|
||||
|
||||
for i, item in enumerate(items):
|
||||
prompt = item["prompt"]
|
||||
gen_kwargs: dict = {
|
||||
"prompt": prompt,
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"guidance_scale": args.guidance_scale,
|
||||
"generator": generator,
|
||||
}
|
||||
|
||||
if args.mode == "i2i":
|
||||
img_path = item.get("image_path")
|
||||
if img_path and os.path.exists(img_path):
|
||||
gen_kwargs["image"] = [Image.open(img_path).convert("RGB")]
|
||||
else:
|
||||
print(f" [{i + 1}] SKIP: no source image")
|
||||
failed += 1
|
||||
continue
|
||||
|
||||
t_start = time.perf_counter()
|
||||
try:
|
||||
result = pipe(**gen_kwargs)
|
||||
image = result.images[0]
|
||||
elapsed = time.perf_counter() - t_start
|
||||
latencies.append(elapsed)
|
||||
success += 1
|
||||
|
||||
out_path = os.path.join(args.output_dir, f"{i:04d}.png")
|
||||
image.save(out_path)
|
||||
print(f" [{i + 1}/{len(items)}] {elapsed:.3f}s -> {out_path}")
|
||||
except Exception as e:
|
||||
elapsed = time.perf_counter() - t_start
|
||||
failed += 1
|
||||
print(f" [{i + 1}/{len(items)}] FAILED ({elapsed:.3f}s): {e}")
|
||||
|
||||
# Report
|
||||
total_gen_time = sum(latencies) if latencies else 0
|
||||
print("\n" + "=" * 60)
|
||||
print("HuggingFace Baseline Results")
|
||||
print("=" * 60)
|
||||
print(f"{'Mode:':<40} {args.mode}")
|
||||
print(f"{'Model:':<40} {args.model_path}")
|
||||
print(f"{'Image size:':<40} {args.height}x{args.width}")
|
||||
print(f"{'Num inference steps:':<40} {args.num_inference_steps}")
|
||||
print("-" * 50)
|
||||
print(f"{'Pipeline init time (s):':<40} {init_time:.2f}")
|
||||
print(f"{'Successful:':<40} {success}/{len(items)}")
|
||||
print(f"{'Failed:':<40} {failed}")
|
||||
print("-" * 50)
|
||||
if latencies:
|
||||
arr = np.array(latencies)
|
||||
print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
|
||||
print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
|
||||
print(f"{'Latency Mean (s):':<40} {arr.mean():.4f}")
|
||||
print(f"{'Latency Median (s):':<40} {np.median(arr):.4f}")
|
||||
print(f"{'Latency P95 (s):':<40} {np.percentile(arr, 95):.4f}")
|
||||
print(f"{'Latency P99 (s):':<40} {np.percentile(arr, 99):.4f}")
|
||||
|
||||
print(f"\n{'Output dir:':<40} {args.output_dir}")
|
||||
print("=" * 60)
|
||||
|
||||
# Save metrics JSON
|
||||
metrics = {
|
||||
"backend": "huggingface",
|
||||
"mode": args.mode,
|
||||
"model": args.model_path,
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"init_time_s": init_time,
|
||||
"completed_requests": success,
|
||||
"failed_requests": failed,
|
||||
"total_gen_time_s": total_gen_time,
|
||||
"throughput_qps": success / total_gen_time if total_gen_time > 0 else 0,
|
||||
"latency_mean": float(np.mean(latencies)) if latencies else 0,
|
||||
"latency_median": float(np.median(latencies)) if latencies else 0,
|
||||
"latency_p95": float(np.percentile(latencies, 95)) if latencies else 0,
|
||||
"latency_p99": float(np.percentile(latencies, 99)) if latencies else 0,
|
||||
}
|
||||
if args.output_file:
|
||||
with open(args.output_file, "w") as f:
|
||||
json.dump(metrics, f, indent=2)
|
||||
print(f"Metrics saved to {args.output_file}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="GLM-Image HuggingFace baseline benchmark")
|
||||
parser.add_argument("--model-path", type=str, default="zai-org/GLM-Image")
|
||||
parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
|
||||
parser.add_argument("--dataset-path", type=str, default=None, help="Path to prompt.json")
|
||||
parser.add_argument("--num-prompts", type=int, default=10)
|
||||
parser.add_argument("--height", type=int, default=HEIGHT)
|
||||
parser.add_argument("--width", type=int, default=WIDTH)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=NUM_INFERENCE_STEPS)
|
||||
parser.add_argument("--guidance-scale", type=float, default=GUIDANCE_SCALE)
|
||||
parser.add_argument("--seed", type=int, default=SEED)
|
||||
parser.add_argument("--output-dir", type=str, default="benchmarks/glm_image/huggingface/outputs")
|
||||
parser.add_argument("--output-file", type=str, default=None, help="JSON file for metrics")
|
||||
args = parser.parse_args()
|
||||
benchmark(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,505 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
vLLM-Omni offline benchmark for GLM-Image.
|
||||
|
||||
Supports T2I and I2I modes with the prompt.json dataset.
|
||||
Downloads source images for I2I from image_url on first run and caches locally.
|
||||
|
||||
Usage:
|
||||
# T2I mode
|
||||
python benchmarks/glm_image/vllm-omni/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode t2i --num-prompts 10
|
||||
|
||||
# I2I mode (downloads source images)
|
||||
python benchmarks/glm_image/vllm-omni/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode i2i --num-prompts 10
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
from PIL import Image
|
||||
from vllm import SamplingParams
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
BENCHMARK_DIR = Path(__file__).resolve().parent.parent
|
||||
DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
|
||||
IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
|
||||
DEFAULT_DEPLOY_CONFIG = "vllm_omni/deploy/glm_image.yaml"
|
||||
|
||||
DATASET_REPO = "JaredforReal/glm-image-bench"
|
||||
DATASET_FILE = "prompt.json"
|
||||
|
||||
|
||||
def _ensure_prompt_json(dataset_path: str | None) -> str:
|
||||
"""Return path to prompt.json, downloading from HuggingFace if needed."""
|
||||
if dataset_path:
|
||||
return dataset_path
|
||||
local = DEFAULT_PROMPT_JSON
|
||||
if local.exists():
|
||||
return str(local)
|
||||
print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
downloaded = hf_hub_download(
|
||||
repo_id=DATASET_REPO,
|
||||
filename=DATASET_FILE,
|
||||
repo_type="dataset",
|
||||
)
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
import shutil
|
||||
|
||||
shutil.copy2(downloaded, local)
|
||||
print(f"Saved to {local}")
|
||||
except ImportError:
|
||||
url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
|
||||
import urllib.request
|
||||
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
urllib.request.urlretrieve(url, local)
|
||||
print(f"Saved to {local}")
|
||||
return str(local)
|
||||
|
||||
|
||||
SEED = 42
|
||||
HEIGHT = 1024
|
||||
WIDTH = 1024
|
||||
NUM_INFERENCE_STEPS = 50
|
||||
GUIDANCE_SCALE = 1.5
|
||||
|
||||
GLM_IMAGE_EOS_TOKEN_ID = 16385
|
||||
GLM_IMAGE_VISION_VOCAB_SIZE = 16512
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataset
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_dataset(
|
||||
dataset_path: str | None,
|
||||
mode: str,
|
||||
num_prompts: int,
|
||||
) -> list[dict]:
|
||||
path = _ensure_prompt_json(dataset_path)
|
||||
with open(path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
|
||||
items = []
|
||||
for entry in raw:
|
||||
prompt_key = "t2i_prompt" if mode == "t2i" else "i2i_prompt"
|
||||
prompt_text = entry.get(prompt_key, "").strip()
|
||||
if not prompt_text:
|
||||
continue
|
||||
|
||||
item = {"prompt": prompt_text}
|
||||
if mode == "i2i":
|
||||
item["image_url"] = entry.get("image_url", "")
|
||||
items.append(item)
|
||||
|
||||
if num_prompts and len(items) > num_prompts:
|
||||
items = items[:num_prompts]
|
||||
return items
|
||||
|
||||
|
||||
def download_image(url: str, cache_dir: Path) -> str:
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
fname = url.rsplit("/", 1)[-1]
|
||||
local_path = cache_dir / fname
|
||||
if local_path.exists():
|
||||
return str(local_path)
|
||||
print(f" Downloading {url} ...")
|
||||
resp = requests.get(url, timeout=30)
|
||||
resp.raise_for_status()
|
||||
local_path.write_bytes(resp.content)
|
||||
return str(local_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def compute_max_tokens(height: int, width: int, is_i2i: bool = False) -> int:
|
||||
factor = 32
|
||||
token_h = height // factor
|
||||
token_w = width // factor
|
||||
large_tokens = token_h * token_w
|
||||
|
||||
# Small preview tokens (half resolution in each dimension)
|
||||
|
||||
ratio = token_h / token_w if token_w > 0 else 1.0
|
||||
small_token_h = max(1, int(math.sqrt(ratio) * (factor // 2)))
|
||||
small_token_w = max(1, int(math.sqrt(1 / ratio) * (factor // 2)))
|
||||
small_tokens = small_token_h * small_token_w
|
||||
|
||||
# Mode-dependent totals:
|
||||
# - t2i: small + large + EOS
|
||||
# - i2i: large + EOS
|
||||
if is_i2i:
|
||||
return large_tokens + 1
|
||||
return small_tokens + large_tokens + 1
|
||||
|
||||
|
||||
def build_prompt_t2i(prompt: str, height: int, width: int, **gen_kw) -> dict:
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"height": height,
|
||||
"width": width,
|
||||
"mm_processor_kwargs": {"target_h": height, "target_w": width},
|
||||
**gen_kw,
|
||||
}
|
||||
|
||||
|
||||
def build_prompt_i2i(prompt: str, image_path: str, height: int, width: int, **gen_kw) -> dict:
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"height": height,
|
||||
"width": width,
|
||||
"mm_processor_kwargs": {"target_h": height, "target_w": width},
|
||||
"multi_modal_data": {"image": Image.open(image_path).convert("RGB")},
|
||||
**gen_kw,
|
||||
}
|
||||
|
||||
|
||||
def resolve_deploy_config(args: argparse.Namespace) -> str:
|
||||
if args.deploy_config:
|
||||
return args.deploy_config
|
||||
if os.path.exists(DEFAULT_DEPLOY_CONFIG):
|
||||
return DEFAULT_DEPLOY_CONFIG
|
||||
fallback = Path(__file__).resolve().parents[3] / DEFAULT_DEPLOY_CONFIG
|
||||
if fallback.exists():
|
||||
return str(fallback)
|
||||
raise FileNotFoundError("Deploy config not found. Specify --deploy-config.")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def benchmark(args: argparse.Namespace) -> None:
|
||||
is_i2i = args.mode == "i2i"
|
||||
|
||||
print("=" * 60)
|
||||
print("GLM-Image vLLM-Omni Benchmark")
|
||||
print(f"Mode: {args.mode} | Model: {args.model_path}")
|
||||
print(f"Size: {args.height}x{args.width} | Steps: {args.num_inference_steps}")
|
||||
print("=" * 60)
|
||||
|
||||
# Load dataset
|
||||
items = load_dataset(args.dataset_path, args.mode, args.num_prompts)
|
||||
if not items:
|
||||
print("No prompts loaded. Exiting.")
|
||||
return
|
||||
print(f"Loaded {len(items)} prompts for {args.mode} mode")
|
||||
|
||||
# Download I2I source images
|
||||
if is_i2i:
|
||||
print("Preparing source images...")
|
||||
for item in items:
|
||||
url = item.get("image_url", "")
|
||||
if url:
|
||||
item["image_path"] = download_image(url, IMAGE_CACHE_DIR)
|
||||
else:
|
||||
item["image_path"] = None
|
||||
|
||||
# Init Omni
|
||||
deploy_config = resolve_deploy_config(args)
|
||||
print(f"\nInitializing vLLM-Omni (deploy config: {deploy_config}) ...")
|
||||
t0 = time.perf_counter()
|
||||
|
||||
omni = Omni(
|
||||
model=args.model_path,
|
||||
deploy_config=deploy_config,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
enable_ar_profiler=args.enable_ar_profiler,
|
||||
)
|
||||
|
||||
init_time = time.perf_counter() - t0
|
||||
print(f"Initialized in {init_time:.2f}s")
|
||||
|
||||
# Sampling params
|
||||
max_tokens = compute_max_tokens(args.height, args.width, is_i2i=is_i2i)
|
||||
ar_params = SamplingParams(
|
||||
temperature=0.9,
|
||||
top_p=0.75,
|
||||
top_k=GLM_IMAGE_VISION_VOCAB_SIZE,
|
||||
max_tokens=max_tokens,
|
||||
stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID],
|
||||
seed=args.seed,
|
||||
detokenize=False,
|
||||
extra_args={"target_h": args.height, "target_w": args.width},
|
||||
)
|
||||
diff_params = OmniDiffusionSamplingParams(
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
seed=args.seed,
|
||||
)
|
||||
sampling_params_list = [ar_params, diff_params]
|
||||
|
||||
# Build all prompts
|
||||
gen_kw = {
|
||||
"seed": args.seed,
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"guidance_scale": args.guidance_scale,
|
||||
}
|
||||
all_prompts = []
|
||||
for item in items:
|
||||
if is_i2i:
|
||||
img_path = item.get("image_path")
|
||||
if not img_path or not os.path.exists(img_path):
|
||||
continue
|
||||
all_prompts.append(build_prompt_i2i(item["prompt"], img_path, args.height, args.width, **gen_kw))
|
||||
else:
|
||||
all_prompts.append(build_prompt_t2i(item["prompt"], args.height, args.width, **gen_kw))
|
||||
|
||||
valid = len(all_prompts)
|
||||
print(f"Valid prompts: {valid}")
|
||||
|
||||
# Create output dir
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# Warmup: run 1 request to prime caches, CUDA graphs, etc.
|
||||
if all_prompts:
|
||||
print("Running warmup request...")
|
||||
try:
|
||||
warmup_prompt = [all_prompts[0]]
|
||||
omni.generate(warmup_prompt, sampling_params_list, py_generator=False)
|
||||
print("Warmup done.\n")
|
||||
except Exception as e:
|
||||
print(f"Warmup failed (continuing): {e}")
|
||||
|
||||
# Run
|
||||
print(f"\nRunning {valid} requests...")
|
||||
print("-" * 60)
|
||||
|
||||
latencies = []
|
||||
all_stage_durations: list[dict[str, float]] = []
|
||||
success = 0
|
||||
failed = 0
|
||||
wall_start = time.perf_counter()
|
||||
|
||||
try:
|
||||
output_idx = 0
|
||||
for stage_outputs in omni.generate(all_prompts, sampling_params_list, py_generator=True):
|
||||
if stage_outputs.final_output_type == "image":
|
||||
request_output = stage_outputs.request_output
|
||||
request_id = getattr(request_output, "request_id", "")
|
||||
|
||||
images = getattr(request_output, "images", [])
|
||||
if not images and hasattr(request_output, "multimodal_output"):
|
||||
mm = request_output.multimodal_output
|
||||
if isinstance(mm, dict):
|
||||
images = mm.get("images", [])
|
||||
|
||||
elapsed = time.perf_counter() - wall_start
|
||||
if images:
|
||||
for img in images:
|
||||
if isinstance(img, Image.Image):
|
||||
out_path = os.path.join(args.output_dir, f"{output_idx:04d}.png")
|
||||
img.save(out_path)
|
||||
success += 1
|
||||
latencies.append(elapsed)
|
||||
stage_durations = getattr(stage_outputs, "stage_durations", {})
|
||||
if stage_durations:
|
||||
all_stage_durations.append(stage_durations)
|
||||
# Show wall-clock elapsed and pipeline breakdown if available
|
||||
preprocess_str = ""
|
||||
if "preprocess_ms" in stage_durations:
|
||||
preprocess_str = f" preprocess={stage_durations['preprocess_ms'] / 1000.0:.2f}s"
|
||||
print(f" [{success}/{valid}] id={request_id[:8]} {elapsed:.2f}s{preprocess_str}")
|
||||
output_idx += 1
|
||||
else:
|
||||
failed += 1
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
failed = valid - success
|
||||
|
||||
total_gen_time = time.perf_counter() - wall_start
|
||||
|
||||
# Diff stage_0_gen_ms with previous request to remove accumulated wait time.
|
||||
# stage_0_gen_ms is measured from submit_ts (same for all requests submitted
|
||||
# at once), so it accumulates queue/scheduling overhead across requests.
|
||||
# Other stages and pipeline timings are per-request already.
|
||||
_TIMING_ORDER = [
|
||||
"preprocess_ms",
|
||||
"stage_0_gen_ms",
|
||||
"ar2diffusion_ms",
|
||||
"stage_1_gen_ms",
|
||||
"queue_wait_ms",
|
||||
]
|
||||
|
||||
per_request_actual: list[dict[str, float]] = []
|
||||
prev_stage_0_ms = 0.0
|
||||
for sd in all_stage_durations:
|
||||
actual = dict(sd)
|
||||
s0 = sd.get("stage_0_gen_ms", 0.0)
|
||||
actual["stage_0_gen_ms"] = s0 - prev_stage_0_ms
|
||||
prev_stage_0_ms = s0
|
||||
per_request_actual.append(actual)
|
||||
|
||||
per_request_e2e_ms: list[float] = []
|
||||
for actual in per_request_actual:
|
||||
e2e_ms = sum(v for k, v in actual.items() if k in _TIMING_ORDER)
|
||||
if e2e_ms > 0:
|
||||
per_request_e2e_ms.append(e2e_ms)
|
||||
|
||||
# Report
|
||||
print("\n" + "=" * 60)
|
||||
print("vLLM-Omni Benchmark Results")
|
||||
print("=" * 60)
|
||||
print(f"{'Mode:':<40} {args.mode}")
|
||||
print(f"{'Model:':<40} {args.model_path}")
|
||||
print(f"{'Image size:':<40} {args.height}x{args.width}")
|
||||
print(f"{'Num inference steps:':<40} {args.num_inference_steps}")
|
||||
print("-" * 50)
|
||||
print(f"{'Init time (s):':<40} {init_time:.2f}")
|
||||
print(f"{'Successful:':<40} {success}/{valid}")
|
||||
print(f"{'Failed:':<40} {failed}")
|
||||
print("-" * 50)
|
||||
|
||||
if per_request_e2e_ms:
|
||||
per_request_s = np.array(per_request_e2e_ms) / 1000.0
|
||||
print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
|
||||
print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
|
||||
print(f"{'Latency Mean (s):':<40} {per_request_s.mean():.4f}")
|
||||
print(f"{'Latency Median (s):':<40} {np.median(per_request_s):.4f}")
|
||||
print(f"{'Latency P95 (s):':<40} {np.percentile(per_request_s, 95):.4f}")
|
||||
print(f"{'Latency P99 (s):':<40} {np.percentile(per_request_s, 99):.4f}")
|
||||
print(f"{'Latency Min (s):':<40} {per_request_s.min():.4f}")
|
||||
print(f"{'Latency Max (s):':<40} {per_request_s.max():.4f}")
|
||||
elif latencies:
|
||||
per_request = np.diff([0.0] + list(latencies))
|
||||
print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
|
||||
print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
|
||||
print(f"{'Latency Mean (s) [wall-clock]:':<40} {per_request.mean():.4f}")
|
||||
print(f"{'Latency Median (s) [wall-clock]:':<40} {np.median(per_request):.4f}")
|
||||
print(f"{'Latency P95 (s) [wall-clock]:':<40} {np.percentile(per_request, 95):.4f}")
|
||||
print(f"{'Latency P99 (s) [wall-clock]:':<40} {np.percentile(per_request, 99):.4f}")
|
||||
print(f"{'Latency Min (s) [wall-clock]:':<40} {per_request.min():.4f}")
|
||||
print(f"{'Latency Max (s) [wall-clock]:':<40} {per_request.max():.4f}")
|
||||
|
||||
if per_request_actual:
|
||||
print("-" * 50)
|
||||
print("Pipeline Timings Mean:")
|
||||
for key in _TIMING_ORDER:
|
||||
vals = [d.get(key, 0.0) for d in per_request_actual]
|
||||
if any(v != 0 for v in vals):
|
||||
unit = "ms" if key.endswith("_ms") else "s"
|
||||
print(f" {key + ':':<38} {np.mean(vals):.4f} ({unit})")
|
||||
# Show any extra keys not in the ordered list
|
||||
ordered_set = set(_TIMING_ORDER)
|
||||
extra_keys = sorted(k for k in per_request_actual[0].keys() if k not in ordered_set)
|
||||
for key in extra_keys:
|
||||
vals = [d.get(key, 0.0) for d in per_request_actual]
|
||||
if any(v != 0 for v in vals):
|
||||
unit = "ms" if key.endswith("_ms") else "s"
|
||||
print(f" {key + ':':<38} {np.mean(vals):.4f} ({unit})")
|
||||
|
||||
print(f"\n{'Output dir:':<40} {args.output_dir}")
|
||||
print("=" * 60)
|
||||
|
||||
# Metrics JSON
|
||||
metrics = {
|
||||
"backend": "vllm-omni",
|
||||
"mode": args.mode,
|
||||
"model": args.model_path,
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"init_time_s": init_time,
|
||||
"completed_requests": success,
|
||||
"failed_requests": failed,
|
||||
"total_gen_time_s": total_gen_time,
|
||||
"throughput_qps": success / total_gen_time if total_gen_time > 0 else 0,
|
||||
}
|
||||
if per_request_e2e_ms:
|
||||
per_request_s = np.array(per_request_e2e_ms) / 1000.0
|
||||
metrics["latency_mean"] = float(per_request_s.mean())
|
||||
metrics["latency_median"] = float(np.median(per_request_s))
|
||||
metrics["latency_p95"] = float(np.percentile(per_request_s, 95))
|
||||
metrics["latency_p99"] = float(np.percentile(per_request_s, 99))
|
||||
elif latencies:
|
||||
per_request = np.diff([0.0] + list(latencies))
|
||||
metrics["latency_mean"] = float(per_request.mean())
|
||||
metrics["latency_median"] = float(np.median(per_request))
|
||||
metrics["latency_p95"] = float(np.percentile(per_request, 95))
|
||||
metrics["latency_p99"] = float(np.percentile(per_request, 99))
|
||||
else:
|
||||
metrics["latency_mean"] = 0
|
||||
metrics["latency_median"] = 0
|
||||
metrics["latency_p95"] = 0
|
||||
metrics["latency_p99"] = 0
|
||||
if per_request_actual:
|
||||
all_keys = list(_TIMING_ORDER) + sorted(k for k in per_request_actual[0].keys() if k not in set(_TIMING_ORDER))
|
||||
stage_metrics = {}
|
||||
for key in all_keys:
|
||||
vals = [d.get(key, 0.0) for d in per_request_actual]
|
||||
stage_metrics[key] = {
|
||||
"mean": float(np.mean(vals)),
|
||||
"median": float(np.median(vals)),
|
||||
"p95": float(np.percentile(vals, 95)),
|
||||
}
|
||||
metrics["stage_durations"] = stage_metrics
|
||||
if args.output_file:
|
||||
with open(args.output_file, "w") as f:
|
||||
json.dump(metrics, f, indent=2)
|
||||
print(f"Metrics saved to {args.output_file}")
|
||||
|
||||
omni.close()
|
||||
print("Done!")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="GLM-Image vLLM-Omni offline benchmark")
|
||||
parser.add_argument("--model-path", type=str, default="zai-org/GLM-Image")
|
||||
parser.add_argument("--deploy-config", type=str, default=None, help="Deploy config YAML")
|
||||
parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
|
||||
parser.add_argument("--dataset-path", type=str, default=None, help="Path to prompt.json")
|
||||
parser.add_argument("--num-prompts", type=int, default=10)
|
||||
parser.add_argument("--height", type=int, default=HEIGHT)
|
||||
parser.add_argument("--width", type=int, default=WIDTH)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=NUM_INFERENCE_STEPS)
|
||||
parser.add_argument("--guidance-scale", type=float, default=GUIDANCE_SCALE)
|
||||
parser.add_argument("--seed", type=int, default=SEED)
|
||||
parser.add_argument("--output-dir", type=str, default="benchmarks/glm_image/vllm-omni/outputs")
|
||||
parser.add_argument("--output-file", type=str, default=None, help="JSON file for metrics")
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=600)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler for stage-level timing",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-ar-profiler",
|
||||
action="store_true",
|
||||
help="Enable AR stage profiler to include AR timing in stage_durations",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
help="Enable detailed per-request pipeline stats logging",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
benchmark(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,86 +0,0 @@
|
||||
# Benchmarks Guide
|
||||
|
||||
This README explains how to (1) prepare benchmark datasets and (2) run the provided Qwen3-Omni benchmarks.
|
||||
|
||||
## 1) Prepare the dataset (SeedTTS top100)
|
||||
|
||||
```bash
|
||||
cd benchmarks/build_dataset
|
||||
pip install gdown
|
||||
|
||||
# Download SeedTTS test set from Google Drive
|
||||
gdown 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
|
||||
|
||||
# Extract
|
||||
tar -xf seedtts_testset.tar
|
||||
|
||||
# Copy metadata and extract top-100 prompts
|
||||
cp seedtts_testset/en/meta.lst meta.lst
|
||||
python extract_prompts.py -i meta.lst -o top100.txt -n 100
|
||||
|
||||
# (Optional) clean up to save space
|
||||
rm -rf seedtts_testset seedtts_testset.tar meta.lst
|
||||
```
|
||||
|
||||
Artifacts:
|
||||
- `benchmarks/build_dataset/top100.txt` — 100 text prompts (one per line).
|
||||
|
||||
## 2) Run benchmarks
|
||||
|
||||
All commands assume repo root (`vllm-omni`).
|
||||
|
||||
### A. Transformers benchmark (offline, HF Transformers)
|
||||
|
||||
```
|
||||
bash benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh
|
||||
```
|
||||
|
||||
What it does:
|
||||
- Runs `qwen3_omni_moe_transformers.py` over `top100.txt` with `--num-prompts 100`.
|
||||
- Outputs to `benchmarks/qwen3-omni/transformers/benchmark_results/`:
|
||||
- `perf_stats.json` — aggregated & per-prompt TPS/latency (thinker/talker/code2wav/overall).
|
||||
- `results.json` — per-prompt outputs and audio paths.
|
||||
- `audio/` — ~100 generated `.wav` files.
|
||||
|
||||
Key checks:
|
||||
- `overall_tps` and `*_tps_avg` should be non-zero and reasonably stable.
|
||||
- Investigate any 0/NaN or unusually low TPS / long-tail latency.
|
||||
|
||||
### B. vLLM Omni end-to-end benchmark (pipeline)
|
||||
|
||||
```
|
||||
bash benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh
|
||||
```
|
||||
|
||||
What it does:
|
||||
- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
|
||||
- Uses `benchmarks/build_dataset/top100.txt` and writes to:
|
||||
- Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
|
||||
- `omni_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
|
||||
- `omni_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
|
||||
- `omni_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
|
||||
- Outputs: `benchmarks/qwen3-omni/vllm_omni/outputs/` — ~100 text and `.wav` files.
|
||||
|
||||
Key checks:
|
||||
- Overall stats: end-to-end latency/TPS should be reasonable.
|
||||
- Orchestrator stats: per-stage latency should be stable; investigate long tails.
|
||||
- Stage logs: ensure no errors and no unusually slow stages.
|
||||
|
||||
|
||||
## Performance snapshot
|
||||
|
||||
The chart below summarizes our measured Qwen3-Omni MoE end-to-end benchmark, comparing vLLM-Omni against HF Transformers. It shows the overall throughput advantage for vLLM-Omni. These are actual experiment results—please refer to this performance when evaluating or reproducing the benchmark.
|
||||
|
||||

|
||||
|
||||
## Directory layout
|
||||
- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
|
||||
- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
|
||||
- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
|
||||
- `benchmarks/<model>/vllm-omni-vs-hf.png` — current performance snapshot (overall throughput comparison).
|
||||
- `benchmarks/<model>/transformers/` — HF Transformers benchmarks (offline reference).
|
||||
|
||||
## Troubleshooting
|
||||
- Make sure GPU/driver/FlashAttention2 requirements are met for the chosen model.
|
||||
- If downloads fail, confirm network access to Google Drive (`gdown`) and Hugging Face.
|
||||
- If audio files are missing, check for errors in stage logs or model generation.***
|
||||
@@ -1,26 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Qwen3-Omni Transformers Benchmark Evaluation Script
|
||||
# This script must be run from the vllm-omni root directory
|
||||
|
||||
# Get the directory where this script is located
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
# Navigate to vllm-omni root directory (4 levels up from script location)
|
||||
VLLM_OMNI_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
|
||||
cd "$VLLM_OMNI_ROOT" || { echo "Error: Failed to navigate to vllm-omni directory"; exit 1; }
|
||||
|
||||
echo "Working directory: $(pwd)"
|
||||
# Verify we're in the correct directory and run benchmark
|
||||
if [[ ! -f "benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py" ]]; then
|
||||
echo "Error: Not in vllm-omni root directory. Please run from vllm-omni folder."
|
||||
else
|
||||
cd benchmarks/qwen3-omni/transformers
|
||||
|
||||
python qwen3_omni_moe_transformers.py --prompts-file ../../build_dataset/top100.txt --num-prompts 100
|
||||
|
||||
echo "Logs and outputs are saved to $(pwd)/benchmark_results:"
|
||||
echo " - perf_stats.json Aggregated/per-prompt TPS and latency (thinker/talker/code2wav/overall)"
|
||||
echo " - results.json Per-prompt outputs and audio paths"
|
||||
echo " - audio/ Generated wav files, there should be 100 wav file generated"
|
||||
echo "Key checks: overall_tps and *_tps_avg should be non-zero and stable; investigate 0/NaN or unusually low TPS/long-tail latency."
|
||||
fi
|
||||
@@ -1,265 +0,0 @@
|
||||
import time
|
||||
|
||||
import torch
|
||||
from transformers import Qwen3OmniMoeForConditionalGeneration
|
||||
|
||||
|
||||
class Qwen3OmniMoeForConditionalGenerationWithLogging(Qwen3OmniMoeForConditionalGeneration):
|
||||
@torch.no_grad()
|
||||
def generate(
|
||||
self,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
speaker: str = "Ethan",
|
||||
use_audio_in_video: bool = False,
|
||||
return_audio: bool | None = None,
|
||||
thinker_max_new_tokens: int = 1024,
|
||||
thinker_eos_token_id: int = 151645,
|
||||
talker_max_new_tokens: int = 4096,
|
||||
talker_do_sample: bool = True,
|
||||
talker_top_k: int = 50,
|
||||
talker_top_p: float = 1.0,
|
||||
talker_temperature: float = 0.9,
|
||||
talker_repetition_penalty: float = 1.05,
|
||||
**kwargs,
|
||||
):
|
||||
total_t0 = time.time()
|
||||
perf_stats = {
|
||||
"thinker_tokens": 0,
|
||||
"thinker_time_s": 0.0,
|
||||
"thinker_tps": 0.0,
|
||||
"talker_tokens": 0,
|
||||
"talker_time_s": 0.0,
|
||||
"talker_tps": 0.0,
|
||||
"code2wav_tokens": 0,
|
||||
"code2wav_time_s": 0.0,
|
||||
"code2wav_tps": 0.0,
|
||||
"total_tokens": 0,
|
||||
"total_time_s": 0.0,
|
||||
"total_tps": 0.0,
|
||||
}
|
||||
if return_audio and not self.has_talker:
|
||||
raise ValueError(
|
||||
"Cannot use talker when talker module not initialized. "
|
||||
"Use `enable_talker` method or set enable_talker in config "
|
||||
"to enable talker."
|
||||
)
|
||||
if return_audio is None:
|
||||
return_audio = self.has_talker
|
||||
|
||||
shared_kwargs = {"use_audio_in_video": use_audio_in_video}
|
||||
thinker_kwargs = {
|
||||
"max_new_tokens": thinker_max_new_tokens,
|
||||
"eos_token_id": thinker_eos_token_id,
|
||||
}
|
||||
|
||||
talker_kwargs = {}
|
||||
token2wav_kwargs = {}
|
||||
if return_audio:
|
||||
speaker_id = self.config.talker_config.speaker_id.get(speaker.lower())
|
||||
if speaker_id is None:
|
||||
raise NotImplementedError(f"Speaker {speaker} not implemented")
|
||||
if input_ids.shape[0] != 1:
|
||||
raise NotImplementedError("Qwen3-Omni currently does not support batched inference with audio output")
|
||||
talker_suppressed_tokens = [
|
||||
i
|
||||
for i in range(
|
||||
self.config.talker_config.text_config.vocab_size - 1024,
|
||||
self.config.talker_config.text_config.vocab_size,
|
||||
)
|
||||
if i != self.config.talker_config.codec_eos_token_id
|
||||
] # Suppress additional special tokens, should not be predicted
|
||||
talker_kwargs = {
|
||||
"max_new_tokens": talker_max_new_tokens,
|
||||
"do_sample": talker_do_sample,
|
||||
"top_k": talker_top_k,
|
||||
"top_p": talker_top_p,
|
||||
"temperature": talker_temperature,
|
||||
"eos_token_id": self.config.talker_config.codec_eos_token_id,
|
||||
"repetition_penalty": talker_repetition_penalty,
|
||||
"suppress_tokens": talker_suppressed_tokens,
|
||||
"output_hidden_states": True,
|
||||
"return_dict_in_generate": True,
|
||||
}
|
||||
token2wav_kwargs = {}
|
||||
|
||||
for key, value in kwargs.items():
|
||||
if key.startswith("thinker_"):
|
||||
thinker_kwargs[key[len("thinker_") :]] = value
|
||||
elif key.startswith("talker_"):
|
||||
talker_kwargs[key[len("talker_") :]] = value
|
||||
elif key.startswith("token2wav_"):
|
||||
token2wav_kwargs[key[len("token2wav_") :]] = value
|
||||
# Process special input values
|
||||
elif key == "feature_attention_mask":
|
||||
thinker_kwargs[key] = value
|
||||
talker_kwargs["audio_feature_lengths"] = torch.sum(value, dim=1)
|
||||
elif key in ("input_features", "attention_mask"):
|
||||
thinker_kwargs[key] = value
|
||||
# Put other key to shared kwargs
|
||||
else:
|
||||
shared_kwargs[key] = value
|
||||
|
||||
# Merge kwargs
|
||||
for key, value in shared_kwargs.items():
|
||||
if key not in thinker_kwargs:
|
||||
thinker_kwargs[key] = value
|
||||
if key not in talker_kwargs and key in ["image_grid_thw", "video_grid_thw", "video_second_per_grid"]:
|
||||
talker_kwargs[key] = value
|
||||
if key not in token2wav_kwargs:
|
||||
token2wav_kwargs[key] = value
|
||||
|
||||
# 1. Generate from thinker module
|
||||
generate_audio = return_audio and self.has_talker
|
||||
if generate_audio:
|
||||
thinker_kwargs["output_hidden_states"] = True
|
||||
thinker_kwargs["return_dict_in_generate"] = True
|
||||
|
||||
t0 = time.time()
|
||||
thinker_result = self.thinker.generate(input_ids=input_ids, **thinker_kwargs)
|
||||
t1 = time.time()
|
||||
perf_stats["thinker_time_s"] = max(0.0, t1 - t0)
|
||||
try:
|
||||
prompt_len = int(input_ids.shape[1]) if input_ids is not None else 0
|
||||
total_len = int(thinker_result.sequences.shape[-1])
|
||||
thinker_out_len = max(0, total_len - prompt_len)
|
||||
except Exception:
|
||||
thinker_out_len = 0
|
||||
perf_stats["thinker_tokens"] = thinker_out_len
|
||||
perf_stats["thinker_tps"] = (
|
||||
(thinker_out_len / perf_stats["thinker_time_s"]) if perf_stats["thinker_time_s"] > 0 else 0.0
|
||||
)
|
||||
|
||||
if not generate_audio:
|
||||
perf_stats["total_tokens"] = perf_stats["thinker_tokens"]
|
||||
perf_stats["total_time_s"] = time.time() - total_t0
|
||||
perf_stats["total_tps"] = (
|
||||
(perf_stats["total_tokens"] / perf_stats["total_time_s"]) if perf_stats["total_time_s"] > 0 else 0.0
|
||||
)
|
||||
# attach stats to self
|
||||
setattr(self, "_perf_stats_last", perf_stats)
|
||||
if not hasattr(self, "_perf_stats_history"):
|
||||
setattr(self, "_perf_stats_history", [])
|
||||
self._perf_stats_history.append(perf_stats)
|
||||
return thinker_result, None
|
||||
|
||||
# 2. Prepare talker input
|
||||
thinker_embed = torch.cat([hidden_states[0] for hidden_states in thinker_result.hidden_states], dim=1).to(
|
||||
self.talker.device
|
||||
) # [1 t d]
|
||||
thinker_hidden = torch.cat(
|
||||
[
|
||||
hidden_states[self.config.talker_config.accept_hidden_layer]
|
||||
for hidden_states in thinker_result.hidden_states
|
||||
],
|
||||
dim=1,
|
||||
).to(self.talker.device) # [1 t d]
|
||||
|
||||
im_start_indexes = torch.cat(
|
||||
(
|
||||
torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(),
|
||||
torch.tensor([thinker_result.sequences.shape[-1]], device=input_ids.device, dtype=input_ids.dtype),
|
||||
),
|
||||
dim=-1,
|
||||
).to(self.talker.device) # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here.
|
||||
multimodal_mask = (
|
||||
(thinker_result.sequences == self.config.thinker_config.audio_token_id) |
|
||||
(thinker_result.sequences == self.config.thinker_config.image_token_id) |
|
||||
(thinker_result.sequences == self.config.thinker_config.video_token_id)
|
||||
).to(self.talker.device) # [1 t] # fmt: skip
|
||||
|
||||
talker_special_tokens = torch.tensor(
|
||||
[[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
|
||||
device=self.thinker.device,
|
||||
dtype=input_ids.dtype,
|
||||
)
|
||||
tts_bos_embed, tts_eos_embed, tts_pad_embed = (
|
||||
self.talker.text_projection(self.thinker.get_input_embeddings()(talker_special_tokens))
|
||||
.to(self.talker.device)
|
||||
.chunk(3, dim=1)
|
||||
) # 3 * [1 1 d]
|
||||
|
||||
talker_input_embeds = [] # [1 t d]
|
||||
talker_input_ids = []
|
||||
# For every chatml parts
|
||||
for i in range(len(im_start_indexes) - 1):
|
||||
im_start_index = im_start_indexes[i]
|
||||
segment_end_index = im_start_indexes[i + 1]
|
||||
role_token = input_ids[0][im_start_index + 1]
|
||||
# Talker should ignore thinker system prompt
|
||||
if role_token == self.config.system_token_id:
|
||||
continue
|
||||
# Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs
|
||||
elif role_token == self.config.user_token_id:
|
||||
talker_user_part = self._get_talker_user_parts(
|
||||
im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed
|
||||
)
|
||||
talker_input_embeds.append(talker_user_part)
|
||||
talker_input_ids.append(thinker_result.sequences[:, im_start_index:segment_end_index])
|
||||
# Take assistant output (for now)
|
||||
elif role_token == self.config.assistant_token_id and i == len(im_start_indexes) - 2:
|
||||
talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts(
|
||||
im_start_index,
|
||||
segment_end_index,
|
||||
speaker_id,
|
||||
thinker_embed,
|
||||
tts_pad_embed,
|
||||
tts_bos_embed,
|
||||
tts_eos_embed,
|
||||
)
|
||||
talker_input_embeds.append(talker_assistant_embeds)
|
||||
talker_input_ids.append(talker_assistant_ids)
|
||||
# History assistant output (ignore for now)
|
||||
elif role_token == self.config.assistant_token_id and i != len(im_start_indexes) - 2:
|
||||
continue
|
||||
else:
|
||||
raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)")
|
||||
talker_input_embed = torch.cat([embed.to(self.talker.device) for embed in talker_input_embeds], dim=1)
|
||||
talker_input_id = torch.cat([embed.to(self.talker.device) for embed in talker_input_ids], dim=1)
|
||||
t2 = time.time()
|
||||
talker_result = self.talker.generate(
|
||||
inputs_embeds=talker_input_embed,
|
||||
trailing_text_hidden=trailing_text_hidden,
|
||||
tts_pad_embed=tts_pad_embed,
|
||||
talker_input_ids=talker_input_id, # Not use input_ids to prevent repetition penalty out of bound
|
||||
**talker_kwargs,
|
||||
)
|
||||
t3 = time.time()
|
||||
perf_stats["talker_time_s"] = max(0.0, t3 - t2)
|
||||
talker_codes = (
|
||||
torch.stack([hid[-1] for hid in talker_result.hidden_states if hid[-1] is not None], dim=1)
|
||||
.transpose(1, 2)
|
||||
.to(self.code2wav.device)
|
||||
)
|
||||
try:
|
||||
# codes shape: (B, num_quantizers, T). We log T as token length.
|
||||
perf_stats["talker_tokens"] = int(talker_codes.shape[-1])
|
||||
except Exception:
|
||||
perf_stats["talker_tokens"] = 0
|
||||
perf_stats["talker_tps"] = (
|
||||
(perf_stats["talker_tokens"] / perf_stats["talker_time_s"]) if perf_stats["talker_time_s"] > 0 else 0.0
|
||||
)
|
||||
t4 = time.time()
|
||||
talker_wavs = self.code2wav.chunked_decode(talker_codes, chunk_size=300, left_context_size=25).float()
|
||||
t5 = time.time()
|
||||
perf_stats["code2wav_time_s"] = max(0.0, t5 - t4)
|
||||
perf_stats["code2wav_tokens"] = perf_stats["talker_tokens"] # same T, not times 16
|
||||
perf_stats["code2wav_tps"] = (
|
||||
(perf_stats["code2wav_tokens"] / perf_stats["code2wav_time_s"])
|
||||
if perf_stats["code2wav_time_s"] > 0
|
||||
else 0.0
|
||||
)
|
||||
perf_stats["total_tokens"] = perf_stats["thinker_tokens"] + perf_stats["talker_tokens"]
|
||||
perf_stats["total_time_s"] = time.time() - total_t0
|
||||
perf_stats["total_tps"] = (
|
||||
(perf_stats["total_tokens"] / perf_stats["total_time_s"]) if perf_stats["total_time_s"] > 0 else 0.0
|
||||
)
|
||||
setattr(self, "_perf_stats_last", perf_stats)
|
||||
if not hasattr(self, "_perf_stats_history"):
|
||||
setattr(self, "_perf_stats_history", [])
|
||||
self._perf_stats_history.append(perf_stats)
|
||||
return thinker_result, talker_wavs.float()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Qwen3OmniMoeForConditionalGenerationWithLogging",
|
||||
]
|
||||
@@ -1,275 +0,0 @@
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
|
||||
import soundfile as sf
|
||||
from qwen3_omni_moe_model import Qwen3OmniMoeForConditionalGenerationWithLogging
|
||||
from qwen_omni_utils import process_mm_info
|
||||
from tqdm import tqdm
|
||||
from transformers import Qwen3OmniMoeProcessor
|
||||
|
||||
MODEL_PATH = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
|
||||
# MODEL_PATH = "Qwen/Qwen3-Omni-30B-A3B-Thinking"
|
||||
|
||||
|
||||
def load_prompts(prompts_file: str) -> list[str]:
|
||||
"""Load prompts from a text file, one prompt per line."""
|
||||
prompts = []
|
||||
with open(prompts_file, encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line:
|
||||
prompts.append(line)
|
||||
return prompts
|
||||
|
||||
|
||||
def run_benchmark(
|
||||
model,
|
||||
processor,
|
||||
prompts: list[str],
|
||||
output_dir: str = "benchmark_results",
|
||||
speaker: str = "Ethan",
|
||||
use_audio_in_video: bool = True,
|
||||
):
|
||||
"""
|
||||
Run benchmark on a list of prompts and collect performance stats.
|
||||
|
||||
Args:
|
||||
model: The Qwen3OmniMoe model
|
||||
processor: The Qwen3OmniMoe processor
|
||||
prompts: List of text prompts to process
|
||||
output_dir: Directory to save results
|
||||
speaker: Speaker voice for audio output
|
||||
use_audio_in_video: Whether to use audio in video
|
||||
|
||||
Returns:
|
||||
tuple: (aggregated_stats, results, audio_outputs)
|
||||
- aggregated_stats: dict with aggregated performance statistics
|
||||
- results: list of dicts with per-prompt results
|
||||
- audio_outputs: list of audio tensors/arrays (or None if no audio)
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
audio_dir = os.path.join(output_dir, "audio")
|
||||
os.makedirs(audio_dir, exist_ok=True)
|
||||
|
||||
all_stats = []
|
||||
results = []
|
||||
audio_outputs = []
|
||||
|
||||
for idx, prompt in enumerate(tqdm(prompts, desc="Processing prompts")):
|
||||
conversation = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": prompt}],
|
||||
},
|
||||
]
|
||||
|
||||
# Preparation for inference
|
||||
text = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
|
||||
audios, images, videos = process_mm_info(conversation, use_audio_in_video=use_audio_in_video)
|
||||
inputs = processor(
|
||||
text=text,
|
||||
audio=audios,
|
||||
images=images,
|
||||
videos=videos,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
)
|
||||
inputs = inputs.to(model.device).to(model.dtype)
|
||||
|
||||
# Inference: Generation of the output text and audio
|
||||
text_ids, audio = model.generate(
|
||||
**inputs, speaker=speaker, thinker_return_dict_in_generate=True, use_audio_in_video=use_audio_in_video
|
||||
)
|
||||
|
||||
# Decode output text
|
||||
output_text = processor.batch_decode(
|
||||
text_ids.sequences[:, inputs["input_ids"].shape[1] :],
|
||||
skip_special_tokens=True,
|
||||
clean_up_tokenization_spaces=False,
|
||||
)[0]
|
||||
|
||||
# Collect performance stats
|
||||
perf_stats = None
|
||||
if hasattr(model, "_perf_stats_last"):
|
||||
perf_stats = model._perf_stats_last.copy()
|
||||
perf_stats["prompt_idx"] = idx
|
||||
perf_stats["prompt"] = prompt
|
||||
all_stats.append(perf_stats)
|
||||
|
||||
# Save audio and collect audio output
|
||||
audio_path = None
|
||||
audio_data = None
|
||||
if audio is not None:
|
||||
audio_data = audio.reshape(-1).detach().cpu().numpy()
|
||||
audio_path = os.path.join(audio_dir, f"output_{idx:04d}.wav")
|
||||
sf.write(
|
||||
audio_path,
|
||||
audio_data,
|
||||
samplerate=24000,
|
||||
)
|
||||
audio_outputs.append(audio_data)
|
||||
else:
|
||||
audio_outputs.append(None)
|
||||
|
||||
# Save result
|
||||
result = {
|
||||
"idx": idx,
|
||||
"prompt": prompt,
|
||||
"output": output_text,
|
||||
"audio_path": audio_path,
|
||||
"perf_stats": perf_stats,
|
||||
}
|
||||
results.append(result)
|
||||
|
||||
# Aggregate statistics
|
||||
aggregated_stats = aggregate_stats(all_stats)
|
||||
|
||||
# Save all results
|
||||
results_path = os.path.join(output_dir, "results.json")
|
||||
with open(results_path, "w", encoding="utf-8") as f:
|
||||
json.dump(results, f, ensure_ascii=False, indent=2)
|
||||
|
||||
# Save aggregated stats
|
||||
stats_path = os.path.join(output_dir, "perf_stats.json")
|
||||
with open(stats_path, "w", encoding="utf-8") as f:
|
||||
json.dump({"aggregated": aggregated_stats, "per_prompt": all_stats}, f, ensure_ascii=False, indent=2)
|
||||
|
||||
# Count saved audio files
|
||||
num_audio_saved = sum(1 for a in audio_outputs if a is not None)
|
||||
print(f"\nSaved {num_audio_saved} audio files to {audio_dir}/")
|
||||
|
||||
return aggregated_stats, results, audio_outputs
|
||||
|
||||
|
||||
def aggregate_stats(all_stats: list[dict]) -> dict:
|
||||
"""Aggregate performance statistics from multiple runs."""
|
||||
if not all_stats:
|
||||
return {}
|
||||
|
||||
keys = [
|
||||
"thinker_tokens",
|
||||
"thinker_time_s",
|
||||
"thinker_tps",
|
||||
"talker_tokens",
|
||||
"talker_time_s",
|
||||
"talker_tps",
|
||||
"code2wav_tokens",
|
||||
"code2wav_time_s",
|
||||
"code2wav_tps",
|
||||
"total_tokens",
|
||||
"total_time_s",
|
||||
"total_tps",
|
||||
]
|
||||
|
||||
aggregated = {
|
||||
"num_samples": len(all_stats),
|
||||
}
|
||||
|
||||
for key in keys:
|
||||
values = [s.get(key, 0) for s in all_stats if key in s]
|
||||
if values:
|
||||
aggregated[f"{key}_sum"] = sum(values)
|
||||
aggregated[f"{key}_avg"] = sum(values) / len(values)
|
||||
aggregated[f"{key}_min"] = min(values)
|
||||
aggregated[f"{key}_max"] = max(values)
|
||||
|
||||
# Calculate overall throughput
|
||||
total_tokens = aggregated.get("total_tokens_sum", 0)
|
||||
total_time = aggregated.get("total_time_s_sum", 0)
|
||||
if total_time > 0:
|
||||
aggregated["overall_tps"] = total_tokens / total_time
|
||||
|
||||
return aggregated
|
||||
|
||||
|
||||
def print_stats(stats: dict):
|
||||
"""Print performance statistics in a formatted way."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Performance Statistics Summary")
|
||||
print("=" * 60)
|
||||
|
||||
print(f"\nNumber of samples: {stats.get('num_samples', 0)}")
|
||||
|
||||
print("\n--- Thinker ---")
|
||||
print(f" Total tokens: {stats.get('thinker_tokens_sum', 0):.0f}")
|
||||
print(f" Total time: {stats.get('thinker_time_s_sum', 0):.2f}s")
|
||||
print(f" Avg TPS: {stats.get('thinker_tps_avg', 0):.2f}")
|
||||
print(f" Min TPS: {stats.get('thinker_tps_min', 0):.2f}")
|
||||
print(f" Max TPS: {stats.get('thinker_tps_max', 0):.2f}")
|
||||
|
||||
print("\n--- Talker ---")
|
||||
print(f" Total tokens: {stats.get('talker_tokens_sum', 0):.0f}")
|
||||
print(f" Total time: {stats.get('talker_time_s_sum', 0):.2f}s")
|
||||
print(f" Avg TPS: {stats.get('talker_tps_avg', 0):.2f}")
|
||||
print(f" Min TPS: {stats.get('talker_tps_min', 0):.2f}")
|
||||
print(f" Max TPS: {stats.get('talker_tps_max', 0):.2f}")
|
||||
|
||||
print("\n--- Code2Wav ---")
|
||||
print(f" Total tokens: {stats.get('code2wav_tokens_sum', 0):.0f}")
|
||||
print(f" Total time: {stats.get('code2wav_time_s_sum', 0):.2f}s")
|
||||
print(f" Avg TPS: {stats.get('code2wav_tps_avg', 0):.2f}")
|
||||
print(f" Min TPS: {stats.get('code2wav_tps_min', 0):.2f}")
|
||||
print(f" Max TPS: {stats.get('code2wav_tps_max', 0):.2f}")
|
||||
|
||||
print("\n--- Overall ---")
|
||||
print(f" Total tokens: {stats.get('total_tokens_sum', 0):.0f}")
|
||||
print(f" Total time: {stats.get('total_time_s_sum', 0):.2f}s")
|
||||
print(f" Overall TPS: {stats.get('overall_tps', 0):.2f}")
|
||||
print(f" Avg TPS: {stats.get('total_tps_avg', 0):.2f}")
|
||||
print(f" Min TPS: {stats.get('total_tps_min', 0):.2f}")
|
||||
print(f" Max TPS: {stats.get('total_tps_max', 0):.2f}")
|
||||
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-Omni Benchmark Script")
|
||||
parser.add_argument(
|
||||
"--prompts-file",
|
||||
type=str,
|
||||
default="benchmark/build_dataset/top100.txt",
|
||||
help="Path to the prompts file (one prompt per line)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir", type=str, default="benchmark_results", help="Directory to save benchmark results"
|
||||
)
|
||||
parser.add_argument("--model-path", type=str, default=MODEL_PATH, help="Path to the model")
|
||||
parser.add_argument("--speaker", type=str, default="Ethan", help="Speaker voice for audio output")
|
||||
parser.add_argument("--num-prompts", type=int, default=None, help="Number of prompts to process (default: all)")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Load model and processor
|
||||
print(f"Loading model from {args.model_path}...")
|
||||
model = Qwen3OmniMoeForConditionalGenerationWithLogging.from_pretrained(
|
||||
args.model_path,
|
||||
dtype="auto",
|
||||
device_map="auto",
|
||||
attn_implementation="flash_attention_2",
|
||||
)
|
||||
processor = Qwen3OmniMoeProcessor.from_pretrained(args.model_path)
|
||||
|
||||
# Benchmark mode
|
||||
print(f"Loading prompts from {args.prompts_file}...")
|
||||
prompts = load_prompts(args.prompts_file)
|
||||
|
||||
if args.num_prompts:
|
||||
prompts = prompts[: args.num_prompts]
|
||||
|
||||
print(f"Running benchmark on {len(prompts)} prompts...")
|
||||
|
||||
aggregated_stats, results, audio_outputs = run_benchmark(
|
||||
model=model,
|
||||
processor=processor,
|
||||
prompts=prompts,
|
||||
output_dir=args.output_dir,
|
||||
speaker=args.speaker,
|
||||
)
|
||||
|
||||
print_stats(aggregated_stats)
|
||||
print(f"\nResults saved to {args.output_dir}/")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 115 KiB |
@@ -1,37 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Qwen3-Omni Benchmark Evaluation Script
|
||||
# This script must be run from the vllm-omni root directory
|
||||
|
||||
# Get the directory where this script is located
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
# Navigate to vllm-omni root directory (4 levels up from script location)
|
||||
VLLM_OMNI_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
|
||||
cd "$VLLM_OMNI_ROOT" || { echo "Error: Failed to navigate to vllm-omni directory"; exit 1; }
|
||||
|
||||
echo "Working directory: $(pwd)"
|
||||
|
||||
# Verify we're in the correct directory and run benchmark
|
||||
if [[ ! -d "benchmarks/qwen3-omni/vllm_omni" ]]; then
|
||||
echo "Error: Not in vllm-omni root directory. Please run from vllm-omni folder."
|
||||
else
|
||||
log_dir=benchmarks/qwen3-omni/vllm_omni/logs
|
||||
outputs_dir=benchmarks/qwen3-omni/vllm_omni/outputs
|
||||
end2end_script_path=examples/offline_inference/qwen3_omni/end2end.py
|
||||
build_dataset_path=benchmarks/build_dataset/top100.txt
|
||||
|
||||
python $end2end_script_path --output-wav $outputs_dir \
|
||||
--query-type text \
|
||||
--txt-prompts $build_dataset_path \
|
||||
--log-stats \
|
||||
--log-dir $log_dir
|
||||
echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
|
||||
echo " - omni_pipeline_text run dir/base name"
|
||||
echo " - omni_pipeline_text.orchestrator.stats.jsonl orchestrator-stage latency stats"
|
||||
echo " - omni_pipeline_text.overall.stats.jsonl overall latency/TPS stats"
|
||||
echo " - omni_pipeline_text.stage0.log per-stage detailed logs"
|
||||
echo " - omni_pipeline_text.stage1.log"
|
||||
echo " - omni_pipeline_text.stage2.log"
|
||||
echo "Key checks: overall.stats.jsonl for end-to-end latency/TPS; orchestrator.stats.jsonl for stable per-stage latency; stage*.log for errors or long tails."
|
||||
echo " - outputs/ Generated txt and wav files, there should be 100 text and wav files generated respectively"
|
||||
fi
|
||||
@@ -1,107 +0,0 @@
|
||||
# Qwen3-TTS Benchmark
|
||||
|
||||
Benchmarks for Qwen3-TTS text-to-speech models, comparing vLLM-Omni streaming serving against HuggingFace Transformers offline inference.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
```bash
|
||||
pip install matplotlib aiohttp soundfile numpy tqdm
|
||||
pip install qwen_tts # for HF baseline
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
Run the full benchmark (vllm-omni + HF baseline) with a single command:
|
||||
|
||||
```bash
|
||||
cd benchmarks/qwen3-tts
|
||||
bash run_benchmark.sh
|
||||
```
|
||||
|
||||
Results (JSON + PNG plots) are saved to `results/`.
|
||||
|
||||
### Common options
|
||||
|
||||
```bash
|
||||
# Only vllm-omni (skip HF baseline)
|
||||
bash run_benchmark.sh --async-only
|
||||
|
||||
# Only HF baseline
|
||||
bash run_benchmark.sh --hf-only
|
||||
|
||||
# Use a different model (e.g. 1.7B)
|
||||
MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
|
||||
|
||||
# Use a Voice Clone model
|
||||
MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-Base TASK_TYPE=Base bash run_benchmark.sh --async-only
|
||||
|
||||
# Use batch size 16 for higher throughput
|
||||
BATCH_SIZE=16 bash run_benchmark.sh --async-only
|
||||
|
||||
# Custom GPU, prompt count, concurrency levels
|
||||
GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
|
||||
```
|
||||
|
||||
## Manual Steps
|
||||
|
||||
### 1) Start the vLLM-Omni server
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
|
||||
"Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
|
||||
--omni --host 127.0.0.1 --port 8000 \
|
||||
--deploy-config vllm_omni/deploy/qwen3_tts.yaml \
|
||||
--stage-overrides '{"0":{"max_num_seqs":1,"gpu_memory_utilization":0.3,"max_num_batched_tokens":512},"1":{"max_num_seqs":1,"gpu_memory_utilization":0.3,"max_num_batched_tokens":8192}}' \
|
||||
--trust-remote-code
|
||||
```
|
||||
|
||||
### 2) Run online serving benchmark
|
||||
|
||||
```bash
|
||||
python benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py \
|
||||
--port 8000 \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 4 10 \
|
||||
--config-name "async_chunk" \
|
||||
--result-dir results/
|
||||
```
|
||||
|
||||
### 3) Run HuggingFace baseline
|
||||
|
||||
```bash
|
||||
python benchmarks/qwen3-tts/transformers/bench_tts_hf.py \
|
||||
--model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
|
||||
--num-prompts 50 \
|
||||
--gpu-device 0 \
|
||||
--result-dir results/
|
||||
```
|
||||
|
||||
### 4) Generate comparison plots
|
||||
|
||||
```bash
|
||||
python benchmarks/qwen3-tts/plot_results.py \
|
||||
--results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
|
||||
--labels "vllm-omni" "hf_transformers" \
|
||||
--output results/comparison.png
|
||||
```
|
||||
|
||||
## Batch-size presets
|
||||
|
||||
The bench script loads the bundled production deploy (`vllm_omni/deploy/qwen3_tts.yaml`) and layers per-stage budgets on top via `--stage-overrides`, driven by the `BATCH_SIZE` env var. Each batch size picks compatible per-stage `max_num_seqs`, `max_num_batched_tokens`, and `gpu_memory_utilization` defaults:
|
||||
|
||||
| `BATCH_SIZE` | Description |
|
||||
|:--:|-------------|
|
||||
| `1` (default) | Single-request processing (lowest latency) |
|
||||
| `4` | Moderate-throughput concurrent processing |
|
||||
| `16` | High-throughput concurrent processing |
|
||||
|
||||
The 2-stage pipeline (Talker -> Code2Wav) runs with `async_chunk` streaming enabled via the prod deploy; the `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
|
||||
|
||||
The model is specified via the CLI `--model` flag (or `MODEL` env var), so the same bench script works for both the 0.6B and 1.7B model variants.
|
||||
|
||||
## Metrics
|
||||
|
||||
- **TTFP (Time to First Audio Packet)**: Time from request to first audio chunk (streaming latency)
|
||||
- **E2E (End-to-End Latency)**: Total time from request to complete audio response
|
||||
- **RTF (Real-Time Factor)**: E2E latency / audio duration. RTF < 1.0 means faster-than-real-time synthesis
|
||||
- **Throughput**: Total audio seconds generated per wall-clock second
|
||||
@@ -1,254 +0,0 @@
|
||||
"""Plot Qwen3-TTS benchmark results.
|
||||
|
||||
Generates comparison bar charts similar to the async_chunk design doc:
|
||||
- TTFP (Time-to-First-Packet) across concurrency levels
|
||||
- E2E latency across concurrency levels
|
||||
- RTF (Real-Time Factor) across concurrency levels
|
||||
|
||||
Usage:
|
||||
# Compare two configs (async_chunk vs no_async_chunk):
|
||||
python plot_results.py \
|
||||
--results results/bench_async_chunk_*.json results/bench_no_async_chunk_*.json \
|
||||
--labels "async_chunk" "no_async_chunk" \
|
||||
--output results/qwen3_tts_benchmark.png
|
||||
|
||||
# Single config:
|
||||
python plot_results.py \
|
||||
--results results/bench_async_chunk_*.json \
|
||||
--labels "async_chunk" \
|
||||
--output results/qwen3_tts_benchmark.png
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
|
||||
def load_results(result_files: list[str]) -> list[list[dict]]:
|
||||
"""Load benchmark results from JSON files."""
|
||||
all_results = []
|
||||
for f in result_files:
|
||||
with open(f) as fh:
|
||||
data = json.load(fh)
|
||||
all_results.append(data)
|
||||
return all_results
|
||||
|
||||
|
||||
def plot_comparison(
|
||||
all_results: list[list[dict]],
|
||||
labels: list[str],
|
||||
output_path: str,
|
||||
title_prefix: str = "Qwen3-TTS",
|
||||
):
|
||||
"""Generate comparison bar charts."""
|
||||
n_configs = len(all_results)
|
||||
|
||||
# Collect concurrency levels present in ALL configs (skip missing data)
|
||||
all_concurrencies = [set(r["concurrency"] for r in results) for results in all_results]
|
||||
concurrencies = sorted(set.union(*all_concurrencies))
|
||||
|
||||
# Build data arrays, using None for missing concurrency levels
|
||||
ttfp_data = {label: [] for label in labels}
|
||||
e2e_data = {label: [] for label in labels}
|
||||
rtf_data = {label: [] for label in labels}
|
||||
throughput_data = {label: [] for label in labels}
|
||||
|
||||
for results, label in zip(all_results, labels):
|
||||
conc_map = {r["concurrency"]: r for r in results}
|
||||
for c in concurrencies:
|
||||
r = conc_map.get(c)
|
||||
ttfp_data[label].append(r["mean_ttfp_ms"] if r else None)
|
||||
e2e_data[label].append(r["mean_e2e_ms"] if r else None)
|
||||
rtf_data[label].append(r["mean_rtf"] if r else None)
|
||||
throughput_data[label].append(r["audio_throughput"] if r else None)
|
||||
|
||||
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
|
||||
fig.suptitle(f"{title_prefix} Performance Benchmark", fontsize=16, fontweight="bold")
|
||||
|
||||
x = np.arange(len(concurrencies))
|
||||
width = 0.35 if n_configs == 2 else 0.5
|
||||
if n_configs > 1:
|
||||
offsets = np.linspace(-width / 2 * (n_configs - 1), width / 2 * (n_configs - 1), n_configs)
|
||||
else:
|
||||
offsets = [0]
|
||||
|
||||
colors = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107"]
|
||||
|
||||
def plot_metric(ax, data_dict, ylabel, title, fmt=".1f"):
|
||||
bars = []
|
||||
for i, (label, values) in enumerate(data_dict.items()):
|
||||
# Replace None with 0 for plotting, but track which are missing
|
||||
plot_values = [v if v is not None else 0 for v in values]
|
||||
color = colors[i % len(colors)]
|
||||
bar = ax.bar(x + offsets[i], plot_values, width, label=label, color=color, alpha=0.85)
|
||||
bars.append(bar)
|
||||
# Add value labels on bars (skip None/missing data)
|
||||
max_val = max((v for v in values if v is not None), default=1)
|
||||
for rect, val in zip(bar, values):
|
||||
if val is not None and val > 0:
|
||||
ax.text(
|
||||
rect.get_x() + rect.get_width() / 2,
|
||||
rect.get_height() + max_val * 0.02,
|
||||
f"{val:{fmt}}",
|
||||
ha="center",
|
||||
va="bottom",
|
||||
fontsize=9,
|
||||
fontweight="bold",
|
||||
)
|
||||
ax.set_xlabel("Concurrency", fontsize=12)
|
||||
ax.set_ylabel(ylabel, fontsize=12)
|
||||
ax.set_title(title, fontsize=13, fontweight="bold")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.legend(fontsize=10)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
plot_metric(axes[0, 0], ttfp_data, "TTFP (ms)", "Time to First Audio Packet (TTFP)")
|
||||
plot_metric(axes[0, 1], e2e_data, "E2E Latency (ms)", "End-to-End Latency (E2E)")
|
||||
plot_metric(axes[1, 0], rtf_data, "RTF", "Real-Time Factor (RTF)", fmt=".3f")
|
||||
plot_metric(axes[1, 1], throughput_data, "Audio-sec / Wall-sec", "Audio Throughput", fmt=".2f")
|
||||
|
||||
plt.tight_layout()
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def plot_single_summary(results: list[dict], label: str, output_path: str):
|
||||
"""Generate a single-config summary with percentile breakdown."""
|
||||
concurrencies = [r["concurrency"] for r in results]
|
||||
|
||||
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
|
||||
fig.suptitle(f"Qwen3-TTS Benchmark - {label}", fontsize=15, fontweight="bold")
|
||||
|
||||
# TTFP breakdown
|
||||
ax = axes[0]
|
||||
means = [r["mean_ttfp_ms"] for r in results]
|
||||
medians = [r["median_ttfp_ms"] for r in results]
|
||||
p90s = [r["p90_ttfp_ms"] for r in results]
|
||||
p99s = [r["p99_ttfp_ms"] for r in results]
|
||||
x = np.arange(len(concurrencies))
|
||||
w = 0.2
|
||||
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
|
||||
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
|
||||
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
|
||||
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_xlabel("Concurrency")
|
||||
ax.set_ylabel("TTFP (ms)")
|
||||
ax.set_title("Time to First Audio Packet")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
|
||||
# E2E breakdown
|
||||
ax = axes[1]
|
||||
means = [r["mean_e2e_ms"] for r in results]
|
||||
medians = [r["median_e2e_ms"] for r in results]
|
||||
p90s = [r["p90_e2e_ms"] for r in results]
|
||||
p99s = [r["p99_e2e_ms"] for r in results]
|
||||
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
|
||||
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
|
||||
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
|
||||
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_xlabel("Concurrency")
|
||||
ax.set_ylabel("E2E Latency (ms)")
|
||||
ax.set_title("End-to-End Latency")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
|
||||
# RTF
|
||||
ax = axes[2]
|
||||
means = [r["mean_rtf"] for r in results]
|
||||
medians = [r["median_rtf"] for r in results]
|
||||
ax.bar(x - 0.15, means, 0.3, label="mean", color="#2196F3")
|
||||
ax.bar(x + 0.15, medians, 0.3, label="median", color="#4CAF50")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_xlabel("Concurrency")
|
||||
ax.set_ylabel("RTF")
|
||||
ax.set_title("Real-Time Factor")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
|
||||
plt.tight_layout()
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def print_comparison_table(all_results: list[list[dict]], labels: list[str]):
|
||||
"""Print a markdown-formatted comparison table."""
|
||||
concurrencies = sorted(set(r["concurrency"] for r in all_results[0]))
|
||||
|
||||
print("\n## Benchmark Results\n")
|
||||
header = "| Metric | Concurrency |"
|
||||
sep = "| --- | --- |"
|
||||
for label in labels:
|
||||
header += f" {label} |"
|
||||
sep += " --- |"
|
||||
print(header)
|
||||
print(sep)
|
||||
|
||||
for metric, key, fmt in [
|
||||
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
|
||||
("E2E (ms)", "mean_e2e_ms", ".1f"),
|
||||
("RTF", "mean_rtf", ".3f"),
|
||||
("Throughput (audio-s/s)", "audio_throughput", ".2f"),
|
||||
]:
|
||||
for c in concurrencies:
|
||||
row = f"| {metric} | {c} |"
|
||||
for results in all_results:
|
||||
conc_map = {r["concurrency"]: r for r in results}
|
||||
val = conc_map.get(c, {}).get(key, 0)
|
||||
row += f" {val:{fmt}} |"
|
||||
print(row)
|
||||
|
||||
# Improvement calculation (only if 2 configs)
|
||||
if len(all_results) == 2:
|
||||
print(f"\n## Improvement ({labels[0]} vs {labels[1]})\n")
|
||||
print("| Metric | Concurrency | Improvement |")
|
||||
print("| --- | --- | --- |")
|
||||
for metric, key in [("TTFP", "mean_ttfp_ms"), ("E2E", "mean_e2e_ms"), ("RTF", "mean_rtf")]:
|
||||
for c in concurrencies:
|
||||
m0 = {r["concurrency"]: r for r in all_results[0]}
|
||||
m1 = {r["concurrency"]: r for r in all_results[1]}
|
||||
v0 = m0.get(c, {}).get(key, 0)
|
||||
v1 = m1.get(c, {}).get(key, 0)
|
||||
if v1 > 0:
|
||||
pct = (v1 - v0) / v1 * 100
|
||||
print(f"| {metric} | {c} | {pct:+.1f}% |")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Plot Qwen3-TTS benchmark results")
|
||||
parser.add_argument(
|
||||
"--results", type=str, nargs="+", required=True, help="Path(s) to result JSON files (one per config)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--labels", type=str, nargs="+", required=True, help="Labels for each config (must match --results count)"
|
||||
)
|
||||
parser.add_argument("--output", type=str, default="results/qwen3_tts_benchmark.png", help="Output image path")
|
||||
parser.add_argument("--title", type=str, default="Qwen3-TTS", help="Title prefix for the plot")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
assert len(args.results) == len(args.labels), "--results and --labels must have the same count"
|
||||
|
||||
all_results = load_results(args.results)
|
||||
print_comparison_table(all_results, args.labels)
|
||||
|
||||
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if len(all_results) == 1:
|
||||
plot_single_summary(all_results[0], args.labels[0], args.output)
|
||||
else:
|
||||
plot_comparison(all_results, args.labels, args.output, title_prefix=args.title)
|
||||
@@ -1,3 +0,0 @@
|
||||
# Benchmark results are machine-specific - do not commit
|
||||
*
|
||||
!.gitignore
|
||||
@@ -1,285 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Qwen3-TTS Benchmark Runner
|
||||
#
|
||||
# Compares vllm-omni streaming serving vs HuggingFace transformers offline inference.
|
||||
# Produces JSON results and comparison plots.
|
||||
#
|
||||
# Usage:
|
||||
# # Full comparison (vllm-omni + HF):
|
||||
# bash run_benchmark.sh
|
||||
#
|
||||
# # Only vllm-omni async_chunk config:
|
||||
# bash run_benchmark.sh --async-only
|
||||
#
|
||||
# # Only HuggingFace baseline:
|
||||
# bash run_benchmark.sh --hf-only
|
||||
#
|
||||
# # vllm-omni only (skip HF):
|
||||
# bash run_benchmark.sh --skip-hf
|
||||
#
|
||||
# # Custom settings:
|
||||
# GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
|
||||
#
|
||||
# # Use 1.7B model:
|
||||
# MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
|
||||
#
|
||||
# # Use Voice Clone model
|
||||
# MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-Base TASK_TYPE=Base bash run_benchmark.sh --async-only
|
||||
#
|
||||
# # Use batch_size=4:
|
||||
# BATCH_SIZE=4 bash run_benchmark.sh --async-only
|
||||
#
|
||||
# Environment variables:
|
||||
# GPU_DEVICE - GPU index to use (default: 0)
|
||||
# NUM_PROMPTS - Number of prompts per concurrency level (default: 50)
|
||||
# CONCURRENCY - Space-separated concurrency levels (default: "1 4 10")
|
||||
# MODEL - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
|
||||
# PORT - Server port (default: 8000)
|
||||
# BATCH_SIZE - Per-stage ``max_num_seqs`` for both talker and code2wav (default: 1)
|
||||
# GPU_MEM_TALKER - gpu_memory_utilization for talker stage (default: 0.3 at bs=1, else 0.2)
|
||||
# GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.3 at bs=1, else 0.2)
|
||||
# TASK_TYPE - Task type: CustomVoice, VoiceDesign, Base (default: CustomVoice)
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
|
||||
|
||||
# Defaults
|
||||
GPU_DEVICE="${GPU_DEVICE:-0}"
|
||||
NUM_PROMPTS="${NUM_PROMPTS:-50}"
|
||||
CONCURRENCY="${CONCURRENCY:-1 4 10}"
|
||||
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
|
||||
PORT="${PORT:-8000}"
|
||||
BATCH_SIZE="${BATCH_SIZE:-1}"
|
||||
DEFAULT_MEM=$([ "${BATCH_SIZE}" = "1" ] && echo "0.3" || echo "0.2")
|
||||
GPU_MEM_TALKER="${GPU_MEM_TALKER:-${DEFAULT_MEM}}"
|
||||
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-${DEFAULT_MEM}}"
|
||||
NUM_WARMUPS="${NUM_WARMUPS:-3}"
|
||||
DEPLOY_CONFIG="vllm_omni/deploy/qwen3_tts.yaml"
|
||||
RESULT_DIR="${SCRIPT_DIR}/results"
|
||||
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
||||
TASK_TYPE="${TASK_TYPE:-CustomVoice}"
|
||||
|
||||
# Build --stage-overrides JSON from BATCH_SIZE + GPU_MEM_*.
|
||||
STAGE_OVERRIDES=$(
|
||||
BATCH_SIZE="${BATCH_SIZE}" \
|
||||
GPU_MEM_TALKER="${GPU_MEM_TALKER}" \
|
||||
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV}" \
|
||||
python - <<'PYEOF'
|
||||
import json, os
|
||||
bs = int(os.environ["BATCH_SIZE"])
|
||||
mem_t = float(os.environ["GPU_MEM_TALKER"])
|
||||
mem_c = float(os.environ["GPU_MEM_CODE2WAV"])
|
||||
# Prefill budget grows with batch size on both stages.
|
||||
talker_batched = 512 if bs <= 4 else 4096
|
||||
code2wav_batched = 8192 if bs <= 4 else 32768
|
||||
print(json.dumps({
|
||||
"0": {"max_num_seqs": bs, "gpu_memory_utilization": mem_t, "max_num_batched_tokens": talker_batched},
|
||||
"1": {"max_num_seqs": bs, "gpu_memory_utilization": mem_c, "max_num_batched_tokens": code2wav_batched},
|
||||
}))
|
||||
PYEOF
|
||||
)
|
||||
|
||||
# Parse args
|
||||
RUN_ASYNC=true
|
||||
RUN_HF=true
|
||||
for arg in "$@"; do
|
||||
case "$arg" in
|
||||
--async-only) RUN_HF=false ;;
|
||||
--hf-only) RUN_ASYNC=false ;;
|
||||
--skip-hf) RUN_HF=false ;;
|
||||
esac
|
||||
done
|
||||
|
||||
mkdir -p "${RESULT_DIR}"
|
||||
|
||||
echo "============================================================"
|
||||
echo " Qwen3-TTS Benchmark"
|
||||
echo "============================================================"
|
||||
echo " GPU: ${GPU_DEVICE}"
|
||||
echo " Model: ${MODEL}"
|
||||
echo " Prompts: ${NUM_PROMPTS}"
|
||||
echo " Concurrency: ${CONCURRENCY}"
|
||||
echo " Port: ${PORT}"
|
||||
echo " Deploy config: ${DEPLOY_CONFIG}"
|
||||
echo " Batch size: ${BATCH_SIZE}"
|
||||
echo " GPU mem T/C: ${GPU_MEM_TALKER} / ${GPU_MEM_CODE2WAV}"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo " Task type: ${TASK_TYPE}"
|
||||
echo "============================================================"
|
||||
|
||||
# Start server and wait for it to be ready
|
||||
start_server() {
|
||||
local config_name="$1"
|
||||
local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log"
|
||||
|
||||
echo ""
|
||||
echo "Starting server with config: ${config_name}"
|
||||
echo " Deploy config: ${DEPLOY_CONFIG}"
|
||||
echo " Stage overrides: ${STAGE_OVERRIDES}"
|
||||
echo " Log file: ${log_file}"
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
CUDA_VISIBLE_DEVICES="${GPU_DEVICE}" \
|
||||
python -m vllm_omni.entrypoints.cli.main serve "${MODEL}" \
|
||||
--omni \
|
||||
--host 127.0.0.1 \
|
||||
--port "${PORT}" \
|
||||
--deploy-config "${DEPLOY_CONFIG}" \
|
||||
--stage-overrides "${STAGE_OVERRIDES}" \
|
||||
--stage-init-timeout 120 \
|
||||
--trust-remote-code \
|
||||
--disable-log-stats \
|
||||
> "${log_file}" 2>&1 &
|
||||
|
||||
SERVER_PID=$!
|
||||
echo " Server PID: ${SERVER_PID}"
|
||||
|
||||
# Wait for server to be ready
|
||||
echo " Waiting for server to be ready..."
|
||||
local max_wait=300
|
||||
local waited=0
|
||||
while [ ${waited} -lt ${max_wait} ]; do
|
||||
if curl -sf "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then
|
||||
echo " Server is ready! (waited ${waited}s)"
|
||||
return 0
|
||||
fi
|
||||
# Check if process is still alive
|
||||
if ! kill -0 ${SERVER_PID} 2>/dev/null; then
|
||||
echo " ERROR: Server process died. Check log: ${log_file}"
|
||||
tail -20 "${log_file}"
|
||||
return 1
|
||||
fi
|
||||
sleep 2
|
||||
waited=$((waited + 2))
|
||||
done
|
||||
|
||||
echo " ERROR: Server did not start within ${max_wait}s. Check log: ${log_file}"
|
||||
kill ${SERVER_PID} 2>/dev/null || true
|
||||
return 1
|
||||
}
|
||||
|
||||
# Stop the server
|
||||
stop_server() {
|
||||
if [ -n "${SERVER_PID:-}" ]; then
|
||||
echo " Stopping server (PID: ${SERVER_PID})..."
|
||||
kill ${SERVER_PID} 2>/dev/null || true
|
||||
wait ${SERVER_PID} 2>/dev/null || true
|
||||
# Kill any remaining child processes on the port
|
||||
local pids
|
||||
pids=$(lsof -ti:${PORT} 2>/dev/null || true)
|
||||
if [ -n "${pids}" ]; then
|
||||
echo " Cleaning up remaining processes on port ${PORT}..."
|
||||
echo "${pids}" | xargs kill -9 2>/dev/null || true
|
||||
fi
|
||||
echo " Server stopped."
|
||||
SERVER_PID=""
|
||||
fi
|
||||
}
|
||||
|
||||
# Cleanup on exit
|
||||
trap 'stop_server' EXIT
|
||||
|
||||
# Run benchmark for a given config
|
||||
run_bench() {
|
||||
local config_name="$1"
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmarking: ${config_name}"
|
||||
echo "============================================================"
|
||||
|
||||
start_server "${config_name}"
|
||||
|
||||
# Convert concurrency string to args
|
||||
local conc_args=""
|
||||
for c in ${CONCURRENCY}; do
|
||||
conc_args="${conc_args} ${c}"
|
||||
done
|
||||
|
||||
cd "${PROJECT_ROOT}"
|
||||
python "${SCRIPT_DIR}/vllm_omni/bench_tts_serve.py" \
|
||||
--host 127.0.0.1 \
|
||||
--port "${PORT}" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${conc_args} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--config-name "${config_name}" \
|
||||
--result-dir "${RESULT_DIR}" \
|
||||
--task-type "${TASK_TYPE}"
|
||||
|
||||
stop_server
|
||||
|
||||
# Allow GPU memory to settle
|
||||
sleep 5
|
||||
}
|
||||
|
||||
# Run vllm-omni benchmark
|
||||
if [ "${RUN_ASYNC}" = true ]; then
|
||||
run_bench "async_chunk"
|
||||
fi
|
||||
|
||||
# Run HuggingFace baseline benchmark
|
||||
if [ "${RUN_HF}" = true ]; then
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmarking: HuggingFace transformers (offline)"
|
||||
echo "============================================================"
|
||||
|
||||
cd "${PROJECT_ROOT}"
|
||||
python "${SCRIPT_DIR}/transformers/bench_tts_hf.py" \
|
||||
--model "${MODEL}" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--gpu-device "${GPU_DEVICE}" \
|
||||
--config-name "hf_transformers" \
|
||||
--result-dir "${RESULT_DIR}" \
|
||||
--task-type "${TASK_TYPE}"
|
||||
|
||||
# Allow GPU memory to settle
|
||||
sleep 5
|
||||
fi
|
||||
|
||||
# Plot results
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Generating plots..."
|
||||
echo "============================================================"
|
||||
|
||||
RESULT_FILES=""
|
||||
LABELS=""
|
||||
|
||||
if [ "${RUN_ASYNC}" = true ]; then
|
||||
ASYNC_FILE=$(ls -t "${RESULT_DIR}"/bench_async_chunk_*.json 2>/dev/null | head -1)
|
||||
if [ -n "${ASYNC_FILE}" ]; then
|
||||
RESULT_FILES="${ASYNC_FILE}"
|
||||
LABELS="async_chunk"
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ "${RUN_HF}" = true ]; then
|
||||
HF_FILE=$(ls -t "${RESULT_DIR}"/bench_hf_transformers_*.json 2>/dev/null | head -1)
|
||||
if [ -n "${HF_FILE}" ]; then
|
||||
if [ -n "${RESULT_FILES}" ]; then
|
||||
RESULT_FILES="${RESULT_FILES} ${HF_FILE}"
|
||||
LABELS="${LABELS} hf_transformers"
|
||||
else
|
||||
RESULT_FILES="${HF_FILE}"
|
||||
LABELS="hf_transformers"
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ -n "${RESULT_FILES}" ]; then
|
||||
python "${SCRIPT_DIR}/plot_results.py" \
|
||||
--results ${RESULT_FILES} \
|
||||
--labels ${LABELS} \
|
||||
--output "${RESULT_DIR}/qwen3_tts_benchmark_${TIMESTAMP}.png"
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmark complete!"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo "============================================================"
|
||||
@@ -1,301 +0,0 @@
|
||||
"""Benchmark Qwen3-TTS using HuggingFace transformers (qwen_tts library).
|
||||
|
||||
Measures E2E latency, RTF, and audio duration for offline (non-serving) inference.
|
||||
Results are saved in the same JSON format as bench_tts_serve.py for unified plotting.
|
||||
|
||||
Usage:
|
||||
python bench_tts_hf.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
|
||||
--num-prompts 50 \
|
||||
--num-warmups 3 \
|
||||
--gpu-device 0 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
|
||||
REF_AUDIO = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
|
||||
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
|
||||
INSTRUCT = "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
config_name: str = ""
|
||||
concurrency: int = 1 # always 1 for offline
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
# TTFP stats - not applicable for HF offline, set to E2E for compatibility
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
# E2E stats (ms)
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
# RTF stats
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
p99_rtf: float = 0.0
|
||||
# Audio stats
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0
|
||||
request_throughput: float = 0.0
|
||||
# Per-request details
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def generate_audio(model, prompt: str, args):
|
||||
if args.task_type == "Base":
|
||||
return model.generate_voice_clone(
|
||||
text=prompt,
|
||||
language=args.language,
|
||||
ref_audio=REF_AUDIO,
|
||||
ref_text=REF_TEXT,
|
||||
)
|
||||
|
||||
if args.task_type == "VoiceDesign":
|
||||
return model.generate_voice_design(
|
||||
text=prompt,
|
||||
language=args.language,
|
||||
instruct=INSTRUCT,
|
||||
)
|
||||
|
||||
return model.generate_custom_voice(
|
||||
text=prompt,
|
||||
language=args.language,
|
||||
speaker=args.voice,
|
||||
)
|
||||
|
||||
|
||||
def run_benchmark(args):
|
||||
from qwen_tts import Qwen3TTSModel
|
||||
|
||||
device = f"cuda:{args.gpu_device}"
|
||||
print(f"Loading model: {args.model} on {device}")
|
||||
model = Qwen3TTSModel.from_pretrained(
|
||||
args.model,
|
||||
device_map=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
print("Model loaded.")
|
||||
|
||||
# Build prompt list
|
||||
prompts = [PROMPTS[i % len(PROMPTS)] for i in range(args.num_prompts)]
|
||||
|
||||
# Warmup
|
||||
if args.num_warmups > 0:
|
||||
print(f"Warming up with {args.num_warmups} requests...")
|
||||
for i in range(args.num_warmups):
|
||||
p = PROMPTS[i % len(PROMPTS)]
|
||||
wavs, sr = generate_audio(model, p, args)
|
||||
# Sync GPU
|
||||
torch.cuda.synchronize(device)
|
||||
print("Warmup done.")
|
||||
|
||||
# Benchmark
|
||||
print(f"Running {args.num_prompts} requests sequentially...")
|
||||
e2e_times = []
|
||||
rtfs = []
|
||||
audio_durations = []
|
||||
per_request = []
|
||||
failed = 0
|
||||
|
||||
audio_dir = None
|
||||
if args.save_audio:
|
||||
audio_dir = Path(args.result_dir) / "audio_hf"
|
||||
audio_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
total_start = time.perf_counter()
|
||||
|
||||
for i, prompt in enumerate(prompts):
|
||||
try:
|
||||
torch.cuda.synchronize(device)
|
||||
st = time.perf_counter()
|
||||
|
||||
wavs, sr = generate_audio(model, prompt, args)
|
||||
|
||||
torch.cuda.synchronize(device)
|
||||
elapsed = time.perf_counter() - st
|
||||
|
||||
# Compute audio duration
|
||||
audio_samples = wavs[0]
|
||||
if isinstance(audio_samples, torch.Tensor):
|
||||
audio_samples = audio_samples.cpu().numpy()
|
||||
audio_dur = len(audio_samples) / sr
|
||||
|
||||
rtf = elapsed / audio_dur if audio_dur > 0 else 0.0
|
||||
|
||||
e2e_times.append(elapsed)
|
||||
rtfs.append(rtf)
|
||||
audio_durations.append(audio_dur)
|
||||
per_request.append(
|
||||
{
|
||||
"e2e_ms": elapsed * 1000,
|
||||
"ttfp_ms": elapsed * 1000, # no streaming, TTFP = E2E
|
||||
"rtf": rtf,
|
||||
"audio_duration_s": audio_dur,
|
||||
"prompt": prompt,
|
||||
}
|
||||
)
|
||||
|
||||
if audio_dir:
|
||||
sf.write(str(audio_dir / f"output_{i:04d}.wav"), audio_samples, sr)
|
||||
|
||||
if (i + 1) % 10 == 0 or i == 0:
|
||||
print(
|
||||
f" [{i + 1}/{args.num_prompts}] e2e={elapsed * 1000:.0f}ms rtf={rtf:.3f} audio={audio_dur:.2f}s"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f" [{i + 1}/{args.num_prompts}] FAILED: {e}")
|
||||
failed += 1
|
||||
|
||||
total_duration = time.perf_counter() - total_start
|
||||
completed = len(e2e_times)
|
||||
|
||||
# Compute stats
|
||||
result = BenchmarkResult(
|
||||
config_name=args.config_name,
|
||||
concurrency=1,
|
||||
num_prompts=args.num_prompts,
|
||||
completed=completed,
|
||||
failed=failed,
|
||||
duration_s=total_duration,
|
||||
)
|
||||
|
||||
if e2e_times:
|
||||
e2e_ms = [t * 1000 for t in e2e_times]
|
||||
|
||||
result.mean_e2e_ms = float(np.mean(e2e_ms))
|
||||
result.median_e2e_ms = float(np.median(e2e_ms))
|
||||
result.std_e2e_ms = float(np.std(e2e_ms))
|
||||
result.p90_e2e_ms = float(np.percentile(e2e_ms, 90))
|
||||
result.p95_e2e_ms = float(np.percentile(e2e_ms, 95))
|
||||
result.p99_e2e_ms = float(np.percentile(e2e_ms, 99))
|
||||
|
||||
# For HF offline, TTFP = E2E (no streaming)
|
||||
result.mean_ttfp_ms = result.mean_e2e_ms
|
||||
result.median_ttfp_ms = result.median_e2e_ms
|
||||
result.std_ttfp_ms = result.std_e2e_ms
|
||||
result.p90_ttfp_ms = result.p90_e2e_ms
|
||||
result.p95_ttfp_ms = result.p95_e2e_ms
|
||||
result.p99_ttfp_ms = result.p99_e2e_ms
|
||||
|
||||
result.mean_rtf = float(np.mean(rtfs))
|
||||
result.median_rtf = float(np.median(rtfs))
|
||||
result.std_rtf = float(np.std(rtfs))
|
||||
result.p99_rtf = float(np.percentile(rtfs, 99))
|
||||
|
||||
result.mean_audio_duration_s = float(np.mean(audio_durations))
|
||||
result.total_audio_duration_s = float(np.sum(audio_durations))
|
||||
result.audio_throughput = result.total_audio_duration_s / total_duration
|
||||
result.request_throughput = completed / total_duration
|
||||
result.per_request = per_request
|
||||
|
||||
# Print summary in standardized performance template
|
||||
W = 50
|
||||
print("")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Serving Benchmark Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Successful requests:':<40}{completed:<10}")
|
||||
print(f"{'Failed requests:':<40}{failed:<10}")
|
||||
print(f"{'Maximum request concurrency:':<40}{1:<10}")
|
||||
print(f"{'Benchmark duration (s):':<40}{total_duration:<10.2f}")
|
||||
print(f"{'Request throughput (req/s):':<40}{result.request_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'End-to-end Latency':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean E2EL (ms):':<40}{result.mean_e2e_ms:<10.2f}")
|
||||
print(f"{'Median E2EL (ms):':<40}{result.median_e2e_ms:<10.2f}")
|
||||
print(f"{'P99 E2EL (ms):':<40}{result.p99_e2e_ms:<10.2f}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Audio Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Total audio duration generated (s):':<40}{result.total_audio_duration_s:<10.2f}")
|
||||
print(f"{'Audio throughput (audio duration/s):':<40}{result.audio_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Time to First Packet':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_TTFP (ms):':<40}{result.mean_ttfp_ms:<10.2f}")
|
||||
print(f"{'Median AUDIO_TTFP (ms):':<40}{result.median_ttfp_ms:<10.2f}")
|
||||
print(f"{'P99 AUDIO_TTFP (ms):':<40}{result.p99_ttfp_ms:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Real Time Factor':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_RTF:':<40}{result.mean_rtf:<10.3f}")
|
||||
print(f"{'Median AUDIO_RTF:':<40}{result.median_rtf:<10.3f}")
|
||||
print(f"{'P99 AUDIO_RTF:':<40}{result.p99_rtf:<10.3f}")
|
||||
print(f"{'=' * W}")
|
||||
print("")
|
||||
|
||||
# Save results (as a list with single concurrency=1 entry, matching serve format)
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump([asdict(result)], f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS HuggingFace Benchmark")
|
||||
parser.add_argument(
|
||||
"--model", type=str, default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", help="HuggingFace model name or path"
|
||||
)
|
||||
parser.add_argument("--num-prompts", type=int, default=50)
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--gpu-device", type=int, default=0)
|
||||
parser.add_argument("--voice", type=str, default="Vivian")
|
||||
parser.add_argument("--language", type=str, default="English")
|
||||
parser.add_argument("--task-type", type=str, default="CustomVoice", choices=["CustomVoice", "VoiceDesign", "Base"])
|
||||
parser.add_argument(
|
||||
"--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
|
||||
)
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
parser.add_argument("--save-audio", action="store_true", help="Save generated audio files")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
run_benchmark(args)
|
||||
@@ -1,301 +0,0 @@
|
||||
"""Benchmark comparing async_chunk on vs off for Qwen3-TTS.
|
||||
|
||||
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across
|
||||
concurrency levels for both async_chunk modes. Saves results as JSON.
|
||||
|
||||
Usage:
|
||||
# Run against a server already serving with a given config:
|
||||
python bench_async_chunk.py \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--config-name async_chunk_on \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 10 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestResult:
|
||||
success: bool = False
|
||||
ttfp: float = 0.0
|
||||
e2e: float = 0.0
|
||||
audio_bytes: int = 0
|
||||
audio_duration: float = 0.0
|
||||
rtf: float = 0.0
|
||||
prompt: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
config_name: str = ""
|
||||
concurrency: int = 0
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0
|
||||
request_throughput: float = 0.0
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
|
||||
return num_bytes / sample_width / sample_rate
|
||||
|
||||
|
||||
async def send_tts_request(
|
||||
session: aiohttp.ClientSession,
|
||||
api_url: str,
|
||||
prompt: str,
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
stream: bool = True,
|
||||
pbar: tqdm | None = None,
|
||||
) -> RequestResult:
|
||||
payload = {
|
||||
"input": prompt,
|
||||
"voice": voice,
|
||||
"language": language,
|
||||
"stream": stream,
|
||||
"response_format": "pcm",
|
||||
}
|
||||
|
||||
result = RequestResult(prompt=prompt)
|
||||
st = time.perf_counter()
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status != 200:
|
||||
result.error = f"HTTP {response.status}: {await response.text()}"
|
||||
return result
|
||||
|
||||
first_chunk = True
|
||||
total_bytes = 0
|
||||
|
||||
async for chunk in response.content.iter_any():
|
||||
if first_chunk and len(chunk) > 0:
|
||||
result.ttfp = time.perf_counter() - st
|
||||
first_chunk = False
|
||||
total_bytes += len(chunk)
|
||||
|
||||
result.e2e = time.perf_counter() - st
|
||||
result.audio_bytes = total_bytes
|
||||
result.audio_duration = pcm_bytes_to_duration(total_bytes)
|
||||
if result.audio_duration > 0:
|
||||
result.rtf = result.e2e / result.audio_duration
|
||||
result.success = True
|
||||
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.e2e = time.perf_counter() - st
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return result
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
host: str,
|
||||
port: int,
|
||||
num_prompts: int,
|
||||
max_concurrency: int,
|
||||
num_warmups: int = 3,
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
stream: bool = True,
|
||||
) -> BenchmarkResult:
|
||||
api_url = f"http://{host}:{port}/v1/audio/speech"
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
|
||||
session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600))
|
||||
|
||||
if num_warmups > 0:
|
||||
print(f" Warming up with {num_warmups} requests...")
|
||||
warmup_tasks = [
|
||||
send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream)
|
||||
for i in range(num_warmups)
|
||||
]
|
||||
await asyncio.gather(*warmup_tasks)
|
||||
print(" Warmup done.")
|
||||
|
||||
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
|
||||
|
||||
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
|
||||
semaphore = asyncio.Semaphore(max_concurrency)
|
||||
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
|
||||
|
||||
async def limited_request(prompt):
|
||||
async with semaphore:
|
||||
return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
|
||||
results: list[RequestResult] = await asyncio.gather(*tasks)
|
||||
duration = time.perf_counter() - start_time
|
||||
pbar.close()
|
||||
|
||||
await session.close()
|
||||
|
||||
successful = [r for r in results if r.success]
|
||||
failed = [r for r in results if not r.success]
|
||||
|
||||
bench = BenchmarkResult(
|
||||
concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
completed=len(successful),
|
||||
failed=len(failed),
|
||||
duration_s=duration,
|
||||
)
|
||||
|
||||
if successful:
|
||||
ttfps = [r.ttfp * 1000 for r in successful]
|
||||
e2es = [r.e2e * 1000 for r in successful]
|
||||
rtfs = [r.rtf for r in successful]
|
||||
audio_durs = [r.audio_duration for r in successful]
|
||||
|
||||
bench.mean_ttfp_ms = float(np.mean(ttfps))
|
||||
bench.median_ttfp_ms = float(np.median(ttfps))
|
||||
bench.std_ttfp_ms = float(np.std(ttfps))
|
||||
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
|
||||
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
|
||||
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
|
||||
|
||||
bench.mean_e2e_ms = float(np.mean(e2es))
|
||||
bench.median_e2e_ms = float(np.median(e2es))
|
||||
bench.std_e2e_ms = float(np.std(e2es))
|
||||
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
|
||||
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
|
||||
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
|
||||
|
||||
bench.mean_rtf = float(np.mean(rtfs))
|
||||
bench.median_rtf = float(np.median(rtfs))
|
||||
bench.std_rtf = float(np.std(rtfs))
|
||||
|
||||
bench.mean_audio_duration_s = float(np.mean(audio_durs))
|
||||
bench.total_audio_duration_s = float(np.sum(audio_durs))
|
||||
bench.audio_throughput = bench.total_audio_duration_s / duration
|
||||
bench.request_throughput = len(successful) / duration
|
||||
|
||||
bench.per_request = [
|
||||
{
|
||||
"ttfp_ms": r.ttfp * 1000,
|
||||
"e2e_ms": r.e2e * 1000,
|
||||
"rtf": r.rtf,
|
||||
"audio_duration_s": r.audio_duration,
|
||||
"prompt": r.prompt,
|
||||
}
|
||||
for r in successful
|
||||
]
|
||||
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f" Concurrency: {max_concurrency} | Completed: {bench.completed} | Failed: {bench.failed}")
|
||||
print(f" Duration: {duration:.2f}s | Throughput: {bench.request_throughput:.2f} req/s")
|
||||
print(
|
||||
f" TTFP (ms): mean={bench.mean_ttfp_ms:.1f} median={bench.median_ttfp_ms:.1f}"
|
||||
f" p90={bench.p90_ttfp_ms:.1f} p99={bench.p99_ttfp_ms:.1f}"
|
||||
)
|
||||
print(
|
||||
f" E2E (ms): mean={bench.mean_e2e_ms:.1f} median={bench.median_e2e_ms:.1f}"
|
||||
f" p90={bench.p90_e2e_ms:.1f} p99={bench.p99_e2e_ms:.1f}"
|
||||
)
|
||||
print(f" RTF: mean={bench.mean_rtf:.3f} median={bench.median_rtf:.3f}")
|
||||
print(f" Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
if failed:
|
||||
for r in failed[:3]:
|
||||
print(f" [ERROR] {r.error[:200]}")
|
||||
|
||||
return bench
|
||||
|
||||
|
||||
async def main(args):
|
||||
all_results = []
|
||||
|
||||
for concurrency in args.max_concurrency:
|
||||
result = await run_benchmark(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
num_prompts=args.num_prompts,
|
||||
max_concurrency=concurrency,
|
||||
num_warmups=args.num_warmups,
|
||||
voice=args.voice,
|
||||
language=args.language,
|
||||
stream=args.stream,
|
||||
)
|
||||
result.config_name = args.config_name
|
||||
all_results.append(asdict(result))
|
||||
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump(all_results, f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client")
|
||||
parser.add_argument("--host", type=str, default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--num-prompts", type=int, default=50)
|
||||
parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10])
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--voice", type=str, default="vivian")
|
||||
parser.add_argument("--language", type=str, default="English")
|
||||
parser.add_argument("--stream", action="store_true", default=True)
|
||||
parser.add_argument("--no-stream", dest="stream", action="store_false")
|
||||
parser.add_argument("--config-name", type=str, default="async_chunk_on")
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
asyncio.run(main(args))
|
||||
@@ -1,371 +0,0 @@
|
||||
"""Benchmark client for Qwen3-TTS via /v1/audio/speech endpoint.
|
||||
|
||||
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF (Real-Time Factor)
|
||||
across configurable concurrency levels. Saves results as JSON for plotting.
|
||||
|
||||
Usage:
|
||||
python bench_tts_serve.py \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 4 10 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
REF_AUDIO = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
|
||||
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
|
||||
INSTRUCT = "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestResult:
|
||||
success: bool = False
|
||||
ttfp: float = 0.0 # Time to first audio packet (seconds)
|
||||
e2e: float = 0.0 # End-to-end latency (seconds)
|
||||
audio_bytes: int = 0 # Total audio bytes received
|
||||
audio_duration: float = 0.0 # Audio duration in seconds (estimated from PCM)
|
||||
rtf: float = 0.0 # Real-time factor = e2e / audio_duration
|
||||
prompt: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
config_name: str = ""
|
||||
concurrency: int = 0
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
# TTFP stats (ms)
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
# E2E stats (ms)
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
# RTF stats
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
p99_rtf: float = 0.0
|
||||
# Audio stats
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0 # audio_duration / wall_time
|
||||
request_throughput: float = 0.0 # requests / second
|
||||
# Per-request details
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
|
||||
"""Convert raw PCM byte count to duration in seconds."""
|
||||
num_samples = num_bytes / sample_width
|
||||
return num_samples / sample_rate
|
||||
|
||||
|
||||
def create_payload(
|
||||
prompt: str, task_type: str = "CustomVoice", voice: str = "vivian", language: str = "English"
|
||||
) -> dict:
|
||||
payload = {
|
||||
"input": prompt,
|
||||
"language": language,
|
||||
"stream": True,
|
||||
"response_format": "pcm",
|
||||
"task_type": task_type,
|
||||
}
|
||||
|
||||
if task_type == "Base":
|
||||
payload["ref_audio"] = REF_AUDIO
|
||||
payload["ref_text"] = REF_TEXT
|
||||
elif task_type == "CustomVoice":
|
||||
payload["voice"] = voice
|
||||
elif task_type == "VoiceDesign":
|
||||
payload["instructions"] = INSTRUCT
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
async def send_tts_request(
|
||||
session: aiohttp.ClientSession,
|
||||
api_url: str,
|
||||
prompt: str,
|
||||
task_type: str = "CustomVoice",
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
pbar: tqdm | None = None,
|
||||
) -> RequestResult:
|
||||
"""Send a streaming TTS request and measure latency metrics."""
|
||||
payload = create_payload(prompt, task_type, voice, language)
|
||||
|
||||
result = RequestResult(prompt=prompt)
|
||||
st = time.perf_counter()
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status != 200:
|
||||
result.error = f"HTTP {response.status}: {await response.text()}"
|
||||
result.success = False
|
||||
return result
|
||||
|
||||
first_chunk = True
|
||||
total_bytes = 0
|
||||
|
||||
async for chunk in response.content.iter_any():
|
||||
if first_chunk and len(chunk) > 0:
|
||||
result.ttfp = time.perf_counter() - st
|
||||
first_chunk = False
|
||||
total_bytes += len(chunk)
|
||||
|
||||
result.e2e = time.perf_counter() - st
|
||||
result.audio_bytes = total_bytes
|
||||
result.audio_duration = pcm_bytes_to_duration(total_bytes)
|
||||
|
||||
if result.audio_duration > 0:
|
||||
result.rtf = result.e2e / result.audio_duration
|
||||
result.success = True
|
||||
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.success = False
|
||||
result.e2e = time.perf_counter() - st
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return result
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
host: str,
|
||||
port: int,
|
||||
num_prompts: int,
|
||||
max_concurrency: int,
|
||||
num_warmups: int = 3,
|
||||
task_type: str = "CustomVoice",
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
) -> BenchmarkResult:
|
||||
"""Run benchmark at a given concurrency level."""
|
||||
api_url = f"http://{host}:{port}/v1/audio/speech"
|
||||
|
||||
connector = aiohttp.TCPConnector(
|
||||
limit=max_concurrency,
|
||||
limit_per_host=max_concurrency,
|
||||
keepalive_timeout=60,
|
||||
)
|
||||
session = aiohttp.ClientSession(
|
||||
connector=connector,
|
||||
timeout=aiohttp.ClientTimeout(total=600),
|
||||
)
|
||||
|
||||
# Warmup
|
||||
if num_warmups > 0:
|
||||
print(f" Warming up with {num_warmups} requests...")
|
||||
warmup_tasks = []
|
||||
for i in range(num_warmups):
|
||||
prompt = PROMPTS[i % len(PROMPTS)]
|
||||
warmup_tasks.append(send_tts_request(session, api_url, prompt, task_type, voice, language))
|
||||
await asyncio.gather(*warmup_tasks)
|
||||
print(" Warmup done.")
|
||||
|
||||
# Build request list
|
||||
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
|
||||
|
||||
# Run benchmark
|
||||
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
|
||||
semaphore = asyncio.Semaphore(max_concurrency)
|
||||
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
|
||||
|
||||
async def limited_request(prompt):
|
||||
async with semaphore:
|
||||
return await send_tts_request(session, api_url, prompt, task_type, voice, language, pbar)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
|
||||
results: list[RequestResult] = await asyncio.gather(*tasks)
|
||||
duration = time.perf_counter() - start_time
|
||||
pbar.close()
|
||||
|
||||
await session.close()
|
||||
|
||||
# Compute stats
|
||||
successful = [r for r in results if r.success]
|
||||
failed = [r for r in results if not r.success]
|
||||
|
||||
bench = BenchmarkResult(
|
||||
concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
completed=len(successful),
|
||||
failed=len(failed),
|
||||
duration_s=duration,
|
||||
)
|
||||
|
||||
if successful:
|
||||
ttfps = [r.ttfp * 1000 for r in successful] # convert to ms
|
||||
e2es = [r.e2e * 1000 for r in successful]
|
||||
rtfs = [r.rtf for r in successful]
|
||||
audio_durs = [r.audio_duration for r in successful]
|
||||
|
||||
bench.mean_ttfp_ms = float(np.mean(ttfps))
|
||||
bench.median_ttfp_ms = float(np.median(ttfps))
|
||||
bench.std_ttfp_ms = float(np.std(ttfps))
|
||||
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
|
||||
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
|
||||
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
|
||||
|
||||
bench.mean_e2e_ms = float(np.mean(e2es))
|
||||
bench.median_e2e_ms = float(np.median(e2es))
|
||||
bench.std_e2e_ms = float(np.std(e2es))
|
||||
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
|
||||
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
|
||||
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
|
||||
|
||||
bench.mean_rtf = float(np.mean(rtfs))
|
||||
bench.median_rtf = float(np.median(rtfs))
|
||||
bench.std_rtf = float(np.std(rtfs))
|
||||
bench.p99_rtf = float(np.percentile(rtfs, 99))
|
||||
|
||||
bench.mean_audio_duration_s = float(np.mean(audio_durs))
|
||||
bench.total_audio_duration_s = float(np.sum(audio_durs))
|
||||
bench.audio_throughput = bench.total_audio_duration_s / duration
|
||||
bench.request_throughput = len(successful) / duration
|
||||
|
||||
bench.per_request = [
|
||||
{
|
||||
"ttfp_ms": r.ttfp * 1000,
|
||||
"e2e_ms": r.e2e * 1000,
|
||||
"rtf": r.rtf,
|
||||
"audio_duration_s": r.audio_duration,
|
||||
"prompt": r.prompt,
|
||||
}
|
||||
for r in successful
|
||||
]
|
||||
|
||||
# Print summary in standardized performance template
|
||||
W = 50
|
||||
print("")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Serving Benchmark Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Successful requests:':<40}{bench.completed:<10}")
|
||||
print(f"{'Failed requests:':<40}{bench.failed:<10}")
|
||||
print(f"{'Maximum request concurrency:':<40}{max_concurrency:<10}")
|
||||
print(f"{'Benchmark duration (s):':<40}{duration:<10.2f}")
|
||||
print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'End-to-end Latency':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
|
||||
print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
|
||||
print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Audio Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
|
||||
print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Time to First Packet':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
|
||||
print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
|
||||
print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Real Time Factor':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
|
||||
print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
|
||||
print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
|
||||
print(f"{'=' * W}")
|
||||
print("")
|
||||
|
||||
if failed:
|
||||
for r in failed[:3]:
|
||||
print(f" [ERROR] {r.error[:200]}")
|
||||
|
||||
return bench
|
||||
|
||||
|
||||
async def main(args):
|
||||
all_results = []
|
||||
|
||||
for concurrency in args.max_concurrency:
|
||||
result = await run_benchmark(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
num_prompts=args.num_prompts,
|
||||
max_concurrency=concurrency,
|
||||
num_warmups=args.num_warmups,
|
||||
task_type=args.task_type,
|
||||
voice=args.voice,
|
||||
language=args.language,
|
||||
)
|
||||
result.config_name = args.config_name
|
||||
all_results.append(asdict(result))
|
||||
|
||||
# Save results
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump(all_results, f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
return all_results
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS Benchmark Client")
|
||||
parser.add_argument("--host", type=str, default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--num-prompts", type=int, default=50, help="Number of prompts per concurrency level")
|
||||
parser.add_argument( # noqa: E501
|
||||
"--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
|
||||
)
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--task-type", type=str, default="CustomVoice", choices=["CustomVoice", "VoiceDesign", "Base"])
|
||||
parser.add_argument("--voice", type=str, default="vivian")
|
||||
parser.add_argument("--language", type=str, default="English")
|
||||
parser.add_argument(
|
||||
"--config-name", type=str, default="async_chunk", help="Label for this config (used in filenames)"
|
||||
)
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
asyncio.run(main(args))
|
||||
@@ -1,249 +0,0 @@
|
||||
"""Plot TTFP comparison: async_chunk off vs on.
|
||||
|
||||
Generates a bar chart with improvement arrows, matching the Qwen3-Omni
|
||||
async_chunk benchmark figure style.
|
||||
|
||||
Usage:
|
||||
python plot_async_chunk.py \
|
||||
--off results/bench_async_chunk_off_*.json \
|
||||
--on results/bench_async_chunk_on_*.json \
|
||||
--output results/qwen3_tts_async_chunk_ttfp.png
|
||||
|
||||
# Also supports E2E and RTF metrics:
|
||||
python plot_async_chunk.py \
|
||||
--off results/bench_async_chunk_off_*.json \
|
||||
--on results/bench_async_chunk_on_*.json \
|
||||
--metric e2e \
|
||||
--output results/qwen3_tts_async_chunk_e2e.png
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
METRIC_CONFIG = {
|
||||
"ttfp": {
|
||||
"key": "mean_ttfp_ms",
|
||||
"ylabel": "TTFP (s)",
|
||||
"title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": True,
|
||||
},
|
||||
"e2e": {
|
||||
"key": "mean_e2e_ms",
|
||||
"ylabel": "E2E (s)",
|
||||
"title": "E2E Latency - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": True,
|
||||
},
|
||||
"rtf": {
|
||||
"key": "mean_rtf",
|
||||
"ylabel": "RTF",
|
||||
"title": "Real-Time Factor - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": False,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def load_results(path: str) -> list[dict]:
|
||||
with open(path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def plot_ttfp_comparison(
|
||||
off_results: list[dict],
|
||||
on_results: list[dict],
|
||||
metric: str,
|
||||
output_path: str,
|
||||
title_override: str | None = None,
|
||||
):
|
||||
cfg = METRIC_CONFIG[metric]
|
||||
key = cfg["key"]
|
||||
to_seconds = cfg["to_seconds"]
|
||||
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
off_vals = []
|
||||
on_vals = []
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c][key]
|
||||
v_on = on_map[c][key]
|
||||
if to_seconds:
|
||||
v_off /= 1000.0
|
||||
v_on /= 1000.0
|
||||
off_vals.append(v_off)
|
||||
on_vals.append(v_on)
|
||||
|
||||
fig, ax = plt.subplots(figsize=(8, 6))
|
||||
|
||||
x = np.arange(len(concurrencies))
|
||||
width = 0.3
|
||||
|
||||
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none")
|
||||
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
|
||||
|
||||
# Draw improvement arrows and labels
|
||||
for i in range(len(concurrencies)):
|
||||
v_off = off_vals[i]
|
||||
v_on = on_vals[i]
|
||||
if v_on > 0:
|
||||
improvement = v_off / v_on
|
||||
else:
|
||||
improvement = float("inf")
|
||||
|
||||
# Arrow from top of off-bar to top of on-bar
|
||||
arrow_start_x = x[i] - width / 2
|
||||
arrow_start_y = v_off * 0.95
|
||||
arrow_end_x = x[i] + width / 2
|
||||
arrow_end_y = v_on * 1.05
|
||||
|
||||
ax.annotate(
|
||||
"",
|
||||
xy=(arrow_end_x, arrow_end_y),
|
||||
xytext=(arrow_start_x, arrow_start_y),
|
||||
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
|
||||
)
|
||||
|
||||
# Improvement label
|
||||
label_x = (arrow_start_x + arrow_end_x) / 2
|
||||
label_y = arrow_start_y + (v_off - v_on) * 0.15
|
||||
ax.text(
|
||||
label_x,
|
||||
label_y,
|
||||
f"{improvement:.1f}x improvement",
|
||||
ha="center",
|
||||
va="bottom",
|
||||
fontsize=10,
|
||||
color="red",
|
||||
fontweight="bold",
|
||||
)
|
||||
|
||||
title = title_override or cfg["title"]
|
||||
ax.set_title(title, fontsize=13, fontweight="bold")
|
||||
ax.set_ylabel(cfg["ylabel"], fontsize=12)
|
||||
ax.set_xlabel("Max concurrency", fontsize=12)
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_yscale("log")
|
||||
ax.legend(loc="upper left", fontsize=11)
|
||||
ax.grid(axis="y", alpha=0.3, linestyle="--")
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
plt.tight_layout()
|
||||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str):
|
||||
"""Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons."""
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
|
||||
fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold")
|
||||
|
||||
for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]):
|
||||
cfg = METRIC_CONFIG[metric]
|
||||
key = cfg["key"]
|
||||
to_seconds = cfg["to_seconds"]
|
||||
|
||||
off_vals = []
|
||||
on_vals = []
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c][key]
|
||||
v_on = on_map[c][key]
|
||||
if to_seconds:
|
||||
v_off /= 1000.0
|
||||
v_on /= 1000.0
|
||||
off_vals.append(v_off)
|
||||
on_vals.append(v_on)
|
||||
|
||||
x = np.arange(len(concurrencies))
|
||||
width = 0.3
|
||||
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB")
|
||||
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
|
||||
|
||||
for i in range(len(concurrencies)):
|
||||
if on_vals[i] > 0:
|
||||
improvement = off_vals[i] / on_vals[i]
|
||||
ax.annotate(
|
||||
"",
|
||||
xy=(x[i] + width / 2, on_vals[i] * 1.05),
|
||||
xytext=(x[i] - width / 2, off_vals[i] * 0.95),
|
||||
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
|
||||
)
|
||||
label_y = off_vals[i] * 0.85
|
||||
ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold")
|
||||
|
||||
ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold")
|
||||
ax.set_ylabel(cfg["ylabel"], fontsize=11)
|
||||
ax.set_xlabel("Max concurrency", fontsize=11)
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
if metric != "rtf":
|
||||
ax.set_yscale("log")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3, linestyle="--")
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
plt.tight_layout()
|
||||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def print_table(off_results: list[dict], on_results: list[dict]):
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
print("\n## Benchmark Results: async_chunk off vs on\n")
|
||||
print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |")
|
||||
print("| --- | --- | --- | --- | --- |")
|
||||
|
||||
for name, key, fmt in [
|
||||
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
|
||||
("E2E (ms)", "mean_e2e_ms", ".1f"),
|
||||
("RTF", "mean_rtf", ".3f"),
|
||||
("Throughput", "audio_throughput", ".2f"),
|
||||
]:
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c].get(key, 0)
|
||||
v_on = on_map[c].get(key, 0)
|
||||
if v_on > 0 and key != "audio_throughput":
|
||||
ratio = f"{v_off / v_on:.1f}x"
|
||||
elif v_off > 0 and key == "audio_throughput":
|
||||
ratio = f"{v_on / v_off:.1f}x"
|
||||
else:
|
||||
ratio = "N/A"
|
||||
print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS")
|
||||
parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off")
|
||||
parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on")
|
||||
parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"])
|
||||
parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png")
|
||||
parser.add_argument("--title", type=str, default=None, help="Custom title override")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
off_results = load_results(args.off)
|
||||
on_results = load_results(args.on)
|
||||
|
||||
print_table(off_results, on_results)
|
||||
|
||||
if args.metric == "all":
|
||||
plot_all_metrics(off_results, on_results, args.output)
|
||||
else:
|
||||
plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title)
|
||||
@@ -1,169 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Qwen3-TTS async_chunk on vs off Benchmark
|
||||
#
|
||||
# Starts two servers (async_chunk on and off), benchmarks both,
|
||||
# and generates comparison plots.
|
||||
#
|
||||
# Usage:
|
||||
# bash run_async_chunk_benchmark.sh
|
||||
#
|
||||
# Environment variables:
|
||||
# GPU_DEVICE - GPU index (default: 0)
|
||||
# MODEL - Model path (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
|
||||
# NUM_PROMPTS - Prompts per concurrency level (default: 50)
|
||||
# CONCURRENCY - Space-separated concurrency levels (default: "1 10")
|
||||
# PORT_ON - Port for async_chunk on server (default: 8000)
|
||||
# PORT_OFF - Port for async_chunk off server (default: 8001)
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
|
||||
cd "$PROJECT_ROOT"
|
||||
|
||||
GPU_DEVICE="${GPU_DEVICE:-0}"
|
||||
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
|
||||
NUM_PROMPTS="${NUM_PROMPTS:-50}"
|
||||
CONCURRENCY="${CONCURRENCY:-1 10}"
|
||||
NUM_WARMUPS="${NUM_WARMUPS:-3}"
|
||||
PORT_ON="${PORT_ON:-8000}"
|
||||
PORT_OFF="${PORT_OFF:-8001}"
|
||||
RESULT_DIR="${SCRIPT_DIR}/results"
|
||||
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
||||
|
||||
# The bundled ``vllm_omni/deploy/qwen3_tts.yaml`` is auto-loaded by the model
|
||||
# registry; no ``--deploy-config`` flag needed on the default (ON) path.
|
||||
# async_chunk OFF is selected by the ``--no-async-chunk`` CLI flag —
|
||||
# the single ``qwen3_tts`` pipeline dispatches to the end-to-end codec
|
||||
# processor when ``deploy.async_chunk`` is false.
|
||||
|
||||
mkdir -p "${RESULT_DIR}"
|
||||
|
||||
echo "============================================================"
|
||||
echo " Qwen3-TTS async_chunk Benchmark"
|
||||
echo "============================================================"
|
||||
echo " GPU: ${GPU_DEVICE}"
|
||||
echo " Model: ${MODEL}"
|
||||
echo " Prompts: ${NUM_PROMPTS}"
|
||||
echo " Concurrency: ${CONCURRENCY}"
|
||||
echo " Port (on/off): ${PORT_ON} / ${PORT_OFF}"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo "============================================================"
|
||||
|
||||
cleanup() {
|
||||
echo "Cleaning up servers..."
|
||||
kill "$PID_ON" 2>/dev/null || true
|
||||
kill "$PID_OFF" 2>/dev/null || true
|
||||
wait "$PID_ON" 2>/dev/null || true
|
||||
wait "$PID_OFF" 2>/dev/null || true
|
||||
}
|
||||
trap cleanup EXIT
|
||||
|
||||
wait_for_server() {
|
||||
local port=$1
|
||||
local name=$2
|
||||
local max_wait=300
|
||||
local elapsed=0
|
||||
echo "Waiting for ${name} server on port ${port}..."
|
||||
while ! curl -s "http://localhost:${port}/health" >/dev/null 2>&1; do
|
||||
sleep 5
|
||||
elapsed=$((elapsed + 5))
|
||||
if [ $elapsed -ge $max_wait ]; then
|
||||
echo "ERROR: ${name} server failed to start within ${max_wait}s"
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
echo "${name} server ready (${elapsed}s)"
|
||||
}
|
||||
|
||||
# ---- Phase 1: Start async_chunk ON server ----
|
||||
echo ""
|
||||
echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..."
|
||||
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
|
||||
--host 0.0.0.0 --port "${PORT_ON}" \
|
||||
--trust-remote-code --enforce-eager --omni \
|
||||
> "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 &
|
||||
PID_ON=$!
|
||||
|
||||
wait_for_server "${PORT_ON}" "async_chunk_on"
|
||||
|
||||
echo "[Phase 1] Benchmarking async_chunk ON..."
|
||||
# shellcheck disable=SC2086
|
||||
python "${SCRIPT_DIR}/bench_async_chunk.py" \
|
||||
--host 127.0.0.1 --port "${PORT_ON}" \
|
||||
--config-name "async_chunk_on" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${CONCURRENCY} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--result-dir "${RESULT_DIR}"
|
||||
|
||||
echo "[Phase 1] Stopping async_chunk ON server..."
|
||||
kill "$PID_ON" 2>/dev/null || true
|
||||
wait "$PID_ON" 2>/dev/null || true
|
||||
sleep 5
|
||||
|
||||
# ---- Phase 2: Start async_chunk OFF server ----
|
||||
echo ""
|
||||
echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..."
|
||||
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
|
||||
--no-async-chunk \
|
||||
--host 0.0.0.0 --port "${PORT_OFF}" \
|
||||
--trust-remote-code --enforce-eager --omni \
|
||||
> "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 &
|
||||
PID_OFF=$!
|
||||
|
||||
wait_for_server "${PORT_OFF}" "async_chunk_off"
|
||||
|
||||
echo "[Phase 2] Benchmarking async_chunk OFF (non-streaming)..."
|
||||
# shellcheck disable=SC2086
|
||||
python "${SCRIPT_DIR}/bench_async_chunk.py" \
|
||||
--host 127.0.0.1 --port "${PORT_OFF}" \
|
||||
--config-name "async_chunk_off" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${CONCURRENCY} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--no-stream \
|
||||
--result-dir "${RESULT_DIR}"
|
||||
|
||||
echo "[Phase 2] Stopping async_chunk OFF server..."
|
||||
kill "$PID_OFF" 2>/dev/null || true
|
||||
wait "$PID_OFF" 2>/dev/null || true
|
||||
|
||||
# ---- Phase 3: Plot results ----
|
||||
echo ""
|
||||
echo "[Phase 3] Generating plots..."
|
||||
|
||||
# Find the latest result files
|
||||
RESULT_ON=$(ls -t "${RESULT_DIR}"/bench_async_chunk_on_*.json 2>/dev/null | head -1)
|
||||
RESULT_OFF=$(ls -t "${RESULT_DIR}"/bench_async_chunk_off_*.json 2>/dev/null | head -1)
|
||||
|
||||
if [ -z "$RESULT_ON" ] || [ -z "$RESULT_OFF" ]; then
|
||||
echo "ERROR: Could not find result files. Check logs in ${RESULT_DIR}/"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo " ON results: ${RESULT_ON}"
|
||||
echo " OFF results: ${RESULT_OFF}"
|
||||
|
||||
# TTFP comparison (main figure)
|
||||
python "${SCRIPT_DIR}/plot_async_chunk.py" \
|
||||
--off "${RESULT_OFF}" \
|
||||
--on "${RESULT_ON}" \
|
||||
--metric ttfp \
|
||||
--output "${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
|
||||
|
||||
# All metrics comparison
|
||||
python "${SCRIPT_DIR}/plot_async_chunk.py" \
|
||||
--off "${RESULT_OFF}" \
|
||||
--on "${RESULT_ON}" \
|
||||
--metric all \
|
||||
--output "${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmark complete!"
|
||||
echo " Results: ${RESULT_DIR}/"
|
||||
echo " Plots:"
|
||||
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
|
||||
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
|
||||
echo "============================================================"
|
||||
@@ -0,0 +1,227 @@
|
||||
# TTS Universal Benchmark
|
||||
|
||||
A model-agnostic serving benchmark for TTS models in vllm-omni. One CLI
|
||||
(`bench_tts.py`) + one YAML registry (`model_configs.yaml`) drive perf and
|
||||
quality runs for every registered checkpoint: **Qwen3-TTS** (Base / CustomVoice)
|
||||
and **VoxCPM2** today, more to come.
|
||||
|
||||
The same three task types — `voice_clone`, `default_voice`, `voice_design` —
|
||||
are wired into both the manual CLI and the DFX nightly CI matrix
|
||||
(`tests/dfx/perf/tests/test_tts.json`).
|
||||
|
||||
## Quick start
|
||||
|
||||
### 1. Start the server
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --port 8000
|
||||
```
|
||||
|
||||
The server auto-loads its Deploy YAML from `vllm_omni/deploy/qwen3_tts.yaml`
|
||||
(Pipeline + Deploy schema introduced in #2383). No `--stage-configs-path` or
|
||||
`--deploy-config` flag is needed for any registered model.
|
||||
|
||||
### 2. Run the benchmark (`vllm bench serve --omni`)
|
||||
|
||||
The primary, directly-controllable path. Copy-paste one of these and tweak
|
||||
any bench flag (sampling params, endpoint, extra body, warmups, etc.):
|
||||
|
||||
#### voice_clone (Qwen3-TTS-Base, seed-tts dataset)
|
||||
|
||||
```bash
|
||||
vllm bench serve --omni \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--backend openai-audio-speech \
|
||||
--endpoint /v1/audio/speech \
|
||||
--dataset-name seed-tts \
|
||||
--dataset-path /path/to/seed-tts-eval \
|
||||
--seed-tts-locale en \
|
||||
--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 \
|
||||
--save-result --result-dir ./results
|
||||
```
|
||||
|
||||
#### default_voice (Qwen3-TTS-CustomVoice, bundled seed_tts_smoke)
|
||||
|
||||
```bash
|
||||
vllm bench serve --omni \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--backend openai-audio-speech \
|
||||
--endpoint /v1/audio/speech \
|
||||
--dataset-name seed-tts-text \
|
||||
--dataset-path benchmarks/build_dataset/seed_tts_smoke \
|
||||
--seed-tts-locale en \
|
||||
--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 \
|
||||
--save-result --result-dir ./results
|
||||
```
|
||||
|
||||
#### voice_design (Qwen3-TTS-CustomVoice, bundled seed_tts_design)
|
||||
|
||||
```bash
|
||||
vllm bench serve --omni \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--backend openai-audio-speech \
|
||||
--endpoint /v1/audio/speech \
|
||||
--dataset-name seed-tts-design \
|
||||
--dataset-path benchmarks/build_dataset/seed_tts_design \
|
||||
--seed-tts-locale en \
|
||||
--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 \
|
||||
--save-result --result-dir ./results
|
||||
```
|
||||
|
||||
#### Add WER / SIM / UTMOS to any of the above
|
||||
|
||||
Append `--seed-tts-wer-eval` (and optionally `SEED_TTS_EVAL_DEVICE=cuda:0`
|
||||
in the env, per PR #2558). This triggers the seed-tts-eval protocol:
|
||||
Whisper-large-v3 ASR → WER, WavLM embeddings → SIM, balacoon/utmos → UTMOS.
|
||||
|
||||
### 3. Convenience wrapper (`bench_tts.py`)
|
||||
|
||||
If you're running the **canonical** configuration for a registered model,
|
||||
`bench_tts.py` loads the right defaults from `model_configs.yaml` and
|
||||
emits the exact `vllm bench serve --omni` command above — useful for
|
||||
concurrency sweeps and multi-task runs:
|
||||
|
||||
```bash
|
||||
# Smallest smoke — 5 prompts, concurrency=1
|
||||
python benchmarks/tts/bench_tts.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--task voice_clone \
|
||||
--dataset-path /path/to/seed-tts-eval \
|
||||
--concurrency 1 --num-prompts 5 \
|
||||
--output-dir ./results
|
||||
|
||||
# Full concurrency sweep
|
||||
python benchmarks/tts/bench_tts.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--task voice_clone \
|
||||
--dataset-path /path/to/seed-tts-eval \
|
||||
--concurrency 1 2 4 8 16 32 \
|
||||
--num-prompts 20 \
|
||||
--output-dir ./results
|
||||
|
||||
# With WER / SIM / UTMOS quality eval (adds ASR + embedding compute)
|
||||
python benchmarks/tts/bench_tts.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--task voice_clone \
|
||||
--dataset-path /path/to/seed-tts-eval \
|
||||
--wer-eval \
|
||||
--concurrency 4 --num-prompts 200 \
|
||||
--output-dir ./results
|
||||
```
|
||||
|
||||
### 4. Plot a sweep
|
||||
|
||||
```bash
|
||||
python benchmarks/tts/plot_results.py \
|
||||
--results ./results/*.json \
|
||||
--output ./results/curve.png
|
||||
```
|
||||
|
||||
Outputs TTFP / RTF / throughput curves (and a markdown table) for every
|
||||
`(task, concurrency)` combination in the result set.
|
||||
|
||||
## Task types
|
||||
|
||||
| Task | Dataset | Request body | Checkpoints that support it |
|
||||
|-----------------|-------------------|-----------------------------------------------------|------------------------------------------|
|
||||
| `voice_clone` | `seed-tts` | `ref_audio` + `ref_text` + `task_type=Base` | `Qwen3-TTS-*-Base`, `VoxCPM2` |
|
||||
| `default_voice` | `seed-tts-text` | `voice=Vivian` + `task_type=CustomVoice` | `Qwen3-TTS-*-CustomVoice` |
|
||||
| `voice_design` | `seed-tts-design` | `instructions=<natural-language description>` + `task_type=VoiceDesign` | `Qwen3-TTS-*-CustomVoice` |
|
||||
|
||||
**`-CustomVoice` checkpoints do NOT ship `speaker_encoder` weights**, so
|
||||
voice_clone requests raise `ValueError` at model runtime. Use `-Base` for
|
||||
voice_clone.
|
||||
|
||||
## Adding a new TTS model
|
||||
|
||||
Drop an entry into `model_configs.yaml` — no Python changes required:
|
||||
|
||||
```yaml
|
||||
models:
|
||||
<org>/<model-id>:
|
||||
supported_tasks: [voice_clone] # or default_voice / voice_design
|
||||
backend: openai-audio-speech # vllm bench serve backend
|
||||
endpoint: /v1/audio/speech # OpenAI-compatible endpoint
|
||||
task_extra_body: # merged into every request's body
|
||||
voice_clone:
|
||||
task_type: Base
|
||||
```
|
||||
|
||||
Then add the model's Deploy YAML under `vllm_omni/deploy/<model>.yaml`
|
||||
(Pipeline + Deploy schema) and it's immediately benchable.
|
||||
|
||||
## Datasets
|
||||
|
||||
| Dataset | Bundled? | Format | Source |
|
||||
|--------------------|----------|-------------------|----------------------------------------------------------------|
|
||||
| `seed-tts-design` | ✅ | 5-field meta.lst | `benchmarks/build_dataset/seed_tts_design/en/meta.lst` (20 prompts) |
|
||||
| `seed_tts_smoke` | ✅ | 4-field meta.lst | `benchmarks/build_dataset/seed_tts_smoke/en/meta.lst` (20 text-only) |
|
||||
| `seed-tts` | ❌ | 4-field meta.lst + WAVs | Google-Drive: [BytedanceSpeech/seed-tts-eval][seedtts] (~1.2 GB) |
|
||||
| `seed-tts-text` | ❌ | 4-field meta.lst | Same archive as `seed-tts` (wav column unused) |
|
||||
|
||||
[seedtts]: https://github.com/BytedanceSpeech/seed-tts-eval
|
||||
|
||||
For manual voice_clone / default_voice runs against the full corpus, follow
|
||||
`benchmarks/build_dataset/download_process_data_seedtts.md` and point
|
||||
`--dataset-path` at the extracted `seedtts_testset` directory.
|
||||
|
||||
## DFX nightly CI
|
||||
|
||||
`tests/dfx/perf/tests/test_tts.json` wires three perf regimes plus quality:
|
||||
|
||||
| eval_phase | concurrency | purpose | Baseline metrics |
|
||||
|---------------|-------------|---------------------------------------------------------|-----------------------------------------|
|
||||
| `latency` | 1 | Single-request TTFP / RTF SLO | `median_audio_ttfp_ms`, `median_audio_rtf` |
|
||||
| `throughput` | 8 | Codec-batching cliff sentinel (PDF #272 concurrency≥8) | `median_audio_ttfp_ms`, `median_audio_rtf` |
|
||||
| `quality` | 4 | WER / SIM / UTMOS regression (disabled in CI by default)| `mean_audio_rtf` |
|
||||
|
||||
Why `median_*` for latency/throughput and `mean_*` for quality: latency
|
||||
distributions have cold-start tails that drag the mean; quality aggregates
|
||||
over 200 prompts so single-request outliers don't matter.
|
||||
|
||||
Quality entries are `enabled: false` in CI because seed-tts-eval is not
|
||||
staged in the Buildkite container (matches the precedent in
|
||||
PR #2558 — quality runs are manual / release-validation, not nightly).
|
||||
|
||||
## Concurrency cliff regression sentinel
|
||||
|
||||
Observed on H20-3e, Qwen3-TTS-1.7B (measured pre-merge on this branch):
|
||||
|
||||
| Task | Model | c=1 | c=4 | **c=8** | c=16 | c=32 |
|
||||
|---------------|---------------|--------|--------|------------|--------|--------|
|
||||
| voice_clone | 1.7B-Base | RTF 0.15 / TTFP 165ms | 0.28 / 412ms | **0.49 / 1701ms** | 0.72 / 3355ms | 0.77 / 3772ms |
|
||||
| voice_design | 1.7B-CustomVoice | RTF 0.08 / TTFP 53ms | 0.11 / 154ms | **0.21 / 872ms** | 0.33 / 1801ms | 0.38 / 1989ms |
|
||||
|
||||
Both models show a **4–6× TTFP jump from c=4 to c=8** while audio throughput
|
||||
saturates around c=4–8 — the codec-bs=1 bottleneck documented in
|
||||
vllm-project/vllm-omni#272. The `throughput` CI regime at c=8 is the
|
||||
sentinel for regressions in this area.
|
||||
|
||||
## File layout
|
||||
|
||||
```
|
||||
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)
|
||||
```
|
||||
|
||||
## Related
|
||||
|
||||
- Upstream seed-tts-eval integration: vllm-project/vllm-omni#2558
|
||||
- Pipeline + Deploy schema: vllm-project/vllm-omni#2383
|
||||
- Concurrency cliff RFC: vllm-project/vllm-omni#272
|
||||
@@ -0,0 +1,308 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Universal TTS benchmark CLI for vllm-omni.
|
||||
|
||||
Runs ``vllm bench serve --omni`` with model-aware defaults loaded from
|
||||
``model_configs.yaml``. Supports Qwen3-TTS, VoxCPM2, and any future TTS
|
||||
model registered in the config file -- no code changes needed to add models.
|
||||
|
||||
Usage::
|
||||
|
||||
python benchmarks/tts/bench_tts.py \\
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \\
|
||||
--task voice_clone \\
|
||||
--locale en \\
|
||||
--concurrency 1 4 \\
|
||||
--num-prompts 20 \\
|
||||
--dataset-path /path/to/seed-tts-eval \\
|
||||
--host localhost --port 8000
|
||||
|
||||
See ``--help`` for full option list.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
def _vllm_omni_bin() -> str:
|
||||
"""Return the vllm-omni (or vllm) binary co-located with the current Python."""
|
||||
bin_dir = Path(sys.executable).parent
|
||||
for candidate in ("vllm-omni", "vllm"):
|
||||
p = bin_dir / candidate
|
||||
if p.is_file():
|
||||
return str(p)
|
||||
return "vllm-omni" # fall back and let the shell resolve it
|
||||
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parent.parent.parent
|
||||
_SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
_DEFAULT_MODEL_CONFIGS = _SCRIPT_DIR / "model_configs.yaml"
|
||||
|
||||
# Maps task name to the dataset_name used with vllm bench serve
|
||||
_TASK_TO_DATASET: dict[str, str] = {
|
||||
"voice_clone": "seed-tts",
|
||||
"default_voice": "seed-tts-text",
|
||||
"voice_design": "seed-tts-design",
|
||||
}
|
||||
|
||||
# Default design dataset path (bundled with the repo)
|
||||
_DEFAULT_DESIGN_DATASET_PATH = str(_REPO_ROOT / "benchmarks" / "build_dataset" / "seed_tts_design")
|
||||
|
||||
|
||||
def load_model_configs(path: Path) -> dict[str, Any]:
|
||||
"""Load model registry from YAML file."""
|
||||
with open(path, encoding="utf-8") as f:
|
||||
data = yaml.safe_load(f)
|
||||
return data.get("models", {})
|
||||
|
||||
|
||||
def build_bench_args(
|
||||
*,
|
||||
host: str,
|
||||
port: int,
|
||||
model: str,
|
||||
task: str,
|
||||
model_cfg: dict[str, Any],
|
||||
locale: str,
|
||||
num_prompts: int,
|
||||
concurrency: int | None,
|
||||
dataset_path: str | None,
|
||||
wer_eval: bool,
|
||||
output_dir: str | None,
|
||||
result_filename: str | None,
|
||||
extra_cli_args: list[str],
|
||||
) -> list[str]:
|
||||
"""Build the ``vllm bench serve --omni`` command for one (task, concurrency) run."""
|
||||
dataset_name = _TASK_TO_DATASET[task]
|
||||
backend: str = model_cfg["backend"]
|
||||
endpoint: str = model_cfg["endpoint"]
|
||||
task_extra_body: dict[str, Any] = (model_cfg.get("task_extra_body") or {}).get(task) or {}
|
||||
|
||||
# Resolve dataset path
|
||||
if dataset_path:
|
||||
resolved_dataset_path = dataset_path
|
||||
elif task == "voice_design":
|
||||
resolved_dataset_path = _DEFAULT_DESIGN_DATASET_PATH
|
||||
else:
|
||||
resolved_dataset_path = None
|
||||
|
||||
cmd = [
|
||||
_vllm_omni_bin(),
|
||||
"bench",
|
||||
"serve",
|
||||
"--omni",
|
||||
"--host",
|
||||
host,
|
||||
"--port",
|
||||
str(port),
|
||||
"--model",
|
||||
model,
|
||||
"--backend",
|
||||
backend,
|
||||
"--endpoint",
|
||||
endpoint,
|
||||
"--dataset-name",
|
||||
dataset_name,
|
||||
"--num-prompts",
|
||||
str(num_prompts),
|
||||
"--num-warmups",
|
||||
"2",
|
||||
"--percentile-metrics",
|
||||
"ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
]
|
||||
|
||||
if resolved_dataset_path:
|
||||
cmd += ["--dataset-path", resolved_dataset_path]
|
||||
|
||||
if locale:
|
||||
cmd += ["--seed-tts-locale", locale]
|
||||
|
||||
if task_extra_body:
|
||||
cmd += ["--extra-body", json.dumps(task_extra_body, separators=(",", ":"))]
|
||||
|
||||
if concurrency is not None:
|
||||
cmd += ["--max-concurrency", str(concurrency), "--request-rate", "inf"]
|
||||
|
||||
if wer_eval:
|
||||
cmd.append("--seed-tts-wer-eval")
|
||||
|
||||
if output_dir or result_filename:
|
||||
out_dir = output_dir or "."
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
cmd += ["--save-result", "--result-dir", out_dir]
|
||||
if result_filename:
|
||||
cmd += ["--result-filename", result_filename]
|
||||
|
||||
cmd += extra_cli_args
|
||||
return cmd
|
||||
|
||||
|
||||
def run_one_benchmark(cmd: list[str]) -> dict[str, Any] | None:
|
||||
"""Run a single benchmark subprocess and return parsed JSON result if available."""
|
||||
print(f"\n{'=' * 60}")
|
||||
print("Running:", " ".join(cmd))
|
||||
print("=" * 60)
|
||||
result = subprocess.run(cmd, check=False)
|
||||
if result.returncode != 0:
|
||||
print(f"[bench_tts] WARNING: benchmark exited with code {result.returncode}")
|
||||
return None
|
||||
# If --save-result was used, find the result file
|
||||
try:
|
||||
result_dir_idx = cmd.index("--result-dir")
|
||||
result_dir = Path(cmd[result_dir_idx + 1])
|
||||
if "--result-filename" in cmd:
|
||||
fname_idx = cmd.index("--result-filename")
|
||||
result_file = result_dir / cmd[fname_idx + 1]
|
||||
else:
|
||||
# find most recently modified json
|
||||
jsons = sorted(result_dir.glob("result_*.json"), key=lambda p: p.stat().st_mtime)
|
||||
result_file = jsons[-1] if jsons else None
|
||||
if result_file and result_file.is_file():
|
||||
return json.loads(result_file.read_text(encoding="utf-8"))
|
||||
except (ValueError, IndexError, OSError):
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def print_summary_table(results: list[dict[str, Any]]) -> None:
|
||||
"""Print a unified metrics table across all (task, concurrency) runs."""
|
||||
if not results:
|
||||
return
|
||||
header = (
|
||||
f"{'Task':<16} {'Concurrency':>11} {'RTF mean':>10} "
|
||||
f"{'TTFP (ms)':>10} {'Throughput':>12} {'WER':>7} {'SIM':>7} {'UTMOS':>7}"
|
||||
)
|
||||
print(f"\n{'=' * len(header)}")
|
||||
print("BENCHMARK SUMMARY")
|
||||
print("=" * len(header))
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
for r in results:
|
||||
task = r.get("_task", "?")
|
||||
conc = r.get("_concurrency", "?")
|
||||
rtf = r.get("mean_audio_rtf", float("nan"))
|
||||
ttfp = r.get("mean_audio_ttfp_ms", float("nan"))
|
||||
throughput = r.get("audio_throughput", float("nan"))
|
||||
wer = r.get("seed_tts_mean_wer", float("nan"))
|
||||
sim = r.get("seed_tts_mean_sim", float("nan"))
|
||||
utmos = r.get("seed_tts_mean_utmos", float("nan"))
|
||||
|
||||
def fmt(v: float, digits: int = 3) -> str:
|
||||
return f"{v:.{digits}f}" if not math.isnan(v) else " n/a"
|
||||
|
||||
print(
|
||||
f"{task:<16} {str(conc):>11} {fmt(rtf):>10} {fmt(ttfp, 0):>10} "
|
||||
f"{fmt(throughput):>12} {fmt(wer):>7} {fmt(sim):>7} {fmt(utmos):>7}"
|
||||
)
|
||||
print("=" * len(header))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Entry point for the universal TTS benchmark CLI."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model", required=True, help="HuggingFace model ID (e.g. Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice)"
|
||||
)
|
||||
parser.add_argument("--task", default="all", help="Task type: voice_clone | default_voice | voice_design | all")
|
||||
parser.add_argument("--locale", default="en", choices=["en", "zh"])
|
||||
parser.add_argument("--concurrency", type=int, nargs="+", default=[1, 4], metavar="N")
|
||||
parser.add_argument(
|
||||
"--num-prompts",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[20],
|
||||
metavar="N",
|
||||
help="Number of prompts per run. If one value, applied to all concurrency levels.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dataset-path", default=None, help="Root of seed-tts-eval dataset (required for voice_clone/default_voice)"
|
||||
)
|
||||
parser.add_argument("--wer-eval", action="store_true", help="Enable WER/SIM/UTMOS quality eval")
|
||||
parser.add_argument("--output-dir", default=None, help="Directory to save result JSON files")
|
||||
parser.add_argument("--host", default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--model-configs", default=str(_DEFAULT_MODEL_CONFIGS), help="Path to model_configs.yaml")
|
||||
parser.add_argument("extra", nargs=argparse.REMAINDER, help="Extra args passed directly to vllm bench serve")
|
||||
args = parser.parse_args()
|
||||
|
||||
model_configs = load_model_configs(Path(args.model_configs))
|
||||
if args.model not in model_configs:
|
||||
known = "\n ".join(model_configs.keys())
|
||||
print(f"[bench_tts] ERROR: model '{args.model}' not in model_configs.yaml.\nKnown models:\n {known}")
|
||||
sys.exit(1)
|
||||
|
||||
model_cfg = model_configs[args.model]
|
||||
supported_tasks: list[str] = model_cfg.get("supported_tasks", [])
|
||||
|
||||
tasks_to_run: list[str]
|
||||
if args.task == "all":
|
||||
tasks_to_run = supported_tasks
|
||||
elif args.task in supported_tasks:
|
||||
tasks_to_run = [args.task]
|
||||
else:
|
||||
print(
|
||||
f"[bench_tts] ERROR: task '{args.task}' not supported by {args.model}.\nSupported tasks: {supported_tasks}"
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
# Align num_prompts list with concurrency list
|
||||
num_prompts_list: list[int] = args.num_prompts
|
||||
if len(num_prompts_list) == 1:
|
||||
num_prompts_list = num_prompts_list * len(args.concurrency)
|
||||
elif len(num_prompts_list) != len(args.concurrency):
|
||||
print(
|
||||
f"[bench_tts] ERROR: --num-prompts ({len(num_prompts_list)} values) must be "
|
||||
f"length 1 or match --concurrency ({len(args.concurrency)} values)."
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
all_results: list[dict[str, Any]] = []
|
||||
|
||||
for task in tasks_to_run:
|
||||
for concurrency, num_prompts in zip(args.concurrency, num_prompts_list):
|
||||
ts = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||||
result_filename = f"bench_tts_{args.model.replace('/', '_')}_{task}_c{concurrency}_{ts}.json"
|
||||
cmd = build_bench_args(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
model=args.model,
|
||||
task=task,
|
||||
model_cfg=model_cfg,
|
||||
locale=args.locale,
|
||||
num_prompts=num_prompts,
|
||||
concurrency=concurrency,
|
||||
dataset_path=args.dataset_path,
|
||||
wer_eval=args.wer_eval,
|
||||
output_dir=args.output_dir,
|
||||
result_filename=result_filename,
|
||||
extra_cli_args=args.extra or [],
|
||||
)
|
||||
result = run_one_benchmark(cmd)
|
||||
if result is not None:
|
||||
result["_task"] = task
|
||||
result["_concurrency"] = concurrency
|
||||
all_results.append(result)
|
||||
# Persist the metadata so plot_results.py can pick it up.
|
||||
if args.output_dir and result_filename:
|
||||
result_path = Path(args.output_dir) / result_filename
|
||||
if result_path.is_file():
|
||||
result_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
|
||||
|
||||
print_summary_table(all_results)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+37
-5
@@ -6,6 +6,27 @@ Supports both:
|
||||
- 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
|
||||
@@ -26,9 +47,23 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import AsyncOmni, Omni
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
|
||||
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"
|
||||
DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "deploy" / "voxcpm.yaml"
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -473,9 +508,6 @@ def parse_args():
|
||||
|
||||
def _is_streaming_stage_config(stage_configs_path: str) -> bool:
|
||||
cfg_name = Path(stage_configs_path).name.lower()
|
||||
# Keep routing purely config-path based:
|
||||
# - voxcpm.yaml => sync
|
||||
# - voxcpm_async_chunk.yaml => streaming
|
||||
return "async_chunk" in cfg_name
|
||||
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
# Universal TTS benchmark model registry.
|
||||
# Maps HuggingFace model ID → supported tasks + per-task extra body fields.
|
||||
# To add a new TTS model: add an entry here. No code changes required.
|
||||
#
|
||||
# The server auto-loads its Deploy YAML from vllm_omni/deploy/<model>.yaml via
|
||||
# the Pipeline + Deploy schema introduced in #2383, so no stage_config path
|
||||
# is tracked here.
|
||||
|
||||
models:
|
||||
# -CustomVoice checkpoints lack speaker_encoder weights, so voice_clone is
|
||||
# NOT supported (an attempt raises ValueError from _extract_speaker_embedding
|
||||
# at model runtime). Use -Base for voice_clone.
|
||||
Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice:
|
||||
supported_tasks: [default_voice, voice_design]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
default_voice:
|
||||
voice: Vivian
|
||||
language: English
|
||||
task_type: CustomVoice
|
||||
voice_design:
|
||||
task_type: VoiceDesign
|
||||
language: English
|
||||
|
||||
Qwen/Qwen3-TTS-12Hz-1.7B-Base:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_clone:
|
||||
task_type: Base
|
||||
|
||||
openbmb/VoxCPM2:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
@@ -0,0 +1,324 @@
|
||||
"""Plot universal TTS benchmark results.
|
||||
|
||||
Reads JSON files saved by ``bench_tts.py`` (via ``vllm bench serve --omni``)
|
||||
and generates comparison bar charts grouped by task type.
|
||||
|
||||
Metrics plotted:
|
||||
- AUDIO_TTFP (mean audio time-to-first-packet, ms)
|
||||
- E2EL (mean end-to-end latency, ms)
|
||||
- Audio RTF (mean real-time factor)
|
||||
- Audio throughput (audio-seconds / wall-second)
|
||||
|
||||
Quality metrics (WER / SIM / UTMOS) are printed in a table when present.
|
||||
|
||||
Usage::
|
||||
|
||||
# Single run — one JSON per task, all in results/
|
||||
python benchmarks/tts/plot_results.py \\
|
||||
--results results/bench_tts_*.json \\
|
||||
--output results/tts_benchmark.png
|
||||
|
||||
# Compare two runs (e.g. async_chunk on vs off)
|
||||
python benchmarks/tts/plot_results.py \\
|
||||
--results run_a/bench_tts_*.json \\
|
||||
--results run_b/bench_tts_*.json \\
|
||||
--labels "async_chunk_on" "async_chunk_off" \\
|
||||
--output results/comparison.png
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# JSON loading
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_run(paths: list[str]) -> list[dict]:
|
||||
"""Load and merge all JSON files for one run into a flat list of records.
|
||||
|
||||
Each record is expected to have at least ``_concurrency`` (int) and
|
||||
``_task`` (str) keys injected by ``bench_tts.py``. Records that come
|
||||
from a file that contains a list are flattened.
|
||||
"""
|
||||
records: list[dict] = []
|
||||
for p in paths:
|
||||
raw = json.loads(Path(p).read_text(encoding="utf-8"))
|
||||
if isinstance(raw, list):
|
||||
records.extend(raw)
|
||||
elif isinstance(raw, dict):
|
||||
records.append(raw)
|
||||
return records
|
||||
|
||||
|
||||
def _get(record: dict, key: str) -> float:
|
||||
v = record.get(key, float("nan"))
|
||||
if v is None or (isinstance(v, float) and math.isnan(v)):
|
||||
return float("nan")
|
||||
try:
|
||||
return float(v)
|
||||
except (TypeError, ValueError):
|
||||
return float("nan")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Plotting helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _bar_group(
|
||||
ax: plt.Axes,
|
||||
x: np.ndarray,
|
||||
data_per_label: dict[str, list[float]],
|
||||
width: float,
|
||||
colors: list[str],
|
||||
ylabel: str,
|
||||
title: str,
|
||||
concurrency_labels: list[str],
|
||||
fmt: str = ".1f",
|
||||
) -> None:
|
||||
n = len(data_per_label)
|
||||
offsets = np.linspace(-(n - 1) * width / 2, (n - 1) * width / 2, n) if n > 1 else [0.0]
|
||||
|
||||
for i, (label, values) in enumerate(data_per_label.items()):
|
||||
plot_vals = [0.0 if math.isnan(v) else v for v in values]
|
||||
bar = ax.bar(x + offsets[i], plot_vals, width, label=label, color=colors[i % len(colors)], alpha=0.85)
|
||||
max_val = max((v for v in values if not math.isnan(v)), default=1.0)
|
||||
for rect, val in zip(bar, values):
|
||||
if not math.isnan(val) and val > 0:
|
||||
ax.text(
|
||||
rect.get_x() + rect.get_width() / 2,
|
||||
rect.get_height() + max_val * 0.02,
|
||||
f"{val:{fmt}}",
|
||||
ha="center",
|
||||
va="bottom",
|
||||
fontsize=8,
|
||||
fontweight="bold",
|
||||
)
|
||||
|
||||
ax.set_xlabel("Concurrency", fontsize=11)
|
||||
ax.set_ylabel(ylabel, fontsize=11)
|
||||
ax.set_title(title, fontsize=12, fontweight="bold")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(concurrency_labels)
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
|
||||
COLORS = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107", "#9C27B0"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Comparison plot (multiple labels / runs)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def plot_comparison(
|
||||
all_runs: list[list[dict]],
|
||||
labels: list[str],
|
||||
output_path: str,
|
||||
task_filter: str | None = None,
|
||||
title_prefix: str = "TTS",
|
||||
) -> None:
|
||||
"""One 2×2 subplot per task found in the data."""
|
||||
# Determine tasks to plot
|
||||
tasks: list[str] = []
|
||||
for run in all_runs:
|
||||
for r in run:
|
||||
t = r.get("_task", "unknown")
|
||||
if t not in tasks:
|
||||
tasks.append(t)
|
||||
if task_filter:
|
||||
tasks = [t for t in tasks if t == task_filter]
|
||||
|
||||
n_tasks = len(tasks)
|
||||
if n_tasks == 0:
|
||||
print("[plot_results] No tasks found in data.")
|
||||
return
|
||||
|
||||
fig, axes_grid = plt.subplots(n_tasks, 4, figsize=(18, 4.5 * n_tasks))
|
||||
fig.suptitle(f"{title_prefix} Benchmark", fontsize=15, fontweight="bold")
|
||||
|
||||
# Ensure axes_grid is always 2D
|
||||
if n_tasks == 1:
|
||||
axes_grid = [axes_grid]
|
||||
|
||||
for row_idx, task in enumerate(tasks):
|
||||
# Collect concurrencies across all runs for this task
|
||||
all_concs: set[int] = set()
|
||||
for run in all_runs:
|
||||
for r in run:
|
||||
if r.get("_task") == task:
|
||||
c = r.get("_concurrency")
|
||||
if c is not None:
|
||||
all_concs.add(int(c))
|
||||
concurrencies = sorted(all_concs)
|
||||
x = np.arange(len(concurrencies))
|
||||
conc_labels = [str(c) for c in concurrencies]
|
||||
|
||||
def _series(run: list[dict], metric_key: str) -> list[float]:
|
||||
conc_map = {int(r["_concurrency"]): r for r in run if r.get("_task") == task and "_concurrency" in r}
|
||||
return [_get(conc_map.get(c, {}), metric_key) for c in concurrencies]
|
||||
|
||||
metrics = [
|
||||
("mean_audio_ttfp_ms", "TTFP (ms)", "Time-to-First-Packet", ".0f"),
|
||||
("mean_e2el_ms", "E2E Latency (ms)", "End-to-End Latency", ".0f"),
|
||||
("mean_audio_rtf", "RTF", "Real-Time Factor (RTF)", ".3f"),
|
||||
("audio_throughput", "audio-s / wall-s", "Audio Throughput", ".2f"),
|
||||
]
|
||||
|
||||
axes_row = axes_grid[row_idx]
|
||||
for col_idx, (key, ylabel, subtitle, fmt) in enumerate(metrics):
|
||||
data_per_label = {lbl: _series(run, key) for lbl, run in zip(labels, all_runs)}
|
||||
_bar_group(
|
||||
axes_row[col_idx],
|
||||
x,
|
||||
data_per_label,
|
||||
width=0.3 if len(labels) > 1 else 0.5,
|
||||
colors=COLORS,
|
||||
ylabel=ylabel,
|
||||
title=f"{task} — {subtitle}",
|
||||
concurrency_labels=conc_labels,
|
||||
fmt=fmt,
|
||||
)
|
||||
|
||||
plt.tight_layout()
|
||||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Markdown comparison table
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def print_comparison_table(all_runs: list[list[dict]], labels: list[str]) -> None:
|
||||
tasks: list[str] = []
|
||||
for run in all_runs:
|
||||
for r in run:
|
||||
t = r.get("_task", "unknown")
|
||||
if t not in tasks:
|
||||
tasks.append(t)
|
||||
|
||||
perf_metrics = [
|
||||
("TTFP (ms)", "mean_audio_ttfp_ms", ".1f"),
|
||||
("E2E (ms)", "mean_e2el_ms", ".1f"),
|
||||
("RTF", "mean_audio_rtf", ".3f"),
|
||||
("Throughput (a-s/s)", "audio_throughput", ".2f"),
|
||||
]
|
||||
quality_metrics = [
|
||||
("WER (%)", "seed_tts_mean_wer", ".1f"),
|
||||
("SIM", "seed_tts_mean_sim", ".3f"),
|
||||
("UTMOS", "seed_tts_mean_utmos", ".2f"),
|
||||
]
|
||||
|
||||
for task in tasks:
|
||||
all_concs: set[int] = set()
|
||||
for run in all_runs:
|
||||
for r in run:
|
||||
if r.get("_task") == task:
|
||||
c = r.get("_concurrency")
|
||||
if c is not None:
|
||||
all_concs.add(int(c))
|
||||
concurrencies = sorted(all_concs)
|
||||
|
||||
print(f"\n## {task}\n")
|
||||
col_header = "| Metric | Concurrency |" + "".join(f" {lbl} |" for lbl in labels)
|
||||
sep = "| --- | --- |" + " --- |" * len(labels)
|
||||
print(col_header)
|
||||
print(sep)
|
||||
|
||||
for metric, key, fmt in perf_metrics + quality_metrics:
|
||||
for c in concurrencies:
|
||||
row = f"| {metric} | {c} |"
|
||||
for run in all_runs:
|
||||
conc_map = {
|
||||
int(r["_concurrency"]): r for r in run if r.get("_task") == task and "_concurrency" in r
|
||||
}
|
||||
val = _get(conc_map.get(c, {}), key)
|
||||
row += f" {val:{fmt}} |" if not math.isnan(val) else " n/a |"
|
||||
print(row)
|
||||
|
||||
# Improvement column (2-run comparison only)
|
||||
if len(all_runs) == 2:
|
||||
print(f"\n### Improvement ({labels[0]} vs {labels[1]})\n")
|
||||
print("| Metric | Concurrency | Change |")
|
||||
print("| --- | --- | --- |")
|
||||
for metric, key, _ in perf_metrics:
|
||||
for c in concurrencies:
|
||||
conc_map0 = {
|
||||
int(r["_concurrency"]): r for r in all_runs[0] if r.get("_task") == task and "_concurrency" in r
|
||||
}
|
||||
conc_map1 = {
|
||||
int(r["_concurrency"]): r for r in all_runs[1] if r.get("_task") == task and "_concurrency" in r
|
||||
}
|
||||
v0 = _get(conc_map0.get(c, {}), key)
|
||||
v1 = _get(conc_map1.get(c, {}), key)
|
||||
if not math.isnan(v0) and not math.isnan(v1) and v1 > 0:
|
||||
pct = (v1 - v0) / v1 * 100
|
||||
print(f"| {metric} | {c} | {pct:+.1f}% |")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--results",
|
||||
type=str,
|
||||
nargs="+",
|
||||
action="append",
|
||||
required=True,
|
||||
metavar="FILE",
|
||||
help="JSON result file(s) for one run. Repeat --results for multiple runs to compare.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--labels",
|
||||
type=str,
|
||||
nargs="+",
|
||||
default=None,
|
||||
help="Label for each --results group (must match the number of --results groups).",
|
||||
)
|
||||
parser.add_argument("--output", type=str, default="results/tts_benchmark.png", help="Output image path.")
|
||||
parser.add_argument("--title", type=str, default="TTS", help="Title prefix for the plot.")
|
||||
parser.add_argument("--task", type=str, default=None, help="Filter to a single task (e.g. voice_clone).")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
# args.results is a list-of-lists due to action="append"
|
||||
all_runs: list[list[dict]] = [load_run(group) for group in args.results]
|
||||
n_runs = len(all_runs)
|
||||
|
||||
labels: list[str]
|
||||
if args.labels:
|
||||
if len(args.labels) != n_runs:
|
||||
raise SystemExit(f"--labels count ({len(args.labels)}) must match --results groups ({n_runs})")
|
||||
labels = args.labels
|
||||
else:
|
||||
labels = [f"run{i + 1}" for i in range(n_runs)]
|
||||
|
||||
print_comparison_table(all_runs, labels)
|
||||
plot_comparison(all_runs, labels, args.output, task_filter=args.task, title_prefix=args.title)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,119 +0,0 @@
|
||||
# VoxCPM Benchmark
|
||||
|
||||
This directory contains both:
|
||||
|
||||
- online serving benchmark through the OpenAI-compatible `/v1/audio/speech` API
|
||||
- offline benchmark for `Omni` / `AsyncOmni`
|
||||
- full offline smoke-matrix orchestration
|
||||
|
||||
Both benchmark paths report:
|
||||
|
||||
- TTFP: time to first PCM packet
|
||||
- E2E latency
|
||||
- RTF: real-time factor (`e2e / audio_duration`)
|
||||
|
||||
## Offline Benchmark
|
||||
|
||||
Single offline benchmark run:
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/bench_tts_offline.py \
|
||||
--model /path/to/voxcpm-model \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm.yaml \
|
||||
--text "This is a split-stage VoxCPM synthesis example running on vLLM Omni." \
|
||||
--warmup-runs 1 \
|
||||
--output-dir benchmarks/voxcpm/results/offline_single
|
||||
```
|
||||
|
||||
Streaming offline benchmark:
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/bench_tts_offline.py \
|
||||
--model /path/to/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." \
|
||||
--warmup-runs 1 \
|
||||
--output-dir benchmarks/voxcpm/results/offline_streaming
|
||||
```
|
||||
|
||||
Full fixed offline matrix, equivalent to the old `examples/offline_inference/voxcpm/test.py`:
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/run_offline_matrix.py \
|
||||
--model /path/to/voxcpm-model \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "The exact transcript spoken in reference.wav." \
|
||||
--output-root benchmarks/voxcpm/results/offline_matrix
|
||||
```
|
||||
|
||||
The full matrix covers both routes:
|
||||
|
||||
- streaming: `voxcpm_async_chunk.yaml`
|
||||
- sync: `voxcpm.yaml`
|
||||
|
||||
And these six scenarios under each route:
|
||||
|
||||
- warmup + single TTS
|
||||
- warmup + single voice cloning
|
||||
- warmup + batch TTS
|
||||
- warmup + batch voice cloning
|
||||
- cold single TTS
|
||||
- cold single voice cloning
|
||||
|
||||
`bench_tts_offline.py` itself no longer writes `summary.json` / `results.json`; it prints TTFP / RTF inline and saves generated WAV files only. The matrix runner keeps only per-case `run.log`.
|
||||
|
||||
## Start the Server
|
||||
|
||||
Async-chunk:
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/voxcpm-model \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
Non-streaming:
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/voxcpm-model \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm.yaml \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
## Run the Benchmark
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/bench_tts_serve.py \
|
||||
--host 127.0.0.1 \
|
||||
--port 8091 \
|
||||
--num-prompts 20 \
|
||||
--max-concurrency 1 \
|
||||
--result-dir /tmp/voxcpm_bench
|
||||
```
|
||||
|
||||
Voice cloning benchmark:
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/bench_tts_serve.py \
|
||||
--host 127.0.0.1 \
|
||||
--port 8091 \
|
||||
--num-prompts 10 \
|
||||
--max-concurrency 1 \
|
||||
--ref-audio https://example.com/reference.wav \
|
||||
--ref-text "The exact transcript spoken in the reference audio." \
|
||||
--result-dir /tmp/voxcpm_clone_bench
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
- The benchmark uses `stream=true` and `response_format=pcm` so TTFP is measured from the first audio packet.
|
||||
- `RTF < 1.0` means the server generates audio faster than real time.
|
||||
- For `voxcpm_async_chunk.yaml`, keep concurrency at `1`. This matches native VoxCPM streaming more closely.
|
||||
- Do not benchmark concurrent online streaming on `voxcpm_async_chunk.yaml`; use `voxcpm.yaml` for multi-request throughput runs.
|
||||
- For the offline matrix mode, `--ref-audio` and `--ref-text` are required because clone cases are part of the fixed coverage set.
|
||||
@@ -1,283 +0,0 @@
|
||||
"""Benchmark VoxCPM via /v1/audio/speech.
|
||||
|
||||
Reports TTFP (time to first packet), E2E latency, and RTF (real-time factor).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
DEFAULT_MODEL = "OpenBMB/VoxCPM1.5"
|
||||
DEFAULT_SAMPLE_RATE = 24000
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the VoxCPM speech benchmark.",
|
||||
"This is a short benchmark prompt for online text-to-speech generation.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"Please remember to bring your identification documents tomorrow morning.",
|
||||
"Learning a new language takes patience, practice, and curiosity.",
|
||||
"This benchmark reports TTFP and RTF for the VoxCPM online serving path.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestResult:
|
||||
success: bool = False
|
||||
ttfp: float = 0.0
|
||||
e2e: float = 0.0
|
||||
audio_bytes: int = 0
|
||||
audio_duration: float = 0.0
|
||||
rtf: float = 0.0
|
||||
prompt: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
concurrency: int = 0
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
p95_rtf: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
request_throughput: float = 0.0
|
||||
per_request: list[dict[str, float | str]] = field(default_factory=list)
|
||||
|
||||
|
||||
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = DEFAULT_SAMPLE_RATE, sample_width: int = 2) -> float:
|
||||
num_samples = num_bytes / sample_width
|
||||
return num_samples / sample_rate
|
||||
|
||||
|
||||
async def send_tts_request(
|
||||
session: aiohttp.ClientSession,
|
||||
api_url: str,
|
||||
*,
|
||||
model: str,
|
||||
prompt: str,
|
||||
ref_audio: str | None,
|
||||
ref_text: str | None,
|
||||
pbar: tqdm | None = None,
|
||||
) -> RequestResult:
|
||||
payload: dict[str, object] = {
|
||||
"model": model,
|
||||
"input": prompt,
|
||||
"stream": True,
|
||||
"response_format": "pcm",
|
||||
}
|
||||
if ref_audio is not None:
|
||||
payload["ref_audio"] = ref_audio
|
||||
if ref_text is not None:
|
||||
payload["ref_text"] = ref_text
|
||||
|
||||
result = RequestResult(prompt=prompt)
|
||||
started_at = time.perf_counter()
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status != 200:
|
||||
result.error = f"HTTP {response.status}: {await response.text()}"
|
||||
return result
|
||||
|
||||
first_chunk = True
|
||||
total_bytes = 0
|
||||
async for chunk in response.content.iter_any():
|
||||
if not chunk:
|
||||
continue
|
||||
if first_chunk:
|
||||
result.ttfp = time.perf_counter() - started_at
|
||||
first_chunk = False
|
||||
total_bytes += len(chunk)
|
||||
|
||||
result.e2e = time.perf_counter() - started_at
|
||||
result.audio_bytes = total_bytes
|
||||
result.audio_duration = pcm_bytes_to_duration(total_bytes)
|
||||
if result.audio_duration > 0:
|
||||
result.rtf = result.e2e / result.audio_duration
|
||||
result.success = True
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.e2e = time.perf_counter() - started_at
|
||||
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
return result
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
*,
|
||||
host: str,
|
||||
port: int,
|
||||
model: str,
|
||||
num_prompts: int,
|
||||
max_concurrency: int,
|
||||
num_warmups: int,
|
||||
ref_audio: str | None,
|
||||
ref_text: str | None,
|
||||
) -> BenchmarkResult:
|
||||
api_url = f"http://{host}:{port}/v1/audio/speech"
|
||||
connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
|
||||
timeout = aiohttp.ClientTimeout(total=600)
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
|
||||
if num_warmups > 0:
|
||||
print(f" Warming up with {num_warmups} requests...")
|
||||
warmup_tasks = [
|
||||
send_tts_request(
|
||||
session,
|
||||
api_url,
|
||||
model=model,
|
||||
prompt=PROMPTS[i % len(PROMPTS)],
|
||||
ref_audio=ref_audio,
|
||||
ref_text=ref_text,
|
||||
)
|
||||
for i in range(num_warmups)
|
||||
]
|
||||
await asyncio.gather(*warmup_tasks)
|
||||
print(" Warmup done.")
|
||||
|
||||
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
|
||||
semaphore = asyncio.Semaphore(max_concurrency)
|
||||
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
|
||||
|
||||
async def limited_request(prompt: str) -> RequestResult:
|
||||
async with semaphore:
|
||||
return await send_tts_request(
|
||||
session,
|
||||
api_url,
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
ref_audio=ref_audio,
|
||||
ref_text=ref_text,
|
||||
pbar=pbar,
|
||||
)
|
||||
|
||||
started_at = time.perf_counter()
|
||||
results = await asyncio.gather(*[asyncio.create_task(limited_request(prompt)) for prompt in request_prompts])
|
||||
duration = time.perf_counter() - started_at
|
||||
pbar.close()
|
||||
|
||||
succeeded = [result for result in results if result.success]
|
||||
bench = BenchmarkResult(
|
||||
concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
completed=len(succeeded),
|
||||
failed=len(results) - len(succeeded),
|
||||
duration_s=duration,
|
||||
)
|
||||
|
||||
if not succeeded:
|
||||
return bench
|
||||
|
||||
ttfps = np.array([result.ttfp * 1000 for result in succeeded], dtype=np.float64)
|
||||
e2es = np.array([result.e2e * 1000 for result in succeeded], dtype=np.float64)
|
||||
rtfs = np.array([result.rtf for result in succeeded], dtype=np.float64)
|
||||
audio_durations = np.array([result.audio_duration for result in succeeded], dtype=np.float64)
|
||||
|
||||
bench.mean_ttfp_ms = float(np.mean(ttfps))
|
||||
bench.median_ttfp_ms = float(np.median(ttfps))
|
||||
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
|
||||
bench.mean_e2e_ms = float(np.mean(e2es))
|
||||
bench.median_e2e_ms = float(np.median(e2es))
|
||||
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
|
||||
bench.mean_rtf = float(np.mean(rtfs))
|
||||
bench.median_rtf = float(np.median(rtfs))
|
||||
bench.p95_rtf = float(np.percentile(rtfs, 95))
|
||||
bench.total_audio_duration_s = float(np.sum(audio_durations))
|
||||
bench.request_throughput = len(succeeded) / duration if duration > 0 else 0.0
|
||||
bench.per_request = [
|
||||
{
|
||||
"prompt": result.prompt,
|
||||
"ttfp_ms": result.ttfp * 1000,
|
||||
"e2e_ms": result.e2e * 1000,
|
||||
"rtf": result.rtf,
|
||||
"audio_duration_s": result.audio_duration,
|
||||
}
|
||||
for result in succeeded
|
||||
]
|
||||
|
||||
return bench
|
||||
|
||||
|
||||
def print_summary(result: BenchmarkResult) -> None:
|
||||
width = 54
|
||||
print("")
|
||||
print("=" * width)
|
||||
print(f"{'VoxCPM Serving Benchmark':^{width}}")
|
||||
print("=" * width)
|
||||
print(f"concurrency : {result.concurrency}")
|
||||
print(f"requests : {result.completed}/{result.num_prompts} succeeded")
|
||||
print(f"wall time (s) : {result.duration_s:.3f}")
|
||||
print(f"mean TTFP (ms) : {result.mean_ttfp_ms:.2f}")
|
||||
print(f"p95 TTFP (ms) : {result.p95_ttfp_ms:.2f}")
|
||||
print(f"mean E2E (ms) : {result.mean_e2e_ms:.2f}")
|
||||
print(f"p95 E2E (ms) : {result.p95_e2e_ms:.2f}")
|
||||
print(f"mean RTF : {result.mean_rtf:.3f}")
|
||||
print(f"p95 RTF : {result.p95_rtf:.3f}")
|
||||
print(f"request throughput : {result.request_throughput:.2f} req/s")
|
||||
print("=" * width)
|
||||
|
||||
|
||||
async def main_async(args) -> None:
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
all_results: list[BenchmarkResult] = []
|
||||
for concurrency in args.max_concurrency:
|
||||
result = await run_benchmark(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
model=args.model,
|
||||
num_prompts=args.num_prompts,
|
||||
max_concurrency=concurrency,
|
||||
num_warmups=args.num_warmups,
|
||||
ref_audio=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
)
|
||||
print_summary(result)
|
||||
all_results.append(result)
|
||||
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"created_at": datetime.utcnow().isoformat() + "Z",
|
||||
"results": [asdict(result) for result in all_results],
|
||||
}
|
||||
result_path = result_dir / "bench_tts_serve.json"
|
||||
result_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
|
||||
print(f"Saved results to: {result_path}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Benchmark VoxCPM via /v1/audio/speech")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="Server host")
|
||||
parser.add_argument("--port", type=int, default=8091, help="Server port")
|
||||
parser.add_argument("--model", default=DEFAULT_MODEL, help="Model name or path")
|
||||
parser.add_argument("--num-prompts", type=int, default=20, help="Number of prompts to send")
|
||||
parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1], help="Concurrency levels to benchmark")
|
||||
parser.add_argument("--num-warmups", type=int, default=3, help="Warmup request count")
|
||||
parser.add_argument("--ref-audio", default=None, help="Reference audio URL or data URL for voice cloning")
|
||||
parser.add_argument("--ref-text", default=None, help="Reference audio transcript for voice cloning")
|
||||
parser.add_argument("--result-dir", default="results", help="Directory to save benchmark JSON")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main_async(parse_args()))
|
||||
@@ -1,303 +0,0 @@
|
||||
"""Run the full offline VoxCPM smoke matrix.
|
||||
|
||||
This script keeps the old `test.py` coverage, but delegates each case to
|
||||
`bench_tts_offline.py` so the benchmark runner itself stays focused on a
|
||||
single execution path.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shlex
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
BENCH_SCRIPT = Path(__file__).with_name("bench_tts_offline.py")
|
||||
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"
|
||||
DEFAULT_OUTPUT_ROOT = BENCH_SCRIPT.parents[1] / "results" / "offline_matrix"
|
||||
|
||||
SINGLE_TTS_TEXT = "This is a single text-to-speech smoke test for VoxCPM on vLLM Omni."
|
||||
SINGLE_CLONE_TEXT = "This sentence is synthesized with the cloned voice for validation."
|
||||
BATCH_TTS_TEXTS = [
|
||||
"The first batch text-to-speech sample validates sequential batch execution.",
|
||||
"The second batch text-to-speech sample checks another prompt in the same file.",
|
||||
"The third batch text-to-speech sample completes the sequential batch path.",
|
||||
]
|
||||
BATCH_CLONE_TEXTS = [
|
||||
"The first cloned sample validates sequential batch voice cloning.",
|
||||
"The second cloned sample checks the same reference voice on another prompt.",
|
||||
"The third cloned sample finishes the shared-reference clone batch path.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ModeSpec:
|
||||
name: str
|
||||
stage_config: Path
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CaseSpec:
|
||||
name: str
|
||||
warmup_runs: int
|
||||
prompt_kind: str
|
||||
voice_clone: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CaseResult:
|
||||
mode: str
|
||||
case: str
|
||||
returncode: int
|
||||
elapsed_s: float
|
||||
output_dir: Path
|
||||
log_path: Path
|
||||
|
||||
@property
|
||||
def ok(self) -> bool:
|
||||
return self.returncode == 0
|
||||
|
||||
|
||||
MODE_SPECS = [
|
||||
ModeSpec(name="streaming", stage_config=DEFAULT_STAGE_ASYNC),
|
||||
ModeSpec(name="sync", stage_config=DEFAULT_STAGE_SYNC),
|
||||
]
|
||||
|
||||
CASE_SPECS = [
|
||||
CaseSpec(name="warmup_single_tts", warmup_runs=1, prompt_kind="single", voice_clone=False),
|
||||
CaseSpec(name="warmup_single_clone", warmup_runs=1, prompt_kind="single", voice_clone=True),
|
||||
CaseSpec(name="warmup_batch_tts", warmup_runs=1, prompt_kind="batch", voice_clone=False),
|
||||
CaseSpec(name="warmup_batch_clone", warmup_runs=1, prompt_kind="batch", voice_clone=True),
|
||||
CaseSpec(name="cold_single_tts", warmup_runs=0, prompt_kind="single", voice_clone=False),
|
||||
CaseSpec(name="cold_single_clone", warmup_runs=0, prompt_kind="single", voice_clone=True),
|
||||
]
|
||||
|
||||
|
||||
def _write_lines(path: Path, lines: list[str]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
def _prepare_batch_inputs(output_root: Path) -> tuple[Path, Path]:
|
||||
input_dir = output_root / "inputs"
|
||||
batch_tts_path = input_dir / "batch_tts_prompts.txt"
|
||||
batch_clone_path = input_dir / "batch_clone_prompts.txt"
|
||||
_write_lines(batch_tts_path, BATCH_TTS_TEXTS)
|
||||
_write_lines(batch_clone_path, BATCH_CLONE_TEXTS)
|
||||
return batch_tts_path, batch_clone_path
|
||||
|
||||
|
||||
def _base_command(args, mode: ModeSpec, output_dir: Path) -> list[str]:
|
||||
cmd = [
|
||||
args.python,
|
||||
str(BENCH_SCRIPT),
|
||||
"--model",
|
||||
args.model,
|
||||
"--stage-configs-path",
|
||||
str(mode.stage_config),
|
||||
"--output-dir",
|
||||
str(output_dir),
|
||||
"--num-runs",
|
||||
str(args.num_runs),
|
||||
"--stage-init-timeout",
|
||||
str(args.stage_init_timeout),
|
||||
]
|
||||
cmd.append("--log-stats" if args.log_stats else "--no-log-stats")
|
||||
cmd.extend(["--cfg-value", str(args.cfg_value)])
|
||||
cmd.extend(["--inference-timesteps", str(args.inference_timesteps)])
|
||||
cmd.extend(["--min-len", str(args.min_len)])
|
||||
cmd.extend(["--max-new-tokens", str(args.max_new_tokens)])
|
||||
if args.streaming_prefix_len is not None:
|
||||
cmd.extend(["--streaming-prefix-len", str(args.streaming_prefix_len)])
|
||||
if args.enable_profiler:
|
||||
profiler_dir = Path(args.profiler_dir) if args.profiler_dir is not None else (output_dir / "profiler")
|
||||
cmd.append("--enable-profiler")
|
||||
cmd.extend(["--profiler-dir", str(profiler_dir)])
|
||||
cmd.extend(["--profiler-wait-seconds", str(args.profiler_wait_seconds)])
|
||||
if args.profiler_stages is not None:
|
||||
cmd.append("--profiler-stages")
|
||||
cmd.extend(str(stage_id) for stage_id in args.profiler_stages)
|
||||
return cmd
|
||||
|
||||
|
||||
def _build_case_command(
|
||||
args,
|
||||
mode: ModeSpec,
|
||||
case: CaseSpec,
|
||||
*,
|
||||
batch_tts_path: Path,
|
||||
batch_clone_path: Path,
|
||||
output_dir: Path,
|
||||
) -> list[str]:
|
||||
cmd = _base_command(args, mode, output_dir)
|
||||
cmd.extend(["--warmup-runs", str(case.warmup_runs)])
|
||||
if case.prompt_kind == "single":
|
||||
cmd.extend(["--text", SINGLE_CLONE_TEXT if case.voice_clone else SINGLE_TTS_TEXT])
|
||||
else:
|
||||
cmd.extend(["--txt-prompts", str(batch_clone_path if case.voice_clone else batch_tts_path)])
|
||||
if case.voice_clone:
|
||||
cmd.extend(["--ref-audio", args.ref_audio, "--ref-text", args.ref_text])
|
||||
return cmd
|
||||
|
||||
|
||||
def _run_case(
|
||||
args,
|
||||
mode: ModeSpec,
|
||||
case: CaseSpec,
|
||||
*,
|
||||
batch_tts_path: Path,
|
||||
batch_clone_path: Path,
|
||||
output_root: Path,
|
||||
) -> CaseResult:
|
||||
case_output_dir = output_root / mode.name / case.name
|
||||
case_output_dir.mkdir(parents=True, exist_ok=True)
|
||||
case_log_path = case_output_dir / "run.log"
|
||||
cmd = _build_case_command(
|
||||
args,
|
||||
mode,
|
||||
case,
|
||||
batch_tts_path=batch_tts_path,
|
||||
batch_clone_path=batch_clone_path,
|
||||
output_dir=case_output_dir,
|
||||
)
|
||||
|
||||
print()
|
||||
print("=" * 80)
|
||||
print(f"[{mode.name}] {case.name}")
|
||||
print(f"Output directory: {case_output_dir}")
|
||||
print(shlex.join(cmd))
|
||||
|
||||
start = time.perf_counter()
|
||||
with case_log_path.open("w", encoding="utf-8") as log_fp:
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
)
|
||||
assert process.stdout is not None
|
||||
for line in process.stdout:
|
||||
print(line, end="")
|
||||
log_fp.write(line)
|
||||
process.wait()
|
||||
|
||||
elapsed_s = time.perf_counter() - start
|
||||
status = "PASS" if (process.returncode or 0) == 0 else f"FAIL({process.returncode})"
|
||||
print(f"[{mode.name}] {case.name} -> {status} ({elapsed_s:.2f}s)")
|
||||
return CaseResult(
|
||||
mode=mode.name,
|
||||
case=case.name,
|
||||
returncode=int(process.returncode or 0),
|
||||
elapsed_s=elapsed_s,
|
||||
output_dir=case_output_dir,
|
||||
log_path=case_log_path,
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Run the full offline VoxCPM smoke matrix.")
|
||||
parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.")
|
||||
parser.add_argument("--ref-audio", type=str, required=True, help="Reference audio path for clone cases.")
|
||||
parser.add_argument("--ref-text", type=str, required=True, help="Exact transcript spoken in --ref-audio.")
|
||||
parser.add_argument("--output-root", type=str, default=str(DEFAULT_OUTPUT_ROOT), help="Root directory for outputs.")
|
||||
parser.add_argument("--python", type=str, default=sys.executable, help="Python executable used to launch cases.")
|
||||
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 measured runs per case.")
|
||||
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="Optional VoxCPM streaming window passed to streaming cases.",
|
||||
)
|
||||
parser.add_argument("--enable-profiler", action="store_true", help="Enable torch profiler for each case.")
|
||||
parser.add_argument(
|
||||
"--profiler-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Profiler output root. Defaults to <case-output-dir>/profiler.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-stages",
|
||||
type=int,
|
||||
nargs="*",
|
||||
default=None,
|
||||
help="Optional stage ids to profile. Defaults to all configured stages.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-wait-seconds",
|
||||
type=float,
|
||||
default=30.0,
|
||||
help="Seconds to wait after stopping profiler for traces to flush.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
if args.num_runs < 1:
|
||||
parser.error("--num-runs must be >= 1")
|
||||
return args
|
||||
|
||||
|
||||
def main(args) -> int:
|
||||
output_root = Path(args.output_root)
|
||||
output_root.mkdir(parents=True, exist_ok=True)
|
||||
batch_tts_path, batch_clone_path = _prepare_batch_inputs(output_root)
|
||||
|
||||
print(f"Model: {args.model}")
|
||||
print(f"Reference audio: {args.ref_audio}")
|
||||
print(f"Reference text: {args.ref_text}")
|
||||
print(f"Python: {args.python}")
|
||||
print(f"Output root: {output_root}")
|
||||
print(f"Cases: {len(MODE_SPECS) * len(CASE_SPECS)}")
|
||||
|
||||
results: list[CaseResult] = []
|
||||
for mode in MODE_SPECS:
|
||||
for case in CASE_SPECS:
|
||||
results.append(
|
||||
_run_case(
|
||||
args,
|
||||
mode,
|
||||
case,
|
||||
batch_tts_path=batch_tts_path,
|
||||
batch_clone_path=batch_clone_path,
|
||||
output_root=output_root,
|
||||
)
|
||||
)
|
||||
|
||||
failed = [result for result in results if not result.ok]
|
||||
print()
|
||||
print("=" * 80)
|
||||
print("Summary:")
|
||||
for result in results:
|
||||
status = "PASS" if result.ok else f"FAIL({result.returncode})"
|
||||
print(f"- [{result.mode}] {result.case}: {status} ({result.elapsed_s:.2f}s)")
|
||||
print(f" output_dir={result.output_dir}")
|
||||
print(f" log={result.log_path}")
|
||||
|
||||
print(f"Passed: {len(results) - len(failed)}/{len(results)}")
|
||||
if failed:
|
||||
print("Failed cases:")
|
||||
for result in failed:
|
||||
print(f"- [{result.mode}] {result.case}: see {result.log_path}")
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main(parse_args()))
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
|
||||
ARG VLLM_BASE_TAG=v0.19.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,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.19.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.20.0
|
||||
FROM ${BASE_IMAGE}
|
||||
|
||||
ARG COMMON_WORKDIR=/app
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
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 && \
|
||||
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" \
|
||||
--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 []
|
||||
@@ -0,0 +1,42 @@
|
||||
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 && \
|
||||
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" \
|
||||
--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,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.19.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.
|
||||
|
||||
+29
-22
@@ -14,8 +14,9 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO
|
||||
RUN apt clean && apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends --fix-missing \
|
||||
curl \
|
||||
espeak-ng \
|
||||
ffmpeg \
|
||||
git \
|
||||
libsndfile1 \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libgl1 \
|
||||
@@ -28,8 +29,20 @@ RUN apt clean && apt-get update -y && \
|
||||
python3.12-dev \
|
||||
python3-pip
|
||||
|
||||
RUN apt update && apt upgrade -y && \
|
||||
apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
|
||||
# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer.
|
||||
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
|
||||
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
|
||||
printf '%s\n' \
|
||||
'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \
|
||||
'Pin: version 2025.3*' \
|
||||
'Pin-Priority: 1001' \
|
||||
> /etc/apt/preferences.d/oneapi-2025.3.pref && \
|
||||
apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
intel-oneapi-compiler-dpcpp-cpp-2025.3 \
|
||||
intel-oneapi-mkl-devel-2025.3 \
|
||||
intel-oneapi-dnnl-devel-2025.3 && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install UMD with fixed version
|
||||
RUN mkdir neo && \
|
||||
@@ -52,9 +65,9 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
|
||||
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
|
||||
|
||||
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2.
|
||||
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh"
|
||||
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \
|
||||
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
|
||||
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
|
||||
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
|
||||
bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
|
||||
rm "${ONECCL_INSTALLER}" && \
|
||||
echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
|
||||
@@ -66,6 +79,7 @@ SHELL ["bash", "-c"]
|
||||
CMD ["bash", "-c", "source /root/.bashrc && exec bash"]
|
||||
|
||||
WORKDIR /workspace/
|
||||
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
# Configure package index for XPU
|
||||
@@ -74,24 +88,18 @@ 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.19.0
|
||||
ARG VLLM_VERSION=v0.20.0
|
||||
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --upgrade pip && \
|
||||
uv pip install -r requirements/xpu.txt
|
||||
|
||||
# used for suffix method speculative decoding
|
||||
# build deps for proto + nanobind-based extensions to set up the build environment
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install grpcio-tools protobuf nanobind
|
||||
# arctic-inference is built from source which needs torch-xpu properly installed first
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install -r requirements/xpu.txt && \
|
||||
uv pip install setuptools_scm &&\
|
||||
uv pip install grpcio-tools protobuf nanobind && \
|
||||
source /opt/intel/oneapi/setvars.sh --force && \
|
||||
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
|
||||
export CMAKE_PREFIX_PATH="$(python -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
|
||||
uv pip install --no-build-isolation arctic-inference==0.1.1
|
||||
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
|
||||
|
||||
ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
|
||||
|
||||
@@ -105,6 +113,9 @@ CMD ["/bin/bash"]
|
||||
|
||||
FROM vllm-base AS vllm-openai
|
||||
|
||||
# install nixl from source code
|
||||
RUN uv pip install nixl
|
||||
|
||||
# install additional dependencies for openai api server
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install accelerate hf_transfer pytest pytest_asyncio lm_eval[api] modelscope
|
||||
@@ -112,10 +123,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
# install development dependencies (for testing)
|
||||
RUN uv pip install -e tests/vllm_test_utils
|
||||
|
||||
# install nixl from source code
|
||||
ENV NIXL_VERSION=0.7.0
|
||||
RUN python /workspace/vllm/tools/install_nixl_from_source_ubuntu.py
|
||||
|
||||
# ensure vllm is properly installed
|
||||
RUN python -c "import vllm, inspect; print(vllm.__file__)"
|
||||
RUN uv pip show vllm
|
||||
@@ -135,7 +142,7 @@ RUN uv pip install --no-cache-dir ".[dev]" --no-build-isolation
|
||||
# FIX triton
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip uninstall triton triton-xpu && \
|
||||
uv pip install triton-xpu==3.6.0
|
||||
uv pip install triton-xpu==3.7.0
|
||||
|
||||
# remove torch bundled oneccl to avoid conflicts
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
|
||||
+18
-13
@@ -10,38 +10,38 @@ nav:
|
||||
- Image Generation: serving/image_generation_api.md
|
||||
- Image Edit: serving/image_edit_api.md
|
||||
- Text to Speech: serving/speech_api.md
|
||||
- Audio Generation: serving/audio_generate_api.md
|
||||
- Streaming Video Input: serving/video_stream_api.md
|
||||
- Examples:
|
||||
- examples/README.md
|
||||
- Offline Inference:
|
||||
- BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
|
||||
- CosyVoice3: user_guide/examples/offline_inference/cosyvoice3.md
|
||||
- Fish Speech S2 Pro: user_guide/examples/offline_inference/fish_speech.md
|
||||
- GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
|
||||
- Helios Video Generation: user_guide/examples/offline_inference/helios.md
|
||||
- HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
|
||||
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
|
||||
- Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
|
||||
- InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
|
||||
- MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
|
||||
- MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
|
||||
- Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
|
||||
- Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
|
||||
- Qwen3-TTS: user_guide/examples/offline_inference/qwen3_tts.md
|
||||
- Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
|
||||
- Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
|
||||
- Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
|
||||
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
|
||||
- Voxtral TTS Offline Inference: user_guide/examples/offline_inference/voxtral_tts.md
|
||||
- X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
|
||||
- Online Serving:
|
||||
- BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
|
||||
- vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
|
||||
- Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md
|
||||
- Diffusers Backend Adapter: user_guide/examples/online_serving/diffusers_pipeline_adapter.md
|
||||
- GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
|
||||
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md
|
||||
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md
|
||||
- Online serving Example of vLLM-Omni for MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
|
||||
- Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
|
||||
- Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
|
||||
- Qwen3-TTS: user_guide/examples/online_serving/qwen3_tts.md
|
||||
- Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
|
||||
- Text-To-Image: user_guide/examples/online_serving/text_to_image.md
|
||||
- Text-To-Video: user_guide/examples/online_serving/text_to_video.md
|
||||
- General:
|
||||
@@ -59,11 +59,7 @@ nav:
|
||||
- Cache Acceleration:
|
||||
- TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
|
||||
- Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
|
||||
- Quantization:
|
||||
- Overview: user_guide/diffusion/quantization/overview.md
|
||||
- FP8: user_guide/diffusion/quantization/fp8.md
|
||||
- Int8: user_guide/diffusion/quantization/int8.md
|
||||
- GGUF: user_guide/diffusion/quantization/gguf.md
|
||||
- Attention Backends: user_guide/diffusion/attention_backends.md
|
||||
- Frame Interpolation: user_guide/diffusion/frame_interpolation.md
|
||||
- Parallelism:
|
||||
- Overview: user_guide/diffusion/parallelism/overview.md
|
||||
@@ -77,11 +73,20 @@ nav:
|
||||
- LoRA: user_guide/diffusion/lora.md
|
||||
- Custom Pipeline: features/custom_pipeline.md
|
||||
- Step Execution: user_guide/diffusion/step_execution.md
|
||||
- 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
|
||||
- ModelOpt: user_guide/quantization/modelopt.md
|
||||
- GGUF: user_guide/quantization/gguf.md
|
||||
- AutoRound: user_guide/quantization/autoround.md
|
||||
- msModelSlim: user_guide/quantization/msmodelslim.md
|
||||
- ComfyUI: features/comfyui.md
|
||||
- Developer Guide:
|
||||
- General:
|
||||
- contributing/README.md
|
||||
- pr_reviewer.md
|
||||
- glob: contributing/*
|
||||
flatten_single_child_sections: true
|
||||
- Model Implementation:
|
||||
@@ -89,7 +94,6 @@ nav:
|
||||
- contributing/model/adding_omni_model.md
|
||||
- contributing/model/adding_tts_model.md
|
||||
- contributing/model/adding_diffusion_model.md
|
||||
- contributing/model/adding_quantization_model.md
|
||||
- CI: contributing/ci
|
||||
- Design Documents:
|
||||
- design/index.md
|
||||
@@ -110,6 +114,7 @@ nav:
|
||||
- design/feature/async_chunk.md
|
||||
- design/feature/vae_parallel.md
|
||||
- design/feature/diffusion_step_execution.md
|
||||
- design/feature/diffusion_continuous_batching.md
|
||||
- Module Design:
|
||||
- design/module/ar_module.md
|
||||
- design/module/dit_module.md
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 96 KiB After Width: | Height: | Size: 182 KiB |
@@ -1,5 +1,61 @@
|
||||
# vllm-omni serve
|
||||
|
||||
## Stage-based CLI quickstart
|
||||
|
||||
The stage-based CLI is designed for deployments that require launching each pipeline stage in an isolated process
|
||||
(e.g., across separate operating system processes, distinct GPUs, or distributed hosts).
|
||||
|
||||
- For **migrated models** that utilize the bundled deployment YAML configurations located in
|
||||
`vllm_omni/deploy/`, the `--deploy-config` flag is only required to override the default configuration. By default, executing `vllm serve MODEL --omni ...`
|
||||
automatically loads the bundled deployment configuration.
|
||||
- For **legacy models** utilizing configuration files located in
|
||||
`vllm_omni/model_executor/stage_configs/`, the `--stage-configs-path` parameter remains mandatory.
|
||||
|
||||
Example: Initializing Stage 0 (Orchestrator and API Server):
|
||||
The commands below show a common device mapping where Stage 0 uses GPU 0 and
|
||||
worker stages use GPU 1 via `CUDA_VISIBLE_DEVICES`.
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--port 8091 \
|
||||
--stage-id 0 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
Example: Initializing a Headless Worker Stage (Stage 1):
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
When utilizing a custom deployment YAML based on the new schema, append `--deploy-config /path/to/override.yaml` to each command execution. Conversely, for legacy models, substitute this parameter with `--stage-configs-path /path/to/stage_configs.yaml`.
|
||||
|
||||
In the standard execution paradigm, the `--stage-overrides` argument is utilized to apply stage-specific configurations from a single CLI command.
|
||||
However, under the **stage-based CLI** paradigm, where each process strictly encapsulates a single stage, it is recommended to specify tuning parameters directly via discrete command-line flags for the respective stage, rather than constructing a composite `--stage-overrides` JSON string.
|
||||
|
||||
For example, as an alternative to the following composite configuration:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
|
||||
```
|
||||
|
||||
the stage-based CLI permits the direct initialization of Stage 1 with explicit parameters:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--gpu-memory-utilization 0.5 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
## JSON CLI Arguments
|
||||
|
||||
--8<-- "docs/cli/json_tip.inc.md"
|
||||
|
||||
@@ -19,7 +19,7 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze
|
||||
| `platforms` | dict | optional | `null` | Keyed by `npu` / `rocm` / `xpu`, each contains a `stages:` list with per-platform overrides applied on top of the CUDA defaults. |
|
||||
| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only`). |
|
||||
| `trust_remote_code` | bool | optional | `true` | **Pipeline-wide.** Trust HF remote code on model load; applies to every stage. |
|
||||
| `distributed_executor_backend` | str | optional | `"mp"` | **Pipeline-wide.** Executor backend (`"mp"` or `"ray"`). |
|
||||
| `distributed_executor_backend` | str \| null | optional | `null` | **Pipeline-wide.** Distributed executor backend forwarded to vLLM (`"mp"`, `"ray"`, `"external_launcher"`). If omitted, vLLM auto-selects backend from runtime topology. |
|
||||
| `dtype` | str \| null | optional | `null` | **Pipeline-wide.** Model dtype for every stage. |
|
||||
| `quantization` | str \| null | optional | `null` | **Pipeline-wide.** Quantization method for every stage. |
|
||||
| `enable_prefix_caching` | bool | optional | `false` | **Pipeline-wide.** Prefix cache toggle applied to every stage. |
|
||||
@@ -27,6 +27,9 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze
|
||||
| `data_parallel_size` | int | optional | `1` | **Pipeline-wide.** DP degree for every stage. |
|
||||
| `pipeline_parallel_size` | int | optional | `1` | **Pipeline-wide.** PP degree for every stage. |
|
||||
|
||||
Note: for diffusion path, `distributed_executor_backend` currently defaults to
|
||||
`mp`, and `ray` / `external_launcher` are not fully supported yet.
|
||||
|
||||
### Stage fields
|
||||
|
||||
Each entry under `stages:` accepts any `StageDeployConfig` field directly (no nested `engine_args:`). Only fields whose value legitimately varies across stages live here; pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level and applied to every stage. Unknown keys fall through to `engine_extras:` and are forwarded to the engine.
|
||||
@@ -88,6 +91,57 @@ stages:
|
||||
| `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. |
|
||||
| `--stage-configs-path` | **Deprecated.** Accepts legacy `stage_args` yamls and (auto-detected) new deploy yamls; emits a deprecation warning. Migrate to `--deploy-config`. To be removed in a follow-up PR. |
|
||||
|
||||
### Stage-Based CLI Paradigm
|
||||
|
||||
The stage-based CLI paradigm facilitates the execution of discrete pipeline stages within isolated processes:
|
||||
|
||||
- **Stage 0** typically encapsulates the orchestrator and the primary API server. Invocation requires `--stage-id 0`,
|
||||
`--omni-master-address`, `--omni-master-port`, and standard port declarations (e.g., `--port`).
|
||||
- **Worker Stages** operate without a distinct API server (i.e., using `--headless`), are assigned sequential `--stage-id` identifiers, and must reference the corresponding
|
||||
`--omni-master-address` and `--omni-master-port` parameters to successfully register with Stage 0.
|
||||
|
||||
For migrated architectures, the system automatically resolves and loads the bundled deployment YAML. Consequently, the primary execution path
|
||||
does **not** necessitate the explicit definition of `--deploy-config`:
|
||||
the example below uses `CUDA_VISIBLE_DEVICES=0` for Stage 0 and
|
||||
`CUDA_VISIBLE_DEVICES=1` for Stage 1.
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--port 8091 \
|
||||
--stage-id 0 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
When instantiating a custom deployment YAML conforming to the updated schema, append the `--deploy-config /path/to/override.yaml` directive
|
||||
to all node invocations. For legacy architectures (e.g., BAGEL) configured via deprecated `stage_args:` schemas, continue to specify the relevant configuration via `--stage-configs-path /path/to/config.yaml`.
|
||||
|
||||
In the context of standard initialization architectures, utilizing the `--stage-overrides` parameter operates as the optimal methodology
|
||||
for delineating stage-specific tuning from the CLI interface:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
|
||||
```
|
||||
|
||||
Conversely, in the context of the **stage-based CLI** paradigm, given that each execution process exclusively instantiates a single pipeline stage, configuration override attributes
|
||||
can be defined uniformly via explicit CLI flags on the corresponding instantiation command, rendering composite `--stage-overrides` JSON strings unnecessary:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--gpu-memory-utilization 0.5 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
### Precedence
|
||||
|
||||
From highest to lowest:
|
||||
@@ -133,6 +187,17 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--stage-overrides '{"0": {"max_num_seqs": 8}}'
|
||||
```
|
||||
|
||||
Within the stage-based CLI paradigm, equivalent configuration parameters can inherently be passed directly
|
||||
as command-line arguments to the designated single-stage process instantiation:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 0 \
|
||||
--max-num-seqs 8 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
Effective config per stage after the merge:
|
||||
|
||||
| Stage | Field | Final value | Source |
|
||||
@@ -153,9 +218,14 @@ Therefore, as a core part of vLLM-Omni, the stage configs for a model have sever
|
||||
- Input and output dependencies for each stage.
|
||||
- Default input parameters.
|
||||
|
||||
If users want to modify some part of it. The custom stage_configs file can be input as input argument in both online and offline. Just like examples below:
|
||||
To override specific parameters, explicitly inject the customized configuration schema
|
||||
in both online and offline instantiation flows. Prioritize the `--deploy-config` flag
|
||||
when loading the new-schema deploy YAML schemas, reserving the `--stage-configs-path` parameter
|
||||
exclusively to maintain compatibility with legacy `stage_args` YAML constructs.
|
||||
|
||||
For offline (Assume necessary dependencies have ben imported):
|
||||
Examples:
|
||||
|
||||
For offline (Assume necessary dependencies have been imported):
|
||||
```python
|
||||
model_name = "Qwen/Qwen2.5-Omni-7B"
|
||||
omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
|
||||
@@ -163,7 +233,13 @@ omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.
|
||||
|
||||
For online serving:
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
|
||||
```
|
||||
|
||||
Legacy online serving:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
|
||||
```
|
||||
!!! important
|
||||
We are actively iterating on the definition of stage configs, and we welcome all feedbacks from both community users and developers to help us shape the development!
|
||||
|
||||
@@ -105,9 +105,11 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<td> Depends on reality </td>
|
||||
<td>
|
||||
<strong>Stability:</strong><br>
|
||||
/tests/dfx/stability/tests/test.json<br>
|
||||
/tests/dfx/stability/tests/test_qwen3_omni.json<br>
|
||||
/tests/dfx/stability/tests/test_wan22.json<br>
|
||||
<strong>Reliability:</strong><br>
|
||||
tests/e2e/reliability/test_{model_name}.py
|
||||
tests/dfx/reliability/test_reliability_{model_key}.py<br>
|
||||
(e.g. <code>test_reliability_qwen3_omni.py</code>, <code>test_reliability_wan22.py</code>)
|
||||
</td>
|
||||
<td>
|
||||
<a href="#chapter-4-l5-level-testing-stability-and-reliability-testing">Chapter 4</a><br>
|
||||
@@ -150,6 +152,8 @@ 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/
|
||||
@@ -230,15 +234,13 @@ 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_t2i_model.py
|
||||
├── test_t2v_model.py
|
||||
├── test_ovis_image.py
|
||||
├── test_z_image.py
|
||||
├── test_wan22.py
|
||||
├── test_zimage_tensor_parallel.py
|
||||
├── test_cache_dit.py
|
||||
├── test_teacache.py
|
||||
@@ -418,13 +420,13 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
|
||||
**Explanation**:
|
||||
|
||||
@pytest.mark.advanced_model: Marks the test as L3 or L4 level, indicating that this test case performs deep validation, using real models for performance, integration, and accuracy testing. This forms a "basic-advanced" correspondence with the core_model mark at the L2 level.
|
||||
@pytest.mark.advanced_model: Marks the test as L3 merge level, indicating deep validation with real models. @pytest.mark.full_model: Marks L4 nightly-only suites (e.g. `test_*_expansion.py`, doc examples).
|
||||
|
||||
@pytest.mark.core_model: Marks the test as L1 or L2 level, indicating that this test case validates the basic functionality of the core model. It uses mock weights and only checks if the relevant interface functions correctly.
|
||||
|
||||
@pytest.mark.parametrize: A parameterization decorator that allows abstracting test data into parameters, enabling reuse of the same test logic across different data configurations. indirect=True indicates that parameters will be passed to the fixture for processing.
|
||||
|
||||
**Notes**: If you believe the test case only needs to execute basic run logic at the PR-level CI, you can mark it only with @pytest.mark.core_model. If you believe it only needs to execute deep validation run logic at the merge or nightly level, you can mark it only with @pytest.mark.advanced_model. If you believe the test case needs to accommodate both basic run and deep validation test logic, you should mark it with both @pytest.mark.core_model and @pytest.mark.advanced_model.
|
||||
**Notes**: If you believe the test case only needs to execute basic run logic at the PR-level CI, you can mark it only with @pytest.mark.core_model. If you believe it only needs to execute deep validation at merge (L3), use @pytest.mark.advanced_model. For L4 nightly-only expansion and doc-example tests, use @pytest.mark.full_model with `--run-level full_model`. If the test case needs both basic run and deep validation, mark with @pytest.mark.core_model and the appropriate L3/L4 marker (`advanced_model` and/or `full_model`).
|
||||
|
||||
**2.4.2 Test Function Definition and Documentation**
|
||||
|
||||
@@ -516,9 +518,11 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
|
||||
**Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n.
|
||||
|
||||
**Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
|
||||
**Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model and full_model perform deep validation.
|
||||
|
||||
- ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
|
||||
- ***Run Command (L3 merge)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
|
||||
|
||||
- ***Run Command (L4 nightly expansion)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}_expansion.py -m full_model --run-level=full_model`
|
||||
|
||||
## Chapter 3: L4 Level Testing - Full Functionality, Performance, and Documentation Testing
|
||||
|
||||
@@ -572,13 +576,18 @@ L5 level testing focuses on the performance of model services under ***long-runn
|
||||
|
||||
### 4.2 Testing Content and Scope
|
||||
|
||||
- ***Long-term Stability (Stability) Testing***: Uses the `tests/dfx/stability/tests/test.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
|
||||
- ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
|
||||
- ***Long-term Stability (Stability) Testing***: Uses JSON under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` and `test_wan22.json`) to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
|
||||
- ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni or `multiprocessing.spawn` for Wan2.2 video workers), **large multimodal chat** or **`/v1/videos`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See the Reliability `<details>` block in Section 4.4 for file-level responsibilities and CI markers (`slow`, `hardware_test`, POSIX-only kill).
|
||||
|
||||
### 4.3 Test Directory and Execution Files
|
||||
|
||||
- ***Stability Test Configuration***: `tests/dfx/stability/tests/test.json`
|
||||
- ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
|
||||
- ***Stability Test Configuration***: `tests/dfx/stability/tests/test_qwen3_omni.json`, `tests/dfx/stability/tests/test_wan22.json` (one JSON per model / runner family)
|
||||
- ***Reliability Test Suite*** (`tests/dfx/reliability/`):
|
||||
- `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default).
|
||||
- `test_reliability_wan22.py` — Wan2.2 T2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
|
||||
- `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`).
|
||||
- `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready.
|
||||
- `README.md` — Short local run commands for this directory.
|
||||
|
||||
### 4.4 Execution Method and Example
|
||||
|
||||
@@ -588,7 +597,7 @@ L5 level testing focuses on the performance of model services under ***long-runn
|
||||
<details>
|
||||
<summary> Test Examples</summary>
|
||||
|
||||
When you want to add L5-level stability test cases, you can refer to the following format for case addition in `tests/dfx/stability/tests/test.json`:
|
||||
When you want to add L5-level stability test cases, add or extend the appropriate JSON file under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` for Omni bench traffic, or `test_wan22.json` for diffusion `/v1/videos` workloads). The following illustrates the Qwen3-Omni shape:
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -659,8 +668,57 @@ All other optional parameters follow the same rules as the in Chapter 3.4.
|
||||
|
||||
</details>
|
||||
|
||||
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_{model_name}.py`
|
||||
- ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py`
|
||||
<details>
|
||||
<summary> Reliability test suite (<code>tests/dfx/reliability</code>)</summary>
|
||||
|
||||
#### Purpose and relationship to stability
|
||||
|
||||
Reliability tests are **short fault-injection** integration runs (L5 **(b)** in `tests/dfx/reliability/README.md`). They complement **stability** JSON-driven long runs: instead of hours of steady traffic, they **perturb** the server (GPU OOM sidecar, fatal signals on selected processes) and check **failure mode** and **latency bounds** (e.g. chat or `/v1/videos` must not hang under concurrent fault-time load).
|
||||
|
||||
#### Directory layout
|
||||
|
||||
| Path | Responsibility |
|
||||
| ---- | -------------- |
|
||||
| `helpers.py` | Shared helpers used by current reliability suites: raw `POST`/`GET` probes (`/v1/chat/completions`, `/health`), OpenAI error parsing (`extract_openai_error_contract_from_bytes`), GPU OOM sidecar lifecycle (`inject_gpu_oom`, `stop_gpu_oom_hogs`), and process-kill injector builder (`make_process_kill_fault_injector`). |
|
||||
| `conftest.py` | Pytest fixtures: indirect `fault_injector`, `omni_server_after_fault` / `omni_server_after_fault_function` (run injector after server is ready, then yield server). |
|
||||
| `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). |
|
||||
| `test_reliability_wan22.py` | Wan2.2 T2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
|
||||
| `README.md` | Minimal run / collect examples. |
|
||||
|
||||
#### Parametrization and markers
|
||||
|
||||
- Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`.
|
||||
- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni paths often require **2** cards; Wan2.2 video paths **1** card).
|
||||
- **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**.
|
||||
|
||||
#### CI trigger
|
||||
|
||||
Weekly Buildkite (`.buildkite/test-weekly.yml`) runs, for example:
|
||||
|
||||
```bash
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
|
||||
```
|
||||
|
||||
#### Local commands
|
||||
|
||||
```bash
|
||||
pytest --collect-only tests/dfx/reliability
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
|
||||
```
|
||||
|
||||
#### Adding a new model suite
|
||||
|
||||
1. Add `test_reliability_<model>.py` under `tests/dfx/reliability/`.
|
||||
2. Define **`RELIABILITY_SCENARIOS`** and pass them through **`create_reliability_omni_server_params()`** with the correct deploy or e2e stage-config directory (same pattern as existing files).
|
||||
3. Reuse **`helpers`** for OOM / kill / raw HTTP; prefer **`assert_fault_exception()`** and **`resolve_oom_device_spec()`** from `tests/dfx/conftest.py` for consistent device selection vs stage YAML.
|
||||
4. Register **`slow`** (and **`hardware_test`** if needed); extend **`.buildkite/test-weekly.yml`** when the suite should run in weekly L5.
|
||||
|
||||
</details>
|
||||
|
||||
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
|
||||
- ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow` and/or `pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow` (add `test_reliability_<suite>.py` for new models)
|
||||
|
||||
## Summary
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
**Scope**
|
||||
|
||||
For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
|
||||
For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/).
|
||||
|
||||
**Test Case Design**
|
||||
|
||||
@@ -37,7 +37,7 @@ Currently all the features are available in online serving mode. Hence, only nee
|
||||
**Code Style**
|
||||
|
||||
- Validation: test that the multimodal output files of your model have the correct shapes. `OpenAIClientHandler.send_diffusion_request` should have taken care of this.
|
||||
- Test marks: always add `advanced_model` and `diffusion`. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/).
|
||||
- Test marks: always add `full_model` and `diffusion` for L4 nightly `test_*_expansion.py` cases. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/).
|
||||
- To maximize code reuse, you may refer to
|
||||
- `tests/conftest.py` for `omni_server` (running server in subprocess) and `openai_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper.
|
||||
- `tests/helpers/mark.py` for `@hardware_test(...)` and `hardware_marks`.
|
||||
|
||||
@@ -42,31 +42,69 @@ Our test scripts use the pytest framework. First, please use `git clone https://
|
||||
```
|
||||
The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
|
||||
|
||||
=== "L3 level & L4 level"
|
||||
=== "L3 level"
|
||||
|
||||
```bash
|
||||
pytest -s -v -m "advanced_model" --run-level=advanced_model
|
||||
```
|
||||
If you only want to run L3 test case, you can use:
|
||||
```bash
|
||||
pytest -s -v e2e/ --ignore-glob='*expansion.py' -m "advanced_model" --run-level=advanced_model
|
||||
```
|
||||
If you only want to run a specific test case, you can use:
|
||||
```bash
|
||||
pytest -s -v test_xxxx.py --run-level=advanced_model
|
||||
```
|
||||
If you only want to run specific test cases on a particular platform, you can use:
|
||||
```bash
|
||||
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
|
||||
pytest -s -v -m "advanced_model and distributed_cuda and L4" --run-level=advanced_model
|
||||
```
|
||||
The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
|
||||
|
||||
|
||||
=== "L4 level"
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v -m "full_model" --run-level=full_model
|
||||
```
|
||||
If you only want to run a specific test case, you can use:
|
||||
```bash
|
||||
pytest -s -v test_xxxx.py --run-level=full_model
|
||||
```
|
||||
If you only want to run specific test cases on a particular platform, you can use:
|
||||
```bash
|
||||
pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model
|
||||
```
|
||||
Note: To run performance tests (defaults to ``test_qwen_omni.json``; use ``--test-config-file tests/dfx/perf/tests/test_tts.json`` for TTS):
|
||||
```bash
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
|
||||
```
|
||||
The latest L4 (nightly) test commands use the `full_model` marker and `--run-level full_model` (see [test-nightly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml) and [test-nightly-diffusion.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly-diffusion.yml)). Example:
|
||||
|
||||
The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v -m "full_model and (omni or tts) and H100" --run-level=full_model
|
||||
```
|
||||
|
||||
The latest L4 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml).
|
||||
=== "L5 level"
|
||||
|
||||
L5 includes stability and reliability testing. Typical commands:
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
|
||||
# Stability: Qwen3-Omni
|
||||
pytest -s -v dfx/stability/scripts/test_stability_qwen3_omni.py -m slow
|
||||
|
||||
# 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
|
||||
|
||||
```
|
||||
|
||||
The latest L5 commands for CI can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
|
||||
|
||||
You can find more information about markers in the documentation: [marker doc](./tests_markers.md)
|
||||
|
||||
|
||||
@@ -8,9 +8,11 @@ Defined in `pyproject.toml`:
|
||||
| Marker | Description |
|
||||
| ------------------ | --------------------------------------------------------- |
|
||||
| `core_model` | L1&L2 tests (run in each PR) |
|
||||
| `advanced_model` | L3&L4 level tests (run in each merge or nightly) |
|
||||
| `advanced_model` | L3 tests (run on each merge to main) |
|
||||
| `full_model` | L4 tests (run nightly) |
|
||||
| `diffusion` | Diffusion model tests |
|
||||
| `omni` | Omni model tests |
|
||||
| `omni` | Omni multimodal model tests |
|
||||
| `tts` | TTS model tests |
|
||||
| `cache` | Cache backend tests |
|
||||
| `parallel` | Parallelism/distributed tests |
|
||||
| `cpu` | Tests that run on CPU |
|
||||
|
||||
@@ -54,6 +54,8 @@ 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/
|
||||
@@ -134,15 +136,13 @@ 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_t2i_model.py
|
||||
├── test_t2v_model.py
|
||||
├── test_ovis_image.py
|
||||
├── test_z_image.py
|
||||
├── test_wan22.py
|
||||
├── test_zimage_tensor_parallel.py
|
||||
├── test_cache_dit.py
|
||||
├── test_teacache.py
|
||||
@@ -213,7 +213,6 @@ 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,17 +115,20 @@ class YourAttentionBlock(nn.Module):
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
|
||||
class YourAttentionBlock(nn.Module):
|
||||
class YourSelfAttentionBlock(nn.Module):
|
||||
def __init__(self, ...):
|
||||
super().__init__()
|
||||
|
||||
# Initialize vLLM-Omni's Attention layer
|
||||
# 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).
|
||||
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, ...):
|
||||
@@ -133,7 +136,6 @@ class YourAttentionBlock(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:**
|
||||
@@ -141,8 +143,35 @@ class YourAttentionBlock(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)
|
||||
|
||||
**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.
|
||||
**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.
|
||||
|
||||
#### 1.3: Replace Imports and Utilities
|
||||
|
||||
@@ -737,6 +766,16 @@ step-scoped phases. The reference implementation is
|
||||
Do not enable `step_execution=True` until those four methods are implemented
|
||||
and validated against the request-level path.
|
||||
|
||||
If you want the pipeline to work with the experimental batched step-wise path
|
||||
(`max_num_seqs > 1`), also see:
|
||||
[Continuous Batching for Step-Wise Diffusion](../../design/feature/diffusion_continuous_batching.md).
|
||||
|
||||
If you expose this in example scripts or recipes, keep it opt-in. Surface
|
||||
runtime features like `step_execution` as optional flags instead of silently
|
||||
turning them on. For Qwen-Image-style serving examples, document
|
||||
`--step-execution` as the feature gate and `--max-num-seqs N` as the
|
||||
companion batching knob.
|
||||
|
||||
### Cache Acceleration
|
||||
|
||||
#### TeaCache
|
||||
@@ -813,16 +852,16 @@ class WanTransformer3DModel(nn.Module):
|
||||
|
||||
---
|
||||
|
||||
### Diffusion Timing (Performance Profiling)
|
||||
### Diffusion Pipeline Profiler (Performance Profiling)
|
||||
When adapting a new diffusion model, it is often useful to analyze the latency of key components such as text encoding, diffusion denoising, and VAE decoding.
|
||||
vLLM-Omni provides a timing utility via `DiffusionPipelineProfilerMixin` to help developers quickly identify performance bottlenecks.
|
||||
|
||||
!!! info
|
||||
`DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md#3-profiling-diffusion-models). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
|
||||
`DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
|
||||
|
||||
This tool automatically measures the execution time of selected pipeline modules and prints the results in the logs.
|
||||
|
||||
**Enabling Diffusion Timing**
|
||||
**Enabling Diffusion Pipeline Profiler**
|
||||
|
||||
|
||||
Enable timing by setting:
|
||||
@@ -843,7 +882,7 @@ If not specified, the default targets are used:
|
||||
**Adding DiffusionPipelineProfilerMixin to a Pipeline**
|
||||
To enable timing support in your pipeline, inherit from DiffusionPipelineProfilerMixin.
|
||||
```python
|
||||
from vllm_omni.diffusion.utils.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
|
||||
from vllm_omni.diffusion.profiler import DiffusionPipelineProfilerMixin
|
||||
|
||||
class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
# Optional: Specify custom timing targets
|
||||
@@ -862,7 +901,9 @@ class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
...
|
||||
|
||||
# initialize timing profiler
|
||||
self.setup_diffusion_pipeline_profiler(enable_diffusion_pipeline_profiler)
|
||||
self.setup_diffusion_pipeline_profiler(
|
||||
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
|
||||
)
|
||||
```
|
||||
The mixin dynamically wraps selected methods and records their execution time during inference.
|
||||
|
||||
@@ -906,9 +947,9 @@ tokenizer.forward
|
||||
|
||||
When enabled, timing logs appear like this:
|
||||
```
|
||||
[DiffusionTiming] text_encoder.forward took 0.018s
|
||||
[DiffusionTiming] diffuse took 2.412s
|
||||
[DiffusionTiming] vae.decode took 0.063s
|
||||
[DiffusionPipelineProfiler] text_encoder.forward took 0.018s
|
||||
[DiffusionPipelineProfiler] diffuse took 2.412s
|
||||
[DiffusionPipelineProfiler] vae.decode took 0.063s
|
||||
```
|
||||
These measurements help identify bottlenecks during model adaptation and optimization
|
||||
|
||||
|
||||
@@ -1,478 +0,0 @@
|
||||
# Support FP8 Quantization
|
||||
|
||||
This section describes how to add FP8 quantization support to a diffusion model. We use the Qwen-Image pipeline as the reference implementation.
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Step-by-Step Implementation](#step-by-step-implementation)
|
||||
- [Testing](#testing)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
- [Reference Implementations](#reference-implementations)
|
||||
- [Summary](#summary)
|
||||
|
||||
---
|
||||
|
||||
## Overview
|
||||
|
||||
### What is FP8 Quantization?
|
||||
|
||||
FP8 quantization reduces model memory footprint by storing weights in 8-bit floating point (`float8_e4m3fn`) instead of 16-bit (`bfloat16`/`float16`). This saves ~50% of weight memory.
|
||||
|
||||
A diffusion pipeline typically has three quantizable components:
|
||||
|
||||
| Component | Layer Types | Quantization Approach |
|
||||
|-----------|-----------|----------------------|
|
||||
| **DiT (transformer)** | `nn.Linear` | Replace with vLLM quantized linear layers |
|
||||
| **Text encoder** | `nn.Linear` | FP8 weight storage via hooks |
|
||||
| **VAE** | `nn.Conv2d`, `nn.Conv3d` | FP8 weight storage via hooks |
|
||||
|
||||
### Architecture
|
||||
|
||||
The quantization system has two mechanisms:
|
||||
|
||||
**1. vLLM Quantized Linear Layers (for DiT)**
|
||||
|
||||
vLLM provides drop-in replacements for `nn.Linear` that perform FP8 W8A8 computation on supported hardware. These are configured via `quant_config` passed to layer constructors.
|
||||
|
||||
Key classes:
|
||||
|
||||
| Class | Location | Purpose |
|
||||
|-------|----------|---------|
|
||||
| `build_quant_config` | `vllm_omni/quantization/__init__.py` | Build quantization config from string/dict/per-component spec |
|
||||
| `ComponentQuantizationConfig` | `vllm_omni/quantization/component_config.py` | Per-component routing by layer prefix |
|
||||
|
||||
**2. FP8 Weight Storage via Hooks (for text encoder and VAE)**
|
||||
|
||||
For models loaded via `from_pretrained()` (outside vLLM's weight pipeline), we use forward hooks to store weights in FP8 and dequantize on-the-fly:
|
||||
|
||||
```
|
||||
At rest: weight.data = fp8_tensor (half memory)
|
||||
Pre-hook: weight.data = fp8_weight.to(bf16) * scale
|
||||
Forward: Normal computation in BF16
|
||||
Post-hook: weight.data = fp8_weight (back to FP8)
|
||||
```
|
||||
|
||||
Key function:
|
||||
|
||||
| Function | Location | Purpose |
|
||||
|----------|----------|---------|
|
||||
| `apply_fp8_weight_storage` | `vllm_omni/diffusion/models/utils.py` | Apply FP8 hooks to Linear/Conv layers (introduced in [#1412](https://github.com/vllm-project/vllm-omni/pull/1412)) |
|
||||
|
||||
---
|
||||
|
||||
## Step-by-Step Implementation
|
||||
|
||||
### Step 1: Add `quant_config` to the Transformer
|
||||
|
||||
Pass `quant_config` to vLLM parallel linear layers in the transformer. This enables FP8 W8A8 for DiT linear layers.
|
||||
|
||||
#### 1.1: Accept `quant_config` in Constructor
|
||||
|
||||
```python
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
||||
|
||||
class YourTransformer2DModel(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
od_config: OmniDiffusionConfig,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
# ... other params
|
||||
):
|
||||
self.quant_config = quant_config
|
||||
# Pass to sub-modules
|
||||
self.blocks = nn.ModuleList([
|
||||
YourTransformerBlock(quant_config=quant_config, ...)
|
||||
for _ in range(num_layers)
|
||||
])
|
||||
```
|
||||
|
||||
#### 1.2: Pass `quant_config` to vLLM Linear Layers
|
||||
|
||||
Replace `nn.Linear` with vLLM parallel linear layers that accept `quant_config`:
|
||||
|
||||
```python
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
|
||||
class YourAttentionBlock(nn.Module):
|
||||
def __init__(self, hidden_size, num_heads, quant_config=None, prefix=""):
|
||||
# QKV projection
|
||||
self.to_qkv = QKVParallelLinear(
|
||||
hidden_size=hidden_size,
|
||||
head_size=hidden_size // num_heads,
|
||||
total_num_heads=num_heads,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.to_qkv",
|
||||
)
|
||||
# Output projection
|
||||
self.to_out = RowParallelLinear(
|
||||
input_size=hidden_size,
|
||||
output_size=hidden_size,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.to_out",
|
||||
)
|
||||
|
||||
class YourFeedForward(nn.Module):
|
||||
def __init__(self, hidden_size, intermediate_size, quant_config=None, prefix=""):
|
||||
self.gate_up_proj = ColumnParallelLinear(
|
||||
input_size=hidden_size,
|
||||
output_size=intermediate_size * 2,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.gate_up_proj",
|
||||
)
|
||||
self.down_proj = RowParallelLinear(
|
||||
input_size=intermediate_size,
|
||||
output_size=hidden_size,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.down_proj",
|
||||
)
|
||||
```
|
||||
|
||||
!!! important
|
||||
Always pass the `prefix` parameter. It is used by `ignored_layers` matching to skip sensitive layers.
|
||||
|
||||
#### 1.3: Identify Sensitive Layers
|
||||
|
||||
Some layers degrade output quality when quantized. Common sensitive layers:
|
||||
|
||||
| Layer | Why Sensitive | Typical Models | Quant Methods Affected |
|
||||
|-------|-------------|---------------|----------------------|
|
||||
| `img_mlp` | Processes denoising latents with shifting dynamic range | Qwen-Image | FP8, Int8 |
|
||||
| `feed_forward` | FFN layers in DiT blocks, large dynamic range | Z-Image | Int8 |
|
||||
| `proj_out` | Final output projection, small quantization errors amplified | Various | FP8, Int8 |
|
||||
| `lm_head` | Output vocabulary projection, precision-critical | Omni (thinker) | NVFP4 |
|
||||
| `mlp.gate` | MoE router gates, routing decisions are precision-critical | MoE models | NVFP4 |
|
||||
|
||||
**Quantization quality benchmarks** (LPIPS: lower = better, <0.01 imperceptible, >0.1 noticeable):
|
||||
|
||||
| Model | Method | All Layers | With `ignored_layers` | Recommendation |
|
||||
|-------|--------|-----------|----------------------|----------------|
|
||||
| Qwen-Image-2512 | Int8 | LPIPS 0.0197 | 0.0027 (skip `img_mlp`) | Skip `img_mlp` |
|
||||
| Z-Image-Turbo | Int8 | LPIPS 0.1597 | 0.0290 (skip `feed_forward`) | Skip `feed_forward` |
|
||||
| Z-Image-Turbo | FP8 | LPIPS ~0.005 | — | All layers OK |
|
||||
|
||||
To identify sensitive layers for a new model:
|
||||
|
||||
1. Run inference with all layers quantized
|
||||
2. Compare output quality with BF16 baseline using LPIPS (see `benchmarks/diffusion/quantization_quality.py`)
|
||||
3. Selectively disable quantization per layer (via `ignored_layers`) until quality matches
|
||||
4. Document the recommended `ignored_layers` in the model's supported models table
|
||||
|
||||
### Step 2: Wire Up Quantization in the Pipeline
|
||||
|
||||
#### 2.1: Create `quant_config` and Pass to Transformer
|
||||
|
||||
In the pipeline's `__init__`, extract the vLLM quant config and pass it to the transformer:
|
||||
|
||||
```python
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
class YourPipeline(nn.Module):
|
||||
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
# ... load scheduler, text_encoder, vae ...
|
||||
|
||||
# Build quant config from OmniDiffusionConfig
|
||||
quant_config = build_quant_config(od_config.quantization_config)
|
||||
|
||||
# Pass to transformer
|
||||
self.transformer = YourTransformer2DModel(
|
||||
od_config=od_config,
|
||||
quant_config=quant_config,
|
||||
# ... other kwargs
|
||||
)
|
||||
```
|
||||
|
||||
#### 2.2: Apply FP8 Weight Storage to Text Encoder and VAE
|
||||
|
||||
For models loaded via `from_pretrained()`, apply hook-based FP8 weight storage.
|
||||
|
||||
!!! note
|
||||
The `apply_fp8_weight_storage` function is provided by `vllm_omni/diffusion/models/utils.py`, introduced in [PR #1412](https://github.com/vllm-project/vllm-omni/pull/1412).
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.models.utils import apply_fp8_weight_storage
|
||||
|
||||
class YourPipeline(nn.Module):
|
||||
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
|
||||
self.text_encoder = SomeTextEncoder.from_pretrained(...)
|
||||
self.vae = SomeVAE.from_pretrained(...).to(self.device)
|
||||
|
||||
# Apply FP8 weight storage to text encoder and VAE
|
||||
if (
|
||||
od_config.quantization_config is not None
|
||||
and getattr(od_config.quantization_config, "quant_method", None) == "fp8"
|
||||
):
|
||||
apply_fp8_weight_storage(self.vae)
|
||||
apply_fp8_weight_storage(self.text_encoder)
|
||||
|
||||
# ... rest of init
|
||||
```
|
||||
|
||||
`apply_fp8_weight_storage` walks all `nn.Linear`, `nn.Conv2d`, and `nn.Conv3d` modules in the model and applies FP8 hooks to each.
|
||||
|
||||
#### 2.3: Fix `load_weights` to Include Pre-loaded Components
|
||||
|
||||
Since text encoders and VAEs are loaded via `from_pretrained()` (not through vLLM's weight pipeline), their parameters won't appear in the `loaded_weights` set. Mark them explicitly:
|
||||
|
||||
```python
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
loader = AutoWeightsLoader(self)
|
||||
loaded_weights = loader.load_weights(weights)
|
||||
# Mark pre-loaded component weights as loaded
|
||||
loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
|
||||
loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
|
||||
return loaded_weights
|
||||
```
|
||||
|
||||
### Step 3: Add to Supported Models Documentation
|
||||
|
||||
Update the FP8 supported models table in `docs/user_guide/diffusion/quantization/fp8.md`:
|
||||
|
||||
```markdown
|
||||
| Your-Model | `org/Your-Model` | ✅ | ✅ | ✅ | `sensitive_layer` |
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Testing
|
||||
|
||||
### Quick Verification
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
# Test with FP8 enabled
|
||||
omni = Omni(model="your-model", quantization="fp8")
|
||||
|
||||
outputs = omni.generate(
|
||||
"A cat sitting on a windowsill",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=50),
|
||||
)
|
||||
```
|
||||
|
||||
### Quality Gate Test (LPIPS)
|
||||
|
||||
We provide a pytest-based quality gate at `tests/diffusion/quantization/test_quantization_quality.py`.
|
||||
It generates outputs with both BF16 and your quantized method using the same seed, computes
|
||||
[LPIPS](https://github.com/richzhang/PerceptualSimilarity) perceptual distance, and **fails if it
|
||||
exceeds a threshold**. This is the recommended way to validate that a quantization method does not
|
||||
unacceptably degrade output quality.
|
||||
|
||||
**Step 1: Add your test config** — edit `QUALITY_CONFIGS` in the test file:
|
||||
|
||||
```python
|
||||
QualityTestConfig(
|
||||
id="int8_z_image", # pytest ID
|
||||
model="Tongyi-MAI/Z-Image-Turbo",
|
||||
quantization="int8", # your quantization method
|
||||
task="t2i", # "t2i" or "t2v"
|
||||
prompt="a cup of coffee on a wooden table, morning light",
|
||||
max_lpips=0.15, # fail threshold (see guide below)
|
||||
num_inference_steps=20,
|
||||
),
|
||||
```
|
||||
|
||||
**Step 2: Run the test:**
|
||||
|
||||
```bash
|
||||
# Install dependency
|
||||
pip install lpips
|
||||
|
||||
# Run all quality tests
|
||||
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m ""
|
||||
|
||||
# Run only your method
|
||||
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "int8"
|
||||
|
||||
# Run a specific model
|
||||
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image"
|
||||
```
|
||||
|
||||
**Step 3: Read the output:**
|
||||
|
||||
```
|
||||
Quantization Quality: int8_z_image
|
||||
Model: Tongyi-MAI/Z-Image-Turbo
|
||||
Method: int8
|
||||
LPIPS: 0.0290 (threshold: 0.15)
|
||||
BF16 memory: 19.15 GiB
|
||||
Quant memory: 14.80 GiB (23% reduction)
|
||||
Result: PASS
|
||||
```
|
||||
|
||||
**LPIPS threshold guide:**
|
||||
|
||||
| LPIPS Range | Meaning | Recommendation |
|
||||
|-------------|---------|----------------|
|
||||
| < 0.01 | Imperceptible | Ideal — no visible difference |
|
||||
| 0.01 – 0.05 | Minor differences | Good — acceptable for most use cases |
|
||||
| 0.05 – 0.15 | Noticeable on close inspection | Acceptable with `ignored_layers` tuning |
|
||||
| 0.15 – 0.30 | Clearly visible | Investigate — likely needs `ignored_layers` |
|
||||
| > 0.30 | Significant degradation | Quantization method may not suit this model |
|
||||
|
||||
!!! tip
|
||||
If your test fails, don't raise the threshold — instead, identify sensitive layers and add
|
||||
them to `ignored_layers` (see next section). Then re-run the test with the lower LPIPS.
|
||||
|
||||
### Quality Benchmark (Detailed)
|
||||
|
||||
For a full multi-prompt benchmark with Markdown output (suitable for PR descriptions), use the
|
||||
standalone benchmark script:
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/quantization_quality.py \
|
||||
--model Tongyi-MAI/Z-Image-Turbo \
|
||||
--quantization fp8 \
|
||||
--prompts \
|
||||
"a cup of coffee on the table" \
|
||||
"an aerial view of a coral reef" \
|
||||
--height 1024 --width 1024 \
|
||||
--num-inference-steps 50 --seed 42
|
||||
```
|
||||
|
||||
This generates BF16 and quantized outputs, computes per-prompt LPIPS, and prints a Markdown
|
||||
summary table ready to paste into your PR.
|
||||
|
||||
### Identifying Sensitive Layers with `ignored_layers`
|
||||
|
||||
Not all layers tolerate quantization equally. **`ignored_layers`** lets you keep specific layers
|
||||
in BF16 while quantizing everything else. This is critical for achieving good quality with
|
||||
aggressive quantization methods (Int8, NVFP4).
|
||||
|
||||
**Why some layers are sensitive:**
|
||||
|
||||
| Layer Type | Why Sensitive | Typical Impact |
|
||||
|------------|--------------|----------------|
|
||||
| `img_mlp` | Processes denoising latents with shifting dynamic range across timesteps. No built-in normalization to absorb quantization error. | Color shifts, blurriness |
|
||||
| `feed_forward` | FFN layers with large dynamic range in DiT blocks. | Artifacts, loss of detail |
|
||||
| `proj_out` | Final output projection — small errors get amplified since there are no subsequent layers to correct them. | Overall quality loss |
|
||||
| `lm_head` | Vocabulary projection in LLMs — precision-critical for token selection. | Garbage text output |
|
||||
| `mlp.gate` | MoE router gates — routing decisions are binary and precision-critical. | Wrong expert selection, quality collapse |
|
||||
|
||||
**How to find sensitive layers for your model:**
|
||||
|
||||
1. Run the quality gate test with all layers quantized:
|
||||
```bash
|
||||
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "your_model"
|
||||
```
|
||||
|
||||
2. If LPIPS exceeds your threshold, try skipping common sensitive layers one at a time:
|
||||
```python
|
||||
QualityTestConfig(
|
||||
id="fp8_your_model_skip_mlp",
|
||||
model="your-org/Your-Model",
|
||||
quantization="fp8",
|
||||
task="t2i",
|
||||
prompt="a cup of coffee on a wooden table",
|
||||
max_lpips=0.05,
|
||||
),
|
||||
```
|
||||
|
||||
And in the test, modify the Omni call:
|
||||
```python
|
||||
omni_qt = Omni(
|
||||
model=config.model,
|
||||
quantization_config={
|
||||
"method": config.quantization,
|
||||
"ignored_layers": ["img_mlp"],
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
3. Compare LPIPS with and without `ignored_layers` to isolate the problematic layers.
|
||||
|
||||
4. Document the recommended `ignored_layers` in the supported models table.
|
||||
|
||||
**Real-world examples:**
|
||||
|
||||
| Model | Method | All Layers (LPIPS) | With `ignored_layers` | Which Layers to Skip |
|
||||
|-------|--------|-------------------|----------------------|---------------------|
|
||||
| Qwen-Image-2512 | Int8 | 0.0197 | 0.0027 | `img_mlp` |
|
||||
| Z-Image-Turbo | Int8 | 0.1597 | 0.0290 | `feed_forward` |
|
||||
| Z-Image-Turbo | FP8 | ~0.005 | — | None needed |
|
||||
|
||||
### What to Verify (Checklist)
|
||||
|
||||
1. **Model loads without errors** — Check logs for quantization initialization messages
|
||||
2. **Quality gate passes** — LPIPS within threshold vs BF16 baseline
|
||||
3. **Memory reduction** — Expect ~30-50% savings on quantized components
|
||||
4. **Sensitive layers identified** — Document which layers (if any) need `ignored_layers`
|
||||
5. **`ignored_layers` pattern works** — Verify that the layer name prefix matches your transformer's module names
|
||||
|
||||
---
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Issue: Output Quality Degradation
|
||||
|
||||
**Symptoms:** Generated images have visible artifacts, color shifts, or blurriness compared to BF16 baseline.
|
||||
|
||||
**Causes & Solutions:**
|
||||
|
||||
- **Sensitive layers being quantized:**
|
||||
|
||||
**Solution:** Identify and skip sensitive layers:
|
||||
```python
|
||||
omni = Omni(
|
||||
model="your-model",
|
||||
quantization_config={
|
||||
"method": "fp8",
|
||||
"ignored_layers": ["img_mlp"],
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
- **All layers sensitive (rare):**
|
||||
|
||||
**Solution:** The model may not be suitable for FP8 quantization. Consider weight-only quantization or other methods.
|
||||
|
||||
### Issue: Missing Weights Warning
|
||||
|
||||
**Symptoms:** Warnings about unloaded weights for `vae.*` or `text_encoder.*` parameters.
|
||||
|
||||
**Solution:** Ensure `load_weights` marks pre-loaded component weights. See [Step 2.3](#23-fix-load_weights-to-include-pre-loaded-components).
|
||||
|
||||
### Issue: CUDA Out of Memory
|
||||
|
||||
**Symptoms:** OOM errors during model loading despite FP8 being enabled.
|
||||
|
||||
**Causes & Solutions:**
|
||||
|
||||
- **Text encoder/VAE loaded in BF16 before FP8 conversion:**
|
||||
|
||||
FP8 weight storage is applied after `from_pretrained()`, so peak load memory is still BF16. This is a one-time cost during initialization.
|
||||
|
||||
---
|
||||
|
||||
## Reference Implementations
|
||||
|
||||
| Model | Pipeline File | Transformer File | Notes |
|
||||
|-------|-------------|-----------------|-------|
|
||||
| **Z-Image** | `models/z_image/pipeline_z_image.py` | `models/z_image/z_image_transformer.py` | DiT + text encoder FP8 |
|
||||
| **Qwen-Image** | `models/qwen_image/pipeline_qwen_image.py` | `models/qwen_image/qwen_image_transformer.py` | DiT + text encoder + VAE FP8 |
|
||||
|
||||
All files are under `vllm_omni/diffusion/`.
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
Adding FP8 quantization support to a new model:
|
||||
|
||||
1. ✅ **Add `quant_config` to transformer** — Accept and pass to vLLM parallel linear layers with proper `prefix`
|
||||
2. ✅ **Identify sensitive layers** — Profile which layers need `ignored_layers` for quality
|
||||
3. ✅ **Apply FP8 weight storage** — Call `apply_fp8_weight_storage()` on text encoder and VAE in pipeline `__init__`
|
||||
4. ✅ **Fix `load_weights`** — Mark pre-loaded component weights as loaded
|
||||
5. ✅ **Update docs** — Add model to the supported models table in `docs/user_guide/diffusion/quantization/fp8.md`
|
||||
6. ✅ **Test** — Verify quality, memory savings, and sensitive layer behavior
|
||||
@@ -1,20 +1,93 @@
|
||||
# Adding a TTS Model
|
||||
|
||||
This guide walks through adding a new TTS model to vLLM-Omni, using **Qwen3-TTS**
|
||||
as a reference. Qwen3-TTS demonstrates the standard two-stage TTS pipeline and the
|
||||
key optimizations all TTS models in this repo should follow.
|
||||
This guide walks through adding a new TTS model to vLLM-Omni. Two patterns are
|
||||
supported:
|
||||
|
||||
- **Two-stage pipeline** (e.g. Qwen3-TTS, Fish Speech): an AR code-predictor stage
|
||||
feeds an audio decoder stage via the `async_chunk` framework. This is the standard
|
||||
pattern for maximum streaming performance.
|
||||
- **Single-stage AR model** (e.g. MOSS-TTS-Nano): the model runs entirely inside one
|
||||
AR worker and streams audio chunks directly from its own `inference_stream()` generator.
|
||||
|
||||
Qwen3-TTS is used as the reference for the two-stage pattern. For the single-stage
|
||||
pattern, refer to MOSS-TTS-Nano.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
1. [Overview](#overview)
|
||||
2. [Directory Structure](#directory-structure)
|
||||
3. [Step-by-Step Implementation](#step-by-step-implementation)
|
||||
4. [Key Components](#key-components)
|
||||
5. [Model Registration](#model-registration)
|
||||
6. [Stage Configuration](#stage-configuration)
|
||||
7. [Stage Input Processors](#stage-input-processors)
|
||||
8. [Testing](#testing)
|
||||
9. [Summary](#summary)
|
||||
2. [Cross-Cutting Invariants](#cross-cutting-invariants)
|
||||
3. [Directory Structure](#directory-structure)
|
||||
4. [Step-by-Step Implementation](#step-by-step-implementation)
|
||||
5. [Key Components](#key-components)
|
||||
6. [Model Registration](#model-registration)
|
||||
7. [Stage Configuration](#stage-configuration)
|
||||
8. [Stage Input Processors](#stage-input-processors)
|
||||
9. [Online Serving Integration](#online-serving-integration)
|
||||
10. [Single-Stage Models](#single-stage-models)
|
||||
11. [Testing](#testing)
|
||||
12. [Pre-commit and DCO](#pre-commit-and-dco)
|
||||
13. [Summary](#summary)
|
||||
|
||||
## Cross-Cutting Invariants
|
||||
|
||||
These rules apply to every TTS model regardless of architecture (AR vs AR+diffusion,
|
||||
single-stage vs two-stage, codec-based vs VAE-based). Each has surfaced as a silent
|
||||
bug in a shipped PR — check them at the end of every phase, not just at the start.
|
||||
|
||||
**I1. Streaming output contract.** Pick one per-step semantics for `forward()` and
|
||||
document it in the docstring:
|
||||
|
||||
- *Delta*: yield only new audio samples produced this step. Preferred — linear cost.
|
||||
- *Cumulative*: re-decode from step 0 every call. O(N²); only acceptable when the
|
||||
codec exposes no streaming decode.
|
||||
|
||||
If you choose delta, audit the full chain: `forward()` returns the new chunk →
|
||||
`_consolidate_multimodal_tensors()` in `vllm_omni/engine/output_processor.py`
|
||||
concatenates the audio key into a single tensor at finish → streaming consumers
|
||||
receive per-step chunks, offline consumers receive the concatenated tensor. A
|
||||
mismatch (consolidator skips the key with `continue`, or consumers expect a list
|
||||
but receive a tensor) is invisible in offline RTF benchmarks — users hear replays
|
||||
or truncation only under live playback.
|
||||
|
||||
**I2. Multimodal output consumer hygiene.** `outputs[0].outputs[0].multimodal_output[key]`
|
||||
can be `Tensor`, `list[Tensor]` (pre-consolidation snapshot), `np.ndarray`, or
|
||||
scalar. In every test, example, and benchmark:
|
||||
|
||||
- Never write `dict.get("a") or dict.get("b")` on tensor values — Python evaluates
|
||||
the tensor's truthiness and raises `Boolean value of Tensor with more than one
|
||||
value is ambiguous`. Use explicit `if x is None` chains.
|
||||
- Defensively handle the list form:
|
||||
`if isinstance(x, list): x = torch.cat([t.reshape(-1) for t in x], dim=0)`.
|
||||
- Assert `shape` / `dtype` / `duration` explicitly — do not rely on truthiness for
|
||||
presence checks.
|
||||
|
||||
**I3. Hot-loop GPU discipline.** Inside any per-step model loop (AR decode,
|
||||
diffusion solver, CFM Euler step, per-frame vocoder):
|
||||
|
||||
- No `tensor.item()`, `.cpu()`, or `.tolist()` — each triggers a GPU→CPU sync; a
|
||||
10-step × 60-frame × 4-op loop creates 2400 syncs per request.
|
||||
- Prefer `dst.copy_(src)` over `dst.fill_(src.item())` for scalar-into-buffer writes.
|
||||
- Whole-model `torch.compile(Model.forward, fullgraph=False)` usually outperforms
|
||||
per-submodule compile — fewer dispatch boundaries, larger fusion regions. Measure
|
||||
before choosing granularity.
|
||||
- No Python control flow that depends on tensor values; use `torch.where` or masking.
|
||||
|
||||
Profile before optimizing.
|
||||
|
||||
**I4. Validation pyramid.** Offline RTF alone is necessary but not sufficient. A
|
||||
new TTS model must pass all three levels:
|
||||
|
||||
| Layer | Catches | Tool |
|
||||
|-------|---------|------|
|
||||
| Offline RTF / duration | Throughput regressions, missing audio, wrong sample rate | `end2end.py`, pytest e2e |
|
||||
| Browser streaming playback | Delta-vs-cumulative bugs, chunk boundary glitches, TTFP regressions | Gradio demo over `/v1/audio/speech?stream=true` |
|
||||
| Concurrent requests | Per-request state leaks, codec window round-robin gaps | `max_num_seqs>1` smoke with 4+ parallel prompts |
|
||||
|
||||
**I5. Per-request state belongs to the request.** If the model caches anything
|
||||
across `forward()` calls (streaming generators, codec buffers, sliding-window pads,
|
||||
CUDA graph state), key it by `info.get("_omni_req_id")` and free the entry on
|
||||
request finish. A shared buffer silently corrupts audio across concurrent requests —
|
||||
the symptom is crosstalk or truncation under load, nothing in single-request tests.
|
||||
|
||||
## Overview
|
||||
|
||||
@@ -113,6 +186,17 @@ vllm_omni/model_executor/stage_configs/
|
||||
your_model_name_async_chunk.yaml # Streaming mode config
|
||||
```
|
||||
|
||||
### Example placement
|
||||
|
||||
TTS examples live in the consolidated text-to-speech hub, **not** in their
|
||||
own top-level directory. Place per-model scripts under
|
||||
`examples/offline_inference/text_to_speech/<your_model>/` and
|
||||
`examples/online_serving/text_to_speech/<your_model>/`, and add a section
|
||||
to the hub `README.md` files (table row + per-model section) instead of a
|
||||
new per-model `README.md`. The mkdocs `generate_examples` hook treats the
|
||||
`text_to_speech/` parent as a single example, so per-model READMEs inside
|
||||
it would not be picked up — the hub README is the documented surface.
|
||||
|
||||
**Qwen3-TTS reference files:**
|
||||
|
||||
| File | Purpose |
|
||||
@@ -556,6 +640,302 @@ Recommended test cases for a new TTS model:
|
||||
|
||||
Reference test: `tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py`
|
||||
|
||||
### E2E Online Serving Tests (`tests/e2e/online_serving/test_<your_model>.py`)
|
||||
|
||||
The `omni_server` fixture in `tests/conftest.py` is **module-scoped**. Each distinct
|
||||
`OmniServerParams` id in the same test file forces the fixture to tear the server
|
||||
down and spawn a new one mid-module. A few rules that save real CI debugging time:
|
||||
|
||||
- **Prefer a single `OmniServerParams` set per file.** If you need to exercise two
|
||||
deploy variants (e.g. `model.yaml` and `model_async_chunk.yaml`), either use one
|
||||
variant and exercise streaming via request args, or split into two test files so
|
||||
each file does exactly one server lifecycle. Mid-module teardown/restart is the
|
||||
fragile path and surfaces startup races first.
|
||||
- **Never depend on server-side fetching of external URLs** for reference audio or
|
||||
other fixture data. CI runners (and China-hosted dev boxes) routinely fail on
|
||||
SSL/DNS for public URLs. Inline the payload as a `data:audio/wav;base64,...`
|
||||
ref_audio value — the serving layer accepts both forms.
|
||||
- **Don't roll your own readiness probe.** The harness already waits for HTTP 200
|
||||
on `/health` before releasing the server to the test. If your model needs extra
|
||||
warmup signals, expose them through `/health` rather than adding `time.sleep(...)`
|
||||
inside the test. (Bare TCP `connect_ex` probes were insufficient; see
|
||||
`tests/conftest.py` `OmniServer.wait_for_ready`.)
|
||||
- **Use `core_model` marker + H100 hardware_test** to match the `test-ready.yml`
|
||||
pipeline so your test is picked up by the `ready` label, not only nightly.
|
||||
|
||||
## Online Serving Integration
|
||||
|
||||
To expose your model through the `/v1/audio/speech` OpenAI-compatible endpoint, add
|
||||
**all five** of the following integration points to
|
||||
`vllm_omni/entrypoints/openai/serving_speech.py` in a **single commit**. Adding them
|
||||
piecemeal causes partial-integration failures that are hard to debug.
|
||||
|
||||
### 1. Stage constant
|
||||
|
||||
Near the top of the file, alongside the other `_*_TTS_MODEL_STAGES` constants:
|
||||
|
||||
```python
|
||||
_YOUR_MODEL_TTS_MODEL_STAGES = {"your_model_stage_key"}
|
||||
```
|
||||
|
||||
### 2. Union into `_TTS_MODEL_STAGES`
|
||||
|
||||
Add to the `_TTS_MODEL_STAGES` set union:
|
||||
|
||||
```python
|
||||
_TTS_MODEL_STAGES: set[str] = (
|
||||
...
|
||||
| _YOUR_MODEL_TTS_MODEL_STAGES
|
||||
)
|
||||
```
|
||||
|
||||
### 3. Model type detection
|
||||
|
||||
In `_detect_tts_model_type()`, add before the final `return None`:
|
||||
|
||||
```python
|
||||
if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES:
|
||||
return "your_model"
|
||||
```
|
||||
|
||||
### 4. Request validation dispatch
|
||||
|
||||
In `_validate_tts_request()`, add before the fallback `return`:
|
||||
|
||||
```python
|
||||
if self._tts_model_type == "your_model":
|
||||
return self._validate_your_model_request(request)
|
||||
```
|
||||
|
||||
### 5. Validation and parameter-builder methods
|
||||
|
||||
Add two new methods:
|
||||
|
||||
```python
|
||||
def _validate_your_model_request(
|
||||
self, request: OpenAICreateSpeechRequest
|
||||
) -> str | None:
|
||||
"""Validate YourModel request. Returns an error string or None."""
|
||||
if not request.input or not request.input.strip():
|
||||
return "Input text cannot be empty"
|
||||
return None
|
||||
|
||||
def _build_your_model_params(
|
||||
self, request: OpenAICreateSpeechRequest
|
||||
) -> dict[str, Any]:
|
||||
"""Build additional_information dict for YourModel."""
|
||||
params: dict[str, Any] = {"text": [request.input]}
|
||||
if request.voice is not None:
|
||||
params["voice"] = [request.voice]
|
||||
# Add any other model-specific fields here
|
||||
return params
|
||||
```
|
||||
|
||||
Then wire `_build_your_model_params` into the request-dispatch block in
|
||||
`_create_tts_request()` (search for the equivalent `_build_*_params` call for an
|
||||
existing model to find the right location). If the model supports voice cloning
|
||||
(`ref_audio` → `prompt_audio_path`, `ref_text` → `prompt_text`), add those mappings
|
||||
here too — follow any existing `_build_<model>_params` in `serving_speech.py` (e.g.
|
||||
`_build_moss_tts_params` for the voice-cloning variant) for the pattern.
|
||||
|
||||
> **Two dispatch patterns coexist:** Fish Speech uses a `self._is_fish_speech` boolean
|
||||
> checked *before* `elif self._is_tts`. All newer models use the `_tts_model_type`
|
||||
> string pattern shown above. For new models, always use the string pattern — do not
|
||||
> add new `_is_*` boolean flags.
|
||||
|
||||
> **Note on unused variables:** Only extract parameters in `_build_your_model_params`
|
||||
> that you actually pass to the model's generate / `inference_stream` call. Extracting
|
||||
> a variable without forwarding it will trigger a `ruff F841` pre-commit failure.
|
||||
|
||||
### Merge conflicts
|
||||
|
||||
`serving_speech.py` is modified by every new model PR and is the most common source of
|
||||
rebase conflicts. When rebasing onto `main` and a conflict appears here, the resolution
|
||||
is always to **keep both** the upstream model's additions and your own — never discard
|
||||
either side. After resolving:
|
||||
|
||||
```bash
|
||||
git add vllm_omni/entrypoints/openai/serving_speech.py
|
||||
git rebase --continue
|
||||
```
|
||||
|
||||
## Single-Stage Models
|
||||
|
||||
Some TTS models (e.g. MOSS-TTS-Nano) do not use a two-stage pipeline. Instead the
|
||||
entire AR LM and audio decoder run inside a single AR worker, streaming audio chunks
|
||||
directly from the model's own generator.
|
||||
|
||||
### Directory structure
|
||||
|
||||
```
|
||||
vllm_omni/model_executor/models/your_model_name/
|
||||
__init__.py
|
||||
modeling_your_model_name.py # unified class: load_weights + forward + streaming
|
||||
|
||||
vllm_omni/model_executor/stage_configs/your_model_name.yaml
|
||||
```
|
||||
|
||||
No stage input processor is needed.
|
||||
|
||||
### Stage config
|
||||
|
||||
Use a single stage with `worker_type: ar`. The `is_comprehension: true` field and the
|
||||
top-level `async_chunk: false` are required — omitting them causes silent
|
||||
misclassification in the serving layer. Set `max_num_seqs` to at least 4 for
|
||||
concurrent production use.
|
||||
|
||||
```yaml
|
||||
# stage_configs/your_model_name.yaml
|
||||
async_chunk: false
|
||||
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_comprehension: true # required for serving_speech.py dispatch
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
model_stage: your_model_stage_key
|
||||
model_arch: YourModelForCausalLM
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
engine_output_type: audio
|
||||
max_num_seqs: 4 # min 4 for concurrent requests; default 1 causes gaps
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
```
|
||||
|
||||
### Generator-based streaming pattern
|
||||
|
||||
This is the MOSS-TTS-Nano pattern, distinct from VoxCPM2's vLLM-native AR pattern
|
||||
(see `plan/voxcpm2_native_ar_design.md` for that variant). Load model weights in
|
||||
`load_weights()` (not `__init__`) so vLLM finishes distributed initialisation before
|
||||
any CUDA allocations. Stream via a per-request generator stored in an instance dict:
|
||||
|
||||
```python
|
||||
class YourModelForCausalLM(nn.Module):
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
self._lm: nn.Module | None = None # populated in load_weights()
|
||||
self._stream_gens: dict[str, Any] = {} # request_key → generator
|
||||
|
||||
def load_weights(self, weights):
|
||||
# Load self._lm here, after vLLM distributed init
|
||||
...
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids,
|
||||
positions,
|
||||
intermediate_tensors=None,
|
||||
inputs_embeds=None,
|
||||
runtime_additional_information: list[dict] | None = None, # one dict per request
|
||||
**kwargs,
|
||||
) -> OmniOutput:
|
||||
infos = runtime_additional_information or [{}]
|
||||
# Return empty output during dummy/profiling calls
|
||||
if not runtime_additional_information or all(i.get("_is_dummy") for i in infos):
|
||||
self._ar_emit_stop_token = True
|
||||
return OmniOutput(...)
|
||||
|
||||
outputs, last_flags = [], []
|
||||
for info in infos:
|
||||
request_key = str(info.get("_omni_req_id", "0")) # set by vLLM, not user code
|
||||
if request_key not in self._stream_gens:
|
||||
self._stream_gens[request_key] = self._create_stream_gen(info)
|
||||
try:
|
||||
chunk, is_last = next(self._stream_gens[request_key])
|
||||
except StopIteration:
|
||||
chunk, is_last = torch.zeros(0), True
|
||||
if is_last:
|
||||
del self._stream_gens[request_key]
|
||||
outputs.append(chunk)
|
||||
last_flags.append(is_last)
|
||||
|
||||
self._ar_emit_stop_token = all(last_flags)
|
||||
return OmniOutput(multimodal_outputs={"model_outputs": outputs, "is_last": last_flags})
|
||||
|
||||
def _create_stream_gen(self, info: dict):
|
||||
"""Yield (waveform_tensor, is_last) from the model's inference_stream().
|
||||
|
||||
Handle both incremental ("audio" events) and batch ("result" event) models:
|
||||
some upstream implementations emit one "result" event with the full waveform
|
||||
instead of incremental "audio" events. Both paths must be covered.
|
||||
"""
|
||||
for event in self._lm.inference_stream(...):
|
||||
if event["type"] == "audio":
|
||||
yield event["waveform"], False
|
||||
elif event["type"] == "result":
|
||||
# Fallback for models that don't emit incremental audio events
|
||||
yield event.get("waveform", torch.zeros(0)), True
|
||||
return
|
||||
yield torch.zeros(0), True
|
||||
|
||||
def compute_logits(self, hidden_states, sampling_metadata):
|
||||
# Emit EOS only when the last chunk has been yielded so the AR
|
||||
# scheduler ends the request at the right time.
|
||||
...
|
||||
```
|
||||
|
||||
For an in-tree reference, look for any single-stage AR model under
|
||||
`vllm_omni/model_executor/models/` (for example
|
||||
`moss_tts_nano/modeling_moss_tts_nano.py` once its integration has landed).
|
||||
|
||||
## Pre-commit and DCO
|
||||
|
||||
All contributions must pass the pre-commit checks and the Developer Certificate of
|
||||
Origin (DCO) sign-off before merging.
|
||||
|
||||
### Running pre-commit
|
||||
|
||||
Install the hooks once with `pre-commit install`. Then run before committing:
|
||||
|
||||
```bash
|
||||
pre-commit run --files \
|
||||
vllm_omni/model_executor/models/your_model_name/*.py \
|
||||
vllm_omni/entrypoints/openai/serving_speech.py \
|
||||
vllm_omni/model_executor/models/registry.py \
|
||||
tests/e2e/offline_inference/test_your_model_name.py \
|
||||
tests/e2e/online_serving/test_your_model_name.py
|
||||
```
|
||||
|
||||
When pre-commit **modifies files**, it exits with a non-zero code but the reformatting
|
||||
is correct. Stage the modified files and commit again — do not revert the changes.
|
||||
|
||||
Common failures and fixes:
|
||||
|
||||
| Check | Cause | Fix |
|
||||
|-------|-------|-----|
|
||||
| `ruff F841` | Local variable assigned but never used | Remove the extraction or forward it to the model call |
|
||||
| `ruff E402` | Module-level import not at top of file | Move import to the top-level import block |
|
||||
| `ruff format` | Line length, spacing, or quote style | Accept the auto-fix, stage, and re-commit |
|
||||
|
||||
### DCO sign-off
|
||||
|
||||
Every commit must carry a `Signed-off-by` trailer. Use the `-s` flag when committing:
|
||||
|
||||
```bash
|
||||
git commit -s -m "feat(your-model): add YourModel TTS support"
|
||||
```
|
||||
|
||||
Or configure git to add it automatically:
|
||||
|
||||
```bash
|
||||
git config format.signOff true
|
||||
```
|
||||
|
||||
To fix a missing sign-off on the most recent commit:
|
||||
|
||||
```bash
|
||||
git commit --amend -s --no-edit
|
||||
git push origin your-branch --force-with-lease
|
||||
```
|
||||
|
||||
> The DCO check verifies that the commit author email matches the `Signed-off-by` email.
|
||||
> Make sure `git config user.email` is set to the address associated with your GitHub
|
||||
> account before committing.
|
||||
|
||||
## Adding a Model Recipe
|
||||
|
||||
After implementing and testing your model, add a model recipe to the
|
||||
@@ -567,15 +947,19 @@ for the expected format.
|
||||
|
||||
Adding a TTS model to vLLM-Omni involves:
|
||||
|
||||
1. **Create model directory** with AR stage, decoder stage, and unified class
|
||||
1. **Create model directory** with AR stage, decoder stage, and unified class (two-stage)
|
||||
or a single unified class with generator-based streaming (single-stage)
|
||||
2. **AR stage** - use vLLM's native decoder layers with fused QKV; do not wrap HF directly
|
||||
3. **Decoder stage** - thin wrapper around your audio decoder; implement `chunked_decode_streaming()`
|
||||
4. **Unified class** - dispatches on `model_stage`; same structure as `Qwen3TTSModelForGeneration`
|
||||
5. **Register** all stage classes in `registry.py`
|
||||
6. **YAML configs** - provide both batch and `async_chunk` variants
|
||||
7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25
|
||||
8. **Tests** - cover single request, batching, and async_chunk streaming
|
||||
9. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
|
||||
6. **YAML configs** - provide both batch and `async_chunk` variants (two-stage), or a single-stage AR config
|
||||
7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25 (two-stage only)
|
||||
8. **Online serving** - add all 5 integration points to `serving_speech.py` in one commit
|
||||
9. **Tests** - cover single request, batching, and streaming
|
||||
10. **Pre-commit + DCO** - run `pre-commit` before pushing; sign every commit with `git commit -s`
|
||||
11. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
|
||||
12. **Invariants** - re-check I1–I5 (streaming contract, consumer hygiene, hot-loop discipline, validation pyramid, per-request state) at the end of every phase
|
||||
|
||||
### Qwen3-TTS Reference Files
|
||||
|
||||
|
||||
+148
-74
@@ -1,15 +1,17 @@
|
||||
# Profiling vLLM-Omni
|
||||
# Profiling Diffusion Models
|
||||
|
||||
> **Warning:** Profiling is for development and debugging only. It adds significant overhead and should not be enabled in production.
|
||||
|
||||
vLLM-Omni supports two profiler backends through `profiler_config`:
|
||||
Diffusion profiling supports two backends through `profiler_config`:
|
||||
|
||||
- `torch`: detailed CPU/CUDA traces written to `torch_profiler_dir`
|
||||
- `torch`: detailed CPU/CUDA traces, operator tables, and optional memory snapshots
|
||||
- `cuda`: low-overhead CUDA range control for NVIDIA Nsight Systems (`nsys`)
|
||||
|
||||
## 1. Configure Profiling
|
||||
## 1. Configure `profiler_config`
|
||||
|
||||
Use the same `profiler_config` shape everywhere:
|
||||
Use `profiler_config` to enable profiling for a diffusion model. For diffusion usage, pass it directly to `Omni(...)` or `vllm serve`.
|
||||
|
||||
Minimal torch-profiler config:
|
||||
|
||||
```yaml
|
||||
profiler_config:
|
||||
@@ -21,54 +23,116 @@ Supported fields:
|
||||
|
||||
| Field | Description |
|
||||
|---|---|
|
||||
| `profiler` | Profiler backend. Supported values: `torch`, `cuda`. |
|
||||
| `torch_profiler_dir` | Output directory for torch traces. Required when `profiler: torch`. |
|
||||
| `profiler` | Profiler backend. Supported values: `torch`, `cuda`. Use `torch` for `trace.json`, Excel operator tables, and optional memory snapshots. Use `cuda` for Nsight Systems only. |
|
||||
| `torch_profiler_dir` | Output directory for torch-profiler artifacts. Required when `profiler: torch`. |
|
||||
| `torch_profiler_use_gzip` | Compress `trace_rank*.json` into `trace_rank*.json.gz`. |
|
||||
| `torch_profiler_record_shapes` | Record input shapes and add a `by_shape` sheet to `ops_rank*.xlsx`. |
|
||||
| `torch_profiler_with_stack` | Record call stacks, add a `by_stack` sheet to `ops_rank*.xlsx`, and export `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt`. |
|
||||
| `torch_profiler_with_memory` | Enable memory profiling and attempt to dump `memory_snapshot_rank*.pickle`. The pickle is only generated when the current backend supports memory history and snapshot APIs. |
|
||||
| `torch_profiler_with_flops` | Enable FLOPs collection in `torch.profiler`. This does not add a separate output file. |
|
||||
| `torch_profiler_dump_cuda_time_total` | Export an additional text summary `profiler_out_<rank>.txt` sorted by `self_cuda_time_total`. |
|
||||
| `delay_iterations` | Number of worker iterations to skip before profiling starts. |
|
||||
| `max_iterations` | Maximum number of worker iterations to capture before auto-stop. |
|
||||
| `wait_iterations` | Torch-profiler wait iterations before warmup. |
|
||||
| `warmup_iterations` | Torch-profiler warmup iterations. |
|
||||
| `active_iterations` | Torch-profiler active iterations. |
|
||||
| `wait_iterations` | Torch-profiler wait iterations before warmup. |
|
||||
|
||||
For multi-stage omni pipelines, put `profiler_config` under the target stage's `engine_args`.
|
||||
For detailed explanations of the fields, please refer to upstream vLLM implementation [vllm/config/profiler.py](https://github.com/vllm-project/vllm/blob/v0.20.1/vllm/config/profiler.py)
|
||||
|
||||
```yaml
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
engine_args:
|
||||
profiler_config:
|
||||
profiler: torch
|
||||
torch_profiler_dir: ./perf
|
||||
```
|
||||
### Minimal configurations by output
|
||||
|
||||
For single-stage diffusion usage, pass `profiler_config` directly to `Omni(...)` or `vllm serve`.
|
||||
|
||||
## 2. Profiling Omni Pipelines
|
||||
|
||||
It is usually best to profile only the stages you need.
|
||||
Only collect trace output:
|
||||
|
||||
```python
|
||||
# Profile all stages.
|
||||
omni.start_profile()
|
||||
|
||||
# Profile selected stages only.
|
||||
omni.start_profile(stages=[0, 2])
|
||||
...
|
||||
omni.stop_profile(stages=[0, 2])
|
||||
profiler_config = {
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": "./perf",
|
||||
}
|
||||
```
|
||||
|
||||
Always stop the same stage set that you started. If only some stages have `profiler_config`, pass an explicit `stages=[...]` list instead of relying on the default "all stages" behavior.
|
||||
Outputs:
|
||||
|
||||
Examples:
|
||||
- `trace_rank*.json`
|
||||
- `ops_rank*.xlsx` with a `summary` sheet
|
||||
|
||||
1. [Qwen2.5-Omni end2end](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
|
||||
2. [Qwen3-Omni end2end](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
|
||||
Collect compressed trace output:
|
||||
|
||||
## 3. Profiling Single-Stage Diffusion
|
||||
```python
|
||||
profiler_config = {
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": "./perf",
|
||||
"torch_profiler_use_gzip": True,
|
||||
}
|
||||
```
|
||||
|
||||
Single-stage diffusion models use the same `start_profile()` / `stop_profile()` controls, but you must provide `profiler_config` explicitly.
|
||||
Outputs:
|
||||
|
||||
### PyTorch profiler
|
||||
- `trace_rank*.json.gz`
|
||||
- `ops_rank*.xlsx` with a `summary` sheet
|
||||
|
||||
Collect trace and full operator tables:
|
||||
|
||||
```python
|
||||
profiler_config = {
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": "./perf",
|
||||
"torch_profiler_record_shapes": True,
|
||||
"torch_profiler_with_stack": True,
|
||||
}
|
||||
```
|
||||
|
||||
Outputs:
|
||||
|
||||
- `trace_rank*.json`
|
||||
- `ops_rank*.xlsx` with `summary`, `by_shape`, and `by_stack`
|
||||
- `stacks_cpu_rank*.txt`
|
||||
- `stacks_cuda_rank*.txt`
|
||||
|
||||
Collect trace, operator tables, and memory snapshots:
|
||||
|
||||
```python
|
||||
profiler_config = {
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": "./perf",
|
||||
"torch_profiler_record_shapes": True,
|
||||
"torch_profiler_with_stack": True,
|
||||
"torch_profiler_with_memory": True,
|
||||
}
|
||||
```
|
||||
|
||||
Outputs:
|
||||
|
||||
- `trace_rank*.json`
|
||||
- `ops_rank*.xlsx` with `summary`, `by_shape`, and `by_stack`
|
||||
- `stacks_cpu_rank*.txt`
|
||||
- `stacks_cuda_rank*.txt`
|
||||
- `memory_snapshot_rank*.pickle` when supported by the current backend
|
||||
|
||||
### Full torch-profiler configuration
|
||||
|
||||
If you want to enable the commonly used torch-profiler options together:
|
||||
|
||||
```python
|
||||
profiler_config = {
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": "./perf",
|
||||
"torch_profiler_use_gzip": False,
|
||||
"torch_profiler_record_shapes": True,
|
||||
"torch_profiler_with_stack": True,
|
||||
"torch_profiler_with_memory": True,
|
||||
"torch_profiler_with_flops": False,
|
||||
"torch_profiler_dump_cuda_time_total": False,
|
||||
"delay_iterations": 0,
|
||||
"max_iterations": 0,
|
||||
"wait_iterations": 0,
|
||||
"warmup_iterations": 0,
|
||||
"active_iterations": 1,
|
||||
}
|
||||
```
|
||||
|
||||
## 2. Profiling Diffusion with PyTorch Profiler
|
||||
|
||||
Single-stage diffusion models use `start_profile()` / `stop_profile()` controls. The profiler only writes artifacts after profiling has been started and then stopped.
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
@@ -86,7 +150,29 @@ omni.start_profile()
|
||||
omni.stop_profile()
|
||||
```
|
||||
|
||||
### Nsight Systems (`nsys`)
|
||||
For diffusion offline example scripts under `examples/offline_inference/`, pass `--profiler-config` as a JSON object. The script enables profiling when this argument is set and wraps generation with `start_profile()` / `stop_profile()`.
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/image_to_video/image_to_video.py \
|
||||
--model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
|
||||
--image input.jpg \
|
||||
--prompt "A cat playing with yarn" \
|
||||
--profiler-config '{
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": "./perf",
|
||||
"torch_profiler_record_shapes": true,
|
||||
"torch_profiler_with_stack": true
|
||||
}'
|
||||
```
|
||||
|
||||
Examples:
|
||||
|
||||
1. [Image edit example](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
|
||||
2. [Image to video example](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
|
||||
|
||||
## 3. Profiling Diffusion with Nsight Systems (`nsys`)
|
||||
|
||||
For Nsight Systems, use `profiler: cuda` and wrap the process with `nsys profile`.
|
||||
|
||||
@@ -103,35 +189,20 @@ nsys profile \
|
||||
The Python process being profiled must create the diffusion engine with:
|
||||
|
||||
```python
|
||||
profiler_config={"profiler": "cuda"}
|
||||
profiler_config = {"profiler": "cuda"}
|
||||
```
|
||||
|
||||
Then call `start_profile()` before the requests you want to capture and `stop_profile()` after them. The diffusion worker processes open and close the CUDA capture range themselves, so `nsys` sees the actual GPU work instead of only the parent process.
|
||||
|
||||
Examples:
|
||||
|
||||
1. [Image edit example](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
|
||||
2. [Image to video example](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
|
||||
|
||||
## 4. Profiling Online Serving
|
||||
|
||||
When any stage has `profiler_config.profiler` set, the server exposes:
|
||||
When `profiler_config.profiler` is set for a diffusion model, the server exposes:
|
||||
|
||||
- `POST /start_profile`
|
||||
- `POST /stop_profile`
|
||||
|
||||
### Start the server
|
||||
|
||||
Multi-stage omni serving:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
(The default deploy config at `vllm_omni/deploy/qwen2_5_omni.yaml` is loaded automatically. Pass `--deploy-config /path/to/custom.yaml` to override.)
|
||||
|
||||
Single-stage diffusion serving with torch profiler:
|
||||
|
||||
```bash
|
||||
@@ -144,21 +215,17 @@ vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers \
|
||||
"torch_profiler_with_stack": true,
|
||||
"torch_profiler_with_flops": false,
|
||||
"torch_profiler_use_gzip": true,
|
||||
"torch_profiler_dump_cuda_time_total": true,
|
||||
"torch_profiler_dump_cuda_time_total": false,
|
||||
"torch_profiler_record_shapes": true,
|
||||
"torch_profiler_with_memory": true,
|
||||
"delay_iterations": 0,
|
||||
"max_iterations": 0
|
||||
"max_iterations": 0,
|
||||
"wait_iterations": 0,
|
||||
"warmup_iterations": 0,
|
||||
"active_iterations": 1
|
||||
}'
|
||||
```
|
||||
|
||||
Useful optional fields:
|
||||
|
||||
- `torch_profiler_with_stack`: export `by_stack` operator views and stack text files.
|
||||
- `torch_profiler_record_shapes`: export `by_shape` operator views.
|
||||
- `torch_profiler_with_memory`: dump `memory_snapshot_rank*.pickle` when the backend supports memory history.
|
||||
- `torch_profiler_use_gzip`: write the trace as `trace_rank*.json.gz`.
|
||||
|
||||
Single-stage diffusion serving with Nsight Systems:
|
||||
|
||||
```bash
|
||||
@@ -176,30 +243,37 @@ nsys profile \
|
||||
|
||||
### Control capture
|
||||
|
||||
Example profiling flow for an online Qwen-Image request:
|
||||
|
||||
```bash
|
||||
# Start profiling on all profiled stages.
|
||||
# Start profiling.
|
||||
curl -X POST http://localhost:8091/start_profile
|
||||
|
||||
# Start profiling on selected stages.
|
||||
curl -X POST http://localhost:8091/start_profile \
|
||||
# Send a Qwen-Image generation request while profiling is active.
|
||||
curl http://localhost:8091/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"stages": [0]}'
|
||||
-d '{
|
||||
"model": "Qwen/Qwen-Image",
|
||||
"prompt": "A red vintage bicycle parked beside a quiet canal at sunset"
|
||||
}'
|
||||
|
||||
# Stop profiling.
|
||||
# Stop profiling and flush profiler artifacts.
|
||||
curl -X POST http://localhost:8091/stop_profile
|
||||
```
|
||||
|
||||
For mixed-stage pipelines, use explicit `stages` and pass the same stage list to both endpoints.
|
||||
## 5. Diffusion Pipeline Profiler
|
||||
|
||||
## 5. Analyze Results
|
||||
For lightweight per-stage pipeline timing such as `vae.decode` or `diffuse`, see [Diffusion Pipeline Profiler](model/adding_diffusion_model.md#diffusion-pipeline-profiler-performance-profiling). That utility logs stage durations only and does not generate torch-profiler artifacts such as `trace.json`, Excel tables, or memory snapshots.
|
||||
|
||||
Torch profiler output:
|
||||
## 6. Analyze Results
|
||||
|
||||
Torch-profiler output:
|
||||
|
||||
- Chrome/Perfetto trace: `trace_rank*.json` or `trace_rank*.json.gz`
|
||||
- Excel workbook: `ops_rank*.xlsx` with `summary`, and optional `by_shape` / `by_stack` sheets
|
||||
- Stack exports: `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt` when stack capture is enabled
|
||||
- Memory snapshot: `memory_snapshot_rank*.pickle` when memory capture is enabled and supported by the backend
|
||||
- Optional aggregated CUDA-time tables under the same session directory
|
||||
- Optional CUDA-time text summary: `profiler_out_<rank>.txt` when `torch_profiler_dump_cuda_time_total` is enabled
|
||||
|
||||
CUDA profiler / Nsight Systems output:
|
||||
|
||||
@@ -211,4 +285,4 @@ Recommended viewers:
|
||||
- `nsys stats <report>.nsys-rep` for CLI summaries
|
||||
- Nsight Systems GUI for CUDA kernel timelines
|
||||
|
||||
vLLM-Omni reuses the vLLM profiling infrastructure where possible. For the upstream reference, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/).
|
||||
For upstream background on the underlying vLLM profiling infrastructure, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/).
|
||||
@@ -0,0 +1,156 @@
|
||||
# Continuous Batching for Step-Wise Diffusion
|
||||
|
||||
!!! warning "Experimental Feature"
|
||||
This feature is experimental. It currently applies only to native
|
||||
diffusion pipelines running with `step_execution=True`.
|
||||
|
||||
This document describes the batching extension built on top of
|
||||
[Diffusion Step Execution](diffusion_step_execution.md). The base
|
||||
step-execution contract is unchanged. The batching work is mainly in the
|
||||
scheduler and runner layers.
|
||||
|
||||
## Why It Helps
|
||||
|
||||
Step-wise execution breaks a long denoise loop into scheduler-visible units.
|
||||
That gives the runtime a place to admit other compatible requests between
|
||||
steps instead of waiting for an entire request to finish.
|
||||
|
||||
This matters most in low-MFU or bursty serving scenarios:
|
||||
|
||||
- one request's denoise step may not fully saturate the GPU
|
||||
- several compatible requests can share the same denoise forward pass
|
||||
- throughput and device utilization can improve without changing request-local
|
||||
scheduler state
|
||||
|
||||
This is **not** a guaranteed single-request latency win. The main benefit is
|
||||
usually higher utilization and better throughput when the workload contains
|
||||
multiple in-flight compatible requests.
|
||||
|
||||
## Overview
|
||||
|
||||
With continuous batching enabled:
|
||||
|
||||
- the scheduler may keep multiple compatible requests active at the same time
|
||||
- the runner packs request-local step state into one `InputBatch`
|
||||
- `denoise_step()` runs on that batch
|
||||
- `step_scheduler()` and `post_decode()` still run per request
|
||||
|
||||
The current implementation is conservative:
|
||||
|
||||
- only compatible requests are batched together
|
||||
- request-mode diffusion still runs with `max_num_seqs=1`
|
||||
- per-request progress and completion remain independent
|
||||
|
||||
## Enablement
|
||||
|
||||
Use `--step-execution` as the feature gate, then increase `--max-num-seqs`
|
||||
above `1` if you want batching:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--step-execution \
|
||||
--max-num-seqs 8
|
||||
```
|
||||
|
||||
`--max-num-seqs 1` keeps the step-wise path without enabling batching.
|
||||
|
||||
For a reproducible replay flow using the bundled serving benchmark, see the
|
||||
Qwen-Image replay commands in
|
||||
[`benchmarks/diffusion/README.md`](gh-file:benchmarks/diffusion/README.md)
|
||||
and
|
||||
[`benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md`](gh-file:benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md).
|
||||
|
||||
## Scheduler
|
||||
|
||||
The scheduler derives its batch capacity from `max_num_seqs` through
|
||||
`max_num_running_reqs`.
|
||||
|
||||
Batch admission is gated by
|
||||
[`SamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py),
|
||||
which is built from shape-sensitive and CFG-sensitive sampling fields. This is
|
||||
the core correctness rule for batching: requests are only co-batched when they
|
||||
share the same denoise tensor contract.
|
||||
|
||||
There are two important details:
|
||||
|
||||
- `num_inference_steps` is not part of the key, so requests with different
|
||||
total step counts can still share a batch
|
||||
- requests also do not need to be at the same current denoise progress; active
|
||||
requests can continue batching even when their current step indices diverge
|
||||
- admission is still FIFO, so an incompatible request at the head of the
|
||||
waiting queue blocks later compatible requests
|
||||
|
||||
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 current batching unit is one `OmniDiffusionRequest`. Requests with
|
||||
multiple prompts do not participate in batching today.
|
||||
|
||||
## Runner
|
||||
|
||||
The runner keeps persistent per-request execution state in
|
||||
[`DiffusionRequestState`](gh-file:vllm_omni/diffusion/worker/utils.py),
|
||||
while the scheduler owns a separate lightweight request state for queueing and
|
||||
lifecycle tracking.
|
||||
|
||||
For each step, the runner builds an
|
||||
[`InputBatch`](gh-file:vllm_omni/diffusion/worker/input_batch.py) from the
|
||||
active request states:
|
||||
|
||||
- prompt embeddings and masks are normalized and padded
|
||||
- dynamic tensors such as `latents` and `timesteps` are gathered each step
|
||||
- buffers are reused when batch composition stays the same
|
||||
|
||||
The step-wise batched path is:
|
||||
|
||||
1. Run `prepare_encode()` for newly admitted requests.
|
||||
2. Build or refresh `InputBatch`.
|
||||
3. Run one batched `denoise_step(input_batch)`.
|
||||
4. Slice the batched `noise_pred` back per request.
|
||||
5. Run per-request `step_scheduler()`.
|
||||
6. Run `post_decode()` only for requests that finished denoising.
|
||||
7. Scatter updated latents back into persistent request state with
|
||||
[`scatter_latents()`](gh-file:vllm_omni/diffusion/worker/input_batch.py).
|
||||
|
||||
This keeps the shared work limited to the denoise forward pass while preserving
|
||||
request-local scheduler state and outputs.
|
||||
|
||||
## Engine
|
||||
|
||||
[`DiffusionEngine`](gh-file:vllm_omni/diffusion/diffusion_engine.py) provides
|
||||
the background loop and async add-request path needed for multiple requests to
|
||||
accumulate in the scheduler.
|
||||
|
||||
This is supporting infrastructure, not the main design point. The batching
|
||||
behavior is defined by scheduler-side compatibility gating and runner-side
|
||||
batch packing.
|
||||
|
||||
## Current Limitations
|
||||
|
||||
- Experimental feature; use `max_num_seqs=1` for the older conservative path.
|
||||
- Only native pipelines that already support `step_execution=True`.
|
||||
- Request-mode diffusion still clamps `max_num_seqs` back to `1`.
|
||||
- Only homogeneous batches keyed by `SamplingParamsKey` are supported.
|
||||
- Multi-prompt requests are not batched.
|
||||
- `cache_backend`, KV transfer, and other request-mode extras are not wired
|
||||
into the batched step-wise path yet.
|
||||
- Future work can relax the current same-shape restriction with richer
|
||||
heterogeneous batching policies such as bucketing or padded execution for
|
||||
different resolutions.
|
||||
|
||||
## Related Files
|
||||
|
||||
- Scheduler base:
|
||||
[`vllm_omni/diffusion/sched/base_scheduler.py`](gh-file:vllm_omni/diffusion/sched/base_scheduler.py)
|
||||
- Scheduler interface:
|
||||
[`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
- Step scheduler:
|
||||
[`vllm_omni/diffusion/sched/step_scheduler.py`](gh-file:vllm_omni/diffusion/sched/step_scheduler.py)
|
||||
- Runner:
|
||||
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
||||
- Input batch:
|
||||
[`vllm_omni/diffusion/worker/input_batch.py`](gh-file:vllm_omni/diffusion/worker/input_batch.py)
|
||||
- Tests:
|
||||
[`tests/diffusion/test_diffusion_scheduler.py`](gh-file:tests/diffusion/test_diffusion_scheduler.py)
|
||||
@@ -1,4 +1,4 @@
|
||||
# Diffusion Step Execution
|
||||
# Adding Step Execution Support for Diffusion Pipelines
|
||||
|
||||
This guide documents vLLM-Omni's stepwise diffusion contract for model authors
|
||||
and contributors implementing `step_execution=True` support for a diffusion
|
||||
@@ -7,12 +7,20 @@ pipeline.
|
||||
For end-user enablement, supported models, and current limitations, see
|
||||
[Step Execution](../../user_guide/diffusion/step_execution.md).
|
||||
|
||||
This document describes the base step-execution contract only. For the
|
||||
experimental batching policy layered on top of the step-wise path, see
|
||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
||||
|
||||
## Current Support Scope
|
||||
|
||||
`step_execution` is **not** a generic diffusion toggle. It only works for
|
||||
pipelines that implement the segmented stateful contract in
|
||||
[`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py).
|
||||
|
||||
This page is intentionally author-facing. Treat runtime enablement
|
||||
(`step_execution=True` in Python or `--step-execution` in serving) as an
|
||||
opt-in user knob layered on top of the implementation contract below.
|
||||
|
||||
Current in-tree support:
|
||||
|
||||
| Pipeline | Example models | Step execution |
|
||||
@@ -22,7 +30,9 @@ Current in-tree support:
|
||||
|
||||
Current engine/runtime limitations:
|
||||
|
||||
- `StepScheduler` only schedules `batch_size=1`.
|
||||
- Continuous batching with `max_num_seqs > 1` is experimental and documented in
|
||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
||||
Keep `max_num_seqs=1` if you want the older conservative behavior.
|
||||
- `cache_backend` is not supported in step mode.
|
||||
- Request-mode extras such as KV transfer are not wired into step mode yet.
|
||||
- Unsupported pipelines now fail early during model loading instead of failing on the first request.
|
||||
|
||||
@@ -10,6 +10,8 @@ This section contains design documents and architecture specifications for vLLM-
|
||||
|
||||
- [Disaggregated Inference](feature/disaggregated_inference.md)
|
||||
- [Ray-based Execution](feature/ray_based_execution.md)
|
||||
- [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
|
||||
- [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
|
||||
|
||||
## Module Design Documents
|
||||
|
||||
|
||||
@@ -201,7 +201,7 @@ class _BaseScheduler(SchedulerInterface):
|
||||
|
||||
- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
|
||||
|
||||
- **Current constraint**: `max_num_running_reqs` remains `1` because the current engine path is still synchronous request-mode execution.
|
||||
- **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.
|
||||
|
||||
#### 2.4 Current `RequestScheduler` Policy
|
||||
|
||||
@@ -217,7 +217,7 @@ class RequestScheduler(_BaseScheduler):
|
||||
|
||||
- **FIFO request scheduling**: Waiting requests are promoted in queue order.
|
||||
|
||||
- **Single-request admission**: The current policy only admits one active request at a time.
|
||||
- **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 scheduler ids.
|
||||
|
||||
@@ -237,7 +237,7 @@ while True:
|
||||
|
||||
- **No scheduler-owned IPC**: Scheduler no longer talks to workers directly.
|
||||
|
||||
- **Conservative concurrency**: The current request-mode implementation still allows only one active request at a time.
|
||||
- **Split concurrency model**: Request-mode diffusion remains single-active-request, while the step-wise path can keep multiple compatible requests running and advance them independently between denoise steps.
|
||||
|
||||
---
|
||||
|
||||
@@ -455,7 +455,7 @@ To learn more about the diffusion pipeline and how to add a new diffusion pipeli
|
||||
|
||||
#### Architecture
|
||||
|
||||
The attention system uses a **backend selector pattern** that automatically chooses the optimal attention implementation based on hardware and model configuration.
|
||||
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.
|
||||
|
||||
#### Backend Selection
|
||||
|
||||
@@ -463,92 +463,90 @@ The attention system uses a **backend selector pattern** that automatically choo
|
||||
|
||||
```python
|
||||
class Attention(nn.Module):
|
||||
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(...)
|
||||
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)
|
||||
```
|
||||
|
||||
**Available Backends**:
|
||||
|
||||
- **FlashAttention**: Optimized CUDA kernel (FA2/FA3) - memory efficient via tiling
|
||||
- **FlashAttention** (`FLASH_ATTN`): Memory-efficient kernel that dispatches to FA2/FA3 on CUDA, AITER on ROCm, and `mindiesd` on Ascend NPU.
|
||||
|
||||
- **SDPA**: PyTorch's scaled dot-product attention - default, cross-platform
|
||||
- **SDPA** (`TORCH_SDPA`): PyTorch's scaled dot-product attention — cross-platform fallback.
|
||||
|
||||
- **SageAttention**: Sparse attention implementation from SageAttention library
|
||||
|
||||
- **AscendAttention**: NPU-optimized attention for Ascend hardware
|
||||
- **SageAttention** (`SAGE_ATTN`): Quantized attention implementation from the SageAttention library.
|
||||
|
||||
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(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
|
||||
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
|
||||
```
|
||||
|
||||
**Selection Priority**:
|
||||
|
||||
1. **Environment Variable**: `DIFFUSION_ATTENTION_BACKEND` for manual override
|
||||
1. **Per-role override** (`--diffusion-attention-config.per_role.<role>.backend`) — finest control; matched on the layer's exact `role` string.
|
||||
|
||||
- Valid values: `FLASH_ATTN`, `TORCH_SDPA`, `SAGE_ATTN`, `ASCEND`
|
||||
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"`).
|
||||
|
||||
- Example: `export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN`
|
||||
3. **Global default** (`--diffusion-attention-backend FLASH_ATTN`, or `--diffusion-attention-config.default.backend`, or `DIFFUSION_ATTENTION_BACKEND` env var).
|
||||
|
||||
2. **Automatic Fallback**: Falls back to SDPA if selected backend unavailable
|
||||
4. **Platform default** — `current_omni_platform.get_diffusion_attn_backend_cls(...)` picks the best-available kernel for the hardware.
|
||||
|
||||
3. **Hardware Detection**: Can select based on device type (NPU, CUDA, etc.)
|
||||
`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).
|
||||
|
||||
**Backend Availability**:
|
||||
|
||||
- **SDPA**: Always available (PyTorch built-in)
|
||||
|
||||
- **FlashAttention**: Requires `flash-attn` package installed
|
||||
- **FlashAttention**: Requires `flash-attn` on CUDA, `aiter` on ROCm, or `mindiesd` on Ascend NPU
|
||||
|
||||
- **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`
|
||||
**Location**: `vllm_omni/diffusion/attention/selector.py` + `vllm_omni/platforms/<device>/platform.py`
|
||||
|
||||
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**:
|
||||
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)`.
|
||||
|
||||
```python
|
||||
# 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",
|
||||
},
|
||||
}
|
||||
@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)
|
||||
```
|
||||
|
||||
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:
|
||||
|
||||
@@ -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, 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).
|
||||
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, and 0.20.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 --torch-backend=auto
|
||||
uv pip install vllm==0.20.0 --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.19.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.19.0 --torch-backend=auto
|
||||
uv pip install vllm==0.20.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
The 0.19.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
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.19.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
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.19.0/vllm-0.19.0+cu130-cp38-abi3-manylinux_2_35_x86_64.whl
|
||||
#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
|
||||
@@ -80,7 +80,7 @@ uv pip install --editable .
|
||||
Install vllm with command below (If you already have PyTorch).
|
||||
```bash
|
||||
python use_existing_torch.py
|
||||
uv pip install -r requirements/build.txt
|
||||
uv pip install -r requirements/build/cuda.txt
|
||||
uv pip install --no-build-isolation --editable .
|
||||
```
|
||||
</details>
|
||||
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
--ipc=host \
|
||||
vllm/vllm-omni:v0.18.0 \
|
||||
vllm/vllm-omni:v0.20.0 \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
|
||||
```
|
||||
|
||||
|
||||
@@ -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.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
|
||||
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.19.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.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
|
||||
uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
|
||||
```
|
||||
|
||||
The pre-built 0.19.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
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.19.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.19.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.19.0
|
||||
vllm/vllm-omni-rocm:v0.20.0
|
||||
```
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
|
||||
source .venv/bin/activate
|
||||
|
||||
# On CUDA
|
||||
uv pip install vllm==0.19.0 --torch-backend=auto
|
||||
uv pip install vllm==0.20.0 --torch-backend=auto
|
||||
|
||||
# On ROCm
|
||||
uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
|
||||
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
|
||||
@@ -31,6 +31,12 @@ uv pip install -e .
|
||||
|
||||
For additional installation methods — please see the [installation guide](installation/README.md).
|
||||
|
||||
|
||||
!!! 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.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
|
||||
|
||||
Text-to-image generation quickstart with vLLM-Omni:
|
||||
|
||||
@@ -121,6 +121,7 @@ def extract_omni_serve_subparser_init():
|
||||
"FlexibleArgumentParser": _FlexibleArgumentParser,
|
||||
"make_arg_parser": lambda parser: parser, # no-op for doc
|
||||
"_ensure_vllm_platform": lambda: None, # no-op for doc
|
||||
"nullify_stage_engine_defaults": lambda parser: None, # no-op for doc
|
||||
"VLLM_SUBCMD_PARSER_EPILOG": "",
|
||||
"logger": logger,
|
||||
"DummySubparsers": DummySubparsers,
|
||||
|
||||
@@ -20,6 +20,7 @@ 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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
@@ -31,11 +32,14 @@ th {
|
||||
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||
| `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
|
||||
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
|
||||
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
@@ -50,17 +54,21 @@ th {
|
||||
| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
|
||||
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
|
||||
Loaded 100 of 973 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user