Compare commits
33
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6838533998 | ||
|
|
d6e338f9d6 | ||
|
|
a90a7690b8 | ||
|
|
8e12007ac1 | ||
|
|
da21e991cc | ||
|
|
a978c60163 | ||
|
|
c1c0a5b3cd | ||
|
|
072647eeb2 | ||
|
|
7007217d4b | ||
|
|
95018c536a | ||
|
|
b33d7637de | ||
|
|
ff254795fa | ||
|
|
6901ba426e | ||
|
|
da7a8f8014 | ||
|
|
199f7832e7 | ||
|
|
fafa248665 | ||
|
|
05d8de813e | ||
|
|
4c2293743b | ||
|
|
6813432f90 | ||
|
|
c2a5420d97 | ||
|
|
b4342fb487 | ||
|
|
81a90d24ef | ||
|
|
6ac7ba692d | ||
|
|
e9fe78e28c | ||
|
|
b766c4702c | ||
|
|
0239057cba | ||
|
|
99cc76690d | ||
|
|
60a8c01a96 | ||
|
|
89fff09628 | ||
|
|
ba5a66b15b | ||
|
|
f9268822cd | ||
|
|
48b3969717 | ||
|
|
b1ceaa7b34 |
No files matched your search
@@ -90,10 +90,18 @@ upload_pipeline() {
|
||||
FAIL_FAST=$(fail_fast)
|
||||
|
||||
cd .buildkite
|
||||
|
||||
# Select test definition file: merge suite for main, ready suite for PRs
|
||||
if [[ $BUILDKITE_BRANCH == "main" ]]; then
|
||||
TEST_YAML="test-amd-merge.yml"
|
||||
else
|
||||
TEST_YAML="test-amd-ready.yaml"
|
||||
fi
|
||||
|
||||
(
|
||||
set -x
|
||||
# Output pipeline.yaml with all blank lines removed
|
||||
minijinja-cli test-template.j2 test-amd.yaml \
|
||||
minijinja-cli test-template.j2 "$TEST_YAML" \
|
||||
-D branch="$BUILDKITE_BRANCH" \
|
||||
-D list_file_diff="$LIST_FILE_DIFF" \
|
||||
-D run_all="$RUN_ALL" \
|
||||
|
||||
@@ -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 \
|
||||
|
||||
@@ -0,0 +1,197 @@
|
||||
steps:
|
||||
|
||||
- label: "Simple Unit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
|
||||
|
||||
- label: "Diffusion Images API LoRA E2E"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
|
||||
# - label: "Audio Generation Model Test"
|
||||
# agent_pool: mi325_1
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
# commands:
|
||||
# - export GPU_ARCHS=gfx942
|
||||
# - export VLLM_LOGGING_LEVEL=DEBUG
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
|
||||
|
||||
- label: "Diffusion Cache Backend Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
|
||||
|
||||
# merge-only tests
|
||||
- label: "Diffusion Tensor Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Benchmark & Engine Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/benchmarks/test_serve_cli.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
- label: "Omni Model Test Qwen2-5-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
|
||||
- label: "Qwen3-TTS E2E Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
|
||||
|
||||
- label: "Diffusion Image Edit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
|
||||
|
||||
# split Bagel Model Test with H100 (Real Weights) into three tests
|
||||
- label: "Bagel Text2Img Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
|
||||
|
||||
- label: "Bagel Img2Img Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
|
||||
|
||||
- label: "Bagel Online Serving Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
|
||||
@@ -0,0 +1,178 @@
|
||||
steps:
|
||||
|
||||
- label: "Simple Unit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
|
||||
# - label: "Audio Generation Model Test"
|
||||
# agent_pool: mi325_1
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
# commands:
|
||||
# - export GPU_ARCHS=gfx942
|
||||
# - export VLLM_LOGGING_LEVEL=DEBUG
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
|
||||
|
||||
- label: "Diffusion Cache Backend Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Benchmark & Engine Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/benchmarks/test_serve_cli.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
- label: "Omni Model Test Qwen2-5-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 17m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- timeout 10m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
|
||||
|
||||
- label: "Qwen3-TTS E2E Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
|
||||
|
||||
- label: "Diffusion Image Edit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
|
||||
|
||||
- label: "Bagel Text2Img Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
|
||||
- label: "Bagel Img2Img Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
|
||||
- label: "Bagel Online Serving Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
@@ -1,7 +1,7 @@
|
||||
steps:
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
timeout_in_minutes: 20
|
||||
timeout_in_minutes: 30
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
@@ -11,7 +11,7 @@ steps:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
|
||||
|
||||
- label: "Diffusion Images API LoRA E2E"
|
||||
timeout_in_minutes: 20
|
||||
timeout_in_minutes: 30
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
@@ -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,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}
|
||||
|
||||
@@ -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}}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
```
|
||||
|
||||
@@ -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:**
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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` |
|
||||
|
||||
@@ -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).
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 = """\
|
||||
|
||||
@@ -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"
|
||||
@@ -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)
|
||||
|
||||
@@ -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
Reference in new issue
Block a user