Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6838533998 | ||
|
|
d6e338f9d6 | ||
|
|
a90a7690b8 | ||
|
|
8e12007ac1 | ||
|
|
da21e991cc | ||
|
|
a978c60163 | ||
|
|
c1c0a5b3cd | ||
|
|
072647eeb2 | ||
|
|
7007217d4b | ||
|
|
95018c536a | ||
|
|
b33d7637de | ||
|
|
ff254795fa | ||
|
|
6901ba426e | ||
|
|
da7a8f8014 | ||
|
|
199f7832e7 | ||
|
|
fafa248665 | ||
|
|
05d8de813e | ||
|
|
4c2293743b | ||
|
|
6813432f90 | ||
|
|
c2a5420d97 | ||
|
|
b4342fb487 | ||
|
|
81a90d24ef | ||
|
|
6ac7ba692d | ||
|
|
e9fe78e28c | ||
|
|
b766c4702c | ||
|
|
0239057cba | ||
|
|
99cc76690d | ||
|
|
60a8c01a96 | ||
|
|
89fff09628 | ||
|
|
ba5a66b15b | ||
|
|
f9268822cd | ||
|
|
48b3969717 | ||
|
|
b1ceaa7b34 | ||
|
|
5699fe7a96 | ||
|
|
81793f1918 | ||
|
|
31c8cb22cf | ||
|
|
fde4808723 | ||
|
|
ae92ef5f00 | ||
|
|
ca6c7ad2be | ||
|
|
e1ac1826cd | ||
|
|
c85acb1773 | ||
|
|
fc5e8f825e | ||
|
|
53ebdd0b1d | ||
|
|
4a44f6ee5c | ||
|
|
61e170c69a | ||
|
|
f52b51533a | ||
|
|
d636c32ff9 | ||
|
|
3158e9d28f | ||
|
|
6a66a1fe36 | ||
|
|
1a5d25a0ba | ||
|
|
fe786f1a19 | ||
|
|
cf328eb8e8 | ||
|
|
4036cd547c | ||
|
|
9718a9e369 | ||
|
|
88caaf182b | ||
|
|
6a5fa588c6 | ||
|
|
1deb4fdf92 | ||
|
|
9aae75c3ee | ||
|
|
284575ab82 | ||
|
|
2472c360aa | ||
|
|
af012babf8 | ||
|
|
4de0c8d41b | ||
|
|
365dfd5f8e | ||
|
|
5377619c0b | ||
|
|
dc72aede40 | ||
|
|
c107f0b336 | ||
|
|
e17130090c | ||
|
|
cebf78ba3e | ||
|
|
3e8cb6c631 | ||
|
|
92f771d455 | ||
|
|
e0b71056ba | ||
|
|
26b3ddee47 | ||
|
|
53a3df7309 | ||
|
|
75056b2db2 | ||
|
|
717dbee5a6 | ||
|
|
c8cf8a799d | ||
|
|
0540af327e | ||
|
|
54a999d031 | ||
|
|
366b3368ea | ||
|
|
36c6cb9634 | ||
|
|
9c6fbd76be | ||
|
|
8ad8988e2b | ||
|
|
b3cced310a | ||
|
|
21e621f84c | ||
|
|
c8c60e1970 | ||
|
|
28dd1a699d | ||
|
|
ff3365f6a2 | ||
|
|
fe7e8770db | ||
|
|
a52867abde | ||
|
|
7c5159037f | ||
|
|
1c7f963a2d | ||
|
|
4dbaa74f2b | ||
|
|
11d9ca8f41 | ||
|
|
3306fa6611 | ||
|
|
0e005817f7 | ||
|
|
c501c2d512 | ||
|
|
5d28b6a40c | ||
|
|
09b66bf6db | ||
|
|
ac17431745 | ||
|
|
e1efcb7560 | ||
|
|
da4a077155 | ||
|
|
acde431668 | ||
|
|
f144f5effe | ||
|
|
a6b5519616 | ||
|
|
0f7c9c755d | ||
|
|
af9784c03f | ||
|
|
74eedd3ef1 | ||
|
|
a19098233b | ||
|
|
b313279af7 | ||
|
|
4d89ebaa07 | ||
|
|
617b1d6053 | ||
|
|
ca37f221a8 | ||
|
|
2b5a3dd6b3 | ||
|
|
7543f2ff3b | ||
|
|
6ac85d4607 | ||
|
|
4725593a4b | ||
|
|
c51b4fd7b8 | ||
|
|
2ee4a07adb | ||
|
|
1d0a97f0d1 | ||
|
|
2fdc647ed7 | ||
|
|
4027534dbf | ||
|
|
e449cc411e | ||
|
|
5d62fcd717 | ||
|
|
745ef312a7 | ||
|
|
7a560492a2 | ||
|
|
761eff930a | ||
|
|
112a9aa237 | ||
|
|
e79ad8d59a | ||
|
|
82f63a4403 | ||
|
|
7271116973 | ||
|
|
5628d41242 |
No files matched your search
@@ -90,10 +90,18 @@ upload_pipeline() {
|
||||
FAIL_FAST=$(fail_fast)
|
||||
|
||||
cd .buildkite
|
||||
|
||||
# Select test definition file: merge suite for main, ready suite for PRs
|
||||
if [[ $BUILDKITE_BRANCH == "main" ]]; then
|
||||
TEST_YAML="test-amd-merge.yml"
|
||||
else
|
||||
TEST_YAML="test-amd-ready.yaml"
|
||||
fi
|
||||
|
||||
(
|
||||
set -x
|
||||
# Output pipeline.yaml with all blank lines removed
|
||||
minijinja-cli test-template.j2 test-amd.yaml \
|
||||
minijinja-cli test-template.j2 "$TEST_YAML" \
|
||||
-D branch="$BUILDKITE_BRANCH" \
|
||||
-D list_file_diff="$LIST_FILE_DIFF" \
|
||||
-D run_all="$RUN_ALL" \
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
#!/bin/bash
|
||||
# vllm-omni Intel bootstrap
|
||||
# Uses static pipeline-intel.yaml for Intel XPU tests
|
||||
# Based on: bootstrap-amd-omni.sh
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
if [[ -z "${DOCS_ONLY_DISABLE:-}" ]]; then
|
||||
DOCS_ONLY_DISABLE=0
|
||||
fi
|
||||
|
||||
upload_pipeline() {
|
||||
echo "--- 🛠 Preparing Intel pipeline"
|
||||
buildkite-agent pipeline upload .buildkite/pipeline-intel.yaml
|
||||
}
|
||||
|
||||
get_diff() {
|
||||
$(git add .)
|
||||
echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
|
||||
}
|
||||
|
||||
get_diff_main() {
|
||||
$(git add .)
|
||||
echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
|
||||
}
|
||||
|
||||
file_diff=$(get_diff)
|
||||
if [[ $BUILDKITE_BRANCH == "main" ]]; then
|
||||
file_diff=$(get_diff_main)
|
||||
fi
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Early exit: skip pipeline if all changed files are under docs/
|
||||
# ----------------------------------------------------------------------
|
||||
if [[ "${DOCS_ONLY_DISABLE}" != "1" ]] && [[ -n "${file_diff:-}" ]]; then
|
||||
docs_only=1
|
||||
while IFS= read -r f; do
|
||||
[[ -z "$f" ]] && continue
|
||||
if [[ "$f" != docs/* ]]; then
|
||||
docs_only=0
|
||||
break
|
||||
fi
|
||||
done < <(printf '%s\n' "$file_diff" | tr ' ' '\n' | tr -d '\r')
|
||||
|
||||
if [[ "$docs_only" -eq 1 ]]; then
|
||||
buildkite-agent annotate ":memo: CI skipped — docs only" --style "info"
|
||||
exit 0
|
||||
fi
|
||||
fi
|
||||
|
||||
upload_pipeline
|
||||
@@ -0,0 +1,16 @@
|
||||
# .buildkite/pipeline-intel.yaml
|
||||
steps:
|
||||
- group: "INTEL Tests"
|
||||
steps:
|
||||
- label: "Intel: XPU Test"
|
||||
command: "bash .buildkite/scripts/hardware_ci/run-xpu-test.sh"
|
||||
agents:
|
||||
queue: "intel-gpu-omni"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
# Buildkite will automatically replace this with the actual commit hash
|
||||
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
|
||||
VLLM_VERSION: "v0.18.0"
|
||||
priority: 100
|
||||
timeout_in_minutes: 60
|
||||
soft_fail: true
|
||||
@@ -87,7 +87,11 @@ HF_CACHE="$(realpath ~)/huggingface"
|
||||
mkdir -p "${HF_CACHE}"
|
||||
HF_MOUNT="/root/.cache/huggingface"
|
||||
|
||||
commands=$@
|
||||
if [[ -n "${TEST_COMMAND:-}" ]]; then
|
||||
commands="$TEST_COMMAND"
|
||||
else
|
||||
commands="$@"
|
||||
fi
|
||||
echo "Commands:$commands"
|
||||
|
||||
PARALLEL_JOB_COUNT=8
|
||||
@@ -102,6 +106,7 @@ if [[ -z "$render_gid" ]]; then
|
||||
fi
|
||||
|
||||
# check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs.
|
||||
# TODO: @tjtanaa reenable to run VLLM_ROCM_USE_AITER=1 when AITER is shipped with prebuilt kernels.
|
||||
if [[ $commands == *"--shard-id="* ]]; then
|
||||
# assign job count as the number of shards used
|
||||
commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g')
|
||||
@@ -118,7 +123,7 @@ if [[ $commands == *"--shard-id="* ]]; then
|
||||
--rm \
|
||||
-e MIOPEN_DEBUG_CONV_DIRECT=0 \
|
||||
-e MIOPEN_DEBUG_CONV_GEMM=0 \
|
||||
-e VLLM_ROCM_USE_AITER=1 \
|
||||
-e VLLM_ROCM_USE_AITER=0 \
|
||||
-e HIP_VISIBLE_DEVICES="${GPU}" \
|
||||
-e HF_TOKEN \
|
||||
-e AWS_ACCESS_KEY_ID \
|
||||
@@ -153,7 +158,7 @@ else
|
||||
--rm \
|
||||
-e MIOPEN_DEBUG_CONV_DIRECT=0 \
|
||||
-e MIOPEN_DEBUG_CONV_GEMM=0 \
|
||||
-e VLLM_ROCM_USE_AITER=1 \
|
||||
-e VLLM_ROCM_USE_AITER=0 \
|
||||
-e HF_TOKEN \
|
||||
-e AWS_ACCESS_KEY_ID \
|
||||
-e AWS_SECRET_ACCESS_KEY \
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
#!/bin/bash
|
||||
|
||||
# This script build the XPU docker image and run the offline inference inside the container.
|
||||
set -ex
|
||||
|
||||
omni_source_dir=$(git rev-parse --show-toplevel)
|
||||
|
||||
base_image_name="xpu/vllm-omni-ci-base:${VLLM_VERSION:?VLLM_VERSION must be set}"
|
||||
image_name="xpu/vllm-omni-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set}"
|
||||
container_name="xpu_${BUILDKITE_COMMIT}_$(
|
||||
tr -dc A-Za-z0-9 </dev/urandom | head -c 10
|
||||
echo
|
||||
)"
|
||||
|
||||
cd "${omni_source_dir}"
|
||||
if [ -z "$(docker images -q "${base_image_name}")" ]; then
|
||||
docker build --target vllm-base -t "${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -f docker/Dockerfile.xpu .
|
||||
fi
|
||||
|
||||
# Try building the docker image
|
||||
docker build --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -t "${image_name}" -f docker/Dockerfile.xpu .
|
||||
|
||||
# Setup cleanup
|
||||
remove_docker_container() {
|
||||
docker rm -f "${container_name}" || true
|
||||
docker image rm -f "${image_name}" || true
|
||||
docker system prune -f || true
|
||||
}
|
||||
trap remove_docker_container EXIT
|
||||
|
||||
HF_CACHE="${HF_CACHE:-$(realpath ~)/.cache/huggingface}"
|
||||
mkdir -p "${HF_CACHE}"
|
||||
HF_MOUNT="/root/.cache/huggingface"
|
||||
|
||||
time timeout -k 30 30m docker run \
|
||||
--device /dev/dri:/dev/dri \
|
||||
--net=host \
|
||||
--ipc=host \
|
||||
-v /dev/dri/by-path:/dev/dri/by-path \
|
||||
-v "${HF_CACHE}:${HF_MOUNT}" \
|
||||
--security-opt seccomp=unconfined \
|
||||
--entrypoint="" \
|
||||
-e VLLM_LOGGING_LEVEL \
|
||||
-e VLLM_OMNI_LOGGING_LEVEL \
|
||||
-e HF_TOKEN \
|
||||
-e ZE_AFFINITY_MASK \
|
||||
--name "${container_name}" \
|
||||
"${image_name}" \
|
||||
bash -c '
|
||||
set -e
|
||||
echo $ZE_AFFINITY_MASK
|
||||
pip install tblib==3.1.0
|
||||
cd /workspace/vllm-omni
|
||||
pytest -v -s -m "core_model and xpu and B60"
|
||||
pytest -v -s -m "advanced_model and xpu and B60"
|
||||
'
|
||||
@@ -0,0 +1,197 @@
|
||||
steps:
|
||||
|
||||
- label: "Simple Unit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
|
||||
|
||||
- label: "Diffusion Images API LoRA E2E"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
|
||||
# - label: "Audio Generation Model Test"
|
||||
# agent_pool: mi325_1
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
# commands:
|
||||
# - export GPU_ARCHS=gfx942
|
||||
# - export VLLM_LOGGING_LEVEL=DEBUG
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
|
||||
|
||||
- label: "Diffusion Cache Backend Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
|
||||
|
||||
# merge-only tests
|
||||
- label: "Diffusion Tensor Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Benchmark & Engine Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/benchmarks/test_serve_cli.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
- label: "Omni Model Test Qwen2-5-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
|
||||
- label: "Qwen3-TTS E2E Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
|
||||
|
||||
- label: "Diffusion Image Edit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
|
||||
|
||||
# split Bagel Model Test with H100 (Real Weights) into three tests
|
||||
- label: "Bagel Text2Img Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
|
||||
|
||||
- label: "Bagel Img2Img Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
|
||||
|
||||
- label: "Bagel Online Serving Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
|
||||
@@ -0,0 +1,178 @@
|
||||
steps:
|
||||
|
||||
- label: "Simple Unit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
|
||||
# - label: "Audio Generation Model Test"
|
||||
# agent_pool: mi325_1
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
# commands:
|
||||
# - export GPU_ARCHS=gfx942
|
||||
# - export VLLM_LOGGING_LEVEL=DEBUG
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
|
||||
|
||||
- label: "Diffusion Cache Backend Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Benchmark & Engine Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/benchmarks/test_serve_cli.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
- label: "Omni Model Test Qwen2-5-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 17m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- timeout 10m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
|
||||
|
||||
- label: "Qwen3-TTS E2E Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
|
||||
|
||||
- label: "Diffusion Image Edit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
|
||||
|
||||
- label: "Bagel Text2Img Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
|
||||
- label: "Bagel Img2Img Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
|
||||
- label: "Bagel Online Serving Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
@@ -1,7 +1,7 @@
|
||||
steps:
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
timeout_in_minutes: 20
|
||||
timeout_in_minutes: 30
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
@@ -11,7 +11,7 @@ steps:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
|
||||
|
||||
- label: "Diffusion Images API LoRA E2E"
|
||||
timeout_in_minutes: 20
|
||||
timeout_in_minutes: 30
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
@@ -101,6 +101,7 @@ steps:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
|
||||
|
||||
|
||||
@@ -12,11 +12,12 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
timeout_in_minutes: 20
|
||||
timeout_in_minutes: 30
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
|
||||
@@ -29,11 +30,12 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Images API LoRA E2E"
|
||||
timeout_in_minutes: 20
|
||||
timeout_in_minutes: 30
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
|
||||
@@ -46,6 +48,7 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -64,6 +67,7 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -81,6 +85,7 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -98,6 +103,7 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -116,6 +122,7 @@ steps:
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -134,6 +141,7 @@ steps:
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -152,21 +160,17 @@ steps:
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Benchmark & Engine Test"
|
||||
- label: "Engine Test"
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 15m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
set +e
|
||||
pytest -s -v tests/benchmarks/test_serve_cli.py
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
EXIT2=$$?
|
||||
exit $$((EXIT1 | EXIT2))
|
||||
'
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
@@ -178,6 +182,7 @@ steps:
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -197,6 +202,7 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -226,6 +232,11 @@ steps:
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
@@ -261,6 +272,11 @@ steps:
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
@@ -272,12 +288,64 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# - label: "Bagel Text2Img Model Test with H100"
|
||||
# timeout_in_minutes: 30
|
||||
- label: "Bagel Model Test with H100 (Real Weights)"
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 55m bash -c '
|
||||
set -e
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory"
|
||||
pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model"
|
||||
pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# - label: "Voxtral TTS E2E Test"
|
||||
# timeout_in_minutes: 20
|
||||
# depends_on: upload-merge-pipeline
|
||||
# commands:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
|
||||
# - |
|
||||
# timeout 20m bash -c '
|
||||
# export VLLM_LOGGING_LEVEL=DEBUG
|
||||
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
# '
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
|
||||
@@ -11,6 +11,9 @@ steps:
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
EXIT2=$$?
|
||||
# TODO: uncomment when mistralai/tts-model is public
|
||||
# pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
# EXIT3=$$?
|
||||
exit $$((EXIT1 | EXIT2))
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
@@ -30,6 +33,11 @@ steps:
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
@@ -54,14 +62,55 @@ steps:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
shm-size: "8gb"
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: ":full_moon: Diffusion Model Test with H100"
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
|
||||
timeout_in_minutes: 30
|
||||
depends_on: upload-nightly-pipeline
|
||||
@@ -81,10 +130,11 @@ steps:
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: ":full_moon: Omni Model Perf Test with H100"
|
||||
- label: ":full_moon: Omni Model Perf Test & Test Case Statistics"
|
||||
key: nightly-performance
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-nightly-pipeline
|
||||
@@ -94,6 +144,8 @@ steps:
|
||||
- export BENCHMARK_DIR=tests
|
||||
- pytest -s -v tests/perf/scripts/run_benchmark.py
|
||||
- buildkite-agent artifact upload "tests/*.json"
|
||||
- python tools/nightly/buildkite_testcase_statistics.py -o tests/buildkite_testcase_statistics.html
|
||||
- buildkite-agent artifact upload "tests/*.html"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -112,6 +164,13 @@ steps:
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
|
||||
value: "120"
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
@@ -123,17 +182,19 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":email: Nightly Perf Collection & Email"
|
||||
- label: ":email: Nightly Collection & Email"
|
||||
key: nightly-perf-distribution
|
||||
depends_on: nightly-performance
|
||||
depends_on:
|
||||
- nightly-performance
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- pip install openpyxl
|
||||
- export DEFAULT_INPUT_DIR=tests
|
||||
- export DEFAULT_OUTPUT_DIR=tests
|
||||
- buildkite-agent artifact download "tests/*.json" . --step nightly-performance
|
||||
- buildkite-agent artifact download "tests/*.html" . --step nightly-performance
|
||||
- python tools/nightly/generate_nightly_perf_excel.py
|
||||
- python tools/nightly/send_nightly_perf_email.py
|
||||
- python tools/nightly/send_nightly_email.py --report-file "tests/*.xlsx, tests/*.html"
|
||||
- buildkite-agent artifact upload "tests/*.xlsx"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
@@ -5,6 +5,24 @@ steps:
|
||||
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Voxtral TTS CUDA Unit Test"
|
||||
timeout_in_minutes: 10
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- "timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
@@ -18,7 +36,7 @@ steps:
|
||||
- label: "Diffusion Model Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
@@ -28,21 +46,14 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
EXIT2=$$?
|
||||
exit $$((EXIT1 | EXIT2))
|
||||
'
|
||||
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
@@ -52,6 +63,7 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -68,6 +80,7 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -84,6 +97,7 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -101,6 +115,7 @@ steps:
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -118,22 +133,18 @@ steps:
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
|
||||
- label: "Benchmark & Engine Test"
|
||||
- label: "Engine Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 15m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
set +e
|
||||
pytest -s -v tests/benchmarks/test_serve_cli.py
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
EXIT2=$$?
|
||||
exit $$((EXIT1 | EXIT2))
|
||||
'
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
@@ -145,6 +156,7 @@ steps:
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -167,6 +179,7 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -229,9 +242,49 @@ steps:
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
# TODO: uncomment when mistralai/tts-model is public
|
||||
# - label: "Voxtral-TTS E2E Test"
|
||||
# depends_on: upload-ready-pipeline
|
||||
# commands:
|
||||
# - |
|
||||
# timeout 20m bash -c '
|
||||
# export VLLM_LOGGING_LEVEL=DEBUG
|
||||
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
# '
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 1
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
# - label: "Diffusion Image Edit Test with H100 (1 GPU)"
|
||||
# depends_on: upload-ready-pipeline
|
||||
# commands:
|
||||
@@ -269,39 +322,132 @@ steps:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
# - label: "Bagel Text2Img Model Test with H100"
|
||||
# depends_on: upload-ready-pipeline
|
||||
# commands:
|
||||
# - |
|
||||
# timeout 30m bash -c '
|
||||
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
|
||||
# '
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 1
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
- label: "Bagel Text2Img Model Test with H100"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model"
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Bagel Img2Img Model Test with H100"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model"
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Bagel Online Serving Test with H100"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model"
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
@@ -40,9 +40,14 @@
|
||||
{% else %}
|
||||
queue: amd_mi325_1
|
||||
{% endif %}
|
||||
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh "(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }} ; {{ step.command or (step.commands | join(" && ")) | safe }}"
|
||||
{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
|
||||
{% set indented_cmd = cmd_body | replace("\n", "\n ") %}
|
||||
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
TEST_COMMAND: |-
|
||||
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
|
||||
{{ indented_cmd | safe }}
|
||||
priority: 100
|
||||
{% if step.grade and step.grade == "Blocking" %}
|
||||
soft_fail: false
|
||||
|
||||
@@ -74,28 +74,21 @@ body:
|
||||
If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. For example:
|
||||
|
||||
```python
|
||||
from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm import SamplingParams
|
||||
|
||||
# Create stage configurations
|
||||
ar_config = create_ar_stage_config(
|
||||
stage_id=0,
|
||||
model_path="Qwen/Qwen3-0.6B",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"]
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
stage_configs_path="/path/to/stage_configs.yaml",
|
||||
)
|
||||
|
||||
dit_config = create_dit_stage_config(
|
||||
stage_id=1,
|
||||
model_path="stabilityai/stable-diffusion-2-1",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["image"]
|
||||
)
|
||||
|
||||
# Initialize OmniLLM
|
||||
omni = OmniLLM([ar_config, dit_config])
|
||||
|
||||
# Generate
|
||||
outputs = omni.generate(prompt="A scenic watercolor painting of a lighthouse at sunset")
|
||||
prompts = [{"prompt": "A scenic watercolor painting of a lighthouse at sunset"}]
|
||||
sampling_params_list = [
|
||||
SamplingParams(max_tokens=1),
|
||||
OmniDiffusionSamplingParams(num_outputs_per_prompt=1),
|
||||
]
|
||||
outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
|
||||
```
|
||||
|
||||
If the code is too long (hopefully, it isn't), feel free to put it in a public gist and link it in the issue: https://gist.github.com.
|
||||
|
||||
@@ -8,6 +8,7 @@ repos:
|
||||
# list of supported hooks: https://pre-commit.com/hooks.html
|
||||
- id: check-yaml
|
||||
args: ["--unsafe"]
|
||||
exclude: "examples/online_serving/chart-helm/templates/"
|
||||
- id: debug-statements
|
||||
- id: end-of-file-fixer
|
||||
- id: mixed-line-ending
|
||||
|
||||
@@ -9,7 +9,7 @@ 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://drive.google.com/file/d/12k_DOo2xoEROvd4xQxYNuEi6C2GECoRx/view?usp=drive_link"><b>Slides</b></a> |
|
||||
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
|
||||
</p>
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
||||
|
||||
*Latest News* 🔥
|
||||
|
||||
- [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.
|
||||
|
||||
@@ -47,6 +47,7 @@ If no, check your shell running the ComfyUI process. There may be some error mes
|
||||
This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
|
||||
|
||||
- **Generate Image** for text-to-image and image-to-image tasks
|
||||
- **Generate Video** for text-to-video and image-to-video tasks
|
||||
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
|
||||
- **TTS** and **TTS Voice Clone** for TTS tasks
|
||||
|
||||
@@ -59,7 +60,7 @@ To build a simple workflow yourself,
|
||||
|
||||
- Drag a generation node onto the canvas.
|
||||
- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
|
||||
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
|
||||
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
|
||||
- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
|
||||
- For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
|
||||
- For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
|
||||
@@ -104,6 +105,20 @@ You can configure per-stage sampling parameters for multi-stage models.
|
||||
> [!TIP]
|
||||
> The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
|
||||
|
||||
### Text-to-video and image-to-video generation (e.g., Wan)
|
||||
|
||||
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video Generation**)
|
||||
|
||||
<p align="center">
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg">
|
||||
<img alt="vLLM-Omni video generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg" width=55%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
> [!TIP]
|
||||
> The node automatically choose text-to-video or image-to-video API endpoints depending on whether you connect an image input or not.
|
||||
|
||||
### TTS (e.g., Qwen TTS series)
|
||||
|
||||
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni TTS**)
|
||||
|
||||
@@ -14,17 +14,21 @@ from .comfyui_vllm_omni.nodes import (
|
||||
VLLMOmniARSampling,
|
||||
VLLMOmniDiffusionSampling,
|
||||
VLLMOmniGenerateImage,
|
||||
VLLMOmniGenerateVideo,
|
||||
VLLMOmniQwenTTSParams,
|
||||
VLLMOmniRemoteLoRA,
|
||||
VLLMOmniSamplingParamsList,
|
||||
VLLMOmniTTS,
|
||||
VLLMOmniUnderstanding,
|
||||
VLLMOmniVoiceClone,
|
||||
VLLMOmniWanParams,
|
||||
)
|
||||
|
||||
# A dictionary that contains all nodes you want to export with their names
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
# === Generation ===
|
||||
"VLLMOmniGenerateImage": VLLMOmniGenerateImage,
|
||||
"VLLMOmniGenerateVideo": VLLMOmniGenerateVideo,
|
||||
"VLLMOmniUnderstanding": VLLMOmniUnderstanding,
|
||||
"VLLMOmniTTS": VLLMOmniTTS,
|
||||
"VLLMOmniVoiceClone": VLLMOmniVoiceClone,
|
||||
@@ -32,13 +36,16 @@ NODE_CLASS_MAPPINGS = {
|
||||
"VLLMOmniARSampling": VLLMOmniARSampling,
|
||||
"VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
|
||||
"VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList,
|
||||
"VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA,
|
||||
"VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
|
||||
"VLLMOmniWanParams": VLLMOmniWanParams,
|
||||
}
|
||||
|
||||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
# === Generation ===
|
||||
"VLLMOmniGenerateImage": "Generate Image",
|
||||
"VLLMOmniGenerateVideo": "Generate Video",
|
||||
"VLLMOmniUnderstanding": "Multimodality Understanding",
|
||||
"VLLMOmniTTS": "TTS (Text to Speech)",
|
||||
"VLLMOmniVoiceClone": "TTS Voice Cloning",
|
||||
@@ -46,7 +53,9 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"VLLMOmniARSampling": "AR Sampling Params",
|
||||
"VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
|
||||
"VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List",
|
||||
"VLLMOmniRemoteLoRA": "LoRA",
|
||||
"VLLMOmniQwenTTSParams": "Qwen TTS Params",
|
||||
"VLLMOmniWanParams": "Wan Video Params",
|
||||
}
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
@@ -11,6 +11,7 @@ from .utils.types import (
|
||||
AutoregressionSamplingParams,
|
||||
DiffusionSamplingParams,
|
||||
QwenTTSModelSpecificParams,
|
||||
WanModelSpecificParams,
|
||||
)
|
||||
from .utils.validators import (
|
||||
add_sampling_parameters_to_stage,
|
||||
@@ -56,6 +57,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
# "video": ("VIDEO",),
|
||||
# "audio": ("AUDIO",),
|
||||
"sampling_params": ("SAMPLING_PARAMS",),
|
||||
"lora": ("REMOTE_LORA",),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -76,9 +78,11 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
audio: AudioInput | None = None, # Hidden & unused
|
||||
video: VideoInput | None = None, # Hidden & unused
|
||||
sampling_params: dict | list[dict] | None = None,
|
||||
lora: dict | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
if kwargs:
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
logger.debug("Got sampling params: %s", sampling_params)
|
||||
validate_model_and_sampling_params_types(model, sampling_params)
|
||||
if image is None and mask is not None:
|
||||
@@ -105,6 +109,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
height=height,
|
||||
negative_prompt=negative_prompt,
|
||||
sampling_params=sampling_params,
|
||||
lora=lora,
|
||||
)
|
||||
return (output,)
|
||||
elif image is not None and audio is None and video is None:
|
||||
@@ -119,6 +124,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
negative_prompt=negative_prompt,
|
||||
mask=mask,
|
||||
sampling_params=sampling_params,
|
||||
lora=lora,
|
||||
)
|
||||
return (output,)
|
||||
|
||||
@@ -136,11 +142,91 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
audio=audio,
|
||||
video=video,
|
||||
sampling_params=sampling_params,
|
||||
lora=lora,
|
||||
)
|
||||
|
||||
return (output,)
|
||||
|
||||
|
||||
class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"url": ("STRING", {"default": "http://localhost:8000/v1"}),
|
||||
"model": ("STRING", {"default": "Wan-AI/Wan2.2-T2V-A14B-Diffusers"}),
|
||||
"prompt": ("STRING", {"multiline": True}),
|
||||
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
|
||||
"width": ("INT", {"default": 832, "min": 1}),
|
||||
"height": ("INT", {"default": 480, "min": 1}),
|
||||
"fps": ("INT", {"default": 16, "min": 1}),
|
||||
"num_frames": ("INT", {"default": 41, "min": 1}),
|
||||
},
|
||||
"optional": {
|
||||
"image": ("IMAGE",),
|
||||
"sampling_params": ("SAMPLING_PARAMS",),
|
||||
"lora": ("REMOTE_LORA",),
|
||||
"model_params": ("VIDEO_PARAMS",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO",)
|
||||
RETURN_NAMES = ("video",)
|
||||
FUNCTION = "generate"
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
url: str,
|
||||
model: str,
|
||||
prompt: str,
|
||||
width: int,
|
||||
height: int,
|
||||
fps: int,
|
||||
num_frames: int,
|
||||
negative_prompt: str | None = None,
|
||||
image: torch.Tensor | None = None,
|
||||
sampling_params: dict | list[dict] | None = None,
|
||||
model_params: dict | None = None,
|
||||
lora: dict | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
if kwargs:
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
logger.debug("Got sampling params: %s", sampling_params)
|
||||
logger.debug("Got model params: %s", model_params)
|
||||
validate_model_and_sampling_params_types(model, sampling_params)
|
||||
|
||||
# Currently, all video generation models are single-stage diffusion models
|
||||
if isinstance(sampling_params, list):
|
||||
if len(sampling_params) != 1:
|
||||
raise ValueError(
|
||||
"Video generation expects a single sampling params group. "
|
||||
"Please provide one Diffusion sampling node."
|
||||
)
|
||||
sampling_params = sampling_params[0]
|
||||
|
||||
if sampling_params is not None:
|
||||
sampling_params.pop("type", None) # internal fields
|
||||
if model_params is not None:
|
||||
model_params.pop("type", None) # internal fields
|
||||
|
||||
client = VLLMOmniClient(url)
|
||||
output = await client.generate_video(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
image=image, # image present => i2v, absent => t2v
|
||||
width=width,
|
||||
height=height,
|
||||
num_frames=num_frames,
|
||||
fps=fps,
|
||||
negative_prompt=negative_prompt,
|
||||
sampling_params=sampling_params,
|
||||
lora=lora,
|
||||
model_params=model_params,
|
||||
)
|
||||
return (output,)
|
||||
|
||||
|
||||
class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -188,7 +274,8 @@ class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
|
||||
use_audio_in_video: bool = True,
|
||||
**kwargs,
|
||||
) -> tuple[str, AudioInput]:
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
if kwargs:
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
logger.debug("Got sampling params: %s", sampling_params)
|
||||
validate_model_and_sampling_params_types(model, sampling_params)
|
||||
|
||||
@@ -547,6 +634,52 @@ class VLLMOmniSamplingParamsList:
|
||||
return (params,)
|
||||
|
||||
|
||||
class VLLMOmniRemoteLoRA:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"local_path": ("STRING", {"default": ""}),
|
||||
"name": ("STRING", {"default": ""}),
|
||||
"scale": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1},
|
||||
),
|
||||
"int_id": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"step": 1,
|
||||
"tooltip": "0 means it is not set and the server can derive it.",
|
||||
},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("REMOTE_LORA",)
|
||||
RETURN_NAMES = ("lora",)
|
||||
FUNCTION = "get_lora"
|
||||
CATEGORY = "vLLM-Omni"
|
||||
|
||||
@classmethod
|
||||
def VALIDATE_INPUTS(cls, local_path, name) -> str | Literal[True]:
|
||||
if not local_path.strip() or not name.strip():
|
||||
return "Both local_path and name must be provided."
|
||||
return True
|
||||
|
||||
def get_lora(self, local_path: str, name: str, scale: float, int_id: int):
|
||||
local_path = local_path.strip()
|
||||
name = name.strip()
|
||||
lora = {
|
||||
"local_path": local_path or None,
|
||||
"name": name or None,
|
||||
"scale": float(scale),
|
||||
"int_id": int(int_id) if int_id > 0 else None,
|
||||
}
|
||||
return (lora,)
|
||||
|
||||
|
||||
class VLLMOmniQwenTTSParams:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -572,3 +705,32 @@ class VLLMOmniQwenTTSParams:
|
||||
|
||||
def get_params(self, **kwargs):
|
||||
return (QwenTTSModelSpecificParams(kwargs),)
|
||||
|
||||
|
||||
class VLLMOmniWanParams:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"guidance_scale_2": (
|
||||
"FLOAT",
|
||||
{"default": 4.0, "min": 0.0, "max": 20.0, "step": 0.1},
|
||||
),
|
||||
"boundary_ratio": (
|
||||
"FLOAT",
|
||||
{"default": 0.875, "min": 0.0, "max": 1.0, "step": 0.001},
|
||||
),
|
||||
"flow_shift": (
|
||||
"FLOAT",
|
||||
{"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO_PARAMS",)
|
||||
RETURN_NAMES = ("Wan video params",)
|
||||
FUNCTION = "get_params"
|
||||
CATEGORY = "vLLM-Omni/Video Params"
|
||||
|
||||
def get_params(self, **kwargs):
|
||||
return (WanModelSpecificParams(kwargs),)
|
||||
@@ -6,6 +6,8 @@ The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@do
|
||||
Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
@@ -18,6 +20,7 @@ from .format import (
|
||||
base64_to_audio,
|
||||
base64_to_image_tensor,
|
||||
bytes_to_audio,
|
||||
bytes_to_video,
|
||||
image_tensor_to_base64,
|
||||
image_tensor_to_png_bytes,
|
||||
video_to_base64,
|
||||
@@ -29,10 +32,43 @@ from .types import AudioFormat
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
async def url_json(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> dict[str, Any]:
|
||||
try:
|
||||
async with getattr(session, verb)(url, **kwargs) as response:
|
||||
if not response.ok:
|
||||
error_text = await response.text()
|
||||
raise (ValueError if response.status < 500 else RuntimeError)(
|
||||
f"vLLM-Omni API returned status {response.status}: {error_text}"
|
||||
)
|
||||
try:
|
||||
return await response.json()
|
||||
except aiohttp.ContentTypeError as e:
|
||||
raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
|
||||
except aiohttp.ClientError as e:
|
||||
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
|
||||
|
||||
|
||||
async def url_bytes(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> bytes:
|
||||
try:
|
||||
async with getattr(session, verb)(url, **kwargs) as response:
|
||||
if not response.ok:
|
||||
error_text = await response.text()
|
||||
raise (ValueError if response.status < 500 else RuntimeError)(
|
||||
f"vLLM-Omni API returned status {response.status}: {error_text}"
|
||||
)
|
||||
return await response.read()
|
||||
except aiohttp.ClientError as e:
|
||||
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
|
||||
|
||||
|
||||
class VLLMOmniClient:
|
||||
def __init__(self, base_url: str, timeout: float = 300.0):
|
||||
def __init__(
|
||||
self, base_url: str, timeout: float | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
|
||||
):
|
||||
self.base_url = base_url
|
||||
self.timeout = aiohttp.ClientTimeout(total=timeout)
|
||||
self.poll_interval = poll_interval
|
||||
self.max_poll_duration = max_poll_duration
|
||||
|
||||
async def generate_image(
|
||||
self,
|
||||
@@ -43,12 +79,13 @@ class VLLMOmniClient:
|
||||
height: int,
|
||||
negative_prompt: str | None = None,
|
||||
sampling_params: dict | None = None,
|
||||
lora: dict | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run text-to-image generation via DALLE API"""
|
||||
await self._check_model_exist(model)
|
||||
|
||||
size = f"{width}x{height}"
|
||||
payload = {
|
||||
payload: dict[str, Any] = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"size": size,
|
||||
@@ -58,6 +95,8 @@ class VLLMOmniClient:
|
||||
payload["negative_prompt"] = negative_prompt
|
||||
if sampling_params is not None:
|
||||
payload.update(sampling_params)
|
||||
if lora is not None:
|
||||
payload["lora"] = lora
|
||||
logger.debug("img gen payload: %s", payload)
|
||||
|
||||
url = self.base_url + "/images/generations"
|
||||
@@ -110,6 +149,7 @@ class VLLMOmniClient:
|
||||
negative_prompt: str | None = None,
|
||||
mask: torch.Tensor | None = None,
|
||||
sampling_params: dict | None = None,
|
||||
lora: dict | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run image editing via DALLE API"""
|
||||
await self._check_model_exist(model)
|
||||
@@ -131,6 +171,8 @@ class VLLMOmniClient:
|
||||
if sampling_params is not None:
|
||||
for k, v in sampling_params.items():
|
||||
form.add_field(k, str(v))
|
||||
if lora is not None:
|
||||
form.add_field("lora", json.dumps(lora, ensure_ascii=False))
|
||||
if mask is not None:
|
||||
mask_filename = "mask.png"
|
||||
form.add_field(
|
||||
@@ -183,6 +225,7 @@ class VLLMOmniClient:
|
||||
audio: AudioInput | None = None,
|
||||
video: VideoInput | None = None,
|
||||
sampling_params: dict | list[dict] | None = None,
|
||||
lora: dict | None = None,
|
||||
) -> torch.Tensor:
|
||||
payload = VLLMOmniClient._prepare_chat_completion_messages(
|
||||
model=model,
|
||||
@@ -193,6 +236,8 @@ class VLLMOmniClient:
|
||||
video=video,
|
||||
sampling_params=sampling_params,
|
||||
modalities=["image"],
|
||||
# === below are additional `extra_body` fields, handled by **kwargs ===
|
||||
lora=lora,
|
||||
)
|
||||
choices = await self._generate_base_chat_completion(model, payload)
|
||||
|
||||
@@ -206,6 +251,87 @@ class VLLMOmniClient:
|
||||
|
||||
return torch.stack(image_tensors, dim=0)
|
||||
|
||||
async def generate_video(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
prompt: str,
|
||||
width: int,
|
||||
height: int,
|
||||
num_frames: int,
|
||||
fps: int,
|
||||
negative_prompt: str | None = None,
|
||||
image: torch.Tensor | None = None,
|
||||
sampling_params: dict | None = None,
|
||||
model_params: dict | None = None,
|
||||
lora: dict | None = None,
|
||||
**extra_body,
|
||||
) -> VideoInput:
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("model", model)
|
||||
form.add_field("prompt", prompt)
|
||||
form.add_field("width", str(width))
|
||||
form.add_field("height", str(height))
|
||||
form.add_field("num_frames", str(num_frames))
|
||||
form.add_field("fps", str(fps))
|
||||
if negative_prompt:
|
||||
form.add_field("negative_prompt", negative_prompt)
|
||||
if sampling_params is not None:
|
||||
for k, v in sampling_params.items():
|
||||
form.add_field(k, str(v))
|
||||
if model_params is not None:
|
||||
for k, v in model_params.items():
|
||||
form.add_field(k, str(v))
|
||||
if lora is not None:
|
||||
form.add_field("lora", json.dumps(lora, ensure_ascii=False))
|
||||
if extra_body:
|
||||
form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False))
|
||||
|
||||
if image is not None:
|
||||
image_filename = "image.png" # Required for multipart form
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
image_tensor_to_png_bytes(image, image_filename),
|
||||
filename=image_filename,
|
||||
content_type="image/png",
|
||||
)
|
||||
|
||||
async with aiohttp.ClientSession(timeout=self.timeout) as session:
|
||||
# Start the video generation job
|
||||
url = f"{self.base_url}/videos"
|
||||
data = await url_json(session, url, "post", data=form)
|
||||
if (job_id := data.get("id", None)) is None:
|
||||
raise RuntimeError("API response missing job 'id' field - expected OpenAI compliant format")
|
||||
if (job_status := data.get("status", None)) is None:
|
||||
raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
|
||||
|
||||
# Poll for video generation job completion
|
||||
deadline = asyncio.get_running_loop().time() + self.max_poll_duration
|
||||
url = f"{self.base_url}/videos/{job_id}"
|
||||
while job_status not in {"completed", "failed"}:
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
|
||||
data = await url_json(session, url)
|
||||
if (job_status := data.get("status", None)) is None:
|
||||
raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
|
||||
if asyncio.get_running_loop().time() >= deadline:
|
||||
raise RuntimeError(f"Timed out waiting for video job {job_id} to complete")
|
||||
|
||||
if job_status == "failed":
|
||||
raise RuntimeError(f"Video job failed: {data}")
|
||||
|
||||
# Retrieve completed content
|
||||
video_bytes = await url_bytes(session, f"{url}/content")
|
||||
|
||||
# Decode video and make a best effort at cleaning up server resources
|
||||
try:
|
||||
return bytes_to_video(video_bytes)
|
||||
finally:
|
||||
try:
|
||||
await url_json(session, url, "delete")
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to clean up video job %s: %s", job_id, exc)
|
||||
|
||||
async def generate_understanding_chat_completion(
|
||||
self,
|
||||
*,
|
||||
|
||||
@@ -6,12 +6,17 @@ Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed und
|
||||
|
||||
import base64
|
||||
import mimetypes
|
||||
from fractions import Fraction
|
||||
from io import BytesIO
|
||||
|
||||
import av
|
||||
import numpy as np
|
||||
import torch
|
||||
from av.audio.frame import AudioFrame
|
||||
from av.audio.resampler import AudioResampler
|
||||
from av.video.frame import VideoFrame
|
||||
from comfy_api.input import AudioInput, VideoInput
|
||||
from comfy_api.latest import InputImpl, Types
|
||||
from comfy_extras import nodes_audio
|
||||
from PIL import Image
|
||||
|
||||
@@ -138,6 +143,75 @@ def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
|
||||
return f"data:{mime_type};base64,{base64_str}"
|
||||
|
||||
|
||||
def bytes_to_video(video_bytes: bytes) -> VideoInput:
|
||||
video_buffer = BytesIO(video_bytes)
|
||||
|
||||
try:
|
||||
with av.open(video_buffer, mode="r") as container:
|
||||
video_stream = next((s for s in container.streams if s.type == "video"), None)
|
||||
if video_stream is None:
|
||||
raise ValueError("No video stream found in decoded payload.")
|
||||
|
||||
frames = []
|
||||
for frame in container.decode(video_stream):
|
||||
if not isinstance(frame, VideoFrame):
|
||||
continue
|
||||
image = frame.to_ndarray(format="rgb24")
|
||||
frames.append(torch.from_numpy(image).float() / 255.0)
|
||||
|
||||
if len(frames) == 0:
|
||||
raise ValueError("No video frames found in decoded payload.")
|
||||
|
||||
images = torch.stack(frames, dim=0)
|
||||
frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
|
||||
|
||||
audio: AudioInput | None = None
|
||||
if len(container.streams.audio):
|
||||
audio_stream = container.streams.audio[-1]
|
||||
audio_frames = []
|
||||
resampler = AudioResampler(format="fltp")
|
||||
for frame in container.decode(audio_stream):
|
||||
if not isinstance(frame, AudioFrame):
|
||||
continue
|
||||
resampled = resampler.resample(frame)
|
||||
if not isinstance(resampled, list):
|
||||
resampled = [resampled]
|
||||
for audio_frame in resampled:
|
||||
if audio_frame is not None:
|
||||
audio_frames.append(audio_frame.to_ndarray())
|
||||
|
||||
if len(audio_frames) > 0:
|
||||
audio_data = np.concatenate(audio_frames, axis=1)
|
||||
sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1
|
||||
audio = {
|
||||
"waveform": torch.from_numpy(audio_data).unsqueeze(0),
|
||||
"sample_rate": sample_rate,
|
||||
}
|
||||
|
||||
components = Types.VideoComponents(
|
||||
images=images,
|
||||
frame_rate=frame_rate,
|
||||
audio=audio,
|
||||
metadata=container.metadata if container.metadata else None,
|
||||
)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to decode video: {e}")
|
||||
|
||||
return InputImpl.VideoFromComponents(components)
|
||||
|
||||
|
||||
def base64_to_video(base64_str: str) -> VideoInput:
|
||||
if base64_str.startswith("data:video"):
|
||||
_, base64_str = base64_str.split(",", 1)
|
||||
|
||||
try:
|
||||
video_bytes = base64.b64decode(base64_str)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid base64 string: {e}")
|
||||
|
||||
return bytes_to_video(video_bytes)
|
||||
|
||||
|
||||
def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO:
|
||||
waveform = audio["waveform"][0] # Shape: (C, T)
|
||||
sample_rate = audio["sample_rate"]
|
||||
|
||||
@@ -23,6 +23,10 @@ class QwenTTSModelSpecificParams(dict):
|
||||
pass
|
||||
|
||||
|
||||
class WanModelSpecificParams(dict):
|
||||
pass
|
||||
|
||||
|
||||
class ModelMode(Enum):
|
||||
IMAGE_GENERATION = auto()
|
||||
VIDEO_GENERATION = auto()
|
||||
|
||||
@@ -16,7 +16,7 @@ def validate_model_and_sampling_params_types(
|
||||
# Skip if no spec or no sampling params
|
||||
pipeline_spec, _ = lookup_model_spec(model_name)
|
||||
if pipeline_spec is None:
|
||||
logger.info("skipping sampling params check because spec is not found")
|
||||
logger.info(f"skipping sampling params check because spec for {model_name} is not found")
|
||||
return
|
||||
if sampling_param_list is None:
|
||||
return
|
||||
@@ -70,13 +70,17 @@ def add_sampling_parameters_to_stage(
|
||||
|
||||
stages = pipeline_spec["stages"]
|
||||
if isinstance(sampling_param_list, dict):
|
||||
sampling_param_list = sampling_param_list.__class__(sampling_param_list)
|
||||
sampling_param_list.update(params_to_add)
|
||||
elif sampling_param_list is None:
|
||||
sampling_param_list = params_to_add.copy()
|
||||
else:
|
||||
for i, stage in enumerate(stages):
|
||||
if stage == stage_type:
|
||||
sampling_param_list[i].update(params_to_add)
|
||||
stage_param = sampling_param_list[i]
|
||||
stage_param = stage_param.__class__(stage_param)
|
||||
stage_param.update(params_to_add)
|
||||
sampling_param_list[i] = stage_param
|
||||
|
||||
return sampling_param_list
|
||||
|
||||
|
||||
|
After Width: | Height: | Size: 129 KiB |
@@ -0,0 +1,513 @@
|
||||
{
|
||||
"id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21",
|
||||
"revision": 0,
|
||||
"last_node_id": 14,
|
||||
"last_link_id": 28,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 10,
|
||||
"type": "MarkdownNote",
|
||||
"pos": [
|
||||
227.99306819575278,
|
||||
-231.24143306069843
|
||||
],
|
||||
"size": [
|
||||
240.3326478922474,
|
||||
136.3505820791881
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [],
|
||||
"title": "Note: Task and Input",
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload depending on whether an input image is present or not."
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#000"
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
496.31859627609606,
|
||||
-229.71277089860084
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
314
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "image",
|
||||
"name": "image",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "image"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "choose file to upload",
|
||||
"name": "upload",
|
||||
"type": "IMAGEUPLOAD",
|
||||
"widget": {
|
||||
"name": "upload"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "IMAGE",
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
25
|
||||
]
|
||||
},
|
||||
{
|
||||
"localized_name": "MASK",
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"cute-cat.jpg",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 8,
|
||||
"type": "VLLMOmniDiffusionSampling",
|
||||
"pos": [
|
||||
465.5140218220927,
|
||||
148.80072646828987
|
||||
],
|
||||
"size": [
|
||||
284.205078125,
|
||||
226
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "n",
|
||||
"name": "n",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "n"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "num_inference_steps",
|
||||
"name": "num_inference_steps",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "num_inference_steps"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "guidance_scale",
|
||||
"name": "guidance_scale",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "guidance_scale"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "true_cfg_scale",
|
||||
"name": "true_cfg_scale",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "true_cfg_scale"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "vae_use_slicing",
|
||||
"name": "vae_use_slicing",
|
||||
"type": "BOOLEAN",
|
||||
"widget": {
|
||||
"name": "vae_use_slicing"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "vae_use_tiling",
|
||||
"name": "vae_use_tiling",
|
||||
"type": "BOOLEAN",
|
||||
"widget": {
|
||||
"name": "vae_use_tiling"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "seed",
|
||||
"name": "seed",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "seed"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "diffusion sampling params",
|
||||
"name": "diffusion sampling params",
|
||||
"type": "SAMPLING_PARAMS",
|
||||
"links": [
|
||||
27
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniDiffusionSampling"
|
||||
},
|
||||
"widgets_values": [
|
||||
1,
|
||||
50,
|
||||
1,
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
1992,
|
||||
"randomize"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 14,
|
||||
"type": "VLLMOmniWanParams",
|
||||
"pos": [
|
||||
518.6005521845747,
|
||||
428.4455663719664
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
106
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "guidance_scale_2",
|
||||
"name": "guidance_scale_2",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "guidance_scale_2"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "boundary_ratio",
|
||||
"name": "boundary_ratio",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "boundary_ratio"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "flow_shift",
|
||||
"name": "flow_shift",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "flow_shift"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "Wan video params",
|
||||
"name": "Wan video params",
|
||||
"type": "VIDEO_PARAMS",
|
||||
"links": [
|
||||
28
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniWanParams"
|
||||
},
|
||||
"widgets_values": [
|
||||
4,
|
||||
0.875,
|
||||
5
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 12,
|
||||
"type": "VLLMOmniGenerateVideo",
|
||||
"pos": [
|
||||
827.2566336087867,
|
||||
-73.77449831827545
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
346
|
||||
],
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "image",
|
||||
"name": "image",
|
||||
"shape": 7,
|
||||
"type": "IMAGE",
|
||||
"link": 25
|
||||
},
|
||||
{
|
||||
"localized_name": "sampling_params",
|
||||
"name": "sampling_params",
|
||||
"shape": 7,
|
||||
"type": "SAMPLING_PARAMS",
|
||||
"link": 27
|
||||
},
|
||||
{
|
||||
"localized_name": "lora",
|
||||
"name": "lora",
|
||||
"shape": 7,
|
||||
"type": "REMOTE_LORA",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "model_params",
|
||||
"name": "model_params",
|
||||
"shape": 7,
|
||||
"type": "VIDEO_PARAMS",
|
||||
"link": 28
|
||||
},
|
||||
{
|
||||
"localized_name": "url",
|
||||
"name": "url",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "url"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "model",
|
||||
"name": "model",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "model"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "prompt",
|
||||
"name": "prompt",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "prompt"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "negative_prompt",
|
||||
"name": "negative_prompt",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "negative_prompt"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "width",
|
||||
"name": "width",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "width"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "height",
|
||||
"name": "height",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "height"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "fps",
|
||||
"name": "fps",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "fps"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "num_frames",
|
||||
"name": "num_frames",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "num_frames"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "video",
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"links": [
|
||||
26
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniGenerateVideo"
|
||||
},
|
||||
"widgets_values": [
|
||||
"http://localhost:8000/v1",
|
||||
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
||||
"Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.",
|
||||
"Cartoonish",
|
||||
900,
|
||||
1200,
|
||||
12,
|
||||
41
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 13,
|
||||
"type": "SaveVideo",
|
||||
"pos": [
|
||||
1257.5441369501657,
|
||||
-44.12956394419863
|
||||
],
|
||||
"size": [
|
||||
374,
|
||||
372.34896826553796
|
||||
],
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "video",
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"link": 26
|
||||
},
|
||||
{
|
||||
"localized_name": "filename_prefix",
|
||||
"name": "filename_prefix",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "filename_prefix"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "format",
|
||||
"name": "format",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "format"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "codec",
|
||||
"name": "codec",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "codec"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"video/ComfyUI",
|
||||
"auto",
|
||||
"auto"
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
25,
|
||||
4,
|
||||
0,
|
||||
12,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
26,
|
||||
12,
|
||||
0,
|
||||
13,
|
||||
0,
|
||||
"VIDEO"
|
||||
],
|
||||
[
|
||||
27,
|
||||
8,
|
||||
0,
|
||||
12,
|
||||
1,
|
||||
"SAMPLING_PARAMS"
|
||||
],
|
||||
[
|
||||
28,
|
||||
14,
|
||||
0,
|
||||
12,
|
||||
3,
|
||||
"VIDEO_PARAMS"
|
||||
]
|
||||
],
|
||||
"groups": [
|
||||
{
|
||||
"id": 1,
|
||||
"title": "Input",
|
||||
"bounding": [
|
||||
213.1706010087147,
|
||||
-313.20095750667554,
|
||||
560.8246471437027,
|
||||
407.4246532472182
|
||||
],
|
||||
"color": "#3f789e",
|
||||
"font_size": 24,
|
||||
"flags": {}
|
||||
}
|
||||
],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"workflowRendererVersion": "LG",
|
||||
"ds": {
|
||||
"scale": 1.1469075819486894,
|
||||
"offset": [
|
||||
101.7632950847089,
|
||||
492.29121889347806
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
@@ -43,6 +43,27 @@ The benchmark supports three dataset modes via `--dataset`:
|
||||
|
||||
### VBench dataset
|
||||
|
||||
`vbench` only provides prompt data (and image paths for i2v/i2i); it does not carry
|
||||
per-request generation fields. In this mode, all requests share CLI values:
|
||||
`--width --height --num-frames --fps --num-inference-steps`
|
||||
(pass `--width` and `--height` together).
|
||||
|
||||
Example (`t2v`):
|
||||
|
||||
```bash
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--base-url http://localhost:8099 \
|
||||
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--task t2v \
|
||||
--dataset vbench \
|
||||
--num-prompts 50 \
|
||||
--width 640 --height 480 \
|
||||
--num-frames 81 --fps 16 \
|
||||
--num-inference-steps 40
|
||||
```
|
||||
|
||||
Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths).
|
||||
|
||||
If you use i2v/i2i bench datasets and need auto-download support, you may need:
|
||||
|
||||
```bash
|
||||
@@ -90,7 +111,7 @@ Common optional flags:
|
||||
Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`.
|
||||
|
||||
- For `vbench` / `random`: these CLI flags act as global defaults for all generated requests.
|
||||
- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`).
|
||||
- For `trace`: requests can carry their own fields (e.g., `width/height/num_frames/num_inference_steps`), with overrides/fallbacks as below.
|
||||
|
||||
Precedence rules for `trace` (i.e., what actually gets sent):
|
||||
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import mimetypes
|
||||
import os
|
||||
@@ -35,6 +36,7 @@ class RequestFuncOutput:
|
||||
error: str = ""
|
||||
start_time: float = 0.0
|
||||
response_body: dict[str, Any] = field(default_factory=dict)
|
||||
stage_durations: dict[str, float] = field(default_factory=dict)
|
||||
peak_memory_mb: float = 0.0
|
||||
slo_achieved: bool | None = None
|
||||
|
||||
@@ -55,6 +57,7 @@ async def async_request_chat_completions(
|
||||
input: RequestFuncInput,
|
||||
session: aiohttp.ClientSession,
|
||||
pbar: tqdm | None = None,
|
||||
enable_diffusion_pipeline_profiler: bool = False,
|
||||
) -> RequestFuncOutput:
|
||||
output = RequestFuncOutput()
|
||||
output.start_time = time.perf_counter()
|
||||
@@ -108,6 +111,18 @@ async def async_request_chat_completions(
|
||||
output.success = True
|
||||
if "peak_memory_mb" in resp_json:
|
||||
output.peak_memory_mb = resp_json["peak_memory_mb"]
|
||||
try:
|
||||
choices = resp_json.get("choices", [])
|
||||
if choices and isinstance(choices, list):
|
||||
msg = choices[0].get("message", {})
|
||||
if isinstance(msg, dict):
|
||||
content = msg.get("content", [])
|
||||
if content and isinstance(content, list) and len(content) > 0:
|
||||
first_item = content[0]
|
||||
if isinstance(first_item, dict):
|
||||
output.stage_durations = first_item.get("stage_durations")
|
||||
except (IndexError, TypeError, AttributeError):
|
||||
pass
|
||||
else:
|
||||
output.error = f"HTTP {response.status}: {await response.text()}"
|
||||
output.success = False
|
||||
@@ -192,7 +207,140 @@ async def async_request_openai_images(
|
||||
return output
|
||||
|
||||
|
||||
async def async_request_v1_videos(
|
||||
input: RequestFuncInput,
|
||||
session: aiohttp.ClientSession,
|
||||
pbar: tqdm | None = None,
|
||||
) -> RequestFuncOutput:
|
||||
output = RequestFuncOutput()
|
||||
output.start_time = time.perf_counter()
|
||||
|
||||
files = dict(input.extra_body)
|
||||
if input.prompt:
|
||||
files.setdefault("prompt", input.prompt)
|
||||
if input.width and input.height:
|
||||
files.setdefault("height", input.height)
|
||||
files.setdefault("width", input.width)
|
||||
if input.num_frames:
|
||||
files.setdefault("num_frames", input.num_frames)
|
||||
if input.num_inference_steps:
|
||||
files.setdefault("num_inference_steps", input.num_inference_steps)
|
||||
if input.seed is not None:
|
||||
files.setdefault("seed", input.seed)
|
||||
if input.fps:
|
||||
files.setdefault("fps", input.fps)
|
||||
|
||||
form = aiohttp.FormData()
|
||||
for k, v in files.items():
|
||||
form.add_field(k, str(v))
|
||||
|
||||
image_file = None
|
||||
if input.image_paths and len(input.image_paths) > 0:
|
||||
image_path = input.image_paths[0]
|
||||
image_file = open(image_path, "rb")
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
image_file,
|
||||
filename=os.path.basename(image_path),
|
||||
content_type="application/octet-stream",
|
||||
)
|
||||
|
||||
job_id = None
|
||||
job_status = None
|
||||
poll_json = {}
|
||||
resp_json = {}
|
||||
|
||||
try:
|
||||
# invoke a post request (POST /v1/videos)
|
||||
async with session.post(input.api_url, data=form) as response:
|
||||
if response.status == 200:
|
||||
resp_json = await response.json()
|
||||
job_id = resp_json.get("id")
|
||||
job_status = resp_json.get("status")
|
||||
if not job_id or not job_status:
|
||||
output.error = "API response missing job 'id' or 'status' field."
|
||||
output.success = False
|
||||
return output
|
||||
else:
|
||||
output.error = f"HTTP {response.status}: {await response.text()}"
|
||||
output.success = False
|
||||
return output
|
||||
|
||||
# invoke a poll request (GET /v1/videos/{video_id})
|
||||
poll_interval = 2.0 # Unit(s)
|
||||
timeout_seconds = 600.0
|
||||
deadline = time.perf_counter() + timeout_seconds
|
||||
job_url = f"{input.api_url}/{job_id}"
|
||||
|
||||
while job_status not in {"completed", "failed"}:
|
||||
await asyncio.sleep(poll_interval)
|
||||
|
||||
async with session.get(job_url) as poll_response:
|
||||
if poll_response.status != 200:
|
||||
output.error = f"Polling failed HTTP {poll_response.status}: {await poll_response.text()}"
|
||||
output.success = False
|
||||
return output
|
||||
|
||||
poll_json = await poll_response.json()
|
||||
job_status = poll_json.get("status")
|
||||
|
||||
if time.perf_counter() >= deadline:
|
||||
output.error = f"Timed out waiting for video job {job_id} to complete."
|
||||
output.success = False
|
||||
return output
|
||||
|
||||
if job_status == "failed":
|
||||
output.error = f"Video job failed: {poll_json}"
|
||||
output.success = False
|
||||
return output
|
||||
|
||||
# invoke a get request (GET /v1/videos/{video_id}/content)
|
||||
content_url = f"{job_url}/content"
|
||||
async with session.get(content_url) as content_response:
|
||||
if content_response.status != 200:
|
||||
output.error = (
|
||||
f"Content retrieval failed HTTP {content_response.status}: {await content_response.text()}"
|
||||
)
|
||||
output.success = False
|
||||
return output
|
||||
|
||||
video_bytes = await content_response.read()
|
||||
output.response_body = video_bytes
|
||||
output.success = True
|
||||
if "peak_memory_mb" in poll_json:
|
||||
output.peak_memory_mb = poll_json["peak_memory_mb"]
|
||||
elif "peak_memory_mb" in resp_json:
|
||||
output.peak_memory_mb = resp_json["peak_memory_mb"]
|
||||
except Exception as e:
|
||||
output.error = str(e)
|
||||
output.success = False
|
||||
finally:
|
||||
if image_file is not None:
|
||||
image_file.close()
|
||||
|
||||
if job_id is not None:
|
||||
try:
|
||||
async with session.delete(f"{input.api_url}/{job_id}") as _:
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"Failed to clean up video job {job_id}: {e}")
|
||||
|
||||
output.latency = time.perf_counter() - output.start_time
|
||||
|
||||
if output.success and input.slo_ms is not None:
|
||||
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return output
|
||||
|
||||
|
||||
backends_function_mapping = {
|
||||
"vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
|
||||
"openai": (async_request_openai_images, "/v1/images/generations"),
|
||||
"2i": {
|
||||
"vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
|
||||
"openai": (async_request_openai_images, "/v1/images/generations"),
|
||||
},
|
||||
"2v": {
|
||||
"v1/videos": (async_request_v1_videos, "/v1/videos"),
|
||||
},
|
||||
}
|
||||
@@ -9,6 +9,7 @@ If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
|
||||
Supports multiple backends:
|
||||
- vllm-omni: Uses /v1/chat/completions endpoint (default)
|
||||
- openai: Uses /v1/images/generations endpoint
|
||||
- v1/videos: Use /v1/videos endpoint
|
||||
|
||||
Usage:
|
||||
# Video (vllm-omni backend)
|
||||
@@ -28,6 +29,16 @@ Usage:
|
||||
--backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
|
||||
--height 1024 --width 1024
|
||||
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni --dataset random --task t2i --num-prompts 1 \
|
||||
--max-concurrency 1 --enable-negative-prompt \
|
||||
--random-request-config '[
|
||||
{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
|
||||
{"width":768,"height":768,"num_inference_steps":20,"weight":0.25},
|
||||
{"width":1024,"height":1024,"num_inference_steps":25,"weight":0.45},
|
||||
{"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
|
||||
]'
|
||||
|
||||
i2i:
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni --dataset vbench --task i2i --num-prompts 10
|
||||
@@ -38,6 +49,16 @@ Usage:
|
||||
--backend openai --dataset vbench --task t2i --num-prompts 10 \
|
||||
--height 1024 --width 1024 --port 3000
|
||||
|
||||
# Video (v1/vedeos)
|
||||
t2v:
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend v1/videos --dataset random --task t2v --num-prompts 1 \
|
||||
--max-concurrency 1 --enable-negative-prompt \
|
||||
--random-request-config '[
|
||||
{"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1}
|
||||
]'
|
||||
|
||||
|
||||
"""
|
||||
|
||||
import argparse
|
||||
@@ -45,7 +66,10 @@ import ast
|
||||
import asyncio
|
||||
import glob
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from abc import ABC, abstractmethod
|
||||
@@ -57,8 +81,11 @@ import aiohttp
|
||||
import numpy as np
|
||||
import requests
|
||||
from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
|
||||
from PIL import Image
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BaseDataset(ABC):
|
||||
def __init__(self, args, api_url: str, model: str):
|
||||
@@ -516,26 +543,66 @@ class TraceDataset(BaseDataset):
|
||||
|
||||
|
||||
class RandomDataset(BaseDataset):
|
||||
def __init__(self, args, api_url: str, model: str):
|
||||
self.args = args
|
||||
self.api_url = api_url
|
||||
self.model = model
|
||||
def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool = False):
|
||||
super().__init__(args, api_url, model)
|
||||
self.num_prompts = args.num_prompts
|
||||
self.enable_negative_prompt = enable_negative_prompt
|
||||
self.random_request_config = getattr(args, "random_request_config", None)
|
||||
if self.random_request_config:
|
||||
self.random_request_config = json.loads(self.random_request_config)
|
||||
self._weights = [p["weight"] for p in self.random_request_config]
|
||||
|
||||
self.random_request_config = [
|
||||
{k: v for k, v in p.items() if k != "weight"} for p in self.random_request_config
|
||||
]
|
||||
|
||||
seed = getattr(args, "random_request_seed", 42)
|
||||
self._rng = random.Random(seed)
|
||||
|
||||
self._sampled_requests = self._rng.choices(
|
||||
self.random_request_config,
|
||||
weights=self._weights,
|
||||
k=self.num_prompts,
|
||||
)
|
||||
else:
|
||||
self._sampled_requests = None
|
||||
|
||||
# Random image generate
|
||||
if self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
|
||||
img = Image.new("RGB", (512, 512), (255, 255, 255))
|
||||
|
||||
image_path = os.path.join(tempfile.gettempdir(), "diffusion_benchmark_random_image.png")
|
||||
self._random_image_path = [image_path]
|
||||
img.save(image_path)
|
||||
else:
|
||||
self._random_image_path = None
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.num_prompts
|
||||
|
||||
def __getitem__(self, idx: int) -> RequestFuncInput:
|
||||
extra_body = {}
|
||||
if self.enable_negative_prompt:
|
||||
extra_body["negative_prompt"] = f"Negative prompt {idx} for benchmarking diffusion models"
|
||||
|
||||
params = {
|
||||
"width": self.args.width,
|
||||
"height": self.args.height,
|
||||
"num_frames": self.args.num_frames,
|
||||
"num_inference_steps": self.args.num_inference_steps,
|
||||
"fps": self.args.fps,
|
||||
}
|
||||
if self._sampled_requests:
|
||||
profile = self._sampled_requests[idx]
|
||||
params.update(profile)
|
||||
return RequestFuncInput(
|
||||
prompt=f"Random prompt {idx} for benchmarking diffusion models",
|
||||
api_url=self.api_url,
|
||||
model=self.model,
|
||||
width=self.args.width,
|
||||
height=self.args.height,
|
||||
num_frames=self.args.num_frames,
|
||||
num_inference_steps=self.args.num_inference_steps,
|
||||
seed=self.args.seed,
|
||||
fps=self.args.fps,
|
||||
extra_body=extra_body,
|
||||
image_paths=self._random_image_path,
|
||||
**params,
|
||||
)
|
||||
|
||||
def get_requests(self) -> list[RequestFuncInput]:
|
||||
@@ -686,6 +753,7 @@ def calculate_metrics(
|
||||
"latency_mean": np.mean(latencies) if latencies else 0,
|
||||
"latency_median": np.median(latencies) if latencies else 0,
|
||||
"latency_p99": np.percentile(latencies, 99) if latencies else 0,
|
||||
"latency_p95": np.percentile(latencies, 95) if latencies else 0,
|
||||
"latency_p50": np.percentile(latencies, 50) if latencies else 0,
|
||||
"peak_memory_mb_max": max(peak_memories) if peak_memories else 0,
|
||||
"peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0,
|
||||
@@ -742,8 +810,32 @@ async def benchmark(args):
|
||||
if args.base_url is None:
|
||||
args.base_url = f"http://{args.host}:{args.port}"
|
||||
|
||||
VIDEO_TASKS = {"t2v", "i2v", "ti2v"}
|
||||
IMAGE_TASKS = {"t2i", "i2i", "ti2i"}
|
||||
|
||||
if args.task in VIDEO_TASKS:
|
||||
task_type = "2v"
|
||||
elif args.task in IMAGE_TASKS:
|
||||
task_type = "2i"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported task: '{args.task}'. "
|
||||
f"Valid video tasks: {sorted(VIDEO_TASKS)}, "
|
||||
f"Valid image tasks: {sorted(IMAGE_TASKS)}"
|
||||
)
|
||||
|
||||
valid_backends = sorted(backends_function_mapping[task_type].keys())
|
||||
|
||||
if args.backend not in valid_backends:
|
||||
logger.error(
|
||||
f"Invalid backend '{args.backend}' for task '{args.task}' (task type: '{task_type}').\n"
|
||||
f"Valid backends for this task type: {valid_backends}\n"
|
||||
f"Example usage: --task {args.task} --backend {valid_backends[0]}"
|
||||
)
|
||||
raise ValueError("Backend validation failed. See log above for valid options.")
|
||||
|
||||
# Setup API URL and request function based on backend
|
||||
request_func, api_url = backends_function_mapping[args.backend]
|
||||
request_func, api_url = backends_function_mapping[task_type][args.backend]
|
||||
api_url = f"{args.base_url}{api_url}"
|
||||
|
||||
if args.dataset == "vbench":
|
||||
@@ -751,7 +843,7 @@ async def benchmark(args):
|
||||
elif args.dataset == "trace":
|
||||
dataset = TraceDataset(args, api_url, args.model)
|
||||
elif args.dataset == "random":
|
||||
dataset = RandomDataset(args, api_url, args.model)
|
||||
dataset = RandomDataset(args, api_url, args.model, args.enable_negative_prompt)
|
||||
else:
|
||||
raise ValueError(f"Unknown dataset: {args.dataset}")
|
||||
|
||||
@@ -847,6 +939,7 @@ async def benchmark(args):
|
||||
print("{:<40} {:<15.4f}".format("Latency Mean (s):", metrics["latency_mean"]))
|
||||
print("{:<40} {:<15.4f}".format("Latency Median (s):", metrics["latency_median"]))
|
||||
print("{:<40} {:<15.4f}".format("Latency P99 (s):", metrics["latency_p99"]))
|
||||
print("{:<40} {:<15.4f}".format("Latency P95 (s):", metrics["latency_p95"]))
|
||||
|
||||
if args.slo:
|
||||
print(f"{'-' * 50}")
|
||||
@@ -883,7 +976,7 @@ if __name__ == "__main__":
|
||||
"--backend",
|
||||
type=str,
|
||||
default="vllm-omni",
|
||||
choices=["vllm-omni", "openai"],
|
||||
choices=["vllm-omni", "openai", "v1/videos"],
|
||||
help="Backend to target the benchmark to.",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -972,6 +1065,25 @@ if __name__ == "__main__":
|
||||
help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).",
|
||||
)
|
||||
parser.add_argument("--disable-tqdm", action="store_true", help="Disable progress bar.")
|
||||
parser.add_argument(
|
||||
"--enable-negative-prompt",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Generate negative prompts when using the random dataset.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--random-request-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"JSON string defining random request profiles. "
|
||||
"Each profile may contain: width, height, num_inference_steps, etc. "
|
||||
"The 'weight' field controls sampling probability (relative weight). "
|
||||
"Example: "
|
||||
'[{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},'
|
||||
'{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
|
||||
),
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
@@ -0,0 +1,169 @@
|
||||
# Qwen-Image Serving Performance Dashboard
|
||||
|
||||
This document describes how to deploy and benchmark **Qwen-Image** using vLLM-Omni. It includes service startup configuration, acceleration-related options, benchmark methodology, dataset settings, and performance results.
|
||||
|
||||
---
|
||||
|
||||
# 1. Overview
|
||||
|
||||
Qwen-Image is a multimodal text-to-image generation model served through the vLLM-Omni infrastructure.
|
||||
|
||||
This document covers:
|
||||
|
||||
* Service launch configuration (including acceleration options)
|
||||
* Benchmark scripts and usage
|
||||
* Dataset and workload settings
|
||||
* Performance measurement results
|
||||
* Reproducibility guidelines
|
||||
|
||||
---
|
||||
|
||||
# 2. Test Environment
|
||||
| Component | Specification |
|
||||
|------------|----------------|
|
||||
| GPU | NVIDIA A100-SXM4-80GB |
|
||||
| Diffusion Attention Backend | FlashAttention |
|
||||
|
||||
# 3. Service Launch Configuration
|
||||
|
||||
## 3.1 Basic Serving Command
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
## 3.2 Key Parameters
|
||||
|
||||
| Parameter | Description |
|
||||
| --------------------- | ------------------------ |
|
||||
| `--cfg-parallel-size` | CFG parallelism degree |
|
||||
| `--ulysses-degree` | Ulysses parallel degree |
|
||||
| `--vae-patch-parallel-size` | VAE parallel degree |
|
||||
| `--tensor-parallel-size` | Tensor parallelism degree |
|
||||
|
||||
Record these parameters when reporting performance results.
|
||||
|
||||
---
|
||||
|
||||
# 4. Benchmark Script
|
||||
|
||||
## 4.1 Benchmark Entry
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni \
|
||||
--dataset <DATASET_NAME> \
|
||||
--task t2i \
|
||||
--num-prompts <N> \
|
||||
--max-concurrency <C> \
|
||||
--enable-negative-prompt \
|
||||
--random-request-config <CFG>
|
||||
```
|
||||
|
||||
## 4.2 Key Benchmark Arguments
|
||||
|
||||
| Parameter | Description |
|
||||
| ---------------------- | --------------------------------- |
|
||||
| `--backend` | Serving backend (use `vllm-omni`) |
|
||||
| `--dataset` | Dataset name (`random` or custom) |
|
||||
| `--task` | Task type (e.g., `t2i`) |
|
||||
| `--num-prompts` | Total number of requests |
|
||||
| `--max-concurrency` | Client-side concurrency |
|
||||
| `--random-request-config`| JSON string defining random request |
|
||||
|
||||
---
|
||||
|
||||
# 5. Dataset & Workload Settings
|
||||
|
||||
## 5.1 Recommended Evaluation Configurations
|
||||
|
||||
### Dataset A ( 512 Resolution)
|
||||
|
||||
* Dataset: `random`
|
||||
* Task: t2i
|
||||
* Concurrency: 1
|
||||
* Mix Resolution
|
||||
```
|
||||
[
|
||||
{"width":512,"height":512,"num_inference_steps":20,"weight":1}
|
||||
]
|
||||
```
|
||||
|
||||
### Dataset B (1536 Resolution)
|
||||
|
||||
* Dataset: `random`
|
||||
* Task: t2i
|
||||
* Concurrency: 1
|
||||
* Mix Resolution
|
||||
```
|
||||
[
|
||||
{"width":1536,"height":1536,"num_inference_steps":35,"weight":1}
|
||||
]
|
||||
```
|
||||
|
||||
### Dataset C (Mix Resolution)
|
||||
|
||||
* Dataset: `random`
|
||||
* Task: t2i
|
||||
* Concurrency: 1
|
||||
* Mix Resolution
|
||||
```
|
||||
[
|
||||
{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
|
||||
{"width":768,"height":768,"num_inference_steps":20,"weight":0.25},
|
||||
{"width":1024,"height":1024,"num_inference_steps":25,"weight":0.45},
|
||||
{"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
|
||||
]
|
||||
```
|
||||
---
|
||||
|
||||
## 5.2 Example Benchmark Command
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni \
|
||||
--dataset random \
|
||||
--task t2i \
|
||||
--num-prompts 1 \
|
||||
--max-concurrency 1 \
|
||||
--enable-negative-prompt \
|
||||
--random-request-config '[
|
||||
{"width":512,"height":512,"num_inference_steps":20,"weight":1}
|
||||
]'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
# 6. Performance Metrics
|
||||
|
||||
The following metrics are collected during benchmarking:
|
||||
|
||||
| Metric | Description | Unit |
|
||||
| ------------------ | ----------------------------- | ------- |
|
||||
| Mean Latency | Mean of latency | seconds |
|
||||
| P99 Latency | P99 of latency | seconds |
|
||||
|
||||
---
|
||||
|
||||
# 7. Performance Results
|
||||
|
||||
| Dataset Configuration | Max Concur. | CFG | Usp | Tp | VAE Parallel | Mean Latency (s) | P99 Latency (s) |
|
||||
|-----------------------|-----|-----|-----|----|--------------|------------------|------------------|
|
||||
| Dataset A | 1 | 2 | 2 | Off | Off | 2.2087 | 2.2087 |
|
||||
| Dataset B | 1 | 2 | 2 | Off | Off | 19.6739 | 19.6739 |
|
||||
| Dataset C | 1 | 2 | 2 | Off | Off | 5.67259 | 18.6234 |
|
||||
---
|
||||
|
||||
# 8. Reproducibility Checklist
|
||||
|
||||
To ensure consistent and comparable benchmark results:
|
||||
|
||||
* Record GPU type
|
||||
* Record parallel configuration
|
||||
* Record benchmark parameters (resolution, concurrency, number of prompts)
|
||||
* Ensure no background workload on GPUs during testing
|
||||
|
||||
---
|
||||
|
||||
This document serves as the official Qwen-Image serving performance reference under vLLM-Omni.
|
||||
@@ -0,0 +1,170 @@
|
||||
# Wan2.2 Serving Performance Dashboard
|
||||
|
||||
This document describes how to deploy and benchmark **Wan-AI/Wan2.2-T2V-A14B-Diffusers** using vLLM-Omni. It includes service startup configuration, acceleration-related options, benchmark methodology, dataset settings, and performance results.
|
||||
|
||||
---
|
||||
|
||||
# 1. Overview
|
||||
|
||||
Wan-AI/Wan2.2-T2V-A14B-Diffusers is a multimodal text-to-video generation model served through the vLLM-Omni infrastructure.
|
||||
|
||||
This document covers:
|
||||
|
||||
* Service launch configuration (including acceleration options)
|
||||
* Benchmark scripts and usage
|
||||
* Dataset and workload settings
|
||||
* Performance measurement results
|
||||
* Reproducibility guidelines
|
||||
|
||||
---
|
||||
|
||||
# 2. Test Environment
|
||||
| Component | Specification |
|
||||
|------------|----------------|
|
||||
| GPU | NVIDIA A100-SXM4-80GB |
|
||||
| Diffusion Attention Backend | FlashAttention |
|
||||
|
||||
# 3. Service Launch Configuration
|
||||
|
||||
## 3.1 Basic Serving Command
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
## 3.2 Key Parameters
|
||||
|
||||
| Parameter | Description |
|
||||
| --------------------- | ------------------------ |
|
||||
| `--cfg-parallel-size` | CFG parallelism degree |
|
||||
| `--ulysses-degree` | Ulysses parallel degree |
|
||||
| `--vae-patch-parallel-size` | VAE parallel degree |
|
||||
| `--tensor-parallel-size` | Tensor parallelism degree |
|
||||
| `--use-hsdp` | Enable Hybrid Sharded Data Parallel to shard model weights across GPUs |
|
||||
|
||||
Record these parameters when reporting performance results.
|
||||
|
||||
---
|
||||
|
||||
# 4. Benchmark Script
|
||||
|
||||
## 4.1 Benchmark Entry
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend v1/videos \
|
||||
--dataset <DATASET_NAME> \
|
||||
--task t2v \
|
||||
--num-prompts <N> \
|
||||
--max-concurrency <C> \
|
||||
--enable-negative-prompt \
|
||||
--random-request-config <CFG>
|
||||
```
|
||||
|
||||
## 4.2 Key Benchmark Arguments
|
||||
|
||||
| Parameter | Description |
|
||||
| ---------------------- | --------------------------------- |
|
||||
| `--backend` | Serving backend (use `v1/videos`) |
|
||||
| `--dataset` | Dataset name (`random` or custom) |
|
||||
| `--task` | Task type (e.g., `t2v`) |
|
||||
| `--num-prompts` | Total number of requests |
|
||||
| `--max-concurrency` | Client-side concurrency |
|
||||
| `--random-request-config`| JSON string defining random request |
|
||||
|
||||
---
|
||||
|
||||
# 5. Dataset & Workload Settings
|
||||
|
||||
## 5.1 Recommended Evaluation Configurations
|
||||
|
||||
### Dataset A (480p)
|
||||
|
||||
* Dataset: `random`
|
||||
* Task: t2v
|
||||
* Concurrency: 1
|
||||
* Mix Resolution
|
||||
```
|
||||
[
|
||||
{"width":854,"height":480,"num_inference_steps":3,"num_frames":80,"fps":16,"weight":1}
|
||||
]
|
||||
```
|
||||
### Dataset B (720p)
|
||||
|
||||
* Dataset: `random`
|
||||
* Task: t2v
|
||||
* Concurrency: 1
|
||||
* Mix Resolution
|
||||
```
|
||||
[
|
||||
{"width":1280,"height":720,"num_inference_steps":6,"num_frames":80,"fps":16,"weight":1}
|
||||
]
|
||||
```
|
||||
### Dataset C (Mix Resolution)
|
||||
|
||||
* Dataset: `random`
|
||||
* Task: t2v
|
||||
* Concurrency: 1
|
||||
* Mix Resolution
|
||||
```
|
||||
[
|
||||
{"width":854,"height":480,"num_inference_steps":3,"num_frames":80,"fps":16,"weight":0.15},
|
||||
{"width":854,"height":480,"num_inference_steps":4,"num_frames":120,"fps":24,"weight":0.25},
|
||||
{"width":1280,"height":720,"num_inference_steps":6,"num_frames":80,"fps":16,"weight":0.6}
|
||||
]
|
||||
```
|
||||
---
|
||||
|
||||
## 5.2 Example Benchmark Command
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend v1/videos \
|
||||
--dataset random \
|
||||
--task t2v \
|
||||
--num-prompts 1 \
|
||||
--max-concurrency 1 \
|
||||
--enable-negative-prompt \
|
||||
--random-request-config '[
|
||||
{"width":854,"height":480,"num_inference_steps":18,"num_frames": 33,"fps":16",weight":1}
|
||||
]'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
# 6. Performance Metrics
|
||||
|
||||
The following metrics are collected during benchmarking:
|
||||
|
||||
| Metric | Description | Unit |
|
||||
| ------------------ | ----------------------------- | ------- |
|
||||
| Mean Latency | Mean of latency | seconds |
|
||||
| P99 Latency | P99 of latency | seconds |
|
||||
|
||||
---
|
||||
|
||||
# 7. Performance Results
|
||||
|
||||
| Dataset Configuration | Max Concur. | CFG | Usp | Tp | Hsdp | VAE Parallel | Mean Latency (s) | P99 Latency (s) |
|
||||
|-----------------------|-----|-----|-----|-----|----|--------------|------------------|------------------|
|
||||
| Dataset A | 1 | 2 | 2 | 1 | On | 1 | 24.6766 | 24.6766 |
|
||||
| Dataset A | 1 | 2 | 2 | 1 | On | 4 | 21.6810 | 21.6810 |
|
||||
| Dataset B | 1 | 2 | 2 | 1 | On | 1 | 124.6639 | 124.6639 |
|
||||
| Dataset B | 1 | 2 | 2 | 1 | On | 4 | 117.44 | 117.44 |
|
||||
| Dataset C | 1 | 2 | 2 | 1 | On | 1 | 79.2175 | 124.2565 |
|
||||
| Dataset C | 1 | 2 | 2 | 1 | On | 4 | 74.4977 | 117.710 |
|
||||
---
|
||||
|
||||
# 8. Reproducibility Checklist
|
||||
|
||||
To ensure consistent and comparable benchmark results:
|
||||
|
||||
* Record GPU type
|
||||
* Record parallel configuration
|
||||
* Record benchmark parameters (resolution, concurrency, number of prompts)
|
||||
* Ensure no background workload on GPUs during testing
|
||||
|
||||
---
|
||||
|
||||
This document serves as the official Wan2.2 serving performance reference under vLLM-Omni.
|
||||
@@ -56,9 +56,9 @@ What it does:
|
||||
- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
|
||||
- Uses `benchmarks/build_dataset/top100.txt` and writes to:
|
||||
- Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
|
||||
- `omni_llm_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
|
||||
- `omni_llm_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
|
||||
- `omni_llm_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
|
||||
- `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:
|
||||
|
||||
@@ -26,12 +26,12 @@ else
|
||||
--log-stats \
|
||||
--log-dir $log_dir
|
||||
echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
|
||||
echo " - omni_llm_pipeline_text run dir/base name"
|
||||
echo " - omni_llm_pipeline_text.orchestrator.stats.jsonl orchestrator-stage latency stats"
|
||||
echo " - omni_llm_pipeline_text.overall.stats.jsonl overall latency/TPS stats"
|
||||
echo " - omni_llm_pipeline_text.stage0.log per-stage detailed logs"
|
||||
echo " - omni_llm_pipeline_text.stage1.log"
|
||||
echo " - omni_llm_pipeline_text.stage2.log"
|
||||
echo " - 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
|
||||
@@ -0,0 +1,100 @@
|
||||
# Qwen3-TTS Benchmark
|
||||
|
||||
Benchmarks for Qwen3-TTS text-to-speech models, comparing vLLM-Omni streaming serving against HuggingFace Transformers offline inference.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
```bash
|
||||
pip install matplotlib aiohttp soundfile numpy tqdm
|
||||
pip install qwen_tts # for HF baseline
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
Run the full benchmark (vllm-omni + HF baseline) with a single command:
|
||||
|
||||
```bash
|
||||
cd benchmarks/qwen3-tts
|
||||
bash run_benchmark.sh
|
||||
```
|
||||
|
||||
Results (JSON + PNG plots) are saved to `results/`.
|
||||
|
||||
### Common options
|
||||
|
||||
```bash
|
||||
# Only vllm-omni (skip HF baseline)
|
||||
bash run_benchmark.sh --async-only
|
||||
|
||||
# Only HF baseline
|
||||
bash run_benchmark.sh --hf-only
|
||||
|
||||
# Use a different model (e.g. 1.7B)
|
||||
MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
|
||||
|
||||
# Use bs16 config for higher throughput
|
||||
STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs16.yaml bash run_benchmark.sh --async-only
|
||||
|
||||
# Custom GPU, prompt count, concurrency levels
|
||||
GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
|
||||
```
|
||||
|
||||
## Manual Steps
|
||||
|
||||
### 1) Start the vLLM-Omni server
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
|
||||
"Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
|
||||
--omni --host 127.0.0.1 --port 8000 \
|
||||
--stage-configs-path benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml \
|
||||
--trust-remote-code
|
||||
```
|
||||
|
||||
### 2) Run online serving benchmark
|
||||
|
||||
```bash
|
||||
python benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py \
|
||||
--port 8000 \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 4 10 \
|
||||
--config-name "async_chunk" \
|
||||
--result-dir results/
|
||||
```
|
||||
|
||||
### 3) Run HuggingFace baseline
|
||||
|
||||
```bash
|
||||
python benchmarks/qwen3-tts/transformers/bench_tts_hf.py \
|
||||
--model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
|
||||
--num-prompts 50 \
|
||||
--gpu-device 0 \
|
||||
--result-dir results/
|
||||
```
|
||||
|
||||
### 4) Generate comparison plots
|
||||
|
||||
```bash
|
||||
python benchmarks/qwen3-tts/plot_results.py \
|
||||
--results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
|
||||
--labels "vllm-omni" "hf_transformers" \
|
||||
--output results/comparison.png
|
||||
```
|
||||
|
||||
## Stage Configs
|
||||
|
||||
| Config | max_num_seqs | Description |
|
||||
|--------|:------------:|-------------|
|
||||
| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lowest latency) |
|
||||
| `vllm_omni/configs/qwen3_tts_bs16.yaml` | 16 | High-throughput concurrent processing |
|
||||
|
||||
All configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
|
||||
|
||||
The model is specified via the CLI `--model` flag (or `MODEL` env var), so the same configs work for both the 0.6B and 1.7B model variants.
|
||||
|
||||
## Metrics
|
||||
|
||||
- **TTFP (Time to First Audio Packet)**: Time from request to first audio chunk (streaming latency)
|
||||
- **E2E (End-to-End Latency)**: Total time from request to complete audio response
|
||||
- **RTF (Real-Time Factor)**: E2E latency / audio duration. RTF < 1.0 means faster-than-real-time synthesis
|
||||
- **Throughput**: Total audio seconds generated per wall-clock second
|
||||
@@ -0,0 +1,254 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,3 @@
|
||||
# Benchmark results are machine-specific - do not commit
|
||||
*
|
||||
!.gitignore
|
||||
@@ -0,0 +1,272 @@
|
||||
#!/bin/bash
|
||||
# Qwen3-TTS Benchmark Runner
|
||||
#
|
||||
# Compares vllm-omni streaming serving vs HuggingFace transformers offline inference.
|
||||
# Produces JSON results and comparison plots.
|
||||
#
|
||||
# Usage:
|
||||
# # Full comparison (vllm-omni + HF):
|
||||
# bash run_benchmark.sh
|
||||
#
|
||||
# # Only vllm-omni async_chunk config:
|
||||
# bash run_benchmark.sh --async-only
|
||||
#
|
||||
# # Only HuggingFace baseline:
|
||||
# bash run_benchmark.sh --hf-only
|
||||
#
|
||||
# # vllm-omni only (skip HF):
|
||||
# bash run_benchmark.sh --skip-hf
|
||||
#
|
||||
# # Custom settings:
|
||||
# GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
|
||||
#
|
||||
# # Use 1.7B model:
|
||||
# MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
|
||||
#
|
||||
# # Use batch_size=4 config:
|
||||
# STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
|
||||
#
|
||||
# Environment variables:
|
||||
# GPU_DEVICE - GPU index to use (default: 0)
|
||||
# NUM_PROMPTS - Number of prompts per concurrency level (default: 50)
|
||||
# CONCURRENCY - Space-separated concurrency levels (default: "1 4 10")
|
||||
# MODEL - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
|
||||
# PORT - Server port (default: 8000)
|
||||
# GPU_MEM_TALKER - gpu_memory_utilization for talker stage (default: 0.3)
|
||||
# GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.2)
|
||||
# STAGE_CONFIG - Path to stage config YAML (default: configs/qwen3_tts_bs1.yaml)
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
|
||||
|
||||
# Defaults
|
||||
GPU_DEVICE="${GPU_DEVICE:-0}"
|
||||
NUM_PROMPTS="${NUM_PROMPTS:-50}"
|
||||
CONCURRENCY="${CONCURRENCY:-1 4 10}"
|
||||
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
|
||||
PORT="${PORT:-8000}"
|
||||
GPU_MEM_TALKER="${GPU_MEM_TALKER:-0.3}"
|
||||
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-0.2}"
|
||||
NUM_WARMUPS="${NUM_WARMUPS:-3}"
|
||||
STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/configs/qwen3_tts_bs1.yaml}"
|
||||
RESULT_DIR="${SCRIPT_DIR}/results"
|
||||
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
||||
|
||||
# Parse args
|
||||
RUN_ASYNC=true
|
||||
RUN_HF=true
|
||||
for arg in "$@"; do
|
||||
case "$arg" in
|
||||
--async-only) RUN_HF=false ;;
|
||||
--hf-only) RUN_ASYNC=false ;;
|
||||
--skip-hf) RUN_HF=false ;;
|
||||
esac
|
||||
done
|
||||
|
||||
mkdir -p "${RESULT_DIR}"
|
||||
|
||||
echo "============================================================"
|
||||
echo " Qwen3-TTS Benchmark"
|
||||
echo "============================================================"
|
||||
echo " GPU: ${GPU_DEVICE}"
|
||||
echo " Model: ${MODEL}"
|
||||
echo " Prompts: ${NUM_PROMPTS}"
|
||||
echo " Concurrency: ${CONCURRENCY}"
|
||||
echo " Port: ${PORT}"
|
||||
echo " Stage config: ${STAGE_CONFIG}"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo "============================================================"
|
||||
|
||||
# Prepare stage config with correct GPU device and memory settings
|
||||
prepare_config() {
|
||||
local config_template="$1"
|
||||
local config_name="$2"
|
||||
local output_path="${RESULT_DIR}/${config_name}_stage_config.yaml"
|
||||
|
||||
# Use sed to patch GPU device and memory utilization
|
||||
sed \
|
||||
-e "s/devices: \"0\"/devices: \"${GPU_DEVICE}\"/g" \
|
||||
-e "s/gpu_memory_utilization: 0.3/gpu_memory_utilization: ${GPU_MEM_TALKER}/g" \
|
||||
-e "s/gpu_memory_utilization: 0.2/gpu_memory_utilization: ${GPU_MEM_CODE2WAV}/g" \
|
||||
"${config_template}" > "${output_path}"
|
||||
|
||||
echo "${output_path}"
|
||||
}
|
||||
|
||||
# Start server and wait for it to be ready
|
||||
start_server() {
|
||||
local stage_config="$1"
|
||||
local config_name="$2"
|
||||
local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log"
|
||||
|
||||
echo ""
|
||||
echo "Starting server with config: ${config_name}"
|
||||
echo " Stage config: ${stage_config}"
|
||||
echo " Log file: ${log_file}"
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
CUDA_VISIBLE_DEVICES="${GPU_DEVICE}" \
|
||||
python -m vllm_omni.entrypoints.cli.main serve "${MODEL}" \
|
||||
--omni \
|
||||
--host 127.0.0.1 \
|
||||
--port "${PORT}" \
|
||||
--stage-configs-path "${stage_config}" \
|
||||
--stage-init-timeout 120 \
|
||||
--trust-remote-code \
|
||||
--disable-log-stats \
|
||||
> "${log_file}" 2>&1 &
|
||||
|
||||
SERVER_PID=$!
|
||||
echo " Server PID: ${SERVER_PID}"
|
||||
|
||||
# Wait for server to be ready
|
||||
echo " Waiting for server to be ready..."
|
||||
local max_wait=300
|
||||
local waited=0
|
||||
while [ ${waited} -lt ${max_wait} ]; do
|
||||
if curl -sf "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then
|
||||
echo " Server is ready! (waited ${waited}s)"
|
||||
return 0
|
||||
fi
|
||||
# Check if process is still alive
|
||||
if ! kill -0 ${SERVER_PID} 2>/dev/null; then
|
||||
echo " ERROR: Server process died. Check log: ${log_file}"
|
||||
tail -20 "${log_file}"
|
||||
return 1
|
||||
fi
|
||||
sleep 2
|
||||
waited=$((waited + 2))
|
||||
done
|
||||
|
||||
echo " ERROR: Server did not start within ${max_wait}s. Check log: ${log_file}"
|
||||
kill ${SERVER_PID} 2>/dev/null || true
|
||||
return 1
|
||||
}
|
||||
|
||||
# Stop the server
|
||||
stop_server() {
|
||||
if [ -n "${SERVER_PID:-}" ]; then
|
||||
echo " Stopping server (PID: ${SERVER_PID})..."
|
||||
kill ${SERVER_PID} 2>/dev/null || true
|
||||
wait ${SERVER_PID} 2>/dev/null || true
|
||||
# Kill any remaining child processes on the port
|
||||
local pids
|
||||
pids=$(lsof -ti:${PORT} 2>/dev/null || true)
|
||||
if [ -n "${pids}" ]; then
|
||||
echo " Cleaning up remaining processes on port ${PORT}..."
|
||||
echo "${pids}" | xargs kill -9 2>/dev/null || true
|
||||
fi
|
||||
echo " Server stopped."
|
||||
SERVER_PID=""
|
||||
fi
|
||||
}
|
||||
|
||||
# Cleanup on exit
|
||||
trap 'stop_server' EXIT
|
||||
|
||||
# Run benchmark for a given config
|
||||
run_bench() {
|
||||
local config_name="$1"
|
||||
local config_template="$2"
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmarking: ${config_name}"
|
||||
echo "============================================================"
|
||||
|
||||
local stage_config
|
||||
stage_config=$(prepare_config "${config_template}" "${config_name}")
|
||||
|
||||
start_server "${stage_config}" "${config_name}"
|
||||
|
||||
# Convert concurrency string to args
|
||||
local conc_args=""
|
||||
for c in ${CONCURRENCY}; do
|
||||
conc_args="${conc_args} ${c}"
|
||||
done
|
||||
|
||||
cd "${PROJECT_ROOT}"
|
||||
python "${SCRIPT_DIR}/vllm_omni/bench_tts_serve.py" \
|
||||
--host 127.0.0.1 \
|
||||
--port "${PORT}" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${conc_args} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--config-name "${config_name}" \
|
||||
--result-dir "${RESULT_DIR}"
|
||||
|
||||
stop_server
|
||||
|
||||
# Allow GPU memory to settle
|
||||
sleep 5
|
||||
}
|
||||
|
||||
# Run vllm-omni benchmark
|
||||
if [ "${RUN_ASYNC}" = true ]; then
|
||||
run_bench "async_chunk" "${SCRIPT_DIR}/${STAGE_CONFIG}"
|
||||
fi
|
||||
|
||||
# Run HuggingFace baseline benchmark
|
||||
if [ "${RUN_HF}" = true ]; then
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmarking: HuggingFace transformers (offline)"
|
||||
echo "============================================================"
|
||||
|
||||
cd "${PROJECT_ROOT}"
|
||||
python "${SCRIPT_DIR}/transformers/bench_tts_hf.py" \
|
||||
--model "${MODEL}" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--gpu-device "${GPU_DEVICE}" \
|
||||
--config-name "hf_transformers" \
|
||||
--result-dir "${RESULT_DIR}"
|
||||
|
||||
# Allow GPU memory to settle
|
||||
sleep 5
|
||||
fi
|
||||
|
||||
# Plot results
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Generating plots..."
|
||||
echo "============================================================"
|
||||
|
||||
RESULT_FILES=""
|
||||
LABELS=""
|
||||
|
||||
if [ "${RUN_ASYNC}" = true ]; then
|
||||
ASYNC_FILE=$(ls -t "${RESULT_DIR}"/bench_async_chunk_*.json 2>/dev/null | head -1)
|
||||
if [ -n "${ASYNC_FILE}" ]; then
|
||||
RESULT_FILES="${ASYNC_FILE}"
|
||||
LABELS="async_chunk"
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ "${RUN_HF}" = true ]; then
|
||||
HF_FILE=$(ls -t "${RESULT_DIR}"/bench_hf_transformers_*.json 2>/dev/null | head -1)
|
||||
if [ -n "${HF_FILE}" ]; then
|
||||
if [ -n "${RESULT_FILES}" ]; then
|
||||
RESULT_FILES="${RESULT_FILES} ${HF_FILE}"
|
||||
LABELS="${LABELS} hf_transformers"
|
||||
else
|
||||
RESULT_FILES="${HF_FILE}"
|
||||
LABELS="hf_transformers"
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ -n "${RESULT_FILES}" ]; then
|
||||
python "${SCRIPT_DIR}/plot_results.py" \
|
||||
--results ${RESULT_FILES} \
|
||||
--labels ${LABELS} \
|
||||
--output "${RESULT_DIR}/qwen3_tts_benchmark_${TIMESTAMP}.png"
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmark complete!"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo "============================================================"
|
||||
@@ -0,0 +1,281 @@
|
||||
"""Benchmark Qwen3-TTS using HuggingFace transformers (qwen_tts library).
|
||||
|
||||
Measures E2E latency, RTF, and audio duration for offline (non-serving) inference.
|
||||
Results are saved in the same JSON format as bench_tts_serve.py for unified plotting.
|
||||
|
||||
Usage:
|
||||
python bench_tts_hf.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
|
||||
--num-prompts 50 \
|
||||
--num-warmups 3 \
|
||||
--gpu-device 0 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
config_name: str = ""
|
||||
concurrency: int = 1 # always 1 for offline
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
# TTFP stats - not applicable for HF offline, set to E2E for compatibility
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
# E2E stats (ms)
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
# RTF stats
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
p99_rtf: float = 0.0
|
||||
# Audio stats
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0
|
||||
request_throughput: float = 0.0
|
||||
# Per-request details
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def run_benchmark(args):
|
||||
from qwen_tts import Qwen3TTSModel
|
||||
|
||||
device = f"cuda:{args.gpu_device}"
|
||||
print(f"Loading model: {args.model} on {device}")
|
||||
model = Qwen3TTSModel.from_pretrained(
|
||||
args.model,
|
||||
device_map=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
print("Model loaded.")
|
||||
|
||||
# Build prompt list
|
||||
prompts = [PROMPTS[i % len(PROMPTS)] for i in range(args.num_prompts)]
|
||||
|
||||
# Warmup
|
||||
if args.num_warmups > 0:
|
||||
print(f"Warming up with {args.num_warmups} requests...")
|
||||
for i in range(args.num_warmups):
|
||||
p = PROMPTS[i % len(PROMPTS)]
|
||||
wavs, sr = model.generate_custom_voice(
|
||||
text=p,
|
||||
language=args.language,
|
||||
speaker=args.voice,
|
||||
)
|
||||
# Sync GPU
|
||||
torch.cuda.synchronize(device)
|
||||
print("Warmup done.")
|
||||
|
||||
# Benchmark
|
||||
print(f"Running {args.num_prompts} requests sequentially...")
|
||||
e2e_times = []
|
||||
rtfs = []
|
||||
audio_durations = []
|
||||
per_request = []
|
||||
failed = 0
|
||||
|
||||
audio_dir = None
|
||||
if args.save_audio:
|
||||
audio_dir = Path(args.result_dir) / "audio_hf"
|
||||
audio_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
total_start = time.perf_counter()
|
||||
|
||||
for i, prompt in enumerate(prompts):
|
||||
try:
|
||||
torch.cuda.synchronize(device)
|
||||
st = time.perf_counter()
|
||||
|
||||
wavs, sr = model.generate_custom_voice(
|
||||
text=prompt,
|
||||
language=args.language,
|
||||
speaker=args.voice,
|
||||
)
|
||||
|
||||
torch.cuda.synchronize(device)
|
||||
elapsed = time.perf_counter() - st
|
||||
|
||||
# Compute audio duration
|
||||
audio_samples = wavs[0]
|
||||
if isinstance(audio_samples, torch.Tensor):
|
||||
audio_samples = audio_samples.cpu().numpy()
|
||||
audio_dur = len(audio_samples) / sr
|
||||
|
||||
rtf = elapsed / audio_dur if audio_dur > 0 else 0.0
|
||||
|
||||
e2e_times.append(elapsed)
|
||||
rtfs.append(rtf)
|
||||
audio_durations.append(audio_dur)
|
||||
per_request.append(
|
||||
{
|
||||
"e2e_ms": elapsed * 1000,
|
||||
"ttfp_ms": elapsed * 1000, # no streaming, TTFP = E2E
|
||||
"rtf": rtf,
|
||||
"audio_duration_s": audio_dur,
|
||||
"prompt": prompt,
|
||||
}
|
||||
)
|
||||
|
||||
if audio_dir:
|
||||
sf.write(str(audio_dir / f"output_{i:04d}.wav"), audio_samples, sr)
|
||||
|
||||
if (i + 1) % 10 == 0 or i == 0:
|
||||
print(
|
||||
f" [{i + 1}/{args.num_prompts}] e2e={elapsed * 1000:.0f}ms rtf={rtf:.3f} audio={audio_dur:.2f}s"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f" [{i + 1}/{args.num_prompts}] FAILED: {e}")
|
||||
failed += 1
|
||||
|
||||
total_duration = time.perf_counter() - total_start
|
||||
completed = len(e2e_times)
|
||||
|
||||
# Compute stats
|
||||
result = BenchmarkResult(
|
||||
config_name=args.config_name,
|
||||
concurrency=1,
|
||||
num_prompts=args.num_prompts,
|
||||
completed=completed,
|
||||
failed=failed,
|
||||
duration_s=total_duration,
|
||||
)
|
||||
|
||||
if e2e_times:
|
||||
e2e_ms = [t * 1000 for t in e2e_times]
|
||||
|
||||
result.mean_e2e_ms = float(np.mean(e2e_ms))
|
||||
result.median_e2e_ms = float(np.median(e2e_ms))
|
||||
result.std_e2e_ms = float(np.std(e2e_ms))
|
||||
result.p90_e2e_ms = float(np.percentile(e2e_ms, 90))
|
||||
result.p95_e2e_ms = float(np.percentile(e2e_ms, 95))
|
||||
result.p99_e2e_ms = float(np.percentile(e2e_ms, 99))
|
||||
|
||||
# For HF offline, TTFP = E2E (no streaming)
|
||||
result.mean_ttfp_ms = result.mean_e2e_ms
|
||||
result.median_ttfp_ms = result.median_e2e_ms
|
||||
result.std_ttfp_ms = result.std_e2e_ms
|
||||
result.p90_ttfp_ms = result.p90_e2e_ms
|
||||
result.p95_ttfp_ms = result.p95_e2e_ms
|
||||
result.p99_ttfp_ms = result.p99_e2e_ms
|
||||
|
||||
result.mean_rtf = float(np.mean(rtfs))
|
||||
result.median_rtf = float(np.median(rtfs))
|
||||
result.std_rtf = float(np.std(rtfs))
|
||||
result.p99_rtf = float(np.percentile(rtfs, 99))
|
||||
|
||||
result.mean_audio_duration_s = float(np.mean(audio_durations))
|
||||
result.total_audio_duration_s = float(np.sum(audio_durations))
|
||||
result.audio_throughput = result.total_audio_duration_s / total_duration
|
||||
result.request_throughput = completed / total_duration
|
||||
result.per_request = per_request
|
||||
|
||||
# Print summary in standardized performance template
|
||||
W = 50
|
||||
print("")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Serving Benchmark Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Successful requests:':<40}{completed:<10}")
|
||||
print(f"{'Failed requests:':<40}{failed:<10}")
|
||||
print(f"{'Maximum request concurrency:':<40}{1:<10}")
|
||||
print(f"{'Benchmark duration (s):':<40}{total_duration:<10.2f}")
|
||||
print(f"{'Request throughput (req/s):':<40}{result.request_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'End-to-end Latency':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean E2EL (ms):':<40}{result.mean_e2e_ms:<10.2f}")
|
||||
print(f"{'Median E2EL (ms):':<40}{result.median_e2e_ms:<10.2f}")
|
||||
print(f"{'P99 E2EL (ms):':<40}{result.p99_e2e_ms:<10.2f}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Audio Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Total audio duration generated (s):':<40}{result.total_audio_duration_s:<10.2f}")
|
||||
print(f"{'Audio throughput (audio duration/s):':<40}{result.audio_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Time to First Packet':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_TTFP (ms):':<40}{result.mean_ttfp_ms:<10.2f}")
|
||||
print(f"{'Median AUDIO_TTFP (ms):':<40}{result.median_ttfp_ms:<10.2f}")
|
||||
print(f"{'P99 AUDIO_TTFP (ms):':<40}{result.p99_ttfp_ms:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Real Time Factor':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_RTF:':<40}{result.mean_rtf:<10.3f}")
|
||||
print(f"{'Median AUDIO_RTF:':<40}{result.median_rtf:<10.3f}")
|
||||
print(f"{'P99 AUDIO_RTF:':<40}{result.p99_rtf:<10.3f}")
|
||||
print(f"{'=' * W}")
|
||||
print("")
|
||||
|
||||
# Save results (as a list with single concurrency=1 entry, matching serve format)
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump([asdict(result)], f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS HuggingFace Benchmark")
|
||||
parser.add_argument(
|
||||
"--model", type=str, default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", help="HuggingFace model name or path"
|
||||
)
|
||||
parser.add_argument("--num-prompts", type=int, default=50)
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--gpu-device", type=int, default=0)
|
||||
parser.add_argument("--voice", type=str, default="Vivian")
|
||||
parser.add_argument("--language", type=str, default="English")
|
||||
parser.add_argument(
|
||||
"--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
|
||||
)
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
parser.add_argument("--save-audio", action="store_true", help="Save generated audio files")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
run_benchmark(args)
|
||||
@@ -0,0 +1,348 @@
|
||||
"""Benchmark client for Qwen3-TTS via /v1/audio/speech endpoint.
|
||||
|
||||
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF (Real-Time Factor)
|
||||
across configurable concurrency levels. Saves results as JSON for plotting.
|
||||
|
||||
Usage:
|
||||
python bench_tts_serve.py \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 4 10 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestResult:
|
||||
success: bool = False
|
||||
ttfp: float = 0.0 # Time to first audio packet (seconds)
|
||||
e2e: float = 0.0 # End-to-end latency (seconds)
|
||||
audio_bytes: int = 0 # Total audio bytes received
|
||||
audio_duration: float = 0.0 # Audio duration in seconds (estimated from PCM)
|
||||
rtf: float = 0.0 # Real-time factor = e2e / audio_duration
|
||||
prompt: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
config_name: str = ""
|
||||
concurrency: int = 0
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
# TTFP stats (ms)
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
# E2E stats (ms)
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
# RTF stats
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
p99_rtf: float = 0.0
|
||||
# Audio stats
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0 # audio_duration / wall_time
|
||||
request_throughput: float = 0.0 # requests / second
|
||||
# Per-request details
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
|
||||
"""Convert raw PCM byte count to duration in seconds."""
|
||||
num_samples = num_bytes / sample_width
|
||||
return num_samples / sample_rate
|
||||
|
||||
|
||||
async def send_tts_request(
|
||||
session: aiohttp.ClientSession,
|
||||
api_url: str,
|
||||
prompt: str,
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
pbar: tqdm | None = None,
|
||||
) -> RequestResult:
|
||||
"""Send a streaming TTS request and measure latency metrics."""
|
||||
payload = {
|
||||
"input": prompt,
|
||||
"voice": voice,
|
||||
"language": language,
|
||||
"stream": True,
|
||||
"response_format": "pcm",
|
||||
}
|
||||
|
||||
result = RequestResult(prompt=prompt)
|
||||
st = time.perf_counter()
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status != 200:
|
||||
result.error = f"HTTP {response.status}: {await response.text()}"
|
||||
result.success = False
|
||||
return result
|
||||
|
||||
first_chunk = True
|
||||
total_bytes = 0
|
||||
|
||||
async for chunk in response.content.iter_any():
|
||||
if first_chunk and len(chunk) > 0:
|
||||
result.ttfp = time.perf_counter() - st
|
||||
first_chunk = False
|
||||
total_bytes += len(chunk)
|
||||
|
||||
result.e2e = time.perf_counter() - st
|
||||
result.audio_bytes = total_bytes
|
||||
result.audio_duration = pcm_bytes_to_duration(total_bytes)
|
||||
|
||||
if result.audio_duration > 0:
|
||||
result.rtf = result.e2e / result.audio_duration
|
||||
result.success = True
|
||||
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.success = False
|
||||
result.e2e = time.perf_counter() - st
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return result
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
host: str,
|
||||
port: int,
|
||||
num_prompts: int,
|
||||
max_concurrency: int,
|
||||
num_warmups: int = 3,
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
) -> BenchmarkResult:
|
||||
"""Run benchmark at a given concurrency level."""
|
||||
api_url = f"http://{host}:{port}/v1/audio/speech"
|
||||
|
||||
connector = aiohttp.TCPConnector(
|
||||
limit=max_concurrency,
|
||||
limit_per_host=max_concurrency,
|
||||
keepalive_timeout=60,
|
||||
)
|
||||
session = aiohttp.ClientSession(
|
||||
connector=connector,
|
||||
timeout=aiohttp.ClientTimeout(total=600),
|
||||
)
|
||||
|
||||
# Warmup
|
||||
if num_warmups > 0:
|
||||
print(f" Warming up with {num_warmups} requests...")
|
||||
warmup_tasks = []
|
||||
for i in range(num_warmups):
|
||||
prompt = PROMPTS[i % len(PROMPTS)]
|
||||
warmup_tasks.append(send_tts_request(session, api_url, prompt, voice, language))
|
||||
await asyncio.gather(*warmup_tasks)
|
||||
print(" Warmup done.")
|
||||
|
||||
# Build request list
|
||||
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
|
||||
|
||||
# Run benchmark
|
||||
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
|
||||
semaphore = asyncio.Semaphore(max_concurrency)
|
||||
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
|
||||
|
||||
async def limited_request(prompt):
|
||||
async with semaphore:
|
||||
return await send_tts_request(session, api_url, prompt, voice, language, pbar)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
|
||||
results: list[RequestResult] = await asyncio.gather(*tasks)
|
||||
duration = time.perf_counter() - start_time
|
||||
pbar.close()
|
||||
|
||||
await session.close()
|
||||
|
||||
# Compute stats
|
||||
successful = [r for r in results if r.success]
|
||||
failed = [r for r in results if not r.success]
|
||||
|
||||
bench = BenchmarkResult(
|
||||
concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
completed=len(successful),
|
||||
failed=len(failed),
|
||||
duration_s=duration,
|
||||
)
|
||||
|
||||
if successful:
|
||||
ttfps = [r.ttfp * 1000 for r in successful] # convert to ms
|
||||
e2es = [r.e2e * 1000 for r in successful]
|
||||
rtfs = [r.rtf for r in successful]
|
||||
audio_durs = [r.audio_duration for r in successful]
|
||||
|
||||
bench.mean_ttfp_ms = float(np.mean(ttfps))
|
||||
bench.median_ttfp_ms = float(np.median(ttfps))
|
||||
bench.std_ttfp_ms = float(np.std(ttfps))
|
||||
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
|
||||
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
|
||||
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
|
||||
|
||||
bench.mean_e2e_ms = float(np.mean(e2es))
|
||||
bench.median_e2e_ms = float(np.median(e2es))
|
||||
bench.std_e2e_ms = float(np.std(e2es))
|
||||
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
|
||||
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
|
||||
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
|
||||
|
||||
bench.mean_rtf = float(np.mean(rtfs))
|
||||
bench.median_rtf = float(np.median(rtfs))
|
||||
bench.std_rtf = float(np.std(rtfs))
|
||||
bench.p99_rtf = float(np.percentile(rtfs, 99))
|
||||
|
||||
bench.mean_audio_duration_s = float(np.mean(audio_durs))
|
||||
bench.total_audio_duration_s = float(np.sum(audio_durs))
|
||||
bench.audio_throughput = bench.total_audio_duration_s / duration
|
||||
bench.request_throughput = len(successful) / duration
|
||||
|
||||
bench.per_request = [
|
||||
{
|
||||
"ttfp_ms": r.ttfp * 1000,
|
||||
"e2e_ms": r.e2e * 1000,
|
||||
"rtf": r.rtf,
|
||||
"audio_duration_s": r.audio_duration,
|
||||
"prompt": r.prompt,
|
||||
}
|
||||
for r in successful
|
||||
]
|
||||
|
||||
# Print summary in standardized performance template
|
||||
W = 50
|
||||
print("")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Serving Benchmark Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Successful requests:':<40}{bench.completed:<10}")
|
||||
print(f"{'Failed requests:':<40}{bench.failed:<10}")
|
||||
print(f"{'Maximum request concurrency:':<40}{max_concurrency:<10}")
|
||||
print(f"{'Benchmark duration (s):':<40}{duration:<10.2f}")
|
||||
print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'End-to-end Latency':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
|
||||
print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
|
||||
print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Audio Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
|
||||
print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Time to First Packet':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
|
||||
print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
|
||||
print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Real Time Factor':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
|
||||
print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
|
||||
print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
|
||||
print(f"{'=' * W}")
|
||||
print("")
|
||||
|
||||
if failed:
|
||||
for r in failed[:3]:
|
||||
print(f" [ERROR] {r.error[:200]}")
|
||||
|
||||
return bench
|
||||
|
||||
|
||||
async def main(args):
|
||||
all_results = []
|
||||
|
||||
for concurrency in args.max_concurrency:
|
||||
result = await run_benchmark(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
num_prompts=args.num_prompts,
|
||||
max_concurrency=concurrency,
|
||||
num_warmups=args.num_warmups,
|
||||
voice=args.voice,
|
||||
language=args.language,
|
||||
)
|
||||
result.config_name = args.config_name
|
||||
all_results.append(asdict(result))
|
||||
|
||||
# Save results
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump(all_results, f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
return all_results
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS Benchmark Client")
|
||||
parser.add_argument("--host", type=str, default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--num-prompts", type=int, default=50, help="Number of prompts per concurrency level")
|
||||
parser.add_argument( # noqa: E501
|
||||
"--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
|
||||
)
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--voice", type=str, default="vivian")
|
||||
parser.add_argument("--language", type=str, default="English")
|
||||
parser.add_argument(
|
||||
"--config-name", type=str, default="async_chunk", help="Label for this config (used in filenames)"
|
||||
)
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
asyncio.run(main(args))
|
||||
@@ -0,0 +1,93 @@
|
||||
# Qwen3-TTS batch_size=1 config (streaming with async_chunk)
|
||||
# 2-stage pipeline: Talker -> Code2Wav
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_comprehension: true
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 1
|
||||
model_stage: qwen3_tts
|
||||
model_arch: Qwen3TTSTalkerForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
enforce_eager: false
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent
|
||||
gpu_memory_utilization: 0.3
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 512
|
||||
max_model_len: 4096
|
||||
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
detokenize: false
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [2150]
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 1
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen3TTSCode2Wav
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio
|
||||
gpu_memory_utilization: 0.3
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 8192
|
||||
max_model_len: 32768
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
tts_args:
|
||||
max_instructions_length: 500
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
detokenize: true
|
||||
repetition_penalty: 1.0
|
||||
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 1
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 25
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
@@ -0,0 +1,94 @@
|
||||
# Qwen3-TTS max_num_seqs=16 config (streaming with async_chunk)
|
||||
# High-throughput concurrent request processing
|
||||
# 2-stage pipeline: Talker -> Code2Wav
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_comprehension: true
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 16
|
||||
model_stage: qwen3_tts
|
||||
model_arch: Qwen3TTSTalkerForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
enforce_eager: false
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent
|
||||
gpu_memory_utilization: 0.3
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 4096
|
||||
max_model_len: 4096
|
||||
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
detokenize: false
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [2150]
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 16
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen3TTSCode2Wav
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio
|
||||
gpu_memory_utilization: 0.2
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 16384
|
||||
max_model_len: 32768
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
tts_args:
|
||||
max_instructions_length: 500
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
detokenize: true
|
||||
repetition_penalty: 1.0
|
||||
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 16
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 25
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
@@ -0,0 +1,94 @@
|
||||
# Qwen3-TTS batch_size=4 config (streaming with async_chunk)
|
||||
# Enables concurrent request processing
|
||||
# 2-stage pipeline: Talker -> Code2Wav
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_comprehension: true
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 4
|
||||
model_stage: qwen3_tts
|
||||
model_arch: Qwen3TTSTalkerForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
enforce_eager: false
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent
|
||||
gpu_memory_utilization: 0.3
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 512
|
||||
max_model_len: 4096
|
||||
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
detokenize: false
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [2150]
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 4
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen3TTSCode2Wav
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio
|
||||
gpu_memory_utilization: 0.2
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 8192
|
||||
max_model_len: 32768
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
tts_args:
|
||||
max_instructions_length: 500
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
detokenize: true
|
||||
repetition_penalty: 1.0
|
||||
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 4
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 25
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
|
||||
ARG VLLM_BASE_TAG=v0.17.0
|
||||
ARG VLLM_BASE_TAG=v0.18.0
|
||||
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
|
||||
ARG APP_DIR=/workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
@@ -7,7 +7,7 @@ COPY . .
|
||||
|
||||
# Install system dependencies
|
||||
RUN apt-get update && \
|
||||
apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
|
||||
@@ -1,17 +1,17 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.14.0rc1
|
||||
ARG VLLM_ASCEND_TAG=v0.17.0rc1
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
WORKDIR /vllm-workspace/vllm-ascend
|
||||
RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
|
||||
RUN pip install -v -e .
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
|
||||
COPY . .
|
||||
|
||||
RUN pip install -v -e .
|
||||
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
|
||||
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
|
||||
@@ -1,17 +1,17 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
|
||||
ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
WORKDIR /vllm-workspace/vllm-ascend
|
||||
RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
|
||||
RUN pip install -v -e .
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
|
||||
COPY . .
|
||||
|
||||
RUN pip install -v -e .
|
||||
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
|
||||
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.17.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.18.0
|
||||
FROM ${BASE_IMAGE} AS base
|
||||
|
||||
# Declare a variable to know if we want to use the nightly build or the stable build.
|
||||
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
|
||||
# we should swap over to use stable release ASAP.
|
||||
# We should avoid relying on custom commits.
|
||||
ARG USE_NIGHTLY_BUILD=0
|
||||
ARG VLLM_VERSION_OR_COMMIT_HASH=2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
|
||||
ARG VLLM_VERSION_OR_COMMIT_HASH=89138b21cc246ae944c741d5c399c148e2b770ab
|
||||
ARG ARG_PYTORCH_ROCM_ARCH
|
||||
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
|
||||
|
||||
@@ -19,7 +19,7 @@ WORKDIR ${COMMON_WORKDIR}
|
||||
|
||||
# Step 1: Setup - Install system dependencies
|
||||
RUN apt-get update && \
|
||||
apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
|
||||
@@ -14,6 +14,7 @@ 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 \
|
||||
@@ -75,7 +76,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
ENV UV_LINK_MODE="copy"
|
||||
|
||||
ARG VLLM_VERSION=v0.16.0
|
||||
ARG VLLM_VERSION=v0.18.0
|
||||
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
|
||||
@@ -13,9 +13,14 @@ nav:
|
||||
- examples/README.md
|
||||
- Offline Inference:
|
||||
- BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
|
||||
- CosyVoice3: user_guide/examples/offline_inference/cosyvoice3.md
|
||||
- Fish Speech S2 Pro: user_guide/examples/offline_inference/fish_speech.md
|
||||
- GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
|
||||
- Helios Video Generation: user_guide/examples/offline_inference/helios.md
|
||||
- HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
|
||||
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
|
||||
- Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
|
||||
- MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
|
||||
- MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
|
||||
- Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
|
||||
- Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
|
||||
@@ -23,8 +28,11 @@ nav:
|
||||
- Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
|
||||
- Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
|
||||
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
|
||||
- X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
|
||||
- Online Serving:
|
||||
- BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
|
||||
- vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
|
||||
- Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md
|
||||
- GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
|
||||
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md
|
||||
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md
|
||||
@@ -50,7 +58,9 @@ nav:
|
||||
- Quantization:
|
||||
- Overview: user_guide/diffusion/quantization/overview.md
|
||||
- FP8: user_guide/diffusion/quantization/fp8.md
|
||||
- Int8: user_guide/diffusion/quantization/int8.md
|
||||
- GGUF: user_guide/diffusion/quantization/gguf.md
|
||||
- Step Execution: user_guide/diffusion/step_execution.md
|
||||
- Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
|
||||
- CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
|
||||
- LoRA: user_guide/diffusion/lora.md
|
||||
@@ -82,6 +92,7 @@ nav:
|
||||
- design/feature/cache_dit.md
|
||||
- design/feature/teacache.md
|
||||
- design/feature/async_chunk_design.md
|
||||
- design/feature/diffusion_step_execution.md
|
||||
- Module Design:
|
||||
- design/module/ar_module.md
|
||||
- design/module/dit_module.md
|
||||
@@ -92,6 +103,7 @@ nav:
|
||||
- api/vllm_omni
|
||||
- CLI Reference: cli
|
||||
- Community:
|
||||
- Governance: community/governance.md
|
||||
- community/*
|
||||
- Slack: https://slack.vllm.ai
|
||||
- Blog: https://blog.vllm.ai
|
||||
|
||||
@@ -60,5 +60,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
|
||||
|
||||
For more information, checkout the following:
|
||||
|
||||
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
|
||||
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1)
|
||||
- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
|
||||
@@ -6,19 +6,13 @@ Main entry points for vLLM-Omni inference and serving.
|
||||
|
||||
- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
|
||||
- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
|
||||
- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
|
||||
- [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
|
||||
- [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
|
||||
- [vllm_omni.entrypoints.omni.Omni][]
|
||||
- [vllm_omni.entrypoints.omni.OmniBase][]
|
||||
- [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]
|
||||
- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
|
||||
- [vllm_omni.entrypoints.omni_stage.OmniStage][]
|
||||
- [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
|
||||
- [vllm_omni.entrypoints.zmq_utils.ZmqQueue][]
|
||||
- [vllm_omni.entrypoints.omni_base.OmniBase][]
|
||||
|
||||
## Inputs
|
||||
|
||||
@@ -48,9 +42,7 @@ Engine classes for offline and online inference.
|
||||
- [vllm_omni.engine.OmniEngineCoreOutputs][]
|
||||
- [vllm_omni.engine.OmniEngineCoreRequest][]
|
||||
- [vllm_omni.engine.PromptEmbedsPayload][]
|
||||
- [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
|
||||
- [vllm_omni.engine.arg_utils.OmniEngineArgs][]
|
||||
- [vllm_omni.engine.input_processor.OmniInputProcessor][]
|
||||
- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
|
||||
- [vllm_omni.engine.output_processor.OmniRequestState][]
|
||||
|
||||
|
||||
|
Before Width: | Height: | Size: 125 KiB After Width: | Height: | Size: 97 KiB |
@@ -0,0 +1,52 @@
|
||||
# Governance
|
||||
|
||||
vLLM-Omni's governance is inspired by the [vLLM governance process](https://docs.vllm.ai/en/latest/governance/process/). We share the same commitment to open source community and meritocratic norms.
|
||||
|
||||
## Values
|
||||
|
||||
vLLM-Omni aims to be the fastest and easiest-to-use omni-modality inference and serving engine. Our values are aligned with [vLLM's values](https://docs.vllm.ai/en/latest/governance/process/#values):
|
||||
|
||||
### Design Values
|
||||
|
||||
1. **Top performance**: System performance is our top priority. We monitor overheads, optimize kernels, and publish benchmarks. We never leave performance on the table.
|
||||
2. **Ease of use**: vLLM-Omni must be simple to install, configure, and operate. We provide clear documentation, fast startup, clean logs, helpful error messages, and monitoring guides. Many users fork our code or study it deeply, so we keep it readable and modular.
|
||||
3. **Wide coverage**: vLLM-Omni supports frontier models and high-performance accelerators. We make it easy to add new models and hardware. vLLM-Omni + PyTorch form a simple interface that avoids complexity.
|
||||
4. **Production ready**: vLLM-Omni runs 24/7 in production. It must be easy to operate and monitor for health issues.
|
||||
5. **Extensibility**: vLLM-Omni serves as fundamental omni-modality infrastructure. Our codebase cannot cover every use case, so we design for easy forking and customization.
|
||||
|
||||
### Collaboration Values
|
||||
|
||||
1. **Tightly Knit and Fast-Moving**: Our maintainer team is aligned on vision, philosophy, and roadmap. We work closely to unblock each other and move quickly.
|
||||
2. **Individual Merit**: No one buys their way into governance. Committer status belongs to individuals, not companies. We reward contribution, maintenance, and project stewardship.
|
||||
|
||||
## Project Maintainers
|
||||
|
||||
### Lead Maintainers
|
||||
|
||||
Lead maintainers are responsible for the overall direction and strategy of the project:
|
||||
|
||||
- [@Gaohan123](https://github.com/Gaohan123)
|
||||
- [@hsliuustc0106](https://github.com/hsliuustc0106)
|
||||
- [@ywang96](https://github.com/ywang96)
|
||||
|
||||
### Active Committers
|
||||
|
||||
Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
|
||||
|
||||
- [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
|
||||
- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
|
||||
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
|
||||
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
|
||||
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
|
||||
- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
|
||||
- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
|
||||
- [@wtomin](https://github.com/wtomin):
|
||||
- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
|
||||
|
||||
## Meetings
|
||||
|
||||
Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
|
||||
|
||||
## Committer Nomination Process
|
||||
|
||||
Every month, any active committer can nominate new committer(s) to the project. Up to **two new committers** will be admitted per month based on the quality and impact of their contributions.
|
||||
@@ -14,6 +14,10 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
|
||||
|
||||
- **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
|
||||
|
||||
## Multi-Stage Recipes
|
||||
|
||||
- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
|
||||
|
||||
## Optimization Features
|
||||
|
||||
- **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
|
||||
|
||||
@@ -158,7 +158,7 @@ The stage initialized but ran out of memory during processing.
|
||||
|
||||
**Solutions:**
|
||||
1. Reduce `max_num_batched_tokens`
|
||||
2. Reduce `max_batch_size` in runtime config
|
||||
2. Reduce `max_num_seqs` in engine_args
|
||||
3. Lower `gpu_memory_utilization` slightly
|
||||
4. Enable quantization if supported
|
||||
|
||||
|
||||
@@ -0,0 +1,177 @@
|
||||
# Prefill-Decode (PD) Disaggregation
|
||||
|
||||
PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
|
||||
stages so prompt processing and token generation can run on different workers.
|
||||
|
||||
This is documented as a stage-config recipe instead of a bundled YAML because the
|
||||
deployment-specific values usually change per environment:
|
||||
|
||||
- GPU placement
|
||||
- `tensor_parallel_size`
|
||||
- connector backend and connector ports
|
||||
- connector IPs or bootstrap addresses
|
||||
|
||||
Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml)
|
||||
and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
|
||||
changes below.
|
||||
|
||||
## Requirements
|
||||
|
||||
- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
|
||||
- A KV connector supported by vLLM, such as `MooncakeConnector`
|
||||
- Matching `tensor_parallel_size` on the prefill and decode thinker stages
|
||||
|
||||
## 1. Split the thinker into prefill and decode stages
|
||||
|
||||
Replace the original thinker stage with two stages:
|
||||
|
||||
```yaml
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_prefill_only: true
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 16
|
||||
model_stage: thinker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.9
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 1
|
||||
kv_transfer_config:
|
||||
kv_connector: "MooncakeConnector"
|
||||
kv_role: "kv_producer"
|
||||
kv_rank: 0
|
||||
kv_parallel_size: 2
|
||||
kv_connector_extra_config:
|
||||
mooncake_bootstrap_port: 25201
|
||||
final_output: false
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
is_decode_only: true
|
||||
runtime:
|
||||
devices: "1"
|
||||
engine_args:
|
||||
max_num_seqs: 64
|
||||
model_stage: thinker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.9
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 1
|
||||
kv_transfer_config:
|
||||
kv_connector: "MooncakeConnector"
|
||||
kv_role: "kv_consumer"
|
||||
kv_rank: 1
|
||||
kv_parallel_size: 2
|
||||
kv_connector_extra_config:
|
||||
mooncake_bootstrap_port: 25202
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
```
|
||||
|
||||
Notes:
|
||||
|
||||
- `is_prefill_only: true` marks the thinker stage that only saves KV.
|
||||
- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
|
||||
- `kv_transfer_config` is required on both stages.
|
||||
- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
|
||||
prefill side only processes the prompt and exports KV.
|
||||
|
||||
## 2. Shift the downstream stages by one index
|
||||
|
||||
After inserting the extra thinker stage, renumber the remaining stages:
|
||||
|
||||
```yaml
|
||||
- stage_id: 2
|
||||
runtime:
|
||||
devices: "2"
|
||||
engine_input_source: [1]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
|
||||
|
||||
- stage_id: 3
|
||||
runtime:
|
||||
devices: "2"
|
||||
engine_args:
|
||||
max_num_seqs: 1
|
||||
engine_input_source: [2]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
|
||||
```
|
||||
|
||||
Compared with the default Qwen3-Omni config:
|
||||
|
||||
- the talker becomes stage `2` instead of stage `1`
|
||||
- the code2wav stage becomes stage `3` instead of stage `2`
|
||||
- the talker now reads from decode stage `1`
|
||||
|
||||
## 3. Add runtime edges for the four-stage pipeline
|
||||
|
||||
```yaml
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 1
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
- from: 1
|
||||
to: 2
|
||||
window_size: -1
|
||||
- from: 2
|
||||
to: 3
|
||||
window_size: -1
|
||||
```
|
||||
|
||||
## 4. Launch with your custom config
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--stage-configs-path /path/to/qwen3_omni_pd.yaml
|
||||
```
|
||||
|
||||
## Operational Notes
|
||||
|
||||
- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
|
||||
pair. Keep prefill and decode at the same `tensor_parallel_size`.
|
||||
- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
|
||||
separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
|
||||
- Choose connector ports and addresses that match your deployment. The values
|
||||
shown above are examples only.
|
||||
@@ -18,7 +18,7 @@ If users want to modify some part of it. The custom stage_configs file can be in
|
||||
For offline (Assume necessary dependencies have ben imported):
|
||||
```python
|
||||
model_name = "Qwen/Qwen2.5-Omni-7B"
|
||||
omni_llm = OmniLLM(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
|
||||
omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
|
||||
```
|
||||
|
||||
For online serving:
|
||||
@@ -30,15 +30,15 @@ vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to
|
||||
|
||||
Below is a specific example of stage_configs.yaml in Qwen2.5-omni.
|
||||
```python
|
||||
# stage config for running qwen2.5-omni with architecture of OmniLLM.
|
||||
# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
|
||||
stage_args:
|
||||
- stage_id: 0 # mark the unique id for each stage
|
||||
runtime: # The disaggregated configuration
|
||||
process: true # Run this stage in a separate process
|
||||
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
|
||||
max_batch_size: 1 # the batch_size for offline inference
|
||||
engine_args: # Engine arguments for a certain engine
|
||||
model_stage: thinker
|
||||
max_num_seqs: 1
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
|
||||
worker_type: ar # The specific worker used
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used
|
||||
@@ -62,9 +62,9 @@ stage_args:
|
||||
runtime:
|
||||
process: true
|
||||
devices: "1"
|
||||
max_batch_size: 3
|
||||
engine_args:
|
||||
model_stage: talker
|
||||
max_num_seqs: 3
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
@@ -88,9 +88,9 @@ stage_args:
|
||||
runtime:
|
||||
process: true
|
||||
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
max_num_seqs: 1
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
@@ -159,9 +159,9 @@ Visible devices for this stage, specified as a string. This controls which GPU d
|
||||
|
||||
Default: `"0"`
|
||||
|
||||
#### `runtime.max_batch_size`
|
||||
#### `engine_args.max_num_seqs`
|
||||
|
||||
The maximum batch size for offline inference in this stage. This limits how many sequences can be processed together in a single batch during offline inference operations.
|
||||
The maximum number of sequences for concurrent processing in this stage. For LLM stages, this controls the vLLM scheduler's maximum concurrent sequences. For all stage types, this also controls how many tasks can be batched together in the task processing loop.
|
||||
|
||||
Default: `1`
|
||||
|
||||
|
||||
@@ -1,12 +1,12 @@
|
||||
# stage config for running qwen2.5-omni with architecture of OmniLLM.
|
||||
# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
runtime:
|
||||
process: true # Run this stage in a separate process
|
||||
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: thinker
|
||||
max_num_seqs: 1
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
@@ -30,9 +30,9 @@ stage_args:
|
||||
runtime:
|
||||
process: true
|
||||
devices: "1"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: talker
|
||||
max_num_seqs: 1
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
@@ -56,9 +56,9 @@ stage_args:
|
||||
runtime:
|
||||
process: true
|
||||
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
max_num_seqs: 1
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
|
||||
@@ -107,7 +107,7 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
|
||||
- `[CI/Build]` for build or continuous integration improvements.
|
||||
- `[Doc]` for documentation fixes and improvements.
|
||||
- `[Model]` for adding a new model or improving an existing model. Model name should appear in the title.
|
||||
- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `OmniLLM` class, etc.)
|
||||
- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `Omni`/`AsyncOmni`, etc.)
|
||||
- `[Kernel]` for changes affecting CUDA kernels or other compute kernels.
|
||||
- `[Core]` for changes in the core vLLM-Omni logic (e.g., `OmniProcessor`, `OmniARScheduler`, etc.)
|
||||
- `[Hardware][Vendor]` for hardware-specific changes. Vendor name should appear in the prefix, such as [Ascend] for Ascend NPUs.
|
||||
|
||||
@@ -67,8 +67,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<td>Online & Offline (multiple deployment scenarios):<br>real model, normal inference function, normal accuracy</td>
|
||||
<td><30min</td>
|
||||
<td>
|
||||
/tests/e2e/online_serving/test_{model_name}_expansion.py<br>
|
||||
/tests/e2e/offline_inference/test_{model_name}_expansion.py
|
||||
/tests/e2e/online_serving/test_{model_name}.py<br>
|
||||
/tests/e2e/offline_inference/test_{model_name}.py
|
||||
</td>
|
||||
<td>
|
||||
<a href="#chapter-2-l3-level-testing-core-integration-performance-and-accuracy-verification">Chapter 2</a><br>
|
||||
@@ -168,11 +168,6 @@ vllm_omni/ tests/
|
||||
│ └── arg_utils.py │ └── test_arg_utils.py ⬜
|
||||
│
|
||||
├── entrypoints/ → ├── entrypoints/
|
||||
│ ├── omni.py │ ├── test_omni.py ⬜ (E2E covered by e2e/offline, e2e/online)
|
||||
│ ├── omni_llm.py │ ├── test_omni_llm.py ✅
|
||||
│ ├── omni_stage.py │ ├── test_omni_stage.py ⬜ (partial in test_omni_stage_diffusion_config.py)
|
||||
│ ├── omni_diffusion.py │ ├── test_omni_diffusion.py ✅
|
||||
│ ├── async_omni.py │ ├── test_async_omni.py ✅ actually in e2e/online_serving/test_async_omni.py
|
||||
│ ├── async_omni_diffusion.py │ ├── test_async_omni_diffusion_config.py ✅
|
||||
│ ├── stage_utils.py │ ├── test_stage_utils.py ✅
|
||||
│ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve)
|
||||
@@ -508,7 +503,9 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
**2.4.5 Request Execution**
|
||||
|
||||
```python
|
||||
openai_client.send_request(request_config, request_num=1)
|
||||
openai_client.send_omni_request(request_config, request_num=1) # for omni-understanding models
|
||||
# or
|
||||
openai_client.send_diffusion_request(request_config, request_num=1) # for diffusion models
|
||||
```
|
||||
|
||||
**Explanation**:
|
||||
@@ -517,7 +514,7 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
|
||||
**Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n.
|
||||
|
||||
**Implicit Validation**: The send_request method internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
|
||||
**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.
|
||||
|
||||
- ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ Defined in `pyproject.toml`:
|
||||
| `gpu` | Tests that run on GPU * |
|
||||
| `cuda` | Tests that run on CUDA * |
|
||||
| `rocm` | Tests that run on AMD/ROCm * |
|
||||
| `xpu` | Tests that run on Intel XPU * |
|
||||
| `npu` | Tests that run on NPU/Ascend * |
|
||||
| `H100` | Tests that require H100 GPU * |
|
||||
| `L4` | Tests that require L4 GPU * |
|
||||
@@ -32,7 +33,7 @@ Defined in `pyproject.toml`:
|
||||
| `slow` | Slow tests (may skip in quick CI) |
|
||||
| `benchmark` | Benchmark tests |
|
||||
|
||||
\* Means those markers are auto-added, and they will be added by the `@hardware_test` decorator.
|
||||
\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility).
|
||||
|
||||
### Example usage for markers
|
||||
|
||||
@@ -49,12 +50,13 @@ from tests.utils import hardware_test
|
||||
def test_video_to_audio()
|
||||
...
|
||||
```
|
||||
|
||||
### Decorator: `@hardware_test`
|
||||
|
||||
This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions:
|
||||
|
||||
1. **Applies platform and resource markers**
|
||||
Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `A2`, `A3`).
|
||||
Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`).
|
||||
```
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.L4
|
||||
@@ -79,14 +81,14 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
|
||||
- Single call for multiple platforms:
|
||||
```python
|
||||
@hardware_test(
|
||||
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
|
||||
num_cards={"cuda": 2, "rocm": 2, "npu": 2},
|
||||
res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
|
||||
num_cards={"cuda": 2, "rocm": 2, "xpu": 2, "npu": 2},
|
||||
)
|
||||
```
|
||||
or
|
||||
```python
|
||||
@hardware_test(
|
||||
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
|
||||
res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
|
||||
num_cards=2,
|
||||
)
|
||||
```
|
||||
@@ -98,6 +100,24 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
|
||||
- ROCm only: `pytest -m "distributed_rocm and MI325"`
|
||||
- NPU only: `pytest -m "distributed_npu"`
|
||||
|
||||
### Function: `hardware_marks`
|
||||
|
||||
`hardware_marks` returns a list of pytest mark objects with the same signature as `@hardware_test`. Use it when you need more flexibility, such as attaching hardware marks to individual `pytest.param` entries rather than an entire test function.
|
||||
|
||||
```python
|
||||
from tests.utils import hardware_marks
|
||||
|
||||
MULTI_CARD_MARKS = hardware_marks(
|
||||
res={"cuda": "H100", "rocm": "MI325", "npu": "A2"}, num_cards=2
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("omni_server", [
|
||||
pytest.param(OmniServerParams(...), id="case_001", marks=MULTI_CARD_MARKS),
|
||||
], indirect=True)
|
||||
def test_feature(omni_server):
|
||||
...
|
||||
```
|
||||
|
||||
## Add Support for a New Platform
|
||||
|
||||
If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps:
|
||||
@@ -132,12 +152,13 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
|
||||
# Optionally: add skipif_tpu when implemented
|
||||
return [test_platform, test_resource, test_distributed]
|
||||
```
|
||||
3. **Update `hardware_test` to recognize your new platform**:
|
||||
In the relevant place (see the `hardware_test` implementation), add:
|
||||
3. **Update `hardware_marks` to recognize your new platform**:
|
||||
In the relevant place (see the `hardware_marks` implementation), add:
|
||||
```python
|
||||
if platform == "tpu":
|
||||
marks = tpu_marks(res=resource, num_cards=cards)
|
||||
```
|
||||
(`hardware_test` calls `hardware_marks` internally, so both will pick up the change.)
|
||||
4. **(Recommended) Add a test using your new markers**:
|
||||
```python
|
||||
@hardware_test(
|
||||
@@ -151,7 +172,7 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
|
||||
**Summary**:
|
||||
- Add pytest markers for your new platform/resources
|
||||
- Implement a marker function (`xxx_marks`)
|
||||
- Plug into `hardware_test`
|
||||
- You're done: tests decorated with `@hardware_test` using your platform now automatically get the correct markers, distribution, and isolation!
|
||||
- Plug into `hardware_marks`
|
||||
- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation!
|
||||
|
||||
See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).
|
||||
@@ -24,8 +24,6 @@ End-to-end tests verify the complete functionality of a system or component. For
|
||||
|
||||
- **`tests/e2e/online_serving/`**: Tests for online serving scenarios (e.g., API server tests)
|
||||
|
||||
**Example:** The test file for `vllm_omni/entrypoints/omni_llm.py` should be located at `tests/entrypoints/test_omni_llm.py`.
|
||||
|
||||
## Test Directory Structure
|
||||
|
||||
The ideal directory structure mirrors the source code organization. Legend: `✅` = test exists, `⬜` = suggested to add.
|
||||
@@ -75,11 +73,6 @@ vllm_omni/ tests/
|
||||
│ └── arg_utils.py │ └── test_arg_utils.py ⬜
|
||||
│
|
||||
├── entrypoints/ → ├── entrypoints/
|
||||
│ ├── omni.py │ ├── test_omni.py ⬜ (E2E covered by e2e/offline, e2e/online)
|
||||
│ ├── omni_llm.py │ ├── test_omni_llm.py ✅
|
||||
│ ├── omni_stage.py │ ├── test_omni_stage.py ⬜ (partial in test_omni_stage_diffusion_config.py)
|
||||
│ ├── omni_diffusion.py │ ├── test_omni_diffusion.py ✅
|
||||
│ ├── async_omni.py │ ├── test_async_omni.py ✅ actually in e2e/online_serving/test_async_omni.py
|
||||
│ ├── async_omni_diffusion.py │ ├── test_async_omni_diffusion_config.py ✅
|
||||
│ ├── stage_utils.py │ ├── test_stage_utils.py ✅
|
||||
│ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve)
|
||||
@@ -158,6 +151,7 @@ vllm_omni/ tests/
|
||||
├── test_diffusion_layerwise_offload.py
|
||||
├── test_diffusion_lora.py
|
||||
├── test_sequence_parallel.py
|
||||
├── test_qwen_image_edit_expansion.py
|
||||
└── stage_configs/
|
||||
├── qwen2_5_omni_ci.yaml
|
||||
├── qwen3_omni_ci.yaml
|
||||
@@ -169,7 +163,7 @@ vllm_omni/ tests/
|
||||
|
||||
### Naming Conventions
|
||||
|
||||
- **Unit Tests**: Use `test_<module_name>.py` format. Example: `omni_llm.py` → `test_omni_llm.py`
|
||||
- **Unit Tests**: Use `test_<module_name>.py` format. Example: `stage_utils.py` → `test_stage_utils.py`
|
||||
|
||||
- **E2E Tests**: Place in `tests/e2e/offline_inference/` or `tests/e2e/online_serving/` with descriptive names. Example: `tests/e2e/offline_inference/test_qwen3_omni.py`, `tests/e2e/offline_inference/test_diffusion_model.py`
|
||||
|
||||
|
||||
@@ -651,8 +651,27 @@ For a fair comparison, keep the same **prompt**, **seed**, **resolution**, **num
|
||||
|
||||
#### 5.2 Functionality Check in CI
|
||||
|
||||
To ensure project maintainability and sustainable development, we encourage contributors to submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
|
||||
For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
|
||||
To ensure project maintainability and sustainable development, please submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
|
||||
|
||||
For comprehensive testing guidelines and the definition of test levels (L1-L5), please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
|
||||
The following tests are required to add:
|
||||
|
||||
- L4 test of the model's full *functionality* (i.e., all the *diffusion features* that are supported by this model), including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
|
||||
- Test cases: Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`. The following test cases shall cover all features:
|
||||
- 1 GPU: TeaCache & GGUF (or fallback to FP8, or disable it) & Layer-wise CPU offloading (or fallback to Module-wise)
|
||||
- 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Ulysses = 2
|
||||
- 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & Ring = 2
|
||||
- 2 GPUs: TeaCache & FP8 (or fallback to GGUF, or disable it) & CFG Parallel = 2
|
||||
- 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Tensor Parallel = 2 & VAE Patch Parallel = 2
|
||||
- 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & HSDP = 2 & VAE Patch Parallel = 2
|
||||
- Validation: test that the multimodal output files of your model have the correct shapes.
|
||||
- Test marks: always add `advanced_model` and `diffusion`. Add `parallel` and GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/)
|
||||
- To maximize code reuse, you may refer to
|
||||
- `tests/conftest.py` for `omni_server` and `openai_client` fixtures, `generate_synthetic_image` and `assert_XXX_valid` helper.
|
||||
- `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
|
||||
- [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
|
||||
- Doc: add a concise dostring for each test function.
|
||||
- Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).
|
||||
|
||||
---
|
||||
|
||||
@@ -720,6 +739,22 @@ See detailed guide: [How to add Sequence Parallel support](../../design/feature/
|
||||
omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2)
|
||||
```
|
||||
|
||||
### Step Execution
|
||||
|
||||
See detailed design guide: [How to add step execution support](../../design/feature/diffusion_step_execution.md)
|
||||
|
||||
Use this only when your pipeline can be split into stable request-scoped and
|
||||
step-scoped phases. The reference implementation is
|
||||
`QwenImagePipeline`, which maps its request-level `forward()` into:
|
||||
|
||||
1. `prepare_encode()` for prompt encoding, latent init, timestep prep, and per-request scheduler setup.
|
||||
2. `denoise_step()` for one transformer/noise prediction.
|
||||
3. `step_scheduler()` for one scheduler update and `step_index` advance.
|
||||
4. `post_decode()` for the final VAE decode.
|
||||
|
||||
Do not enable `step_execution=True` until those four methods are implemented
|
||||
and validated against the request-level path.
|
||||
|
||||
### Cache Acceleration
|
||||
|
||||
#### TeaCache
|
||||
@@ -796,6 +831,107 @@ class WanTransformer3DModel(nn.Module):
|
||||
|
||||
---
|
||||
|
||||
### Diffusion Timing (Performance Profiling)
|
||||
When adapting a new diffusion model, it is often useful to analyze the latency of key components such as text encoding, diffusion denoising, and VAE decoding.
|
||||
vLLM-Omni provides a timing utility via `DiffusionPipelineProfilerMixin` to help developers quickly identify performance bottlenecks.
|
||||
|
||||
!!! info
|
||||
`DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md#3-profiling-diffusion-models). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
|
||||
|
||||
This tool automatically measures the execution time of selected pipeline modules and prints the results in the logs.
|
||||
|
||||
**Enabling Diffusion Timing**
|
||||
|
||||
|
||||
Enable timing by setting:
|
||||
```
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --enable-diffusion-pipeline-profiler
|
||||
```
|
||||
You can optionally specify which modules to profile:
|
||||
```
|
||||
class YourPipeline(xxx, DiffusionPipelineProfilerMixin):
|
||||
def __init__(self, xxx):
|
||||
...
|
||||
self.setup_diffusion_pipeline_profiler(profiler_targets=["diffuse"], enable_diffusion_pipeline_profiler)
|
||||
```
|
||||
If not specified, the default targets are used:
|
||||
```
|
||||
["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"]
|
||||
```
|
||||
**Adding DiffusionPipelineProfilerMixin to a Pipeline**
|
||||
To enable timing support in your pipeline, inherit from DiffusionPipelineProfilerMixin.
|
||||
```python
|
||||
from vllm_omni.diffusion.utils.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
|
||||
|
||||
class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
# Optional: Specify custom timing targets
|
||||
_PROFILER_TARGETS = ["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
od_config: OmniDiffusionConfig,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
self.od_config = od_config
|
||||
self.parallel_config = od_config.parallel_config
|
||||
# initialize pipeline components
|
||||
...
|
||||
|
||||
# initialize timing profiler
|
||||
self.setup_diffusion_pipeline_profiler(enable_diffusion_pipeline_profiler)
|
||||
```
|
||||
The mixin dynamically wraps selected methods and records their execution time during inference.
|
||||
|
||||
If you need to fetch the execution time of different modules, you will need to pass `self.stage_durations` to `DiffusionOutput`, as shown below:
|
||||
|
||||
```diff
|
||||
- return DiffusionOutput(output=img)
|
||||
+ return DiffusionOutput(
|
||||
output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
|
||||
)
|
||||
```
|
||||
|
||||
**Pipeline Design for Timing**
|
||||
The current diffusion timing utility is function-based, meaning it measures the execution time of individual methods.
|
||||
|
||||
When implementing a new pipeline, avoid putting all logic inside a single function (e.g., forward). Instead, structure the pipeline in a modular way by separating key stages into independent methods, such as the diffusion loop.
|
||||
|
||||
For example:
|
||||
```
|
||||
def forward(self, req: OmniDiffusionRequest):
|
||||
prompt_embeds = self.encode_prompt(req)
|
||||
latents = self.diffuse(prompt_embeds, req)
|
||||
images = self.vae.decode(latents)
|
||||
return DiffusionOutput(output=images)
|
||||
```
|
||||
This allows the timing utility to measure each stage (e.g., encode_prompt, diffuse, vae.decode) separately and helps identify performance bottlenecks more easily.
|
||||
|
||||
|
||||
**Default Profiled Modules**
|
||||
|
||||
By default, the following pipeline modules are timed:
|
||||
```
|
||||
vae.encode
|
||||
vae.decode
|
||||
diffuse
|
||||
text_encoder.forward
|
||||
tokenizer.forward
|
||||
```
|
||||
|
||||
**Example Output**
|
||||
|
||||
When enabled, timing logs appear like this:
|
||||
```
|
||||
[DiffusionTiming] text_encoder.forward took 0.018s
|
||||
[DiffusionTiming] diffuse took 2.412s
|
||||
[DiffusionTiming] vae.decode took 0.063s
|
||||
```
|
||||
These measurements help identify bottlenecks during model adaptation and optimization
|
||||
|
||||
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
|
||||
@@ -847,7 +983,7 @@ hidden_states = hidden_states.reshape(batch_size, seq_len, -1)
|
||||
|
||||
1. **Reduce batch size:**
|
||||
```python
|
||||
omni.generate(prompts=[...], max_batch_size=2)
|
||||
omni.generate(prompts=[...], max_num_seqs=2)
|
||||
```
|
||||
|
||||
2. **Use smaller image size:**
|
||||
|
||||
@@ -330,54 +330,46 @@ Stage transitions are the mechanism by which outputs from one stage are converte
|
||||
|
||||
### Where Stage Transitions Are Called
|
||||
|
||||
Stage transitions happen automatically in the orchestrator (`OmniLLM` class) during the generation loop. Here's the detailed flow:
|
||||
Stage transitions happen automatically in the runtime orchestrator. Here's the detailed flow:
|
||||
|
||||
1. **Location**: `vllm_omni/entrypoints/omni_llm.py` in the `_run_generation()` method
|
||||
1. **Location**: `vllm_omni/engine/orchestrator.py` in `_forward_to_next_stage()`
|
||||
2. **Trigger**: When a stage completes processing and produces outputs
|
||||
3. **Execution Flow**:
|
||||
```python
|
||||
# In omni_llm.py, _run_generation() method (around line 345-460)
|
||||
# In orchestrator.py
|
||||
next_stage_id = stage_id + 1
|
||||
next_client = self.stage_clients[next_stage_id]
|
||||
params = req_state.sampling_params_list[next_stage_id]
|
||||
|
||||
# Main orchestrator loop polls each stage for completed requests
|
||||
for stage_id, stage in enumerate(self.stage_list):
|
||||
result = stage.try_collect() # Get completed request
|
||||
if result is None:
|
||||
continue
|
||||
# Save current stage outputs so stage_input_processors can consume them.
|
||||
self.stage_clients[stage_id].set_engine_outputs([output])
|
||||
|
||||
# Store outputs from this stage
|
||||
engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
|
||||
stage.set_engine_outputs(engine_outputs)
|
||||
# THIS IS WHERE STAGE TRANSITION HAPPENS
|
||||
next_inputs = next_client.process_engine_inputs(
|
||||
stage_list=self.stage_clients,
|
||||
prompt=req_state.prompt,
|
||||
)
|
||||
|
||||
# Check if there's a next stage to forward to
|
||||
next_stage_id = stage_id + 1
|
||||
if next_stage_id < len(self.stage_list):
|
||||
next_stage: OmniStage = self.stage_list[next_stage_id]
|
||||
|
||||
# THIS IS WHERE STAGE TRANSITION HAPPENS
|
||||
next_inputs = next_stage.process_engine_inputs(
|
||||
self.stage_list,
|
||||
[request_id_to_prompt[req_id]]
|
||||
)
|
||||
|
||||
# Submit to next stage
|
||||
task = {
|
||||
"type": OmniStageTaskType.GENERATE,
|
||||
"request_id": req_id,
|
||||
"engine_inputs": next_inputs[0],
|
||||
"sampling_params": sampling_params_list[next_stage_id],
|
||||
}
|
||||
next_stage.submit(task)
|
||||
# Build and submit request(s) to the next stage.
|
||||
for next_input in next_inputs:
|
||||
request = build_engine_core_request_from_tokens(
|
||||
request_id=req_id,
|
||||
prompt=next_input,
|
||||
params=params,
|
||||
model_config=self.stage_vllm_configs[next_stage_id].model_config,
|
||||
)
|
||||
await next_client.add_request_async(request)
|
||||
```
|
||||
|
||||
### How Stage Transitions Work
|
||||
|
||||
The stage transition process follows these steps:
|
||||
|
||||
1. **Stage Completion**: When a stage finishes processing a request, it stores outputs via `stage.set_engine_outputs(engine_outputs)`
|
||||
1. **Stage Completion**: When a stage finishes processing a request, the orchestrator stores outputs via `stage_client.set_engine_outputs(...)`
|
||||
|
||||
2. **Transition Detection**: The orchestrator checks if there's a next stage and calls `process_engine_inputs()` on it
|
||||
|
||||
3. **Input Processing**: The `process_engine_inputs()` method in `OmniStage` (`omni_stage.py`) handles the transition:
|
||||
3. **Input Processing**: The stage input processor configured in stage YAML (under `vllm_omni/model_executor/stage_input_processors/`) handles the transition:
|
||||
```python
|
||||
def process_engine_inputs(
|
||||
self, stage_list: list[Any], prompt: OmniTokensPrompt | TextPrompt = None
|
||||
|
||||
@@ -361,9 +361,9 @@ stage_args:
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 64
|
||||
engine_args:
|
||||
model_stage: ar_stage
|
||||
max_num_seqs: 64
|
||||
model_arch: YourTTSModelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
@@ -403,9 +403,9 @@ stage_args:
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 64
|
||||
engine_args:
|
||||
model_stage: ar_stage
|
||||
max_num_seqs: 64
|
||||
model_arch: YourTTSModelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
|
||||
@@ -23,30 +23,32 @@ export VLLM_PROFILER_MAX_ITERS=1
|
||||
The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
|
||||
```python
|
||||
# Profile all stages
|
||||
omni_llm.start_profile()
|
||||
omni.start_profile()
|
||||
|
||||
# Only profile Stage 1
|
||||
omni_llm.start_profile(stages=[1])
|
||||
omni.start_profile(stages=[1])
|
||||
```
|
||||
|
||||
```python
|
||||
# Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
|
||||
omni_llm.start_profile(stages=[0, 2])
|
||||
omni.start_profile(stages=[0, 2])
|
||||
```
|
||||
|
||||
**Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
|
||||
|
||||
```python
|
||||
from vllm_omni import omni_llm
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
|
||||
|
||||
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
|
||||
|
||||
# 1. Start profiling if enabled
|
||||
if profiler_enabled:
|
||||
omni_llm.start_profile(stages=[0])
|
||||
omni.start_profile(stages=[0])
|
||||
|
||||
# Initialize generator
|
||||
omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
|
||||
omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
|
||||
|
||||
total_requests = len(prompts)
|
||||
processed_count = 0
|
||||
@@ -57,21 +59,21 @@ for stage_outputs in omni_generator:
|
||||
# ... [Output processing logic for text/audio would go here] ...
|
||||
|
||||
# Update count to track when to stop profiling
|
||||
processed_count += len(stage_outputs.request_output)
|
||||
processed_count += 1
|
||||
|
||||
# 2. Check if all requests are done to stop the profiler safely
|
||||
if profiler_enabled and processed_count >= total_requests:
|
||||
print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
|
||||
|
||||
# Stop the profiler while workers are still active
|
||||
omni_llm.stop_profile()
|
||||
omni.stop_profile()
|
||||
|
||||
# Wait for traces to flush to disk
|
||||
print("[Info] Waiting 30s for workers to write trace files to disk...")
|
||||
time.sleep(30)
|
||||
print("[Info] Trace export wait time finished.")
|
||||
|
||||
omni_llm.close()
|
||||
omni.close()
|
||||
```
|
||||
|
||||
|
||||
|
||||
@@ -67,12 +67,12 @@ According to analysis for current popular open-source models, most of them have
|
||||
| Component | Description |
|
||||
| ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| **OmniRouter** | provide an intelligent router for Omni-modality requests dispatch |
|
||||
| **EntryPoints** | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni) and provide the OmniStage abstraction for different AR/DiT stages |
|
||||
| **EntryPoints** | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni), while `AsyncOmniEngine` and `Orchestrator` coordinate multi-stage AR/DiT execution |
|
||||
| **AR** | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management |
|
||||
| **Diffusion** | natively implemented and optimized using acceleration components |
|
||||
| **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages |
|
||||
|
||||
Disaggregated stages are managed through configuration, such as in the Qwen3-Omni example, where stages like Thinker, Talker, and Code2wav are defined as separate OmniStage instances with specific resources and input/output type.
|
||||
Disaggregated stages are managed through stage configuration. In Qwen3-Omni, Thinker/Talker/Code2wav are declared as separate configured stages, and runtime routing is handled by `Orchestrator` over `StageEngineCoreClient` / `StageDiffusionClient`.
|
||||
|
||||
## Main features
|
||||
|
||||
@@ -127,10 +127,10 @@ Taking **Qwen3-Omni** as an example:
|
||||
The **Omni** class provides a Python interface for offline batched inference. Users initialize the Omni class with a Hugging Face model name and use the generate method, passing inputs that include both text prompts and multi-modal data:
|
||||
|
||||
```
|
||||
# Create an omni_lm with HF model name.
|
||||
# Create an omni runtime with HF model name.
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni_lm = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
|
||||
omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
|
||||
|
||||
# Example prompts.
|
||||
om_inputs = {"prompt": prompt,
|
||||
@@ -140,7 +140,7 @@ om_inputs = {"prompt": prompt,
|
||||
}}
|
||||
|
||||
# Generate texts and audio from the multi-modality inputs.
|
||||
outputs = omni_lm.generate(om_inputs, sampling_params_list)
|
||||
outputs = omni.generate(om_inputs, sampling_params_list)
|
||||
```
|
||||
|
||||
## Online Serving
|
||||
|
||||
@@ -19,7 +19,7 @@ The `async_chunk` feature enables asynchronous, chunked processing of data acros
|
||||
|
||||
For qwen3-omni:
|
||||
- **Thinker → Talker**: Per decode step (typically chunk_size=1)
|
||||
- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. Set `initial_codec_chunk_frames` to emit smaller chunks during the initial phase for reduced TTFA
|
||||
- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. During the initial phase, a dynamic initial chunk size (IC) is automatically selected based on server load to reduce TTFA. Use the per-request `initial_codec_chunk_frames` API field to override.
|
||||
- **Code2Wav**: Streaming decode with code2wav chunk_size
|
||||
|
||||
With `async_chunk`:
|
||||
@@ -67,6 +67,13 @@ Enabling **async_chunk** (False→True) sharply reduces time-to-first-audio (TTF
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png">
|
||||
<img alt="E2E Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png" width=100%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
## Architecture
|
||||
### Data Flow
|
||||
|
||||
@@ -164,7 +171,7 @@ stage_args:
|
||||
- `custom_process_next_stage_input_func: str`: Path to custom chunk processing function; receives `(transfer_manager, pooling_output, request)`. For qwen3-omni: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
|
||||
- `stage_connector_config: dict`: Connector configuration
|
||||
- `worker_type: str`: Model type, e.g. `"ar"` or `"generation"` (used by OmniChunkTransferAdapter for mode-specific payload handling)
|
||||
- `max_batch_size: int`: Maximum batch size for the stage
|
||||
- `max_num_seqs: int`: Maximum number of sequences for concurrent processing in the stage
|
||||
|
||||
|
||||
### Connector Configuration
|
||||
@@ -188,9 +195,9 @@ stage_args:
|
||||
- stage_id: 2 # code2wav stage
|
||||
runtime:
|
||||
devices: "1"
|
||||
max_batch_size: 64 # Enables batched audio generation
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
max_num_seqs: 64 # Enables batched audio generation
|
||||
```
|
||||
|
||||
## Related Files
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
# Diffusion Step Execution
|
||||
|
||||
This guide documents vLLM-Omni's stepwise diffusion contract for model authors
|
||||
and contributors implementing `step_execution=True` support for a diffusion
|
||||
pipeline.
|
||||
|
||||
For end-user enablement, supported models, and current limitations, see
|
||||
[Step Execution](../../user_guide/diffusion/step_execution.md).
|
||||
|
||||
## Current Support Scope
|
||||
|
||||
`step_execution` is **not** a generic diffusion toggle. It only works for
|
||||
pipelines that implement the segmented stateful contract in
|
||||
[`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py).
|
||||
|
||||
Current in-tree support:
|
||||
|
||||
| Pipeline | Example models | Step execution |
|
||||
|----------|----------------|----------------|
|
||||
| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
|
||||
| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
|
||||
|
||||
Current engine/runtime limitations:
|
||||
|
||||
- `StepScheduler` only schedules `batch_size=1`.
|
||||
- `cache_backend` is not supported in step mode.
|
||||
- Request-mode extras such as KV transfer are not wired into step mode yet.
|
||||
- Unsupported pipelines now fail early during model loading instead of failing on the first request.
|
||||
|
||||
## Execution Contract
|
||||
|
||||
Step mode is driven by four pipeline methods plus the shared mutable request
|
||||
state object:
|
||||
|
||||
- `prepare_encode(state)`: one-time request preparation.
|
||||
- `denoise_step(state)`: compute the noise prediction for the current step.
|
||||
- `step_scheduler(state, noise_pred)`: mutate latents and advance step state.
|
||||
- `post_decode(state)`: decode the final output after denoising is complete.
|
||||
|
||||
The state lives in
|
||||
[`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
|
||||
as `DiffusionRequestState`. Store request-scoped tensors there, or use
|
||||
`state.extra` for model-specific fields that do not justify extending the core
|
||||
dataclass.
|
||||
|
||||
The worker-side step loop lives in
|
||||
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py):
|
||||
|
||||
1. `prepare_encode()` runs once for a new request.
|
||||
2. `denoise_step()` runs every scheduler tick.
|
||||
3. `step_scheduler()` mutates `state.latents` and advances `state.step_index`.
|
||||
4. `post_decode()` runs exactly once after `state.denoise_completed` becomes true.
|
||||
|
||||
## Recommended Split
|
||||
|
||||
When converting an existing request-level `forward()` pipeline, keep the split
|
||||
strict and mechanical:
|
||||
|
||||
| Request-level phase | Stepwise method | What belongs there |
|
||||
|---------------------|-----------------|--------------------|
|
||||
| Input validation, prompt encoding, latent init, timestep prep, per-request scheduler creation | `prepare_encode()` | Anything that should happen once per request |
|
||||
| Transformer forward / noise prediction | `denoise_step()` | Pure denoise computation for the current timestep |
|
||||
| `scheduler.step(...)` and `step_index += 1` | `step_scheduler()` | Only latent/state mutation for one step |
|
||||
| VAE decode / postprocess | `post_decode()` | Final decode only |
|
||||
|
||||
Keep the stepwise path reusing the same helpers as the request-level path
|
||||
whenever possible. Reimplementing the denoise loop from scratch is the easiest
|
||||
way to introduce behavioral drift.
|
||||
|
||||
## Qwen-Image Reference
|
||||
|
||||
[`pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)
|
||||
is the reference implementation and is split correctly for the current
|
||||
contract:
|
||||
|
||||
- `prepare_encode()` reuses `_prepare_generation_context()` so prompt encoding,
|
||||
latent init, timestep creation, CFG setup, and shape bookkeeping stay aligned
|
||||
with `forward()`.
|
||||
- `prepare_encode()` deep-copies `self.scheduler` **after**
|
||||
`prepare_timesteps()` so request-specific scheduler state is isolated.
|
||||
- `denoise_step()` reuses `_build_denoise_kwargs()` plus
|
||||
`predict_noise_maybe_with_cfg()`, so sequential CFG, CFG-parallel, and
|
||||
non-CFG behavior stay identical to the request-level path.
|
||||
- `step_scheduler()` only calls
|
||||
`scheduler_step_maybe_with_cfg(..., per_request_scheduler=state.scheduler)`
|
||||
and increments `state.step_index`.
|
||||
- `post_decode()` reuses `_decode_latents()`, so the final image decode matches
|
||||
the normal `forward()` path.
|
||||
|
||||
That decomposition is the target pattern for future models.
|
||||
|
||||
## Rules For New Pipelines
|
||||
|
||||
- Do not keep request-scoped scheduler state on `self.scheduler`. Copy it into
|
||||
`state.scheduler` during `prepare_encode()`.
|
||||
- Do not mutate `state.step_index` inside `denoise_step()`. Only
|
||||
`step_scheduler()` should advance the step.
|
||||
- Do not decode partial outputs in `denoise_step()` or `step_scheduler()`.
|
||||
- If the request-level pipeline has condition latents, masks, or edit-specific
|
||||
tensors, store them in `state` or `state.extra`, not in global pipeline
|
||||
attributes.
|
||||
- Preserve CFG behavior by sharing the same helper path used by `forward()`.
|
||||
- Keep `post_decode()` equivalent to the tail of `forward()`.
|
||||
|
||||
## Validation Checklist
|
||||
|
||||
Before marking a pipeline as `supports_step_execution = True`, verify:
|
||||
|
||||
- Stepwise output matches request-level output for the same seed and sampling params.
|
||||
- Per-request scheduler state is isolated across concurrent requests.
|
||||
- Abort during denoise does not leak cached state.
|
||||
- `step_index` reported by `RunnerOutput` matches the scheduler progress.
|
||||
- CFG-parallel and non-CFG paths both work if the request-level pipeline supports them.
|
||||
|
||||
## Related Files
|
||||
|
||||
- Contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py)
|
||||
- State: [`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
|
||||
- Runner loop: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
||||
- Scheduler transport: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
- Reference pipeline: [`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)
|
||||
@@ -72,3 +72,241 @@ Parameters:
|
||||
|
||||
For more details, refer to the
|
||||
[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
|
||||
|
||||
---
|
||||
|
||||
## Design
|
||||
|
||||
### 1. Overview
|
||||
|
||||
`MooncakeStoreConnector` is the storage-oriented remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake distributed store APIs and provides a uniform `put()` / `get()` abstraction for transferring stage payloads across nodes.
|
||||
|
||||
Compared with `SharedMemoryConnector`, this connector is intended for multi-node deployments. Compared with `MooncakeTransferEngineConnector`, it is simpler and more object-oriented: data is written into a distributed store and fetched back by key, rather than transferred through an explicit remote-memory write protocol.
|
||||
|
||||
Its primary role is to provide a general-purpose remote connector for stage payloads when:
|
||||
|
||||
- the pipeline spans multiple hosts
|
||||
- a shared-memory connector is not possible
|
||||
- a simpler remote backend is preferred over the RDMA-oriented transfer engine
|
||||
|
||||
### 2. Relationship with the OmniConnector System
|
||||
|
||||
`MooncakeStoreConnector` implements `OmniConnectorBase`, so it integrates with the same connector lifecycle as the other backends:
|
||||
|
||||
- `OmniConnectorFactory` instantiates it from `ConnectorSpec`
|
||||
- `load_omni_transfer_config()` maps edge config to the connector spec
|
||||
- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
|
||||
|
||||
This means the rest of the pipeline does not need store-specific logic. It only relies on the generic connector contract.
|
||||
|
||||
### 3. Design Goals
|
||||
|
||||
The connector is designed around the following goals:
|
||||
|
||||
- **Cross-node object transfer** using a single key-based abstraction
|
||||
- **Minimal connector-specific control plane**, since the store itself is the shared medium
|
||||
- **Compatibility with arbitrary Python payloads** through the Omni serializer
|
||||
- **Operational simplicity** compared with the transfer-engine-based remote connector
|
||||
|
||||
The connector is intentionally not optimized for zero-copy tensor movement or direct remote-memory access.
|
||||
|
||||
### 4. Core Design
|
||||
|
||||
#### 4.1 Object-Oriented Transfer Model
|
||||
|
||||
`MooncakeStoreConnector` treats the transport backend as a distributed object store:
|
||||
|
||||
1. Serialize the Python object into bytes.
|
||||
2. Generate a unique connector key.
|
||||
3. Store the bytes in Mooncake.
|
||||
4. Fetch the bytes from Mooncake on the receiver side.
|
||||
5. Deserialize the bytes back into the original object.
|
||||
|
||||
This design makes the connector conceptually close to a distributed key-value transport.
|
||||
|
||||
#### 4.2 Key Construction
|
||||
|
||||
The connector uses `OmniConnectorBase._make_key()` to derive the internal store key:
|
||||
|
||||
```text
|
||||
{request_key}@{from_stage}_{to_stage}
|
||||
```
|
||||
|
||||
This adds stage routing information to the logical request key and avoids collisions between different pipeline edges.
|
||||
|
||||
#### 4.3 No Extra Metadata Path
|
||||
|
||||
Unlike `SharedMemoryConnector` and `MooncakeTransferEngineConnector`, this connector returns:
|
||||
|
||||
```python
|
||||
metadata = None
|
||||
```
|
||||
|
||||
from `put()`.
|
||||
|
||||
This is a meaningful design difference:
|
||||
|
||||
- the store itself is the shared rendezvous point
|
||||
- consumers only need the same key
|
||||
- no transport handle, address, or side-channel metadata needs to be propagated
|
||||
|
||||
As a result, the control plane is simpler for this connector than for the SHM and RDMA variants.
|
||||
|
||||
### 5. Initialization
|
||||
|
||||
#### 5.1 Required Mooncake Components
|
||||
|
||||
The connector requires the Mooncake Python bindings to expose:
|
||||
|
||||
- `MooncakeDistributedStore`
|
||||
- `ReplicateConfig`
|
||||
|
||||
If these symbols are unavailable, construction fails immediately with `ImportError`. This keeps startup failures explicit and avoids silent fallback behavior.
|
||||
|
||||
#### 5.2 Store Setup
|
||||
|
||||
During `_init_store()`, the connector:
|
||||
|
||||
1. creates `MooncakeDistributedStore()`
|
||||
2. calls `store.setup(...)`
|
||||
3. validates the return code
|
||||
4. creates a `ReplicateConfig`
|
||||
5. enables `with_soft_pin = True`
|
||||
|
||||
The setup step is completed during connector construction, not lazily at first use.
|
||||
|
||||
### 6. Put / Get Flow
|
||||
|
||||
#### 6.1 Producer Flow: `put()`
|
||||
|
||||
The producer-side flow is:
|
||||
|
||||
1. Validate that the store has been initialized.
|
||||
2. Serialize the payload into bytes.
|
||||
3. Build the internal key with stage routing information.
|
||||
4. Call `store.put(key, serialized_data, self.pin)`.
|
||||
5. Update metrics and return success.
|
||||
|
||||
The returned tuple is:
|
||||
|
||||
```python
|
||||
(True, len(serialized_data), None)
|
||||
```
|
||||
|
||||
#### 6.2 Consumer Flow: `get()`
|
||||
|
||||
The consumer-side flow is:
|
||||
|
||||
1. Validate that the store has been initialized.
|
||||
2. Build the same internal key.
|
||||
3. Poll `store.get(key)` for up to 20 retries.
|
||||
4. Sleep 50 ms between retries.
|
||||
5. If data is found, deserialize it and return `(data, payload_size)`.
|
||||
6. If all retries are exhausted, record a timeout and return `None`.
|
||||
|
||||
This design gives the connector a bounded waiting model rather than an indefinite blocking get.
|
||||
|
||||
### 7. Integration with Stage Communication
|
||||
|
||||
All callers use the connector through the same `put()` / `get()` contract:
|
||||
|
||||
- the sender calls `put()` to serialize and store the payload
|
||||
- the receiver calls `get()` to retrieve and deserialize it
|
||||
- no connector-specific metadata is required, since the store key is the rendezvous point
|
||||
|
||||
Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, which makes the connector functional but not the highest-performance option for large payloads such as KV cache blocks.
|
||||
|
||||
### 8. Data Flow in the Pipeline
|
||||
|
||||
The end-to-end transfer model is:
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant SenderStage
|
||||
participant MooncakeStoreConnector
|
||||
participant MooncakeStore
|
||||
participant ReceiverStage
|
||||
|
||||
SenderStage->>MooncakeStoreConnector: put(from_stage, to_stage, put_key, data)
|
||||
MooncakeStoreConnector->>MooncakeStoreConnector: serialize object
|
||||
MooncakeStoreConnector->>MooncakeStore: put(internal_key, bytes)
|
||||
MooncakeStoreConnector-->>SenderStage: (success, size, None)
|
||||
|
||||
ReceiverStage->>MooncakeStoreConnector: get(from_stage, to_stage, get_key)
|
||||
MooncakeStoreConnector->>MooncakeStore: get(internal_key)
|
||||
MooncakeStoreConnector->>MooncakeStoreConnector: deserialize bytes
|
||||
MooncakeStoreConnector-->>ReceiverStage: (data, size)
|
||||
```
|
||||
|
||||
This is a store-mediated remote transfer model rather than a direct peer-to-peer transport model.
|
||||
|
||||
### 9. Strengths and Trade-offs
|
||||
|
||||
#### Strengths
|
||||
|
||||
- Simple conceptual model: store by key, fetch by key.
|
||||
- No connector-specific metadata handoff is required.
|
||||
- Works naturally for remote multi-node stage transfer.
|
||||
- Easy to integrate into the existing connector abstraction.
|
||||
|
||||
#### Trade-offs
|
||||
|
||||
- Full serialization and deserialization are always required.
|
||||
- Large payloads are more expensive than in direct-memory transports.
|
||||
- Runtime behavior depends on external Mooncake services being available.
|
||||
- Cleanup semantics are weaker than request-scoped local buffer management.
|
||||
|
||||
### 10. Important Implementation Characteristics
|
||||
|
||||
#### 10.1 Cleanup Is a No-op
|
||||
|
||||
`cleanup()` only logs a debug message and does not actively delete remote data. The current design assumes Mooncake-side lifecycle management rather than explicit per-request removal.
|
||||
|
||||
This keeps the connector implementation simple but means request-level reclamation is not modeled inside the connector itself.
|
||||
|
||||
#### 10.2 Close Releases the Store Handle
|
||||
|
||||
Unlike `SharedMemoryConnector`, where `close()` is a no-op, `MooncakeStoreConnector.close()` performs a meaningful shutdown:
|
||||
|
||||
1. Calls `self.store.close()` to release the Mooncake store handle.
|
||||
2. Sets `self.store = None` to prevent further operations.
|
||||
|
||||
This ensures the connector releases its connection to the external Mooncake service on shutdown. Errors during close are logged but do not propagate.
|
||||
|
||||
#### 10.3 Health Output Uses a Shared Schema
|
||||
|
||||
`health()` reports:
|
||||
|
||||
- `host`
|
||||
- `metadata_server`
|
||||
- `master`
|
||||
- metrics
|
||||
|
||||
and also includes placeholder fields such as `protocol`, `pool_device`, and `pool_size` to keep a more uniform shape across connector health outputs. This reflects a system-level consistency choice rather than a store-specific design need.
|
||||
|
||||
#### 10.4 Get Is Retry-Based, Not Event-Driven
|
||||
|
||||
The receiver side polls the store with a short retry loop. This is simple and robust, but it means the connector is latency-sensitive to:
|
||||
|
||||
- store visibility delay
|
||||
- network jitter
|
||||
- payload size
|
||||
|
||||
If the payload is not visible within the retry window, the connector reports a timeout.
|
||||
|
||||
### 11. Summary
|
||||
|
||||
`MooncakeStoreConnector` is the remote, store-based connector in the OmniConnector stack. Its design is straightforward:
|
||||
|
||||
- serialize the payload
|
||||
- store it in Mooncake under a stage-qualified key
|
||||
- fetch it by the same key on the receiver side
|
||||
- deserialize it back into the original object
|
||||
|
||||
This connector fills an important gap in the system:
|
||||
|
||||
1. It enables cross-node transfer without requiring shared memory.
|
||||
2. It keeps the stage communication model uniform.
|
||||
3. It provides a simpler operational alternative to the transfer-engine-based remote connector.
|
||||
|
||||
Its simplicity is also its main trade-off: it favors a clean object-transfer model over the fast-path and peer-to-peer optimizations implemented by `MooncakeTransferEngineConnector`.
|
||||
@@ -205,3 +205,589 @@ for test-specific setup instructions.
|
||||
|
||||
For more details on the underlying engine, refer to the
|
||||
[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
|
||||
|
||||
---
|
||||
|
||||
## Design
|
||||
|
||||
### 1. Overview
|
||||
|
||||
`MooncakeTransferEngineConnector` is the high-performance remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake `TransferEngine` and combines:
|
||||
|
||||
- a **direct data plane** for remote memory writes
|
||||
- a **ZMQ side channel** for metadata lookup, handshake, and completion signaling
|
||||
- a **managed local memory pool** for both send and receive buffers
|
||||
|
||||
Unlike `MooncakeStoreConnector`, which treats the backend as a distributed store, `MooncakeTransferEngineConnector` is designed as a peer-to-peer transport. Its goal is to move large stage payloads efficiently while still fitting the common `put()` / `get()` API defined by `OmniConnectorBase`.
|
||||
|
||||
It is the most performance-oriented connector in the current OmniConnector family and is intended for large remote payloads such as:
|
||||
|
||||
- KV cache transfer
|
||||
- stage hidden-state payloads
|
||||
- streaming chunk payloads
|
||||
- other binary-heavy inter-stage artifacts
|
||||
|
||||
### 2. Relationship with the OmniConnector System
|
||||
|
||||
`MooncakeTransferEngineConnector` implements the same connector contract as the other backends:
|
||||
|
||||
- `put(from_stage, to_stage, put_key, data)`
|
||||
- `get(from_stage, to_stage, get_key, metadata=None)`
|
||||
- `cleanup(request_id, ...)`
|
||||
- `health()`
|
||||
- `close()`
|
||||
|
||||
It is integrated into the system through the standard connector plumbing:
|
||||
|
||||
- `OmniConnectorFactory` constructs the connector from `ConnectorSpec`
|
||||
- `load_omni_transfer_config()` resolves the edge-level connector configuration
|
||||
- `get_connectors_config_for_stage()` and `resolve_omni_kv_config_for_stage()` inject the connector role
|
||||
- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
|
||||
|
||||
The key system-level distinction is that this connector is **role-aware**:
|
||||
|
||||
- sender instances expose data and listen for pull requests
|
||||
- receiver instances allocate buffers and actively pull data from the sender
|
||||
|
||||
### 3. Design Goals
|
||||
|
||||
The connector is designed around four primary goals:
|
||||
|
||||
1. **High-throughput remote transfer**
|
||||
Avoid store-mediated round trips and write directly into the receiver memory region.
|
||||
|
||||
2. **Fast path for raw payloads**
|
||||
Support `torch.Tensor`, `bytes`, and `ManagedBuffer` without forcing all traffic through full object serialization.
|
||||
|
||||
3. **Unified connector abstraction**
|
||||
Preserve the same `put()` / `get()` interface used by the rest of the OmniConnector stack.
|
||||
|
||||
4. **Safe lifecycle management**
|
||||
Manage allocation, reuse, cleanup, and failure recovery for a registered memory pool.
|
||||
|
||||
### 4. Architecture Overview
|
||||
|
||||
At a high level, the connector is composed of four main subsystems:
|
||||
|
||||
```mermaid
|
||||
classDiagram
|
||||
class OmniConnectorBase {
|
||||
<<abstract>>
|
||||
+put(from_stage, to_stage, put_key, data)
|
||||
+get(from_stage, to_stage, get_key, metadata)
|
||||
+cleanup(request_id)
|
||||
+health()
|
||||
+close()
|
||||
}
|
||||
|
||||
class MooncakeTransferEngineConnector {
|
||||
+supports_raw_data: bool
|
||||
-engine: TransferEngine
|
||||
-allocator: BufferAllocator
|
||||
-pool: torch.Tensor
|
||||
-zmq_ctx: zmq.Context
|
||||
-_local_buffers: dict
|
||||
-_sender_executor: ThreadPoolExecutor
|
||||
-_listener_thread: threading.Thread
|
||||
+put(...)
|
||||
+get(...)
|
||||
+update_sender_info(sender_host, sender_zmq_port)
|
||||
+get_connection_info()
|
||||
+cleanup(request_id, from_stage, to_stage)
|
||||
+close()
|
||||
}
|
||||
|
||||
class BufferAllocator {
|
||||
-total_size: int
|
||||
-alignment: int
|
||||
-free_blocks: list
|
||||
+alloc(size) int
|
||||
+free(offset, size)
|
||||
}
|
||||
|
||||
class ManagedBuffer {
|
||||
-allocator: BufferAllocator
|
||||
-offset: int
|
||||
-size: int
|
||||
-pool_tensor: torch.Tensor
|
||||
+tensor
|
||||
+as_tensor(dtype, shape) torch.Tensor
|
||||
+to_bytes() bytes
|
||||
+release()
|
||||
}
|
||||
|
||||
class TransferEngine {
|
||||
+initialize(host, handshake, protocol, device_name)
|
||||
+register_memory(base_ptr, size)
|
||||
+batch_transfer_sync_write(remote_session, src_addrs, dst_addrs, lengths)
|
||||
+unregister_memory(base_ptr)
|
||||
+get_rpc_port() int
|
||||
}
|
||||
|
||||
class QueryRequest {
|
||||
+request_id: str
|
||||
}
|
||||
|
||||
class QueryResponse {
|
||||
+request_id: str
|
||||
+data_size: int
|
||||
+is_fast_path: bool
|
||||
}
|
||||
|
||||
class MooncakeAgentMetadata {
|
||||
+remote_hostname: str
|
||||
+remote_port: int
|
||||
+request_id: str
|
||||
+dst_addrs: list[int]
|
||||
+lengths: list[int]
|
||||
}
|
||||
|
||||
OmniConnectorBase <|-- MooncakeTransferEngineConnector
|
||||
MooncakeTransferEngineConnector *-- BufferAllocator
|
||||
MooncakeTransferEngineConnector *-- TransferEngine
|
||||
ManagedBuffer --> BufferAllocator : releases to
|
||||
ManagedBuffer --> "1" torch.Tensor : views
|
||||
MooncakeTransferEngineConnector ..> ManagedBuffer : returns / retains
|
||||
MooncakeTransferEngineConnector ..> QueryRequest : decodes
|
||||
MooncakeTransferEngineConnector ..> QueryResponse : encodes
|
||||
MooncakeTransferEngineConnector ..> MooncakeAgentMetadata : exchanges
|
||||
```
|
||||
|
||||
#### 4.1 Transfer Engine
|
||||
|
||||
Mooncake `TransferEngine` is responsible for the actual data-plane transfer. It registers local memory and performs synchronous remote writes through:
|
||||
|
||||
```python
|
||||
batch_transfer_sync_write(...)
|
||||
```
|
||||
|
||||
#### 4.2 Managed Memory Pool
|
||||
|
||||
Each connector instance owns a large pre-registered memory pool:
|
||||
|
||||
- CPU pinned memory when `memory_pool_device == "cpu"`
|
||||
- GPU memory when `memory_pool_device == "cuda"`
|
||||
|
||||
This avoids repeated memory registration and allows each transfer to allocate subranges from one long-lived pool.
|
||||
|
||||
#### 4.3 Buffer Manager
|
||||
|
||||
Two helper classes control local memory ownership:
|
||||
|
||||
- `BufferAllocator`
|
||||
Manages aligned subrange allocation and free-list merging.
|
||||
|
||||
- `ManagedBuffer`
|
||||
Represents one live slice of the pool and exposes:
|
||||
- `.tensor`
|
||||
- `.as_tensor(dtype, shape)`
|
||||
- `.to_bytes()`
|
||||
- `.release()`
|
||||
|
||||
#### 4.4 ZMQ Side Channel
|
||||
|
||||
ZMQ is used for transport coordination, not for the data payload itself. It handles:
|
||||
|
||||
- metadata query from receiver to sender
|
||||
- pull request submission
|
||||
- completion or error signaling
|
||||
- internal notification from worker threads back to the listener thread
|
||||
|
||||
This split makes the control plane lightweight while keeping the bulk payload on the transfer engine data plane.
|
||||
|
||||
### 5. Role Model
|
||||
|
||||
#### 5.1 Sender Role
|
||||
|
||||
A sender connector:
|
||||
|
||||
- accepts `put()` calls
|
||||
- stores live transfer-ready buffers in `_local_buffers`
|
||||
- starts a ZMQ listener thread
|
||||
- responds to metadata queries and pull requests from receivers
|
||||
|
||||
#### 5.2 Receiver Role
|
||||
|
||||
A receiver connector:
|
||||
|
||||
- does not bind the sender-side ZMQ listener
|
||||
- accepts `get()` calls
|
||||
- allocates receive buffers from its own pool
|
||||
- requests metadata or transfer service from the sender
|
||||
|
||||
The role is not inferred dynamically. It is injected by the stage configuration layer:
|
||||
|
||||
- incoming edge for a stage -> `role="receiver"`
|
||||
- outgoing edge for a stage -> `role="sender"`
|
||||
|
||||
This is important because incorrect role assignment would break initialization semantics.
|
||||
|
||||
#### 5.3 Host Auto-Detection
|
||||
|
||||
The `host` configuration field supports the special value `"auto"`. When set, the connector auto-detects the local IP address that would be used for external communication (via a UDP socket probe to `8.8.8.8`). If that fails, it falls back to hostname resolution, and ultimately to `127.0.0.1`.
|
||||
|
||||
This is useful in environments where the operator does not want to hard-code IP addresses in the connector config.
|
||||
|
||||
#### 5.4 RDMA Device Filtering
|
||||
|
||||
The `device_name` configuration field allows the operator to specify which RDMA NICs to use (comma-separated, e.g. `"mlx5_0,mlx5_1"`). If not set in config, the connector also checks the `RDMA_DEVICE_NAME` environment variable.
|
||||
|
||||
This is important in environments with mixed InfiniBand/RoCE NICs, where not all devices are suitable for the transfer engine.
|
||||
|
||||
### 6. Local Memory Management
|
||||
|
||||
#### 6.1 Memory Pool Registration
|
||||
|
||||
During initialization, the connector:
|
||||
|
||||
1. allocates a large pool tensor
|
||||
2. records its base pointer
|
||||
3. registers that memory with Mooncake
|
||||
4. creates a `BufferAllocator` for subrange management
|
||||
|
||||
This means every later transfer only allocates offsets inside the pre-registered pool rather than registering memory per request.
|
||||
|
||||
#### 6.2 BufferAllocator
|
||||
|
||||
`BufferAllocator` maintains a sorted free list of `(offset, size)` blocks and enforces alignment. Its responsibilities include:
|
||||
|
||||
- aligned allocation
|
||||
- freeing previously allocated blocks
|
||||
- adjacent block merging
|
||||
- double-free detection
|
||||
- overlap detection to catch corruption
|
||||
|
||||
This is a critical piece of the connector because both sender and receiver depend on long-lived pool reuse.
|
||||
|
||||
#### 6.3 ManagedBuffer
|
||||
|
||||
`ManagedBuffer` is the main fast-path data wrapper. It can:
|
||||
|
||||
- expose the pool slice as a zero-copy 1D `uint8` tensor
|
||||
- reinterpret that slice as a typed tensor
|
||||
- copy out the contents as Python `bytes`
|
||||
- release the slice back to the allocator
|
||||
|
||||
The connector uses `ManagedBuffer` in two different ways:
|
||||
|
||||
- as a send-side holder to keep the pool slice alive
|
||||
- as a receive-side return type when `is_fast_path=True`
|
||||
|
||||
### 7. Put Flow
|
||||
|
||||
#### 7.1 High-Level Behavior
|
||||
|
||||
`put()` is only valid in sender mode. Its job is to expose a payload for later remote pull by the receiver.
|
||||
|
||||
The high-level flow is:
|
||||
|
||||
1. validate connector state and role
|
||||
2. convert the input into a pool-backed transferable representation
|
||||
3. store the transfer metadata in `_local_buffers`
|
||||
4. return lightweight metadata describing how the receiver can fetch the data
|
||||
|
||||
#### 7.2 Payload Type Handling
|
||||
|
||||
`put()` supports three payload classes:
|
||||
|
||||
**A. `ManagedBuffer`**
|
||||
|
||||
If the buffer belongs to the same pool, the connector can use it directly without copying. This is the most efficient path.
|
||||
|
||||
If the buffer comes from a different pool, the connector falls back to a copy path.
|
||||
|
||||
**B. `torch.Tensor` or `bytes`**
|
||||
|
||||
These payloads are copied into the local pool and marked as fast-path data:
|
||||
|
||||
- no Omni object serialization is required
|
||||
- receiver can get a `ManagedBuffer` back
|
||||
|
||||
**C. Generic Python object**
|
||||
|
||||
Any other payload is serialized via `OmniSerializer.serialize(...)` and then copied into the pool.
|
||||
|
||||
In this case:
|
||||
|
||||
- `is_fast_path=False`
|
||||
- the receiver will deserialize back into a Python object
|
||||
|
||||
#### 7.3 Sender Metadata
|
||||
|
||||
The sender returns:
|
||||
|
||||
```python
|
||||
{
|
||||
"source_host": self.host,
|
||||
"source_port": self.zmq_port,
|
||||
"data_size": size,
|
||||
"is_fast_path": is_fast_path,
|
||||
}
|
||||
```
|
||||
|
||||
This metadata is intentionally lightweight. It tells the receiver:
|
||||
|
||||
- where the sender-side control plane lives
|
||||
- how large the remote transfer will be
|
||||
- whether the payload should be returned as a `ManagedBuffer` or a deserialized object
|
||||
|
||||
#### 7.4 Sender Buffer Table
|
||||
|
||||
The sender stores each live payload in `_local_buffers` under the stage-qualified key. Each entry contains:
|
||||
|
||||
- source addresses
|
||||
- lengths
|
||||
- the holder object
|
||||
- ownership information (`should_release`)
|
||||
- `is_fast_path`
|
||||
- creation time
|
||||
|
||||
This table is the sender-side truth source for both metadata queries and pull requests.
|
||||
|
||||
### 8. Get Flow
|
||||
|
||||
#### 8.1 High-Level Behavior
|
||||
|
||||
`get()` runs on the receiver side and performs four steps:
|
||||
|
||||
1. resolve metadata
|
||||
2. allocate a destination buffer in the local pool
|
||||
3. request the sender to write into that destination buffer
|
||||
4. return either a `ManagedBuffer` or a deserialized object
|
||||
|
||||
#### 8.2 Metadata Resolution Paths
|
||||
|
||||
The `metadata` parameter in `get()` is optional. The connector supports two resolution modes depending on whether the caller supplies it.
|
||||
|
||||
**With metadata**
|
||||
|
||||
When the caller passes metadata, the connector uses it directly. The metadata carries:
|
||||
|
||||
- `source_host` / `source_port` — sender ZMQ endpoint
|
||||
- `data_size` — payload byte count
|
||||
- `is_fast_path` — whether the receiver gets a `ManagedBuffer` or a deserialized object
|
||||
|
||||
This mode is suitable when the control plane already forwards the sender's `put()` output to the receiver.
|
||||
|
||||
**Without metadata**
|
||||
|
||||
When `get(metadata=None)` is called, the connector queries the sender over ZMQ to discover the same fields (`data_size`, `is_fast_path`). The caller must first call:
|
||||
|
||||
```python
|
||||
update_sender_info(sender_host, sender_zmq_port)
|
||||
```
|
||||
|
||||
so that the connector knows where to send the query.
|
||||
|
||||
This mode is suitable for polling-based flows (e.g. KV transfer, async chunk transfer) where the receiver does not have metadata from the control plane.
|
||||
|
||||
#### 8.3 Destination Allocation
|
||||
|
||||
Once metadata is resolved, the receiver:
|
||||
|
||||
1. allocates a subrange from its own local pool
|
||||
2. wraps it in a `ManagedBuffer`
|
||||
3. builds a `MooncakeAgentMetadata` request containing:
|
||||
- receiver hostname
|
||||
- receiver RPC port
|
||||
- request ID
|
||||
- destination addresses
|
||||
- transfer lengths
|
||||
|
||||
This tells the sender exactly where to write the incoming data.
|
||||
|
||||
#### 8.4 Transfer Completion
|
||||
|
||||
The receiver then sends the pull request over ZMQ and waits for:
|
||||
|
||||
- `TRANS_DONE`
|
||||
- or `TRANS_ERROR`
|
||||
|
||||
If the transfer succeeds:
|
||||
|
||||
- for `is_fast_path=True`, the receiver returns `(ManagedBuffer, size)`
|
||||
- for `is_fast_path=False`, the receiver copies to bytes, deserializes, releases the buffer, and returns `(object, size)`
|
||||
|
||||
### 9. Sender-Side Listener Design
|
||||
|
||||
#### 9.1 Listener Thread
|
||||
|
||||
In sender mode, the connector starts `_zmq_listener_loop()` after initialization. The listener:
|
||||
|
||||
- binds `tcp://{host}:{zmq_port}`
|
||||
- receives incoming requests
|
||||
- uses a poller for socket events and internal notifications
|
||||
- periodically reclaims stale buffers
|
||||
|
||||
If the bind fails, initialization fails immediately. The code does not silently downgrade the connector role.
|
||||
|
||||
#### 9.2 Worker Thread Pool
|
||||
|
||||
The listener hands work to `_sender_executor` so that the listener thread does not block on transfer work.
|
||||
|
||||
There are two request types:
|
||||
|
||||
- metadata query -> `_handle_query_request(...)`
|
||||
- transfer request -> `_handle_pull_request(...)`
|
||||
|
||||
#### 9.3 Query Handling
|
||||
|
||||
For metadata queries, the sender looks up the request ID in `_local_buffers` and returns:
|
||||
|
||||
- data size
|
||||
- fast-path flag
|
||||
|
||||
This supports consumers that only know the sender endpoint but not the original sender metadata.
|
||||
|
||||
#### 9.4 Pull Handling
|
||||
|
||||
For a transfer request, the sender:
|
||||
|
||||
1. locates the source addresses in `_local_buffers`
|
||||
2. constructs the remote session identifier
|
||||
3. calls `batch_transfer_sync_write(...)`
|
||||
4. replies `TRANS_DONE` or `TRANS_ERROR`
|
||||
|
||||
On success, the sender immediately calls `cleanup(meta.request_id)` and frees the producer-side buffer if it owns it.
|
||||
|
||||
This choice is important: it makes the connector effectively a **single-consumer transfer model** for each successful put/get pair.
|
||||
|
||||
### 10. Fast Path Semantics
|
||||
|
||||
This connector explicitly advertises:
|
||||
|
||||
```python
|
||||
supports_raw_data = True
|
||||
```
|
||||
|
||||
That means it can move raw payloads without forcing everything through the Omni object serializer.
|
||||
|
||||
#### Fast Path
|
||||
|
||||
For `torch.Tensor`, `bytes`, or pool-local `ManagedBuffer`:
|
||||
|
||||
- sender returns `is_fast_path=True`
|
||||
- receiver returns a `ManagedBuffer`
|
||||
- caller is responsible for calling `release()`
|
||||
|
||||
This avoids an unnecessary copy on the receiver side.
|
||||
|
||||
#### Serialized Object Path
|
||||
|
||||
For arbitrary Python objects:
|
||||
|
||||
- sender serializes the payload
|
||||
- receiver converts the receive buffer to bytes
|
||||
- receiver deserializes the object
|
||||
- receive buffer is released internally
|
||||
|
||||
This preserves a uniform object-oriented API while still allowing optimized raw-data transport when possible.
|
||||
|
||||
### 11. Failure Handling and Cleanup
|
||||
|
||||
#### 11.1 Timeouts and Socket Recovery
|
||||
|
||||
The receiver caches ZMQ REQ sockets per thread, but invalidates them after failures. This avoids reusing sockets that may be stuck in a bad state after timeout or receive errors.
|
||||
|
||||
Timeout is scaled based on payload size:
|
||||
|
||||
- a base timeout
|
||||
- plus additional time for large payloads
|
||||
|
||||
This is intended to reduce false timeouts for large remote writes.
|
||||
|
||||
#### 11.2 Stale Buffer Reclamation
|
||||
|
||||
The sender periodically reclaims old entries from `_local_buffers` using a TTL policy. This protects the memory pool from permanent leaks if a receiver crashes or never consumes a prepared payload.
|
||||
|
||||
This is a practical recovery mechanism, although the code notes that TTL cleanup can still race with very long-running in-flight transfers.
|
||||
|
||||
#### 11.3 Connector Shutdown
|
||||
|
||||
`close()` is a full resource teardown routine. It:
|
||||
|
||||
- stops the listener thread
|
||||
- shuts down the worker executor
|
||||
- releases all pending buffers
|
||||
- closes cached sockets
|
||||
- unregisters memory from the engine when supported
|
||||
- terminates the ZMQ context
|
||||
- drops the pool reference
|
||||
|
||||
This makes `MooncakeTransferEngineConnector` the most lifecycle-aware connector in the current connector family.
|
||||
|
||||
### 12. Current Implementation Constraints
|
||||
|
||||
The current code documents several important topology constraints.
|
||||
|
||||
#### 12.1 One Sender to One Receiver per Successful Transfer
|
||||
|
||||
After a successful transfer, the sender-side buffer is cleaned up immediately. This means the same prepared payload is not retained for multiple independent receivers.
|
||||
|
||||
#### 12.2 One Receiver to One Active Sender Endpoint
|
||||
|
||||
The receiver only stores one `(sender_host, sender_zmq_port)` pair through `update_sender_info(...)`. So the metadata-query mode is currently single-sender at a time.
|
||||
|
||||
#### 12.3 Explicit Buffer Ownership Matters
|
||||
|
||||
When the connector allocates a pool slice internally, it is responsible for releasing it. When a caller passes an externally owned `ManagedBuffer`, the connector keeps it alive for transfer but does not assume ownership of its eventual release.
|
||||
|
||||
These constraints are consistent with the current implementation and should be treated as design assumptions rather than incidental behavior.
|
||||
|
||||
### 13. Data Flow in the Pipeline
|
||||
|
||||
The end-to-end sender/receiver interaction is:
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant SenderStage
|
||||
participant SenderConnector
|
||||
participant ReceiverConnector
|
||||
participant ReceiverStage
|
||||
|
||||
SenderStage->>SenderConnector: put(from_stage, to_stage, put_key, data)
|
||||
SenderConnector->>SenderConnector: place payload in local memory pool
|
||||
SenderConnector-->>SenderStage: metadata(source_host, source_port, data_size, is_fast_path)
|
||||
|
||||
ReceiverStage->>ReceiverConnector: get(..., metadata)
|
||||
ReceiverConnector->>ReceiverConnector: allocate destination buffer
|
||||
ReceiverConnector->>SenderConnector: ZMQ pull request with dst addr
|
||||
SenderConnector->>ReceiverConnector: TransferEngine remote write
|
||||
SenderConnector-->>ReceiverConnector: TRANS_DONE
|
||||
ReceiverConnector-->>ReceiverStage: ManagedBuffer or deserialized object
|
||||
```
|
||||
|
||||
For metadata-less polling, the flow simply adds a metadata query step before the pull request.
|
||||
|
||||
### 14. Strengths and Trade-offs
|
||||
|
||||
#### Strengths
|
||||
|
||||
- Best remote-transfer design in the current connector stack for large payloads.
|
||||
- Supports raw-data fast path.
|
||||
- Keeps stage communication under the same connector abstraction.
|
||||
- Includes real lifecycle and memory-pool management.
|
||||
- Works for both stage payload transfer and KV transfer scenarios.
|
||||
|
||||
#### Trade-offs
|
||||
|
||||
- More complex than the store-based connector.
|
||||
- Correctness depends on role injection and endpoint coordination.
|
||||
- Caller must release fast-path receive buffers.
|
||||
- Current implementation is optimized for single-consumer transfer semantics.
|
||||
|
||||
### 15. Summary
|
||||
|
||||
`MooncakeTransferEngineConnector` is the high-performance peer-to-peer transport in the OmniConnector system. Its design combines:
|
||||
|
||||
- a registered memory pool
|
||||
- a safe subrange allocator
|
||||
- a ZMQ control plane
|
||||
- a Mooncake transfer-engine data plane
|
||||
|
||||
This allows the connector to support both:
|
||||
|
||||
1. a **fast path** for raw tensors and bytes
|
||||
2. a **generic object path** for arbitrary Python payloads
|
||||
|
||||
Within vLLM-Omni, it is the connector that most directly targets performance-sensitive remote transfer, especially for large payloads and KV cache movement. Its additional complexity is deliberate: it is the connector that turns the generic OmniConnector abstraction into a transport capable of efficient remote memory movement rather than simple object storage.
|
||||
@@ -7,9 +7,10 @@ auto-configured when no explicit connector is specified for an edge.
|
||||
|
||||
## How It Works
|
||||
|
||||
- Small payloads (< threshold): serialized and passed inline in metadata.
|
||||
- Large payloads (>= threshold): stored in shared memory; the SHM name is
|
||||
returned in metadata.
|
||||
All payloads are serialized and stored in shared memory (`/dev/shm`); the SHM
|
||||
segment name is returned in metadata. The configuration exposes a
|
||||
`shm_threshold_bytes` field for a future inline-vs-SHM split, but the current
|
||||
implementation always uses shared memory regardless of payload size.
|
||||
|
||||
## Configuration
|
||||
|
||||
@@ -25,3 +26,234 @@ runtime:
|
||||
## Notes
|
||||
|
||||
- Auto-mode uses SharedMemoryConnector if no connector is declared for an edge.
|
||||
|
||||
---
|
||||
|
||||
## Design
|
||||
|
||||
### 1. Overview
|
||||
|
||||
`SharedMemoryConnector` is the default same-node connector in `vllm_omni/distributed/omni_connectors`. It is designed for stage-to-stage transfer when producer and consumer processes run on the same host and can share `/dev/shm`.
|
||||
|
||||
The connector provides a unified `put()` / `get()` API for arbitrary Python objects while keeping the control plane lightweight:
|
||||
|
||||
- The payload is serialized by the connector.
|
||||
- The serialized bytes are placed in shared memory.
|
||||
- The queue/control plane only carries a small metadata handle.
|
||||
|
||||
This makes `SharedMemoryConnector` the simplest connector in the OmniConnector family and the default fallback when an edge does not explicitly configure another backend.
|
||||
|
||||
### 2. Relationship with the OmniConnector System
|
||||
|
||||
`SharedMemoryConnector` implements `OmniConnectorBase`, so it follows the same lifecycle and API contract as the other connectors:
|
||||
|
||||
- `put(from_stage, to_stage, put_key, data)`
|
||||
- `get(from_stage, to_stage, get_key, metadata=None)`
|
||||
- `cleanup(request_id)`
|
||||
- `health()`
|
||||
- `close()`
|
||||
|
||||
Within the larger system:
|
||||
|
||||
- `load_omni_transfer_config()` automatically fills missing edges with `SharedMemoryConnector`.
|
||||
- Callers interact with the connector exclusively through the `put()` / `get()` / `cleanup()` contract — the connector does not require caller-specific logic.
|
||||
|
||||
Compared with the remote Mooncake-based connectors, `SharedMemoryConnector` is intentionally minimal and local-only.
|
||||
|
||||
### 3. Design Goals
|
||||
|
||||
The connector is built around the following goals:
|
||||
|
||||
- **Low-friction local transfer** for single-node multi-process pipelines.
|
||||
- **Unified object semantics** for arbitrary Python payloads.
|
||||
- **Small control-plane overhead** by passing only metadata through queues.
|
||||
- **Zero external dependencies** beyond Python shared memory and the existing stage utilities.
|
||||
|
||||
It is not intended to provide cross-node transfer, RDMA, or raw tensor zero-copy semantics across processes.
|
||||
|
||||
### 4. Core Design
|
||||
|
||||
#### 4.1 Serialization Model
|
||||
|
||||
`SharedMemoryConnector` always starts from a Python object and serializes it through the shared Omni serializer:
|
||||
|
||||
```python
|
||||
payload = self.serialize_obj(data)
|
||||
```
|
||||
|
||||
This keeps the connector behavior consistent with the rest of the connector stack:
|
||||
|
||||
- producer code does not need connector-specific serialization logic
|
||||
- consumer code always receives the original object after deserialization
|
||||
- the connector can reuse the same serializer used by other backends
|
||||
|
||||
#### 4.2 Shared Memory as the Data Plane
|
||||
|
||||
The actual data plane is a shared-memory segment created by:
|
||||
|
||||
- `shm_write_bytes(...)`
|
||||
- `shm_read_bytes(...)`
|
||||
|
||||
The connector stores a small metadata object such as:
|
||||
|
||||
```python
|
||||
{
|
||||
"shm": {"name": ..., "size": ...},
|
||||
"size": ...
|
||||
}
|
||||
```
|
||||
|
||||
This metadata is passed over the control plane and allows the downstream stage to locate the shared-memory segment.
|
||||
|
||||
#### 4.3 Locking Model
|
||||
|
||||
To avoid races between the producer and consumer, the connector uses a lock file per request:
|
||||
|
||||
```text
|
||||
/dev/shm/shm_{put_key}_lockfile.lock
|
||||
```
|
||||
|
||||
Locking is done with `fcntl.flock`:
|
||||
|
||||
- producer uses `LOCK_EX`
|
||||
- consumer uses `LOCK_EX`
|
||||
|
||||
Both sides acquire an exclusive lock. This ensures that the shared-memory segment is not read while it is still being written and makes the handoff safer in a multi-process environment.
|
||||
|
||||
### 5. Put / Get Flow
|
||||
|
||||
#### 5.1 Producer Flow: `put()`
|
||||
|
||||
The producer-side flow is:
|
||||
|
||||
1. Serialize the input object to bytes.
|
||||
2. Compute the payload size.
|
||||
3. Acquire the per-request lock file.
|
||||
4. Write the bytes into shared memory.
|
||||
5. Return lightweight metadata to the caller.
|
||||
|
||||
The returned tuple is:
|
||||
|
||||
```python
|
||||
(success, serialized_size, metadata)
|
||||
```
|
||||
|
||||
where `metadata` contains the shared-memory handle needed by the consumer.
|
||||
|
||||
#### 5.2 Consumer Flow: `get(metadata=...)`
|
||||
|
||||
The primary consumer path is metadata-driven:
|
||||
|
||||
1. Extract the shared-memory handle from `metadata`.
|
||||
2. Acquire the exclusive lock.
|
||||
3. Read the raw bytes from shared memory.
|
||||
4. Deserialize the bytes back into the original Python object.
|
||||
5. Remove the lock file if it still exists.
|
||||
|
||||
This is the path used by the current stage-to-stage connector flow.
|
||||
|
||||
#### 5.3 Compatibility Flow: `get(metadata=None)`
|
||||
|
||||
The connector also keeps a compatibility path for callers that only know the key:
|
||||
|
||||
1. Attempt to open the shared-memory segment by name via `SharedMemory(name=get_key)`.
|
||||
2. If the segment exists and has non-zero size, acquire the exclusive lock and read the bytes.
|
||||
3. Deserialize the bytes and return the object.
|
||||
|
||||
If the segment does not exist or any exception occurs, the call returns `None` immediately. There is no retry loop in this path -- it is a single-attempt open.
|
||||
|
||||
This path is mainly for older code paths and is not the preferred mode for the current connector pipeline.
|
||||
|
||||
### 6. Key Implementation Characteristics
|
||||
|
||||
#### 6.1 Threshold Exists, but the Current Code Always Uses SHM
|
||||
|
||||
The class keeps a `shm_threshold_bytes` field and still exposes metrics for inline writes. However, the current implementation uses:
|
||||
|
||||
```python
|
||||
if True:
|
||||
...
|
||||
```
|
||||
|
||||
inside `put()`, which means the current code path always writes to shared memory.
|
||||
|
||||
So the design still suggests a future split between:
|
||||
|
||||
- small payloads inline
|
||||
- large payloads in shared memory
|
||||
|
||||
but the current behavior is effectively:
|
||||
|
||||
- all payloads go through shared memory
|
||||
|
||||
This should be documented because it affects real runtime behavior.
|
||||
|
||||
#### 6.2 Cleanup Is Currently Passive
|
||||
|
||||
`cleanup()` is currently a no-op. The intended assumption is:
|
||||
|
||||
- the consumer reads the segment
|
||||
- the underlying shared-memory helpers unlink it
|
||||
|
||||
If the consumer never executes `get()`, the shared-memory segment may remain allocated. This means the connector relies on the normal success path for resource reclamation.
|
||||
|
||||
#### 6.3 Close Is Currently Minimal
|
||||
|
||||
`close()` is also a no-op. There is no connector-owned background thread, socket, or memory pool to tear down, so the lifecycle is simple. The trade-off is that `close()` does not scan or recover leaked shared-memory resources.
|
||||
|
||||
### 7. Data Flow in the Pipeline
|
||||
|
||||
The typical flow with `SharedMemoryConnector` is:
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant SenderStage
|
||||
participant SharedMemoryConnector
|
||||
participant QueueOrControlPlane
|
||||
participant ReceiverStage
|
||||
|
||||
SenderStage->>SharedMemoryConnector: put(from_stage, to_stage, put_key, data)
|
||||
SharedMemoryConnector->>SharedMemoryConnector: serialize object
|
||||
SharedMemoryConnector->>SharedMemoryConnector: write bytes to /dev/shm
|
||||
SharedMemoryConnector-->>SenderStage: metadata {shm, size}
|
||||
SenderStage->>QueueOrControlPlane: forward connector metadata
|
||||
QueueOrControlPlane->>ReceiverStage: task + connector metadata
|
||||
ReceiverStage->>SharedMemoryConnector: get(from_stage, to_stage, get_key, metadata)
|
||||
SharedMemoryConnector->>SharedMemoryConnector: read bytes from /dev/shm
|
||||
SharedMemoryConnector->>SharedMemoryConnector: deserialize object
|
||||
SharedMemoryConnector-->>ReceiverStage: (data, size)
|
||||
```
|
||||
|
||||
This is a classic split-control-plane / data-plane design, but constrained to a single host.
|
||||
|
||||
### 8. Strengths and Trade-offs
|
||||
|
||||
#### Strengths
|
||||
|
||||
- Very simple deployment model.
|
||||
- No external service dependency.
|
||||
- Fits naturally into the existing queue-driven orchestration flow.
|
||||
- Good default for local multi-process pipelines.
|
||||
|
||||
#### Trade-offs
|
||||
|
||||
- Same-node only.
|
||||
- Full object serialization and deserialization are still required.
|
||||
- Resource cleanup depends on the normal consumer path.
|
||||
- Shared memory capacity is limited by host configuration.
|
||||
|
||||
### 9. Summary
|
||||
|
||||
`SharedMemoryConnector` is the baseline local transport for the OmniConnector system. Its design is intentionally straightforward:
|
||||
|
||||
- serialize object
|
||||
- place bytes in shared memory
|
||||
- pass metadata through the control plane
|
||||
- deserialize on the receiving side
|
||||
|
||||
It plays two important roles in vLLM-Omni:
|
||||
|
||||
1. It is the simplest production-ready connector for same-node stage pipelines.
|
||||
2. It serves as the automatic fallback connector when no explicit edge transport is configured.
|
||||
|
||||
Although the current implementation is deliberately minimal, it provides the foundation for reliable local connector semantics and keeps the stage communication model uniform across the system.
|
||||
@@ -94,8 +94,265 @@ stage_args:
|
||||
Parameters:
|
||||
|
||||
- host: datasystem worker host.
|
||||
- port: datasystem worker port.
|
||||
- port: datasystem worker port (default: `35001` if omitted; the example above uses `31501` to match the worker startup command).
|
||||
- get_sub_timeout_ms: get timeout in milliseconds (0 for no timeout).
|
||||
|
||||
For more details, refer to the
|
||||
[Yuanrong Datasystem repository](https://atomgit.com/openeuler/yuanrong-datasystem).
|
||||
|
||||
---
|
||||
|
||||
## Design
|
||||
|
||||
### 1. Overview
|
||||
|
||||
`YuanrongConnector` is the Datasystem-based remote connector in `vllm_omni/distributed/omni_connectors`. It uses Yuanrong Datasystem's distributed key-value client as the transport backend and exposes the same `put()` / `get()` interface as the other OmniConnectors.
|
||||
|
||||
Like `MooncakeStoreConnector`, it is a store-oriented remote connector rather than a direct peer-to-peer transport. Its role is to let stage payloads move across nodes through a deterministic key-based storage abstraction while keeping the rest of the pipeline on the common connector API.
|
||||
|
||||
It is intended for deployments that already use Yuanrong Datasystem and want a remote connector that integrates with the existing OmniConnector configuration and orchestration model.
|
||||
|
||||
### 2. Relationship with the OmniConnector System
|
||||
|
||||
`YuanrongConnector` implements `OmniConnectorBase`, so it participates in the same connector lifecycle as the other backends:
|
||||
|
||||
- `OmniConnectorFactory` constructs it from a `ConnectorSpec`
|
||||
- stage edge configuration is resolved by `load_omni_transfer_config()`
|
||||
- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
|
||||
|
||||
This means the connector is not exposed directly to stage logic. Stages only interact with the generic connector contract, and the backend choice remains a configuration concern.
|
||||
|
||||
### 3. Design Goals
|
||||
|
||||
The connector is built around the following goals:
|
||||
|
||||
1. **Cross-node payload transfer through Datasystem**
|
||||
Reuse Yuanrong Datasystem as the remote exchange medium for stage data.
|
||||
|
||||
2. **Uniform object transfer semantics**
|
||||
Allow arbitrary Python objects to be transmitted through the shared Omni serializer.
|
||||
|
||||
3. **Minimal connector-specific control plane**
|
||||
Use deterministic keys so that consumers can retrieve data without an extra transport metadata handoff.
|
||||
|
||||
4. **Operational reuse of existing infrastructure**
|
||||
Fit into environments that already deploy Yuanrong Datasystem workers and etcd.
|
||||
|
||||
The connector is not designed for direct remote-memory writes or tensor-specific fast-path transfer.
|
||||
|
||||
### 4. Core Design
|
||||
|
||||
#### 4.1 Store-Oriented Transfer Model
|
||||
|
||||
`YuanrongConnector` treats the transport backend as a distributed object store:
|
||||
|
||||
1. serialize the Python payload
|
||||
2. build a deterministic connector key
|
||||
3. write the serialized bytes into Datasystem
|
||||
4. read the bytes back on the receiver side
|
||||
5. deserialize them into the original object
|
||||
|
||||
This is the same broad architectural class as `MooncakeStoreConnector`, but implemented on top of Yuanrong Datasystem APIs instead of Mooncake store APIs.
|
||||
|
||||
#### 4.2 Deterministic Keying
|
||||
|
||||
Unlike the default `_make_key()` in `OmniConnectorBase`, `YuanrongConnector` defines its own key format:
|
||||
|
||||
```text
|
||||
{request_id}:{from_stage}_{to_stage}
|
||||
```
|
||||
|
||||
This has two design implications:
|
||||
|
||||
- the request identifier remains the primary lookup handle
|
||||
- stage routing information is embedded in the key so that the same logical request ID can safely appear on different edges
|
||||
|
||||
The explicit override also makes the key format easier to align with Datasystem-side debugging and operational inspection.
|
||||
|
||||
#### 4.3 No Extra Metadata Hand-off
|
||||
|
||||
`put()` returns:
|
||||
|
||||
```python
|
||||
(success, serialized_size, None)
|
||||
```
|
||||
|
||||
and does not generate connector-specific metadata.
|
||||
|
||||
This design works because the receiver can reconstruct the exact same key from:
|
||||
|
||||
- `get_key`
|
||||
- `from_stage`
|
||||
- `to_stage`
|
||||
|
||||
As a result, the connector does not require a separate side-channel metadata handoff.
|
||||
|
||||
### 5. Initialization
|
||||
|
||||
#### 5.1 Datasystem Client Dependency
|
||||
|
||||
The connector requires the Datasystem Python bindings to expose:
|
||||
|
||||
- `KVClient`
|
||||
- `SetParam`
|
||||
- `WriteMode`
|
||||
|
||||
If any of these symbols are unavailable, connector construction fails immediately with `ImportError`. This keeps configuration errors explicit and avoids a partially initialized runtime.
|
||||
|
||||
#### 5.2 Client Setup
|
||||
|
||||
During `_init_client()`, the connector:
|
||||
|
||||
1. reads `host` and `port`
|
||||
2. creates `KVClient(host, port)`
|
||||
3. calls `client.init()`
|
||||
|
||||
At construction time it also creates a `SetParam` and fixes:
|
||||
|
||||
```python
|
||||
self.set_param.write_mode = WriteMode.NONE_L2_CACHE_EVICT
|
||||
```
|
||||
|
||||
This means the connector has a stable write policy for all writes and does not currently expose write-mode selection as a higher-level connector option.
|
||||
|
||||
### 6. Put / Get Flow
|
||||
|
||||
#### 6.1 Producer Flow: `put()`
|
||||
|
||||
The producer-side flow is:
|
||||
|
||||
1. verify that the Datasystem client has been initialized
|
||||
2. serialize the input object with the shared Omni serializer
|
||||
3. build the Datasystem key using the connector-specific `_make_key()`
|
||||
4. call `client.set(key, serialized_data, self.set_param.write_mode)`
|
||||
5. update metrics and return success
|
||||
|
||||
The returned metadata is always `None`, because the Datasystem key itself is the lookup contract between producer and consumer.
|
||||
|
||||
#### 6.2 Consumer Flow: `get()`
|
||||
|
||||
The consumer-side flow is:
|
||||
|
||||
1. verify that the Datasystem client has been initialized
|
||||
2. rebuild the same key with `from_stage`, `to_stage`, and `get_key`
|
||||
3. call:
|
||||
|
||||
```python
|
||||
client.get([key], False, self.get_sub_timeout_ms)
|
||||
```
|
||||
|
||||
4. take the first returned element if present
|
||||
5. deserialize it and return `(data, payload_size)`
|
||||
|
||||
If the returned list is empty or contains no data for the key, `get()` returns `None`.
|
||||
|
||||
### 7. Timeout and Retrieval Semantics
|
||||
|
||||
The connector uses `get_sub_timeout_ms` as its read timeout. Unlike `MooncakeStoreConnector`, which performs an explicit retry loop in Python, `YuanrongConnector` delegates waiting behavior more directly to the Datasystem client call.
|
||||
|
||||
This leads to a slightly different retrieval model:
|
||||
|
||||
- `MooncakeStoreConnector`: retry-oriented polling in connector code
|
||||
- `YuanrongConnector`: single client call with backend-managed timeout behavior
|
||||
|
||||
From the connector API perspective the result is the same, but operationally the waiting behavior is more dependent on Datasystem client semantics.
|
||||
|
||||
### 8. Integration with Stage Communication
|
||||
|
||||
All callers use the connector through the same `put()` / `get()` contract:
|
||||
|
||||
- the sender calls `put()` to serialize and store the payload
|
||||
- the receiver calls `get()` to retrieve and deserialize it
|
||||
- no connector-specific metadata is required, since the Datasystem key is the rendezvous point
|
||||
|
||||
Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, and there is no raw tensor fast path, which makes the connector functional but not optimized for the largest payloads.
|
||||
|
||||
### 9. Data Flow in the Pipeline
|
||||
|
||||
The end-to-end transfer model is:
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant SenderStage
|
||||
participant YuanrongConnector
|
||||
participant Datasystem
|
||||
participant ReceiverStage
|
||||
|
||||
SenderStage->>YuanrongConnector: put(from_stage, to_stage, put_key, data)
|
||||
YuanrongConnector->>YuanrongConnector: serialize object
|
||||
YuanrongConnector->>Datasystem: set(key, bytes)
|
||||
YuanrongConnector-->>SenderStage: (success, size, None)
|
||||
|
||||
ReceiverStage->>YuanrongConnector: get(from_stage, to_stage, get_key)
|
||||
YuanrongConnector->>Datasystem: get([key], timeout)
|
||||
YuanrongConnector->>YuanrongConnector: deserialize bytes
|
||||
YuanrongConnector-->>ReceiverStage: (data, size)
|
||||
```
|
||||
|
||||
This is a store-mediated remote connector design with deterministic key lookup and no explicit side-channel metadata exchange.
|
||||
|
||||
### 10. Strengths and Trade-offs
|
||||
|
||||
#### Strengths
|
||||
|
||||
- Reuses existing Yuanrong Datasystem infrastructure.
|
||||
- Keeps the connector contract simple and uniform.
|
||||
- Requires no connector-specific metadata handoff.
|
||||
- Suitable for remote stage transfer in Datasystem-based deployments.
|
||||
|
||||
#### Trade-offs
|
||||
|
||||
- Always pays full serialization and deserialization cost.
|
||||
- Does not support raw bytes / tensor fast-path semantics.
|
||||
- Depends on external Datasystem worker availability.
|
||||
- Retrieval and timeout behavior are partly delegated to the backend client.
|
||||
|
||||
### 11. Important Implementation Characteristics
|
||||
|
||||
#### 11.1 Cleanup Is a No-op
|
||||
|
||||
`cleanup()` only logs a debug message and does not explicitly remove data from Datasystem. The current design assumes backend-side lifecycle or garbage collection rather than request-scoped delete semantics inside the connector.
|
||||
|
||||
This mirrors the same design trade-off seen in other store-based connectors: simplicity over explicit per-request data reclamation.
|
||||
|
||||
#### 11.2 Close Only Releases the Local Client Handle
|
||||
|
||||
`close()` does not perform a remote shutdown. It simply clears the local `client` reference and marks the connector as closed from the process perspective:
|
||||
|
||||
```python
|
||||
self.client = None
|
||||
```
|
||||
|
||||
This is appropriate for a client-based store connector, but it also means that backend resource lifecycle remains outside the connector's control.
|
||||
|
||||
#### 11.3 Error and Timeout Metrics Are Coarse-Grained
|
||||
|
||||
The connector tracks:
|
||||
|
||||
- `puts`
|
||||
- `gets`
|
||||
- `bytes_transferred`
|
||||
- `errors`
|
||||
- `timeouts`
|
||||
|
||||
In the current code, failed `put()` increments `errors`, while `get()` exceptions increment `timeouts`. This is operationally useful, but it does not distinguish between:
|
||||
|
||||
- backend timeout
|
||||
- not-found result
|
||||
- transport failure
|
||||
- deserialization failure
|
||||
|
||||
So the metrics should be read as high-level health indicators, not detailed root-cause diagnostics.
|
||||
|
||||
### 12. Summary
|
||||
|
||||
`YuanrongConnector` is the Datasystem-backed remote connector in the OmniConnector family. Its design is straightforward:
|
||||
|
||||
- serialize payloads
|
||||
- store them under a deterministic stage-qualified key
|
||||
- retrieve them by the same key
|
||||
- deserialize them on the receiving side
|
||||
|
||||
Within vLLM-Omni, it provides a clean integration point for Yuanrong Datasystem-based deployments while preserving the same connector abstraction used by the rest of the pipeline.
|
||||
|
||||
Its design priorities are simplicity, infrastructure reuse, and API consistency rather than direct-memory transport optimization.
|
||||
@@ -55,6 +55,8 @@ The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The
|
||||
- Tensor operations happen at `nn.Module` boundaries
|
||||
- Predictable sharding/gathering patterns
|
||||
|
||||
This is the ideal approach for integrating sequence parallelism into new models, as it is easier to maintain and ensure compatibility with other types of acceleration.
|
||||
|
||||
**How it works:**
|
||||
1. Declare `_sp_plan` dict in your transformer class
|
||||
2. Framework automatically applies hooks when `sequence_parallel_size > 1`
|
||||
@@ -201,6 +203,36 @@ class TransformerWithRoPE(nn.Module):
|
||||
}
|
||||
```
|
||||
|
||||
**Pattern 3: Shard RoPE for Dual Stream Attention**
|
||||
In some cases, different streams in attention may need to handle sequence parallelism differently. For example, we may want to shard the image embeddings, while replicating the text embeddings to correctly configure joint attention.
|
||||
|
||||
```python
|
||||
class DualStreamTransformer(nn.Module):
|
||||
"""
|
||||
Dual-stream model where we need to replicate the text components, but shard
|
||||
the image components to correctly handle sequence parallelism.
|
||||
"""
|
||||
_sp_plan = {
|
||||
# In this case, the rope_preparer returns a tuple of len 4, where the
|
||||
# first 2 items correspond to the text, and the second 2 correspond to
|
||||
# visual inputs, so we only shard the second.
|
||||
"rope_preparer": {
|
||||
# Outputs 0, 1 (text) - NOT sharded (replicated)
|
||||
# Outputs 2, 3 (image) - sharded
|
||||
2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), # img_cos
|
||||
3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), # img_sin
|
||||
},
|
||||
# Shard transformer block INPUT
|
||||
"transformer_blocks.0": {
|
||||
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
|
||||
},
|
||||
# Gather at output
|
||||
"proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
|
||||
}
|
||||
```
|
||||
|
||||
NOTE: be careful to test adequately when refactoring classes that take this style of plan, as changing the order of the return values will break sequence parallelism.
|
||||
|
||||
### API Reference
|
||||
|
||||
**SequenceParallelInput Parameters:**
|
||||
@@ -240,7 +272,7 @@ class TransformerWithRoPE(nn.Module):
|
||||
|
||||
## Approach 2: Intrusive Modification (For Complex Cases)
|
||||
|
||||
For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls. Importantly, when taking this approach, be careful to ensure that you correctly manage the `_sp_shard_depth`; if the sequence parallel shard depth is 0, Ulysses will not be used.
|
||||
For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls.
|
||||
|
||||
|
||||
**When to use:**
|
||||
@@ -253,18 +285,15 @@ from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
|
||||
|
||||
def forward(self, hidden_states, ...):
|
||||
if self.parallel_config.sequence_parallel_size > 1:
|
||||
# <Increment the _sp_shard_depth on the fwd context>
|
||||
hidden_states = sp_shard(hidden_states, dim=1)
|
||||
|
||||
# ... computation ...
|
||||
|
||||
if self.parallel_config.sequence_parallel_size > 1:
|
||||
output = sp_gather(output, dim=1)
|
||||
# <Decrement the _sp_shard_depth on the fwd context>
|
||||
|
||||
return output
|
||||
```
|
||||
Note that currently, `sp_shard` / `sp_gather` do *not* automatically manage the `_sp_shard_depth`; you need to be careful to manage it yourself.
|
||||
|
||||
---
|
||||
|
||||
@@ -439,6 +468,7 @@ Complete examples in the codebase:
|
||||
|
||||
| Model | Path | Pattern | Notes |
|
||||
|-------|------|---------|-------|
|
||||
| **LongCat** | `vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py` | Dual-stream | Text components replicated, image components sharded |
|
||||
| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Dual-stream + preprocessing | auto_pad, separate RoPE |
|
||||
| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Dual-Transformer + RoPE | Video transformer |
|
||||
| **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Unified sequence | Concatenated input |
|
||||
@@ -453,7 +483,7 @@ Complete examples in the codebase:
|
||||
Adding Sequence Parallel support to a transformer:
|
||||
|
||||
1. ✅ **Choose approach** - Use `_sp_plan` for standard cases, intrusive modification for complex cases
|
||||
2. ✅ **Identify sharding boundaries** - Where should tensors be split/gathered?
|
||||
2. ✅ **Identify sharding boundaries** - Where should tensors be split/gathered? And which module boundaries need to be moved to facilitate this?
|
||||
3. ✅ **Extract inline operations** - Move `torch.cat`, `pad_sequence`, etc. to submodules
|
||||
4. ✅ **Define `_sp_plan`** - Declare shard/gather points as class attribute
|
||||
5. ✅ **Use `auto_pad` for variable lengths** - Support non-uniform sequences
|
||||
|
||||
@@ -85,9 +85,6 @@ classDiagram
|
||||
class InputProcessor {
|
||||
+process_inputs() EngineCoreRequest
|
||||
}
|
||||
class OmniInputProcessor {
|
||||
+process_inputs() OmniEngineCoreRequest
|
||||
}
|
||||
|
||||
class VLLMOutputProcessor {
|
||||
+process_outputs() OutputProcessorOutput
|
||||
@@ -96,7 +93,6 @@ classDiagram
|
||||
+process_outputs() OutputProcessorOutput
|
||||
+_route_and_normalize()
|
||||
}
|
||||
InputProcessor <|-- OmniInputProcessor
|
||||
VLLMOutputProcessor <|-- MultimodalOutputProcessor
|
||||
```
|
||||
|
||||
@@ -105,7 +101,7 @@ classDiagram
|
||||
- **Scheduler**: `OmniARScheduler` extends `vllm.v1.core.sched.scheduler.Scheduler` to enrich scheduled requests with omni-specific payloads
|
||||
- **Worker**: `GPUARWorker` extends `vllm.v1.worker.gpu_worker.Worker` to initialize AR-specific model runners
|
||||
- **ModelRunner**: `GPUARModelRunner` extends `OmniGPUModelRunner` → `vllm.v1.worker.gpu_model_runner.GPUModelRunner` to expose hidden states and handle multimodal outputs
|
||||
- **InputProcessor**: `OmniInputProcessor` extends `vllm.v1.engine.input_processor.InputProcessor` to serialize prompt embeddings and additional information
|
||||
- **InputProcessor**: Stage-0 uses upstream `vllm.v1.engine.input_processor.InputProcessor`; `AsyncOmniEngine` then restores omni-specific payloads (for example `additional_information` and `prompt_embeds`) when building `OmniEngineCoreRequest`
|
||||
- **OutputProcessor**: `MultimodalOutputProcessor` extends `vllm.v1.engine.output_processor.OutputProcessor` to route and accumulate multimodal outputs
|
||||
|
||||
## 3. Scheduler Design
|
||||
@@ -118,7 +114,7 @@ The following diagram illustrates the request flow through the AR module compone
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
A[OmniInputProcessor] -->|OmniEngineCoreRequest| B[OmniARScheduler]
|
||||
A[InputProcessor stage-0 in AsyncOmniEngine] -->|EngineCoreRequest then upgraded to OmniEngineCoreRequest| B[OmniARScheduler]
|
||||
B -->|schedule: OmniNewRequestData| C[GPUARWorker]
|
||||
C -->|SchedulerOutput| D[GPUARModelRunner]
|
||||
D -->|execute_model: None| E[Model Forward Pass]
|
||||
@@ -297,15 +293,17 @@ The input/output processing pipeline handles serialization, routing, and accumul
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant Client
|
||||
participant OmniInputProcessor
|
||||
participant AsyncOmniEngine
|
||||
participant InputProcessor
|
||||
participant Scheduler
|
||||
participant ModelRunner
|
||||
participant MultimodalOutputProcessor
|
||||
participant Client
|
||||
|
||||
Client->>OmniInputProcessor: prompt + prompt_embeds + additional_info
|
||||
OmniInputProcessor->>OmniInputProcessor: Serialize tensors to payloads
|
||||
OmniInputProcessor->>Scheduler: OmniEngineCoreRequest (with payloads)
|
||||
Client->>AsyncOmniEngine: prompt + prompt_embeds + additional_info
|
||||
AsyncOmniEngine->>InputProcessor: process_inputs()
|
||||
InputProcessor->>Scheduler: EngineCoreRequest
|
||||
AsyncOmniEngine->>AsyncOmniEngine: _upgrade_to_omni_request() + serialize_additional_information()
|
||||
Scheduler->>ModelRunner: OmniNewRequestData (with payloads)
|
||||
ModelRunner->>ModelRunner: Decode payloads → CPU tensors
|
||||
ModelRunner->>ModelRunner: Overlay prompt_embeds on inputs_embeds
|
||||
@@ -318,43 +316,18 @@ sequenceDiagram
|
||||
MultimodalOutputProcessor->>Client: RequestOutput (with multimodal_output)
|
||||
```
|
||||
|
||||
### OmniInputProcessor
|
||||
### Stage-0 Input Processing
|
||||
|
||||
`OmniInputProcessor` extends the base input processor to serialize prompt embeddings and additional information for inter-stage transfer.
|
||||
|
||||
#### Payload Serialization
|
||||
|
||||
Converts PyTorch tensors to serialized payloads:
|
||||
Stage-0 now uses upstream `InputProcessor` directly, and `AsyncOmniEngine` upgrades the request to `OmniEngineCoreRequest` while restoring omni-specific payloads.
|
||||
|
||||
```python
|
||||
def process_inputs(self, ...) -> OmniEngineCoreRequest:
|
||||
# Serialize prompt_embeds
|
||||
if "prompt_embeds" in decoder_inputs:
|
||||
pe_cpu = decoder_inputs["prompt_embeds"].detach().to("cpu").contiguous()
|
||||
prompt_embeds_payload = PromptEmbedsPayload(
|
||||
data=pe_cpu.numpy().tobytes(),
|
||||
shape=[seq_len, hidden_size],
|
||||
dtype=dtype_str,
|
||||
)
|
||||
|
||||
# Serialize additional_information
|
||||
if "additional_information" in decoder_inputs:
|
||||
entries = {}
|
||||
for key, value in raw_info.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
entry = AdditionalInformationEntry(
|
||||
tensor_data=value.numpy().tobytes(),
|
||||
tensor_shape=list(value.shape),
|
||||
tensor_dtype=dtype_str,
|
||||
)
|
||||
entries[key] = entry
|
||||
additional_information_payload = AdditionalInformationPayload(entries=entries)
|
||||
|
||||
return OmniEngineCoreRequest(
|
||||
# ... standard fields ...
|
||||
prompt_embeds=prompt_embeds_payload,
|
||||
additional_information=additional_information_payload,
|
||||
)
|
||||
request = self.input_processor.process_inputs(
|
||||
request_id=request_id,
|
||||
prompt=prompt,
|
||||
params=params,
|
||||
supported_tasks=self.supported_tasks,
|
||||
)
|
||||
request = _upgrade_to_omni_request(request, prompt)
|
||||
```
|
||||
|
||||
### MultimodalOutputProcessor
|
||||
@@ -400,13 +373,13 @@ The AR module of vLLM-Omni extends vLLM through strategic inheritance and minima
|
||||
### Key Design Patterns
|
||||
|
||||
1. **Inheritance over composition**: Extends vLLM classes to preserve compatibility with existing scheduling, batching, and execution mechanisms
|
||||
2. **Payload serialization**: Uses `PromptEmbedsPayload` and `AdditionalInformationPayload` for efficient inter-stage data transfer
|
||||
2. **Payload serialization**: Uses serialized `additional_information` payloads together with prompt-embedding handoff for efficient inter-stage transfer
|
||||
3. **Two-phase execution**: Maintains vLLM's execute/sample separation for AR models while supporting single-phase execution for generation models
|
||||
4. **Multimodal routing**: Routes outputs by `output_type` and accumulates tensors incrementally to support streaming
|
||||
|
||||
### Differences from vLLM
|
||||
|
||||
- **Payload support**: Serialized prompt embeddings and additional information enable direct transfer between pipeline stages
|
||||
- **Payload support**: Serialized additional information and prompt embeddings enable direct transfer between pipeline stages
|
||||
- **Multimodal handling**: Extended input/output processors support images, audio, and other modalities alongside text
|
||||
- **Hidden state exposure**: AR model runners expose per-request hidden states via `pooler_output` for downstream consumption
|
||||
- **Generation scheduler**: Fast-path scheduling for basic heterogeneous architectures that complete in one step
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
# AsyncOmni Architecture (Qwen3-Omni Example)
|
||||
|
||||
## 1. System Architecture
|
||||
|
||||
```text
|
||||
• ┌─────────────────────────────────────────────────────────────────────────────────┐
|
||||
│ API Layer │
|
||||
│ ┌─────────────────────────────────────┐ ┌──────────────────────────────────┐ │
|
||||
│ │ AsyncOmni (EngineClient) │ │ Omni │ │
|
||||
│ │ • generate() / abort() / shutdown() │ │ • generate() │ │
|
||||
│ │ • _final_output_handler() │ │ | │
|
||||
│ └─────────────────────────────────────┘ └──────────────────────────────────┘ │
|
||||
├─────────────────────────────────────────────────────────────────────────────────┤
|
||||
│ Engine Layer (Proxy) │
|
||||
│ ┌───────────────────────────────────────────────────────────────────────────┐ │
|
||||
│ │ AsyncOmniEngine │ │
|
||||
│ │ • _bootstrap_orchestrator() & _initialize_stages() │ │
|
||||
│ │ • add_request() / add_request_async() -> input_processor.process_inputs() │ │
|
||||
│ │ • try_get_output() / try_get_output_async() │ │
|
||||
│ └───────────────────┬─────────────────────────────────▲─────────────────────┘ │
|
||||
│ request_queue (janus.Queue) output_queue (janus.Queue) │
|
||||
├──────────────────────┼─────────────────────────────────┼────────────────────────┤
|
||||
│ ▼ Orchestration Layer │ │
|
||||
│ ┌───────────────────────────────────────────────────────────────────────────┐ │
|
||||
│ │ Orchestrator [background thread] │ │
|
||||
│ │ • _request_handler() │ │
|
||||
│ │ - stage_client.add_request_async() & _prewarm_async_chunk_stages() │ │
|
||||
│ │ • _orchestration_output_handler() │ │
|
||||
│ │ - _process_stage_outputs() -> output_processors[i].process_outputs() │ │
|
||||
│ │ - _route_output() & _forward_to_next_stage() │ │
|
||||
│ └──────────┬─────────────────────────┬────────────────────────┬─────────────┘ │
|
||||
├─────────────┼─────────────────────────┼────────────────────────┼────────────────┤
|
||||
│ │ Communication Layer │ │
|
||||
│ ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐ │
|
||||
│ │ StageEngineCoreClient │ │ StageEngineCoreClient │ │ StageDiffusionClient │ │
|
||||
│ │ • ZMQ ROUTER / PULL │ │ • ZMQ ROUTER / PULL │ │ • ZMQ ROUTER / PULL │ │
|
||||
│ │ • Msgpack codec │ │ • Msgpack codec │ │ • Msgpack codec │ │
|
||||
│ └──────────┬────────────┘ └──────────┬────────────┘ └──────────┬────────────┘ │
|
||||
│ ▼ ZMQ IPC ▼ ZMQ IPC ▼ ZMQ IPC │
|
||||
├─────────────────────────────────────────────────────────────────────────────────┤
|
||||
│ Execution Layer │
|
||||
│ ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐ │
|
||||
│ │ StageCoreProc │ │ StageCoreProc │ │ DiffusionEngine │ │
|
||||
│ │ [background process] │ │ [background process] │ │ [background process] │ │
|
||||
│ └───────────────────────┘ └───────────────────────┘ └───────────────────────┘ │
|
||||
└─────────────────────────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
## 2. Execution Flow (Arrow Steps, one generate request)
|
||||
|
||||
```text
|
||||
[1] App
|
||||
-> AsyncOmni.generate(prompt, request_id)
|
||||
|
||||
[2] AsyncOmni
|
||||
-> _final_output_handler() (started on first request)
|
||||
-> AsyncOmniEngine.add_request(stage_id=0, ...)
|
||||
|
||||
[3] AsyncOmniEngine.add_request
|
||||
-> (if stage-0 is llm and input is not EngineCoreRequest)
|
||||
InputProcessor.process_inputs()
|
||||
OutputProcessor[0].add_request()
|
||||
-> request_queue.put(add_request_msg)
|
||||
|
||||
[4] Orchestrator._request_handler
|
||||
-> _handle_add_request(msg)
|
||||
-> stage_clients[0].add_request_async(...)
|
||||
|
||||
[5] Orchestrator._orchestration_loop (loop)
|
||||
-> poll stage output
|
||||
- llm stage: await get_output_async()
|
||||
- diffusion stage: get_diffusion_output_async()
|
||||
-> (llm stage) output_processors[i].process_outputs(...)
|
||||
-> _route_output(...)
|
||||
-> if finished and not final_stage and non-async-chunk:
|
||||
_forward_to_next_stage(...)
|
||||
-> next_stage.add_request_async(...)
|
||||
-> output_queue.put(output)
|
||||
|
||||
[6] AsyncOmni._final_output_loop (background coroutine)
|
||||
-> AsyncOmniEngine.try_get_output_async()
|
||||
-> route by request_id to ClientRequestState.queue
|
||||
|
||||
[7] AsyncOmni._process_orchestrator_results
|
||||
-> read from ClientRequestState.queue
|
||||
-> _process_single_result(...)
|
||||
-> yield OmniRequestOutput
|
||||
|
||||
[8] Exit condition
|
||||
-> receive result["finished"] == True
|
||||
-> generate() ends
|
||||
```
|
||||
|
||||
## 3. Runtime Sequence (one generate request)
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant APP as App
|
||||
participant AO as AsyncOmni
|
||||
participant ENG as AsyncOmniEngine
|
||||
participant ORCH as Orchestrator
|
||||
participant S0 as Stage-0 Client
|
||||
participant SN as Next Stage Client
|
||||
|
||||
APP->>AO: generate
|
||||
AO->>AO: start output_handler once
|
||||
AO->>ENG: add_request(stage_id=0, ...)
|
||||
ENG->>ENG: input_processor.process_inputs()
|
||||
ENG->>ORCH: request_queue.put(add_request)
|
||||
|
||||
ORCH->>ORCH: _handle_add_request
|
||||
ORCH->>S0: add_request_async
|
||||
|
||||
loop poll route forward
|
||||
ORCH->>S0: get_output_async / get_diffusion_output_async
|
||||
ORCH->>ORCH: _route_output
|
||||
alt need forward to next stage
|
||||
ORCH->>SN: add_request_async
|
||||
end
|
||||
ORCH-->>ENG: output_queue.put
|
||||
end
|
||||
|
||||
AO->>ENG: try_get_output_async
|
||||
ENG-->>AO: message
|
||||
AO-->>APP: yield OmniRequestOutput
|
||||
```
|
||||
|
||||
## 4. Comparison
|
||||
|
||||
Previous topology (reference):
|
||||
|
||||
```text
|
||||
┌────────────────────────────────────────────────────────────────────────────┐
|
||||
│ Main Process │
|
||||
│ ┌──────────────────────┐ ┌────────────────────────────────────────────┐ │
|
||||
│ │ generate() │ │ final_output_handler() │ │
|
||||
│ └──────────────────────┘ └────────────────────────────────────────────┘ │
|
||||
└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
|
||||
mp.Queue (in_q/out_q) mp.Queue (in_q/out_q) mp.Queue (in_q/out_q)
|
||||
▼▲ ▼▲ ▼▲
|
||||
┌───────────────────────┐ ┌───────────────────────┐ ┌──────────────────────┐
|
||||
│ Worker Proc-0 │ │ Worker Proc-1 │ │ Worker Proc-2 │
|
||||
│ (Thinker LLM) │ │ (Talker LLM) │ │ (Vocoder) │
|
||||
│ ┌────────────────┐ │ │ ┌────────────────┐ │ │ ┌────────────────┐ │
|
||||
│ │_stage_worker │ │ │ │_stage_worker │ │ │ │_stage_worker │ │
|
||||
│ │_async() │ │ │ │_async() │ │ │ │_async() │ │
|
||||
│ └────────────────┘ │ │ └────────────────┘ │ │ └────────────────┘ │
|
||||
│ ┌────────────────┐ │ │ ┌────────────────┐ │ │ ┌────────────────┐ │
|
||||
│ │output_handler()│ │ │ │output_handler()│ │ │ │output_handler()│ │
|
||||
│ └────────────────┘ │ │ └────────────────┘ │ │ └────────────────┘ │
|
||||
└──────────┬────────────┘ └──────────┬────────────┘ └──────────┬───────────┘
|
||||
ZMQ ▼ ▲ ZMQ ZMQ ▼ ▲ ZMQ ZMQ ▼ ▲ ZMQ
|
||||
┌──────────────────────┐ ┌──────────────────────┐ ┌──────────────────────┐
|
||||
│ EngineCore Proc-0 │ │ EngineCore Proc-1 │ │ EngineCore Proc-2 │
|
||||
│ (Thinker) │ │ (Talker) │ │ (Vocoder) │
|
||||
└──────────────────────┘ └──────────────────────┘ └──────────────────────┘
|
||||
```
|
||||
|
||||
Current topology:
|
||||
|
||||
```text
|
||||
┌────────────────────────────────────────────────────────────────────────────┐
|
||||
│ Main Process │
|
||||
│ ┌──────────────────────────────────────────────────────────────────────┐ │
|
||||
│ │ Main Thread │ │
|
||||
│ │ ┌──────────────────────┐ ┌─────────────────────────────────────┐ │ │
|
||||
│ │ │ generate() │ │ final_output_handler() │ │ │
|
||||
│ │ └──────────────────────┘ └─────────────────────────────────────┘ │ │
|
||||
│ └──────────────────────────────────────────────────────────────────────┘ │
|
||||
│ janus.Queue (request_queue) ▼ ▲ janus.Queue (output_queue) │
|
||||
│ ┌──────────────────────────────────────────────────────────────────────┐ │
|
||||
│ │ Orchestrator Thread │ │
|
||||
│ │ ┌──────────────────────┐ ┌──────────────────────────────────────┐ │ │
|
||||
│ │ │ _request_handler() │ │ _orchestration_output_handler() │ │ │
|
||||
│ │ └──────────────────────┘ └──────────────────────────────────────┘ │ │
|
||||
│ │ ┌────────────────────────────────────────────────────────────────┐ │ │
|
||||
│ │ │ _orchestration_loop(): poll/process/route outputs for all stages│ │ │
|
||||
│ │ └────────────────────────────────────────────────────────────────┘ │ │
|
||||
│ └───────┬─────────────────────────┬─────────────────────────┬──────────┘ │
|
||||
└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
|
||||
ZMQ ▼ ▲ ZMQ ZMQ ▼ ▲ ZMQ ZMQ ▼ ▲ ZMQ
|
||||
┌──────────────────────┐ ┌──────────────────────┐ ┌──────────────────────┐
|
||||
│ EngineCore Proc-0 │ │ EngineCore Proc-1 │ │ EngineCore Proc-2 │
|
||||
│ (Thinker) │ │ (Talker) │ │ (Vocoder) │
|
||||
└──────────────────────┘ └──────────────────────┘ └──────────────────────┘
|
||||
```
|
||||
|
||||
|
||||
Test scripts:
|
||||
```bash
|
||||
# enter offline inference folder.
|
||||
cd examples/offline_inference/qwen2_5_omni
|
||||
python end2end.py --output-dir output_audio --query-type use_mixed_modalities
|
||||
|
||||
cd ../qwen3_omni
|
||||
python end2end.py --output-dir output_audio --query-type text --async-chunk --enable-stats
|
||||
|
||||
cd ../bagel
|
||||
python end2end.py --prompts "A cute cat"
|
||||
|
||||
cd ../text_to_image
|
||||
python text_to_image.py --prompt "a cup of coffee on the table" --output output.png
|
||||
```
|
||||
@@ -47,7 +47,8 @@ If no, check your shell running the ComfyUI process. There may be some error mes
|
||||
This extension offers the following nodes based on the output modalities:
|
||||
|
||||
- **Generate Image** for text-to-image and image-to-image tasks
|
||||
- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
|
||||
- **Generate Video** for text-to-video and image-to-video tasks
|
||||
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
|
||||
- **TTS** and **TTS Voice Clone** for TTS tasks
|
||||
|
||||
This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
|
||||
@@ -59,7 +60,7 @@ To build a simple workflow yourself,
|
||||
|
||||
- Drag a generation node onto the canvas.
|
||||
- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
|
||||
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
|
||||
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
|
||||
- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
|
||||
- For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
|
||||
- For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
|
||||
|
||||
@@ -97,7 +97,7 @@ outputs = omni.generate(
|
||||
)
|
||||
|
||||
# Access custom trajectory data
|
||||
output = outputs[0].request_output[0]
|
||||
output = outputs[0].request_output
|
||||
print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
|
||||
print(f"Trajectory latents shape: {output.latents.shape}")
|
||||
```
|
||||
@@ -115,7 +115,7 @@ python image_edit.py \
|
||||
--image cherry_blossom.jpg \
|
||||
--prompt "Let this mascot dance under the moon, surrounded by floating stars" \
|
||||
--output output_image_edit.png \
|
||||
--num_inference_steps 10
|
||||
--num-inference-steps 10
|
||||
```
|
||||
|
||||
## Advanced Usage
|
||||
|
||||
@@ -30,7 +30,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
|
||||
|
||||
### Pre-built wheels
|
||||
|
||||
Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
|
||||
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, and 0.18.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
|
||||
|
||||
=== "NVIDIA CUDA"
|
||||
|
||||
|
||||
@@ -34,13 +34,13 @@ uv pip install vllm-omni
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.16.0 release version of the library
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.16.0 --torch-backend=auto
|
||||
uv pip install vllm==0.18.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
The release 0.16.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
|
||||
The 0.18.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -56,11 +56,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.16.0
|
||||
git checkout v0.18.0
|
||||
```
|
||||
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
|
||||
```bash
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.16.0/vllm-0.16.0-cp38-abi3-manylinux_2_31_x86_64.whl
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.18.0/vllm-0.18.0+cu129-cp38-abi3-manylinux_2_35_x86_64.whl
|
||||
```
|
||||
Install vllm with command below (If you have no existing PyTorch).
|
||||
```bash
|
||||
@@ -91,7 +91,7 @@ docker run --runtime nvidia --gpus 2 \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
--ipc=host \
|
||||
vllm/vllm-omni:v0.16.0 \
|
||||
vllm/vllm-omni:v0.18.0 \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
|
||||
```
|
||||
|
||||
|
||||
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
|
||||
|
||||
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||
```bash
|
||||
uv pip install vllm==0.16.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
|
||||
uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
|
||||
```
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.16.0 release version of the library
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.16.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
|
||||
uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
|
||||
```
|
||||
|
||||
The release 0.16.0 of vLLM requires ROCm 7.0 environment.
|
||||
The pre-built 0.18.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.16.0
|
||||
git checkout v0.18.0
|
||||
python3 -m pip install -r requirements/rocm.txt
|
||||
python3 setup.py develop
|
||||
```
|
||||
@@ -130,7 +130,7 @@ docker run --rm \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
vllm/vllm-omni-rocm:v0.16.0 \
|
||||
vllm/vllm-omni-rocm:v0.18.0 \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
@@ -149,7 +149,7 @@ docker run --rm -it \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
--entrypoint bash \
|
||||
vllm/vllm-omni-rocm:v0.16.0
|
||||
vllm/vllm-omni-rocm:v0.18.0
|
||||
```
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
@@ -33,18 +33,25 @@ docker run --rm \
|
||||
-p 8000:8000 \
|
||||
-it $IMAGE bash
|
||||
|
||||
cd /vllm-workspace/vllm
|
||||
git pull origin main
|
||||
git fetch origin --tags
|
||||
git checkout v0.16.0
|
||||
|
||||
# Because vllm-ascend will release v0.16.0rc1 after vllm-omni 0.16.0,
|
||||
# we have to pin vllm-ascend at the current commit.
|
||||
cd /vllm-workspace/vllm-ascend
|
||||
git pull origin main
|
||||
git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
|
||||
pip install -v -e .
|
||||
|
||||
# Inside the container, install vLLM-Omni from source
|
||||
cd /vllm-workspace
|
||||
git clone -b v0.16.0 https://github.com/vllm-project/vllm-omni.git
|
||||
|
||||
cd vllm-omni
|
||||
pip install -v -e .
|
||||
pip install -v -e . --no-build-isolation
|
||||
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
|
||||
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
```
|
||||
|
||||
@@ -61,22 +68,22 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
|
||||
You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
|
||||
|
||||
```bash
|
||||
# Pin vLLM version to 0.16.0
|
||||
# Pin vLLM version to 0.18.0
|
||||
cd /vllm-workspace/vllm
|
||||
git pull origin main
|
||||
git fetch origin --tags
|
||||
git checkout v0.16.0
|
||||
git checkout v0.18.0
|
||||
VLLM_TARGET_DEVICE=empty pip install -v -e .
|
||||
|
||||
# Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
|
||||
cd ../vllm-ascend
|
||||
cd /vllm-workspace/vllm-ascend
|
||||
git pull origin main
|
||||
git fetch origin --tags
|
||||
git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
|
||||
git checkout 1e05c4908f31737bc4eef865a9f351d030a77c9d
|
||||
pip install -v -e .
|
||||
|
||||
# Install vLLM-Omni from the latest main branch
|
||||
cd ../vllm-omni
|
||||
cd /vllm-workspace/vllm-omni
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
pip install -v -e . --no-build-isolation
|
||||
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
|
||||
|
||||
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
|
||||
source .venv/bin/activate
|
||||
|
||||
# On CUDA
|
||||
uv pip install vllm==0.16.0 --torch-backend=auto
|
||||
uv pip install vllm==0.18.0 --torch-backend=auto
|
||||
|
||||
# On ROCm
|
||||
uv pip install vllm==0.16.0 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
|
||||
uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
|
||||
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
@@ -42,7 +42,7 @@ if __name__ == "__main__":
|
||||
omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
|
||||
prompt = "a cup of coffee on the table"
|
||||
outputs = omni.generate(prompt)
|
||||
images = outputs[0].request_output[0].images
|
||||
images = outputs[0].request_output.images
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
@@ -70,7 +70,7 @@ if __name__ == "__main__":
|
||||
]
|
||||
omni_outputs = omni.generate(prompts)
|
||||
for i_prompt, prompt_output in enumerate(omni_outputs):
|
||||
this_request_output = prompt_output.request_output[0]
|
||||
this_request_output = prompt_output.request_output
|
||||
this_images = this_request_output.images
|
||||
for i_image, image in enumerate(this_images):
|
||||
image.save(f"p{i_prompt}-img{i_image}.jpg")
|
||||
@@ -82,7 +82,7 @@ if __name__ == "__main__":
|
||||
|
||||
!!! info
|
||||
|
||||
For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
|
||||
For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
|
||||
list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
|
||||
batched inputs, you can [modify this configuration](../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
|
||||
|
||||
|
||||
@@ -141,6 +141,7 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]:
|
||||
"vllm_omni.diffusion.models.qwen_image",
|
||||
"vllm_omni.entrypoints.async_diffusion",
|
||||
"vllm_omni.entrypoints.openai",
|
||||
"vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts",
|
||||
]
|
||||
if any(module_path.startswith(prefix) for prefix in excluded_prefixes):
|
||||
continue
|
||||
|
||||
@@ -132,9 +132,19 @@ class Example:
|
||||
".dat",
|
||||
".db",
|
||||
".sqlite", # Data files
|
||||
".pyc",
|
||||
".pyo",
|
||||
".pyd", # Python compiled
|
||||
".npy",
|
||||
".npz",
|
||||
".pkl",
|
||||
".pickle", # Serialized data
|
||||
}
|
||||
excluded_dirs = {"__pycache__", ".git", "node_modules", ".tox", ".mypy_cache"}
|
||||
|
||||
def is_other_file(file: Path) -> bool:
|
||||
if any(part in excluded_dirs for part in file.parts):
|
||||
return False
|
||||
return file.is_file() and file != self.main_file and file.suffix.lower() not in binary_extensions
|
||||
|
||||
return [file for file in self.path.rglob("*") if is_other_file(file)]
|
||||
|
||||
@@ -28,7 +28,7 @@ th {
|
||||
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
|
||||
| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` |
|
||||
|`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
|
||||
| `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
|
||||
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
|
||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
|
||||
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` |
|
||||
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` |
|
||||
@@ -49,6 +49,12 @@ th {
|
||||
|`GlmImageForConditionalGeneration` | GLM-Image | `zai-org/GLM-Image` |
|
||||
|`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` |
|
||||
|`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` |
|
||||
|`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
|
||||
|`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` |
|
||||
|`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` |
|
||||
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` |
|
||||
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` |
|
||||
|`VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` |
|
||||
|
||||
|
||||
## List of Supported Models for NPU
|
||||
@@ -64,6 +70,7 @@ th {
|
||||
|--------------|--------|-------------------|
|
||||
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
|
||||
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
|
||||
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |
|
||||
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
|
||||
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
|
||||
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
|
||||
@@ -71,7 +78,7 @@ th {
|
||||
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
|
||||
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` |
|
||||
|`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
|
||||
| `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
|
||||
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
|
||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
|
||||
|`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
|
||||
|`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
# Speech API
|
||||
|
||||
vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation using Qwen3-TTS models.
|
||||
vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. Supported TTS models include:
|
||||
|
||||
- **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
|
||||
- **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
|
||||
- **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
|
||||
|
||||
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
|
||||
|
||||
@@ -9,13 +13,30 @@ Each server instance runs a single model (specified at startup via `vllm serve <
|
||||
### Start the Server
|
||||
|
||||
```bash
|
||||
# CustomVoice model (predefined speakers)
|
||||
# Qwen3-TTS: CustomVoice model (predefined speakers)
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
|
||||
# Fish Speech S2 Pro
|
||||
vllm-omni serve fishaudio/s2-pro \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.9
|
||||
|
||||
# Voxtral TTS
|
||||
vllm serve mistralai/Voxtral-4B-TTS-2603 \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
### Generate Speech
|
||||
@@ -96,7 +117,7 @@ Content-Type: application/json
|
||||
| `language` | string | "Auto" | Language (see supported languages below) |
|
||||
| `instructions` | string | "" | Voice style/emotion instructions |
|
||||
| `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
|
||||
| `initial_codec_chunk_frames` | integer | null | Initial chunk size for reduced TTFA (overrides stage config) |
|
||||
| `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
|
||||
|
||||
**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
|
||||
|
||||
@@ -164,6 +185,42 @@ curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
-F "name=custom_voice_1"
|
||||
```
|
||||
|
||||
## Streaming Text Input (WebSocket)
|
||||
|
||||
The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally and generates audio per sentence as boundaries are detected.
|
||||
|
||||
> Note: text input is always streamed incrementally. Audio output remains sentence-scoped:
|
||||
> use `stream_audio=false` for one binary frame per sentence, or `stream_audio=true` for one or more PCM chunks per sentence.
|
||||
|
||||
### WebSocket Protocol
|
||||
|
||||
Client -> Server:
|
||||
|
||||
| Message | Description |
|
||||
|---------|-------------|
|
||||
| `{"type": "session.config", ...}` | Session configuration (sent once, first message) |
|
||||
| `{"type": "input.text", "text": "..."}` | Text chunk |
|
||||
| `{"type": "input.done"}` | End of input, flushes remaining buffer |
|
||||
|
||||
Server -> Client:
|
||||
|
||||
| Message | Description |
|
||||
|---------|-------------|
|
||||
| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for a sentence |
|
||||
| Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
|
||||
| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for a sentence |
|
||||
| `{"type": "session.done", "total_sentences": N}` | Session complete |
|
||||
| `{"type": "error", "message": "..."}` | Non-fatal error |
|
||||
|
||||
### Session Config Parameters
|
||||
|
||||
All REST API parameters are supported, plus:
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `stream_audio` | bool | false | Stream one or more PCM chunks per sentence over WebSocket |
|
||||
| `split_granularity` | string | "sentence" | Text splitting granularity |
|
||||
|
||||
|
||||
```bash
|
||||
DELETE /v1/audio/voices/{name}
|
||||
@@ -281,6 +338,8 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
|
||||
## Supported Models
|
||||
|
||||
### Qwen3-TTS
|
||||
|
||||
| Model | Task Type | Description |
|
||||
|-------|-----------|-------------|
|
||||
| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
|
||||
@@ -289,6 +348,20 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
|
||||
| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
|
||||
|
||||
### Fish Speech S2 Pro
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| `fishaudio/s2-pro` | 4B dual-AR TTS with DAC codec (44.1 kHz). Supports text-to-speech and voice cloning. |
|
||||
|
||||
Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech online serving example](../user_guide/examples/online_serving/fish_speech.md) for details.
|
||||
|
||||
### Voxtral TTS
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| `mistralai/Voxtral-4B-TTS-2603` | 3B AR + FlowMatching TTS. Supports text-to-speech with preset voices. |
|
||||
|
||||
## Error Responses
|
||||
|
||||
### 400 Bad Request
|
||||
|
||||
|
After Width: | Height: | Size: 34 KiB |
|
Before Width: | Height: | Size: 36 KiB After Width: | Height: | Size: 37 KiB |
|
Before Width: | Height: | Size: 33 KiB After Width: | Height: | Size: 45 KiB |
@@ -18,23 +18,28 @@ The following parallelism methods are currently supported in vLLM-Omni:
|
||||
|
||||
6. [HSDP](#hsdp): Hybrid Sharded Data Parallel shards model weights across GPUs using PyTorch FSDP2. This reduces per-GPU memory usage, enabling inference of large models on GPUs with limited memory.
|
||||
|
||||
7. [Expert Parallel](#expert-parallelism): Expert Parallelism shards the Experts of a Mixture-of-Experts (MoE) layer across multiple devices. During the forward, a gating mechanism routes tokens to their designated experts, necessitating cross-cards communication(all-to-all) to dispatch tokens to the correct ranks and combine the results. This parallelism allows for massive scaling of model parameters without a proportional increase in the computational load per device.
|
||||
|
||||
The following table shows which models are currently supported by parallelism method:
|
||||
|
||||
### ImageGen
|
||||
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel |
|
||||
|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|
|
||||
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ |
|
||||
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ |
|
||||
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ |
|
||||
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ |
|
||||
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel | HSDP |
|
||||
|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|:----:|
|
||||
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ❌ |
|
||||
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ❌ |
|
||||
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | N/A | ❌ |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | N/A | ❌ |
|
||||
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ |
|
||||
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ |
|
||||
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ |
|
||||
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | N/A | ❌ |
|
||||
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ | N/A | ❌ |
|
||||
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ✅ |
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | N/A | ✅ |
|
||||
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | ✅ |
|
||||
| **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ |
|
||||
| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ✅ | ❌ | ❌ | N/A | ❌ |
|
||||
|
||||
!!! note "TP Limitations for Diffusion Models"
|
||||
We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
|
||||
@@ -53,6 +58,8 @@ The following table shows which models are currently supported by parallelism me
|
||||
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | VAE-Patch-Parallel |
|
||||
|-------|------------------|:----------:|:--------------:|:---------------:|:----:| :----:|
|
||||
| **Wan2.1** | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Wan2.1** | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **LTX-2** | `Lightricks/LTX-2` | ✅ | ✅ | ✅ | ❌ | ❌ |
|
||||
|
||||
@@ -414,3 +421,28 @@ vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --usp
|
||||
#### Adding HSDP Support to New Models
|
||||
|
||||
For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../design/feature/hsdp.md).
|
||||
|
||||
### Expert Parallelism
|
||||
|
||||
Unlike Tensor Parallelism which shards every layer's weights, EP only shards the MoE expert MLP blocks. This significantly reduces the memory footprint of MoE models (e.g., HunyuanImage3.0) while maintaining constant dense-equivalent compute efficiency. Expert Parallelism is enabled via `DiffusionParallelConfig.enable_expert_parallel`. And `self.ep = tp * sp * cfg * dp` for now, so at least one of TP/SP/CFG/DP should set when EP enabled.
|
||||
|
||||
#### Offline Inference
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
|
||||
omni = Omni(
|
||||
model="tencent/HunyuanImage-3.0",
|
||||
parallel_config=DiffusionParallelConfig(tensor_parallel_size=8, enable_expert_parallel=True),
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
"A brown and white dog is running on the grass",
|
||||
OmniDiffusionSamplingParams(
|
||||
num_inference_steps=50,
|
||||
width=1024,
|
||||
height=1024,
|
||||
),
|
||||
)
|
||||
```
|
||||
@@ -63,6 +63,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
|
||||
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
|
||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Skip sensitive layers | `img_mlp` |
|
||||
| Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
|
||||
| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None |
|
||||
|
||||
## Combining with Other Features
|
||||
|
||||
|
||||
@@ -22,10 +22,16 @@
|
||||
CLI (examples/offline_inference/text_to_image/text_to_image.py)
|
||||
|
|
||||
v
|
||||
Omni (vllm_omni/entrypoints/omni.py)
|
||||
Omni (vllm_omni/entrypoints/__init__.py)
|
||||
|
|
||||
v
|
||||
OmniStage (diffusion)
|
||||
AsyncOmniEngine
|
||||
|
|
||||
v
|
||||
Orchestrator
|
||||
|
|
||||
v
|
||||
StageDiffusionClient
|
||||
|
|
||||
v
|
||||
DiffusionWorker
|
||||
@@ -64,7 +70,13 @@ _generate_with_async_omni
|
||||
AsyncOmni
|
||||
|
|
||||
v
|
||||
DiffusionEngine
|
||||
AsyncOmniEngine
|
||||
|
|
||||
v
|
||||
Orchestrator
|
||||
|
|
||||
v
|
||||
StageDiffusionClient
|
||||
|
|
||||
v
|
||||
OmniRequestOutput
|
||||
@@ -136,9 +148,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--height 768 \
|
||||
--width 1360 \
|
||||
--seed 42 \
|
||||
--cfg_scale 4.0 \
|
||||
--num_images_per_prompt 1 \
|
||||
--num_inference_steps 4 \
|
||||
--cfg-scale 4.0 \
|
||||
--num-images-per-prompt 1 \
|
||||
--num-inference-steps 4 \
|
||||
--output outputs/flux2_klein_4b.png
|
||||
```
|
||||
|
||||
@@ -152,9 +164,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--height 768 \
|
||||
--width 1360 \
|
||||
--seed 42 \
|
||||
--cfg_scale 4.0 \
|
||||
--num_images_per_prompt 1 \
|
||||
--num_inference_steps 4 \
|
||||
--cfg-scale 4.0 \
|
||||
--num-images-per-prompt 1 \
|
||||
--num-inference-steps 4 \
|
||||
--output outputs/flux2_klein_4b_gguf.png
|
||||
```
|
||||
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
# Int8 Quantization
|
||||
|
||||
## Overview
|
||||
|
||||
Int8 quantization converts BF16/FP16 weights to Int8 at model load time. No calibration or pre-quantized checkpoint needed.
|
||||
|
||||
Depending on the model, either all layers can be quantized, or some sensitive layers should stay in BF16/FP16. See the [per-model table](#supported-models) for which case applies.
|
||||
|
||||
## Configuration
|
||||
|
||||
1. **Python API**: set `quantization="int8"`. To skip sensitive layers, use `quantization_config` with `ignored_layers`.
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
# All layers quantized
|
||||
omni = Omni(model="<your-model>", quantization="int8")
|
||||
|
||||
# Skip sensitive layers
|
||||
omni = Omni(
|
||||
model="<your-model>",
|
||||
quantization_config={
|
||||
"method": "int8",
|
||||
"ignored_layers": ["<layer-name>"],
|
||||
},
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
"A cat sitting on a windowsill",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=50),
|
||||
)
|
||||
```
|
||||
|
||||
2. **CLI**: pass `--quantization int8` and optionally `--ignored-layers`.
|
||||
|
||||
```bash
|
||||
# All layers
|
||||
python text_to_image.py --model <your-model> --quantization int8
|
||||
|
||||
# Skip sensitive layers
|
||||
python text_to_image.py --model <your-model> --quantization int8 --ignored-layers "img_mlp"
|
||||
|
||||
# Online serving
|
||||
vllm serve <your-model> --omni --quantization int8
|
||||
```
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `method` | str | — | Quantization method (`"int8"`) |
|
||||
| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16/FP16 |
|
||||
| `activation_scheme` | str | `"dynamic"` | `"dynamic"` (no calibration) |
|
||||
|
||||
|
||||
The available `ignored_layers` names depend on the model architecture (e.g., `to_qkv`, `to_out`, `img_mlp`, `txt_mlp`). Consult the transformer source for your target model.
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | HF Models | Recommendation | `ignored_layers` |
|
||||
|-------|-----------|---------------|------------------|
|
||||
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
|
||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | All layers | None |
|
||||
|
||||
## Combining with Other Features
|
||||
|
||||
Int8 quantization can be combined with cache acceleration:
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="<your-model>",
|
||||
quantization="int8",
|
||||
cache_backend="tea_cache",
|
||||
cache_config={"rel_l1_thresh": 0.2},
|
||||
)
|
||||
```
|
||||
@@ -7,12 +7,20 @@ vLLM-Omni supports quantization of DiT linear layers to reduce memory usage and
|
||||
| Method | Guide |
|
||||
|--------|-------|
|
||||
| FP8 | [FP8](fp8.md) |
|
||||
| Int8 | [Int8](int8.md) |
|
||||
| GGUF | [GGUF](gguf.md) |
|
||||
|
||||
## Device Compatibility
|
||||
## Device Compatibility for FP8
|
||||
|
||||
| GPU Generation | Example GPUs | FP8 Mode |
|
||||
|---------------|-------------------|----------|
|
||||
| Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
|
||||
|
||||
Kernel selection is automatic.
|
||||
|
||||
## Device Compatibility for Int8
|
||||
|
||||
| Device Type | Generation | Example | Int8 Mode |
|
||||
|-------------|---------------|-------------------|----------|
|
||||
| NVIDIA GPU | Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
|
||||
| Ascend NPU | Atlas A2/Atlas A3 | Atlas 800T A2/Atlas 900 A3 | Full W8A8 with native hardware |
|
||||
@@ -0,0 +1,61 @@
|
||||
# Step Execution
|
||||
|
||||
Step execution is an opt-in diffusion execution mode enabled with
|
||||
`step_execution=True` when constructing `Omni`.
|
||||
|
||||
It is not a generic diffusion toggle for every pipeline. Only pipelines that
|
||||
implement the stepwise contract support it today.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
step_execution=True,
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
"A cat sitting on a windowsill",
|
||||
OmniDiffusionSamplingParams(
|
||||
num_inference_steps=50,
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
## Supported Pipelines
|
||||
|
||||
| Pipeline | Example models | Step execution |
|
||||
|----------|----------------|----------------|
|
||||
| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
|
||||
| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
|
||||
|
||||
## Current Limitations
|
||||
|
||||
- `step_execution` currently supports `batch_size=1` only.
|
||||
- `cache_backend` is not supported together with step execution.
|
||||
- Unsupported pipelines fail early during model loading.
|
||||
- Request-mode extras such as KV transfer are not wired into step mode yet.
|
||||
|
||||
## When To Use It
|
||||
|
||||
Use step execution only when you specifically need the pipeline to run through
|
||||
its stepwise request state machine. For normal diffusion inference, leave it
|
||||
disabled unless your workflow depends on this mode.
|
||||
|
||||
If you are looking for general diffusion speedups, see
|
||||
[Diffusion Acceleration Overview](../diffusion_acceleration.md).
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
If model loading fails with a message mentioning `prepare_encode()`,
|
||||
`denoise_step()`, `step_scheduler()`, and `post_decode()`, the selected
|
||||
pipeline does not support step execution.
|
||||
|
||||
## For Model Authors
|
||||
|
||||
If you want to add step execution support to a new diffusion pipeline, see the
|
||||
implementation guide:
|
||||
[Diffusion Step Execution Design](../../design/feature/diffusion_step_execution.md).
|
||||
@@ -16,7 +16,7 @@ Both methods can provide significant speedups (typically **1.5x-2.0x**) while ma
|
||||
|
||||
vLLM-Omni also supports quantization methods:
|
||||
|
||||
3. **[FP8 Quantization](diffusion/quantization/overview.md)** - Reduces DiT linear layers from BF16 to FP8, providing ~1.28x speedup with minimal quality loss. Supports per-layer skip for sensitive layers.
|
||||
3. **[Quantization](diffusion/quantization/overview.md)** - Reduces DiT linear layers from BF16 to FP8 or Int8, providing ~1.28x speedup with minimal quality loss. Supports per-layer skip for sensitive layers.
|
||||
|
||||
vLLM-Omni also supports parallelism methods for diffusion models, including:
|
||||
|
||||
@@ -46,6 +46,7 @@ vLLM-Omni also supports parallelism methods for diffusion models, including:
|
||||
| Method | Configuration | Description | Best For |
|
||||
|--------|--------------|-------------|----------|
|
||||
| **FP8** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup |
|
||||
| **Int8** | `quantization="int8"` | Int8 W8A8 | Memory reduction, inference speedup |
|
||||
|
||||
## Supported Models
|
||||
|
||||
@@ -69,21 +70,25 @@ The following table shows which models are currently supported by each accelerat
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
|
||||
### VideoGen
|
||||
|
||||
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | VAE-Patch-Parallel |
|
||||
|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
|
||||
| **Wan2.1-T2V** | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Wan2.1-T2V** | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
|
||||
| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ |
|
||||
|
||||
### Quantization
|
||||
|
||||
| Model | Model Identifier | FP8 |
|
||||
|-------|------------------|:---:|
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ |
|
||||
| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ |
|
||||
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ |
|
||||
| Model | Model Identifier | FP8 | Int8 |
|
||||
|-------|------------------|:---:|:---:|
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ |
|
||||
| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ |
|
||||
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ |
|
||||
|
||||
|
||||
## Performance Benchmarks
|
||||
@@ -336,13 +341,30 @@ outputs = omni.generate(
|
||||
)
|
||||
```
|
||||
|
||||
### Using Int8 Quantization
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
omni = Omni(
|
||||
model="<your-model>",
|
||||
quantization="int8",
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
"A cat sitting on a windowsill",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=50),
|
||||
)
|
||||
```
|
||||
|
||||
## Documentation
|
||||
|
||||
For detailed information on each acceleration method:
|
||||
|
||||
- **[TeaCache Guide](diffusion/teacache.md)** - Complete TeaCache documentation, configuration options, and best practices
|
||||
- **[Cache-DiT Acceleration Guide](diffusion/cache_dit_acceleration.md)** - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters
|
||||
- **[FP8 Quantization Guide](diffusion/quantization/overview.md)** - FP8 quantization for DiT models with per-layer control
|
||||
- **[Quantization Guide](diffusion/quantization/overview.md)** - Quantization for DiT models with per-layer control
|
||||
- **[Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism)** - Guidance on how to enable TP for diffusion models.
|
||||
- **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
|
||||
- **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.
|
||||
|
||||
@@ -130,7 +130,7 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
|
||||
| :------------------------------- | :------------------------------ | :----------------------- |
|
||||
| `stage_type` | `llm` | Stage type |
|
||||
| `devices` | `"0"` | GPU device ID |
|
||||
| `max_batch_size` | `1` | Maximum batch size |
|
||||
| `max_num_seqs` | `1` | Maximum batch size |
|
||||
| `model_stage` | `thinker` | Model stage identifier |
|
||||
| `model_arch` | `BagelForConditionalGeneration` | Model architecture |
|
||||
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
|
||||
@@ -146,7 +146,7 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
|
||||
| :------------------------------- | :---------- | :-------------------------- |
|
||||
| `stage_type` | `diffusion` | Stage type |
|
||||
| `devices` | `"0"` | GPU device ID |
|
||||
| `max_batch_size` | `1` | Maximum batch size |
|
||||
| `max_num_seqs` | `1` | Maximum batch size |
|
||||
| `model_stage` | `dit` | Model stage identifier |
|
||||
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
|
||||
| `omni_kv_config.need_recv_cache` | `true` | Whether to receive KV cache |
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
# CosyVoice3
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/cosyvoice3>.
|
||||
|
||||
|
||||
## Setup
|
||||
|
||||
Install dependencies:
|
||||
```
|
||||
uv pip install -e .
|
||||
```
|
||||
|
||||
> **Note:** This includes required libraries such as `librosa`, `soundfile`,
|
||||
> `onnxruntime`, `x-transformers`, and `einops` via
|
||||
> `requirements/common.txt` and platform-specific requirements files.
|
||||
|
||||
Download the model snapshot:
|
||||
```
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', local_dir='pretrained_models/Fun-CosyVoice3-0.5B')
|
||||
```
|
||||
|
||||
Add `config.json` in `pretrained_models/Fun-CosyVoice3-0.5B/`:
|
||||
```json
|
||||
{
|
||||
"model_type": "cosyvoice3",
|
||||
"architectures": [
|
||||
"CosyVoice3Model"
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
> **Why `config.json` is required:**
|
||||
> `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers a class mapping.
|
||||
> The loader still needs `model_type: "cosyvoice3"` from `config.json` to select that class.
|
||||
> If no `config.json` is present, model type cannot be inferred automatically.
|
||||
> If your downloaded checkpoint already includes a valid `config.json` with
|
||||
> `model_type: "cosyvoice3"`, this manual step can be skipped.
|
||||
|
||||
Run the offline verification script:
|
||||
```
|
||||
python examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py \
|
||||
--model pretrained_models/Fun-CosyVoice3-0.5B \
|
||||
--tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN
|
||||
```
|
||||
|
||||
## Implementation Overview
|
||||
|
||||
CosyVoice3 runs as a **2-stage Omni pipeline**.
|
||||
|
||||
- **Stage 0** (`talker`) converts text + prompt audio to speech tokens.
|
||||
- **Stage 1** (`code2wav`) converts speech tokens + prompt features to waveform via flow matching and HiFiGAN.
|
||||
|
||||
Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py`.
|
||||
|
||||
- `CosyVoice3MultiModalProcessor` builds the multimodal inputs:
|
||||
- Tokenizes `prompt` and `prompt_text`.
|
||||
- Extracts speech tokens and mel features from the prompt audio.
|
||||
- Extracts a speaker embedding.
|
||||
- `CosyVoice3Model` implements both stages:
|
||||
- Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features.
|
||||
- Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform.
|
||||
|
||||
Stage wiring is configured in `vllm_omni/model_executor/stage_configs/cosyvoice3.yaml`.
|
||||
|
||||
- Stage 0 emits latent speech tokens.
|
||||
- Stage 1 consumes them via `custom_process_input_func` and outputs audio.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "verify_e2e_cosyvoice.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py"
|
||||
``````
|
||||
@@ -0,0 +1,97 @@
|
||||
# Fish Speech S2 Pro
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/fish_speech>.
|
||||
|
||||
|
||||
This directory contains an offline demo for running Fish Speech S2 Pro (`fishaudio/s2-pro`) with vLLM Omni. It supports text-only synthesis and voice cloning with reference audio.
|
||||
|
||||
## Model Overview
|
||||
|
||||
[Fish Speech S2 Pro](https://huggingface.co/fishaudio/s2-pro) is a 4B dual-AR text-to-speech model by FishAudio. It uses:
|
||||
|
||||
- **Slow AR**: Generates semantic tokens autoregressively (Qwen3-based backbone)
|
||||
- **Fast AR**: Predicts residual codebook tokens from semantic tokens
|
||||
- **DAC Decoder**: Converts codec codes to 44.1 kHz audio waveform
|
||||
|
||||
The model produces high-quality speech with voice cloning capabilities.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Install the `fish-speech` package for the DAC codec:
|
||||
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
## Setup
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Text-Only Synthesis
|
||||
|
||||
```bash
|
||||
python end2end.py --text "Hello, this is a test of the Fish Speech text to speech system."
|
||||
```
|
||||
|
||||
Generated audio files are saved to `output_audio/` by default.
|
||||
|
||||
### Voice Cloning
|
||||
|
||||
Provide a reference audio file and its transcript to clone a voice:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--text "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio."
|
||||
```
|
||||
|
||||
## Streaming Mode
|
||||
|
||||
Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
|
||||
|
||||
```bash
|
||||
python end2end.py --text "Hello, this is a test." --streaming
|
||||
```
|
||||
|
||||
Each DAC decoder chunk is logged as it arrives. The final WAV file is written once generation completes.
|
||||
|
||||
## Configuration
|
||||
|
||||
| Argument | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `--text` | `"Hello, this is a test..."` | Text to synthesize |
|
||||
| `--ref-audio` | None | Path to reference audio for voice cloning |
|
||||
| `--ref-text` | None | Transcript of the reference audio |
|
||||
| `--model` | `fishaudio/s2-pro` | HuggingFace model path |
|
||||
| `--stage-configs-path` | `fish_speech_s2_pro.yaml` | Path to stage configs YAML |
|
||||
| `--output-dir` | `output_audio` | Output directory for WAV files |
|
||||
| `--streaming` | False | Enable streaming via AsyncOmni |
|
||||
| `--stage-init-timeout` | 600 | Stage initialization timeout (seconds) |
|
||||
|
||||
## Notes
|
||||
|
||||
- Output audio is 44.1 kHz mono WAV.
|
||||
- The DAC codec weights (`codec.pth`) are loaded lazily from the model directory.
|
||||
- Voice cloning encodes the reference audio using the DAC codec to extract semantic codes, then prepends them as a system message.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "diag_codebook.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/fish_speech/diag_codebook.py"
|
||||
``````
|
||||
??? abstract "diag_reference_compare.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/fish_speech/diag_reference_compare.py"
|
||||
``````
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/fish_speech/end2end.py"
|
||||
``````
|
||||
??? abstract "standalone_test.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/fish_speech/standalone_test.py"
|
||||
``````
|
||||
@@ -99,14 +99,15 @@ python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_str
|
||||
```
|
||||
|
||||
Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames`
|
||||
and `initial_codec_chunk_frames` in the stage config). The final WAV file is written once generation
|
||||
in the stage config). The initial chunk size is dynamically selected based on server load for reduced
|
||||
TTFA, and can be overridden per-request via the `initial_codec_chunk_frames` API field. The final WAV file is written once generation
|
||||
completes. This demonstrates that audio data is available progressively rather than only at the end.
|
||||
|
||||
> **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.
|
||||
|
||||
## Batched Decoding
|
||||
|
||||
The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
|
||||
The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_num_seqs > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
|
||||
|
||||
```
|
||||
python end2end.py --query-type CustomVoice \
|
||||
@@ -115,7 +116,7 @@ python end2end.py --query-type CustomVoice \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
|
||||
```
|
||||
|
||||
**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_batch_size >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_batch_size`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
|
||||
**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_num_seqs >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_num_seqs`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
|
||||
|
||||
## Notes
|
||||
|
||||
|
||||
@@ -19,7 +19,7 @@ if __name__ == "__main__":
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
prompt = "a cup of coffee on the table"
|
||||
outputs = omni.generate(prompt)
|
||||
images = outputs[0].request_output[0].images
|
||||
images = outputs[0].request_output.images
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
@@ -37,7 +37,7 @@ if __name__ == "__main__":
|
||||
]
|
||||
outputs = omni.generate(prompts)
|
||||
for i, output in enumerate(outputs):
|
||||
image = output.request_output[0].images[0].save(f"{i}.jpg")
|
||||
image = output.request_output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
!!! info
|
||||
@@ -49,7 +49,7 @@ if __name__ == "__main__":
|
||||
|
||||
!!! info
|
||||
|
||||
For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
|
||||
For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
|
||||
list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
|
||||
batched inputs, you can [modify this configuration](https://github.com/vllm-project/vllm-omni/tree/main/configuration/stage_configs.md) to let the model accept a longer batch of prompts.
|
||||
|
||||
@@ -72,7 +72,7 @@ if __name__ == "__main__":
|
||||
}
|
||||
])
|
||||
for i, output in enumerate(outputs):
|
||||
image = output.request_output[0].images[0].save(f"{i}.jpg")
|
||||
image = output.request_output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
@@ -3,7 +3,10 @@
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_video>.
|
||||
|
||||
|
||||
The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts.
|
||||
The Wan T2V pipeline generates short videos from text prompts. Supported models include:
|
||||
|
||||
- `Wan-AI/Wan2.2-T2V-A14B-Diffusers` (14B, requires ~48GB VRAM)
|
||||
- `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` (1.3B, requires ~16GB VRAM, use `--flow-shift 3.0 --boundary-ratio 0.0`)
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
|
||||