Compare commits

...
33 Commits
Author SHA1 Message Date
6838533998 [Model] Add HunyuanVideo-1.5 T2V and I2V support (#1516)
Signed-off-by: lishunyang <lishunyang12@163.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-03-21 21:26:39 +08:00
Zhou TaichangandGao Han d6e338f9d6 [Doc] Update docs and dockerfiles for rebase of vllm v0.18.0 (#2038)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-03-21 13:21:33 +00:00
Zhou Taichangandgcanlin a90a7690b8 [Rebase] Rebase to vllm v0.18.0 (#2037)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: gcanlin <canlinguosdu@gmail.com>
2026-03-21 20:37:38 +08:00
Juan Pablo ZuluagaandJuanPZuluaga 8e12007ac1 [BugFix][Qwen3TTS] CodePredictor CudaGraph Pool (#2059)
Signed-off-by: JuanPZuluaga <juanz9312@gmal.com>
Co-authored-by: JuanPZuluaga <juanz9312@gmal.com>
2026-03-21 19:11:34 +08:00
Zeyu Huang | 黃澤宇andSamit da21e991cc [bugfix] /chat/completion doesn't read extra_body for diffusion model (#2042)
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Samit <285365963@qq.com>
2026-03-21 15:53:04 +08:00
Gao Han a978c60163 [CI] Uncomment condition for nightly build in YAML (#2057)
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
2026-03-21 13:55:42 +08:00
Patrick von Platen c1c0a5b3cd [Voxtral] Improve example (#2045)
Signed-off-by: Patrick von Platen <patrick.v.platen@gmail.com>
2026-03-21 05:27:43 +00:00
Sy03andsy0307 072647eeb2 [Perf] Improve Fish Speech S2 Pro inference performance (#1859)
Signed-off-by: sy0307 <sy0307@users.noreply.github.com>
Signed-off-by: Sy03 <1370724210@qq.com>
Co-authored-by: sy0307 <sy0307@users.noreply.github.com>
2026-03-21 01:15:02 -04:00
NumberWan 7007217d4b [Bugfix]: fixed ServerDisconnectedError in benchmark test (reapply #1683, fixes #1374) (#1841)
Signed-off-by: NumberWan <wantszkin2003@gmail.com>
2026-03-21 11:41:02 +08:00
Chen-Yo Sun 95018c536a [Voxtral TTS] Remove redundant yaml (#2056)
Signed-off-by: Chen-Yo Sun <chenyo.sun@mistral.ai>
2026-03-20 14:02:35 -07:00
Lancer b33d7637de [Bugfix] Z-Image CFG threshold should be > 0 instead of > 1 (#1634)
Signed-off-by: Lancer <maruixiang6688@gmail.com>
2026-03-20 20:13:10 +08:00
ff254795fa [Benchmark] [Diffusion] [Enhancement] Random dataset (#1657)
Signed-off-by: dengyunyang <584797741@qq.com>
Signed-off-by: Samit <285365963@qq.com>
Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: Samit <285365963@qq.com>
Co-authored-by: bjf-frz <frz123db@gmail.com>
2026-03-20 20:12:49 +08:00
Junhong Liu 6901ba426e [Fixbug][Perf] Qwen3-omni: code predictor with re-prefill + SDPA and eliminate decode hot-path CPU round-trips (#2012)
Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>
2026-03-20 19:33:52 +08:00
Wu JIAZHEN da7a8f8014 reafator pipeline stage/step pipeline (#1368)
Signed-off-by: asukaqaq-s <1311722138@qq.com>
2026-03-20 19:21:50 +08:00
Alicia 199f7832e7 [CI] Fix test. (#2031)
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-03-20 16:41:45 +08:00
wangyu fafa248665 [Test] Implement mock HTTP request handling in benchmark CLI tests (#2014)
Signed-off-by: yenuo26 <410167048@qq.com>
Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com>
2026-03-20 16:13:39 +08:00
Zhou Taichang 05d8de813e [Bugfix] Remove duplicated config keyword max batch size (#1851)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
2026-03-20 16:13:24 +08:00
Samit 4c2293743b [Docs] Add Wan2.1-T2V as supported video generation models (#1920)
Signed-off-by: samithuang <285365963@qq.com>
2026-03-20 14:14:20 +08:00
Chen-Yo Sun 6813432f90 [BugFix][Doc]Update voxtral_tts end2end.py & README.md (#2026)
Signed-off-by: Chen-Yo Sun <chenyo.sun@mistral.ai>
2026-03-20 05:05:39 +00:00
汪志鹏 c2a5420d97 [CI] Change Bagel online test environment variable VLLM_TEST_CLEAN_GPU_MEMORY to 0 (#2032)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
2026-03-20 12:13:21 +08:00
Juan Pablo ZuluagaandJuanPZuluaga b4342fb487 [Optim][Qwen3TTS][CodePredictor] support torch.compile with reduce-overhead and dynamic False (#1913)
Signed-off-by: JuanPZuluaga <juanz9312@gmal.com>
Signed-off-by: JuanPZuluaga <juanz9312@gmail.com>
Co-authored-by: JuanPZuluaga <juanz9312@gmal.com>
2026-03-20 12:12:12 +08:00
81a90d24ef [Feat][Qwen3TTS][Code2wav] triton SnakeBeta and Cuda Graph (#1797)
Signed-off-by: pablo <pablo@agigo.ai>
Signed-off-by: JuanPZuluaga <juanz9312@gmal.com>
Co-authored-by: pablo <pablo@agigo.ai>
Co-authored-by: JuanPZuluaga <juanz9312@gmal.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-03-20 12:11:37 +08:00
Yuanheng Zhao 6ac7ba692d [Feat] Support T5 Tensor Parallelism (#1881)
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
2026-03-20 08:49:01 +08:00
Chen-Yo Sun e9fe78e28c [Model] Add Voxtral TTS model (#1803)
Signed-off-by: Chen-Yo Sun <chenyo.sun@mistral.ai>
2026-03-19 16:02:11 -04:00
b766c4702c Int8 Quantization Support for DiT (Z-Image & Qwen-Image) (#1470)
Signed-off-by: juboyu <767868009@qq.com>
Signed-off-by: JuboYu <767868009@qq.com>
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-03-19 22:58:29 +08:00
TJian 0239057cba [CI] [ROCm] Setup test-ready.yml and test-merge.yml (#2017)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-03-19 14:26:31 +00:00
Hui.andAlvisGong 99cc76690d [Bugfix]Fix bug of online server can not return mutli images (#2007)
Signed-off-by: Hui <1779066624@qq.com>
Co-authored-by: AlvisGong <gwly0401@163.com>
2026-03-19 03:57:04 -07:00
rein yang 60a8c01a96 [Bugfix] revert PR#1758 which introduced the accuracy problem of qwen3-omni (#2009)
Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
2026-03-19 17:34:01 +08:00
Ziming Huang 89fff09628 [Fix] Fix slow hasattr in CUDAGraphWrapper.__getattr__ (#1982)
Signed-off-by: ZeldaHuang <hzm414167@alibaba-inc.com>
2026-03-19 17:10:08 +08:00
Chendi.XueandHongsheng Liu ba5a66b15b [XPU] update bagel modeling to remove cuda hardcode, add xpu stage_config (#1931)
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-03-19 16:55:16 +08:00
rongfu.leng f9268822cd [Bugfix] fix helios video generate use cpu device (#1915)
Signed-off-by: rongfu.leng <lenronfu@gmail.com>
2026-03-19 16:51:45 +08:00
Dnoob 48b3969717 [Bugfix] Fix Ovis Image crash when guidance_scale is set without negative_prompt (#1956)
Signed-off-by: Dnoob <dxpouo@gmail.com>
2026-03-19 16:51:02 +08:00
Kevin H. Luu b1ceaa7b34 Add HF token to H100 jobs (#2008)
Signed-off-by: khluu <khluu000@gmail.com>
2026-03-19 01:25:14 -07:00
202 changed files with 13793 additions and 1108 deletions

No files matched your search

+9 -1
View File
@@ -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" \
+1 -1
View File
@@ -10,7 +10,7 @@ steps:
DOCKER_BUILDKIT: "1"
# Buildkite will automatically replace this with the actual commit hash
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
VLLM_VERSION: "v0.17.0"
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 \
+197
View File
@@ -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"
+178
View File
@@ -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"
+2 -2
View File
@@ -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]
+68 -8
View File
@@ -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:
@@ -304,6 +320,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:
@@ -314,3 +335,42 @@ steps:
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
# - label: "Voxtral TTS E2E Test"
# timeout_in_minutes: 20
# depends_on: upload-merge-pipeline
# commands:
# - |
# timeout 20m bash -c '
# export VLLM_LOGGING_LEVEL=DEBUG
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
# pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
# '
# 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
+24 -2
View File
@@ -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:
@@ -58,13 +66,14 @@ steps:
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"
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"
@@ -86,6 +95,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:
@@ -116,6 +130,7 @@ steps:
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -149,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:
@@ -163,7 +185,7 @@ steps:
- label: ":email: Nightly Collection & Email"
key: nightly-perf-distribution
depends_on:
- nightly-performance
- nightly-performance
if: build.env("NIGHTLY") == "1"
commands:
- pip install openpyxl
+83 -7
View File
@@ -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,6 +46,7 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -44,6 +63,7 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -60,6 +80,7 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -76,6 +97,7 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -93,6 +115,7 @@ steps:
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -110,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
@@ -137,6 +156,7 @@ steps:
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -159,6 +179,7 @@ steps:
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
@@ -221,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:
@@ -288,6 +349,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:
@@ -326,6 +392,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:
@@ -365,6 +436,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:
+6 -1
View File
@@ -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
+136 -2
View File
@@ -1,3 +1,4 @@
import asyncio
import base64
import mimetypes
import os
@@ -206,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.
+11 -11
View File
@@ -32,8 +32,8 @@ 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 batch_size=4 config for higher throughput
STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml 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
@@ -54,7 +54,7 @@ CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
### 2) Run online serving benchmark
```bash
python vllm_omni/bench_tts_serve.py \
python benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py \
--port 8000 \
--num-prompts 50 \
--max-concurrency 1 4 10 \
@@ -65,7 +65,7 @@ python vllm_omni/bench_tts_serve.py \
### 3) Run HuggingFace baseline
```bash
python transformers/bench_tts_hf.py \
python benchmarks/qwen3-tts/transformers/bench_tts_hf.py \
--model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
--num-prompts 50 \
--gpu-device 0 \
@@ -75,7 +75,7 @@ python transformers/bench_tts_hf.py \
### 4) Generate comparison plots
```bash
python plot_results.py \
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
@@ -83,14 +83,14 @@ python plot_results.py \
## Stage Configs
| Config | Batch Size | Description |
|--------|:----------:|-------------|
| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lower latency) |
| `vllm_omni/configs/qwen3_tts_bs4.yaml` | 4 | Concurrent request processing (higher throughput) |
| 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 |
Both 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.
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 in `run_benchmark.sh`), so the same configs work for both the 0.6B and 1.7B model variants.
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
@@ -4,17 +4,18 @@ async_chunk: true
stage_args:
- stage_id: 0
stage_type: llm
is_comprehension: true
runtime:
devices: "0"
max_batch_size: 1
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: false
async_scheduling: true
enable_prefix_caching: false
engine_output_type: latent
gpu_memory_utilization: 0.3
@@ -37,18 +38,18 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 1
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: false
async_scheduling: true
enable_prefix_caching: false
engine_output_type: audio
gpu_memory_utilization: 0.2
gpu_memory_utilization: 0.3
distributed_executor_backend: "mp"
max_num_batched_tokens: 8192
max_model_len: 32768
@@ -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
@@ -1,21 +1,22 @@
# Qwen3-TTS batch_size=4 config (streaming with async_chunk)
# Enables concurrent request processing with max_inflight=4
# 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"
max_batch_size: 4
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: false
async_scheduling: true
enable_prefix_caching: false
engine_output_type: latent
gpu_memory_utilization: 0.3
@@ -38,15 +39,15 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 4
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: false
async_scheduling: true
enable_prefix_caching: false
engine_output_type: audio
gpu_memory_utilization: 0.2
+1 -1
View File
@@ -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}
+2 -2
View File
@@ -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}}
+1 -1
View File
@@ -76,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.17.0
ARG VLLM_VERSION=v0.18.0
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
WORKDIR /workspace/vllm
+3
View File
@@ -58,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
@@ -90,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
+1 -1
View File
@@ -158,7 +158,7 @@ The stage initialized but ran out of memory during processing.
**Solutions:**
1. Reduce `max_num_batched_tokens`
2. Reduce `max_batch_size` in runtime config
2. Reduce `max_num_seqs` in engine_args
3. Lower `gpu_memory_utilization` slightly
4. Enable quantization if supported
+4 -3
View File
@@ -32,8 +32,8 @@ stage_args:
is_prefill_only: true
runtime:
devices: "0"
max_batch_size: 16
engine_args:
max_num_seqs: 16
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
@@ -70,8 +70,8 @@ stage_args:
is_decode_only: true
runtime:
devices: "1"
max_batch_size: 64
engine_args:
max_num_seqs: 64
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
@@ -128,7 +128,8 @@ After inserting the extra thinker stage, renumber the remaining stages:
- stage_id: 3
runtime:
devices: "2"
max_batch_size: 1
engine_args:
max_num_seqs: 1
engine_input_source: [2]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
```
+5 -5
View File
@@ -36,9 +36,9 @@ stage_args:
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`
@@ -4,9 +4,9 @@ stage_args:
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
@@ -739,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
@@ -967,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:**
+2 -2
View File
@@ -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
+2 -2
View File
@@ -171,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
@@ -195,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)
+1 -1
View File
@@ -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.17.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.17.0 --torch-backend=auto
uv pip install vllm==0.18.0 --torch-backend=auto
```
The release 0.17.0 of vLLM is based on PyTorch 2.10.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.17.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.17.0/vllm-0.17.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.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.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.17.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.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.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.17.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.17.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]
@@ -68,18 +68,18 @@ 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.17.0
# Pin vLLM version to 0.18.0
cd /vllm-workspace/vllm
git pull origin main
git fetch origin --tags
git checkout v0.17.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-workspace/vllm-ascend
git pull origin main
git fetch origin --tags
git checkout v0.17.0
git checkout 1e05c4908f31737bc4eef865a9f351d030a77c9d
pip install -v -e .
# Install vLLM-Omni from the latest main branch
+3 -3
View File
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
source .venv/bin/activate
# On CUDA
uv pip install vllm==0.17.0 --torch-backend=auto
uv pip install vllm==0.18.0 --torch-backend=auto
# On ROCm
uv pip install vllm==0.17.0 --extra-index-url https://wheels.vllm.ai/rocm/0.17.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
@@ -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.
+1
View File
@@ -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
+5 -2
View 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` |
@@ -52,6 +52,9 @@ th {
|`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
@@ -75,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` |
+15
View File
@@ -4,6 +4,7 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
- **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
- **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
- **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
@@ -28,6 +29,14 @@ vllm-omni serve fishaudio/s2-pro \
--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
@@ -347,6 +356,12 @@ curl -X POST http://localhost:8091/v1/audio/speech \
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
@@ -58,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` | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -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
@@ -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).
+27 -7
View File
@@ -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
@@ -75,17 +76,19 @@ The following table shows which models are currently supported by each accelerat
| 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
@@ -338,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 |
@@ -107,7 +107,7 @@ completes. This demonstrates that audio data is available progressively rather t
## 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 \
@@ -116,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
@@ -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.
@@ -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
+2 -2
View File
@@ -127,7 +127,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 |
@@ -143,7 +143,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 |
+11 -3
View File
@@ -34,7 +34,14 @@ def parse_args():
help="Path to input image for img2img.",
)
# Omni runtime init args
parser.add_argument(
"--output",
type=str,
default=".",
help="Output directory to save images.",
)
# OmniLLM init args
parser.add_argument("--log-stats", action="store_true", default=False)
parser.add_argument("--init-sleep-seconds", type=int, default=20)
parser.add_argument("--batch-timeout", type=int, default=5)
@@ -70,6 +77,7 @@ def parse_args():
def main():
args = parse_args()
os.makedirs(args.output, exist_ok=True)
model_name = args.model
prompts: list[OmniPromptType] = []
try:
@@ -172,13 +180,13 @@ def main():
img_idx = 0
for req_output in omni_outputs:
images = getattr(req_output, "images", None)
if not images:
continue
for j, img in enumerate(images):
save_path = f"output_{img_idx}_{j}.png"
save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
img.save(save_path)
print(f"[Info] Saved image to {save_path}")
img_idx += 1
print(omni_outputs)
@@ -2,15 +2,16 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Image-to-Video generation example using Wan2.2 I2V/TI2V models or LTX2.
Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, or HunyuanVideo-1.5.
Supports:
- Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
- Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B)
- LTX2 image-to-video pipeline
- HunyuanVideo-1.5 I2V: SigLIP + VAE dual image conditioning
Usage:
# I2V-A14B (MoE)
# Wan I2V-A14B (MoE)
python image_to_video.py --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
--image input.jpg --prompt "A cat playing with yarn"
@@ -24,6 +25,11 @@ Usage:
--image input.jpg --prompt "A cinematic dolly shot of a boat" \
--num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
--frame-rate 24 --fps 24 --output ltx2_i2v.mp4
# HunyuanVideo-1.5 I2V (480p)
python image_to_video.py --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \
--image input.jpg --prompt "A cat playing with yarn" \
--flow-shift 5.0 --guidance-scale 6.0
"""
import argparse
@@ -43,11 +49,11 @@ from vllm_omni.platforms import current_omni_platform
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate a video from an image with Wan2.2 or LTX2.")
parser = argparse.ArgumentParser(description="Generate a video from an image (Wan2.2, LTX2, HunyuanVideo-1.5).")
parser.add_argument(
"--model",
default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
help="Diffusers Wan2.2 I2V model ID or local path.",
help="Diffusers I2V model ID or local path (Wan2.2 or HunyuanVideo-1.5).",
)
parser.add_argument(
"--model-class-name",
@@ -104,7 +104,7 @@ completes. This demonstrates that audio data is available progressively rather t
## 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 \
@@ -113,7 +113,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
@@ -46,7 +46,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.
@@ -131,12 +131,10 @@ def parse_args() -> argparse.Namespace:
"--quantization",
type=str,
default=None,
choices=["fp8", "gguf"],
help=(
"Quantization method for the transformer. "
"Options: 'fp8' (FP8 W8A8), 'gguf' (GGUF quantized weights). "
"Default: None (no quantization, uses BF16)."
),
choices=["fp8", "int8", "gguf"],
help="Quantization method for the transformer. "
"Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8), 'gguf' (GGUF quantized weights). "
"Default: None (no quantization, uses BF16).",
)
parser.add_argument(
"--gguf-model",
@@ -1,10 +1,19 @@
# Text-To-Video
The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts. This script can also be used
for `Lightricks/LTX-2` to generate video+audio.
A unified script for text-to-video generation. Supports multiple models with model-aware defaults.
## Supported Models
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|---|---|---|---|---|---|
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
## Local CLI Usage
### Wan2.2 (default)
```bash
python text_to_video.py \
--prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
@@ -36,22 +45,100 @@ python text_to_video.py \
--output ltx2_out.mp4
```
Key arguments:
### HunyuanVideo-1.5 (480p)
```bash
python text_to_video.py \
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
--prompt "A cat walks through a sunlit garden, flowers swaying gently in the breeze." \
--height 480 \
--width 832 \
--num-frames 121 \
--guidance-scale 6.0 \
--flow-shift 5.0 \
--num-inference-steps 50 \
--fps 24 \
--output hunyuan_video_15_output.mp4
```
### HunyuanVideo-1.5 (720p)
```bash
python text_to_video.py \
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v \
--prompt "A serene lakeside sunrise with mist over the water." \
--height 720 \
--width 1280 \
--num-frames 121 \
--guidance-scale 6.0 \
--flow-shift 9.0 \
--num-inference-steps 50 \
--fps 24 \
--output hunyuan_720p.mp4
```
### HunyuanVideo-1.5 with FP8 Quantization
```bash
python text_to_video.py \
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
--prompt "A dog running across a field of golden wheat." \
--quantization fp8 \
--height 480 --width 832 --num-frames 121 \
--guidance-scale 6.0 --flow-shift 5.0 \
--output hunyuan_fp8.mp4
```
Quick test (smaller resolution, fewer frames):
```bash
python text_to_video.py \
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
--prompt "A serene lakeside sunrise with mist over the water." \
--height 320 --width 576 --num-frames 17 --num-inference-steps 30 \
--flow-shift 5.0 \
--output quick_test.mp4
```
## Key Arguments
### Common
- `--model`: Diffusers model ID or local path.
- `--prompt`: text description (string).
- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
- `--num-frames`: Number of frames (Wan default is 81).
- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high).
- `--negative-prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
- `--boundary-ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
- `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
- `--height/--width`: output resolution. Default depends on model.
- `--num-frames`: number of frames. Default depends on model.
- `--guidance-scale`: CFG scale. Default depends on model.
- `--num-inference-steps`: sampling steps. Default depends on model.
- `--fps`: frames per second for the saved MP4.
- `--output`: path to save the generated video.
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
- `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
- `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio).
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
- `--flow-shift`: scheduler flow_shift parameter.
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
### Wan2.2-specific
- `--negative-prompt`: artifacts to suppress.
- `--guidance-scale-high`: separate CFG scale for high-noise stage.
- `--boundary-ratio`: boundary split for low/high DiT (default 0.875).
- `--flow-shift`: scheduler flow_shift (5.0 for 720p, 12.0 for 480p).
- `--cache-backend`: `cache_dit` for acceleration.
### HunyuanVideo-1.5 Optimal Configs
| Variant | flow_shift | guidance_scale | steps |
|---------|-----------|----------------|-------|
| 480p T2V | 5.0 | 6.0 | 50 |
| 720p T2V | 9.0 | 6.0 | 50 |
| 480p I2V | 5.0 | 6.0 | 50 |
| 720p I2V | 7.0 | 6.0 | 50 |
| CFG-distilled | (same) | 1.0 | 50 |
> If you encounter OOM errors, try `--vae-use-slicing`, `--vae-use-tiling`, `--enable-cpu-offload`, or `--quantization fp8`.
@@ -15,56 +15,92 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
_MODEL_PRESETS = {
"wan": {
"height": 720,
"width": 1280,
"num_frames": 81,
"num_inference_steps": 40,
"guidance_scale": 4.0,
"fps": 24,
"output": "wan22_output.mp4",
},
"hunyuan": {
"height": 480,
"width": 832,
"num_frames": 121,
"num_inference_steps": 50,
"guidance_scale": 6.0,
"fps": 24,
"output": "hunyuan_video_15_output.mp4",
},
}
def _detect_preset(model: str) -> dict:
model_lower = model.lower()
if "hunyuan" in model_lower:
return _MODEL_PRESETS["hunyuan"]
return _MODEL_PRESETS["wan"]
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate a video with Wan2.2 T2V.")
parser = argparse.ArgumentParser(
description="Generate a video from a text prompt. "
"Supports Wan2.2, HunyuanVideo-1.5, and other text-to-video models."
)
parser.add_argument(
"--model",
default="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
help="Diffusers Wan2.2 model ID or local path.",
help="Diffusers model ID or local path. "
"Examples: Wan-AI/Wan2.2-T2V-A14B-Diffusers, "
"hunyuanvideo-community/HunyuanVideo-1.5-480p_t2v",
)
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
parser.add_argument("--negative-prompt", default="", help="Negative prompt (Wan2.2 only).")
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (applied to low/high).")
parser.add_argument("--guidance-scale-high", type=float, default=None, help="Optional separate CFG for high-noise.")
parser.add_argument("--height", type=int, default=720, help="Video height.")
parser.add_argument("--width", type=int, default=1280, help="Video width.")
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames (Wan default is 81).")
parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
parser.add_argument(
"--guidance-scale-high", type=float, default=None, help="Separate CFG for high-noise stage (Wan2.2 only)."
)
parser.add_argument("--height", type=int, default=None, help="Video height. Default: model-specific.")
parser.add_argument("--width", type=int, default=None, help="Video width. Default: model-specific.")
parser.add_argument("--num-frames", type=int, default=None, help="Number of frames. Default: model-specific.")
parser.add_argument(
"--frame-rate",
type=float,
default=None,
help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
)
parser.add_argument("--num-inference-steps", type=int, default=40, help="Sampling steps.")
parser.add_argument(
"--num-inference-steps", type=int, default=None, help="Sampling steps. Default: model-specific."
)
parser.add_argument(
"--boundary-ratio",
type=float,
default=0.875,
help="Boundary split ratio for low/high DiT. Default 0.875 uses both transformers for best quality. Set to 1.0 to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited).",
default=None,
help="(Wan2.2) Boundary split ratio for low/high DiT. Default 0.875.",
)
parser.add_argument(
"--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
"--flow-shift",
type=float,
default=None,
help="Scheduler flow_shift. Wan2.2: 5.0(720p)/12.0(480p). HunyuanVideo-1.5: 5.0(480p)/9.0(720p).",
)
parser.add_argument(
"--cache-backend",
type=str,
default=None,
choices=["cache_dit"],
help=(
"Cache backend to use for acceleration. "
"Options: 'cache_dit' (DBCache + SCM + TaylorSeer). "
"Default: None (no cache acceleration)."
),
help="Cache backend for acceleration (Wan2.2). Default: None.",
)
parser.add_argument(
"--enable-cache-dit-summary",
action="store_true",
help="Enable cache-dit summary logging after diffusion forward passes.",
)
parser.add_argument("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).")
parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.")
parser.add_argument("--output", type=str, default=None, help="Output path (mp4). Default: model-specific.")
parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.")
parser.add_argument(
"--vae-use-slicing",
action="store_true",
@@ -137,32 +173,41 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
parser.add_argument(
"--quantization",
type=str,
default=None,
choices=["fp8", "gguf"],
help="Quantization method for the transformer (fp8 for online FP8 quantization).",
)
return parser.parse_args()
def main():
args = parse_args()
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
frame_rate = args.frame_rate if args.frame_rate is not None else float(args.fps)
# Wan2.2 cache-dit tuning (from cache-dit examples and cache_alignment).
preset = _detect_preset(args.model)
for key, default_val in preset.items():
if getattr(args, key.replace("-", "_"), None) is None:
setattr(args, key.replace("-", "_"), default_val)
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
# Cache-dit config (Wan2.2 only)
cache_config = None
if args.cache_backend == "cache_dit":
cache_config = {
# DBCache parameters [cache-dit only]
"Fn_compute_blocks": 1, # Optimized for single-transformer models
"Bn_compute_blocks": 0, # Number of backward compute blocks
"max_warmup_steps": 4, # Maximum warmup steps (works for few-step models)
"Fn_compute_blocks": 1,
"Bn_compute_blocks": 0,
"max_warmup_steps": 4,
"max_cached_steps": 20,
"residual_diff_threshold": 0.24, # Higher threshold for more aggressive caching
"max_continuous_cached_steps": 3, # Limit to prevent precision degradation
# TaylorSeer parameters [cache-dit only]
"enable_taylorseer": False, # Disabled by default (not suitable for few-step models)
"taylorseer_order": 1, # TaylorSeer polynomial order
# SCM (Step Computation Masking) parameters [cache-dit only]
"scm_steps_mask_policy": None, # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
"scm_steps_policy": "dynamic", # SCM steps policy: "dynamic" or "static"
"residual_diff_threshold": 0.24,
"max_continuous_cached_steps": 3,
"enable_taylorseer": False,
"taylorseer_order": 1,
"scm_steps_mask_policy": None,
"scm_steps_policy": "dynamic",
}
# Configure parallel settings
parallel_config = DiffusionParallelConfig(
ulysses_degree=args.ulysses_degree,
@@ -176,14 +221,11 @@ def main():
# Check if profiling is requested via environment variable
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
omni = Omni(
omni_kwargs = dict(
model=args.model,
enable_layerwise_offload=args.enable_layerwise_offload,
vae_use_slicing=args.vae_use_slicing,
vae_use_tiling=args.vae_use_tiling,
boundary_ratio=args.boundary_ratio,
flow_shift=args.flow_shift,
enable_cache_dit_summary=args.enable_cache_dit_summary,
enable_cpu_offload=args.enable_cpu_offload,
parallel_config=parallel_config,
enforce_eager=args.enforce_eager,
@@ -191,6 +233,18 @@ def main():
cache_config=cache_config,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
)
if args.boundary_ratio is not None:
omni_kwargs["boundary_ratio"] = args.boundary_ratio
if args.flow_shift is not None:
omni_kwargs["flow_shift"] = args.flow_shift
if args.quantization is not None:
omni_kwargs["quantization"] = args.quantization
if args.cache_backend is not None:
omni_kwargs["cache_backend"] = args.cache_backend
omni_kwargs["cache_config"] = cache_config
omni_kwargs["enable_cache_dit_summary"] = args.enable_cache_dit_summary
omni = Omni(**omni_kwargs)
if profiler_enabled:
print("[Profiler] Starting profiling...")
@@ -210,22 +264,25 @@ def main():
print(f" Video size: {args.width}x{args.height}")
print(f"{'=' * 60}\n")
prompt_dict = {"prompt": args.prompt}
if args.negative_prompt:
prompt_dict["negative_prompt"] = args.negative_prompt
sampling_kwargs = dict(
height=args.height,
width=args.width,
generator=generator,
guidance_scale=args.guidance_scale,
num_inference_steps=args.num_inference_steps,
num_frames=args.num_frames,
)
if args.guidance_scale_high is not None:
sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
generation_start = time.perf_counter()
frames = omni.generate(
{
"prompt": args.prompt,
"negative_prompt": args.negative_prompt,
},
OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
generator=generator,
guidance_scale=args.guidance_scale,
guidance_scale_2=args.guidance_scale_high,
num_inference_steps=args.num_inference_steps,
num_frames=args.num_frames,
frame_rate=frame_rate,
),
prompt_dict,
OmniDiffusionSamplingParams(**sampling_kwargs),
)
generation_end = time.perf_counter()
generation_time = generation_end - generation_start
@@ -0,0 +1,64 @@
# Voxtral TTS Offline Inference
`end2end.py` runs Voxtral TTS end-to-end offline inference using vLLM. It supports both blocking (`Omni`) and streaming (`AsyncOmni`) generation, batched prompts with configurable concurrency, and voice selection via preset name or reference audio file.
When `mistral_common` has `SpeechRequest` support, prompt token IDs are built via `encode_speech_request`. Otherwise, the script falls back to manual token construction.
## Usage Examples
```bash
# Basic single-prompt with cheerful_female voice preset
python3 examples/offline_inference/voxtral_tts/end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
--write-audio --voice cheerful_female \
--model mistralai/Voxtral-4B-TTS-2603 \
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
# 32 replicate prompts with cheerful_female voice preset
python3 examples/offline_inference/voxtral_tts/end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
--num-prompts 32 --write-audio --voice cheerful_female \
--model mistralai/Voxtral-4B-TTS-2603 \
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
# Streaming with neutral_female voice preset
python3 examples/offline_inference/voxtral_tts/end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
--streaming --write-audio --voice neutral_female \
--model mistralai/Voxtral-4B-TTS-2603 \
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
# 32 prompts, 8 concurrent requests per wave, streaming with neutral_female voice
python3 examples/offline_inference/voxtral_tts/end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
--num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \
--model mistralai/Voxtral-4B-TTS-2603 \
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
# Short debug prompt with reference audio
# Note: Reference audio capability is not yet released.
python3 examples/offline_inference/voxtral_tts/end2end.py \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
--write-audio \
--model mistralai/Voxtral-4B-TTS-2603 \
--text "This is a test message." \
--audio-path path/to/reference_audio.wav
```
## Arguments
| Argument | Description |
|---|---|
| `--model PATH` | HuggingFace repo ID or local directory path (default: `mistralai/Voxtral-4B-TTS-2603`) |
| `--text TEXT` | Text to synthesize (default: `"This is a test message."`) |
| `--audio-path PATH` | Path to reference audio file for voice cloning |
| `--output-dir DIR` | Directory to write output WAV files (default: `output_audio`) |
| `--stage-configs-path PATH` | Path to stage configs YAML (currently it must be set for VoxtralTTS) |
| `--num-prompts N` | Number of replicate prompts to run for measuring performance (default: 1) |
| `--streaming` | Use streaming generation via `AsyncOmni` (default: blocking `Omni`) |
| `--concurrency N` | Max concurrent requests per wave (must be used with `--streaming`, must evenly divide `--num-prompts`) |
| `--voice NAME` | Voice preset to use instead of reference audio file (e.g., casual_female, casual_male, cheerful_female, neutral_female, neutral_male) |
| `--write-audio` | Write generated audio to WAV files |
| `--profiling-mode` | Enable profiling mode (reduces max tokens to 50) |
| `--log-stats` | Enable detailed statistics logging |
@@ -0,0 +1,377 @@
"""
This example shows how to use vLLM for running Voxtral TTS
"""
import asyncio
import gc
import logging
import os
import time
import uuid
from argparse import Namespace
from pathlib import Path
from typing import Any
import numpy as np
import soundfile as sf
import torch
from mistral_common.protocol.instruct.chunk import TextChunk
try:
from mistral_common.protocol.speech.request import SpeechRequest
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
except ImportError:
raise ImportError(
"Could not import SpeechRequest or MistralTokenizer from mistral_common. "
"Please pull the latest mistral-common code and install it with: "
"pip install -e ."
)
from vllm import SamplingParams
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import AsyncOmni
from vllm_omni.entrypoints.omni import Omni
logger = logging.getLogger(__name__)
# ---- Streaming version ----
async def run_streaming(inputs, sampling_params_list, model_name, args, output_dir):
async_omni = AsyncOmni(
model=model_name,
stage_configs_path=args.stage_configs_path,
log_stats=args.log_stats,
)
# Normalize to a list so batch and single-input share the same code path
if isinstance(inputs, list):
inputs_list = inputs
else:
inputs_list = [inputs]
total_audio_dur = 0.0
total_gen_time = 0.0
total_ttfa = 0.0
ttfa_count = 0
total_waits = 0
total_chunks = 0
results_lock = asyncio.Lock()
async def _generate_one(batch_idx, single_input):
nonlocal total_audio_dur, total_gen_time, total_ttfa, ttfa_count
nonlocal total_waits, total_chunks
request_id = str(uuid.uuid4())
all_audio_chunks = []
chunk_arrival_times = []
chunk_durations = []
chunk_idx = 0
gen_start = time.time()
ttfa = None
accumulated_sample = 0
async for stage_output in async_omni.generate(single_input, request_id, sampling_params_list):
mm_output = stage_output.multimodal_output
finished = stage_output.finished
if not mm_output or "audio" not in mm_output:
continue
now = time.time()
if ttfa is None:
ttfa = now - gen_start
audio_chunk = mm_output["audio"]
# audio_chunk is a 1-D tensor from decode_helper_async
if isinstance(audio_chunk, torch.Tensor):
if finished:
# Last chunk may return whole audio instead of chunk delta.
# Cut accumulated samples from previous chunks.
audio_numpy = audio_chunk[accumulated_sample:].float().detach().cpu().numpy()
else:
audio_numpy = audio_chunk.float().detach().cpu().numpy()
elif isinstance(audio_chunk, list):
# Audio_chunk list contain all previous chunks. Use chunk_idx to index
audio_numpy = audio_chunk[chunk_idx].float().detach().cpu().numpy()
else:
audio_numpy = audio_chunk
accumulated_sample += len(audio_numpy)
all_audio_chunks.append(audio_numpy)
chunk_arrival_times.append(now)
chunk_durations.append(len(audio_numpy) / 24000)
chunk_idx += 1
gen_elapsed = time.time() - gen_start
# Analyze wait/no-wait per chunk:
# A "wait" means the client has finished playing all previous audio
# before this chunk arrived, so there is a playback gap.
# buffer_time > 0: client still has audio to play (no wait)
# buffer_time <= 0: client ran out of audio and had to wait
chunk_labels = []
accumulated_audio_dur = 0.0
if chunk_arrival_times:
first_arrival = chunk_arrival_times[0]
for i in range(len(chunk_arrival_times)):
if i == 0:
chunk_labels.append("no_wait")
else:
playback_elapsed = chunk_arrival_times[i] - first_arrival
buffer_time = accumulated_audio_dur - playback_elapsed
if buffer_time <= 0:
chunk_labels.append("wait")
else:
chunk_labels.append("no_wait")
accumulated_audio_dur += chunk_durations[i]
req_wait_count = sum(1 for chunk_label in chunk_labels if chunk_label == "wait")
req_wait_rate = req_wait_count / len(chunk_labels) if chunk_labels else 0.0
# Concatenate all chunks for this request's audio
if all_audio_chunks:
full_audio = np.concatenate(all_audio_chunks)
output_audio_dur = len(full_audio) / 24000
output_path = os.path.join(output_dir, f"tts_output_{batch_idx}.wav")
if args.write_audio:
sf.write(output_path, full_audio, 24000)
print(
f"Request {batch_idx}: saved {len(full_audio)} samples ({output_audio_dur:.2f}s) to {output_path}"
)
# Per-chunk wait details
for i, label in enumerate(chunk_labels):
dur_ms = chunk_durations[i] * 1000
arrival_ms = (chunk_arrival_times[i] - first_arrival) * 1000 if i > 0 else 0.0
print(
f" Request {batch_idx} chunk {i}: {label} | arrived={arrival_ms:.1f}ms | chunk_dur={dur_ms:.1f}ms"
)
print(
f"Request {batch_idx}: TTFA={ttfa:.4f}s | "
f"Generation={gen_elapsed:.4f}s | "
f"Audio={output_audio_dur:.2f}s | "
f"RTF={output_audio_dur / gen_elapsed:.4f} | "
f"WaitRate={req_wait_rate:.2%} ({req_wait_count}/{len(chunk_labels)})"
)
async with results_lock:
total_audio_dur += output_audio_dur
total_gen_time += gen_elapsed
total_waits += req_wait_count
total_chunks += len(chunk_labels)
if ttfa is not None:
total_ttfa += ttfa
ttfa_count += 1
# Launch requests in waves of `concurrency` size
concurrency = args.concurrency or len(inputs_list)
gen_start_all = time.time()
for wave_start in range(0, len(inputs_list), concurrency):
wave = inputs_list[wave_start : wave_start + concurrency]
wave_idx_offset = wave_start
print(f"\n--- Wave {wave_start // concurrency + 1} (requests {wave_start}-{wave_start + len(wave) - 1}) ---")
await asyncio.gather(*[_generate_one(wave_idx_offset + i, inp) for i, inp in enumerate(wave)])
generation_time = time.time() - gen_start_all
avg_ttfa = total_ttfa / ttfa_count if ttfa_count else float("nan")
overall_wait_rate = total_waits / total_chunks if total_chunks else float("nan")
print(
f"\nAll requests: Generation={generation_time:.4f}s | "
f"TotalAudio={total_audio_dur:.2f}s | "
f"Concurrency={concurrency} | "
f"AvgTTFA={avg_ttfa:.4f}s | "
f"RTF(total)={total_audio_dur / generation_time:.4f} | "
f"RTF(per-request)={total_audio_dur / total_gen_time:.4f} | "
f"WaitRate={overall_wait_rate:.2%} ({total_waits}/{total_chunks})"
)
async_omni.shutdown()
torch.cuda.empty_cache()
gc.collect()
# ---- Non-streaming version ----
def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir):
llm = Omni(
model=model_name,
log_stats=args.log_stats,
stage_configs_path=args.stage_configs_path,
)
if args.profiling_mode:
llm.start_profile()
vllm_start = time.time()
outputs = llm.generate(inputs, sampling_params_list)
vllm_elapsed = time.time() - vllm_start
if args.profiling_mode:
llm.stop_profile()
time.sleep(10)
print(f"vLLM run time: {vllm_elapsed:.4f}s")
output_audio_dur = 0.0
for batch_idx, o in enumerate(outputs):
audio_tensor = torch.cat(o.multimodal_output["audio"])
audio_array = audio_tensor.tolist()
output_audio_dur += float(len(audio_array)) / 24000
if args.write_audio:
output_path = os.path.join(output_dir, f"tts_output_{batch_idx}.wav")
sf.write(output_path, audio_array, 24000)
print(f"Audio saved to {output_path}")
print(f"Total audio duration: {output_audio_dur:.2f}s")
print(f"RTF: {output_audio_dur / vllm_elapsed:.4f}")
del llm
torch.cuda.empty_cache()
gc.collect()
def parse_args() -> Namespace:
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
parser.add_argument(
"--model",
type=str,
default="mistralai/Voxtral-4B-TTS-2603",
help="Model name or path.",
)
parser.add_argument(
"--text",
type=str,
default="This is a test message.",
help="Text to synthesize.",
)
parser.add_argument(
"--audio-path",
type=str,
default=None,
help="Path to reference audio file for voice cloning.",
)
parser.add_argument(
"--output-dir",
type=str,
default="output_audio",
help="Directory to write output wav files.",
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="Path to stage configs YAML. Auto-resolved from model if not set.",
)
parser.add_argument(
"--num-prompts", type=int, default=1, help="Number of replicate prompts to run for measuring performance"
)
parser.add_argument(
"--profiling-mode",
default=False,
action="store_true",
help="Set max_num_tokens=2 to reduce profiling time",
)
parser.add_argument(
"--log-stats",
action="store_true",
default=False,
help="Enable writing detailed statistics (default: disabled)",
)
parser.add_argument(
"--write-audio",
action="store_true",
default=False,
help="Write audio output to files.",
)
parser.add_argument(
"--streaming",
action="store_true",
default=False,
help="Use streaming generation via AsyncOmni instead of blocking Omni.",
)
parser.add_argument(
"--concurrency",
type=int,
default=None,
help="Max concurrent requests per wave (default: all at once). Must evenly divide --num-prompts.",
)
parser.add_argument(
"--voice",
type=str,
default=None,
help="Voice to use instead of audio file.",
)
return parser.parse_args()
def compose_request(
model_name: str,
text_chunk: TextChunk,
audio_prompt_file: str,
args: Any,
) -> dict:
"""Build the full TTS input dict (prompt_token_ids, multi_modal_data or additional_information)."""
inputs: dict[str, Any] = {}
if Path(model_name).is_dir():
mistral_tokenizer = MistralTokenizer.from_file(str(Path(model_name) / "tekken.json"))
else:
mistral_tokenizer = MistralTokenizer.from_hf_hub(model_name)
instruct_tokenizer = mistral_tokenizer.instruct_tokenizer
if args.voice is not None:
tokenized = instruct_tokenizer.encode_speech_request(SpeechRequest(input=text_chunk.text, voice=args.voice))
inputs["additional_information"] = {"voice": [args.voice]}
inputs["prompt_token_ids"] = tokenized.tokens
else:
with open(audio_prompt_file, "rb") as f:
ref_audio_bytes = f.read()
tokenized = instruct_tokenizer.encode_speech_request(
SpeechRequest(input=text_chunk.text, ref_audio=ref_audio_bytes)
)
audio = tokenized.audios[0]
inputs["multi_modal_data"] = {"audio": [(audio.audio_array, audio.sampling_rate)]}
inputs["prompt_token_ids"] = tokenized.tokens
return inputs
def main(args: Any) -> None:
max_num_tokens = 50 if args.profiling_mode else 2500
model_name = args.model
output_dir = args.output_dir
if args.voice is None and args.audio_path is None:
raise ValueError("Either --voice or --audio-path must be provided.")
audio_prompt_file = args.audio_path
text_chunk = TextChunk(text=args.text)
if args.write_audio:
os.makedirs(output_dir, exist_ok=True)
inputs = compose_request(model_name, text_chunk, audio_prompt_file, args)
sampling_params = SamplingParams(
max_tokens=max_num_tokens,
)
sampling_params_list = [
sampling_params,
sampling_params,
]
if args.num_prompts > 1:
inputs = [inputs] * args.num_prompts
if args.concurrency is not None:
assert args.streaming, "--concurrency must be used with --streaming on AsyncOmni"
assert args.num_prompts % args.concurrency == 0, (
f"--num-prompts ({args.num_prompts}) must be divisible by --concurrency ({args.concurrency})"
)
torch.cuda.empty_cache()
gc.collect()
if args.streaming:
asyncio.run(run_streaming(inputs, sampling_params_list, model_name, args, output_dir))
else:
run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir)
if __name__ == "__main__":
args = parse_args()
main(args)
@@ -0,0 +1,68 @@
#!/bin/bash
# HunyuanVideo-1.5 image-to-video curl example using the async video job API.
set -euo pipefail
INPUT_IMAGE="${INPUT_IMAGE:-test_input.jpg}"
BASE_URL="${BASE_URL:-http://localhost:8099}"
OUTPUT_PATH="${OUTPUT_PATH:-hunyuan_video_15_i2v.mp4}"
POLL_INTERVAL="${POLL_INTERVAL:-2}"
if [ ! -f "$INPUT_IMAGE" ]; then
echo "Input image not found: $INPUT_IMAGE"
echo "Provide an image via INPUT_IMAGE env var."
exit 1
fi
create_response=$(
curl -sS -X POST "${BASE_URL}/v1/videos" \
-H "Accept: application/json" \
-F "prompt=The camera follows the puppy as it runs forward on the grass, its four legs alternating steps, its tail held high and wagging side to side." \
-F "input_reference=@${INPUT_IMAGE}" \
-F "size=832x480" \
-F "num_frames=33" \
-F "fps=24" \
-F "num_inference_steps=30" \
-F "guidance_scale=6.0" \
-F "flow_shift=5.0" \
-F "seed=42"
)
video_id="$(echo "${create_response}" | jq -r '.id')"
if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then
echo "Failed to create video job:"
echo "${create_response}" | jq .
exit 1
fi
echo "Created video job ${video_id}"
echo "${create_response}" | jq .
while true; do
status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
status="$(echo "${status_response}" | jq -r '.status')"
case "${status}" in
queued|in_progress)
echo "Video job ${video_id} status: ${status}"
sleep "${POLL_INTERVAL}"
;;
completed)
echo "${status_response}" | jq .
break
;;
failed)
echo "Video generation failed:"
echo "${status_response}" | jq .
exit 1
;;
*)
echo "Unexpected status response:"
echo "${status_response}" | jq .
exit 1
;;
esac
done
curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
echo "Saved video to ${OUTPUT_PATH}"
@@ -0,0 +1,31 @@
#!/bin/bash
# HunyuanVideo-1.5 image-to-video online serving startup script
#
# 480p: ~35 GB VRAM (BF16), fits 1x A100 80GB
# 720p: needs FP8 + VAE tiling, ~35 GB VRAM
MODEL="${MODEL:-hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v}"
PORT="${PORT:-8099}"
FLOW_SHIFT="${FLOW_SHIFT:-5.0}"
QUANTIZATION="${QUANTIZATION:-}"
CACHE_BACKEND="${CACHE_BACKEND:-none}"
echo "Starting HunyuanVideo-1.5 I2V server..."
echo "Model: $MODEL"
echo "Port: $PORT"
echo "Flow shift: $FLOW_SHIFT"
echo "Quantization: ${QUANTIZATION:-none}"
echo "Cache backend: $CACHE_BACKEND"
EXTRA_FLAGS=""
if [ -n "$QUANTIZATION" ]; then
EXTRA_FLAGS="$EXTRA_FLAGS --quantization $QUANTIZATION"
fi
if [ "$CACHE_BACKEND" != "none" ]; then
EXTRA_FLAGS="$EXTRA_FLAGS --cache-backend $CACHE_BACKEND"
fi
vllm serve "$MODEL" --omni \
--port "$PORT" \
--flow-shift "$FLOW_SHIFT" \
$EXTRA_FLAGS
@@ -9,9 +9,9 @@ stage_args:
- stage_id: 0
runtime:
devices: "0,1"
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -0,0 +1,60 @@
#!/bin/bash
# HunyuanVideo-1.5 text-to-video curl example using the async video job API.
set -euo pipefail
BASE_URL="${BASE_URL:-http://localhost:8098}"
OUTPUT_PATH="${OUTPUT_PATH:-hunyuan_video_15_t2v.mp4}"
POLL_INTERVAL="${POLL_INTERVAL:-2}"
create_response=$(
curl -sS -X POST "${BASE_URL}/v1/videos" \
-H "Accept: application/json" \
-F "prompt=A little girl wearing a straw hat runs through a summer meadow full of wildflowers. A wide shot is used, with the camera panning right to follow her." \
-F "size=832x480" \
-F "num_frames=33" \
-F "fps=24" \
-F "num_inference_steps=30" \
-F "guidance_scale=6.0" \
-F "flow_shift=5.0" \
-F "seed=42"
)
video_id="$(echo "${create_response}" | jq -r '.id')"
if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then
echo "Failed to create video job:"
echo "${create_response}" | jq .
exit 1
fi
echo "Created video job ${video_id}"
echo "${create_response}" | jq .
while true; do
status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
status="$(echo "${status_response}" | jq -r '.status')"
case "${status}" in
queued|in_progress)
echo "Video job ${video_id} status: ${status}"
sleep "${POLL_INTERVAL}"
;;
completed)
echo "${status_response}" | jq .
break
;;
failed)
echo "Video generation failed:"
echo "${status_response}" | jq .
exit 1
;;
*)
echo "Unexpected status response:"
echo "${status_response}" | jq .
exit 1
;;
esac
done
curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
echo "Saved video to ${OUTPUT_PATH}"
@@ -0,0 +1,31 @@
#!/bin/bash
# HunyuanVideo-1.5 text-to-video online serving startup script
#
# 480p: ~35 GB VRAM (BF16), fits 1x A100 80GB
# 720p: needs FP8 + VAE tiling, ~35 GB VRAM
MODEL="${MODEL:-hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v}"
PORT="${PORT:-8098}"
FLOW_SHIFT="${FLOW_SHIFT:-5.0}"
QUANTIZATION="${QUANTIZATION:-}"
CACHE_BACKEND="${CACHE_BACKEND:-none}"
echo "Starting HunyuanVideo-1.5 T2V server..."
echo "Model: $MODEL"
echo "Port: $PORT"
echo "Flow shift: $FLOW_SHIFT"
echo "Quantization: ${QUANTIZATION:-none}"
echo "Cache backend: $CACHE_BACKEND"
EXTRA_FLAGS=""
if [ -n "$QUANTIZATION" ]; then
EXTRA_FLAGS="$EXTRA_FLAGS --quantization $QUANTIZATION"
fi
if [ "$CACHE_BACKEND" != "none" ]; then
EXTRA_FLAGS="$EXTRA_FLAGS --cache-backend $CACHE_BACKEND"
fi
vllm serve "$MODEL" --omni \
--port "$PORT" \
--flow-shift "$FLOW_SHIFT" \
$EXTRA_FLAGS
@@ -0,0 +1,269 @@
"""
- Make sure to install the following for this example to function correctly:
- `pip install -e .`
- `pip install gradio==5.50 mistral_common=1.10.0`
Example use case:
python examples/online_serving/voxtral_tts/gradio_demo.py --host slurm-199-077 --port 8000
"""
import argparse
import io
import json
import logging
import socket
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any
import gradio as gr
import httpx
import numpy as np
import soundfile as sf
logger = logging.getLogger()
LOGFORMAT = "%(asctime)s - %(levelname)s - %(message)s"
TIMEFORMAT = "%Y-%m-%d %H:%M:%S"
logging.basicConfig(level=logging.INFO, force=True, format=LOGFORMAT, datefmt=TIMEFORMAT)
logger.setLevel(logging.INFO)
_BASE_URL = f"{socket.gethostname()}:7860"
def run_inference(
voice_name: str,
text_prompt: str,
base_url: str,
model: str,
) -> tuple[int, np.ndarray]:
"""Call /v1/audio/speech and return (sample_rate, audio_array)."""
text_prompt = text_prompt.strip()
if not text_prompt:
raise gr.Error("Please enter a text prompt.")
payload: dict[str, Any] = {
"input": text_prompt,
"model": model,
"response_format": "wav",
"voice": voice_name,
}
response = httpx.post(
f"{base_url}/audio/speech",
json=payload,
timeout=120.0,
)
response.raise_for_status()
audio_array, sr = sf.read(io.BytesIO(response.content), dtype="float32")
return sr, audio_array
def _save_example(
outputs_dir: Path,
voice_name: str,
text_prompt: str,
sr: int,
audio_array: np.ndarray,
) -> tuple[str, str]:
"""
Save inputs/outputs for sharing.
Returns (share_id, saved_audio_path)
"""
share_id = uuid.uuid4().hex
# Save generated audio
saved_audio_path = outputs_dir / f"{share_id}_gen.wav"
sf.write(str(saved_audio_path), audio_array, sr)
meta = {
"id": share_id,
"created_at": datetime.utcnow().isoformat(),
"voice_name": voice_name,
"text_prompt": text_prompt,
"generated_audio_path": str(saved_audio_path),
}
with open(outputs_dir / f"{share_id}.json", "w") as f:
json.dump(meta, f, ensure_ascii=False)
return share_id, str(saved_audio_path)
def _load_from_share(
outputs_dir: Path | None,
request: gr.Request,
) -> tuple[str | None, str, str | None, dict[str, Any], str]:
"""
Called on page load. If ?share_id=... is present, load stored example.
Returns: (voice_name, text_prompt, output_audio, submit_btn_update, share_link_text)
"""
if outputs_dir is None:
return "Neutral_Male", "", None, gr.update(interactive=False), ""
share_id = None
if request and request.query_params:
share_id = request.query_params.get("share_id")
if not share_id:
return "Neutral_Male", "", None, gr.update(interactive=False), ""
meta_path = outputs_dir / f"{share_id}.json"
if not meta_path.exists():
logger.warning("No stored example for share_id=%s", share_id)
return "Neutral_Male", "", None, gr.update(interactive=False), ""
with open(meta_path) as f:
meta = json.load(f)
voice = meta.get("voice_name", "Neutral_Male")
text_prompt = meta.get("text_prompt", "")
gen_path = meta.get("generated_audio_path")
if _BASE_URL:
share_link = f"http://{_BASE_URL}?share_id={share_id}"
else:
share_link = f"Use this query on your current URL: ?share_id={share_id}"
return voice, text_prompt, gen_path, gr.update(interactive=True), share_link
def main(
model: str,
host: str,
port: str,
output_dir: str | None = None,
) -> None:
base_url = f"http://{host}:{port}/v1"
logger.info(f"Using speech API at: {base_url}/audio/speech")
outputs_dir: Path | None = None
if output_dir is not None:
outputs_dir = Path(output_dir)
outputs_dir.mkdir(parents=True, exist_ok=True)
with gr.Blocks(title="Voxtral TTS", fill_height=True) as demo:
gr.Markdown("## Voxtral TTS")
with gr.Row():
with gr.Column():
voice_name = gr.Dropdown(
["Neutral_Male", "Neutral_Female", "Cheerful_Female", "Casual_Male", "Casual_Female"],
label="Voice",
value="Neutral_Male",
)
text_prompt = gr.Textbox(
label="Text prompt",
placeholder="Enter the text you want to synthesize...",
lines=4,
)
with gr.Row():
reset_btn = gr.Button("Clear")
submit_btn = gr.Button("Generate audio", interactive=False)
with gr.Column():
output_audio = gr.Audio(
label="Generated audio",
show_download_button=True,
interactive=False,
autoplay=True,
type="filepath",
)
share_link_box = gr.Textbox(
label="Shareable link",
interactive=False,
show_copy_button=True,
visible=outputs_dir is not None,
)
# --- UI logic: disable submit until text is non-empty ---
def _toggle_submit(text: str):
enabled = bool(text.strip())
return gr.update(interactive=enabled)
text_prompt.change(
fn=_toggle_submit,
inputs=[text_prompt],
outputs=submit_btn,
)
# --- Wiring inference + persistence to the button ---
def _on_submit(voice: str, text: str):
assert text.strip() != ""
sr, audio_array = run_inference(voice, text, base_url, model)
if outputs_dir is not None:
share_id, saved_audio_path = _save_example(
outputs_dir,
voice_name=voice,
text_prompt=text,
sr=sr,
audio_array=audio_array,
)
share_link = f"{_BASE_URL}?share_id={share_id}"
return saved_audio_path, share_link
return (sr, audio_array), ""
submit_btn.click(
fn=_on_submit,
inputs=[voice_name, text_prompt],
outputs=[output_audio, share_link_box],
)
# --- Clear everything and disable submit again ---
def _on_reset():
return (
"Neutral_Male", # voice_name
"", # text_prompt
None, # output_audio
gr.update(interactive=False), # submit_btn
"", # share_link_box
)
reset_btn.click(
fn=_on_reset,
inputs=[],
outputs=[voice_name, text_prompt, output_audio, submit_btn, share_link_box],
)
demo.load(
fn=lambda request: _load_from_share(outputs_dir, request),
inputs=[],
outputs=[voice_name, text_prompt, output_audio, submit_btn, share_link_box],
)
launch_kwargs: dict[str, Any] = {
"server_name": "0.0.0.0",
"share": True,
}
if outputs_dir is not None:
launch_kwargs["allowed_paths"] = [str(outputs_dir)]
demo.launch(**launch_kwargs)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Voxtral TTS Gradio Demo")
parser.add_argument("--model", type=str, default="mistralai/Voxtral-4B-TTS-2603", help="Name of model repo on HF")
parser.add_argument("--host", type=str, default="localhost", help="Name of host")
parser.add_argument("--port", type=str, default="8091", help="port number")
parser.add_argument(
"--output-dir",
type=str,
default=None,
help="Directory to save generated audio and share links. "
"If not provided, save/share functionality is disabled.",
)
args = parser.parse_args()
main(
model=args.model,
host=args.host,
port=args.port,
output_dir=args.output_dir,
)
+116 -26
View File
@@ -1,38 +1,101 @@
import json
import os
import subprocess
import textwrap
from pathlib import Path
import pytest
from tests.conftest import OmniServerParams
from tests.utils import hardware_test
from vllm_omni.platforms import current_omni_platform
models = ["Qwen/Qwen2.5-Omni-7B"]
stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_ci.yaml")]
if current_omni_platform.is_xpu():
stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "xpu" / "qwen2_5_omni_ci.yaml")]
# Create parameter combinations for model and stage config
test_params = [
OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
]
@pytest.mark.core_model
@pytest.mark.benchmark
@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=3)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_bench_serve_chat(omni_server):
command = [
@pytest.mark.cpu
def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
num_prompts = 5
port = 18000
result_filename = "bench-result.json"
calls_filename = "http-post-calls.json"
result_path = tmp_path / result_filename
calls_path = tmp_path / calls_filename
sitecustomize_path = tmp_path / "sitecustomize.py"
sitecustomize_path.write_text(
textwrap.dedent(
"""
import atexit
import json
import os
POST_CALLS = []
SSE_CHUNKS = [
b'data: {"choices":[{"delta":{"content":"hi"}}],"modality":"text"}\\n\\n',
b'data: {"choices":[{"delta":{"content":" there"}}],"modality":"text","metrics":{"num_tokens_out":4,"num_tokens_in":5}}\\n\\n',
b"data: [DONE]\\n\\n",
]
class _Content:
def __init__(self, chunks):
self._chunks = chunks
async def iter_any(self):
for chunk in self._chunks:
yield chunk
class MockResponse:
def __init__(self):
self.status = 200
self.reason = "OK"
self.content = _Content(SSE_CHUNKS)
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
class MockClientSession:
def __init__(self, *args, **kwargs):
pass
def post(self, url=None, *args, **kwargs):
if url is not None:
POST_CALLS.append(url)
return MockResponse()
async def close(self):
return None
import vllm_omni.benchmarks.patch.patch as patch_mod
patch_mod.aiohttp.ClientSession = MockClientSession
patch_mod.aiohttp.TCPConnector = lambda *args, **kwargs: object()
calls_file = os.environ.get("VLLM_OMNI_TEST_POST_CALLS_FILE")
if calls_file:
def _write_calls():
with open(calls_file, "w", encoding="utf-8") as f:
json.dump({"requested_urls": POST_CALLS}, f)
atexit.register(_write_calls)
"""
),
encoding="utf-8",
)
env = os.environ.copy()
env["PYTHONPATH"] = str(tmp_path) + os.pathsep + env.get("PYTHONPATH", "")
env["VLLM_OMNI_TEST_POST_CALLS_FILE"] = str(calls_path)
cmd = [
"vllm",
"bench",
"serve",
"--omni",
"--model",
omni_server.model,
"Qwen/Qwen2.5-Omni-7B",
"--port",
str(omni_server.port),
str(port),
"--dataset-name",
"random",
"--random-input-len",
@@ -40,14 +103,41 @@ def test_bench_serve_chat(omni_server):
"--random-output-len",
"4",
"--num-prompts",
"5",
str(num_prompts),
"--endpoint",
"/v1/chat/completions",
"--backend",
"openai-chat-omni",
"--disable-tqdm",
"--num-warmups",
"0",
"--ready-check-timeout-sec",
"0",
"--save-result",
"--result-dir",
str(tmp_path),
"--result-filename",
result_filename,
]
result = subprocess.run(command, capture_output=True, text=True)
print(result.stdout)
print(result.stderr)
proc = subprocess.run(
cmd,
cwd=str(Path(__file__).resolve().parents[2]),
env=env,
capture_output=True,
text=True,
)
assert result.returncode == 0, f"Benchmark failed: {result.stderr}"
assert proc.returncode == 0, f"CLI failed: stdout={proc.stdout}\nstderr={proc.stderr}"
print(f"CLI output: {proc.stdout}")
assert result_path.exists()
result = json.loads(result_path.read_text(encoding="utf-8"))
assert calls_path.exists()
calls = json.loads(calls_path.read_text(encoding="utf-8"))
expected_url = f"http://127.0.0.1:{port}/v1/chat/completions"
requested_urls = calls["requested_urls"]
sent_requests = len(requested_urls)
assert result["completed"] == sent_requests == num_prompts
assert requested_urls
assert all(url == expected_url for url in requested_urls), f"Unexpected target URLs: {requested_urls}"
+25 -10
View File
@@ -843,7 +843,7 @@ def modify_stage_config(
'async_chunk': True,
'stage_args': {
0: {'engine_args.max_model_len': 5800},
1: {'runtime.max_batch_size': 2}
1: {'engine_args.max_num_seqs': 2}
}
}
deletes: Dictionary containing configurations to delete.
@@ -1236,8 +1236,6 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st
port = params.port
stage_config_path = params.stage_config_path
if run_level == "advanced_model" and stage_config_path is not None:
# Dynamically detect stages from config to avoid KeyError
# for models with fewer stages (e.g., BAGEL has 2, Qwen3-Omni has 3)
with open(stage_config_path, encoding="utf-8") as f:
_cfg = yaml.safe_load(f) or {}
_stage_ids = [s["stage_id"] for s in _cfg.get("stage_args", []) if "stage_id" in s]
@@ -1912,6 +1910,7 @@ class OmniRunner:
def _cleanup_process(self):
try:
keywords = ["enginecore"]
matched = []
for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]):
try:
@@ -1924,16 +1923,32 @@ class OmniRunner:
if is_process:
print(f"Found vllm process: PID={proc.pid}, cmd={cmdline[:100]}")
try:
proc.terminate()
time.sleep(2)
except Exception:
proc.kill()
matched.append(proc)
except (psutil.NoSuchProcess, psutil.AccessDenied):
pass
for proc in matched:
try:
proc.terminate()
except (psutil.NoSuchProcess, psutil.AccessDenied):
pass
_, still_alive = psutil.wait_procs(matched, timeout=5)
for proc in still_alive:
try:
proc.kill()
except (psutil.NoSuchProcess, psutil.AccessDenied):
pass
if still_alive:
_, stubborn = psutil.wait_procs(still_alive, timeout=3)
if stubborn:
print(f"Warning: failed to kill residual vllm pids: {[p.pid for p in stubborn]}")
else:
print(f"Force-killed residual vllm pids: {[p.pid for p in still_alive]}")
elif matched:
print(f"Terminated vllm pids: {[p.pid for p in matched]}")
except Exception as e:
print(f"Error in psutil vllm cleanup: {e}")
@@ -0,0 +1,161 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from transformers import T5Config
from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
from vllm_omni.diffusion.models.t5_encoder.t5_encoder import (
T5EncoderModel,
T5SelfAttention,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
_T5_MODULE = "vllm_omni.diffusion.models.t5_encoder.t5_encoder"
SMALL_T5_CONFIG = dict(
d_model=64,
d_kv=8,
d_ff=128,
num_heads=8,
num_layers=2,
vocab_size=256,
relative_attention_num_buckets=32,
relative_attention_max_distance=128,
is_gated_act=True,
dense_act_fn="gelu_new",
layer_norm_epsilon=1e-6,
feed_forward_proj="gated-gelu",
)
@pytest.fixture(scope="module")
def t5_config() -> T5Config:
return T5Config(**SMALL_T5_CONFIG)
@pytest.fixture(scope="function", autouse=True)
def setup_tp_group(monkeypatch, mocker):
"""Set up TP=2, rank=0, VllmConfig, and mock activation for all tests."""
device_config = DeviceConfig(device="cpu")
# TP world size
monkeypatch.setattr("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", lambda: 2)
monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_world_size", lambda: 2)
monkeypatch.setattr(
"vllm.model_executor.layers.vocab_parallel_embedding.get_tensor_model_parallel_world_size",
lambda: 2,
)
monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_rank", lambda: 0)
monkeypatch.setattr(
"vllm.model_executor.layers.vocab_parallel_embedding.get_tensor_model_parallel_rank",
lambda: 0,
)
# TP group
mock_tp_group = mocker.MagicMock()
mock_tp_group.world_size = 2
mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_tp_group)
monkeypatch.setattr(f"{_T5_MODULE}.get_act_fn", lambda _: torch.nn.GELU())
with set_current_vllm_config(VllmConfig(device_config=device_config)):
yield
class TestRelativePositionBiasTPSlicing:
"""Verify compute_bias slices heads correctly per TP rank."""
def test_compute_bias_shape(self, t5_config):
attn = T5SelfAttention(t5_config, has_relative_attention_bias=True)
seq_len = 6
bias = attn.compute_bias(seq_len, seq_len, device=torch.device("cpu"))
local_heads = t5_config.num_heads // 2
assert bias.shape == (1, local_heads, seq_len, seq_len)
def test_all_ranks_cover_all_heads(self, t5_config, monkeypatch):
seq_len = 4
biases = []
ref_weight = None
for rank in range(2):
monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_rank", lambda r=rank: r)
attn = T5SelfAttention(t5_config, has_relative_attention_bias=True)
if rank > 0:
attn.relative_attention_bias.weight.data.copy_(ref_weight)
else:
ref_weight = attn.relative_attention_bias.weight.data.clone()
biases.append(attn.compute_bias(seq_len, seq_len, device=torch.device("cpu")))
full_bias = torch.cat(biases, dim=1)
assert full_bias.shape == (1, t5_config.num_heads, seq_len, seq_len)
class TestT5EncoderModelWeightLoading:
"""Test weight loading at the top-level T5EncoderModel."""
def test_model_instantiation(self, t5_config):
model = T5EncoderModel(t5_config, prefix="text_encoder")
assert model.config is t5_config
assert model.encoder is not None
assert len(model.encoder.block) == t5_config.num_layers
def test_embedding_shape(self, t5_config):
model = T5EncoderModel(t5_config, prefix="text_encoder")
assert model.shared.embedding_dim == t5_config.d_model
def test_embed_input_ids(self, t5_config, monkeypatch):
# Verify method and output shape
model = T5EncoderModel(t5_config, prefix="text_encoder")
# Mock all-reduce to be identity (no actual TP communication)
monkeypatch.setattr(
"vllm.model_executor.layers.vocab_parallel_embedding.tensor_model_parallel_all_reduce",
lambda x: x,
)
input_ids = torch.randint(0, t5_config.vocab_size, (2, 8))
embeddings = model.embed_input_ids(input_ids)
assert embeddings.shape == (2, 8, t5_config.d_model)
def test_qkv_weights_loaded_through_blocks(self):
# Verify that HF-style separate Q/K/V weights can be loaded through
# the block hierarchy
config = T5Config(**{**SMALL_T5_CONFIG, "num_layers": 1})
model = T5EncoderModel(config, prefix="text_encoder")
inner_dim = config.num_heads * config.d_kv
prefix = "encoder.block.0.layer.0.SelfAttention."
loaded = model.load_weights(
[
(prefix + "q.weight", torch.randn(inner_dim, config.d_model)),
(prefix + "k.weight", torch.randn(inner_dim, config.d_model)),
(prefix + "v.weight", torch.randn(inner_dim, config.d_model)),
]
)
assert len(loaded) > 0
attn = model.encoder.block[0].layer[0].SelfAttention
expected_qkv_dim = 3 * (config.num_heads // 2) * config.d_kv
assert attn.qkv_proj.weight.shape == (expected_qkv_dim, config.d_model)
class TestTPConstraints:
"""Verify that invalid TP configurations raise clear errors."""
def test_num_heads_not_divisible_by_tp(self):
config = T5Config(**{**SMALL_T5_CONFIG, "num_heads": 7})
with pytest.raises(AssertionError, match=r"num_heads.*must be divisible by tp_size"):
T5SelfAttention(config)
def test_num_heads_divisible_by_tp(self, t5_config):
attn = T5SelfAttention(t5_config)
assert attn.n_heads_per_partition == 4
@@ -0,0 +1,463 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for Int8 quantization config."""
from unittest.mock import MagicMock, patch
import pytest
import torch
from pytest_mock import MockerFixture
from torch.nn import Module, Parameter
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
from vllm_omni.diffusion.quantization import (
get_diffusion_quant_config,
get_vllm_quant_config_for_layers,
)
from vllm_omni.platforms import current_omni_platform
npu_available = pytest.mark.skipif(not current_omni_platform.is_npu(), reason="NPU platform not available.")
cuda_available = pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="GPU platform not available.")
def test_int8_config_creation():
"""Test that Int8 config can be created."""
config = get_diffusion_quant_config("int8")
assert config is not None
assert config.get_name() == "int8"
def test_vllm_config_extraction():
"""Test that vLLM config can be extracted from diffusion config."""
diff_config = get_diffusion_quant_config("int8")
vllm_config = get_vllm_quant_config_for_layers(diff_config)
assert vllm_config is not None
assert vllm_config.activation_scheme == "dynamic"
def test_none_quantization():
"""Test that None quantization returns None config."""
config = get_diffusion_quant_config(None)
assert config is None
vllm_config = get_vllm_quant_config_for_layers(config)
assert vllm_config is None
def test_invalid_quantization():
"""Test that invalid quantization method raises error."""
with pytest.raises(ValueError, match="Unknown quantization method"):
get_diffusion_quant_config("invalid_method")
def test_int8_config_with_custom_params():
"""Test Int8 config with custom parameters."""
config = get_diffusion_quant_config(
"int8",
activation_scheme="dynamic",
ignored_layers=["proj_out"],
)
assert config is not None
assert config.activation_scheme == "dynamic"
assert "proj_out" in config.ignored_layers
def test_supported_methods():
"""Test that supported methods list is correct."""
from vllm_omni.diffusion.quantization import SUPPORTED_QUANTIZATION_METHODS
assert "int8" in SUPPORTED_QUANTIZATION_METHODS
def test_quantization_integration():
"""Test end-to-end quantization flow through OmniDiffusionConfig."""
from vllm_omni.diffusion.data import OmniDiffusionConfig
# Test with quantization string only
config = OmniDiffusionConfig(model="test", quantization="int8")
assert config.quantization_config is not None
assert config.quantization_config.get_name() == "int8"
# Test with quantization_config dict
config2 = OmniDiffusionConfig(
model="test",
quantization_config={"method": "int8", "activation_scheme": "dynamic"},
)
assert config2.quantization_config is not None
assert config2.quantization_config.get_name() == "int8"
assert config2.quantization_config.activation_scheme == "dynamic"
# Test that vLLM config can be extracted
vllm_config = config.quantization_config.get_vllm_quant_config()
assert vllm_config is not None
def test_quantization_dict_not_mutated():
"""Test that passing a dict to quantization_config doesn't mutate it."""
from vllm_omni.diffusion.data import OmniDiffusionConfig
original_dict = {"method": "int8", "activation_scheme": "dynamic"}
dict_copy = original_dict.copy()
OmniDiffusionConfig(model="test", quantization_config=original_dict)
# Original dict should be unchanged
assert original_dict == dict_copy
def test_quantization_conflicting_methods_warning(caplog):
"""Test warning when quantization and quantization_config['method'] conflict."""
import logging
from vllm_omni.diffusion.data import OmniDiffusionConfig
with caplog.at_level(logging.WARNING):
config = OmniDiffusionConfig(
model="test",
quantization="int8", # This should be overridden
quantization_config={"method": "int8", "activation_scheme": "dynamic"},
)
# No warning when methods match
assert config.quantization_config is not None
def test_get_quant_method(mocker: MockerFixture):
"""Test for get_quant_method method for GPU"""
from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
diff_config = get_diffusion_quant_config("int8")
vllm_config = get_vllm_quant_config_for_layers(diff_config)
def _fake_init(self, quant_config):
pass
layer = MagicMock(spec=LinearBase)
mocker.patch.object(Int8OnlineLinearMethod, "__init__", _fake_init)
prefix = "test_layer"
# Mock the platform to be GPU
with (
patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=True),
patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=False),
):
method = vllm_config.get_quant_method(layer, prefix)
assert isinstance(method, Int8OnlineLinearMethod)
# Test skipping quantization for a layer
vllm_config.ignored_layers = [prefix]
method = vllm_config.get_quant_method(layer, prefix)
assert isinstance(method, UnquantizedLinearMethod)
def test_get_npu_quant_method():
"""Test for get_quant_method method for NPU"""
from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
diff_config = get_diffusion_quant_config("int8")
vllm_config = get_vllm_quant_config_for_layers(diff_config)
layer = MagicMock(spec=LinearBase)
prefix = "test_layer"
# Mock the platform to be NPU
with (
patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=False),
patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=True),
):
method = vllm_config.get_quant_method(layer, prefix)
assert isinstance(method, NPUInt8OnlineLinearMethod)
# Test skipping quantization for a layer
vllm_config.ignored_layers = [prefix]
method = vllm_config.get_quant_method(layer, prefix)
assert isinstance(method, UnquantizedLinearMethod)
class TestInt8LinearMethod:
@pytest.fixture
def mock_quant_config(self, mocker):
return mocker.Mock()
@pytest.fixture
def mock_kernel(self, mocker):
kernel = mocker.Mock()
kernel.process_weights_after_loading = mocker.Mock()
kernel.apply_weights = mocker.Mock(return_value=torch.randn(1, 10))
return kernel
@pytest.fixture
def patch_deps(self, mocker, mock_kernel):
# mock init_int8_linear_kernel
mocker.patch("vllm_omni.diffusion.quantization.int8.init_int8_linear_kernel", return_value=mock_kernel)
return mock_kernel
def test_init(self, patch_deps, mock_quant_config):
# test for Int8LinearMethod init
from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod, init_int8_linear_kernel
method = Int8LinearMethod(mock_quant_config)
assert method.quant_config == mock_quant_config
init_int8_linear_kernel.assert_called_once_with(
is_channelwise=False, is_static_input_scheme=False, input_symmetric=True, module_name="Int8LinearMethod"
)
assert method.int8_linear == patch_deps
def test_process_weights_after_loading(self, patch_deps, mock_quant_config):
from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
method = Int8LinearMethod(mock_quant_config)
layer = Module()
method.process_weights_after_loading(layer)
patch_deps.process_weights_after_loading.assert_called_once_with(layer)
def test_apply(self, patch_deps, mock_quant_config):
from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
method = Int8LinearMethod(mock_quant_config)
layer = Module()
x = torch.randn(1, 128)
bias = torch.randn(128)
output = method.apply(layer, x, bias)
patch_deps.apply_weights.assert_called_once_with(layer, x, bias)
assert isinstance(output, torch.Tensor)
class TestInt8OnlineLinearMethod:
@pytest.fixture
def mock_quant_config(self, mocker):
return mocker.Mock()
@pytest.fixture
def mock_deps(self, mocker):
# mock kernel
mock_kernel = mocker.Mock()
mock_kernel.layer_param_names = ("weight", "weight_scale", "input_scale", "input_zero_point", "azp_adj")
mocker.patch("vllm_omni.diffusion.quantization.int8.init_int8_linear_kernel", return_value=mock_kernel)
mocker.patch("vllm_omni.diffusion.quantization.int8.replace_parameter")
# mock scaled_int8_quant return value
mock_qweight = torch.ones((128, 64), dtype=torch.int8)
mock_scale = torch.randn(128)
mock_quant = mocker.patch(
"vllm_omni.diffusion.quantization.int8.ops.scaled_int8_quant", return_value=(mock_qweight, mock_scale, None)
)
return {"kernel": mock_kernel, "quant": mock_quant, "mock_qweight": mock_qweight, "mock_scale": mock_scale}
def test_process_weights_after_loading(self, mock_deps, mock_quant_config):
from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
method = Int8OnlineLinearMethod(mock_quant_config)
layer = Module()
layer.weight = Parameter(torch.randn(128, 64))
method.process_weights_after_loading(layer)
mock_deps["quant"].assert_called_once_with(layer.weight, scale=None)
@npu_available
class TestNPUInt8LinearMethod:
qweight_mock = torch.randn((128, 64)).to(dtype=torch.int8)
scale_mock = torch.randn(128)
out_mock = torch.randn((16, 128))
@pytest.fixture
def mock_torch_npu(self, mocker):
torch_npu = MagicMock()
mocker.patch("vllm_omni.diffusion.quantization.int8.torch_npu", return_value=torch_npu)
mocker.patch(
"vllm_omni.diffusion.quantization.int8.torch_npu.npu_dynamic_quant",
return_value=(self.qweight_mock, self.scale_mock),
)
mocker.patch("vllm_omni.diffusion.quantization.int8.torch_npu.npu_quant_matmul", return_value=self.out_mock)
return torch_npu
@pytest.fixture
def mock_quant_config(self, mocker):
return mocker.Mock()
@pytest.fixture
def mock_layer(self, mocker):
layer = torch.nn.Module()
layer.weight = torch.nn.Parameter(self.qweight_mock, requires_grad=False)
layer.weight_scale = torch.nn.Parameter(self.scale_mock, requires_grad=False)
return layer
def test_npu_int8_process_weights_after_loading(self, mock_layer, mock_quant_config, mock_torch_npu):
from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
method = NPUInt8LinearMethod(mock_quant_config)
ori_weight_shape = mock_layer.weight.shape
method.process_weights_after_loading(mock_layer)
assert mock_layer.weight.shape == ori_weight_shape[::-1]
assert mock_layer.weight.is_contiguous()
def test_npu_int8_apply(self, mock_layer, mock_quant_config, mock_torch_npu):
from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
method = NPUInt8LinearMethod(mock_quant_config)
x = torch.randn(1, 16, 64)
output = method.apply(mock_layer, x)
assert output.shape == (1, 16, 128)
def test_npu_int8_online_process_weights(self, mock_layer, mock_quant_config, mock_torch_npu):
from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
method = NPUInt8OnlineLinearMethod(mock_quant_config)
method.process_weights_after_loading(mock_layer)
assert mock_layer.weight.shape == (64, 128)
assert torch.equal(mock_layer.weight_scale, self.scale_mock)
@pytest.fixture
def quant_config():
"""Shared quant config fixture for smoke tests."""
from vllm_omni.diffusion.quantization.int8 import Int8Config
return Int8Config(
is_checkpoint_int8_serialized=False,
activation_scheme="dynamic",
)
@npu_available
class TestNPUInt8Smoke:
"""Smoke tests using real torch_npu, only run on NPU."""
@pytest.fixture
def real_layer(self):
"""Create a real linear layer with fp16 weights on NPU"""
layer = torch.nn.Module()
layer.weight = torch.nn.Parameter(
torch.randn(128, 64, dtype=torch.float16, device="npu"),
requires_grad=False,
)
layer.logical_widths = [128]
layer.input_size_per_partition = 64
layer.output_size_per_partition = 128
layer.orig_dtype = torch.float16
return layer
def test_real_npu_dynamic_quant_shape_contract(self, quant_config, real_layer):
"""Smoke test: verify npu_dynamic_quant returns correct shapes."""
import torch_npu
# Call real torch_npu.npu_dynamic_quant
weight = real_layer.weight
qweight, scale = torch_npu.npu_dynamic_quant(weight)
assert qweight.shape == weight.shape
assert qweight.dtype == torch.int8
assert scale.shape == (weight.shape[0],)
def test_real_npu_online_process_weights_after_loading(self, quant_config, real_layer):
"""Smoke test: full process_weights_after_loading with real torch_npu."""
from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
method = NPUInt8OnlineLinearMethod(quant_config)
method.process_weights_after_loading(real_layer)
assert real_layer.weight.shape == (64, 128)
assert real_layer.weight.dtype == torch.int8
assert hasattr(real_layer, "weight_scale")
assert real_layer.weight_scale.shape == (128,)
def test_real_npu_int8_apply_forward(self, quant_config):
"""Smoke test: forward pass with real npu_quant_matmul."""
import torch_npu
from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
method = NPUInt8LinearMethod(quant_config)
# Create layer with pre-processed weights on NPU
layer = torch.nn.Module()
weight_fp16 = torch.randn(128, 64, dtype=torch.float16, device="npu")
qweight, scale = torch_npu.npu_dynamic_quant(weight_fp16)
layer.weight = torch.nn.Parameter(qweight.t().contiguous(), requires_grad=False)
layer.weight_scale = torch.nn.Parameter(scale.squeeze(), requires_grad=False)
# Forward pass on NPU
x = torch.randn(2, 16, 64, dtype=torch.float16, device="npu")
output = method.apply(layer, x)
assert output.shape == (2, 16, 128)
assert output.dtype == torch.float16
@cuda_available
class TestCudaInt8Smoke:
"""Smoke tests using real CUDA kernels, only on CUDA"""
@pytest.fixture
def real_layer(self):
"""Create a real linear layer with fp16 weights on CUDA"""
layer = torch.nn.Module()
layer.weight = torch.nn.Parameter(
torch.randn(128, 64, dtype=torch.float16, device="cuda"),
requires_grad=False,
)
layer.logical_widths = [128]
layer.input_size_per_partition = 64
layer.output_size_per_partition = 128
layer.orig_dtype = torch.float16
return layer
def test_real_cuda_scaled_int8_quant_shape_contract(self, quant_config):
"""Smoke test: verify scaled_int8_quant returns correct shapes."""
from vllm import _custom_ops as ops
weight = torch.randn(128, 64, dtype=torch.float16, device="cuda")
qweight, scale, _ = ops.scaled_int8_quant(weight, scale=None)
assert qweight.shape == weight.shape
assert qweight.dtype == torch.int8
assert scale.shape == (weight.shape[0], 1)
def test_real_cuda_online_process_weights_after_loading(self, quant_config, real_layer):
"""Smoke test: full process_weights_after_loading with real CUDA ops."""
from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
method = Int8OnlineLinearMethod(quant_config)
method.process_weights_after_loading(real_layer)
assert real_layer.weight.shape == (64, 128)
assert real_layer.weight.dtype == torch.int8
assert hasattr(real_layer, "weight_scale")
def test_real_cuda_int8_apply_forward(self, quant_config):
"""Smoke test: forward pass with real CUDA int8 kernel."""
from vllm import _custom_ops as ops
from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
method = Int8LinearMethod(quant_config)
# Create layer with pre-processed weights
layer = torch.nn.Module()
weight_fp16 = torch.randn(128, 64, dtype=torch.float16, device="cuda")
qweight, scale, _ = ops.scaled_int8_quant(weight_fp16, scale=None)
layer.weight = torch.nn.Parameter(qweight.t(), requires_grad=False)
layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False)
# Set required attributes for kernel
layer.input_scale = None
layer.input_zero_point = None
layer.azp_adj = None
# Forward pass
x = torch.randn(2, 16, 64, dtype=torch.float16, device="cuda")
output = method.apply(layer, x)
assert output.shape == (2, 16, 128)
assert output.dtype == torch.float16
@@ -0,0 +1,507 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for step-level diffusion runner and worker execution."""
import os
from contextlib import contextmanager
from types import SimpleNamespace
import pytest
import torch
import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
from tests.utils import hardware_test
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D
from vllm_omni.diffusion.distributed.parallel_state import (
destroy_distributed_env,
get_sp_group,
init_distributed_environment,
initialize_model_parallel,
)
from vllm_omni.diffusion.ipc import (
pack_diffusion_output_shm,
unpack_diffusion_output_shm,
)
from vllm_omni.diffusion.sched.interface import (
DiffusionRequestState as SchedulerRequestState,
)
from vllm_omni.diffusion.sched.interface import (
DiffusionSchedulerOutput,
)
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
from vllm_omni.diffusion.worker.utils import RunnerOutput
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
# ---------------------------------------------------------------------------
# Helpers & fixtures
# ---------------------------------------------------------------------------
@contextmanager
def _noop_forward_context(*args, **kwargs):
del args, kwargs
yield
def _update_environment_variables(envs_dict: dict[str, str]) -> None:
for key, value in envs_dict.items():
os.environ[key] = value
class _StepPipeline:
"""Minimal pipeline stub that supports step-wise execution."""
supports_step_execution = True
def __init__(self):
self.prepare_calls = 0
self.denoise_calls = 0
self.scheduler_calls = 0
self.decode_calls = 0
def prepare_encode(self, state, **kwargs):
del kwargs
self.prepare_calls += 1
state.timesteps = [torch.tensor(10), torch.tensor(5)]
state.latents = torch.tensor([0.0])
return state
def denoise_step(self, state, **kwargs):
del state, kwargs
self.denoise_calls += 1
return torch.tensor([1.0])
def step_scheduler(self, state, noise_pred, **kwargs):
del noise_pred, kwargs
self.scheduler_calls += 1
state.step_index += 1
def post_decode(self, state, **kwargs):
del kwargs
self.decode_calls += 1
return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
class _IdentityNoiseTransformer(torch.nn.Module):
def forward(self, x: torch.Tensor, **kwargs):
del kwargs
return (x,)
class _AdditiveScheduler:
def step(self, noise_pred: torch.Tensor, t: torch.Tensor, latents: torch.Tensor, return_dict: bool = False):
del t, return_dict
return (latents + noise_pred,)
class _DistributedStepPipeline(CFGParallelMixin):
supports_step_execution = True
def __init__(self, mode: str, device: torch.device):
self.mode = mode
self.device = device
self._interrupt = False
self.scheduler = _AdditiveScheduler()
self.transformer = _IdentityNoiseTransformer()
@property
def interrupt(self):
return self._interrupt
def prepare_encode(self, state, **kwargs):
del kwargs
state.timesteps = [torch.tensor(1.0, device=self.device)]
state.latents = torch.ones((1, 1), device=self.device)
state.step_index = 0
state.scheduler = self.scheduler
state.do_true_cfg = self.mode == "cfg"
return state
def denoise_step(self, state, **kwargs):
del kwargs
if self.mode == "ulysses":
sp_group = get_sp_group().ulysses_group
seq_world_size = torch.distributed.get_world_size(sp_group)
input_tensor = torch.randn(1, 2, 2 * seq_world_size, 2, device=self.device)
original = input_tensor.clone()
intermediate = SeqAllToAll4D.apply(sp_group, input_tensor, 2, 1, False)
output = SeqAllToAll4D.apply(sp_group, intermediate, 1, 2, False)
torch.testing.assert_close(output, original, rtol=1e-5, atol=1e-5)
return torch.ones_like(state.latents)
if self.mode == "ring":
ring_group = get_sp_group().ring_group
rank = torch.distributed.get_rank(ring_group)
world_size = torch.distributed.get_world_size(ring_group)
comm = RingComm(ring_group)
input_tensor = torch.full((1, 2, 2), float(rank + 1), device=self.device)
recv_tensor = comm.send_recv(input_tensor)
comm.commit()
comm.wait()
expected = torch.full_like(recv_tensor, float(((rank - 1) % world_size) + 1))
torch.testing.assert_close(recv_tensor, expected, rtol=1e-5, atol=1e-5)
return torch.ones_like(state.latents)
positive_kwargs = {"x": state.latents + 1}
negative_kwargs = {"x": state.latents - 1}
return self.predict_noise_maybe_with_cfg(
do_true_cfg=True,
true_cfg_scale=1.0,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=False,
)
def step_scheduler(self, state, noise_pred, **kwargs):
del kwargs
if self.mode == "cfg":
state.latents = self.scheduler_step_maybe_with_cfg(
noise_pred,
state.current_timestep,
state.latents,
do_true_cfg=True,
per_request_scheduler=state.scheduler,
)
else:
state.latents = state.latents + noise_pred
state.step_index += 1
def post_decode(self, state, **kwargs):
del kwargs
return DiffusionOutput(output=state.latents.detach().cpu())
def _make_step_request(num_inference_steps: int = 2):
return SimpleNamespace(
prompts=["a prompt"],
request_ids=["req-1"],
sampling_params=SimpleNamespace(
generator=None,
seed=None,
generator_device=None,
num_inference_steps=num_inference_steps,
),
)
def _make_runner():
runner = object.__new__(DiffusionModelRunner)
runner.vllm_config = object()
runner.od_config = SimpleNamespace(
cache_backend=None,
parallel_config=SimpleNamespace(use_hsdp=False),
)
runner.device = torch.device("cpu")
runner.pipeline = _StepPipeline()
runner.cache_backend = None
runner.offload_backend = None
runner.state_cache = {}
runner.kv_transfer_manager = SimpleNamespace()
return runner
def _make_distributed_runner(mode: str, device: torch.device):
runner = object.__new__(DiffusionModelRunner)
runner.vllm_config = object()
runner.od_config = SimpleNamespace(
cache_backend=None,
parallel_config=SimpleNamespace(use_hsdp=False),
)
runner.device = device
runner.pipeline = _DistributedStepPipeline(mode=mode, device=device)
runner.cache_backend = None
runner.offload_backend = None
runner.state_cache = {}
runner.kv_transfer_manager = SimpleNamespace()
return runner
def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_ids=None):
return DiffusionSchedulerOutput(
step_id=step_id,
req_states=[SchedulerRequestState(sched_req_id=sched_req_id, req=req)],
finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
num_running_reqs=1,
num_waiting_reqs=0,
)
def _expected_output_for_mode(mode: str) -> torch.Tensor:
if mode == "cfg":
return torch.tensor([[3.0]])
return torch.tensor([[2.0]])
def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master_port: str):
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
_update_environment_variables(
{
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
"WORLD_SIZE": str(world_size),
"MASTER_ADDR": "localhost",
"MASTER_PORT": master_port,
}
)
model_runner_module.set_forward_context = _noop_forward_context
try:
init_distributed_environment()
if mode == "ulysses":
initialize_model_parallel(ulysses_degree=world_size)
elif mode == "ring":
initialize_model_parallel(ring_degree=world_size)
elif mode == "cfg":
initialize_model_parallel(cfg_parallel_size=world_size)
else:
raise ValueError(f"Unsupported distributed test mode: {mode}")
runner = _make_distributed_runner(mode, device)
output = DiffusionModelRunner.execute_stepwise(
runner,
_make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0),
)
assert output.finished is True
assert output.result is not None
torch.testing.assert_close(output.result.output, _expected_output_for_mode(mode), rtol=1e-5, atol=1e-5)
assert "req-1" not in runner.state_cache
finally:
destroy_distributed_env()
# ---------------------------------------------------------------------------
# Runner / Worker
# ---------------------------------------------------------------------------
@pytest.mark.cpu
class TestRunner:
"""DiffusionModelRunner.execute_stepwise"""
def test_completes_request_and_clears_state(self, monkeypatch):
runner = _make_runner()
req = _make_step_request()
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
first = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
assert first.req_id == "req-1"
assert first.step_index == 1
assert first.finished is False
assert first.result is None
assert "req-1" in runner.state_cache
second = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=1))
assert second.req_id == "req-1"
assert second.step_index == 2
assert second.finished is True
assert second.result is not None
assert second.result.error is None
assert torch.equal(second.result.output, torch.tensor([2.0]))
assert "req-1" not in runner.state_cache
assert runner.pipeline.prepare_calls == 1
assert runner.pipeline.denoise_calls == 2
assert runner.pipeline.scheduler_calls == 2
assert runner.pipeline.decode_calls == 1
def test_load_model_rejects_unsupported_step_execution(self, monkeypatch):
class _RequestOnlyPipeline:
pass
class _FakeLoader:
def __init__(self, *args, **kwargs):
del args, kwargs
def load_model(self, **kwargs):
del kwargs
return _RequestOnlyPipeline()
class _FakeProfiler:
consumed_memory = 0
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb):
del exc_type, exc, tb
return False
runner = object.__new__(DiffusionModelRunner)
runner.vllm_config = object()
runner.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
enforce_eager=True,
cache_backend=None,
cache_config=None,
step_execution=True,
model_class_name="RequestOnlyPipeline",
parallel_config=SimpleNamespace(use_hsdp=False),
)
runner.device = torch.device("cpu")
runner.pipeline = None
runner.cache_backend = None
runner.offload_backend = None
runner.state_cache = {}
runner.kv_transfer_manager = SimpleNamespace()
monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeLoader)
monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeProfiler)
monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None)
monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None)
with pytest.raises(ValueError, match="RequestOnlyPipeline"):
DiffusionModelRunner.load_model(runner)
@pytest.mark.cpu
class TestWorker:
"""DiffusionWorker.execute_stepwise"""
def test_delegates_to_model_runner(self):
worker = object.__new__(DiffusionWorker)
expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None)
scheduler_output = SimpleNamespace(
req_states=[
SimpleNamespace(
req=SimpleNamespace(
sampling_params=SimpleNamespace(lora_request=None),
)
)
]
)
worker.lora_manager = None
worker.model_runner = SimpleNamespace(
execute_stepwise=lambda arg: expected if arg is scheduler_output else None
)
output = DiffusionWorker.execute_stepwise(worker, scheduler_output)
assert output is expected
def test_clears_active_lora_before_stepwise_execution(self):
worker = object.__new__(DiffusionWorker)
scheduler_output = SimpleNamespace(
req_states=[
SimpleNamespace(
req=SimpleNamespace(
sampling_params=SimpleNamespace(lora_request=None),
)
)
]
)
calls: list[object | None] = []
class _FakeLoRAManager:
def set_active_adapter(self, adapter):
calls.append(adapter)
worker.lora_manager = _FakeLoRAManager()
worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1"))
DiffusionWorker.execute_stepwise(worker, scheduler_output)
assert calls == [None]
def test_rejects_lora_requests_in_step_mode(self):
worker = object.__new__(DiffusionWorker)
scheduler_output = SimpleNamespace(
req_states=[
SimpleNamespace(
req=SimpleNamespace(
sampling_params=SimpleNamespace(lora_request=object()),
)
)
]
)
worker.lora_manager = None
worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1"))
with pytest.raises(ValueError, match="does not support LoRA"):
DiffusionWorker.execute_stepwise(worker, scheduler_output)
@pytest.mark.cpu
class TestIPC:
def test_pack_unpack_runner_output_shm(self):
tensor = torch.zeros(300_000, dtype=torch.float32)
output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
packed = pack_diffusion_output_shm(output)
assert isinstance(packed.result.output, dict)
assert packed.result.output["__tensor_shm__"] is True
unpacked = unpack_diffusion_output_shm(packed)
assert isinstance(unpacked.result.output, torch.Tensor)
torch.testing.assert_close(unpacked.result.output, tensor)
@pytest.mark.cpu
class TestSupportedPipelines:
"""Step-execution protocol checks for supported pipelines."""
def test_qwen_image_supports_step_execution(self):
from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImagePipeline
# Avoid loading model weights; protocol membership depends on the class contract.
pipeline = object.__new__(QwenImagePipeline)
assert pipeline.supports_step_execution is True
assert supports_step_execution(pipeline) is True
assert isinstance(pipeline, SupportsStepExecution) is True
@hardware_test(
res={"cuda": "L4"},
num_cards=2,
)
def test_execute_stepwise_with_ulysses_parallel():
world_size = 2
if current_omni_platform.get_device_count() < world_size:
pytest.skip(f"Test requires {world_size} devices")
torch.multiprocessing.spawn(
_distributed_step_worker,
args=(world_size, "ulysses", "29540"),
nprocs=world_size,
)
@hardware_test(
res={"cuda": "L4"},
num_cards=2,
)
def test_execute_stepwise_with_ring_parallel():
world_size = 2
if current_omni_platform.get_device_count() < world_size:
pytest.skip(f"Test requires {world_size} devices")
torch.multiprocessing.spawn(
_distributed_step_worker,
args=(world_size, "ring", "29541"),
nprocs=world_size,
)
@hardware_test(
res={"cuda": "L4"},
num_cards=2,
)
def test_execute_stepwise_with_cfg_parallel():
world_size = 2
if current_omni_platform.get_device_count() < world_size:
pytest.skip(f"Test requires {world_size} devices")
torch.multiprocessing.spawn(
_distributed_step_worker,
args=(world_size, "cfg", "29542"),
nprocs=world_size,
)
@@ -7,9 +7,9 @@ stage_args:
prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: BagelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -44,9 +44,9 @@ stage_args:
cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: dit
max_num_seqs: 1
gpu_memory_utilization: 0.45
enforce_eager: true
trust_remote_code: true
@@ -7,9 +7,9 @@ stage_args:
prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: OmniBagelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -43,9 +43,9 @@ stage_args:
cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: dit
max_num_seqs: 1
gpu_memory_utilization: 0.45
enforce_eager: true
trust_remote_code: true
@@ -6,7 +6,6 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0"
max_batch_size: 1
engine_args:
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
@@ -37,7 +36,6 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
@@ -67,9 +65,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
@@ -5,7 +5,7 @@
End-to-end test for Bagel img2img generation.
This test validates that the Bagel model generates images from an input image
and text prompt that match expected reference pixel values within a ±5 tolerance.
and text prompt that match expected reference pixel values within a ±10 tolerance.
Equivalent to running:
python3 examples/offline_inference/bagel/end2end.py \
@@ -25,6 +25,7 @@ from vllm.assets.image import ImageAsset
from tests.conftest import modify_stage_config
from tests.utils import hardware_test
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.platforms import current_omni_platform
# Reference pixel data extracted from the known-good output image
# Generated with seed=52, num_inference_steps=15,
@@ -43,7 +44,21 @@ REFERENCE_PIXELS = [
{"position": (256, 256), "rgb": (181, 202, 222)},
]
PIXEL_TOLERANCE = 5
if current_omni_platform.is_rocm():
REFERENCE_PIXELS = [
{"position": (100, 100), "rgb": (156, 172, 215)},
{"position": (400, 50), "rgb": (106, 144, 216)},
{"position": (700, 100), "rgb": (118, 158, 231)},
{"position": (150, 400), "rgb": (183, 23, 48)},
{"position": (512, 336), "rgb": (218, 215, 191)},
{"position": (700, 400), "rgb": (194, 14, 42)},
{"position": (100, 600), "rgb": (105, 10, 16)},
{"position": (400, 600), "rgb": (167, 33, 46)},
{"position": (700, 600), "rgb": (102, 86, 92)},
{"position": (256, 256), "rgb": (181, 201, 220)},
]
PIXEL_TOLERANCE = 10
DEFAULT_PROMPT = "<|fim_middle|><|im_start|>Change the grass color to red<|im_end|>"
@@ -191,7 +206,7 @@ def _resolve_stage_config(config_path: str, run_level: str) -> str:
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"})
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
def test_bagel_img2img_shared_memory_connector(run_level):
"""Test Bagel img2img with shared memory connector."""
input_image = _load_input_image()
@@ -31,6 +31,7 @@ from PIL import Image
from tests.conftest import modify_stage_config
from tests.utils import hardware_test
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.platforms import current_omni_platform
# Reference pixel data extracted from the known-good output image
# Each entry contains (x, y) position and expected (R, G, B) values
@@ -49,6 +50,20 @@ REFERENCE_PIXELS = [
{"position": (256, 256), "rgb": (171, 160, 153)},
]
if current_omni_platform.is_rocm():
REFERENCE_PIXELS = [
{"position": (100, 100), "rgb": (123, 119, 100)},
{"position": (400, 50), "rgb": (162, 161, 142)},
{"position": (700, 100), "rgb": (171, 156, 127)},
{"position": (150, 400), "rgb": (131, 128, 112)},
{"position": (512, 512), "rgb": (134, 61, 59)},
{"position": (700, 400), "rgb": (204, 107, 43)},
{"position": (100, 700), "rgb": (201, 180, 165)},
{"position": (400, 700), "rgb": (140, 108, 87)},
{"position": (700, 700), "rgb": (247, 205, 145)},
{"position": (256, 256), "rgb": (171, 160, 153)},
]
# Maximum allowed difference per color channel
PIXEL_TOLERANCE = 5
@@ -181,7 +196,7 @@ def _resolve_stage_config(config_path: str, run_level: str) -> str:
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"})
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
def test_bagel_text2img_shared_memory_connector(run_level):
"""Test Bagel text2img with shared memory connector."""
config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
@@ -24,6 +24,7 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
# This test is specific to Z-Image LoRA behavior. Keep it focused on a single
# model to reduce runtime and avoid extra downloads.
models = ["Tongyi-MAI/Z-Image-Turbo"]
DIFFUSION_INIT_TIMEOUT_S = 600
@pytest.mark.parametrize("model_name", models)
@@ -76,7 +77,11 @@ def test_diffusion_model(model_name: str, tmp_path: Path):
)
return str(adapter_dir)
m = Omni(model=model_name)
m = Omni(
model=model_name,
stage_init_timeout=DIFFUSION_INIT_TIMEOUT_S,
init_timeout=DIFFUSION_INIT_TIMEOUT_S,
)
try:
# high resolution may cause OOM on L4
height = 256
@@ -32,7 +32,7 @@ from tests.conftest import OmniServerParams
from tests.utils import hardware_test
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
STAGE_CONFIGS_PATH = str(
@@ -97,7 +97,7 @@ def test_bagel_text2img_online(omni_server, openai_client) -> None:
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"})
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_bagel_img2img_online(omni_server, openai_client) -> None:
"""Test Bagel img2img via OpenAI-compatible chat completions API."""
@@ -0,0 +1,98 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Tests of common diffusion feature combinations in online serving mode
for HunyuanVideo-1.5-T2V (480p).
Coverage (H100, since model cannot fit L4):
- TeaCache + Layerwise CPU offloading (1 GPU)
- CacheDiT + TP=2 + VAE patch parallel=2 (2 GPUs)
"""
import pytest
from tests.conftest import (
OmniServer,
OmniServerParams,
OpenAIClientHandler,
dummy_messages_from_mix_data,
)
from tests.utils import hardware_marks
PROMPT = "A cat walking across a sunlit garden, cinematic lighting, slow motion."
NEGATIVE_PROMPT = "low quality, blurry, distorted"
MODEL = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
def _get_diffusion_feature_cases(model: str):
"""Return L4 diffusion feature cases for HunyuanVideo-1.5.
Designed for 2x H100 environment per issue #1832.
"""
return [
# (1 GPU) TeaCache + Layerwise CPU offloading
pytest.param(
OmniServerParams(
model=model,
server_args=[
"--cache-backend",
"tea_cache",
"--enable-layerwise-offload",
],
),
id="single_card_teacache_layerwise",
marks=SINGLE_CARD_MARKS,
),
# (2 GPUs) CacheDiT + TP=2 + VAE patch parallel=2
pytest.param(
OmniServerParams(
model=model,
server_args=[
"--cache-backend",
"cache_dit",
"--tensor-parallel-size",
"2",
"--vae-patch-parallel-size",
"2",
"--vae-use-tiling",
],
),
id="parallel_cachedit_tp2_vae2",
marks=PARALLEL_MARKS,
),
]
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parametrize(
"omni_server",
_get_diffusion_feature_cases(MODEL),
indirect=True,
)
def test_hunyuan_video_15_t2v(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
):
"""L4 diffusion feature coverage for HunyuanVideo-1.5-T2V on H100."""
messages = dummy_messages_from_mix_data(content_text=PROMPT)
request_config = {
"model": omni_server.model,
"messages": messages,
"extra_body": {
"height": 480,
"width": 640,
"num_frames": 5,
"num_inference_steps": 2,
"guidance_scale": 6.0,
"negative_prompt": NEGATIVE_PROMPT,
"seed": 42,
},
}
openai_client.send_diffusion_request(request_config)
@@ -28,6 +28,7 @@ from tests.utils import hardware_test
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
MODEL = "Tongyi-MAI/Z-Image-Turbo"
DIFFUSION_INIT_TIMEOUT_S = 600
PROMPT = "a photo of a cat sitting on a laptop keyboard"
@@ -37,7 +38,17 @@ SEED = 42
@pytest.fixture(scope="module")
def omni_server():
with OmniServer(MODEL, ["--num-gpus", "1"]) as server:
with OmniServer(
MODEL,
[
"--num-gpus",
"1",
"--stage-init-timeout",
str(DIFFUSION_INIT_TIMEOUT_S),
"--init-timeout",
str(DIFFUSION_INIT_TIMEOUT_S),
],
) as server:
yield server
+2 -6
View File
@@ -44,13 +44,9 @@ def get_chunk_config():
return path
# CI stage config for 2xH100-80G GPUs or AMD GPU MI325
if current_omni_platform.is_rocm():
# ROCm stage config optimized for MI325 GPU
stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
elif current_omni_platform.is_xpu():
if current_omni_platform.is_xpu():
stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")]
else:
else: # MI325 GPU should share the same config as H100
stage_configs = [get_chunk_config()]
# Create parameter combinations for model and stage config
@@ -0,0 +1,192 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
E2E Online tests for Voxtral TTS model with text input and audio output.
These tests verify the /v1/audio/speech endpoint works correctly with
actual model inference, not mocks.
"""
import os
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
from pathlib import Path
import httpx
import pytest
from tests.conftest import OmniServerParams
from tests.utils import hardware_test
MODEL = "mistralai/Voxtral-4B-TTS-2603"
STAGE_CONFIG = str(
Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "voxtral_tts.yaml"
)
EXTRA_ARGS = ["--trust-remote-code", "--enforce-eager", "--disable-log-stats"]
TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG, server_args=EXTRA_ARGS)]
def make_speech_request(
host: str,
port: int,
text: str,
voice: str = "casual_female",
language: str = "English",
timeout: float = 120.0,
) -> httpx.Response:
"""Make a request to the /v1/audio/speech endpoint."""
url = f"http://{host}:{port}/v1/audio/speech"
payload = {
"input": text,
"voice": voice,
"language": language,
}
with httpx.Client(timeout=timeout) as client:
return client.post(url, json=payload)
def verify_wav_audio(content: bytes) -> bool:
"""Verify that content is valid WAV audio data."""
# WAV files start with "RIFF" header
if len(content) < 44: # Minimum WAV header size
return False
return content[:4] == b"RIFF" and content[8:12] == b"WAVE"
# Minimum expected audio size for a short sentence (~1 second of 24kHz 16-bit mono WAV)
MIN_AUDIO_BYTES = 10000
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
class TestVoxtralTTSFixedVoice:
"""E2E tests for Voxtral TTS model."""
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_speech_english_basic(self, omni_server) -> None:
"""Test basic English TTS generation."""
response = make_speech_request(
host=omni_server.host,
port=omni_server.port,
text="Hello, how are you?",
voice="casual_female",
language="English",
)
assert response.status_code == 200, f"Request failed: {response.text}"
assert response.headers.get("content-type") == "audio/wav"
assert verify_wav_audio(response.content), "Response is not valid WAV audio"
assert len(response.content) > MIN_AUDIO_BYTES, (
f"Audio content too small ({len(response.content)} bytes), expected at least {MIN_AUDIO_BYTES} bytes"
)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_speech_different_voices(self, omni_server) -> None:
"""Test TTS with different voice presets."""
voices = ["casual_female", "neutral_male"]
for voice in voices:
response = make_speech_request(
host=omni_server.host,
port=omni_server.port,
text="Testing voice selection.",
voice=voice,
language="English",
)
assert response.status_code == 200, f"Request failed for voice {voice}: {response.text}"
assert verify_wav_audio(response.content), f"Invalid WAV for voice {voice}"
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_speech_invalid_voice_rejected(self, omni_server) -> None:
"""Request with a non-existent voice should return an error response."""
response = make_speech_request(
host=omni_server.host,
port=omni_server.port,
text="This voice does not exist.",
voice="nonexistent_voice_xyz",
language="English",
)
# vLLM returns ErrorResponse as JSON with HTTP 200;
# verify the body is a JSON error, not valid WAV audio.
data = response.json()
assert "error" in data or "message" in data, f"Expected error response for invalid voice, got: {data}"
assert not verify_wav_audio(response.content), "Should not return valid audio for invalid voice"
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_speech_binary_response_not_utf8_error(self, omni_server) -> None:
"""
Regression test: Verify binary audio is returned, not UTF-8 error.
This test ensures the multimodal_output property correctly retrieves
audio from completion outputs, preventing the "TTS model did not
produce audio output" error.
"""
response = make_speech_request(
host=omni_server.host,
port=omni_server.port,
text="This should return binary audio, not a JSON error.",
voice="casual_female",
language="English",
)
# Should NOT be a JSON error response
assert response.status_code == 200, f"Request failed: {response.text}"
# Verify it's binary audio, not JSON
try:
# If this succeeds and starts with {"error", it's a bug
text = response.content.decode("utf-8")
assert not text.startswith('{"error"'), f"Got error response instead of audio: {text}"
except UnicodeDecodeError:
# This is expected - binary audio can't be decoded as UTF-8
pass
assert verify_wav_audio(response.content), "Response is not valid WAV audio"
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
class TestVoxtralTTSAPIEndpoints:
"""Test API endpoint functionality."""
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_list_voices_endpoint(self, omni_server) -> None:
"""Test the /v1/audio/voices endpoint returns available voices."""
url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/voices"
with httpx.Client(timeout=30.0) as client:
response = client.get(url)
assert response.status_code == 200
data = response.json()
assert "voices" in data
assert isinstance(data["voices"], list)
assert len(data["voices"]) > 0
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_models_endpoint(self, omni_server) -> None:
"""Test the /v1/models endpoint returns loaded model."""
url = f"http://{omni_server.host}:{omni_server.port}/v1/models"
with httpx.Client(timeout=30.0) as client:
response = client.get(url)
assert response.status_code == 200
data = response.json()
assert "data" in data
assert len(data["data"]) > 0
+3 -5
View File
@@ -7,7 +7,6 @@ stage_args:
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
model_arch: Qwen2_5OmniForConditionalGeneration
@@ -38,7 +37,6 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
@@ -47,7 +45,7 @@ stage_args:
max_model_len: 16384
max_num_batched_tokens: 16384
max_num_seqs: 1
gpu_memory_utilization: 0.9
gpu_memory_utilization: 0.4
skip_mm_profiling: true
enforce_eager: true
trust_remote_code: true
@@ -68,13 +66,13 @@ stage_args:
runtime:
process: true
devices: "2" # 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
gpu_memory_utilization: 0.9 #increase the gpu memory utilization to enable the test on H800
gpu_memory_utilization: 0.5 #increase the gpu memory utilization to enable the test on H800
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
@@ -3,7 +3,6 @@ stage_args:
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
model_arch: Qwen2_5OmniForConditionalGeneration
+3 -3
View File
@@ -8,9 +8,9 @@ stage_args:
- stage_id: 0
runtime:
devices: "0"
max_batch_size: 5
engine_args:
model_stage: thinker
max_num_seqs: 5
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -42,9 +42,9 @@ stage_args:
- stage_id: 1
runtime:
devices: "1"
max_batch_size: 5
engine_args:
model_stage: talker
max_num_seqs: 5
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -72,9 +72,9 @@ stage_args:
- stage_id: 2
runtime:
devices: "1"
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
@@ -7,14 +7,13 @@ stage_args:
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
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 2400
max_num_batched_tokens: 2400
max_model_len: 16384
max_num_batched_tokens: 16384
max_num_seqs: 1
gpu_memory_utilization: 0.8
skip_mm_profiling: true
@@ -38,14 +37,13 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 2400
max_num_batched_tokens: 2400
max_model_len: 16384
max_num_batched_tokens: 16384
max_num_seqs: 1
gpu_memory_utilization: 0.8
skip_mm_profiling: true
@@ -59,7 +57,7 @@ stage_args:
temperature: 0.9
top_p: 0.8
top_k: 40
max_tokens: 128
max_tokens: 4096
seed: 42
detokenize: True
repetition_penalty: 1.05
@@ -68,9 +66,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
@@ -79,6 +77,8 @@ stage_args:
trust_remote_code: true
enable_prefix_caching: false
engine_output_type: audio
max_num_batched_tokens: 4096
max_model_len: 4096
engine_input_source: [1]
final_output: true
final_output_type: audio
@@ -86,7 +86,7 @@ stage_args:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 128
max_tokens: 4096
seed: 42
detokenize: True
repetition_penalty: 1.1
@@ -8,9 +8,9 @@ stage_args:
- stage_id: 0
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: thinker
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -39,9 +39,9 @@ stage_args:
- stage_id: 1
runtime:
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -70,9 +70,9 @@ stage_args:
- stage_id: 2
runtime:
devices: "1"
max_batch_size: 1
engine_args:
model_stage: code2wav
max_num_seqs: 1
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
@@ -7,8 +7,8 @@ stage_args:
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
@@ -39,8 +39,8 @@ stage_args:
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
@@ -71,8 +71,8 @@ stage_args:
runtime:
process: true
devices: "2"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: code2wav
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
@@ -9,8 +9,8 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "0,1,2,3"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
@@ -44,8 +44,8 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "4"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: talker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
@@ -78,8 +78,8 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "5"
max_batch_size: 1
engine_args:
max_num_seqs: 1
model_stage: code2wav
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
@@ -123,6 +123,7 @@ def make_mock_model(hidden: int = 8):
cfg.video_token_index = VIDEO_TOKEN_ID
cfg.audio_token_index = AUDIO_TOKEN_ID
model.config = cfg
model._has_oov_mm_tokens = False
def fake_lm_embed(ids: torch.Tensor) -> torch.Tensor:
# Use .clone() so the tensor is contiguous (expand() creates a strided
@@ -137,13 +138,12 @@ def make_mock_model(hidden: int = 8):
model._embed_text_input_ids = lambda *a, **kw: SupportsMultiModal._embed_text_input_ids(model, *a, **kw)
def fake_super_embed(ids, mm_embs=None, *, is_multimodal=None, handle_oov_mm_token=False):
def fake_super_embed(ids, mm_embs=None, *, is_multimodal=None):
return SupportsMultiModal.embed_input_ids(
model,
ids,
mm_embs,
is_multimodal=is_multimodal,
handle_oov_mm_token=handle_oov_mm_token,
)
model.embed_input_ids = lambda *a, **kw: Qwen2_5OmniThinkerForConditionalGeneration.embed_input_ids(model, *a, **kw)
@@ -270,3 +270,63 @@ def test_deterministic_across_calls(decoder, wrapper):
out1 = wrapper.decode(codes)
out2 = wrapper.decode(codes)
torch.testing.assert_close(out1, out2, atol=0, rtol=0)
# ──────────────────────────────────────────────────────────────────
# 6. compute_capture_sizes
# ──────────────────────────────────────────────────────────────────
@pytest.mark.parametrize(
"kwargs,expected_in,not_expected",
[
({}, [2, 4, 8, 16, 32, 64, 128, 256, 325], [512]),
(
{"codec_chunk_frames": 33, "codec_left_context_frames": 25},
[2, 4, 8, 16, 32, 33, 58, 64, 128, 256, 325],
[512],
),
(
{"codec_chunk_frames": 25, "codec_left_context_frames": 25},
[2, 4, 8, 16, 25, 32, 50, 64, 128, 256, 325],
[512],
),
],
ids=["default", "streaming_c33", "streaming_c25"],
)
def test_compute_capture_sizes(kwargs, expected_in, not_expected):
"""compute_capture_sizes produces expected sizes capped by max useful size."""
sizes = CUDAGraphDecoderWrapper.compute_capture_sizes(**kwargs)
for val in expected_in:
assert val in sizes, f"{val} not in {sizes}"
for val in not_expected:
assert val not in sizes, f"{val} should not be in {sizes}"
# ──────────────────────────────────────────────────────────────────
# 7. SnakeBeta Triton kernel vs eager equivalence
# ──────────────────────────────────────────────────────────────────
@pytest.mark.parametrize(
"batch,channels,seq_len",
[(2, 64, 1000), (1, 32, 1), (1, 32, 7), (1, 32, 128), (1, 32, 1024), (1, 32, 4096)],
)
def test_snakebeta_triton_vs_eager(batch, channels, seq_len):
"""Fused Triton SnakeBeta kernel must match eager PyTorch output."""
from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
SnakeBeta,
)
if not SnakeBeta._init_triton():
pytest.skip("Triton not available")
torch.manual_seed(42)
snake = SnakeBeta(in_features=channels).to(DEVICE).eval()
x = torch.randn(batch, channels, seq_len, device=DEVICE)
with torch.no_grad():
eager_out = snake._eager_forward(x)
triton_out = snake._triton_forward(x)
torch.testing.assert_close(triton_out, eager_out, atol=1e-5, rtol=1e-5)
Whitespace-only changes.
@@ -0,0 +1,161 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Tests for parse_batched_audio_input and apply_ctx_frames_cutting.
These functions parse the [ctx_frames, context_length, ...tokens] wire format
produced by generator2tokenizer_async_chunk and cut leading context samples
from decoded audio arrays.
"""
import pytest
import torch
from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts import (
apply_ctx_frames_cutting,
parse_batched_audio_input,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
NUM_CODEBOOKS = 37
# ──────────────────────────────────────────────────────────────────
# parse_batched_audio_input tests
# ──────────────────────────────────────────────────────────────────
def _make_input_ids(*requests: tuple[int, int, torch.Tensor]) -> torch.Tensor:
"""Build a flat input_ids tensor from (ctx_frames, context_length, tokens) tuples."""
parts = []
for ctx_frames, context_length, tokens in requests:
parts.append(torch.tensor([ctx_frames, context_length]))
parts.append(tokens.flatten())
return torch.cat(parts)
def test_parse_single_request():
"""Single request is parsed correctly."""
ctx_frames = 3
context_length = 5
total_frames = ctx_frames + context_length
tokens = torch.arange(total_frames * NUM_CODEBOOKS).reshape(total_frames, NUM_CODEBOOKS)
input_ids = _make_input_ids((ctx_frames, context_length, tokens))
all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS)
assert len(all_audio_tokens) == 1
assert len(all_ctx_frames) == 1
assert all_ctx_frames[0] == ctx_frames
assert all_audio_tokens[0].shape == (total_frames, NUM_CODEBOOKS)
torch.testing.assert_close(all_audio_tokens[0], tokens)
def test_parse_multiple_requests():
"""Multiple requests in a batch are parsed correctly."""
req1 = (2, 4, torch.ones((6, NUM_CODEBOOKS), dtype=torch.long))
req2 = (0, 3, torch.full((3, NUM_CODEBOOKS), 7, dtype=torch.long))
input_ids = _make_input_ids(req1, req2)
all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS)
assert len(all_audio_tokens) == 2
assert all_ctx_frames == [2, 0]
assert all_audio_tokens[0].shape == (6, NUM_CODEBOOKS)
assert all_audio_tokens[1].shape == (3, NUM_CODEBOOKS)
assert torch.all(all_audio_tokens[0] == 1)
assert torch.all(all_audio_tokens[1] == 7)
def test_parse_zero_ctx_frames():
"""ctx_frames=0 means no context, just chunk data."""
ctx_frames = 0
context_length = 10
tokens = torch.randint(0, 100, (context_length, NUM_CODEBOOKS))
input_ids = _make_input_ids((ctx_frames, context_length, tokens))
all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS)
assert all_ctx_frames[0] == 0
assert all_audio_tokens[0].shape == (context_length, NUM_CODEBOOKS)
def test_parse_fails_on_misaligned_input():
"""Input not divisible by num_codebooks should raise AssertionError."""
# 2 header elements + 5 data elements (not divisible by 37)
input_ids = torch.tensor([0, 1] + [0] * 5)
with pytest.raises(AssertionError, match="divisible by"):
parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS)
def test_parse_custom_num_codebooks():
"""Works with a non-default num_codebooks value."""
cb = 4
ctx_frames = 1
context_length = 2
total_frames = ctx_frames + context_length
tokens = torch.arange(total_frames * cb).reshape(total_frames, cb)
input_ids = _make_input_ids((ctx_frames, context_length, tokens))
all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=cb)
assert all_audio_tokens[0].shape == (total_frames, cb)
torch.testing.assert_close(all_audio_tokens[0], tokens)
# ──────────────────────────────────────────────────────────────────
# apply_ctx_frames_cutting tests
# ──────────────────────────────────────────────────────────────────
def test_cut_removes_leading_samples():
"""Context frames are removed from the front of the audio array."""
downsample_factor = 240
ctx_frames = 5
total_samples = 2400 # 10 frames * 240
audio = torch.arange(total_samples, dtype=torch.float32)
result = apply_ctx_frames_cutting([audio], [ctx_frames], downsample_factor)
expected_cut = ctx_frames * downsample_factor # 1200
assert len(result) == 1
assert result[0].shape[0] == total_samples - expected_cut
torch.testing.assert_close(result[0], audio[expected_cut:])
def test_cut_zero_ctx_frames_unchanged():
"""ctx_frames=0 leaves the audio unchanged."""
downsample_factor = 240
audio = torch.randn(2400)
result = apply_ctx_frames_cutting([audio], [0], downsample_factor)
assert result[0].shape == audio.shape
torch.testing.assert_close(result[0], audio)
def test_cut_multiple_requests():
"""Each request in the batch gets its own ctx_frames cut."""
downsample_factor = 100
audio1 = torch.arange(1000, dtype=torch.float32) # 10 frames
audio2 = torch.arange(500, dtype=torch.float32) # 5 frames
result = apply_ctx_frames_cutting([audio1, audio2], [3, 0], downsample_factor)
assert len(result) == 2
assert result[0].shape[0] == 1000 - 3 * 100 # 700
assert result[1].shape[0] == 500 # unchanged
torch.testing.assert_close(result[0], audio1[300:])
torch.testing.assert_close(result[1], audio2)
def test_cut_all_frames_returns_empty():
"""Cutting all frames returns an empty tensor."""
downsample_factor = 100
ctx_frames = 5
audio = torch.randn(500) # exactly 5 frames
result = apply_ctx_frames_cutting([audio], [ctx_frames], downsample_factor)
assert result[0].numel() == 0
@@ -0,0 +1,330 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Tests for CUDA Graph acoustic transformer wrapper numerical equivalence.
Verifies that CUDA Graph-accelerated decoding produces results equivalent
to eager mode, with special attention to padding cases where zero-padding
may introduce small numerical differences due to the flow matching ODE.
Architecture note: the real FlowMatchingAudioTransformer uses a multi-step
Euler ODE with CFG. The synthetic model here uses simple linear layers
to exercise the wrapper mechanism while keeping the test lightweight.
"""
import importlib.util
import os
import pytest
import torch
import torch.nn as nn
pytestmark = [
pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
pytest.mark.core_model,
pytest.mark.cuda,
pytest.mark.L4,
]
DEVICE = torch.device("cuda:0")
HIDDEN_DIM = 64
N_ACOUSTIC_CODEBOOK = 7
SEMANTIC_CODEBOOK_SIZE = 128
ACOUSTIC_EMBEDDINGS_LEVELS = 1024
# Load CUDAGraphAcousticTransformerWrapper: try package import first, fall back to direct file load
try:
from vllm_omni.model_executor.models.voxtral_tts.cuda_graph_acoustic_transformer_wrapper import (
CUDAGraphAcousticTransformerWrapper,
)
from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation import (
AudioSpecialTokens,
)
except Exception:
_WRAPPER_PATH = os.path.join(
os.path.dirname(__file__),
os.pardir,
os.pardir,
os.pardir,
os.pardir,
"vllm_omni",
"model_executor",
"models",
"voxtral_tts",
"cuda_graph_acoustic_transformer_wrapper.py",
)
_spec = importlib.util.spec_from_file_location(
"cuda_graph_acoustic_transformer_wrapper", os.path.abspath(_WRAPPER_PATH)
)
_mod = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_mod)
CUDAGraphAcousticTransformerWrapper = _mod.CUDAGraphAcousticTransformerWrapper
_AUDIO_GEN_PATH = os.path.join(
os.path.dirname(__file__),
os.pardir,
os.pardir,
os.pardir,
os.pardir,
"vllm_omni",
"model_executor",
"models",
"voxtral_tts",
"voxtral_tts_audio_generation.py",
)
_spec2 = importlib.util.spec_from_file_location("voxtral_tts_audio_generation", os.path.abspath(_AUDIO_GEN_PATH))
_mod2 = importlib.util.module_from_spec(_spec2)
_spec2.loader.exec_module(_mod2)
AudioSpecialTokens = _mod2.AudioSpecialTokens
class SyntheticModelArgs:
"""Mimics MultimodalAudioModelArgs interface."""
def __init__(self):
self.semantic_codebook_size = SEMANTIC_CODEBOOK_SIZE
self.n_acoustic_codebook = N_ACOUSTIC_CODEBOOK
class SyntheticAcousticTransformer(nn.Module):
"""A small acoustic transformer mimicking FlowMatchingAudioTransformer's interface.
Uses simple linear layers so the test stays lightweight while exercising
the full wrapper mechanism (semantic logit, flow matching ODE, quantization).
"""
def __init__(self):
super().__init__()
self.model_args = SyntheticModelArgs()
self.acoustic_embeddings_levels = ACOUSTIC_EMBEDDINGS_LEVELS
# semantic_codebook_output: hidden_dim -> padded_codebook_size
padded_semantic_size = len(AudioSpecialTokens) + SEMANTIC_CODEBOOK_SIZE
# Round up to multiple of 128 like the real model
padded_semantic_size = 128 * ((padded_semantic_size + 127) // 128)
self.semantic_codebook_output = nn.Linear(HIDDEN_DIM, padded_semantic_size, bias=False)
# time_embedding: (B, 1) -> (B, dim)
self.time_embedding = nn.Linear(1, HIDDEN_DIM, bias=False)
# _predict_velocity: takes x_t (B, C), llm_output (B, D), t_emb (B, D) -> (B, C)
self._velocity_proj = nn.Linear(N_ACOUSTIC_CODEBOOK + HIDDEN_DIM + HIDDEN_DIM, N_ACOUSTIC_CODEBOOK, bias=False)
def _predict_velocity(self, x_t: torch.Tensor, llm_output: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
combined = torch.cat([x_t, llm_output, t_emb], dim=-1)
return self._velocity_proj(combined)
class SyntheticModel(nn.Module):
"""Mimics VoxtralTTSAudioGenerationForConditionalGeneration interface."""
def __init__(self):
super().__init__()
self.acoustic_transformer = SyntheticAcousticTransformer()
end_audio_id = AudioSpecialTokens.id(AudioSpecialTokens.end_audio)
empty_audio_id = AudioSpecialTokens.id(AudioSpecialTokens.empty_audio)
self.end_audio_id = end_audio_id
self.empty_audio_id = empty_audio_id
def compute_mm_logits(
self,
hidden_states: torch.Tensor,
mm_sampling_tensors=None,
):
"""Eager fallback path: replicate what the wrapper does."""
at = self.acoustic_transformer
B = hidden_states.shape[0]
empty_audio_id = self.empty_audio_id
end_audio_id = self.end_audio_id
semantic_mask_start = len(AudioSpecialTokens) + at.model_args.semantic_codebook_size
# Semantic logits
semantic_logit = at.semantic_codebook_output(hidden_states).float()
semantic_logit[:, empty_audio_id] = -float("inf")
semantic_logit[:, semantic_mask_start:] = -float("inf")
semantic_code = semantic_logit.argmax(dim=-1, keepdim=True)
# Flow matching Euler ODE
should_decode = semantic_code.squeeze(1) != end_audio_id
x = torch.randn(B, at.model_args.n_acoustic_codebook, device=hidden_states.device, dtype=hidden_states.dtype)
hidden_zero = torch.zeros_like(hidden_states)
timesteps = torch.linspace(0, 1, 16, device=hidden_states.device, dtype=hidden_states.dtype)
for i in range(len(timesteps) - 1):
t = timesteps[i]
dt = timesteps[i + 1] - timesteps[i]
t_emb = at.time_embedding(t.view(-1, 1).repeat(B, 1)).to(hidden_states.dtype)
x_batched = torch.cat([x, x], dim=0)
llm_batched = torch.cat([hidden_states, hidden_zero], dim=0)
t_emb_batched = t_emb.repeat(2, 1)
v_all = at._predict_velocity(x_t=x_batched, llm_output=llm_batched, t_emb=t_emb_batched)
v_t, uncond_v_t = v_all[:B], v_all[B:]
cfg_alpha = 1.2
v_t = cfg_alpha * v_t + (1 - cfg_alpha) * uncond_v_t
x = x + v_t * dt
sampled = torch.clamp(x, -1, 1)
scaled_x = ((sampled + 1) / 2) * (at.acoustic_embeddings_levels - 1)
output_codes = scaled_x.round().long()
output_codes[~should_decode] = empty_audio_id
acoustic_codes = output_codes + len(AudioSpecialTokens)
audio_codes = torch.cat([semantic_code, acoustic_codes], dim=1)
fake_eos = torch.where(
audio_codes[:, 0] == end_audio_id,
torch.tensor(1.0, dtype=hidden_states.dtype, device=hidden_states.device),
torch.tensor(0.0, dtype=hidden_states.dtype, device=hidden_states.device),
)
audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0))
mm_tokens = {"audio": audio_list}
return fake_eos, mm_tokens
@pytest.fixture(scope="module")
def model():
"""Create a synthetic model on CUDA with fixed weights in bfloat16."""
torch.manual_seed(42)
return SyntheticModel().to(device=DEVICE, dtype=torch.bfloat16).eval()
@pytest.fixture(scope="module")
def wrapper(model):
"""Create a warmed-up CUDAGraphAcousticTransformerWrapper."""
w = CUDAGraphAcousticTransformerWrapper(
model=model,
capture_sizes=[1, 2, 4, 8, 16, 32],
)
w._warmup_and_capture(
device=DEVICE,
dtype=torch.bfloat16,
hidden_dim=HIDDEN_DIM,
)
return w
def _random_hidden(batch_size, device=DEVICE, dtype=torch.bfloat16):
return torch.randn(batch_size, HIDDEN_DIM, device=device, dtype=dtype)
def _unpack_audio_codes(result):
"""Unpack (fake_eos, {"audio": [list of tensors]}) into (fake_eos, audio_codes)."""
fake_eos, mm_tokens = result
audio_list = mm_tokens["audio"]
# Each element is (1, 1, 1+n_acoustic_codebook), concat along batch dim
audio_codes = torch.cat(audio_list, dim=0).squeeze(1) # (B, 1+n_acoustic)
return fake_eos, audio_codes
# ──────────────────────────────────────────────────────────────────
# 1. Exact-size inputs (no padding) → correct format and bounds
# ──────────────────────────────────────────────────────────────────
@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16, 32])
def test_exact_size_output_format(model, wrapper, batch_size):
"""Graph path returns correctly shaped and bounded outputs."""
hidden = _random_hidden(batch_size)
with torch.no_grad():
graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
assert graph_eos.shape == (batch_size,)
assert graph_codes.shape == (batch_size, 1 + N_ACOUSTIC_CODEBOOK)
# fake_eos should be 0.0 or 1.0
assert torch.all((graph_eos == 0.0) | (graph_eos == 1.0))
# Audio codes should be non-negative
assert torch.all(graph_codes >= 0)
@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16, 32])
def test_exact_size_deterministic(model, wrapper, batch_size):
"""Same input + same RNG state produces identical CUDA graph output."""
hidden = _random_hidden(batch_size)
with torch.no_grad():
torch.manual_seed(42)
eos1, codes1 = _unpack_audio_codes(wrapper(hidden))
torch.manual_seed(42)
eos2, codes2 = _unpack_audio_codes(wrapper(hidden))
torch.testing.assert_close(eos1, eos2, atol=0, rtol=0)
torch.testing.assert_close(codes1, codes2, atol=0, rtol=0)
# ──────────────────────────────────────────────────────────────────
# 2. Padded inputs (zero-padding to nearest capture size)
# ──────────────────────────────────────────────────────────────────
@pytest.mark.parametrize("batch_size", [3, 5, 7, 10, 24])
def test_padded_output_shape(model, wrapper, batch_size):
"""Padded decode must return output trimmed to actual batch size."""
hidden = _random_hidden(batch_size)
with torch.no_grad():
graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
assert graph_eos.shape == (batch_size,)
assert graph_codes.shape == (batch_size, 1 + N_ACOUSTIC_CODEBOOK)
@pytest.mark.parametrize("batch_size", [3, 5, 7, 10, 24])
def test_padded_output_bounded(model, wrapper, batch_size):
"""Padded output audio codes should be non-negative integers."""
hidden = _random_hidden(batch_size)
with torch.no_grad():
graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
# fake_eos should be 0.0 or 1.0
assert torch.all((graph_eos == 0.0) | (graph_eos == 1.0))
# Audio codes should be non-negative
assert torch.all(graph_codes >= 0)
# ──────────────────────────────────────────────────────────────────
# 3. Batch size exceeds all capture sizes so fallback to eager
# ──────────────────────────────────────────────────────────────────
@pytest.mark.parametrize("batch_size", [33, 48, 64])
def test_fallback_eager_exact_match(model, wrapper, batch_size):
"""Cudagraph fallback to eager. Two eager runs should produce identical results."""
hidden = _random_hidden(batch_size)
with torch.no_grad():
torch.manual_seed(100)
eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden))
torch.manual_seed(100)
graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
torch.testing.assert_close(graph_eos, eager_eos, atol=0, rtol=0)
torch.testing.assert_close(graph_codes, eager_codes, atol=0, rtol=0)
# ──────────────────────────────────────────────────────────────────
# 4. Disabled wrapper so fallback to eager
# ──────────────────────────────────────────────────────────────────
def test_disabled_wrapper_matches_eager(model, wrapper):
"""Cudagraph fallback to eager. Two eager runs should produce identical results."""
hidden = _random_hidden(4)
wrapper.enabled = False
with torch.no_grad():
torch.manual_seed(200)
eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden))
torch.manual_seed(200)
graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
wrapper.enabled = True
torch.testing.assert_close(graph_eos, eager_eos, atol=0, rtol=0)
torch.testing.assert_close(graph_codes, eager_codes, atol=0, rtol=0)
# ──────────────────────────────────────────────────────────────────
# 5. Determinism - same input + same RNG state on cudagraph runs
# ──────────────────────────────────────────────────────────────────
def test_deterministic_across_calls(model, wrapper):
"""Same input + same RNG state. Two cudagraph runs should produce identical results."""
hidden = _random_hidden(4)
with torch.no_grad():
torch.manual_seed(300)
eos1, codes1 = _unpack_audio_codes(wrapper(hidden))
torch.manual_seed(300)
eos2, codes2 = _unpack_audio_codes(wrapper(hidden))
torch.testing.assert_close(eos1, eos2, atol=0, rtol=0)
torch.testing.assert_close(codes1, codes2, atol=0, rtol=0)
@@ -0,0 +1,281 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections import defaultdict
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.model_executor.stage_input_processors.voxtral_tts import (
generator2tokenizer,
generator2tokenizer_async_chunk,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
# ---- Helpers for generator2tokenizer (non-streaming) ----
def _make_stage(engine_outputs):
return SimpleNamespace(engine_outputs=engine_outputs)
def _make_output(audio_tensors):
"""Build a single generator output with multimodal_output["audio"] = audio_tensors."""
return SimpleNamespace(
outputs=[SimpleNamespace(multimodal_output={"audio": audio_tensors})],
)
# ---- Tests for generator2tokenizer ----
class TestGenerator2Tokenizer:
def test_single_output_single_chunk(self):
"""Single output with one audio tensor produces correct prompt_token_ids."""
audio = torch.tensor([10, 20, 30, 40])
stage = _make_stage([_make_output([audio])])
result = generator2tokenizer([stage], engine_input_source=[0])
assert len(result) == 1
assert result[0]["prompt_token_ids"] == [10, 20, 30, 40]
assert result[0]["multi_modal_data"] is None
def test_2d_audio_tensors_flattened(self):
"""2D audio tensors (e.g., multi-codebook frames) are flattened."""
audio = torch.tensor([[1, 2, 3], [4, 5, 6]]) # (2, 3)
stage = _make_stage([_make_output([audio])])
result = generator2tokenizer([stage], engine_input_source=[0])
assert result[0]["prompt_token_ids"] == [1, 2, 3, 4, 5, 6]
def test_empty_engine_input_source_raises(self):
"""Empty engine_input_source raises ValueError."""
stage = _make_stage([_make_output([torch.tensor([1])])])
with pytest.raises(ValueError, match="cannot be empty"):
generator2tokenizer([stage], engine_input_source=[])
def test_invalid_stage_id_raises(self):
"""Stage id beyond stage_list length raises IndexError."""
stage = _make_stage([_make_output([torch.tensor([1])])])
with pytest.raises(IndexError, match="Invalid stage_id"):
generator2tokenizer([stage], engine_input_source=[5])
def test_no_outputs_yet_raises(self):
"""Stage with engine_outputs=None raises RuntimeError."""
stage = _make_stage(engine_outputs=None)
with pytest.raises(RuntimeError, match="no outputs yet"):
generator2tokenizer([stage], engine_input_source=[0])
# ---- Helpers for generator2tokenizer_async_chunk (streaming) ----
def _req(external_req_id: str, *, finished: bool):
return SimpleNamespace(
external_req_id=external_req_id,
is_finished=lambda: finished,
)
def _make_transfer_manager(codec_chunk_frames=25, codec_left_context_frames=25, codec_chunk_frames_at_begin=5):
return SimpleNamespace(
code_prompt_token_ids=defaultdict(list),
connector=SimpleNamespace(
config={
"extra": {
"codec_chunk_frames": codec_chunk_frames,
"codec_left_context_frames": codec_left_context_frames,
"codec_chunk_frames_at_begin": codec_chunk_frames_at_begin,
}
}
),
)
def test_empty_chunk_when_not_finished():
"""Returns None when no audio data and not finished."""
transfer_manager = _make_transfer_manager()
request = _req("rid-empty", finished=False)
payload = generator2tokenizer_async_chunk(
transfer_manager=transfer_manager,
pooling_output={"audio": torch.zeros((0,))},
request=request,
)
assert payload is None
def test_flush_tail_when_finished():
"""Emits remaining frames on finish with finished=True."""
transfer_manager = _make_transfer_manager()
request_id = "rid-tail"
# Pre-populate with 24 frames of 8-element codes
transfer_manager.code_prompt_token_ids[request_id] = [[1, 2, 3, 4, 5, 6, 7, 8] for _ in range(24)]
request = _req(request_id, finished=True)
payload = generator2tokenizer_async_chunk(
transfer_manager=transfer_manager,
pooling_output=None, # e.g. EOS step with no audio
request=request,
)
assert payload is not None
assert payload["finished"].item() is True
codes = payload["code_predictor_codes"]
# Format: [ctx_frames, context_length, ...flat_codes]
assert len(codes) >= 2 # At least ctx_frames + context_length header
ctx_frames = codes[0]
context_length = codes[1]
flat_codes = codes[2:]
# Total frames in window = ctx_frames + context_length
total_window_frames = ctx_frames + context_length
assert len(flat_codes) == total_window_frames * 8
def test_eof_marker_when_finished_with_no_frames():
"""Emits EOF marker with empty codes when finished with no accumulated frames."""
transfer_manager = _make_transfer_manager()
request = _req("rid-eof", finished=True)
payload = generator2tokenizer_async_chunk(
transfer_manager=transfer_manager,
pooling_output=None,
request=request,
)
assert payload == {
"code_predictor_codes": [],
"finished": torch.tensor(True, dtype=torch.bool),
}
def test_normal_chunk_emission():
"""Emits chunks at codec_chunk_frames boundary."""
transfer_manager = _make_transfer_manager(
codec_chunk_frames=25,
codec_left_context_frames=25,
codec_chunk_frames_at_begin=5,
)
request_id = "rid-chunk"
request = _req(request_id, finished=False)
# Feed 25 frames one by one
for i in range(25):
pooling_output = {"audio": torch.tensor([float(i)] * 4)}
payload = generator2tokenizer_async_chunk(
transfer_manager=transfer_manager,
pooling_output=pooling_output,
request=request,
)
# A chunk should be emitted
assert payload is not None
codes = payload["code_predictor_codes"]
ctx_frames = codes[0]
context_length = codes[1]
assert ctx_frames == 20 # 25 - 5(chunk_size_at_begin)
assert context_length == 5 # 25 - 20(chunk_size_at_begin * 4)
def test_small_initial_chunks():
"""Uses codec_chunk_frames_at_begin for initial frames when length <= chunk_size."""
transfer_manager = _make_transfer_manager(
codec_chunk_frames=25,
codec_left_context_frames=25,
codec_chunk_frames_at_begin=5,
)
request_id = "rid-begin"
request = _req(request_id, finished=False)
# Feed 5 frames (should trigger emission because codec_chunk_frames_at_begin=5)
for i in range(5):
pooling_output = {"audio": torch.tensor([float(i + 1)] * 3)}
payload = generator2tokenizer_async_chunk(
transfer_manager=transfer_manager,
pooling_output=pooling_output,
request=request,
)
assert payload is not None
codes = payload["code_predictor_codes"]
ctx_frames = codes[0]
context_length = codes[1]
assert ctx_frames == 0
assert context_length == 5 # chunk_size_at_begin
def test_no_emission_between_boundaries():
"""No chunk emitted when not at chunk boundary and not finished."""
transfer_manager = _make_transfer_manager(
codec_chunk_frames=25,
codec_left_context_frames=25,
codec_chunk_frames_at_begin=5,
)
request_id = "rid-mid"
request = _req(request_id, finished=False)
# Feed 3 frames (3 % 5 != 0, should not emit)
for i in range(3):
pooling_output = {"audio": torch.tensor([float(i)] * 4)}
payload = generator2tokenizer_async_chunk(
transfer_manager=transfer_manager,
pooling_output=pooling_output,
request=request,
)
assert payload is None
def test_context_handling_format():
"""Verifies [ctx_frames, context_length, ...codes] format."""
transfer_manager = _make_transfer_manager(
codec_chunk_frames=10,
codec_left_context_frames=5,
codec_chunk_frames_at_begin=5,
)
request_id = "rid-ctx"
request = _req(request_id, finished=False)
# Feed 5 frames to trigger a chunk (chunk_size_at_begin=5)
for i in range(5):
pooling_output = {"audio": torch.tensor([float(i + 10)] * 2)} # codebook_dim=2
payload = generator2tokenizer_async_chunk(
transfer_manager=transfer_manager,
pooling_output=pooling_output,
request=request,
)
assert payload is not None
codes = payload["code_predictor_codes"]
# First two elements are ctx_frames and context_length
ctx_frames = codes[0]
context_length = codes[1]
flat_codes = codes[2:]
# The window has ctx_frames + context_length frames
assert isinstance(ctx_frames, int)
assert isinstance(context_length, int)
assert ctx_frames >= 0
assert context_length > 0
# flat_codes = total_window_frames * codebook_dim
total_window_frames = ctx_frames + context_length
assert len(flat_codes) == total_window_frames * 2 # codebook_dim=2
def test_none_pooling_output_not_finished_returns_none():
"""None pooling_output when not finished returns None."""
transfer_manager = _make_transfer_manager()
request = _req("rid-none", finished=False)
payload = generator2tokenizer_async_chunk(
transfer_manager=transfer_manager,
pooling_output=None,
request=request,
)
assert payload is None
+3 -3
View File
@@ -10,9 +10,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "0"
max_batch_size: 64
engine_args:
model_stage: thinker
max_num_seqs: 64
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -41,9 +41,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "1"
max_batch_size: 64
engine_args:
model_stage: talker
max_num_seqs: 64
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -72,9 +72,9 @@ stage_args:
stage_type: llm # Use llm stage type for AR stages
runtime:
devices: "1"
max_batch_size: 64
engine_args:
model_stage: code2wav
max_num_seqs: 64
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+2 -2
View File
@@ -10,8 +10,8 @@ stage_args:
is_comprehension: true
runtime:
devices: "0"
max_batch_size: 4
engine_args:
max_num_seqs: 4
model_stage: qwen3_tts
model_arch: Qwen3TTSTalkerForConditionalGeneration
worker_type: ar
@@ -41,8 +41,8 @@ stage_args:
stage_type: llm
runtime:
devices: "0"
max_batch_size: 4
engine_args:
max_num_seqs: 4
model_stage: code2wav
model_arch: Qwen3TTSCode2Wav
worker_type: generation
+16 -8
View File
@@ -12,11 +12,13 @@
"endpoint": "/v1/chat/completions",
"num_prompts": [
10,
40
40,
100
],
"max_concurrency": [
1,
4
4,
10
],
"random_input_len": 100,
"random_output_len": 100,
@@ -34,11 +36,13 @@
"endpoint": "/v1/chat/completions",
"num_prompts": [
10,
40
40,
100
],
"max_concurrency": [
1,
4
4,
10
],
"random_input_len": 100,
"random_output_len": 100,
@@ -96,11 +100,13 @@
"endpoint": "/v1/chat/completions",
"num_prompts": [
10,
40
40,
100
],
"max_concurrency": [
1,
4
4,
10
],
"random_input_len": 100,
"random_output_len": 100,
@@ -118,11 +124,13 @@
"endpoint": "/v1/chat/completions",
"num_prompts": [
10,
40
40,
100
],
"max_concurrency": [
1,
4
4,
10
],
"random_input_len": 100,
"random_output_len": 100,
+56 -2
View File
@@ -83,6 +83,8 @@ class TestStageConfig:
def test_to_omegaconf_with_runtime_overrides(self):
"""Test that runtime overrides are applied to OmegaConf output."""
import warnings
config = StageConfig(
stage_id=0,
model_stage="thinker",
@@ -93,12 +95,41 @@ class TestStageConfig:
"max_batch_size": 64,
},
)
omega_config = config.to_omegaconf()
with warnings.catch_warnings():
warnings.simplefilter("ignore", FutureWarning)
omega_config = config.to_omegaconf()
assert omega_config.engine_args.gpu_memory_utilization == 0.9
assert omega_config.engine_args.tensor_parallel_size == 2
assert omega_config.runtime.devices == "0,1"
assert omega_config.runtime.max_batch_size == 64
# max_batch_size is migrated to engine_args.max_num_seqs
assert omega_config.engine_args.max_num_seqs == 64
def test_to_omegaconf_max_batch_size_deprecation(self):
"""Test that runtime.max_batch_size emits a FutureWarning."""
import warnings
config = StageConfig(
stage_id=0,
model_stage="thinker",
runtime_overrides={"max_batch_size": 8},
)
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
config.to_omegaconf()
deprecation_warnings = [x for x in w if issubclass(x.category, FutureWarning)]
assert len(deprecation_warnings) == 1
assert "max_batch_size" in str(deprecation_warnings[0].message)
def test_to_omegaconf_max_num_seqs_in_engine_args(self):
"""Test that max_num_seqs in yaml_engine_args takes precedence."""
config = StageConfig(
stage_id=0,
model_stage="thinker",
yaml_engine_args={"max_num_seqs": 32},
)
omega_config = config.to_omegaconf()
assert omega_config.engine_args.max_num_seqs == 32
class TestModelPipeline:
@@ -425,6 +456,29 @@ stages:
errors = pipeline.validate_pipeline()
assert errors == [], f"Unexpected validation errors: {errors}"
def test_parse_yaml_migrates_runtime_max_batch_size(self, tmp_path):
"""Test that runtime.max_batch_size is migrated to engine_args.max_num_seqs."""
yaml_content = """\
model_type: legacy_batch
stages:
- stage_id: 0
model_stage: entry
stage_type: llm
input_sources: []
runtime:
devices: "0"
max_batch_size: 16
engine_args:
model_arch: SomeModel
"""
yaml_file = tmp_path / "legacy_batch.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_batch")
s0 = pipeline.stages[0]
assert "max_batch_size" not in s0.yaml_runtime
assert s0.yaml_engine_args.get("max_num_seqs") == 16
def test_parse_diffusion_stage_type(self, tmp_path):
"""Test parsing a diffusion stage type from YAML."""
yaml_content = """\
+185
View File
@@ -0,0 +1,185 @@
"""Tests for CUDAGraphWrapper.__getattr__ performance optimization.
This module tests that the patched CUDAGraphWrapper avoids expensive __repr__
calls when hasattr() is used for non-existent attributes. The original vLLM
implementation includes {self.runnable} in the AttributeError message, which
triggers model tree traversal and can take ~6ms on large models.
"""
import time
import pytest
import torch
import torch.nn as nn
from vllm_omni.worker.gpu_model_runner import CUDAGraphWrapper
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class SlowReprModel(nn.Module):
"""A mock model with artificially slow __repr__ to detect unwanted calls."""
def __init__(self, repr_delay_ms: float = 10.0):
super().__init__()
self.linear = nn.Linear(16, 16)
self.repr_delay_ms = repr_delay_ms
self.repr_call_count = 0
def forward(self, x):
return self.linear(x)
def __repr__(self):
self.repr_call_count += 1
# Simulate expensive repr by sleeping
time.sleep(self.repr_delay_ms / 1000.0)
return f"SlowReprModel(delay={self.repr_delay_ms}ms)"
class MockCUDAGraphWrapper:
"""A minimal mock that mimics CUDAGraphWrapper structure for CPU testing."""
def __init__(self, runnable):
# Store in __dict__ directly to avoid triggering __getattr__
object.__setattr__(self, "runnable", runnable)
def __getattr__(self, key: str):
# This is the optimized implementation we're testing
runnable = object.__getattribute__(self, "runnable")
if hasattr(runnable, key):
return getattr(runnable, key)
# Key optimization: DO NOT include {self.runnable} in error message
# as it triggers expensive __repr__ on large models
raise AttributeError(f"Attribute {key} not exists in the runnable of cudagraph wrapper")
def test_hasattr_nonexistent_does_not_trigger_repr():
"""Verify that hasattr for non-existent attributes doesn't call __repr__."""
model = SlowReprModel(repr_delay_ms=100.0) # Very slow repr
wrapper = MockCUDAGraphWrapper(model)
# Reset counter
model.repr_call_count = 0
# Call hasattr for non-existent attribute multiple times
for _ in range(10):
result = hasattr(wrapper, "nonexistent_attribute_xyz")
assert result is False
# __repr__ should never have been called
assert model.repr_call_count == 0, (
f"__repr__ was called {model.repr_call_count} times when checking "
"for non-existent attributes. This indicates the AttributeError "
"message contains {self.runnable} which triggers expensive repr."
)
def test_hasattr_nonexistent_is_fast():
"""Verify that hasattr for non-existent attributes is fast (<1ms per call)."""
model = SlowReprModel(repr_delay_ms=100.0)
wrapper = MockCUDAGraphWrapper(model)
num_iterations = 100
start = time.perf_counter()
for _ in range(num_iterations):
hasattr(wrapper, "nonexistent_attribute_xyz")
elapsed_ms = (time.perf_counter() - start) * 1000
avg_ms = elapsed_ms / num_iterations
# If __repr__ were being called, each would take ~100ms
# We expect <1ms per call with the fix
assert avg_ms < 1.0, (
f"hasattr for non-existent attribute took {avg_ms:.2f}ms on average. "
"Expected <1ms. This suggests __repr__ is being triggered."
)
def test_hasattr_existing_attribute_works():
"""Verify that hasattr for existing attributes returns True and works correctly."""
model = SlowReprModel()
wrapper = MockCUDAGraphWrapper(model)
# 'forward' exists on nn.Module
assert hasattr(wrapper, "forward") is True
# 'linear' exists on our model
assert hasattr(wrapper, "linear") is True
# Can actually access the attribute
linear = wrapper.linear
assert isinstance(linear, nn.Linear)
def test_getattr_existing_attribute_returns_value():
"""Verify that getattr for existing attributes returns the correct value."""
model = SlowReprModel()
wrapper = MockCUDAGraphWrapper(model)
# Access forward method
forward_method = wrapper.forward
assert callable(forward_method)
# Access linear layer
linear = wrapper.linear
assert isinstance(linear, nn.Linear)
assert linear.in_features == 16
assert linear.out_features == 16
def test_getattr_nonexistent_raises_attribute_error():
"""Verify that getattr for non-existent attributes raises AttributeError."""
model = SlowReprModel()
wrapper = MockCUDAGraphWrapper(model)
with pytest.raises(AttributeError) as exc_info:
_ = wrapper.nonexistent_attribute
# Verify error message format (should NOT contain model repr)
error_msg = str(exc_info.value)
assert "nonexistent_attribute" in error_msg
assert "cudagraph wrapper" in error_msg
# Should NOT contain the slow repr output
assert "SlowReprModel(delay=" not in error_msg
def test_attribute_error_message_does_not_contain_runnable_repr():
"""Explicitly verify the error message doesn't trigger runnable repr."""
model = SlowReprModel(repr_delay_ms=100.0)
wrapper = MockCUDAGraphWrapper(model)
model.repr_call_count = 0
try:
_ = wrapper.nonexistent_attr
except AttributeError:
pass
# __repr__ should not have been called during error construction
assert model.repr_call_count == 0, (
"AttributeError message construction triggered __repr__. The error message should not include {self.runnable}."
)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_real_cudagraph_wrapper_hasattr_performance():
"""Test the actual CUDAGraphWrapper from vllm_omni (requires CUDA)."""
from vllm.config import CUDAGraphMode
model = SlowReprModel(repr_delay_ms=50.0).cuda()
model.repr_call_count = 0
# Create actual CUDAGraphWrapper
try:
wrapper = CUDAGraphWrapper(model, runtime_mode=CUDAGraphMode.NONE)
except Exception:
pytest.skip("Could not create CUDAGraphWrapper")
# Test hasattr performance
num_iterations = 50
start = time.perf_counter()
for _ in range(num_iterations):
hasattr(wrapper, "nonexistent_xyz")
elapsed_ms = (time.perf_counter() - start) * 1000
avg_ms = elapsed_ms / num_iterations
assert avg_ms < 1.0, f"Real CUDAGraphWrapper hasattr took {avg_ms:.2f}ms avg. Expected <1ms with the optimization."
assert model.repr_call_count == 0, f"__repr__ called {model.repr_call_count} times"
+117 -81
View File
@@ -114,91 +114,127 @@ async def async_request_openai_chat_omni_completions(
output = MixRequestFuncOutput()
output.prompt_len = request_func_input.prompt_len
generated_text = ""
generated_audio = None
ttft = 0.0
st = time.perf_counter()
output.start_time = st
most_recent_timestamp = st
timestamp = st
audio_generate_time = 0.0
try:
async with session.post(url=api_url, json=payload, headers=headers) as response:
if response.status == 200:
handler = StreamedResponseHandler()
async for chunk_bytes in response.content.iter_any():
chunk_bytes = chunk_bytes.strip()
if not chunk_bytes:
continue
messages = handler.add_chunk(chunk_bytes)
for message in messages:
if type(message) is bytes:
message = message.decode("utf-8")
# NOTE: SSE comments (often used as pings) start with
# a colon. These are not JSON data payload and should
# be skipped.
if message.startswith(":"):
max_retries = 3
retry_delay = 0.1
for attempt in range(max_retries + 1):
# Reset per-attempt state so that retries do not mix partial
# outputs or metrics from previous attempts.
generated_text = ""
generated_audio = None
ttft = 0.0
st = time.perf_counter()
output.start_time = st
most_recent_timestamp = st
timestamp = st
audio_generate_time = 0.0
output.itl = []
output.generated_text = ""
output.ttft = 0.0
output.audio_ttfp = 0.0
output.audio_duration = 0.0
output.audio_frames = 0
output.audio_rtf = 0.0
output.text_latency = 0.0
output.output_tokens = 0
output.error = ""
output.success = False
try:
async with session.post(url=api_url, json=payload, headers=headers) as response:
if response.status == 200:
handler = StreamedResponseHandler()
async for chunk_bytes in response.content.iter_any():
chunk_bytes = chunk_bytes.strip()
if not chunk_bytes:
continue
chunk = message.removeprefix("data: ")
if chunk != "[DONE]":
timestamp = time.perf_counter()
data = json.loads(chunk)
if choices := data.get("choices"):
modality = data.get("modality")
content = choices[0]["delta"].get("content")
if modality == "text":
# First token
if ttft == 0.0:
ttft = timestamp - st
output.ttft = ttft
else:
output.itl.append(timestamp - most_recent_timestamp)
generated_text += content or ""
most_recent_timestamp = timestamp
output.text_latency = timestamp - st
elif modality == "audio":
if output.audio_ttfp == 0.0:
output.audio_ttfp = timestamp - st
audio_generate_time = timestamp - st
if content != "":
audio_bytes = base64.b64decode(content)
seg = AudioSegment.from_file(io.BytesIO(audio_bytes))
if seg is not None:
if generated_audio is None:
generated_audio = seg
else:
generated_audio = generated_audio + seg
messages = handler.add_chunk(chunk_bytes)
for message in messages:
if type(message) is bytes:
message = message.decode("utf-8")
# NOTE: SSE comments (often used as pings) start with
# a colon. These are not JSON data payload and should
# be skipped.
if message.startswith(":"):
continue
if metrics := data.get("metrics"):
output.output_tokens = metrics.get("num_tokens_out", 0)
chunk = message.removeprefix("data: ")
if chunk != "[DONE]":
timestamp = time.perf_counter()
data = json.loads(chunk)
if choices := data.get("choices"):
modality = data.get("modality")
content = choices[0]["delta"].get("content")
if modality == "text":
# First token
if ttft == 0.0:
ttft = timestamp - st
output.ttft = ttft
else:
output.itl.append(timestamp - most_recent_timestamp)
generated_text += content or ""
most_recent_timestamp = timestamp
output.text_latency = timestamp - st
elif modality == "audio":
if output.audio_ttfp == 0.0:
output.audio_ttfp = timestamp - st
audio_generate_time = timestamp - st
if content != "":
audio_bytes = base64.b64decode(content)
seg = AudioSegment.from_file(io.BytesIO(audio_bytes))
if seg is not None:
if generated_audio is None:
generated_audio = seg
else:
generated_audio = generated_audio + seg
output.latency = timestamp - st
output.generated_text = generated_text
if generated_audio is not None:
output.audio_duration = len(generated_audio) / 1000.0
frame_width = generated_audio.frame_width
if frame_width > 0:
output.audio_frames = len(generated_audio.raw_data) // frame_width
else:
output.audio_frames = 0
logger.warning("Audio frame width is zero")
audio_duration = output.audio_duration
if audio_duration > 0:
output.audio_rtf = audio_generate_time / output.audio_duration
else:
output.audio_rtf = 0
logger.warning("Audio duration is zero")
output.success = True
else:
output.error = response.reason or ""
output.success = False
except Exception:
output.success = False
output.error = traceback.format_exc()
logger.error(f"ERROR: send request failed, reason is: {output.error}")
if metrics := data.get("metrics"):
output.output_tokens = metrics.get("num_tokens_out", 0)
output.latency = timestamp - st
output.generated_text = generated_text
if generated_audio is not None:
output.audio_duration = len(generated_audio) / 1000.0
frame_width = generated_audio.frame_width
if frame_width > 0:
output.audio_frames = len(generated_audio.raw_data) // frame_width
else:
output.audio_frames = 0
logger.warning("Audio frame width is zero")
audio_duration = output.audio_duration
if audio_duration > 0:
output.audio_rtf = audio_generate_time / output.audio_duration
else:
output.audio_rtf = 0
logger.warning("Audio duration is zero")
output.success = True
else:
output.error = response.reason or ""
output.success = False
break
except aiohttp.ClientError as e:
# transient transport error: may retry
output.success = False
output.error = traceback.format_exc()
if attempt < max_retries:
logger.warning(
"ClientError in omni benchmark request (will retry): attempt=%d/%d delay=%.2fs: %s",
attempt + 1,
max_retries + 1,
retry_delay,
str(e),
)
await asyncio.sleep(retry_delay)
continue
logger.error(
"ClientError in omni benchmark request (giving up):\n%s",
output.error,
)
break
except Exception:
output.success = False
output.error = traceback.format_exc()
logger.error(f"ERROR: send request failed, reason is: {output.error}")
break
if pbar:
pbar.update(1)
+29 -6
View File
@@ -11,6 +11,7 @@ Runtime parameters (gpu_memory_utilization, tp_size, etc.) come from CLI.
from __future__ import annotations
import re
import warnings
from dataclasses import asdict, dataclass, field
from enum import Enum
from pathlib import Path
@@ -74,7 +75,7 @@ class StageConfig:
# Per-stage engine args from pipeline YAML (defaults)
yaml_engine_args: dict[str, Any] = field(default_factory=dict)
# Per-stage runtime config from pipeline YAML (devices, max_batch_size)
# Per-stage runtime config from pipeline YAML (devices, etc.)
yaml_runtime: dict[str, Any] = field(default_factory=dict)
# Pass-through fields from pipeline YAML (default_sampling_params,
# output_connectors, input_connectors, tts_args, etc.)
@@ -109,13 +110,23 @@ class StageConfig:
# Build runtime config from YAML defaults + CLI overrides
runtime: dict[str, Any] = dict(self.yaml_runtime)
runtime.setdefault("process", True)
runtime.setdefault("max_batch_size", self.runtime_overrides.get("max_batch_size", 1))
if "devices" in self.runtime_overrides:
runtime["devices"] = self.runtime_overrides["devices"]
# Inject max_num_seqs from max_batch_size (legacy compat)
max_batch_size = int(runtime.get("max_batch_size", 1) or 1)
engine_args.setdefault("max_num_seqs", max_batch_size)
# Legacy compat: migrate runtime.max_batch_size → engine_args.max_num_seqs
legacy_mbs = runtime.pop("max_batch_size", None)
cli_mbs = self.runtime_overrides.get("max_batch_size")
if legacy_mbs is not None or cli_mbs is not None:
warnings.warn(
"runtime.max_batch_size is deprecated and will be removed in a "
"future release. Use engine_args.max_num_seqs instead.",
FutureWarning,
stacklevel=2,
)
effective_mbs = int(cli_mbs or legacy_mbs or 1)
engine_args.setdefault("max_num_seqs", effective_mbs)
engine_args.setdefault("max_num_seqs", 1)
# Build full config dict
config_dict: dict[str, Any] = {
@@ -222,6 +233,7 @@ class StageConfigFactory:
"qwen2_5_omni": "qwen2_5_omni",
"bagel": "bagel",
"qwen3_tts": "qwen3_tts",
"voxtral_tts": "voxtral_tts",
"mimo_audio": "mimo_audio",
"glm-image": "glm_image",
"cosyvoice3": "cosyvoice3",
@@ -321,13 +333,14 @@ class StageConfigFactory:
if "dtype" in engine_args:
engine_args["dtype"] = str(engine_args["dtype"])
engine_args.setdefault("max_num_seqs", 1)
config_dict: dict[str, Any] = {
"stage_id": 0,
"stage_type": StageType.DIFFUSION.value,
"runtime": {
"process": True,
"devices": devices,
"max_batch_size": 1,
},
"engine_args": create_config(engine_args),
"final_output": True,
@@ -425,6 +438,16 @@ class StageConfigFactory:
raw_rt = stage_data.get("runtime", None)
yaml_runtime = to_dict(raw_rt) if raw_rt is not None else {}
# Migrate legacy runtime.max_batch_size → engine_args.max_num_seqs
if "max_batch_size" in yaml_runtime:
mbs = yaml_runtime.pop("max_batch_size")
yaml_engine_args.setdefault("max_num_seqs", int(mbs))
logger.debug(
"Stage %s: migrated runtime.max_batch_size=%s to engine_args.max_num_seqs",
stage_data.get("stage_id", "?"),
mbs,
)
# Topology-level fields that also live inside engine_args in legacy
# YAMLs (worker_type, scheduler_cls, etc.) — read from both places.
worker_type = stage_data.get("worker_type", None) or yaml_engine_args.pop("worker_type", None)
Loaded 100 of 202 files, more files were not shown because too many files have changed in this diff. Show more