Compare commits
200
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
963ba1ab7a | ||
|
|
e25239c2ec | ||
|
|
2414b6f94d | ||
|
|
b3dd0a4baf | ||
|
|
6cefe91370 | ||
|
|
84569280b1 | ||
|
|
1e43656ab1 | ||
|
|
3875b81c18 | ||
|
|
b63b36ded3 | ||
|
|
95d56cf062 | ||
|
|
6e8ef213d8 | ||
|
|
31f882935d | ||
|
|
fccd95b12b | ||
|
|
469c1dc87e | ||
|
|
519cf3ecf1 | ||
|
|
bfc7cedcf1 | ||
|
|
e10aca30ec | ||
|
|
c6e8d78eda | ||
|
|
e5030c7524 | ||
|
|
a1cc583665 | ||
|
|
a57b8f858b | ||
|
|
9a9c8f7503 | ||
|
|
af02713952 | ||
|
|
d8dc1defa7 | ||
|
|
800c2b8251 | ||
|
|
221584a489 | ||
|
|
4e1c13236e | ||
|
|
706bad253d | ||
|
|
8949fdc440 | ||
|
|
3818ba4c65 | ||
|
|
c33ca60208 | ||
|
|
0fc4032a71 | ||
|
|
2e108a4109 | ||
|
|
ce69eb35e9 | ||
|
|
07737e6c61 | ||
|
|
d2e65a92f1 | ||
|
|
3842970a82 | ||
|
|
9bb4974198 | ||
|
|
0e2d55f46d | ||
|
|
40b2959183 | ||
|
|
3329f1b674 | ||
|
|
23484d8a11 | ||
|
|
bdfc77147d | ||
|
|
e35cd213d7 | ||
|
|
35ee3c7543 | ||
|
|
3a7c7f1422 | ||
|
|
12e75c6618 | ||
|
|
bd37f3c27d | ||
|
|
7c729e19d7 | ||
|
|
b550709bba | ||
|
|
da9c45e1de | ||
|
|
1fb423e5f6 | ||
|
|
2a2033df02 | ||
|
|
bc794e625f | ||
|
|
2b7249fbb3 | ||
|
|
29343540f7 | ||
|
|
c067485030 | ||
|
|
8d83fb59e9 | ||
|
|
41a795b5d6 | ||
|
|
ee33601571 | ||
|
|
b76291d6a1 | ||
|
|
2dcfcc0519 | ||
|
|
634af60fba | ||
|
|
ac4106900b | ||
|
|
9afdab4939 | ||
|
|
ef71d74c81 | ||
|
+1 |
f92d84fcf9 | ||
|
|
d4c139508f | ||
|
|
57e5f0fe69 | ||
|
|
5dfdf584dd | ||
|
|
6f4bd3e45c | ||
|
|
412bf42a75 | ||
|
|
ffd9771c54 | ||
|
|
11c4fcedd3 | ||
|
|
b0f3672aec | ||
|
|
2289478e3a | ||
|
|
8d13b54383 | ||
|
|
802a6242f0 | ||
|
|
ea6d811b58 | ||
|
|
5994a0bbf5 | ||
|
|
914e92cf80 | ||
|
|
3b317285c4 | ||
|
|
a3456752c5 | ||
|
|
f43bb124bc | ||
|
|
272e626100 | ||
|
|
5377108438 | ||
|
|
9fdb3bcb89 | ||
|
|
f1c2caa3c7 | ||
|
|
b690a10e3a | ||
|
|
eba5617c95 | ||
|
|
2696751003 | ||
|
|
8c4a42b16c | ||
|
|
3457758114 | ||
|
|
67bb9b4a12 | ||
|
|
4d9262a6eb | ||
|
|
c69ed4037d | ||
|
|
68a1d23e4d | ||
|
|
593c4c06aa | ||
|
|
4987bec425 | ||
|
|
5dd1d0a772 | ||
|
|
834abdaf8b | ||
|
|
5dcf18a984 | ||
|
|
900b4187a1 | ||
|
|
8935439428 | ||
|
|
bb9f21d0db | ||
|
|
5eac50f22c | ||
|
|
126fb4b4a9 | ||
|
|
8d9a4afed9 | ||
|
|
1acd27c205 | ||
|
|
740ec608f4 | ||
|
|
74258b18b7 | ||
|
|
1c876f0cd1 | ||
|
|
ee212d12b1 | ||
|
|
c20afebce4 | ||
|
|
0e3bf4de8e | ||
|
|
b72ffb82bb | ||
|
|
168033cac0 | ||
|
|
6fa9dbbe42 | ||
|
|
beb00fce31 | ||
|
|
3a347b10b0 | ||
|
|
0f3264a9c0 | ||
|
|
dab596e468 | ||
|
|
2b999429c4 | ||
|
|
6031c7d393 | ||
|
|
373d3c7476 | ||
|
|
1903cee56a | ||
|
|
39705f9ee8 | ||
|
|
8766c0c59f | ||
|
|
6ea66f8ed0 | ||
|
|
220ce0b4ea | ||
|
|
5c3c81920b | ||
|
|
1f3a3e3848 | ||
|
|
d37912b071 | ||
|
|
37eebff674 | ||
|
|
a698269e47 | ||
|
|
36dc32ff81 | ||
|
|
9abf9545df | ||
|
|
c33ef34654 | ||
|
|
c349f62839 | ||
|
|
41cb6378a2 | ||
|
|
753508fe02 | ||
|
|
a3802d1e5c | ||
|
|
3c48257f8d | ||
|
|
283873411a | ||
|
|
1f1c82ecdd | ||
|
|
e459fdb30a | ||
|
|
9322cc6176 | ||
|
|
abc31731ab | ||
|
|
385df0359e | ||
|
|
36eb8afba5 | ||
|
|
ec94e83a25 | ||
|
|
cad25fa20c | ||
|
|
11a951b962 | ||
|
|
410a6a20b0 | ||
|
|
52020c7af7 | ||
|
|
8f45e68bcd | ||
|
|
80dea88007 | ||
|
|
e996723bc9 | ||
|
|
5f9aee193e | ||
|
|
9a37ca83d4 | ||
|
|
b52a5ca401 | ||
|
|
c890b3a99c | ||
|
|
6c1bb27b69 | ||
|
|
cc0e15ba57 | ||
|
|
28ce618f0e | ||
|
|
748470c1fb | ||
|
|
d01eaca39a | ||
|
|
4e4458e7bf | ||
|
|
e7644daa7f | ||
|
|
d714b4ce04 | ||
|
|
23e243378f | ||
|
|
5799d85a51 | ||
|
|
89e389cf37 | ||
|
|
90ed618b5a | ||
|
|
28467dc19a | ||
|
|
f0af3ab200 | ||
|
|
0ed3186db1 | ||
|
|
8e937400dc | ||
|
|
076c34abfb | ||
|
|
7d8816b46b | ||
|
|
1de15bdf4d | ||
|
|
e949ccf087 | ||
|
|
bd1617ce2b | ||
|
|
2bc9cb7f13 | ||
|
|
8297570c2b | ||
|
|
474710dda9 | ||
|
|
7d9d8aee79 | ||
|
|
083b5e363f | ||
|
|
f74cba84ac | ||
|
|
2aafe0bf94 | ||
|
|
5cf4605bf9 | ||
|
|
e1c6bc7597 | ||
|
|
fc8486cf20 | ||
|
|
1c530d984d | ||
|
|
9dd36e3afd | ||
|
|
109bba8cdf | ||
|
|
f99e8b2d3c | ||
|
|
06ab723afc | ||
|
|
12bcc32c3d | ||
|
|
d4e90bbba8 |
No files matched your search
@@ -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.21.0"
|
||||
VLLM_VERSION: "v0.22.0"
|
||||
priority: 100
|
||||
timeout_in_minutes: 60
|
||||
soft_fail: false
|
||||
@@ -52,6 +52,7 @@ time timeout -k 30 30m docker run \
|
||||
pip install tblib==3.1.0
|
||||
cd /workspace/vllm-omni
|
||||
pytest -v -s -m "core_model and xpu and B60"
|
||||
pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py
|
||||
export VLLM_XPU_USE_SAMPLER_KERNEL=0 # NOTE: Remove this after vLLM v0.21.1 is merged. Fixes Qwen2-5 omni-expansion tests.
|
||||
pytest -v -s -m "advanced_model and xpu and B60"
|
||||
'
|
||||
@@ -1,15 +1,28 @@
|
||||
steps:
|
||||
|
||||
- label: "Simple Unit Test"
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Simple · Other Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Diffusion Images API LoRA E2E"
|
||||
agent_pool: mi325_1
|
||||
@@ -49,16 +62,6 @@ steps:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
|
||||
|
||||
- label: "Diffusion Cache Backend Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- 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 (Need 4 GPUs)"
|
||||
agent_pool: mi325_4
|
||||
depends_on: amd-build
|
||||
@@ -123,7 +126,7 @@ steps:
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
|
||||
- label: "Qwen3-TTS CustomVoice E2E Test"
|
||||
agent_pool: mi325_1
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
steps:
|
||||
|
||||
- label: "Simple Unit Test"
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
@@ -9,7 +9,18 @@ steps:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Simple · Other Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Voxtral TTS CUDA Unit Test"
|
||||
agent_pool: mi325_1
|
||||
@@ -19,13 +30,13 @@ steps:
|
||||
commands:
|
||||
- timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
|
||||
|
||||
- label: "Diffusion Batching Test"
|
||||
- label: "Diffusion · Batch Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/batching/test_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
|
||||
- timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
|
||||
|
||||
- label: "Custom Pipeline Test"
|
||||
agent_pool: mi325_1
|
||||
|
||||
@@ -107,7 +107,7 @@ steps:
|
||||
depends_on: upload-merge-pipeline
|
||||
steps:
|
||||
- label: "Entrypoint Test with H100"
|
||||
timeout_in_minutes: 35
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
@@ -248,6 +248,47 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "TTS · GLM-TTS Test"
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
pytest -s -v tests/e2e/online_serving/test_glm_tts.py tests/e2e/offline_inference/test_glm_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni · Qwen3-Omni Test"
|
||||
timeout_in_minutes: 30
|
||||
@@ -348,6 +389,54 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · VoxCPM2 Online Test"
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni · Higgs-Audio V2 Test"
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_USE_DEEP_GEMM=0
|
||||
export VLLM_MOE_USE_DEEP_GEMM=0
|
||||
pytest -s -v tests/e2e/online_serving/test_higgs_audio_v2.py tests/e2e/offline_inference/test_higgs_audio_v2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion · Qwen Image Edit Test"
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
|
||||
+300
-25
@@ -305,6 +305,51 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Omni · Perf Test · Multi-Replica"
|
||||
key: nightly-omni-performance-multi-replicas
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 3
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
|
||||
- label: ":full_moon: Omni · Multi-Replica Startup Test with 4x H100"
|
||||
timeout_in_minutes: 45
|
||||
@@ -349,11 +394,11 @@ steps:
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
|
||||
steps:
|
||||
- label: ":full_moon: TTS · Function Test"
|
||||
- label: ":full_moon: TTS · Function Test with L4"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
@@ -368,6 +413,69 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
# - label: ":full_moon: TTS · Function Test with H100"
|
||||
# timeout_in_minutes: 90
|
||||
# commands:
|
||||
# - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
# - pytest -s -v tests/e2e/ -m "full_model and H100 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 1
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# - name: HF_TOKEN
|
||||
# valueFrom:
|
||||
# secretKeyRef:
|
||||
# name: hf-token-secret
|
||||
# key: token
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: TTS · Accuracy Test · VoxCPM2"
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
|
||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts_expansion.py -m 'full_model and cuda and tts' --run-level 'full_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: ":full_moon: TTS · Perf Test"
|
||||
key: nightly-tts-performance
|
||||
timeout_in_minutes: 180
|
||||
@@ -414,6 +522,52 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: TTS · VoxCPM2 · Perf Test"
|
||||
key: nightly-tts-performance-voxcpm2
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_voxcpm2.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below.
|
||||
- group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
|
||||
key: nightly-diffusion-x2iat-group
|
||||
@@ -525,6 +679,7 @@ steps:
|
||||
tests/e2e/online_serving/test_longcat_image_expansion.py
|
||||
tests/e2e/online_serving/test_longcat_image_edit_expansion.py
|
||||
tests/e2e/online_serving/test_flux_2_dev_expansion.py
|
||||
tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
|
||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
agents:
|
||||
@@ -578,7 +733,88 @@ steps:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 3
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU DreamZero"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_dreamzero_expansion.py
|
||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · GLM-Image Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- pytest -sv tests/e2e/online_serving/test_glm_image_expansion.py -m "full_model and diffusion and H100" --run-level "full_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
@@ -701,6 +937,46 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: HunyuanImage3-DIT · Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export HUNYUAN_IMAGE3_MODEL="tencent/HunyuanImage-3.0-Instruct"
|
||||
- export HUNYUAN_IMAGE3_DEVICES="0,1,2,3"
|
||||
- pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m full_model --run-level full_model
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 4
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image
|
||||
timeout_in_minutes: 180
|
||||
@@ -797,8 +1073,8 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2509"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2509
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
@@ -806,7 +1082,7 @@ steps:
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
@@ -893,26 +1169,20 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · HunyuanImage3 · DiT Perf Test"
|
||||
key: nightly-hunyuan-image3-performance
|
||||
soft_fail: true
|
||||
timeout_in_minutes: 120
|
||||
if: build.env("RUN_HUNYUAN_IMAGE3_PERF") == "1"
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · BAGEL"
|
||||
key: nightly-diffusion-x2iat-performance-bagel
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp4_fp8.json
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_sp2.json
|
||||
EXIT2=$$?
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_cfgp2.json
|
||||
EXIT3=$$?
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$((EXIT1 | EXIT2 | EXIT3))
|
||||
exit $$EXIT
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -922,7 +1192,7 @@ steps:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 4
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
@@ -959,7 +1229,7 @@ steps:
|
||||
- label: ":full_moon: Diffusion X2V · Function Test"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and cuda" --run-level "full_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -998,6 +1268,8 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
|
||||
agents:
|
||||
@@ -1086,9 +1358,8 @@ steps:
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "diffusion-x2iat-test"
|
||||
steps:
|
||||
- label: ":full_moon: Quantization Test with H100"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- timeout 30m pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
|
||||
- timeout 60m pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -1189,14 +1460,16 @@ steps:
|
||||
- nightly-omni-performance-no-async-chunk
|
||||
- nightly-omni-performance-async-chunk
|
||||
- nightly-omni-performance-vllm-text
|
||||
- nightly-omni-performance-multi-replicas
|
||||
- nightly-tts-performance
|
||||
- nightly-diffusion-x2iat-performance-qwen-image
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit-2509
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
- nightly-diffusion-x2iat-performance-bagel
|
||||
- nightly-diffusion-x2v-performance
|
||||
- nightly-testcase-statistics
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
|
||||
commands:
|
||||
- pip install openpyxl
|
||||
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
|
||||
@@ -1205,10 +1478,12 @@ steps:
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2509
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
|
||||
- python tools/nightly/generate_nightly_perf_excel.py
|
||||
|
||||
@@ -183,7 +183,7 @@ steps:
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 15m bash -c "
|
||||
timeout 20m bash -c "
|
||||
pytest -sv tests/distributed/ -m 'core_model and cuda' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
@@ -273,6 +273,42 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni · Qwen3-Omni Realtime WebSocket"
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_streaming_audio_input_pcm_output_async_chunk_smoke -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
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: "Omni · VoxCPM2 Test"
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
@@ -296,6 +332,29 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · VoxCPM2 Online Test"
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model and tts' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · Qwen3-TTS CustomVoice Test"
|
||||
commands:
|
||||
- |
|
||||
|
||||
@@ -53,6 +53,9 @@
|
||||
{% endif %}
|
||||
{{ indented_cmd | safe }}
|
||||
priority: 100
|
||||
{% if step.timeout_in_minutes %}
|
||||
timeout_in_minutes: {{ step.timeout_in_minutes }}
|
||||
{% endif %}
|
||||
{% if step.grade and step.grade == "Blocking" %}
|
||||
soft_fail: false
|
||||
{% else %}
|
||||
|
||||
@@ -21,6 +21,9 @@ Before proposing changes, ask for the optimization scene if it is not already kn
|
||||
- Current enabled strategies: USP/SP, CFG parallel, HSDP/FSDP, VAE patch parallel, torch.compile, profiler options.
|
||||
- Optimization target: latency, throughput, memory, cost, or quality-preserving speed.
|
||||
- Precision/quality tolerance: bf16/fp8/quantization/sparsity/approximate attention allowed or not.
|
||||
- Quality validation method: exact-output comparison when feasible, or metrics
|
||||
such as SSIM, PSNR, LPIPS, cosine similarity, MAE, MSE, temporal flicker
|
||||
checks, baseline self-run variance, and reviewable output artifacts.
|
||||
|
||||
## Workflow
|
||||
|
||||
@@ -79,7 +82,19 @@ Before proposing changes, ask for the optimization scene if it is not already kn
|
||||
- Record latency, stage timings, memory, output quality, and logs.
|
||||
- Report comparison tables in ms. Include at least end-to-end client time, server `inference_time`, server stage generation time if available, `vae.encode`, `diffuse`, `vae.decode`, and peak memory.
|
||||
|
||||
6. **Collect diagnostic trace only after narrowing hypotheses.**
|
||||
6. **Enforce the quality and precision gate for every optimization.**
|
||||
- Do not mark an optimization ready until quality is checked on the same
|
||||
workload: model, prompt/media, seed, shape, steps, scheduler, dtype,
|
||||
backend, parallelism, and output encoding.
|
||||
- For math-preserving changes, require exact/near-exact agreement or show
|
||||
that differences are within baseline self-run variance.
|
||||
- For precision, quantization, approximate attention, sparse/custom kernels,
|
||||
or backend changes, include reviewable artifacts and metrics such as SSIM,
|
||||
PSNR, LPIPS, cosine similarity, MAE, MSE, and temporal flicker checks.
|
||||
- If quality tolerance is not stated, default to quality-preserving behavior.
|
||||
Failed or inconclusive quality validation blocks a ready-to-merge claim.
|
||||
|
||||
7. **Collect diagnostic trace only after narrowing hypotheses.**
|
||||
- Use torch profiler for a small number of requests.
|
||||
- Run two separate diagnostic traces instead of mixing concerns:
|
||||
- **Operator/shape trace**: enable `torch_profiler_record_shapes=True` and keep stack collection disabled. Use this to rank CUDA kernels, NCCL collectives, attention/MLP/norm/RoPE work, and shape-specific hot operators.
|
||||
@@ -91,15 +106,16 @@ Before proposing changes, ask for the optimization scene if it is not already kn
|
||||
- Diagnostic reports must be written to disk and preserve: server command, profiler config, warmup/request/polling commands, trace artifact paths, rank analyzed, analyzer output or summary, and the decision about whether additional ranks are necessary.
|
||||
- Analyze both rank-level balance and device-level free bubbles when additional ranks are opened.
|
||||
|
||||
7. **Analyze host, communication, and operators.**
|
||||
8. **Analyze host, communication, and operators.**
|
||||
- Find GPU idle/free intervals and map each large gap to the enclosing CPU/Python code.
|
||||
- Separate real GPU idle from profiler overhead such as CUPTI `Command Buffer Full`.
|
||||
- Compare NCCL kernel time to user annotations; annotations can overcount nested intervals.
|
||||
- Rank operator work by total CUDA time and by repeated small-kernel launch count.
|
||||
|
||||
8. **Produce an optimization plan.**
|
||||
9. **Produce an optimization plan.**
|
||||
- Classify candidates as P0/P1/P2.
|
||||
- For each candidate, state necessity, expected benefit, implementation path, validation plan, and quality risk.
|
||||
- Include a concrete quality/precision validation gate for each candidate.
|
||||
- Do not implement high-risk operator rewrites before proving the operator is a bottleneck for the target shapes.
|
||||
- End the plan with a user-facing candidate selection table. The assistant
|
||||
should not automatically choose a risky optimization just because it is
|
||||
@@ -124,7 +140,13 @@ Before proposing changes, ask for the optimization scene if it is not already kn
|
||||
- **P1:** meaningful code changes with contained risk. Examples: cross-attention KV caching, VAE gather/broadcast reduction, AdaLayerNorm/RMSNorm/RoPE fusion after trace evidence.
|
||||
- **P2:** high implementation or quality risk. Examples: FA to LA replacement, custom Triton/CUDA fused kernels, FP8/quantization, sparsity/Rainfusion-style acceleration.
|
||||
|
||||
Every implemented optimization needs A/B validation and quality regression. A/B means same workload and hardware before/after. Quality regression means checking generated image/video stability, artifacts, temporal flicker, and seed behavior when precision or approximate kernels change.
|
||||
Every implemented optimization needs A/B validation and a passing
|
||||
quality/precision gate, including math-preserving changes such as padding trim,
|
||||
layout cleanup, cache reuse, or host/runtime cleanup. Use objective metrics
|
||||
when available, keep reviewable artifacts, and compare against baseline
|
||||
self-run variance for precision, quantization, approximate kernels, or backend
|
||||
changes. If quality validation fails or is inconclusive, do not present the
|
||||
optimization as ready to merge.
|
||||
|
||||
## Optimization Layers
|
||||
|
||||
|
||||
@@ -270,3 +270,6 @@ vllm_omni/_version.py
|
||||
.worktrees/
|
||||
# CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
|
||||
tests/.ci_generated/
|
||||
|
||||
# Local Lance precision-matching dump scripts (not part of the repo)
|
||||
lance_compare/
|
||||
@@ -10,6 +10,11 @@ build:
|
||||
jobs:
|
||||
post_checkout:
|
||||
- git fetch --unshallow || true
|
||||
pre_install:
|
||||
# Pre-install CPU torch so the subsequent `pip install .[docs]` does not
|
||||
# pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
|
||||
# 1 GB memory cap and time out the mkdocs build.
|
||||
- python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
|
||||
|
||||
mkdocs:
|
||||
configuration: mkdocs.yml
|
||||
|
||||
+51
-35
@@ -1,44 +1,60 @@
|
||||
# Benchmarks Overview and Architecture
|
||||
# Benchmarks
|
||||
|
||||
This document explains the benchmark architecture across all benchmark assets in this repo. It describes what we measure, and where to find or plug in new scenarios. Per-task details remain in subfolder READMEs (e.g., `benchmarks/<model>/README.md`).
|
||||
This directory contains benchmark suites for evaluating different model families and infrastructure components in vLLM-Omni. Each subfolder targets a different benchmark family with its own scripts, configs, and metrics. See the per-subfolder READMEs for detailed usage.
|
||||
|
||||
## Scope and goals
|
||||
- Establish repeatable latency/throughput measurements for multimodal LLM pipelines.
|
||||
- Provide both HF Transformers (offline) and vLLM-Omni (multi-stage/pipeline) baselines.
|
||||
- Make it easy to plug in new datasets and models with minimal changes to the runner scripts.
|
||||
## Benchmark families
|
||||
|
||||
## Dataset and inputs
|
||||
- Default example: SeedTTS top-100 prompts (`benchmarks/build_dataset/top100.txt`) via `benchmarks/build_dataset/`.
|
||||
- Extensible: drop in new prompt files or modality-aligned payloads; keep the expected format for the consuming scripts (e.g., one prompt per line).
|
||||
- If you add a new dataset, document it under `benchmarks/<model>/README.md` and point scripts to your data path.
|
||||
### [TTS](tts/README.md) — Text-to-Speech
|
||||
|
||||
## Directory layout
|
||||
- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
|
||||
- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
|
||||
- `benchmarks/accuracy/` — accuracy benchmark integrations that adapt external
|
||||
benchmark suites to vLLM-Omni serving and evaluation flows.
|
||||
- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
|
||||
Model-agnostic serving benchmarks for TTS models, including Qwen3-TTS and VoxCPM2.
|
||||
|
||||
## Reference workflows
|
||||
- **HF Transformers (offline, single process)**
|
||||
Script (example): `benchmarks/<model>/transformers/eval_qwen3_moe_omni_transformers.sh`
|
||||
Outputs: `benchmark_results/perf_stats.json`, `benchmark_results/results.json`, `benchmark_results/audio/` (if audio is produced).
|
||||
- **Layout**: `tts/bench_tts.py` (serving benchmark driver), `tts/model_configs.yaml` (model registry), `tts/plot_results.py` (result plotting)
|
||||
- **Dataset**: Seed-TTS full or text-only datasets, plus bundled smoke/design prompts under `build_dataset/`
|
||||
- **Key metrics**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor), throughput (audio seconds / wall-clock second)
|
||||
|
||||
- **vLLM-Omni end-to-end pipeline**
|
||||
Script (example): `benchmarks/<model>/vllm_omni/eval_qwen3_moe_omni.sh`
|
||||
Outputs: `vllm_omni/logs/*.stats.jsonl` (per-stage/overall latency & TPS), `vllm_omni/logs/stage*.log`, `vllm_omni/outputs/` (text/audio artifacts).
|
||||
### [Diffusion](diffusion/README.md) — Image and Video Generation
|
||||
|
||||
- **Adding a new task/model**
|
||||
1) Create `benchmarks/<model>/transformers/` and/or `benchmarks/<model>/vllm_omni/` with scripts referencing your model and dataset.
|
||||
2) Add a task README describing dataset, configs, and expected outputs.
|
||||
3) Keep the output/log structure similar for easy comparison (perf_stats/results/audio or text outputs; stats.jsonl/logs for pipeline).
|
||||
Online-serving benchmark for diffusion image/video models, sending requests to the configured vLLM serving endpoint (`/v1/chat/completions`, `/v1/images/generations`, or `/v1/videos`, depending on backend/task).
|
||||
|
||||
## Metrics to watch
|
||||
- **Throughput**: `overall_tps`, `*_tps_avg` per stage.
|
||||
- **Latency distribution**: look for long tails in `*.stats.jsonl`.
|
||||
- **Quality/completeness**: missing outputs or errors in stage logs indicate pipeline failures or misconfigurations.
|
||||
- **Tasks**: text-to-image, text-to-video, image-to-image, image-to-video, text+image-to-image, text+image-to-video
|
||||
- **Datasets**: `vbench`, `trace`, `random`
|
||||
- **Key metrics**: request throughput, latency percentiles, optional SLO attainment
|
||||
|
||||
## Troubleshooting
|
||||
- Verify GPU/driver/FlashAttention2 requirements for your chosen model/config.
|
||||
- Ensure network access for dataset/model downloads (Google Drive, Hugging Face, etc.).
|
||||
- If outputs are missing or slow, inspect per-stage logs and `*.stats.jsonl` for errors, stragglers, or contention.
|
||||
### [GLM-Image](glm_image/README.md) — Text-to-Image and Image-to-Image
|
||||
|
||||
Benchmarks for GLM-Image performance across HuggingFace baseline, vLLM-Omni offline inference, and vLLM-Omni online serving.
|
||||
|
||||
- **Layout**: `glm_image/huggingface/` (HF baseline), `glm_image/vllm-omni/` (offline inference), `glm_image/benchmark_glm_image.py` (online serving)
|
||||
- **Tasks**: text-to-image and image-to-image
|
||||
- **Key metrics**: request/image throughput, latency percentiles, optional per-stage pipeline timings
|
||||
|
||||
### [Distributed](distributed/omni_connectors/README.md) — RDMA Connector Testing
|
||||
|
||||
RDMA environment setup and transfer tests for `MooncakeTransferEngineConnector`, including pytest-based single-node checks and manual cross-node benchmarks.
|
||||
|
||||
- **Transfer modes**: `copy`, `zerocopy`, `gpu` (GPUDirect)
|
||||
- **Supports**: single-node pytest suites and manual multi-node/cross-node transfer testing
|
||||
|
||||
### [Accuracy](accuracy/README.md) — Image Generation and Editing Quality
|
||||
|
||||
Accuracy benchmarks for image generation/editing models, adapting external suites to vLLM-Omni serving and local judge-evaluation flows.
|
||||
|
||||
- **Layout**: `accuracy/text_to_image/` (GEBench), `accuracy/image_to_image/` (GEdit-Bench)
|
||||
- **Method**: generation and judge scoring both run through local `vllm-omni serve` endpoints
|
||||
|
||||
### Common serving metrics framework
|
||||
|
||||
`vllm_omni/benchmarks/` extends `vllm bench serve --omni` with Omni-specific datasets, backends, and multimodal metrics. Key metrics include:
|
||||
|
||||
- **Text output**: TTFT (time to first token), TPOT (time per output token), ITL (inter-token latency)
|
||||
- **Audio output**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor)
|
||||
- **Throughput**: request throughput, output token throughput, total token throughput, audio throughput
|
||||
|
||||
See `vllm_omni/benchmarks/serve.py` for the `vllm bench serve --omni` runner wrapper and `vllm_omni/benchmarks/metrics/` for Omni metric definitions.
|
||||
|
||||
## Adding a new benchmark
|
||||
|
||||
1. Create a subfolder under `benchmarks/<name>/` with scripts, configs if needed, and a `README.md`.
|
||||
2. If comparing against another runtime, use clear backend subfolders where applicable, such as `huggingface/` and `vllm-omni/`, or follow the shared TTS serving benchmark pattern in `tts/`.
|
||||
3. Add reusable dataset or prompt-building utilities to `build_dataset/` if applicable.
|
||||
4. Update this README with a link to the new benchmark family.
|
||||
@@ -151,6 +151,9 @@ class VllmOmniImageClient:
|
||||
seed: int | None = None,
|
||||
negative_prompt: str | None = None,
|
||||
output_compression: int | None = None,
|
||||
bot_task: str | None = None,
|
||||
sys_type: str | None = None,
|
||||
system_prompt: str | None = None,
|
||||
) -> Image.Image:
|
||||
if not isinstance(images, list):
|
||||
images = [images]
|
||||
@@ -170,6 +173,12 @@ class VllmOmniImageClient:
|
||||
data["negative_prompt"] = negative_prompt
|
||||
if output_compression is not None:
|
||||
data["output_compression"] = str(output_compression)
|
||||
if bot_task is not None:
|
||||
data["bot_task"] = bot_task
|
||||
if sys_type is not None:
|
||||
data["sys_type"] = sys_type
|
||||
if system_prompt is not None:
|
||||
data["system_prompt"] = system_prompt
|
||||
|
||||
files = [
|
||||
(
|
||||
|
||||
@@ -64,6 +64,47 @@ def resolve_model_name(*, model_name: str | None, model: str | None = None, outp
|
||||
raise ValueError("model-name is required when it cannot be inferred from model or output-root")
|
||||
|
||||
|
||||
def _coerce_score_int(value: Any) -> int:
|
||||
if isinstance(value, bool):
|
||||
raise TypeError(f"Expected numeric score, got bool: {value!r}")
|
||||
if isinstance(value, (int, float)):
|
||||
return int(value)
|
||||
if isinstance(value, str):
|
||||
stripped = value.strip()
|
||||
if stripped.lstrip("-").isdigit():
|
||||
return int(stripped)
|
||||
raise TypeError(f"Expected numeric score, got {type(value).__name__}: {value!r}")
|
||||
|
||||
|
||||
def _expand_score_entry(entry: Any) -> list[int]:
|
||||
if isinstance(entry, dict):
|
||||
if "score" in entry:
|
||||
return [_coerce_score_int(entry["score"])]
|
||||
if "value" in entry:
|
||||
return [_coerce_score_int(entry["value"])]
|
||||
numeric_values = [_coerce_score_int(value) for value in entry.values()]
|
||||
if numeric_values:
|
||||
return numeric_values
|
||||
raise TypeError(f"Score dict has no numeric values: {entry!r}")
|
||||
return [_coerce_score_int(entry)]
|
||||
|
||||
|
||||
def _normalize_score_list(score: Any) -> list[int]:
|
||||
if isinstance(score, dict):
|
||||
entries: list[Any] = [score]
|
||||
elif isinstance(score, list):
|
||||
entries = score
|
||||
else:
|
||||
entries = [score]
|
||||
|
||||
normalized: list[int] = []
|
||||
for entry in entries:
|
||||
normalized.extend(_expand_score_entry(entry))
|
||||
if not normalized:
|
||||
raise ValueError(f"Empty score list after normalization: {score!r}")
|
||||
return normalized
|
||||
|
||||
|
||||
def parse_score_payload(raw_text: str) -> dict[str, Any]:
|
||||
try:
|
||||
parsed = extract_json_object(raw_text)
|
||||
@@ -73,10 +114,7 @@ def parse_score_payload(raw_text: str) -> dict[str, Any]:
|
||||
parsed = {"score": json.loads(stripped), "reasoning": ""}
|
||||
else:
|
||||
raise
|
||||
score = parsed.get("score", [])
|
||||
if not isinstance(score, list):
|
||||
score = [score]
|
||||
parsed["score"] = [int(value) for value in score]
|
||||
parsed["score"] = _normalize_score_list(parsed.get("score", []))
|
||||
return parsed
|
||||
|
||||
|
||||
@@ -391,6 +429,9 @@ class GEditBenchRunner:
|
||||
num_inference_steps: int = 20,
|
||||
guidance_scale: float | None = None,
|
||||
seed: int | None = 42,
|
||||
bot_task: str | None = None,
|
||||
sys_type: str | None = None,
|
||||
system_prompt: str | None = None,
|
||||
):
|
||||
self.dataset_ref = dataset_ref
|
||||
self.output_root = output_root
|
||||
@@ -400,6 +441,9 @@ class GEditBenchRunner:
|
||||
self.num_inference_steps = num_inference_steps
|
||||
self.guidance_scale = guidance_scale
|
||||
self.seed = seed
|
||||
self.bot_task = bot_task
|
||||
self.sys_type = sys_type
|
||||
self.system_prompt = system_prompt
|
||||
self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
|
||||
|
||||
def generate(
|
||||
@@ -478,6 +522,9 @@ class GEditBenchRunner:
|
||||
num_inference_steps=self.num_inference_steps,
|
||||
guidance_scale=self.guidance_scale,
|
||||
seed=self.seed,
|
||||
bot_task=self.bot_task,
|
||||
sys_type=self.sys_type,
|
||||
system_prompt=self.system_prompt,
|
||||
)
|
||||
edited_image.save(output_path)
|
||||
return {
|
||||
@@ -633,6 +680,25 @@ def build_parser() -> argparse.ArgumentParser:
|
||||
generate.add_argument("--num-inference-steps", type=int, default=20)
|
||||
generate.add_argument("--guidance-scale", type=float, default=None)
|
||||
generate.add_argument("--seed", type=int, default=42)
|
||||
generate.add_argument(
|
||||
"--bot-task",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["think", "recaption", "think_recaption", "vanilla"],
|
||||
help="HunyuanImage-3 prompting mode for /v1/images/edits (optional).",
|
||||
)
|
||||
generate.add_argument(
|
||||
"--sys-type",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Hunyuan system-prompt type (e.g. en_unified). Defaults from --bot-task when omitted.",
|
||||
)
|
||||
generate.add_argument(
|
||||
"--system-prompt",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Custom Hunyuan system prompt body (sys_type=custom).",
|
||||
)
|
||||
generate.add_argument("--workers", type=int, default=1)
|
||||
generate.add_argument("--max-samples", type=int, default=None)
|
||||
generate.add_argument("--samples-per-group", type=int, default=None)
|
||||
@@ -675,6 +741,9 @@ def main(argv: list[str] | None = None) -> int:
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
seed=args.seed,
|
||||
bot_task=args.bot_task,
|
||||
sys_type=args.sys_type,
|
||||
system_prompt=args.system_prompt,
|
||||
)
|
||||
records = runner.generate(
|
||||
model_name=model_name,
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
se001|ocean waves crashing on a rocky beach|10.0
|
||||
se002|heavy rain on a metal roof at night|8.0
|
||||
se003|a dog barking in the distance|6.0
|
||||
se004|a crackling campfire in a quiet forest|10.0
|
||||
se005|footsteps on dry leaves|6.0
|
||||
se006|thunder rolling across a wide valley|8.0
|
||||
se007|a busy coffee shop with chatter and espresso machines|10.0
|
||||
se008|wind howling through tall pine trees|8.0
|
||||
se009|a small stream flowing over pebbles|10.0
|
||||
se010|a creaking wooden door slowly opening|6.0
|
||||
@@ -0,0 +1,10 @@
|
||||
ttsd001|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Hello and welcome to the show today. [S2] Thanks for having me, it's a pleasure to be here. [S1] Let's dive right in.
|
||||
ttsd002|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you finish the report? [S2] Almost — I'm just polishing the conclusion. [S1] Great, send it over when you're done.
|
||||
ttsd003|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have you tried that new restaurant downtown? [S2] Yes, the pasta is amazing. [S1] I'll book a table for Friday.
|
||||
ttsd004|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] What did you think of yesterday's meeting? [S2] Productive, but a bit long. [S1] Agreed, we should set an hour cap next time.
|
||||
ttsd005|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] The weather is beautiful today. [S2] Perfect for a hike. [S1] Want to drive out to the trails this afternoon?
|
||||
ttsd006|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] How did the demo go? [S2] Better than expected — the client signed off on the prototype. [S1] Fantastic news!
|
||||
ttsd007|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Coffee or tea this morning? [S2] Coffee please, with a splash of milk. [S1] Coming right up.
|
||||
ttsd008|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you read the email about the deadline change? [S2] I did, we now have until next Friday. [S1] That should give us enough buffer.
|
||||
ttsd009|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Are you joining the team dinner tonight? [S2] I'll be there around seven. [S1] See you then.
|
||||
ttsd010|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have we backed up the database? [S2] Last backup ran two hours ago. [S1] Good, I'll trigger a fresh snapshot before the migration.
|
||||
Whitespace-only changes.
@@ -10,6 +10,8 @@ from typing import Any
|
||||
import aiohttp
|
||||
from tqdm import tqdm
|
||||
|
||||
DEFAULT_EDITS_BOT_TASK = "think"
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestFuncInput:
|
||||
@@ -27,6 +29,7 @@ class RequestFuncInput:
|
||||
extra_body: dict[str, Any] = field(default_factory=dict)
|
||||
image_paths: list[str] | None = None
|
||||
request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
default_bot_task: str | None = DEFAULT_EDITS_BOT_TASK
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -53,6 +56,95 @@ def _encode_image_as_data_url(path: str) -> str:
|
||||
return f"data:{mime};base64,{encoded}"
|
||||
|
||||
|
||||
async def async_request_image_edits(
|
||||
input: RequestFuncInput,
|
||||
session: aiohttp.ClientSession,
|
||||
pbar: tqdm | None = None,
|
||||
enable_diffusion_pipeline_profiler: bool = False,
|
||||
) -> RequestFuncOutput:
|
||||
"""POST /v1/images/edits (multipart)."""
|
||||
del enable_diffusion_pipeline_profiler
|
||||
output = RequestFuncOutput()
|
||||
output.start_time = time.perf_counter()
|
||||
|
||||
extra_body = dict(input.extra_body)
|
||||
width = input.width or extra_body.get("width") or 1024
|
||||
height = input.height or extra_body.get("height") or 1024
|
||||
edits_url = input.api_url
|
||||
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("model", input.model)
|
||||
form.add_field("prompt", input.prompt)
|
||||
form.add_field("size", f"{width}x{height}")
|
||||
form.add_field("response_format", "b64_json")
|
||||
|
||||
if input.num_inference_steps is not None:
|
||||
form.add_field("num_inference_steps", str(input.num_inference_steps))
|
||||
elif extra_body.get("num_inference_steps") is not None:
|
||||
form.add_field("num_inference_steps", str(extra_body["num_inference_steps"]))
|
||||
|
||||
if input.seed is not None:
|
||||
form.add_field("seed", str(input.seed))
|
||||
elif extra_body.get("seed") is not None:
|
||||
form.add_field("seed", str(extra_body["seed"]))
|
||||
|
||||
if extra_body.get("guidance_scale") is not None:
|
||||
form.add_field("guidance_scale", str(extra_body["guidance_scale"]))
|
||||
if extra_body.get("negative_prompt") is not None:
|
||||
form.add_field("negative_prompt", str(extra_body["negative_prompt"]))
|
||||
if extra_body.get("true_cfg_scale") is not None:
|
||||
form.add_field("true_cfg_scale", str(extra_body["true_cfg_scale"]))
|
||||
if extra_body.get("sys_type") is not None:
|
||||
form.add_field("sys_type", str(extra_body["sys_type"]))
|
||||
if extra_body.get("system_prompt") is not None:
|
||||
form.add_field("system_prompt", str(extra_body["system_prompt"]))
|
||||
|
||||
bot_task = extra_body.get("bot_task")
|
||||
if bot_task is None and input.default_bot_task is not None:
|
||||
bot_task = input.default_bot_task
|
||||
if bot_task is not None:
|
||||
form.add_field("bot_task", str(bot_task))
|
||||
|
||||
assert input.image_paths is not None
|
||||
for img_path in input.image_paths:
|
||||
if not os.path.exists(img_path):
|
||||
output.error = f"Image file not found: {img_path}"
|
||||
output.success = False
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return output
|
||||
with open(img_path, "rb") as img_f:
|
||||
image_bytes = img_f.read()
|
||||
form.add_field(
|
||||
"image",
|
||||
image_bytes,
|
||||
filename=os.path.basename(img_path),
|
||||
content_type=_guess_mime_type(img_path),
|
||||
)
|
||||
|
||||
try:
|
||||
async with session.post(edits_url, data=form) as response:
|
||||
if response.status == 200:
|
||||
resp_json = await response.json()
|
||||
output.response_body = resp_json
|
||||
output.success = True
|
||||
else:
|
||||
output.error = f"HTTP {response.status}: {await response.text()}"
|
||||
output.success = False
|
||||
except Exception as e:
|
||||
output.error = str(e)
|
||||
output.success = False
|
||||
|
||||
output.latency = time.perf_counter() - output.start_time
|
||||
|
||||
if output.success and input.slo_ms is not None:
|
||||
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return output
|
||||
|
||||
|
||||
async def async_request_chat_completions(
|
||||
input: RequestFuncInput,
|
||||
session: aiohttp.ClientSession,
|
||||
@@ -378,6 +470,7 @@ backends_function_mapping = {
|
||||
"2i": {
|
||||
"/v1/chat/completions": (async_request_chat_completions, "/v1/chat/completions"),
|
||||
"/v1/images/generations": (async_request_openai_image_generations, "/v1/images/generations"),
|
||||
"/v1/images/edits": (async_request_image_edits, "/v1/images/edits"),
|
||||
},
|
||||
"2v": {
|
||||
"/v1/videos": (async_request_v1_videos, "/v1/videos"),
|
||||
|
||||
@@ -0,0 +1,328 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Standalone diagnostic for the diffusion attention backends added in #3079.
|
||||
|
||||
Exercises the same synthetic attention shape through each backend and each
|
||||
torch SDPA sub-kernel, with and without an attn_mask, so we can pinpoint why
|
||||
CUDNN_ATTN underperforms on some SKUs. Two typical causes:
|
||||
|
||||
(a) cuDNN lacks a tuned kernel for this (SM, head_dim, seq) combination
|
||||
and silently falls back to MATH.
|
||||
(b) attn_mask is non-None and the cuDNN SDPA dispatch rejects it,
|
||||
walking CUDNN -> FLASH -> MATH.
|
||||
|
||||
Run:
|
||||
python benchmarks/diffusion/bench_attention_backends.py --preset hv15
|
||||
python benchmarks/diffusion/bench_attention_backends.py --preset wan22
|
||||
python benchmarks/diffusion/bench_attention_backends.py \
|
||||
--batch 1 --heads 24 --seq 14336 --head-dim 128
|
||||
|
||||
Optional extras:
|
||||
pip install --pre flash-attn-4 # FA4 is currently pre-release only
|
||||
pip install -U flashinfer # latest FlashInfer (0.6.9)
|
||||
|
||||
The table at the end is the data we want on the PR — surface the row where a
|
||||
backend is >1.5x the SDPA baseline, that's the one to gate off the auto-route.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
|
||||
import torch
|
||||
from torch.nn.attention import SDPBackend, sdpa_kernel
|
||||
|
||||
# Presets mirror the DiT hot-path attention call (single-stream block) at the
|
||||
# shapes we run in the PR's validation configs.
|
||||
_PRESETS = {
|
||||
# HunyuanVideo-1.5 480p/33f: (30 * 52 * 9) latent + ~256 text tokens.
|
||||
"hv15": {"batch": 1, "heads": 24, "seq": 14336, "head_dim": 128},
|
||||
# Wan 2.2 480p/33f rough estimate; adjust when we measure real shapes.
|
||||
"wan22": {"batch": 1, "heads": 40, "seq": 16384, "head_dim": 128},
|
||||
# Smaller image-gen shape for quick smoke tests.
|
||||
"flux": {"batch": 1, "heads": 24, "seq": 4096, "head_dim": 128},
|
||||
}
|
||||
|
||||
_SDPA_BACKENDS = [
|
||||
("CUDNN_ATTENTION", [SDPBackend.CUDNN_ATTENTION]),
|
||||
("FLASH_ATTENTION", [SDPBackend.FLASH_ATTENTION]),
|
||||
("EFFICIENT_ATTENTION", [SDPBackend.EFFICIENT_ATTENTION]),
|
||||
("MATH", [SDPBackend.MATH]),
|
||||
# The PR's CUDNN_ATTN impl uses this priority chain; keep it here to show
|
||||
# what torch actually picks when multiple are allowed.
|
||||
("CUDNN_ATTN_CHAIN", [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]),
|
||||
]
|
||||
|
||||
|
||||
def _env_report() -> None:
|
||||
if not torch.cuda.is_available():
|
||||
raise SystemExit("CUDA not available; this diagnostic needs a GPU.")
|
||||
cc = torch.cuda.get_device_capability()
|
||||
print("=" * 72)
|
||||
print(f"GPU : {torch.cuda.get_device_name()}")
|
||||
print(f"SM capability : sm_{cc[0]}{cc[1]}")
|
||||
print(f"torch : {torch.__version__}")
|
||||
print(f"cuDNN version : {torch.backends.cudnn.version()}")
|
||||
print("flashinfer : ", end="")
|
||||
try:
|
||||
import flashinfer
|
||||
|
||||
print(getattr(flashinfer, "__version__", "present"))
|
||||
except Exception as e:
|
||||
print(f"not installed ({type(e).__name__})")
|
||||
print("=" * 72)
|
||||
|
||||
|
||||
def _make_qkv(batch: int, heads: int, seq: int, head_dim: int, device: str, dtype: torch.dtype):
|
||||
# Layout: (B, S, H, D) — what the PR's backends accept as input.
|
||||
q = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
|
||||
k = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
|
||||
v = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
|
||||
return q, k, v
|
||||
|
||||
|
||||
def _make_mask(batch: int, seq: int, device: str, dtype: torch.dtype, pad_tokens: int = 128) -> torch.Tensor:
|
||||
# Mirrors a text-encoder padding mask: last `pad_tokens` positions masked.
|
||||
mask = torch.zeros(batch, 1, seq, seq, device=device, dtype=dtype)
|
||||
if pad_tokens > 0:
|
||||
mask[..., -pad_tokens:] = float("-inf")
|
||||
return mask
|
||||
|
||||
|
||||
def _time_call(fn, *args, warmup: int = 3, iters: int = 10) -> tuple[float, str]:
|
||||
"""Return (median ms, error string). Skips on any exception so that a
|
||||
backend that rejects our args (wrong dtype, missing JIT module, unsupported
|
||||
kwarg value) doesn't abort the whole sweep."""
|
||||
try:
|
||||
for _ in range(warmup):
|
||||
fn(*args)
|
||||
torch.accelerator.synchronize()
|
||||
times = []
|
||||
for _ in range(iters):
|
||||
torch.accelerator.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
fn(*args)
|
||||
torch.accelerator.synchronize()
|
||||
times.append((time.perf_counter() - t0) * 1000.0)
|
||||
times.sort()
|
||||
return times[len(times) // 2], ""
|
||||
except Exception as e: # noqa: BLE001 — probe script, keep the whole table going
|
||||
msg = str(e).split("\n", 1)[0][:60]
|
||||
return float("nan"), f"{type(e).__name__}: {msg}" if msg else type(e).__name__
|
||||
|
||||
|
||||
def _run_sdpa_variants(q, k, v, attn_mask, scale: float) -> list[tuple[str, float, str]]:
|
||||
# q/k/v come in (B, S, H, D); F.sdpa wants (B, H, S, D).
|
||||
q_t = q.permute(0, 2, 1, 3)
|
||||
k_t = k.permute(0, 2, 1, 3)
|
||||
v_t = v.permute(0, 2, 1, 3)
|
||||
rows: list[tuple[str, float, str]] = []
|
||||
for name, backends in _SDPA_BACKENDS:
|
||||
|
||||
def _call(mask=attn_mask):
|
||||
with sdpa_kernel(backends):
|
||||
return torch.nn.functional.scaled_dot_product_attention(
|
||||
q_t, k_t, v_t, attn_mask=mask, dropout_p=0.0, is_causal=False, scale=scale
|
||||
)
|
||||
|
||||
ms, err = _time_call(_call)
|
||||
rows.append((name, ms, err))
|
||||
return rows
|
||||
|
||||
|
||||
def _run_flashinfer(q, k, v, scale: float, backend: str | None = None, attn_mask=None) -> tuple[float, str]:
|
||||
"""Call FlashInfer's dense single-prefill.
|
||||
|
||||
``backend`` hints at cutlass/fa3/trtllm if the installed version exposes
|
||||
that kwarg. ``attn_mask`` is the same additive 4D mask we pass to SDPA —
|
||||
we convert it to the 2D boolean ``custom_mask`` form FlashInfer accepts.
|
||||
Per FlashInfer docs, ``custom_mask`` only applies when ``causal=False``.
|
||||
"""
|
||||
try:
|
||||
import inspect
|
||||
|
||||
from flashinfer.prefill import single_prefill_with_kv_cache
|
||||
except Exception as e:
|
||||
return float("nan"), f"import-{type(e).__name__}"
|
||||
|
||||
kwargs: dict = {"sm_scale": scale, "causal": False, "return_lse": False}
|
||||
if backend is not None:
|
||||
sig = inspect.signature(single_prefill_with_kv_cache)
|
||||
if "backend" not in sig.parameters:
|
||||
return float("nan"), "no-backend-kwarg"
|
||||
kwargs["backend"] = backend
|
||||
|
||||
if attn_mask is not None:
|
||||
# Collapse (B, 1, S, S) additive float mask to (S, S) boolean.
|
||||
# FlashInfer expects True = keep, False = masked out.
|
||||
mask_2d = attn_mask[0, 0]
|
||||
kwargs["custom_mask"] = mask_2d != float("-inf")
|
||||
|
||||
def _call():
|
||||
out = single_prefill_with_kv_cache(q[0], k[0], v[0], **kwargs)
|
||||
return out.unsqueeze(0)
|
||||
|
||||
return _time_call(_call)
|
||||
|
||||
|
||||
def _run_fa4(q, k, v, scale: float) -> tuple[float, str]:
|
||||
"""Call FlashAttention-4 directly (``pip install flash-attn-4``). FA4
|
||||
ships a Blackwell-native kernel via CuTe-DSL; on sm_120 it should beat
|
||||
cuDNN by ~20%. API lives under ``flash_attn.cute``; the older
|
||||
``flash_attn.flash_attn_func`` path is FA2/FA3 only."""
|
||||
try:
|
||||
from flash_attn.cute import flash_attn_func
|
||||
except Exception as e:
|
||||
return float("nan"), f"import-{type(e).__name__}"
|
||||
|
||||
def _call():
|
||||
# FA4 accepts (B, S, H, D) directly — same layout the PR's backends use.
|
||||
return flash_attn_func(q, k, v, softmax_scale=scale, causal=False)
|
||||
|
||||
return _time_call(_call)
|
||||
|
||||
|
||||
def _run_flashinfer_cudnn_batch(q, k, v, scale: float) -> tuple[float, str]:
|
||||
"""Call FlashInfer's direct cuDNN wrapper. Bypasses PyTorch SDPA dispatch
|
||||
(which has a few hundred ns of overhead per call) and talks to cuDNN FMHA
|
||||
straight. Useful as a ceiling for 'pure cuDNN, no SDPA' on Blackwell."""
|
||||
try:
|
||||
from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
|
||||
except Exception as e:
|
||||
return float("nan"), f"import-{type(e).__name__}"
|
||||
|
||||
b, s, h, d = q.shape
|
||||
qo_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
|
||||
kv_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
|
||||
|
||||
def _call():
|
||||
return cudnn_batch_prefill_with_kv_cache(
|
||||
q.reshape(b * s, h, d),
|
||||
k.reshape(b * s, h, d),
|
||||
v.reshape(b * s, h, d),
|
||||
qo_indptr=qo_indptr,
|
||||
kv_indptr=kv_indptr,
|
||||
max_qo_len=s,
|
||||
max_kv_len=s,
|
||||
sm_scale=scale,
|
||||
causal=False,
|
||||
)
|
||||
|
||||
return _time_call(_call)
|
||||
|
||||
|
||||
def _print_table(title: str, rows: list[tuple[str, float, str]], baseline_name: str | None = None) -> None:
|
||||
print(f"\n{title}")
|
||||
print("-" * 88)
|
||||
header = f"{'backend':<24} {'median (ms)':>14} {'vs baseline':>12} status"
|
||||
print(header)
|
||||
print("-" * 88)
|
||||
baseline_ms = None
|
||||
if baseline_name is not None:
|
||||
for name, ms, _ in rows:
|
||||
if name == baseline_name and ms == ms: # not NaN
|
||||
baseline_ms = ms
|
||||
break
|
||||
for name, ms, err in rows:
|
||||
ms_str = f"{ms:>14.3f}" if ms == ms else f"{'n/a':>14}"
|
||||
if baseline_ms is not None and ms == ms:
|
||||
ratio = baseline_ms / ms
|
||||
ratio_str = f"{ratio:>11.2f}x"
|
||||
else:
|
||||
ratio_str = f"{'—':>12}"
|
||||
status = "ok" if not err else f"FAILED ({err})"
|
||||
print(f"{name:<24} {ms_str} {ratio_str} {status}")
|
||||
|
||||
|
||||
def _pick_winner(rows: list[tuple[str, float, str]]) -> tuple[str, float] | None:
|
||||
"""Return (backend, ms) of the fastest non-failing row, or None."""
|
||||
ok_rows = [(n, m) for n, m, err in rows if not err and m == m]
|
||||
if not ok_rows:
|
||||
return None
|
||||
return min(ok_rows, key=lambda x: x[1])
|
||||
|
||||
|
||||
def _bench_one_shape(shape: dict, dtype: torch.dtype, device: str) -> tuple[list, list]:
|
||||
torch.manual_seed(0)
|
||||
q, k, v = _make_qkv(**shape, device=device, dtype=dtype)
|
||||
scale = 1.0 / (shape["head_dim"] ** 0.5)
|
||||
|
||||
rows_nomask = _run_sdpa_variants(q, k, v, attn_mask=None, scale=scale)
|
||||
rows_nomask.append(("FLASHINFER (default)", *_run_flashinfer(q, k, v, scale)))
|
||||
# `trtllm-gen` has no sm_120 cubins (NVIDIA/TensorRT-LLM#11799) — skipped.
|
||||
for fi_backend in ("fa2", "fa3", "cutlass", "auto"):
|
||||
rows_nomask.append((f"FLASHINFER ({fi_backend})", *_run_flashinfer(q, k, v, scale, backend=fi_backend)))
|
||||
rows_nomask.append(("FLASHINFER (cudnn-batch)", *_run_flashinfer_cudnn_batch(q, k, v, scale)))
|
||||
rows_nomask.append(("FA4 (direct)", *_run_fa4(q, k, v, scale)))
|
||||
|
||||
mask = _make_mask(shape["batch"], shape["seq"], device=device, dtype=dtype, pad_tokens=256)
|
||||
rows_mask = _run_sdpa_variants(q, k, v, attn_mask=mask, scale=scale)
|
||||
rows_mask.append(("FLASHINFER (dense)", *_run_flashinfer(q, k, v, scale, attn_mask=mask)))
|
||||
rows_mask.append(("FA4 (direct)", float("nan"), "mask-not-supported"))
|
||||
|
||||
return rows_nomask, rows_mask
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--preset", choices=list(_PRESETS.keys()), default="hv15")
|
||||
parser.add_argument("--sweep", action="store_true", help="Run all presets and print a ranking")
|
||||
parser.add_argument("--batch", type=int, default=None)
|
||||
parser.add_argument("--heads", type=int, default=None)
|
||||
parser.add_argument("--seq", type=int, default=None)
|
||||
parser.add_argument("--head-dim", type=int, default=None)
|
||||
parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16")
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
args = parser.parse_args()
|
||||
|
||||
dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
|
||||
_env_report()
|
||||
print(f"dtype : {args.dtype}")
|
||||
|
||||
presets = list(_PRESETS.keys()) if args.sweep else [args.preset]
|
||||
summary: dict[str, tuple[str, float] | None] = {}
|
||||
|
||||
for preset_name in presets:
|
||||
shape = dict(_PRESETS[preset_name])
|
||||
if not args.sweep:
|
||||
for k in ("batch", "heads", "seq", "head_dim"):
|
||||
v = getattr(args, k if k != "head_dim" else "head_dim")
|
||||
if v is not None:
|
||||
shape[k] = v
|
||||
|
||||
print("\n" + "=" * 88)
|
||||
print(
|
||||
f"Preset: {preset_name} | batch={shape['batch']} heads={shape['heads']} "
|
||||
f"seq={shape['seq']} head_dim={shape['head_dim']}"
|
||||
)
|
||||
print("=" * 88)
|
||||
|
||||
rows_nomask, rows_mask = _bench_one_shape(shape, dtype, args.device)
|
||||
_print_table(f"[{preset_name}] No attention mask", rows_nomask, baseline_name="CUDNN_ATTENTION")
|
||||
_print_table(
|
||||
f"[{preset_name}] With attention mask (pad 256 tokens)", rows_mask, baseline_name="CUDNN_ATTENTION"
|
||||
)
|
||||
|
||||
summary[preset_name] = _pick_winner(rows_nomask)
|
||||
|
||||
if args.sweep or len(presets) > 1:
|
||||
print("\n" + "=" * 88)
|
||||
print("Winners per preset (no-mask path)")
|
||||
print("=" * 88)
|
||||
for preset_name, winner in summary.items():
|
||||
if winner is None:
|
||||
print(f" {preset_name:<10} — no successful backend")
|
||||
else:
|
||||
name, ms = winner
|
||||
print(f" {preset_name:<10} {name:<24} {ms:>8.3f} ms")
|
||||
|
||||
print("\nNotes:")
|
||||
print(" * Ratios are relative to CUDNN_ATTENTION. >1.0x means faster than cuDNN.")
|
||||
print(" * Mask-path winner inherits CUDNN_ATTN's fallback in the PR's backends.")
|
||||
print(" * `trtllm-gen` and FA4 4.0.0b10 are known-broken on sm_120 as of Apr 2026.")
|
||||
print(" * For e2e timings, run text_to_video.py / text_to_image.py while varying DIFFUSION_ATTENTION_BACKEND.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,108 @@
|
||||
#!/usr/bin/env bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# 2x2 timing matrix (backend x compile mode) for SAGE_ATTN_3 vs TORCH_SDPA,
|
||||
# on HunyuanVideo-1.5 480p T2V and Wan 2.2 TI2V-5B. Run from the repo root:
|
||||
#
|
||||
# bash benchmarks/diffusion/bench_sage_attn3_matrix.sh
|
||||
#
|
||||
# Requires: sageattn3 installed (for SAGE_ATTN_3 rows) and a Blackwell GPU.
|
||||
|
||||
set -u
|
||||
|
||||
STEPS=${STEPS:-50}
|
||||
SEED=${SEED:-42}
|
||||
PROMPT=${PROMPT:-"A dog running across a field of golden wheat."}
|
||||
OUT_DIR=${OUT_DIR:-/tmp/bench_sage3}
|
||||
mkdir -p "$OUT_DIR"
|
||||
|
||||
declare -A TOTAL PER_STEP TEXT_ENC TRANSFORMER VAE
|
||||
|
||||
run() {
|
||||
local name="$1"; shift
|
||||
local model="$1"; shift
|
||||
local h="$1"; shift
|
||||
local w="$1"; shift
|
||||
local f="$1"; shift
|
||||
local g="$1"; shift
|
||||
local log="$OUT_DIR/${name}.log"
|
||||
echo "[$(date +%T)] Running $name ..."
|
||||
env "$@" python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model "$model" \
|
||||
--prompt "$PROMPT" \
|
||||
--height "$h" --width "$w" --num-frames "$f" \
|
||||
--num-inference-steps "$STEPS" --seed "$SEED" --guidance-scale "$g" \
|
||||
--enable-diffusion-pipeline-profiler \
|
||||
--output "$OUT_DIR/${name}.mp4" > "$log" 2>&1 || {
|
||||
echo " FAILED — see $log"
|
||||
TOTAL[$name]="FAIL"; PER_STEP[$name]="FAIL"
|
||||
TEXT_ENC[$name]="-"; TRANSFORMER[$name]="-"; VAE[$name]="-"
|
||||
return
|
||||
}
|
||||
|
||||
local wait_ms
|
||||
wait_ms=$(grep "add_req_and_wait=" "$log" | tail -1 | sed -nE 's/.*add_req_and_wait=([0-9.]+) ms.*/\1/p')
|
||||
local total_s
|
||||
total_s=$(grep -oE "Total generation time: [0-9.]+ seconds" "$log" | tail -1 | sed -nE 's/.*: ([0-9.]+) seconds.*/\1/p')
|
||||
local text_enc_s
|
||||
text_enc_s=$(grep -oE "text_encoder.forward took [0-9.]+s" "$log" | awk '{sub(/s$/,"",$NF); sum+=$NF} END {printf "%.3f", sum+0}')
|
||||
local vae_s
|
||||
vae_s=$(grep -oE "vae.decode took [0-9.]+s" "$log" | awk '{sub(/s$/,"",$NF); sum+=$NF} END {printf "%.3f", sum+0}')
|
||||
local pipeline_s
|
||||
pipeline_s=$(grep -oE "Pipeline\.forward took [0-9.]+s" "$log" | tail -1 | sed -nE 's/.*took ([0-9.]+)s.*/\1/p')
|
||||
local transformer_s
|
||||
transformer_s=$(awk "BEGIN {printf \"%.3f\", ${pipeline_s:-0} - ${vae_s:-0} - ${text_enc_s:-0}}")
|
||||
|
||||
TOTAL[$name]="${total_s:-0}"
|
||||
PER_STEP[$name]=$(awk "BEGIN {printf \"%.3f\", ${wait_ms:-0} / 1000 / $STEPS}")
|
||||
TEXT_ENC[$name]="$text_enc_s"
|
||||
TRANSFORMER[$name]="$transformer_s"
|
||||
VAE[$name]="$vae_s"
|
||||
|
||||
echo " total=${TOTAL[$name]}s per-step=${PER_STEP[$name]}s/it transformer=${TRANSFORMER[$name]}s vae=${VAE[$name]}s"
|
||||
}
|
||||
|
||||
# ========== HunyuanVideo-1.5 480p T2V (480x832x33, CFG=6.0) ==========
|
||||
HV_MODEL=hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
|
||||
run hv15_sdpa_eager "$HV_MODEL" 480 832 33 6.0 TORCH_COMPILE_DISABLE=1
|
||||
run hv15_sdpa_compiled "$HV_MODEL" 480 832 33 6.0
|
||||
run hv15_sage3_eager "$HV_MODEL" 480 832 33 6.0 TORCH_COMPILE_DISABLE=1 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
|
||||
run hv15_sage3_compiled "$HV_MODEL" 480 832 33 6.0 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
|
||||
|
||||
# ========== Wan 2.2 TI2V-5B (704x1280x33, CFG=5.0) ==========
|
||||
WAN_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers
|
||||
run wan22_sdpa_eager "$WAN_MODEL" 704 1280 33 5.0 TORCH_COMPILE_DISABLE=1
|
||||
run wan22_sdpa_compiled "$WAN_MODEL" 704 1280 33 5.0
|
||||
run wan22_sage3_eager "$WAN_MODEL" 704 1280 33 5.0 TORCH_COMPILE_DISABLE=1 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
|
||||
run wan22_sage3_compiled "$WAN_MODEL" 704 1280 33 5.0 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
|
||||
|
||||
print_table() {
|
||||
local title="$1"; shift
|
||||
echo
|
||||
echo "$title (${STEPS} steps, seed=${SEED})"
|
||||
echo
|
||||
printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
|
||||
"Config" "Total (s)" "Per-step (s/it)" "Text Enc (s)" "Transformer (s)" "VAE Decode (s)"
|
||||
printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
|
||||
"-----------------------" "---------" "---------------" "------------" "---------------" "--------------"
|
||||
for row in "$@"; do
|
||||
case "$row" in
|
||||
*_sdpa_eager) label="SDPA + Eager" ;;
|
||||
*_sdpa_compiled) label="SDPA + Compiled" ;;
|
||||
*_sage3_eager) label="SAGE_ATTN_3 + Eager" ;;
|
||||
*_sage3_compiled) label="SAGE_ATTN_3 + Compiled" ;;
|
||||
esac
|
||||
printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
|
||||
"$label" "${TOTAL[$row]}" "${PER_STEP[$row]}" "${TEXT_ENC[$row]}" "${TRANSFORMER[$row]}" "${VAE[$row]}"
|
||||
done
|
||||
}
|
||||
|
||||
print_table "HunyuanVideo-1.5 480p T2V, 832x480x33" \
|
||||
hv15_sdpa_eager hv15_sdpa_compiled hv15_sage3_eager hv15_sage3_compiled
|
||||
|
||||
print_table "Wan 2.2 TI2V-5B, 1280x704x33" \
|
||||
wan22_sdpa_eager wan22_sdpa_compiled wan22_sage3_eager wan22_sage3_compiled
|
||||
|
||||
echo
|
||||
echo "Logs: $OUT_DIR/*.log"
|
||||
@@ -7,10 +7,13 @@ Benchmark online serving for diffusion models (Image/Video Generation).
|
||||
If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
|
||||
|
||||
Supports multiple endpoints:
|
||||
- /v1/chat/completions: OpenAI chat-compatible image requests
|
||||
- /v1/chat/completions: OpenAI chat-compatible image requests (e.g. t2i, Qwen i2i)
|
||||
- /v1/images/edits: OpenAI image edit / IT2I (multipart; e.g. Hunyuan --bot-task think)
|
||||
- /v1/images/generations: OpenAI image generation requests
|
||||
- /v1/videos: Async video jobs
|
||||
|
||||
Legacy --backend vllm-omni and openai are aliases for chat/completions and images/generations.
|
||||
|
||||
Usage:
|
||||
# Video (/v1/videos endpoint)
|
||||
t2v:
|
||||
@@ -39,7 +42,19 @@ Usage:
|
||||
{"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
|
||||
]'
|
||||
|
||||
i2i:
|
||||
ti2i (Hunyuan / OpenAI image edit API):
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--endpoint /v1/images/edits --dataset random --task ti2i --num-prompts 10 \
|
||||
--bot-task think
|
||||
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--endpoint /v1/images/edits \
|
||||
--dataset custom \
|
||||
--dataset-path custom_requests.jsonl \
|
||||
--task ti2i \
|
||||
--bot-task think
|
||||
|
||||
i2i (chat-based models such as Qwen-Image-Edit):
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--endpoint /v1/chat/completions --dataset vbench --task i2i --num-prompts 10
|
||||
|
||||
@@ -132,7 +147,7 @@ class VBenchDataset(BaseDataset):
|
||||
def _load_data(self) -> list[dict[str, Any]]:
|
||||
if self.args.task == "t2v":
|
||||
return self._load_t2v_prompts()
|
||||
elif self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
|
||||
elif self.args.task in ["i2v", "ti2v", "ti2i", "i2i", "it2i"]:
|
||||
return self._load_i2v_data()
|
||||
else:
|
||||
return self._load_t2v_prompts()
|
||||
@@ -524,7 +539,7 @@ class TraceDataset(BaseDataset):
|
||||
single = row.get("image_path")
|
||||
image_paths = [single] if single else None
|
||||
|
||||
if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i"]:
|
||||
if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i", "it2i"]:
|
||||
raise ValueError(
|
||||
f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row."
|
||||
)
|
||||
@@ -558,6 +573,135 @@ class TraceDataset(BaseDataset):
|
||||
return [self[i] for i in range(len(self))]
|
||||
|
||||
|
||||
class CustomDataset(BaseDataset):
|
||||
"""
|
||||
Custom dataset that loads requests from a JSONL file.
|
||||
|
||||
Each line in the JSONL file should be a JSON object with the following fields:
|
||||
- prompt (required): The text prompt for the request
|
||||
- width (optional): Image/video width
|
||||
- height (optional): Image height
|
||||
- num_inference_steps (optional): Number of diffusion steps
|
||||
- seed (optional): Random seed
|
||||
- image_paths (optional): List of input image paths for i2i/i2v/ti2i tasks
|
||||
- image_urls (optional): List of input image URLs (alternative to image_paths)
|
||||
|
||||
Example JSONL for ti2i:
|
||||
{"prompt": "Add sunset lighting", "width": 1024, "height": 1024, "image_urls": ["https://example.com/image.jpg"]}
|
||||
"""
|
||||
|
||||
def __init__(self, args, api_url: str, model: str):
|
||||
super().__init__(args, api_url, model)
|
||||
self.dataset_path = args.dataset_path
|
||||
if not self.dataset_path:
|
||||
raise ValueError("--dataset-path must be provided when using 'custom' dataset")
|
||||
self.load_data()
|
||||
|
||||
def load_data(self) -> None:
|
||||
"""Load data from JSONL file."""
|
||||
import pandas as pd
|
||||
|
||||
if not self.dataset_path.endswith(".jsonl"):
|
||||
raise ValueError("Custom dataset must be a JSONL file")
|
||||
|
||||
try:
|
||||
df = pd.read_json(path_or_buf=self.dataset_path, lines=True)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to load JSONL file {self.dataset_path}: {e}")
|
||||
|
||||
if "prompt" not in df.columns:
|
||||
raise ValueError("JSONL file must contain a 'prompt' column")
|
||||
|
||||
self.data = df.to_dict("records")
|
||||
print(f"Loaded {len(self.data)} requests from {self.dataset_path}")
|
||||
|
||||
def _resolve_image_paths(self, item: dict) -> list[str] | None:
|
||||
# Handle local file paths
|
||||
if "image_paths" in item and item["image_paths"]:
|
||||
paths = item["image_paths"]
|
||||
if isinstance(paths, str):
|
||||
paths = [paths]
|
||||
|
||||
# Verify all paths exist
|
||||
valid_paths = []
|
||||
for path in paths:
|
||||
if os.path.exists(path):
|
||||
valid_paths.append(path)
|
||||
else:
|
||||
raise ValueError(f"Image file not found: {path}")
|
||||
|
||||
return valid_paths if valid_paths else None
|
||||
|
||||
# Handle URLs - download to temp files
|
||||
if "image_urls" in item and item["image_urls"]:
|
||||
urls = item["image_urls"]
|
||||
if isinstance(urls, str):
|
||||
urls = [urls]
|
||||
|
||||
downloaded_paths = []
|
||||
for url in urls:
|
||||
try:
|
||||
response = requests.get(url, timeout=30)
|
||||
response.raise_for_status()
|
||||
|
||||
# Create temp file with appropriate extension
|
||||
suffix = os.path.splitext(url)[1] or ".png"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
|
||||
tmp.write(response.content)
|
||||
downloaded_paths.append(tmp.name)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to download image from {url}: {e}")
|
||||
|
||||
return downloaded_paths if downloaded_paths else None
|
||||
|
||||
return None
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.data)
|
||||
|
||||
def __getitem__(self, idx: int) -> RequestFuncInput:
|
||||
item = self.data[idx]
|
||||
|
||||
# Extract fields with fallback to args defaults
|
||||
prompt = item["prompt"]
|
||||
width = item.get("width", self.args.width)
|
||||
height = item.get("height", self.args.height)
|
||||
num_inference_steps = item.get("num_inference_steps", self.args.num_inference_steps)
|
||||
seed = item.get("seed", self.args.seed)
|
||||
reserved_keys = {"prompt", "width", "height", "num_inference_steps", "seed", "image_paths", "image_urls"}
|
||||
extra_body = {k: v for k, v in item.items() if k not in reserved_keys}
|
||||
|
||||
# Handle image paths/URLs
|
||||
image_paths = self._resolve_image_paths(item)
|
||||
|
||||
return RequestFuncInput(
|
||||
prompt=prompt,
|
||||
api_url=self.api_url,
|
||||
model=self.model,
|
||||
seed=seed,
|
||||
image_paths=image_paths,
|
||||
extra_body=extra_body,
|
||||
width=width,
|
||||
height=height,
|
||||
num_inference_steps=num_inference_steps,
|
||||
)
|
||||
|
||||
def get_requests(self) -> list[RequestFuncInput]:
|
||||
"""Get all requests, cycling through data if num_prompts > dataset size."""
|
||||
num_requests = getattr(self.args, "num_prompts", len(self.data))
|
||||
|
||||
if num_requests <= len(self.data):
|
||||
return [self[i] for i in range(num_requests)]
|
||||
|
||||
# Cycle through the dataset to reach num_prompts
|
||||
requests = []
|
||||
for i in range(num_requests):
|
||||
idx = i % len(self.data)
|
||||
requests.append(self[idx])
|
||||
|
||||
return requests
|
||||
|
||||
|
||||
class RandomDataset(BaseDataset):
|
||||
def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool = False):
|
||||
super().__init__(args, api_url, model)
|
||||
@@ -585,7 +729,7 @@ class RandomDataset(BaseDataset):
|
||||
self._sampled_requests = None
|
||||
|
||||
# Random image generate
|
||||
if self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
|
||||
if self.args.task in ["i2v", "ti2v", "ti2i", "i2i", "it2i"]:
|
||||
self._random_image_path = self._generate_random_image_paths()
|
||||
else:
|
||||
self._random_image_path = None
|
||||
@@ -891,7 +1035,9 @@ def wait_for_service(base_url: str, timeout: int = 120) -> None:
|
||||
def _default_endpoint_for_task(task: str) -> str:
|
||||
if task in {"t2v", "i2v", "ti2v"}:
|
||||
return "/v1/videos"
|
||||
if task in {"t2i", "i2i", "ti2i"}:
|
||||
if task in {"i2i", "ti2i", "it2i"}:
|
||||
return "/v1/images/edits"
|
||||
if task == "t2i":
|
||||
return "/v1/chat/completions"
|
||||
raise ValueError(f"Unsupported task for endpoint resolution: {task}")
|
||||
|
||||
@@ -902,7 +1048,7 @@ async def benchmark(args):
|
||||
args.base_url = f"http://{args.host}:{args.port}"
|
||||
|
||||
VIDEO_TASKS = {"t2v", "i2v", "ti2v"}
|
||||
IMAGE_TASKS = {"t2i", "i2i", "ti2i"}
|
||||
IMAGE_TASKS = {"t2i", "i2i", "ti2i", "it2i"}
|
||||
|
||||
if args.task in VIDEO_TASKS:
|
||||
task_type = "2v"
|
||||
@@ -940,6 +1086,8 @@ async def benchmark(args):
|
||||
dataset = TraceDataset(args, api_url, args.model)
|
||||
elif args.dataset == "random":
|
||||
dataset = RandomDataset(args, api_url, args.model, args.enable_negative_prompt)
|
||||
elif args.dataset == "custom":
|
||||
dataset = CustomDataset(args, api_url, args.model)
|
||||
else:
|
||||
raise ValueError(f"Unknown dataset: {args.dataset}")
|
||||
|
||||
@@ -947,6 +1095,10 @@ async def benchmark(args):
|
||||
requests_list = dataset.get_requests()
|
||||
print(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.")
|
||||
|
||||
if args.endpoint == "/v1/images/edits":
|
||||
for req in requests_list:
|
||||
req.default_bot_task = args.bot_task
|
||||
|
||||
# Limit concurrency
|
||||
if args.max_concurrency is not None:
|
||||
semaphore = asyncio.Semaphore(args.max_concurrency)
|
||||
@@ -998,6 +1150,8 @@ async def benchmark(args):
|
||||
metrics["model"] = args.model
|
||||
metrics["dataset"] = args.dataset
|
||||
metrics["task"] = args.task
|
||||
if args.endpoint == "/v1/images/edits":
|
||||
metrics["bot_task"] = args.bot_task
|
||||
|
||||
print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))
|
||||
|
||||
@@ -1081,14 +1235,14 @@ if __name__ == "__main__":
|
||||
"--dataset",
|
||||
type=str,
|
||||
default="vbench",
|
||||
choices=["vbench", "trace", "random"],
|
||||
choices=["vbench", "trace", "random", "custom"],
|
||||
help="Dataset to use.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--task",
|
||||
type=str,
|
||||
default="t2v",
|
||||
choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "t2i"],
|
||||
choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "it2i", "t2i"],
|
||||
help="Task type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -1195,9 +1349,15 @@ if __name__ == "__main__":
|
||||
default=1,
|
||||
help=(
|
||||
"Number of synthetic input images to attach for image-conditioned tasks "
|
||||
"(i2v, ti2v, ti2i, i2i) when using random dataset."
|
||||
"(i2v, ti2v, ti2i, i2i, it2i) when using random dataset."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--bot-task",
|
||||
type=str,
|
||||
default="think",
|
||||
help=("bot_task form field for --endpoint /v1/images/edits (think, recaption, think_recaption, vanilla)."),
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
asyncio.run(benchmark(args))
|
||||
@@ -34,6 +34,16 @@ Video example (text-to-video):
|
||||
--height 720 --width 1280 \
|
||||
--num-frames 81 --num-inference-steps 40 --seed 42
|
||||
|
||||
LTX-2 example (text-to-video; audio output is generated but not scored by LPIPS):
|
||||
python benchmarks/diffusion/quantization_quality.py \
|
||||
--model Lightricks/LTX-Video-2 \
|
||||
--task t2v \
|
||||
--quantization fp8 int8 \
|
||||
--prompts \
|
||||
"A serene lakeside sunrise with mist over the water" \
|
||||
--height 704 --width 1216 \
|
||||
--num-frames 121 --num-inference-steps 40 --seed 42
|
||||
|
||||
Multiple quantization methods:
|
||||
python benchmarks/diffusion/quantization_quality.py \
|
||||
--model Tongyi-MAI/Z-Image-Turbo \
|
||||
@@ -204,10 +214,19 @@ def _generate_video(omni, args, prompt, seed):
|
||||
else:
|
||||
frames = inner
|
||||
elif hasattr(first, "images") and first.images:
|
||||
frames = first.images
|
||||
frames = first.images[0]
|
||||
else:
|
||||
raise ValueError("Could not extract video frames from output.")
|
||||
|
||||
# LTX-2 (and similar audio+video models) may surface a dict or (video, audio) tuple
|
||||
if isinstance(frames, dict):
|
||||
frames = frames.get("video") or frames.get("frames")
|
||||
elif isinstance(frames, tuple) and len(frames) == 2:
|
||||
frames = frames[0]
|
||||
|
||||
if frames is None:
|
||||
raise ValueError("Could not extract video frames from output.")
|
||||
|
||||
if isinstance(frames, torch.Tensor):
|
||||
video = frames.detach().cpu()
|
||||
if video.dim() == 5:
|
||||
@@ -225,9 +244,11 @@ def _generate_video(omni, args, prompt, seed):
|
||||
return frames_array, elapsed, peak_mem
|
||||
|
||||
|
||||
def _unload_omni(omni):
|
||||
"""Delete Omni instance and free GPU memory."""
|
||||
del omni
|
||||
def _free_gpu_memory():
|
||||
"""Force GC and release cached GPU memory.
|
||||
|
||||
Must be called AFTER the caller has dropped (i.e., via `del`)
|
||||
every reference to the Omni instance"""
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.accelerator.empty_cache()
|
||||
@@ -267,7 +288,9 @@ def run_benchmark(args):
|
||||
|
||||
bl_avg_time = np.mean([v[1] for v in baseline_outputs.values()])
|
||||
bl_mem = baseline_outputs[prompts[0]][2] # use first prompt's memory
|
||||
_unload_omni(omni_bl)
|
||||
omni_bl.shutdown()
|
||||
del omni_bl
|
||||
_free_gpu_memory()
|
||||
|
||||
# Save baseline outputs
|
||||
bl_dir = output_dir / "baseline"
|
||||
@@ -307,7 +330,9 @@ def run_benchmark(args):
|
||||
|
||||
qt_avg_time = np.mean([v[1] for v in qt_outputs.values()])
|
||||
qt_mem = qt_outputs[prompts[0]][2]
|
||||
_unload_omni(omni_qt)
|
||||
omni_qt.shutdown()
|
||||
del omni_qt
|
||||
_free_gpu_memory()
|
||||
|
||||
# Save quantized outputs
|
||||
qt_dir = output_dir / config_label.replace(" ", "_")
|
||||
|
||||
@@ -40,7 +40,7 @@ vllm bench serve --omni \
|
||||
--num-prompts 20 --num-warmups 2 \
|
||||
--extra-body '{"task_type":"Base"}' \
|
||||
--max-concurrency 1 --request-rate inf \
|
||||
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
|
||||
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \
|
||||
--save-result --result-dir ./results
|
||||
```
|
||||
|
||||
@@ -58,7 +58,7 @@ vllm bench serve --omni \
|
||||
--num-prompts 20 --num-warmups 2 \
|
||||
--extra-body '{"voice":"Vivian","language":"English","task_type":"CustomVoice"}' \
|
||||
--max-concurrency 1 --request-rate inf \
|
||||
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
|
||||
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \
|
||||
--save-result --result-dir ./results
|
||||
```
|
||||
|
||||
@@ -76,10 +76,28 @@ vllm bench serve --omni \
|
||||
--num-prompts 20 --num-warmups 2 \
|
||||
--extra-body '{"task_type":"VoiceDesign","language":"English"}' \
|
||||
--max-concurrency 1 --request-rate inf \
|
||||
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
|
||||
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \
|
||||
--save-result --result-dir ./results
|
||||
```
|
||||
|
||||
#### Streaming continuity (`audio_underrun`)
|
||||
|
||||
`audio_underrun` (seconds) is the per-request worst-case buffer deficit
|
||||
under a realtime-rate player simulation. The audio-speech backend records
|
||||
chunk arrival times during the SSE stream and surfaces:
|
||||
|
||||
- `Mean / Median / P{50,99} AUDIO_UNDERRUN (s)` per the standard percentile
|
||||
output - any positive value means at least one inter-chunk gap was longer
|
||||
than the chunk's audio duration.
|
||||
- `Streaming continuity OK rate` - fraction of requests whose worst-case
|
||||
underrun stayed under the threshold (default 100 ms, the commonly cited
|
||||
"audible gap" budget). Override the threshold with
|
||||
`VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S=<float>`.
|
||||
|
||||
This captures the failure mode where `RTF_p50 < 1` (server keeps up in
|
||||
aggregate) but per-stream chunk arrival is bursty enough that listeners
|
||||
still hear gaps - common at high concurrency on streaming TTS pipelines.
|
||||
|
||||
#### Add WER / SIM / UTMOS to any of the above
|
||||
|
||||
Append `--seed-tts-wer-eval` (and optionally `SEED_TTS_EVAL_DEVICE=cuda:0`
|
||||
@@ -215,7 +233,6 @@ sentinel for regressions in this area.
|
||||
benchmarks/tts/
|
||||
├── README.md (this file)
|
||||
├── bench_tts.py CLI — serve-mode benchmark driver
|
||||
├── bench_voxcpm_offline.py CLI — offline VoxCPM benchmark (sync + streaming)
|
||||
├── plot_results.py Generate per-task / per-concurrency curves
|
||||
└── model_configs.yaml Model registry (supported tasks + extra body)
|
||||
```
|
||||
|
||||
@@ -53,6 +53,8 @@ _TASK_TO_DATASET: dict[str, str] = {
|
||||
"voice_clone": "seed-tts",
|
||||
"default_voice": "seed-tts-text",
|
||||
"voice_design": "seed-tts-design",
|
||||
"dialogue": "ttsd",
|
||||
"sound_effect": "sound-effect",
|
||||
}
|
||||
|
||||
# Default design dataset path (bundled with the repo)
|
||||
@@ -81,6 +83,7 @@ def build_bench_args(
|
||||
output_dir: str | None,
|
||||
result_filename: str | None,
|
||||
extra_cli_args: list[str],
|
||||
output_len: int | None = None,
|
||||
) -> list[str]:
|
||||
"""Build the ``vllm bench serve --omni`` command for one (task, concurrency) run."""
|
||||
dataset_name = _TASK_TO_DATASET[task]
|
||||
@@ -118,7 +121,7 @@ def build_bench_args(
|
||||
"--num-warmups",
|
||||
"2",
|
||||
"--percentile-metrics",
|
||||
"ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun",
|
||||
]
|
||||
|
||||
if resolved_dataset_path:
|
||||
@@ -130,12 +133,18 @@ def build_bench_args(
|
||||
if task_extra_body:
|
||||
cmd += ["--extra-body", json.dumps(task_extra_body, separators=(",", ":"))]
|
||||
|
||||
if model_cfg.get("trust_remote_code"):
|
||||
cmd.append("--trust-remote-code")
|
||||
|
||||
if concurrency is not None:
|
||||
cmd += ["--max-concurrency", str(concurrency), "--request-rate", "inf"]
|
||||
|
||||
if wer_eval:
|
||||
cmd.append("--seed-tts-wer-eval")
|
||||
|
||||
if output_len is not None:
|
||||
cmd += ["--hf-output-len", str(output_len)]
|
||||
|
||||
if output_dir or result_filename:
|
||||
out_dir = output_dir or "."
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
@@ -235,6 +244,15 @@ def main() -> None:
|
||||
parser.add_argument("--host", default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--model-configs", default=str(_DEFAULT_MODEL_CONFIGS), help="Path to model_configs.yaml")
|
||||
parser.add_argument(
|
||||
"--output-len",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Override per-request max_new_tokens forwarded to the server "
|
||||
"(seed-tts datasets default to 2048; the MOSS-TTS full talker takes "
|
||||
"this as 'max_new_frames' so smaller values cap audio length per "
|
||||
"request and keep bench runtime tractable).",
|
||||
)
|
||||
parser.add_argument("extra", nargs=argparse.REMAINDER, help="Extra args passed directly to vllm bench serve")
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -289,6 +307,7 @@ def main() -> None:
|
||||
output_dir=args.output_dir,
|
||||
result_filename=result_filename,
|
||||
extra_cli_args=args.extra or [],
|
||||
output_len=args.output_len,
|
||||
)
|
||||
result = run_one_benchmark(cmd)
|
||||
if result is not None:
|
||||
|
||||
@@ -1,922 +0,0 @@
|
||||
"""Offline VoxCPM benchmark for vLLM Omni.
|
||||
|
||||
Supports both:
|
||||
- sync one-shot (Omni.generate)
|
||||
- streaming (AsyncOmni.generate with async_chunk config)
|
||||
- text-only synthesis
|
||||
- voice cloning
|
||||
- text/clone batch inputs from txt or jsonl
|
||||
|
||||
Usage::
|
||||
|
||||
# Sync (default voice)
|
||||
python benchmarks/tts/bench_voxcpm_offline.py \\
|
||||
--model /path/to/VoxCPM \\
|
||||
--text "Hello world" \\
|
||||
--output-dir results/audio/
|
||||
|
||||
# Streaming (async_chunk)
|
||||
python benchmarks/tts/bench_voxcpm_offline.py \\
|
||||
--model /path/to/VoxCPM \\
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \\
|
||||
--txt-prompts prompts.txt \\
|
||||
--output-dir results/audio/
|
||||
|
||||
# Voice cloning batch via JSONL
|
||||
python benchmarks/tts/bench_voxcpm_offline.py \\
|
||||
--model /path/to/VoxCPM \\
|
||||
--jsonl-prompts prompts.jsonl \\
|
||||
--output-dir results/audio/
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import AsyncOmni, Omni
|
||||
|
||||
|
||||
def _find_repo_root(start: Path) -> Path:
|
||||
"""Walk up from ``start`` until a repo marker is found.
|
||||
|
||||
Falls back to ``parents[2]`` for backwards compatibility if no marker hits
|
||||
(which can only happen in unusual checkouts — the tree should always have
|
||||
pyproject.toml + vllm_omni/ at the top level).
|
||||
"""
|
||||
for candidate in [start, *start.parents]:
|
||||
if (candidate / "pyproject.toml").is_file() and (candidate / "vllm_omni").is_dir():
|
||||
return candidate
|
||||
return start.parents[2]
|
||||
|
||||
|
||||
REPO_ROOT = _find_repo_root(Path(__file__).resolve())
|
||||
DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml"
|
||||
DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "deploy" / "voxcpm.yaml"
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PromptSpec:
|
||||
text: str
|
||||
label: str
|
||||
ref_audio: str | None = None
|
||||
ref_text: str | None = None
|
||||
|
||||
|
||||
def _require_soundfile():
|
||||
try:
|
||||
import soundfile as sf # type: ignore
|
||||
except ModuleNotFoundError as exc:
|
||||
raise RuntimeError(
|
||||
"soundfile is required to write VoxCPM benchmark WAV outputs. Install it with: pip install soundfile"
|
||||
) from exc
|
||||
return sf
|
||||
|
||||
|
||||
def _build_prompt(
|
||||
args,
|
||||
*,
|
||||
text: str,
|
||||
ref_audio: str | None = None,
|
||||
ref_text: str | None = None,
|
||||
global_request_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
additional_information: dict[str, list[Any]] = {
|
||||
"text": [text],
|
||||
"cfg_value": [args.cfg_value],
|
||||
"inference_timesteps": [args.inference_timesteps],
|
||||
"min_len": [args.min_len],
|
||||
"max_new_tokens": [args.max_new_tokens],
|
||||
}
|
||||
if args.streaming_prefix_len is not None:
|
||||
additional_information["streaming_prefix_len"] = [args.streaming_prefix_len]
|
||||
|
||||
if ref_audio:
|
||||
additional_information["ref_audio"] = [ref_audio]
|
||||
if ref_text:
|
||||
additional_information["ref_text"] = [ref_text]
|
||||
if global_request_id is not None:
|
||||
additional_information["global_request_id"] = [global_request_id]
|
||||
|
||||
return {
|
||||
"prompt_token_ids": [1],
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
|
||||
|
||||
def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor:
|
||||
audio = mm.get("audio", mm.get("model_outputs"))
|
||||
if audio is None:
|
||||
raise ValueError("No audio output found in multimodal output.")
|
||||
if isinstance(audio, list):
|
||||
parts = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio]
|
||||
audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0)
|
||||
if not isinstance(audio, torch.Tensor):
|
||||
audio = torch.as_tensor(audio)
|
||||
return audio.float().cpu().reshape(-1)
|
||||
|
||||
|
||||
def _extract_sample_rate(mm: dict[str, Any]) -> int:
|
||||
sr_raw = mm.get("sr", 24000)
|
||||
if isinstance(sr_raw, list) and sr_raw:
|
||||
sr_raw = sr_raw[-1]
|
||||
if hasattr(sr_raw, "item"):
|
||||
return int(sr_raw.item())
|
||||
return int(sr_raw)
|
||||
|
||||
|
||||
def _emit_offline_metrics(
|
||||
*,
|
||||
request_id: str,
|
||||
elapsed_s: float,
|
||||
first_audio_elapsed: float | None,
|
||||
audio_duration_s: float,
|
||||
) -> None:
|
||||
metrics = {
|
||||
"request_id": request_id,
|
||||
"ttfp_ms": round(first_audio_elapsed * 1000.0, 3) if first_audio_elapsed is not None else None,
|
||||
"audio_duration_s": round(audio_duration_s, 6),
|
||||
"rtf": round(elapsed_s / audio_duration_s, 6) if audio_duration_s > 0 else None,
|
||||
}
|
||||
print(f"[OfflineMetrics] {metrics}")
|
||||
|
||||
|
||||
def _write_audio_tensor(output_path: Path, audio_tensor: Any, sample_rate: int) -> None:
|
||||
sf = _require_soundfile()
|
||||
if isinstance(audio_tensor, torch.Tensor):
|
||||
audio_np = audio_tensor.float().cpu().clamp(-1.0, 1.0).numpy()
|
||||
else:
|
||||
audio_np = torch.as_tensor(audio_tensor).float().cpu().clamp(-1.0, 1.0).numpy()
|
||||
sf.write(
|
||||
output_path,
|
||||
audio_np,
|
||||
sample_rate,
|
||||
format="WAV",
|
||||
subtype="PCM_16",
|
||||
)
|
||||
|
||||
|
||||
def _save_wav(mm: dict[str, Any], output_dir: Path, request_id: str) -> Path:
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
output_path = output_dir / f"output_{request_id}.wav"
|
||||
_write_audio_tensor(output_path, _extract_audio_tensor(mm), _extract_sample_rate(mm))
|
||||
return output_path
|
||||
|
||||
|
||||
def _iter_request_multimodal_outputs(request_output: Any):
|
||||
outputs = getattr(request_output, "outputs", None)
|
||||
if outputs:
|
||||
for output in outputs:
|
||||
mm = getattr(output, "multimodal_output", None)
|
||||
if isinstance(mm, dict):
|
||||
yield mm
|
||||
|
||||
mm = getattr(request_output, "multimodal_output", None)
|
||||
if isinstance(mm, dict):
|
||||
yield mm
|
||||
|
||||
|
||||
def _read_non_empty_lines(path: str) -> list[str]:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
return [line.strip() for line in f if line.strip()]
|
||||
|
||||
|
||||
def _load_prompt_specs(args) -> list[PromptSpec]:
|
||||
specs: list[PromptSpec] = []
|
||||
|
||||
if args.txt_prompts is not None:
|
||||
texts = _read_non_empty_lines(args.txt_prompts)
|
||||
if not texts:
|
||||
raise ValueError(f"No prompts found in {args.txt_prompts}")
|
||||
for idx, text in enumerate(texts, start=1):
|
||||
specs.append(
|
||||
PromptSpec(
|
||||
text=text,
|
||||
label=f"item{idx:03d}",
|
||||
ref_audio=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
)
|
||||
)
|
||||
return specs
|
||||
|
||||
if args.jsonl_prompts is not None:
|
||||
with open(args.jsonl_prompts, encoding="utf-8") as f:
|
||||
for line_no, raw_line in enumerate(f, start=1):
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
item = json.loads(line)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"{args.jsonl_prompts}:{line_no} is not valid JSON: {exc}") from exc
|
||||
if not isinstance(item, dict):
|
||||
raise ValueError(f"{args.jsonl_prompts}:{line_no} must be a JSON object")
|
||||
|
||||
text = item.get("text")
|
||||
if not isinstance(text, str) or not text.strip():
|
||||
raise ValueError(f"{args.jsonl_prompts}:{line_no} requires non-empty string field 'text'")
|
||||
|
||||
ref_audio = item.get("ref_audio", args.ref_audio)
|
||||
ref_text = item.get("ref_text", args.ref_text)
|
||||
if (ref_audio is None) != (ref_text is None):
|
||||
raise ValueError(
|
||||
f"{args.jsonl_prompts}:{line_no} must provide both 'ref_audio' and 'ref_text' together"
|
||||
)
|
||||
|
||||
specs.append(
|
||||
PromptSpec(
|
||||
text=text.strip(),
|
||||
label=f"item{len(specs) + 1:03d}",
|
||||
ref_audio=ref_audio,
|
||||
ref_text=ref_text,
|
||||
)
|
||||
)
|
||||
|
||||
if not specs:
|
||||
raise ValueError(f"No prompts found in {args.jsonl_prompts}")
|
||||
return specs
|
||||
|
||||
specs.append(
|
||||
PromptSpec(
|
||||
text=args.text,
|
||||
label="item001",
|
||||
ref_audio=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
)
|
||||
)
|
||||
return specs
|
||||
|
||||
|
||||
def _build_prompt_for_spec(args, spec: PromptSpec, *, global_request_id: str | None = None) -> dict[str, Any]:
|
||||
return _build_prompt(
|
||||
args,
|
||||
text=spec.text,
|
||||
ref_audio=spec.ref_audio,
|
||||
ref_text=spec.ref_text,
|
||||
global_request_id=global_request_id,
|
||||
)
|
||||
|
||||
|
||||
def _count_voice_clone_prompts(prompt_specs: list[PromptSpec]) -> int:
|
||||
return sum(1 for spec in prompt_specs if spec.ref_audio is not None)
|
||||
|
||||
|
||||
def _get_warmup_specs(prompt_specs: list[PromptSpec]) -> list[PromptSpec]:
|
||||
return prompt_specs[:1]
|
||||
|
||||
|
||||
def _extract_stream_finished(stage_output: Any) -> bool:
|
||||
request_output = getattr(stage_output, "request_output", None)
|
||||
request_finished = getattr(request_output, "finished", None)
|
||||
if request_finished is not None:
|
||||
return bool(request_finished)
|
||||
return bool(getattr(stage_output, "finished", False))
|
||||
|
||||
|
||||
def _build_profiled_stage_config(
|
||||
stage_configs_path: str,
|
||||
profiler_dir: str,
|
||||
) -> str:
|
||||
stage_config_path = Path(stage_configs_path)
|
||||
yaml_text = stage_config_path.read_text(encoding="utf-8")
|
||||
injected_lines: list[str] = []
|
||||
injected_count = 0
|
||||
|
||||
for line in yaml_text.splitlines():
|
||||
injected_lines.append(line)
|
||||
if line.strip() != "engine_args:":
|
||||
continue
|
||||
indent = line[: len(line) - len(line.lstrip())]
|
||||
child_indent = indent + " "
|
||||
grandchild_indent = child_indent + " "
|
||||
injected_lines.extend(
|
||||
[
|
||||
f"{child_indent}profiler_config:",
|
||||
f'{grandchild_indent}profiler: "torch"',
|
||||
f'{grandchild_indent}torch_profiler_dir: "{profiler_dir}"',
|
||||
f"{grandchild_indent}torch_profiler_with_stack: true",
|
||||
]
|
||||
)
|
||||
injected_count += 1
|
||||
|
||||
if injected_count == 0:
|
||||
raise ValueError(f"No engine_args block found in stage config: {stage_configs_path}")
|
||||
|
||||
tmp = tempfile.NamedTemporaryFile(
|
||||
mode="w",
|
||||
encoding="utf-8",
|
||||
delete=False,
|
||||
suffix=".yaml",
|
||||
prefix=f"{stage_config_path.stem}_profile_",
|
||||
)
|
||||
tmp.write("\n".join(injected_lines) + "\n")
|
||||
tmp.close()
|
||||
return tmp.name
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(
|
||||
description="Offline split-stage VoxCPM inference with vLLM Omni (auto sync/streaming by stage config)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default=os.environ.get("VOXCPM_MODEL"),
|
||||
help="Local VoxCPM model directory. Defaults to $VOXCPM_MODEL.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
type=str,
|
||||
default="This is a split-stage VoxCPM synthesis example running on vLLM Omni.",
|
||||
help="Text to synthesize. Ignored when --txt-prompts or --jsonl-prompts is used.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--txt-prompts",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a .txt file with one synthesis text per line.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--jsonl-prompts",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Path to a .jsonl file. Each line must contain at least {'text': ...}; "
|
||||
"clone rows can also set ref_audio/ref_text, and ref_text must be the "
|
||||
"real transcript of ref_audio."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Optional reference audio path for voice cloning. With --txt-prompts, "
|
||||
"the same reference is applied to every line."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-text",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Real transcript of the reference audio. Placeholder text or mismatched "
|
||||
"text will usually produce noisy/electronic clone audio."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=str(DEFAULT_STAGE_SYNC),
|
||||
help="Stage config YAML path. Routing is selected only from this path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg-value",
|
||||
type=float,
|
||||
default=2.0,
|
||||
help="Classifier-free guidance value for VoxCPM.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--inference-timesteps",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Number of inference timesteps.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min-len",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Minimum generated token length.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-new-tokens",
|
||||
type=int,
|
||||
default=4096,
|
||||
help="Maximum generated token length.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--streaming-prefix-len",
|
||||
type=int,
|
||||
default=None,
|
||||
help="VoxCPM streaming window (optional, streaming mode only).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory for output WAV files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-init-timeout",
|
||||
type=int,
|
||||
default=600,
|
||||
help="Stage initialization timeout in seconds.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
dest="log_stats",
|
||||
action="store_true",
|
||||
help="Enable vLLM Omni stats logging.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-log-stats",
|
||||
dest="log_stats",
|
||||
action="store_false",
|
||||
help="Disable vLLM Omni stats logging.",
|
||||
)
|
||||
parser.set_defaults(log_stats=True)
|
||||
parser.add_argument(
|
||||
"--num-runs",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of full inference runs (same prompt each time). Default 1.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup-runs",
|
||||
type=int,
|
||||
default=0,
|
||||
help=(
|
||||
"Optional number of warmup passes before measured runs. Warmup uses only "
|
||||
"the first prompt and does not save outputs."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-profiler",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Enable torch profiler for the configured stages. A temporary profiled "
|
||||
"stage config is generated automatically."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory for profiler traces. Defaults to <output-dir>/profiler when profiling is enabled.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-stages",
|
||||
type=int,
|
||||
nargs="*",
|
||||
default=None,
|
||||
help="Optional stage ids to profile. Defaults to all stages that have profiler_config.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-wait-seconds",
|
||||
type=float,
|
||||
default=30.0,
|
||||
help="Seconds to wait after stop_profile for trace files to flush.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.model:
|
||||
parser.error("--model is required unless $VOXCPM_MODEL is set")
|
||||
if args.txt_prompts is not None and args.jsonl_prompts is not None:
|
||||
parser.error("--txt-prompts and --jsonl-prompts are mutually exclusive")
|
||||
if (args.ref_audio is None) != (args.ref_text is None):
|
||||
parser.error("--ref-audio and --ref-text must be provided together")
|
||||
if args.num_runs < 1:
|
||||
parser.error("--num-runs must be >= 1")
|
||||
if args.warmup_runs < 0:
|
||||
parser.error("--warmup-runs must be >= 0")
|
||||
if args.output_dir is None:
|
||||
args.output_dir = (
|
||||
"output_audio_streaming" if _is_streaming_stage_config(args.stage_configs_path) else "output_audio"
|
||||
)
|
||||
if args.enable_profiler and args.profiler_dir is None:
|
||||
args.profiler_dir = str(Path(args.output_dir) / "profiler")
|
||||
try:
|
||||
args.prompt_specs = _load_prompt_specs(args)
|
||||
except ValueError as exc:
|
||||
parser.error(str(exc))
|
||||
|
||||
return args
|
||||
|
||||
|
||||
def _is_streaming_stage_config(stage_configs_path: str) -> bool:
|
||||
cfg_name = Path(stage_configs_path).name.lower()
|
||||
return "async_chunk" in cfg_name
|
||||
|
||||
|
||||
async def _collect_streaming_audio(
|
||||
omni: AsyncOmni,
|
||||
args: Any,
|
||||
spec: PromptSpec,
|
||||
request_id: str,
|
||||
*,
|
||||
phase_label: str,
|
||||
prompt_index: int,
|
||||
prompt_count: int,
|
||||
print_prompt: bool = False,
|
||||
) -> tuple[torch.Tensor, int, float, float | None]:
|
||||
prompt = _build_prompt_for_spec(args, spec, global_request_id=request_id)
|
||||
delta_chunks: list[torch.Tensor] = []
|
||||
sample_rate = 24000
|
||||
chunk_i = 0
|
||||
prev_total_samples = 0
|
||||
t_start = time.perf_counter()
|
||||
first_audio_elapsed: float | None = None
|
||||
|
||||
if print_prompt:
|
||||
print(f"---prompt---:{prompt}")
|
||||
|
||||
async for stage_output in omni.generate(prompt, request_id=request_id):
|
||||
mm = getattr(stage_output, "multimodal_output", None)
|
||||
if not isinstance(mm, dict):
|
||||
ro = getattr(stage_output, "request_output", None)
|
||||
if ro is None:
|
||||
continue
|
||||
mm = getattr(ro, "multimodal_output", None)
|
||||
if not isinstance(mm, dict) and getattr(ro, "outputs", None):
|
||||
seq = ro.outputs[0]
|
||||
mm = getattr(seq, "multimodal_output", None)
|
||||
if not isinstance(mm, dict):
|
||||
continue
|
||||
sample_rate = _extract_sample_rate(mm)
|
||||
try:
|
||||
w = _extract_audio_tensor(mm)
|
||||
n = int(w.numel())
|
||||
if n == 0:
|
||||
continue
|
||||
finished = _extract_stream_finished(stage_output)
|
||||
if n > prev_total_samples:
|
||||
delta = w.reshape(-1)[prev_total_samples:]
|
||||
prev_total_samples = n
|
||||
elif finished and n == prev_total_samples:
|
||||
delta = w.reshape(-1)[:0]
|
||||
else:
|
||||
delta = w.reshape(-1)
|
||||
prev_total_samples += int(delta.numel())
|
||||
if int(delta.numel()) > 0:
|
||||
delta_chunks.append(delta)
|
||||
if first_audio_elapsed is None and int(delta.numel()) > 0:
|
||||
first_audio_elapsed = time.perf_counter() - t_start
|
||||
logger.info(
|
||||
"%s prompt=%d/%d chunk=%d delta_samples=%d buf_len=%d finished=%s",
|
||||
phase_label,
|
||||
prompt_index + 1,
|
||||
prompt_count,
|
||||
chunk_i,
|
||||
int(delta.numel()),
|
||||
n,
|
||||
finished,
|
||||
)
|
||||
chunk_i += 1
|
||||
except ValueError:
|
||||
if not _extract_stream_finished(stage_output):
|
||||
logger.debug("skip non-audio partial output chunk=%d", chunk_i)
|
||||
|
||||
if not delta_chunks:
|
||||
raise RuntimeError("No audio chunks received; check stage config and logs.")
|
||||
|
||||
audio_cat = torch.cat([c.reshape(-1) for c in delta_chunks], dim=0)
|
||||
elapsed = time.perf_counter() - t_start
|
||||
return audio_cat, sample_rate, elapsed, first_audio_elapsed
|
||||
|
||||
|
||||
async def _abort_streaming_residual_work(
|
||||
omni: AsyncOmni,
|
||||
request_id: str,
|
||||
*,
|
||||
settle_seconds: float = 0.1,
|
||||
) -> None:
|
||||
"""Stop any late stage-0 work once the final audio has been collected."""
|
||||
await omni.engine.abort_async([request_id])
|
||||
if settle_seconds > 0:
|
||||
await asyncio.sleep(settle_seconds)
|
||||
|
||||
|
||||
async def _run_streaming_single(
|
||||
omni: AsyncOmni,
|
||||
args: Any,
|
||||
spec: PromptSpec,
|
||||
output_dir: Path,
|
||||
request_id: str,
|
||||
*,
|
||||
run_index: int,
|
||||
num_runs: int,
|
||||
prompt_index: int,
|
||||
prompt_count: int,
|
||||
) -> Path:
|
||||
audio_cat, sample_rate, elapsed, first_audio_elapsed = await _collect_streaming_audio(
|
||||
omni,
|
||||
args,
|
||||
spec,
|
||||
request_id,
|
||||
phase_label=f"run={run_index + 1}/{num_runs}",
|
||||
prompt_index=prompt_index,
|
||||
prompt_count=prompt_count,
|
||||
print_prompt=(run_index == 0 and prompt_index == 0),
|
||||
)
|
||||
await _abort_streaming_residual_work(omni, request_id)
|
||||
output_path = output_dir / f"output_run{run_index + 1}_{spec.label}.wav"
|
||||
_write_audio_tensor(output_path, audio_cat, sample_rate)
|
||||
audio_duration_s = float(audio_cat.numel()) / float(sample_rate) if sample_rate > 0 else 0.0
|
||||
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
|
||||
rtf_text = f", rtf={elapsed / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
|
||||
print(
|
||||
f"Saved (streaming) run {run_index + 1}/{num_runs}, "
|
||||
f"prompt {prompt_index + 1}/{prompt_count}: {output_path} ({elapsed:.2f}s{ttfp_text}{rtf_text})"
|
||||
)
|
||||
_emit_offline_metrics(
|
||||
request_id=request_id,
|
||||
elapsed_s=elapsed,
|
||||
first_audio_elapsed=first_audio_elapsed,
|
||||
audio_duration_s=audio_duration_s,
|
||||
)
|
||||
return output_path
|
||||
|
||||
|
||||
async def _run_streaming_warmup(args, omni: AsyncOmni) -> None:
|
||||
if args.warmup_runs == 0:
|
||||
return
|
||||
|
||||
warmup_specs = _get_warmup_specs(args.prompt_specs)
|
||||
print(
|
||||
f"Warmup: {args.warmup_runs} run(s) using the first prompt "
|
||||
f"({len(warmup_specs)} prompt(s)); outputs will be discarded."
|
||||
)
|
||||
for warmup_index in range(args.warmup_runs):
|
||||
t_warmup = time.perf_counter()
|
||||
tasks = []
|
||||
request_ids: list[str] = []
|
||||
for prompt_index, spec in enumerate(warmup_specs):
|
||||
request_id = f"warmup_stream_{warmup_index + 1}_{spec.label}_{uuid.uuid4().hex[:8]}"
|
||||
request_ids.append(request_id)
|
||||
tasks.append(
|
||||
_collect_streaming_audio(
|
||||
omni,
|
||||
args,
|
||||
spec,
|
||||
request_id,
|
||||
phase_label=f"warmup={warmup_index + 1}/{args.warmup_runs}",
|
||||
prompt_index=prompt_index,
|
||||
prompt_count=len(warmup_specs),
|
||||
)
|
||||
)
|
||||
results = await asyncio.gather(*tasks)
|
||||
for request_id in request_ids:
|
||||
await _abort_streaming_residual_work(omni, request_id)
|
||||
total_samples = sum(int(audio.numel()) for audio, _, _, _ in results)
|
||||
warmup_ttfps = [ttfp for _, _, _, ttfp in results if ttfp is not None]
|
||||
ttfp_text = f", ttfp={min(warmup_ttfps):.2f}s" if warmup_ttfps else ""
|
||||
print(
|
||||
f"Warmup (streaming) {warmup_index + 1}/{args.warmup_runs} finished: "
|
||||
f"{len(results)} prompt(s), {total_samples} sample(s) "
|
||||
f"({time.perf_counter() - t_warmup:.2f}s{ttfp_text})"
|
||||
)
|
||||
|
||||
|
||||
async def _run_streaming(args) -> list[Path]:
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
omni = AsyncOmni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
await _run_streaming_warmup(args, omni)
|
||||
profiler_started = False
|
||||
if args.enable_profiler:
|
||||
profile_prefix = f"voxcpm_streaming_{int(time.time())}"
|
||||
stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured"
|
||||
print(f"Starting profiler (streaming): stages={stages_text}, dir={args.profiler_dir}")
|
||||
await omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages)
|
||||
profiler_started = True
|
||||
t_total = time.perf_counter()
|
||||
total_elapsed = 0.0
|
||||
paths: list[Path] = []
|
||||
prompt_specs: list[PromptSpec] = args.prompt_specs
|
||||
try:
|
||||
for run in range(args.num_runs):
|
||||
for prompt_index, spec in enumerate(prompt_specs):
|
||||
request_id = f"stream_{run + 1}_{spec.label}_{uuid.uuid4().hex[:8]}"
|
||||
paths.append(
|
||||
await _run_streaming_single(
|
||||
omni,
|
||||
args,
|
||||
spec,
|
||||
output_dir,
|
||||
request_id,
|
||||
run_index=run,
|
||||
num_runs=args.num_runs,
|
||||
prompt_index=prompt_index,
|
||||
prompt_count=len(prompt_specs),
|
||||
)
|
||||
)
|
||||
total_elapsed = time.perf_counter() - t_total
|
||||
finally:
|
||||
if profiler_started:
|
||||
print("Stopping profiler (streaming)...")
|
||||
await omni.stop_profile(stages=args.profiler_stages)
|
||||
if args.profiler_wait_seconds > 0:
|
||||
print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...")
|
||||
await asyncio.sleep(args.profiler_wait_seconds)
|
||||
|
||||
print(
|
||||
f"All streaming runs finished: {args.num_runs} run(s), "
|
||||
f"{len(prompt_specs)} prompt(s), {len(paths)} file(s) in {total_elapsed:.2f}s total"
|
||||
)
|
||||
return paths
|
||||
|
||||
|
||||
def _run_sync(args) -> list[Path]:
|
||||
output_dir = Path(args.output_dir)
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
def _run_sync_single(
|
||||
spec: PromptSpec,
|
||||
*,
|
||||
request_prefix: str,
|
||||
save_outputs: bool,
|
||||
run_index: int | None = None,
|
||||
) -> tuple[list[Path], int, float | None, float, float, str]:
|
||||
global_request_id = f"{request_prefix}_{spec.label}"
|
||||
prompt = _build_prompt_for_spec(args, spec, global_request_id=global_request_id)
|
||||
if save_outputs and run_index == 0 and spec.label == "item001":
|
||||
print(f"---prompt---:{prompt}")
|
||||
|
||||
saved_paths: list[Path] = []
|
||||
output_count = 0
|
||||
first_audio_elapsed: float | None = None
|
||||
total_audio_duration_s = 0.0
|
||||
metrics_request_id = global_request_id
|
||||
t_start = time.perf_counter()
|
||||
for stage_outputs in omni.generate(prompt):
|
||||
request_output = stage_outputs.request_output
|
||||
if request_output is None:
|
||||
continue
|
||||
request_output_id = getattr(request_output, "request_id", None)
|
||||
if isinstance(request_output_id, str) and request_output_id:
|
||||
metrics_request_id = request_output_id
|
||||
for j, mm in enumerate(_iter_request_multimodal_outputs(request_output)):
|
||||
output_count += 1
|
||||
if first_audio_elapsed is None:
|
||||
try:
|
||||
audio_tensor = _extract_audio_tensor(mm)
|
||||
if int(audio_tensor.numel()) > 0:
|
||||
first_audio_elapsed = time.perf_counter() - t_start
|
||||
total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm))
|
||||
except ValueError:
|
||||
pass
|
||||
else:
|
||||
try:
|
||||
audio_tensor = _extract_audio_tensor(mm)
|
||||
total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm))
|
||||
except ValueError:
|
||||
pass
|
||||
if not save_outputs:
|
||||
continue
|
||||
save_stem = f"run{run_index + 1}_{spec.label}" if j == 0 else f"run{run_index + 1}_{spec.label}_{j}"
|
||||
saved_paths.append(_save_wav(mm, output_dir, save_stem))
|
||||
|
||||
if output_count == 0:
|
||||
raise RuntimeError("No output from Omni.generate")
|
||||
elapsed_s = time.perf_counter() - t_start
|
||||
return saved_paths, output_count, first_audio_elapsed, elapsed_s, total_audio_duration_s, metrics_request_id
|
||||
|
||||
if args.warmup_runs:
|
||||
warmup_specs = _get_warmup_specs(args.prompt_specs)
|
||||
print(
|
||||
f"Warmup: {args.warmup_runs} run(s) using the first prompt "
|
||||
f"({len(warmup_specs)} prompt(s)); outputs will be discarded."
|
||||
)
|
||||
for warmup_index in range(args.warmup_runs):
|
||||
t_warmup = time.perf_counter()
|
||||
_, output_count, first_audio_elapsed, elapsed_s, audio_duration_s, _ = _run_sync_single(
|
||||
warmup_specs[0],
|
||||
request_prefix=f"warmup_sync{warmup_index + 1}",
|
||||
save_outputs=False,
|
||||
)
|
||||
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
|
||||
rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
|
||||
print(
|
||||
f"Warmup (sync) {warmup_index + 1}/{args.warmup_runs} finished: "
|
||||
f"{output_count} output(s) ({time.perf_counter() - t_warmup:.2f}s{ttfp_text}{rtf_text})"
|
||||
)
|
||||
|
||||
profiler_started = False
|
||||
if args.enable_profiler:
|
||||
profile_prefix = f"voxcpm_sync_{int(time.time())}"
|
||||
stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured"
|
||||
print(f"Starting profiler (sync): stages={stages_text}, dir={args.profiler_dir}")
|
||||
omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages)
|
||||
profiler_started = True
|
||||
|
||||
t_total = time.perf_counter()
|
||||
total_elapsed = 0.0
|
||||
saved_paths: list[Path] = []
|
||||
prompt_specs: list[PromptSpec] = args.prompt_specs
|
||||
try:
|
||||
for run in range(args.num_runs):
|
||||
t_run = time.perf_counter()
|
||||
run_paths: list[Path] = []
|
||||
for prompt_index, spec in enumerate(prompt_specs):
|
||||
prompt_paths, _, first_audio_elapsed, elapsed_s, audio_duration_s, metrics_request_id = (
|
||||
_run_sync_single(
|
||||
spec,
|
||||
request_prefix=f"sync_run{run + 1}_{prompt_index + 1:03d}",
|
||||
save_outputs=True,
|
||||
run_index=run,
|
||||
)
|
||||
)
|
||||
run_paths.extend(prompt_paths)
|
||||
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
|
||||
rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
|
||||
print(
|
||||
f"Saved (sync) run {run + 1}/{args.num_runs}, "
|
||||
f"prompt {prompt_index + 1}/{len(prompt_specs)}: {len(prompt_paths)} file(s){ttfp_text}{rtf_text}"
|
||||
)
|
||||
_emit_offline_metrics(
|
||||
request_id=metrics_request_id,
|
||||
elapsed_s=elapsed_s,
|
||||
first_audio_elapsed=first_audio_elapsed,
|
||||
audio_duration_s=audio_duration_s,
|
||||
)
|
||||
|
||||
saved_paths.extend(run_paths)
|
||||
print(
|
||||
f"Run {run + 1}/{args.num_runs} finished: {len(run_paths)} file(s) ({time.perf_counter() - t_run:.2f}s)"
|
||||
)
|
||||
for path in run_paths:
|
||||
print(f" {path}")
|
||||
|
||||
total_elapsed = time.perf_counter() - t_total
|
||||
finally:
|
||||
if profiler_started:
|
||||
print("Stopping profiler (sync)...")
|
||||
omni.stop_profile(stages=args.profiler_stages)
|
||||
if args.profiler_wait_seconds > 0:
|
||||
print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...")
|
||||
time.sleep(args.profiler_wait_seconds)
|
||||
|
||||
print(
|
||||
f"All sync runs finished: {args.num_runs} run(s), "
|
||||
f"{len(prompt_specs)} prompt(s), {len(saved_paths)} file(s) in {total_elapsed:.2f}s total"
|
||||
)
|
||||
return saved_paths
|
||||
|
||||
|
||||
def main(args) -> int:
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
profiled_stage_config_path: str | None = None
|
||||
original_stage_config_path = args.stage_configs_path
|
||||
if args.enable_profiler:
|
||||
Path(args.profiler_dir).mkdir(parents=True, exist_ok=True)
|
||||
profiled_stage_config_path = _build_profiled_stage_config(
|
||||
args.stage_configs_path,
|
||||
str(Path(args.profiler_dir).resolve()),
|
||||
)
|
||||
args.stage_configs_path = profiled_stage_config_path
|
||||
|
||||
is_streaming = _is_streaming_stage_config(args.stage_configs_path)
|
||||
voice_clone_count = _count_voice_clone_prompts(args.prompt_specs)
|
||||
print(f"Model: {args.model}")
|
||||
print(f"Stage config: {original_stage_config_path}")
|
||||
print(f"Route: {'streaming' if is_streaming else 'sync'} (from stage-configs-path)")
|
||||
print(f"Prompt count: {len(args.prompt_specs)}")
|
||||
print("Batch mode: sequential (aligned with native VoxCPM)")
|
||||
print(f"Warmup runs: {args.warmup_runs}")
|
||||
print(f"Voice cloning prompts: {voice_clone_count}/{len(args.prompt_specs)}")
|
||||
if args.enable_profiler:
|
||||
print(f"Profiler: enabled (dir={args.profiler_dir}, stages={args.profiler_stages or 'all-configured'})")
|
||||
print(f"Profiled stage config: {args.stage_configs_path}")
|
||||
if voice_clone_count:
|
||||
print("Voice cloning note: --ref-text/ref_text must match the spoken content of the reference audio.")
|
||||
print(f"Num runs: {args.num_runs}")
|
||||
try:
|
||||
if is_streaming:
|
||||
asyncio.run(_run_streaming(args))
|
||||
else:
|
||||
_run_sync(args)
|
||||
finally:
|
||||
if profiled_stage_config_path is not None and os.path.exists(profiled_stage_config_path):
|
||||
os.unlink(profiled_stage_config_path)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
raise SystemExit(main(parse_args()))
|
||||
@@ -14,6 +14,7 @@ models:
|
||||
supported_tasks: [default_voice, voice_design]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
trust_remote_code: true
|
||||
task_extra_body:
|
||||
default_voice:
|
||||
voice: Vivian
|
||||
@@ -27,13 +28,76 @@ models:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
trust_remote_code: true
|
||||
task_extra_body:
|
||||
voice_clone:
|
||||
task_type: Base
|
||||
|
||||
openbmb/VoxCPM2:
|
||||
supported_tasks: [voice_clone, default_voice]
|
||||
trust_remote_code: true
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
trust_remote_code: true
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
default_voice:
|
||||
voice: default
|
||||
|
||||
# MOSS-TTS family. After the serving_speech.py variant-detection refactor
|
||||
# all five variants now go through OpenAI /v1/audio/speech with the
|
||||
# variant-specific request fields (ambient_sound, ref_audio_2, instructions).
|
||||
OpenMOSS-Team/MOSS-TTS:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
|
||||
# MOSS-TTS-v1.5: continued-training upgrade of MOSS-TTS 1.0 — same
|
||||
# MossTTSDelay-8B architecture and API, so identical wiring. Adds 31-language
|
||||
# synthesis via the `language` field and inline `[pause Xs]` markers.
|
||||
OpenMOSS-Team/MOSS-TTS-v1.5:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
|
||||
OpenMOSS-Team/MOSS-TTS-Realtime:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
|
||||
# MOSS-TTSD synthesizes [S1]/[S2] dialogue. The seed-tts voice_clone dataset
|
||||
# has single-speaker rows; bench piggybacks on it (both speakers share the
|
||||
# same ref_audio). For real dialogue evaluation use the ttsd dataset.
|
||||
OpenMOSS-Team/MOSS-TTSD-v1.0:
|
||||
supported_tasks: [voice_clone, dialogue]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
dialogue: {}
|
||||
|
||||
# MOSS-SoundEffect ignores text/ref_audio; expects ambient_sound +
|
||||
# duration_seconds. The sound_effect dataset module emits those.
|
||||
OpenMOSS-Team/MOSS-SoundEffect:
|
||||
supported_tasks: [sound_effect]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
sound_effect:
|
||||
duration_seconds: 10
|
||||
|
||||
# MOSS-VoiceGenerator: instructions-driven (no ref_audio). seed-tts-design
|
||||
# dataset already emits {instructions, target_text}, so voice_design works
|
||||
# out of the box once the serving validator accepts instructions.
|
||||
OpenMOSS-Team/MOSS-VoiceGenerator:
|
||||
supported_tasks: [voice_design]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_design: {}
|
||||
@@ -1,5 +1,8 @@
|
||||
# Target vllm version: v0.21.0 (releases/v0.21.0 branch).
|
||||
# The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must
|
||||
# be updated on every rebase to match the new upstream vllm commit/tag.
|
||||
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
|
||||
ARG VLLM_BASE_TAG=v0.21.0
|
||||
ARG VLLM_BASE_TAG=v0.22.0
|
||||
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
|
||||
ARG APP_DIR=/workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.20.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.22.0
|
||||
FROM ${BASE_IMAGE}
|
||||
|
||||
ARG COMMON_WORKDIR=/app
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.22.0
|
||||
FROM ${BASE_IMAGE} AS base
|
||||
|
||||
# Declare a variable to know if we want to use the nightly build or the stable build.
|
||||
|
||||
@@ -88,7 +88,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.21.0
|
||||
ARG VLLM_VERSION=v0.22.0
|
||||
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
@@ -133,6 +133,11 @@ ENTRYPOINT []
|
||||
|
||||
FROM ${VLLM_BASE} AS vllm-omni
|
||||
|
||||
RUN apt-get update && \
|
||||
apt-get install -y espeak-ng ffmpeg git jq && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /workspace/vllm-omni
|
||||
COPY . .
|
||||
|
||||
|
||||
@@ -80,6 +80,8 @@ nav:
|
||||
- Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
|
||||
- FP8 W8A8: user_guide/quantization/fp8.md
|
||||
- Int8 W8A8: user_guide/quantization/int8.md
|
||||
- MXFP8 W8A8: user_guide/quantization/mxfp8.md
|
||||
- MXFP4 W4A4: user_guide/quantization/mxfp4.md
|
||||
- ModelOpt: user_guide/quantization/modelopt.md
|
||||
- GGUF: user_guide/quantization/gguf.md
|
||||
- AutoRound: user_guide/quantization/autoround.md
|
||||
|
||||
@@ -121,7 +121,6 @@ Configuration classes.
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.voxcpm.VoxCPMConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][]
|
||||
|
||||
## Workers
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 261 KiB After Width: | Height: | Size: 408 KiB |
+134
-26
@@ -2,29 +2,29 @@
|
||||
The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
|
||||
|
||||
## Notes
|
||||
We currently only support using the "openai-chat-omni" backend.
|
||||
We currently only support using the "openai-chat-omni" and "openai-image-edits-omni" backend.
|
||||
|
||||
## Basic Parameter Description
|
||||
You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
|
||||
- `--omni`
|
||||
- `--omni`
|
||||
Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio.
|
||||
|
||||
- `--backend`
|
||||
- `--backend`
|
||||
Specify the backend adapter as openai-chat-omni, using OpenAI Chat compatible API behavior as the protocol. Currently only openai-chat-omni is supported.
|
||||
|
||||
- `--model`
|
||||
- `--model`
|
||||
The model identifier to load, filled according to the models supported by vLLM-Omni.
|
||||
|
||||
- `--endpoint`
|
||||
- `--endpoint`
|
||||
The API endpoint exposed externally, to which clients send their requests.
|
||||
|
||||
- `--dataset-name`
|
||||
- `--dataset-name`
|
||||
The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio).
|
||||
|
||||
- `--num-prompts`
|
||||
- `--num-prompts`
|
||||
The total number of requests to send, an integer.
|
||||
|
||||
- `--max-concurrency`
|
||||
- `--max-concurrency`
|
||||
"Maximum number of concurrent requests. This can be used "
|
||||
"to help simulate an environment where a higher level component "
|
||||
"is enforcing a maximum number of concurrent requests. While the "
|
||||
@@ -34,69 +34,73 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
|
||||
"actual request rate may be lower than specified with --request-rate, "
|
||||
"if the server is not processing requests fast enough to keep up."
|
||||
|
||||
- `--request-rate`
|
||||
- `--request-rate`
|
||||
"Number of requests per second. If this is inf, "
|
||||
"then all the requests are sent at time 0. "
|
||||
"Otherwise, we use Poisson process or gamma distribution "
|
||||
"to synthesize the request arrival times."
|
||||
|
||||
- `--ignore-eos`
|
||||
- `--ignore-eos`
|
||||
"Set ignore_eos flag when sending the benchmark request."
|
||||
|
||||
- `--metric-percentiles`
|
||||
- `--metric-percentiles`
|
||||
Comma-separated list of percentiles for selected metrics. "
|
||||
"To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". "
|
||||
"Default value is \"99\"."
|
||||
"Use \"--percentile-metrics\" to select metrics.
|
||||
|
||||
- `--percentile-metrics`
|
||||
- `--percentile-metrics`
|
||||
"Comma-separated list of selected metrics to report percentiles."
|
||||
"This argument specifies the metrics to report percentiles."
|
||||
'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
|
||||
'Allowed metric names are "ttft", "tpot", "itl", "ttfc", "tpoc", "icl", '
|
||||
'"tpop", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
|
||||
|
||||
- `--save-result`
|
||||
- `-print-stage`
|
||||
Print per-stage benchmark metrics for --omni serving when stage metrics are returned by the server. Disabled by default.
|
||||
|
||||
- `--save-result`
|
||||
Specify to save benchmark results to a json file
|
||||
|
||||
- `--save-detailed`
|
||||
- `--save-detailed`
|
||||
"When saving the results, whether to include per request "
|
||||
"information such as response, error, ttfs, tpots, etc."
|
||||
|
||||
- `--result-dir`
|
||||
- `--result-dir`
|
||||
"Specify directory to save benchmark json results."
|
||||
"If not specified, results are saved in the current directory."
|
||||
|
||||
- `--result-filename`
|
||||
- `--result-filename`
|
||||
"Specify the filename to save benchmark json results."
|
||||
"If not specified, results will be saved in "
|
||||
"{label}-{args.request_rate}qps-{base_model_id}-{current_dt}.json"
|
||||
|
||||
- `--random-prefix-len`
|
||||
- `--random-prefix-len`
|
||||
Number of fixed prefix tokens before the random context in a request.
|
||||
The total input length is the sum of random-prefix-len and a random
|
||||
context length sampled from [input_len * (1 - range_ratio),
|
||||
input_len * (1 + range_ratio)].Only the random and random-mm modes
|
||||
support this parameter.
|
||||
|
||||
- `--random-input-len`
|
||||
- `--random-input-len`
|
||||
Number of input tokens per request.Only the random and random-mm modes support this parameter.
|
||||
|
||||
- `--random-output-len`
|
||||
- `--random-output-len`
|
||||
Number of output tokens per request.Only the random and random-mm modes support this parameter.
|
||||
|
||||
- `--random-range-ratio`
|
||||
- `--random-range-ratio`
|
||||
Range ratio for sampling input/output length,
|
||||
used only for random sampling. Must be in the range [0, 1) to define
|
||||
a symmetric sampling range
|
||||
[length * (1 - range_ratio), length * (1 + range_ratio)].
|
||||
Only the random and random-mm modes support this parameter.
|
||||
|
||||
- `--random-mm-base-items-per-request`
|
||||
- `--random-mm-base-items-per-request`
|
||||
Base number of multimodal items per request for random-mm.
|
||||
Actual per-request count is sampled around this base using
|
||||
--random-mm-num-mm-items-range-ratio.
|
||||
Only the random-mm mode supports this parameter.
|
||||
|
||||
- `--random-mm-limit-mm-per-prompt`
|
||||
- `--random-mm-limit-mm-per-prompt`
|
||||
Per-modality hard caps for items attached per request, e.g.
|
||||
'{"image": 3, "video": 1, "audio": 1}'. The sampled per-request item
|
||||
count is clamped to the sum of these limits. When a modality
|
||||
@@ -104,7 +108,7 @@ Specify to save benchmark results to a json file
|
||||
renormalized.
|
||||
Only the random-mm mode supports this parameter.
|
||||
|
||||
- `--random-mm-num-mm-items-range-ratio`
|
||||
- `--random-mm-num-mm-items-range-ratio`
|
||||
Range ratio r in [0, 1] for sampling items per request.
|
||||
We sample uniformly from the closed integer range
|
||||
[floor(n*(1-r)), ceil(n*(1+r))]
|
||||
@@ -115,7 +119,7 @@ Specify to save benchmark results to a json file
|
||||
An error is raised if the computed min exceeds the max.
|
||||
Only the random-mm mode supports this parameter.
|
||||
|
||||
- `--random-mm-bucket-config`
|
||||
- `--random-mm-bucket-config`
|
||||
The bucket config is a dictionary mapping a multimodal item
|
||||
sampling configuration to a probability.
|
||||
Currently allows for 3 modalities: audio, images and videos.
|
||||
@@ -276,7 +280,9 @@ Generate synthetic image、video、audio inputs alongside random text prompts to
|
||||
|
||||
Notes:
|
||||
|
||||
- Works only with online benchmark via the OpenAI backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
|
||||
- Works only with online benchmark via:
|
||||
- the OpenAI chat backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
|
||||
- the OpenAI edit image backend (`--backend openai-image-edits-omni`) and endpoint `/v1/images/edits`.
|
||||
|
||||
Start the server (example):
|
||||
|
||||
@@ -357,3 +363,105 @@ How sampling works:
|
||||
This should be seen as an edge case, and if this behavior can be avoided by setting `--random-mm-limit-mm-per-prompt` to a large number. Note that this might result in errors due to engine config `--limit-mm-per-prompt`.
|
||||
- The resulting request contains synthetic image data in `multi_modal_data` (OpenAI Chat format). When `random-mm` is used with the OpenAI Chat backend, prompts remain text and MM content is attached via `multi_modal_data`.
|
||||
</details>
|
||||
|
||||
### Multi-Stage Benchmark
|
||||
|
||||
<details class="admonition abstract" markdown="1">
|
||||
<summary>Show more</summary>
|
||||
|
||||
Start the server (change the port and model):
|
||||
|
||||
```bash
|
||||
vllm serve --omni --port 29999 --model ~/Qwen3-Omni-30B-A3B-Instruct
|
||||
```
|
||||
|
||||
Then run the benchmarking script (remember to add the flag `--print-stage`):
|
||||
```bash
|
||||
vllm bench serve --omni \
|
||||
--dataset-name random \
|
||||
--port 29999 \
|
||||
--max-concurrency 1 \
|
||||
--num-warmups 2 \
|
||||
--model ~/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--endpoint /v1/chat/completions \
|
||||
--backend openai-chat-omni \
|
||||
--num-prompts 3 \
|
||||
--random-input-len 2500 \
|
||||
--ignore-eos \
|
||||
--percentile-metrics ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,ttfc,tpoc,icl \
|
||||
--random-output-len 900 \
|
||||
--extra_body '{"modalities": ["text","audio"]}' \
|
||||
--print-stage
|
||||
```
|
||||
|
||||
Besides the "Serving Benchmark Result", there will be "Stage Benchmark Result" below:
|
||||
```text
|
||||
============= Stage Benchmark Result =============
|
||||
=============== Stage 0 (thinker) ================
|
||||
-------------------Stage Timing-------------------
|
||||
Mean stage_gen_time (ms): 18243.16
|
||||
Median stage_gen_time (ms): 18220.83
|
||||
P99 stage_gen_time (ms): 18296.75
|
||||
================== Text Result ===================
|
||||
Stage generated tokens: 2700
|
||||
-----------Serving Time to First Token------------
|
||||
Mean Serving TTFT (ms): 816.31
|
||||
Median Serving TTFT (ms): 818.94
|
||||
P99 Serving TTFT (ms): 829.76
|
||||
-----Time per Output Token (excl. 1st token)------
|
||||
Mean TPOP (ms): 19.41
|
||||
Median TPOP (ms): 19.40
|
||||
P99 TPOP (ms): 19.47
|
||||
---------------Inter-token Latency----------------
|
||||
Mean ITL (ms): 19.41
|
||||
Median ITL (ms): 19.12
|
||||
P99 ITL (ms): 23.58
|
||||
================ Stage 1 (talker) ================
|
||||
-------------------Stage Timing-------------------
|
||||
Mean stage_gen_time (ms): 18849.64
|
||||
Median stage_gen_time (ms): 13641.15
|
||||
P99 stage_gen_time (ms): 33655.58
|
||||
============= Internal Stream Result =============
|
||||
-----------Serving Time to First Chunk------------
|
||||
Mean Serving TTFC (ms): 1029.14
|
||||
Median Serving TTFC (ms): 1027.98
|
||||
P99 Serving TTFC (ms): 1032.59
|
||||
-----Time per Output Chunk (excl. 1st chunk)------
|
||||
Mean TPOP (ms): 17.53
|
||||
Median TPOP (ms): 18.70
|
||||
P99 TPOP (ms): 18.88
|
||||
---------------Inter-chunk Latency----------------
|
||||
Mean ICL (ms): 16.27
|
||||
Median ICL (ms): 18.19
|
||||
P99 ICL (ms): 25.86
|
||||
=============== Stage 2 (code2wav) ===============
|
||||
-------------------Stage Timing-------------------
|
||||
Mean stage_gen_time (ms): 18860.74
|
||||
Median stage_gen_time (ms): 13655.84
|
||||
P99 stage_gen_time (ms): 33659.79
|
||||
================== Audio Result ==================
|
||||
Stage audio duration generated(s): 260.11
|
||||
Stage audio frames generated: 6242730
|
||||
-----------Serving Time to First Packet-----------
|
||||
Mean Serving AUDIO_TTFP (ms): 1366.16
|
||||
Median Serving AUDIO_TTFP (ms): 1367.20
|
||||
P99 Serving AUDIO_TTFP (ms): 1368.26
|
||||
-----------------Real Time Factor-----------------
|
||||
Mean AUDIO_RTF: 0.24
|
||||
Median AUDIO_RTF: 0.26
|
||||
P99 AUDIO_RTF: 0.27
|
||||
```
|
||||
Explanation:
|
||||
- stage_gen_time: Time from submitting a request to a specific stage to that stage finishing generation.
|
||||
|
||||
- Serving TTFC (Time to First Chunk): Time from the HTTP request being accepted by the serving frontend to
|
||||
the stage producing its first non-empty output chunk. Keep the name TTFT in text stage, TTFP in audio stage
|
||||
for easier mapping.
|
||||
|
||||
- TPOP (Time per Output Chunk): Average time from the first output chunk to stage completion, divided by the
|
||||
number of remaining output chunks. The TPOP abbreviation follows Qwen3.5-Omni Technical Report.
|
||||
|
||||
- ICL (Inter-Chunk Latency): Time between two consecutive output chunks produced by the same stage. Keep the
|
||||
name ITL in text stage for easier mapping.
|
||||
|
||||
</details>
|
||||
@@ -39,11 +39,34 @@ Committers have write access and merge rights. They typically have deep expertis
|
||||
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
|
||||
- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
|
||||
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
|
||||
- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
|
||||
- [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs
|
||||
- [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache
|
||||
- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
|
||||
- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism
|
||||
- [@wtomin](https://github.com/wtomin): Diffusion models, parallel, and docs
|
||||
- [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
|
||||
- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
|
||||
- [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
|
||||
- [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
|
||||
|
||||
## Diffusion Workload Division
|
||||
|
||||
This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review.
|
||||
|
||||
### Runtime-related
|
||||
|
||||
- Diffusion models: @RuixiangMa, @wtomin
|
||||
|
||||
### Optimization
|
||||
|
||||
- Parallel: @RuixiangMa, @wtomin
|
||||
- Attention backend: @ZJY0516
|
||||
- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
|
||||
- Offload: @yuanheng-zhao
|
||||
- Quantization: @lishunyang12, @david6666666
|
||||
- Kernel fusion / communication-computation: @ZJY0516
|
||||
|
||||
### Docs & Test
|
||||
|
||||
- Diffusion docs: @RuixiangMa, @wtomin
|
||||
|
||||
## Meetings
|
||||
|
||||
|
||||
@@ -83,7 +83,10 @@ There are two important details:
|
||||
|
||||
Today that compatibility rule is still shape-sensitive. `height`, `width`,
|
||||
`num_frames`, and CFG-related fields remain part of the key, so different
|
||||
resolutions or incompatible guidance settings do **not** co-batch yet.
|
||||
resolutions or incompatible guidance settings do **not** co-batch yet. The
|
||||
key also covers LoRA identity (`lora_int_id`, `lora_scale`), so requests
|
||||
targeting different adapters or scales run in separate batches and the
|
||||
worker can activate exactly one adapter per step.
|
||||
|
||||
The current batching unit is one `OmniDiffusionRequest`. Requests with
|
||||
multiple prompts do not participate in batching today.
|
||||
|
||||
@@ -9,6 +9,7 @@ Backend-specific setup lives in separate docs:
|
||||
- [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
|
||||
- [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
|
||||
- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
|
||||
- [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md)
|
||||
|
||||
## Overview
|
||||
|
||||
@@ -23,6 +24,7 @@ Current connectors operate in D2H2D (device to host to device) mode.
|
||||
| Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
|
||||
| Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
|
||||
| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
|
||||
| Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. |
|
||||
|
||||
## Core API
|
||||
|
||||
|
||||
@@ -0,0 +1,211 @@
|
||||
# YuanrongTransferEngineConnector
|
||||
|
||||
## When to Use
|
||||
|
||||
Use `YuanrongTransferEngineConnector` for high-performance peer-to-peer KV cache
|
||||
transfer on Ascend NPU. It is different from `YuanrongConnector`:
|
||||
|
||||
- `YuanrongConnector` uses Yuanrong Datasystem as a distributed KV store and
|
||||
requires Datasystem workers plus etcd.
|
||||
- `YuanrongTransferEngineConnector` uses Yuanrong TransferEngine directly. It
|
||||
registers local memory in each stage worker and lets the receiver pull data
|
||||
from the sender with TransferEngine.
|
||||
|
||||
The implementation currently lives under
|
||||
`vllm_omni/platforms/npu/omni_connectors/` because the supported backend is
|
||||
Ascend NPU only. The connector is still registered with the generic
|
||||
OmniConnector factory under the name `YuanrongTransferEngineConnector`.
|
||||
|
||||
For Ascend P2P transfer, the connector should use NPU memory as its transfer
|
||||
pool. The Mooncake-style CPU pool is not supported for this connector.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Install Yuanrong Datasystem/TransferEngine Python bindings in the runtime
|
||||
environment. The packaged install is:
|
||||
|
||||
```bash
|
||||
pip install openyuanrong-datasystem
|
||||
```
|
||||
|
||||
To build from source, follow the upstream project instructions:
|
||||
|
||||
- Source repository: [openeuler/yuanrong-datasystem](https://atomgit.com/openeuler/yuanrong-datasystem)
|
||||
- Installation guide: [Yuanrong Datasystem Linux installation](https://pages.openeuler.openatom.cn/openyuanrong-datasystem/docs/zh-cn/latest/installation/installation_linux.html)
|
||||
|
||||
The Ascend driver, CANN runtime, HCCL/HCCP, and NPU device network must also be
|
||||
configured on every machine participating in the transfer.
|
||||
|
||||
## Check NPU Device IPv4
|
||||
|
||||
Yuanrong TransferEngine with `protocol: "ascend"` uses the NPU device network,
|
||||
not only the host management IP. Each participating NPU must have an IPv4 HCCN
|
||||
device IP.
|
||||
|
||||
Check the device IP, link, and health state:
|
||||
|
||||
```bash
|
||||
for i in {0..7}; do
|
||||
echo "===== NPU $i ====="
|
||||
hccn_tool -i "$i" -ip -g
|
||||
hccn_tool -i "$i" -link -g
|
||||
hccn_tool -i "$i" -net_health -g
|
||||
done
|
||||
```
|
||||
|
||||
Expected output for each NPU should include an IPv4 address and healthy link:
|
||||
|
||||
```text
|
||||
ipaddr:10.86.1.80
|
||||
netmask:255.255.255.0
|
||||
link status: UP
|
||||
net health status: Success
|
||||
```
|
||||
|
||||
If the output says `Get ipconf failed, because no ip was preset there!`, the NPU
|
||||
device IP is not configured and Ascend P2P will fail. If the address is IPv6,
|
||||
configure IPv4 for this connector path or validate that the underlying Yuanrong
|
||||
TransferEngine build supports the IPv6 HCCN path before using it.
|
||||
|
||||
## Check Device-to-Device Connectivity
|
||||
|
||||
After collecting the target-stage NPU IPv4 addresses, ping them from each
|
||||
source-stage NPU:
|
||||
|
||||
```bash
|
||||
# Example: AR uses NPU 0..3 and DiT uses NPU 4..7.
|
||||
# Replace the destination IPs with the values reported by hccn_tool -i 4..7 -ip -g.
|
||||
for src in 0 1 2 3; do
|
||||
for dst_ip in 10.86.1.78 10.86.1.79 10.86.1.80 10.86.1.81; do
|
||||
echo "===== NPU $src -> $dst_ip ====="
|
||||
hccn_tool -i "$src" -ping -g address "$dst_ip" || true
|
||||
done
|
||||
done
|
||||
```
|
||||
|
||||
The expected result is `0.00% packet loss` for every required source/target pair.
|
||||
This verifies basic HCCN IPv4 connectivity. TransferEngine can still fail later
|
||||
if RDMA/QP/HCCL setup is broken, but failed ping means the device network must be
|
||||
fixed first.
|
||||
|
||||
## YAML Configuration
|
||||
|
||||
For Ascend P2P, configure the connector with `protocol: "ascend"` and an NPU
|
||||
memory pool:
|
||||
|
||||
```yaml
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 1
|
||||
connectors:
|
||||
yuanrong_te_connector:
|
||||
name: YuanrongTransferEngineConnector
|
||||
extra:
|
||||
host: "auto"
|
||||
zmq_port: 50051
|
||||
rpc_port: "auto"
|
||||
protocol: "ascend"
|
||||
device_name: "auto"
|
||||
memory_pool_size: 1073741824
|
||||
memory_pool_device: "npu"
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
```
|
||||
|
||||
Important fields:
|
||||
|
||||
| Parameter | Recommended Value | Notes |
|
||||
|---|---|---|
|
||||
| `host` | `"auto"` | Advertises a routable local host IP for ZMQ metadata exchange. |
|
||||
| `zmq_port` | `50051` | Base ZMQ port. The runtime applies stage/rank offsets. |
|
||||
| `rpc_port` | `"auto"` | Lets TransferEngine choose an available RPC port. |
|
||||
| `protocol` | `"ascend"` | Uses the Ascend TransferEngine backend. |
|
||||
| `device_name` | `"auto"` | Each TP worker resolves to its local logical NPU, for example `npu:0`, `npu:1`, etc. |
|
||||
| `memory_pool_device` | `"npu"` | Required for the Ascend fast path. This registers NPU memory with TransferEngine. |
|
||||
| `memory_pool_size` | `1073741824` | 1 GiB per worker is a practical starting point for KV-cache transfer validation. Increase if the transfer pool is exhausted. |
|
||||
|
||||
`memory_pool_device: "cpu"` is rejected for
|
||||
`YuanrongTransferEngineConnector` with `protocol: "ascend"`. CPU pool registration
|
||||
can lead to Ascend P2P failures such as HCCL RA init errors, QP timeouts, or
|
||||
FFTS/SDMA runtime errors. `memory_pool_device: "cpu"` is common for
|
||||
`MooncakeTransferEngineConnector`, but Yuanrong TE on Ascend only supports
|
||||
`"npu"`.
|
||||
|
||||
## 4-Card AR-to-DiT TP Example
|
||||
|
||||
For AR 4-way TP sending KV cache to DiT 4-way TP, set the stage TP topology in
|
||||
both stage `omni_kv_config` blocks:
|
||||
|
||||
```yaml
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
runtime:
|
||||
devices: "0,1,2,3"
|
||||
engine_args:
|
||||
tensor_parallel_size: 4
|
||||
omni_kv_config:
|
||||
need_send_cache: true
|
||||
rank_mapping:
|
||||
from_tp: 4
|
||||
to_tp: 4
|
||||
output_connectors:
|
||||
to_stage_1: yuanrong_te_connector
|
||||
|
||||
- stage_id: 1
|
||||
runtime:
|
||||
devices: "4,5,6,7"
|
||||
engine_args:
|
||||
parallel_config:
|
||||
tensor_parallel_size: 4
|
||||
omni_kv_config:
|
||||
need_recv_cache: true
|
||||
rank_mapping:
|
||||
from_tp: 4
|
||||
to_tp: 4
|
||||
input_connectors:
|
||||
from_stage_0: yuanrong_te_connector
|
||||
```
|
||||
|
||||
In this topology each DiT rank receives the corresponding AR rank shard:
|
||||
|
||||
```text
|
||||
DiT rank 0 <- AR rank 0
|
||||
DiT rank 1 <- AR rank 1
|
||||
DiT rank 2 <- AR rank 2
|
||||
DiT rank 3 <- AR rank 3
|
||||
```
|
||||
|
||||
## Validate a Successful Run
|
||||
|
||||
Search the inference log for Yuanrong TE receive-side messages:
|
||||
|
||||
```bash
|
||||
grep -E "\[YR GET\]|Successfully received KV cache|transfer_engine get failed" inference.log
|
||||
```
|
||||
|
||||
A successful Ascend fast-path transfer contains lines like:
|
||||
|
||||
```text
|
||||
[YR GET] ... fast_path, zero-copy
|
||||
Successfully received KV cache ... across 1 key(s)
|
||||
Applied CFG KV caches
|
||||
[KV Reuse] ...
|
||||
```
|
||||
|
||||
Sender-side `KV transfer OK` only means that the sender registered the payload
|
||||
and metadata. The actual P2P data transfer happens on the receiver side during
|
||||
`[YR GET]`.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
| Symptom | Likely Cause | Action |
|
||||
|---|---|---|
|
||||
| `IPv4 device IP not found` | NPU HCCN IPv4 is missing. | Run `hccn_tool -i <id> -ip -g` and configure IPv4 for every participating NPU. |
|
||||
| `P2PCommInitRootInfo failed`, `ra init failed`, `RA_QP_STATUS_TIMEOUT` | Ascend P2P/RDMA/QP setup failed. | First verify device IPv4, link health, and device-to-device ping. Then test a minimal Yuanrong TE pair outside vLLM. |
|
||||
| `fftsplus sdma error`, `context is abort` | Runtime SDMA/FFTS failure, often after an invalid or failed P2P transfer. | Use `memory_pool_device: "npu"` for Ascend TE, reduce `memory_pool_size` if OOM occurs, and isolate the failing NPU pair with a Yuanrong TE demo. |
|
||||
| Sender logs `KV transfer OK` but receiver never logs `[YR GET]` | Sender `put()` succeeded, receiver P2P read failed or timed out. | Check receiver logs and search for `transfer_engine get failed`. |
|
||||
| Pool allocation fails or OOM | Transfer pool is too small or too large for available NPU memory. | Start with 1 GiB (`1073741824`), reduce to 512 MiB for validation, or increase for larger KV payloads. |
|
||||
@@ -13,6 +13,10 @@ This section contains design documents and architecture specifications for vLLM-
|
||||
- [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
|
||||
- [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
|
||||
|
||||
## Infrastructure Design Documents
|
||||
|
||||
- [Prometheus Metrics](metrics.md)
|
||||
|
||||
## Module Design Documents
|
||||
|
||||
- [AR Module](module/ar_module.md)
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
# Prometheus Metrics Design
|
||||
|
||||
This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. This iteration covers the text and audio surface; diffusion / image / video families land in a follow-up.
|
||||
|
||||
## Objectives
|
||||
|
||||
- Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope).
|
||||
- Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically.
|
||||
- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss.
|
||||
- Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable.
|
||||
- Keep the metrics collection overhead low enough that it does not regress TTFA or throughput.
|
||||
|
||||
## Background
|
||||
|
||||
### Upstream vLLM Metrics
|
||||
|
||||
Upstream vLLM defines ~37 Prometheus metric families under the `vllm:` prefix. These are registered by `PrometheusStatLogger` and cover engine-level state: KV cache usage, running/waiting request counts, token throughput, TTFT, inter-token latency, e2e latency, and so on. They are served via the `/metrics` HTTP endpoint provided by `prometheus_fastapi_instrumentator` and the default `prometheus_client` WSGI handler.
|
||||
|
||||
vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` collector whose `_name` attribute contains the substring `"vllm"`. This runs during engine initialization to clean up stale collectors from prior instantiations within the same process.
|
||||
|
||||
### The Problem
|
||||
|
||||
vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm:omni_` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
|
||||
|
||||
Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below).
|
||||
|
||||
## Architecture
|
||||
|
||||
### Component Overview
|
||||
|
||||
```text
|
||||
+------------------------+
|
||||
| API Server (FastAPI) |
|
||||
| GET /metrics |
|
||||
+-----------+------------+
|
||||
|
|
||||
prometheus_client default registry
|
||||
|
|
||||
+--------+--------+--------+--------+--------+
|
||||
| |
|
||||
vllm:omni_* vllm:*
|
||||
collectors collectors
|
||||
| |
|
||||
+----+--------+ +-----------+ +----------+ +-----------+
|
||||
| OmniPromet- | | OmniMod- | | OmniTra- | | OmniProm- |
|
||||
| heusMetrics | | alityMet- | | nsferMe- | | etheusSt- |
|
||||
| | | rics | | trics | | atLogger |
|
||||
+----+--------+ +-----+-----+ +----+-----+ +----+------+
|
||||
| | | |
|
||||
OmniBase OmniBase Orchestrator Orchestrator
|
||||
(request life- (finalize + (record_trans- (per-(stage,
|
||||
cycle, success/ streaming fer_tx/rx replica)
|
||||
fail counter) hooks via hooks via scheduler/
|
||||
observe_* emit hook in iteration
|
||||
APIs) OrchestratorAg- stats)
|
||||
gregator)
|
||||
```
|
||||
|
||||
### Data Flow
|
||||
|
||||
There are four independent paths for metric collection.
|
||||
|
||||
**Path 1: Pipeline-level metrics (`vllm:omni_*`)**
|
||||
|
||||
`OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress:
|
||||
|
||||
- `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`.
|
||||
- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm:omni_requests_success_total{finished_reason}`.
|
||||
- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` so a single Counter family covers both natural and aborted completion.
|
||||
|
||||
**Path 2: Audio modality metrics (`vllm:omni_audio_*`)**
|
||||
|
||||
`OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites:
|
||||
|
||||
- `observe_modality_at_finalize(...)` — called from `omni_base._process_single_result` inside the existing `e2e_done` finalize guard. For `output_type == "audio"` it emits `audio_frames_total`, `audio_duration_s`, `audio_rtf` (or `audio_skipped_requests_total{reason="no_audio_data"}` when no audio was produced). Sample rate is resolved from `engine_outputs.multimodal_output` via `definitions.resolve_audio_sample_rate(...)` (fallback chain mirrors `serving_chat.py`'s audio response path).
|
||||
- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry.
|
||||
- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`.
|
||||
|
||||
**Path 3: Cross-stage transfer metrics (`vllm:omni_transfer_*`)**
|
||||
|
||||
`OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution.
|
||||
|
||||
The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transfer_rx`. After the existing `TransferEdgeStats` accumulation, the aggregator calls `_emit_transfer_tx` / `_emit_transfer_rx`. Those:
|
||||
|
||||
1. Resolve `from_replica` / `to_replica` via a `replica_resolver` callback supplied by `async_omni.py`. The resolver delegates to `stage_pool.get_bound_replica_id(request_id)` — i.e. the orchestrator's existing sticky-routing binding is the source of truth.
|
||||
2. Convert the underlying `_ms` accumulators to seconds and call the `_s`-suffixed observe methods on `OmniTransferMetrics`.
|
||||
|
||||
Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected).
|
||||
|
||||
**Path 4: Per-engine metrics (`vllm:*`, stage/replica wrap)**
|
||||
|
||||
The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics.
|
||||
|
||||
### Shared State Between Threads
|
||||
|
||||
The Orchestrator runs in a background thread. The API server (OmniBase) runs in the asyncio event loop thread. `OmniRequestCounter` bridges them — a plain Python object with an `int` field. The Orchestrator increments/decrements it; the entrypoint reads it for gauge updates. No lock is needed because the counter is advisory (a stale read by one Prometheus scrape interval is acceptable). It is created by `AsyncOmniEngine.__init__()` and passed to the Orchestrator at construction time.
|
||||
|
||||
### Metric Registration and Lifecycle
|
||||
|
||||
All `vllm:omni_*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
|
||||
|
||||
The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm:omni_*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
|
||||
|
||||
## OmniPrometheusStatLogger Wrap
|
||||
|
||||
Upstream `PrometheusStatLogger.__init__` hard-codes `labelnames = ["model_name", "engine"]` as a local variable, references it across ~37 metric-family construction sites, and uses the `engine` label value in five different `.labels()` call shapes (kwarg with int engine, kwarg with str engine, positional with str engine in the middle, plus a `metrics_info["engine"] = str(...)` dict pattern). To reshape `engine` into `stage` + `replica` without forking the entire upstream `__init__`, the wrap uses three coordinated mechanisms:
|
||||
|
||||
1. **Class-level metric class slot overrides.** `OmniPrometheusStatLogger` overrides `_gauge_cls`, `_counter_cls`, `_histogram_cls` (which upstream calls via `self._gauge_cls(...)` etc.) with `_RelabelGauge` / `_RelabelCounter` / `_RelabelHistogram` wrapper classes. These intercept the `labelnames` kwarg at metric family creation time and replace `engine` with `("stage", "replica")`.
|
||||
2. **Property descriptor for `per_engine_labelvalues`.** Upstream builds `self.per_engine_labelvalues = {idx: [model_name, str(idx)]}` inside `__init__` and then captures it into a local variable for `create_metric_per_engine` calls. By making `per_engine_labelvalues` a Python property on the subclass, the setter intercepts upstream's assignment and rewrites each 2-tuple into a 3-tuple `[model_name, stage, replica]` using the `stage_replica_map` supplied at construction time. The captured local then sees the rewritten dict.
|
||||
3. **Override of `.labels()` on the wrapper classes.** For the five call sites that pass `engine` directly (kwarg or positional, int or str), `_RelabelMixin.labels()` translates the engine value back to `(stage, replica)` via a process-level `_ENGINE_INDEX_MAP` populated by `OmniPrometheusStatLogger.__init__`. This handles `gauge_engine_sleep_state.labels(engine=idx, ...)`, `counter_request_success_base.labels(model_name, str(idx), str(reason))`, `info_gauge.labels(**metrics_info)`, etc.
|
||||
|
||||
The three sub-helpers that upstream `PrometheusStatLogger.__init__` constructs (`spec_decoding_prom` / `kv_connector_prom` / `perf_metrics_prom`) use their own `_counter_cls` / `_gauge_cls` / `_histogram_cls` slots and would otherwise build families with the raw 2-element labelnames. `_OmniPerfMetricsProm` / `_OmniSpecDecodingProm` / `_OmniKVConnectorProm` subclass each helper to route the same relabel mixin through their internal family construction.
|
||||
|
||||
The `Orchestrator` constructs `stage_replica_map` from the static `stage_pools` configuration at startup:
|
||||
|
||||
```python
|
||||
stage_replica_map = {
|
||||
flat_idx: (str(stage_id), str(replica_id))
|
||||
for flat_idx, (stage_id, replica_id) in enumerate(
|
||||
(s, r)
|
||||
for s, pool in enumerate(stage_pools)
|
||||
for r in range(pool.num_replicas)
|
||||
)
|
||||
}
|
||||
```
|
||||
|
||||
A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestrator so the per-replica `record(engine_idx=...)` call site can look up the right flat index.
|
||||
|
||||
> Dynamic add/remove of replicas at runtime is intentionally out of scope — the upstream `PrometheusStatLogger` materializes per-engine_idx child metrics at init time, and supporting hot-add would require non-trivial intervention into upstream's per-family child dictionaries.
|
||||
|
||||
## Gating: `--log-stats`
|
||||
|
||||
All metrics — both the 15 `vllm:omni_*` families and the ~65 upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
|
||||
|
||||
Behavior with `--log-stats=off` (default):
|
||||
|
||||
- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm:omni_*` stay materialized but never have data written to them.
|
||||
- `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all.
|
||||
- The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate.
|
||||
|
||||
Behavior with `--log-stats=on`: all metric paths fire normally; the orchestrator's per-replica recording is bounded only by `OmniSchedulerMixin.make_stats()`'s per-scheduler 1 Hz throttle (see next section).
|
||||
|
||||
The overhead with the flag on is small enough that an A/B benchmark on Qwen3-Omni-30B single replica (30 sequential audio requests) showed a mean latency delta of +0.6% (Welch's t = 0.318, n=30, not statistically significant at α=0.05); the `/metrics` line count drops from 1358 to 124 lines when the flag is off.
|
||||
|
||||
## Throttling: `make_stats()` Override
|
||||
|
||||
Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, returning a SchedulerStats object for the orchestrator. Under vLLM's architecture, this is fine. But since vLLM-Omni requires that the object be serialized and transferred over ZMQ, receiving a SchedulerStats object on every step can introduce unacceptable overhead to the system.
|
||||
|
||||
`OmniSchedulerMixin.make_stats()` (in `vllm_omni/core/sched/omni_scheduler_mixin.py`) throttles stats emission to at most once per second **per scheduler** — i.e. per `(stage, replica)` since each replica owns its own scheduler instance. Between intervals it returns `None`, which the engine core skips serializing. This keeps gauges fresh enough for Prometheus scrapes (typically 15-30s intervals) while eliminating the per-step overhead.
|
||||
|
||||
The orchestrator side does not add its own throttle on top: the per-replica recording loop gates only on `raw_outputs.scheduler_stats is not None` (i.e. this replica's scheduler passed its own 1 Hz gate). A previous global `_last_stats_ts` on the orchestrator was removed because it starved every replica other than the first to emit within each second.
|
||||
|
||||
## Metric Definitions
|
||||
|
||||
### Pipeline (4)
|
||||
|
||||
| Metric | Type | Labels | Description |
|
||||
|--------|------|--------|-------------|
|
||||
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
|
||||
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
|
||||
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
|
||||
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
|
||||
|
||||
### Audio (7)
|
||||
|
||||
Labels: `{model_name, stage, replica}` plus the listed extra label.
|
||||
|
||||
| Metric | Type | Extra label | Description |
|
||||
|--------|------|-------------|-------------|
|
||||
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
|
||||
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
|
||||
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
|
||||
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frames generated |
|
||||
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
|
||||
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
|
||||
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
|
||||
|
||||
### Cross-stage transfer (4)
|
||||
|
||||
Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`.
|
||||
|
||||
| Metric | Type | Description |
|
||||
|--------|------|-------------|
|
||||
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
|
||||
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
|
||||
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
|
||||
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
|
||||
|
||||
### LLM stage-level (wrapped `vllm:*`)
|
||||
|
||||
After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — carries `{model_name, stage, replica}` labels. For the full upstream catalog see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md); note that metrics depending on features unsupported in vLLM-Omni (e.g. speculative decoding, LoRA) will not be available.
|
||||
|
||||
## Naming Convention
|
||||
|
||||
- All time-bearing metrics use the `_s` suffix (values in seconds). Two bucket families are used:
|
||||
- `SECONDS_BUCKETS` (0.05 s – 300 s) for e2e / generation / TTFP style values.
|
||||
- `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution.
|
||||
- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
|
||||
- Sizes use the `_bytes` suffix.
|
||||
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
|
||||
|
||||
## Logging vs. Prometheus
|
||||
|
||||
`OrchestratorAggregator` (in `vllm_omni/metrics/stats.py`) is the logging-oriented metrics path. It collects detailed per-request, per-stage, and per-transfer statistics and prints formatted tables to the `INFO` log. This is designed for development and debugging — individual request traces, transfer bandwidth, inter-stage timing.
|
||||
|
||||
`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics` form the Prometheus-oriented path. They record aggregate counters, gauges, and histograms suitable for time-series monitoring and alerting. Both paths share the same source data (`StageRequestStats`, `TransferEdgeStats`) — `OrchestratorAggregator.record_transfer_tx/rx` in particular calls both the existing accumulator code and the Prometheus emit hook in the same method body. The two consumption models can run simultaneously without coupling.
|
||||
|
||||
The separation follows upstream vLLM's pattern of `LoggingStatLogger` vs. `PrometheusStatLogger` — same underlying data, different consumption models.
|
||||
@@ -169,15 +169,13 @@ class DiffusionRequestStatus(enum.IntEnum):
|
||||
|
||||
@dataclass
|
||||
class DiffusionRequestState:
|
||||
sched_req_id: str
|
||||
request_id: str
|
||||
req: OmniDiffusionRequest
|
||||
status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
|
||||
```
|
||||
|
||||
**Design Features**:
|
||||
|
||||
- **Scheduler-owned ID**: Each `OmniDiffusionRequest` is tracked by an internal `sched_req_id`, separated from public `request_id` values.
|
||||
|
||||
- **Explicit lifecycle**: Requests move through waiting, running, optional preemption, and terminal states.
|
||||
|
||||
- **Centralized error handling**: Completion, abort, and error states are all normalized in the scheduler layer.
|
||||
@@ -188,7 +186,6 @@ class DiffusionRequestState:
|
||||
class _BaseScheduler(SchedulerInterface):
|
||||
def __init__(self) -> None:
|
||||
self._request_states = {}
|
||||
self._request_id_to_sched_req_id = {}
|
||||
self._waiting = deque()
|
||||
self._running = []
|
||||
self._finished_req_ids = set()
|
||||
@@ -199,7 +196,8 @@ class _BaseScheduler(SchedulerInterface):
|
||||
|
||||
- **Common state storage**: Shared request maps and waiting/running sets live in the base class.
|
||||
|
||||
- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
|
||||
- **Shared cleanup logic**: Duplicate-request checks, finish handling, and state removal are centralized instead of
|
||||
duplicated in each policy.
|
||||
|
||||
- **Current constraint boundary**: `_BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`, but request-mode diffusion is still clamped back to `1` by the engine. The step-wise path can keep this above `1` for compatible-request batching.
|
||||
|
||||
@@ -219,12 +217,12 @@ class RequestScheduler(_BaseScheduler):
|
||||
|
||||
- **Single-request admission**: `RequestScheduler` still admits one active request at a time because request-mode execution completes a whole request per dispatch.
|
||||
|
||||
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished scheduler ids.
|
||||
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished request ids.
|
||||
|
||||
#### 2.5 Engine-Driven Execution Loop
|
||||
|
||||
```python
|
||||
sched_req_id = scheduler.add_request(request)
|
||||
request_id = scheduler.add_request(request)
|
||||
while True:
|
||||
sched_output = scheduler.schedule()
|
||||
output = executor.add_req(req)
|
||||
|
||||
@@ -411,6 +411,8 @@ Notes:
|
||||
- `runtime.max_batch_size` controls stage-level batching.
|
||||
- Thinker/Talker commonly use `enforce_eager: false` for CUDA Graph paths.
|
||||
- Code2Wav often remains eager (`enforce_eager: true`) depending on runtime behavior.
|
||||
- Qwen3-Omni defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. The Triton has been more stable & faster
|
||||
than the FlashInfer CUTLASS unquantized MoE backend on recent vLLM rebases.
|
||||
|
||||
#### 2) Enable async chunk
|
||||
|
||||
|
||||
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
|
||||
|
||||
### Pre-built wheels
|
||||
|
||||
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, and 0.21.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
|
||||
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, and 0.22.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
|
||||
|
||||
=== "NVIDIA CUDA"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
|
||||
|
||||
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||
```bash
|
||||
uv pip install vllm==0.21.0 --torch-backend=auto
|
||||
uv pip install vllm==0.22.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.22.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.21.0 --torch-backend=auto
|
||||
uv pip install vllm==0.22.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
The 0.21.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
The 0.22.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.21.0
|
||||
git checkout v0.22.0
|
||||
```
|
||||
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
|
||||
```bash
|
||||
#For CUDA 13.0 (the default for v0.21.0; the wheel filename has no `+cu130` suffix)
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.21.0/vllm-0.21.0-cp38-abi3-manylinux_2_35_x86_64.whl
|
||||
#For CUDA 13.0 (the default for v0.22.0; the wheel filename has no `+cu130` suffix)
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_35_x86_64.whl
|
||||
```
|
||||
Install vllm with command below (If you have no existing PyTorch).
|
||||
```bash
|
||||
@@ -102,11 +102,11 @@ docker run --runtime nvidia --gpus 2 \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
--ipc=host \
|
||||
vllm/vllm-omni:v0.21.0 \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
|
||||
vllm/vllm-omni:v0.22.0 \
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
!!! tip
|
||||
You can use this docker image to serve models the same way you would with in vLLM! To do so, make sure you overwrite the default entrypoint (`vllm serve --omni`) which works only for models supported in the vLLM-Omni project.
|
||||
The CUDA image does not define a default entrypoint, so include `vllm serve ... --omni` after the image name.
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
@@ -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.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
|
||||
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
|
||||
```
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.22.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
|
||||
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
|
||||
```
|
||||
|
||||
The pre-built 0.21.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
The pre-built 0.22.0 vLLM wheel targets ROCm 7.2.2. 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.21.0
|
||||
git checkout v0.22.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.21.0 \
|
||||
vllm/vllm-omni-rocm:v0.22.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.21.0
|
||||
vllm/vllm-omni-rocm:v0.22.0
|
||||
```
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
|
||||
source .venv/bin/activate
|
||||
|
||||
# On CUDA
|
||||
uv pip install vllm==0.21.0 --torch-backend=auto
|
||||
uv pip install vllm==0.22.0 --torch-backend=auto
|
||||
|
||||
# On ROCm
|
||||
uv pip install vllm==0.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721
|
||||
uv pip install vllm==0.22.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm721
|
||||
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
@@ -35,7 +35,7 @@ For additional installation methods — please see the [installation guide](inst
|
||||
!!! note
|
||||
It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
|
||||
|
||||
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.21.0` and vLLM Omni `0.21.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
|
||||
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.22.0` and vLLM Omni `0.22.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
|
||||
|
||||
## Offline Inference
|
||||
|
||||
|
||||
@@ -119,9 +119,10 @@ def extract_omni_serve_subparser_init():
|
||||
exec_globals = {
|
||||
"_FlexibleArgumentParser": _FlexibleArgumentParser,
|
||||
"FlexibleArgumentParser": _FlexibleArgumentParser,
|
||||
"TrackingArgumentParser": _FlexibleArgumentParser, # only needs add_argument, not tracking
|
||||
"TrackingNamespace": None, # type annotation only
|
||||
"make_arg_parser": lambda parser: parser, # no-op for doc
|
||||
"_ensure_vllm_platform": lambda: None, # no-op for doc
|
||||
"nullify_stage_engine_defaults": lambda parser: None, # no-op for doc
|
||||
"VLLM_SUBCMD_PARSER_EPILOG": "",
|
||||
"logger": logger,
|
||||
"DummySubparsers": DummySubparsers,
|
||||
|
||||
@@ -16,6 +16,7 @@ ROOT_DIR_RELATIVE = "../../../../.."
|
||||
EXAMPLE_DIR = ROOT_DIR / "examples"
|
||||
EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
|
||||
NAV_FILE = ROOT_DIR / "docs/.nav.yml"
|
||||
MAX_INLINE_MATERIAL_SIZE = 8 * 1024
|
||||
|
||||
|
||||
def fix_case(text: str) -> str:
|
||||
@@ -191,18 +192,21 @@ class Example:
|
||||
|
||||
return re.sub(link_pattern, replace_link, content)
|
||||
|
||||
def github_url(self, path: Path) -> str:
|
||||
url = "https://github.com/vllm-project/vllm-omni/"
|
||||
url += "tree/main" if path.is_dir() else "blob/main"
|
||||
return f"{url}/{path.relative_to(ROOT_DIR).as_posix()}"
|
||||
|
||||
def generate(self) -> str:
|
||||
content = f"# {self.title}\n\n"
|
||||
url = "https://github.com/vllm-project/vllm-omni/"
|
||||
url += "tree/main" if self.path.is_dir() else "blob/main"
|
||||
content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"
|
||||
content += f"Source <{self.github_url(self.path)}>.\n\n"
|
||||
|
||||
# Use long code fence to avoid issues with
|
||||
# included files containing code fences too
|
||||
code_fence = "``````"
|
||||
|
||||
if self.is_code:
|
||||
main_file_rel = self.main_file.relative_to(ROOT_DIR)
|
||||
main_file_rel = self.main_file.relative_to(ROOT_DIR).as_posix()
|
||||
content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel}"\n{code_fence}\n'
|
||||
else:
|
||||
with open(self.main_file, encoding="utf-8") as f:
|
||||
@@ -216,10 +220,15 @@ class Example:
|
||||
|
||||
content += "## Example materials\n\n"
|
||||
for file in sorted(self.other_files):
|
||||
content += f'??? abstract "{file.relative_to(self.path)}"\n'
|
||||
content += f'??? abstract "{file.relative_to(self.path).as_posix()}"\n'
|
||||
if file.stat().st_size > MAX_INLINE_MATERIAL_SIZE:
|
||||
content += (
|
||||
f" Large file omitted from the rendered docs. View it on GitHub: <{self.github_url(file)}>.\n\n"
|
||||
)
|
||||
continue
|
||||
if file.suffix != ".md":
|
||||
content += f" {code_fence}{file.suffix[1:]}\n"
|
||||
content += f' --8<-- "{file.relative_to(ROOT_DIR)}"\n'
|
||||
content += f' --8<-- "{file.relative_to(ROOT_DIR).as_posix()}"\n'
|
||||
if file.suffix != ".md":
|
||||
content += f" {code_fence}\n"
|
||||
|
||||
@@ -291,7 +300,7 @@ def update_nav_file(examples: list[Example]):
|
||||
for example in category_examples:
|
||||
doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
|
||||
rel_path = doc_path.relative_to(ROOT_DIR / "docs")
|
||||
category_items.append({example.title: str(rel_path)})
|
||||
category_items.append({example.title: rel_path.as_posix()})
|
||||
|
||||
if category_items:
|
||||
# Format category name (e.g., "offline_inference" -> "Offline Inference")
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Reduce browser search-index memory for generated API reference pages."""
|
||||
|
||||
from material.plugins.search.plugin import SearchIndex
|
||||
|
||||
_add_entry_from_context = SearchIndex.add_entry_from_context
|
||||
|
||||
|
||||
def add_entry_from_context(self: SearchIndex, page) -> None:
|
||||
search = page.meta.get("search") or {}
|
||||
if search.get("exclude"):
|
||||
return
|
||||
|
||||
if str(page.url).startswith("api/"):
|
||||
entry = {
|
||||
"location": page.url,
|
||||
"title": str(page.meta.get("title", page.title)),
|
||||
"text": "",
|
||||
}
|
||||
if "boost" in search:
|
||||
entry["boost"] = search["boost"]
|
||||
self.entries.append(entry)
|
||||
return
|
||||
|
||||
_add_entry_from_context(self, page)
|
||||
|
||||
|
||||
SearchIndex.add_entry_from_context = add_entry_from_context
|
||||
@@ -1,20 +1,44 @@
|
||||
// Enables MathJax rendering
|
||||
// Enables MathJax rendering only on pages that contain math.
|
||||
window.MathJax = {
|
||||
tex: {
|
||||
inlineMath: [["\\(", "\\)"]],
|
||||
displayMath: [["\\[", "\\]"]],
|
||||
processEscapes: true,
|
||||
processEnvironments: true
|
||||
},
|
||||
options: {
|
||||
ignoreHtmlClass: ".*|",
|
||||
processHtmlClass: "arithmatex"
|
||||
}
|
||||
};
|
||||
tex: {
|
||||
inlineMath: [["\\(", "\\)"]],
|
||||
displayMath: [["\\[", "\\]"]],
|
||||
processEscapes: true,
|
||||
processEnvironments: true,
|
||||
},
|
||||
options: {
|
||||
ignoreHtmlClass: ".*|",
|
||||
processHtmlClass: "arithmatex",
|
||||
},
|
||||
};
|
||||
|
||||
document$.subscribe(() => {
|
||||
MathJax.startup.output.clearCache()
|
||||
MathJax.typesetClear()
|
||||
MathJax.texReset()
|
||||
MathJax.typesetPromise()
|
||||
})
|
||||
let mathJaxLoading;
|
||||
|
||||
function loadMathJax() {
|
||||
if (window.MathJax.typesetPromise) {
|
||||
return Promise.resolve();
|
||||
}
|
||||
if (!mathJaxLoading) {
|
||||
mathJaxLoading = new Promise((resolve, reject) => {
|
||||
const script = document.createElement("script");
|
||||
script.src = "https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js";
|
||||
script.async = true;
|
||||
script.onload = resolve;
|
||||
script.onerror = reject;
|
||||
document.head.appendChild(script);
|
||||
});
|
||||
}
|
||||
return mathJaxLoading;
|
||||
}
|
||||
|
||||
document$.subscribe(() => {
|
||||
if (!document.querySelector(".arithmatex")) {
|
||||
return;
|
||||
}
|
||||
loadMathJax().then(() => {
|
||||
MathJax.startup.output.clearCache();
|
||||
MathJax.typesetClear();
|
||||
MathJax.texReset();
|
||||
MathJax.typesetPromise();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,2 @@
|
||||
// Opt out of bfcache so large generated API pages are released after navigation.
|
||||
window.addEventListener("unload", () => {});
|
||||
@@ -1,21 +1,42 @@
|
||||
// Initialize Mermaid for diagram rendering
|
||||
mermaid.initialize({
|
||||
startOnLoad: false,
|
||||
theme: 'default',
|
||||
securityLevel: 'loose',
|
||||
flowchart: {
|
||||
useMaxWidth: true,
|
||||
htmlLabels: true
|
||||
}
|
||||
});
|
||||
// Load Mermaid only on pages that contain diagrams.
|
||||
let mermaidLoading;
|
||||
|
||||
// Render Mermaid diagrams when page content is ready
|
||||
document$.subscribe(() => {
|
||||
const mermaidElements = document.querySelectorAll('.mermaid');
|
||||
if (mermaidElements.length > 0) {
|
||||
mermaid.run({
|
||||
querySelector: '.mermaid',
|
||||
nodes: mermaidElements
|
||||
function loadMermaid() {
|
||||
if (window.mermaid) {
|
||||
return Promise.resolve();
|
||||
}
|
||||
if (!mermaidLoading) {
|
||||
mermaidLoading = new Promise((resolve, reject) => {
|
||||
const script = document.createElement("script");
|
||||
script.src = "https://unpkg.com/mermaid@10/dist/mermaid.min.js";
|
||||
script.async = true;
|
||||
script.onload = () => {
|
||||
mermaid.initialize({
|
||||
startOnLoad: false,
|
||||
theme: "default",
|
||||
securityLevel: "loose",
|
||||
flowchart: {
|
||||
useMaxWidth: true,
|
||||
htmlLabels: true,
|
||||
},
|
||||
});
|
||||
resolve();
|
||||
};
|
||||
script.onerror = reject;
|
||||
document.head.appendChild(script);
|
||||
});
|
||||
}
|
||||
return mermaidLoading;
|
||||
}
|
||||
|
||||
document$.subscribe(() => {
|
||||
if (!document.querySelector(".mermaid")) {
|
||||
return;
|
||||
}
|
||||
loadMermaid().then(() => {
|
||||
mermaid.run({
|
||||
querySelector: ".mermaid",
|
||||
nodes: document.querySelectorAll(".mermaid"),
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -19,6 +19,7 @@ th {
|
||||
|--------------|--------|-------------------|------------|---------|-----|-----------|
|
||||
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | |
|
||||
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
|
||||
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
@@ -32,6 +33,7 @@ th {
|
||||
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
|
||||
| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
@@ -40,13 +42,14 @@ th {
|
||||
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||
| `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | | | |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
|
||||
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
|
||||
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | |
|
||||
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
|
||||
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
|
||||
@@ -58,6 +61,7 @@ th {
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
|
||||
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
|
||||
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
|
||||
@@ -71,4 +75,6 @@ th {
|
||||
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
|
||||
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
|
||||
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
|
||||
|
||||
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
|
||||
@@ -13,7 +13,7 @@ For example...
|
||||
```bash
|
||||
# Qwen-Image
|
||||
vllm serve Qwen/Qwen-Image-Edit-2511 --omni --port 8000
|
||||
|
||||
```
|
||||
|
||||
### Generate Images
|
||||
|
||||
@@ -53,7 +53,7 @@ result = client.images.edit(
|
||||
output_format='jpeg',
|
||||
# url格式
|
||||
extra_body={
|
||||
"url": [input_image_url1,input_image_url],
|
||||
"url": [input_image_url,input_image_url],
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1,
|
||||
"seed": 777,
|
||||
@@ -92,7 +92,8 @@ Content-Type: multipart/form-data
|
||||
| `user` | string | null | User identifier for tracking |
|
||||
| `output_format` | string | "png" | The format in which the generated images are returned. Must be one of "png", "jpg", "jpeg", "webp". |
|
||||
| `output_compression` | integer | 100 | The compression level (0-100%) for the generated images. |
|
||||
| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s).
|
||||
| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s). |
|
||||
| `stream` | boolean | false | Return Server-Sent Events instead of a single JSON response. Currently supported only by multi-stage image edit pipelines that produce AR text before the final image, such as HunyuanImage3 IT2I AR+DiT. |
|
||||
|
||||
#### vllm-omni Extension Parameters
|
||||
|
||||
@@ -109,6 +110,8 @@ Content-Type: multipart/form-data
|
||||
|
||||
### Response Format
|
||||
|
||||
When `stream=false` or omitted, the endpoint returns the standard image edit response:
|
||||
|
||||
```json
|
||||
{
|
||||
"created": 1701234567,
|
||||
@@ -120,10 +123,56 @@ Content-Type: multipart/form-data
|
||||
}
|
||||
],
|
||||
"output_format": null,
|
||||
"size": null,
|
||||
"size": null
|
||||
}
|
||||
```
|
||||
|
||||
### Streaming Response Format
|
||||
|
||||
Set `stream=true` when you want to receive the HunyuanImage3 IT2I AR
|
||||
recaption text before the final edited image is ready. Streaming is only
|
||||
available for multi-stage image edit pipelines where stage 0 produces AR text
|
||||
and a later diffusion stage produces the image. Single-stage diffusion image
|
||||
edit pipelines reject `stream=true` with HTTP 400.
|
||||
|
||||
The response uses Server-Sent Events with `Content-Type: text/event-stream`.
|
||||
Each event is sent as one `data:` line. The stream order is:
|
||||
|
||||
1. One or more AR text delta chunks.
|
||||
2. One final image chunk.
|
||||
3. `data: [DONE]`.
|
||||
|
||||
AR text chunks expose only the generated text delta and the AR completion
|
||||
`index`. They do not expose token ids, ratio tokens, KV cache data, or internal
|
||||
prompt token ids.
|
||||
|
||||
Example AR delta event:
|
||||
|
||||
```text
|
||||
data: {"object":"image.edit.chunk","type":"ar_delta","delta":"A close-up product photo...","index":0,"created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct"}
|
||||
```
|
||||
|
||||
Example final image event:
|
||||
|
||||
```text
|
||||
data: {"object":"image.edit.chunk","type":"image","data":[{"b64_json":"<base64-encoded PNG>","url":null,"revised_prompt":null}],"output_format":"png","size":"1024x1024","created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct"}
|
||||
```
|
||||
|
||||
Terminal event:
|
||||
|
||||
```text
|
||||
data: [DONE]
|
||||
```
|
||||
|
||||
If the engine fails after the SSE response has started, the stream emits an
|
||||
error chunk followed by `[DONE]`:
|
||||
|
||||
```text
|
||||
data: {"object":"error","created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct","error":{"message":"<error message>","type":"server_error","code":500}}
|
||||
|
||||
data: [DONE]
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
### Multiple Images input
|
||||
@@ -134,12 +183,28 @@ curl -s -D >(grep -i x-request-id >&2) \
|
||||
-X POST "http://localhost:8000/v1/images/edits" \
|
||||
-F "model=xxx" \
|
||||
-F "image=@xx.png" \
|
||||
-F "image=@xx.png"
|
||||
-F "image=@xx.png" \
|
||||
-F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
|
||||
-F "size=1024x1024" \
|
||||
-F "output_format=png"
|
||||
```
|
||||
|
||||
### Streaming HunyuanImage3 IT2I AR Text
|
||||
|
||||
```bash
|
||||
curl -N -X POST "http://localhost:8000/v1/images/edits" \
|
||||
-F "model=tencent/HunyuanImage-3.0-Instruct" \
|
||||
-F "image=@./input.png" \
|
||||
-F "prompt=Turn the product into a clean studio advertisement." \
|
||||
-F "size=1024x1024" \
|
||||
-F "output_format=png" \
|
||||
-F "stream=true"
|
||||
```
|
||||
|
||||
Use the `ar_delta` events for progressive display of the AR-generated
|
||||
recaption. Decode the `b64_json` field from the final `image` event to get the
|
||||
edited image.
|
||||
|
||||
|
||||
## Parameter Handling
|
||||
|
||||
|
||||
@@ -5,6 +5,9 @@ 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.
|
||||
- **CosyVoice3** (`FunAudioLLM/Fun-CosyVoice3-0.5B-2512`) -- 2-stage talker + flow-matching code2wav. Voice cloning via `ref_audio` + `ref_text` (no presets). Output: 24 kHz.
|
||||
|
||||
See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, and MOSS-TTS-Nano.
|
||||
|
||||
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
|
||||
|
||||
@@ -26,6 +29,10 @@ vllm serve fishaudio/s2-pro --omni --port 8091
|
||||
|
||||
# Voxtral TTS
|
||||
vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
|
||||
|
||||
# CosyVoice3 (voice cloning only — supply ref_audio + ref_text per request)
|
||||
vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 \
|
||||
--omni --port 8091 --trust-remote-code
|
||||
```
|
||||
|
||||
### Generate Speech
|
||||
@@ -374,6 +381,44 @@ are cached in-process with a shared LRU so repeated requests with the same
|
||||
all TTS model types; deleting a voice invalidates every model-type slot at
|
||||
once.
|
||||
|
||||
### Precomputed Custom Voices
|
||||
|
||||
Qwen3-TTS Base and VoxCPM2 can load offline-precomputed voices at startup.
|
||||
Generate a directory containing `custom_voice_manifest.json` plus one
|
||||
`.safetensors` file per voice, then set the pipeline-wide deploy config field:
|
||||
|
||||
```yaml
|
||||
custom_voice_dir: /path/to/custom_voices
|
||||
```
|
||||
|
||||
Qwen3-TTS profiles are created with:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/qwen3_tts/precompute_custom_voice.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--voice-name alice \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Original transcript of the reference audio" \
|
||||
--mode icl \
|
||||
--output-dir /path/to/custom_voices
|
||||
```
|
||||
|
||||
VoxCPM2 profiles are created with:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/voxcpm2/precompute_custom_voice.py \
|
||||
--model openbmb/VoxCPM2 \
|
||||
--voice-name alice \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--mode ref_continuation \
|
||||
--prompt-text "Original transcript of the reference audio" \
|
||||
--output-dir /path/to/custom_voices
|
||||
```
|
||||
|
||||
Only profiles whose safetensors payload can be loaded and validated are exposed
|
||||
by `GET /v1/audio/voices`. Valid precomputed voices can be used in
|
||||
`POST /v1/audio/speech` by passing `voice="alice"` without `ref_audio`.
|
||||
|
||||
**Configuration (environment variables):**
|
||||
|
||||
| Variable | Default | Description |
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
# Production Metrics
|
||||
|
||||
vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. This page covers the text and audio surface; diffusion / image / video metrics are tracked in a follow-up PR.
|
||||
|
||||
```bash
|
||||
vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats
|
||||
curl http://localhost:8000/metrics
|
||||
```
|
||||
|
||||
**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm:omni_*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
|
||||
|
||||
## Metric Namespaces
|
||||
|
||||
| Prefix | Source | Present when |
|
||||
|--------|--------|--------------|
|
||||
| `vllm:omni_` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
|
||||
| `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage |
|
||||
| `http_` / `process_` | Uvicorn / Python runtime | Always |
|
||||
|
||||
## Pipeline-Level Metrics (`vllm:omni_`)
|
||||
|
||||
Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline.
|
||||
|
||||
### Request counts
|
||||
|
||||
| Metric | Type | Labels | Description |
|
||||
|--------|------|--------|-------------|
|
||||
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
|
||||
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
|
||||
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
|
||||
|
||||
### Latency
|
||||
|
||||
| Metric | Type | Labels | Description |
|
||||
|--------|------|--------|-------------|
|
||||
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
|
||||
|
||||
## Audio Modality Metrics (`vllm:omni_`)
|
||||
|
||||
Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label.
|
||||
|
||||
| Metric | Type | Extra label | Description |
|
||||
|--------|------|-------------|-------------|
|
||||
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
|
||||
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
|
||||
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
|
||||
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
|
||||
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
|
||||
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
|
||||
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
|
||||
|
||||
The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition.
|
||||
|
||||
## Cross-Stage Transfer Metrics (`vllm:omni_`)
|
||||
|
||||
Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed.
|
||||
|
||||
| Metric | Type | Description |
|
||||
|--------|------|-------------|
|
||||
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
|
||||
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
|
||||
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
|
||||
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
|
||||
|
||||
## vLLM Engine Metrics (`vllm:`)
|
||||
|
||||
When the pipeline includes an LLM stage, the upstream vLLM engine exposes its full set of ~37 metric families under the `vllm:` prefix.
|
||||
|
||||
vLLM-Omni wraps the upstream `vllm.v1.metrics.loggers.PrometheusStatLogger` with `OmniPrometheusStatLogger` so that the original `engine` single label is reshaped into `stage` + `replica`. Every `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — therefore gains per-`(stage, replica)` visibility automatically. No omni-side duplicate is needed for the text path.
|
||||
|
||||
```text
|
||||
# Before wrap:
|
||||
vllm:num_requests_running{model_name="...", engine="1"} 3.0
|
||||
|
||||
# After wrap:
|
||||
vllm:num_requests_running{model_name="...", stage="1", replica="0"} 2.0
|
||||
vllm:num_requests_running{model_name="...", stage="1", replica="1"} 1.0
|
||||
```
|
||||
|
||||
For the full list of upstream metrics, see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md).
|
||||
|
||||
## Metric Availability by Pipeline Type
|
||||
|
||||
| Metric group | Multi-stage LLM (Qwen3-Omni) |
|
||||
|---|---|
|
||||
| `vllm:omni_` request tracking + latency | With `--log-stats` |
|
||||
| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage |
|
||||
| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
|
||||
| `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` |
|
||||
| `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` |
|
||||
|
||||
## Naming Convention
|
||||
|
||||
- All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values.
|
||||
- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
|
||||
- Sizes use the `_bytes` suffix.
|
||||
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
|
||||
- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label.
|
||||
@@ -1,28 +1,26 @@
|
||||
# Diffusion Attention Backends
|
||||
|
||||
This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, and how to use SageAttention.
|
||||
This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, the per-platform defaults, and how to use SageAttention.
|
||||
|
||||
## Overview
|
||||
|
||||
Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured.
|
||||
|
||||
This backend is used by diffusion attention layers such as the DiT attention in video and image generation models.
|
||||
This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector.
|
||||
|
||||
On CUDA, the practical choices today are:
|
||||
|
||||
- `FLASH_ATTN`: FlashAttention backend. This is the default on supported CUDA systems when FlashAttention is installed.
|
||||
- `TORCH_SDPA`: PyTorch `scaled_dot_product_attention`.
|
||||
- `SAGE_ATTN`: SageAttention backend, if `sageattention` is installed.
|
||||
|
||||
If no attention backend is configured, vLLM-Omni asks the current platform to choose the default. On CUDA, that normally means `FLASH_ATTN` when available, otherwise `TORCH_SDPA`.
|
||||
The full set of backends and their platform defaults is in the **Backend Options** and **Platform Defaults** sections below. If no attention backend is configured, vLLM-Omni asks the current platform to choose the default.
|
||||
|
||||
## Backend Options
|
||||
|
||||
| Value | Notes |
|
||||
|---|---|
|
||||
| `FLASH_ATTN` | Default on CUDA when FlashAttention is available. Good default for most diffusion workloads. |
|
||||
| `TORCH_SDPA` | Most conservative fallback. Useful for debugging or compatibility. |
|
||||
| `SAGE_ATTN` | Requires `sageattention`. Can improve performance on some workloads, but output quality must be validated model-by-model. |
|
||||
| `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. |
|
||||
| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
|
||||
| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
|
||||
| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
|
||||
| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
|
||||
| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
|
||||
|
||||
|
||||
## Configuration
|
||||
|
||||
@@ -95,6 +93,80 @@ config = OmniDiffusionConfig(
|
||||
|
||||
A plain dict is also accepted and normalized to `AttentionConfig`.
|
||||
|
||||
## Platform Defaults
|
||||
|
||||
### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
|
||||
|
||||
Auto-route preference, in order:
|
||||
|
||||
1. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
|
||||
2. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
|
||||
3. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
|
||||
4. `TORCH_SDPA` — last resort
|
||||
|
||||
The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
|
||||
|
||||
**Why CUDNN_ATTN by default**: on mask-heavy diffusion models (HunyuanVideo-1.5, Qwen-Image), cuDNN's pinned FMHA kernel sidesteps a PyTorch SDPA dispatch quirk where the unpinned dispatcher picks `EFFICIENT_ATTENTION` (~25 ms) for masked calls instead of cuDNN (~11 ms). The pin gives 2× e2e on HV-1.5 with no regression on lighter models.
|
||||
|
||||
### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86)
|
||||
|
||||
Auto-route preference:
|
||||
|
||||
1. `FLASH_ATTN` — when `flash-attn` is installed
|
||||
2. `TORCH_SDPA` — fallback
|
||||
|
||||
`CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware.
|
||||
|
||||
## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
|
||||
|
||||
Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
|
||||
|
||||
| Model | Shape | TORCH_SDPA | CUDNN_ATTN | FLASHINFER_ATTN |
|
||||
|---|---|---|---|---|
|
||||
| HunyuanVideo-1.5 (T2V) | 480p / 33f / 50 steps | 147.05 s | **73.02 s** | 127.84 s |
|
||||
| Wan 2.2 14B (T2V) | 480p / 33f / 40 steps | 117.75 s | 117.17 s | **115.07 s** |
|
||||
| Qwen-Image (T2I) | 1024² / 50 steps | 17.41 s | **15.14 s** | 16.02 s |
|
||||
| FLUX.2-dev (T2I) | 1024² / 50 steps, TP=2 | 53.62 s | **53.30 s** | 54.94 s |
|
||||
|
||||
Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
|
||||
|
||||
## Known Limitations
|
||||
|
||||
### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
|
||||
|
||||
LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
|
||||
|
||||
**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Lightricks/LTX-2 ...
|
||||
```
|
||||
|
||||
### FA4 not yet integrated
|
||||
|
||||
FlashAttention-4 (released March 2026) targets Blackwell natively and reportedly beats cuDNN by ~20% on B200. As of this writing the `flash-attn-4 4.0.0b10` wheel crashes with `AttributeError: 'NoneType' object has no attribute '_trait'` during JIT on sm_120. Not yet wired into vLLM-Omni; revisit when stable lands.
|
||||
|
||||
## Choosing a Backend Manually
|
||||
|
||||
### When to override the default
|
||||
|
||||
- **Quality validation**: compare a new backend against `TORCH_SDPA` as the reference, since SDPA's default dispatcher is the most extensively tested.
|
||||
- **Lossy speedup hunting**: try `SAGE_ATTN` (INT8 quantized) on diffusion outputs — typically indistinguishable visually but always validate.
|
||||
- **Workaround for known issues**: see Known Limitations above.
|
||||
|
||||
### Verifying which backend is in use
|
||||
|
||||
The startup log prints one of:
|
||||
|
||||
```
|
||||
Using diffusion attention backend 'CUDNN_ATTN' # explicit override
|
||||
Defaulting to diffusion attention backend CUDNN_ATTN ... # auto-route
|
||||
Defaulting to diffusion attention backend SDPA # nothing else available
|
||||
```
|
||||
|
||||
If you don't see one of these, the model didn't reach diffusion stage init — check earlier logs for failures.
|
||||
|
||||
## SageAttention Installation
|
||||
|
||||
vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
|
||||
@@ -115,11 +187,56 @@ Quick check:
|
||||
python -c "import sageattention; print(sageattention.__file__)"
|
||||
```
|
||||
|
||||
## Usage
|
||||
## SageAttention3 Installation
|
||||
|
||||
### Enable SageAttention
|
||||
vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni.
|
||||
|
||||
Example: HunyuanVideo-1.5 text-to-video
|
||||
Build from source:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/thu-ml/SageAttention.git
|
||||
cd SageAttention/sageattention3_blackwell
|
||||
python setup.py install
|
||||
```
|
||||
|
||||
Quick check:
|
||||
|
||||
```bash
|
||||
python -c "import sageattn3; print(sageattn3.__file__)"
|
||||
```
|
||||
|
||||
Notes:
|
||||
|
||||
- `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class.
|
||||
- SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness.
|
||||
|
||||
## Usage Examples
|
||||
|
||||
### Default (auto-route)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 50 --seed 42 --guidance-scale 6.0 \
|
||||
--output hv15.mp4
|
||||
```
|
||||
|
||||
On Blackwell this picks `CUDNN_ATTN` automatically. Check the log for the `Defaulting to ...` line.
|
||||
|
||||
### Explicit backend selection
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Lightricks/LTX-2 \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 40 --seed 42 --guidance-scale 4.0 \
|
||||
--output ltx2.mp4
|
||||
```
|
||||
|
||||
### SageAttention (lossy)
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
@@ -128,7 +245,7 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output ../tmp/hv15_modelopt_sage.mp4
|
||||
--output hv15_sage.mp4
|
||||
```
|
||||
|
||||
Example: Wan2.2 TI2V 5B
|
||||
@@ -143,6 +260,20 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
|
||||
--output outputs/wan22_sage.mp4
|
||||
```
|
||||
|
||||
### Enable SageAttention3
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/hv15_sage3.mp4
|
||||
```
|
||||
|
||||
### Mixed backends across roles
|
||||
|
||||
Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention:
|
||||
@@ -173,13 +304,13 @@ python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
|
||||
## Validation Guidance
|
||||
|
||||
Do not assume that a faster attention backend is numerically interchangeable with `FLASH_ATTN`.
|
||||
Don't assume a faster attention backend is numerically interchangeable with `TORCH_SDPA`.
|
||||
|
||||
Always compare:
|
||||
|
||||
- End-to-end runtime
|
||||
- DiT / diffusion stage runtime
|
||||
- Output quality against a known-good baseline
|
||||
- Diffusion-stage runtime (`add_req_and_wait` line in DiffusionEngine.step breakdown)
|
||||
- Output quality against a known-good baseline (CLIP similarity, frame-level diff, or visual review)
|
||||
|
||||
At minimum, keep the same:
|
||||
|
||||
@@ -187,6 +318,19 @@ At minimum, keep the same:
|
||||
- prompt
|
||||
- seed
|
||||
- resolution
|
||||
- frame count
|
||||
- inference steps
|
||||
- parallel config
|
||||
- frame count / step count
|
||||
- parallel config (TP / CFG-parallel / Ulysses degrees)
|
||||
|
||||
## Reproducing the Benchmark Table
|
||||
|
||||
The end-to-end numbers above were collected by running `text_to_video.py` /
|
||||
`text_to_image.py` with the same prompt and seed while varying
|
||||
`DIFFUSION_ATTENTION_BACKEND`. For a quick kernel-level comparison of the
|
||||
backends without loading a model:
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/bench_attention_backends.py --preset hv15
|
||||
```
|
||||
|
||||
It runs all three BF16 backends on representative DiT attention shapes and
|
||||
prints a ranking table at the end.
|
||||
@@ -60,6 +60,9 @@ batching only becomes relevant when `--max-num-seqs > 1`.
|
||||
- `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.
|
||||
- LoRA is supported in step mode, but each batch must use a single adapter:
|
||||
requests with different `lora_request` or `lora_scale` are scheduled into
|
||||
separate batches.
|
||||
|
||||
## When To Use It
|
||||
|
||||
|
||||
@@ -113,8 +113,9 @@ The following tables show which models support each feature:
|
||||
| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.2-dev** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Hidream-I1-Full** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
@@ -128,26 +129,31 @@ The following tables show which models support each feature:
|
||||
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
|
||||
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **SenseNova-U1** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **SenseNova-U1** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
|
||||
| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
|
||||
|
||||
> Notes:
|
||||
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
|
||||
> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
|
||||
> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload is not supported; use layerwise offload.
|
||||
|
||||
### VideoGen
|
||||
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.2-S2V** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
|
||||
|
||||
**Frame Interpolation Support**
|
||||
|
||||
@@ -184,7 +190,7 @@ The following tables show which models support each feature:
|
||||
| **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | | |
|
||||
| **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | | |
|
||||
| **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | |
|
||||
| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ❌ | |
|
||||
| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ✅ | |
|
||||
|
||||
!!! info
|
||||
|
||||
@@ -193,7 +199,7 @@ The following tables show which models support each feature:
|
||||
3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
|
||||
4. CPU Offloading (Layerwise) supports single-card for now.
|
||||
5. Using FP8-Quant as an example of qunatization methods.
|
||||
6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT) or LoRA.
|
||||
6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT). LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
|
||||
|
||||
|
||||
## Multi-Thread Weight Loading
|
||||
|
||||
@@ -17,12 +17,13 @@ list of supported architectures across all modalities, see
|
||||
| Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
|
||||
|---|---|---|---|---|---|---|
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | — | continuation (`--ref-audio` + `--ref-text`) | 48 kHz |
|
||||
| VoxCPM | local model dir | split | ✓ | ✓ (`voxcpm_async_chunk.yaml`) | — | 24 kHz |
|
||||
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ (`async_chunk: true` default) | — | 22.05 kHz |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ (`--streaming`) | — | 44.1 kHz |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design (`--instruct`), language (`--lang`) | 24 kHz |
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants (`--query-type`) | 24 kHz |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets (`--voice`) | 24 kHz |
|
||||
| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` (+ `k2-fsa/OmniVoice/audio_tokenizer/` codec) | 2 (talker + code2wav, DualFFN) | ✓ (`--ref-audio` + `--ref-text`) | — | — | 24 kHz |
|
||||
|
||||
## Common Quick Start
|
||||
|
||||
@@ -72,61 +73,6 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
|
||||
|
||||
---
|
||||
|
||||
## VoxCPM
|
||||
|
||||
Split-stage TTS. The hub example covers single TTS, single voice cloning, and streaming. Use `benchmarks/voxcpm/` for warmup, batch JSONL prompts, profiler injection, and offline TTFP / RTF measurement.
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
pip install voxcpm soundfile
|
||||
# or use a local source tree:
|
||||
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
|
||||
```
|
||||
|
||||
If the native VoxCPM `config.json` does not contain HF metadata such as `model_type`, prepare a persistent HF-compatible config directory and point the stage configs to it via `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH`:
|
||||
|
||||
```bash
|
||||
export VOXCPM_MODEL=/path/to/voxcpm-model
|
||||
export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
|
||||
mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
|
||||
cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
|
||||
cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
|
||||
python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
|
||||
```
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
|
||||
--model "$VOXCPM_MODEL" \
|
||||
--text "This is a split-stage VoxCPM synthesis example running on vLLM Omni."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
|
||||
--model "$VOXCPM_MODEL" \
|
||||
--text "This sentence is synthesized with a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "The exact transcript spoken in reference.wav."
|
||||
```
|
||||
|
||||
### Streaming
|
||||
Pass the async-chunk stage config:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
|
||||
--model "$VOXCPM_MODEL" \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
|
||||
--text "This is a split-stage VoxCPM streaming example running on vLLM Omni."
|
||||
```
|
||||
|
||||
### Notes
|
||||
- `voxcpm.yaml` is the default non-streaming stage config; `voxcpm_async_chunk.yaml` enables streaming.
|
||||
- Streaming is currently single-request oriented.
|
||||
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades quality.
|
||||
- For online serving, see [`examples/online_serving/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/voxcpm/README.md). For benchmark reporting, see [`benchmarks/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/benchmarks/voxcpm/README.md).
|
||||
|
||||
---
|
||||
|
||||
## CosyVoice3
|
||||
|
||||
2-stage TTS pipeline (`talker` + `code2wav`) at 22.05 kHz.
|
||||
@@ -214,6 +160,35 @@ Streaming requires `async_chunk: true` in the stage config.
|
||||
|
||||
---
|
||||
|
||||
## GLM-TTS
|
||||
|
||||
2-stage TTS pipeline (AR + DiT flow-matching) at 24 kHz. Every request requires reference audio and its transcript for zero-shot voice cloning.
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
|
||||
--model zai-org/GLM-TTS \
|
||||
--text "你好,这是语音合成测试。" \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "这是参考音频的文本内容。" \
|
||||
--output-dir ./output
|
||||
```
|
||||
|
||||
### Architecture
|
||||
```
|
||||
Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz)
|
||||
(Llama-based) (32k vocab) (Flow Matching)
|
||||
```
|
||||
|
||||
### Notes
|
||||
- `--ref-audio` and `--ref-text` are **required** together; GLM-TTS does not support text-only synthesis.
|
||||
- Reference audio should be 3-10 seconds.
|
||||
- First run may be slow due to lazy loading of WhisperVQ tokenizer and CampPlus ONNX speaker embedder.
|
||||
- Default sampling: temperature=1.0, top_k=25, top_p=0.8 (RAS method).
|
||||
- The `--model` path should point to the repository root (not `llm/` subdirectory).
|
||||
|
||||
---
|
||||
|
||||
## OmniVoice
|
||||
|
||||
Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design).
|
||||
@@ -359,6 +334,46 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
|
||||
- `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
|
||||
- Run `--help` for the full argument surface.
|
||||
|
||||
---
|
||||
|
||||
## higgs-audio v2
|
||||
|
||||
2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) over the shared-memory connector. Stage 1 builds on the HiggsAudio decoder kernel at `vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py`.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Voice clone uses HF's `HiggsAudioV2TokenizerModel`, instantiated from `k2-fsa/OmniVoice/audio_tokenizer/` — the boson-ai standalone tokenizer Hub repo's `model.safetensors` is actually the 3B talker LM, so we point HF at k2's repackaged codec weights instead. Only the `audio_tokenizer/` subdir (~806 MB) is downloaded.
|
||||
|
||||
```bash
|
||||
pip install -U "transformers>=5.3.0"
|
||||
```
|
||||
|
||||
### Quick start (plain TTS)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
|
||||
--texts "Hello world." "The quick brown fox jumps over the lazy dog." \
|
||||
--output-dir results/wavs
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
|
||||
Pass both `--ref-audio` and `--ref-text` together:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
|
||||
--texts "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Exact transcript spoken in reference.wav." \
|
||||
--output-dir results/wavs
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
|
||||
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
|
||||
- Out of scope (rejected with 4xx by the request validator): multi-speaker `[SPEAKERn]` dialogue, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "cosyvoice3/end2end.py"
|
||||
@@ -369,6 +384,13 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/fish_speech/end2end.py"
|
||||
``````
|
||||
??? abstract "higgs_audio_v2/end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py"
|
||||
??? abstract "glm_tts/end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/glm_tts/end2end.py"
|
||||
``````
|
||||
??? abstract "omnivoice/end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/omnivoice/end2end.py"
|
||||
@@ -381,10 +403,6 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/qwen3_tts/end2end.py"
|
||||
``````
|
||||
??? abstract "voxcpm/end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/voxcpm/end2end.py"
|
||||
``````
|
||||
??? abstract "voxcpm2/end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/voxcpm2/end2end.py"
|
||||
|
||||
@@ -17,6 +17,9 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
|
||||
The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded
|
||||
automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies.
|
||||
The bundled Qwen3-Omni setup defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. This keeps the Thinker & Talker on vLLM's
|
||||
Triton unquantized MoE path and avoids the performance regression observed with the FlashInfer CUTLASS unquantized MoE
|
||||
backend.
|
||||
Asynchronous chunk streaming is **enabled by default** within the bundled configuration.
|
||||
|
||||
To explicitly utilize a custom deployment YAML, specify the configuration path:
|
||||
@@ -72,6 +75,12 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
|
||||
```
|
||||
|
||||
To experiment with the FlashInfer FP16 MoE path, set `VLLM_USE_FLASHINFER_MOE_FP16=1` before launching the server:
|
||||
```bash
|
||||
VLLM_USE_FLASHINFER_MOE_FP16=1 \
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
For the stage-based CLI, you usually do **not** need `--stage-overrides` for
|
||||
that kind of change. Since each command launches one stage, just pass the knob
|
||||
directly on that stage command:
|
||||
|
||||
@@ -17,14 +17,16 @@ For the full list of supported architectures across all modalities, see
|
||||
|
||||
| Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
|
||||
|---|---|---|---|---|---|
|
||||
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
|
||||
| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
|
||||
| VoxCPM | local model dir | ✓ | ✓ (PCM stream) | — | — |
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
|
||||
|
||||
CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3) instead.
|
||||
For offline inference of any of these models, see the [offline TTS hub](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md).
|
||||
|
||||
## Common Quick Start
|
||||
|
||||
@@ -96,6 +98,54 @@ For full request-shape documentation (all parameters, response formats, error co
|
||||
|
||||
---
|
||||
|
||||
## CosyVoice3
|
||||
|
||||
2-stage TTS (`talker` + flow-matching `code2wav`) at 24 kHz. Voice cloning only — every request needs `ref_audio` + `ref_text`; there are no built-in voice presets.
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
huggingface-cli download FunAudioLLM/Fun-CosyVoice3-0.5B-2512
|
||||
```
|
||||
|
||||
If your downloaded checkpoint lacks `config.json`, add one with `{"model_type": "cosyvoice3", "architectures": ["CosyVoice3Model"]}` (the loader reads `model_type` to select the class).
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code
|
||||
# or:
|
||||
./cosyvoice3/run_server.sh
|
||||
```
|
||||
|
||||
Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path.
|
||||
|
||||
### CLI client
|
||||
The client defaults to the official upstream zero-shot prompt, so it runs without extra flags:
|
||||
```bash
|
||||
cd examples/online_serving/text_to_speech/cosyvoice3
|
||||
python speech_client.py --text "收到好友从远方寄来的生日礼物。"
|
||||
```
|
||||
|
||||
Pass your own reference clip and transcript for a different voice:
|
||||
```bash
|
||||
python speech_client.py --text "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio."
|
||||
```
|
||||
|
||||
Stream PCM instead of WAV:
|
||||
```bash
|
||||
python speech_client.py --text "Hello world" --stream --output output.pcm
|
||||
```
|
||||
|
||||
The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-text`, `--response-format`, `--stream`, `--output`.
|
||||
|
||||
### Notes
|
||||
- Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
|
||||
- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
|
||||
- For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3).
|
||||
|
||||
---
|
||||
|
||||
## Fish Speech S2 Pro
|
||||
|
||||
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
|
||||
@@ -144,6 +194,84 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server
|
||||
|
||||
---
|
||||
|
||||
## higgs-audio v2
|
||||
|
||||
2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) that streams chunks back to the client.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Voice clone uses HF's `HiggsAudioV2TokenizerModel` loaded from `k2-fsa/OmniVoice/audio_tokenizer/` (~806 MB subdir; the boson-ai standalone tokenizer Hub repo's `model.safetensors` is the 3B talker LM, not the codec):
|
||||
|
||||
```bash
|
||||
pip install -U "transformers>=5.3.0"
|
||||
```
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
|
||||
```
|
||||
Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
|
||||
|
||||
### Sending requests
|
||||
```bash
|
||||
# Plain TTS
|
||||
python higgs_audio_v2/batch_speech_client.py \
|
||||
--base-url http://localhost:8094 \
|
||||
--output-dir /tmp/higgs_out \
|
||||
--prompts "Hello world." "The quick brown fox jumps over the lazy dog."
|
||||
|
||||
# Voice cloning — pass a reference clip and its transcript together
|
||||
python higgs_audio_v2/batch_speech_client.py \
|
||||
--base-url http://localhost:8094 \
|
||||
--output-dir /tmp/higgs_clone \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Exact transcript spoken in reference.wav." \
|
||||
--prompts "Hello, this is a cloned voice."
|
||||
```
|
||||
|
||||
### Notes
|
||||
- Output: 24 kHz mono.
|
||||
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
|
||||
- Out of scope (rejected with explicit 4xx): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
|
||||
## GLM-TTS
|
||||
|
||||
2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
|
||||
# or:
|
||||
bash examples/online_serving/text_to_speech/glm_tts/run_server.sh /path/to/GLM-TTS
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
```bash
|
||||
# Voice cloning (required)
|
||||
python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
|
||||
--text "你好,这是语音克隆测试。" \
|
||||
--ref-audio file:///path/to/ref.wav \
|
||||
--ref-text "这是参考音频的文本内容。"
|
||||
|
||||
# Custom format
|
||||
python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
|
||||
--text "Hello, this is a voice cloning test." \
|
||||
--ref-audio file:///path/to/ref.wav \
|
||||
--ref-text "Transcript of the reference audio." \
|
||||
--response-format mp3 -o output.mp3
|
||||
```
|
||||
|
||||
### Gradio demo
|
||||
```bash
|
||||
bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
|
||||
```
|
||||
|
||||
### Notes
|
||||
- Output: 24 kHz mono WAV via HiFT vocoder.
|
||||
- `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds.
|
||||
- Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer.
|
||||
|
||||
---
|
||||
|
||||
## OmniVoice
|
||||
|
||||
Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline.
|
||||
@@ -196,10 +324,20 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
|
||||
./qwen3_tts/run_server.sh Base
|
||||
```
|
||||
|
||||
### Choosing an executor backend (uniproc vs mp)
|
||||
Stage configs ship with the chunked-streaming default. To use the uniproc executor (lower IPC overhead for the Base cloning task), pass `--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml`. See [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) and [#2604](https://github.com/vllm-project/vllm-omni/pull/2604) for the full investigation.
|
||||
### Executor backend
|
||||
Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override.
|
||||
|
||||
To opt out of chunked streaming, pass `--no-async-chunk` instead — the pipeline auto-dispatches to the end-to-end codec processor.
|
||||
To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.
|
||||
|
||||
### Tuning stage 1 `max_num_seqs` per task type
|
||||
The bundled `qwen3_tts.yaml` ships stage 1 (Code2Wav) at `max_num_seqs: 10`, tuned for Base voice cloning: stage-1 lifetimes are long (~3 s/req), so admitting up to 10 concurrent codec sequences lets requests progress in parallel in the scheduler — ~2× TTFA p95 at c=4 / c=8 (1× H100, 1.7B-Base, seed-tts) at an 8–12 % audio-throughput cost.
|
||||
|
||||
CustomVoice / VoiceDesign have much shorter stage-1 lifetimes (~50–200 ms) and are TTFA-optimal at `max_num_seqs: 1`. Override the default when serving those task types:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \
|
||||
--stage-overrides '{"1": {"max_num_seqs": 1}}'
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
```bash
|
||||
@@ -242,6 +380,31 @@ curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
```
|
||||
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
|
||||
|
||||
### Precomputed custom voices
|
||||
For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
|
||||
```bash
|
||||
python qwen3_tts/precompute_custom_voice.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--voice-name alice \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Original transcript of the reference audio" \
|
||||
--mode icl \
|
||||
--output-dir /path/to/custom_voices
|
||||
```
|
||||
`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
|
||||
```yaml
|
||||
custom_voice_dir: /path/to/custom_voices
|
||||
```
|
||||
Then start the server with that config and call the Speech API with only the voice name:
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
|
||||
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"input":"Hello from a precomputed voice.","voice":"alice","task_type":"Base"}' \
|
||||
--output alice.wav
|
||||
```
|
||||
|
||||
### Streaming PCM
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
@@ -282,70 +445,12 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
|
||||
|
||||
### Notes
|
||||
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
|
||||
- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
|
||||
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
|
||||
- Under vocoder-bound overload (single-stream `rtf_p99 ≥ 1` at the target concurrency), set `active_stream_window: 2` at the top of the deploy yaml to cap simultaneously active Stage 1 streams. Off by default; trades TTFP for streaming continuity. See [#3592](https://github.com/vllm-project/vllm-omni/pull/3592) for the mechanism and tradeoff numbers.
|
||||
|
||||
---
|
||||
|
||||
## VoxCPM
|
||||
|
||||
Split-stage TTS at 24 kHz.
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
pip install voxcpm
|
||||
# or use a local source tree:
|
||||
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
|
||||
```
|
||||
|
||||
If the native VoxCPM `config.json` lacks HF `model_type`, set up an HF-compatible config dir:
|
||||
```bash
|
||||
export VOXCPM_MODEL=/path/to/voxcpm-model
|
||||
export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
|
||||
mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
|
||||
cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
|
||||
cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
|
||||
python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
|
||||
```
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
export VOXCPM_MODEL=/path/to/voxcpm-model
|
||||
./voxcpm/run_server.sh # async-chunk streaming (default)
|
||||
./voxcpm/run_server.sh sync # non-streaming
|
||||
```
|
||||
Or directly:
|
||||
```bash
|
||||
vllm serve "$VOXCPM_MODEL" \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
|
||||
--trust-remote-code --enforce-eager --omni --port 8091
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
```bash
|
||||
# Basic TTS
|
||||
python voxcpm/openai_speech_client.py \
|
||||
--model "$VOXCPM_MODEL" \
|
||||
--text "This is a VoxCPM online text-to-speech example."
|
||||
|
||||
# Voice cloning
|
||||
python voxcpm/openai_speech_client.py \
|
||||
--model "$VOXCPM_MODEL" \
|
||||
--text "This sentence is synthesized with a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "The exact transcript spoken in reference.wav."
|
||||
|
||||
# Streaming PCM
|
||||
python voxcpm/openai_speech_client.py \
|
||||
--model "$VOXCPM_MODEL" \
|
||||
--text "This is a streaming VoxCPM request." \
|
||||
--stream --output voxcpm_stream.pcm
|
||||
```
|
||||
|
||||
### Notes
|
||||
- `voxcpm.yaml` for one-shot decode; `voxcpm_async_chunk.yaml` for single-request streaming. Do not use the async-chunk config for concurrent requests or `/v1/audio/speech/batch`.
|
||||
- Generic TTS fields not supported by VoxCPM: `voice`, `instructions`, `language`, `speaker_embedding`, `x_vector_only_mode`.
|
||||
- For benchmark measurement, see [`benchmarks/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/benchmarks/voxcpm/README.md).
|
||||
|
||||
---
|
||||
|
||||
## VoxCPM2
|
||||
@@ -370,6 +475,23 @@ python voxcpm2/openai_speech_client.py \
|
||||
```
|
||||
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
|
||||
|
||||
### Precomputed custom voices
|
||||
For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
|
||||
```bash
|
||||
python voxcpm2/precompute_custom_voice.py \
|
||||
--model openbmb/VoxCPM2 \
|
||||
--voice-name alice \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--mode ref_continuation \
|
||||
--prompt-text "Original transcript of the reference audio" \
|
||||
--output-dir /path/to/custom_voices
|
||||
```
|
||||
Add the output directory to the deploy config:
|
||||
```yaml
|
||||
custom_voice_dir: /path/to/custom_voices
|
||||
```
|
||||
After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
|
||||
|
||||
### Gradio demo (gapless streaming via AudioWorklet)
|
||||
```bash
|
||||
python voxcpm2/gradio_demo.py
|
||||
@@ -407,6 +529,14 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "cosyvoice3/run_server.sh"
|
||||
|
||||
--8<-- "examples/online_serving/text_to_speech/cosyvoice3/run_server.sh"
|
||||
|
||||
??? abstract "cosyvoice3/speech_client.py"
|
||||
|
||||
--8<-- "examples/online_serving/text_to_speech/cosyvoice3/speech_client.py"
|
||||
|
||||
??? abstract "fish_speech/gradio_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/fish_speech/gradio_demo.py"
|
||||
@@ -423,6 +553,29 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/fish_speech/speech_client.py"
|
||||
``````
|
||||
??? abstract "higgs_audio_v2/batch_speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py"
|
||||
``````
|
||||
??? abstract "higgs_audio_v2/run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh"
|
||||
??? abstract "glm_tts/gradio_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/glm_tts/gradio_demo.py"
|
||||
``````
|
||||
??? abstract "glm_tts/openai_speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py"
|
||||
``````
|
||||
??? abstract "glm_tts/run_gradio_demo.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh"
|
||||
``````
|
||||
??? abstract "glm_tts/run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_speech/glm_tts/run_server.sh"
|
||||
``````
|
||||
??? abstract "omnivoice/run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_speech/omnivoice/run_server.sh"
|
||||
@@ -463,14 +616,6 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/qwen3_tts/tts_common.py"
|
||||
``````
|
||||
??? abstract "voxcpm/openai_speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/voxcpm/openai_speech_client.py"
|
||||
``````
|
||||
??? abstract "voxcpm/run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_speech/voxcpm/run_server.sh"
|
||||
``````
|
||||
??? abstract "voxcpm2/gradio_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py"
|
||||
|
||||
@@ -32,7 +32,9 @@ guide. AutoRound is Intel-supported.
|
||||
|-------|------------|-------|--------|---------|
|
||||
| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
|
||||
| Qwen-Image | Not listed | Diffusion transformer | W4A16 | Not validated |
|
||||
| Wan2.2 | Not listed | Diffusion transformer | W4A16 | Not validated |
|
||||
| Wan2.2-I2V | `Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
|
||||
| Wan2.2-T2V | `Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
|
||||
| Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
@@ -50,8 +52,8 @@ vLLM-Omni's runtime module names.
|
||||
|
||||
| Model | Scope | Status | Notes |
|
||||
|-------|-------|--------|-------|
|
||||
| GLM-Image | Diffusion transformer | ✅ | `Intel/GLM-Image-int4-AutoRound` |
|
||||
| BAGEL | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
|
||||
| GLM-Image | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
|
||||
|
||||
## Configuration
|
||||
|
||||
|
||||
@@ -74,7 +74,7 @@ Online:
|
||||
vllm serve Qwen/Qwen-Image \
|
||||
--omni \
|
||||
--port 8000 \
|
||||
--quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
|
||||
--diffusion-quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
@@ -6,34 +6,46 @@ ModelOpt quantization loads checkpoints produced by NVIDIA ModelOpt. The
|
||||
quantized weights and scale tensors are generated before serving, so inference
|
||||
does not run online calibration or convert a BF16 checkpoint at startup.
|
||||
|
||||
vLLM-Omni currently validates the ModelOpt FP8 checkpoint path for diffusion
|
||||
transformers. The loader auto-detects supported ModelOpt FP8 checkpoint configs
|
||||
and keeps non-transformer components, such as the tokenizer, scheduler, text
|
||||
vLLM-Omni validates ModelOpt FP8, ModelOpt NVFP4, and ModelOpt mixed
|
||||
FP8/NVFP4 checkpoint loading for diffusion transformer stages. The loader
|
||||
auto-detects supported ModelOpt checkpoint configs and keeps non-transformer
|
||||
components, such as the tokenizer, scheduler, text encoder, vision/audio
|
||||
encoder, and VAE, on the base checkpoint unless a model-specific guide says
|
||||
otherwise.
|
||||
|
||||
!!! note
|
||||
`--force-cutlass-fp8` is an explicit runtime override for diffusion
|
||||
checkpoints that already carry a supported ModelOpt FP8 config. It does not
|
||||
quantize BF16 checkpoints and it does not apply to online `--quantization
|
||||
fp8`. The flag only takes effect for ModelOpt FP8 diffusion stages on CUDA
|
||||
SM89+ devices; other platforms and non-ModelOpt FP8 paths fall back to the
|
||||
normal vLLM kernel selection.
|
||||
ModelOpt checkpoints are pre-quantized checkpoints. Do not pass
|
||||
`--quantization fp8` for these checkpoints. The checkpoint
|
||||
`quantization_config` selects the ModelOpt path.
|
||||
|
||||
!!! note
|
||||
`--force-cutlass-fp8`, `--linear-backend cutlass`, and
|
||||
`--moe-backend cutlass` are runtime backend selections for checkpoints that
|
||||
already carry supported ModelOpt quantized weights and scales. They do not
|
||||
quantize BF16 checkpoints at startup.
|
||||
|
||||
## Supported ModelOpt Checkpoint Formats
|
||||
|
||||
vLLM-Omni treats ModelOpt checkpoints as pre-quantized checkpoints. The
|
||||
checkpoint config must identify ModelOpt as the quantization method or producer,
|
||||
and the quantization algorithm must be one of the validated FP8 algorithms.
|
||||
and the quantization algorithm must be one of the validated algorithms below.
|
||||
|
||||
| Checkpoint field | Supported value |
|
||||
|------------------|-----------------|
|
||||
| `method` / `quant_method` | `modelopt` |
|
||||
| `method` / `quant_method` | `modelopt`, `modelopt_fp4`, `modelopt_mixed` |
|
||||
| `producer.name` | `modelopt` |
|
||||
| `quant_algo` | `FP8`, `FP8_PER_CHANNEL_PER_TOKEN` |
|
||||
| `quant_algo` | `FP8`, `FP8_PER_CHANNEL_PER_TOKEN`, `NVFP4`, `MIXED_PRECISION` |
|
||||
|
||||
Other ModelOpt algorithms, such as NVFP4, are not enabled by this diffusion
|
||||
FP8 path until they have separate model and quality validation.
|
||||
| `quant_algo` | Runtime method | Typical use |
|
||||
|--------------|----------------|-------------|
|
||||
| `FP8`, `FP8_PER_CHANNEL_PER_TOKEN` | `modelopt` | FP8 diffusion transformer checkpoints |
|
||||
| `NVFP4` | `modelopt_fp4` | NVFP4 diffusion transformer checkpoints |
|
||||
| `MIXED_PRECISION` | `modelopt_mixed` | Mixed FP8/NVFP4 checkpoints with a ModelOpt per-layer policy |
|
||||
|
||||
For multi-component diffusion or omni models, only the checkpoint component
|
||||
that contains ModelOpt quantized weights should use the ModelOpt quantization
|
||||
method. Encoders, decoders, tokenizers, schedulers, and other BF16 components
|
||||
stay unquantized unless the model-specific recipe validates otherwise.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
@@ -47,7 +59,10 @@ FP8 path until they have separate model and quality validation.
|
||||
| Ascend NPU | ❌ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
|
||||
The optional CUTLASS FP8 runtime override requires CUDA SM89+.
|
||||
The optional CUTLASS FP8 runtime override requires CUDA SM89+. ModelOpt NVFP4
|
||||
and mixed FP8/NVFP4 diffusion checkpoints are currently validated on Blackwell
|
||||
CUDA systems in the recipes below; other CUDA generations require separate
|
||||
backend and quality validation.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
@@ -56,12 +71,19 @@ The optional CUTLASS FP8 runtime override requires CUDA SM89+.
|
||||
| Model | HF checkpoint | Scope | Status |
|
||||
|-------|---------------|-------|--------|
|
||||
| Qwen-Image 2512 | `feizhai123/qwen-image-2512-modelopt-fp8-dynamic-all` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
|
||||
| Qwen-Image 2512 | `feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy` | Diffusion transformer | Validated for ModelOpt mixed FP8/NVFP4 checkpoints |
|
||||
| Z-Image | `feizhai123/z-image-modelopt-fp8-conservative` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
|
||||
| FLUX.2-dev | `feizhai123/flux2-dev-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
|
||||
| FLUX.2-klein 4B | `feizhai123/flux2-klein-4b-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
|
||||
| HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-fp8` | MoE diffusion transformer | Validated for ModelOpt FP8 checkpoints |
|
||||
| HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8` | MoE diffusion transformer | Validated for ModelOpt mixed FP8/NVFP4 checkpoints |
|
||||
| Wan2.2 | Not available | Diffusion transformer | Not validated |
|
||||
|
||||
For full serving commands and benchmark context, see
|
||||
[`recipes/Qwen/Qwen-Image.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image.md)
|
||||
and
|
||||
[`recipes/Tencent/HunyuanImage-3.0-Instruct.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md).
|
||||
|
||||
### Multi-Stage Omni/TTS Model
|
||||
|
||||
| Model | Scope | Status |
|
||||
@@ -80,15 +102,28 @@ ModelOpt targets yet.
|
||||
|
||||
## Configuration
|
||||
|
||||
For pre-quantized ModelOpt FP8 checkpoints, no `--quantization fp8` flag is
|
||||
needed. The checkpoint config selects the ModelOpt path.
|
||||
For pre-quantized ModelOpt checkpoints, no `--quantization fp8` flag is needed.
|
||||
The checkpoint config selects the ModelOpt path.
|
||||
|
||||
Online serving:
|
||||
|
||||
```bash
|
||||
vllm serve <modelopt-fp8-checkpoint> \
|
||||
vllm serve <modelopt-checkpoint> \
|
||||
--omni \
|
||||
--tensor-parallel-size <N> \
|
||||
--linear-backend cutlass \
|
||||
--force-cutlass-fp8
|
||||
```
|
||||
|
||||
For mixed FP8/NVFP4 MoE checkpoints, also select the validated MoE backend:
|
||||
|
||||
```bash
|
||||
vllm serve <modelopt-mixed-moe-checkpoint> \
|
||||
--omni \
|
||||
--tensor-parallel-size <N> \
|
||||
--enable-expert-parallel \
|
||||
--linear-backend cutlass \
|
||||
--moe-backend cutlass \
|
||||
--force-cutlass-fp8
|
||||
```
|
||||
|
||||
@@ -96,14 +131,14 @@ Offline inference:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model <modelopt-fp8-checkpoint> \
|
||||
--model <modelopt-checkpoint> \
|
||||
--tensor-parallel-size <N> \
|
||||
--prompt "a red ceramic teapot on a wooden table" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--num-inference-steps 20 \
|
||||
--seed 42 \
|
||||
--output outputs/modelopt_fp8.png
|
||||
--output outputs/modelopt.png
|
||||
```
|
||||
|
||||
Python API:
|
||||
@@ -112,7 +147,7 @@ Python API:
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="<modelopt-fp8-checkpoint>",
|
||||
model="<modelopt-checkpoint>",
|
||||
tensor_parallel_size=2,
|
||||
force_cutlass_fp8=True,
|
||||
)
|
||||
@@ -123,16 +158,34 @@ omni = Omni(
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `force_cutlass_fp8` / `--force-cutlass-fp8` | bool | `False` | Force CUTLASS FP8 linear kernels for supported ModelOpt FP8 diffusion stages on CUDA SM89+ |
|
||||
| `--linear-backend cutlass` | str | auto | Select the validated CUTLASS linear backend for supported ModelOpt NVFP4 or mixed FP8/NVFP4 diffusion stages |
|
||||
| `--moe-backend cutlass` | str | auto | Select the validated CUTLASS MoE backend for supported ModelOpt mixed MoE checkpoints |
|
||||
|
||||
## Validation and Notes
|
||||
|
||||
1. Compare the ModelOpt FP8 checkpoint against the BF16 baseline with the same
|
||||
1. Compare the ModelOpt checkpoint against the BF16 baseline with the same
|
||||
prompt, resolution, seed, and inference steps.
|
||||
2. Use `tests/diffusion/quantization/test_quantization_quality.py` with
|
||||
`VLLM_OMNI_QUALITY_CONFIGS` to validate local baseline and quantized model
|
||||
paths.
|
||||
3. Report LPIPS, PSNR, MAE, throughput, latency, and peak memory when adding a
|
||||
new validated ModelOpt diffusion checkpoint.
|
||||
4. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this
|
||||
ModelOpt path only when the checkpoint already contains ModelOpt FP8 weights
|
||||
and scales.
|
||||
3. For HunyuanImage-3.0 quantized DiT checkpoints, the opt-in accuracy check is:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=2,3 \
|
||||
HUNYUAN_IMAGE3_RUN_QUANT_ACCURACY=1 \
|
||||
HUNYUAN_IMAGE3_QUANT_DEVICES=0,1 \
|
||||
HUNYUAN_IMAGE3_QUANT_TP=2 \
|
||||
HUNYUAN_IMAGE3_BF16_MODEL=/path/to/hunyuan-image3-bf16 \
|
||||
HUNYUAN_IMAGE3_FP8_MODEL=/path/to/hunyuan-image3-modelopt-fp8 \
|
||||
HUNYUAN_IMAGE3_NVFP4_MODEL=/path/to/hunyuan-image3-modelopt-mixed-experts-nvfp4-dense-fp8 \
|
||||
PYTHONPATH=/path/to/vllm-omni:${PYTHONPATH:-} \
|
||||
python -m pytest -s -v \
|
||||
tests/e2e/accuracy/test_hunyuan_image3.py \
|
||||
-k quantized_dit_matches_bf16_accuracy
|
||||
```
|
||||
|
||||
4. Report CLIP score deltas, SSIM, PSNR, throughput, latency, and peak memory
|
||||
when adding a new validated ModelOpt diffusion checkpoint.
|
||||
5. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this
|
||||
ModelOpt path only when the checkpoint already contains ModelOpt quantized
|
||||
weights and scales.
|
||||
@@ -0,0 +1,479 @@
|
||||
# W4A4 MXFP4 Quantization
|
||||
|
||||
## Overview
|
||||
|
||||
W4A4 MXFP4 (Microscaling FP4) quantizes both weights and activations to FP4
|
||||
(`float4_e2m1fn_x2`, packed 2 values per byte) using the OCP MX format: groups
|
||||
of 32 K-dimension elements share a single `float8_e8m0fnu` exponent scale.
|
||||
|
||||
vLLM-Omni provides two quantization methods with different scale structures:
|
||||
|
||||
| Method | Scale structure | Mode | Use case |
|
||||
|--------|----------------|------|----------|
|
||||
| `mxfp4` | Single-scale (per-32 fine only) | Online only | Quick accuracy baseline; no checkpoint prep needed |
|
||||
| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | Production; better accuracy; offline recommended |
|
||||
|
||||
!!! tip "Recommended: `mxfp4_dualscale` offline"
|
||||
For production deployments, use the `mxfp4_dualscale` offline mode with a
|
||||
pre-quantized checkpoint produced by msModelSlim. Offline checkpoints load
|
||||
calibrated `mul_scale` tensors from disk, providing measurably better accuracy
|
||||
than any online method. The one-time preprocessing cost amortises across all
|
||||
subsequent inference runs.
|
||||
|
||||
Use `mxfp4` online only for quick experimentation where preprocessing time
|
||||
is not acceptable and accuracy loss is tolerable.
|
||||
|
||||
!!! warning "Online single-scale ≠ Offline dual-scale"
|
||||
`mxfp4_dualscale` offline mode uses `NPUMxfp4DualScaleLinearMethod`:
|
||||
fine scale (per-32 K), coarse scale (per-512 K), and per-input-channel
|
||||
`mul_scale` from calibration — all loaded from the checkpoint.
|
||||
`mxfp4_dualscale` online mode uses `NPUMxfp4DualScaleOnlineLinearMethod`:
|
||||
dual-level scales computed on the fly from BF16 weights; no calibration
|
||||
`mul_scale` is available. Loading an offline checkpoint with the online
|
||||
method (or vice versa) will produce incorrect results or shape errors.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | Support |
|
||||
|--------|---------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ⭕ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ⭕ |
|
||||
| AMD ROCm | ⭕ |
|
||||
| Intel XPU | ⭕ |
|
||||
| Ascend NPU (Atlas 950 A5) | ✅ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Wan2.2)
|
||||
|
||||
| Model | Online | Offline | Notes |
|
||||
|-------|--------|---------|-------|
|
||||
| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade (`transformer` + `transformer_2`); both transformers quantized with the same config |
|
||||
| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade; same scheme as T2V-A14B |
|
||||
| Wan2.2-TI2V-5B | ❌ | ❌ | Parameter count too small; W4A4 causes unacceptable accuracy loss |
|
||||
|
||||
The choice between `mxfp4` and `mxfp4_dualscale` in **online mode** is about
|
||||
quantization quality, not model compatibility — both work on cascade (A14B) and
|
||||
single-transformer models alike, the same as `mxfp8` online:
|
||||
|
||||
- `mxfp4`: single-scale, lower overhead, simpler compute, online only
|
||||
- `mxfp4_dualscale`: dual-scale + optional BF16 fallback, better accuracy, online **and** offline
|
||||
|
||||
**Offline** checkpoints for A14B are always in `mxfp4_dualscale` format (produced
|
||||
by the merge script); there is no offline `mxfp4` single-scale format.
|
||||
|
||||
!!! note "Per-layer BF16 fallback in offline cascade models"
|
||||
The A14B offline checkpoint uses `quant_method: mxfp4_dualscale`. Most
|
||||
linear layers are stored as W4A4 MXFP4 DualScale; precision-sensitive layers
|
||||
retain their original BF16 weights and are listed in `ignored_layers` inside
|
||||
each transformer's `config.json`. The two transformers may have different
|
||||
`ignored_layers` sets — the pipeline reads each transformer's own `config.json`
|
||||
and rebuilds the config locally when they differ, so routing is always
|
||||
per-transformer-accurate.
|
||||
|
||||
!!! warning "TI2V-5B not supported"
|
||||
Wan2.2-TI2V-5B is excluded from W4A4 quantization. Its smaller parameter
|
||||
count makes it significantly more sensitive to 4-bit quantization noise,
|
||||
resulting in unacceptable accuracy loss. Use [MXFP8](mxfp8.md) for TI2V-5B.
|
||||
|
||||
## Configuration
|
||||
|
||||
### `mxfp4` — Single-Scale Online Mode
|
||||
|
||||
Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to
|
||||
MXFP4 at load time using `npu_dynamic_mx_quant`. A single block scale
|
||||
(`float8_e8m0fnu`, one per 32 K elements) is computed on the fly; no
|
||||
calibration `mul_scale` is available. Applies equally to single-transformer
|
||||
and cascade (A14B) models — both transformers in a cascade receive the same
|
||||
quantization config automatically.
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
omni = Omni(model="<your-model>", quantization="mxfp4")
|
||||
outputs = omni.generate(
|
||||
"A cat sitting on a windowsill",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=50),
|
||||
)
|
||||
```
|
||||
|
||||
```bash
|
||||
# Single-transformer or cascade model — same command
|
||||
python text_to_video.py --model <your-model> --quantization mxfp4
|
||||
|
||||
# Online serving
|
||||
vllm serve <your-model> --omni --quantization mxfp4
|
||||
```
|
||||
|
||||
### `mxfp4_dualscale` — DualScale Online Mode
|
||||
|
||||
Online DualScale mode computes both fine and coarse scales on the fly from BF16
|
||||
weights using `npu_dynamic_dual_level_mx_quant`. Applies equally to
|
||||
single-transformer and cascade (A14B) models. Compared to `mxfp4` online,
|
||||
DualScale provides better quantization accuracy at higher compute cost.
|
||||
|
||||
The default configuration keeps the leading 5 transformer blocks in BF16
|
||||
(`num_bf16_fallback_layers=5`). Accuracy evaluation on Wan2.2-A14B shows this
|
||||
is sufficient to meet quality requirements and is the recommended setting.
|
||||
|
||||
```python
|
||||
omni = Omni(model="<your-model>", quantization="mxfp4_dualscale")
|
||||
```
|
||||
|
||||
```bash
|
||||
python text_to_video.py --model <your-model> --quantization mxfp4_dualscale
|
||||
```
|
||||
|
||||
If accuracy debugging identifies additional precision-sensitive layers, they
|
||||
can be pinned to BF16 via the Python API:
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="<your-model>",
|
||||
quantization_config={
|
||||
"method": "mxfp4_dualscale",
|
||||
"ignored_layers": ["blocks.10.attn1.to_q"], # explicit per-layer override
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
BF16 fallback routing in online mode applies two rules in priority order:
|
||||
|
||||
1. **`ignored_layers`** (explicit per-layer override): any layer whose prefix
|
||||
matches is kept in BF16 regardless of block index.
|
||||
2. **`num_bf16_fallback_layers`** (coarse leading-block rule): the first N
|
||||
transformer blocks (`blocks.0` … `blocks.N-1`) fall back to BF16. Defaults
|
||||
to `5` (recommended). Layers outside `blocks.N.*`
|
||||
(e.g. `condition_embedder`) are always quantized.
|
||||
|
||||
### `mxfp4_dualscale` — DualScale Offline Mode (Recommended)
|
||||
|
||||
Offline mode loads a pre-quantized DualScale checkpoint from msModelSlim. A
|
||||
preprocessing step converts the raw quantized output to the diffusers format
|
||||
expected by vLLM-Omni and injects the quantization config into each
|
||||
`transformer/config.json` so that vLLM-Omni auto-detects the offline path
|
||||
without a `--quantization` flag.
|
||||
|
||||
BF16 fallback layers may be interleaved anywhere in the transformer — they are
|
||||
not restricted to leading blocks. The merge script detects them from
|
||||
`quant_model_description.json` and writes their prefixes into `ignored_layers`
|
||||
inside `config.json`. At runtime, each layer's prefix is matched against
|
||||
`ignored_layers` to decide BF16 vs. MXFP4 DualScale.
|
||||
|
||||
#### Checkpoint tensor layout
|
||||
|
||||
Each quantized linear layer stores four tensors:
|
||||
|
||||
| Tensor | Shape | dtype | Description |
|
||||
|--------|-------|-------|-------------|
|
||||
| `weight` | `(N, K)` | float8_e4m3fn | FP4 packed (2 values per byte) |
|
||||
| `weight_scale` | `(N, K//32)` | uint8 | Fine block scale (`float8_e8m0fnu` bit pattern) |
|
||||
| `weight_dual_scale` | `(N, K//512, 1)` | float32 | Coarse block scale |
|
||||
| `mul_scale` | `(K,)` | float32 | Per-input-channel smooth pre-scale (from calibration) |
|
||||
|
||||
BF16 fallback layers have no quantization tensors; only the original `weight`
|
||||
(and optional `bias`) are present, loaded directly from the base checkpoint.
|
||||
|
||||
#### Step 1 — Quantize with msModelSlim
|
||||
|
||||
```bash
|
||||
msmodelslim quant \
|
||||
--model_path /path/to/Wan2.2-T2V-A14B-Diffusers \
|
||||
--save_path /path/to/wan2_2_t2v_quantized_raw \
|
||||
--device npu \
|
||||
--model_type Wan2_2 \
|
||||
--config_path /path/to/wan2_2_w4a4_mxfp4_dualscale.yaml \
|
||||
--trust_remote_code True
|
||||
```
|
||||
|
||||
After this step, `--save_path` contains raw quantized safetensors files,
|
||||
scale files, and a metadata JSON (`quant_model_description*.json`).
|
||||
|
||||
For cascade MoE models (T2V-A14B, I2V-A14B), msModelSlim outputs two
|
||||
subdirectories: `high_noise_model/` (transformer) and `low_noise_model/`
|
||||
(transformer_2).
|
||||
|
||||
#### Step 2 — Preprocess with merge_mxfp4_dualscale_checkpoint.py
|
||||
|
||||
The script (`vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py`):
|
||||
|
||||
1. Copies the original diffusers model to `--output-path` (VAE, text encoder,
|
||||
scheduler, etc. are preserved).
|
||||
2. Remaps tensor names from msModelSlim convention to diffusers convention and
|
||||
strips `.linear.` / `.div.` wrappers added by the quantization tool.
|
||||
3. Overlays MXFP4 tensors (weight, fine/coarse scales, `mul_scale`) onto the
|
||||
BF16 base checkpoint. Non-quantized layers keep their original BF16 weights.
|
||||
4. Detects all linear layers that remain in BF16 and writes their prefixes into
|
||||
`ignored_layers` in `config.json`.
|
||||
5. Injects `quantization_config` so vLLM-Omni auto-detects offline MXFP4
|
||||
DualScale.
|
||||
|
||||
For cascade MoE models, steps 2–5 run separately for each transformer.
|
||||
|
||||
```bash
|
||||
python vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py \
|
||||
--model-type Wan2.2-T2V-A14B \
|
||||
--original-model /path/to/Wan2.2-T2V-A14B-Diffusers \
|
||||
--quant-path /path/to/wan2_2_t2v_quantized_raw \
|
||||
--output-path /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale
|
||||
```
|
||||
|
||||
| Argument | Description |
|
||||
|----------|-------------|
|
||||
| `--model-type` | Model variant: `Wan2.2-T2V-A14B` or `Wan2.2-I2V-A14B` |
|
||||
| `--original-model` | Root directory of the original BF16 diffusers model |
|
||||
| `--quant-path` | Root directory of the msModelSlim quantized output |
|
||||
| `--output-path` | Output directory for the merged model (created by the script) |
|
||||
|
||||
The script outputs a complete diffusers model directory at `--output-path`,
|
||||
with each transformer subfolder containing:
|
||||
|
||||
- `diffusion_pytorch_model.safetensors` — MXFP4 weights + scale tensors, with BF16 fallback layers from the base checkpoint
|
||||
- `config.json` — original transformer config with `quantization_config` injected
|
||||
- `quant_model_description.json` — quantization metadata (reference only)
|
||||
|
||||
The `quantization_config` injected into `config.json` for each transformer:
|
||||
|
||||
```json
|
||||
{
|
||||
"quant_method": "mxfp4_dualscale",
|
||||
"is_checkpoint_serialized": true,
|
||||
"ignored_layers": [
|
||||
"blocks.0.attn1.to_qkv",
|
||||
"blocks.0.attn1.to_out",
|
||||
"proj_out"
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
`ignored_layers` lists every linear layer that retains its original BF16 weight,
|
||||
using vllm-omni model parameter names (QKV-fused, FFN underscored, `to_out`
|
||||
unindexed). The exact entries are determined by the quantization tool (msModelSlim)
|
||||
and may differ between `transformer` and `transformer_2` in a cascade model.
|
||||
|
||||
#### Step 3 — Serve
|
||||
|
||||
```bash
|
||||
python text_to_video.py --model /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale
|
||||
|
||||
# Online serving
|
||||
vllm serve /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale --omni
|
||||
```
|
||||
|
||||
```python
|
||||
omni = Omni(model="/path/to/Wan2.2-T2V-A14B-MXFP4-DualScale")
|
||||
```
|
||||
|
||||
!!! note
|
||||
No `--quantization` flag is needed for offline mode. The preprocessing
|
||||
script injects `quantization_config` into each `transformer/config.json`,
|
||||
which vLLM-Omni reads automatically to activate the correct offline path.
|
||||
|
||||
## Parameters
|
||||
|
||||
### `mxfp4` (single-scale, online only)
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `method` | str | — | `"mxfp4"` |
|
||||
| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16 |
|
||||
|
||||
### `mxfp4_dualscale` (dual-scale, online + offline)
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `method` | str | — | `"mxfp4_dualscale"` |
|
||||
| `is_checkpoint_serialized` | bool | `False` | `True` for offline DualScale checkpoints; auto-set from `config.json` when using the preprocessing script |
|
||||
| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16. **Works in both modes**: offline — populated by the merge script for interleaved sensitive layers; online — user-supplied for explicit per-layer precision override |
|
||||
| `num_bf16_fallback_layers` | int | `5` | **Online mode only**: leading N transformer blocks (`blocks.0` … `blocks.N-1`) kept in BF16. Applied after `ignored_layers`; ignored in offline mode. Default of `5` is the evaluated recommended value for Wan2.2-A14B |
|
||||
|
||||
#### BF16 fallback priority (online mode)
|
||||
|
||||
```
|
||||
for each linear layer:
|
||||
if prefix in ignored_layers → BF16 (explicit override, highest priority)
|
||||
elif block_idx < num_bf16_fallback_layers → BF16 (coarse leading-block rule)
|
||||
else → MXFP4 DualScale online
|
||||
```
|
||||
|
||||
Layers outside `blocks.N.*` (e.g. `condition_embedder.*`) are always quantized
|
||||
unless they appear in `ignored_layers`.
|
||||
|
||||
## Validation and Notes
|
||||
|
||||
1. **Online single-scale (`mxfp4`)** quantizes BF16 weights at load time using
|
||||
`npu_dynamic_mx_quant` (single-scale). No calibration `mul_scale` is
|
||||
available — all output partitions receive an identity pre-scale. No offline
|
||||
checkpoint format exists for this method.
|
||||
|
||||
2. **Online dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=False`)**
|
||||
quantizes BF16 weights using `npu_dynamic_dual_level_mx_quant` (fine + coarse
|
||||
scales computed on the fly). No calibration `mul_scale`; leading blocks or
|
||||
explicit `ignored_layers` stay in BF16 for accuracy.
|
||||
|
||||
3. **Offline dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=True`)** —
|
||||
**recommended for production** — loads four tensors per quantized layer: FP4
|
||||
weight, fine scale (`uint8` reinterpreted as `float8_e8m0fnu`), coarse scale
|
||||
(`float32`), and per-input-channel `mul_scale` (`float32`). BF16 fallback
|
||||
layers have no quantization tensors and are routed via `ignored_layers`.
|
||||
|
||||
4. **Scale dtype**: fine scales are stored as `uint8` in safetensors (same bit
|
||||
layout as `float8_e8m0fnu`) and reinterpreted at load time without a lossy
|
||||
float32 round-trip.
|
||||
|
||||
5. **Cascade model config propagation**: in a cascade model (transformer +
|
||||
transformer_2), vLLM-Omni reads each transformer's own `config.json` and
|
||||
rebuilds the quant config locally when `ignored_layers` differs between
|
||||
transformers, ensuring per-layer routing is accurate for each. The first
|
||||
transformer's config is propagated to `od_config` so the second transformer
|
||||
can reuse it as a starting point.
|
||||
|
||||
6. **Self-attention QKV fusion**: Q, K, V projection weights are fused into a
|
||||
single `QKVParallelLinear` layer at runtime. `ignored_layers` entries use the
|
||||
fused name (`attn1.to_qkv`), written automatically by the merge script.
|
||||
|
||||
7. W4A4 carries higher quantization noise than W8A8 (16 vs 256 levels). The
|
||||
DualScale offline method mitigates this with calibrated `mul_scale` smooth
|
||||
quantization. Use `ignored_layers` and `num_bf16_fallback_layers` to trade
|
||||
off compression vs. accuracy for precision-sensitive layers.
|
||||
|
||||
## Adapting MXFP4 for a New Model
|
||||
|
||||
This section is aimed at developers who want to add MXFP4 support to a model
|
||||
other than Wan2.2. The three integration points are: (1) discovering the correct
|
||||
runtime layer names, (2) wiring `ignored_layers` into the model, and (3) writing
|
||||
a merge script for offline checkpoints.
|
||||
|
||||
### Step 1 — Discover runtime layer names
|
||||
|
||||
`ignored_layers` entries must match the **runtime parameter names** used inside
|
||||
vllm-omni, which may differ from the names stored in the diffusers checkpoint.
|
||||
The canonical source of truth is the model's own `named_parameters()`.
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
# Load the model without quantization to inspect parameter names.
|
||||
omni = Omni(model="/path/to/your-model") # no --quantization flag
|
||||
for name, _ in omni.pipeline.transformer.named_parameters():
|
||||
if "weight" in name and "scale" not in name:
|
||||
print(name)
|
||||
```
|
||||
|
||||
Compare the printed names against the diffusers checkpoint keys
|
||||
(`safetensors.safe_open` or `torch.load`) to identify any renames your model
|
||||
applies. Common patterns that differ in Wan2.2 (and may appear in other
|
||||
models):
|
||||
|
||||
| Diffusers checkpoint name | vllm-omni runtime name | Reason |
|
||||
|---------------------------|------------------------|--------|
|
||||
| `attn1.to_q`, `attn1.to_k`, `attn1.to_v` | `attn1.to_qkv` | Self-attention Q/K/V fused into `QKVParallelLinear` |
|
||||
| `ffn.net.0.proj` | `ffn.net_0.proj` | Dots in sub-module names replaced with underscores |
|
||||
| `ffn.net.2` | `ffn.net_2` | Same underscore rule |
|
||||
| `to_out.0` | `to_out` | Sequential index stripped |
|
||||
|
||||
If your model has different fusion patterns, inspect `packed_modules_mapping`
|
||||
on the model class — this dict records how checkpoint keys are mapped to
|
||||
fused runtime parameters.
|
||||
|
||||
!!! warning "Partial QKV fallback is not allowed"
|
||||
If your model fuses Q, K, V into a single layer, `ignored_layers` must
|
||||
include **all three or none**. A partial fallback (e.g. `to_q` in BF16 but
|
||||
`to_k`, `to_v` quantized) cannot be expressed at runtime because they share
|
||||
one `QKVParallelLinear`. The merge script enforces this and raises an error
|
||||
if only some of the trio appear as non-quantized.
|
||||
|
||||
### Step 2 — Add ignored_layers to the model
|
||||
|
||||
#### Online mode
|
||||
|
||||
Pass `ignored_layers` directly in the quantization config using the **runtime
|
||||
names** discovered in Step 1. No code changes to the model are required.
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="/path/to/your-model",
|
||||
quantization={
|
||||
"method": "mxfp4_dualscale",
|
||||
"ignored_layers": [
|
||||
"blocks.0.attn1.to_qkv", # runtime name, not diffusers name
|
||||
"blocks.0.attn1.to_out",
|
||||
"blocks.0.ffn.net_0.proj",
|
||||
],
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
```bash
|
||||
# CLI does not support list-typed ignored_layers directly.
|
||||
# Use the Python API or set ignored_layers in config.json (offline).
|
||||
python your_script.py --model /path/to/your-model --quantization mxfp4_dualscale
|
||||
```
|
||||
|
||||
The `num_bf16_fallback_layers` coarse rule is an alternative to listing layers
|
||||
individually: set it to N to keep all linear layers in blocks 0 … N-1 in BF16.
|
||||
The right value depends on the model's sensitivity; evaluate on a validation
|
||||
set and pick the smallest N that meets your accuracy target.
|
||||
|
||||
#### Offline mode
|
||||
|
||||
For offline checkpoints, `ignored_layers` is written into each transformer's
|
||||
`config.json` by the merge script (see Step 3). No manual editing is needed if
|
||||
the merge script is correct. The injected block:
|
||||
|
||||
```json
|
||||
{
|
||||
"quant_method": "mxfp4_dualscale",
|
||||
"is_checkpoint_serialized": true,
|
||||
"ignored_layers": [
|
||||
"blocks.0.attn1.to_qkv",
|
||||
"blocks.0.attn1.to_out"
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
To add a layer manually (e.g. to pin an additional layer to BF16 without
|
||||
re-running the merge script), edit `config.json` inside the transformer
|
||||
subfolder. Use runtime names, not diffusers checkpoint names.
|
||||
|
||||
### Step 3 — Write a merge script for offline mode
|
||||
|
||||
The merge script for a new model mirrors
|
||||
`vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py`. The four
|
||||
things it must do:
|
||||
|
||||
1. **Remap tensor names** from the quantization tool convention to diffusers
|
||||
convention (strip wrappers like `.linear.`, `.div.`; fix any prefix
|
||||
differences).
|
||||
|
||||
2. **Collect ignored_layers**: after loading, enumerate all `*.weight` keys that
|
||||
have no corresponding `*.weight_scale` (i.e. layers the tool left in BF16).
|
||||
Convert diffusers names to vllm-omni runtime names (fuse QKV, rename FFN
|
||||
sub-modules, etc.). Write the result to `config.json`.
|
||||
|
||||
3. **Inject `quantization_config`** into `config.json`:
|
||||
```python
|
||||
config["quantization_config"] = {
|
||||
"quant_method": "mxfp4_dualscale",
|
||||
"is_checkpoint_serialized": True,
|
||||
"ignored_layers": ignored_layers, # runtime names
|
||||
}
|
||||
```
|
||||
|
||||
4. **Save** the merged safetensors and the updated `config.json`.
|
||||
|
||||
The key helper to implement is the diffusers-to-runtime name translator
|
||||
(equivalent to `_diffusers_to_vllm_ignored` in the Wan2.2 merge script).
|
||||
For each non-quantized diffusers weight key, apply your model's specific
|
||||
renaming rules and collect the results.
|
||||
|
||||
!!! tip "Validate before serving"
|
||||
After producing the offline checkpoint, load it without a `--quantization`
|
||||
flag and verify that vLLM-Omni auto-detects the correct method. Check that
|
||||
the layer count reported in the startup log matches expectations: quantized
|
||||
layer count + `ignored_layers` count should equal total linear layer count.
|
||||
Any mismatch indicates a name-mapping bug in the merge script.
|
||||
@@ -7,26 +7,28 @@ using the OCP MX format: groups of 32 K-dimension elements share a single
|
||||
`float8_e8m0fnu` exponent scale. This gives better accuracy than channel-wise
|
||||
FP8 while keeping the same 8-bit weight footprint.
|
||||
|
||||
This method supports two modes:
|
||||
This method supports three modes:
|
||||
|
||||
| Mode | Description |
|
||||
|------|-------------|
|
||||
| **Online** | BF16 weights are quantized to MXFP8 at load time — no pre-processing needed |
|
||||
| **Offline** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
|
||||
| **Offline (Native)** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
|
||||
| **Offline (AutoRound)** | AutoRound MXFP8 checkpoints with `data_type="mx_fp"` — auto-detected from `config.json` |
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | Support |
|
||||
|--------|---------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ⭕ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ⭕ |
|
||||
| AMD ROCm | ⭕ |
|
||||
| Intel XPU | ⭕ |
|
||||
| Ascend NPU (Atlas 950 A5) | ✅ |
|
||||
| Device | Online | Offline (Native) | Offline (AutoRound) |
|
||||
|--------|--------|------------------|---------------------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ⭕ | ⭕ | ⭕ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | ⭕ | ⭕ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ⭕ | ⭕ | ⭕ |
|
||||
| AMD ROCm | ⭕ | ⭕ | ⭕ |
|
||||
| Intel XPU | ✅ | ❌ | ✅ |
|
||||
| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ⭕ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide.
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
|
||||
|
||||
**Note**: Intel XPU only supports AutoRound MXFP8 for offline mode. Use AutoRound quantized checkpoints or online mode for XPU.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
@@ -82,9 +84,9 @@ python text_to_video.py --model <your-model> --quantization mxfp8
|
||||
vllm serve <your-model> --omni --quantization mxfp8
|
||||
```
|
||||
|
||||
### Offline Mode
|
||||
### Offline Mode (Native)
|
||||
|
||||
Offline mode loads a pre-quantized checkpoint from msModelSlim. A preprocessing
|
||||
Native offline mode loads a pre-quantized checkpoint from msModelSlim. A preprocessing
|
||||
step converts the raw quantized output to the diffusers format expected by
|
||||
vLLM-Omni and injects the quantization config into `transformer/config.json` so
|
||||
that vLLM-Omni auto-detects the offline path without a `--quantization` flag.
|
||||
@@ -159,10 +161,41 @@ omni = Omni(model="/path/to/Wan2.2-TI2V-5B-MXFP8")
|
||||
```
|
||||
|
||||
!!! note
|
||||
No `--quantization` flag is needed for offline mode. The preprocessing
|
||||
No `--quantization` flag is needed for native offline mode. The preprocessing
|
||||
script injects `quantization_config` into each `transformer/config.json`,
|
||||
which vLLM-Omni reads automatically to activate the offline MXFP8 method.
|
||||
|
||||
### Offline Mode (AutoRound)
|
||||
|
||||
AutoRound MXFP8 checkpoints declare `quant_method="auto-round"` with `data_type="mx_fp"` in their `config.json`. These are automatically detected and use the `IncMxfp8OfflineLinearMethod` backend.
|
||||
|
||||
To use an AutoRound MXFP8 checkpoint:
|
||||
|
||||
```bash
|
||||
python text_to_video.py --model <autoround-mxfp8-model>
|
||||
|
||||
# Online serving
|
||||
vllm serve <autoround-mxfp8-model> --omni
|
||||
```
|
||||
|
||||
Python API:
|
||||
|
||||
```python
|
||||
omni = Omni(model="<autoround-mxfp8-model>")
|
||||
```
|
||||
|
||||
!!! note
|
||||
AutoRound MXFP8 checkpoints are auto-detected from `config.json` and do not require a `--quantization` flag. The config must include:
|
||||
```json
|
||||
{
|
||||
"quantization_config": {
|
||||
"quant_method": "auto-round",
|
||||
"data_type": "mx_fp",
|
||||
...
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|
||||
@@ -9,20 +9,24 @@ separate quantized checkpoint.
|
||||
This mode is different from pre-quantized checkpoint formats such as GGUF,
|
||||
AutoRound, msModelSlim, or serialized Int8 checkpoints. Those formats are
|
||||
prepared before serving and are documented in their method-specific guides.
|
||||
For MXFP8 and MXFP4, use this page for load-time quantization from BF16
|
||||
checkpoints, and use the method-specific pages for offline checkpoints produced
|
||||
by msModelSlim and the merge tools.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | FP8 W8A8 | Int8 W8A8 |
|
||||
|--------|----------|-----------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ |
|
||||
| AMD ROCm | ⭕ | ⭕ |
|
||||
| Intel XPU | ⭕ | ⭕ |
|
||||
| Ascend NPU | ❌ | ✅ |
|
||||
| Device | FP8 W8A8 | Int8 W8A8 | MXFP8 W8A8 | MXFP4 W4A4 |
|
||||
|--------|----------|-----------|------------|------------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ⭕ | ⭕ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ |
|
||||
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ |
|
||||
| Intel XPU | ⭕ | ⭕ | ✅ | ⭕ |
|
||||
| Ascend NPU | ❌ | ✅ | ✅ | ✅ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide. FP8 on Ampere may use a weight-only path where available.
|
||||
guide. FP8 on Ampere may use a weight-only path where available. MXFP8 and
|
||||
MXFP4 are documented for the Ascend NPU path.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
@@ -32,6 +36,8 @@ guide. FP8 on Ampere may use a weight-only path where available.
|
||||
|--------|-------|----------------|--------|
|
||||
| FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
|
||||
| Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
|
||||
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Validated on Ascend NPU and Intel XPU |
|
||||
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B | Ascend NPU only; TI2V-5B is not supported |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
@@ -53,6 +59,9 @@ from vllm_omni import Omni
|
||||
|
||||
omni_fp8 = Omni(model="<your-model>", quantization="fp8")
|
||||
omni_int8 = Omni(model="<your-model>", quantization="int8")
|
||||
omni_mxfp8 = Omni(model="<your-model>", quantization="mxfp8")
|
||||
omni_mxfp4 = Omni(model="<your-model>", quantization="mxfp4")
|
||||
omni_mxfp4_dualscale = Omni(model="<your-model>", quantization="mxfp4_dualscale")
|
||||
```
|
||||
|
||||
CLI:
|
||||
@@ -60,6 +69,9 @@ CLI:
|
||||
```bash
|
||||
vllm serve <your-model> --omni --quantization fp8
|
||||
vllm serve <your-model> --omni --quantization int8
|
||||
vllm serve <your-model> --omni --quantization mxfp8
|
||||
vllm serve <your-model> --omni --quantization mxfp4
|
||||
vllm serve <your-model> --omni --quantization mxfp4_dualscale
|
||||
```
|
||||
|
||||
Per-component routing:
|
||||
@@ -77,10 +89,11 @@ config = build_quant_config({
|
||||
|
||||
| Parameter | Methods | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `method` | FP8, Int8 | Quantization method, usually `"fp8"` or `"int8"` |
|
||||
| `ignored_layers` | FP8, Int8 | Layer name patterns to keep in BF16/FP16 |
|
||||
| `method` | FP8, Int8, MXFP8, MXFP4 | Quantization method: `"fp8"`, `"int8"`, `"mxfp8"`, `"mxfp4"`, or `"mxfp4_dualscale"` |
|
||||
| `ignored_layers` | FP8, Int8, MXFP8, MXFP4 | Layer name patterns to keep in BF16/FP16 |
|
||||
| `activation_scheme` | FP8, Int8 | The runtime value `"dynamic"` selects online activation scaling |
|
||||
| `weight_block_size` | FP8 | Optional block-wise FP8 weight quantization size |
|
||||
| `num_bf16_fallback_layers` | MXFP4 DualScale | Leading transformer blocks to keep in BF16 for online `mxfp4_dualscale`; defaults to `5` |
|
||||
|
||||
## Validation and Notes
|
||||
|
||||
@@ -91,3 +104,5 @@ config = build_quant_config({
|
||||
model as supported.
|
||||
4. If a model already ships quantized weights, use the matching pre-quantized
|
||||
method guide instead of online quantization.
|
||||
5. For Ascend MXFP4 deployments, prefer offline `mxfp4_dualscale` checkpoints
|
||||
when production quality is more important than avoiding preprocessing.
|
||||
@@ -9,20 +9,20 @@ type has a different quantization scope.
|
||||
|
||||
| Mode | Guide | Description | Methods |
|
||||
|------|-------|-------------|---------|
|
||||
| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8 |
|
||||
| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8, MXFP4 W4A4 |
|
||||
| Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA |
|
||||
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8 |
|
||||
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale |
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | GGUF | AutoRound | msModelSlim |
|
||||
|--------|----------|-----------|----------|------------|------|-----------|-------------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ✅ | ✅ | ❌ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ✅ | ✅ | ❌ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
|
||||
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
|
||||
| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
|
||||
| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ |
|
||||
| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | GGUF | AutoRound | msModelSlim |
|
||||
|--------|----------|-----------|----------|------------|------------|------|-----------|-------------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
|
||||
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
|
||||
| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
|
||||
| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide. FP8 on Ampere may use a weight-only path where available.
|
||||
@@ -42,6 +42,7 @@ otherwise.
|
||||
| Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
|
||||
| ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints |
|
||||
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Online W8A8 or offline pre-quantized | Wan2.2-T2V-A14B, I2V-A14B, TI2V-5B | Ascend NPU only; validated for Wan2.2 |
|
||||
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | `mxfp4`: online single-scale only; `mxfp4_dualscale`: online or offline dual-scale (offline recommended) | Wan2.2-T2V-A14B, I2V-A14B | Ascend NPU only; validated for Wan2.2 A14B cascade models; TI2V-5B not supported; offline `mxfp4_dualscale` uses calibrated `mul_scale` for best accuracy |
|
||||
| GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists |
|
||||
| AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven |
|
||||
| msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path |
|
||||
@@ -58,6 +59,7 @@ in BF16 unless the model guide explicitly adds support.
|
||||
| ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
|
||||
| Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints |
|
||||
| msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
@@ -73,6 +75,7 @@ attached to the intended stage rather than applied globally.
|
||||
| Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
|
||||
| ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation |
|
||||
| MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
|
||||
| MXFP4 | [MXFP4](mxfp4.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
|
||||
| GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed |
|
||||
| AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed |
|
||||
| msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
|
||||
@@ -100,7 +103,7 @@ config = build_quant_config({
|
||||
|
||||
| Component | Default quantized? | Notes |
|
||||
|-----------|--------------------|-------|
|
||||
| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, GGUF, AutoRound, and msModelSlim |
|
||||
| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, MXFP8, MXFP4, GGUF, AutoRound, and msModelSlim |
|
||||
| Text encoder | No | Keep BF16 unless a method-specific guide documents support |
|
||||
| VAE | No | Keep BF16; storage-only paths are method-specific |
|
||||
| Scheduler/tokenizer | No | Loaded from the base model repository |
|
||||
|
||||
@@ -195,6 +195,101 @@ stages:
|
||||
devices: "0,1"
|
||||
```
|
||||
|
||||
### VAE Patch Parallelism
|
||||
|
||||
[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) splits Bagel VAE **decode/encode** tiles across multiple GPUs on the **DiT stage**, reducing **per-GPU peak memory during VAE decode**. Use it when high-resolution `text2img` or `img2img` hits VAE OOM or large decode spikes.
|
||||
|
||||
**Bagel-specific notes:**
|
||||
|
||||
- Implemented in `BagelPipeline` via `DistributedAutoEncoder` (DiT stage only).
|
||||
- **Single-stage** is the simplest path: one DiT process with TP + VAE patch parallel.
|
||||
- **Two-stage**: enable on **stage 1 (DiT)** only; stage 0 (Thinker) keeps encoder-only `VAEEncoder` and does not use VAE patch parallel.
|
||||
- You need a DiT `world_size` ≥ `vae_patch_parallel_size` (typically `tensor_parallel_size=2` on that stage). VAE PP reuses the DiT process group; it is not a standalone second-GPU VAE worker.
|
||||
|
||||
**Single-stage via deploy YAML** (recommended for `end2end.py`):
|
||||
|
||||
```yaml
|
||||
pipeline: bagel_single_stage
|
||||
async_chunk: false
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_batched_tokens: 32768
|
||||
max_num_seqs: 1
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
devices: "0,1"
|
||||
vae_use_tiling: true
|
||||
parallel_config:
|
||||
tensor_parallel_size: 2
|
||||
vae_patch_parallel_size: 2
|
||||
default_sampling_params:
|
||||
seed: 52
|
||||
```
|
||||
|
||||
```bash
|
||||
cd examples/offline_inference/bagel
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1 python end2end.py \
|
||||
--model /path/to/BAGEL-7B-MoT \
|
||||
--deploy-config /path/to/bagel_single_stage_vae_pp.yaml \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--steps 10 \
|
||||
--output ./out_vae_pp
|
||||
```
|
||||
|
||||
**Single-stage via `Omni` kwargs** (same flags as online serving):
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="ByteDance-Seed/BAGEL-7B-MoT",
|
||||
deploy_config="vllm_omni/deploy/bagel_single_stage.yaml",
|
||||
tensor_parallel_size=2,
|
||||
vae_patch_parallel_size=2,
|
||||
vae_use_tiling=True,
|
||||
)
|
||||
# Then call omni.generate(...) as in end2end.py
|
||||
```
|
||||
|
||||
**Two-stage (VAE PP on DiT only):**
|
||||
|
||||
```yaml
|
||||
stages:
|
||||
- stage_id: 0
|
||||
devices: "0"
|
||||
# AR Thinker — no vae_patch_parallel here
|
||||
|
||||
- stage_id: 1
|
||||
devices: "0,1"
|
||||
vae_use_tiling: true
|
||||
parallel_config:
|
||||
tensor_parallel_size: 2
|
||||
vae_patch_parallel_size: 2
|
||||
```
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--deploy-config /path/to/bagel_vae_pp.yaml \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat"
|
||||
```
|
||||
|
||||
**Startup log checks:**
|
||||
|
||||
```text
|
||||
INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
|
||||
```
|
||||
|
||||
| Setting | Role |
|
||||
| :------ | :--- |
|
||||
| `parallel_config.tensor_parallel_size` | DiT world size / TP (must be ≥ `vae_patch_parallel_size`) |
|
||||
| `parallel_config.vae_patch_parallel_size` | Number of ranks for distributed VAE tiles (`1` = off) |
|
||||
| `vae_use_tiling` | Enable spatial tiling (auto-enabled when `vae_patch_parallel_size > 1`) |
|
||||
|
||||
#### Hybrid Sharded Data Parallel (HSDP)
|
||||
|
||||
For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
|
||||
|
||||
@@ -121,9 +121,6 @@ def parse_args():
|
||||
help="Temperature for text generation sampling (default: 0.3).",
|
||||
)
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
@@ -154,6 +151,7 @@ def main():
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni_kwargs = vars(args).copy()
|
||||
omni_kwargs.pop("model", None)
|
||||
deploy_config = args.deploy_config
|
||||
if args.think and deploy_config is None:
|
||||
deploy_config = "vllm_omni/deploy/bagel_think.yaml"
|
||||
@@ -164,9 +162,7 @@ def main():
|
||||
if args.quantization:
|
||||
omni_kwargs["quantization_config"] = args.quantization
|
||||
|
||||
# Override CLI --model with the derived model_name.
|
||||
omni_kwargs["model"] = model_name
|
||||
omni = Omni(**omni_kwargs)
|
||||
omni = Omni(model=model_name, **omni_kwargs)
|
||||
|
||||
formatted_prompts = []
|
||||
for p in prompts:
|
||||
|
||||
@@ -0,0 +1,304 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
WORKER_EXTENSION = "vllm_omni.diffusion.models.dreamzero.video_export_worker.DreamZeroVideoExportWorkerExtension"
|
||||
ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
|
||||
DEFAULT_SESSION_PREFIX = "dreamzero-export"
|
||||
RELATIVE_OFFSETS = [-23, -16, -8, 0]
|
||||
ACTION_HORIZON = 24
|
||||
CAMERA_FILES = {
|
||||
"observation/exterior_image_0_left": "exterior_image_1_left.mp4",
|
||||
"observation/exterior_image_1_left": "exterior_image_2_left.mp4",
|
||||
"observation/wrist_image_left": "wrist_image_left.mp4",
|
||||
}
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
repo_root = Path(__file__).resolve().parents[3]
|
||||
parser = argparse.ArgumentParser(description="Export DreamZero prediction video from downloaded example inputs.")
|
||||
parser.add_argument("--model", default="GEAR-Dreams/DreamZero-DROID")
|
||||
parser.add_argument("--deploy-config", type=Path, required=True)
|
||||
parser.add_argument(
|
||||
"--video-dir",
|
||||
type=Path,
|
||||
default=repo_root / "outputs" / "dreamzero" / "assets",
|
||||
help="Directory containing the three camera MP4 files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir", type=Path, default=repo_root / "outputs" / "dreamzero" / "generated_predictions"
|
||||
)
|
||||
parser.add_argument("--output-stem", default="dreamzero_prediction")
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
|
||||
)
|
||||
parser.add_argument("--session-id", default=None)
|
||||
parser.add_argument("--save-input-video", action="store_true")
|
||||
parser.add_argument("--save-gif", action="store_true")
|
||||
parser.add_argument("--save-actions", action="store_true")
|
||||
parser.add_argument("--fps", type=int, default=5)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _load_all_frames(video_path: Path) -> np.ndarray:
|
||||
cap = cv2.VideoCapture(str(video_path))
|
||||
frames = []
|
||||
while True:
|
||||
ok, frame = cap.read()
|
||||
if not ok:
|
||||
break
|
||||
frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
cap.release()
|
||||
if not frames:
|
||||
raise RuntimeError(f"No frames loaded from {video_path}")
|
||||
return np.stack(frames, axis=0)
|
||||
|
||||
|
||||
def _load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
|
||||
camera_frames: dict[str, np.ndarray] = {}
|
||||
for camera_key, file_name in CAMERA_FILES.items():
|
||||
video_path = video_dir / file_name
|
||||
if not video_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"Missing DreamZero example asset: {video_path}. "
|
||||
"Download the example videos with: "
|
||||
f"`hf download {ASSET_REPO_ID} --repo-type dataset --local-dir {video_dir}`"
|
||||
)
|
||||
camera_frames[camera_key] = _load_all_frames(video_path)
|
||||
return camera_frames
|
||||
|
||||
|
||||
def _build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
|
||||
chunks: list[list[int]] = []
|
||||
current_frame = 23
|
||||
for _ in range(num_chunks):
|
||||
indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
|
||||
if indices[-1] >= total_frames:
|
||||
break
|
||||
chunks.append(indices)
|
||||
current_frame += ACTION_HORIZON
|
||||
return chunks
|
||||
|
||||
|
||||
def _make_obs_from_video(
|
||||
camera_frames: dict[str, np.ndarray],
|
||||
frame_indices: list[int],
|
||||
*,
|
||||
prompt: str,
|
||||
session_id: str,
|
||||
) -> dict:
|
||||
obs: dict = {}
|
||||
for camera_key, all_frames in camera_frames.items():
|
||||
selected = all_frames[frame_indices]
|
||||
obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
|
||||
|
||||
obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
|
||||
obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
|
||||
obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
|
||||
obs["prompt"] = prompt
|
||||
obs["session_id"] = session_id
|
||||
return obs
|
||||
|
||||
|
||||
def _build_observations(video_dir: Path, prompt: str, session_id: str) -> tuple[dict[str, np.ndarray], list[dict]]:
|
||||
camera_frames = _load_camera_frames(video_dir)
|
||||
total_frames = min(frames.shape[0] for frames in camera_frames.values())
|
||||
chunks = _build_frame_schedule(total_frames, 1)
|
||||
observations = [
|
||||
_make_obs_from_video(camera_frames, [0], prompt=prompt, session_id=session_id),
|
||||
]
|
||||
if chunks:
|
||||
observations.append(
|
||||
_make_obs_from_video(
|
||||
camera_frames,
|
||||
chunks[0],
|
||||
prompt=prompt,
|
||||
session_id=session_id,
|
||||
)
|
||||
)
|
||||
if len(observations) < 2:
|
||||
raise RuntimeError("Need at least two DreamZero example observations to export a prediction video.")
|
||||
return camera_frames, observations[:2]
|
||||
|
||||
|
||||
def _extract_latents(output: OmniRequestOutput) -> torch.Tensor:
|
||||
if not isinstance(output, OmniRequestOutput):
|
||||
raise TypeError(f"Expected OmniRequestOutput, got {type(output)!r}")
|
||||
if not output.images:
|
||||
raise RuntimeError("DreamZero output does not contain video latents in `images`.")
|
||||
|
||||
latents = output.images[0]
|
||||
if not isinstance(latents, torch.Tensor):
|
||||
raise TypeError(f"Expected tensor latents, got {type(latents)!r}")
|
||||
|
||||
latents = latents.detach().cpu()
|
||||
if latents.dim() == 4:
|
||||
latents = latents.unsqueeze(0)
|
||||
if latents.dim() != 5:
|
||||
raise ValueError(f"Unexpected latent shape: {tuple(latents.shape)}")
|
||||
|
||||
if latents.shape[1] < latents.shape[2]:
|
||||
latents = latents.transpose(1, 2).contiguous()
|
||||
return latents
|
||||
|
||||
|
||||
def _write_mp4(path: Path, frames: np.ndarray, fps: int) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
height, width = frames.shape[1:3]
|
||||
writer = cv2.VideoWriter(
|
||||
str(path),
|
||||
cv2.VideoWriter_fourcc(*"mp4v"),
|
||||
float(fps),
|
||||
(width, height),
|
||||
)
|
||||
if not writer.isOpened():
|
||||
raise RuntimeError(f"Failed to open video writer for {path}")
|
||||
try:
|
||||
for frame in frames:
|
||||
writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
|
||||
finally:
|
||||
writer.release()
|
||||
|
||||
|
||||
def _write_gif(path: Path, frames: np.ndarray, fps: int) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
images = [Image.fromarray(frame) for frame in frames]
|
||||
duration_ms = max(int(round(1000 / max(fps, 1))), 1)
|
||||
images[0].save(
|
||||
path,
|
||||
save_all=True,
|
||||
append_images=images[1:],
|
||||
duration=duration_ms,
|
||||
loop=0,
|
||||
)
|
||||
|
||||
|
||||
def _stitch_input_frames(camera_frames: dict[str, np.ndarray]) -> np.ndarray:
|
||||
total_frames = min(frames.shape[0] for frames in camera_frames.values())
|
||||
stitched = []
|
||||
for frame_index in range(total_frames):
|
||||
left = camera_frames["observation/exterior_image_0_left"][frame_index]
|
||||
right = camera_frames["observation/exterior_image_1_left"][frame_index]
|
||||
wrist = camera_frames["observation/wrist_image_left"][frame_index]
|
||||
pad = np.zeros((left.shape[0], left.shape[1], 3), dtype=np.uint8)
|
||||
canvas = np.concatenate([left, right], axis=1)
|
||||
bottom = np.concatenate([wrist, pad], axis=1)
|
||||
stitched.append(np.concatenate([canvas, bottom], axis=0))
|
||||
return np.stack(stitched, axis=0)
|
||||
|
||||
|
||||
def _run_generation(
|
||||
model: str, deploy_config_path: Path, observations: list[dict]
|
||||
) -> tuple[Omni, list[OmniRequestOutput]]:
|
||||
omni = Omni(
|
||||
model=model,
|
||||
deploy_config=str(deploy_config_path),
|
||||
enforce_eager=True,
|
||||
worker_extension_cls=WORKER_EXTENSION,
|
||||
)
|
||||
|
||||
outputs: list[OmniRequestOutput] = []
|
||||
for index, obs in enumerate(observations):
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
extra_args={
|
||||
"reset": index == 0,
|
||||
"session_id": obs["session_id"],
|
||||
"robot_obs": obs,
|
||||
}
|
||||
)
|
||||
result = omni.generate(obs["prompt"], sampling_params_list=[sampling_params])
|
||||
if not result:
|
||||
raise RuntimeError(f"No output returned for DreamZero request {index}")
|
||||
outputs.append(result[0])
|
||||
return omni, outputs
|
||||
|
||||
|
||||
def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray:
|
||||
stage_client = omni.engine.stage_clients[0]
|
||||
engine = getattr(stage_client, "_engine", None)
|
||||
if engine is None:
|
||||
raise RuntimeError("DreamZero export requires inline diffusion stage access.")
|
||||
|
||||
decoded = engine.executor.collective_rpc(
|
||||
"decode_video_latents_to_uint8",
|
||||
args=(full_latents,),
|
||||
unique_reply_rank=0,
|
||||
exec_all_ranks=True,
|
||||
)
|
||||
if isinstance(decoded, torch.Tensor):
|
||||
decoded = decoded.numpy()
|
||||
if not isinstance(decoded, np.ndarray):
|
||||
raise TypeError(f"Unexpected decoded output type: {type(decoded)!r}")
|
||||
return decoded
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _parse_args()
|
||||
session_id = args.session_id or f"{DEFAULT_SESSION_PREFIX}-{uuid.uuid4()}"
|
||||
|
||||
camera_frames, observations = _build_observations(
|
||||
video_dir=args.video_dir,
|
||||
prompt=args.prompt,
|
||||
session_id=session_id,
|
||||
)
|
||||
|
||||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if args.save_input_video:
|
||||
input_frames = _stitch_input_frames(camera_frames)
|
||||
_write_mp4(args.output_dir / f"{args.output_stem}_input.mp4", input_frames, fps=15)
|
||||
if args.save_gif:
|
||||
_write_gif(args.output_dir / f"{args.output_stem}_input.gif", input_frames[::3], fps=5)
|
||||
|
||||
omni = None
|
||||
try:
|
||||
omni, outputs = _run_generation(
|
||||
model=args.model,
|
||||
deploy_config_path=args.deploy_config,
|
||||
observations=observations,
|
||||
)
|
||||
latent_steps = [_extract_latents(output) for output in outputs]
|
||||
full_latents = torch.cat(latent_steps, dim=2)
|
||||
frames = _decode_with_worker(omni, full_latents)
|
||||
finally:
|
||||
if omni is not None:
|
||||
omni.close()
|
||||
|
||||
mp4_path = args.output_dir / f"{args.output_stem}.mp4"
|
||||
|
||||
_write_mp4(mp4_path, frames, fps=args.fps)
|
||||
print(f"SAVED_MP4={mp4_path}")
|
||||
|
||||
if args.save_gif:
|
||||
gif_path = args.output_dir / f"{args.output_stem}.gif"
|
||||
_write_gif(gif_path, frames, fps=args.fps)
|
||||
print(f"SAVED_GIF={gif_path}")
|
||||
|
||||
if args.save_actions:
|
||||
npz_path = args.output_dir / f"{args.output_stem}_actions.npz"
|
||||
np.savez(
|
||||
npz_path,
|
||||
step0=np.asarray(outputs[0].multimodal_output.get("actions")),
|
||||
step1=np.asarray(outputs[1].multimodal_output.get("actions")),
|
||||
)
|
||||
print(f"SAVED_ACTIONS={npz_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -42,7 +42,7 @@ pip install \
|
||||
## 3. Hardware and VRAM Requirements
|
||||
|
||||
This example uses a 3-stage pipeline on one GPU by default
|
||||
([`dynin_omni.yaml`](../../../vllm_omni/model_executor/stage_configs/dynin_omni.yaml)):
|
||||
([`dynin_omni.yaml`](../../../vllm_omni/deploy/dynin_omni.yaml)):
|
||||
|
||||
- Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
|
||||
- Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`
|
||||
|
||||
@@ -18,8 +18,6 @@ import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
TASK_CHOICES = ("t2t", "t2i", "t2s", "i2i", "i2t", "s2t", "v2t")
|
||||
|
||||
TASK_DEFAULT_RUNTIME = {
|
||||
@@ -880,8 +878,8 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
|
||||
parser.add_argument(
|
||||
"--stage-config-path",
|
||||
type=str,
|
||||
default=str(repo_root / "vllm_omni/model_executor/stage_configs/dynin_omni.yaml"),
|
||||
help="Path to stage config yaml.",
|
||||
default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
|
||||
help="Path to deploy config yaml.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dynin-config-path",
|
||||
@@ -972,8 +970,6 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
|
||||
parser.add_argument("--vq-model-audio-local-files-only", action=argparse.BooleanOptionalAction, default=None)
|
||||
|
||||
parser.add_argument("--disable-hf-xet", action=argparse.BooleanOptionalAction, default=True)
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -17,7 +17,7 @@ from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniPromptType
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3.yaml")
|
||||
_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image_3_moe.yaml")
|
||||
_DEFAULT_AR_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml")
|
||||
|
||||
_MODALITY_TASK_MAP: dict[str, tuple[str, str | None]] = {
|
||||
@@ -61,8 +61,8 @@ def parse_args():
|
||||
parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=5.0, help="Classifier-free guidance scale.")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--height", type=int, default=1024, help="Output image height.")
|
||||
parser.add_argument("--width", type=int, default=1024, help="Output image width.")
|
||||
parser.add_argument("--height", type=int, default=None, help="Output image height.")
|
||||
parser.add_argument("--width", type=int, default=None, help="Output image width.")
|
||||
parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
|
||||
parser.add_argument(
|
||||
"--bot-task",
|
||||
@@ -107,9 +107,6 @@ def parse_args():
|
||||
),
|
||||
)
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -229,7 +226,21 @@ def main():
|
||||
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
ar_stop_token_ids = resolve_stop_token_ids(task=task, bot_task=bot_task, tokenizer=tokenizer)
|
||||
if (args.height is None) != (args.width is None):
|
||||
raise ValueError("--height and --width must both be specified or both omitted.")
|
||||
user_specified_size = args.height is not None and args.width is not None
|
||||
if args.modality in ("img2text", "text2text"):
|
||||
ar_image_size = "auto"
|
||||
elif user_specified_size:
|
||||
ar_image_size = f"{args.width}x{args.height}"
|
||||
else:
|
||||
ar_image_size = None
|
||||
ar_stop_token_ids = resolve_stop_token_ids(
|
||||
task=task, bot_task=bot_task, tokenizer=tokenizer, image_size=ar_image_size
|
||||
)
|
||||
print(
|
||||
f"[AR Config] task={task}, bot_task={bot_task}, image_size={ar_image_size}, stop_token_ids={ar_stop_token_ids}"
|
||||
)
|
||||
for sp in params_list:
|
||||
if isinstance(sp, OmniDiffusionSamplingParams):
|
||||
sp.num_inference_steps = args.steps
|
||||
|
||||
@@ -173,8 +173,8 @@ def parse_args() -> argparse.Namespace:
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["fp8", "mxfp8", "int8", "gguf"],
|
||||
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
|
||||
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
|
||||
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
|
||||
)
|
||||
|
||||
# Distributed and parallel execution
|
||||
|
||||
@@ -127,7 +127,12 @@ def run_one_path(
|
||||
),
|
||||
args.device,
|
||||
)
|
||||
pred = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
pred = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-sample-{index}",
|
||||
)
|
||||
pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
|
||||
results.append(
|
||||
{
|
||||
@@ -150,6 +155,7 @@ def run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs: dict[str, torch.Tensor],
|
||||
noise: torch.Tensor,
|
||||
request_id: str,
|
||||
) -> torch.Tensor:
|
||||
output = pipeline.forward(
|
||||
OmniDiffusionRequest(
|
||||
@@ -161,6 +167,7 @@ def run_pipeline_forward(
|
||||
"decode_image": False,
|
||||
}
|
||||
),
|
||||
request_id=request_id,
|
||||
)
|
||||
)
|
||||
if output.error:
|
||||
@@ -212,23 +219,38 @@ def benchmark_forward(
|
||||
|
||||
_synchronize(args.device)
|
||||
cold_start_begin = time.perf_counter()
|
||||
pred = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
pred = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-benchmark-{index}-cold",
|
||||
)
|
||||
_synchronize(args.device)
|
||||
cold_start_ms = (time.perf_counter() - cold_start_begin) * 1000.0
|
||||
|
||||
warmup_ms: list[float] = []
|
||||
for _ in range(args.warmup_iters):
|
||||
for iter_idx in range(args.warmup_iters):
|
||||
_synchronize(args.device)
|
||||
begin = time.perf_counter()
|
||||
_ = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
_ = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-benchmark-{index}-warmup-{iter_idx}",
|
||||
)
|
||||
_synchronize(args.device)
|
||||
warmup_ms.append((time.perf_counter() - begin) * 1000.0)
|
||||
|
||||
benchmark_ms: list[float] = []
|
||||
for _ in range(args.benchmark_iters):
|
||||
for iter_idx in range(args.benchmark_iters):
|
||||
_synchronize(args.device)
|
||||
begin = time.perf_counter()
|
||||
_ = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
_ = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-benchmark-{index}-iter-{iter_idx}",
|
||||
)
|
||||
_synchronize(args.device)
|
||||
benchmark_ms.append((time.perf_counter() - begin) * 1000.0)
|
||||
|
||||
@@ -295,7 +317,12 @@ def main() -> None:
|
||||
dataset=dataset,
|
||||
train_meta=train_meta,
|
||||
collate_samples=collate_open_loop_samples,
|
||||
run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(policy, batch_inputs, noise),
|
||||
run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(
|
||||
policy,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id="internvla-a1-open-loop",
|
||||
),
|
||||
num_episodes=args.num_episodes,
|
||||
seed=args.seed,
|
||||
device=args.device,
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
# Lance (ByteDance) — offline inference
|
||||
|
||||
[Lance](https://huggingface.co/bytedance-research/Lance) is a 3B unified
|
||||
autoregressive + diffusion multimodal model on a Qwen2.5-VL backbone. It is
|
||||
**BAGEL-lineage** (ByteDance Mixture-of-Transformers): the released `Lance_3B`
|
||||
checkpoint uses the same `*_moe_gen` MoT weight layout as BAGEL, so vLLM-Omni
|
||||
implements it by reusing the BAGEL transformer core and specializing only the
|
||||
ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the checkpoint layout.
|
||||
|
||||
This example covers all six Lance modalities from the upstream HF model card:
|
||||
`t2i`, `t2v`, `image_edit`, `video_edit`, `x2t_image` (image understanding) and
|
||||
`x2t_video` (video understanding).
|
||||
|
||||
## Hardware
|
||||
|
||||
Single NVIDIA GPU with 16 GB+ VRAM in BF16 (we test on B300 / A100). CUDA ≥ 12.4.
|
||||
|
||||
## Run
|
||||
|
||||
```bash
|
||||
# Text-to-image
|
||||
python examples/offline_inference/lance/end2end.py \
|
||||
--model bytedance-research/Lance \
|
||||
--prompts "a corgi astronaut on the moon, cinematic" \
|
||||
--steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
|
||||
--height 1024 --width 1024 \
|
||||
--output ./out
|
||||
|
||||
# Text-to-video (uses the Lance_3B_Video subfolder; see ``--modality``
|
||||
# choices for all six task variants)
|
||||
python examples/offline_inference/lance/end2end.py \
|
||||
--model bytedance-research/Lance/Lance_3B_Video --modality text2video \
|
||||
--num-frames 25 --video-height 480 --video-width 768 \
|
||||
--prompts "a cat playing piano, cinematic" \
|
||||
--steps 30 --fps 8 --output ./out
|
||||
```
|
||||
|
||||
`video_edit` requires `--model bytedance-research/Lance/Lance_3B_Video` so
|
||||
the 3-D `latent_pos_embed` table is loaded; the other paths can point at
|
||||
the top-level `bytedance-research/Lance` repo and resolve the right
|
||||
sub-checkpoint automatically.
|
||||
|
||||
The HF repo bundles everything (`Lance_3B/`, `Lance_3B_Video/`,
|
||||
`Qwen2.5-VL-ViT/`, `Wan2.2_VAE.pth`); no separate downloads are required.
|
||||
|
||||
## Defaults
|
||||
|
||||
Matches upstream `inference_lance.sh`: 30 denoising steps, timestep-shift 3.5,
|
||||
text CFG 4.0, seed 42, 1024×1024 (override with `--height` / `--width`). For
|
||||
the understanding paths (`img2text` / `video2text`), sampling is enabled by
|
||||
default at `--text-temperature 0.8` because Lance's greedy decoder emits an
|
||||
immediate EOS for many prompts.
|
||||
@@ -0,0 +1,339 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Lance-3B offline inference (single-stage).
|
||||
|
||||
Lance is BAGEL-lineage; this example covers the single-stage paths wired in
|
||||
this PR:
|
||||
|
||||
text2img — image generation from text
|
||||
img2text — image understanding (caption / VQA)
|
||||
|
||||
Examples:
|
||||
|
||||
python examples/offline_inference/lance/end2end.py \
|
||||
--model bytedance-research/Lance \
|
||||
--prompts "a corgi astronaut on the moon, cinematic" \
|
||||
--output ./out
|
||||
|
||||
python examples/offline_inference/lance/end2end.py \
|
||||
--model bytedance-research/Lance --modality img2text \
|
||||
--image-path /path/to/photo.jpg \
|
||||
--prompts "Describe this image in detail."
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--model", default="bytedance-research/Lance", help="HF repo or local path.")
|
||||
p.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
|
||||
p.add_argument("--txt-prompts", type=str, default=None, help="File with one prompt per line.")
|
||||
p.add_argument(
|
||||
"--modality",
|
||||
default="text2img",
|
||||
choices=["text2img", "img2text", "text2video", "video2text", "img2img", "video2video", "image2video"],
|
||||
help="Lance single-stage modality.",
|
||||
)
|
||||
p.add_argument("--image-path", type=str, default=None, help="Input image for img2text.")
|
||||
p.add_argument("--video-path", type=str, default=None, help="Input video file for video2text.")
|
||||
p.add_argument("--num-frames", type=int, default=25, help="Number of RGB frames for text2video (max 121).")
|
||||
p.add_argument("--video-height", type=int, default=480, help="Video frame height.")
|
||||
p.add_argument("--video-width", type=int, default=768, help="Video frame width.")
|
||||
p.add_argument("--height", type=int, default=None, help="Image height (t2i). Default = max_hw (1024).")
|
||||
p.add_argument("--width", type=int, default=None, help="Image width (t2i). Default = max_hw (1024).")
|
||||
p.add_argument(
|
||||
"--fps", type=int, default=12, help="Output video FPS when saving MP4 (matches upstream Lance's save_fps=12)."
|
||||
)
|
||||
p.add_argument("--output", type=str, default=".", help="Output directory.")
|
||||
p.add_argument("--steps", type=int, default=30, help="Denoising steps (Lance default 30).")
|
||||
p.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (Lance default 4.0).")
|
||||
p.add_argument("--timestep-shift", type=float, default=3.5, help="Flow-match timestep shift (Lance default 3.5).")
|
||||
p.add_argument("--negative-prompt", type=str, default=None)
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument("--max-text-tokens", type=int, default=512, help="Max tokens for img2text generation.")
|
||||
p.add_argument(
|
||||
"--do-sample",
|
||||
action="store_true",
|
||||
default=True,
|
||||
help="Sample (vs greedy) for x2t generation; greedy frequently emits an immediate EOS for Lance.",
|
||||
)
|
||||
p.add_argument("--no-sample", dest="do_sample", action="store_false", help="Disable sampling (greedy decoding).")
|
||||
p.add_argument(
|
||||
"--text-temperature",
|
||||
type=float,
|
||||
default=0.8,
|
||||
help="Sampling temperature for x2t generation. Lance frequently emits an immediate EOS below ~0.7; 0.8 is a good default.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--system-prompt",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Override the task system prompt (used by x2t_image / x2t_video for "
|
||||
'per-example QA instructions like "Look at the image carefully and '
|
||||
'answer the question."; without it x2t falls back to a '
|
||||
"caption-style default and the model describes instead of answering)."
|
||||
),
|
||||
)
|
||||
# Lance is single-stage diffusion — no deploy YAML needed. Required
|
||||
# engine knobs (``pipeline``, ``enforce_eager``, ``trust_remote_code``,
|
||||
# ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
|
||||
# and ``create_default_diffusion`` materializes the stage config.
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(p)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
from vllm_omni.diffusion.models.lance.prompts import ( # noqa: E402
|
||||
VIDEO_PAD,
|
||||
VISION_END,
|
||||
VISION_START,
|
||||
render_lance_prompt,
|
||||
)
|
||||
|
||||
# Visual block matching upstream Lance: <|vision_start|><|video_pad|><|vision_end|>
|
||||
# (Lance's renderer uses ``<|video_pad|>`` even for image inputs by default.)
|
||||
_VISION_BLOCK = f"{VISION_START}{VIDEO_PAD}{VISION_END}"
|
||||
|
||||
|
||||
def _format_text2img_prompts(prompts):
|
||||
return [{"prompt": render_lance_prompt("t2i", p), "modalities": ["image"]} for p in prompts]
|
||||
|
||||
|
||||
def _format_text2video_prompts(prompts, num_frames, video_h, video_w):
|
||||
extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
|
||||
return [{"prompt": render_lance_prompt("t2v", p), "modalities": ["video"], "extra_args": extra} for p in prompts]
|
||||
|
||||
|
||||
def _format_image2video_prompts(prompts, image_path, num_frames, video_h, video_w):
|
||||
"""Image-to-Video (no first-frame pin): image + text → long video.
|
||||
|
||||
Passes the input image via ``multi_modal_data.first_frame``. The pipeline
|
||||
treats it as a 1-frame reference (VAE+ViT prefill) and generates a fresh
|
||||
multi-frame video at the requested ``num_frames`` × ``video_h`` ×
|
||||
``video_w`` shape.
|
||||
"""
|
||||
import os as _os
|
||||
|
||||
if not image_path or not _os.path.exists(image_path):
|
||||
raise ValueError(f"image2video requires --image-path pointing to an existing file, got: {image_path}")
|
||||
from PIL import Image as _Image
|
||||
|
||||
img = _Image.open(image_path).convert("RGB")
|
||||
extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
|
||||
return [
|
||||
{
|
||||
"prompt": render_lance_prompt("i2v", p, vision_token=_VISION_BLOCK),
|
||||
"multi_modal_data": {"first_frame": img},
|
||||
"modalities": ["video"],
|
||||
"extra_args": extra,
|
||||
}
|
||||
for p in prompts
|
||||
]
|
||||
|
||||
|
||||
def _format_img2img_prompts(prompts, image_path):
|
||||
import os as _os
|
||||
|
||||
if not image_path or not _os.path.exists(image_path):
|
||||
raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {image_path}")
|
||||
from PIL import Image as _Image
|
||||
|
||||
img = _Image.open(image_path).convert("RGB")
|
||||
return [
|
||||
{
|
||||
"prompt": render_lance_prompt("image_edit", p, vision_token=_VISION_BLOCK),
|
||||
"multi_modal_data": {"img2img": img},
|
||||
"modalities": ["image"],
|
||||
}
|
||||
for p in prompts
|
||||
]
|
||||
|
||||
|
||||
def _format_video2video_prompts(prompts, video_path, num_frames, video_h, video_w):
|
||||
import os as _os
|
||||
|
||||
if not video_path or not _os.path.exists(video_path):
|
||||
raise ValueError(f"video2video requires --video-path pointing to an existing file, got: {video_path}")
|
||||
|
||||
extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
|
||||
return [
|
||||
{
|
||||
"prompt": render_lance_prompt("video_edit", p, vision_token=_VISION_BLOCK),
|
||||
"multi_modal_data": {"video": video_path},
|
||||
"modalities": ["video"],
|
||||
"extra_args": extra,
|
||||
}
|
||||
for p in prompts
|
||||
]
|
||||
|
||||
|
||||
def _format_video2text_prompts(prompts, video_path, system_prompt=None):
|
||||
import os as _os
|
||||
|
||||
if not video_path or not _os.path.exists(video_path):
|
||||
raise ValueError(f"video2text requires --video-path pointing to an existing file, got: {video_path}")
|
||||
import imageio.v3 as iio
|
||||
|
||||
video = iio.imread(video_path)
|
||||
return [
|
||||
{
|
||||
"prompt": render_lance_prompt("x2t_video", p, vision_token=_VISION_BLOCK, system_prompt=system_prompt),
|
||||
"multi_modal_data": {"video": video},
|
||||
"modalities": ["text"],
|
||||
}
|
||||
for p in prompts
|
||||
]
|
||||
|
||||
|
||||
def _format_img2text_prompts(prompts, image_path, system_prompt=None):
|
||||
from PIL import Image
|
||||
|
||||
if not image_path or not os.path.exists(image_path):
|
||||
raise ValueError(f"img2text requires --image-path pointing to an existing file, got: {image_path}")
|
||||
img = Image.open(image_path).convert("RGB")
|
||||
return [
|
||||
{
|
||||
"prompt": render_lance_prompt("x2t_image", p, vision_token=_VISION_BLOCK, system_prompt=system_prompt),
|
||||
"multi_modal_data": {"image": img},
|
||||
"modalities": ["text"],
|
||||
}
|
||||
for p in prompts
|
||||
]
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
os.makedirs(args.output, exist_ok=True)
|
||||
|
||||
# Lance aligns more naturally with upstream when using FLASH_ATTN
|
||||
# (upstream uses flash_attn_varlen_func; SDPA accumulates ~5x more
|
||||
# numerical drift through the 36-layer Qwen2 stack on B300). Default
|
||||
# to flash_attn when the user hasn't explicitly set a backend.
|
||||
os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
|
||||
|
||||
if args.txt_prompts:
|
||||
with open(args.txt_prompts, encoding="utf-8") as f:
|
||||
prompts = [ln.strip() for ln in f if ln.strip()]
|
||||
elif args.modality == "img2text":
|
||||
prompts = args.prompts or ["Describe this image in detail."]
|
||||
else:
|
||||
prompts = args.prompts or ["A cute cat, studio lighting, 4k"]
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
# Video modalities require the Lance_3B_Video checkpoint (3D
|
||||
# latent_pos_embed shaped (126976, 2048)). Without it, vllm-omni silently
|
||||
# loads Lance_3B (4096, 2048) — the image-only table — and any t_lat >= 1
|
||||
# immediately indexes out of bounds. Rewrite the model path so
|
||||
# LancePipeline._select_video_variant picks the video subfolder.
|
||||
resolved_model = args.model
|
||||
if args.modality in ("text2video", "video2video", "video2text", "image2video"):
|
||||
if not str(resolved_model).rstrip("/").endswith("Lance_3B_Video"):
|
||||
candidate = os.path.join(resolved_model, "Lance_3B_Video")
|
||||
if os.path.isdir(candidate):
|
||||
resolved_model = candidate
|
||||
|
||||
omni_kwargs = vars(args).copy()
|
||||
omni_kwargs["model"] = resolved_model
|
||||
# Lance single-stage defaults (formerly in vllm_omni/deploy/lance.yaml):
|
||||
omni_kwargs.setdefault("pipeline", "lance")
|
||||
omni_kwargs.setdefault("max_num_batched_tokens", 32768)
|
||||
omni_kwargs.setdefault("max_num_seqs", 1)
|
||||
omni_kwargs.setdefault("enforce_eager", True)
|
||||
omni_kwargs.setdefault("trust_remote_code", True)
|
||||
omni_kwargs.setdefault("enable_prefix_caching", False)
|
||||
omni_kwargs.setdefault("async_chunk", False)
|
||||
omni = Omni(**omni_kwargs)
|
||||
|
||||
if args.modality == "img2text":
|
||||
formatted = _format_img2text_prompts(prompts, args.image_path, system_prompt=args.system_prompt)
|
||||
elif args.modality == "text2video":
|
||||
formatted = _format_text2video_prompts(prompts, args.num_frames, args.video_height, args.video_width)
|
||||
elif args.modality == "video2text":
|
||||
formatted = _format_video2text_prompts(prompts, args.video_path, system_prompt=args.system_prompt)
|
||||
elif args.modality == "img2img":
|
||||
formatted = _format_img2img_prompts(prompts, args.image_path)
|
||||
elif args.modality == "video2video":
|
||||
formatted = _format_video2video_prompts(
|
||||
prompts, args.video_path, args.num_frames, args.video_height, args.video_width
|
||||
)
|
||||
elif args.modality == "image2video":
|
||||
formatted = _format_image2video_prompts(
|
||||
prompts, args.image_path, args.num_frames, args.video_height, args.video_width
|
||||
)
|
||||
else:
|
||||
formatted = _format_text2img_prompts(prompts)
|
||||
|
||||
params_list = omni.default_sampling_params_list
|
||||
diffusion_params = params_list[0] # single-stage: one param set
|
||||
diffusion_params.num_inference_steps = args.steps # type: ignore
|
||||
if args.seed is not None:
|
||||
diffusion_params.seed = args.seed # type: ignore
|
||||
if args.height is not None:
|
||||
diffusion_params.height = args.height # type: ignore
|
||||
if args.width is not None:
|
||||
diffusion_params.width = args.width # type: ignore
|
||||
extra = getattr(diffusion_params, "extra_args", {}) or {}
|
||||
extra["cfg_text_scale"] = args.cfg_text_scale
|
||||
extra["timestep_shift"] = args.timestep_shift
|
||||
if args.modality in ("img2text", "video2text"):
|
||||
extra["max_think_tokens"] = args.max_text_tokens
|
||||
extra["do_sample"] = args.do_sample
|
||||
extra["text_temperature"] = args.text_temperature
|
||||
if args.negative_prompt is not None:
|
||||
extra["negative_prompt"] = args.negative_prompt
|
||||
diffusion_params.extra_args = extra # type: ignore
|
||||
|
||||
outputs = list(omni.generate(prompts=formatted, sampling_params_list=params_list))
|
||||
|
||||
if args.modality in ("img2text", "video2text"):
|
||||
for i, req_output in enumerate(outputs):
|
||||
text = getattr(req_output, "output", None) or getattr(req_output, "text", None)
|
||||
if text is None and hasattr(req_output, "custom_output"):
|
||||
text = (req_output.custom_output or {}).get("text_output")
|
||||
print(f"[Output {i}] {text}")
|
||||
return
|
||||
|
||||
if args.modality in ("text2video", "video2video", "image2video"):
|
||||
for i, req_output in enumerate(outputs):
|
||||
custom = getattr(req_output, "custom_output", None) or {}
|
||||
frames = custom.get("video_frames")
|
||||
if not frames:
|
||||
print(f"[Output {i}] no video frames returned")
|
||||
continue
|
||||
mp4_path = os.path.join(args.output, f"lance_{i}.mp4")
|
||||
try:
|
||||
import imageio.v3 as iio
|
||||
import numpy as np
|
||||
|
||||
arr = np.stack([np.asarray(f) for f in frames], axis=0)
|
||||
iio.imwrite(mp4_path, arr, fps=args.fps, codec="libx264", quality=8)
|
||||
print(f"[Output {i}] Saved MP4 ({len(frames)} frames, {arr.shape[1]}x{arr.shape[2]}) to {mp4_path}")
|
||||
except Exception as exc:
|
||||
print(f"[Output {i}] MP4 encode failed ({exc!r}); writing per-frame ONGs instead")
|
||||
frame_dir = os.path.join(args.output, f"lance_{i}_frames")
|
||||
os.makedirs(frame_dir, exist_ok=True)
|
||||
for j, frame in enumerate(frames):
|
||||
frame.save(os.path.join(frame_dir, f"{j:04d}.png"))
|
||||
print(f" → {frame_dir} ({len(frames)} frames)")
|
||||
return
|
||||
|
||||
idx = 0
|
||||
for req_output in outputs:
|
||||
images = getattr(req_output, "images", None)
|
||||
if not images:
|
||||
continue
|
||||
for j, img in enumerate(images):
|
||||
path = os.path.join(args.output, f"lance_{idx}_{j}.png")
|
||||
img.save(path)
|
||||
print(f"[Output] Saved image to {path}")
|
||||
idx += 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large.
Load diff
@@ -19,7 +19,6 @@ from vllm import SamplingParams
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
DEFAULT_SYSTEM = "You are a helpful assistant."
|
||||
DEFAULT_QUESTION = "Please summarize the content of this image."
|
||||
@@ -49,7 +48,6 @@ def parse_args() -> argparse.Namespace:
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -29,7 +29,6 @@ from PIL import Image
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -118,7 +117,6 @@ def parse_args() -> argparse.Namespace:
|
||||
)
|
||||
p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
|
||||
p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
|
||||
nullify_stage_engine_defaults(p)
|
||||
args = p.parse_args()
|
||||
if not args.prompt:
|
||||
args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]
|
||||
|
||||
@@ -21,11 +21,10 @@ from message_convert import (
|
||||
to_prompt,
|
||||
)
|
||||
from vllm import SamplingParams
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniTokensPrompt
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
SEED = 42
|
||||
|
||||
@@ -343,7 +342,7 @@ def main(args):
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser.add_argument(
|
||||
"--model-name",
|
||||
"-m",
|
||||
@@ -439,7 +438,6 @@ def parse_args():
|
||||
"vllm_omni/deploy/mimo_audio.yaml based on the HF model_type.",
|
||||
)
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -2,12 +2,13 @@
|
||||
|
||||
[Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with text and speech outputs.
|
||||
|
||||
vLLM-Omni supports two deployment modes:
|
||||
vLLM-Omni supports three deployment modes:
|
||||
|
||||
| Mode | Deploy config | Output |
|
||||
|------|--------------|--------|
|
||||
| Thinker + Talker (omni-speech, default) | `vllm_omni/deploy/ming_flash_omni.yaml` | Text + Audio |
|
||||
| Thinker only (multimodal understanding) | `vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
|
||||
| Thinker + Imagegen (text-to-image / img2img) | `vllm_omni/deploy/ming_flash_omni_image.yaml` | Image (online-serving only at the moment) |
|
||||
|
||||
For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
|
||||
|
||||
@@ -52,7 +53,7 @@ python examples/offline_inference/ming_flash_omni/end2end.py \
|
||||
--image-path ./3_0.png
|
||||
```
|
||||
|
||||
### Omni-Speech (Thinker + Talker)
|
||||
### Omni-Speech (Thinker + Talker) {#omni-speech-thinker--talker}
|
||||
|
||||
The default deploy YAML already runs thinker+talker, so spoken output only requires requesting `audio` (or `text,audio`) modalities.
|
||||
The thinker processes your multimodal input, generates text, then the talker synthesises the response as speech.
|
||||
@@ -87,6 +88,17 @@ The default deploy YAML allocates thinker on GPUs 0–3 and talker on GPU 3 for
|
||||
|
||||
Pass `--deploy-config /path/to/your_deploy.yaml` to any of the commands above to override the bundled deploy config.
|
||||
|
||||
### Image generation (text-to-image / img2img)
|
||||
|
||||
The diffusion-stage image-generation path is currently only wired through
|
||||
the online OpenAI-compatible chat endpoint (`/v1/chat/completions` with
|
||||
`"modalities": ["image"]`). For payloads, optional knobs
|
||||
(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`), and the img2img
|
||||
reference-image flow, see the
|
||||
[image-gen section in the online-serving README](../../online_serving/ming_flash_omni/README.md#image-generation-text-to-image).
|
||||
`end2end.py` does not yet exercise the imagegen stage; that is tracked as
|
||||
follow-up work.
|
||||
|
||||
## Online serving
|
||||
|
||||
For online serving via the OpenAI-compatible API, see [examples/online_serving/ming_flash_omni/README.md](../../online_serving/ming_flash_omni/README.md).
|
||||
@@ -17,14 +17,13 @@ from vllm.assets.image import ImageAsset
|
||||
from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
import vllm_omni
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
# Imports the processor also registers itself
|
||||
from vllm_omni.transformers_utils.processors.ming import MingFlashOmniProcessor # noqa: F401
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
SEED = 42
|
||||
MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0"
|
||||
@@ -397,7 +396,7 @@ def main(args):
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
|
||||
parser = TrackingArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
@@ -496,7 +495,6 @@ def parse_args():
|
||||
help="Output directory for results.",
|
||||
)
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -19,10 +19,9 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
SEED = 42
|
||||
|
||||
@@ -432,7 +431,7 @@ def main(args):
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
@@ -554,7 +553,6 @@ def parse_args():
|
||||
default=False,
|
||||
help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
|
||||
)
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -19,10 +19,9 @@ from vllm.assets.image import ImageAsset
|
||||
from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
SEED = 42
|
||||
|
||||
@@ -422,7 +421,7 @@ def main(args):
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
@@ -561,7 +560,6 @@ def parse_args():
|
||||
help="Model dtype (auto, half, float16, bfloat16, float, float32).",
|
||||
)
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -39,10 +39,9 @@ from vllm.assets.image import ImageAsset
|
||||
from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -383,7 +382,9 @@ async def run_all(args):
|
||||
print(f"[Info] Creating AsyncOmni with deploy_config={args.deploy_config}")
|
||||
async_omni = None
|
||||
try:
|
||||
async_omni = AsyncOmni(**vars(args))
|
||||
# ``from_cli_args`` forwards only explicitly-passed CLI args so
|
||||
# argparse defaults do not silently override deploy YAML values.
|
||||
async_omni = AsyncOmni.from_cli_args(args, model=args.model)
|
||||
|
||||
# Use default sampling params from stage config (they are pre-configured
|
||||
# in the YAML for each stage).
|
||||
@@ -453,7 +454,7 @@ async def run_all(args):
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(
|
||||
parser = TrackingArgumentParser(
|
||||
description=(
|
||||
"Offline inference with async_chunk enabled via AsyncOmni. "
|
||||
"Downstream stages start before upstream stages finish, "
|
||||
@@ -589,7 +590,6 @@ def parse_args():
|
||||
default=16000,
|
||||
help="Sampling rate for audio loading.",
|
||||
)
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -37,6 +37,10 @@ python end2end.py --modality text2text \
|
||||
# Custom resolution (image generation)
|
||||
python end2end.py --prompt "A futuristic cityscape at sunset" \
|
||||
--width 2048 --height 1024 --think
|
||||
|
||||
# Text-to-image with Cache-DiT acceleration
|
||||
python end2end.py --prompt "A cute cat" \
|
||||
--cache-backend cache_dit \
|
||||
```
|
||||
|
||||
> **Note**: Default configuration works on a single **NVIDIA A100 (80GB)** or **H100** GPU.
|
||||
@@ -135,6 +139,8 @@ python end2end.py \
|
||||
| `--tensor-parallel-size` | 1 | Number of GPUs for tensor parallelism |
|
||||
| `--enforce-eager` | False | Disable torch.compile |
|
||||
| `--enable-cpu-offload` | False | Enable module-wise (sequential) CPU offload to reduce peak VRAM |
|
||||
| `--cache-backend` | None | Set to `cache_dit` for Cache-DiT acceleration |
|
||||
| `--enable-cache-dit-summary` | False | Print Cache-DiT cache statistics after generation |
|
||||
|
||||
## Reducing GPU Memory Usage
|
||||
|
||||
|
||||
@@ -29,7 +29,6 @@ import os
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
@@ -155,8 +154,19 @@ def parse_args():
|
||||
"of extra CPU<->GPU transfers."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cache-backend",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["cache_dit"],
|
||||
help="Cache backend to use for image generation acceleration.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-cache-dit-summary",
|
||||
action="store_true",
|
||||
help="Enable cache-dit summary logging after diffusion forward passes.",
|
||||
)
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -182,6 +192,8 @@ def main():
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
enforce_eager=args.enforce_eager,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
cache_backend=args.cache_backend,
|
||||
enable_cache_dit_summary=args.enable_cache_dit_summary,
|
||||
)
|
||||
|
||||
extra_args = {
|
||||
@@ -225,6 +237,7 @@ def main():
|
||||
print(f" Seed : {args.seed}")
|
||||
print(f" Think mode : {args.think}")
|
||||
print(f" TP size : {args.tensor_parallel_size}")
|
||||
print(f" Cache backend : {args.cache_backend or 'none'}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
# Build prompt dict
|
||||
|
||||
@@ -16,9 +16,10 @@ Usage:
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
@@ -29,6 +30,16 @@ from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def parse_profiler_config(value: str) -> dict[str, Any]:
|
||||
try:
|
||||
config = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
|
||||
if not isinstance(config, dict):
|
||||
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
|
||||
return config
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate a talking-head video from a reference image and audio (Wan2.2 S2V)."
|
||||
@@ -87,6 +98,12 @@ def parse_args() -> argparse.Namespace:
|
||||
default=3.0,
|
||||
help="Scheduler flow shift (default: 3.0).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--boundary-ratio",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Boundary split ratio for low/high DiT stages (S2V). Default varies by model.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
@@ -178,30 +195,10 @@ def parse_args() -> argparse.Namespace:
|
||||
help="Enable cache-dit summary logging after diffusion forward passes.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-dir",
|
||||
type=str,
|
||||
"--profiler-config",
|
||||
type=parse_profiler_config,
|
||||
default=None,
|
||||
help="Enable torch profiler and save traces to this directory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-record-shapes",
|
||||
action="store_true",
|
||||
help="Record tensor shapes in profiler (increases trace size).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-with-stack",
|
||||
action="store_true",
|
||||
help="Record stack traces in profiler (increases overhead).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-with-memory",
|
||||
action="store_true",
|
||||
help="Profile memory usage.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-with-flops",
|
||||
action="store_true",
|
||||
help="Estimate FLOPs for operations.",
|
||||
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
@@ -239,23 +236,9 @@ def main():
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
)
|
||||
|
||||
# Check if profiling is requested via CLI args or environment variable
|
||||
profile_dir = args.profile_dir or os.getenv("VLLM_TORCH_PROFILER_DIR")
|
||||
profiler_enabled = bool(profile_dir)
|
||||
profiler_config = None
|
||||
if profiler_enabled:
|
||||
from vllm.config import ProfilerConfig
|
||||
profiler_enabled = args.profiler_config is not None
|
||||
|
||||
profiler_config = ProfilerConfig(
|
||||
profiler="torch",
|
||||
torch_profiler_dir=profile_dir,
|
||||
torch_profiler_record_shapes=args.profile_record_shapes,
|
||||
torch_profiler_with_stack=args.profile_with_stack,
|
||||
torch_profiler_with_memory=args.profile_with_memory,
|
||||
torch_profiler_with_flops=args.profile_with_flops,
|
||||
)
|
||||
|
||||
omni = Omni(
|
||||
omni_kwargs = dict(
|
||||
model=args.model,
|
||||
model_class_name="WanS2VPipeline",
|
||||
flow_shift=args.flow_shift,
|
||||
@@ -269,8 +252,12 @@ def main():
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
enable_cache_dit_summary=args.enable_cache_dit_summary,
|
||||
profiler_config=profiler_config,
|
||||
profiler_config=args.profiler_config,
|
||||
)
|
||||
if args.boundary_ratio is not None:
|
||||
omni_kwargs["boundary_ratio"] = args.boundary_ratio
|
||||
|
||||
omni = Omni(**omni_kwargs)
|
||||
|
||||
# Print generation configuration
|
||||
print(f"\n{'=' * 60}")
|
||||
@@ -325,13 +312,16 @@ def main():
|
||||
if profiler_enabled:
|
||||
print("\n[Profiler] Stopping profiler and collecting results...")
|
||||
profile_results = omni.stop_profile()
|
||||
if profile_results and isinstance(profile_results, list):
|
||||
if profile_results and isinstance(profile_results, dict):
|
||||
traces = profile_results.get("traces", [])
|
||||
print("\n" + "=" * 60)
|
||||
print("PROFILING RESULTS:")
|
||||
for i, result in enumerate(profile_results):
|
||||
print(f"\nStage {i}:")
|
||||
if result:
|
||||
print(f" {result}")
|
||||
for rank, trace in enumerate(traces):
|
||||
print(f"\nRank {rank}:")
|
||||
if trace:
|
||||
print(f" • Trace: {trace}")
|
||||
if not traces:
|
||||
print(" No traces collected.")
|
||||
print("=" * 60)
|
||||
else:
|
||||
print("[Profiler] No valid profiling data returned.")
|
||||
|
||||
@@ -35,6 +35,7 @@ This folder provides several entrypoints for experimenting with text-to-image di
|
||||
| `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 |
|
||||
| `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
|
||||
| `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3) | 160 |
|
||||
| `HiDream-I1-Full` | 1024 x 1024 | 63.7 | 57.7 |
|
||||
|
||||
!!! info
|
||||
*Peak VRAM: based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU.
|
||||
@@ -94,6 +95,7 @@ python text_to_image.py \
|
||||
| `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
|
||||
| `--use-system-prompt` | str | `None` | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text. Recommended: `en_unified`. Only for HunyuanImage-3.0.|
|
||||
| `--system-prompt` | str | `None` | Custom system prompt text. Only used when `--use-system-prompt` is set to `custom`. Only for HunyuanImage-3.0.|
|
||||
| `--auxiliary-text-encoder` | str | `None` | Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full). |
|
||||
|
||||
**NextStep-1.1 specific arguments:**
|
||||
|
||||
@@ -177,6 +179,24 @@ python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--output flux2-dev.png
|
||||
```
|
||||
|
||||
|
||||
### HiDream-I1-Full Models
|
||||
|
||||
The `--auxiliary-text-encoder` parameter is required when running HiDream‑I1‑Full:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model HiDream-ai/HiDream-I1-Full \
|
||||
--prompt "The setting sun of late autumn dyes the riverside with a warm orange hue" \
|
||||
--seed 42 \
|
||||
--guidance-scale 5.0 \
|
||||
--tensor-parallel-size 1 \
|
||||
--num-images-per-prompt 1 \
|
||||
--num-inference-steps 50 \
|
||||
--auxiliary-text-encoder meta-llama/Meta-Llama-3.1-8B-Instruct \
|
||||
--output /output.png
|
||||
```
|
||||
|
||||
### Batch Requests (Multiple Prompts)
|
||||
|
||||
You can pass multiple prompts in a single `generate` call.
|
||||
|
||||
@@ -300,10 +300,13 @@ def parse_args() -> argparse.Namespace:
|
||||
default=None,
|
||||
help=("Custom system prompt. Used when --use-system-prompt is custom. "),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auxiliary-text-encoder",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
|
||||
)
|
||||
current_omni_platform.pre_register_and_update(parser)
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -392,6 +395,7 @@ def main():
|
||||
"profiler_config": args.profiler_config,
|
||||
"init_timeout": args.init_timeout,
|
||||
"stage_init_timeout": args.stage_init_timeout,
|
||||
"auxiliary_text_encoder": args.auxiliary_text_encoder,
|
||||
**lora_args,
|
||||
**quant_kwargs,
|
||||
}
|
||||
|
||||
@@ -14,13 +14,13 @@ list of supported architectures across all modalities, see
|
||||
|
||||
| Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
|
||||
|---|---|---|---|---|---|---|
|
||||
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 22.05 kHz |
|
||||
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
|
||||
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
|
||||
| MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | single (AR + codec) | ✓ (required) | ✓ | voice_clone, continuation | 48 kHz |
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
|
||||
| VoxCPM | `openbmb/VoxCPM-0.5B` | split | ✓ | ✓ | — | 24 kHz |
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
|
||||
|
||||
@@ -41,7 +41,7 @@ python examples/offline_inference/text_to_speech/<model>/end2end.py \
|
||||
|
||||
## CosyVoice3
|
||||
|
||||
2-stage TTS pipeline (`talker` + `code2wav`) at 22.05 kHz.
|
||||
2-stage TTS pipeline (`talker` + `code2wav`) at 24 kHz.
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
@@ -93,6 +93,35 @@ Streaming is enabled by default via `async_chunk: true` in `vllm_omni/deploy/cos
|
||||
|
||||
---
|
||||
|
||||
## GLM-TTS
|
||||
|
||||
2-stage TTS pipeline (AR + DiT flow-matching) at 24 kHz. Every request requires reference audio and its transcript for zero-shot voice cloning.
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
|
||||
--model zai-org/GLM-TTS \
|
||||
--text "你好,这是语音合成测试。" \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "这是参考音频的文本内容。" \
|
||||
--output-dir ./output
|
||||
```
|
||||
|
||||
### Architecture
|
||||
```
|
||||
Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz)
|
||||
(Llama-based) (32k vocab) (Flow Matching)
|
||||
```
|
||||
|
||||
### Notes
|
||||
- `--ref-audio` and `--ref-text` are **required** together; GLM-TTS does not support text-only synthesis.
|
||||
- Reference audio should be 3-10 seconds.
|
||||
- First run may be slow due to lazy loading of WhisperVQ tokenizer and CampPlus ONNX speaker embedder.
|
||||
- Default sampling: temperature=1.0, top_k=25, top_p=0.8 (RAS method).
|
||||
- The `--model` path should point to the repository root (not `llm/` subdirectory).
|
||||
|
||||
---
|
||||
|
||||
## Fish Speech S2 Pro
|
||||
|
||||
4B dual-AR text-to-speech model from FishAudio with the DAC codec at 44.1 kHz.
|
||||
@@ -242,6 +271,14 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--lang zh
|
||||
```
|
||||
|
||||
### Seed for Reproducibility
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
--text "Hello, this is a test." \
|
||||
--seed 42
|
||||
```
|
||||
|
||||
### Notes
|
||||
- Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking.
|
||||
- Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz.
|
||||
@@ -308,56 +345,6 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
|
||||
---
|
||||
|
||||
## VoxCPM
|
||||
|
||||
Split-stage TTS. The hub example covers single TTS, single voice cloning, and streaming. Use `benchmarks/voxcpm/` for warmup, batch JSONL prompts, profiler injection, and offline TTFP / RTF measurement.
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
pip install voxcpm soundfile
|
||||
# or use a local source tree:
|
||||
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
|
||||
```
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
|
||||
--model openbmb/VoxCPM-0.5B \
|
||||
--text "This is a split-stage VoxCPM synthesis example running on vLLM Omni."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
|
||||
--model openbmb/VoxCPM-0.5B \
|
||||
--text "This sentence is synthesized with a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "The exact transcript spoken in reference.wav."
|
||||
```
|
||||
|
||||
### Streaming
|
||||
Pass `--streaming` together with the legacy async-chunk stage config:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
|
||||
--model openbmb/VoxCPM-0.5B \
|
||||
--streaming \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
|
||||
--text "This is a split-stage VoxCPM streaming example running on vLLM Omni."
|
||||
```
|
||||
|
||||
### Persistent HF config (optional)
|
||||
The engine auto-patches a missing `model_type` in the checkpoint's `config.json` in-memory at startup (see `vllm_omni/engine/arg_utils.py`), so no manual setup is required.
|
||||
|
||||
If you'd rather maintain a persistent patched config dir (e.g. to share across runs or to add fields the auto-patcher doesn't set), export `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH` to its location — both the bundled deploy and async-chunk yamls read it via `${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,}`.
|
||||
|
||||
### Notes
|
||||
- Non-streaming auto-loads `vllm_omni/deploy/voxcpm.yaml` from the model directory name / HF `model_type`. Streaming uses the legacy `voxcpm_async_chunk.yaml` passed via `--stage-configs-path`.
|
||||
- Streaming is currently single-request oriented.
|
||||
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades quality.
|
||||
- For online serving, see the [VoxCPM section in the online hub](../../online_serving/text_to_speech/README.md#voxcpm). For benchmark reporting, see [`benchmarks/voxcpm`](../../../benchmarks/voxcpm/README.md).
|
||||
|
||||
---
|
||||
|
||||
## VoxCPM2
|
||||
|
||||
Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`.
|
||||
|
||||
@@ -10,7 +10,6 @@ import soundfile as sf
|
||||
from vllm import SamplingParams
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
|
||||
from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
|
||||
@@ -65,7 +64,6 @@ def run_e2e():
|
||||
required=True,
|
||||
help="Path to tokenizer directory (e.g., <model_path>/CosyVoice-BlankEN).",
|
||||
)
|
||||
nullify_stage_engine_defaults(parser)
|
||||
args = parser.parse_args()
|
||||
# Ensure tokenizer directory exists
|
||||
if not os.path.exists(args.tokenizer):
|
||||
|
||||
@@ -26,16 +26,14 @@ import torch
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import AsyncOmni, Omni
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.model_executor.models.fish_speech.dac_utils import DAC_HOP_LENGTH, DAC_SAMPLE_RATE
|
||||
from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
|
||||
build_fish_text_only_prompt_ids,
|
||||
estimate_fish_voice_clone_prompt_len_from_normalized,
|
||||
normalize_fish_voice_clone_texts,
|
||||
)
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -211,7 +209,7 @@ async def main_streaming(args):
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(
|
||||
parser = TrackingArgumentParser(
|
||||
description="Fish Speech S2 Pro offline inference with vLLM Omni",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -267,7 +265,6 @@ def parse_args():
|
||||
default=False,
|
||||
help="Stream audio chunks as they arrive via AsyncOmni.",
|
||||
)
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,194 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""GLM-TTS End-to-End Offline Inference Example.
|
||||
|
||||
GLM-TTS is a two-stage TTS system:
|
||||
- Stage 0 (AR): Llama-based model generates speech tokens from text
|
||||
- Stage 1 (DiT): Flow matching model converts speech tokens to audio
|
||||
|
||||
Usage:
|
||||
# Sync two-stage (default)
|
||||
python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
|
||||
--model /path/to/GLM-TTS \
|
||||
--text "你好,这是一个语音合成测试。" \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "参考音频的转录文本。" \
|
||||
--output-dir ./output
|
||||
|
||||
# Async chunk mode (streaming DiT)
|
||||
python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
|
||||
--model /path/to/GLM-TTS --async-chunk \
|
||||
--text "你好,这是一个语音合成测试。" \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "参考音频的转录文本。" \
|
||||
--output-dir ./output
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
from typing import Any
|
||||
from urllib.request import urlopen
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.model_executor.models.glm_tts.glm_tts import build_glm_tts_prefill_metadata
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_DEPLOY_CONFIG = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
"..",
|
||||
"..",
|
||||
"..",
|
||||
"..",
|
||||
"vllm_omni",
|
||||
"deploy",
|
||||
"glm_tts.yaml",
|
||||
)
|
||||
SAMPLE_RATE = 24000
|
||||
|
||||
|
||||
def _load_ref_audio(ref_audio: str) -> tuple[torch.Tensor, int]:
|
||||
"""Load reference audio from file path, URL, or data URI."""
|
||||
if ref_audio.startswith(("http://", "https://")):
|
||||
with urlopen(ref_audio, timeout=60) as response:
|
||||
audio_obj: Any = io.BytesIO(response.read())
|
||||
elif ref_audio.startswith("data:"):
|
||||
_, _, encoded = ref_audio.partition(",")
|
||||
audio_obj = io.BytesIO(base64.b64decode(encoded))
|
||||
else:
|
||||
audio_obj = ref_audio
|
||||
wav_np, sr = sf.read(audio_obj, dtype="float32")
|
||||
if wav_np.ndim > 1:
|
||||
wav_np = wav_np.mean(axis=1)
|
||||
return torch.from_numpy(wav_np), int(sr)
|
||||
|
||||
|
||||
def _concat_audio(audio_val: Any) -> np.ndarray:
|
||||
"""Concatenate audio tensors from multimodal output."""
|
||||
if isinstance(audio_val, list):
|
||||
tensors = [torch.as_tensor(t).float().reshape(-1) for t in audio_val if t is not None]
|
||||
if not tensors:
|
||||
return np.zeros((0,), dtype=np.float32)
|
||||
return torch.cat(tensors, dim=-1).cpu().numpy().astype(np.float32, copy=False)
|
||||
if isinstance(audio_val, torch.Tensor):
|
||||
return audio_val.float().cpu().numpy().reshape(-1)
|
||||
return np.asarray(audio_val, dtype=np.float32).reshape(-1)
|
||||
|
||||
|
||||
def _extract_sample_rate(audio_mm: dict) -> int:
|
||||
"""Extract sample rate from multimodal output dict."""
|
||||
sr_raw = audio_mm.get("sr", SAMPLE_RATE)
|
||||
if isinstance(sr_raw, list):
|
||||
sr_raw = sr_raw[-1] if sr_raw else SAMPLE_RATE
|
||||
if hasattr(sr_raw, "item"):
|
||||
return int(sr_raw.item())
|
||||
return int(sr_raw)
|
||||
|
||||
|
||||
def _modify_deploy_config(base_path: str, async_chunk: bool) -> str:
|
||||
"""Build deploy config with explicit sync/async mode and eager execution.
|
||||
|
||||
Mirrors the logic in ``tests/e2e/offline_inference/test_glm_tts.py``
|
||||
(``_get_deploy_config``) so that example runs match CI behavior.
|
||||
"""
|
||||
with open(base_path) as f:
|
||||
cfg = yaml.safe_load(f)
|
||||
cfg["async_chunk"] = async_chunk
|
||||
for stage in cfg.get("stages", []):
|
||||
stage["enforce_eager"] = True
|
||||
if stage.get("stage_id") == 0:
|
||||
stage["async_scheduling"] = bool(async_chunk)
|
||||
tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False, prefix="glm_tts_")
|
||||
yaml.dump(cfg, tmp)
|
||||
tmp.close()
|
||||
return tmp.name
|
||||
|
||||
|
||||
def main(args):
|
||||
"""Run offline GLM-TTS inference."""
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
base_deploy_config = args.deploy_config or DEFAULT_DEPLOY_CONFIG
|
||||
deploy_config_path = _modify_deploy_config(base_deploy_config, args.async_chunk)
|
||||
|
||||
ref_audio_wav, ref_audio_sr = _load_ref_audio(args.ref_audio)
|
||||
if not args.ref_text:
|
||||
raise ValueError("GLM-TTS requires --ref-audio and --ref-text.")
|
||||
|
||||
inputs = [
|
||||
{
|
||||
"prompt": args.text,
|
||||
"multi_modal_data": {
|
||||
"audio": (ref_audio_wav.float().cpu().numpy(), ref_audio_sr),
|
||||
},
|
||||
"modalities": ["audio"],
|
||||
"mm_processor_kwargs": {"prompt_text": args.ref_text},
|
||||
"additional_information": build_glm_tts_prefill_metadata(
|
||||
args.model,
|
||||
args.text,
|
||||
args.ref_text,
|
||||
),
|
||||
}
|
||||
]
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
stage_configs_path=deploy_config_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
t_start = time.perf_counter()
|
||||
outputs = omni.generate(inputs)
|
||||
elapsed = (time.perf_counter() - t_start) * 1000
|
||||
|
||||
assert outputs, "No outputs returned"
|
||||
audio_mm = outputs[0].multimodal_output
|
||||
assert "audio" in audio_mm, "No audio output found"
|
||||
|
||||
audio = _concat_audio(audio_mm["audio"])
|
||||
sr = _extract_sample_rate(audio_mm)
|
||||
out_path = os.path.join(args.output_dir, "output.wav")
|
||||
sf.write(out_path, audio, samplerate=sr, format="WAV")
|
||||
|
||||
logger.info("Saved %s (%.2fs @ %dHz)", out_path, len(audio) / sr, sr)
|
||||
logger.info("Total inference: %.1f ms", elapsed)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="GLM-TTS Text-to-Speech Example")
|
||||
parser.add_argument("--model", type=str, required=True, help="Model path")
|
||||
parser.add_argument("--text", type=str, default="你好,这是一个语音合成测试。")
|
||||
parser.add_argument("--output-dir", type=str, default="./output")
|
||||
parser.add_argument("--ref-audio", type=str, required=True, help="Reference WAV path/URL")
|
||||
parser.add_argument("--ref-text", type=str, required=True, help="Transcript of ref audio")
|
||||
parser.add_argument("--deploy-config", type=str, default=None)
|
||||
parser.add_argument(
|
||||
"--async-chunk",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable async_chunk mode (streaming DiT). Default: sync two-stage.",
|
||||
)
|
||||
parser.add_argument("--log-stats", action="store_true")
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=600)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
main(parse_args())
|
||||
@@ -0,0 +1,38 @@
|
||||
# higgs-audio v2 — offline example
|
||||
|
||||
Drives Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) for `bosonai/higgs-audio-v2-generation-3B-base` end-to-end through the vLLM-Omni engine and writes a 24 kHz mono WAV per prompt.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Voice clone needs `transformers>=5.3.0` — vllm-omni loads the audio codec via HF's `HiggsAudioV2TokenizerModel`, instantiated from the `k2-fsa/OmniVoice/audio_tokenizer/` subdirectory (only that ~806 MB subdir is downloaded). The boson-ai standalone tokenizer repo's `model.safetensors` is actually a copy of the 3B talker LM, so HF can't load it directly; the k2 bundle ships the same codec weights repackaged with HF-compatible key naming.
|
||||
|
||||
```bash
|
||||
pip install -U "transformers>=5.3.0"
|
||||
```
|
||||
|
||||
## Quick start
|
||||
|
||||
Plain TTS:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
|
||||
--texts "Hello world." "The quick brown fox jumps over the lazy dog." \
|
||||
--output-dir results/wavs
|
||||
```
|
||||
|
||||
## Voice cloning
|
||||
|
||||
Pass both `--ref-audio` and `--ref-text` together:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
|
||||
--texts "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Exact transcript spoken in reference.wav." \
|
||||
--output-dir results/wavs
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
- Output: 24 kHz mono WAV.
|
||||
- Deploy config: `vllm_omni/deploy/higgs_audio_v2.yaml` (auto-loaded by HF `model_type`).
|
||||
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
|
||||
- For online serving, see [`examples/online_serving/text_to_speech/higgs_audio_v2/`](../../../online_serving/text_to_speech/higgs_audio_v2/).
|
||||
@@ -0,0 +1,203 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline higgs-audio v2 inference example.
|
||||
|
||||
Runs Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) end-to-end
|
||||
through the vLLM-Omni engine without going through the HTTP server, and
|
||||
saves a 24 kHz mono WAV per prompt.
|
||||
|
||||
Example:
|
||||
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \\
|
||||
--texts "Hello world." \\
|
||||
"The quick brown fox jumps over the lazy dog." \\
|
||||
--output-dir results/wavs
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
# DeepGEMM FP8 kernels require an optional backend that may not be installed.
|
||||
# Disable the warmup before importing vLLM so engine startup falls back to the
|
||||
# regular gemm path. Users with deep_gemm installed can override these.
|
||||
os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
|
||||
os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
SAMPLE_RATE = 24_000
|
||||
DEFAULT_TEXTS = (
|
||||
"Hello world.",
|
||||
"The quick brown fox jumps over the lazy dog.",
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="Offline higgs-audio v2 inference")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="bosonai/higgs-audio-v2-generation-3B-base",
|
||||
help="Stage-0 talker model id or path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--texts",
|
||||
type=str,
|
||||
nargs="+",
|
||||
default=list(DEFAULT_TEXTS),
|
||||
help="One or more plain-text prompts to synthesize.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default="results/wavs",
|
||||
help="Directory to write per-prompt WAV files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override the deploy config path. Auto-loads "
|
||||
"vllm_omni/deploy/higgs_audio_v2.yaml from the HF model_type by default.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-new-tokens",
|
||||
type=int,
|
||||
default=500,
|
||||
help="Cap on Stage-0 codec frames per request.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Reference clip for voice clone (path to a WAV file). Paired with --ref-text.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-text",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Transcript of the reference clip. Required when --ref-audio is set.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _slugify(text: str) -> str:
|
||||
import re
|
||||
|
||||
slug = re.sub(r"\s+", "_", text.strip().lower())
|
||||
slug = re.sub(r"[^a-z0-9_]+", "", slug)
|
||||
return slug[:48] or "prompt"
|
||||
|
||||
|
||||
def _extract_pcm(multimodal_output: dict) -> torch.Tensor:
|
||||
"""Pull the final concatenated PCM tensor out of a request's multimodal_output."""
|
||||
audio = multimodal_output.get("model_outputs")
|
||||
if audio is None:
|
||||
audio = multimodal_output.get("audio")
|
||||
if audio is None:
|
||||
raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
|
||||
if isinstance(audio, list):
|
||||
valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
|
||||
if not valid:
|
||||
raise ValueError("audio list is empty")
|
||||
return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
|
||||
return torch.as_tensor(audio).float().cpu().reshape(-1)
|
||||
|
||||
|
||||
def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray:
|
||||
arr = pcm.numpy()
|
||||
if arr.dtype.kind == "f":
|
||||
arr = np.clip(arr, -1.0, 1.0)
|
||||
arr = (arr * 32767.0).astype(np.int16)
|
||||
else:
|
||||
arr = arr.astype(np.int16)
|
||||
return arr
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
if (args.ref_audio is None) != (args.ref_text is None):
|
||||
raise SystemExit("--ref-audio and --ref-text must be supplied together")
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
engine = Omni(model=args.model, deploy_config=args.deploy_config)
|
||||
|
||||
# Build prompt_token_ids using the same path serving_speech.py takes online.
|
||||
from transformers import AutoProcessor
|
||||
|
||||
from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
|
||||
build_plain_text_prompt,
|
||||
build_voice_clone_prompt,
|
||||
input_ids_to_python_list,
|
||||
)
|
||||
|
||||
processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
|
||||
|
||||
# Voice-clone path: load the reference clip once. The HF processor will
|
||||
# encode it via the bundled HiggsAudioV2TokenizerModel each time we build
|
||||
# a prompt below. This is cheap on CPU for a few-second clip.
|
||||
ref_wav: np.ndarray | None = None
|
||||
ref_sr: int | None = None
|
||||
if args.ref_audio is not None:
|
||||
ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
|
||||
if ref_wav.ndim == 2:
|
||||
ref_wav = ref_wav.mean(axis=1)
|
||||
|
||||
print(f"Model : {args.model}")
|
||||
print(f"Prompts : {len(args.texts)}")
|
||||
print(f"Output dir : {output_dir}")
|
||||
print(f"Voice clone : {'yes' if ref_wav is not None else 'no'}")
|
||||
|
||||
# Run one prompt at a time. The Stage-0 talker's per-slot audio state is
|
||||
# request-scoped; submitting multiple prompts in the same engine.generate()
|
||||
# call would batch them in the AR runner and exercise a code path that is
|
||||
# not validated for this model yet.
|
||||
total_elapsed = 0.0
|
||||
total_dur = 0.0
|
||||
for text in args.texts:
|
||||
if ref_wav is not None:
|
||||
out = build_voice_clone_prompt(processor, text, ref_wav, int(ref_sr), args.ref_text)
|
||||
prompt = {
|
||||
"prompt_token_ids": out["prompt_token_ids"],
|
||||
# Bare tensors (NOT list-wrapped): the msgspec serializer in
|
||||
# vllm_omni.data_entry_keys routes torch.Tensor → tensor_data
|
||||
# and list[Tensor] → list_data (which silently strips tensors).
|
||||
"additional_information": {
|
||||
"audio_input_ids": out["audio_input_ids"],
|
||||
"audio_input_ids_mask": out["audio_input_ids_mask"],
|
||||
},
|
||||
}
|
||||
else:
|
||||
inputs = build_plain_text_prompt(processor, text)
|
||||
prompt = {"prompt_token_ids": input_ids_to_python_list(inputs)}
|
||||
t_start = time.perf_counter()
|
||||
outputs = engine.generate([prompt])
|
||||
elapsed = time.perf_counter() - t_start
|
||||
total_elapsed += elapsed
|
||||
|
||||
mm = outputs[0].outputs[0].multimodal_output
|
||||
pcm = _extract_pcm(mm)
|
||||
slug = _slugify(text)
|
||||
out_path = output_dir / f"{slug}.wav"
|
||||
sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16")
|
||||
dur = pcm.numel() / SAMPLE_RATE
|
||||
total_dur += dur
|
||||
print(f" {slug:<50} dur={dur:5.2f}s -> {out_path}")
|
||||
|
||||
rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
|
||||
print(f"Total infer : {total_elapsed:.2f}s total audio: {total_dur:.2f}s RTF: {rtf:.3f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
main()
|
||||
@@ -5,11 +5,11 @@ from typing import Any
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniTokensPrompt
|
||||
from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
|
||||
@@ -68,7 +68,7 @@ def save_audio(mm: dict[str, Any], output_path: str) -> None:
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
|
||||
parser = TrackingArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
|
||||
parser.add_argument("--model", type=str, default=MODEL_NAME, help="Model name or local path.")
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
@@ -89,7 +89,6 @@ def parse_args():
|
||||
parser.add_argument("--init-timeout", type=int, default=600, help="Engine init timeout in seconds.")
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=300, help="Single stage init timeout in seconds.")
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
Loaded 100 of 755 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user