Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
202961e56c |
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.22.0"
|
||||
VLLM_VERSION: "v0.21.0"
|
||||
priority: 100
|
||||
timeout_in_minutes: 60
|
||||
soft_fail: false
|
||||
@@ -52,7 +52,6 @@ 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,28 +1,15 @@
|
||||
steps:
|
||||
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
- label: "Simple Unit 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: "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"
|
||||
- "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Diffusion Images API LoRA E2E"
|
||||
agent_pool: mi325_1
|
||||
@@ -62,6 +49,16 @@ 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
|
||||
@@ -126,7 +123,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 -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
|
||||
|
||||
- label: "Qwen3-TTS CustomVoice E2E Test"
|
||||
agent_pool: mi325_1
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
steps:
|
||||
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
- label: "Simple Unit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
@@ -9,18 +9,7 @@ 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 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"
|
||||
- "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Voxtral TTS CUDA Unit Test"
|
||||
agent_pool: mi325_1
|
||||
@@ -30,13 +19,13 @@ steps:
|
||||
commands:
|
||||
- timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
|
||||
|
||||
- label: "Diffusion · Batch Test"
|
||||
- label: "Diffusion Batching Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
|
||||
- timeout 20m pytest -s -v tests/diffusion/batching/test_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
|
||||
|
||||
- label: "Custom Pipeline Test"
|
||||
agent_pool: mi325_1
|
||||
|
||||
@@ -107,7 +107,7 @@ steps:
|
||||
depends_on: upload-merge-pipeline
|
||||
steps:
|
||||
- label: "Entrypoint Test with H100"
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 35
|
||||
commands:
|
||||
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
@@ -389,54 +389,6 @@ 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:
|
||||
|
||||
+21
-297
@@ -305,51 +305,6 @@ 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
|
||||
@@ -394,11 +349,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 with L4"
|
||||
- label: ":full_moon: TTS · Function Test"
|
||||
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 --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
|
||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
@@ -413,69 +368,6 @@ 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
|
||||
@@ -522,52 +414,6 @@ 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
|
||||
@@ -679,7 +525,6 @@ 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:
|
||||
@@ -733,88 +578,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: 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
|
||||
nvidia.com/gpu: 3
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
@@ -937,46 +701,6 @@ 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
|
||||
@@ -1073,8 +797,8 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2509"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2509
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
@@ -1082,7 +806,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_2511_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_2509_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"
|
||||
@@ -1169,20 +893,26 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · BAGEL"
|
||||
key: nightly-diffusion-x2iat-performance-bagel
|
||||
timeout_in_minutes: 180
|
||||
- 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"
|
||||
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_bagel_vllm_omni.json
|
||||
EXIT=$$?
|
||||
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=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
exit $$((EXIT1 | EXIT2 | EXIT3))
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -1192,7 +922,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: 1
|
||||
nvidia.com/gpu: 4
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
@@ -1268,8 +998,6 @@ 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:
|
||||
@@ -1460,13 +1188,11 @@ 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-2511
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit-2509
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
- nightly-diffusion-x2iat-performance-bagel
|
||||
- nightly-diffusion-x2v-performance
|
||||
- nightly-testcase-statistics
|
||||
if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
|
||||
@@ -1478,12 +1204,10 @@ 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-2511
|
||||
- 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-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 20m bash -c "
|
||||
timeout 15m bash -c "
|
||||
pytest -sv tests/distributed/ -m 'core_model and cuda' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
@@ -273,42 +273,6 @@ 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:
|
||||
@@ -332,29 +296,6 @@ 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:
|
||||
- |
|
||||
@@ -399,29 +340,6 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · Higgs-Audio-V3 Test"
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
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_v3.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_4_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"
|
||||
commands:
|
||||
- |
|
||||
|
||||
@@ -53,9 +53,6 @@
|
||||
{% 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,9 +21,6 @@ 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
|
||||
|
||||
@@ -82,19 +79,7 @@ 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. **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.**
|
||||
6. **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.
|
||||
@@ -106,16 +91,15 @@ 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.
|
||||
|
||||
8. **Analyze host, communication, and operators.**
|
||||
7. **Analyze host, communication, and operators.**
|
||||
- Find GPU idle/free intervals and map each large gap to the enclosing CPU/Python code.
|
||||
- Separate real GPU idle from profiler overhead such as CUPTI `Command Buffer Full`.
|
||||
- Compare NCCL kernel time to user annotations; annotations can overcount nested intervals.
|
||||
- Rank operator work by total CUDA time and by repeated small-kernel launch count.
|
||||
|
||||
9. **Produce an optimization plan.**
|
||||
8. **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
|
||||
@@ -140,13 +124,7 @@ 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 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.
|
||||
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.
|
||||
|
||||
## Optimization Layers
|
||||
|
||||
|
||||
@@ -1,97 +0,0 @@
|
||||
---
|
||||
name: precheck-pr
|
||||
description: Self-check your branch before creating a PR — catch dead code, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub.
|
||||
---
|
||||
|
||||
# PR Pre-Check
|
||||
|
||||
Self-review your branch before creating a PR against `vllm-project/vllm-omni`. Two modes: **quick** catches showstoppers, **full** does a thorough maintainer-grade review. Never posts to GitHub; the report is for the contributor's terminal only.
|
||||
|
||||
## Mode Selection
|
||||
|
||||
| Mode | When | Time |
|
||||
|------|------|------|
|
||||
| **Quick** | About to push, final sanity check | ~3 min |
|
||||
| **Full** | Ready for review, want maintainer-level scan | ~10 min |
|
||||
|
||||
Default to quick if unsure. Run full before marking a PR "ready for review."
|
||||
|
||||
## Workflow
|
||||
|
||||
### Step 1: Detect Base Branch
|
||||
|
||||
```bash
|
||||
BASE_SHA=$(git merge-base HEAD origin/main 2>/dev/null \
|
||||
|| git merge-base HEAD main 2>/dev/null \
|
||||
|| echo origin/main)
|
||||
echo "diffing against ${BASE_SHA}"
|
||||
git diff --name-only ${BASE_SHA}...HEAD
|
||||
```
|
||||
|
||||
### Step 2: Validate PR Title
|
||||
|
||||
Check the most recent commit message (or branch name if no commit yet) against the project convention. Valid prefixes:
|
||||
|
||||
| Prefix | Applies to |
|
||||
|--------|-----------|
|
||||
| `[Bugfix]` | Bug fixes |
|
||||
| `[CI/Build]` | Build or CI improvements |
|
||||
| `[Doc]` | Documentation changes |
|
||||
| `[Model]` | New/improved models (include model name) |
|
||||
| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
|
||||
| `[Kernel]` | CUDA/kernel changes |
|
||||
| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
|
||||
| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
|
||||
| `[Misc]` | Other changes (use sparingly) |
|
||||
|
||||
✗ if: missing prefix, wrong case (`[bugfix]`), or WIP/Draft in title.
|
||||
⚠ if: `[Model]` prefix without the model identifier (e.g., `[Model] Add new model` — should be `[Model] Add <ModelName> ...`).
|
||||
|
||||
### Step 3: Categorize the PR
|
||||
|
||||
| Diff contains | PR type |
|
||||
|---------------|---------|
|
||||
| New files under `vllm_omni/model_executor/models/<name>/` | **New Model** |
|
||||
| Changes to `vllm_omni/diffusion/` | **Diffusion Model** |
|
||||
| `[Bugfix]` prefix or single-file fix | **Bug Fix** |
|
||||
| Perf/benchmark/throughput claims in commit msg or diff | **Performance** |
|
||||
| Everything else | **General** |
|
||||
|
||||
If multiple rows apply (e.g., a diffusion model is also a new model), union the checklists.
|
||||
|
||||
### Step 4: Run Checklist
|
||||
|
||||
Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type from [references/checklists.md](references/checklists.md). Each item produces ✓, ✗, or ⚠.
|
||||
|
||||
### Step 5: Print Report
|
||||
|
||||
```
|
||||
Pre-check report for <branch>
|
||||
|
||||
Mode: quick | full
|
||||
Type: <new-model | diffusion-model | bug-fix | perf | general>
|
||||
|
||||
Dimension Result
|
||||
───────────────── ──────
|
||||
PR title format ✓
|
||||
PR desc integrity ✓
|
||||
Registry/config ✓
|
||||
Dead code ⚠ 2 warnings
|
||||
Accuracy ✓
|
||||
Benchmark ✗ missing software versions
|
||||
|
||||
Verdict: 1 blocking | 2 warnings | recommend fixing ✗ before PR
|
||||
```
|
||||
|
||||
**Severity:**
|
||||
|
||||
| Mark | Meaning |
|
||||
|------|---------|
|
||||
| ✗ | Blocking — fix before opening PR |
|
||||
| ⚠ | Warning — consider fixing |
|
||||
| ✓ | Pass |
|
||||
| — | Skipped (not applicable) |
|
||||
|
||||
### Stop Here
|
||||
|
||||
Do not post comments, open PRs, or modify files. The report is for the contributor's terminal only.
|
||||
@@ -1,135 +0,0 @@
|
||||
# Pre-Check Checklists
|
||||
|
||||
Each item is a concrete check to run against the diff. ✗ = fix before PR. ⚠ = consider fixing.
|
||||
|
||||
---
|
||||
|
||||
## All PRs: Title Format (check first)
|
||||
|
||||
PR title must follow the project convention documented in `docs/contributing/README.md`:
|
||||
|
||||
| Prefix | Applies to |
|
||||
|--------|-----------|
|
||||
| `[Bugfix]` | Bug fixes |
|
||||
| `[CI/Build]` | Build or CI improvements |
|
||||
| `[Doc]` | Documentation changes |
|
||||
| `[Model]` | New/improved models (include model name) |
|
||||
| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
|
||||
| `[Kernel]` | CUDA/kernel changes |
|
||||
| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
|
||||
| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
|
||||
| `[Misc]` | Other changes (use sparingly) |
|
||||
|
||||
**Examples:**
|
||||
```
|
||||
✓ [Model] Add MiniCPM-o 4.5 support
|
||||
✓ [Bugfix] Fix Ovis image text encoder dtype
|
||||
✓ [Core] Qwen3-Omni performance optimization
|
||||
✓ [Hardware][Ascend] Fix VoxCPM2 on Ascend
|
||||
✗ Fix Ovis image text encoder dtype ← missing prefix
|
||||
✗ [bugfix] Fix Ovis image text encoder ← wrong case
|
||||
✗ [WIP] Add model ← WIP in title
|
||||
```
|
||||
|
||||
- [ ] **Title has valid prefix** — ✗ if missing prefix, wrong case, or WIP/Draft in title. For `[Model]` PRs, the title MUST also include the model identifier (e.g., `[Model] Add <ModelName> ...`). ⚠ if missing the model name.
|
||||
|
||||
---
|
||||
|
||||
## Bug Fix PRs
|
||||
|
||||
### Quick
|
||||
|
||||
- [ ] **PR body has 4 sections:** what broke (error message), repro steps, root cause, fix description
|
||||
- [ ] **Fix is minimal:** diff touches only files directly related to the bug. No unrelated refactors or style changes.
|
||||
- [ ] **Repro is runnable:** the PR body includes a copy-paste command that reproduces the bug
|
||||
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
|
||||
|
||||
### Full (adds)
|
||||
|
||||
- [ ] **Regression test exists:** `git diff --name-only` includes at least one `tests/` file
|
||||
- [ ] **Regression test reproduces the original error:** the test would fail on main, pass on this branch
|
||||
- [ ] **Fix matches root cause exactly:** no "fix the symptom + something else" — if the root cause is dtype, the fix is dtype, not dtype + reformatting
|
||||
- [ ] **No silent failure risk:** no bare `except: pass`, no `try: ... except: return None`, no empty fallback added
|
||||
- [ ] **Upstream pattern match:** the fix follows the same pattern as existing code for similar cases (grep for analogous `from_pretrained` calls, etc.)
|
||||
- [ ] **Environment documented:** torch/transformers/vllm versions listed if the bug is version-dependent
|
||||
|
||||
---
|
||||
|
||||
## Performance PRs
|
||||
|
||||
### Quick
|
||||
|
||||
- [ ] **Before/after numbers in PR body:** at minimum, one metric (RTF, RPS, latency, VRAM) with main vs PR values
|
||||
- [ ] **Hardware specified:** GPU model, count, VRAM. "L4" or "H20" is sufficient. "tested on GPU" is not.
|
||||
- [ ] **Concurrency scaling:** if the optimization claims better throughput, numbers at 2+ concurrency levels
|
||||
- [ ] **No unexplained regressions:** if any metric goes down at any concurrency level, explain why
|
||||
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
|
||||
|
||||
### Full (adds)
|
||||
|
||||
- [ ] **Software versions:** torch, CUDA, vllm, vllm-omni versions in PR body
|
||||
- [ ] **Warmup stated:** "1 warmup + 3 measured" or equivalent
|
||||
- [ ] **Stddev or range:** not just a single best run. "average of 2 rounds" is minimum.
|
||||
- [ ] **VRAM measured:** peak GPU memory before and after
|
||||
- [ ] **Benchmark script checked in:** under `tests/` or `examples/`, runnable by reviewer
|
||||
- [ ] **Exact command line in PR:** not pseudocode, not "run the benchmark"
|
||||
- [ ] **Pytest summary included:** `36 passed, 3 skipped, 17 warnings in 2897.66s` — proves the whole suite passes
|
||||
- [ ] **Config changes documented:** if `max_num_seqs`, `gpu_memory_utilization`, or deploy YAML changed, old and new values stated with reasoning
|
||||
|
||||
---
|
||||
|
||||
## New Model PRs
|
||||
|
||||
### Quick
|
||||
|
||||
- [ ] **PR body file list matches `git diff --name-only`:** every file path in the description exists in the diff
|
||||
- [ ] **Registry entries match claimed count:** count `_OMNI_MODELS` or `_DIFFUSION_MODELS` entries added vs claimed in body
|
||||
- [ ] **Pipeline config `custom_process_input_func` paths resolve:** each referenced module.function exists
|
||||
- [ ] **`__init__.py` exports match registry:** each registered class appears in `__init__.py` and `__all__`
|
||||
- [ ] **Shell script modes match README:** `run_curl.sh` supports every mode the README documents
|
||||
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
|
||||
|
||||
### Full (adds)
|
||||
|
||||
- [ ] **Dead code scan (5 patterns):**
|
||||
1. Dead `forward()` — defined but never called through the pipeline path
|
||||
2. Dead factory/builder functions — instantiated nowhere
|
||||
3. Dead wrappers — thin wrappers with no callers
|
||||
4. Dead branch guards — `if version > X` where X is always satisfied
|
||||
5. Unused parameters — `__init__` args never accessed
|
||||
- [ ] **Copy-paste detection:** no string constant defined in 3+ files, no cross-module validation duplication, no near-identical shape coercion fns
|
||||
- [ ] **Import hygiene:** no re-export-only imports, no module-level side effects without docs, no `import os` in function body when already at top. Function-body imports are acceptable when the import is only reachable from a closure or factory return value (i.e., `import` would otherwise execute at module load even though the user may never trigger that path). Either lift to top-of-file or annotate with `# noqa: PLC0415`.
|
||||
- [ ] **Accuracy:** at least one test compares output against reference; audio = valid WAV at correct sample rate; determinism verified if seed claimed. ⚠ if missing and not explicitly deferred to a follow-up PR or RFC section; — if explicitly deferred (the deferral target SHOULD be linked in the PR body).
|
||||
- [ ] **Performance:** hardware, software versions, warmup, and metrics (RTF, VRAM, RPS, latency) stated. ⚠ if missing and not explicitly deferred; — if explicitly deferred.
|
||||
- [ ] **Benchmark settings:** model config (TP, PP, max_model_len, enforce_eager, quant), runtime config (batch, input/output spec), environment versions
|
||||
- [ ] **Benchmark script checked in + exact command line + pytest summary**
|
||||
|
||||
---
|
||||
|
||||
## Diffusion Model PRs
|
||||
|
||||
Diffusion models live under `vllm_omni/diffusion/`. In addition to the New Model checklist above, verify:
|
||||
|
||||
- [ ] **Transformer adapter follows contract:** inherits from `nn.Module`, implements `load_weights()`, uses `vllm_omni.diffusion.attention.layer.Attention`
|
||||
- [ ] **Pipeline has both offline and online paths:** `forward()` for batch generation, streaming path for online serving
|
||||
- [ ] **Cache-DiT integration:** at least one acceleration (TeaCache, FBCache, etc.) is wired up. ✗ for new **top-level** pipelines. For **subclass** pipelines that inherit `__call__` from a parent where Cache-DiT is already absent, downgrade to ⚠ and recommend filing a follow-up issue against the parent pipeline family.
|
||||
- [ ] **Parallelism config:** TP/SP/USP/CFG-Parallel options exposed in pipeline config
|
||||
- [ ] **Docs table updated:** model added to `docs/models/supported_models.md` with correct metadata
|
||||
- [ ] **E2E test exists:** at least one test exercises the full generate path
|
||||
|
||||
---
|
||||
|
||||
## General PRs
|
||||
|
||||
### Quick
|
||||
|
||||
- [ ] **PR body matches diff:** description doesn't claim changes not in the diff
|
||||
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
|
||||
- [ ] **CI gates passing:** DCO, pre-commit, build — check `gh pr view --json statusCheckRollup` if a PR already exists
|
||||
|
||||
### Full (adds)
|
||||
|
||||
- [ ] **No new dead code in changed files:** scan for unreachable paths, unused imports
|
||||
- [ ] **Test coverage:** if core code changed (`engine/`, `stages/`, `connectors/`), tests are added or existing tests cover the path
|
||||
- [ ] **Import hygiene:** no redundant imports added
|
||||
- [ ] **No unrelated changes:** diff scope matches PR description
|
||||
@@ -25,12 +25,7 @@ include:
|
||||
- `add-tts-model`: covers integration of new TTS models and related serving
|
||||
workflows
|
||||
- `generate-release-note`: helps prepare release notes for repository changes
|
||||
- `precheck-pr`: self-check a branch before creating a PR — validates PR title
|
||||
format, catches dead code, verifies accuracy/perf claims, and confirms merge
|
||||
readiness
|
||||
- `review-pr`: provides a structured workflow for reviewing pull requests
|
||||
- `vllm-omni-npu-model-runner-upgrade`: upgrades NPU model runners to align with the
|
||||
latest vllm-ascend NPUModelRunner
|
||||
|
||||
## Maintenance Guidelines
|
||||
|
||||
|
||||
@@ -1,6 +1,4 @@
|
||||
<!-- markdownlint-disable -->
|
||||
**Before submitting:** run the precheck-pr skill with code agent for a self-check against project conventions — catches dead code, missing benchmarks, and title format issues.
|
||||
|
||||
PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
|
||||
|
||||
## Purpose
|
||||
@@ -20,5 +18,4 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
|
||||
- [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
|
||||
</details>
|
||||
|
||||
**BEFORE SUBMITTING,** please read <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>
|
||||
(anything written below this line will be removed by GitHub Actions)
|
||||
**BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)
|
||||
@@ -270,6 +270,3 @@ 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,11 +10,6 @@ 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
|
||||
|
||||
@@ -16,7 +16,6 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
||||
---
|
||||
|
||||
*Latest News* 🔥
|
||||
- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
|
||||
- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
|
||||
- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
|
||||
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
|
||||
|
||||
+35
-51
@@ -1,60 +1,44 @@
|
||||
# Benchmarks
|
||||
# Benchmarks Overview and Architecture
|
||||
|
||||
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.
|
||||
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`).
|
||||
|
||||
## Benchmark families
|
||||
## 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.
|
||||
|
||||
### [TTS](tts/README.md) — Text-to-Speech
|
||||
## 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.
|
||||
|
||||
Model-agnostic serving benchmarks for TTS models, including Qwen3-TTS and VoxCPM2.
|
||||
## 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.
|
||||
|
||||
- **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)
|
||||
## 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).
|
||||
|
||||
### [Diffusion](diffusion/README.md) — Image and Video Generation
|
||||
- **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).
|
||||
|
||||
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).
|
||||
- **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).
|
||||
|
||||
- **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
|
||||
## 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.
|
||||
|
||||
### [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.
|
||||
## 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.
|
||||
@@ -151,9 +151,6 @@ 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]
|
||||
@@ -173,12 +170,6 @@ 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,47 +64,6 @@ 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)
|
||||
@@ -114,7 +73,10 @@ def parse_score_payload(raw_text: str) -> dict[str, Any]:
|
||||
parsed = {"score": json.loads(stripped), "reasoning": ""}
|
||||
else:
|
||||
raise
|
||||
parsed["score"] = _normalize_score_list(parsed.get("score", []))
|
||||
score = parsed.get("score", [])
|
||||
if not isinstance(score, list):
|
||||
score = [score]
|
||||
parsed["score"] = [int(value) for value in score]
|
||||
return parsed
|
||||
|
||||
|
||||
@@ -429,9 +391,6 @@ 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
|
||||
@@ -441,9 +400,6 @@ 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(
|
||||
@@ -522,9 +478,6 @@ 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 {
|
||||
@@ -680,25 +633,6 @@ 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)
|
||||
@@ -741,9 +675,6 @@ 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,
|
||||
|
||||
@@ -1,10 +0,0 @@
|
||||
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
|
||||
@@ -1,10 +0,0 @@
|
||||
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.
@@ -1,328 +0,0 @@
|
||||
# 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()
|
||||
@@ -1,108 +0,0 @@
|
||||
#!/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"
|
||||
@@ -47,13 +47,6 @@ Usage:
|
||||
--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
|
||||
@@ -573,135 +566,6 @@ 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)
|
||||
@@ -1086,8 +950,6 @@ 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}")
|
||||
|
||||
@@ -1235,7 +1097,7 @@ if __name__ == "__main__":
|
||||
"--dataset",
|
||||
type=str,
|
||||
default="vbench",
|
||||
choices=["vbench", "trace", "random", "custom"],
|
||||
choices=["vbench", "trace", "random"],
|
||||
help="Dataset to use.",
|
||||
)
|
||||
parser.add_argument(
|
||||
|
||||
@@ -53,8 +53,6 @@ _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)
|
||||
@@ -83,7 +81,6 @@ 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]
|
||||
@@ -133,18 +130,12 @@ 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)
|
||||
@@ -244,15 +235,6 @@ 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()
|
||||
|
||||
@@ -307,7 +289,6 @@ 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:
|
||||
|
||||
@@ -14,7 +14,6 @@ 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
|
||||
@@ -28,76 +27,13 @@ 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,8 +1,5 @@
|
||||
# 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.22.0
|
||||
ARG VLLM_BASE_TAG=v0.21.0
|
||||
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
|
||||
ARG APP_DIR=/workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.22.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.21.0
|
||||
FROM ${BASE_IMAGE}
|
||||
|
||||
ARG COMMON_WORKDIR=/app
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.22.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0
|
||||
FROM ${BASE_IMAGE} AS base
|
||||
|
||||
# Declare a variable to know if we want to use the nightly build or the stable build.
|
||||
|
||||
@@ -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.22.0
|
||||
ARG VLLM_VERSION=v0.21.0
|
||||
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
|
||||
@@ -92,6 +92,7 @@ Core scheduling and caching components.
|
||||
Configuration classes.
|
||||
|
||||
- [vllm_omni.config.model.OmniModelConfig][]
|
||||
- [vllm_omni.config.stage_config.ModelPipeline][]
|
||||
- [vllm_omni.config.stage_config.StageConfig][]
|
||||
- [vllm_omni.config.stage_config.StageConfigFactory][]
|
||||
- [vllm_omni.config.stage_config.StageType][]
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 403 KiB After Width: | Height: | Size: 402 KiB |
+26
-134
@@ -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" and "openai-image-edits-omni" backend.
|
||||
We currently only support using the "openai-chat-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,73 +34,69 @@ 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", "ttfc", "tpoc", "icl", '
|
||||
'"tpop", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
|
||||
'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
|
||||
|
||||
- `-print-stage`
|
||||
Print per-stage benchmark metrics for --omni serving when stage metrics are returned by the server. Disabled by default.
|
||||
|
||||
- `--save-result`
|
||||
- `--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
|
||||
@@ -108,7 +104,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))]
|
||||
@@ -119,7 +115,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.
|
||||
@@ -280,9 +276,7 @@ Generate synthetic image、video、audio inputs alongside random text prompts to
|
||||
|
||||
Notes:
|
||||
|
||||
- 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`.
|
||||
- Works only with online benchmark via the OpenAI backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
|
||||
|
||||
Start the server (example):
|
||||
|
||||
@@ -363,105 +357,3 @@ 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>
|
||||
@@ -116,20 +116,6 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
|
||||
!!! note
|
||||
If the PR spans more than one category, please include all relevant prefixes.
|
||||
|
||||
### Pre-Check Before Submitting
|
||||
|
||||
Before submitting a PR, run the precheck-pr skill with code agent for a self-review against project conventions:
|
||||
|
||||
```
|
||||
/precheck-pr
|
||||
```
|
||||
|
||||
The skill offers two modes:
|
||||
- **Quick (~3 min):** catches showstoppers — PR title format, missing benchmark claims, rebase status
|
||||
- **Full (~10 min):** thorough maintainer-grade review — dead code scan, copy-paste detection, import hygiene
|
||||
|
||||
The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Model, and General. Each type has a tailored checklist that validates evidence quality (repro steps, A/B benchmarks, registry entries, etc.). See the skill in `.claude/skills/precheck-pr/` for the full checklist.
|
||||
|
||||
### Local Test
|
||||
Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
|
||||
|
||||
|
||||
@@ -13,10 +13,6 @@ 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)
|
||||
|
||||
@@ -1,208 +0,0 @@
|
||||
# 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,13 +169,15 @@ class DiffusionRequestStatus(enum.IntEnum):
|
||||
|
||||
@dataclass
|
||||
class DiffusionRequestState:
|
||||
request_id: str
|
||||
sched_req_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.
|
||||
@@ -186,6 +188,7 @@ 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()
|
||||
@@ -196,8 +199,7 @@ class _BaseScheduler(SchedulerInterface):
|
||||
|
||||
- **Common state storage**: Shared request maps and waiting/running sets live in the base class.
|
||||
|
||||
- **Shared cleanup logic**: Duplicate-request checks, finish handling, and state removal are centralized instead of
|
||||
duplicated in each policy.
|
||||
- **Shared cleanup logic**: Request-id registration, 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.
|
||||
|
||||
@@ -217,12 +219,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 request ids.
|
||||
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished scheduler ids.
|
||||
|
||||
#### 2.5 Engine-Driven Execution Loop
|
||||
|
||||
```python
|
||||
request_id = scheduler.add_request(request)
|
||||
sched_req_id = scheduler.add_request(request)
|
||||
while True:
|
||||
sched_output = scheduler.schedule()
|
||||
output = executor.add_req(req)
|
||||
@@ -305,8 +307,8 @@ class GPUWorker:
|
||||
)
|
||||
|
||||
# Load model
|
||||
model_loader = DiffusersPipelineLoader(load_config, self.od_config)
|
||||
self.pipeline = model_loader.load_model(load_device=f"cuda:{rank}")
|
||||
model_loader = DiffusersPipelineLoader(load_config)
|
||||
self.pipeline = model_loader.load_model(od_config, load_device=f"cuda:{rank}")
|
||||
|
||||
# Setup cache backend
|
||||
from vllm_omni.diffusion.cache.selector import get_cache_backend
|
||||
|
||||
@@ -411,8 +411,6 @@ 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, 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).
|
||||
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).
|
||||
|
||||
=== "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.22.0 --torch-backend=auto
|
||||
uv pip install vllm==0.21.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.22.0 release version of the library
|
||||
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
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.22.0 --torch-backend=auto
|
||||
uv pip install vllm==0.21.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
The 0.22.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
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.
|
||||
|
||||
#### 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.22.0
|
||||
git checkout v0.21.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 12.9 (the default for v0.22.0; the wheel filename has no `+cu129` 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_28_x86_64.whl
|
||||
#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
|
||||
```
|
||||
Install vllm with command below (If you have no existing PyTorch).
|
||||
```bash
|
||||
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
--ipc=host \
|
||||
vllm/vllm-omni:v0.22.0 \
|
||||
vllm/vllm-omni:v0.21.0 \
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
|
||||
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
|
||||
|
||||
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||
```bash
|
||||
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
|
||||
uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.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.22.0 release version of the library
|
||||
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
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
|
||||
uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
|
||||
```
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
#### 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.22.0
|
||||
git checkout v0.21.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.22.0 \
|
||||
vllm/vllm-omni-rocm:v0.21.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.22.0
|
||||
vllm/vllm-omni-rocm:v0.21.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.22.0 --torch-backend=auto
|
||||
uv pip install vllm==0.21.0 --torch-backend=auto
|
||||
|
||||
# On ROCm
|
||||
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
|
||||
uv pip install vllm==0.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.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.22.0` and vLLM Omni `0.22.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.21.0` and vLLM Omni `0.21.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
|
||||
|
||||
## Offline Inference
|
||||
|
||||
|
||||
@@ -119,10 +119,9 @@ 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,7 +16,6 @@ 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:
|
||||
@@ -192,21 +191,18 @@ 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"
|
||||
content += f"Source <{self.github_url(self.path)}>.\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"
|
||||
|
||||
# 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).as_posix()
|
||||
main_file_rel = self.main_file.relative_to(ROOT_DIR)
|
||||
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:
|
||||
@@ -220,15 +216,10 @@ class Example:
|
||||
|
||||
content += "## Example materials\n\n"
|
||||
for file in sorted(self.other_files):
|
||||
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
|
||||
content += f'??? abstract "{file.relative_to(self.path)}"\n'
|
||||
if file.suffix != ".md":
|
||||
content += f" {code_fence}{file.suffix[1:]}\n"
|
||||
content += f' --8<-- "{file.relative_to(ROOT_DIR).as_posix()}"\n'
|
||||
content += f' --8<-- "{file.relative_to(ROOT_DIR)}"\n'
|
||||
if file.suffix != ".md":
|
||||
content += f" {code_fence}\n"
|
||||
|
||||
@@ -300,7 +291,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: rel_path.as_posix()})
|
||||
category_items.append({example.title: str(rel_path)})
|
||||
|
||||
if category_items:
|
||||
# Format category name (e.g., "offline_inference" -> "Offline Inference")
|
||||
|
||||
@@ -1,29 +0,0 @@
|
||||
# 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,44 +1,20 @@
|
||||
// Enables MathJax rendering only on pages that contain math.
|
||||
// Enables MathJax rendering
|
||||
window.MathJax = {
|
||||
tex: {
|
||||
inlineMath: [["\\(", "\\)"]],
|
||||
displayMath: [["\\[", "\\]"]],
|
||||
processEscapes: true,
|
||||
processEnvironments: true,
|
||||
},
|
||||
options: {
|
||||
ignoreHtmlClass: ".*|",
|
||||
processHtmlClass: "arithmatex",
|
||||
},
|
||||
};
|
||||
tex: {
|
||||
inlineMath: [["\\(", "\\)"]],
|
||||
displayMath: [["\\[", "\\]"]],
|
||||
processEscapes: true,
|
||||
processEnvironments: true
|
||||
},
|
||||
options: {
|
||||
ignoreHtmlClass: ".*|",
|
||||
processHtmlClass: "arithmatex"
|
||||
}
|
||||
};
|
||||
|
||||
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();
|
||||
});
|
||||
});
|
||||
document$.subscribe(() => {
|
||||
MathJax.startup.output.clearCache()
|
||||
MathJax.typesetClear()
|
||||
MathJax.texReset()
|
||||
MathJax.typesetPromise()
|
||||
})
|
||||
@@ -1,2 +0,0 @@
|
||||
// Opt out of bfcache so large generated API pages are released after navigation.
|
||||
window.addEventListener("unload", () => {});
|
||||
@@ -1,42 +1,21 @@
|
||||
// Load Mermaid only on pages that contain diagrams.
|
||||
let mermaidLoading;
|
||||
// Initialize Mermaid for diagram rendering
|
||||
mermaid.initialize({
|
||||
startOnLoad: false,
|
||||
theme: 'default',
|
||||
securityLevel: 'loose',
|
||||
flowchart: {
|
||||
useMaxWidth: true,
|
||||
htmlLabels: true
|
||||
}
|
||||
});
|
||||
|
||||
function loadMermaid() {
|
||||
if (window.mermaid) {
|
||||
return Promise.resolve();
|
||||
// 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
|
||||
});
|
||||
}
|
||||
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,7 +19,6 @@ 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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
@@ -33,7 +32,6 @@ 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, V2V, 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` | ✅︎ | ✅︎ | | |
|
||||
@@ -42,16 +40,13 @@ 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` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | |
|
||||
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | |
|
||||
| `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` | ✅︎ | ✅︎ | | |
|
||||
@@ -64,10 +59,6 @@ th {
|
||||
| `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` | ✅︎ | | | |
|
||||
| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
|
||||
| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
|
||||
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
|
||||
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
|
||||
| `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` | ✅︎ | ✅︎ | | |
|
||||
@@ -75,13 +66,11 @@ th {
|
||||
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
|
||||
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
|
||||
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
|
||||
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | |
|
||||
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
|
||||
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
|
||||
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
|
||||
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
|
||||
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
|
||||
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | |
|
||||
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
|
||||
|
||||
|
||||
@@ -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_url,input_image_url],
|
||||
"url": [input_image_url1,input_image_url],
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1,
|
||||
"seed": 777,
|
||||
@@ -92,8 +92,7 @@ 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). |
|
||||
| `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. |
|
||||
| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s).
|
||||
|
||||
#### vllm-omni Extension Parameters
|
||||
|
||||
@@ -110,8 +109,6 @@ Content-Type: multipart/form-data
|
||||
|
||||
### Response Format
|
||||
|
||||
When `stream=false` or omitted, the endpoint returns the standard image edit response:
|
||||
|
||||
```json
|
||||
{
|
||||
"created": 1701234567,
|
||||
@@ -123,56 +120,10 @@ When `stream=false` or omitted, the endpoint returns the standard image edit res
|
||||
}
|
||||
],
|
||||
"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
|
||||
@@ -183,28 +134,12 @@ 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
|
||||
|
||||
|
||||
@@ -116,7 +116,7 @@ Content-Type: application/json
|
||||
| `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
|
||||
| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
|
||||
|
||||
**Supported languages:** Only applicable to Qwen3-TTS. Derived from the model configuration (`talker_config.codec_language_id` in the checkpoint's `config.json`), plus `Auto`, which is always accepted. Official Qwen3-TTS checkpoints support: Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian.
|
||||
**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
|
||||
|
||||
#### Voice Clone Parameters (Base task)
|
||||
|
||||
@@ -381,44 +381,6 @@ 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 |
|
||||
|
||||
@@ -67,9 +67,8 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `input_reference` | file | null | Uploaded reference image or video for image-to-video/video-to-video requests |
|
||||
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` or `video_reference` |
|
||||
| `video_reference` | string | null | JSON-encoded reference video payload; do not combine with `input_reference` or `image_reference` |
|
||||
| `input_reference` | file | null | Uploaded reference image for image-to-video requests |
|
||||
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` |
|
||||
| `width` | integer | model default | Output video width |
|
||||
| `height` | integer | model default | Output video height |
|
||||
| `num_frames` | integer | model default | Number of generated frames |
|
||||
@@ -81,8 +80,6 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
|
||||
| `flow_shift` | number | null | Scheduler flow-shift value |
|
||||
| `true_cfg_scale` | number | null | True CFG scale when supported by the model |
|
||||
| `seed` | integer | null | Random seed for reproducibility |
|
||||
| `generate_sound` | boolean | false | Request model-generated audio for video models that support sound generation |
|
||||
| `sound_duration` | number | null | Duration in seconds for generated audio; defaults to generated video duration |
|
||||
| `negative_prompt` | string | null | Text describing what to avoid in the generated video |
|
||||
| `enable_frame_interpolation` | boolean | null | Enable post-generation frame interpolation |
|
||||
| `frame_interpolation_exp` | integer | null | Interpolation exponent; `1=2x`, `2=4x`, and so on |
|
||||
@@ -126,41 +123,6 @@ curl -s http://localhost:8091/v1/videos \
|
||||
-F "fps=16"
|
||||
```
|
||||
|
||||
### Video-to-Video
|
||||
|
||||
For models that support video conditioning, upload the reference video with
|
||||
`input_reference`:
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/videos \
|
||||
-F "prompt=continue this motion with consistent subjects and lighting" \
|
||||
-F "input_reference=@input.mp4;type=video/mp4" \
|
||||
-F "width=1280" \
|
||||
-F "height=720" \
|
||||
-F "num_frames=80" \
|
||||
-F "fps=16"
|
||||
```
|
||||
|
||||
You can also pass a JSON-safe video URL or `data:video/...;base64,...` payload
|
||||
through `video_reference`. Do not send `video_reference` together with
|
||||
`input_reference` or `image_reference`.
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/videos \
|
||||
-F "prompt=continue this motion with consistent subjects and lighting" \
|
||||
-F 'video_reference={"video_url":"https://example.com/input.mp4"}' \
|
||||
-F "width=1280" \
|
||||
-F "height=720" \
|
||||
-F "num_frames=80" \
|
||||
-F "fps=16"
|
||||
```
|
||||
|
||||
JSON references currently support `image_url`/`video_url`; `file_id` references
|
||||
are not implemented yet. Models may expose additional V2V controls through
|
||||
`extra_params`. For example, Cosmos3 supports
|
||||
`condition_frame_indexes_vision` and `condition_video_keep` to select which
|
||||
decoded reference frames are used as clean conditioning.
|
||||
|
||||
### Synchronous Generation
|
||||
|
||||
```bash
|
||||
@@ -184,10 +146,7 @@ export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
|
||||
|
||||
## Model-Specific Examples
|
||||
|
||||
For complete text-to-video, image-to-video, and model-specific video-to-video
|
||||
walkthroughs, see:
|
||||
For complete text-to-video and image-to-video walkthroughs, see:
|
||||
|
||||
- [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md)
|
||||
- [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md)
|
||||
- [Cosmos3 recipes](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md)
|
||||
for model-specific video-to-video examples and conditioning controls
|
||||
@@ -1,98 +0,0 @@
|
||||
# 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,26 +1,28 @@
|
||||
# Diffusion Attention Backends
|
||||
|
||||
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.
|
||||
This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, 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. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector.
|
||||
This backend is used by diffusion attention layers such as the DiT attention in video and image generation models.
|
||||
|
||||
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.
|
||||
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`.
|
||||
|
||||
## Backend Options
|
||||
|
||||
| Value | Notes |
|
||||
|---|---|
|
||||
| `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. |
|
||||
|
||||
| `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. |
|
||||
|
||||
## Configuration
|
||||
|
||||
@@ -93,80 +95,6 @@ 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.
|
||||
@@ -187,56 +115,11 @@ Quick check:
|
||||
python -c "import sageattention; print(sageattention.__file__)"
|
||||
```
|
||||
|
||||
## SageAttention3 Installation
|
||||
## Usage
|
||||
|
||||
vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni.
|
||||
### Enable SageAttention
|
||||
|
||||
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)
|
||||
Example: HunyuanVideo-1.5 text-to-video
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
@@ -245,7 +128,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 hv15_sage.mp4
|
||||
--output ../tmp/hv15_modelopt_sage.mp4
|
||||
```
|
||||
|
||||
Example: Wan2.2 TI2V 5B
|
||||
@@ -260,20 +143,6 @@ 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:
|
||||
@@ -304,13 +173,13 @@ python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
|
||||
## Validation Guidance
|
||||
|
||||
Don't assume a faster attention backend is numerically interchangeable with `TORCH_SDPA`.
|
||||
Do not assume that a faster attention backend is numerically interchangeable with `FLASH_ATTN`.
|
||||
|
||||
Always compare:
|
||||
|
||||
- End-to-end runtime
|
||||
- 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)
|
||||
- DiT / diffusion stage runtime
|
||||
- Output quality against a known-good baseline
|
||||
|
||||
At minimum, keep the same:
|
||||
|
||||
@@ -318,19 +187,6 @@ At minimum, keep the same:
|
||||
- prompt
|
||||
- seed
|
||||
- resolution
|
||||
- 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.
|
||||
- frame count
|
||||
- inference steps
|
||||
- parallel config
|
||||
@@ -113,7 +113,7 @@ 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** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
@@ -129,31 +129,26 @@ 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** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
|
||||
|
||||
**Frame Interpolation Support**
|
||||
|
||||
|
||||
@@ -23,7 +23,6 @@ list of supported architectures across all modalities, see
|
||||
| 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
|
||||
|
||||
@@ -334,46 +333,6 @@ 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"
|
||||
@@ -384,9 +343,6 @@ python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
|
||||
``````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"
|
||||
|
||||
@@ -17,9 +17,6 @@ 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:
|
||||
@@ -75,12 +72,6 @@ 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:
|
||||
|
||||
@@ -19,7 +19,6 @@ For the full list of supported architectures across all modalities, see
|
||||
|---|---|---|---|---|---|
|
||||
| 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) |
|
||||
@@ -194,45 +193,6 @@ 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`.
|
||||
@@ -329,16 +289,6 @@ Single-GPU serves now default to the uniproc executor (lower IPC overhead, the B
|
||||
|
||||
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
|
||||
# CustomVoice with a predefined speaker
|
||||
@@ -380,31 +330,6 @@ 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 \
|
||||
@@ -444,10 +369,8 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
|
||||
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
|
||||
|
||||
### 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.
|
||||
- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
|
||||
- `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.
|
||||
|
||||
---
|
||||
|
||||
@@ -475,23 +398,6 @@ 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
|
||||
@@ -553,13 +459,6 @@ 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"
|
||||
|
||||
@@ -52,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
|
||||
|
||||
|
||||
@@ -28,50 +28,6 @@ in deep DiT blocks.
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide. FP8 on Ampere may use a weight-only path where available.
|
||||
|
||||
### Faster FP8 GEMM on Blackwell (quack)
|
||||
|
||||
On Blackwell (SM 100+), vLLM runs FP8 linears through the FlashInfer kernel, which
|
||||
applies the bias as a separate kernel after the GEMM. On the small GEMMs in video
|
||||
DiTs this bias add is a significant overhead. Installing the optional `quack` kernel
|
||||
lets vLLM-Omni fuse `alpha * (A @ B) + bias` into a single CuteDSL GEMM, recovering
|
||||
that overhead (e.g. HunyuanVideo-1.5 FP8 goes from slower-than-BF16 to faster).
|
||||
|
||||
```bash
|
||||
# CUDA 12.9
|
||||
pip install vllm-omni[quack]
|
||||
|
||||
# CUDA 13.x
|
||||
pip install 'quack-kernels[cu13]' --extra-index-url https://download.pytorch.org/whl/cu130
|
||||
```
|
||||
|
||||
It is enabled automatically once installed (no flag needed) and is **Blackwell-only**:
|
||||
on Hopper/Ada the CUTLASS FP8 kernel already fuses bias, so quack is not used there.
|
||||
Set `VLLM_OMNI_USE_QUACK_FP8=0` to force the FlashInfer path. If `quack-kernels` is
|
||||
not installed, FP8 still works — it just keeps the unfused FlashInfer path.
|
||||
|
||||
#### Compile cache and warmup
|
||||
|
||||
quack JIT-compiles its kernel once per distinct GEMM shape (tens of seconds, longer
|
||||
the first time across all autotuned configs). The compiled `.o` files are cached on
|
||||
disk and reused on later runs, so this is a one-time cost — **not** per request.
|
||||
|
||||
vLLM-Omni points that cache at `~/.cache/vllm_omni/quack` (override with
|
||||
`QUACK_CACHE_DIR`) instead of quack's default under `/tmp`, so it survives restarts.
|
||||
In containers, set `QUACK_CACHE_DIR` to a mounted/persistent path — or bake it into
|
||||
the image — so the first cold start does not recompile. The engine's startup dummy
|
||||
run already exercises the kernels, so with a warm cache the first real request is fast.
|
||||
|
||||
To pre-warm specific shapes (e.g. at image build time):
|
||||
|
||||
```python
|
||||
from vllm_omni.quantization.quack_fp8 import warmup_quack_fp8
|
||||
# (M, K, N) per linear; M = number of tokens for your resolution/frame count
|
||||
warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
|
||||
```
|
||||
|
||||
> The PyPI package is `quack-kernels` (imported as `quack`); plain `pip install
|
||||
> quack` is an unrelated statistics library. Requires CUDA 12.9+ and Python 3.12.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
|
||||
@@ -74,7 +74,7 @@ Online:
|
||||
vllm serve Qwen/Qwen-Image \
|
||||
--omni \
|
||||
--port 8000 \
|
||||
--diffusion-quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
|
||||
--quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
@@ -7,28 +7,26 @@ 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 three modes:
|
||||
This method supports two modes:
|
||||
|
||||
| Mode | Description |
|
||||
|------|-------------|
|
||||
| **Online** | BF16 weights are quantized to MXFP8 at load time — no pre-processing needed |
|
||||
| **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` |
|
||||
| **Offline** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| 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) | ✅ | ✅ | ⭕ |
|
||||
| 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.
|
||||
|
||||
**Note**: Intel XPU only supports AutoRound MXFP8 for offline mode. Use AutoRound quantized checkpoints or online mode for XPU.
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
@@ -84,9 +82,9 @@ python text_to_video.py --model <your-model> --quantization mxfp8
|
||||
vllm serve <your-model> --omni --quantization mxfp8
|
||||
```
|
||||
|
||||
### Offline Mode (Native)
|
||||
### Offline Mode
|
||||
|
||||
Native offline mode loads a pre-quantized checkpoint from msModelSlim. A preprocessing
|
||||
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.
|
||||
@@ -161,41 +159,10 @@ omni = Omni(model="/path/to/Wan2.2-TI2V-5B-MXFP8")
|
||||
```
|
||||
|
||||
!!! note
|
||||
No `--quantization` flag is needed for native offline mode. The preprocessing
|
||||
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 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 |
|
||||
|
||||
@@ -21,7 +21,7 @@ by msModelSlim and the merge tools.
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ |
|
||||
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ |
|
||||
| Intel XPU | ⭕ | ⭕ | ✅ | ⭕ |
|
||||
| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ |
|
||||
| Ascend NPU | ❌ | ✅ | ✅ | ✅ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
@@ -36,7 +36,7 @@ MXFP4 are documented for the Ascend NPU path.
|
||||
|--------|-------|----------------|--------|
|
||||
| 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 |
|
||||
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Ascend NPU only |
|
||||
| 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)
|
||||
|
||||
@@ -195,101 +195,6 @@ 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,6 +121,9 @@ 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
|
||||
|
||||
@@ -151,7 +154,6 @@ 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"
|
||||
@@ -162,7 +164,9 @@ def main():
|
||||
if args.quantization:
|
||||
omni_kwargs["quantization_config"] = args.quantization
|
||||
|
||||
omni = Omni(model=model_name, **omni_kwargs)
|
||||
# Override CLI --model with the derived model_name.
|
||||
omni_kwargs["model"] = model_name
|
||||
omni = Omni(**omni_kwargs)
|
||||
|
||||
formatted_prompts = []
|
||||
for p in prompts:
|
||||
|
||||
@@ -1,14 +0,0 @@
|
||||
{
|
||||
"prompt": "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. The robotic arm is in the process of pouring a transparent liquid into the glass. To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. The lighting is bright, casting soft shadows on the table. The robotic arm's movements are smooth and controlled, demonstrating precision in its task. As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. The other robotic arm on the right side also stays stationary throughout the video. The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper.",
|
||||
"negative_prompt": "",
|
||||
"vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4",
|
||||
"height": 720,
|
||||
"width": 1280,
|
||||
"num_frames": 189,
|
||||
"num_inference_steps": 35,
|
||||
"guidance_scale": 6.0,
|
||||
"flow_shift": 10.0,
|
||||
"fps": 24,
|
||||
"condition_frame_indexes_vision": [0, 1],
|
||||
"condition_video_keep": "first"
|
||||
}
|
||||
@@ -1,304 +0,0 @@
|
||||
#!/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/deploy/dynin_omni.yaml)):
|
||||
([`dynin_omni.yaml`](../../../vllm_omni/model_executor/stage_configs/dynin_omni.yaml)):
|
||||
|
||||
- Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
|
||||
- Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`
|
||||
|
||||
@@ -18,6 +18,8 @@ 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 = {
|
||||
@@ -878,8 +880,8 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
|
||||
parser.add_argument(
|
||||
"--stage-config-path",
|
||||
type=str,
|
||||
default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
|
||||
help="Path to deploy config yaml.",
|
||||
default=str(repo_root / "vllm_omni/model_executor/stage_configs/dynin_omni.yaml"),
|
||||
help="Path to stage config yaml.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dynin-config-path",
|
||||
@@ -970,6 +972,8 @@ 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_image_3_moe.yaml")
|
||||
_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3.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=None, help="Output image height.")
|
||||
parser.add_argument("--width", type=int, default=None, help="Output image width.")
|
||||
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("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
|
||||
parser.add_argument(
|
||||
"--bot-task",
|
||||
@@ -107,6 +107,9 @@ def parse_args():
|
||||
),
|
||||
)
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -226,21 +229,7 @@ def main():
|
||||
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
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}"
|
||||
)
|
||||
ar_stop_token_ids = resolve_stop_token_ids(task=task, bot_task=bot_task, tokenizer=tokenizer)
|
||||
for sp in params_list:
|
||||
if isinstance(sp, OmniDiffusionSamplingParams):
|
||||
sp.num_inference_steps = args.steps
|
||||
|
||||
@@ -127,12 +127,7 @@ def run_one_path(
|
||||
),
|
||||
args.device,
|
||||
)
|
||||
pred = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-sample-{index}",
|
||||
)
|
||||
pred = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
|
||||
results.append(
|
||||
{
|
||||
@@ -155,7 +150,6 @@ def run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs: dict[str, torch.Tensor],
|
||||
noise: torch.Tensor,
|
||||
request_id: str,
|
||||
) -> torch.Tensor:
|
||||
output = pipeline.forward(
|
||||
OmniDiffusionRequest(
|
||||
@@ -167,7 +161,6 @@ def run_pipeline_forward(
|
||||
"decode_image": False,
|
||||
}
|
||||
),
|
||||
request_id=request_id,
|
||||
)
|
||||
)
|
||||
if output.error:
|
||||
@@ -219,38 +212,23 @@ def benchmark_forward(
|
||||
|
||||
_synchronize(args.device)
|
||||
cold_start_begin = time.perf_counter()
|
||||
pred = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-benchmark-{index}-cold",
|
||||
)
|
||||
pred = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
_synchronize(args.device)
|
||||
cold_start_ms = (time.perf_counter() - cold_start_begin) * 1000.0
|
||||
|
||||
warmup_ms: list[float] = []
|
||||
for iter_idx in range(args.warmup_iters):
|
||||
for _ in range(args.warmup_iters):
|
||||
_synchronize(args.device)
|
||||
begin = time.perf_counter()
|
||||
_ = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-benchmark-{index}-warmup-{iter_idx}",
|
||||
)
|
||||
_ = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
_synchronize(args.device)
|
||||
warmup_ms.append((time.perf_counter() - begin) * 1000.0)
|
||||
|
||||
benchmark_ms: list[float] = []
|
||||
for iter_idx in range(args.benchmark_iters):
|
||||
for _ in range(args.benchmark_iters):
|
||||
_synchronize(args.device)
|
||||
begin = time.perf_counter()
|
||||
_ = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-benchmark-{index}-iter-{iter_idx}",
|
||||
)
|
||||
_ = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
_synchronize(args.device)
|
||||
benchmark_ms.append((time.perf_counter() - begin) * 1000.0)
|
||||
|
||||
@@ -317,12 +295,7 @@ 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,
|
||||
request_id="internvla-a1-open-loop",
|
||||
),
|
||||
run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(policy, batch_inputs, noise),
|
||||
num_episodes=args.num_episodes,
|
||||
seed=args.seed,
|
||||
device=args.device,
|
||||
|
||||
@@ -1,52 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,339 +0,0 @@
|
||||
# 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,6 +19,7 @@ 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."
|
||||
@@ -48,6 +49,7 @@ 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,6 +29,7 @@ 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__)
|
||||
@@ -117,6 +118,7 @@ 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,10 +21,11 @@ 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
|
||||
|
||||
@@ -342,7 +343,7 @@ def main(args):
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser.add_argument(
|
||||
"--model-name",
|
||||
"-m",
|
||||
@@ -438,6 +439,7 @@ 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,13 +2,12 @@
|
||||
|
||||
[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 three deployment modes:
|
||||
vLLM-Omni supports two 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).
|
||||
|
||||
@@ -53,7 +52,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.
|
||||
@@ -88,17 +87,6 @@ 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,13 +17,14 @@ 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"
|
||||
@@ -396,7 +397,7 @@ def main(args):
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
|
||||
parser = FlexibleArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
@@ -495,6 +496,7 @@ def parse_args():
|
||||
help="Output directory for results.",
|
||||
)
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -19,9 +19,10 @@ 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
|
||||
|
||||
@@ -431,7 +432,7 @@ def main(args):
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
@@ -553,6 +554,7 @@ 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,9 +19,10 @@ 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
|
||||
|
||||
@@ -421,7 +422,7 @@ def main(args):
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
@@ -560,6 +561,7 @@ def parse_args():
|
||||
help="Model dtype (auto, half, float16, bfloat16, float, float32).",
|
||||
)
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -39,9 +39,10 @@ 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__)
|
||||
|
||||
@@ -382,9 +383,7 @@ async def run_all(args):
|
||||
print(f"[Info] Creating AsyncOmni with deploy_config={args.deploy_config}")
|
||||
async_omni = None
|
||||
try:
|
||||
# ``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)
|
||||
async_omni = AsyncOmni(**vars(args))
|
||||
|
||||
# Use default sampling params from stage config (they are pre-configured
|
||||
# in the YAML for each stage).
|
||||
@@ -454,7 +453,7 @@ async def run_all(args):
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(
|
||||
parser = FlexibleArgumentParser(
|
||||
description=(
|
||||
"Offline inference with async_chunk enabled via AsyncOmni. "
|
||||
"Downstream stages start before upstream stages finish, "
|
||||
@@ -590,6 +589,7 @@ def parse_args():
|
||||
default=16000,
|
||||
help="Sampling rate for audio loading.",
|
||||
)
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -37,10 +37,6 @@ 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.
|
||||
@@ -139,8 +135,6 @@ 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,6 +29,7 @@ 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
|
||||
|
||||
@@ -154,19 +155,8 @@ 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()
|
||||
|
||||
|
||||
@@ -192,8 +182,6 @@ 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 = {
|
||||
@@ -237,7 +225,6 @@ 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,10 +16,9 @@ Usage:
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
@@ -30,16 +29,6 @@ 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)."
|
||||
@@ -98,12 +87,6 @@ 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,
|
||||
@@ -195,10 +178,30 @@ def parse_args() -> argparse.Namespace:
|
||||
help="Enable cache-dit summary logging after diffusion forward passes.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-config",
|
||||
type=parse_profiler_config,
|
||||
"--profile-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
|
||||
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.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
@@ -236,9 +239,23 @@ def main():
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
)
|
||||
|
||||
profiler_enabled = args.profiler_config is not None
|
||||
# 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
|
||||
|
||||
omni_kwargs = dict(
|
||||
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(
|
||||
model=args.model,
|
||||
model_class_name="WanS2VPipeline",
|
||||
flow_shift=args.flow_shift,
|
||||
@@ -252,12 +269,8 @@ def main():
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
enable_cache_dit_summary=args.enable_cache_dit_summary,
|
||||
profiler_config=args.profiler_config,
|
||||
profiler_config=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}")
|
||||
@@ -312,16 +325,13 @@ 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, dict):
|
||||
traces = profile_results.get("traces", [])
|
||||
if profile_results and isinstance(profile_results, list):
|
||||
print("\n" + "=" * 60)
|
||||
print("PROFILING RESULTS:")
|
||||
for rank, trace in enumerate(traces):
|
||||
print(f"\nRank {rank}:")
|
||||
if trace:
|
||||
print(f" • Trace: {trace}")
|
||||
if not traces:
|
||||
print(" No traces collected.")
|
||||
for i, result in enumerate(profile_results):
|
||||
print(f"\nStage {i}:")
|
||||
if result:
|
||||
print(f" {result}")
|
||||
print("=" * 60)
|
||||
else:
|
||||
print("[Profiler] No valid profiling data returned.")
|
||||
|
||||
@@ -307,6 +307,9 @@ def parse_args() -> argparse.Namespace:
|
||||
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()
|
||||
|
||||
|
||||
|
||||
@@ -271,14 +271,6 @@ 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.
|
||||
|
||||
@@ -10,6 +10,7 @@ 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
|
||||
@@ -64,6 +65,7 @@ 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,14 +26,16 @@ 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__)
|
||||
|
||||
@@ -209,7 +211,7 @@ async def main_streaming(args):
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(
|
||||
parser = FlexibleArgumentParser(
|
||||
description="Fish Speech S2 Pro offline inference with vLLM Omni",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -265,6 +267,7 @@ def parse_args():
|
||||
default=False,
|
||||
help="Stream audio chunks as they arrive via AsyncOmni.",
|
||||
)
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -1,38 +0,0 @@
|
||||
# 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/).
|
||||
@@ -1,203 +0,0 @@
|
||||
# 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()
|
||||
@@ -1,204 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline higgs-audio v3 inference example.
|
||||
|
||||
Runs Stage 0 (Qwen3 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 (plain TTS):
|
||||
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
|
||||
--texts "Hello world." \\
|
||||
--output-dir results/higgs_v3_wavs
|
||||
|
||||
Example (voice clone):
|
||||
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
|
||||
--texts "Hello world." \\
|
||||
--ref-audio path/to/reference.wav \\
|
||||
--ref-text "Transcript of the reference clip." \\
|
||||
--output-dir results/higgs_v3_wavs
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
|
||||
os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
|
||||
|
||||
import re
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
from vllm_omni import Omni
|
||||
|
||||
SAMPLE_RATE = 24_000
|
||||
DEFAULT_TEXTS = (
|
||||
"Hello world.",
|
||||
"The quick brown fox jumps over the lazy dog.",
|
||||
"Today is a beautiful day for a walk in the park.",
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Offline higgs-audio v3 inference")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="bosonai/higgs-audio-v3-tts-4b",
|
||||
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/higgs_v3_wavs",
|
||||
help="Directory to write per-prompt WAV files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override the deploy config path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-new-tokens",
|
||||
type=int,
|
||||
default=2048,
|
||||
help="Cap on Stage-0 codec frames per request.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Reference audio for voice clone (WAV/FLAC/MP3 path).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-text",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Transcript of the reference audio. Optional but improves fidelity.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _slugify(text: str) -> str:
|
||||
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:
|
||||
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()
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
|
||||
HiggsAudioV3TokenizerAdapter,
|
||||
apply_delay_pattern,
|
||||
encode_reference_audio,
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
|
||||
adapter = HiggsAudioV3TokenizerAdapter(tokenizer)
|
||||
|
||||
# Load and encode reference audio once (if voice cloning)
|
||||
ref_codes_delayed: torch.Tensor | 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)
|
||||
ref_codes_raw = encode_reference_audio(ref_wav, int(ref_sr))
|
||||
ref_codes_delayed = apply_delay_pattern(ref_codes_raw)
|
||||
print(f"Reference : {args.ref_audio}")
|
||||
print(f"Ref codes : {ref_codes_raw.shape[0]} frames -> {ref_codes_delayed.shape[0]} delayed")
|
||||
if args.ref_text:
|
||||
print(f"Ref text : {args.ref_text}")
|
||||
|
||||
print(f"Model : {args.model}")
|
||||
print(f"Prompts : {len(args.texts)}")
|
||||
print(f"Output dir : {output_dir}")
|
||||
print(f"Voice clone : {'yes' if ref_codes_delayed is not None else 'no'}")
|
||||
|
||||
total_elapsed = 0.0
|
||||
total_dur = 0.0
|
||||
for text in args.texts:
|
||||
if ref_codes_delayed is not None:
|
||||
prompt_ids = adapter.build_prompt(
|
||||
text,
|
||||
num_ref_tokens=int(ref_codes_delayed.shape[0]),
|
||||
reference_text=args.ref_text,
|
||||
)
|
||||
prompt = {
|
||||
"prompt_token_ids": prompt_ids,
|
||||
"additional_information": {
|
||||
"audio_input_ids": ref_codes_delayed.to(torch.long),
|
||||
"audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool),
|
||||
},
|
||||
}
|
||||
else:
|
||||
prompt_ids = adapter.build_prompt(text)
|
||||
prompt = {"prompt_token_ids": prompt_ids}
|
||||
|
||||
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__":
|
||||
main()
|
||||
@@ -5,11 +5,11 @@ from typing import Any
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
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 = TrackingArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
|
||||
parser = FlexibleArgumentParser(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,6 +89,7 @@ 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()
|
||||
|
||||
|
||||
|
||||
@@ -1,377 +0,0 @@
|
||||
"""Offline inference example for MOSS-TTS variants.
|
||||
|
||||
Covers all five models:
|
||||
- OpenMOSS-Team/MOSS-TTS (8B, general TTS)
|
||||
- OpenMOSS-Team/MOSS-TTSD-v1.0 (8B, dialogue TTS)
|
||||
- OpenMOSS-Team/MOSS-SoundEffect (8B, sound effect synthesis)
|
||||
- OpenMOSS-Team/MOSS-TTS-Realtime (1.7B, low-latency streaming)
|
||||
- OpenMOSS-Team/MOSS-VoiceGenerator(1.7B, voice design)
|
||||
|
||||
Usage
|
||||
-----
|
||||
# Standard TTS with voice cloning:
|
||||
python end2end.py \\
|
||||
--model OpenMOSS-Team/MOSS-TTS \\
|
||||
--text "Hello, this is a MOSS-TTS test." \\
|
||||
--ref-audio path/to/reference.wav \\
|
||||
--output output.wav
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import sys
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
|
||||
def _load_ref_audio(path: str, target_sr: int = 24000) -> torch.Tensor:
|
||||
"""Load and resample reference audio to (1, T) at target_sr."""
|
||||
data, sr = sf.read(path, always_2d=True) # (T, C)
|
||||
wav = torch.from_numpy(data.T.astype("float32")) # (C, T)
|
||||
if wav.shape[0] > 1:
|
||||
wav = wav.mean(dim=0, keepdim=True)
|
||||
if sr != target_sr:
|
||||
import torchaudio
|
||||
|
||||
wav = torchaudio.functional.resample(wav, sr, target_sr)
|
||||
return wav # (1, T)
|
||||
|
||||
|
||||
def _build_unified_codes(
|
||||
model_id: str,
|
||||
text: str,
|
||||
ref_audio: torch.Tensor | None,
|
||||
*,
|
||||
instruction: str | None = None,
|
||||
tokens: int | None = None,
|
||||
quality: str | None = None,
|
||||
sound_event: str | None = None,
|
||||
ambient_sound: str | None = None,
|
||||
language: str | None = None,
|
||||
) -> tuple[list[int], torch.Tensor, int]:
|
||||
"""Build the upstream MossTTSDelay unified-codes prompt.
|
||||
|
||||
Returns (text_token_ids, audio_codes_LxNQ, n_vq) where
|
||||
``text_token_ids`` is the prefill prompt for the talker and
|
||||
``audio_codes_LxNQ`` is the matching delay-pattern audio code grid
|
||||
(``audio_pad_code`` everywhere except inside the reference-audio block).
|
||||
"""
|
||||
from transformers import AutoProcessor
|
||||
|
||||
proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
|
||||
|
||||
reference: list[torch.Tensor] | None = None
|
||||
if ref_audio is not None:
|
||||
# Pre-encode so the processor doesn't need torchaudio.load (which
|
||||
# now requires torchcodec on recent torchaudio releases).
|
||||
# encode_audios_from_wav expects (C, T) per item.
|
||||
wav_2d = ref_audio if ref_audio.dim() == 2 else ref_audio.unsqueeze(0)
|
||||
codes_list = proc.encode_audios_from_wav([wav_2d], sampling_rate=24000, n_vq=proc.model_config.n_vq)
|
||||
reference = [codes_list[0]]
|
||||
|
||||
user_msg_kwargs: dict = {"text": text}
|
||||
if reference is not None:
|
||||
user_msg_kwargs["reference"] = reference
|
||||
if instruction is not None:
|
||||
user_msg_kwargs["instruction"] = instruction
|
||||
if tokens is not None:
|
||||
user_msg_kwargs["tokens"] = tokens
|
||||
if quality is not None:
|
||||
user_msg_kwargs["quality"] = quality
|
||||
if sound_event is not None:
|
||||
user_msg_kwargs["sound_event"] = sound_event
|
||||
if ambient_sound is not None:
|
||||
user_msg_kwargs["ambient_sound"] = ambient_sound
|
||||
if language is not None:
|
||||
user_msg_kwargs["language"] = language
|
||||
|
||||
user_msg = proc.build_user_message(**user_msg_kwargs)
|
||||
batch = proc(conversations=[[user_msg]], mode="generation")
|
||||
unified = batch["input_ids"][0] # (L, 1+n_vq)
|
||||
n_vq = int(unified.shape[-1] - 1)
|
||||
text_ids = unified[:, 0].tolist()
|
||||
audio_codes = unified[:, 1:].contiguous().to(torch.int64) # (L, n_vq)
|
||||
|
||||
# Free the processor (incl. ~1.6B-param audio_tokenizer on CPU) before
|
||||
# we instantiate Omni — we only needed it for the one-shot encode.
|
||||
del proc
|
||||
gc.collect()
|
||||
|
||||
return text_ids, audio_codes, n_vq
|
||||
|
||||
|
||||
def _build_realtime_prompt(
|
||||
model_id: str,
|
||||
text: str,
|
||||
ref_audio: torch.Tensor | None,
|
||||
) -> tuple[list[int], torch.Tensor, int]:
|
||||
"""Build the prefill prompt for MOSS-TTS-Realtime.
|
||||
|
||||
Realtime uses a different prompt format than the delay variants: a system
|
||||
prompt + (optional) voice-clone context with the encoded reference audio
|
||||
inlined, followed by the user text and an assistant header. The format is
|
||||
a ``(L, channels+1=17)`` int grid where col 0 is text/special tokens and
|
||||
cols 1..16 are RVQ codebook entries (``audio_channel_pad`` outside the
|
||||
ref-audio block).
|
||||
|
||||
The upstream ``MossTTSRealtimeProcessor`` is shipped via
|
||||
``trust_remote_code`` but is not auto-discovered by ``AutoProcessor``
|
||||
(no ``processor_config.json``) — we import its module directly from the
|
||||
snapshot directory and use its ``make_ensemble``/``make_user_prompt``.
|
||||
"""
|
||||
import importlib.util
|
||||
import sys as _sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from huggingface_hub import snapshot_download
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
snap_dir = Path(snapshot_download(repo_id=model_id))
|
||||
proc_module_path = snap_dir / "processing_mossttsrealtime.py"
|
||||
if not proc_module_path.exists():
|
||||
raise FileNotFoundError(f"realtime processor module missing at {proc_module_path}")
|
||||
|
||||
# Load the processor module without polluting sys.path globally.
|
||||
spec = importlib.util.spec_from_file_location("_moss_tts_realtime_proc", proc_module_path)
|
||||
mod = importlib.util.module_from_spec(spec) # type: ignore[arg-type]
|
||||
_sys.modules[spec.name] = mod # type: ignore[union-attr]
|
||||
spec.loader.exec_module(mod) # type: ignore[union-attr]
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
|
||||
realtime_processor = mod.MossTTSRealtimeProcessor(tokenizer=tokenizer)
|
||||
|
||||
# Encode reference audio via the standalone MOSS Audio Tokenizer (CPU is
|
||||
# fine — it runs once per request).
|
||||
audio_tokens = None
|
||||
if ref_audio is not None:
|
||||
codec = AutoModel.from_pretrained("OpenMOSS-Team/MOSS-Audio-Tokenizer", trust_remote_code=True).eval()
|
||||
wav_2d = ref_audio if ref_audio.dim() == 2 else ref_audio.unsqueeze(0)
|
||||
with torch.no_grad():
|
||||
enc = codec.batch_encode([wav_2d.squeeze(0)], num_quantizers=16)
|
||||
# enc.audio_codes: (NQ, B, T); we want (T, NQ) numpy.
|
||||
codes = enc.audio_codes[:, 0, : int(enc.audio_codes_lengths[0].item())]
|
||||
audio_tokens = codes.transpose(0, 1).contiguous().cpu().numpy() # (T, 16)
|
||||
del codec
|
||||
gc.collect()
|
||||
|
||||
# Build the inference-style prefill (mirrors upstream
|
||||
# ``MossTTSRealtimeInference._build_prefill_batch``):
|
||||
# system_prompt [+ voice_clone_block]
|
||||
# <|im_start|>assistant\n
|
||||
# <text_tokens...> (col 0 = text, cols 1..16 = audio_channel_pad)
|
||||
# <last text token> (col 1 = audio_bos at the last text position)
|
||||
# Generation continues from here; the local transformer emits real audio
|
||||
# codes for codebook 0 directly (no BOS frame at gen_step=0).
|
||||
system_grid = realtime_processor.make_ensemble(prompt_audio_tokens=audio_tokens) # (L_sys, 17)
|
||||
|
||||
# Append "<|im_start|>assistant\n" header (the snapshot's make_ensemble
|
||||
# omits it; inferencer.py's version includes it. We add it here.)
|
||||
assistant_header_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
|
||||
assistant_grid = np.full((len(assistant_header_ids), 17), 1024, dtype=np.int64)
|
||||
assistant_grid[:, 0] = assistant_header_ids
|
||||
|
||||
# Append the FIRST 12 text tokens with audio_bos (1025) at the last
|
||||
# included text position (audio column 0 == grid col 1). The remaining
|
||||
# text tokens are fed one-per-step during generation via the talker's
|
||||
# compute_logits → vLLM sampler pipeline (see ``ids.all`` in the
|
||||
# additional_information). This matches upstream
|
||||
# ``MossTTSRealtimeInference._build_prefill_batch`` (prefill_max_text=12).
|
||||
text_ids_only = tokenizer.encode(text, add_special_tokens=False)
|
||||
if not text_ids_only:
|
||||
raise ValueError("Empty text after tokenisation")
|
||||
PREFILL_MAX_TEXT = 12
|
||||
cur_len = min(len(text_ids_only), PREFILL_MAX_TEXT)
|
||||
text_grid = np.full((cur_len, 17), 1024, dtype=np.int64)
|
||||
text_grid[:, 0] = text_ids_only[:cur_len]
|
||||
text_grid[-1, 1] = 1025 # audio_bos_token at last *prefilled* text token
|
||||
|
||||
grid = np.concatenate([system_grid, assistant_grid, text_grid], axis=0)
|
||||
text_ids = grid[:, 0].tolist()
|
||||
audio_codes = torch.from_numpy(grid[:, 1:].astype(np.int64).copy()) # (L, 16)
|
||||
|
||||
# ``remaining_text_ids`` will be fed one-per-step by the talker.
|
||||
remaining_text_ids = list(text_ids_only[cur_len:])
|
||||
return text_ids, audio_codes, 16, remaining_text_ids
|
||||
|
||||
|
||||
def run_tts(args: argparse.Namespace) -> None:
|
||||
from vllm import SamplingParams
|
||||
|
||||
from vllm_omni import Omni
|
||||
|
||||
deploy_map = {
|
||||
"OpenMOSS-Team/MOSS-TTS": "moss_tts",
|
||||
"OpenMOSS-Team/MOSS-TTSD-v1.0": "moss_ttsd",
|
||||
"OpenMOSS-Team/MOSS-SoundEffect": "moss_sound_effect",
|
||||
"OpenMOSS-Team/MOSS-TTS-Realtime": "moss_tts_realtime",
|
||||
"OpenMOSS-Team/MOSS-VoiceGenerator": "moss_voice_generator",
|
||||
}
|
||||
deploy_config = deploy_map.get(args.model, "moss_tts")
|
||||
|
||||
is_sound_effect = "SoundEffect" in args.model
|
||||
is_voice_gen = "VoiceGenerator" in args.model
|
||||
is_realtime = "Realtime" in args.model
|
||||
|
||||
ref_audio = None
|
||||
builder_kwargs: dict = {}
|
||||
|
||||
remaining_text_ids: list[int] | None = None
|
||||
if is_realtime:
|
||||
# MOSS-TTS-Realtime uses a different prompt format and has its own
|
||||
# per-step depth transformer; the delay-style ``MossTTSDelayProcessor``
|
||||
# path doesn't apply.
|
||||
if not args.ref_audio:
|
||||
sys.exit("MOSS-TTS-Realtime currently requires --ref-audio (voice clone).")
|
||||
ref_audio = _load_ref_audio(args.ref_audio)
|
||||
text = args.text or ""
|
||||
text_ids, audio_codes, n_vq, remaining_text_ids = _build_realtime_prompt(args.model, text, ref_audio)
|
||||
print(
|
||||
f"Prefill prompt: {len(text_ids)} tokens, {n_vq}-quantizer audio block (realtime); "
|
||||
f"{len(remaining_text_ids)} text tokens to stream"
|
||||
)
|
||||
elif is_sound_effect:
|
||||
if not args.ambient_sound:
|
||||
sys.exit("MOSS-SoundEffect requires --ambient-sound")
|
||||
builder_kwargs["ambient_sound"] = args.ambient_sound
|
||||
# MOSS-SoundEffect emits ~12.5 codec frames per second; ``tokens``
|
||||
# tells the AR loop the target duration.
|
||||
if args.duration_tokens:
|
||||
builder_kwargs["tokens"] = args.duration_tokens
|
||||
elif args.duration_seconds:
|
||||
builder_kwargs["tokens"] = max(1, int(float(args.duration_seconds) * 12.5))
|
||||
text = args.text or ""
|
||||
elif is_voice_gen:
|
||||
# MOSS-VoiceGenerator needs an instruction (voice description) plus
|
||||
# the text to synthesise; no reference audio.
|
||||
if not args.instruction:
|
||||
sys.exit("MOSS-VoiceGenerator requires --instruction (voice description)")
|
||||
builder_kwargs["instruction"] = args.instruction
|
||||
text = args.text or ""
|
||||
else:
|
||||
if not args.ref_audio:
|
||||
sys.exit("Voice cloning requires --ref-audio")
|
||||
ref_audio = _load_ref_audio(args.ref_audio)
|
||||
text = args.text or ""
|
||||
|
||||
if not is_realtime:
|
||||
text_ids, audio_codes, n_vq = _build_unified_codes(args.model, text, ref_audio, **builder_kwargs)
|
||||
print(f"Prefill prompt: {len(text_ids)} tokens, {n_vq}-quantizer audio block")
|
||||
|
||||
omni = Omni(model=args.model, deploy_config=deploy_config, stage_init_timeout=600)
|
||||
|
||||
talker_sp = SamplingParams(
|
||||
temperature=1.7,
|
||||
top_p=0.8,
|
||||
top_k=25,
|
||||
max_tokens=2048,
|
||||
seed=42,
|
||||
detokenize=False,
|
||||
)
|
||||
codec_sp = SamplingParams(
|
||||
temperature=0.0,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
max_tokens=65536,
|
||||
seed=42,
|
||||
detokenize=True,
|
||||
)
|
||||
|
||||
additional_info: dict = {
|
||||
# Reference-audio codes consumed by the talker prefill (additive
|
||||
# embedding); shape is (L, n_vq) where pad_code marks non-audio
|
||||
# positions. The OmniPayload schema reserves ``codes.ref`` for
|
||||
# exactly this kind of conditioning input.
|
||||
"codes": {"ref": audio_codes},
|
||||
}
|
||||
if remaining_text_ids:
|
||||
# Realtime variant: feed the remaining text tokens (everything past
|
||||
# the first ``PREFILL_MAX_TEXT``) one-per-step during decode via the
|
||||
# talker's compute_logits → vLLM sampler one-hot trick.
|
||||
additional_info["ids"] = {"all": remaining_text_ids}
|
||||
|
||||
request = {
|
||||
"prompt_token_ids": text_ids,
|
||||
"additional_information": additional_info,
|
||||
}
|
||||
|
||||
import time as _time
|
||||
|
||||
chunks: list[torch.Tensor] = []
|
||||
t_gen_start = _time.monotonic()
|
||||
t_first_audio: float | None = None
|
||||
for stage_outputs in omni.generate(request, [talker_sp, codec_sp]):
|
||||
# Each yielded ``stage_outputs`` is an OmniRequestOutput; the codec
|
||||
# stage attaches its waveform under ``multimodal_output["model_outputs"]``.
|
||||
# Avoid ``a or b`` when ``a`` may be a tensor (truthiness raises).
|
||||
mm = stage_outputs.multimodal_output or {}
|
||||
wav = mm.get("audio")
|
||||
if wav is None:
|
||||
wav = mm.get("model_outputs")
|
||||
if wav is None:
|
||||
continue
|
||||
if isinstance(wav, list):
|
||||
non_empty = [w.reshape(-1) for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0]
|
||||
if non_empty and t_first_audio is None:
|
||||
t_first_audio = _time.monotonic()
|
||||
chunks.extend(non_empty)
|
||||
elif isinstance(wav, torch.Tensor) and wav.numel() > 0:
|
||||
if t_first_audio is None:
|
||||
t_first_audio = _time.monotonic()
|
||||
chunks.append(wav.reshape(-1))
|
||||
|
||||
t_gen_end = _time.monotonic()
|
||||
if not chunks:
|
||||
print("No audio in output — check model loading logs.")
|
||||
return
|
||||
|
||||
wav_t = torch.cat(chunks).float().cpu()
|
||||
sf.write(args.output, wav_t.numpy(), 24000)
|
||||
print(f"Saved {len(wav_t) / 24000:.2f}s of audio to {args.output}")
|
||||
ttfa_ms = (t_first_audio - t_gen_start) * 1000.0 if t_first_audio is not None else float("nan")
|
||||
gen_total_ms = (t_gen_end - t_gen_start) * 1000.0
|
||||
print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={len(wav_t) / 24000:.2f}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="MOSS-TTS offline inference example")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="OpenMOSS-Team/MOSS-TTS",
|
||||
help="HuggingFace model ID",
|
||||
)
|
||||
parser.add_argument("--text", default="", help="Text to synthesise (TTS variants)")
|
||||
parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning")
|
||||
parser.add_argument(
|
||||
"--ambient-sound",
|
||||
default=None,
|
||||
help="Sound description for MOSS-SoundEffect",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--duration-tokens",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Target duration in tokens (1 s ≈ 12.5 tokens) for MOSS-SoundEffect",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--duration-seconds",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Target duration in seconds (converted via ~12.5 tokens/s) for MOSS-SoundEffect",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--instruction",
|
||||
default=None,
|
||||
help="Voice description for MOSS-VoiceGenerator (e.g. 'a young woman with a clear voice')",
|
||||
)
|
||||
parser.add_argument("--output", default="output.wav", help="Output WAV path")
|
||||
args = parser.parse_args()
|
||||
run_tts(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -38,8 +38,7 @@ from pathlib import Path
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from vllm import SamplingParams
|
||||
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
# Prevent multiprocessing from re-importing CUDA in the wrong context.
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
@@ -156,7 +155,7 @@ def main(args) -> None:
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="MOSS-TTS-Nano offline inference")
|
||||
parser = FlexibleArgumentParser(description="MOSS-TTS-Nano offline inference")
|
||||
parser.add_argument("--text", default="Hello, this is MOSS-TTS-Nano speaking.", help="Text to synthesize.")
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
|
||||
@@ -21,6 +21,7 @@ import os
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
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
|
||||
|
||||
@@ -79,12 +80,7 @@ def run_e2e():
|
||||
default=600,
|
||||
help="Stage initialization timeout in seconds",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Random seed for generation",
|
||||
)
|
||||
nullify_stage_engine_defaults(parser)
|
||||
args = parser.parse_args()
|
||||
|
||||
if not os.path.exists(args.stage_config):
|
||||
@@ -126,7 +122,7 @@ def run_e2e():
|
||||
if mm_processor_kwargs:
|
||||
prompts["mm_processor_kwargs"] = mm_processor_kwargs
|
||||
|
||||
sampling_params_list = [OmniDiffusionSamplingParams(extra_args={"seed": args.seed})]
|
||||
sampling_params_list = [OmniDiffusionSamplingParams()]
|
||||
|
||||
print(f"Generating speech for: {args.text}")
|
||||
|
||||
|
||||
@@ -15,8 +15,10 @@ 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.utils.tracking_parser import TrackingArgumentParser
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -41,8 +43,8 @@ def _estimate_prompt_len(
|
||||
"""
|
||||
try:
|
||||
from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
|
||||
from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
|
||||
Qwen3TTSPromptEmbedsBuilder,
|
||||
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
|
||||
Qwen3TTSTalkerForConditionalGeneration,
|
||||
)
|
||||
|
||||
if model_name not in _cache:
|
||||
@@ -122,7 +124,7 @@ def _estimate_prompt_len(
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
|
||||
return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
|
||||
additional_information=additional_information,
|
||||
task_type=task_type,
|
||||
tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
|
||||
@@ -433,7 +435,7 @@ async def main_streaming(args):
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
@@ -546,6 +548,7 @@ def parse_args():
|
||||
help="Number of prompts per batch (default: 1, sequential).",
|
||||
)
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -12,16 +12,17 @@ from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
SAMPLE_RATE = 48_000
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="Offline VoxCPM2 native AR inference")
|
||||
parser = FlexibleArgumentParser(description="Offline VoxCPM2 native AR inference")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
@@ -59,6 +60,7 @@ def parse_args():
|
||||
default=None,
|
||||
help="Optional transcript of --ref-audio (enables continuation mode).",
|
||||
)
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -27,10 +27,11 @@ except ImportError:
|
||||
"pip install -e ."
|
||||
)
|
||||
from vllm import SamplingParams
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import AsyncOmni
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -229,7 +230,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
|
||||
|
||||
|
||||
def parse_args() -> Namespace:
|
||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
|
||||
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
@@ -312,6 +313,7 @@ def parse_args() -> Namespace:
|
||||
default=None,
|
||||
help="Quantization method (e.g. 'fp8'). Applied to the language model only.",
|
||||
)
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
Loaded 100 of 678 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user