Compare commits
420
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4bbe81308b | ||
|
|
74258b18b7 | ||
|
|
1c876f0cd1 | ||
|
|
ee212d12b1 | ||
|
|
5ab2932ce6 | ||
|
|
2b4a76502f | ||
|
|
414d14f37e | ||
|
|
8593781a07 | ||
|
|
c20afebce4 | ||
|
|
eb0c4b1de4 | ||
|
|
0e3bf4de8e | ||
|
|
b72ffb82bb | ||
|
|
f035107ba0 | ||
|
|
65784f8199 | ||
|
|
168033cac0 | ||
|
|
6fa9dbbe42 | ||
|
|
beb00fce31 | ||
|
|
3a347b10b0 | ||
|
|
0f3264a9c0 | ||
|
|
dab596e468 | ||
|
|
2b999429c4 | ||
|
|
6031c7d393 | ||
|
|
373d3c7476 | ||
|
|
1903cee56a | ||
|
|
39705f9ee8 | ||
|
|
8766c0c59f | ||
|
|
6ea66f8ed0 | ||
|
|
220ce0b4ea | ||
|
|
5c3c81920b | ||
|
|
1f3a3e3848 | ||
|
|
d37912b071 | ||
|
|
37eebff674 | ||
|
|
a698269e47 | ||
|
|
36dc32ff81 | ||
|
|
9abf9545df | ||
|
|
c33ef34654 | ||
|
|
c349f62839 | ||
|
|
41cb6378a2 | ||
|
|
753508fe02 | ||
|
|
a3802d1e5c | ||
|
|
3c48257f8d | ||
|
|
283873411a | ||
|
|
1f1c82ecdd | ||
|
|
e459fdb30a | ||
|
|
9322cc6176 | ||
|
|
abc31731ab | ||
|
|
385df0359e | ||
|
|
36eb8afba5 | ||
|
|
ec94e83a25 | ||
|
|
cad25fa20c | ||
|
|
11a951b962 | ||
|
|
410a6a20b0 | ||
|
|
52020c7af7 | ||
|
|
8f45e68bcd | ||
|
|
80dea88007 | ||
|
|
e996723bc9 | ||
|
|
5f9aee193e | ||
|
|
9a37ca83d4 | ||
|
|
b52a5ca401 | ||
|
|
c890b3a99c | ||
|
|
6c1bb27b69 | ||
|
|
cc0e15ba57 | ||
|
|
28ce618f0e | ||
|
|
748470c1fb | ||
|
|
d01eaca39a | ||
|
|
4e4458e7bf | ||
|
|
e7644daa7f | ||
|
|
d714b4ce04 | ||
|
|
23e243378f | ||
|
|
5799d85a51 | ||
|
|
89e389cf37 | ||
|
|
90ed618b5a | ||
|
|
28467dc19a | ||
|
|
f0af3ab200 | ||
|
|
0ed3186db1 | ||
|
|
8e937400dc | ||
|
|
076c34abfb | ||
|
|
7d8816b46b | ||
|
|
1de15bdf4d | ||
|
|
e949ccf087 | ||
|
|
bd1617ce2b | ||
|
|
2bc9cb7f13 | ||
|
|
8297570c2b | ||
|
|
474710dda9 | ||
|
|
7d9d8aee79 | ||
|
|
083b5e363f | ||
|
|
f74cba84ac | ||
|
|
2aafe0bf94 | ||
|
|
5cf4605bf9 | ||
|
|
e1c6bc7597 | ||
|
|
fc8486cf20 | ||
|
|
1c530d984d | ||
|
|
9dd36e3afd | ||
|
|
109bba8cdf | ||
|
|
f99e8b2d3c | ||
|
|
06ab723afc | ||
|
|
12bcc32c3d | ||
|
|
d4e90bbba8 | ||
|
|
29179599fc | ||
|
|
4e82192fd6 | ||
|
|
a3d4ed809d | ||
|
|
3c58868c9a | ||
|
|
da53618793 | ||
|
|
18186db216 | ||
|
|
14e5baceaf | ||
|
|
2af2a50e0e | ||
|
|
bd83ac9b4b | ||
|
|
e277feacaf | ||
|
|
ca9fd0b71c | ||
|
|
3ac739817f | ||
|
|
1fa734419e | ||
|
|
2c6b1bb0c0 | ||
|
|
e2ed1c4574 | ||
|
|
89f8819525 | ||
|
|
475a4002b0 | ||
|
|
ab59673f21 | ||
|
|
821286794f | ||
|
|
309e5c38c6 | ||
|
|
f4115bd771 | ||
|
|
dbc589dbca | ||
|
|
990566aef1 | ||
|
|
6d37e77fb2 | ||
|
|
4ba8e14981 | ||
|
|
c99df1ebd9 | ||
|
|
0a395f9de1 | ||
|
|
c0e132d973 | ||
|
|
471ddfe025 | ||
|
|
8cfc9179e6 | ||
|
|
0da9ffdb0d | ||
|
|
4e88053750 | ||
|
|
768943b879 | ||
|
|
220db62b3f | ||
|
|
5549b7f44a | ||
|
|
bc26cad19a | ||
|
|
9c5e35f485 | ||
|
|
a64ebf103b | ||
|
|
c08959ee04 | ||
|
|
c5ac295e3c | ||
|
|
ea35a0cc4a | ||
|
|
0f4853ff86 | ||
|
|
b5e163cfca | ||
|
|
d647e7e4cf | ||
|
|
33220b1e39 | ||
|
+20 |
eb4e60ee64 | ||
|
|
5e1986206f | ||
|
|
d1c65bdffa | ||
|
|
d18168ccbb | ||
|
|
779bf3118b | ||
|
|
e7ee5de09f | ||
|
|
440c718d2a | ||
|
|
82a0b3a467 | ||
|
|
c7178d89bb | ||
|
|
fdb0efea94 | ||
|
|
90f5b3c3a1 | ||
|
|
bbc00f9f86 | ||
|
|
adb2291c27 | ||
|
|
4f13b871f9 | ||
|
|
94254e015f | ||
|
|
f7161b07d0 | ||
|
|
c0b7509f07 | ||
|
|
3f63aaf982 | ||
|
|
c4f859bf56 | ||
|
|
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
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
@@ -1,20 +0,0 @@
|
||||
steps:
|
||||
- label: "Build and upload wheel"
|
||||
key: "build-wheel"
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "curl -LsSf https://astral.sh/uv/install.sh | sh"
|
||||
- 'export PATH="$HOME/.local/bin:$PATH"'
|
||||
- "uv venv --python=3.12 && source .venv/bin/activate"
|
||||
- "uv pip install --upgrade build"
|
||||
- "python3 -m build"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
|
||||
- label: "Generate and upload wheel indices"
|
||||
depends_on: "build-wheel"
|
||||
allow_dependency_failure: true
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
|
||||
@@ -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.21.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,112 @@
|
||||
steps:
|
||||
- label: ":buildkit: Build and Push NPU Test Image - A2"
|
||||
key: image-build
|
||||
agents:
|
||||
queue: "ascend-a2b3"
|
||||
resource_class: "npu-1"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
metadata:
|
||||
annotations:
|
||||
vault.hashicorp.com/agent-init-first: "true"
|
||||
vault.hashicorp.com/agent-inject: "true"
|
||||
vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
|
||||
vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
|
||||
vault.hashicorp.com/agent-inject-perms-config.json: "0400"
|
||||
vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
|
||||
vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
|
||||
vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
|
||||
vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
|
||||
vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
|
||||
vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
|
||||
vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
|
||||
vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
|
||||
vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
|
||||
vault.hashicorp.com/agent-pre-populate-only: "true"
|
||||
vault.hashicorp.com/agent-run-as-group: "1000"
|
||||
vault.hashicorp.com/agent-run-as-user: "1000"
|
||||
vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
|
||||
vault.hashicorp.com/role: ascend-gha-runners
|
||||
vault.hashicorp.com/secret-volume-path: /home/user/.docker/
|
||||
vault.hashicorp.com/tls-skip-verify: "true"
|
||||
podSpecPatch:
|
||||
volumes:
|
||||
- name: token-vol
|
||||
projected:
|
||||
defaultMode: 420
|
||||
sources:
|
||||
- serviceAccountToken:
|
||||
audience: api
|
||||
expirationSeconds: 600
|
||||
path: token
|
||||
env:
|
||||
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
|
||||
IMAGE_NAME: "vllm-omni-ci-npu"
|
||||
IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
|
||||
BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
|
||||
command: |
|
||||
set -ex
|
||||
|
||||
echo "--- Building and pushing NPU Test Image"
|
||||
echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
|
||||
echo "buildkitd address: $${BUILDKITD_ADDR}"
|
||||
|
||||
if ! command -v buildctl &> /dev/null; then
|
||||
echo "Installing buildctl..."
|
||||
mkdir -p /tmp/buildkit
|
||||
BUILDKIT_VERSION="v0.29.0"
|
||||
wget -q "https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
|
||||
tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
|
||||
cp /tmp/buildkit/bin/buildctl /usr/local/bin/
|
||||
fi
|
||||
|
||||
export DOCKER_CONFIG=/home/user/.docker
|
||||
buildctl \
|
||||
--addr="$${BUILDKITD_ADDR}" \
|
||||
--tlscacert=/home/user/.docker/ca.pem \
|
||||
--tlscert=/home/user/.docker/cert.pem \
|
||||
--tlskey=/home/user/.docker/key.pem \
|
||||
build \
|
||||
--frontend dockerfile.v0 \
|
||||
--local context=. \
|
||||
--local dockerfile=./docker \
|
||||
--opt filename=Dockerfile.npu.ci \
|
||||
--secret id=dockerconfig,src=/home/user/.docker/config.json \
|
||||
--output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
|
||||
--progress=plain
|
||||
|
||||
echo "--- Image pushed successfully"
|
||||
echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
|
||||
|
||||
- label: "Upload Ready Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
if: >-
|
||||
build.branch != "main" && (
|
||||
build.pull_request.labels includes "npu-test" ||
|
||||
build.pull_request.labels includes "ready"
|
||||
)
|
||||
agents:
|
||||
queue: "ascend-a2b3"
|
||||
resource_class: "npu-1"
|
||||
command: |
|
||||
buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml
|
||||
|
||||
# L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
|
||||
- label: "Upload Nightly Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-nightly-pipeline
|
||||
if: >-
|
||||
(build.branch == "main" && build.env("NIGHTLY") == "1") ||
|
||||
(build.branch != "main" && (
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "omni-test" ||
|
||||
build.pull_request.labels includes "tts-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2iat-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2v-test"
|
||||
))
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-npu-nightly.yaml
|
||||
agents:
|
||||
queue: "ascend-a2b3"
|
||||
resource_class: "npu-1"
|
||||
+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"
|
||||
@@ -0,0 +1,38 @@
|
||||
steps:
|
||||
- label: ":docker: Build image"
|
||||
key: image-build
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
|
||||
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L2 Test
|
||||
- label: "Upload Ready Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L3 Test
|
||||
- label: "Upload Merge Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L4 Test
|
||||
- label: "Upload Nightly Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
@@ -1,38 +1,110 @@
|
||||
steps:
|
||||
- label: ":docker: Build image"
|
||||
key: image-build
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
|
||||
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
- input: "Provide Release version here"
|
||||
id: input-release-version
|
||||
fields:
|
||||
- text: "What is the release version?"
|
||||
key: release-version
|
||||
|
||||
# L2 Test
|
||||
- label: "Upload Ready Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
- label: "Build and upload wheel"
|
||||
depends_on: ~
|
||||
key: "build-wheel"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "curl -LsSf https://astral.sh/uv/install.sh | sh"
|
||||
- 'export PATH="$HOME/.local/bin:$PATH"'
|
||||
- "uv venv --python=3.12 && source .venv/bin/activate"
|
||||
- "uv pip install --upgrade build"
|
||||
- "python3 -m build"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
|
||||
# L3 Test
|
||||
- label: "Upload Merge Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
- label: "Generate and upload wheel indices"
|
||||
depends_on: "build-wheel"
|
||||
allow_dependency_failure: true
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
|
||||
|
||||
# L4 Test
|
||||
- label: "Upload Nightly Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
|
||||
- block: "Upload release wheels to PyPI"
|
||||
depends_on: "build-wheel"
|
||||
key: block-upload-release-wheels
|
||||
|
||||
- label: "Upload release wheels to PyPI"
|
||||
depends_on:
|
||||
- block-upload-release-wheels
|
||||
- input-release-version
|
||||
key: upload-release-wheels
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/publish-release-wheels-pypi.sh"
|
||||
|
||||
- block: "Unblock to build release Docker images"
|
||||
depends_on: ~
|
||||
key: block-build-release-images
|
||||
|
||||
- group: "Build release Docker images"
|
||||
key: "build-release-images"
|
||||
depends_on: block-build-release-images
|
||||
allow_dependency_failure: true
|
||||
steps:
|
||||
- label: "Build release image - x86_64"
|
||||
key: build-release-image-x86
|
||||
depends_on: ~
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
--build-arg max_jobs=16 \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m) \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile.cuda .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
|
||||
# re-tag to default image tag and push, just in case arm64 build fails
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- label: "Build release image - aarch64"
|
||||
key: build-release-image-arm64
|
||||
depends_on: ~
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
--build-arg max_jobs=16 \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m) \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile.cuda .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- block: "Publish release images to DockerHub"
|
||||
key: block-publish-release-images
|
||||
depends_on:
|
||||
- build-release-image-x86
|
||||
- build-release-image-arm64
|
||||
- input-release-version
|
||||
|
||||
- label: "Publish release images to DockerHub"
|
||||
depends_on:
|
||||
- block-publish-release-images
|
||||
key: publish-release-images-dockerhub
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/publish-release-images.sh"
|
||||
plugins:
|
||||
- docker-login#v3.0.0:
|
||||
username: vllmbot
|
||||
password-env: DOCKERHUB_TOKEN
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
DOCKERHUB_USERNAME: "vllmbot"
|
||||
@@ -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.
|
||||
|
||||
@@ -52,5 +52,7 @@ time timeout -k 30 30m docker run \
|
||||
pip install tblib==3.1.0
|
||||
cd /workspace/vllm-omni
|
||||
pytest -v -s -m "core_model and xpu and B60"
|
||||
pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py
|
||||
export VLLM_XPU_USE_SAMPLER_KERNEL=0 # NOTE: Remove this after vLLM v0.21.1 is merged. Fixes Qwen2-5 omni-expansion tests.
|
||||
pytest -v -s -m "advanced_model and xpu and B60"
|
||||
'
|
||||
@@ -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
|
||||
'
|
||||
Executable
+50
@@ -0,0 +1,50 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Publish release Docker images from ECR to DockerHub.
|
||||
# Pulls per-arch images, tags with `latest` and versioned tags, pushes them,
|
||||
# then creates and pushes a multi-arch manifest.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
|
||||
if [ -z "${RELEASE_VERSION}" ]; then
|
||||
echo "ERROR: release-version metadata not set"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
COMMIT="$BUILDKITE_COMMIT"
|
||||
DOCKERHUB_REPO="vllm/vllm-omni"
|
||||
ECR_REPO="public.ecr.aws/q9t5s3a7/vllm-omni-release-repo"
|
||||
|
||||
echo "========================================"
|
||||
echo "Publishing release images v${RELEASE_VERSION}"
|
||||
echo " Commit: ${COMMIT}"
|
||||
echo "========================================"
|
||||
|
||||
# Login to ECR to pull staging images
|
||||
aws ecr-public get-login-password --region us-east-1 | \
|
||||
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
|
||||
|
||||
docker pull "${ECR_REPO}:${COMMIT}-x86_64"
|
||||
docker pull "${ECR_REPO}:${COMMIT}-aarch64"
|
||||
|
||||
docker tag "${ECR_REPO}:${COMMIT}-x86_64" "${DOCKERHUB_REPO}:latest-x86_64"
|
||||
docker tag "${ECR_REPO}:${COMMIT}-x86_64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64"
|
||||
docker push "${DOCKERHUB_REPO}:latest-x86_64"
|
||||
docker push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64"
|
||||
|
||||
docker tag "${ECR_REPO}:${COMMIT}-aarch64" "${DOCKERHUB_REPO}:latest-aarch64"
|
||||
docker tag "${ECR_REPO}:${COMMIT}-aarch64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
|
||||
docker push "${DOCKERHUB_REPO}:latest-aarch64"
|
||||
docker push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
|
||||
|
||||
docker manifest rm "${DOCKERHUB_REPO}:latest" || true
|
||||
docker manifest rm "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" || true
|
||||
docker manifest create "${DOCKERHUB_REPO}:latest" "${DOCKERHUB_REPO}:latest-x86_64" "${DOCKERHUB_REPO}:latest-aarch64"
|
||||
docker manifest create "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
|
||||
docker manifest push "${DOCKERHUB_REPO}:latest"
|
||||
docker manifest push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}"
|
||||
|
||||
echo ""
|
||||
echo "Successfully published release images for v${RELEASE_VERSION}"
|
||||
+52
@@ -0,0 +1,52 @@
|
||||
#!/usr/bin/env bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
set -e
|
||||
|
||||
BUCKET="vllm-wheels"
|
||||
SUBPATH="omni/$BUILDKITE_COMMIT"
|
||||
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
|
||||
|
||||
RELEASE_VERSION=$(buildkite-agent meta-data get release-version | sed 's/^v//')
|
||||
if [[ -z "$RELEASE_VERSION" ]]; then
|
||||
echo "[FATAL] release-version metadata not set."
|
||||
exit 1
|
||||
fi
|
||||
echo "Release version: $RELEASE_VERSION"
|
||||
|
||||
if [[ -z "$PYPI_TOKEN" ]]; then
|
||||
echo "[FATAL] PYPI_TOKEN is not set."
|
||||
exit 1
|
||||
else
|
||||
export TWINE_USERNAME="__token__"
|
||||
export TWINE_PASSWORD="$PYPI_TOKEN"
|
||||
fi
|
||||
|
||||
set -x
|
||||
|
||||
if ! command -v uv &> /dev/null; then
|
||||
curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
export PATH="$HOME/.local/bin:$PATH"
|
||||
fi
|
||||
|
||||
uv venv --python 3.12 /tmp/vllm-omni-release-env
|
||||
source /tmp/vllm-omni-release-env/bin/activate
|
||||
uv pip install twine
|
||||
|
||||
DIST_DIR=/tmp/vllm-omni-release-dist
|
||||
mkdir -p "$DIST_DIR"
|
||||
|
||||
echo "Downloading wheels from S3:"
|
||||
aws s3 ls "$S3_COMMIT_PREFIX"
|
||||
aws s3 cp --recursive --exclude "*" --include "vllm_omni-${RELEASE_VERSION}*.whl" --exclude "*dev*" "$S3_COMMIT_PREFIX" "$DIST_DIR"
|
||||
ls -la "$DIST_DIR"
|
||||
|
||||
PYPI_WHEEL_FILES=$(find "$DIST_DIR" -name "vllm_omni-${RELEASE_VERSION}*.whl" -not -name "*+*")
|
||||
if [[ -z "$PYPI_WHEEL_FILES" ]]; then
|
||||
echo "[FATAL] No wheels found for version ${RELEASE_VERSION}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
python3 -m twine check $PYPI_WHEEL_FILES
|
||||
python3 -m twine upload --non-interactive --verbose $PYPI_WHEEL_FILES
|
||||
echo "Wheels uploaded to PyPI"
|
||||
@@ -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"
|
||||
+752
-388
File diff suppressed because it is too large.
Load diff
+780
-116
File diff suppressed because it is too large.
Load diff
@@ -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"
|
||||
@@ -0,0 +1,21 @@
|
||||
steps:
|
||||
- group: ":card_index_dividers: TTS Model Test"
|
||||
key: ready-tts-group
|
||||
depends_on: upload-ready-pipeline
|
||||
steps:
|
||||
- label: ":sound: TTS · Qwen3-TTS CustomVoice Test"
|
||||
agents:
|
||||
queue: "ascend-a2b3"
|
||||
resource_class: "npu-1"
|
||||
image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpecPatch:
|
||||
imagePullSecrets:
|
||||
- name: swr-secret
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model' --run-level 'core_model'
|
||||
+452
-549
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,148 @@
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_USE_FLASHINFER_MOE_FP16: 0
|
||||
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
|
||||
|
||||
# Invalid HTTP / WS payloads vs live Omni servers (Issue #3649 tracks known skips).
|
||||
- group: "Reliability Test - Invalid parameters Test"
|
||||
depends_on: upload-weekly-pipeline
|
||||
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
|
||||
steps:
|
||||
- label: "Invalid parameters Test · H100"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# Stable Audio Open, Qwen3-TTS speech/batch/voices paths marked ``L4``.
|
||||
- label: "Invalid parameters Test · L4"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
@@ -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`.
|
||||
@@ -0,0 +1,510 @@
|
||||
---
|
||||
name: diffusion-perf-opt
|
||||
description: Diagnose and optimize vLLM Omni diffusion workloads, especially Wan/Qwen/Flux-style image and video generation. Use when Codex is asked to analyze profiling traces, choose parallel strategies, inspect torch profiler trace.json or trace.json.gz timelines, estimate optimization ROI, investigate GPU idle/free bubbles, compare USP/CFG/HSDP/VAE parallelism, or design operator/host/quantization optimizations for vLLM Omni.
|
||||
---
|
||||
|
||||
# vLLM Omni Performance Optimization
|
||||
|
||||
Use this skill to run a disciplined optimization loop for vLLM Omni diffusion workloads. Keep two ideas separate: real performance baselines are collected with low overhead, while torch profiler traces are diagnostic artifacts and may distort latency.
|
||||
|
||||
## First Questions
|
||||
|
||||
Before proposing changes, ask for the optimization scene if it is not already known:
|
||||
|
||||
- GPU model, card count, topology, and whether NVLink is present.
|
||||
- Model and pipeline, for example Wan2.2 I2V A14B.
|
||||
- User workload: resolution, frames, steps, batch/concurrency, CFG scales, prompt/image inputs.
|
||||
- Three runnable commands for each model/workload, or scripts that generate them:
|
||||
- **Server startup command**: exact `vllm serve` command, environment variables, model path, port, parallelism flags, profiler flags, and precision/compile settings.
|
||||
- **User request command**: exact single-request client command (for example curl or Python client) with fixed prompt/media/seed/size/steps. Use this to validate correctness and collect per-request stage timings.
|
||||
- **User benchmark command**: exact repeatable benchmark command or script with warmup count, measured iteration count, concurrency/batch policy, output directory, and summary format.
|
||||
- Current enabled strategies: USP/SP, CFG parallel, HSDP/FSDP, VAE patch parallel, torch.compile, profiler options.
|
||||
- Optimization target: latency, throughput, memory, cost, or quality-preserving speed.
|
||||
- Precision/quality tolerance: bf16/fp8/quantization/sparsity/approximate attention allowed or not.
|
||||
- Quality validation method: exact-output comparison when feasible, or metrics
|
||||
such as SSIM, PSNR, LPIPS, cosine similarity, MAE, MSE, temporal flicker
|
||||
checks, baseline self-run variance, and reviewable output artifacts.
|
||||
|
||||
## Workflow
|
||||
|
||||
1. **Freeze the measurement protocol and commands.**
|
||||
- Before analyzing parallel strategies or traces, establish the three commands for the target model/workload: server startup, single user request, and user benchmark.
|
||||
- Prefer checked-in scripts or one-off shell scripts over free-form commands in chat. The scripts should make workload variables explicit, including model path, port/base URL, prompt/media inputs, resolution, frames, steps, seed, warmup, measured iterations, concurrency, and output path.
|
||||
- Measurement reports must preserve the concrete process, not just the final numbers. For each tested configuration, write down the server startup command, the user request command, the user benchmark command or polling loop, the final server response metrics, and the client-side timing/HTTP result.
|
||||
- Present all duration and latency values in **milliseconds (ms)** in tables and summaries. If an API returns seconds, convert to ms before reporting. Keep raw units only inside quoted raw logs or source snippets.
|
||||
- Keep profiler modes separate in the command set:
|
||||
- Baseline/benchmark commands should avoid torch profiler and stack collection.
|
||||
- Diagnostic commands may enable torch profiler after a baseline identifies a bottleneck.
|
||||
- If any of the three commands is missing, ask for it or create a proposed script before proceeding.
|
||||
- Make clear which command is authoritative for correctness validation and which command is authoritative for performance numbers.
|
||||
|
||||
2. **Collect a real baseline.**
|
||||
- Disable PyTorch profiler and stack collection.
|
||||
- Disable or fix torch.compile state so A/B is fair.
|
||||
- Avoid `--enforce-eager` for production-speed baselines unless eager is the target.
|
||||
- Prefer `--log-stats` and `--enable-diffusion-pipeline-profiler` for low-overhead stage timing. PR 3069 has the relevant metrics/log-stats changes; if local code does not include them, fetch or cherry-pick the minimal metrics changes rather than merging unrelated PR drift.
|
||||
- Run warmup requests and exclude first-request lazy init.
|
||||
|
||||
3. **Model the parallel strategy before testing.**
|
||||
- Estimate compute and communication for self-attention, cross-attention, FFN, CFG, VAE encode/decode, and HSDP.
|
||||
- Select a small candidate matrix rather than exhaustively testing everything.
|
||||
- Typical 2-card candidates: `USP=2`, `CFG=2`, `USP=1/HSDP on-off`, VAE parallel on-off if memory allows.
|
||||
- Typical 4-card candidates: `USP=4`, `CFG=2 x USP=2`, `USP=2 x HSDP`, VAE parallel on-off.
|
||||
- Typical 8-card candidates for official CFG-enabled video diffusion models:
|
||||
- Primary candidate: `CFG=2 x USP=4` with VAE patch parallel across all 8 ranks.
|
||||
- Compare against: `CFG=1 x USP=8` to test whether larger sequence parallel groups beat CFG branch parallelism.
|
||||
- Isolate VAE parallelism: keep the DiT strategy fixed and compare VAE patch parallel on/off or different VAE patch sizes.
|
||||
- If long-video `diffuse` remains dominant and Ulysses all-to-all is suspected, test a hybrid sequence strategy such as `CFG=2 x USP=2 x Ring=2`.
|
||||
- Test HSDP off only if the model fits without it; treat HSDP primarily as a memory strategy until A/B proves otherwise.
|
||||
- Prefer USP/SP for long video token sequences; prefer CFG parallel when CFG doubles transformer forwards and sequence length is modest.
|
||||
- Convert workload shape to latent/patch token counts before choosing candidates. For video models, record frames, latent frames, latent HxW, patch size, approximate token count, and which stages should scale with the token count.
|
||||
|
||||
4. **Search for the best parallel configuration.**
|
||||
- Start with a small matrix that answers one question per comparison:
|
||||
- `CFG parallel vs larger USP`: compare `CFG=2 x USP=world/2` against `CFG=1 x USP=world` for CFG-enabled workloads.
|
||||
- `VAE patch parallel value`: compare the best DiT strategy with VAE patch parallel enabled and disabled.
|
||||
- `HSDP cost`: compare HSDP on/off only when both configurations fit memory.
|
||||
- `Ulysses vs Ring`: test a Ring/Ulysses hybrid only after long-sequence `diffuse` is confirmed dominant.
|
||||
- For each configuration, create or record a stable config id, for example `A_cfg2_usp4_vaepp8_hsdp_tiling`.
|
||||
- For each config id, capture:
|
||||
- exact server command and environment variables,
|
||||
- observed distributed setup from logs, such as SP groups, CFG groups, HSDP shard/replicate sizes, VAE patch size,
|
||||
- exact request command for every scenario,
|
||||
- final server response metrics and client-side elapsed time,
|
||||
- output artifact paths and any failed/empty responses.
|
||||
- Use one warmup request per scenario, then at least three measured repeats for the shortlisted config. For exploratory matrix pruning, one measured request is acceptable only if the margin is large; label it as one-shot.
|
||||
- Compare configurations by stage, not only end-to-end latency. A configuration can improve `diffuse` while hurting `vae.decode`; record both effects.
|
||||
- Select the best config only after checking the target metric, dominant stages, memory headroom, and output correctness.
|
||||
|
||||
5. **Run targeted A/B tests.**
|
||||
- Change one variable per test.
|
||||
- Keep model, input, seed, request parameters, GPU placement, and warmup policy fixed.
|
||||
- Record latency, stage timings, memory, output quality, and logs.
|
||||
- Report comparison tables in ms. Include at least end-to-end client time, server `inference_time`, server stage generation time if available, `vae.encode`, `diffuse`, `vae.decode`, and peak memory.
|
||||
|
||||
6. **Enforce the quality and precision gate for every optimization.**
|
||||
- Do not mark an optimization ready until quality is checked on the same
|
||||
workload: model, prompt/media, seed, shape, steps, scheduler, dtype,
|
||||
backend, parallelism, and output encoding.
|
||||
- For math-preserving changes, require exact/near-exact agreement or show
|
||||
that differences are within baseline self-run variance.
|
||||
- For precision, quantization, approximate attention, sparse/custom kernels,
|
||||
or backend changes, include reviewable artifacts and metrics such as SSIM,
|
||||
PSNR, LPIPS, cosine similarity, MAE, MSE, and temporal flicker checks.
|
||||
- If quality tolerance is not stated, default to quality-preserving behavior.
|
||||
Failed or inconclusive quality validation blocks a ready-to-merge claim.
|
||||
|
||||
7. **Collect diagnostic trace only after narrowing hypotheses.**
|
||||
- Use torch profiler for a small number of requests.
|
||||
- Run two separate diagnostic traces instead of mixing concerns:
|
||||
- **Operator/shape trace**: enable `torch_profiler_record_shapes=True` and keep stack collection disabled. Use this to rank CUDA kernels, NCCL collectives, attention/MLP/norm/RoPE work, and shape-specific hot operators.
|
||||
- **Host-stack trace**: enable `torch_profiler_with_stack=True` and normally keep shape collection disabled. Use this to map CPU/Python host gaps, synchronization points, scheduler paths, and request handling overhead.
|
||||
- Keep both trace commands and reports separate from baseline/benchmark commands. Torch profiler latency is diagnostic only and must not be used as the final latency claim.
|
||||
- Prefer profiling only the narrowed dominant scenario, for example the highest-resolution/video-length case where `diffuse` dominates.
|
||||
- If profiler endpoints are available, run one warmup request first, then call `/start_profile`, run one profiled request, and call `/stop_profile`. This keeps model initialization and warmup out of the diagnostic trace.
|
||||
- Start by analyzing rank 0 only. Expand to more ranks only if rank 0 suggests imbalance, unclear GPU idle/free bubbles, high NCCL wait, server timing mismatch, CFG branch imbalance, or USP group stragglers.
|
||||
- Diagnostic reports must be written to disk and preserve: server command, profiler config, warmup/request/polling commands, trace artifact paths, rank analyzed, analyzer output or summary, and the decision about whether additional ranks are necessary.
|
||||
- Analyze both rank-level balance and device-level free bubbles when additional ranks are opened.
|
||||
|
||||
8. **Analyze host, communication, and operators.**
|
||||
- Find GPU idle/free intervals and map each large gap to the enclosing CPU/Python code.
|
||||
- Separate real GPU idle from profiler overhead such as CUPTI `Command Buffer Full`.
|
||||
- Compare NCCL kernel time to user annotations; annotations can overcount nested intervals.
|
||||
- Rank operator work by total CUDA time and by repeated small-kernel launch count.
|
||||
|
||||
9. **Produce an optimization plan.**
|
||||
- Classify candidates as P0/P1/P2.
|
||||
- For each candidate, state necessity, expected benefit, implementation path, validation plan, and quality risk.
|
||||
- Include a concrete quality/precision validation gate for each candidate.
|
||||
- Do not implement high-risk operator rewrites before proving the operator is a bottleneck for the target shapes.
|
||||
- End the plan with a user-facing candidate selection table. The assistant
|
||||
should not automatically choose a risky optimization just because it is
|
||||
technically possible. Present the options clearly and let the user decide
|
||||
which item is worth implementing based on latency target, engineering
|
||||
budget, memory headroom, and quality tolerance.
|
||||
- Organize the plan by optimization layer, not by a flat list of ideas:
|
||||
- host/runtime optimization,
|
||||
- measurement/benchmark reliability,
|
||||
- parallelism and communication,
|
||||
- VAE encode/decode and media pre/post processing,
|
||||
- operator fusion and layout cleanup,
|
||||
- attention main-path optimization,
|
||||
- algorithmic, precision, or approximation changes.
|
||||
- For each layer, tie every candidate back to evidence from baseline
|
||||
metrics, diagnostic trace, source code, or output quality requirements.
|
||||
Do not list generic optimizations without a trace or workload reason.
|
||||
|
||||
## Priority Rules
|
||||
|
||||
- **P0:** low risk, likely useful, or required for trustworthy measurement. Examples: real baseline, warmup, targeted parallel A/B, disabling avoidable `empty_cache`, scheduler coefficient caching.
|
||||
- **P1:** meaningful code changes with contained risk. Examples: cross-attention KV caching, VAE gather/broadcast reduction, AdaLayerNorm/RMSNorm/RoPE fusion after trace evidence.
|
||||
- **P2:** high implementation or quality risk. Examples: FA to LA replacement, custom Triton/CUDA fused kernels, FP8/quantization, sparsity/Rainfusion-style acceleration.
|
||||
|
||||
Every implemented optimization needs A/B validation and a passing
|
||||
quality/precision gate, including math-preserving changes such as padding trim,
|
||||
layout cleanup, cache reuse, or host/runtime cleanup. Use objective metrics
|
||||
when available, keep reviewable artifacts, and compare against baseline
|
||||
self-run variance for precision, quantization, approximate kernels, or backend
|
||||
changes. If quality validation fails or is inconclusive, do not present the
|
||||
optimization as ready to merge.
|
||||
|
||||
## Optimization Layers
|
||||
|
||||
After baseline, parallel-search, and diagnostic traces, summarize optimization
|
||||
opportunities by layer. This is the core of the performance analysis: the goal
|
||||
is to connect evidence to a scoped implementation and a validation plan.
|
||||
|
||||
### Host and Runtime Optimization
|
||||
|
||||
Purpose: remove CPU/Python stalls, synchronization points, allocator overhead,
|
||||
and request-path overhead that leave GPU lanes empty.
|
||||
|
||||
Evidence to look for:
|
||||
|
||||
- High `idle_pct` or large `GAP` blocks in `trace_analyzer.py`.
|
||||
- Host-stack trace lines such as `torch.cuda.empty_cache`,
|
||||
`cudaStreamSynchronize`, `cudaDeviceSynchronize`, Python locks, scheduler
|
||||
waits, image/video preprocessing, or repeated small allocation paths.
|
||||
- Difference between client wall-clock time and server `inference_time_s`.
|
||||
|
||||
Typical candidates:
|
||||
|
||||
- Make avoidable `torch.cuda.empty_cache()` optional or guard it by memory
|
||||
headroom.
|
||||
- Cache scheduler coefficients, timesteps, masks, or other tiny repeated CPU
|
||||
computations when the request shape/steps are fixed.
|
||||
- Remove avoidable host-device synchronizations and blocking logging/stat calls.
|
||||
- Move expensive preprocessing out of the critical path or cache fixed prompt,
|
||||
image, and transform work for benchmark scenarios.
|
||||
- Ensure benchmark scripts record client-side elapsed time, HTTP status, output
|
||||
path, and server response metrics.
|
||||
|
||||
Priority guidance:
|
||||
|
||||
- Usually P0 when the change is measurement reliability or an obvious removable
|
||||
synchronization.
|
||||
- Usually P1 when it changes scheduling, memory lifetime, or request execution
|
||||
order.
|
||||
|
||||
Validation:
|
||||
|
||||
- Re-run non-profiler baseline with same workload and seed.
|
||||
- Confirm peak memory headroom if disabling cache cleanup.
|
||||
- Confirm generated output exists and quality/seed behavior is unchanged.
|
||||
|
||||
### Parallelism and Communication Optimization
|
||||
|
||||
Purpose: choose the right decomposition for CFG branches, sequence tokens, model
|
||||
weights, VAE tiles, and rank topology.
|
||||
|
||||
Evidence to look for:
|
||||
|
||||
- Baseline A/B across `CFG`, `USP/SP`, `Ring`, `HSDP/FSDP`, and VAE patch
|
||||
parallelism.
|
||||
- Stage timing shifts: `diffuse`, `vae.encode`, `vae.decode`, and server
|
||||
end-to-end.
|
||||
- NCCL kernel time from trace, not only `user_annotation` time.
|
||||
- Rank imbalance across SP group ranks or CFG branch ranks.
|
||||
- Memory headroom and OOM risk.
|
||||
|
||||
Typical candidates:
|
||||
|
||||
- `CFG=2 x USP=world/2` versus `CFG=1 x USP=world` for CFG-enabled models.
|
||||
- VAE patch parallel on/off or patch size tuning.
|
||||
- HSDP/FSDP on/off only if both configurations fit memory.
|
||||
- Ulysses versus Ulysses+Ring only after long-sequence `diffuse` is confirmed
|
||||
dominant and all-to-all is suspected.
|
||||
- Rank mapping/topology changes if all-rank traces show stragglers or NCCL wait.
|
||||
- Buffer reuse or preallocation for FSDP/HSDP all-gather paths.
|
||||
|
||||
Priority guidance:
|
||||
|
||||
- P0/P1 for configuration-only changes with strong measured wins.
|
||||
- P1 for buffer reuse or rank mapping changes.
|
||||
- P2 for invasive distributed algorithm changes.
|
||||
|
||||
Validation:
|
||||
|
||||
- Measure by stage and memory, not only end-to-end.
|
||||
- Use one-variable A/B with identical prompt/media/seed/shape/steps.
|
||||
- When communication is suspected, compare rank0-3 in one USP group and rank0
|
||||
versus rank4 across CFG branches for `CFG=2 x USP=4`.
|
||||
|
||||
### VAE and Media Pipeline Optimization
|
||||
|
||||
Purpose: reduce encode/decode, tiling, split/gather, and media conversion time.
|
||||
|
||||
Evidence to look for:
|
||||
|
||||
- Large `vae.encode` or `vae.decode` in low-overhead stage timings.
|
||||
- Host-stack gaps in VAE tile split, gather, merge, broadcast, or image/video
|
||||
transforms.
|
||||
- VAE kernels or cuDNN convolution in operator trace.
|
||||
- Whether every rank needs the final decoded tensor.
|
||||
|
||||
Typical candidates:
|
||||
|
||||
- Keep VAE patch parallel enabled when it has clear measured benefit.
|
||||
- Reduce VAE gather/broadcast to only ranks that need the final media output.
|
||||
- Reuse tile metadata, split buffers, or gather buffers.
|
||||
- Evaluate bf16/autocast behavior for VAE only with visual quality checks.
|
||||
- Avoid redundant image conversion, resize, or tensor construction in repeated
|
||||
benchmark runs.
|
||||
|
||||
Priority guidance:
|
||||
|
||||
- P0/P1 if VAE is a large share of the target workload or if a host gap is
|
||||
obvious and low risk.
|
||||
- Lower priority when `diffuse` dominates and VAE is already patch-parallelized.
|
||||
|
||||
Validation:
|
||||
|
||||
- Compare `vae.encode`, `vae.decode`, server end-to-end, and peak memory.
|
||||
- Check output video integrity, artifacts, flicker, and seed stability.
|
||||
|
||||
### Operator Fusion and Layout Cleanup
|
||||
|
||||
Purpose: reduce high-frequency small kernels, memory bandwidth pressure, layout
|
||||
conversions, and launch overhead in transformer and VAE blocks.
|
||||
|
||||
Evidence to look for:
|
||||
|
||||
- Top operator tables showing many `aten::copy_`, `aten::cat`,
|
||||
`split_with_sizes_copy`, `aten::add`, `aten::mul`, `aten::div`, norm,
|
||||
activation, RoPE, or reshape/layout kernels.
|
||||
- `ops_rankN.xlsx` `by_shape` sheet showing repeated small shapes inside the
|
||||
same block path.
|
||||
- Trace lanes showing many short kernels between larger GEMM/attention kernels.
|
||||
- Source code patterns with repeated elementwise chains or layout conversions.
|
||||
|
||||
Typical fusion targets:
|
||||
|
||||
- AdaLayerNorm / RMSNorm / LayerNorm plus scale/shift fusion.
|
||||
- RoPE fusion with Q/K layout preparation when shapes are stable.
|
||||
- Residual add, scale, gate, and elementwise chains.
|
||||
- MLP gate/up/down cleanup, such as fusing activation and multiply around
|
||||
GEMM outputs when feasible.
|
||||
- QKV projection and reshape/split/cat path cleanup.
|
||||
- Attention pre/post layout cleanup to avoid unnecessary copies, cats, and
|
||||
splits around sequence parallel all-to-all.
|
||||
|
||||
Priority guidance:
|
||||
|
||||
- P1 when implemented with existing PyTorch/Triton/local helper patterns and
|
||||
validated against exact outputs or tolerances.
|
||||
- P2 when it requires custom CUDA/Triton kernels, changes numerics, or touches
|
||||
attention math directly.
|
||||
|
||||
Validation:
|
||||
|
||||
- First prove the operator family is material for the target shape.
|
||||
- Use non-profiler A/B for latency and stage timing.
|
||||
- Use quality regression checks for generated video stability.
|
||||
- Check compile behavior and graph breaks if using `torch.compile`.
|
||||
|
||||
### Attention Main-Path Optimization
|
||||
|
||||
Purpose: address the dominant self-attention cost when FlashAttention or other
|
||||
attention kernels dominate CUDA time.
|
||||
|
||||
Evidence to look for:
|
||||
|
||||
- Operator trace where FlashAttention/SDPA kernels dominate total CUDA time.
|
||||
- Attention shape from `ops_rankN.xlsx` `by_shape`, model code, or trace
|
||||
metadata.
|
||||
- Whether attention cost scales with latent frames, latent H/W, patch size, or
|
||||
CFG duplication.
|
||||
- Layout/copy/all-to-all work around attention.
|
||||
|
||||
Typical candidates:
|
||||
|
||||
- Verify the attention backend and shape are on the intended fast kernel path.
|
||||
- Compare supported attention backends only with identical workload and quality
|
||||
settings.
|
||||
- Reduce attention input size by safe model/config choices when allowed:
|
||||
latent resolution, frame count, patching, boundary ratio, or windowing.
|
||||
- Remove avoidable layout conversions before/after attention.
|
||||
- Reuse condition-side KV or other static inputs if the model structure allows.
|
||||
- Consider custom kernels, sparse/window/linear attention, or approximation only
|
||||
after quality risk is accepted.
|
||||
|
||||
Priority guidance:
|
||||
|
||||
- P1 for backend/config/layout changes with preserved math.
|
||||
- P2 for approximate attention, sparsity, custom kernels, or any change that can
|
||||
alter quality/temporal consistency.
|
||||
|
||||
Validation:
|
||||
|
||||
- Always include output quality and seed behavior checks.
|
||||
- Compare `diffuse`, server end-to-end, peak memory, and attention kernel time
|
||||
in diagnostic traces if needed.
|
||||
|
||||
### Algorithmic, Precision, and Approximation Optimization
|
||||
|
||||
Purpose: reduce mathematical work or precision cost beyond local code cleanup.
|
||||
|
||||
Evidence to look for:
|
||||
|
||||
- A single operator family dominates even after low-risk runtime, parallel, and
|
||||
fusion work.
|
||||
- Memory bandwidth or compute utilization suggests precision or quantization
|
||||
could matter.
|
||||
- The user explicitly allows quality-preserving or approximate methods.
|
||||
|
||||
Typical candidates:
|
||||
|
||||
- FP8/quantization for transformer or selected projections.
|
||||
- Sparsity or Rainfusion-style acceleration.
|
||||
- Reduced steps, scheduler changes, distillation, or caching across frames.
|
||||
- Approximate attention or linear attention.
|
||||
|
||||
Priority guidance:
|
||||
|
||||
- Usually P2 because quality, numerics, and implementation risk are high.
|
||||
|
||||
Validation:
|
||||
|
||||
- Requires strict A/B, visual quality review, temporal flicker checks, seed
|
||||
stability, and possibly human evaluation.
|
||||
|
||||
### Interpolation, Super-Resolution, and E2E Pipeline Optimization
|
||||
|
||||
Purpose: optimize the whole user-visible video product, not only the base
|
||||
diffusion invocation. Some deployments trade base-model latency against
|
||||
post-processing, interpolation, or super-resolution stages.
|
||||
|
||||
Evidence to look for:
|
||||
|
||||
- E2E latency breakdown across base generation, interpolation, super-resolution,
|
||||
encoding, storage, and response streaming.
|
||||
- Fast/slow GPU or fast/slow stage analysis across multiple cards and pipeline
|
||||
stages.
|
||||
- User quality target: resolution, FPS, temporal smoothness, and acceptable
|
||||
post-processing artifacts.
|
||||
|
||||
Typical candidates:
|
||||
|
||||
- Add or optimize a frame interpolation stage when it reduces required base
|
||||
model frames for the same perceived FPS.
|
||||
- Add or optimize a super-resolution model when generating lower base
|
||||
resolution plus SR is faster for the target quality.
|
||||
- Analyze E2E2 pipeline behavior: client request, service scheduling, diffusion,
|
||||
VAE/media, post-process, file write, and response.
|
||||
- Identify fast/slow cards or stages and rebalance pipeline placement.
|
||||
|
||||
Priority guidance:
|
||||
|
||||
- P1 when using proven interpolation/SR components without changing diffusion
|
||||
math.
|
||||
- P2 when quality risk is high or the pipeline adds significant operational
|
||||
complexity.
|
||||
|
||||
Validation:
|
||||
|
||||
- Measure E2E wall-clock, per-stage server timings, output FPS/resolution,
|
||||
artifacts, flicker, and user-visible quality.
|
||||
|
||||
### Optimization Candidate Library
|
||||
|
||||
Use this table as a compact menu, not as automatic recommendations. Pick items
|
||||
only when baseline metrics, trace evidence, source inspection, or quality
|
||||
tolerance supports them.
|
||||
|
||||
| Layer | Candidate | Evidence | Priority | Validation focus |
|
||||
|---|---|---|---|---|
|
||||
| Measurement | Freeze server/request/benchmark commands | Missing or drifting commands | P0 | Repeatable non-profiler A/B |
|
||||
| Measurement | Separate baseline and diagnostic profiler runs | Profiler used for latency claims | P0 | Low-overhead stage timings |
|
||||
| Host/runtime | Guard or remove avoidable `empty_cache` | Host-stack gaps or sync stalls | P0/P1 | Latency, peak memory, OOM safety |
|
||||
| Host/runtime | Cache scheduler coefficients/timesteps | Repeated tiny CPU/GPU work | P0/P1 | Same seed/output, stage timing |
|
||||
| Host/runtime | Reduce framework scheduling overhead | Client time exceeds server time | P1 | E2E latency, throughput |
|
||||
| Parallel | `CFG=2 x USP=world/2` vs `USP=world` | CFG doubles forward work | P0/P1 | `diffuse`, NCCL, memory |
|
||||
| Parallel | Tune VAE patch parallelism | VAE encode/decode is material | P0/P1 | VAE time, output correctness |
|
||||
| Parallel | HSDP on/off or buffer reuse | HSDP affects memory/all-gather | P1 | Memory, latency, OOM risk |
|
||||
| Parallel | Ulysses vs Ulysses+Ring | Long sequence all-to-all suspected | P1/P2 | Rank balance, NCCL kernels |
|
||||
| Cross-attn | Disable SP for short condition tokens | Cross-attn comm exceeds compute | P1 | `diffuse`, correctness |
|
||||
| VAE/media | Reduce VAE gather/broadcast | Rank traces show VAE wait | P1 | Rank balance, output file |
|
||||
| VAE/media | Reuse tile metadata/buffers | Tile split/merge host gaps | P1 | `vae.encode/decode`, memory |
|
||||
| VAE/media | VAE bf16/autocast | VAE float kernels are slow | P1/P2 | Artifacts, flicker, seed stability |
|
||||
| Operator fusion | AdaLayerNorm/LayerNorm fusion | Norm plus scale/shift kernels | P1 | Numeric tolerance, latency |
|
||||
| Operator fusion | RMSNorm fusion | Many small RMSNorm kernels | P1 | Numeric tolerance, latency |
|
||||
| Operator fusion | RoPE cache/fuse/layout cleanup | RoPE copy/reshape kernels | P1/P2 | Kernel count, correctness |
|
||||
| Operator layout | QKV or attention layout cleanup | Copy/cat/split around attention | P1 | Copy kernels, compile behavior |
|
||||
| Attention | Verify backend fast path | FA/SDPA dominates trace | P1 | `diffuse`, attention kernels |
|
||||
| Attention | FA to LA or selected-head LA | Attention remains dominant | P2 | Quality, temporal stability |
|
||||
| Precision | Transformer FP8/quantization | Compute/bandwidth bound and allowed | P2 | Quality, speed, stability |
|
||||
| Sparsity | Rainfusion-style acceleration | DiT compute remains dominant | P2 | Prompt diversity, quality |
|
||||
| Pipeline | Frame interpolation | Fewer base frames can meet FPS | P1/P2 | E2E latency, motion artifacts |
|
||||
| Pipeline | Super-resolution | Lower base res plus SR may win | P1/P2 | Detail quality, artifacts |
|
||||
| E2E | Fast/slow-card analysis | Multi-card stragglers | P0/P1 | Per-rank/stage wall-clock |
|
||||
|
||||
### Optimization Plan Template
|
||||
|
||||
Use this table shape when reporting the next work items:
|
||||
|
||||
| Priority | Layer | Candidate | Evidence | Expected benefit | Implementation path | Validation | Quality risk |
|
||||
|---|---|---|---|---|---|---|---|
|
||||
| P0 | Host/runtime | Guard `empty_cache` | Host-stack gap points to `torch.cuda.empty_cache` | Small latency reduction, less idle | Add config/env guard | Non-profiler A/B, memory check | Low |
|
||||
| P1 | Operator fusion | RMSNorm/AdaLayerNorm fusion | High-frequency norm/elementwise kernels | Lower launch/bandwidth overhead | Use existing fusion helper or targeted Triton | A/B + output check | Medium |
|
||||
| P1/P2 | Attention | Attention layout/backend investigation | FA kernel dominates CUDA time | Potentially large | Inspect shapes/backend and remove layout copies | A/B + trace + quality | Medium/high |
|
||||
|
||||
Then present a short selection prompt using the same rows:
|
||||
|
||||
```text
|
||||
Which candidate should we implement next?
|
||||
|
||||
1. P0 Host/runtime: guard empty_cache
|
||||
- Expected benefit: small but low-risk latency reduction.
|
||||
- Risk: possible memory increase/OOM if memory headroom is insufficient.
|
||||
|
||||
2. P1 Operator fusion: inspect by_shape and implement first norm/RoPE/layout fusion
|
||||
- Expected benefit: medium if high-frequency small kernels are confirmed.
|
||||
- Risk: numerical/compile/quality validation needed.
|
||||
|
||||
3. P1/P2 Attention: FA/LA/backend/layout investigation
|
||||
- Expected benefit: potentially large.
|
||||
- Risk: high quality and implementation risk.
|
||||
```
|
||||
|
||||
If the user has not chosen an item, default to explaining tradeoffs and asking
|
||||
which candidate to execute. Only proceed autonomously on low-risk P0 measurement
|
||||
or instrumentation fixes.
|
||||
|
||||
## Analysis Helpers
|
||||
|
||||
PyTorch profiler traces are Chrome/Perfetto-compatible JSON files, usually
|
||||
`trace_rankN.json` or `trace_rankN.json.gz`. They normally contain a top-level
|
||||
`traceEvents` list, though some exporters emit the raw event list directly.
|
||||
|
||||
Use the checked-in analyzer from the repository root:
|
||||
|
||||
```bash
|
||||
python3 .claude/skills/diffusion-perf-opt/scripts/trace_analyzer.py \
|
||||
vllm_profile/.../trace_rank0.json.gz \
|
||||
--min-gap-ms 5 \
|
||||
--topn 20
|
||||
```
|
||||
|
||||
For rank imbalance or communication questions, pass all relevant rank traces in
|
||||
one command. For host gaps, lower `--min-gap-ms` to `1` and use a host-stack
|
||||
trace. Read `gpu_span_s`, `busy_union_s`, `idle_union_s`, `idle_pct`, `GAP`
|
||||
blocks, `Top GPU/operator events by total duration`, and `Top NCCL-like events
|
||||
by category`. Treat `cat=user_annotation` NCCL ranges as enclosing annotations;
|
||||
prefer `cat=kernel` or `cat=gpu_user_annotation` for real device work.
|
||||
|
||||
The analyzer summarizes timing only. It does not parse tensor shapes, attribute
|
||||
overlap to individual streams, prove quality, or provide final latency claims.
|
||||
Use `ops_rankN.xlsx` or PyTorch key averages for shape analysis, and re-test any
|
||||
optimization with non-profiler baseline commands.
|
||||
|
||||
Read `references/optimization-playbook.md` when drafting the optimization table or comparing candidate techniques.
|
||||
|
||||
## vLLM Omni Heuristics
|
||||
|
||||
- Cross-attention usually should not use USP/SP when text/image condition token count is much smaller than latent video tokens. Confirm via trace; in Wan2.2 I2V, self-attention dominates cross-attention.
|
||||
- VAE bf16/autocast is often worthwhile but requires visual quality checks.
|
||||
- VAE patch parallel can help decode/encode but may add gather/merge/broadcast overhead. Check whether all ranks need the final decoded tensor.
|
||||
- HSDP/FSDP is primarily a memory strategy. If the model fits without it, run an on/off latency comparison.
|
||||
- Scheduler work can create small host/device gaps; cache tiny solve coefficients when timesteps/order are known.
|
||||
- `torch.cuda.empty_cache()` can prevent OOM but creates synchronization/idle. Make it optional if memory headroom is sufficient.
|
||||
- `Command Buffer Full` in profiler output is profiler overhead, not a model optimization target.
|
||||
@@ -0,0 +1,4 @@
|
||||
interface:
|
||||
display_name: "vLLM Omni Perf Opt"
|
||||
short_description: "Profile and optimize vLLM Omni diffusion performance."
|
||||
default_prompt: "Analyze and optimize a vLLM Omni diffusion workload using baselines, traces, parallel strategy modeling, and prioritized optimization plans."
|
||||
@@ -0,0 +1,82 @@
|
||||
# Optimization Playbook
|
||||
|
||||
Use this reference when turning trace evidence into prioritized vLLM Omni optimizations.
|
||||
|
||||
## Parallel Strategy
|
||||
|
||||
Start from workload shape:
|
||||
|
||||
- Long video tokens, high resolution, many frames: self-attention and FFN dominate. USP/SP is usually valuable.
|
||||
- CFG enabled with positive and negative prompts: transformer forward count doubles. CFG parallel may beat USP on small/moderate sequence lengths.
|
||||
- Cross-attention: skip SP/USP when condition tokens are much smaller than latent tokens; extra all-to-all can exceed compute saved.
|
||||
- HSDP/FSDP: use to fit memory first. If memory allows, compare without HSDP because all-gathers and lazy init add overhead.
|
||||
- VAE parallel: useful when encode/decode is material. Check gather/merge/broadcast and whether every rank needs final decoded output.
|
||||
|
||||
Candidate matrices:
|
||||
|
||||
- 2 cards: `USP=2`, `CFG=2`, HSDP on/off, VAE PP on/off.
|
||||
- 4 cards: `USP=4`, `CFG=2 x USP=2`, `USP=2 x HSDP`, VAE PP on/off.
|
||||
- 8 cards: add pipeline/stage parallel only after measuring orchestration and inter-stage transfer overhead.
|
||||
|
||||
## Host And Free Bubble Analysis
|
||||
|
||||
Look for device intervals where no kernel, memcpy, or memset is active.
|
||||
|
||||
Common causes:
|
||||
|
||||
- First request FSDP/HSDP lazy init.
|
||||
- Python preprocessing: image load, resize, `to_tensor`, tokenizer, request marshaling.
|
||||
- `torch.cuda.empty_cache()` or explicit synchronizations.
|
||||
- Scheduler small linear algebra such as tiny `torch.linalg.solve`.
|
||||
- VAE tile split/merge Python loops.
|
||||
- Profiler overhead such as CUPTI command buffer flush/full.
|
||||
|
||||
Also inspect trace-external time:
|
||||
|
||||
- API request parsing and multipart upload.
|
||||
- Engine/orchestrator dispatch.
|
||||
- Worker RPC and queue wait.
|
||||
- Output serialization, video write, base64, HTTP response.
|
||||
|
||||
## Operator Optimization Candidates
|
||||
|
||||
Evaluate necessity before implementation:
|
||||
|
||||
- **FA to LA / LA for head:** high potential only if attention is the main CUDA hotspot and target shapes fit LA well.
|
||||
- **Cross-attention BSND/no transpose:** do only if transpose/copy around cross-attn is measurable.
|
||||
- **AdaLayerNorm fusion:** fuse norm, scale/shift, gate, residual, and dtype conversion if repeated elementwise kernels are significant.
|
||||
- **RMSNorm fusion:** useful for Q/K norms and VAE norms when small-kernel launch count is high.
|
||||
- **FP32 LayerNorm to bf16:** quality-sensitive; run visual regression.
|
||||
- **LA preprocess:** integrate only if preprocess overhead is cached or amortized.
|
||||
- **RoPE move/cache/fuse:** useful when RoPE small kernels are visible around every self-attn block.
|
||||
- **VAE RMSNorm replacement:** only after VAE internal trace shows norm as a real component of decode/encode.
|
||||
- **Scheduler fusion/cache:** cache coefficient solves and fuse CFG combine with scheduler arithmetic when possible.
|
||||
|
||||
Before writing a custom kernel, check:
|
||||
|
||||
- Local vLLM Omni implementations.
|
||||
- Upstream vLLM kernels and attention backends.
|
||||
- SGLang, FlashInfer, FlashAttention, and Triton examples.
|
||||
- Shape constraints and layout expectations.
|
||||
|
||||
## Quantization And Sparsity
|
||||
|
||||
Treat these as high-risk optimizations:
|
||||
|
||||
- Rainfusion/sparse diffusion can be high reward but needs model-specific quality validation.
|
||||
- VAE decode quantization can reduce decode time but usually has smaller end-to-end upside than transformer attention/FFN work.
|
||||
- FP8 or approximate kernels require seed-stable visual regression, temporal flicker checks, and prompt diversity.
|
||||
|
||||
## Output Format For Recommendations
|
||||
|
||||
Return a table with:
|
||||
|
||||
- Priority: P0/P1/P2.
|
||||
- Optimization.
|
||||
- Evidence from trace or code.
|
||||
- Expected benefit.
|
||||
- Risk/quality impact.
|
||||
- Implementation path.
|
||||
- Validation command or A/B test.
|
||||
|
||||
State clearly when a recommendation is a hypothesis requiring new trace with shapes.
|
||||
+144
@@ -0,0 +1,144 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Summarize torch profiler Chrome traces for vLLM Omni diffusion optimization."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import collections
|
||||
import gzip
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
GPU_CATS = {"kernel", "gpu_memcpy", "gpu_memset"}
|
||||
CPU_CATS = {"python_function", "user_annotation", "cpu_op", "cuda_runtime", "cuda_driver"}
|
||||
|
||||
|
||||
def open_trace(path: Path):
|
||||
if path.suffix == ".gz":
|
||||
return gzip.open(path, "rt")
|
||||
return path.open("rt")
|
||||
|
||||
|
||||
def event_name(event: dict[str, Any]) -> str:
|
||||
return str(event.get("name", ""))
|
||||
|
||||
|
||||
def summarize_trace(path: Path, min_gap_us: float, topn: int) -> None:
|
||||
with open_trace(path) as f:
|
||||
data = json.load(f)
|
||||
|
||||
events = data if isinstance(data, list) else data.get("traceEvents", [])
|
||||
gpu: list[tuple[float, float, float, str, str, Any, Any]] = []
|
||||
cpu: list[tuple[float, float, float, str, str, Any, Any]] = []
|
||||
by_gpu_name: dict[str, list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
|
||||
by_name: dict[str, list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
|
||||
nccl: dict[tuple[str, str], list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
|
||||
|
||||
for event in events:
|
||||
dur = event.get("dur")
|
||||
ts = event.get("ts")
|
||||
if dur is None or ts is None or dur <= 0:
|
||||
continue
|
||||
cat = str(event.get("cat", ""))
|
||||
name = event_name(event)
|
||||
row = (float(ts), float(ts + dur), float(dur), name, cat, event.get("pid"), event.get("tid"))
|
||||
stat = by_name[name]
|
||||
stat[0] += 1
|
||||
stat[1] += dur
|
||||
stat[2] = max(stat[2], dur)
|
||||
if cat in GPU_CATS:
|
||||
gpu.append(row)
|
||||
gstat = by_gpu_name[name]
|
||||
gstat[0] += 1
|
||||
gstat[1] += dur
|
||||
gstat[2] = max(gstat[2], dur)
|
||||
elif cat in CPU_CATS:
|
||||
cpu.append(row)
|
||||
if "nccl" in name.lower():
|
||||
nstat = nccl[(cat, name)]
|
||||
nstat[0] += 1
|
||||
nstat[1] += dur
|
||||
nstat[2] = max(nstat[2], dur)
|
||||
|
||||
print(f"\n== {path}")
|
||||
print(f"events={len(events)} gpu_events={len(gpu)} cpu_events={len(cpu)}")
|
||||
if not gpu:
|
||||
print("No GPU events found.")
|
||||
return
|
||||
|
||||
gpu.sort()
|
||||
merged: list[list[Any]] = []
|
||||
for start, end, dur, name, cat, pid, tid in gpu:
|
||||
if not merged or start > merged[-1][1]:
|
||||
merged.append([start, end, [(start, end, dur, name, cat, pid, tid)]])
|
||||
else:
|
||||
merged[-1][1] = max(merged[-1][1], end)
|
||||
merged[-1][2].append((start, end, dur, name, cat, pid, tid))
|
||||
|
||||
span = max(end for _, end, *_ in gpu) - min(start for start, *_ in gpu)
|
||||
busy = sum(end - start for start, end, _ in merged)
|
||||
idle = span - busy
|
||||
print(
|
||||
f"gpu_span_s={span / 1e6:.3f} "
|
||||
f"busy_union_s={busy / 1e6:.3f} "
|
||||
f"idle_union_s={idle / 1e6:.3f} "
|
||||
f"idle_pct={idle / span * 100:.2f}"
|
||||
)
|
||||
|
||||
interesting_cpu = [
|
||||
row
|
||||
for row in cpu
|
||||
if row[2] >= 1000
|
||||
and (
|
||||
row[4] in {"python_function", "user_annotation"}
|
||||
or "cudaStreamSynchronize" in row[3]
|
||||
or "cudaDeviceSynchronize" in row[3]
|
||||
or "cudaLaunch" in row[3]
|
||||
or "cudaMemcpy" in row[3]
|
||||
)
|
||||
]
|
||||
|
||||
gaps = []
|
||||
for idx in range(1, len(merged)):
|
||||
gap_start = merged[idx - 1][1]
|
||||
gap_end = merged[idx][0]
|
||||
gap_dur = gap_end - gap_start
|
||||
if gap_dur >= min_gap_us:
|
||||
prev_event = max(merged[idx - 1][2], key=lambda x: x[1])
|
||||
next_event = min(merged[idx][2], key=lambda x: x[0])
|
||||
mid = (gap_start + gap_end) / 2
|
||||
containers = [row for row in interesting_cpu if row[0] <= mid <= row[1]]
|
||||
containers = sorted(containers, key=lambda x: x[2])[:8]
|
||||
gaps.append((gap_dur, gap_start, gap_end, prev_event, next_event, containers))
|
||||
|
||||
print(f"gaps_ge_{min_gap_us / 1000:.3f}ms count={len(gaps)} sum_s={sum(g[0] for g in gaps) / 1e6:.3f}")
|
||||
for gap_dur, gap_start, gap_end, prev_event, next_event, containers in sorted(gaps, reverse=True)[:topn]:
|
||||
print(f"\nGAP {gap_dur / 1000:.3f} ms ts={gap_start:.0f}->{gap_end:.0f}")
|
||||
print(f" prev {prev_event[4]} {prev_event[2] / 1000:.3f} ms {prev_event[3][:160]}")
|
||||
print(f" next {next_event[4]} {next_event[2] / 1000:.3f} ms {next_event[3][:160]}")
|
||||
for row in containers:
|
||||
print(f" in {row[4]} {row[2] / 1000:.3f} ms {row[3][:180]}")
|
||||
|
||||
print("\nTop GPU/operator events by total duration:")
|
||||
for name, (count, total, max_dur) in sorted(by_gpu_name.items(), key=lambda kv: kv[1][1], reverse=True)[:topn]:
|
||||
print(f" {int(count):8d} total={total / 1e6:9.3f}s max={max_dur / 1000:9.3f}ms {name[:180]}")
|
||||
|
||||
print("\nTop NCCL-like events by category:")
|
||||
for (cat, name), (count, total, max_dur) in sorted(nccl.items(), key=lambda kv: kv[1][1], reverse=True)[:topn]:
|
||||
print(f" {int(count):8d} total={total / 1e6:9.3f}s max={max_dur / 1000:9.3f}ms cat={cat} {name[:160]}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("traces", nargs="+", type=Path, help="trace.json or trace.json.gz files")
|
||||
parser.add_argument("--min-gap-ms", type=float, default=5.0, help="minimum GPU idle gap to print")
|
||||
parser.add_argument("--topn", type=int, default=20, help="number of gaps/hotspots to print")
|
||||
args = parser.parse_args()
|
||||
|
||||
for trace in args.traces:
|
||||
summarize_trace(trace, args.min_gap_ms * 1000.0, args.topn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -18,6 +18,9 @@ include:
|
||||
|
||||
- `add-diffusion-model`: guides integration of a new diffusion model into
|
||||
`vllm-omni`
|
||||
- `diffusion-perf-opt`: guides diffusion model performance optimization,
|
||||
including profiling traces, parallel strategies, stage timing analysis, and
|
||||
benchmark-driven tuning
|
||||
- `add-omni-model`: covers addition of new omni-modality model support
|
||||
- `add-tts-model`: covers integration of new TTS models and related serving
|
||||
workflows
|
||||
|
||||
@@ -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/
|
||||
@@ -10,6 +10,11 @@ build:
|
||||
jobs:
|
||||
post_checkout:
|
||||
- git fetch --unshallow || true
|
||||
pre_install:
|
||||
# Pre-install CPU torch so the subsequent `pip install .[docs]` does not
|
||||
# pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
|
||||
# 1 GB memory cap and time out the mkdocs build.
|
||||
- python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
|
||||
|
||||
mkdocs:
|
||||
configuration: mkdocs.yml
|
||||
|
||||
@@ -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://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/111-L8zF7A1j_YI_cR8JsblofdScdRr2f/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
|
||||
</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.
|
||||
|
||||
---
|
||||
|
||||
+51
-35
@@ -1,44 +1,60 @@
|
||||
# Benchmarks Overview and Architecture
|
||||
# Benchmarks
|
||||
|
||||
This document explains the benchmark architecture across all benchmark assets in this repo. It describes what we measure, and where to find or plug in new scenarios. Per-task details remain in subfolder READMEs (e.g., `benchmarks/<model>/README.md`).
|
||||
This directory contains benchmark suites for evaluating different model families and infrastructure components in vLLM-Omni. Each subfolder targets a different benchmark family with its own scripts, configs, and metrics. See the per-subfolder READMEs for detailed usage.
|
||||
|
||||
## Scope and goals
|
||||
- Establish repeatable latency/throughput measurements for multimodal LLM pipelines.
|
||||
- Provide both HF Transformers (offline) and vLLM-Omni (multi-stage/pipeline) baselines.
|
||||
- Make it easy to plug in new datasets and models with minimal changes to the runner scripts.
|
||||
## Benchmark families
|
||||
|
||||
## Dataset and inputs
|
||||
- Default example: SeedTTS top-100 prompts (`benchmarks/build_dataset/top100.txt`) via `benchmarks/build_dataset/`.
|
||||
- Extensible: drop in new prompt files or modality-aligned payloads; keep the expected format for the consuming scripts (e.g., one prompt per line).
|
||||
- If you add a new dataset, document it under `benchmarks/<model>/README.md` and point scripts to your data path.
|
||||
### [TTS](tts/README.md) — Text-to-Speech
|
||||
|
||||
## Directory layout
|
||||
- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
|
||||
- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
|
||||
- `benchmarks/accuracy/` — accuracy benchmark integrations that adapt external
|
||||
benchmark suites to vLLM-Omni serving and evaluation flows.
|
||||
- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
|
||||
Model-agnostic serving benchmarks for TTS models, including Qwen3-TTS and VoxCPM2.
|
||||
|
||||
## Reference workflows
|
||||
- **HF Transformers (offline, single process)**
|
||||
Script (example): `benchmarks/<model>/transformers/eval_qwen3_moe_omni_transformers.sh`
|
||||
Outputs: `benchmark_results/perf_stats.json`, `benchmark_results/results.json`, `benchmark_results/audio/` (if audio is produced).
|
||||
- **Layout**: `tts/bench_tts.py` (serving benchmark driver), `tts/model_configs.yaml` (model registry), `tts/plot_results.py` (result plotting)
|
||||
- **Dataset**: Seed-TTS full or text-only datasets, plus bundled smoke/design prompts under `build_dataset/`
|
||||
- **Key metrics**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor), throughput (audio seconds / wall-clock second)
|
||||
|
||||
- **vLLM-Omni end-to-end pipeline**
|
||||
Script (example): `benchmarks/<model>/vllm_omni/eval_qwen3_moe_omni.sh`
|
||||
Outputs: `vllm_omni/logs/*.stats.jsonl` (per-stage/overall latency & TPS), `vllm_omni/logs/stage*.log`, `vllm_omni/outputs/` (text/audio artifacts).
|
||||
### [Diffusion](diffusion/README.md) — Image and Video Generation
|
||||
|
||||
- **Adding a new task/model**
|
||||
1) Create `benchmarks/<model>/transformers/` and/or `benchmarks/<model>/vllm_omni/` with scripts referencing your model and dataset.
|
||||
2) Add a task README describing dataset, configs, and expected outputs.
|
||||
3) Keep the output/log structure similar for easy comparison (perf_stats/results/audio or text outputs; stats.jsonl/logs for pipeline).
|
||||
Online-serving benchmark for diffusion image/video models, sending requests to the configured vLLM serving endpoint (`/v1/chat/completions`, `/v1/images/generations`, or `/v1/videos`, depending on backend/task).
|
||||
|
||||
## Metrics to watch
|
||||
- **Throughput**: `overall_tps`, `*_tps_avg` per stage.
|
||||
- **Latency distribution**: look for long tails in `*.stats.jsonl`.
|
||||
- **Quality/completeness**: missing outputs or errors in stage logs indicate pipeline failures or misconfigurations.
|
||||
- **Tasks**: text-to-image, text-to-video, image-to-image, image-to-video, text+image-to-image, text+image-to-video
|
||||
- **Datasets**: `vbench`, `trace`, `random`
|
||||
- **Key metrics**: request throughput, latency percentiles, optional SLO attainment
|
||||
|
||||
## Troubleshooting
|
||||
- Verify GPU/driver/FlashAttention2 requirements for your chosen model/config.
|
||||
- Ensure network access for dataset/model downloads (Google Drive, Hugging Face, etc.).
|
||||
- If outputs are missing or slow, inspect per-stage logs and `*.stats.jsonl` for errors, stragglers, or contention.
|
||||
### [GLM-Image](glm_image/README.md) — Text-to-Image and Image-to-Image
|
||||
|
||||
Benchmarks for GLM-Image performance across HuggingFace baseline, vLLM-Omni offline inference, and vLLM-Omni online serving.
|
||||
|
||||
- **Layout**: `glm_image/huggingface/` (HF baseline), `glm_image/vllm-omni/` (offline inference), `glm_image/benchmark_glm_image.py` (online serving)
|
||||
- **Tasks**: text-to-image and image-to-image
|
||||
- **Key metrics**: request/image throughput, latency percentiles, optional per-stage pipeline timings
|
||||
|
||||
### [Distributed](distributed/omni_connectors/README.md) — RDMA Connector Testing
|
||||
|
||||
RDMA environment setup and transfer tests for `MooncakeTransferEngineConnector`, including pytest-based single-node checks and manual cross-node benchmarks.
|
||||
|
||||
- **Transfer modes**: `copy`, `zerocopy`, `gpu` (GPUDirect)
|
||||
- **Supports**: single-node pytest suites and manual multi-node/cross-node transfer testing
|
||||
|
||||
### [Accuracy](accuracy/README.md) — Image Generation and Editing Quality
|
||||
|
||||
Accuracy benchmarks for image generation/editing models, adapting external suites to vLLM-Omni serving and local judge-evaluation flows.
|
||||
|
||||
- **Layout**: `accuracy/text_to_image/` (GEBench), `accuracy/image_to_image/` (GEdit-Bench)
|
||||
- **Method**: generation and judge scoring both run through local `vllm-omni serve` endpoints
|
||||
|
||||
### Common serving metrics framework
|
||||
|
||||
`vllm_omni/benchmarks/` extends `vllm bench serve --omni` with Omni-specific datasets, backends, and multimodal metrics. Key metrics include:
|
||||
|
||||
- **Text output**: TTFT (time to first token), TPOT (time per output token), ITL (inter-token latency)
|
||||
- **Audio output**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor)
|
||||
- **Throughput**: request throughput, output token throughput, total token throughput, audio throughput
|
||||
|
||||
See `vllm_omni/benchmarks/serve.py` for the `vllm bench serve --omni` runner wrapper and `vllm_omni/benchmarks/metrics/` for Omni metric definitions.
|
||||
|
||||
## Adding a new benchmark
|
||||
|
||||
1. Create a subfolder under `benchmarks/<name>/` with scripts, configs if needed, and a `README.md`.
|
||||
2. If comparing against another runtime, use clear backend subfolders where applicable, such as `huggingface/` and `vllm-omni/`, or follow the shared TTS serving benchmark pattern in `tts/`.
|
||||
3. Add reusable dataset or prompt-building utilities to `build_dataset/` if applicable.
|
||||
4. Update this README with a link to the new benchmark family.
|
||||
@@ -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,30 @@
|
||||
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.
|
||||
vd021|||Please fasten your seatbelts; we will be experiencing some turbulence shortly.|A calm, professional female flight attendant voice, reassuring and clear.
|
||||
vd022|||The fossil record suggests these creatures lived more than sixty million years ago.|A measured male documentary narrator voice, scientific and curious.
|
||||
vd023|||Hey kiddo, did you finish your homework before dinner today?|A gentle, caring father voice, warm and slightly tired after a long day.
|
||||
vd024|||Live from courtside, the home team has just taken a commanding lead.|A high-energy male sports broadcaster, urgent and engaged.
|
||||
vd025|||Please remain on the line while I transfer you to the next available agent.|A polite, neutral female call-center voice, professional and patient.
|
||||
vd026|||I have a dream that one day this nation will rise up and live out the true meaning of its creed.|A powerful, resonant male orator voice, slow and impassioned.
|
||||
vd027|||Mama always said life was like a box of chocolates; you never know what you're gonna get.|A folksy male voice with a Southern American accent, warm and reflective.
|
||||
vd028|||Ahoy there, matey! Hoist the colors and prepare to set sail.|A theatrical pirate male voice, gruff and boisterous.
|
||||
vd029|||Ladies and gentlemen, please welcome to the stage our keynote speaker.|A polished female event MC voice, gracious and articulate.
|
||||
vd030|||Once you've completed the form, click submit and we'll review it within two business days.|A friendly female product-onboarding voice, encouraging and clear.
|
||||
@@ -0,0 +1,30 @@
|
||||
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.
|
||||
smoke021|||A long line of cars stretched along the coastal highway after the lighthouse opened.
|
||||
smoke022|||The chef recommended the seared scallops with a lemon butter sauce as the special.
|
||||
smoke023|||Our quarterly engineering review begins promptly at two in the afternoon on Thursday.
|
||||
smoke024|||Despite the heavy rain, the marathon runners kept a steady pace through downtown.
|
||||
smoke025|||She finished the crossword puzzle in under fifteen minutes while sipping her tea.
|
||||
smoke026|||The new release adds streaming output, faster startup, and better error reporting.
|
||||
smoke027|||His grandfather taught him how to whittle small wooden birds from cedar offcuts.
|
||||
smoke028|||Travelers were advised to check the timetable since trains were running on a holiday schedule.
|
||||
smoke029|||The startup launched a public beta after eighteen months of closed development.
|
||||
smoke030|||At night the desert sky filled with stars so bright the Milky Way was clearly visible.
|
||||
@@ -30,7 +30,7 @@ python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
|
||||
**Notes**
|
||||
|
||||
- The benchmark talks to `http://<host>:<port>/v1/chat/completions`.
|
||||
- By default, image tasks talk to `http://<host>:<port>/v1/chat/completions`; video tasks talk to `/v1/videos`.
|
||||
- If you run the server on another host or port, pass `--base-url` accordingly.
|
||||
|
||||
## 2. Supported Datasets
|
||||
@@ -95,8 +95,9 @@ You can point to your own trace using `--dataset-path`.
|
||||
|
||||
### Basic flags
|
||||
|
||||
- `--base-url`: Server address (the script calls `.../v1/chat/completions`).
|
||||
- `--base-url`: Server address; `--endpoint` selects the path appended to this base URL.
|
||||
- `--model`: The OpenAI-compatible `model` field.
|
||||
- `--endpoint`: API endpoint path. Leading `/` is optional, e.g. `/v1/videos` or `v1/videos`.
|
||||
- `--task`: Task type (e.g., `t2i`, `t2v`, `i2i`, `i2v`).
|
||||
- `--dataset`: Dataset mode (`vbench` / `trace` / `random`).
|
||||
- `--num-prompts`: Number of requests to send.
|
||||
@@ -130,9 +131,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).
|
||||
@@ -10,6 +10,8 @@ from typing import Any
|
||||
import aiohttp
|
||||
from tqdm import tqdm
|
||||
|
||||
DEFAULT_EDITS_BOT_TASK = "think"
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestFuncInput:
|
||||
@@ -27,6 +29,7 @@ class RequestFuncInput:
|
||||
extra_body: dict[str, Any] = field(default_factory=dict)
|
||||
image_paths: list[str] | None = None
|
||||
request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
default_bot_task: str | None = DEFAULT_EDITS_BOT_TASK
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -53,6 +56,95 @@ def _encode_image_as_data_url(path: str) -> str:
|
||||
return f"data:{mime};base64,{encoded}"
|
||||
|
||||
|
||||
async def async_request_image_edits(
|
||||
input: RequestFuncInput,
|
||||
session: aiohttp.ClientSession,
|
||||
pbar: tqdm | None = None,
|
||||
enable_diffusion_pipeline_profiler: bool = False,
|
||||
) -> RequestFuncOutput:
|
||||
"""POST /v1/images/edits (multipart)."""
|
||||
del enable_diffusion_pipeline_profiler
|
||||
output = RequestFuncOutput()
|
||||
output.start_time = time.perf_counter()
|
||||
|
||||
extra_body = dict(input.extra_body)
|
||||
width = input.width or extra_body.get("width") or 1024
|
||||
height = input.height or extra_body.get("height") or 1024
|
||||
edits_url = input.api_url
|
||||
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("model", input.model)
|
||||
form.add_field("prompt", input.prompt)
|
||||
form.add_field("size", f"{width}x{height}")
|
||||
form.add_field("response_format", "b64_json")
|
||||
|
||||
if input.num_inference_steps is not None:
|
||||
form.add_field("num_inference_steps", str(input.num_inference_steps))
|
||||
elif extra_body.get("num_inference_steps") is not None:
|
||||
form.add_field("num_inference_steps", str(extra_body["num_inference_steps"]))
|
||||
|
||||
if input.seed is not None:
|
||||
form.add_field("seed", str(input.seed))
|
||||
elif extra_body.get("seed") is not None:
|
||||
form.add_field("seed", str(extra_body["seed"]))
|
||||
|
||||
if extra_body.get("guidance_scale") is not None:
|
||||
form.add_field("guidance_scale", str(extra_body["guidance_scale"]))
|
||||
if extra_body.get("negative_prompt") is not None:
|
||||
form.add_field("negative_prompt", str(extra_body["negative_prompt"]))
|
||||
if extra_body.get("true_cfg_scale") is not None:
|
||||
form.add_field("true_cfg_scale", str(extra_body["true_cfg_scale"]))
|
||||
if extra_body.get("sys_type") is not None:
|
||||
form.add_field("sys_type", str(extra_body["sys_type"]))
|
||||
if extra_body.get("system_prompt") is not None:
|
||||
form.add_field("system_prompt", str(extra_body["system_prompt"]))
|
||||
|
||||
bot_task = extra_body.get("bot_task")
|
||||
if bot_task is None and input.default_bot_task is not None:
|
||||
bot_task = input.default_bot_task
|
||||
if bot_task is not None:
|
||||
form.add_field("bot_task", str(bot_task))
|
||||
|
||||
assert input.image_paths is not None
|
||||
for img_path in input.image_paths:
|
||||
if not os.path.exists(img_path):
|
||||
output.error = f"Image file not found: {img_path}"
|
||||
output.success = False
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return output
|
||||
with open(img_path, "rb") as img_f:
|
||||
image_bytes = img_f.read()
|
||||
form.add_field(
|
||||
"image",
|
||||
image_bytes,
|
||||
filename=os.path.basename(img_path),
|
||||
content_type=_guess_mime_type(img_path),
|
||||
)
|
||||
|
||||
try:
|
||||
async with session.post(edits_url, data=form) as response:
|
||||
if response.status == 200:
|
||||
resp_json = await response.json()
|
||||
output.response_body = resp_json
|
||||
output.success = True
|
||||
else:
|
||||
output.error = f"HTTP {response.status}: {await response.text()}"
|
||||
output.success = False
|
||||
except Exception as e:
|
||||
output.error = str(e)
|
||||
output.success = False
|
||||
|
||||
output.latency = time.perf_counter() - output.start_time
|
||||
|
||||
if output.success and input.slo_ms is not None:
|
||||
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return output
|
||||
|
||||
|
||||
async def async_request_chat_completions(
|
||||
input: RequestFuncInput,
|
||||
session: aiohttp.ClientSession,
|
||||
@@ -122,6 +214,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
|
||||
@@ -139,7 +243,7 @@ async def async_request_chat_completions(
|
||||
return output
|
||||
|
||||
|
||||
async def async_request_openai_images(
|
||||
async def async_request_openai_image_generations(
|
||||
input: RequestFuncInput,
|
||||
session: aiohttp.ClientSession,
|
||||
pbar: tqdm | None = None,
|
||||
@@ -336,12 +440,39 @@ async def async_request_v1_videos(
|
||||
return output
|
||||
|
||||
|
||||
LEGACY_BACKEND_ENDPOINT_ALIASES = {
|
||||
"vllm-omni": "/v1/chat/completions",
|
||||
"openai": "/v1/images/generations",
|
||||
}
|
||||
|
||||
|
||||
def normalize_endpoint(value: str) -> str:
|
||||
endpoint = str(value).strip()
|
||||
if not endpoint:
|
||||
raise ValueError("endpoint must not be empty.")
|
||||
endpoint = LEGACY_BACKEND_ENDPOINT_ALIASES.get(
|
||||
endpoint,
|
||||
LEGACY_BACKEND_ENDPOINT_ALIASES.get(endpoint.lstrip("/"), endpoint),
|
||||
)
|
||||
if not endpoint.startswith("/"):
|
||||
endpoint = f"/{endpoint}"
|
||||
return endpoint
|
||||
|
||||
|
||||
def endpoint_filename_token(value: str) -> str:
|
||||
token = normalize_endpoint(value).lstrip("/")
|
||||
for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"):
|
||||
token = token.replace(bad, "_")
|
||||
return token or "endpoint"
|
||||
|
||||
|
||||
backends_function_mapping = {
|
||||
"2i": {
|
||||
"vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
|
||||
"openai": (async_request_openai_images, "/v1/images/generations"),
|
||||
"/v1/chat/completions": (async_request_chat_completions, "/v1/chat/completions"),
|
||||
"/v1/images/generations": (async_request_openai_image_generations, "/v1/images/generations"),
|
||||
"/v1/images/edits": (async_request_image_edits, "/v1/images/edits"),
|
||||
},
|
||||
"2v": {
|
||||
"v1/videos": (async_request_v1_videos, "/v1/videos"),
|
||||
"/v1/videos": (async_request_v1_videos, "/v1/videos"),
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,328 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Standalone diagnostic for the diffusion attention backends added in #3079.
|
||||
|
||||
Exercises the same synthetic attention shape through each backend and each
|
||||
torch SDPA sub-kernel, with and without an attn_mask, so we can pinpoint why
|
||||
CUDNN_ATTN underperforms on some SKUs. Two typical causes:
|
||||
|
||||
(a) cuDNN lacks a tuned kernel for this (SM, head_dim, seq) combination
|
||||
and silently falls back to MATH.
|
||||
(b) attn_mask is non-None and the cuDNN SDPA dispatch rejects it,
|
||||
walking CUDNN -> FLASH -> MATH.
|
||||
|
||||
Run:
|
||||
python benchmarks/diffusion/bench_attention_backends.py --preset hv15
|
||||
python benchmarks/diffusion/bench_attention_backends.py --preset wan22
|
||||
python benchmarks/diffusion/bench_attention_backends.py \
|
||||
--batch 1 --heads 24 --seq 14336 --head-dim 128
|
||||
|
||||
Optional extras:
|
||||
pip install --pre flash-attn-4 # FA4 is currently pre-release only
|
||||
pip install -U flashinfer # latest FlashInfer (0.6.9)
|
||||
|
||||
The table at the end is the data we want on the PR — surface the row where a
|
||||
backend is >1.5x the SDPA baseline, that's the one to gate off the auto-route.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
|
||||
import torch
|
||||
from torch.nn.attention import SDPBackend, sdpa_kernel
|
||||
|
||||
# Presets mirror the DiT hot-path attention call (single-stream block) at the
|
||||
# shapes we run in the PR's validation configs.
|
||||
_PRESETS = {
|
||||
# HunyuanVideo-1.5 480p/33f: (30 * 52 * 9) latent + ~256 text tokens.
|
||||
"hv15": {"batch": 1, "heads": 24, "seq": 14336, "head_dim": 128},
|
||||
# Wan 2.2 480p/33f rough estimate; adjust when we measure real shapes.
|
||||
"wan22": {"batch": 1, "heads": 40, "seq": 16384, "head_dim": 128},
|
||||
# Smaller image-gen shape for quick smoke tests.
|
||||
"flux": {"batch": 1, "heads": 24, "seq": 4096, "head_dim": 128},
|
||||
}
|
||||
|
||||
_SDPA_BACKENDS = [
|
||||
("CUDNN_ATTENTION", [SDPBackend.CUDNN_ATTENTION]),
|
||||
("FLASH_ATTENTION", [SDPBackend.FLASH_ATTENTION]),
|
||||
("EFFICIENT_ATTENTION", [SDPBackend.EFFICIENT_ATTENTION]),
|
||||
("MATH", [SDPBackend.MATH]),
|
||||
# The PR's CUDNN_ATTN impl uses this priority chain; keep it here to show
|
||||
# what torch actually picks when multiple are allowed.
|
||||
("CUDNN_ATTN_CHAIN", [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]),
|
||||
]
|
||||
|
||||
|
||||
def _env_report() -> None:
|
||||
if not torch.cuda.is_available():
|
||||
raise SystemExit("CUDA not available; this diagnostic needs a GPU.")
|
||||
cc = torch.cuda.get_device_capability()
|
||||
print("=" * 72)
|
||||
print(f"GPU : {torch.cuda.get_device_name()}")
|
||||
print(f"SM capability : sm_{cc[0]}{cc[1]}")
|
||||
print(f"torch : {torch.__version__}")
|
||||
print(f"cuDNN version : {torch.backends.cudnn.version()}")
|
||||
print("flashinfer : ", end="")
|
||||
try:
|
||||
import flashinfer
|
||||
|
||||
print(getattr(flashinfer, "__version__", "present"))
|
||||
except Exception as e:
|
||||
print(f"not installed ({type(e).__name__})")
|
||||
print("=" * 72)
|
||||
|
||||
|
||||
def _make_qkv(batch: int, heads: int, seq: int, head_dim: int, device: str, dtype: torch.dtype):
|
||||
# Layout: (B, S, H, D) — what the PR's backends accept as input.
|
||||
q = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
|
||||
k = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
|
||||
v = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
|
||||
return q, k, v
|
||||
|
||||
|
||||
def _make_mask(batch: int, seq: int, device: str, dtype: torch.dtype, pad_tokens: int = 128) -> torch.Tensor:
|
||||
# Mirrors a text-encoder padding mask: last `pad_tokens` positions masked.
|
||||
mask = torch.zeros(batch, 1, seq, seq, device=device, dtype=dtype)
|
||||
if pad_tokens > 0:
|
||||
mask[..., -pad_tokens:] = float("-inf")
|
||||
return mask
|
||||
|
||||
|
||||
def _time_call(fn, *args, warmup: int = 3, iters: int = 10) -> tuple[float, str]:
|
||||
"""Return (median ms, error string). Skips on any exception so that a
|
||||
backend that rejects our args (wrong dtype, missing JIT module, unsupported
|
||||
kwarg value) doesn't abort the whole sweep."""
|
||||
try:
|
||||
for _ in range(warmup):
|
||||
fn(*args)
|
||||
torch.accelerator.synchronize()
|
||||
times = []
|
||||
for _ in range(iters):
|
||||
torch.accelerator.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
fn(*args)
|
||||
torch.accelerator.synchronize()
|
||||
times.append((time.perf_counter() - t0) * 1000.0)
|
||||
times.sort()
|
||||
return times[len(times) // 2], ""
|
||||
except Exception as e: # noqa: BLE001 — probe script, keep the whole table going
|
||||
msg = str(e).split("\n", 1)[0][:60]
|
||||
return float("nan"), f"{type(e).__name__}: {msg}" if msg else type(e).__name__
|
||||
|
||||
|
||||
def _run_sdpa_variants(q, k, v, attn_mask, scale: float) -> list[tuple[str, float, str]]:
|
||||
# q/k/v come in (B, S, H, D); F.sdpa wants (B, H, S, D).
|
||||
q_t = q.permute(0, 2, 1, 3)
|
||||
k_t = k.permute(0, 2, 1, 3)
|
||||
v_t = v.permute(0, 2, 1, 3)
|
||||
rows: list[tuple[str, float, str]] = []
|
||||
for name, backends in _SDPA_BACKENDS:
|
||||
|
||||
def _call(mask=attn_mask):
|
||||
with sdpa_kernel(backends):
|
||||
return torch.nn.functional.scaled_dot_product_attention(
|
||||
q_t, k_t, v_t, attn_mask=mask, dropout_p=0.0, is_causal=False, scale=scale
|
||||
)
|
||||
|
||||
ms, err = _time_call(_call)
|
||||
rows.append((name, ms, err))
|
||||
return rows
|
||||
|
||||
|
||||
def _run_flashinfer(q, k, v, scale: float, backend: str | None = None, attn_mask=None) -> tuple[float, str]:
|
||||
"""Call FlashInfer's dense single-prefill.
|
||||
|
||||
``backend`` hints at cutlass/fa3/trtllm if the installed version exposes
|
||||
that kwarg. ``attn_mask`` is the same additive 4D mask we pass to SDPA —
|
||||
we convert it to the 2D boolean ``custom_mask`` form FlashInfer accepts.
|
||||
Per FlashInfer docs, ``custom_mask`` only applies when ``causal=False``.
|
||||
"""
|
||||
try:
|
||||
import inspect
|
||||
|
||||
from flashinfer.prefill import single_prefill_with_kv_cache
|
||||
except Exception as e:
|
||||
return float("nan"), f"import-{type(e).__name__}"
|
||||
|
||||
kwargs: dict = {"sm_scale": scale, "causal": False, "return_lse": False}
|
||||
if backend is not None:
|
||||
sig = inspect.signature(single_prefill_with_kv_cache)
|
||||
if "backend" not in sig.parameters:
|
||||
return float("nan"), "no-backend-kwarg"
|
||||
kwargs["backend"] = backend
|
||||
|
||||
if attn_mask is not None:
|
||||
# Collapse (B, 1, S, S) additive float mask to (S, S) boolean.
|
||||
# FlashInfer expects True = keep, False = masked out.
|
||||
mask_2d = attn_mask[0, 0]
|
||||
kwargs["custom_mask"] = mask_2d != float("-inf")
|
||||
|
||||
def _call():
|
||||
out = single_prefill_with_kv_cache(q[0], k[0], v[0], **kwargs)
|
||||
return out.unsqueeze(0)
|
||||
|
||||
return _time_call(_call)
|
||||
|
||||
|
||||
def _run_fa4(q, k, v, scale: float) -> tuple[float, str]:
|
||||
"""Call FlashAttention-4 directly (``pip install flash-attn-4``). FA4
|
||||
ships a Blackwell-native kernel via CuTe-DSL; on sm_120 it should beat
|
||||
cuDNN by ~20%. API lives under ``flash_attn.cute``; the older
|
||||
``flash_attn.flash_attn_func`` path is FA2/FA3 only."""
|
||||
try:
|
||||
from flash_attn.cute import flash_attn_func
|
||||
except Exception as e:
|
||||
return float("nan"), f"import-{type(e).__name__}"
|
||||
|
||||
def _call():
|
||||
# FA4 accepts (B, S, H, D) directly — same layout the PR's backends use.
|
||||
return flash_attn_func(q, k, v, softmax_scale=scale, causal=False)
|
||||
|
||||
return _time_call(_call)
|
||||
|
||||
|
||||
def _run_flashinfer_cudnn_batch(q, k, v, scale: float) -> tuple[float, str]:
|
||||
"""Call FlashInfer's direct cuDNN wrapper. Bypasses PyTorch SDPA dispatch
|
||||
(which has a few hundred ns of overhead per call) and talks to cuDNN FMHA
|
||||
straight. Useful as a ceiling for 'pure cuDNN, no SDPA' on Blackwell."""
|
||||
try:
|
||||
from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
|
||||
except Exception as e:
|
||||
return float("nan"), f"import-{type(e).__name__}"
|
||||
|
||||
b, s, h, d = q.shape
|
||||
qo_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
|
||||
kv_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
|
||||
|
||||
def _call():
|
||||
return cudnn_batch_prefill_with_kv_cache(
|
||||
q.reshape(b * s, h, d),
|
||||
k.reshape(b * s, h, d),
|
||||
v.reshape(b * s, h, d),
|
||||
qo_indptr=qo_indptr,
|
||||
kv_indptr=kv_indptr,
|
||||
max_qo_len=s,
|
||||
max_kv_len=s,
|
||||
sm_scale=scale,
|
||||
causal=False,
|
||||
)
|
||||
|
||||
return _time_call(_call)
|
||||
|
||||
|
||||
def _print_table(title: str, rows: list[tuple[str, float, str]], baseline_name: str | None = None) -> None:
|
||||
print(f"\n{title}")
|
||||
print("-" * 88)
|
||||
header = f"{'backend':<24} {'median (ms)':>14} {'vs baseline':>12} status"
|
||||
print(header)
|
||||
print("-" * 88)
|
||||
baseline_ms = None
|
||||
if baseline_name is not None:
|
||||
for name, ms, _ in rows:
|
||||
if name == baseline_name and ms == ms: # not NaN
|
||||
baseline_ms = ms
|
||||
break
|
||||
for name, ms, err in rows:
|
||||
ms_str = f"{ms:>14.3f}" if ms == ms else f"{'n/a':>14}"
|
||||
if baseline_ms is not None and ms == ms:
|
||||
ratio = baseline_ms / ms
|
||||
ratio_str = f"{ratio:>11.2f}x"
|
||||
else:
|
||||
ratio_str = f"{'—':>12}"
|
||||
status = "ok" if not err else f"FAILED ({err})"
|
||||
print(f"{name:<24} {ms_str} {ratio_str} {status}")
|
||||
|
||||
|
||||
def _pick_winner(rows: list[tuple[str, float, str]]) -> tuple[str, float] | None:
|
||||
"""Return (backend, ms) of the fastest non-failing row, or None."""
|
||||
ok_rows = [(n, m) for n, m, err in rows if not err and m == m]
|
||||
if not ok_rows:
|
||||
return None
|
||||
return min(ok_rows, key=lambda x: x[1])
|
||||
|
||||
|
||||
def _bench_one_shape(shape: dict, dtype: torch.dtype, device: str) -> tuple[list, list]:
|
||||
torch.manual_seed(0)
|
||||
q, k, v = _make_qkv(**shape, device=device, dtype=dtype)
|
||||
scale = 1.0 / (shape["head_dim"] ** 0.5)
|
||||
|
||||
rows_nomask = _run_sdpa_variants(q, k, v, attn_mask=None, scale=scale)
|
||||
rows_nomask.append(("FLASHINFER (default)", *_run_flashinfer(q, k, v, scale)))
|
||||
# `trtllm-gen` has no sm_120 cubins (NVIDIA/TensorRT-LLM#11799) — skipped.
|
||||
for fi_backend in ("fa2", "fa3", "cutlass", "auto"):
|
||||
rows_nomask.append((f"FLASHINFER ({fi_backend})", *_run_flashinfer(q, k, v, scale, backend=fi_backend)))
|
||||
rows_nomask.append(("FLASHINFER (cudnn-batch)", *_run_flashinfer_cudnn_batch(q, k, v, scale)))
|
||||
rows_nomask.append(("FA4 (direct)", *_run_fa4(q, k, v, scale)))
|
||||
|
||||
mask = _make_mask(shape["batch"], shape["seq"], device=device, dtype=dtype, pad_tokens=256)
|
||||
rows_mask = _run_sdpa_variants(q, k, v, attn_mask=mask, scale=scale)
|
||||
rows_mask.append(("FLASHINFER (dense)", *_run_flashinfer(q, k, v, scale, attn_mask=mask)))
|
||||
rows_mask.append(("FA4 (direct)", float("nan"), "mask-not-supported"))
|
||||
|
||||
return rows_nomask, rows_mask
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--preset", choices=list(_PRESETS.keys()), default="hv15")
|
||||
parser.add_argument("--sweep", action="store_true", help="Run all presets and print a ranking")
|
||||
parser.add_argument("--batch", type=int, default=None)
|
||||
parser.add_argument("--heads", type=int, default=None)
|
||||
parser.add_argument("--seq", type=int, default=None)
|
||||
parser.add_argument("--head-dim", type=int, default=None)
|
||||
parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16")
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
args = parser.parse_args()
|
||||
|
||||
dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
|
||||
_env_report()
|
||||
print(f"dtype : {args.dtype}")
|
||||
|
||||
presets = list(_PRESETS.keys()) if args.sweep else [args.preset]
|
||||
summary: dict[str, tuple[str, float] | None] = {}
|
||||
|
||||
for preset_name in presets:
|
||||
shape = dict(_PRESETS[preset_name])
|
||||
if not args.sweep:
|
||||
for k in ("batch", "heads", "seq", "head_dim"):
|
||||
v = getattr(args, k if k != "head_dim" else "head_dim")
|
||||
if v is not None:
|
||||
shape[k] = v
|
||||
|
||||
print("\n" + "=" * 88)
|
||||
print(
|
||||
f"Preset: {preset_name} | batch={shape['batch']} heads={shape['heads']} "
|
||||
f"seq={shape['seq']} head_dim={shape['head_dim']}"
|
||||
)
|
||||
print("=" * 88)
|
||||
|
||||
rows_nomask, rows_mask = _bench_one_shape(shape, dtype, args.device)
|
||||
_print_table(f"[{preset_name}] No attention mask", rows_nomask, baseline_name="CUDNN_ATTENTION")
|
||||
_print_table(
|
||||
f"[{preset_name}] With attention mask (pad 256 tokens)", rows_mask, baseline_name="CUDNN_ATTENTION"
|
||||
)
|
||||
|
||||
summary[preset_name] = _pick_winner(rows_nomask)
|
||||
|
||||
if args.sweep or len(presets) > 1:
|
||||
print("\n" + "=" * 88)
|
||||
print("Winners per preset (no-mask path)")
|
||||
print("=" * 88)
|
||||
for preset_name, winner in summary.items():
|
||||
if winner is None:
|
||||
print(f" {preset_name:<10} — no successful backend")
|
||||
else:
|
||||
name, ms = winner
|
||||
print(f" {preset_name:<10} {name:<24} {ms:>8.3f} ms")
|
||||
|
||||
print("\nNotes:")
|
||||
print(" * Ratios are relative to CUDNN_ATTENTION. >1.0x means faster than cuDNN.")
|
||||
print(" * Mask-path winner inherits CUDNN_ATTN's fallback in the PR's backends.")
|
||||
print(" * `trtllm-gen` and FA4 4.0.0b10 are known-broken on sm_120 as of Apr 2026.")
|
||||
print(" * For e2e timings, run text_to_video.py / text_to_image.py while varying DIFFUSION_ATTENTION_BACKEND.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -6,31 +6,34 @@
|
||||
Benchmark online serving for diffusion models (Image/Video Generation).
|
||||
If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
|
||||
|
||||
Supports multiple backends:
|
||||
- vllm-omni: Uses /v1/chat/completions endpoint (default)
|
||||
- openai: Uses /v1/images/generations endpoint
|
||||
- v1/videos: Use /v1/videos endpoint
|
||||
Supports multiple endpoints:
|
||||
- /v1/chat/completions: OpenAI chat-compatible image requests (e.g. t2i, Qwen i2i)
|
||||
- /v1/images/edits: OpenAI image edit / IT2I (multipart; e.g. Hunyuan --bot-task think)
|
||||
- /v1/images/generations: OpenAI image generation requests
|
||||
- /v1/videos: Async video jobs
|
||||
|
||||
Legacy --backend vllm-omni and openai are aliases for chat/completions and images/generations.
|
||||
|
||||
Usage:
|
||||
# Video (vllm-omni backend)
|
||||
# Video (/v1/videos endpoint)
|
||||
t2v:
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni --dataset vbench --task t2v --num-prompts 10 \
|
||||
--endpoint /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
|
||||
--endpoint /v1/videos --dataset vbench --task i2v --num-prompts 10
|
||||
|
||||
|
||||
# Image (vllm-omni backend)
|
||||
# Image (/v1/chat/completions endpoint)
|
||||
t2i:
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
|
||||
--endpoint /v1/chat/completions --dataset vbench --task t2i --num-prompts 10 \
|
||||
--height 1024 --width 1024
|
||||
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni --dataset random --task t2i --num-prompts 1 \
|
||||
--endpoint /v1/chat/completions --dataset random --task t2i --num-prompts 1 \
|
||||
--max-concurrency 1 --enable-negative-prompt \
|
||||
--random-request-config '[
|
||||
{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
|
||||
@@ -39,20 +42,32 @@ Usage:
|
||||
{"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
|
||||
]'
|
||||
|
||||
i2i:
|
||||
ti2i (Hunyuan / OpenAI image edit API):
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni --dataset vbench --task i2i --num-prompts 10
|
||||
--endpoint /v1/images/edits --dataset random --task ti2i --num-prompts 10 \
|
||||
--bot-task think
|
||||
|
||||
# Image (openai backend)
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--endpoint /v1/images/edits \
|
||||
--dataset custom \
|
||||
--dataset-path custom_requests.jsonl \
|
||||
--task ti2i \
|
||||
--bot-task think
|
||||
|
||||
i2i (chat-based models such as Qwen-Image-Edit):
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--endpoint /v1/chat/completions --dataset vbench --task i2i --num-prompts 10
|
||||
|
||||
# Image (/v1/images/generations endpoint)
|
||||
t2i:
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend openai --dataset vbench --task t2i --num-prompts 10 \
|
||||
--endpoint /v1/images/generations --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 \
|
||||
--endpoint /v1/videos --dataset random --task t2v --num-prompts 1 \
|
||||
--max-concurrency 1 --enable-negative-prompt \
|
||||
--random-request-config '[
|
||||
{"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1}
|
||||
@@ -80,7 +95,12 @@ from typing import Any
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
import requests
|
||||
from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
|
||||
from backends import (
|
||||
RequestFuncInput,
|
||||
RequestFuncOutput,
|
||||
backends_function_mapping,
|
||||
normalize_endpoint,
|
||||
)
|
||||
from PIL import Image
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
@@ -127,7 +147,7 @@ class VBenchDataset(BaseDataset):
|
||||
def _load_data(self) -> list[dict[str, Any]]:
|
||||
if self.args.task == "t2v":
|
||||
return self._load_t2v_prompts()
|
||||
elif self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
|
||||
elif self.args.task in ["i2v", "ti2v", "ti2i", "i2i", "it2i"]:
|
||||
return self._load_i2v_data()
|
||||
else:
|
||||
return self._load_t2v_prompts()
|
||||
@@ -519,7 +539,7 @@ class TraceDataset(BaseDataset):
|
||||
single = row.get("image_path")
|
||||
image_paths = [single] if single else None
|
||||
|
||||
if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i"]:
|
||||
if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i", "it2i"]:
|
||||
raise ValueError(
|
||||
f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row."
|
||||
)
|
||||
@@ -553,6 +573,135 @@ class TraceDataset(BaseDataset):
|
||||
return [self[i] for i in range(len(self))]
|
||||
|
||||
|
||||
class CustomDataset(BaseDataset):
|
||||
"""
|
||||
Custom dataset that loads requests from a JSONL file.
|
||||
|
||||
Each line in the JSONL file should be a JSON object with the following fields:
|
||||
- prompt (required): The text prompt for the request
|
||||
- width (optional): Image/video width
|
||||
- height (optional): Image height
|
||||
- num_inference_steps (optional): Number of diffusion steps
|
||||
- seed (optional): Random seed
|
||||
- image_paths (optional): List of input image paths for i2i/i2v/ti2i tasks
|
||||
- image_urls (optional): List of input image URLs (alternative to image_paths)
|
||||
|
||||
Example JSONL for ti2i:
|
||||
{"prompt": "Add sunset lighting", "width": 1024, "height": 1024, "image_urls": ["https://example.com/image.jpg"]}
|
||||
"""
|
||||
|
||||
def __init__(self, args, api_url: str, model: str):
|
||||
super().__init__(args, api_url, model)
|
||||
self.dataset_path = args.dataset_path
|
||||
if not self.dataset_path:
|
||||
raise ValueError("--dataset-path must be provided when using 'custom' dataset")
|
||||
self.load_data()
|
||||
|
||||
def load_data(self) -> None:
|
||||
"""Load data from JSONL file."""
|
||||
import pandas as pd
|
||||
|
||||
if not self.dataset_path.endswith(".jsonl"):
|
||||
raise ValueError("Custom dataset must be a JSONL file")
|
||||
|
||||
try:
|
||||
df = pd.read_json(path_or_buf=self.dataset_path, lines=True)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to load JSONL file {self.dataset_path}: {e}")
|
||||
|
||||
if "prompt" not in df.columns:
|
||||
raise ValueError("JSONL file must contain a 'prompt' column")
|
||||
|
||||
self.data = df.to_dict("records")
|
||||
print(f"Loaded {len(self.data)} requests from {self.dataset_path}")
|
||||
|
||||
def _resolve_image_paths(self, item: dict) -> list[str] | None:
|
||||
# Handle local file paths
|
||||
if "image_paths" in item and item["image_paths"]:
|
||||
paths = item["image_paths"]
|
||||
if isinstance(paths, str):
|
||||
paths = [paths]
|
||||
|
||||
# Verify all paths exist
|
||||
valid_paths = []
|
||||
for path in paths:
|
||||
if os.path.exists(path):
|
||||
valid_paths.append(path)
|
||||
else:
|
||||
raise ValueError(f"Image file not found: {path}")
|
||||
|
||||
return valid_paths if valid_paths else None
|
||||
|
||||
# Handle URLs - download to temp files
|
||||
if "image_urls" in item and item["image_urls"]:
|
||||
urls = item["image_urls"]
|
||||
if isinstance(urls, str):
|
||||
urls = [urls]
|
||||
|
||||
downloaded_paths = []
|
||||
for url in urls:
|
||||
try:
|
||||
response = requests.get(url, timeout=30)
|
||||
response.raise_for_status()
|
||||
|
||||
# Create temp file with appropriate extension
|
||||
suffix = os.path.splitext(url)[1] or ".png"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
|
||||
tmp.write(response.content)
|
||||
downloaded_paths.append(tmp.name)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to download image from {url}: {e}")
|
||||
|
||||
return downloaded_paths if downloaded_paths else None
|
||||
|
||||
return None
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.data)
|
||||
|
||||
def __getitem__(self, idx: int) -> RequestFuncInput:
|
||||
item = self.data[idx]
|
||||
|
||||
# Extract fields with fallback to args defaults
|
||||
prompt = item["prompt"]
|
||||
width = item.get("width", self.args.width)
|
||||
height = item.get("height", self.args.height)
|
||||
num_inference_steps = item.get("num_inference_steps", self.args.num_inference_steps)
|
||||
seed = item.get("seed", self.args.seed)
|
||||
reserved_keys = {"prompt", "width", "height", "num_inference_steps", "seed", "image_paths", "image_urls"}
|
||||
extra_body = {k: v for k, v in item.items() if k not in reserved_keys}
|
||||
|
||||
# Handle image paths/URLs
|
||||
image_paths = self._resolve_image_paths(item)
|
||||
|
||||
return RequestFuncInput(
|
||||
prompt=prompt,
|
||||
api_url=self.api_url,
|
||||
model=self.model,
|
||||
seed=seed,
|
||||
image_paths=image_paths,
|
||||
extra_body=extra_body,
|
||||
width=width,
|
||||
height=height,
|
||||
num_inference_steps=num_inference_steps,
|
||||
)
|
||||
|
||||
def get_requests(self) -> list[RequestFuncInput]:
|
||||
"""Get all requests, cycling through data if num_prompts > dataset size."""
|
||||
num_requests = getattr(self.args, "num_prompts", len(self.data))
|
||||
|
||||
if num_requests <= len(self.data):
|
||||
return [self[i] for i in range(num_requests)]
|
||||
|
||||
# Cycle through the dataset to reach num_prompts
|
||||
requests = []
|
||||
for i in range(num_requests):
|
||||
idx = i % len(self.data)
|
||||
requests.append(self[idx])
|
||||
|
||||
return requests
|
||||
|
||||
|
||||
class RandomDataset(BaseDataset):
|
||||
def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool = False):
|
||||
super().__init__(args, api_url, model)
|
||||
@@ -580,7 +729,7 @@ class RandomDataset(BaseDataset):
|
||||
self._sampled_requests = None
|
||||
|
||||
# Random image generate
|
||||
if self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
|
||||
if self.args.task in ["i2v", "ti2v", "ti2i", "i2i", "it2i"]:
|
||||
self._random_image_path = self._generate_random_image_paths()
|
||||
else:
|
||||
self._random_image_path = None
|
||||
@@ -751,6 +900,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,
|
||||
@@ -837,13 +1032,23 @@ def wait_for_service(base_url: str, timeout: int = 120) -> None:
|
||||
time.sleep(1)
|
||||
|
||||
|
||||
def _default_endpoint_for_task(task: str) -> str:
|
||||
if task in {"t2v", "i2v", "ti2v"}:
|
||||
return "/v1/videos"
|
||||
if task in {"i2i", "ti2i", "it2i"}:
|
||||
return "/v1/images/edits"
|
||||
if task == "t2i":
|
||||
return "/v1/chat/completions"
|
||||
raise ValueError(f"Unsupported task for endpoint resolution: {task}")
|
||||
|
||||
|
||||
async def benchmark(args):
|
||||
# Construct base_url if not provided
|
||||
if args.base_url is None:
|
||||
args.base_url = f"http://{args.host}:{args.port}"
|
||||
|
||||
VIDEO_TASKS = {"t2v", "i2v", "ti2v"}
|
||||
IMAGE_TASKS = {"t2i", "i2i", "ti2i"}
|
||||
IMAGE_TASKS = {"t2i", "i2i", "ti2i", "it2i"}
|
||||
|
||||
if args.task in VIDEO_TASKS:
|
||||
task_type = "2v"
|
||||
@@ -856,18 +1061,23 @@ async def benchmark(args):
|
||||
f"Valid image tasks: {sorted(IMAGE_TASKS)}"
|
||||
)
|
||||
|
||||
valid_backends = sorted(backends_function_mapping[task_type].keys())
|
||||
raw_endpoint = args.endpoint if args.endpoint is not None else args.backend
|
||||
if raw_endpoint is None:
|
||||
raw_endpoint = _default_endpoint_for_task(args.task)
|
||||
args.endpoint = normalize_endpoint(raw_endpoint)
|
||||
|
||||
if args.backend not in valid_backends:
|
||||
valid_endpoints = sorted(backends_function_mapping[task_type].keys())
|
||||
|
||||
if args.endpoint not in valid_endpoints:
|
||||
logger.error(
|
||||
f"Invalid backend '{args.backend}' for task '{args.task}' (task type: '{task_type}').\n"
|
||||
f"Valid backends for this task type: {valid_backends}\n"
|
||||
f"Example usage: --task {args.task} --backend {valid_backends[0]}"
|
||||
f"Invalid endpoint '{args.endpoint}' for task '{args.task}' (task type: '{task_type}').\n"
|
||||
f"Valid endpoints for this task type: {valid_endpoints}\n"
|
||||
f"Example usage: --task {args.task} --endpoint {valid_endpoints[0]}"
|
||||
)
|
||||
raise ValueError("Backend validation failed. See log above for valid options.")
|
||||
raise ValueError("Endpoint validation failed. See log above for valid options.")
|
||||
|
||||
# Setup API URL and request function based on backend
|
||||
request_func, api_url = backends_function_mapping[task_type][args.backend]
|
||||
# Setup API URL and request function based on endpoint.
|
||||
request_func, api_url = backends_function_mapping[task_type][args.endpoint]
|
||||
api_url = f"{args.base_url}{api_url}"
|
||||
|
||||
if args.dataset == "vbench":
|
||||
@@ -876,6 +1086,8 @@ async def benchmark(args):
|
||||
dataset = TraceDataset(args, api_url, args.model)
|
||||
elif args.dataset == "random":
|
||||
dataset = RandomDataset(args, api_url, args.model, args.enable_negative_prompt)
|
||||
elif args.dataset == "custom":
|
||||
dataset = CustomDataset(args, api_url, args.model)
|
||||
else:
|
||||
raise ValueError(f"Unknown dataset: {args.dataset}")
|
||||
|
||||
@@ -883,6 +1095,10 @@ async def benchmark(args):
|
||||
requests_list = dataset.get_requests()
|
||||
print(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.")
|
||||
|
||||
if args.endpoint == "/v1/images/edits":
|
||||
for req in requests_list:
|
||||
req.default_bot_task = args.bot_task
|
||||
|
||||
# Limit concurrency
|
||||
if args.max_concurrency is not None:
|
||||
semaphore = asyncio.Semaphore(args.max_concurrency)
|
||||
@@ -900,23 +1116,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.
|
||||
@@ -941,15 +1146,17 @@ async def benchmark(args):
|
||||
metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo)
|
||||
|
||||
# Add configuration info to metrics for JSON output
|
||||
metrics["backend"] = args.backend
|
||||
metrics["endpoint"] = args.endpoint
|
||||
metrics["model"] = args.model
|
||||
metrics["dataset"] = args.dataset
|
||||
metrics["task"] = args.task
|
||||
if args.endpoint == "/v1/images/edits":
|
||||
metrics["bot_task"] = args.bot_task
|
||||
|
||||
print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="="))
|
||||
print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))
|
||||
|
||||
# Section 1: Configuration
|
||||
print("{:<40} {:<15}".format("Backend:", args.backend))
|
||||
print("{:<40} {:<15}".format("Endpoint:", args.endpoint))
|
||||
print("{:<40} {:<15}".format("Model:", args.model))
|
||||
print("{:<40} {:<15}".format("Dataset:", args.dataset))
|
||||
print("{:<40} {:<15}".format("Task:", args.task))
|
||||
@@ -1012,25 +1219,30 @@ if __name__ == "__main__":
|
||||
parser.add_argument("--host", type=str, default="localhost", help="Server host.")
|
||||
parser.add_argument("--port", type=int, default=8091, help="Server port.")
|
||||
parser.add_argument("--model", type=str, default="default", help="Model name.")
|
||||
parser.add_argument(
|
||||
"--endpoint",
|
||||
type=str,
|
||||
default=None,
|
||||
help=("Endpoint path to target. Leading '/' is optional, e.g. /v1/videos or v1/videos."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--backend",
|
||||
type=str,
|
||||
default="vllm-omni",
|
||||
choices=["vllm-omni", "openai", "v1/videos"],
|
||||
help="Backend to target the benchmark to.",
|
||||
default=None,
|
||||
help=argparse.SUPPRESS,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dataset",
|
||||
type=str,
|
||||
default="vbench",
|
||||
choices=["vbench", "trace", "random"],
|
||||
choices=["vbench", "trace", "random", "custom"],
|
||||
help="Dataset to use.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--task",
|
||||
type=str,
|
||||
default="t2v",
|
||||
choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "t2i"],
|
||||
choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "it2i", "t2i"],
|
||||
help="Task type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -1072,6 +1284,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).")
|
||||
@@ -1130,10 +1349,15 @@ if __name__ == "__main__":
|
||||
default=1,
|
||||
help=(
|
||||
"Number of synthetic input images to attach for image-conditioned tasks "
|
||||
"(i2v, ti2v, ti2i, i2i) when using random dataset."
|
||||
"(i2v, ti2v, ti2i, i2i, it2i) when using random dataset."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--bot-task",
|
||||
type=str,
|
||||
default="think",
|
||||
help=("bot_task form field for --endpoint /v1/images/edits (think, recaption, think_recaption, vanilla)."),
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
asyncio.run(benchmark(args))
|
||||
@@ -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 |
|
||||
@@ -52,7 +58,7 @@ Record these parameters when reporting performance results.
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni \
|
||||
--endpoint /v1/chat/completions \
|
||||
--dataset <DATASET_NAME> \
|
||||
--task t2i \
|
||||
--num-prompts <N> \
|
||||
@@ -65,11 +71,12 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
|
||||
| Parameter | Description |
|
||||
| ---------------------- | --------------------------------- |
|
||||
| `--backend` | Serving backend (use `vllm-omni`) |
|
||||
| `--endpoint` | API endpoint (use `/v1/chat/completions`; leading `/` is optional) |
|
||||
| `--dataset` | Dataset name (`random` or custom) |
|
||||
| `--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 |
|
||||
|
||||
---
|
||||
@@ -122,7 +129,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni \
|
||||
--endpoint /v1/chat/completions \
|
||||
--dataset random \
|
||||
--task t2i \
|
||||
--num-prompts 1 \
|
||||
@@ -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
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -53,7 +53,7 @@ Record these parameters when reporting performance results.
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend v1/videos \
|
||||
--endpoint /v1/videos \
|
||||
--dataset <DATASET_NAME> \
|
||||
--task t2v \
|
||||
--num-prompts <N> \
|
||||
@@ -66,7 +66,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
|
||||
| Parameter | Description |
|
||||
| ---------------------- | --------------------------------- |
|
||||
| `--backend` | Serving backend (use `v1/videos`) |
|
||||
| `--endpoint` | API endpoint (use `/v1/videos`; leading `/` is optional) |
|
||||
| `--dataset` | Dataset name (`random` or custom) |
|
||||
| `--task` | Task type (e.g., `t2v`) |
|
||||
| `--num-prompts` | Total number of requests |
|
||||
@@ -120,7 +120,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend v1/videos \
|
||||
--endpoint /v1/videos \
|
||||
--dataset random \
|
||||
--task t2v \
|
||||
--num-prompts 1 \
|
||||
|
||||
@@ -34,6 +34,16 @@ Video example (text-to-video):
|
||||
--height 720 --width 1280 \
|
||||
--num-frames 81 --num-inference-steps 40 --seed 42
|
||||
|
||||
LTX-2 example (text-to-video; audio output is generated but not scored by LPIPS):
|
||||
python benchmarks/diffusion/quantization_quality.py \
|
||||
--model Lightricks/LTX-Video-2 \
|
||||
--task t2v \
|
||||
--quantization fp8 int8 \
|
||||
--prompts \
|
||||
"A serene lakeside sunrise with mist over the water" \
|
||||
--height 704 --width 1216 \
|
||||
--num-frames 121 --num-inference-steps 40 --seed 42
|
||||
|
||||
Multiple quantization methods:
|
||||
python benchmarks/diffusion/quantization_quality.py \
|
||||
--model Tongyi-MAI/Z-Image-Turbo \
|
||||
@@ -148,10 +158,11 @@ def _build_omni_kwargs(args, quantization=None):
|
||||
def _generate_image(omni, args, prompt, seed):
|
||||
"""Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
torch.accelerator.reset_peak_memory_stats()
|
||||
start = time.perf_counter()
|
||||
outputs = omni.generate(
|
||||
{"prompt": prompt},
|
||||
@@ -163,10 +174,11 @@ 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
|
||||
req_out = OmniRequestOutput.unwrap_result(outputs)
|
||||
if not req_out.images:
|
||||
raise ValueError("Could not extract image output from result.")
|
||||
img = req_out.images[0]
|
||||
return img, elapsed, peak_mem
|
||||
|
||||
@@ -178,7 +190,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 +204,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):
|
||||
@@ -202,10 +214,19 @@ def _generate_video(omni, args, prompt, seed):
|
||||
else:
|
||||
frames = inner
|
||||
elif hasattr(first, "images") and first.images:
|
||||
frames = first.images
|
||||
frames = first.images[0]
|
||||
else:
|
||||
raise ValueError("Could not extract video frames from output.")
|
||||
|
||||
# LTX-2 (and similar audio+video models) may surface a dict or (video, audio) tuple
|
||||
if isinstance(frames, dict):
|
||||
frames = frames.get("video") or frames.get("frames")
|
||||
elif isinstance(frames, tuple) and len(frames) == 2:
|
||||
frames = frames[0]
|
||||
|
||||
if frames is None:
|
||||
raise ValueError("Could not extract video frames from output.")
|
||||
|
||||
if isinstance(frames, torch.Tensor):
|
||||
video = frames.detach().cpu()
|
||||
if video.dim() == 5:
|
||||
@@ -223,13 +244,15 @@ def _generate_video(omni, args, prompt, seed):
|
||||
return frames_array, elapsed, peak_mem
|
||||
|
||||
|
||||
def _unload_omni(omni):
|
||||
"""Delete Omni instance and free GPU memory."""
|
||||
del omni
|
||||
def _free_gpu_memory():
|
||||
"""Force GC and release cached GPU memory.
|
||||
|
||||
Must be called AFTER the caller has dropped (i.e., via `del`)
|
||||
every reference to the Omni instance"""
|
||||
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):
|
||||
@@ -265,7 +288,9 @@ def run_benchmark(args):
|
||||
|
||||
bl_avg_time = np.mean([v[1] for v in baseline_outputs.values()])
|
||||
bl_mem = baseline_outputs[prompts[0]][2] # use first prompt's memory
|
||||
_unload_omni(omni_bl)
|
||||
omni_bl.shutdown()
|
||||
del omni_bl
|
||||
_free_gpu_memory()
|
||||
|
||||
# Save baseline outputs
|
||||
bl_dir = output_dir / "baseline"
|
||||
@@ -305,7 +330,9 @@ def run_benchmark(args):
|
||||
|
||||
qt_avg_time = np.mean([v[1] for v in qt_outputs.values()])
|
||||
qt_mem = qt_outputs[prompts[0]][2]
|
||||
_unload_omni(omni_qt)
|
||||
omni_qt.shutdown()
|
||||
del omni_qt
|
||||
_free_gpu_memory()
|
||||
|
||||
# Save quantized outputs
|
||||
qt_dir = output_dir / config_label.replace(" ", "_")
|
||||
|
||||
@@ -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,244 @@
|
||||
# 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,audio_underrun \
|
||||
--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,audio_underrun \
|
||||
--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,audio_underrun \
|
||||
--save-result --result-dir ./results
|
||||
```
|
||||
|
||||
#### Streaming continuity (`audio_underrun`)
|
||||
|
||||
`audio_underrun` (seconds) is the per-request worst-case buffer deficit
|
||||
under a realtime-rate player simulation. The audio-speech backend records
|
||||
chunk arrival times during the SSE stream and surfaces:
|
||||
|
||||
- `Mean / Median / P{50,99} AUDIO_UNDERRUN (s)` per the standard percentile
|
||||
output - any positive value means at least one inter-chunk gap was longer
|
||||
than the chunk's audio duration.
|
||||
- `Streaming continuity OK rate` - fraction of requests whose worst-case
|
||||
underrun stayed under the threshold (default 100 ms, the commonly cited
|
||||
"audible gap" budget). Override the threshold with
|
||||
`VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S=<float>`.
|
||||
|
||||
This captures the failure mode where `RTF_p50 < 1` (server keeps up in
|
||||
aggregate) but per-stream chunk arrival is bursty enough that listeners
|
||||
still hear gaps - common at high concurrency on streaming TTS pipelines.
|
||||
|
||||
#### Add WER / SIM / UTMOS to any of the above
|
||||
|
||||
Append `--seed-tts-wer-eval` (and optionally `SEED_TTS_EVAL_DEVICE=cuda:0`
|
||||
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
|
||||
├── 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,audio_underrun",
|
||||
]
|
||||
|
||||
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()
|
||||
@@ -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,890 +0,0 @@
|
||||
"""Offline VoxCPM benchmark for vLLM Omni.
|
||||
|
||||
Supports both:
|
||||
- sync one-shot (Omni.generate)
|
||||
- streaming (AsyncOmni.generate with async_chunk config)
|
||||
- text-only synthesis
|
||||
- voice cloning
|
||||
- text/clone batch inputs from txt or jsonl
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import AsyncOmni, Omni
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml"
|
||||
DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml"
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PromptSpec:
|
||||
text: str
|
||||
label: str
|
||||
ref_audio: str | None = None
|
||||
ref_text: str | None = None
|
||||
|
||||
|
||||
def _require_soundfile():
|
||||
try:
|
||||
import soundfile as sf # type: ignore
|
||||
except ModuleNotFoundError as exc:
|
||||
raise RuntimeError(
|
||||
"soundfile is required to write VoxCPM benchmark WAV outputs. Install it with: pip install soundfile"
|
||||
) from exc
|
||||
return sf
|
||||
|
||||
|
||||
def _build_prompt(
|
||||
args,
|
||||
*,
|
||||
text: str,
|
||||
ref_audio: str | None = None,
|
||||
ref_text: str | None = None,
|
||||
global_request_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
additional_information: dict[str, list[Any]] = {
|
||||
"text": [text],
|
||||
"cfg_value": [args.cfg_value],
|
||||
"inference_timesteps": [args.inference_timesteps],
|
||||
"min_len": [args.min_len],
|
||||
"max_new_tokens": [args.max_new_tokens],
|
||||
}
|
||||
if args.streaming_prefix_len is not None:
|
||||
additional_information["streaming_prefix_len"] = [args.streaming_prefix_len]
|
||||
|
||||
if ref_audio:
|
||||
additional_information["ref_audio"] = [ref_audio]
|
||||
if ref_text:
|
||||
additional_information["ref_text"] = [ref_text]
|
||||
if global_request_id is not None:
|
||||
additional_information["global_request_id"] = [global_request_id]
|
||||
|
||||
return {
|
||||
"prompt_token_ids": [1],
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
|
||||
|
||||
def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor:
|
||||
audio = mm.get("audio", mm.get("model_outputs"))
|
||||
if audio is None:
|
||||
raise ValueError("No audio output found in multimodal output.")
|
||||
if isinstance(audio, list):
|
||||
parts = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio]
|
||||
audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0)
|
||||
if not isinstance(audio, torch.Tensor):
|
||||
audio = torch.as_tensor(audio)
|
||||
return audio.float().cpu().reshape(-1)
|
||||
|
||||
|
||||
def _extract_sample_rate(mm: dict[str, Any]) -> int:
|
||||
sr_raw = mm.get("sr", 24000)
|
||||
if isinstance(sr_raw, list) and sr_raw:
|
||||
sr_raw = sr_raw[-1]
|
||||
if hasattr(sr_raw, "item"):
|
||||
return int(sr_raw.item())
|
||||
return int(sr_raw)
|
||||
|
||||
|
||||
def _emit_offline_metrics(
|
||||
*,
|
||||
request_id: str,
|
||||
elapsed_s: float,
|
||||
first_audio_elapsed: float | None,
|
||||
audio_duration_s: float,
|
||||
) -> None:
|
||||
metrics = {
|
||||
"request_id": request_id,
|
||||
"ttfp_ms": round(first_audio_elapsed * 1000.0, 3) if first_audio_elapsed is not None else None,
|
||||
"audio_duration_s": round(audio_duration_s, 6),
|
||||
"rtf": round(elapsed_s / audio_duration_s, 6) if audio_duration_s > 0 else None,
|
||||
}
|
||||
print(f"[OfflineMetrics] {metrics}")
|
||||
|
||||
|
||||
def _write_audio_tensor(output_path: Path, audio_tensor: Any, sample_rate: int) -> None:
|
||||
sf = _require_soundfile()
|
||||
if isinstance(audio_tensor, torch.Tensor):
|
||||
audio_np = audio_tensor.float().cpu().clamp(-1.0, 1.0).numpy()
|
||||
else:
|
||||
audio_np = torch.as_tensor(audio_tensor).float().cpu().clamp(-1.0, 1.0).numpy()
|
||||
sf.write(
|
||||
output_path,
|
||||
audio_np,
|
||||
sample_rate,
|
||||
format="WAV",
|
||||
subtype="PCM_16",
|
||||
)
|
||||
|
||||
|
||||
def _save_wav(mm: dict[str, Any], output_dir: Path, request_id: str) -> Path:
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
output_path = output_dir / f"output_{request_id}.wav"
|
||||
_write_audio_tensor(output_path, _extract_audio_tensor(mm), _extract_sample_rate(mm))
|
||||
return output_path
|
||||
|
||||
|
||||
def _iter_request_multimodal_outputs(request_output: Any):
|
||||
outputs = getattr(request_output, "outputs", None)
|
||||
if outputs:
|
||||
for output in outputs:
|
||||
mm = getattr(output, "multimodal_output", None)
|
||||
if isinstance(mm, dict):
|
||||
yield mm
|
||||
|
||||
mm = getattr(request_output, "multimodal_output", None)
|
||||
if isinstance(mm, dict):
|
||||
yield mm
|
||||
|
||||
|
||||
def _read_non_empty_lines(path: str) -> list[str]:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
return [line.strip() for line in f if line.strip()]
|
||||
|
||||
|
||||
def _load_prompt_specs(args) -> list[PromptSpec]:
|
||||
specs: list[PromptSpec] = []
|
||||
|
||||
if args.txt_prompts is not None:
|
||||
texts = _read_non_empty_lines(args.txt_prompts)
|
||||
if not texts:
|
||||
raise ValueError(f"No prompts found in {args.txt_prompts}")
|
||||
for idx, text in enumerate(texts, start=1):
|
||||
specs.append(
|
||||
PromptSpec(
|
||||
text=text,
|
||||
label=f"item{idx:03d}",
|
||||
ref_audio=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
)
|
||||
)
|
||||
return specs
|
||||
|
||||
if args.jsonl_prompts is not None:
|
||||
with open(args.jsonl_prompts, encoding="utf-8") as f:
|
||||
for line_no, raw_line in enumerate(f, start=1):
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
item = json.loads(line)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"{args.jsonl_prompts}:{line_no} is not valid JSON: {exc}") from exc
|
||||
if not isinstance(item, dict):
|
||||
raise ValueError(f"{args.jsonl_prompts}:{line_no} must be a JSON object")
|
||||
|
||||
text = item.get("text")
|
||||
if not isinstance(text, str) or not text.strip():
|
||||
raise ValueError(f"{args.jsonl_prompts}:{line_no} requires non-empty string field 'text'")
|
||||
|
||||
ref_audio = item.get("ref_audio", args.ref_audio)
|
||||
ref_text = item.get("ref_text", args.ref_text)
|
||||
if (ref_audio is None) != (ref_text is None):
|
||||
raise ValueError(
|
||||
f"{args.jsonl_prompts}:{line_no} must provide both 'ref_audio' and 'ref_text' together"
|
||||
)
|
||||
|
||||
specs.append(
|
||||
PromptSpec(
|
||||
text=text.strip(),
|
||||
label=f"item{len(specs) + 1:03d}",
|
||||
ref_audio=ref_audio,
|
||||
ref_text=ref_text,
|
||||
)
|
||||
)
|
||||
|
||||
if not specs:
|
||||
raise ValueError(f"No prompts found in {args.jsonl_prompts}")
|
||||
return specs
|
||||
|
||||
specs.append(
|
||||
PromptSpec(
|
||||
text=args.text,
|
||||
label="item001",
|
||||
ref_audio=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
)
|
||||
)
|
||||
return specs
|
||||
|
||||
|
||||
def _build_prompt_for_spec(args, spec: PromptSpec, *, global_request_id: str | None = None) -> dict[str, Any]:
|
||||
return _build_prompt(
|
||||
args,
|
||||
text=spec.text,
|
||||
ref_audio=spec.ref_audio,
|
||||
ref_text=spec.ref_text,
|
||||
global_request_id=global_request_id,
|
||||
)
|
||||
|
||||
|
||||
def _count_voice_clone_prompts(prompt_specs: list[PromptSpec]) -> int:
|
||||
return sum(1 for spec in prompt_specs if spec.ref_audio is not None)
|
||||
|
||||
|
||||
def _get_warmup_specs(prompt_specs: list[PromptSpec]) -> list[PromptSpec]:
|
||||
return prompt_specs[:1]
|
||||
|
||||
|
||||
def _extract_stream_finished(stage_output: Any) -> bool:
|
||||
request_output = getattr(stage_output, "request_output", None)
|
||||
request_finished = getattr(request_output, "finished", None)
|
||||
if request_finished is not None:
|
||||
return bool(request_finished)
|
||||
return bool(getattr(stage_output, "finished", False))
|
||||
|
||||
|
||||
def _build_profiled_stage_config(
|
||||
stage_configs_path: str,
|
||||
profiler_dir: str,
|
||||
) -> str:
|
||||
stage_config_path = Path(stage_configs_path)
|
||||
yaml_text = stage_config_path.read_text(encoding="utf-8")
|
||||
injected_lines: list[str] = []
|
||||
injected_count = 0
|
||||
|
||||
for line in yaml_text.splitlines():
|
||||
injected_lines.append(line)
|
||||
if line.strip() != "engine_args:":
|
||||
continue
|
||||
indent = line[: len(line) - len(line.lstrip())]
|
||||
child_indent = indent + " "
|
||||
grandchild_indent = child_indent + " "
|
||||
injected_lines.extend(
|
||||
[
|
||||
f"{child_indent}profiler_config:",
|
||||
f'{grandchild_indent}profiler: "torch"',
|
||||
f'{grandchild_indent}torch_profiler_dir: "{profiler_dir}"',
|
||||
f"{grandchild_indent}torch_profiler_with_stack: true",
|
||||
]
|
||||
)
|
||||
injected_count += 1
|
||||
|
||||
if injected_count == 0:
|
||||
raise ValueError(f"No engine_args block found in stage config: {stage_configs_path}")
|
||||
|
||||
tmp = tempfile.NamedTemporaryFile(
|
||||
mode="w",
|
||||
encoding="utf-8",
|
||||
delete=False,
|
||||
suffix=".yaml",
|
||||
prefix=f"{stage_config_path.stem}_profile_",
|
||||
)
|
||||
tmp.write("\n".join(injected_lines) + "\n")
|
||||
tmp.close()
|
||||
return tmp.name
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(
|
||||
description="Offline split-stage VoxCPM inference with vLLM Omni (auto sync/streaming by stage config)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default=os.environ.get("VOXCPM_MODEL"),
|
||||
help="Local VoxCPM model directory. Defaults to $VOXCPM_MODEL.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
type=str,
|
||||
default="This is a split-stage VoxCPM synthesis example running on vLLM Omni.",
|
||||
help="Text to synthesize. Ignored when --txt-prompts or --jsonl-prompts is used.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--txt-prompts",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a .txt file with one synthesis text per line.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--jsonl-prompts",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Path to a .jsonl file. Each line must contain at least {'text': ...}; "
|
||||
"clone rows can also set ref_audio/ref_text, and ref_text must be the "
|
||||
"real transcript of ref_audio."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Optional reference audio path for voice cloning. With --txt-prompts, "
|
||||
"the same reference is applied to every line."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-text",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Real transcript of the reference audio. Placeholder text or mismatched "
|
||||
"text will usually produce noisy/electronic clone audio."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=str(DEFAULT_STAGE_SYNC),
|
||||
help="Stage config YAML path. Routing is selected only from this path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg-value",
|
||||
type=float,
|
||||
default=2.0,
|
||||
help="Classifier-free guidance value for VoxCPM.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--inference-timesteps",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Number of inference timesteps.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min-len",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Minimum generated token length.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-new-tokens",
|
||||
type=int,
|
||||
default=4096,
|
||||
help="Maximum generated token length.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--streaming-prefix-len",
|
||||
type=int,
|
||||
default=None,
|
||||
help="VoxCPM streaming window (optional, streaming mode only).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory for output WAV files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-init-timeout",
|
||||
type=int,
|
||||
default=600,
|
||||
help="Stage initialization timeout in seconds.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
dest="log_stats",
|
||||
action="store_true",
|
||||
help="Enable vLLM Omni stats logging.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-log-stats",
|
||||
dest="log_stats",
|
||||
action="store_false",
|
||||
help="Disable vLLM Omni stats logging.",
|
||||
)
|
||||
parser.set_defaults(log_stats=True)
|
||||
parser.add_argument(
|
||||
"--num-runs",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of full inference runs (same prompt each time). Default 1.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup-runs",
|
||||
type=int,
|
||||
default=0,
|
||||
help=(
|
||||
"Optional number of warmup passes before measured runs. Warmup uses only "
|
||||
"the first prompt and does not save outputs."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-profiler",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Enable torch profiler for the configured stages. A temporary profiled "
|
||||
"stage config is generated automatically."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory for profiler traces. Defaults to <output-dir>/profiler when profiling is enabled.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-stages",
|
||||
type=int,
|
||||
nargs="*",
|
||||
default=None,
|
||||
help="Optional stage ids to profile. Defaults to all stages that have profiler_config.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-wait-seconds",
|
||||
type=float,
|
||||
default=30.0,
|
||||
help="Seconds to wait after stop_profile for trace files to flush.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.model:
|
||||
parser.error("--model is required unless $VOXCPM_MODEL is set")
|
||||
if args.txt_prompts is not None and args.jsonl_prompts is not None:
|
||||
parser.error("--txt-prompts and --jsonl-prompts are mutually exclusive")
|
||||
if (args.ref_audio is None) != (args.ref_text is None):
|
||||
parser.error("--ref-audio and --ref-text must be provided together")
|
||||
if args.num_runs < 1:
|
||||
parser.error("--num-runs must be >= 1")
|
||||
if args.warmup_runs < 0:
|
||||
parser.error("--warmup-runs must be >= 0")
|
||||
if args.output_dir is None:
|
||||
args.output_dir = (
|
||||
"output_audio_streaming" if _is_streaming_stage_config(args.stage_configs_path) else "output_audio"
|
||||
)
|
||||
if args.enable_profiler and args.profiler_dir is None:
|
||||
args.profiler_dir = str(Path(args.output_dir) / "profiler")
|
||||
try:
|
||||
args.prompt_specs = _load_prompt_specs(args)
|
||||
except ValueError as exc:
|
||||
parser.error(str(exc))
|
||||
|
||||
return args
|
||||
|
||||
|
||||
def _is_streaming_stage_config(stage_configs_path: str) -> bool:
|
||||
cfg_name = Path(stage_configs_path).name.lower()
|
||||
# Keep routing purely config-path based:
|
||||
# - voxcpm.yaml => sync
|
||||
# - voxcpm_async_chunk.yaml => streaming
|
||||
return "async_chunk" in cfg_name
|
||||
|
||||
|
||||
async def _collect_streaming_audio(
|
||||
omni: AsyncOmni,
|
||||
args: Any,
|
||||
spec: PromptSpec,
|
||||
request_id: str,
|
||||
*,
|
||||
phase_label: str,
|
||||
prompt_index: int,
|
||||
prompt_count: int,
|
||||
print_prompt: bool = False,
|
||||
) -> tuple[torch.Tensor, int, float, float | None]:
|
||||
prompt = _build_prompt_for_spec(args, spec, global_request_id=request_id)
|
||||
delta_chunks: list[torch.Tensor] = []
|
||||
sample_rate = 24000
|
||||
chunk_i = 0
|
||||
prev_total_samples = 0
|
||||
t_start = time.perf_counter()
|
||||
first_audio_elapsed: float | None = None
|
||||
|
||||
if print_prompt:
|
||||
print(f"---prompt---:{prompt}")
|
||||
|
||||
async for stage_output in omni.generate(prompt, request_id=request_id):
|
||||
mm = getattr(stage_output, "multimodal_output", None)
|
||||
if not isinstance(mm, dict):
|
||||
ro = getattr(stage_output, "request_output", None)
|
||||
if ro is None:
|
||||
continue
|
||||
mm = getattr(ro, "multimodal_output", None)
|
||||
if not isinstance(mm, dict) and getattr(ro, "outputs", None):
|
||||
seq = ro.outputs[0]
|
||||
mm = getattr(seq, "multimodal_output", None)
|
||||
if not isinstance(mm, dict):
|
||||
continue
|
||||
sample_rate = _extract_sample_rate(mm)
|
||||
try:
|
||||
w = _extract_audio_tensor(mm)
|
||||
n = int(w.numel())
|
||||
if n == 0:
|
||||
continue
|
||||
finished = _extract_stream_finished(stage_output)
|
||||
if n > prev_total_samples:
|
||||
delta = w.reshape(-1)[prev_total_samples:]
|
||||
prev_total_samples = n
|
||||
elif finished and n == prev_total_samples:
|
||||
delta = w.reshape(-1)[:0]
|
||||
else:
|
||||
delta = w.reshape(-1)
|
||||
prev_total_samples += int(delta.numel())
|
||||
if int(delta.numel()) > 0:
|
||||
delta_chunks.append(delta)
|
||||
if first_audio_elapsed is None and int(delta.numel()) > 0:
|
||||
first_audio_elapsed = time.perf_counter() - t_start
|
||||
logger.info(
|
||||
"%s prompt=%d/%d chunk=%d delta_samples=%d buf_len=%d finished=%s",
|
||||
phase_label,
|
||||
prompt_index + 1,
|
||||
prompt_count,
|
||||
chunk_i,
|
||||
int(delta.numel()),
|
||||
n,
|
||||
finished,
|
||||
)
|
||||
chunk_i += 1
|
||||
except ValueError:
|
||||
if not _extract_stream_finished(stage_output):
|
||||
logger.debug("skip non-audio partial output chunk=%d", chunk_i)
|
||||
|
||||
if not delta_chunks:
|
||||
raise RuntimeError("No audio chunks received; check stage config and logs.")
|
||||
|
||||
audio_cat = torch.cat([c.reshape(-1) for c in delta_chunks], dim=0)
|
||||
elapsed = time.perf_counter() - t_start
|
||||
return audio_cat, sample_rate, elapsed, first_audio_elapsed
|
||||
|
||||
|
||||
async def _abort_streaming_residual_work(
|
||||
omni: AsyncOmni,
|
||||
request_id: str,
|
||||
*,
|
||||
settle_seconds: float = 0.1,
|
||||
) -> None:
|
||||
"""Stop any late stage-0 work once the final audio has been collected."""
|
||||
await omni.engine.abort_async([request_id])
|
||||
if settle_seconds > 0:
|
||||
await asyncio.sleep(settle_seconds)
|
||||
|
||||
|
||||
async def _run_streaming_single(
|
||||
omni: AsyncOmni,
|
||||
args: Any,
|
||||
spec: PromptSpec,
|
||||
output_dir: Path,
|
||||
request_id: str,
|
||||
*,
|
||||
run_index: int,
|
||||
num_runs: int,
|
||||
prompt_index: int,
|
||||
prompt_count: int,
|
||||
) -> Path:
|
||||
audio_cat, sample_rate, elapsed, first_audio_elapsed = await _collect_streaming_audio(
|
||||
omni,
|
||||
args,
|
||||
spec,
|
||||
request_id,
|
||||
phase_label=f"run={run_index + 1}/{num_runs}",
|
||||
prompt_index=prompt_index,
|
||||
prompt_count=prompt_count,
|
||||
print_prompt=(run_index == 0 and prompt_index == 0),
|
||||
)
|
||||
await _abort_streaming_residual_work(omni, request_id)
|
||||
output_path = output_dir / f"output_run{run_index + 1}_{spec.label}.wav"
|
||||
_write_audio_tensor(output_path, audio_cat, sample_rate)
|
||||
audio_duration_s = float(audio_cat.numel()) / float(sample_rate) if sample_rate > 0 else 0.0
|
||||
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
|
||||
rtf_text = f", rtf={elapsed / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
|
||||
print(
|
||||
f"Saved (streaming) run {run_index + 1}/{num_runs}, "
|
||||
f"prompt {prompt_index + 1}/{prompt_count}: {output_path} ({elapsed:.2f}s{ttfp_text}{rtf_text})"
|
||||
)
|
||||
_emit_offline_metrics(
|
||||
request_id=request_id,
|
||||
elapsed_s=elapsed,
|
||||
first_audio_elapsed=first_audio_elapsed,
|
||||
audio_duration_s=audio_duration_s,
|
||||
)
|
||||
return output_path
|
||||
|
||||
|
||||
async def _run_streaming_warmup(args, omni: AsyncOmni) -> None:
|
||||
if args.warmup_runs == 0:
|
||||
return
|
||||
|
||||
warmup_specs = _get_warmup_specs(args.prompt_specs)
|
||||
print(
|
||||
f"Warmup: {args.warmup_runs} run(s) using the first prompt "
|
||||
f"({len(warmup_specs)} prompt(s)); outputs will be discarded."
|
||||
)
|
||||
for warmup_index in range(args.warmup_runs):
|
||||
t_warmup = time.perf_counter()
|
||||
tasks = []
|
||||
request_ids: list[str] = []
|
||||
for prompt_index, spec in enumerate(warmup_specs):
|
||||
request_id = f"warmup_stream_{warmup_index + 1}_{spec.label}_{uuid.uuid4().hex[:8]}"
|
||||
request_ids.append(request_id)
|
||||
tasks.append(
|
||||
_collect_streaming_audio(
|
||||
omni,
|
||||
args,
|
||||
spec,
|
||||
request_id,
|
||||
phase_label=f"warmup={warmup_index + 1}/{args.warmup_runs}",
|
||||
prompt_index=prompt_index,
|
||||
prompt_count=len(warmup_specs),
|
||||
)
|
||||
)
|
||||
results = await asyncio.gather(*tasks)
|
||||
for request_id in request_ids:
|
||||
await _abort_streaming_residual_work(omni, request_id)
|
||||
total_samples = sum(int(audio.numel()) for audio, _, _, _ in results)
|
||||
warmup_ttfps = [ttfp for _, _, _, ttfp in results if ttfp is not None]
|
||||
ttfp_text = f", ttfp={min(warmup_ttfps):.2f}s" if warmup_ttfps else ""
|
||||
print(
|
||||
f"Warmup (streaming) {warmup_index + 1}/{args.warmup_runs} finished: "
|
||||
f"{len(results)} prompt(s), {total_samples} sample(s) "
|
||||
f"({time.perf_counter() - t_warmup:.2f}s{ttfp_text})"
|
||||
)
|
||||
|
||||
|
||||
async def _run_streaming(args) -> list[Path]:
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
omni = AsyncOmni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
await _run_streaming_warmup(args, omni)
|
||||
profiler_started = False
|
||||
if args.enable_profiler:
|
||||
profile_prefix = f"voxcpm_streaming_{int(time.time())}"
|
||||
stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured"
|
||||
print(f"Starting profiler (streaming): stages={stages_text}, dir={args.profiler_dir}")
|
||||
await omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages)
|
||||
profiler_started = True
|
||||
t_total = time.perf_counter()
|
||||
total_elapsed = 0.0
|
||||
paths: list[Path] = []
|
||||
prompt_specs: list[PromptSpec] = args.prompt_specs
|
||||
try:
|
||||
for run in range(args.num_runs):
|
||||
for prompt_index, spec in enumerate(prompt_specs):
|
||||
request_id = f"stream_{run + 1}_{spec.label}_{uuid.uuid4().hex[:8]}"
|
||||
paths.append(
|
||||
await _run_streaming_single(
|
||||
omni,
|
||||
args,
|
||||
spec,
|
||||
output_dir,
|
||||
request_id,
|
||||
run_index=run,
|
||||
num_runs=args.num_runs,
|
||||
prompt_index=prompt_index,
|
||||
prompt_count=len(prompt_specs),
|
||||
)
|
||||
)
|
||||
total_elapsed = time.perf_counter() - t_total
|
||||
finally:
|
||||
if profiler_started:
|
||||
print("Stopping profiler (streaming)...")
|
||||
await omni.stop_profile(stages=args.profiler_stages)
|
||||
if args.profiler_wait_seconds > 0:
|
||||
print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...")
|
||||
await asyncio.sleep(args.profiler_wait_seconds)
|
||||
|
||||
print(
|
||||
f"All streaming runs finished: {args.num_runs} run(s), "
|
||||
f"{len(prompt_specs)} prompt(s), {len(paths)} file(s) in {total_elapsed:.2f}s total"
|
||||
)
|
||||
return paths
|
||||
|
||||
|
||||
def _run_sync(args) -> list[Path]:
|
||||
output_dir = Path(args.output_dir)
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
def _run_sync_single(
|
||||
spec: PromptSpec,
|
||||
*,
|
||||
request_prefix: str,
|
||||
save_outputs: bool,
|
||||
run_index: int | None = None,
|
||||
) -> tuple[list[Path], int, float | None, float, float, str]:
|
||||
global_request_id = f"{request_prefix}_{spec.label}"
|
||||
prompt = _build_prompt_for_spec(args, spec, global_request_id=global_request_id)
|
||||
if save_outputs and run_index == 0 and spec.label == "item001":
|
||||
print(f"---prompt---:{prompt}")
|
||||
|
||||
saved_paths: list[Path] = []
|
||||
output_count = 0
|
||||
first_audio_elapsed: float | None = None
|
||||
total_audio_duration_s = 0.0
|
||||
metrics_request_id = global_request_id
|
||||
t_start = time.perf_counter()
|
||||
for stage_outputs in omni.generate(prompt):
|
||||
request_output = stage_outputs.request_output
|
||||
if request_output is None:
|
||||
continue
|
||||
request_output_id = getattr(request_output, "request_id", None)
|
||||
if isinstance(request_output_id, str) and request_output_id:
|
||||
metrics_request_id = request_output_id
|
||||
for j, mm in enumerate(_iter_request_multimodal_outputs(request_output)):
|
||||
output_count += 1
|
||||
if first_audio_elapsed is None:
|
||||
try:
|
||||
audio_tensor = _extract_audio_tensor(mm)
|
||||
if int(audio_tensor.numel()) > 0:
|
||||
first_audio_elapsed = time.perf_counter() - t_start
|
||||
total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm))
|
||||
except ValueError:
|
||||
pass
|
||||
else:
|
||||
try:
|
||||
audio_tensor = _extract_audio_tensor(mm)
|
||||
total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm))
|
||||
except ValueError:
|
||||
pass
|
||||
if not save_outputs:
|
||||
continue
|
||||
save_stem = f"run{run_index + 1}_{spec.label}" if j == 0 else f"run{run_index + 1}_{spec.label}_{j}"
|
||||
saved_paths.append(_save_wav(mm, output_dir, save_stem))
|
||||
|
||||
if output_count == 0:
|
||||
raise RuntimeError("No output from Omni.generate")
|
||||
elapsed_s = time.perf_counter() - t_start
|
||||
return saved_paths, output_count, first_audio_elapsed, elapsed_s, total_audio_duration_s, metrics_request_id
|
||||
|
||||
if args.warmup_runs:
|
||||
warmup_specs = _get_warmup_specs(args.prompt_specs)
|
||||
print(
|
||||
f"Warmup: {args.warmup_runs} run(s) using the first prompt "
|
||||
f"({len(warmup_specs)} prompt(s)); outputs will be discarded."
|
||||
)
|
||||
for warmup_index in range(args.warmup_runs):
|
||||
t_warmup = time.perf_counter()
|
||||
_, output_count, first_audio_elapsed, elapsed_s, audio_duration_s, _ = _run_sync_single(
|
||||
warmup_specs[0],
|
||||
request_prefix=f"warmup_sync{warmup_index + 1}",
|
||||
save_outputs=False,
|
||||
)
|
||||
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
|
||||
rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
|
||||
print(
|
||||
f"Warmup (sync) {warmup_index + 1}/{args.warmup_runs} finished: "
|
||||
f"{output_count} output(s) ({time.perf_counter() - t_warmup:.2f}s{ttfp_text}{rtf_text})"
|
||||
)
|
||||
|
||||
profiler_started = False
|
||||
if args.enable_profiler:
|
||||
profile_prefix = f"voxcpm_sync_{int(time.time())}"
|
||||
stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured"
|
||||
print(f"Starting profiler (sync): stages={stages_text}, dir={args.profiler_dir}")
|
||||
omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages)
|
||||
profiler_started = True
|
||||
|
||||
t_total = time.perf_counter()
|
||||
total_elapsed = 0.0
|
||||
saved_paths: list[Path] = []
|
||||
prompt_specs: list[PromptSpec] = args.prompt_specs
|
||||
try:
|
||||
for run in range(args.num_runs):
|
||||
t_run = time.perf_counter()
|
||||
run_paths: list[Path] = []
|
||||
for prompt_index, spec in enumerate(prompt_specs):
|
||||
prompt_paths, _, first_audio_elapsed, elapsed_s, audio_duration_s, metrics_request_id = (
|
||||
_run_sync_single(
|
||||
spec,
|
||||
request_prefix=f"sync_run{run + 1}_{prompt_index + 1:03d}",
|
||||
save_outputs=True,
|
||||
run_index=run,
|
||||
)
|
||||
)
|
||||
run_paths.extend(prompt_paths)
|
||||
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
|
||||
rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
|
||||
print(
|
||||
f"Saved (sync) run {run + 1}/{args.num_runs}, "
|
||||
f"prompt {prompt_index + 1}/{len(prompt_specs)}: {len(prompt_paths)} file(s){ttfp_text}{rtf_text}"
|
||||
)
|
||||
_emit_offline_metrics(
|
||||
request_id=metrics_request_id,
|
||||
elapsed_s=elapsed_s,
|
||||
first_audio_elapsed=first_audio_elapsed,
|
||||
audio_duration_s=audio_duration_s,
|
||||
)
|
||||
|
||||
saved_paths.extend(run_paths)
|
||||
print(
|
||||
f"Run {run + 1}/{args.num_runs} finished: {len(run_paths)} file(s) ({time.perf_counter() - t_run:.2f}s)"
|
||||
)
|
||||
for path in run_paths:
|
||||
print(f" {path}")
|
||||
|
||||
total_elapsed = time.perf_counter() - t_total
|
||||
finally:
|
||||
if profiler_started:
|
||||
print("Stopping profiler (sync)...")
|
||||
omni.stop_profile(stages=args.profiler_stages)
|
||||
if args.profiler_wait_seconds > 0:
|
||||
print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...")
|
||||
time.sleep(args.profiler_wait_seconds)
|
||||
|
||||
print(
|
||||
f"All sync runs finished: {args.num_runs} run(s), "
|
||||
f"{len(prompt_specs)} prompt(s), {len(saved_paths)} file(s) in {total_elapsed:.2f}s total"
|
||||
)
|
||||
return saved_paths
|
||||
|
||||
|
||||
def main(args) -> int:
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
profiled_stage_config_path: str | None = None
|
||||
original_stage_config_path = args.stage_configs_path
|
||||
if args.enable_profiler:
|
||||
Path(args.profiler_dir).mkdir(parents=True, exist_ok=True)
|
||||
profiled_stage_config_path = _build_profiled_stage_config(
|
||||
args.stage_configs_path,
|
||||
str(Path(args.profiler_dir).resolve()),
|
||||
)
|
||||
args.stage_configs_path = profiled_stage_config_path
|
||||
|
||||
is_streaming = _is_streaming_stage_config(args.stage_configs_path)
|
||||
voice_clone_count = _count_voice_clone_prompts(args.prompt_specs)
|
||||
print(f"Model: {args.model}")
|
||||
print(f"Stage config: {original_stage_config_path}")
|
||||
print(f"Route: {'streaming' if is_streaming else 'sync'} (from stage-configs-path)")
|
||||
print(f"Prompt count: {len(args.prompt_specs)}")
|
||||
print("Batch mode: sequential (aligned with native VoxCPM)")
|
||||
print(f"Warmup runs: {args.warmup_runs}")
|
||||
print(f"Voice cloning prompts: {voice_clone_count}/{len(args.prompt_specs)}")
|
||||
if args.enable_profiler:
|
||||
print(f"Profiler: enabled (dir={args.profiler_dir}, stages={args.profiler_stages or 'all-configured'})")
|
||||
print(f"Profiled stage config: {args.stage_configs_path}")
|
||||
if voice_clone_count:
|
||||
print("Voice cloning note: --ref-text/ref_text must match the spoken content of the reference audio.")
|
||||
print(f"Num runs: {args.num_runs}")
|
||||
try:
|
||||
if is_streaming:
|
||||
asyncio.run(_run_streaming(args))
|
||||
else:
|
||||
_run_sync(args)
|
||||
finally:
|
||||
if profiled_stage_config_path is not None and os.path.exists(profiled_stage_config_path):
|
||||
os.unlink(profiled_stage_config_path)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
raise SystemExit(main(parse_args()))
|
||||
@@ -1,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.21.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.21.0
|
||||
FROM ${BASE_IMAGE}
|
||||
|
||||
ARG COMMON_WORKDIR=/app
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=nightly-main
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
COPY . .
|
||||
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
ffmpeg \
|
||||
libgl1 \
|
||||
libglib2.0-0 \
|
||||
libgomp1 \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libxrender1 \
|
||||
libxcb1 \
|
||||
espeak-ng && \
|
||||
rm -rf /var/cache/apt/* && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
python3 -m pip install -U pip && \
|
||||
python3 -m pip install -v \
|
||||
"pytest>=7.0.0" \
|
||||
"pytest-asyncio>=0.21.0" \
|
||||
"pytest-cov>=4.0.0" \
|
||||
"pytest-mock>=3.10.0" \
|
||||
"openpyxl>=3.0.0" \
|
||||
"psutil>=7.2.0" \
|
||||
"soundfile>=0.13.1" \
|
||||
"pyttsx3>=2.99" \
|
||||
"opencv-python>=4.12.0.88" \
|
||||
"imageio[ffmpeg]>=0.6.0" \
|
||||
--index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
|
||||
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
|
||||
--default-timeout=60 && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \
|
||||
--index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
|
||||
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
|
||||
--default-timeout=60
|
||||
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
ENTRYPOINT []
|
||||
@@ -0,0 +1,48 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=nightly-main-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
COPY . .
|
||||
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
ffmpeg \
|
||||
libgl1 \
|
||||
libglib2.0-0 \
|
||||
libgomp1 \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libxrender1 \
|
||||
libxcb1 \
|
||||
espeak-ng && \
|
||||
rm -rf /var/cache/apt/* && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
python3 -m pip install -U pip && \
|
||||
python3 -m pip install -v \
|
||||
"pytest>=7.0.0" \
|
||||
"pytest-asyncio>=0.21.0" \
|
||||
"pytest-cov>=4.0.0" \
|
||||
"pytest-mock>=3.10.0" \
|
||||
"openpyxl>=3.0.0" \
|
||||
"psutil>=7.2.0" \
|
||||
"soundfile>=0.13.1" \
|
||||
"pyttsx3>=2.99" \
|
||||
"opencv-python>=4.12.0.88" \
|
||||
"imageio[ffmpeg]>=0.6.0" \
|
||||
--index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
|
||||
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
|
||||
--default-timeout=60 && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \
|
||||
--index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
|
||||
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
|
||||
--default-timeout=60
|
||||
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
ENTRYPOINT []
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.19.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0
|
||||
FROM ${BASE_IMAGE} AS base
|
||||
|
||||
# Declare a variable to know if we want to use the nightly build or the stable build.
|
||||
|
||||
+34
-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.21.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
|
||||
@@ -126,6 +133,11 @@ ENTRYPOINT []
|
||||
|
||||
FROM ${VLLM_BASE} AS vllm-omni
|
||||
|
||||
RUN apt-get update && \
|
||||
apt-get install -y espeak-ng ffmpeg git jq && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /workspace/vllm-omni
|
||||
COPY . .
|
||||
|
||||
@@ -135,7 +147,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 \
|
||||
|
||||
+22
-14
@@ -6,42 +6,43 @@ nav:
|
||||
- getting_started/installation/*
|
||||
- Serving:
|
||||
- OpenAI-Compatible API:
|
||||
- Diffusion Chat API: serving/diffusion_chat_api.md
|
||||
- Chat Completions API: serving/chat_completions_api.md
|
||||
- Image Generation: serving/image_generation_api.md
|
||||
- Image Edit: serving/image_edit_api.md
|
||||
- Text to Speech: serving/speech_api.md
|
||||
- Videos: serving/videos_api.md
|
||||
- Audio Generation: serving/audio_generate_api.md
|
||||
- Streaming Video Input: serving/video_stream_api.md
|
||||
- Examples:
|
||||
- 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 +60,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 +74,22 @@ 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
|
||||
- MXFP8 W8A8: user_guide/quantization/mxfp8.md
|
||||
- MXFP4 W4A4: user_guide/quantization/mxfp4.md
|
||||
- ModelOpt: user_guide/quantization/modelopt.md
|
||||
- GGUF: user_guide/quantization/gguf.md
|
||||
- AutoRound: user_guide/quantization/autoround.md
|
||||
- msModelSlim: user_guide/quantization/msmodelslim.md
|
||||
- ComfyUI: features/comfyui.md
|
||||
- Developer Guide:
|
||||
- General:
|
||||
- contributing/README.md
|
||||
- pr_reviewer.md
|
||||
- glob: contributing/*
|
||||
flatten_single_child_sections: true
|
||||
- Model Implementation:
|
||||
@@ -89,7 +97,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 +117,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
|
||||
|
||||
+3
-1
@@ -99,7 +99,6 @@ Configuration classes.
|
||||
- [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
|
||||
- [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
|
||||
- [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.config.CosyVoice3Config][]
|
||||
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechConfig][]
|
||||
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechFastARConfig][]
|
||||
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechSlowARConfig][]
|
||||
@@ -116,10 +115,13 @@ Configuration classes.
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderConfig][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.cosyvoice3.CosyVoice3Config][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][]
|
||||
|
||||
## Workers
|
||||
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 96 KiB After Width: | Height: | Size: 402 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"
|
||||
|
||||
@@ -39,11 +39,34 @@ Committers have write access and merge rights. They typically have deep expertis
|
||||
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
|
||||
- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
|
||||
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
|
||||
- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
|
||||
- [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs
|
||||
- [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache
|
||||
- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
|
||||
- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism
|
||||
- [@wtomin](https://github.com/wtomin): Diffusion models, parallel, and docs
|
||||
- [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
|
||||
- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
|
||||
- [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
|
||||
- [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
|
||||
|
||||
## Diffusion Workload Division
|
||||
|
||||
This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review.
|
||||
|
||||
### Runtime-related
|
||||
|
||||
- Diffusion models: @RuixiangMa, @wtomin
|
||||
|
||||
### Optimization
|
||||
|
||||
- Parallel: @RuixiangMa, @wtomin
|
||||
- Attention backend: @ZJY0516
|
||||
- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
|
||||
- Offload: @yuanheng-zhao
|
||||
- Quantization: @lishunyang12, @david6666666
|
||||
- Kernel fusion / communication-computation: @ZJY0516
|
||||
|
||||
### Docs & Test
|
||||
|
||||
- Diffusion docs: @RuixiangMa, @wtomin
|
||||
|
||||
## Meetings
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Loaded 100 of 1234 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user