Merge branch 'main' into feat/moss-tts-nano
Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
No files matched your search
@@ -44,11 +44,19 @@ steps:
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L4 Test — main+NIGHTLY=1 (scheduled), or PR with label nightly-test (e.g. add label then Rebuild)
|
||||
# L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
|
||||
- label: "Upload Nightly Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-nightly-pipeline
|
||||
if: '(build.branch == "main" && build.env("NIGHTLY") == "1") || (build.branch != "main" && build.pull_request.labels includes "nightly-test")'
|
||||
if: >-
|
||||
(build.branch == "main" && build.env("NIGHTLY") == "1") ||
|
||||
(build.branch != "main" && (
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "omni-test" ||
|
||||
build.pull_request.labels includes "tts-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2iat-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2v-test"
|
||||
))
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
|
||||
agents:
|
||||
|
||||
@@ -64,12 +64,13 @@ while true; do
|
||||
done
|
||||
|
||||
echo "--- Pulling container"
|
||||
## Temporary change to use AMD Docker Hub to store the vllm-ci image
|
||||
## Temporary change to use AMD Docker Hub to store the vllm-omni image
|
||||
# to bypass the rate limit issue with ECR Public Gallery.
|
||||
# Images are now stored in a separate repository for vllm-omni, instead of vllm-ci.
|
||||
# TODO: @tjtanaa point back to ECR Public Gallery
|
||||
# once the amd agents are configured to use ECR Public Gallery.
|
||||
# image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
|
||||
image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}-rocm-omni"
|
||||
image_name="rocm/vllm-omni:${BUILDKITE_COMMIT}"
|
||||
container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
|
||||
|
||||
# TODO: @tjtanaa uncomment this once the amd agents are configured to use ECR Public Gallery.
|
||||
|
||||
@@ -117,3 +117,17 @@ steps:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
|
||||
|
||||
|
||||
- label: "Omni Sleep Mode Test"
|
||||
timeout_in_minutes: 40
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- pytest -s -v tests/e2e/offline_inference/test_omni_sleep_mode.py -m "advanced_model and omni and MI325" --run-level "advanced_model"
|
||||
@@ -382,6 +382,46 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni Sleep Mode Test with H100"
|
||||
timeout_in_minutes: 30
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- pytest -s -v tests/e2e/offline_inference/test_omni_sleep_mode.py -m "advanced_model and H100 and omni" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Voxtral-TTS E2E Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-merge-pipeline
|
||||
|
||||
@@ -1,377 +0,0 @@
|
||||
# Nightly diffusion GPU tests — appended to the main nightly build via
|
||||
# buildkite-agent pipeline upload .buildkite/test-nightly-diffusion.yml
|
||||
# from test-nightly.yml (step key: nightly-diffusion-model-test). Top-level groups are
|
||||
# foldable in the Buildkite UI (Other / Wan / Qwen-Image).
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HF_HUB_DOWNLOAD_TIMEOUT: 300
|
||||
HF_HUB_ETAG_TIMEOUT: 60
|
||||
|
||||
steps:
|
||||
- group: ":card_index_dividers: Other Model Test"
|
||||
key: nightly-other-model-test-group
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion · Other · Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
# Shared nightly vs PR label conditional; referenced below as *nightly_or_pr_label
|
||||
if: &nightly_or_pr_label 'build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"'
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not test_qwen_image" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion · Other · Function Test with L4"
|
||||
timeout_in_minutes: 60
|
||||
if: *nightly_or_pr_label
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
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: Diffusion · Other · Doc Test"
|
||||
timeout_in_minutes: 60
|
||||
if: *nightly_or_pr_label
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- pytest -s -v tests/examples/online_serving/test_text_to_image.py tests/examples/offline_inference/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- group: ":card_index_dividers: Wan Series Model Test"
|
||||
key: nightly-wan-model-test-group
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion · Wan · Function Test"
|
||||
timeout_in_minutes: 90
|
||||
if: *nightly_or_pr_label
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py -m "advanced_model" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 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 · Wan · Accuracy Test"
|
||||
key: nightly-wan22-i2v-accuracy
|
||||
timeout_in_minutes: 180
|
||||
if: *nightly_or_pr_label
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --run-level advanced_model
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- group: ":card_index_dividers: Qwen-Image Series Model Test"
|
||||
key: nightly-qwen-image-edit-group
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion · Qwen-Image · Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
if: *nightly_or_pr_label
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen_image*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion · Qwen-Image · GEBench Accuracy Test"
|
||||
key: nightly-gebench-accuracy
|
||||
timeout_in_minutes: 60
|
||||
if: *nightly_or_pr_label
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
|
||||
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: Diffusion · Qwen-Image · GEdit-Bench Accuracy Test"
|
||||
key: nightly-gedit-bench-accuracy
|
||||
timeout_in_minutes: 60
|
||||
if: *nightly_or_pr_label
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
|
||||
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: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
|
||||
value: "120"
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion · Qwen-Image · Perf Test"
|
||||
key: nightly-qwen-image-performance
|
||||
timeout_in_minutes: 180
|
||||
if: *nightly_or_pr_label
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
# [HACK]: run upload in the same command block as pytest.
|
||||
# Because `exit` aborts the entire commands list.
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
|
||||
EXIT2=$$?
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json
|
||||
EXIT3=$$?
|
||||
if [ $$EXIT1 -eq 0 ] || [ $$EXIT2 -eq 0 ] || [ $$EXIT3 -eq 0 ]; then
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
fi
|
||||
exit $$((EXIT1 | EXIT2 | EXIT3))
|
||||
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
|
||||
@@ -7,12 +7,11 @@ steps:
|
||||
# Group: collapses under one heading in the Buildkite UI; child steps still run in parallel.
|
||||
- group: ":card_index_dividers: Omni Model Test"
|
||||
key: nightly-omni-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test"
|
||||
steps:
|
||||
- label: ":full_moon: Omni · Function Test with H100"
|
||||
- label: ":full_moon: Omni · Function Test"
|
||||
timeout_in_minutes: 90
|
||||
depends_on: upload-nightly-pipeline
|
||||
# Shared nightly vs PR label conditional; referenced below as *nightly_or_pr_label
|
||||
if: &nightly_or_pr_label 'build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"'
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and H100 and omni" --run-level "advanced_model"
|
||||
agents:
|
||||
@@ -49,31 +48,8 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Omni · Function Test with L4"
|
||||
timeout_in_minutes: 90
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: *nightly_or_pr_label
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
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: Omni · Doc Test with L4"
|
||||
timeout_in_minutes: 90
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: *nightly_or_pr_label
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/examples/ -m "advanced_model and omni and L4" --run-level "advanced_model"
|
||||
@@ -93,8 +69,6 @@ steps:
|
||||
|
||||
- label: ":full_moon: Omni · Doc Test with H100"
|
||||
timeout_in_minutes: 90
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: *nightly_or_pr_label
|
||||
commands:
|
||||
- pytest -s -v tests/examples/ -m "advanced_model and omni and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
@@ -134,13 +108,14 @@ steps:
|
||||
- label: ":full_moon: Omni · Perf Test"
|
||||
key: nightly-omni-performance
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: *nightly_or_pr_label
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -175,23 +150,448 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# Dynamically appends steps from test-nightly-diffusion.yml into this build (same mechanism as
|
||||
# pipeline.yml → test-ready.yml / test-merge.yml / test-nightly.yml). Foldable groups stay in the
|
||||
# uploaded YAML (Other / Wan / Qwen-Image).
|
||||
- label: ":card_index_dividers: Diffusion Model Test"
|
||||
key: nightly-diffusion-model-test
|
||||
|
||||
- group: ":card_index_dividers: TTS Model Test"
|
||||
key: nightly-tts-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: *nightly_or_pr_label
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly-diffusion.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
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"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
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
|
||||
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_tts.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
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: >-
|
||||
build.env("NIGHTLY") == "1" ||
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2iat-test"
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not hunyuan" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not hunyuan" -m "advanced_model and diffusion and L4" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
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: Diffusion X2I(&A&T) · Doc Test"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- pytest -s -v tests/examples/*/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · GEBench Accuracy Test"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
|
||||
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: Diffusion X2I(&A&T) · GEdit-Bench Accuracy Test"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
|
||||
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: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
|
||||
value: "120"
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level advanced_model
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
|
||||
value: "120"
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test"
|
||||
key: nightly-diffusion-x2iat-performance
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
# [HACK]: run upload in the same command block as pytest.
|
||||
# Because `exit` aborts the entire commands list.
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
|
||||
EXIT2=$$?
|
||||
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
|
||||
EXIT3=$$?
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
|
||||
EXIT4=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$((EXIT1 | EXIT2 | EXIT3 | EXIT4))
|
||||
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
|
||||
|
||||
# Diffusion x2v only (Wan, HunyuanVideo, …). x2i/x2a/x2t live in the X2I group above, not here.
|
||||
- group: ":card_index_dividers: Diffusion X2V Model Test"
|
||||
key: nightly-diffusion-x2v-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: >-
|
||||
build.env("NIGHTLY") == "1" ||
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2v-test"
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion X2V · Function Test"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "advanced_model" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 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 X2V · Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m advanced_model --run-level advanced_model
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 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: ":bar_chart: Testcase Statistics"
|
||||
key: nightly-testcase-statistics
|
||||
timeout_in_minutes: 120
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: *nightly_or_pr_label
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
|
||||
@@ -234,15 +634,17 @@ steps:
|
||||
key: nightly-perf-distribution
|
||||
depends_on:
|
||||
- nightly-omni-performance
|
||||
- nightly-qwen-image-performance
|
||||
- nightly-tts-performance
|
||||
- nightly-diffusion-x2iat-performance
|
||||
- nightly-testcase-statistics
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- pip install openpyxl
|
||||
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
|
||||
- export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-qwen-image-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
|
||||
- python tools/nightly/generate_nightly_perf_excel.py
|
||||
- python tools/nightly/generate_nightly_perf_html.py
|
||||
|
||||
@@ -295,6 +295,31 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "VoxCPM E2E Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
pip install voxcpm
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxcpm.py -m "core_model" --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: "VoxCPM2 Native AR E2E Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-ready-pipeline
|
||||
@@ -342,6 +367,33 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Qwen3-TTS Base E2E Test (ModelRunner V2)"
|
||||
depends_on: upload-ready-pipeline
|
||||
soft_fail:
|
||||
- exit_status: 1
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
export VLLM_OMNI_USE_V2_RUNNER="1"
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m "core_model" --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: "Voxtral-TTS E2E Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-ready-pipeline
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
Last synced: 2025-12-15
|
||||
Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests
|
||||
#}
|
||||
{% set docker_image_amd = "rocm/vllm-ci:$BUILDKITE_COMMIT-rocm-omni" %}
|
||||
{% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %}
|
||||
{% set default_working_dir = "/app/vllm-omni" %}
|
||||
|
||||
- group: "AMD Tests"
|
||||
|
||||
@@ -0,0 +1,300 @@
|
||||
---
|
||||
name: vllm-omni-npu-model-runner-upgrade
|
||||
description: "Upgrade vllm-omni NPU model runners (OmniNPUModelRunner, NPUARModelRunner, NPUGenerationModelRunner) to align with the latest vllm-ascend NPUModelRunner while preserving omni-specific logic."
|
||||
---
|
||||
|
||||
# vLLM-Omni NPU Model Runner Upgrade Skill
|
||||
|
||||
## Overview
|
||||
|
||||
This skill guides the process of upgrading vllm-omni's NPU model runners to align with the latest vllm-ascend codebase while preserving omni-specific enhancements. The NPU runners are designed to run omni multimodal models (like Qwen3-Omni, Bagel, MiMoAudio) on Ascend NPUs.
|
||||
|
||||
## File Structure
|
||||
|
||||
### NPU Model Runner Files
|
||||
```
|
||||
vllm-omni/vllm_omni/platforms/npu/worker/
|
||||
├── __init__.py
|
||||
├── npu_model_runner.py # OmniNPUModelRunner (base class)
|
||||
├── npu_ar_model_runner.py # NPUARModelRunner (autoregressive)
|
||||
├── npu_ar_worker.py # AR worker
|
||||
├── npu_generation_model_runner.py # NPUGenerationModelRunner (diffusion/non-AR)
|
||||
└── npu_generation_worker.py # Generation worker
|
||||
```
|
||||
|
||||
### GPU Reference Files (for omni-specific logic sync)
|
||||
```
|
||||
vllm-omni/vllm_omni/worker/
|
||||
├── __init__.py
|
||||
├── gpu_model_runner.py # OmniGPUModelRunner
|
||||
├── gpu_ar_model_runner.py # GPUARModelRunner
|
||||
├── gpu_ar_worker.py
|
||||
├── gpu_generation_model_runner.py
|
||||
├── gpu_generation_worker.py
|
||||
├── mixins.py
|
||||
├── base.py
|
||||
└── gpu_memory_utils.py
|
||||
```
|
||||
|
||||
### vllm-ascend Reference Files
|
||||
```
|
||||
vllm-ascend/vllm_ascend/worker/
|
||||
├── model_runner_v1.py # NPUModelRunner (base class to copy from)
|
||||
├── npu_input_batch.py
|
||||
├── block_table.py
|
||||
├── pcp_utils.py
|
||||
└── worker.py
|
||||
```
|
||||
|
||||
## Inheritance Hierarchy
|
||||
|
||||
```
|
||||
GPUModelRunner (vllm)
|
||||
|
|
||||
+----------------+----------------+
|
||||
| |
|
||||
OmniGPUModelRunner NPUModelRunner (vllm-ascend)
|
||||
(vllm_omni/worker) (vllm_ascend/worker)
|
||||
| |
|
||||
+----------- OmniNPUModelRunner --+
|
||||
(multiple inheritance)
|
||||
|
|
||||
+---------------+---------------+
|
||||
| |
|
||||
NPUARModelRunner NPUGenerationModelRunner
|
||||
(autoregressive) (non-autoregressive/diffusion)
|
||||
```
|
||||
|
||||
## Omni-Specific Comment Markers
|
||||
|
||||
Omni-specific logic is marked with comment blocks:
|
||||
```python
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
# ... omni-specific code ...
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
```
|
||||
|
||||
Or simpler variations:
|
||||
```python
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
```
|
||||
|
||||
**Important**:
|
||||
- Always preserve and add these markers when modifying code.
|
||||
- **The reference documents (`references/omni-specific-blocks.md`) may not be up-to-date.** Always grep for `Omni-new` in the GPU implementations to find the authoritative list of omni-specific blocks.
|
||||
- When you discover new omni-specific code that is not documented in the references, please update the reference files.
|
||||
|
||||
## Key Methods Requiring Attention
|
||||
|
||||
### OmniNPUModelRunner (npu_model_runner.py)
|
||||
|
||||
| Method | Description | Omni-Specific Logic |
|
||||
|--------|-------------|---------------------|
|
||||
| `load_model` | Load model and initialize talker_mtp | Uses `ACLGraphWrapper` instead of `CUDAGraphWrapper`, initializes talker buffers |
|
||||
| `_dummy_run` | Warmup/profiling run | talker_mtp dummy forward, `extract_multimodal_outputs` |
|
||||
| `_model_forward` | Forward pass wrapper | Injects `model_kwargs_extra`, wraps with `OmniOutput`, NPU-specific graph updates |
|
||||
| `_talker_mtp_forward` | Talker MTP forward for Qwen3-Omni | Uses `set_ascend_forward_context` |
|
||||
|
||||
### NPUARModelRunner (npu_ar_model_runner.py)
|
||||
|
||||
| Method | Description | Omni-Specific Logic |
|
||||
|--------|-------------|---------------------|
|
||||
| `__init__` | Initialize with KV transfer manager | `OmniKVTransferManager` setup |
|
||||
| `execute_model` | Main inference entry | KV transfer handling, `_update_states` override, `extract_multimodal_outputs` |
|
||||
| `sample_tokens` | Token sampling | Hidden states extraction, multimodal outputs processing, `OmniModelRunnerOutput` |
|
||||
| `_resolve_global_request_id` | Request ID resolution | For disaggregated inference |
|
||||
|
||||
### NPUGenerationModelRunner (npu_generation_model_runner.py)
|
||||
|
||||
| Method | Description | Omni-Specific Logic |
|
||||
|--------|-------------|---------------------|
|
||||
| `_update_request_states` | Update request states for async chunk | async_chunk handling |
|
||||
| `execute_model` | Generation forward | async_chunk, `seq_token_counts`, `_run_generation_model` |
|
||||
| `sample_tokens` | Output processing | multimodal output packaging to `OmniModelRunnerOutput` |
|
||||
| `_dummy_run` | Dummy run override | model_kwargs initialization, multimodal extraction |
|
||||
| `_run_generation_model` | Run generation model | Calls `_model_forward` with sampler |
|
||||
|
||||
## Upgrade Workflow
|
||||
|
||||
### Step 1: Preparation
|
||||
|
||||
1. **Identify target versions**(Use gh cli to check):
|
||||
- We're using vllm-omni main branch
|
||||
- Check the last release of vllm-omni
|
||||
- Target vllm-ascend version(Just directly use the local latest vllm-ascend code)
|
||||
|
||||
2. **Check GPU-side changes** (since last release):
|
||||
```bash
|
||||
cd /root/vllm-workspace/vllm-omni
|
||||
git log --oneline --since="<last-release-date>" -- vllm_omni/worker/
|
||||
```
|
||||
|
||||
3. **Read latest vllm-ascend code**:
|
||||
- We don't track vllm-ascend changes - just directly use the latest code from `/root/vllm-workspace/vllm-ascend/vllm_ascend/worker/model_runner_v1.py`
|
||||
- Copy the relevant methods and re-insert omni-specific blocks
|
||||
|
||||
### Step 2: Analyze Omni-Specific Logic
|
||||
|
||||
For each NPU model runner file:
|
||||
|
||||
1. **Extract existing omni-specific blocks**:
|
||||
```bash
|
||||
grep -n "Omni-new" vllm_omni/platforms/npu/worker/npu_model_runner.py
|
||||
```
|
||||
|
||||
2. **Document each omni block**:
|
||||
- Which method it belongs to
|
||||
- What functionality it provides
|
||||
- Dependencies on other omni code
|
||||
|
||||
### Step 3: Update Base Class (OmniNPUModelRunner)
|
||||
|
||||
**Note**: Always check the GPU implementation `gpu_model_runner.py` for any new omni logic not yet documented in references.
|
||||
|
||||
1. **Read the latest vllm-ascend `NPUModelRunner.load_model`**
|
||||
2. **Copy the method, keeping the structure**
|
||||
3. **Re-insert omni-specific logic** (check GPU `gpu_model_runner.py` for authoritative list):
|
||||
- Replace `CUDAGraphWrapper` with `ACLGraphWrapper`
|
||||
- Keep talker_mtp initialization
|
||||
- Preserve buffer allocations for talker
|
||||
- Check for any new omni blocks added since last sync
|
||||
|
||||
4. **Update `_dummy_run`**:
|
||||
- Copy from vllm-ascend
|
||||
- Compare with GPU `_dummy_run` for omni-specific blocks
|
||||
- Re-insert all `Omni-new` marked code from GPU version
|
||||
|
||||
5. **Update `_model_forward`**:
|
||||
- Keep the omni wrapper logic
|
||||
- Update NPU-specific parts (graph params, SP all-gather)
|
||||
- Check GPU version for any new omni logic
|
||||
|
||||
### Step 4: Update AR Model Runner
|
||||
|
||||
1. **Compare with GPU `gpu_ar_model_runner.py`** for any new omni features
|
||||
2. **Copy `execute_model` from vllm-ascend**
|
||||
3. **Re-insert omni blocks** (reference `references/omni-specific-blocks.md`, but note it may be incomplete):
|
||||
- **IMPORTANT**: Always check the GPU implementation `gpu_ar_model_runner.py` for all `Omni-new` marked code blocks
|
||||
- The reference doc may not include newly added omni logic - treat it as a starting point, not exhaustive
|
||||
- When discovering new omni code blocks, please update `references/omni-specific-blocks.md`
|
||||
- Common omni blocks include but are not limited to: KV transfer, multimodal outputs, sampling_metadata handling, etc.
|
||||
|
||||
4. **Update `sample_tokens`** (also compare with GPU implementation):
|
||||
- Compare with `gpu_ar_model_runner.py`'s `sample_tokens` method
|
||||
- Identify all `Omni-new` marked code blocks
|
||||
- Ensure NPU version includes all omni-specific logic
|
||||
|
||||
### Step 5: Update Generation Model Runner
|
||||
|
||||
**Note**: Generation model runner may have unique omni logic for diffusion/non-AR models.
|
||||
|
||||
1. **Compare with GPU `gpu_generation_model_runner.py`** - grep for all `Omni-new` blocks
|
||||
2. **Update `execute_model`**:
|
||||
- Check GPU version for all omni-specific blocks
|
||||
- Keep async_chunk handling
|
||||
- Keep `seq_token_counts` injection
|
||||
- Update forward/context setup from vllm-ascend
|
||||
- Look for any new omni logic not documented in references
|
||||
|
||||
3. **Update `_dummy_run`**:
|
||||
- Copy from vllm-ascend base
|
||||
- Compare with GPU `_dummy_run` if exists
|
||||
- Re-insert all omni-specific logic
|
||||
|
||||
### Step 6: Update Imports
|
||||
|
||||
Check and update imports at the top of each file:
|
||||
|
||||
```python
|
||||
# Common vllm-ascend imports
|
||||
from vllm_ascend.ascend_forward_context import get_forward_context, set_ascend_forward_context
|
||||
from vllm_ascend.attention.attention_v1 import AscendAttentionState
|
||||
from vllm_ascend.attention.utils import using_paged_attention
|
||||
from vllm_ascend.compilation.acl_graph import ACLGraphWrapper, update_full_graph_params
|
||||
from vllm_ascend.ops.rotary_embedding import update_cos_sin
|
||||
from vllm_ascend.utils import enable_sp, lmhead_tp_enable
|
||||
from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE, NPUModelRunner
|
||||
|
||||
# Omni-specific imports
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
|
||||
from vllm_omni.outputs import OmniModelRunnerOutput
|
||||
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
|
||||
```
|
||||
|
||||
### Step 7: Sync GPU-Side Omni Changes
|
||||
|
||||
1. **Check recent GPU worker changes**:
|
||||
```bash
|
||||
git diff <from-tag>..<to-tag> -- vllm_omni/worker/gpu_model_runner.py
|
||||
git diff <from-tag>..<to-tag> -- vllm_omni/worker/gpu_ar_model_runner.py
|
||||
```
|
||||
|
||||
2. **Identify new omni features** that need to be ported to NPU
|
||||
|
||||
3. **Apply corresponding changes** to NPU runners
|
||||
|
||||
### Step 8: Validation
|
||||
|
||||
1. **Run type checking**:
|
||||
```bash
|
||||
cd /root/vllm-workspace/vllm-omni
|
||||
python -m py_compile vllm_omni/platforms/npu/worker/npu_model_runner.py
|
||||
python -m py_compile vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
|
||||
python -m py_compile vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
|
||||
```
|
||||
|
||||
2. **Run import test**:
|
||||
```bash
|
||||
python -c "from vllm_omni.platforms.npu.worker import *"
|
||||
```
|
||||
|
||||
3. **Run model serving test** (if hardware available):
|
||||
```bash
|
||||
vllm serve <model-path> --trust-remote-code
|
||||
```
|
||||
|
||||
## Common Pitfalls
|
||||
|
||||
### 1. Forward Context Differences
|
||||
- GPU uses `set_forward_context`
|
||||
- NPU uses `set_ascend_forward_context`
|
||||
- Parameters may differ slightly
|
||||
|
||||
### 2. Graph Wrapper Differences
|
||||
- GPU: `CUDAGraphWrapper`
|
||||
- NPU: `ACLGraphWrapper`
|
||||
- Constructor parameters may differ
|
||||
|
||||
### 3. Buffer Creation
|
||||
- GPU: `_make_buffer` returns different structure
|
||||
- NPU: May need numpy=True/False parameter
|
||||
|
||||
### 4. Attention Metadata
|
||||
- GPU: Uses vllm attention metadata builders
|
||||
- NPU: Uses `AscendCommonAttentionMetadata`
|
||||
|
||||
### 5. Sampling
|
||||
- GPU: Uses vllm sampler
|
||||
- NPU: Uses `AscendSampler`
|
||||
|
||||
## Checklist Before Commit
|
||||
|
||||
- [ ] All omni-specific comment markers preserved
|
||||
- [ ] New omni logic from GPU side synced
|
||||
- [ ] Imports updated to latest vllm-ascend
|
||||
- [ ] No `CUDAGraphWrapper` references in NPU code
|
||||
- [ ] `set_ascend_forward_context` used instead of `set_forward_context`
|
||||
- [ ] `ACLGraphWrapper` used for talker_mtp wrapping
|
||||
- [ ] Type hints match vllm-ascend signatures
|
||||
- [ ] No duplicate code blocks
|
||||
- [ ] Python syntax valid (py_compile passes)
|
||||
|
||||
## Reference Files for Comparison
|
||||
|
||||
When upgrading, keep these files open for reference:
|
||||
|
||||
1. **vllm-ascend NPUModelRunner**: `/root/vllm-workspace/vllm-ascend/vllm_ascend/worker/model_runner_v1.py`
|
||||
2. **vllm GPUModelRunner**: `/root/vllm-workspace/vllm/vllm/v1/worker/gpu_model_runner.py`
|
||||
3. **vllm-omni OmniGPUModelRunner**: `/root/vllm-workspace/vllm-omni/vllm_omni/worker/gpu_model_runner.py`
|
||||
@@ -0,0 +1,335 @@
|
||||
# GPU to NPU Translation Patterns
|
||||
|
||||
This document provides a quick reference for translating GPU code patterns to NPU equivalents when porting omni-specific logic.
|
||||
|
||||
## Import Translations
|
||||
|
||||
### Forward Context
|
||||
```python
|
||||
# GPU
|
||||
from vllm.forward_context import set_forward_context
|
||||
|
||||
# NPU
|
||||
from vllm_ascend.ascend_forward_context import set_ascend_forward_context
|
||||
```
|
||||
|
||||
### Graph Wrapper
|
||||
```python
|
||||
# GPU
|
||||
from vllm.compilation.cuda_graph import CUDAGraphWrapper
|
||||
|
||||
# NPU
|
||||
from vllm_ascend.compilation.acl_graph import ACLGraphWrapper
|
||||
```
|
||||
|
||||
### Attention State
|
||||
```python
|
||||
# GPU (no equivalent - uses FlashAttention states directly)
|
||||
|
||||
# NPU
|
||||
from vllm_ascend.attention.attention_v1 import AscendAttentionState
|
||||
```
|
||||
|
||||
### Utilities
|
||||
```python
|
||||
# GPU
|
||||
# (directly use torch.cuda functions)
|
||||
|
||||
# NPU
|
||||
from vllm_ascend.utils import enable_sp, lmhead_tp_enable
|
||||
from vllm_ascend.ops.rotary_embedding import update_cos_sin
|
||||
```
|
||||
|
||||
## Context Manager Translations
|
||||
|
||||
### Forward Context Setup
|
||||
```python
|
||||
# GPU
|
||||
with set_forward_context(
|
||||
attn_metadata,
|
||||
self.vllm_config,
|
||||
num_tokens=num_tokens_padded,
|
||||
num_tokens_across_dp=num_tokens_across_dp,
|
||||
cudagraph_runtime_mode=cudagraph_mode,
|
||||
batch_descriptor=batch_desc,
|
||||
):
|
||||
# forward pass
|
||||
|
||||
# NPU
|
||||
with set_ascend_forward_context(
|
||||
attn_metadata,
|
||||
self.vllm_config,
|
||||
num_tokens=num_tokens_padded,
|
||||
num_tokens_across_dp=num_tokens_across_dp,
|
||||
aclgraph_runtime_mode=cudagraph_mode, # Note: 'aclgraph' not 'cudagraph'
|
||||
batch_descriptor=batch_desc,
|
||||
num_actual_tokens=scheduler_output.total_num_scheduled_tokens,
|
||||
model_instance=self.model,
|
||||
):
|
||||
# forward pass
|
||||
```
|
||||
|
||||
### Graph Capture Context
|
||||
```python
|
||||
# GPU
|
||||
from vllm.compilation.cuda_graph import graph_capture as cuda_graph_capture
|
||||
with cuda_graph_capture(self.device):
|
||||
# capture
|
||||
|
||||
# NPU
|
||||
from vllm_ascend.worker.model_runner_v1 import graph_capture
|
||||
with graph_capture(self.device):
|
||||
# capture
|
||||
```
|
||||
|
||||
## Graph Wrapper Usage
|
||||
|
||||
### Creating Graph Wrapper
|
||||
```python
|
||||
# GPU
|
||||
if cudagraph_mode.has_full_cudagraphs() and has_separate_talker:
|
||||
self.talker_mtp = CUDAGraphWrapper(
|
||||
talker_mtp,
|
||||
self.vllm_config,
|
||||
runtime_mode=CUDAGraphMode.FULL
|
||||
)
|
||||
|
||||
# NPU
|
||||
if cudagraph_mode.has_full_cudagraphs() and has_separate_talker:
|
||||
self.talker_mtp = ACLGraphWrapper(
|
||||
talker_mtp,
|
||||
self.vllm_config,
|
||||
runtime_mode=CUDAGraphMode.FULL
|
||||
)
|
||||
```
|
||||
|
||||
### Checking Graph Wrapper Type
|
||||
```python
|
||||
# GPU
|
||||
if not isinstance(self.talker_mtp, CUDAGraphWrapper):
|
||||
_cudagraph_mode = CUDAGraphMode.NONE
|
||||
|
||||
# NPU
|
||||
if not isinstance(self.talker_mtp, ACLGraphWrapper):
|
||||
_cudagraph_mode = CUDAGraphMode.NONE
|
||||
```
|
||||
|
||||
## Device Operations
|
||||
|
||||
### Synchronization
|
||||
```python
|
||||
# GPU
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# NPU
|
||||
torch.npu.synchronize()
|
||||
```
|
||||
|
||||
### Stream Operations
|
||||
```python
|
||||
# GPU
|
||||
stream = torch.cuda.Stream(device=device)
|
||||
torch.cuda.current_stream()
|
||||
|
||||
# NPU
|
||||
stream = torch.npu.Stream(device=device)
|
||||
torch.npu.current_stream()
|
||||
```
|
||||
|
||||
## Attention Metadata
|
||||
|
||||
### State Setting (NPU-specific)
|
||||
```python
|
||||
# GPU - handled internally by attention backends
|
||||
|
||||
# NPU - explicit state setting required
|
||||
self.attn_state = AscendAttentionState.DecodeOnly
|
||||
if self.speculative_config and self.speculative_config.method == "mtp":
|
||||
if self.vllm_config.model_config.use_mla:
|
||||
self.attn_state = AscendAttentionState.SpecDecoding
|
||||
else:
|
||||
self.attn_state = AscendAttentionState.ChunkedPrefill
|
||||
```
|
||||
|
||||
### Building Attention Metadata
|
||||
```python
|
||||
# GPU - uses vllm attention builders
|
||||
|
||||
# NPU - may need additional parameters
|
||||
(attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata(
|
||||
num_tokens=num_tokens_unpadded,
|
||||
num_tokens_padded=num_tokens_padded,
|
||||
num_reqs=num_reqs,
|
||||
num_reqs_padded=num_reqs_padded,
|
||||
max_query_len=max_num_scheduled_tokens,
|
||||
ubatch_slices=ubatch_slices_attn,
|
||||
logits_indices=logits_indices,
|
||||
use_spec_decode=use_spec_decode,
|
||||
num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
|
||||
num_scheduled_tokens_np=num_scheduled_tokens_np,
|
||||
cascade_attn_prefix_lens=cascade_attn_prefix_lens,
|
||||
)
|
||||
```
|
||||
|
||||
## Rotary Embedding
|
||||
|
||||
### Update Cos/Sin Cache
|
||||
```python
|
||||
# GPU - typically handled inside attention
|
||||
|
||||
# NPU - explicit update required before forward
|
||||
from vllm_ascend.ops.rotary_embedding import update_cos_sin
|
||||
update_cos_sin(positions)
|
||||
```
|
||||
|
||||
## Sequence Parallelism
|
||||
|
||||
### Enable SP Check
|
||||
```python
|
||||
# GPU - use vllm distributed utilities
|
||||
|
||||
# NPU - use vllm-ascend wrapper
|
||||
from vllm_ascend.utils import enable_sp
|
||||
|
||||
if enable_sp():
|
||||
# sequence parallelism enabled
|
||||
```
|
||||
|
||||
## Sampler
|
||||
|
||||
### Sampler Type
|
||||
```python
|
||||
# GPU - uses vllm sampler
|
||||
self.sampler = Sampler()
|
||||
|
||||
# NPU - uses AscendSampler
|
||||
from vllm_ascend.sample.sampler import AscendSampler
|
||||
self.sampler = AscendSampler()
|
||||
```
|
||||
|
||||
## Input Batch
|
||||
|
||||
### Batch Class
|
||||
```python
|
||||
# GPU
|
||||
from vllm.v1.worker.gpu_input_batch import InputBatch
|
||||
|
||||
# NPU
|
||||
from vllm_ascend.worker.npu_input_batch import NPUInputBatch
|
||||
```
|
||||
|
||||
## Graph Parameter Updates
|
||||
|
||||
### Full Graph Params Update (NPU-specific)
|
||||
```python
|
||||
# GPU - not needed
|
||||
|
||||
# NPU - required for FULL graph mode
|
||||
from vllm_ascend.compilation.acl_graph import update_full_graph_params
|
||||
|
||||
forward_context = get_forward_context()
|
||||
if (
|
||||
forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL
|
||||
and not forward_context.capturing
|
||||
and not self.use_sparse
|
||||
):
|
||||
update_full_graph_params(
|
||||
self.attn_backend,
|
||||
self.update_stream,
|
||||
forward_context,
|
||||
num_tokens_padded,
|
||||
self.vllm_config,
|
||||
self.speculative_config,
|
||||
positions.shape[0],
|
||||
)
|
||||
```
|
||||
|
||||
## Paged Attention Check
|
||||
|
||||
```python
|
||||
# GPU - not typically needed
|
||||
|
||||
# NPU
|
||||
from vllm_ascend.attention.utils import using_paged_attention
|
||||
|
||||
if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config):
|
||||
seq_lens = SEQ_LEN_WITH_MAX_PA_WORKSPACE
|
||||
```
|
||||
|
||||
## Common Method Signature Differences
|
||||
|
||||
### _dummy_run Parameters
|
||||
```python
|
||||
# GPU (v0.17.0)
|
||||
def _dummy_run(
|
||||
self,
|
||||
num_tokens: int,
|
||||
cudagraph_runtime_mode: CUDAGraphMode | None = None,
|
||||
force_attention: bool = False,
|
||||
uniform_decode: bool = False,
|
||||
allow_microbatching: bool = True,
|
||||
skip_eplb: bool = False,
|
||||
is_profile: bool = False,
|
||||
create_mixed_batch: bool = False,
|
||||
remove_lora: bool = True,
|
||||
is_graph_capturing: bool = False,
|
||||
num_active_loras: int = 0,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
|
||||
# NPU (v0.17.0) - adds with_prefill, activate_lora
|
||||
def _dummy_run(
|
||||
self,
|
||||
num_tokens: int,
|
||||
with_prefill: bool = False,
|
||||
cudagraph_runtime_mode: CUDAGraphMode | None = None,
|
||||
force_attention: bool = False,
|
||||
uniform_decode: bool = False,
|
||||
is_profile: bool = False,
|
||||
create_mixed_batch: bool = False,
|
||||
allow_microbatching: bool = True,
|
||||
skip_eplb: bool = False,
|
||||
remove_lora: bool = True,
|
||||
activate_lora: bool = False,
|
||||
is_graph_capturing: bool = False,
|
||||
num_active_loras: int = 0,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
```
|
||||
|
||||
### _model_forward Parameters
|
||||
```python
|
||||
# GPU - no num_tokens_padded
|
||||
def _model_forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
positions: torch.Tensor | None = None,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**model_kwargs: dict[str, Any],
|
||||
):
|
||||
|
||||
# NPU - has num_tokens_padded as first parameter
|
||||
def _model_forward(
|
||||
self,
|
||||
num_tokens_padded: int,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
positions: torch.Tensor | None = None,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**model_kwargs: dict[str, Any],
|
||||
):
|
||||
```
|
||||
|
||||
## Quick Reference Table
|
||||
|
||||
| Feature | GPU | NPU |
|
||||
|---------|-----|-----|
|
||||
| Graph wrapper | `CUDAGraphWrapper` | `ACLGraphWrapper` |
|
||||
| Forward context | `set_forward_context` | `set_ascend_forward_context` |
|
||||
| Runtime mode param | `cudagraph_runtime_mode` | `aclgraph_runtime_mode` |
|
||||
| Device sync | `torch.cuda.synchronize()` | `torch.npu.synchronize()` |
|
||||
| Stream | `torch.cuda.Stream` | `torch.npu.Stream` |
|
||||
| Current stream | `torch.cuda.current_stream()` | `torch.npu.current_stream()` |
|
||||
| Input batch | `InputBatch` | `NPUInputBatch` |
|
||||
| Sampler | `Sampler` | `AscendSampler` |
|
||||
| Attention state | N/A | `AscendAttentionState` |
|
||||
| RoPE update | N/A | `update_cos_sin()` |
|
||||
@@ -0,0 +1,374 @@
|
||||
# Omni-Specific Code Blocks Reference
|
||||
|
||||
This document catalogs omni-specific code blocks in the NPU model runners, making it easier to identify what needs to be preserved during upgrades.
|
||||
|
||||
> **IMPORTANT**: This document may not be complete or up-to-date!
|
||||
>
|
||||
> - Always grep for `Omni-new` in the GPU implementations (`vllm_omni/worker/`) to find the authoritative list
|
||||
> - New omni features may be added that are not yet documented here
|
||||
> - When you discover new omni-specific blocks during an upgrade, please update this document
|
||||
> - Last verified: Check git history for this file
|
||||
|
||||
## OmniNPUModelRunner (npu_model_runner.py)
|
||||
|
||||
### load_model - Talker MTP Initialization
|
||||
|
||||
```python
|
||||
def load_model(self, *args, **kwargs) -> None:
|
||||
NPUModelRunner.load_model(self, *args, **kwargs)
|
||||
# Initialize enable_sp cache to avoid get_current_vllm_config() error
|
||||
# in _pad_for_sequence_parallelism during execute_model.
|
||||
# This is a workaround for vllm-ascend not passing vllm_config to enable_sp().
|
||||
enable_sp(self.vllm_config)
|
||||
# TODO move this model specific logic to a separate class
|
||||
# TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS.
|
||||
talker_mtp = getattr(self.model, "talker_mtp", None)
|
||||
if talker_mtp is not None:
|
||||
self.talker_mtp = talker_mtp # type: ignore[assignment]
|
||||
cudagraph_mode = self.compilation_config.cudagraph_mode
|
||||
assert cudagraph_mode is not None
|
||||
# Only wrap talker_mtp in CUDAGraphWrapper for Omni models that
|
||||
# have a separate .talker sub-module. TTS models' code predictor
|
||||
# has internal AR loops / torch.multinomial — not graph-safe.
|
||||
has_separate_talker = getattr(self.model, "talker", None) is not None
|
||||
if cudagraph_mode.has_full_cudagraphs() and has_separate_talker:
|
||||
# NOTE: Use ACLGraphWrapper on NPU, not CUDAGraphWrapper
|
||||
self.talker_mtp = ACLGraphWrapper(talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL)
|
||||
# TTS exposes mtp_hidden_size; Omni uses hf_text_config.hidden_size.
|
||||
hidden_size = int(
|
||||
getattr(self.model, "mtp_hidden_size", 0) or getattr(self.model_config.hf_text_config, "hidden_size")
|
||||
)
|
||||
max_batch_size = max(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size)
|
||||
self.talker_mtp_input_ids = self._make_buffer(max_batch_size, dtype=torch.int32)
|
||||
self.talker_mtp_inputs_embeds = self._make_buffer(
|
||||
max_batch_size, hidden_size, dtype=self.dtype, numpy=False
|
||||
)
|
||||
self.last_talker_hidden = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
|
||||
self.text_step = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
|
||||
```
|
||||
|
||||
### _dummy_run - Talker MTP Dummy Forward
|
||||
|
||||
Location: Inside `set_ascend_forward_context` block, before main model forward
|
||||
|
||||
```python
|
||||
# ---------------------------------------Omni-new----------------------------------------------
|
||||
if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"):
|
||||
num_tokens_padded_talker_mtp = num_tokens_padded
|
||||
if num_tokens_padded_talker_mtp == self.max_num_tokens:
|
||||
num_tokens_padded_talker_mtp = self.talker_mtp_input_ids.gpu.shape[0]
|
||||
outputs = self.talker_mtp(
|
||||
self.talker_mtp_input_ids.gpu[:num_tokens_padded_talker_mtp],
|
||||
self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded_talker_mtp],
|
||||
self.last_talker_hidden.gpu[:num_tokens_padded_talker_mtp],
|
||||
self.text_step.gpu[:num_tokens_padded_talker_mtp],
|
||||
)
|
||||
self.compilation_config.cache_dir = None
|
||||
# ---------------------------------------Omni-new----------------------------------------------
|
||||
```
|
||||
|
||||
### _dummy_run - Extract Multimodal Outputs
|
||||
|
||||
Location: After model forward, before dummy_compute_logits
|
||||
|
||||
```python
|
||||
# ---------------------------------------Omni-new----------------------------------------------
|
||||
hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
|
||||
# ---------------------------------------Omni-new----------------------------------------------
|
||||
```
|
||||
|
||||
### _model_forward - Omni Output Wrapping
|
||||
|
||||
```python
|
||||
def _model_forward(
|
||||
self,
|
||||
num_tokens_padded: int,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
positions: torch.Tensor | None = None,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**model_kwargs: dict[str, Any],
|
||||
):
|
||||
"""Override to combine NPUModelRunner's signature with OmniGPUModelRunner's logic."""
|
||||
# Omni-specific: build and inject extra model kwargs
|
||||
model_kwargs_extra = self._build_model_kwargs_extra()
|
||||
|
||||
# Call the model forward (same as NPUModelRunner)
|
||||
assert self.model is not None
|
||||
model_output = self.model(
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
intermediate_tensors=intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds,
|
||||
**model_kwargs,
|
||||
**model_kwargs_extra,
|
||||
)
|
||||
|
||||
# Omni-specific: wrap output if needed
|
||||
if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"):
|
||||
model_output = self.model.make_omni_output(model_output, **model_kwargs_extra)
|
||||
|
||||
# Omni-specific: cache model output for later sample_tokens
|
||||
self._omni_last_model_output = model_output
|
||||
|
||||
# NPU-specific: update full graph params (keep from vllm-ascend)
|
||||
forward_context = get_forward_context()
|
||||
# ... NPU graph update logic ...
|
||||
|
||||
# NPU-specific: all-gather for sequence parallelism (keep from vllm-ascend)
|
||||
if get_forward_context().sp_enabled and not isinstance(model_output, IntermediateTensors):
|
||||
model_output = self._all_gather_hidden_states_and_aux(model_output)
|
||||
|
||||
return model_output
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## NPUARModelRunner (npu_ar_model_runner.py)
|
||||
|
||||
### __init__ - KV Transfer Manager
|
||||
|
||||
```python
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
|
||||
# each model stage has their own hidden size
|
||||
self.hidden_size = self.model_config.hf_text_config.hidden_size
|
||||
self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
|
||||
# Initialize KV cache manager (preserve vllm_config fallback behavior)
|
||||
self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
|
||||
```
|
||||
|
||||
### execute_model - KV Transfer Before Update States
|
||||
|
||||
Location: At the very beginning of execute_model
|
||||
|
||||
```python
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
# [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
|
||||
self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer(
|
||||
finished_reqs=getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}),
|
||||
kv_caches=self.kv_caches,
|
||||
block_size=self.cache_config.block_size,
|
||||
cache_dtype=str(self.cache_config.cache_dtype),
|
||||
request_id_resolver=self._resolve_global_request_id,
|
||||
)
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
```
|
||||
|
||||
### execute_model - Custom _update_states Call
|
||||
|
||||
Location: Inside synchronize_input_prep context
|
||||
|
||||
```python
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
self._update_states(scheduler_output)
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
```
|
||||
|
||||
### execute_model - Extract Multimodal Outputs
|
||||
|
||||
Location: In post process section, after hidden_states assignment
|
||||
|
||||
```python
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
|
||||
|
||||
if multimodal_outputs is not None:
|
||||
keys_or_type = (
|
||||
list(multimodal_outputs.keys())
|
||||
if isinstance(multimodal_outputs, dict)
|
||||
else type(multimodal_outputs)
|
||||
)
|
||||
logger.debug(f"[AR] execute_model: multimodal_outputs keys = {keys_or_type}")
|
||||
else:
|
||||
logger.debug("[AR] execute_model: multimodal_outputs is None")
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
```
|
||||
|
||||
### execute_model - Compute Logits with sampling_metadata
|
||||
|
||||
Location: In both broadcast_pp_output True and False branches
|
||||
|
||||
```python
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
# Try with sampling_metadata first; fall back to without for models that don't support it
|
||||
try:
|
||||
logits = self.model.compute_logits(
|
||||
sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
|
||||
)
|
||||
except TypeError:
|
||||
logits = self.model.compute_logits(sample_hidden_states)
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
```
|
||||
|
||||
### sample_tokens - KV Extracted Req IDs
|
||||
|
||||
Location: At the beginning of sample_tokens
|
||||
|
||||
```python
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
kv_extracted_req_ids = getattr(self, "kv_extracted_req_ids", None)
|
||||
self.kv_extracted_req_ids = None
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
```
|
||||
|
||||
### sample_tokens - Process Additional Information and Build Output
|
||||
|
||||
Location: After bookkeeping sync, replacing the original output construction
|
||||
|
||||
```python
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
|
||||
num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None)
|
||||
if num_scheduled_tokens_np is None:
|
||||
req_ids = self.input_batch.req_ids
|
||||
num_scheduled_tokens_np = np.array(
|
||||
[scheduler_output.num_scheduled_tokens[rid] for rid in req_ids],
|
||||
dtype=np.int32,
|
||||
)
|
||||
|
||||
self._process_additional_information_updates(
|
||||
hidden_states, multimodal_outputs, num_scheduled_tokens_np, scheduler_output
|
||||
)
|
||||
|
||||
pooler_output: list[dict[str, object]] = []
|
||||
for rid in req_ids_output_copy:
|
||||
idx = req_id_to_index_output_copy[rid]
|
||||
start = int(self.query_start_loc.cpu[idx])
|
||||
sched = int(num_scheduled_tokens_np[idx])
|
||||
end = start + sched
|
||||
hidden_slice = hidden_states_cpu[start:end]
|
||||
payload: dict[str, object] = {"hidden": hidden_slice}
|
||||
if isinstance(multimodal_outputs, dict) and multimodal_outputs:
|
||||
# ... multimodal output slicing logic ...
|
||||
pooler_output.append(payload)
|
||||
|
||||
model_runner_output = OmniModelRunnerOutput(
|
||||
req_ids=req_ids_output_copy,
|
||||
req_id_to_index=req_id_to_index_output_copy,
|
||||
sampled_token_ids=valid_sampled_token_ids,
|
||||
logprobs=logprobs_lists,
|
||||
prompt_logprobs_dict=prompt_logprobs_dict,
|
||||
pooler_output=(pooler_output if self.vllm_config.model_config.engine_output_type != "text" else None),
|
||||
kv_connector_output=kv_connector_output,
|
||||
)
|
||||
model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## NPUGenerationModelRunner (npu_generation_model_runner.py)
|
||||
|
||||
### execute_model - Async Chunk Update
|
||||
|
||||
Location: Inside prepare input section, before synchronize_input_prep
|
||||
|
||||
```python
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
if self.model_config.async_chunk and num_scheduled_tokens:
|
||||
self._update_request_states(scheduler_output)
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
```
|
||||
|
||||
### execute_model - Seq Token Counts
|
||||
|
||||
Location: After _preprocess call
|
||||
|
||||
```python
|
||||
# [Omni] Pass token counts per request for code2wav output slicing
|
||||
model_kwargs["seq_token_counts"] = tokens
|
||||
```
|
||||
|
||||
### execute_model - Run Generation Model
|
||||
|
||||
Location: Inside forward context
|
||||
|
||||
```python
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
outputs = self._run_generation_model(
|
||||
num_tokens_padded=num_tokens_padded,
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
intermediate_tensors=intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds,
|
||||
model_kwargs=model_kwargs,
|
||||
logits_indices=logits_indices,
|
||||
)
|
||||
_, multimodal_outputs = self.extract_multimodal_outputs(outputs)
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
```
|
||||
|
||||
### sample_tokens - Multimodal Output Processing
|
||||
|
||||
The entire sample_tokens method body is omni-specific for generation models:
|
||||
|
||||
```python
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
pooler_output: list[object] = []
|
||||
if isinstance(multimodal_outputs, torch.Tensor):
|
||||
# ... tensor handling ...
|
||||
elif isinstance(multimodal_outputs, list):
|
||||
# ... list handling ...
|
||||
elif isinstance(multimodal_outputs, dict):
|
||||
# ... dict handling per request ...
|
||||
else:
|
||||
raise RuntimeError("Unsupported diffusion output type")
|
||||
# [Omni] Copy req_id mappings to avoid async scheduling mutation.
|
||||
req_ids_output_copy = self.input_batch.req_ids.copy()
|
||||
req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
|
||||
output = OmniModelRunnerOutput(
|
||||
req_ids=req_ids_output_copy,
|
||||
req_id_to_index=req_id_to_index_output_copy,
|
||||
sampled_token_ids=[],
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={},
|
||||
pooler_output=pooler_output,
|
||||
kv_connector_output=kv_connector_output,
|
||||
num_nans_in_logits={},
|
||||
ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
|
||||
)
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
```
|
||||
|
||||
### _dummy_run - Model Kwargs Init and Multimodal Extract
|
||||
|
||||
Location: Before model forward and after
|
||||
|
||||
```python
|
||||
model_kwargs = self._init_model_kwargs() # Before forward
|
||||
|
||||
# ... forward ...
|
||||
|
||||
# -------------------------------------- Omni-new -------------------------------------------------
|
||||
hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
|
||||
# -------------------------------------------------------------------------------------------------
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## ExecuteModelState Extension
|
||||
|
||||
The `ExecuteModelState` NamedTuple is extended for omni:
|
||||
|
||||
```python
|
||||
class ExecuteModelState(NamedTuple):
|
||||
"""Ephemeral cached state transferred between execute_model() and
|
||||
sample_tokens(), after execute_model() returns None."""
|
||||
|
||||
scheduler_output: SchedulerOutput
|
||||
logits: torch.Tensor
|
||||
spec_decode_metadata: SpecDecodeMetadata | None
|
||||
spec_decode_common_attn_metadata: AscendCommonAttentionMetadata | None
|
||||
hidden_states: torch.Tensor
|
||||
sample_hidden_states: torch.Tensor
|
||||
aux_hidden_states: list[torch.Tensor] | None
|
||||
attn_metadata: PerLayerAttnMetadata
|
||||
positions: torch.Tensor
|
||||
ec_connector_output: ECConnectorOutput | None
|
||||
cudagraph_stats: CUDAGraphStat | None
|
||||
multimodal_outputs: Any # <-- Omni extension
|
||||
```
|
||||
|
||||
This extended state must be imported from `npu_ar_model_runner` in `npu_generation_model_runner`.
|
||||
@@ -0,0 +1,222 @@
|
||||
# NPU Model Runner Upgrade Workflow Checklist
|
||||
|
||||
> **Note**: Reference documents (`omni-specific-blocks.md`) may not be complete. Always grep for `Omni-new` in GPU implementations to find all omni-specific code blocks. Update the reference docs when discovering new blocks.
|
||||
|
||||
## Pre-Upgrade Preparation
|
||||
|
||||
### 1. Version Information
|
||||
- [ ] Identify current vllm-omni version: `_________`
|
||||
- [ ] Identify target vllm-ascend version: `_________`
|
||||
- [ ] Identify target vllm version: `_________`
|
||||
- [ ] Last release date for GPU worker changes: `_________`
|
||||
|
||||
### 2. Gather Git History
|
||||
```bash
|
||||
# GPU-side omni changes since last release
|
||||
cd /root/vllm-workspace/vllm-omni
|
||||
git log --oneline --since="YYYY-MM-DD" -- vllm_omni/worker/
|
||||
|
||||
# vllm-ascend NPUModelRunner changes
|
||||
cd /root/vllm-workspace/vllm-ascend
|
||||
git log --oneline <from-tag>..<to-tag> -- vllm_ascend/worker/model_runner_v1.py
|
||||
```
|
||||
|
||||
### 3. Backup Current Files
|
||||
- [ ] Create backup of current NPU runners:
|
||||
```bash
|
||||
cp -r vllm_omni/platforms/npu/worker vllm_omni/platforms/npu/worker.backup
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## OmniNPUModelRunner (npu_model_runner.py)
|
||||
|
||||
### Read and Understand
|
||||
- [ ] Read current `npu_model_runner.py`
|
||||
- [ ] Read latest `vllm_ascend/worker/model_runner_v1.py`
|
||||
- [ ] Read latest `vllm_omni/worker/gpu_model_runner.py`
|
||||
|
||||
### Method: load_model
|
||||
- [ ] Document existing omni-specific logic
|
||||
- [ ] Copy latest NPUModelRunner.load_model structure
|
||||
- [ ] Re-insert: `enable_sp(self.vllm_config)` call
|
||||
- [ ] Re-insert: talker_mtp detection and setup
|
||||
- [ ] Replace: `CUDAGraphWrapper` → `ACLGraphWrapper`
|
||||
- [ ] Re-insert: Buffer allocations (talker_mtp_input_ids, etc.)
|
||||
|
||||
### Method: _dummy_run
|
||||
- [ ] Document existing omni-specific logic locations
|
||||
- [ ] Copy latest NPUModelRunner._dummy_run
|
||||
- [ ] Re-insert: talker_mtp dummy forward block (inside context)
|
||||
- [ ] Re-insert: `extract_multimodal_outputs` call
|
||||
- [ ] Verify: Comment markers are present
|
||||
|
||||
### Method: _model_forward
|
||||
- [ ] Copy latest NPUModelRunner._model_forward structure
|
||||
- [ ] Re-insert: `_build_model_kwargs_extra()` call
|
||||
- [ ] Re-insert: OmniOutput wrapping logic
|
||||
- [ ] Re-insert: `_omni_last_model_output` caching
|
||||
- [ ] Keep: NPU graph params update
|
||||
- [ ] Keep: SP all-gather logic
|
||||
|
||||
### Method: _talker_mtp_forward
|
||||
- [ ] Verify: Uses `set_ascend_forward_context`
|
||||
- [ ] Verify: Uses `ACLGraphWrapper` check
|
||||
- [ ] Sync any changes from GPU `_talker_mtp_forward`
|
||||
|
||||
### Imports
|
||||
- [ ] Update vllm-ascend imports to latest paths
|
||||
- [ ] Verify all omni imports are present
|
||||
- [ ] Remove any deprecated imports
|
||||
|
||||
---
|
||||
|
||||
## NPUARModelRunner (npu_ar_model_runner.py)
|
||||
|
||||
### Read and Understand
|
||||
- [ ] Read current `npu_ar_model_runner.py`
|
||||
- [ ] Read latest `vllm_ascend/worker/model_runner_v1.py` execute_model
|
||||
- [ ] Read latest `vllm_omni/worker/gpu_ar_model_runner.py`
|
||||
|
||||
### Method: __init__
|
||||
- [ ] Sync any new initialization from GPU side
|
||||
- [ ] Keep: `OmniKVTransferManager` setup
|
||||
- [ ] Keep: Custom buffer allocations
|
||||
|
||||
### Method: execute_model
|
||||
- [ ] Document all omni blocks with line numbers
|
||||
- [ ] Copy latest NPUModelRunner.execute_model structure
|
||||
- [ ] Re-insert: KV transfer handling (beginning)
|
||||
- [ ] Re-insert: Custom `_update_states` call
|
||||
- [ ] Re-insert: `extract_multimodal_outputs`
|
||||
- [ ] Re-insert: `compute_logits` with sampling_metadata try/except
|
||||
- [ ] Update: ExecuteModelState to include multimodal_outputs
|
||||
|
||||
### Method: sample_tokens
|
||||
- [ ] Document all omni blocks
|
||||
- [ ] Copy latest NPUModelRunner.sample_tokens structure
|
||||
- [ ] Re-insert: `kv_extracted_req_ids` handling
|
||||
- [ ] Re-insert: Hidden states CPU copy
|
||||
- [ ] Re-insert: `_process_additional_information_updates`
|
||||
- [ ] Re-insert: `OmniModelRunnerOutput` construction
|
||||
|
||||
### ExecuteModelState
|
||||
- [ ] Verify: `multimodal_outputs` field is present
|
||||
- [ ] Verify: Imported/used correctly in execute_model
|
||||
|
||||
### Imports
|
||||
- [ ] Update all vllm-ascend imports
|
||||
- [ ] Keep omni-specific imports
|
||||
|
||||
---
|
||||
|
||||
## NPUGenerationModelRunner (npu_generation_model_runner.py)
|
||||
|
||||
### Read and Understand
|
||||
- [ ] Read current `npu_generation_model_runner.py`
|
||||
- [ ] Read latest GPU `gpu_generation_model_runner.py`
|
||||
|
||||
### Method: _update_request_states
|
||||
- [ ] Verify: async_chunk handling is correct
|
||||
- [ ] Sync any changes from GPU side
|
||||
|
||||
### Method: execute_model
|
||||
- [ ] Document all omni blocks
|
||||
- [ ] Copy latest NPUModelRunner.execute_model base structure
|
||||
- [ ] Re-insert: async_chunk update logic
|
||||
- [ ] Re-insert: `seq_token_counts` injection
|
||||
- [ ] Re-insert: `_run_generation_model` call
|
||||
- [ ] Re-insert: `extract_multimodal_outputs`
|
||||
- [ ] Use: ExecuteModelState from npu_ar_model_runner
|
||||
|
||||
### Method: sample_tokens
|
||||
- [ ] Keep: Entire omni multimodal output processing
|
||||
- [ ] Update: Any new output fields needed
|
||||
- [ ] Keep: `OmniModelRunnerOutput` construction
|
||||
|
||||
### Method: _run_generation_model
|
||||
- [ ] Sync any changes from GPU side
|
||||
- [ ] Keep: `_model_forward` call with sampler
|
||||
|
||||
### Method: _dummy_run
|
||||
- [ ] Copy latest NPUModelRunner._dummy_run
|
||||
- [ ] Re-insert: `model_kwargs = self._init_model_kwargs()`
|
||||
- [ ] Re-insert: `extract_multimodal_outputs` at end
|
||||
|
||||
### Imports
|
||||
- [ ] Import ExecuteModelState from npu_ar_model_runner
|
||||
- [ ] Update vllm-ascend imports
|
||||
|
||||
---
|
||||
|
||||
## Post-Upgrade Validation
|
||||
|
||||
### Syntax Validation
|
||||
- [ ] `python -m py_compile vllm_omni/platforms/npu/worker/npu_model_runner.py`
|
||||
- [ ] `python -m py_compile vllm_omni/platforms/npu/worker/npu_ar_model_runner.py`
|
||||
- [ ] `python -m py_compile vllm_omni/platforms/npu/worker/npu_generation_model_runner.py`
|
||||
|
||||
### Import Validation
|
||||
- [ ] `python -c "from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner"`
|
||||
- [ ] `python -c "from vllm_omni.platforms.npu.worker.npu_ar_model_runner import NPUARModelRunner"`
|
||||
- [ ] `python -c "from vllm_omni.platforms.npu.worker.npu_generation_model_runner import NPUGenerationModelRunner"`
|
||||
|
||||
### Comment Markers
|
||||
- [ ] Grep for "Omni-new" in all three files
|
||||
- [ ] Verify all omni blocks have closing markers
|
||||
|
||||
### Code Review
|
||||
- [ ] No `CUDAGraphWrapper` references
|
||||
- [ ] All `set_forward_context` replaced with `set_ascend_forward_context`
|
||||
- [ ] Parameter names correct (`aclgraph_runtime_mode` not `cudagraph_runtime_mode`)
|
||||
- [ ] No duplicate code blocks
|
||||
- [ ] No missing imports
|
||||
|
||||
---
|
||||
|
||||
## Git Commit
|
||||
|
||||
### Commit Message Template
|
||||
```
|
||||
[NPU] Upgrade model runners to align with vllm-ascend vX.Y.Z
|
||||
|
||||
- Update OmniNPUModelRunner with latest NPUModelRunner base
|
||||
- Update NPUARModelRunner execute_model and sample_tokens
|
||||
- Update NPUGenerationModelRunner for async_chunk changes
|
||||
- Sync GPU-side omni changes from vX.Y.Z release
|
||||
- Preserve all omni-specific logic (marked with Omni-new comments)
|
||||
|
||||
Changes from vllm-ascend:
|
||||
- <list key changes>
|
||||
|
||||
Changes synced from GPU:
|
||||
- <list key GPU-side omni changes>
|
||||
```
|
||||
|
||||
### Files to Stage
|
||||
- [ ] `vllm_omni/platforms/npu/worker/npu_model_runner.py`
|
||||
- [ ] `vllm_omni/platforms/npu/worker/npu_ar_model_runner.py`
|
||||
- [ ] `vllm_omni/platforms/npu/worker/npu_generation_model_runner.py`
|
||||
- [ ] Any other modified files
|
||||
|
||||
---
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Import Errors
|
||||
- Check if vllm-ascend module paths have changed
|
||||
- Verify PYTHONPATH includes both vllm-ascend and vllm-omni
|
||||
|
||||
### Type Errors
|
||||
- Check method signatures match between GPU and NPU
|
||||
- Verify NamedTuple fields match expected structure
|
||||
|
||||
### Runtime Errors
|
||||
- Enable debug logging: `export VLLM_LOGGING_LEVEL=DEBUG`
|
||||
- Check graph capture issues: try `--enforce-eager`
|
||||
- Check attention issues: verify AscendAttentionState usage
|
||||
|
||||
### Performance Regression
|
||||
- Compare with previous version on same model
|
||||
- Check if graph capture is working: look for ACLGraph logs
|
||||
- Verify SP/EP configurations are correct
|
||||
@@ -1,629 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
PR Reviewer using GLM API for vllm-omni project.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, TypedDict
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
# Type definitions for API responses
|
||||
class PRDetails(TypedDict):
|
||||
"""Type definition for GitHub PR details response."""
|
||||
|
||||
title: str
|
||||
body: str
|
||||
number: int
|
||||
state: str
|
||||
user: dict[str, Any]
|
||||
|
||||
|
||||
class GLMMessage(TypedDict):
|
||||
"""Type definition for GLM API message."""
|
||||
|
||||
role: str
|
||||
content: str
|
||||
|
||||
|
||||
class GLMChoice(TypedDict):
|
||||
"""Type definition for GLM API choice."""
|
||||
|
||||
message: GLMMessage
|
||||
finish_reason: str
|
||||
|
||||
|
||||
class GLMResponse(TypedDict):
|
||||
"""Type definition for GLM API response."""
|
||||
|
||||
choices: list[GLMChoice]
|
||||
usage: dict[str, int] | None
|
||||
|
||||
|
||||
class GitHubComment(TypedDict):
|
||||
"""Type definition for GitHub comment."""
|
||||
|
||||
id: int
|
||||
body: str
|
||||
created_at: str
|
||||
user: dict[str, Any]
|
||||
|
||||
|
||||
# Configuration
|
||||
TRIGGER_PHRASE: str = "@vllm-omni-reviewer"
|
||||
DEFAULT_GLM_API_URL: str = "https://open.bigmodel.cn/api/paas/v4/chat/completions" # noqa: E501
|
||||
DEFAULT_GLM_MODEL: str = "glm-5"
|
||||
DEFAULT_COOLDOWN_MINUTES: int = 5
|
||||
DEFAULT_MAX_RETRIES: int = 3
|
||||
DEFAULT_RETRY_DELAY: float = 1.0
|
||||
MAX_DIFF_SIZE: int = 100_000 # Maximum diff size in characters
|
||||
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
"""Configuration for the PR reviewer."""
|
||||
|
||||
glm_api_url: str
|
||||
glm_model: str
|
||||
cooldown_minutes: int
|
||||
max_retries: int
|
||||
retry_delay: float
|
||||
max_diff_size: int
|
||||
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="[PR Reviewer] %(asctime)s - %(levelname)s - %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
)
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_config() -> Config:
|
||||
"""Load configuration from environment variables with defaults."""
|
||||
return Config(
|
||||
glm_api_url=os.getenv("GLM_API_URL", DEFAULT_GLM_API_URL),
|
||||
glm_model=os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL),
|
||||
cooldown_minutes=int(
|
||||
os.getenv(
|
||||
"PR_REVIEWER_COOLDOWN_MINUTES",
|
||||
str(DEFAULT_COOLDOWN_MINUTES),
|
||||
)
|
||||
),
|
||||
max_retries=int(
|
||||
os.getenv(
|
||||
"PR_REVIEWER_MAX_RETRIES",
|
||||
str(DEFAULT_MAX_RETRIES),
|
||||
)
|
||||
),
|
||||
retry_delay=float(os.getenv("PR_REVIEWER_RETRY_DELAY", str(DEFAULT_RETRY_DELAY))),
|
||||
max_diff_size=int(os.getenv("PR_REVIEWER_MAX_DIFF_SIZE", str(MAX_DIFF_SIZE))), # noqa: E501
|
||||
)
|
||||
|
||||
|
||||
def get_env_var(name: str) -> str:
|
||||
"""
|
||||
Get an environment variable or raise an error.
|
||||
|
||||
Args:
|
||||
name: Name of the environment variable.
|
||||
|
||||
Returns:
|
||||
The value of the environment variable.
|
||||
|
||||
Raises:
|
||||
SystemExit: If the environment variable is not set.
|
||||
"""
|
||||
value = os.environ.get(name)
|
||||
if not value:
|
||||
logger.error(f"Environment variable {name} is not set")
|
||||
sys.exit(1)
|
||||
return value
|
||||
|
||||
|
||||
def check_trigger(comment_body: str) -> bool:
|
||||
"""
|
||||
Check if the comment contains the trigger phrase.
|
||||
|
||||
Args:
|
||||
comment_body: The body of the comment to check.
|
||||
|
||||
Returns:
|
||||
True if the trigger phrase is found, False otherwise.
|
||||
"""
|
||||
return TRIGGER_PHRASE in comment_body
|
||||
|
||||
|
||||
def fetch_pr_diff(
|
||||
repo_name: str,
|
||||
pr_number: int,
|
||||
token: str,
|
||||
max_size: int = MAX_DIFF_SIZE,
|
||||
) -> str | None:
|
||||
"""
|
||||
Fetch the diff for a pull request.
|
||||
|
||||
Args:
|
||||
repo_name: The repository name in format "owner/repo".
|
||||
pr_number: The pull request number.
|
||||
token: GitHub authentication token.
|
||||
max_size: Maximum diff size in characters.
|
||||
|
||||
Returns:
|
||||
The diff content as a string, or None if fetching failed.
|
||||
Returns empty string if diff is larger than max_size.
|
||||
"""
|
||||
url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/vnd.github.v3.diff",
|
||||
}
|
||||
|
||||
logger.info(f"Fetching PR diff from {url}")
|
||||
response = requests.get(url, headers=headers, timeout=30)
|
||||
|
||||
if response.status_code == 200:
|
||||
diff: str = response.text
|
||||
if len(diff) > max_size:
|
||||
logger.warning(
|
||||
f"Diff size ({len(diff)} bytes) exceeds maximum "
|
||||
f"({max_size} bytes), truncating to first "
|
||||
f"{max_size} bytes"
|
||||
)
|
||||
return diff[:max_size] + "\n\n... [Diff truncated due to size] ..."
|
||||
logger.info(f"Successfully fetched diff ({len(diff)} bytes)")
|
||||
return diff
|
||||
else:
|
||||
logger.error(f"Failed to fetch PR diff: {response.status_code}")
|
||||
logger.error(f"Response: {response.text}")
|
||||
return None
|
||||
|
||||
|
||||
def fetch_pr_details(
|
||||
repo_name: str,
|
||||
pr_number: int,
|
||||
token: str,
|
||||
) -> PRDetails | None:
|
||||
"""
|
||||
Fetch PR details including title and description.
|
||||
|
||||
Args:
|
||||
repo_name: The repository name in format "owner/repo".
|
||||
pr_number: The pull request number.
|
||||
token: GitHub authentication token.
|
||||
|
||||
Returns:
|
||||
A dictionary containing PR details, or None if fetching failed.
|
||||
"""
|
||||
url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/vnd.github.v3+json",
|
||||
}
|
||||
|
||||
logger.info(f"Fetching PR details from {url}")
|
||||
response = requests.get(url, headers=headers, timeout=30)
|
||||
|
||||
if response.status_code == 200:
|
||||
return response.json()
|
||||
else:
|
||||
logger.error(f"Failed to fetch PR details: {response.status_code}")
|
||||
return None
|
||||
|
||||
|
||||
def build_review_prompt(pr_title: str, pr_description: str, diff: str) -> str:
|
||||
"""
|
||||
Build the prompt for the GLM-4.7 API.
|
||||
|
||||
Args:
|
||||
pr_title: The title of the pull request.
|
||||
pr_description: The description/body of the pull request.
|
||||
diff: The diff content of the pull request.
|
||||
|
||||
Returns:
|
||||
The formatted prompt string for the API.
|
||||
"""
|
||||
return f"""You are an expert code reviewer for the VLLM-Omni project. \
|
||||
Please review the following pull request:
|
||||
|
||||
## Pull Request Details
|
||||
**Title:** {pr_title}
|
||||
|
||||
**Description:**
|
||||
{pr_description if pr_description else "No description provided."}
|
||||
|
||||
## Code Changes (Diff)
|
||||
{diff}
|
||||
|
||||
## Review Guidelines
|
||||
|
||||
Please provide a comprehensive code review with the following sections:
|
||||
|
||||
### 1. Overview
|
||||
- Brief summary of the changes
|
||||
- Overall assessment (positive, neutral, or concerns)
|
||||
|
||||
### 2. Code Quality
|
||||
- Code style and consistency
|
||||
- Potential bugs or edge cases
|
||||
- Performance considerations
|
||||
- Error handling
|
||||
|
||||
### 3. Architecture & Design
|
||||
- Integration with existing codebase
|
||||
- Design patterns and best practices
|
||||
- Potential improvements
|
||||
|
||||
### 4. Security & Safety
|
||||
- Security concerns (if any)
|
||||
- Resource management
|
||||
- Input validation
|
||||
|
||||
### 5. Testing & Documentation
|
||||
- Test coverage considerations
|
||||
- Documentation completeness
|
||||
- Examples and usage clarity
|
||||
|
||||
### 6. Specific Suggestions
|
||||
- Line-by-line specific feedback (use `file:line` format)
|
||||
- Concrete actionable suggestions
|
||||
- Code examples for improvements (if applicable)
|
||||
|
||||
### 7. Approval Status
|
||||
- **LGTM** (Looks Good To Me) if the PR is ready to merge
|
||||
- **LGTM with suggestions** if the PR is good but has minor suggestions
|
||||
- **Changes requested** if significant changes are needed
|
||||
|
||||
## Important Notes
|
||||
- Be constructive and helpful
|
||||
- Focus on objective technical feedback
|
||||
- Acknowledge good practices when you see them
|
||||
- Prioritize critical issues over nitpicks
|
||||
- If the diff is empty or minimal, acknowledge this and provide
|
||||
any relevant context-specific guidance
|
||||
|
||||
Please format your response in Markdown with clear section headers.
|
||||
"""
|
||||
|
||||
|
||||
def validate_glm_response(data: dict[str, Any]) -> str | None:
|
||||
"""
|
||||
Validate and extract content from GLM API response.
|
||||
|
||||
Args:
|
||||
data: The response data from GLM API.
|
||||
|
||||
Returns:
|
||||
The review content string if valid, None otherwise.
|
||||
"""
|
||||
# Check if choices exists and is a non-empty list
|
||||
if "choices" not in data:
|
||||
logger.error("GLM API response missing 'choices' field")
|
||||
logger.error(f"Response structure: {json.dumps(data, indent=2)}")
|
||||
return None
|
||||
|
||||
choices = data["choices"]
|
||||
if not isinstance(choices, list):
|
||||
logger.error(f"GLM API 'choices' is not a list: {type(choices)}")
|
||||
return None
|
||||
|
||||
if len(choices) == 0:
|
||||
logger.error("GLM API 'choices' is an empty list")
|
||||
return None
|
||||
|
||||
# Check if first choice has message
|
||||
try:
|
||||
first_choice = choices[0]
|
||||
if not isinstance(first_choice, dict):
|
||||
logger.error(f"GLM API choice is not a dict: {type(first_choice)}")
|
||||
return None
|
||||
|
||||
if "message" not in first_choice:
|
||||
logger.error("GLM API choice missing 'message' field")
|
||||
logger.error(f"Choice structure: {json.dumps(first_choice, indent=2)}") # noqa: E501
|
||||
return None
|
||||
|
||||
message = first_choice["message"]
|
||||
if not isinstance(message, dict):
|
||||
logger.error(f"GLM API message is not a dict: {type(message)}")
|
||||
return None
|
||||
|
||||
if "content" not in message:
|
||||
logger.error("GLM API message missing 'content' field")
|
||||
logger.error(f"Message structure: {json.dumps(message, indent=2)}")
|
||||
return None
|
||||
|
||||
content = message["content"]
|
||||
if not isinstance(content, str):
|
||||
logger.error(f"GLM API content is not a string: {type(content)}")
|
||||
return None
|
||||
|
||||
return content
|
||||
|
||||
except (KeyError, IndexError, TypeError) as e:
|
||||
logger.error(f"Failed to parse GLM API response: {e}")
|
||||
logger.error(f"Response: {json.dumps(data, indent=2)}")
|
||||
return None
|
||||
|
||||
|
||||
def call_glm_api(prompt: str, api_key: str, config: Config) -> str | None:
|
||||
"""
|
||||
Call the GLM-4.7 API to get code review with retry logic.
|
||||
|
||||
Args:
|
||||
prompt: The prompt to send to the API.
|
||||
api_key: The GLM API key.
|
||||
config: Configuration object.
|
||||
|
||||
Returns:
|
||||
The review content as a string, or None if all retries failed.
|
||||
"""
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"model": config.glm_model,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"temperature": 0.3,
|
||||
"max_tokens": 32000,
|
||||
"top_p": 0.9,
|
||||
}
|
||||
|
||||
last_error: str | None = None
|
||||
|
||||
for attempt in range(config.max_retries):
|
||||
try:
|
||||
logger.info(f"Calling GLM API ({config.glm_model}) - Attempt {attempt + 1}/{config.max_retries}")
|
||||
response = requests.post(
|
||||
config.glm_api_url,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=120,
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
review = validate_glm_response(data)
|
||||
if review:
|
||||
logger.info(f"Successfully received review ({len(review)} chars)") # noqa: E501
|
||||
return review
|
||||
else:
|
||||
last_error = "Failed to validate API response structure"
|
||||
logger.error(last_error)
|
||||
else:
|
||||
last_error = f"GLM API request failed: {response.status_code} - {response.text}"
|
||||
logger.error(last_error)
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
last_error = f"GLM API request timed out (attempt {attempt + 1})"
|
||||
logger.error(last_error)
|
||||
except requests.exceptions.RequestException as e:
|
||||
last_error = f"GLM API request exception: {e}"
|
||||
logger.error(last_error)
|
||||
except json.JSONDecodeError as e:
|
||||
last_error = f"Failed to decode GLM API response as JSON: {e}"
|
||||
logger.error(last_error)
|
||||
|
||||
# Exponential backoff before retry
|
||||
if attempt < config.max_retries - 1:
|
||||
wait_time: float = config.retry_delay * (2**attempt)
|
||||
logger.info(f"Waiting {wait_time}s before retry...") # noqa: E501
|
||||
time.sleep(wait_time)
|
||||
|
||||
logger.error(
|
||||
f"All {config.max_retries} attempts failed. Last error: {last_error}" # noqa: E501
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def check_cooldown( # noqa: E501
|
||||
repo_name: str,
|
||||
pr_number: int,
|
||||
token: str,
|
||||
cooldown_minutes: int,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if the PR is within the cooldown period.
|
||||
|
||||
Args:
|
||||
repo_name: The repository name in format "owner/repo".
|
||||
pr_number: The pull request number.
|
||||
token: GitHub authentication token.
|
||||
cooldown_minutes: Cooldown period in minutes.
|
||||
|
||||
Returns:
|
||||
True if within cooldown period (should skip), False otherwise.
|
||||
"""
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
url: str = (
|
||||
f"https://api.github.com/repos/{repo_name}/issues/"
|
||||
f"{pr_number}/comments" # noqa: E501
|
||||
)
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/vnd.github.v3+json",
|
||||
}
|
||||
|
||||
logger.info(f"Checking cooldown period ({cooldown_minutes} minutes)")
|
||||
response = requests.get(url, headers=headers, timeout=30)
|
||||
|
||||
if response.status_code != 200:
|
||||
logger.warning(f"Failed to check cooldown: {response.status_code}, proceeding with review")
|
||||
return False
|
||||
|
||||
comments: list[dict[str, Any]] = response.json()
|
||||
cutoff_time: datetime = datetime.utcnow() - timedelta(minutes=cooldown_minutes) # noqa: E501
|
||||
|
||||
for comment in reversed(comments):
|
||||
# Check if this is a bot comment
|
||||
body: str = comment.get("body", "")
|
||||
if "VLLM-Omni PR Review" in body or "PR Reviewer Bot" in body:
|
||||
created_at_str: str = comment.get("created_at", "")
|
||||
try:
|
||||
# Parse GitHub timestamp format
|
||||
created_at = datetime.fromisoformat(created_at_str.replace("Z", "+00:00"))
|
||||
created_at = created_at.replace(tzinfo=None)
|
||||
if created_at > cutoff_time:
|
||||
logger.info(f"PR is within cooldown period (last review: {created_at_str})")
|
||||
return True
|
||||
except ValueError:
|
||||
logger.warning(f"Failed to parse comment timestamp: {created_at_str}") # noqa: E501
|
||||
continue
|
||||
|
||||
logger.info("PR is outside cooldown period, proceeding with review")
|
||||
return False
|
||||
|
||||
|
||||
def post_review_comment( # noqa: E501
|
||||
repo_name: str,
|
||||
pr_number: int,
|
||||
token: str,
|
||||
review: str,
|
||||
) -> bool:
|
||||
"""
|
||||
Post the review as a comment on the PR.
|
||||
|
||||
Args:
|
||||
repo_name: The repository name in format "owner/repo".
|
||||
pr_number: The pull request number.
|
||||
token: GitHub authentication token.
|
||||
review: The review content to post.
|
||||
|
||||
Returns:
|
||||
True if posting succeeded, False otherwise.
|
||||
"""
|
||||
url: str = (
|
||||
f"https://api.github.com/repos/{repo_name}/issues/"
|
||||
f"{pr_number}/comments" # noqa: E501
|
||||
)
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/vnd.github.v3+json",
|
||||
}
|
||||
|
||||
# Format the review comment
|
||||
comment_body: str = f"""## 🤖 VLLM-Omni PR Review
|
||||
|
||||
{review}
|
||||
|
||||
---
|
||||
*This review was generated automatically by the VLLM-Omni PR Reviewer Bot
|
||||
using {os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL)}.*
|
||||
"""
|
||||
|
||||
payload: dict[str, str] = {"body": comment_body}
|
||||
|
||||
logger.info(f"Posting review comment to PR #{pr_number}")
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=30)
|
||||
|
||||
if response.status_code == 201:
|
||||
logger.info("Successfully posted review comment")
|
||||
return True
|
||||
else:
|
||||
logger.error(f"Failed to post comment: {response.status_code}")
|
||||
logger.error(f"Response: {response.text}")
|
||||
return False
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""
|
||||
Main entry point for the PR reviewer bot.
|
||||
|
||||
Returns:
|
||||
0 on success, 1 on error.
|
||||
"""
|
||||
logger.info("VLLM-Omni PR Reviewer Bot starting...")
|
||||
|
||||
# Load configuration
|
||||
config: Config = get_config()
|
||||
logger.info(
|
||||
f"Configuration: model={config.glm_model}, "
|
||||
f"cooldown={config.cooldown_minutes}min, "
|
||||
f"max_retries={config.max_retries}"
|
||||
)
|
||||
|
||||
# Get environment variables
|
||||
token: str = get_env_var("GITHUB_TOKEN")
|
||||
api_key: str = get_env_var("GLM_API_KEY")
|
||||
repo_name: str = get_env_var("REPO_NAME")
|
||||
pr_number_str: str = get_env_var("PR_NUMBER")
|
||||
comment_body: str = get_env_var("COMMENT_BODY")
|
||||
|
||||
try:
|
||||
pr_number: int = int(pr_number_str)
|
||||
except ValueError:
|
||||
logger.error(f"Invalid PR number: {pr_number_str}")
|
||||
return 1
|
||||
|
||||
logger.info(f"Repository: {repo_name}")
|
||||
logger.info(f"PR Number: {pr_number}")
|
||||
|
||||
# Check if the comment contains the trigger phrase
|
||||
if not check_trigger(comment_body):
|
||||
logger.info(
|
||||
f"Comment does not contain trigger phrase '{TRIGGER_PHRASE}', exiting" # noqa: E501
|
||||
)
|
||||
return 0
|
||||
|
||||
logger.info("Trigger phrase detected! Starting review process...")
|
||||
|
||||
# Check cooldown period
|
||||
if check_cooldown(repo_name, pr_number, token, config.cooldown_minutes):
|
||||
logger.info("Skipping review due to cooldown period")
|
||||
return 0
|
||||
|
||||
# Fetch PR details
|
||||
logger.info("Step 1/4: Fetching PR details...")
|
||||
pr_details: PRDetails | None = fetch_pr_details(repo_name, pr_number, token) # noqa: E501
|
||||
if not pr_details:
|
||||
logger.error("Failed to fetch PR details")
|
||||
return 1
|
||||
|
||||
pr_title: str = pr_details.get("title", "Unknown")
|
||||
pr_description: str = pr_details.get("body", "")
|
||||
|
||||
logger.info(f"PR Title: {pr_title}")
|
||||
|
||||
# Fetch PR diff
|
||||
logger.info("Step 2/4: Fetching PR diff...")
|
||||
diff: str | None = fetch_pr_diff(repo_name, pr_number, token, config.max_diff_size)
|
||||
if diff is None:
|
||||
logger.error("Failed to fetch PR diff")
|
||||
return 1
|
||||
|
||||
if not diff:
|
||||
logger.warning("Warning: Empty diff - this might be a draft PR or no code changes")
|
||||
|
||||
# Build prompt
|
||||
logger.info("Step 3/4: Building review prompt...")
|
||||
prompt: str = build_review_prompt(pr_title, pr_description, diff)
|
||||
|
||||
# Call GLM API
|
||||
logger.info("Step 4/4: Calling GLM API...")
|
||||
review: str | None = call_glm_api(prompt, api_key, config)
|
||||
if not review:
|
||||
logger.error("Failed to get review from GLM API")
|
||||
return 1
|
||||
|
||||
# Post review comment
|
||||
logger.info("Posting review comment...")
|
||||
if not post_review_comment(repo_name, pr_number, token, review):
|
||||
logger.error("Failed to post review comment")
|
||||
return 1
|
||||
|
||||
logger.info("PR review completed successfully!")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -1,62 +0,0 @@
|
||||
name: VLLM-Omni PR Reviewer
|
||||
|
||||
on:
|
||||
issue_comment:
|
||||
types: [created]
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
issues: write
|
||||
|
||||
jobs:
|
||||
pr-reviewer:
|
||||
name: Review Pull Request
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
# Only run when the comment is from a collaborator/owner/member
|
||||
if: |
|
||||
github.event_name == 'issue_comment' &&
|
||||
github.event.issue.pull_request != null &&
|
||||
contains(github.event.comment.body, '@vllm-omni-reviewer') &&
|
||||
(github.event.comment.author_association == 'MEMBER' ||
|
||||
github.event.comment.author_association == 'COLLABORATOR' ||
|
||||
github.event.comment.author_association == 'OWNER')
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6.0.2
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6.2.0
|
||||
with:
|
||||
python-version: '3.11'
|
||||
cache: 'pip'
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip==24.0
|
||||
pip install requests==2.31.0 pyyaml==6.0.1
|
||||
|
||||
- name: Run PR Reviewer
|
||||
id: reviewer
|
||||
env:
|
||||
GLM_API_KEY: ${{ secrets.GLM_API_KEY }}
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
PR_NUMBER: ${{ github.event.issue.number || github.event.pull_request.number }}
|
||||
COMMENT_BODY: ${{ github.event.comment.body }}
|
||||
REPO_NAME: ${{ github.repository }}
|
||||
PR_REVIEWER_COOLDOWN_MINUTES: 5
|
||||
PR_REVIEWER_MAX_RETRIES: 3
|
||||
run: |
|
||||
python .github/scripts/pr_reviewer.py 2>&1 | tee "/tmp/pr_review_${PR_NUMBER}.log"
|
||||
|
||||
- name: Upload review logs
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v7.0.0
|
||||
with:
|
||||
name: pr-review-logs-${{ github.event.issue.number || github.event.pull_request.number }}
|
||||
path: /tmp/pr_review_*.log
|
||||
retention-days: 7
|
||||
if-no-files-found: ignore
|
||||
@@ -174,6 +174,7 @@ CLAUDE.md
|
||||
|
||||
# Codex
|
||||
AGENTS.md
|
||||
.codex
|
||||
.codex/
|
||||
|
||||
# cursor
|
||||
@@ -203,6 +204,7 @@ checkpoints/
|
||||
# Cache directories
|
||||
cache/
|
||||
!vllm_omni/diffusion/cache/
|
||||
!tests/diffusion/cache/
|
||||
.cache/
|
||||
diffusion_cache/
|
||||
kv_cache/
|
||||
@@ -262,3 +264,6 @@ tmp_test
|
||||
vllm_omni/_version.py
|
||||
# output files
|
||||
*.wav
|
||||
.worktrees/
|
||||
# CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
|
||||
tests/.ci_generated/
|
||||
@@ -35,8 +35,8 @@ MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
|
||||
# Use a Voice Clone model
|
||||
MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-Base TASK_TYPE=Base bash run_benchmark.sh --async-only
|
||||
|
||||
# Use bs16 config for higher throughput
|
||||
STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs16.yaml bash run_benchmark.sh --async-only
|
||||
# Use batch size 16 for higher throughput
|
||||
BATCH_SIZE=16 bash run_benchmark.sh --async-only
|
||||
|
||||
# Custom GPU, prompt count, concurrency levels
|
||||
GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
|
||||
@@ -50,7 +50,8 @@ GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
|
||||
CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
|
||||
"Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
|
||||
--omni --host 127.0.0.1 --port 8000 \
|
||||
--stage-configs-path benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml \
|
||||
--deploy-config vllm_omni/deploy/qwen3_tts.yaml \
|
||||
--stage-overrides '{"0":{"max_num_seqs":1,"gpu_memory_utilization":0.3,"max_num_batched_tokens":512},"1":{"max_num_seqs":1,"gpu_memory_utilization":0.3,"max_num_batched_tokens":8192}}' \
|
||||
--trust-remote-code
|
||||
```
|
||||
|
||||
@@ -84,16 +85,19 @@ python benchmarks/qwen3-tts/plot_results.py \
|
||||
--output results/comparison.png
|
||||
```
|
||||
|
||||
## Stage Configs
|
||||
## Batch-size presets
|
||||
|
||||
| Config | max_num_seqs | Description |
|
||||
|--------|:------------:|-------------|
|
||||
| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lowest latency) |
|
||||
| `vllm_omni/configs/qwen3_tts_bs16.yaml` | 16 | High-throughput concurrent processing |
|
||||
The bench script loads the bundled production deploy (`vllm_omni/deploy/qwen3_tts.yaml`) and layers per-stage budgets on top via `--stage-overrides`, driven by the `BATCH_SIZE` env var. Each batch size picks compatible per-stage `max_num_seqs`, `max_num_batched_tokens`, and `gpu_memory_utilization` defaults:
|
||||
|
||||
All configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
|
||||
| `BATCH_SIZE` | Description |
|
||||
|:--:|-------------|
|
||||
| `1` (default) | Single-request processing (lowest latency) |
|
||||
| `4` | Moderate-throughput concurrent processing |
|
||||
| `16` | High-throughput concurrent processing |
|
||||
|
||||
The model is specified via the CLI `--model` flag (or `MODEL` env var), so the same configs work for both the 0.6B and 1.7B model variants.
|
||||
The 2-stage pipeline (Talker -> Code2Wav) runs with `async_chunk` streaming enabled via the prod deploy; the `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
|
||||
|
||||
The model is specified via the CLI `--model` flag (or `MODEL` env var), so the same bench script works for both the 0.6B and 1.7B model variants.
|
||||
|
||||
## Metrics
|
||||
|
||||
|
||||
@@ -26,8 +26,8 @@
|
||||
# # Use Voice Clone model
|
||||
# MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-Base TASK_TYPE=Base bash run_benchmark.sh --async-only
|
||||
#
|
||||
# # Use batch_size=4 config:
|
||||
# STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
|
||||
# # Use batch_size=4:
|
||||
# BATCH_SIZE=4 bash run_benchmark.sh --async-only
|
||||
#
|
||||
# Environment variables:
|
||||
# GPU_DEVICE - GPU index to use (default: 0)
|
||||
@@ -35,9 +35,9 @@
|
||||
# CONCURRENCY - Space-separated concurrency levels (default: "1 4 10")
|
||||
# MODEL - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
|
||||
# PORT - Server port (default: 8000)
|
||||
# GPU_MEM_TALKER - gpu_memory_utilization for talker stage (default: 0.3)
|
||||
# GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.2)
|
||||
# STAGE_CONFIG - Path to stage config YAML (default: configs/qwen3_tts_bs1.yaml)
|
||||
# BATCH_SIZE - Per-stage ``max_num_seqs`` for both talker and code2wav (default: 1)
|
||||
# GPU_MEM_TALKER - gpu_memory_utilization for talker stage (default: 0.3 at bs=1, else 0.2)
|
||||
# GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.3 at bs=1, else 0.2)
|
||||
# TASK_TYPE - Task type: CustomVoice, VoiceDesign, Base (default: CustomVoice)
|
||||
|
||||
set -euo pipefail
|
||||
@@ -51,14 +51,36 @@ NUM_PROMPTS="${NUM_PROMPTS:-50}"
|
||||
CONCURRENCY="${CONCURRENCY:-1 4 10}"
|
||||
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
|
||||
PORT="${PORT:-8000}"
|
||||
GPU_MEM_TALKER="${GPU_MEM_TALKER:-0.3}"
|
||||
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-0.2}"
|
||||
BATCH_SIZE="${BATCH_SIZE:-1}"
|
||||
DEFAULT_MEM=$([ "${BATCH_SIZE}" = "1" ] && echo "0.3" || echo "0.2")
|
||||
GPU_MEM_TALKER="${GPU_MEM_TALKER:-${DEFAULT_MEM}}"
|
||||
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-${DEFAULT_MEM}}"
|
||||
NUM_WARMUPS="${NUM_WARMUPS:-3}"
|
||||
STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/configs/qwen3_tts_bs1.yaml}"
|
||||
DEPLOY_CONFIG="vllm_omni/deploy/qwen3_tts.yaml"
|
||||
RESULT_DIR="${SCRIPT_DIR}/results"
|
||||
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
||||
TASK_TYPE="${TASK_TYPE:-CustomVoice}"
|
||||
|
||||
# Build --stage-overrides JSON from BATCH_SIZE + GPU_MEM_*.
|
||||
STAGE_OVERRIDES=$(
|
||||
BATCH_SIZE="${BATCH_SIZE}" \
|
||||
GPU_MEM_TALKER="${GPU_MEM_TALKER}" \
|
||||
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV}" \
|
||||
python - <<'PYEOF'
|
||||
import json, os
|
||||
bs = int(os.environ["BATCH_SIZE"])
|
||||
mem_t = float(os.environ["GPU_MEM_TALKER"])
|
||||
mem_c = float(os.environ["GPU_MEM_CODE2WAV"])
|
||||
# Prefill budget grows with batch size on both stages.
|
||||
talker_batched = 512 if bs <= 4 else 4096
|
||||
code2wav_batched = 8192 if bs <= 4 else 32768
|
||||
print(json.dumps({
|
||||
"0": {"max_num_seqs": bs, "gpu_memory_utilization": mem_t, "max_num_batched_tokens": talker_batched},
|
||||
"1": {"max_num_seqs": bs, "gpu_memory_utilization": mem_c, "max_num_batched_tokens": code2wav_batched},
|
||||
}))
|
||||
PYEOF
|
||||
)
|
||||
|
||||
# Parse args
|
||||
RUN_ASYNC=true
|
||||
RUN_HF=true
|
||||
@@ -75,41 +97,27 @@ mkdir -p "${RESULT_DIR}"
|
||||
echo "============================================================"
|
||||
echo " Qwen3-TTS Benchmark"
|
||||
echo "============================================================"
|
||||
echo " GPU: ${GPU_DEVICE}"
|
||||
echo " Model: ${MODEL}"
|
||||
echo " Prompts: ${NUM_PROMPTS}"
|
||||
echo " Concurrency: ${CONCURRENCY}"
|
||||
echo " Port: ${PORT}"
|
||||
echo " Stage config: ${STAGE_CONFIG}"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo " Task type: ${TASK_TYPE}"
|
||||
echo " GPU: ${GPU_DEVICE}"
|
||||
echo " Model: ${MODEL}"
|
||||
echo " Prompts: ${NUM_PROMPTS}"
|
||||
echo " Concurrency: ${CONCURRENCY}"
|
||||
echo " Port: ${PORT}"
|
||||
echo " Deploy config: ${DEPLOY_CONFIG}"
|
||||
echo " Batch size: ${BATCH_SIZE}"
|
||||
echo " GPU mem T/C: ${GPU_MEM_TALKER} / ${GPU_MEM_CODE2WAV}"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo " Task type: ${TASK_TYPE}"
|
||||
echo "============================================================"
|
||||
|
||||
# Prepare stage config with correct GPU device and memory settings
|
||||
prepare_config() {
|
||||
local config_template="$1"
|
||||
local config_name="$2"
|
||||
local output_path="${RESULT_DIR}/${config_name}_stage_config.yaml"
|
||||
|
||||
# Use sed to patch GPU device and memory utilization
|
||||
sed \
|
||||
-e "s/devices: \"0\"/devices: \"${GPU_DEVICE}\"/g" \
|
||||
-e "s/gpu_memory_utilization: 0.3/gpu_memory_utilization: ${GPU_MEM_TALKER}/g" \
|
||||
-e "s/gpu_memory_utilization: 0.2/gpu_memory_utilization: ${GPU_MEM_CODE2WAV}/g" \
|
||||
"${config_template}" > "${output_path}"
|
||||
|
||||
echo "${output_path}"
|
||||
}
|
||||
|
||||
# Start server and wait for it to be ready
|
||||
start_server() {
|
||||
local stage_config="$1"
|
||||
local config_name="$2"
|
||||
local config_name="$1"
|
||||
local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log"
|
||||
|
||||
echo ""
|
||||
echo "Starting server with config: ${config_name}"
|
||||
echo " Stage config: ${stage_config}"
|
||||
echo " Deploy config: ${DEPLOY_CONFIG}"
|
||||
echo " Stage overrides: ${STAGE_OVERRIDES}"
|
||||
echo " Log file: ${log_file}"
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
@@ -118,7 +126,8 @@ start_server() {
|
||||
--omni \
|
||||
--host 127.0.0.1 \
|
||||
--port "${PORT}" \
|
||||
--stage-configs-path "${stage_config}" \
|
||||
--deploy-config "${DEPLOY_CONFIG}" \
|
||||
--stage-overrides "${STAGE_OVERRIDES}" \
|
||||
--stage-init-timeout 120 \
|
||||
--trust-remote-code \
|
||||
--disable-log-stats \
|
||||
@@ -175,17 +184,13 @@ trap 'stop_server' EXIT
|
||||
# Run benchmark for a given config
|
||||
run_bench() {
|
||||
local config_name="$1"
|
||||
local config_template="$2"
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmarking: ${config_name}"
|
||||
echo "============================================================"
|
||||
|
||||
local stage_config
|
||||
stage_config=$(prepare_config "${config_template}" "${config_name}")
|
||||
|
||||
start_server "${stage_config}" "${config_name}"
|
||||
start_server "${config_name}"
|
||||
|
||||
# Convert concurrency string to args
|
||||
local conc_args=""
|
||||
@@ -212,7 +217,7 @@ run_bench() {
|
||||
|
||||
# Run vllm-omni benchmark
|
||||
if [ "${RUN_ASYNC}" = true ]; then
|
||||
run_bench "async_chunk" "${SCRIPT_DIR}/${STAGE_CONFIG}"
|
||||
run_bench "async_chunk"
|
||||
fi
|
||||
|
||||
# Run HuggingFace baseline benchmark
|
||||
|
||||
@@ -1,94 +0,0 @@
|
||||
# Qwen3-TTS max_num_seqs=16 config (streaming with async_chunk)
|
||||
# High-throughput concurrent request processing
|
||||
# 2-stage pipeline: Talker -> Code2Wav
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_comprehension: true
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 16
|
||||
model_stage: qwen3_tts
|
||||
model_arch: Qwen3TTSTalkerForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
enforce_eager: false
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent
|
||||
gpu_memory_utilization: 0.3
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 4096
|
||||
max_model_len: 4096
|
||||
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
detokenize: false
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [2150]
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 16
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen3TTSCode2Wav
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio
|
||||
gpu_memory_utilization: 0.2
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 16384
|
||||
max_model_len: 32768
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
tts_args:
|
||||
max_instructions_length: 500
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
detokenize: true
|
||||
repetition_penalty: 1.0
|
||||
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 16
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 25
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
@@ -1,94 +0,0 @@
|
||||
# Qwen3-TTS batch_size=4 config (streaming with async_chunk)
|
||||
# Enables concurrent request processing
|
||||
# 2-stage pipeline: Talker -> Code2Wav
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_comprehension: true
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 4
|
||||
model_stage: qwen3_tts
|
||||
model_arch: Qwen3TTSTalkerForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
enforce_eager: false
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent
|
||||
gpu_memory_utilization: 0.3
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 512
|
||||
max_model_len: 4096
|
||||
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
detokenize: false
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [2150]
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 4
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen3TTSCode2Wav
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio
|
||||
gpu_memory_utilization: 0.2
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 8192
|
||||
max_model_len: 32768
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
tts_args:
|
||||
max_instructions_length: 500
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
detokenize: true
|
||||
repetition_penalty: 1.0
|
||||
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 4
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 25
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
@@ -31,8 +31,11 @@ PORT_OFF="${PORT_OFF:-8001}"
|
||||
RESULT_DIR="${SCRIPT_DIR}/results"
|
||||
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
||||
|
||||
STAGE_CONFIG_ON="vllm_omni/model_executor/stage_configs/qwen3_tts.yaml"
|
||||
STAGE_CONFIG_OFF="vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml"
|
||||
# The bundled ``vllm_omni/deploy/qwen3_tts.yaml`` is auto-loaded by the model
|
||||
# registry; no ``--deploy-config`` flag needed on the default (ON) path.
|
||||
# async_chunk OFF is selected by the ``--no-async-chunk`` CLI flag —
|
||||
# the single ``qwen3_tts`` pipeline dispatches to the end-to-end codec
|
||||
# processor when ``deploy.async_chunk`` is false.
|
||||
|
||||
mkdir -p "${RESULT_DIR}"
|
||||
|
||||
@@ -77,7 +80,6 @@ wait_for_server() {
|
||||
echo ""
|
||||
echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..."
|
||||
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
|
||||
--stage-configs-path "${STAGE_CONFIG_ON}" \
|
||||
--host 0.0.0.0 --port "${PORT_ON}" \
|
||||
--trust-remote-code --enforce-eager --omni \
|
||||
> "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 &
|
||||
@@ -104,7 +106,7 @@ sleep 5
|
||||
echo ""
|
||||
echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..."
|
||||
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
|
||||
--stage-configs-path "${STAGE_CONFIG_OFF}" \
|
||||
--no-async-chunk \
|
||||
--host 0.0.0.0 --port "${PORT_OFF}" \
|
||||
--trust-remote-code --enforce-eager --omni \
|
||||
> "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 &
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
# VoxCPM Benchmark
|
||||
|
||||
This directory contains both:
|
||||
|
||||
- online serving benchmark through the OpenAI-compatible `/v1/audio/speech` API
|
||||
- offline benchmark for `Omni` / `AsyncOmni`
|
||||
- full offline smoke-matrix orchestration
|
||||
|
||||
Both benchmark paths report:
|
||||
|
||||
- TTFP: time to first PCM packet
|
||||
- E2E latency
|
||||
- RTF: real-time factor (`e2e / audio_duration`)
|
||||
|
||||
## Offline Benchmark
|
||||
|
||||
Single offline benchmark run:
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/bench_tts_offline.py \
|
||||
--model /path/to/voxcpm-model \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm.yaml \
|
||||
--text "This is a split-stage VoxCPM synthesis example running on vLLM Omni." \
|
||||
--warmup-runs 1 \
|
||||
--output-dir benchmarks/voxcpm/results/offline_single
|
||||
```
|
||||
|
||||
Streaming offline benchmark:
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/bench_tts_offline.py \
|
||||
--model /path/to/voxcpm-model \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
|
||||
--text "This is a split-stage VoxCPM streaming example running on vLLM Omni." \
|
||||
--warmup-runs 1 \
|
||||
--output-dir benchmarks/voxcpm/results/offline_streaming
|
||||
```
|
||||
|
||||
Full fixed offline matrix, equivalent to the old `examples/offline_inference/voxcpm/test.py`:
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/run_offline_matrix.py \
|
||||
--model /path/to/voxcpm-model \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "The exact transcript spoken in reference.wav." \
|
||||
--output-root benchmarks/voxcpm/results/offline_matrix
|
||||
```
|
||||
|
||||
The full matrix covers both routes:
|
||||
|
||||
- streaming: `voxcpm_async_chunk.yaml`
|
||||
- sync: `voxcpm.yaml`
|
||||
|
||||
And these six scenarios under each route:
|
||||
|
||||
- warmup + single TTS
|
||||
- warmup + single voice cloning
|
||||
- warmup + batch TTS
|
||||
- warmup + batch voice cloning
|
||||
- cold single TTS
|
||||
- cold single voice cloning
|
||||
|
||||
`bench_tts_offline.py` itself no longer writes `summary.json` / `results.json`; it prints TTFP / RTF inline and saves generated WAV files only. The matrix runner keeps only per-case `run.log`.
|
||||
|
||||
## Start the Server
|
||||
|
||||
Async-chunk:
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/voxcpm-model \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
Non-streaming:
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/voxcpm-model \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm.yaml \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
## Run the Benchmark
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/bench_tts_serve.py \
|
||||
--host 127.0.0.1 \
|
||||
--port 8091 \
|
||||
--num-prompts 20 \
|
||||
--max-concurrency 1 \
|
||||
--result-dir /tmp/voxcpm_bench
|
||||
```
|
||||
|
||||
Voice cloning benchmark:
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/bench_tts_serve.py \
|
||||
--host 127.0.0.1 \
|
||||
--port 8091 \
|
||||
--num-prompts 10 \
|
||||
--max-concurrency 1 \
|
||||
--ref-audio https://example.com/reference.wav \
|
||||
--ref-text "The exact transcript spoken in the reference audio." \
|
||||
--result-dir /tmp/voxcpm_clone_bench
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
- The benchmark uses `stream=true` and `response_format=pcm` so TTFP is measured from the first audio packet.
|
||||
- `RTF < 1.0` means the server generates audio faster than real time.
|
||||
- For `voxcpm_async_chunk.yaml`, keep concurrency at `1`. This matches native VoxCPM streaming more closely.
|
||||
- Do not benchmark concurrent online streaming on `voxcpm_async_chunk.yaml`; use `voxcpm.yaml` for multi-request throughput runs.
|
||||
- For the offline matrix mode, `--ref-audio` and `--ref-text` are required because clone cases are part of the fixed coverage set.
|
||||
@@ -0,0 +1,890 @@
|
||||
"""Offline VoxCPM benchmark for vLLM Omni.
|
||||
|
||||
Supports both:
|
||||
- sync one-shot (Omni.generate)
|
||||
- streaming (AsyncOmni.generate with async_chunk config)
|
||||
- text-only synthesis
|
||||
- voice cloning
|
||||
- text/clone batch inputs from txt or jsonl
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import AsyncOmni, Omni
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml"
|
||||
DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml"
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PromptSpec:
|
||||
text: str
|
||||
label: str
|
||||
ref_audio: str | None = None
|
||||
ref_text: str | None = None
|
||||
|
||||
|
||||
def _require_soundfile():
|
||||
try:
|
||||
import soundfile as sf # type: ignore
|
||||
except ModuleNotFoundError as exc:
|
||||
raise RuntimeError(
|
||||
"soundfile is required to write VoxCPM benchmark WAV outputs. Install it with: pip install soundfile"
|
||||
) from exc
|
||||
return sf
|
||||
|
||||
|
||||
def _build_prompt(
|
||||
args,
|
||||
*,
|
||||
text: str,
|
||||
ref_audio: str | None = None,
|
||||
ref_text: str | None = None,
|
||||
global_request_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
additional_information: dict[str, list[Any]] = {
|
||||
"text": [text],
|
||||
"cfg_value": [args.cfg_value],
|
||||
"inference_timesteps": [args.inference_timesteps],
|
||||
"min_len": [args.min_len],
|
||||
"max_new_tokens": [args.max_new_tokens],
|
||||
}
|
||||
if args.streaming_prefix_len is not None:
|
||||
additional_information["streaming_prefix_len"] = [args.streaming_prefix_len]
|
||||
|
||||
if ref_audio:
|
||||
additional_information["ref_audio"] = [ref_audio]
|
||||
if ref_text:
|
||||
additional_information["ref_text"] = [ref_text]
|
||||
if global_request_id is not None:
|
||||
additional_information["global_request_id"] = [global_request_id]
|
||||
|
||||
return {
|
||||
"prompt_token_ids": [1],
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
|
||||
|
||||
def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor:
|
||||
audio = mm.get("audio", mm.get("model_outputs"))
|
||||
if audio is None:
|
||||
raise ValueError("No audio output found in multimodal output.")
|
||||
if isinstance(audio, list):
|
||||
parts = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio]
|
||||
audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0)
|
||||
if not isinstance(audio, torch.Tensor):
|
||||
audio = torch.as_tensor(audio)
|
||||
return audio.float().cpu().reshape(-1)
|
||||
|
||||
|
||||
def _extract_sample_rate(mm: dict[str, Any]) -> int:
|
||||
sr_raw = mm.get("sr", 24000)
|
||||
if isinstance(sr_raw, list) and sr_raw:
|
||||
sr_raw = sr_raw[-1]
|
||||
if hasattr(sr_raw, "item"):
|
||||
return int(sr_raw.item())
|
||||
return int(sr_raw)
|
||||
|
||||
|
||||
def _emit_offline_metrics(
|
||||
*,
|
||||
request_id: str,
|
||||
elapsed_s: float,
|
||||
first_audio_elapsed: float | None,
|
||||
audio_duration_s: float,
|
||||
) -> None:
|
||||
metrics = {
|
||||
"request_id": request_id,
|
||||
"ttfp_ms": round(first_audio_elapsed * 1000.0, 3) if first_audio_elapsed is not None else None,
|
||||
"audio_duration_s": round(audio_duration_s, 6),
|
||||
"rtf": round(elapsed_s / audio_duration_s, 6) if audio_duration_s > 0 else None,
|
||||
}
|
||||
print(f"[OfflineMetrics] {metrics}")
|
||||
|
||||
|
||||
def _write_audio_tensor(output_path: Path, audio_tensor: Any, sample_rate: int) -> None:
|
||||
sf = _require_soundfile()
|
||||
if isinstance(audio_tensor, torch.Tensor):
|
||||
audio_np = audio_tensor.float().cpu().clamp(-1.0, 1.0).numpy()
|
||||
else:
|
||||
audio_np = torch.as_tensor(audio_tensor).float().cpu().clamp(-1.0, 1.0).numpy()
|
||||
sf.write(
|
||||
output_path,
|
||||
audio_np,
|
||||
sample_rate,
|
||||
format="WAV",
|
||||
subtype="PCM_16",
|
||||
)
|
||||
|
||||
|
||||
def _save_wav(mm: dict[str, Any], output_dir: Path, request_id: str) -> Path:
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
output_path = output_dir / f"output_{request_id}.wav"
|
||||
_write_audio_tensor(output_path, _extract_audio_tensor(mm), _extract_sample_rate(mm))
|
||||
return output_path
|
||||
|
||||
|
||||
def _iter_request_multimodal_outputs(request_output: Any):
|
||||
outputs = getattr(request_output, "outputs", None)
|
||||
if outputs:
|
||||
for output in outputs:
|
||||
mm = getattr(output, "multimodal_output", None)
|
||||
if isinstance(mm, dict):
|
||||
yield mm
|
||||
|
||||
mm = getattr(request_output, "multimodal_output", None)
|
||||
if isinstance(mm, dict):
|
||||
yield mm
|
||||
|
||||
|
||||
def _read_non_empty_lines(path: str) -> list[str]:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
return [line.strip() for line in f if line.strip()]
|
||||
|
||||
|
||||
def _load_prompt_specs(args) -> list[PromptSpec]:
|
||||
specs: list[PromptSpec] = []
|
||||
|
||||
if args.txt_prompts is not None:
|
||||
texts = _read_non_empty_lines(args.txt_prompts)
|
||||
if not texts:
|
||||
raise ValueError(f"No prompts found in {args.txt_prompts}")
|
||||
for idx, text in enumerate(texts, start=1):
|
||||
specs.append(
|
||||
PromptSpec(
|
||||
text=text,
|
||||
label=f"item{idx:03d}",
|
||||
ref_audio=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
)
|
||||
)
|
||||
return specs
|
||||
|
||||
if args.jsonl_prompts is not None:
|
||||
with open(args.jsonl_prompts, encoding="utf-8") as f:
|
||||
for line_no, raw_line in enumerate(f, start=1):
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
item = json.loads(line)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"{args.jsonl_prompts}:{line_no} is not valid JSON: {exc}") from exc
|
||||
if not isinstance(item, dict):
|
||||
raise ValueError(f"{args.jsonl_prompts}:{line_no} must be a JSON object")
|
||||
|
||||
text = item.get("text")
|
||||
if not isinstance(text, str) or not text.strip():
|
||||
raise ValueError(f"{args.jsonl_prompts}:{line_no} requires non-empty string field 'text'")
|
||||
|
||||
ref_audio = item.get("ref_audio", args.ref_audio)
|
||||
ref_text = item.get("ref_text", args.ref_text)
|
||||
if (ref_audio is None) != (ref_text is None):
|
||||
raise ValueError(
|
||||
f"{args.jsonl_prompts}:{line_no} must provide both 'ref_audio' and 'ref_text' together"
|
||||
)
|
||||
|
||||
specs.append(
|
||||
PromptSpec(
|
||||
text=text.strip(),
|
||||
label=f"item{len(specs) + 1:03d}",
|
||||
ref_audio=ref_audio,
|
||||
ref_text=ref_text,
|
||||
)
|
||||
)
|
||||
|
||||
if not specs:
|
||||
raise ValueError(f"No prompts found in {args.jsonl_prompts}")
|
||||
return specs
|
||||
|
||||
specs.append(
|
||||
PromptSpec(
|
||||
text=args.text,
|
||||
label="item001",
|
||||
ref_audio=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
)
|
||||
)
|
||||
return specs
|
||||
|
||||
|
||||
def _build_prompt_for_spec(args, spec: PromptSpec, *, global_request_id: str | None = None) -> dict[str, Any]:
|
||||
return _build_prompt(
|
||||
args,
|
||||
text=spec.text,
|
||||
ref_audio=spec.ref_audio,
|
||||
ref_text=spec.ref_text,
|
||||
global_request_id=global_request_id,
|
||||
)
|
||||
|
||||
|
||||
def _count_voice_clone_prompts(prompt_specs: list[PromptSpec]) -> int:
|
||||
return sum(1 for spec in prompt_specs if spec.ref_audio is not None)
|
||||
|
||||
|
||||
def _get_warmup_specs(prompt_specs: list[PromptSpec]) -> list[PromptSpec]:
|
||||
return prompt_specs[:1]
|
||||
|
||||
|
||||
def _extract_stream_finished(stage_output: Any) -> bool:
|
||||
request_output = getattr(stage_output, "request_output", None)
|
||||
request_finished = getattr(request_output, "finished", None)
|
||||
if request_finished is not None:
|
||||
return bool(request_finished)
|
||||
return bool(getattr(stage_output, "finished", False))
|
||||
|
||||
|
||||
def _build_profiled_stage_config(
|
||||
stage_configs_path: str,
|
||||
profiler_dir: str,
|
||||
) -> str:
|
||||
stage_config_path = Path(stage_configs_path)
|
||||
yaml_text = stage_config_path.read_text(encoding="utf-8")
|
||||
injected_lines: list[str] = []
|
||||
injected_count = 0
|
||||
|
||||
for line in yaml_text.splitlines():
|
||||
injected_lines.append(line)
|
||||
if line.strip() != "engine_args:":
|
||||
continue
|
||||
indent = line[: len(line) - len(line.lstrip())]
|
||||
child_indent = indent + " "
|
||||
grandchild_indent = child_indent + " "
|
||||
injected_lines.extend(
|
||||
[
|
||||
f"{child_indent}profiler_config:",
|
||||
f'{grandchild_indent}profiler: "torch"',
|
||||
f'{grandchild_indent}torch_profiler_dir: "{profiler_dir}"',
|
||||
f"{grandchild_indent}torch_profiler_with_stack: true",
|
||||
]
|
||||
)
|
||||
injected_count += 1
|
||||
|
||||
if injected_count == 0:
|
||||
raise ValueError(f"No engine_args block found in stage config: {stage_configs_path}")
|
||||
|
||||
tmp = tempfile.NamedTemporaryFile(
|
||||
mode="w",
|
||||
encoding="utf-8",
|
||||
delete=False,
|
||||
suffix=".yaml",
|
||||
prefix=f"{stage_config_path.stem}_profile_",
|
||||
)
|
||||
tmp.write("\n".join(injected_lines) + "\n")
|
||||
tmp.close()
|
||||
return tmp.name
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(
|
||||
description="Offline split-stage VoxCPM inference with vLLM Omni (auto sync/streaming by stage config)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default=os.environ.get("VOXCPM_MODEL"),
|
||||
help="Local VoxCPM model directory. Defaults to $VOXCPM_MODEL.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
type=str,
|
||||
default="This is a split-stage VoxCPM synthesis example running on vLLM Omni.",
|
||||
help="Text to synthesize. Ignored when --txt-prompts or --jsonl-prompts is used.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--txt-prompts",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a .txt file with one synthesis text per line.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--jsonl-prompts",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Path to a .jsonl file. Each line must contain at least {'text': ...}; "
|
||||
"clone rows can also set ref_audio/ref_text, and ref_text must be the "
|
||||
"real transcript of ref_audio."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Optional reference audio path for voice cloning. With --txt-prompts, "
|
||||
"the same reference is applied to every line."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-text",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Real transcript of the reference audio. Placeholder text or mismatched "
|
||||
"text will usually produce noisy/electronic clone audio."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=str(DEFAULT_STAGE_SYNC),
|
||||
help="Stage config YAML path. Routing is selected only from this path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg-value",
|
||||
type=float,
|
||||
default=2.0,
|
||||
help="Classifier-free guidance value for VoxCPM.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--inference-timesteps",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Number of inference timesteps.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min-len",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Minimum generated token length.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-new-tokens",
|
||||
type=int,
|
||||
default=4096,
|
||||
help="Maximum generated token length.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--streaming-prefix-len",
|
||||
type=int,
|
||||
default=None,
|
||||
help="VoxCPM streaming window (optional, streaming mode only).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory for output WAV files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-init-timeout",
|
||||
type=int,
|
||||
default=600,
|
||||
help="Stage initialization timeout in seconds.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
dest="log_stats",
|
||||
action="store_true",
|
||||
help="Enable vLLM Omni stats logging.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-log-stats",
|
||||
dest="log_stats",
|
||||
action="store_false",
|
||||
help="Disable vLLM Omni stats logging.",
|
||||
)
|
||||
parser.set_defaults(log_stats=True)
|
||||
parser.add_argument(
|
||||
"--num-runs",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of full inference runs (same prompt each time). Default 1.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup-runs",
|
||||
type=int,
|
||||
default=0,
|
||||
help=(
|
||||
"Optional number of warmup passes before measured runs. Warmup uses only "
|
||||
"the first prompt and does not save outputs."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-profiler",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Enable torch profiler for the configured stages. A temporary profiled "
|
||||
"stage config is generated automatically."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory for profiler traces. Defaults to <output-dir>/profiler when profiling is enabled.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-stages",
|
||||
type=int,
|
||||
nargs="*",
|
||||
default=None,
|
||||
help="Optional stage ids to profile. Defaults to all stages that have profiler_config.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-wait-seconds",
|
||||
type=float,
|
||||
default=30.0,
|
||||
help="Seconds to wait after stop_profile for trace files to flush.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.model:
|
||||
parser.error("--model is required unless $VOXCPM_MODEL is set")
|
||||
if args.txt_prompts is not None and args.jsonl_prompts is not None:
|
||||
parser.error("--txt-prompts and --jsonl-prompts are mutually exclusive")
|
||||
if (args.ref_audio is None) != (args.ref_text is None):
|
||||
parser.error("--ref-audio and --ref-text must be provided together")
|
||||
if args.num_runs < 1:
|
||||
parser.error("--num-runs must be >= 1")
|
||||
if args.warmup_runs < 0:
|
||||
parser.error("--warmup-runs must be >= 0")
|
||||
if args.output_dir is None:
|
||||
args.output_dir = (
|
||||
"output_audio_streaming" if _is_streaming_stage_config(args.stage_configs_path) else "output_audio"
|
||||
)
|
||||
if args.enable_profiler and args.profiler_dir is None:
|
||||
args.profiler_dir = str(Path(args.output_dir) / "profiler")
|
||||
try:
|
||||
args.prompt_specs = _load_prompt_specs(args)
|
||||
except ValueError as exc:
|
||||
parser.error(str(exc))
|
||||
|
||||
return args
|
||||
|
||||
|
||||
def _is_streaming_stage_config(stage_configs_path: str) -> bool:
|
||||
cfg_name = Path(stage_configs_path).name.lower()
|
||||
# Keep routing purely config-path based:
|
||||
# - voxcpm.yaml => sync
|
||||
# - voxcpm_async_chunk.yaml => streaming
|
||||
return "async_chunk" in cfg_name
|
||||
|
||||
|
||||
async def _collect_streaming_audio(
|
||||
omni: AsyncOmni,
|
||||
args: Any,
|
||||
spec: PromptSpec,
|
||||
request_id: str,
|
||||
*,
|
||||
phase_label: str,
|
||||
prompt_index: int,
|
||||
prompt_count: int,
|
||||
print_prompt: bool = False,
|
||||
) -> tuple[torch.Tensor, int, float, float | None]:
|
||||
prompt = _build_prompt_for_spec(args, spec, global_request_id=request_id)
|
||||
delta_chunks: list[torch.Tensor] = []
|
||||
sample_rate = 24000
|
||||
chunk_i = 0
|
||||
prev_total_samples = 0
|
||||
t_start = time.perf_counter()
|
||||
first_audio_elapsed: float | None = None
|
||||
|
||||
if print_prompt:
|
||||
print(f"---prompt---:{prompt}")
|
||||
|
||||
async for stage_output in omni.generate(prompt, request_id=request_id):
|
||||
mm = getattr(stage_output, "multimodal_output", None)
|
||||
if not isinstance(mm, dict):
|
||||
ro = getattr(stage_output, "request_output", None)
|
||||
if ro is None:
|
||||
continue
|
||||
mm = getattr(ro, "multimodal_output", None)
|
||||
if not isinstance(mm, dict) and getattr(ro, "outputs", None):
|
||||
seq = ro.outputs[0]
|
||||
mm = getattr(seq, "multimodal_output", None)
|
||||
if not isinstance(mm, dict):
|
||||
continue
|
||||
sample_rate = _extract_sample_rate(mm)
|
||||
try:
|
||||
w = _extract_audio_tensor(mm)
|
||||
n = int(w.numel())
|
||||
if n == 0:
|
||||
continue
|
||||
finished = _extract_stream_finished(stage_output)
|
||||
if n > prev_total_samples:
|
||||
delta = w.reshape(-1)[prev_total_samples:]
|
||||
prev_total_samples = n
|
||||
elif finished and n == prev_total_samples:
|
||||
delta = w.reshape(-1)[:0]
|
||||
else:
|
||||
delta = w.reshape(-1)
|
||||
prev_total_samples += int(delta.numel())
|
||||
if int(delta.numel()) > 0:
|
||||
delta_chunks.append(delta)
|
||||
if first_audio_elapsed is None and int(delta.numel()) > 0:
|
||||
first_audio_elapsed = time.perf_counter() - t_start
|
||||
logger.info(
|
||||
"%s prompt=%d/%d chunk=%d delta_samples=%d buf_len=%d finished=%s",
|
||||
phase_label,
|
||||
prompt_index + 1,
|
||||
prompt_count,
|
||||
chunk_i,
|
||||
int(delta.numel()),
|
||||
n,
|
||||
finished,
|
||||
)
|
||||
chunk_i += 1
|
||||
except ValueError:
|
||||
if not _extract_stream_finished(stage_output):
|
||||
logger.debug("skip non-audio partial output chunk=%d", chunk_i)
|
||||
|
||||
if not delta_chunks:
|
||||
raise RuntimeError("No audio chunks received; check stage config and logs.")
|
||||
|
||||
audio_cat = torch.cat([c.reshape(-1) for c in delta_chunks], dim=0)
|
||||
elapsed = time.perf_counter() - t_start
|
||||
return audio_cat, sample_rate, elapsed, first_audio_elapsed
|
||||
|
||||
|
||||
async def _abort_streaming_residual_work(
|
||||
omni: AsyncOmni,
|
||||
request_id: str,
|
||||
*,
|
||||
settle_seconds: float = 0.1,
|
||||
) -> None:
|
||||
"""Stop any late stage-0 work once the final audio has been collected."""
|
||||
await omni.engine.abort_async([request_id])
|
||||
if settle_seconds > 0:
|
||||
await asyncio.sleep(settle_seconds)
|
||||
|
||||
|
||||
async def _run_streaming_single(
|
||||
omni: AsyncOmni,
|
||||
args: Any,
|
||||
spec: PromptSpec,
|
||||
output_dir: Path,
|
||||
request_id: str,
|
||||
*,
|
||||
run_index: int,
|
||||
num_runs: int,
|
||||
prompt_index: int,
|
||||
prompt_count: int,
|
||||
) -> Path:
|
||||
audio_cat, sample_rate, elapsed, first_audio_elapsed = await _collect_streaming_audio(
|
||||
omni,
|
||||
args,
|
||||
spec,
|
||||
request_id,
|
||||
phase_label=f"run={run_index + 1}/{num_runs}",
|
||||
prompt_index=prompt_index,
|
||||
prompt_count=prompt_count,
|
||||
print_prompt=(run_index == 0 and prompt_index == 0),
|
||||
)
|
||||
await _abort_streaming_residual_work(omni, request_id)
|
||||
output_path = output_dir / f"output_run{run_index + 1}_{spec.label}.wav"
|
||||
_write_audio_tensor(output_path, audio_cat, sample_rate)
|
||||
audio_duration_s = float(audio_cat.numel()) / float(sample_rate) if sample_rate > 0 else 0.0
|
||||
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
|
||||
rtf_text = f", rtf={elapsed / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
|
||||
print(
|
||||
f"Saved (streaming) run {run_index + 1}/{num_runs}, "
|
||||
f"prompt {prompt_index + 1}/{prompt_count}: {output_path} ({elapsed:.2f}s{ttfp_text}{rtf_text})"
|
||||
)
|
||||
_emit_offline_metrics(
|
||||
request_id=request_id,
|
||||
elapsed_s=elapsed,
|
||||
first_audio_elapsed=first_audio_elapsed,
|
||||
audio_duration_s=audio_duration_s,
|
||||
)
|
||||
return output_path
|
||||
|
||||
|
||||
async def _run_streaming_warmup(args, omni: AsyncOmni) -> None:
|
||||
if args.warmup_runs == 0:
|
||||
return
|
||||
|
||||
warmup_specs = _get_warmup_specs(args.prompt_specs)
|
||||
print(
|
||||
f"Warmup: {args.warmup_runs} run(s) using the first prompt "
|
||||
f"({len(warmup_specs)} prompt(s)); outputs will be discarded."
|
||||
)
|
||||
for warmup_index in range(args.warmup_runs):
|
||||
t_warmup = time.perf_counter()
|
||||
tasks = []
|
||||
request_ids: list[str] = []
|
||||
for prompt_index, spec in enumerate(warmup_specs):
|
||||
request_id = f"warmup_stream_{warmup_index + 1}_{spec.label}_{uuid.uuid4().hex[:8]}"
|
||||
request_ids.append(request_id)
|
||||
tasks.append(
|
||||
_collect_streaming_audio(
|
||||
omni,
|
||||
args,
|
||||
spec,
|
||||
request_id,
|
||||
phase_label=f"warmup={warmup_index + 1}/{args.warmup_runs}",
|
||||
prompt_index=prompt_index,
|
||||
prompt_count=len(warmup_specs),
|
||||
)
|
||||
)
|
||||
results = await asyncio.gather(*tasks)
|
||||
for request_id in request_ids:
|
||||
await _abort_streaming_residual_work(omni, request_id)
|
||||
total_samples = sum(int(audio.numel()) for audio, _, _, _ in results)
|
||||
warmup_ttfps = [ttfp for _, _, _, ttfp in results if ttfp is not None]
|
||||
ttfp_text = f", ttfp={min(warmup_ttfps):.2f}s" if warmup_ttfps else ""
|
||||
print(
|
||||
f"Warmup (streaming) {warmup_index + 1}/{args.warmup_runs} finished: "
|
||||
f"{len(results)} prompt(s), {total_samples} sample(s) "
|
||||
f"({time.perf_counter() - t_warmup:.2f}s{ttfp_text})"
|
||||
)
|
||||
|
||||
|
||||
async def _run_streaming(args) -> list[Path]:
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
omni = AsyncOmni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
await _run_streaming_warmup(args, omni)
|
||||
profiler_started = False
|
||||
if args.enable_profiler:
|
||||
profile_prefix = f"voxcpm_streaming_{int(time.time())}"
|
||||
stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured"
|
||||
print(f"Starting profiler (streaming): stages={stages_text}, dir={args.profiler_dir}")
|
||||
await omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages)
|
||||
profiler_started = True
|
||||
t_total = time.perf_counter()
|
||||
total_elapsed = 0.0
|
||||
paths: list[Path] = []
|
||||
prompt_specs: list[PromptSpec] = args.prompt_specs
|
||||
try:
|
||||
for run in range(args.num_runs):
|
||||
for prompt_index, spec in enumerate(prompt_specs):
|
||||
request_id = f"stream_{run + 1}_{spec.label}_{uuid.uuid4().hex[:8]}"
|
||||
paths.append(
|
||||
await _run_streaming_single(
|
||||
omni,
|
||||
args,
|
||||
spec,
|
||||
output_dir,
|
||||
request_id,
|
||||
run_index=run,
|
||||
num_runs=args.num_runs,
|
||||
prompt_index=prompt_index,
|
||||
prompt_count=len(prompt_specs),
|
||||
)
|
||||
)
|
||||
total_elapsed = time.perf_counter() - t_total
|
||||
finally:
|
||||
if profiler_started:
|
||||
print("Stopping profiler (streaming)...")
|
||||
await omni.stop_profile(stages=args.profiler_stages)
|
||||
if args.profiler_wait_seconds > 0:
|
||||
print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...")
|
||||
await asyncio.sleep(args.profiler_wait_seconds)
|
||||
|
||||
print(
|
||||
f"All streaming runs finished: {args.num_runs} run(s), "
|
||||
f"{len(prompt_specs)} prompt(s), {len(paths)} file(s) in {total_elapsed:.2f}s total"
|
||||
)
|
||||
return paths
|
||||
|
||||
|
||||
def _run_sync(args) -> list[Path]:
|
||||
output_dir = Path(args.output_dir)
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
def _run_sync_single(
|
||||
spec: PromptSpec,
|
||||
*,
|
||||
request_prefix: str,
|
||||
save_outputs: bool,
|
||||
run_index: int | None = None,
|
||||
) -> tuple[list[Path], int, float | None, float, float, str]:
|
||||
global_request_id = f"{request_prefix}_{spec.label}"
|
||||
prompt = _build_prompt_for_spec(args, spec, global_request_id=global_request_id)
|
||||
if save_outputs and run_index == 0 and spec.label == "item001":
|
||||
print(f"---prompt---:{prompt}")
|
||||
|
||||
saved_paths: list[Path] = []
|
||||
output_count = 0
|
||||
first_audio_elapsed: float | None = None
|
||||
total_audio_duration_s = 0.0
|
||||
metrics_request_id = global_request_id
|
||||
t_start = time.perf_counter()
|
||||
for stage_outputs in omni.generate(prompt):
|
||||
request_output = stage_outputs.request_output
|
||||
if request_output is None:
|
||||
continue
|
||||
request_output_id = getattr(request_output, "request_id", None)
|
||||
if isinstance(request_output_id, str) and request_output_id:
|
||||
metrics_request_id = request_output_id
|
||||
for j, mm in enumerate(_iter_request_multimodal_outputs(request_output)):
|
||||
output_count += 1
|
||||
if first_audio_elapsed is None:
|
||||
try:
|
||||
audio_tensor = _extract_audio_tensor(mm)
|
||||
if int(audio_tensor.numel()) > 0:
|
||||
first_audio_elapsed = time.perf_counter() - t_start
|
||||
total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm))
|
||||
except ValueError:
|
||||
pass
|
||||
else:
|
||||
try:
|
||||
audio_tensor = _extract_audio_tensor(mm)
|
||||
total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm))
|
||||
except ValueError:
|
||||
pass
|
||||
if not save_outputs:
|
||||
continue
|
||||
save_stem = f"run{run_index + 1}_{spec.label}" if j == 0 else f"run{run_index + 1}_{spec.label}_{j}"
|
||||
saved_paths.append(_save_wav(mm, output_dir, save_stem))
|
||||
|
||||
if output_count == 0:
|
||||
raise RuntimeError("No output from Omni.generate")
|
||||
elapsed_s = time.perf_counter() - t_start
|
||||
return saved_paths, output_count, first_audio_elapsed, elapsed_s, total_audio_duration_s, metrics_request_id
|
||||
|
||||
if args.warmup_runs:
|
||||
warmup_specs = _get_warmup_specs(args.prompt_specs)
|
||||
print(
|
||||
f"Warmup: {args.warmup_runs} run(s) using the first prompt "
|
||||
f"({len(warmup_specs)} prompt(s)); outputs will be discarded."
|
||||
)
|
||||
for warmup_index in range(args.warmup_runs):
|
||||
t_warmup = time.perf_counter()
|
||||
_, output_count, first_audio_elapsed, elapsed_s, audio_duration_s, _ = _run_sync_single(
|
||||
warmup_specs[0],
|
||||
request_prefix=f"warmup_sync{warmup_index + 1}",
|
||||
save_outputs=False,
|
||||
)
|
||||
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
|
||||
rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
|
||||
print(
|
||||
f"Warmup (sync) {warmup_index + 1}/{args.warmup_runs} finished: "
|
||||
f"{output_count} output(s) ({time.perf_counter() - t_warmup:.2f}s{ttfp_text}{rtf_text})"
|
||||
)
|
||||
|
||||
profiler_started = False
|
||||
if args.enable_profiler:
|
||||
profile_prefix = f"voxcpm_sync_{int(time.time())}"
|
||||
stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured"
|
||||
print(f"Starting profiler (sync): stages={stages_text}, dir={args.profiler_dir}")
|
||||
omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages)
|
||||
profiler_started = True
|
||||
|
||||
t_total = time.perf_counter()
|
||||
total_elapsed = 0.0
|
||||
saved_paths: list[Path] = []
|
||||
prompt_specs: list[PromptSpec] = args.prompt_specs
|
||||
try:
|
||||
for run in range(args.num_runs):
|
||||
t_run = time.perf_counter()
|
||||
run_paths: list[Path] = []
|
||||
for prompt_index, spec in enumerate(prompt_specs):
|
||||
prompt_paths, _, first_audio_elapsed, elapsed_s, audio_duration_s, metrics_request_id = (
|
||||
_run_sync_single(
|
||||
spec,
|
||||
request_prefix=f"sync_run{run + 1}_{prompt_index + 1:03d}",
|
||||
save_outputs=True,
|
||||
run_index=run,
|
||||
)
|
||||
)
|
||||
run_paths.extend(prompt_paths)
|
||||
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
|
||||
rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
|
||||
print(
|
||||
f"Saved (sync) run {run + 1}/{args.num_runs}, "
|
||||
f"prompt {prompt_index + 1}/{len(prompt_specs)}: {len(prompt_paths)} file(s){ttfp_text}{rtf_text}"
|
||||
)
|
||||
_emit_offline_metrics(
|
||||
request_id=metrics_request_id,
|
||||
elapsed_s=elapsed_s,
|
||||
first_audio_elapsed=first_audio_elapsed,
|
||||
audio_duration_s=audio_duration_s,
|
||||
)
|
||||
|
||||
saved_paths.extend(run_paths)
|
||||
print(
|
||||
f"Run {run + 1}/{args.num_runs} finished: {len(run_paths)} file(s) ({time.perf_counter() - t_run:.2f}s)"
|
||||
)
|
||||
for path in run_paths:
|
||||
print(f" {path}")
|
||||
|
||||
total_elapsed = time.perf_counter() - t_total
|
||||
finally:
|
||||
if profiler_started:
|
||||
print("Stopping profiler (sync)...")
|
||||
omni.stop_profile(stages=args.profiler_stages)
|
||||
if args.profiler_wait_seconds > 0:
|
||||
print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...")
|
||||
time.sleep(args.profiler_wait_seconds)
|
||||
|
||||
print(
|
||||
f"All sync runs finished: {args.num_runs} run(s), "
|
||||
f"{len(prompt_specs)} prompt(s), {len(saved_paths)} file(s) in {total_elapsed:.2f}s total"
|
||||
)
|
||||
return saved_paths
|
||||
|
||||
|
||||
def main(args) -> int:
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
profiled_stage_config_path: str | None = None
|
||||
original_stage_config_path = args.stage_configs_path
|
||||
if args.enable_profiler:
|
||||
Path(args.profiler_dir).mkdir(parents=True, exist_ok=True)
|
||||
profiled_stage_config_path = _build_profiled_stage_config(
|
||||
args.stage_configs_path,
|
||||
str(Path(args.profiler_dir).resolve()),
|
||||
)
|
||||
args.stage_configs_path = profiled_stage_config_path
|
||||
|
||||
is_streaming = _is_streaming_stage_config(args.stage_configs_path)
|
||||
voice_clone_count = _count_voice_clone_prompts(args.prompt_specs)
|
||||
print(f"Model: {args.model}")
|
||||
print(f"Stage config: {original_stage_config_path}")
|
||||
print(f"Route: {'streaming' if is_streaming else 'sync'} (from stage-configs-path)")
|
||||
print(f"Prompt count: {len(args.prompt_specs)}")
|
||||
print("Batch mode: sequential (aligned with native VoxCPM)")
|
||||
print(f"Warmup runs: {args.warmup_runs}")
|
||||
print(f"Voice cloning prompts: {voice_clone_count}/{len(args.prompt_specs)}")
|
||||
if args.enable_profiler:
|
||||
print(f"Profiler: enabled (dir={args.profiler_dir}, stages={args.profiler_stages or 'all-configured'})")
|
||||
print(f"Profiled stage config: {args.stage_configs_path}")
|
||||
if voice_clone_count:
|
||||
print("Voice cloning note: --ref-text/ref_text must match the spoken content of the reference audio.")
|
||||
print(f"Num runs: {args.num_runs}")
|
||||
try:
|
||||
if is_streaming:
|
||||
asyncio.run(_run_streaming(args))
|
||||
else:
|
||||
_run_sync(args)
|
||||
finally:
|
||||
if profiled_stage_config_path is not None and os.path.exists(profiled_stage_config_path):
|
||||
os.unlink(profiled_stage_config_path)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
raise SystemExit(main(parse_args()))
|
||||
@@ -0,0 +1,283 @@
|
||||
"""Benchmark VoxCPM via /v1/audio/speech.
|
||||
|
||||
Reports TTFP (time to first packet), E2E latency, and RTF (real-time factor).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
DEFAULT_MODEL = "OpenBMB/VoxCPM1.5"
|
||||
DEFAULT_SAMPLE_RATE = 24000
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the VoxCPM speech benchmark.",
|
||||
"This is a short benchmark prompt for online text-to-speech generation.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"Please remember to bring your identification documents tomorrow morning.",
|
||||
"Learning a new language takes patience, practice, and curiosity.",
|
||||
"This benchmark reports TTFP and RTF for the VoxCPM online serving path.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestResult:
|
||||
success: bool = False
|
||||
ttfp: float = 0.0
|
||||
e2e: float = 0.0
|
||||
audio_bytes: int = 0
|
||||
audio_duration: float = 0.0
|
||||
rtf: float = 0.0
|
||||
prompt: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
concurrency: int = 0
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
p95_rtf: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
request_throughput: float = 0.0
|
||||
per_request: list[dict[str, float | str]] = field(default_factory=list)
|
||||
|
||||
|
||||
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = DEFAULT_SAMPLE_RATE, sample_width: int = 2) -> float:
|
||||
num_samples = num_bytes / sample_width
|
||||
return num_samples / sample_rate
|
||||
|
||||
|
||||
async def send_tts_request(
|
||||
session: aiohttp.ClientSession,
|
||||
api_url: str,
|
||||
*,
|
||||
model: str,
|
||||
prompt: str,
|
||||
ref_audio: str | None,
|
||||
ref_text: str | None,
|
||||
pbar: tqdm | None = None,
|
||||
) -> RequestResult:
|
||||
payload: dict[str, object] = {
|
||||
"model": model,
|
||||
"input": prompt,
|
||||
"stream": True,
|
||||
"response_format": "pcm",
|
||||
}
|
||||
if ref_audio is not None:
|
||||
payload["ref_audio"] = ref_audio
|
||||
if ref_text is not None:
|
||||
payload["ref_text"] = ref_text
|
||||
|
||||
result = RequestResult(prompt=prompt)
|
||||
started_at = time.perf_counter()
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status != 200:
|
||||
result.error = f"HTTP {response.status}: {await response.text()}"
|
||||
return result
|
||||
|
||||
first_chunk = True
|
||||
total_bytes = 0
|
||||
async for chunk in response.content.iter_any():
|
||||
if not chunk:
|
||||
continue
|
||||
if first_chunk:
|
||||
result.ttfp = time.perf_counter() - started_at
|
||||
first_chunk = False
|
||||
total_bytes += len(chunk)
|
||||
|
||||
result.e2e = time.perf_counter() - started_at
|
||||
result.audio_bytes = total_bytes
|
||||
result.audio_duration = pcm_bytes_to_duration(total_bytes)
|
||||
if result.audio_duration > 0:
|
||||
result.rtf = result.e2e / result.audio_duration
|
||||
result.success = True
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.e2e = time.perf_counter() - started_at
|
||||
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
return result
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
*,
|
||||
host: str,
|
||||
port: int,
|
||||
model: str,
|
||||
num_prompts: int,
|
||||
max_concurrency: int,
|
||||
num_warmups: int,
|
||||
ref_audio: str | None,
|
||||
ref_text: str | None,
|
||||
) -> BenchmarkResult:
|
||||
api_url = f"http://{host}:{port}/v1/audio/speech"
|
||||
connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
|
||||
timeout = aiohttp.ClientTimeout(total=600)
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
|
||||
if num_warmups > 0:
|
||||
print(f" Warming up with {num_warmups} requests...")
|
||||
warmup_tasks = [
|
||||
send_tts_request(
|
||||
session,
|
||||
api_url,
|
||||
model=model,
|
||||
prompt=PROMPTS[i % len(PROMPTS)],
|
||||
ref_audio=ref_audio,
|
||||
ref_text=ref_text,
|
||||
)
|
||||
for i in range(num_warmups)
|
||||
]
|
||||
await asyncio.gather(*warmup_tasks)
|
||||
print(" Warmup done.")
|
||||
|
||||
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
|
||||
semaphore = asyncio.Semaphore(max_concurrency)
|
||||
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
|
||||
|
||||
async def limited_request(prompt: str) -> RequestResult:
|
||||
async with semaphore:
|
||||
return await send_tts_request(
|
||||
session,
|
||||
api_url,
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
ref_audio=ref_audio,
|
||||
ref_text=ref_text,
|
||||
pbar=pbar,
|
||||
)
|
||||
|
||||
started_at = time.perf_counter()
|
||||
results = await asyncio.gather(*[asyncio.create_task(limited_request(prompt)) for prompt in request_prompts])
|
||||
duration = time.perf_counter() - started_at
|
||||
pbar.close()
|
||||
|
||||
succeeded = [result for result in results if result.success]
|
||||
bench = BenchmarkResult(
|
||||
concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
completed=len(succeeded),
|
||||
failed=len(results) - len(succeeded),
|
||||
duration_s=duration,
|
||||
)
|
||||
|
||||
if not succeeded:
|
||||
return bench
|
||||
|
||||
ttfps = np.array([result.ttfp * 1000 for result in succeeded], dtype=np.float64)
|
||||
e2es = np.array([result.e2e * 1000 for result in succeeded], dtype=np.float64)
|
||||
rtfs = np.array([result.rtf for result in succeeded], dtype=np.float64)
|
||||
audio_durations = np.array([result.audio_duration for result in succeeded], dtype=np.float64)
|
||||
|
||||
bench.mean_ttfp_ms = float(np.mean(ttfps))
|
||||
bench.median_ttfp_ms = float(np.median(ttfps))
|
||||
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
|
||||
bench.mean_e2e_ms = float(np.mean(e2es))
|
||||
bench.median_e2e_ms = float(np.median(e2es))
|
||||
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
|
||||
bench.mean_rtf = float(np.mean(rtfs))
|
||||
bench.median_rtf = float(np.median(rtfs))
|
||||
bench.p95_rtf = float(np.percentile(rtfs, 95))
|
||||
bench.total_audio_duration_s = float(np.sum(audio_durations))
|
||||
bench.request_throughput = len(succeeded) / duration if duration > 0 else 0.0
|
||||
bench.per_request = [
|
||||
{
|
||||
"prompt": result.prompt,
|
||||
"ttfp_ms": result.ttfp * 1000,
|
||||
"e2e_ms": result.e2e * 1000,
|
||||
"rtf": result.rtf,
|
||||
"audio_duration_s": result.audio_duration,
|
||||
}
|
||||
for result in succeeded
|
||||
]
|
||||
|
||||
return bench
|
||||
|
||||
|
||||
def print_summary(result: BenchmarkResult) -> None:
|
||||
width = 54
|
||||
print("")
|
||||
print("=" * width)
|
||||
print(f"{'VoxCPM Serving Benchmark':^{width}}")
|
||||
print("=" * width)
|
||||
print(f"concurrency : {result.concurrency}")
|
||||
print(f"requests : {result.completed}/{result.num_prompts} succeeded")
|
||||
print(f"wall time (s) : {result.duration_s:.3f}")
|
||||
print(f"mean TTFP (ms) : {result.mean_ttfp_ms:.2f}")
|
||||
print(f"p95 TTFP (ms) : {result.p95_ttfp_ms:.2f}")
|
||||
print(f"mean E2E (ms) : {result.mean_e2e_ms:.2f}")
|
||||
print(f"p95 E2E (ms) : {result.p95_e2e_ms:.2f}")
|
||||
print(f"mean RTF : {result.mean_rtf:.3f}")
|
||||
print(f"p95 RTF : {result.p95_rtf:.3f}")
|
||||
print(f"request throughput : {result.request_throughput:.2f} req/s")
|
||||
print("=" * width)
|
||||
|
||||
|
||||
async def main_async(args) -> None:
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
all_results: list[BenchmarkResult] = []
|
||||
for concurrency in args.max_concurrency:
|
||||
result = await run_benchmark(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
model=args.model,
|
||||
num_prompts=args.num_prompts,
|
||||
max_concurrency=concurrency,
|
||||
num_warmups=args.num_warmups,
|
||||
ref_audio=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
)
|
||||
print_summary(result)
|
||||
all_results.append(result)
|
||||
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"created_at": datetime.utcnow().isoformat() + "Z",
|
||||
"results": [asdict(result) for result in all_results],
|
||||
}
|
||||
result_path = result_dir / "bench_tts_serve.json"
|
||||
result_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
|
||||
print(f"Saved results to: {result_path}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Benchmark VoxCPM via /v1/audio/speech")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="Server host")
|
||||
parser.add_argument("--port", type=int, default=8091, help="Server port")
|
||||
parser.add_argument("--model", default=DEFAULT_MODEL, help="Model name or path")
|
||||
parser.add_argument("--num-prompts", type=int, default=20, help="Number of prompts to send")
|
||||
parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1], help="Concurrency levels to benchmark")
|
||||
parser.add_argument("--num-warmups", type=int, default=3, help="Warmup request count")
|
||||
parser.add_argument("--ref-audio", default=None, help="Reference audio URL or data URL for voice cloning")
|
||||
parser.add_argument("--ref-text", default=None, help="Reference audio transcript for voice cloning")
|
||||
parser.add_argument("--result-dir", default="results", help="Directory to save benchmark JSON")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main_async(parse_args()))
|
||||
@@ -0,0 +1,303 @@
|
||||
"""Run the full offline VoxCPM smoke matrix.
|
||||
|
||||
This script keeps the old `test.py` coverage, but delegates each case to
|
||||
`bench_tts_offline.py` so the benchmark runner itself stays focused on a
|
||||
single execution path.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shlex
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
BENCH_SCRIPT = Path(__file__).with_name("bench_tts_offline.py")
|
||||
DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml"
|
||||
DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml"
|
||||
DEFAULT_OUTPUT_ROOT = BENCH_SCRIPT.parents[1] / "results" / "offline_matrix"
|
||||
|
||||
SINGLE_TTS_TEXT = "This is a single text-to-speech smoke test for VoxCPM on vLLM Omni."
|
||||
SINGLE_CLONE_TEXT = "This sentence is synthesized with the cloned voice for validation."
|
||||
BATCH_TTS_TEXTS = [
|
||||
"The first batch text-to-speech sample validates sequential batch execution.",
|
||||
"The second batch text-to-speech sample checks another prompt in the same file.",
|
||||
"The third batch text-to-speech sample completes the sequential batch path.",
|
||||
]
|
||||
BATCH_CLONE_TEXTS = [
|
||||
"The first cloned sample validates sequential batch voice cloning.",
|
||||
"The second cloned sample checks the same reference voice on another prompt.",
|
||||
"The third cloned sample finishes the shared-reference clone batch path.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ModeSpec:
|
||||
name: str
|
||||
stage_config: Path
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CaseSpec:
|
||||
name: str
|
||||
warmup_runs: int
|
||||
prompt_kind: str
|
||||
voice_clone: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CaseResult:
|
||||
mode: str
|
||||
case: str
|
||||
returncode: int
|
||||
elapsed_s: float
|
||||
output_dir: Path
|
||||
log_path: Path
|
||||
|
||||
@property
|
||||
def ok(self) -> bool:
|
||||
return self.returncode == 0
|
||||
|
||||
|
||||
MODE_SPECS = [
|
||||
ModeSpec(name="streaming", stage_config=DEFAULT_STAGE_ASYNC),
|
||||
ModeSpec(name="sync", stage_config=DEFAULT_STAGE_SYNC),
|
||||
]
|
||||
|
||||
CASE_SPECS = [
|
||||
CaseSpec(name="warmup_single_tts", warmup_runs=1, prompt_kind="single", voice_clone=False),
|
||||
CaseSpec(name="warmup_single_clone", warmup_runs=1, prompt_kind="single", voice_clone=True),
|
||||
CaseSpec(name="warmup_batch_tts", warmup_runs=1, prompt_kind="batch", voice_clone=False),
|
||||
CaseSpec(name="warmup_batch_clone", warmup_runs=1, prompt_kind="batch", voice_clone=True),
|
||||
CaseSpec(name="cold_single_tts", warmup_runs=0, prompt_kind="single", voice_clone=False),
|
||||
CaseSpec(name="cold_single_clone", warmup_runs=0, prompt_kind="single", voice_clone=True),
|
||||
]
|
||||
|
||||
|
||||
def _write_lines(path: Path, lines: list[str]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
def _prepare_batch_inputs(output_root: Path) -> tuple[Path, Path]:
|
||||
input_dir = output_root / "inputs"
|
||||
batch_tts_path = input_dir / "batch_tts_prompts.txt"
|
||||
batch_clone_path = input_dir / "batch_clone_prompts.txt"
|
||||
_write_lines(batch_tts_path, BATCH_TTS_TEXTS)
|
||||
_write_lines(batch_clone_path, BATCH_CLONE_TEXTS)
|
||||
return batch_tts_path, batch_clone_path
|
||||
|
||||
|
||||
def _base_command(args, mode: ModeSpec, output_dir: Path) -> list[str]:
|
||||
cmd = [
|
||||
args.python,
|
||||
str(BENCH_SCRIPT),
|
||||
"--model",
|
||||
args.model,
|
||||
"--stage-configs-path",
|
||||
str(mode.stage_config),
|
||||
"--output-dir",
|
||||
str(output_dir),
|
||||
"--num-runs",
|
||||
str(args.num_runs),
|
||||
"--stage-init-timeout",
|
||||
str(args.stage_init_timeout),
|
||||
]
|
||||
cmd.append("--log-stats" if args.log_stats else "--no-log-stats")
|
||||
cmd.extend(["--cfg-value", str(args.cfg_value)])
|
||||
cmd.extend(["--inference-timesteps", str(args.inference_timesteps)])
|
||||
cmd.extend(["--min-len", str(args.min_len)])
|
||||
cmd.extend(["--max-new-tokens", str(args.max_new_tokens)])
|
||||
if args.streaming_prefix_len is not None:
|
||||
cmd.extend(["--streaming-prefix-len", str(args.streaming_prefix_len)])
|
||||
if args.enable_profiler:
|
||||
profiler_dir = Path(args.profiler_dir) if args.profiler_dir is not None else (output_dir / "profiler")
|
||||
cmd.append("--enable-profiler")
|
||||
cmd.extend(["--profiler-dir", str(profiler_dir)])
|
||||
cmd.extend(["--profiler-wait-seconds", str(args.profiler_wait_seconds)])
|
||||
if args.profiler_stages is not None:
|
||||
cmd.append("--profiler-stages")
|
||||
cmd.extend(str(stage_id) for stage_id in args.profiler_stages)
|
||||
return cmd
|
||||
|
||||
|
||||
def _build_case_command(
|
||||
args,
|
||||
mode: ModeSpec,
|
||||
case: CaseSpec,
|
||||
*,
|
||||
batch_tts_path: Path,
|
||||
batch_clone_path: Path,
|
||||
output_dir: Path,
|
||||
) -> list[str]:
|
||||
cmd = _base_command(args, mode, output_dir)
|
||||
cmd.extend(["--warmup-runs", str(case.warmup_runs)])
|
||||
if case.prompt_kind == "single":
|
||||
cmd.extend(["--text", SINGLE_CLONE_TEXT if case.voice_clone else SINGLE_TTS_TEXT])
|
||||
else:
|
||||
cmd.extend(["--txt-prompts", str(batch_clone_path if case.voice_clone else batch_tts_path)])
|
||||
if case.voice_clone:
|
||||
cmd.extend(["--ref-audio", args.ref_audio, "--ref-text", args.ref_text])
|
||||
return cmd
|
||||
|
||||
|
||||
def _run_case(
|
||||
args,
|
||||
mode: ModeSpec,
|
||||
case: CaseSpec,
|
||||
*,
|
||||
batch_tts_path: Path,
|
||||
batch_clone_path: Path,
|
||||
output_root: Path,
|
||||
) -> CaseResult:
|
||||
case_output_dir = output_root / mode.name / case.name
|
||||
case_output_dir.mkdir(parents=True, exist_ok=True)
|
||||
case_log_path = case_output_dir / "run.log"
|
||||
cmd = _build_case_command(
|
||||
args,
|
||||
mode,
|
||||
case,
|
||||
batch_tts_path=batch_tts_path,
|
||||
batch_clone_path=batch_clone_path,
|
||||
output_dir=case_output_dir,
|
||||
)
|
||||
|
||||
print()
|
||||
print("=" * 80)
|
||||
print(f"[{mode.name}] {case.name}")
|
||||
print(f"Output directory: {case_output_dir}")
|
||||
print(shlex.join(cmd))
|
||||
|
||||
start = time.perf_counter()
|
||||
with case_log_path.open("w", encoding="utf-8") as log_fp:
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
)
|
||||
assert process.stdout is not None
|
||||
for line in process.stdout:
|
||||
print(line, end="")
|
||||
log_fp.write(line)
|
||||
process.wait()
|
||||
|
||||
elapsed_s = time.perf_counter() - start
|
||||
status = "PASS" if (process.returncode or 0) == 0 else f"FAIL({process.returncode})"
|
||||
print(f"[{mode.name}] {case.name} -> {status} ({elapsed_s:.2f}s)")
|
||||
return CaseResult(
|
||||
mode=mode.name,
|
||||
case=case.name,
|
||||
returncode=int(process.returncode or 0),
|
||||
elapsed_s=elapsed_s,
|
||||
output_dir=case_output_dir,
|
||||
log_path=case_log_path,
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Run the full offline VoxCPM smoke matrix.")
|
||||
parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.")
|
||||
parser.add_argument("--ref-audio", type=str, required=True, help="Reference audio path for clone cases.")
|
||||
parser.add_argument("--ref-text", type=str, required=True, help="Exact transcript spoken in --ref-audio.")
|
||||
parser.add_argument("--output-root", type=str, default=str(DEFAULT_OUTPUT_ROOT), help="Root directory for outputs.")
|
||||
parser.add_argument("--python", type=str, default=sys.executable, help="Python executable used to launch cases.")
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=600, help="Stage initialization timeout in seconds.")
|
||||
parser.add_argument("--log-stats", dest="log_stats", action="store_true", help="Enable vLLM Omni stats logging.")
|
||||
parser.add_argument(
|
||||
"--no-log-stats",
|
||||
dest="log_stats",
|
||||
action="store_false",
|
||||
help="Disable vLLM Omni stats logging.",
|
||||
)
|
||||
parser.set_defaults(log_stats=True)
|
||||
parser.add_argument("--num-runs", type=int, default=1, help="Number of measured runs per case.")
|
||||
parser.add_argument("--cfg-value", type=float, default=2.0, help="Classifier-free guidance value for VoxCPM.")
|
||||
parser.add_argument("--inference-timesteps", type=int, default=10, help="Number of inference timesteps.")
|
||||
parser.add_argument("--min-len", type=int, default=2, help="Minimum generated token length.")
|
||||
parser.add_argument("--max-new-tokens", type=int, default=4096, help="Maximum generated token length.")
|
||||
parser.add_argument(
|
||||
"--streaming-prefix-len",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Optional VoxCPM streaming window passed to streaming cases.",
|
||||
)
|
||||
parser.add_argument("--enable-profiler", action="store_true", help="Enable torch profiler for each case.")
|
||||
parser.add_argument(
|
||||
"--profiler-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Profiler output root. Defaults to <case-output-dir>/profiler.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-stages",
|
||||
type=int,
|
||||
nargs="*",
|
||||
default=None,
|
||||
help="Optional stage ids to profile. Defaults to all configured stages.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-wait-seconds",
|
||||
type=float,
|
||||
default=30.0,
|
||||
help="Seconds to wait after stopping profiler for traces to flush.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
if args.num_runs < 1:
|
||||
parser.error("--num-runs must be >= 1")
|
||||
return args
|
||||
|
||||
|
||||
def main(args) -> int:
|
||||
output_root = Path(args.output_root)
|
||||
output_root.mkdir(parents=True, exist_ok=True)
|
||||
batch_tts_path, batch_clone_path = _prepare_batch_inputs(output_root)
|
||||
|
||||
print(f"Model: {args.model}")
|
||||
print(f"Reference audio: {args.ref_audio}")
|
||||
print(f"Reference text: {args.ref_text}")
|
||||
print(f"Python: {args.python}")
|
||||
print(f"Output root: {output_root}")
|
||||
print(f"Cases: {len(MODE_SPECS) * len(CASE_SPECS)}")
|
||||
|
||||
results: list[CaseResult] = []
|
||||
for mode in MODE_SPECS:
|
||||
for case in CASE_SPECS:
|
||||
results.append(
|
||||
_run_case(
|
||||
args,
|
||||
mode,
|
||||
case,
|
||||
batch_tts_path=batch_tts_path,
|
||||
batch_clone_path=batch_clone_path,
|
||||
output_root=output_root,
|
||||
)
|
||||
)
|
||||
|
||||
failed = [result for result in results if not result.ok]
|
||||
print()
|
||||
print("=" * 80)
|
||||
print("Summary:")
|
||||
for result in results:
|
||||
status = "PASS" if result.ok else f"FAIL({result.returncode})"
|
||||
print(f"- [{result.mode}] {result.case}: {status} ({result.elapsed_s:.2f}s)")
|
||||
print(f" output_dir={result.output_dir}")
|
||||
print(f" log={result.log_path}")
|
||||
|
||||
print(f"Passed: {len(results) - len(failed)}/{len(results)}")
|
||||
if failed:
|
||||
print("Failed cases:")
|
||||
for result in failed:
|
||||
print(f"- [{result.mode}] {result.case}: see {result.log_path}")
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main(parse_args()))
|
||||
@@ -7,7 +7,7 @@ COPY . .
|
||||
|
||||
# Install system dependencies
|
||||
RUN apt-get update && \
|
||||
apt-get install -y espeak-ng git sox libsox-fmt-all jq && \
|
||||
apt-get install -y espeak-ng git jq && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@ WORKDIR ${COMMON_WORKDIR}
|
||||
|
||||
# Step 1: Setup - Install system dependencies
|
||||
RUN apt-get update && \
|
||||
apt-get install -y git sox libsox-fmt-all jq && \
|
||||
apt-get install -y git jq && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
|
||||
@@ -21,7 +21,7 @@ WORKDIR ${COMMON_WORKDIR}
|
||||
# Need to include ffmpeg because vllm rocm upstream docker image
|
||||
# does not include it.
|
||||
RUN apt-get update && \
|
||||
apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get install -y espeak-ng ffmpeg git jq && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
|
||||
@@ -64,6 +64,7 @@ nav:
|
||||
- FP8: user_guide/diffusion/quantization/fp8.md
|
||||
- Int8: user_guide/diffusion/quantization/int8.md
|
||||
- GGUF: user_guide/diffusion/quantization/gguf.md
|
||||
- Frame Interpolation: user_guide/diffusion/frame_interpolation.md
|
||||
- Parallelism:
|
||||
- Overview: user_guide/diffusion/parallelism/overview.md
|
||||
- CFG Parallel: user_guide/diffusion/parallelism/cfg_parallel.md
|
||||
@@ -80,7 +81,6 @@ nav:
|
||||
- Developer Guide:
|
||||
- General:
|
||||
- contributing/README.md
|
||||
- pr_reviewer.md
|
||||
- glob: contributing/*
|
||||
flatten_single_child_sections: true
|
||||
- Model Implementation:
|
||||
@@ -97,6 +97,7 @@ nav:
|
||||
- design/feature/disaggregated_inference.md
|
||||
- design/feature/ray_based_execution.md
|
||||
- design/feature/omni_connectors/
|
||||
- design/feature/prefix_caching.md
|
||||
- design/feature/cfg_parallel.md
|
||||
- design/feature/expert_parallel.md
|
||||
- design/feature/sequence_parallel.md
|
||||
@@ -105,7 +106,7 @@ nav:
|
||||
- design/feature/hsdp.md
|
||||
- design/feature/cache_dit.md
|
||||
- design/feature/teacache.md
|
||||
- design/feature/async_chunk_design.md
|
||||
- design/feature/async_chunk.md
|
||||
- design/feature/vae_parallel.md
|
||||
- design/feature/diffusion_step_execution.md
|
||||
- Module Design:
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
Main entry points for vLLM-Omni inference and serving.
|
||||
|
||||
- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
|
||||
- [vllm_omni.entrypoints.cfg_companion_tracker.CfgCompanionTracker][]
|
||||
- [vllm_omni.engine.cfg_companion_tracker.CfgCompanionTracker][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
|
||||
|
||||
|
Before Width: | Height: | Size: 97 KiB After Width: | Height: | Size: 96 KiB |
@@ -6,7 +6,7 @@ For options within a vLLM Engine. Please refer to [vLLM Configuration](https://d
|
||||
|
||||
Currently, the main options are maintained by stage configs for each model.
|
||||
|
||||
For specific example, please refer to [Qwen2.5-omni stage config](stage_configs/qwen2_5_omni.yaml)
|
||||
For a specific example, see the [Qwen2.5-Omni deploy config](gh-file:vllm_omni/deploy/qwen2_5_omni.yaml). The matching frozen pipeline topology lives at [vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py](gh-file:vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py).
|
||||
|
||||
For introduction, please check [Introduction for stage config](./stage_configs.md)
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ deployment-specific values usually change per environment:
|
||||
- connector backend and connector ports
|
||||
- connector IPs or bootstrap addresses
|
||||
|
||||
Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml)
|
||||
Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
|
||||
and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
|
||||
changes below.
|
||||
|
||||
@@ -145,19 +145,13 @@ Compared with the default Qwen3-Omni config:
|
||||
```yaml
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 1
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
- from: 1
|
||||
to: 2
|
||||
window_size: -1
|
||||
- from: 2
|
||||
to: 3
|
||||
window_size: -1
|
||||
```
|
||||
|
||||
## 4. Launch with your custom config
|
||||
|
||||
@@ -3,7 +3,147 @@
|
||||
In vLLM-Omni, the target model is separated into multiple stages, which are processed by different LLMEngines, DiffusionEngines or other types of engines. Depending on different types of stages, such as Autoregressive (AR) stage or Diffusion transformer (DiT) stage, each can choose corresponding schedulers, model workers to load with the Engines in a plug-in fashion.
|
||||
|
||||
!!! note
|
||||
Default stage config YAMLs (for example, `vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml` and `vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml`) are bundled and loaded automatically when `stage_configs_path` is not provided. They have been verified to work on 1xH100 for Qwen2.5-Omni and 2xH100 for Qwen3-Omni.
|
||||
Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when neither `--stage-configs-path` nor `--deploy-config` is provided — the model registry resolves the right pipeline + deploy YAML by `model_type`. The bundled defaults have been verified on 1xH100 for Qwen2.5-Omni and 2xH100 for Qwen3-Omni. Models that have not yet migrated to the new schema continue to use the legacy `vllm_omni/model_executor/stage_configs/<model>.yaml` files via `--stage-configs-path`.
|
||||
|
||||
## New deploy schema reference
|
||||
|
||||
The new deploy schema lives under `vllm_omni/deploy/` and is paired with a frozen `PipelineConfig` registered by the model's `pipeline.py`. Each deploy YAML has these top-level fields:
|
||||
|
||||
| Field | Type | Required | Default | Description |
|
||||
|-------|------|----------|---------|-------------|
|
||||
| `base_config` | str (path) | optional | — | Overlay parent (relative or absolute). `stages:` / `platforms:` deep-merged by stage_id; other scalars overlay-wins. Intended for user-authored overlays; prod yamls stay flat. |
|
||||
| `async_chunk` | bool | optional | `true` | Enable chunked streaming between stages. Pin to `false` if the pipeline runs end-to-end. |
|
||||
| `connectors` | dict | optional | `null` | Named connector specs (`{name, extra}`). Referenced by each stage's `input_connectors` / `output_connectors`. See [Connector schema](#connector-schema). |
|
||||
| `edges` | list | optional | `null` | Explicit edge list for the KV transfer graph. Auto-derived from stage inputs if omitted. |
|
||||
| `stages` | list | required | — | Per-stage engine args + wiring (see [Stage fields](#stage-fields)). |
|
||||
| `platforms` | dict | optional | `null` | Keyed by `npu` / `rocm` / `xpu`, each contains a `stages:` list with per-platform overrides applied on top of the CUDA defaults. |
|
||||
| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only`). |
|
||||
| `trust_remote_code` | bool | optional | `true` | **Pipeline-wide.** Trust HF remote code on model load; applies to every stage. |
|
||||
| `distributed_executor_backend` | str | optional | `"mp"` | **Pipeline-wide.** Executor backend (`"mp"` or `"ray"`). |
|
||||
| `dtype` | str \| null | optional | `null` | **Pipeline-wide.** Model dtype for every stage. |
|
||||
| `quantization` | str \| null | optional | `null` | **Pipeline-wide.** Quantization method for every stage. |
|
||||
| `enable_prefix_caching` | bool | optional | `false` | **Pipeline-wide.** Prefix cache toggle applied to every stage. |
|
||||
| `enable_chunked_prefill` | bool \| null | optional | `null` | **Pipeline-wide.** Chunked prefill toggle applied to every stage. |
|
||||
| `data_parallel_size` | int | optional | `1` | **Pipeline-wide.** DP degree for every stage. |
|
||||
| `pipeline_parallel_size` | int | optional | `1` | **Pipeline-wide.** PP degree for every stage. |
|
||||
|
||||
### Stage fields
|
||||
|
||||
Each entry under `stages:` accepts any `StageDeployConfig` field directly (no nested `engine_args:`). Only fields whose value legitimately varies across stages live here; pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level and applied to every stage. Unknown keys fall through to `engine_extras:` and are forwarded to the engine.
|
||||
|
||||
| Field | Type | Required | Default | Description |
|
||||
|-------|------|----------|---------|-------------|
|
||||
| `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. |
|
||||
| `max_num_seqs` | int | optional | `64` | Max concurrent sequences per stage. |
|
||||
| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage memory budget. |
|
||||
| `tensor_parallel_size` | int | optional | `1` | TP degree for this stage. |
|
||||
| `enforce_eager` | bool | optional | `false` | Disable CUDA graphs. |
|
||||
| `max_num_batched_tokens` | int | optional | `32768` | Prefill budget. |
|
||||
| `max_model_len` | int \| null | optional | `null` | Per-stage context length (auto-sets `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1` when larger than HF default). |
|
||||
| `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. |
|
||||
| `devices` | str | optional | `"0"` | `CUDA_VISIBLE_DEVICES`-style device list. |
|
||||
| `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_<n>`; values are names registered under top-level `connectors:`. |
|
||||
| `input_connectors` | dict \| null | optional | `null` | Keyed by `from_stage_<n>`; values are names registered under top-level `connectors:`. |
|
||||
| `default_sampling_params` | dict \| null | optional | `null` | Baseline sampling params. Deep-merged with pipeline `sampling_constraints` (pipeline wins). |
|
||||
| `engine_extras` | dict | optional | `{}` | Catch-all for keys not listed above; deep-merged across overlays. Also carries per-stage overrides of pipeline-wide settings (e.g. stage-specific `dtype`). |
|
||||
|
||||
### Connector schema
|
||||
|
||||
Each entry under top-level `connectors:` follows this shape:
|
||||
|
||||
```yaml
|
||||
connectors:
|
||||
<connector_name>:
|
||||
name: <ConnectorClassName> # required — class registered in vllm_omni.distributed
|
||||
extra: # optional — forwarded to the connector's __init__
|
||||
<key>: <value>
|
||||
...
|
||||
```
|
||||
|
||||
| Connector class | Use case | `extra` keys |
|
||||
|-----------------|----------|--------------|
|
||||
| `SharedMemoryConnector` | Same-host KV transfer between stages (default for bundled YAMLs). | `shm_threshold_bytes` (int, default `65536`). |
|
||||
| `MooncakeStoreConnector` | Cross-host KV transfer over TCP. Required for multi-node deployments. | `host`, `metadata_server`, `master`, `segment` (int bytes), `localbuf` (int bytes), `proto` (`"tcp"` / `"rdma"`). |
|
||||
|
||||
A stage references a connector by name in its `input_connectors` / `output_connectors`:
|
||||
|
||||
```yaml
|
||||
connectors:
|
||||
shm:
|
||||
name: SharedMemoryConnector
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
output_connectors: {to_stage_1: shm}
|
||||
- stage_id: 1
|
||||
input_connectors: {from_stage_0: shm}
|
||||
```
|
||||
|
||||
### CLI flags introduced in this refactor
|
||||
|
||||
| Flag | Description |
|
||||
|------|-------------|
|
||||
| `--deploy-config PATH` | Load a new-schema deploy YAML. Takes precedence over `--stage-configs-path`. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
|
||||
| `--stage-overrides JSON` | Per-stage JSON overrides, e.g. `'{"0":{"gpu_memory_utilization":0.5}}'`. Per-stage values always win over global flags. |
|
||||
| `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. |
|
||||
| `--stage-configs-path` | **Deprecated.** Accepts legacy `stage_args` yamls and (auto-detected) new deploy yamls; emits a deprecation warning. Migrate to `--deploy-config`. To be removed in a follow-up PR. |
|
||||
|
||||
### Precedence
|
||||
|
||||
From highest to lowest:
|
||||
|
||||
1. Per-stage flags (`--stage-overrides` JSON, `--stage-<id>-<key>` if registered)
|
||||
2. Explicit global CLI flags (`--gpu-memory-utilization 0.85`, etc.)
|
||||
3. Platform section (`platforms.npu.stages`, etc.) on top of the base `stages:`
|
||||
4. Overlay YAML (via `base_config:`) on top of the base YAML
|
||||
5. Parser defaults
|
||||
|
||||
### Worked override example
|
||||
|
||||
Starting from the bundled `vllm_omni/deploy/qwen3_omni_moe.yaml`:
|
||||
|
||||
```yaml
|
||||
# vllm_omni/deploy/qwen3_omni_moe.yaml (excerpt)
|
||||
async_chunk: true
|
||||
stages:
|
||||
- stage_id: 0
|
||||
gpu_memory_utilization: 0.9
|
||||
max_num_seqs: 32
|
||||
- stage_id: 1
|
||||
gpu_memory_utilization: 0.7
|
||||
max_num_seqs: 16
|
||||
```
|
||||
|
||||
A user-authored overlay that inherits the base and overrides only stage 1:
|
||||
|
||||
```yaml
|
||||
# my_overrides.yaml
|
||||
base_config: /path/to/vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
stages:
|
||||
- stage_id: 1
|
||||
gpu_memory_utilization: 0.5 # smaller GPU
|
||||
```
|
||||
|
||||
Launched with both an explicit global flag and a per-stage override:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--deploy-config my_overrides.yaml \
|
||||
--max-model-len 16384 \
|
||||
--stage-overrides '{"0": {"max_num_seqs": 8}}'
|
||||
```
|
||||
|
||||
Effective config per stage after the merge:
|
||||
|
||||
| Stage | Field | Final value | Source |
|
||||
|-------|-------|-------------|--------|
|
||||
| 0 | `gpu_memory_utilization` | `0.9` | base YAML (overlay didn't touch stage 0) |
|
||||
| 0 | `max_num_seqs` | `8` | per-stage CLI (`--stage-overrides`) — wins over base `32` |
|
||||
| 0 | `max_model_len` | `16384` | global CLI |
|
||||
| 1 | `gpu_memory_utilization` | `0.5` | overlay YAML — wins over base `0.7` |
|
||||
| 1 | `max_num_seqs` | `16` | base YAML (overlay didn't touch this field) |
|
||||
| 1 | `max_model_len` | `16384` | global CLI |
|
||||
| 2 | (all defaults) | — | base YAML (no overrides apply) |
|
||||
|
||||
Therefore, as a core part of vLLM-Omni, the stage configs for a model have several main functions:
|
||||
|
||||
@@ -35,7 +175,7 @@ stage_args:
|
||||
- stage_id: 0 # mark the unique id for each stage
|
||||
runtime: # The disaggregated configuration
|
||||
process: true # Run this stage in a separate process
|
||||
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
|
||||
devices: "0" # Logical device index for this stage (mapped through CUDA_VISIBLE_DEVICES / ASCEND_RT_VISIBLE_DEVICES if set)
|
||||
engine_args: # Engine arguments for a certain engine
|
||||
model_stage: thinker
|
||||
max_num_seqs: 1
|
||||
@@ -114,16 +254,12 @@ stage_args:
|
||||
# Top-level runtime config (concise): default windows and stage edges
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1 # Simplified: trigger downstream only after full upstream completion
|
||||
max_inflight: 1 # Simplified: process serially within each stage
|
||||
|
||||
edges:
|
||||
- from: 0 # thinker → talker: trigger only after receiving full input (-1)
|
||||
to: 1
|
||||
window_size: -1
|
||||
- from: 1 # talker → code2wav: trigger only after receiving full input (-1)
|
||||
to: 2
|
||||
window_size: -1
|
||||
|
||||
```
|
||||
|
||||
@@ -155,7 +291,9 @@ Default: `true`
|
||||
|
||||
#### `runtime.devices`
|
||||
|
||||
Visible devices for this stage, specified as a string. This controls which GPU devices are available to the stage process, similar to setting `CUDA_VISIBLE_DEVICES` or using `torch.cuda.set_device()`. For example, `"0"` uses GPU 0, `"1"` uses GPU 1, and `"0,1"` makes both GPUs 0 and 1 visible.
|
||||
Logical device indices for this stage, specified as a string. Values are **logical indices** (`0`, `1`, `2`, ...) — not physical GPU IDs — and are mapped through the platform's visibility env var (`CUDA_VISIBLE_DEVICES` on CUDA, `ASCEND_RT_VISIBLE_DEVICES` on NPU) before being applied via `torch.cuda.set_device()` (or the equivalent).
|
||||
|
||||
Example: if `CUDA_VISIBLE_DEVICES=0,2,4` is set in the environment, then `devices: "0"` selects physical GPU 0 (the first visible), `devices: "1"` selects physical GPU 2, and `devices: "0,1"` makes physical GPUs 0 and 2 available to the stage. If no visibility env var is set, logical and physical IDs coincide.
|
||||
|
||||
Default: `"0"`
|
||||
|
||||
|
||||
@@ -1,94 +0,0 @@
|
||||
# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
runtime:
|
||||
process: true # Run this stage in a separate process
|
||||
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
|
||||
engine_args:
|
||||
model_stage: thinker
|
||||
max_num_seqs: 1
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.8
|
||||
enforce_eager: true # Now we only support eager mode
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
enable_prefix_caching: false
|
||||
is_comprehension: true
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.1
|
||||
- stage_id: 1
|
||||
runtime:
|
||||
process: true
|
||||
devices: "1"
|
||||
engine_args:
|
||||
model_stage: talker
|
||||
max_num_seqs: 1
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.8
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent
|
||||
engine_input_source: [0]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_p: 0.8
|
||||
top_k: 40
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [8294]
|
||||
- stage_id: 2
|
||||
runtime:
|
||||
process: true
|
||||
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
max_num_seqs: 1
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
gpu_memory_utilization: 0.15
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio
|
||||
engine_input_source: [1]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.1
|
||||
|
||||
# Top-level runtime config (concise): default windows and stage edges
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1 # Simplified: trigger downstream only after full upstream completion
|
||||
max_inflight: 1 # Simplified: process serially within each stage
|
||||
edges:
|
||||
- from: 0 # thinker → talker: trigger only after receiving full input (-1)
|
||||
to: 1
|
||||
window_size: -1
|
||||
- from: 1 # talker → code2wav: trigger only after receiving full input (-1)
|
||||
to: 2
|
||||
window_size: -1
|
||||
@@ -86,7 +86,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
/tests/e2e/online_serving/test_{model_name}_expansion.py<br>
|
||||
/tests/e2e/offline_inference/test_{model_name}_expansion.py<br>
|
||||
<strong>Performance:</strong><br>
|
||||
/tests/dfx/perf/tests/test.json<br>
|
||||
/tests/dfx/perf/tests/test_qwen_omni.json (Omni), test_tts.json (TTS),<br>
|
||||
and /tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)<br>
|
||||
<strong>Doc Test:</strong><br>
|
||||
tests/example/online_serving/test_{model_name}.py<br>
|
||||
tests/example/offline_inference/test_{model_name}.py
|
||||
@@ -230,8 +231,7 @@ vllm_omni/ tests/
|
||||
│ ├── test_qwen3_omni_expansion.py
|
||||
│ ├── test_mimo_audio.py
|
||||
│ ├── test_image_gen_edit.py
|
||||
│ ├── test_images_generations_lora.py
|
||||
│ └── stage_configs/
|
||||
│ └── test_images_generations_lora.py
|
||||
└── offline_inference/ ✅
|
||||
├── test_qwen2_5_omni.py
|
||||
├── test_qwen3_omni.py
|
||||
@@ -247,11 +247,12 @@ vllm_omni/ tests/
|
||||
├── test_diffusion_layerwise_offload.py
|
||||
├── test_diffusion_lora.py
|
||||
├── test_sequence_parallel.py
|
||||
└── stage_configs/
|
||||
├── qwen2_5_omni_ci.yaml
|
||||
├── qwen3_omni_ci.yaml
|
||||
├── bagel_*.yaml
|
||||
└── npu/, rocm/, etc.
|
||||
└── stage_configs/ (legacy schema, still
|
||||
├── bagel_*.yaml present for unmigrated
|
||||
└── npu/, rocm/, etc. models)
|
||||
|
||||
# Migrated models (qwen3_omni_moe, qwen2_5_omni, qwen3_tts) live under
|
||||
# vllm_omni/deploy/ instead — see docs/configuration/stage_configs.md.
|
||||
```
|
||||
|
||||
|
||||
@@ -530,13 +531,13 @@ L4 level testing is a comprehensive quality audit before a version release. It e
|
||||
### 3.2 Testing Content and Scope
|
||||
|
||||
- ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling.
|
||||
- ***Performance Testing***: Uses the `tests/dfx/perf/tests/test.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
|
||||
- ***Performance Testing***: Uses `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, and diffusion configs in the form `tests/dfx/perf/tests/test_*_vllm_omni.json` (passed to `run_benchmark.py` via `--test-config-file`) to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
|
||||
- ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description.
|
||||
|
||||
### 3.3 Test Directory and Execution Files
|
||||
|
||||
- ***Functional Testing***: Same directories as L3.
|
||||
- ***Performance Test Configuration***: `tests/dfx/perf/tests/test.json`
|
||||
- ***Performance Test Configuration***: `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (e.g. `test_qwen_image_vllm_omni.json`)
|
||||
- ***Documentation Example Tests***:
|
||||
- - `tests/example/online_serving/test_{model_name}.py`
|
||||
- `tests/example/offline_inference/test_{model_name}.py`
|
||||
|
||||
@@ -40,7 +40,7 @@ Currently all the features are available in online serving mode. Hence, only nee
|
||||
- Test marks: always add `advanced_model` and `diffusion`. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/).
|
||||
- To maximize code reuse, you may refer to
|
||||
- `tests/conftest.py` for `omni_server` (running server in subprocess) and `openai_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper.
|
||||
- `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
|
||||
- `tests/helpers/mark.py` for `@hardware_test(...)` and `hardware_marks`.
|
||||
- [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
|
||||
- Doc: add a concise docstring for each test function.
|
||||
- Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).
|
||||
@@ -1,4 +1,4 @@
|
||||
When you want to add L4-level ***performance test*** cases, you can refer to the following format for case addition in tests/dfx/perf/tests/test.json:
|
||||
When you want to add L4-level ***performance test*** cases, you can refer to the following format for case addition in `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, or diffusion configs such as `tests/dfx/perf/tests/test_*_vllm_omni.json` (selected via `pytest ... run_benchmark.py --test-config-file <path>`):
|
||||
|
||||
```JSON
|
||||
{
|
||||
|
||||
@@ -45,7 +45,6 @@ Our test scripts use the pytest framework. First, please use `git clone https://
|
||||
=== "L3 level & L4 level"
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v -m "advanced_model" --run-level=advanced_model
|
||||
```
|
||||
If you only want to run L3 test case, you can use:
|
||||
@@ -60,9 +59,9 @@ Our test scripts use the pytest framework. First, please use `git clone https://
|
||||
```bash
|
||||
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
|
||||
```
|
||||
Note: To run performance tests, use:
|
||||
Note: To run performance tests (defaults to ``test_qwen_omni.json``; use ``--test-config-file tests/dfx/perf/tests/test_tts.json`` for TTS):
|
||||
```bash
|
||||
pytest -s -v perf/scripts/run_benchmark.py
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
|
||||
```
|
||||
|
||||
The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
|
||||
|
||||
@@ -38,7 +38,7 @@ Defined in `pyproject.toml`:
|
||||
### Example usage for markers
|
||||
|
||||
```python
|
||||
from tests.utils import hardware_test
|
||||
from tests.helpers.mark import hardware_test
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@@ -53,7 +53,7 @@ def test_video_to_audio()
|
||||
|
||||
### Decorator: `@hardware_test`
|
||||
|
||||
This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions:
|
||||
This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/helpers/mark.py` performs the following actions:
|
||||
|
||||
1. **Applies platform and resource markers**
|
||||
Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`).
|
||||
@@ -105,7 +105,7 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
|
||||
`hardware_marks` returns a list of pytest mark objects with the same signature as `@hardware_test`. Use it when you need more flexibility, such as attaching hardware marks to individual `pytest.param` entries rather than an entire test function.
|
||||
|
||||
```python
|
||||
from tests.utils import hardware_marks
|
||||
from tests.helpers.mark import hardware_marks
|
||||
|
||||
MULTI_CARD_MARKS = hardware_marks(
|
||||
res={"cuda": "H100", "rocm": "MI325", "npu": "A2"}, num_cards=2
|
||||
@@ -133,9 +133,9 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
|
||||
"distributed_tpu: Tests that require multiple TPU devices",
|
||||
]
|
||||
```
|
||||
2. **Implement a marker construction function for your platform** in `vllm-omni/tests/utils.py`:
|
||||
2. **Implement a marker construction function for your platform** in `vllm-omni/tests/helpers/mark.py`:
|
||||
```python
|
||||
# In vllm-omni/tests/utils.py
|
||||
# In vllm-omni/tests/helpers/mark.py
|
||||
|
||||
def tpu_marks(*, res: str, num_cards: int):
|
||||
test_platform = pytest.mark.tpu
|
||||
@@ -175,4 +175,4 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
|
||||
- Plug into `hardware_marks`
|
||||
- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation!
|
||||
|
||||
See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).
|
||||
See code in `vllm-omni/tests/helpers/mark.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).
|
||||
@@ -135,8 +135,7 @@ vllm_omni/ tests/
|
||||
│ ├── test_qwen3_omni_expansion.py
|
||||
│ ├── test_mimo_audio.py
|
||||
│ ├── test_image_gen_edit.py
|
||||
│ ├── test_images_generations_lora.py
|
||||
│ └── stage_configs/
|
||||
│ └── test_images_generations_lora.py
|
||||
└── offline_inference/ ✅
|
||||
├── test_qwen2_5_omni.py
|
||||
├── test_qwen3_omni.py
|
||||
@@ -153,11 +152,12 @@ vllm_omni/ tests/
|
||||
├── test_diffusion_lora.py
|
||||
├── test_sequence_parallel.py
|
||||
├── test_qwen_image_edit_expansion.py
|
||||
└── stage_configs/
|
||||
├── qwen2_5_omni_ci.yaml
|
||||
├── qwen3_omni_ci.yaml
|
||||
├── bagel_*.yaml
|
||||
└── stage_configs/ (legacy schema, still present
|
||||
├── bagel_*.yaml for unmigrated models)
|
||||
└── npu/, rocm/, etc.
|
||||
|
||||
# Migrated models (qwen3_omni_moe, qwen2_5_omni, qwen3_tts) live under
|
||||
# vllm_omni/deploy/ instead — see docs/configuration/stage_configs.md.
|
||||
examples/ tests
|
||||
│ └── examples
|
||||
├── online_serving/ → ├── online_serving/
|
||||
@@ -221,14 +221,13 @@ from pathlib import Path
|
||||
import openai
|
||||
import pytest
|
||||
|
||||
from tests.conftest import (
|
||||
OmniServer,
|
||||
convert_audio_to_text,
|
||||
from tests.helpers.media import (
|
||||
convert_audio_bytes_to_text,
|
||||
cosine_similarity_text,
|
||||
dummy_messages_from_mix_data,
|
||||
generate_synthetic_video,
|
||||
merge_base64_and_convert_to_text,
|
||||
)
|
||||
from tests.helpers.runtime import OmniServer, dummy_messages_from_mix_data
|
||||
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
# Edit: model name and stage config path
|
||||
@@ -236,7 +235,7 @@ models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
|
||||
|
||||
#If you use the default configuration file, you can directly use the following address.
|
||||
def get_default_config():
|
||||
return str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
|
||||
return get_deploy_config_path("ci/qwen3_omni_moe.yaml")
|
||||
|
||||
#If you need to modify the configuration file, you can use modify_stage_config.
|
||||
def get_chunk_config():
|
||||
@@ -405,7 +404,7 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
|
||||
# PURPOSE: Verify text and audio outputs convey the same information
|
||||
# CUSTOMIZATION: Adjust similarity threshold (0.9) based on accuracy requirements
|
||||
assert audio_data is not None, "No audio output is generated"
|
||||
audio_content = merge_base64_and_convert_to_text(audio_data)
|
||||
audio_content = convert_audio_bytes_to_text(audio_data)
|
||||
print(f"text content is: {text_content}")
|
||||
print(f"audio content is: {audio_content}")
|
||||
similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
|
||||
@@ -428,7 +427,7 @@ from pathlib import Path
|
||||
import pytest
|
||||
from vllm.assets.video import VideoAsset
|
||||
|
||||
from tests.utils import hardware_test
|
||||
from tests.helpers.mark import hardware_test
|
||||
from ..multi_stages.conftest import OmniRunner
|
||||
|
||||
# Optional: set process start method for workers
|
||||
|
||||
@@ -313,7 +313,7 @@ The registry uses lazy loading, so the model class is imported only when needed.
|
||||
|
||||
## Stage Configuration
|
||||
|
||||
Create a YAML configuration file in `vllm_omni/model_executor/stage_configs/`. For a complete example, see the [Qwen3-Omni configuration file](gh-file:vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml).
|
||||
Create a YAML configuration file in `vllm_omni/deploy/`. For a complete example, see the [Qwen3-Omni configuration file](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml).
|
||||
|
||||
### Key Configuration Fields
|
||||
|
||||
@@ -408,18 +408,17 @@ Understanding the data structures is crucial for implementing stage transitions:
|
||||
|
||||
**Input to your function:**
|
||||
- `stage_list[source_stage_id].engine_outputs`: List of `EngineCoreOutput` objects
|
||||
- Each contains `outputs`: List of `RequestOutput` objects
|
||||
- Each `RequestOutput` has:
|
||||
- `token_ids`: Generated token IDs
|
||||
- `multimodal_output`: Dict with keys like `"code_predictor_codes"`, etc.
|
||||
- These are the hidden states or intermediate outputs from the model's forward pass
|
||||
- `prompt_token_ids`: Original prompt token IDs
|
||||
- - Each contains `outputs`: List of `RequestOutput` objects
|
||||
- Each `RequestOutput` has:
|
||||
- - - `token_ids`: Generated token IDs
|
||||
- `multimodal_output`: Dict with keys like `"code_predictor_codes"`, etc.These are the hidden states or intermediate outputs from the model's forward pass
|
||||
- `prompt_token_ids`: Original prompt token IDs
|
||||
|
||||
**Output from your function:**
|
||||
- Must return `list[OmniTokensPrompt]` where each `OmniTokensPrompt` contains:
|
||||
- `prompt_token_ids`: List[int] - Token IDs for the next stage
|
||||
- `additional_information`: Dict[str, Any] - Optional metadata (e.g., embeddings, hidden states)
|
||||
- `multi_modal_data`: Optional multimodal data if needed
|
||||
- - `prompt_token_ids`: List[int] - Token IDs for the next stage
|
||||
- `additional_information`: Dict[str, Any] - Optional metadata (e.g., embeddings, hidden states)
|
||||
- `multi_modal_data`: Optional multimodal data if needed
|
||||
|
||||
### How Model Outputs Are Stored
|
||||
|
||||
@@ -614,7 +613,7 @@ For a complete reference implementation, see:
|
||||
- **Thinker**: `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py`
|
||||
- **Talker**: `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py`
|
||||
- **Code2Wav**: `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py`
|
||||
- **Stage config**: `vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml`
|
||||
- **Stage config**: `vllm_omni/deploy/qwen3_omni_moe.yaml`
|
||||
- **Input processors**: `vllm_omni/model_executor/stage_input_processors/qwen3_omni.py`
|
||||
- **Registry**: `vllm_omni/model_executor/models/registry.py`
|
||||
- **Testing**: `vllm_omni/tests/e2e/offline_inference/test_qwen3_omni.py`
|
||||
|
||||
@@ -28,7 +28,7 @@ and can be placed on different devices. Qwen3-TTS has two stages:
|
||||
|
||||
Each stage is a separate model class configured independently via YAML. The two stages
|
||||
are connected by the `async_chunk` framework, which enables inter-stage streaming for
|
||||
low first-packet latency (see [Async Chunk Design](../../design/feature/async_chunk_design.md)).
|
||||
low first-packet latency (see [Async Chunk Design](../../design/feature/async_chunk.md)).
|
||||
|
||||
### Without async_chunk (batch mode)
|
||||
|
||||
@@ -120,8 +120,18 @@ vllm_omni/model_executor/stage_configs/
|
||||
| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
|
||||
| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | Stage 0 - optimized AR |
|
||||
| `models/qwen3_tts/qwen3_tts_code2wav.py` | Stage 1 - decoder |
|
||||
| `stage_configs/qwen3_tts.yaml` | Stage config (async_chunk enabled) |
|
||||
| `stage_configs/qwen3_tts_batch.yaml` | Batch mode config |
|
||||
| `deploy/qwen3_tts.yaml` (new schema) | Deploy config (async_chunk enabled) — paired with `models/qwen3_tts/pipeline.py` for the frozen topology |
|
||||
|
||||
> **Chunked vs end-to-end modes**: `qwen3_tts` registers a single
|
||||
> pipeline whose stage 1 declares alternate processor functions — an
|
||||
> `async_chunk_process_next_stage_input_func` (per-chunk streaming, used
|
||||
> when `deploy.async_chunk=True`) and a `sync_process_input_func`
|
||||
> (batch-end, used when `deploy.async_chunk=False`). The loader selects
|
||||
> one at merge time based on the bool, so `--no-async-chunk` alone
|
||||
> switches modes — no variant yaml or variant pipeline registration is
|
||||
> needed. Pipelines that only make sense in one mode (e.g.
|
||||
> `qwen3_omni_moe` is always chunked) can keep using the unconditional
|
||||
> `custom_process_*` fields.
|
||||
| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
|
||||
|
||||
## Step-by-Step Implementation
|
||||
@@ -574,11 +584,12 @@ Adding a TTS model to vLLM-Omni involves:
|
||||
| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
|
||||
| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | AR stage with vLLM fused ops |
|
||||
| `models/qwen3_tts/qwen3_tts_code2wav.py` | Decoder stage with `chunked_decode_streaming()` |
|
||||
| `stage_configs/qwen3_tts.yaml` | Stage configuration |
|
||||
| `models/qwen3_tts/pipeline.py` | Frozen pipeline topology (registered at import time) |
|
||||
| `deploy/qwen3_tts.yaml` | Deploy config (user-editable, async_chunk + SharedMemoryConnector) |
|
||||
| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
|
||||
|
||||
For more information, see:
|
||||
|
||||
- [Architecture Overview](../../design/architecture_overview.md)
|
||||
- [Async Chunk Design](../../design/feature/async_chunk_design.md)
|
||||
- [Async Chunk Design](../../design/feature/async_chunk.md)
|
||||
- [Stage Configuration Guide](../../configuration/stage_configs.md)
|
||||
@@ -127,19 +127,38 @@ Multi-stage omni serving:
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B \
|
||||
--omni \
|
||||
--stage-configs-path qwen2_5_omni.yaml \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
(The default deploy config at `vllm_omni/deploy/qwen2_5_omni.yaml` is loaded automatically. Pass `--deploy-config /path/to/custom.yaml` to override.)
|
||||
|
||||
Single-stage diffusion serving with torch profiler:
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--profiler-config '{"profiler": "torch", "torch_profiler_dir": "./vllm_profile"}'
|
||||
--profiler-config '{
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": "/tmp/vllm_profile_wan22_i2v",
|
||||
"torch_profiler_with_stack": true,
|
||||
"torch_profiler_with_flops": false,
|
||||
"torch_profiler_use_gzip": true,
|
||||
"torch_profiler_dump_cuda_time_total": true,
|
||||
"torch_profiler_record_shapes": true,
|
||||
"torch_profiler_with_memory": true,
|
||||
"delay_iterations": 0,
|
||||
"max_iterations": 0
|
||||
}'
|
||||
```
|
||||
|
||||
Useful optional fields:
|
||||
|
||||
- `torch_profiler_with_stack`: export `by_stack` operator views and stack text files.
|
||||
- `torch_profiler_record_shapes`: export `by_shape` operator views.
|
||||
- `torch_profiler_with_memory`: dump `memory_snapshot_rank*.pickle` when the backend supports memory history.
|
||||
- `torch_profiler_use_gzip`: write the trace as `trace_rank*.json.gz`.
|
||||
|
||||
Single-stage diffusion serving with Nsight Systems:
|
||||
|
||||
```bash
|
||||
@@ -176,8 +195,11 @@ For mixed-stage pipelines, use explicit `stages` and pass the same stage list to
|
||||
|
||||
Torch profiler output:
|
||||
|
||||
- Chrome/Perfetto traces under `torch_profiler_dir`
|
||||
- Optional aggregated CUDA-time tables under the same directory
|
||||
- Chrome/Perfetto trace: `trace_rank*.json` or `trace_rank*.json.gz`
|
||||
- Excel workbook: `ops_rank*.xlsx` with `summary`, and optional `by_shape` / `by_stack` sheets
|
||||
- Stack exports: `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt` when stack capture is enabled
|
||||
- Memory snapshot: `memory_snapshot_rank*.pickle` when memory capture is enabled and supported by the backend
|
||||
- Optional aggregated CUDA-time tables under the same session directory
|
||||
|
||||
CUDA profiler / Nsight Systems output:
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
# Async Chunk Design
|
||||
# Async Chunk
|
||||
|
||||
## Table of Contents
|
||||
|
||||
@@ -88,8 +88,9 @@ The following diagram illustrates the **Async Chunk Architecture** for multi-sta
|
||||
</p>
|
||||
|
||||
**Diagram Legend:**
|
||||
|
||||
| Step | Stage Type | Description |
|
||||
|:------:|:-----------:|:------------|
|
||||
|------|-----------|------------|
|
||||
| `prefill` | Initialization | Context processing, KV cache initialization |
|
||||
| `decode` | Autoregressive | Token-by-token generation in AR stages |
|
||||
| `codes` | Audio Encoding | RVQ codec codes from Talker stage |
|
||||
@@ -207,9 +207,9 @@ Complete examples in the codebase:
|
||||
|
||||
| Model | Path | Pattern | Notes |
|
||||
|-------|------|---------|-------|
|
||||
| **HunyuanImage3.0** | `vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py` | Standard EP | Full implementation with validation |
|
||||
| **HunyuanImage3.0** | `vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py` | Standard EP | Full implementation with validation |
|
||||
| **EP Tests** | `vllm-omni/tests/e2e/offline_inference/test_expert_parallel.py` | E2E testing | EP correctness and performance |
|
||||
| **Constraint Tests** | `vllm-omni/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py` | Unit testing | Validation logic |
|
||||
| **Constraint Tests** | `vllm-omni/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py` | Unit testing | Validation logic |
|
||||
|
||||
---
|
||||
## Summary
|
||||
|
||||
@@ -0,0 +1,164 @@
|
||||
# Automatic Prefix Caching in Omni Models
|
||||
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [High-Level Approach](#high-level-approach)
|
||||
- [Example](#example)
|
||||
- [What About Multimodal Inputs?](#what-about-multimodal-inputs)
|
||||
|
||||
---
|
||||
|
||||
### Overview
|
||||
|
||||
Prefix caching in the context of kv-cache management is a useful optimization for avoiding redundant computations. The main idea is that we store portions of the kv-cache from processed requests, so that we can reuse them if incoming requests have the same prefix as previous requests.
|
||||
|
||||
vLLM manages the kv-cache as blocks, which represent a span of tokens of a fixed length. Blocks are hashable by the content that they contain, which typically means the tokens within the span, but also could be influenced by other factors, e.g., LoRA and multimodal data.
|
||||
|
||||
vLLM implements automatic prefix caching for managing its kv-cache, which is best understood by reading the design document [here](https://docs.vllm.ai/en/latest/design/prefix_caching/). vLLM-Omni builds on top of the prefix caching mechanism in a noninvasive way to allow caching between stages in Omni pipelines. This typically means for a given stage we aim to support caching for the following:
|
||||
|
||||
- The last hidden states produced by the stage
|
||||
- Model / stage specific multimodal data
|
||||
|
||||
!!! note "Note 1"
|
||||
This document describes vLLM-Omni's mechanism for caching tensor outputs that are meant to be passed between stages, when requests have common prefixes, similar to the way in which vLLM has prefix caching for the kv-cache. This works in conjunction with vLLM's multimodal encoder caching, but is distinct. See the final section for a concrete example for how they tie together in practice.
|
||||
|
||||
### High-Level Approach
|
||||
!!! note "Note 2"
|
||||
Prior to reading this section, it's recommended to take a look at the design documents in vLLM for [Automatic Prefix Caching](https://docs.vllm.ai/en/latest/features/automatic_prefix_caching/), which will make some of the concepts more clear.
|
||||
|
||||
The main focus of vLLM-Omni's approach to prefix caching stage outputs is to build on vLLM's prefix caching in the least invasive way possible while minimizing impact for cache misses, and consuming a minimal amount of GPU memory. To understand the implementation, there are a few important things to note:
|
||||
|
||||
- Between stages, device tensors are generally moved to CPU; this is important since we're just caching the outputs of stages, so it is okay to keep the entire cache on the CPU.
|
||||
|
||||
- For a tensor to be considered cacheable, the first dimension (currently) needs to be the same as the token count, as it allows us to reuse block/slot mappings for our externally maintained tensor caches. This allows us to dynamically discover the tensors to be marked as cacheable outputs in each Omni model without having to explicitly specify cacheable output field names in every model.
|
||||
|
||||
With this in mind, consider the set of blocks in a 2D layout, where the row represents the index of blocks being considered, and the columns represent the slots corresponding to tokens within each block. Since we know the `num_blocks` and `block_size` from our kv cache config, if we want to cache a tensor with feature size `D`, we can preallocate a CPU tensor of size `(num_blocks, block_size, D)`, and use the same block index and slot mapping to retrieve the corresponding feature vector.
|
||||
|
||||
|
||||
### Example
|
||||
!!! note "Note 3"
|
||||
Prefix caching in vLLM-Omni currently is only supported on AutoRegressive stages with one kv-cache group. It can be enabled/disabled per-stage via the `enable_prefix_caching` parameter in the model's stage config.
|
||||
|
||||
The way in which vLLM-Omni ties into vLLM's prefix caching is best understood by example. Say that we have the following:
|
||||
|
||||
- `num_blocks=8`
|
||||
- `block_size=4`
|
||||
- `hidden_size=2`
|
||||
- A stage specific multimodal output tensor named `mm_feature` with feature dimension `16`
|
||||
|
||||
The prefix cache flow is then outlined below.
|
||||
|
||||
1. When the model is initialized, we can determine the `hidden_size` from the `ModelConfig`, and allocate a cache of size `(num_blocks, block_size, hidden_size)`.
|
||||
|
||||
2. Say we process the request `The quick brown fox was tired and slept beneath the shady tree`, which is 12 tokens and evenly divides into 3 blocks as shown below.
|
||||
|
||||
```
|
||||
[ The quick brown fox ] [ was tired and slept ] [beneath the shady tree ]
|
||||
Block 1: |<--- block tokens ---->|
|
||||
Block 2: |<------- prefix ------>| |<--- block tokens --->|
|
||||
Block 3: |<------------------ prefix -------------------->| |<--- block tokens ---->|
|
||||
```
|
||||
|
||||
When the request processes, we inspect the multimodal outputs and identify the `mm_feature` tensor, which will be of shape `(seq_len, feature_dim)`, i.e., `(12, 16)` in this example. We note that the first axis is dependent on the `seq_len` and add a new cache_tensor of shape `(num_blocks, block_size, feature_dim)` to our multimodal cache for tensors.
|
||||
|
||||
|
||||
3. If we lay out the cache as a 2D tensor of shape (`num_blocks`, `block_size`), we'll have something like the following:
|
||||
|
||||
```
|
||||
0: [ The quick brown fox ]
|
||||
1: [ was tired and slept ]
|
||||
2: [beneath the shady tree ]
|
||||
3: [EMPTY]
|
||||
...
|
||||
7: [EMPTY]
|
||||
```
|
||||
|
||||
Or, if we flatten it down to 1D,
|
||||
```
|
||||
0: The
|
||||
1: quick
|
||||
2: brown
|
||||
3: fox
|
||||
...
|
||||
11: tree
|
||||
12: [EMPTY]
|
||||
...
|
||||
```
|
||||
|
||||
which we can think of as row indices into the hidden states tensor if we view it as the 2D shape `(num_blocks x block_size, feature_dim)`. That is, the analogous flattened (from 3D -> 2D) mapping of the cache for hidden states becomes the following.
|
||||
```
|
||||
0: <hidden states vector of len 2 corresponding to 'The'>
|
||||
1: <hidden states vector of len 2 corresponding to 'quick'>
|
||||
2: <hidden states vector of len 2 corresponding to 'brown'>
|
||||
3: <hidden states vector of len 2 corresponding to 'fox'>
|
||||
...
|
||||
11: <hidden states vector of len 2 corresponding to 'tree'>
|
||||
12: [EMPTY]
|
||||
...
|
||||
```
|
||||
|
||||
Similarly, for the multimodal outputs cache, the flattened coordinates are the same, but the `mm_feature` maps to vectors of length `16` instead of the hidden size of `2`. Note that in practice, we may have multiple multimodal output tensors per forward pass, which may have different names and different feature dimensions.
|
||||
|
||||
|
||||
4. Now, say that we receive a new request `The quick brown fox jumped over the dog`.
|
||||
|
||||
```
|
||||
[ The quick brown fox ] [ jumped over the dog ]
|
||||
Block 1: |<--- block tokens ---->|
|
||||
Block 2: |<------- prefix ------>| |<--- block tokens --->|
|
||||
```
|
||||
|
||||
Here, we will have a cache hit for `Block 1` which will be detected by vLLM based on the hash of the first block when it's handling the prefix caching on the kv-cache. As a result, when we get the output from the scheduler, we will see that `num_computed_tokens=4` (corresponding to the cached first block), and we only need to process the remaining 4 new tokens in the new prefill.
|
||||
|
||||
Since we have the block indices / slot mappings from the kv cache manager, we can simply mirror the mappings and leverage the same indices for the cached hidden states and multimodal outputs. This allows us to look up the correct tensors from our externally maintained 3D caches.
|
||||
|
||||
```
|
||||
0: [ The quick brown fox ] < already in the cache
|
||||
1: [ was tired and slept ]
|
||||
2: [beneath the shady tree ]
|
||||
3: [ jumped over the dog ] < added on the second request
|
||||
4: [EMPTY]
|
||||
...
|
||||
7: [EMPTY]
|
||||
...
|
||||
```
|
||||
|
||||
Finally, to pass the full hidden states and multimodal outputs to the next stage, we simply concatenate the cached contents with the corresponding new tensors computed from the current forward call.
|
||||
|
||||
|
||||
### What About Multimodal Inputs?
|
||||
It's also useful to consider the case about how Omni prefix caching is handled when we have multimodal inputs that don't cleanly end on block boundaries, as well as how this works with multimodal encoder caching in vLLM. For example:
|
||||
|
||||
```
|
||||
[ Im0 Im1 Im2 Im3 ] [ Im4 Im5 foo <empty> ]
|
||||
Block 1: |<--- block tokens ---->|
|
||||
Block 2: |<------- prefix ------>| |<--- block tokens --->|
|
||||
```
|
||||
|
||||
In this case, only `Block 1` will have outputs stored in the prefix tensor cache, because vLLM does not store partial blocks. This may appear to be a problem at first glance, because the multimodal input is fragmented across a new block that wasn't cached.
|
||||
|
||||
In reality, this isn't a big problem for correctness, because vLLM also maintains an encoder cache for multimodal inputs. In other words, after the first pass, we'll have the following:
|
||||
|
||||
- The Block 1 hash, which is used for prefix caching
|
||||
- The hash describing the image data starting at position 0 and with length 6
|
||||
- In vLLM's encoder cache, a mapping from the image hash above to the encoder output
|
||||
|
||||
|
||||
To understand what happens, say we get the following input as a second request:
|
||||
```
|
||||
[ Im0 Im1 Im2 Im3 ] [ Im4 Im5 bar baz ]
|
||||
Block 1: |<--- block tokens ---->|
|
||||
Block 2: |<------- prefix ------>| |<--- block tokens --->|
|
||||
```
|
||||
|
||||
First, the scheduler will check for a prefix cache hit, which we will see on `Block 1`. As a result, we will have 4 tokens marked as precomputed, and only see the remaining 4 tokens in the following prefill.
|
||||
|
||||
Because we have multimodal data in a scheduled span that isn't fully precomputed, we still need to call the visual encoder. However, since we have the image hash and encoder cache, we will retrieve the encoder outputs for `Im4` and `Im5` as we create the multimodal embeddings.
|
||||
|
||||
When we pass our multimodal tensors to the language model component in the same stage, we'll then expect the same outputs, because the prefix caching behaviors in vLLM-Omni / vLLM match, so the LLM will use vLLM's KV cache manager's prefix caching to correctly handle the attention information for `Block 1` while calculating the outputs for `Block 2`, giving us the correct results for processing `Block 2` with the context of `Block 1`.
|
||||
|
||||
Finally, we look up the output hidden states/multimodal tensors corresponding to the prefix cache hit `Block 1` and concatenate it with the forward pass result to get the final result, which is expected to be identical to the full hidden states when prefix caching is disabled.
|
||||
@@ -326,9 +326,41 @@ for prompt in tqdm(prompts, desc="Collecting data"):
|
||||
|
||||
# Estimate coefficients
|
||||
coeffs = estimator.estimate(poly_order=4)
|
||||
print(f"Estimated coefficients: {coeffs.tolist()}")
|
||||
print(f"Estimated coefficients: {coeffs}")
|
||||
```
|
||||
|
||||
Note: some models may require the vLLM context and config to be initialized to initialize vLLM modules. To this end, you may need a workaround like the following to be able to run coefficient estimation.
|
||||
```python
|
||||
from vllm_omni.diffusion.forward_context import set_forward_context
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm.config import VllmConfig
|
||||
...
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ["MASTER_ADDR"] = "localhost"
|
||||
os.environ["MASTER_PORT"] = "8192"
|
||||
os.environ["LOCAL_RANK"] = "0"
|
||||
os.environ["RANK"] = "0"
|
||||
os.environ["WORLD_SIZE"] = "1"
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
init_distributed_environment()
|
||||
initialize_model_parallel()
|
||||
|
||||
# NOTE: you may have to pass an initialized OmniDiffusionConfig as a kwarg
|
||||
# here to make current sp checks happy; if this is the case, just create one
|
||||
# .from_kwargs() with the model name to get around this check for now,
|
||||
# since your estimator subclass should handle the actual model configuration.
|
||||
#
|
||||
# This will be cleaned up in the future
|
||||
with set_forward_context(vllm_config):
|
||||
<create the estimator + run estimation here>
|
||||
```
|
||||
|
||||
|
||||
**Data Statistics Guide:**
|
||||
|
||||
| Metric | Good Range | Warning Signs |
|
||||
|
||||
|
After Width: | Height: | Size: 62 KiB |
|
After Width: | Height: | Size: 79 KiB |
|
After Width: | Height: | Size: 103 KiB |
|
After Width: | Height: | Size: 82 KiB |
|
After Width: | Height: | Size: 97 KiB |
|
After Width: | Height: | Size: 130 KiB |
|
After Width: | Height: | Size: 120 KiB |
|
After Width: | Height: | Size: 90 KiB |
|
After Width: | Height: | Size: 125 KiB |
|
After Width: | Height: | Size: 111 KiB |
|
After Width: | Height: | Size: 77 KiB |
|
After Width: | Height: | Size: 119 KiB |
|
After Width: | Height: | Size: 119 KiB |
|
After Width: | Height: | Size: 134 KiB |
|
After Width: | Height: | Size: 72 KiB |
|
After Width: | Height: | Size: 37 KiB |
|
After Width: | Height: | Size: 36 KiB |
|
After Width: | Height: | Size: 40 KiB |
|
After Width: | Height: | Size: 48 KiB |
|
After Width: | Height: | Size: 53 KiB |
|
After Width: | Height: | Size: 56 KiB |
|
After Width: | Height: | Size: 42 KiB |
|
After Width: | Height: | Size: 50 KiB |
|
After Width: | Height: | Size: 56 KiB |
|
After Width: | Height: | Size: 52 KiB |
|
After Width: | Height: | Size: 56 KiB |
|
After Width: | Height: | Size: 59 KiB |
|
After Width: | Height: | Size: 91 KiB |
|
After Width: | Height: | Size: 90 KiB |
|
After Width: | Height: | Size: 98 KiB |
@@ -0,0 +1,539 @@
|
||||
# Speech Generation on vLLM-Omni: Performance Optimizations for Qwen3-Omni and Qwen3-TTS
|
||||
|
||||
## Summary
|
||||
|
||||
vLLM-Omni supports end-to-end serving for speech-generating models, including both **Qwen3-Omni** (multimodal understanding + speech) and **Qwen3-TTS** (text-to-speech). Despite their different architectures, both models share the same multi-stage pipeline design and benefit from the same set of stacked optimizations:
|
||||
|
||||
1. **Batching** improves GPU utilization stage by stage and increases overall throughput.
|
||||
2. **CUDA Graph** reduces CPU launch overhead and decode-time jitter on stable shapes.
|
||||
3. **Async Chunk and Streaming Output** overlap compute and communication across stages and emit audio incrementally, improving both TTFP and E2E.
|
||||
|
||||
### Model architectures
|
||||
|
||||
**Qwen3-Omni** is a native multimodal model that understands text, audio, image, and video inputs, and generates both text and speech outputs. Its pipeline has three stages:
|
||||
|
||||
- **Thinker**: multimodal understanding and text generation
|
||||
- **Talker (+ Talker-MTP / code predictor path)**: converts semantic/text representations into codec tokens
|
||||
- **Code2Wav**: decodes codec tokens into waveform audio
|
||||
|
||||
**Qwen3-TTS** is a lightweight, high-quality text-to-speech model. Its pipeline has two stages:
|
||||
|
||||
- **Talker (AR decoder)**: auto-regressively generates codec tokens from text input
|
||||
- **Code2Wav (vocoder)**: decodes codec tokens into waveform audio
|
||||
|
||||
The optimizations described in this post apply to both models. We present results for each side by side.
|
||||
|
||||
### vLLM-Omni vs HF Transformers
|
||||
|
||||
Compared with **HF Transformers** (offline, single request), vLLM-Omni with the full optimization stack delivers dramatically lower latency and higher efficiency for both models.
|
||||
|
||||
**Qwen3-Omni** (A100):
|
||||
|
||||
<table><tr>
|
||||
<td><img src="figures/omni/E2EL_s_vllm_omni_vs_transformers.png" alt="Qwen3-Omni E2EL: vLLM vs HF" width="100%"/></td>
|
||||
<td><img src="figures/omni/TTFP_s_vllm_omni_vs_transformers.png" alt="Qwen3-Omni TTFP: vLLM vs HF" width="100%"/></td>
|
||||
<td><img src="figures/omni/RTF_vllm_omni_vs_transformers.png" alt="Qwen3-Omni RTF: vLLM vs HF" width="100%"/></td>
|
||||
</tr></table>
|
||||
|
||||
| Metric | vLLM-Omni | HF Transformers | Improvement |
|
||||
| --- | --- | --- | --- |
|
||||
| E2E latency (s) | 23.78 | 336.10 | ~93% reduction |
|
||||
| TTFP (s) | 0.934 | 336.10 | ~99.7% reduction |
|
||||
| RTF | 0.32 | 3.776 | ~91% reduction (~12× faster) |
|
||||
|
||||
- **E2E latency**: 23.78 s vs 336.10 s - **~93%** reduction
|
||||
- **TTFP**: 0.934 s vs 336.10 s - **~99.7%** reduction
|
||||
- **RTF**: 0.32 vs 3.776 - **~91%** reduction (~12x faster)
|
||||
|
||||
**Qwen3-TTS** (H200, concurrency 1):
|
||||
|
||||
<table><tr>
|
||||
<td><img src="figures/tts/Mean_E2EL_(ms)_vllm_omni_vs_transformers.png" alt="Qwen3-TTS E2EL: vLLM vs HF" width="100%"/></td>
|
||||
<td><img src="figures/tts/Mean_AUDIO_TTFP_(ms)_vllm_omni_vs_transformers.png" alt="Qwen3-TTS TTFP: vLLM vs HF" width="100%"/></td>
|
||||
<td><img src="figures/tts/Mean_AUDIO_RTF_vllm_omni_vs_transformers.png" alt="Qwen3-TTS RTF: vLLM vs HF" width="100%"/></td>
|
||||
</tr></table>
|
||||
|
||||
| Metric | vLLM-Omni | HF Transformers | Improvement |
|
||||
| --- | --- | --- | --- |
|
||||
| E2E latency (ms) | 941 | 15,513 | ~94% reduction |
|
||||
| TTFP (ms) | 64 | 15,513 | ~99.6% reduction (242× faster) |
|
||||
| RTF | 0.16 | 2.64 | ~94% reduction (~16.5× faster) |
|
||||
|
||||
- **E2E latency**: 941 ms vs 15,513 ms - **~94%** reduction
|
||||
- **TTFP**: 64 ms vs 15,513 ms - **~99.6%** reduction (242x faster)
|
||||
- **RTF**: 0.16 vs 2.64 - **~94%** reduction (~16.5x faster)
|
||||
|
||||
### Stacked optimization summary
|
||||
|
||||
Each optimization stacks on the previous one. The summary plots below show the cumulative effect at each step, with one line per concurrency level (1, 4, 10).
|
||||
|
||||
**Qwen3-Omni** (A100):
|
||||
|
||||
<table><tr>
|
||||
<td><img src="figures/omni/Summary_E2EL_ms_vs_features.png" alt="Qwen3-Omni E2EL: stacked optimization" width="100%"/></td>
|
||||
<td><img src="figures/omni/Summary_TTFP_ms_vs_features.png" alt="Qwen3-Omni TTFP: stacked optimization" width="100%"/></td>
|
||||
<td><img src="figures/omni/Summary_RTF_vs_features.png" alt="Qwen3-Omni RTF: stacked optimization" width="100%"/></td>
|
||||
</tr></table>
|
||||
|
||||
- **E2EL reduction**: ~74% at concurrency 10 (410,054 ms -> 104,901 ms); ~90% at concurrency 1 (426,529 ms -> 41,216 ms)
|
||||
- **TTFP reduction**: ~96% at concurrency 10 (409,705 ms -> 16,482 ms); ~99.7% at concurrency 1 (426,078 ms -> 1,164 ms)
|
||||
- **RTF reduction**: ~74% at concurrency 10 (2.83 -> 0.74); ~90% at concurrency 1 (2.08 -> 0.21)
|
||||
|
||||
**Qwen3-TTS** (H200):
|
||||
|
||||
<table><tr>
|
||||
<td><img src="figures/tts/Summary_mean_e2e_ms_vs_features.png" alt="Qwen3-TTS E2EL: stacked optimization" width="100%"/></td>
|
||||
<td><img src="figures/tts/Summary_mean_ttfp_ms_vs_features.png" alt="Qwen3-TTS TTFP: stacked optimization" width="100%"/></td>
|
||||
<td><img src="figures/tts/Summary_mean_rtf_vs_features.png" alt="Qwen3-TTS RTF: stacked optimization" width="100%"/></td>
|
||||
</tr></table>
|
||||
|
||||
- **E2EL reduction**: ~85% at concurrency 10 (12,141 ms -> 1,767 ms); ~29% at concurrency 1 (1,323 ms -> 941 ms)
|
||||
- **TTFP reduction**: ~96.5% at concurrency 10 (12,141 ms -> 425 ms); ~95% at concurrency 1 (1,323 ms -> 64 ms)
|
||||
- **RTF reduction**: ~86% at concurrency 10 (2.19 -> 0.31); ~30% at concurrency 1 (0.23 -> 0.16)
|
||||
|
||||
**Benchmark environment:**
|
||||
|
||||
| | Qwen3-Omni | Qwen3-TTS |
|
||||
| --- |-----------------------------| --- |
|
||||
| **GPU** | A100 | H200 |
|
||||
| **Model** | Qwen3-Omni-30B-A3B-Instruct | Qwen3-TTS-12Hz-1.7B-CustomVoice |
|
||||
| **vLLM** | v0.17.0 | v0.18.0 |
|
||||
| **vllm-omni** | commit 199f7832 | v0.18.0rc2 |
|
||||
| **CUDA** | 12.9 | 12.8 |
|
||||
|
||||
This post walks through each optimization in the same order they are typically enabled in practice, then ends with deployment playbooks for both models.
|
||||
|
||||
---
|
||||
|
||||
## Pipeline Batching
|
||||
|
||||
### How stage-wise batching works
|
||||
|
||||
For both Qwen3-Omni and Qwen3-TTS, batching is a pipeline-level optimization:
|
||||
|
||||
- Requests are grouped per stage using `runtime.max_batch_size`
|
||||
- Each stage executes batch inference with its own scheduler/worker
|
||||
- Stage outputs are routed to downstream stages with per-request mapping preserved
|
||||
|
||||
**Batching strategy by stage:** The understanding and decode stages (Thinker for Omni, Talker for both) use **continuous batching**: requests can join and leave the batch over time. Code2Wav uses **static batching**: once a batch is formed, the stage runs the whole batch before starting the next. This matches the decode pattern of Code2Wav and keeps implementation simple while still improving throughput.
|
||||
|
||||
### Batching results (Baseline vs. Batch)
|
||||
|
||||
Batching alone greatly reduces E2EL and RTF across all concurrencies. The biggest gains appear at high concurrency where requests share GPU resources.
|
||||
|
||||
**Qwen3-Omni** (A100):
|
||||
|
||||
<table><tr>
|
||||
<td><img src="figures/omni/Mean_E2EL_ms_Baseline_vs_Batch.png" alt="Qwen3-Omni E2EL: Baseline vs Batch" width="100%"/></td>
|
||||
<td><img src="figures/omni/Mean_AUDIO_TTFP_ms_Baseline_vs_Batch.png" alt="Qwen3-Omni TTFP: Baseline vs Batch" width="100%"/></td>
|
||||
<td><img src="figures/omni/Mean_AUDIO_RTF_Baseline_vs_Batch.png" alt="Qwen3-Omni RTF: Baseline vs Batch" width="100%"/></td>
|
||||
</tr></table>
|
||||
|
||||
| Metric | Concurrency | Baseline | + Batch | Improvement |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| E2EL (ms) | 1 | 426,529 | 307,719 | 1.4× |
|
||||
| E2EL (ms) | 4 | 407,213 | 376,934 | 1.1× |
|
||||
| E2EL (ms) | 10 | 410,054 | 234,844 | 1.7× |
|
||||
| TTFP (ms) | 1 | 426,078 | 307,262 | 1.4× |
|
||||
| TTFP (ms) | 4 | 406,843 | 376,466 | 1.1× |
|
||||
| TTFP (ms) | 10 | 409,705 | 234,557 | 1.7× |
|
||||
| RTF | 1 | 2.08 | 1.51 | 1.4× |
|
||||
| RTF | 4 | 2.55 | 1.83 | 1.4× |
|
||||
| RTF | 10 | 2.83 | 2.28 | 1.2× |
|
||||
|
||||
At concurrency 10, E2EL drops from ~410 s to ~235 s; at concurrency 1, from ~427 s to ~308 s.
|
||||
|
||||
**Qwen3-TTS** (H200):
|
||||
|
||||
<table><tr>
|
||||
<td><img src="figures/tts/Mean_mean_e2e_ms_baseline_vs_batch.png" alt="Qwen3-TTS E2EL: Baseline vs Batch" width="100%"/></td>
|
||||
<td><img src="figures/tts/Mean_mean_ttfp_ms_baseline_vs_batch.png" alt="Qwen3-TTS TTFP: Baseline vs Batch" width="100%"/></td>
|
||||
<td><img src="figures/tts/Mean_mean_rtf_baseline_vs_batch.png" alt="Qwen3-TTS RTF: Baseline vs Batch" width="100%"/></td>
|
||||
</tr></table>
|
||||
|
||||
| Metric | Concurrency | Baseline | + Batch | Improvement |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| E2EL (ms) | 1 | 1,323 | 1,339 | 1.0× |
|
||||
| E2EL (ms) | 4 | 5,171 | 1,471 | 3.5× |
|
||||
| E2EL (ms) | 10 | 12,141 | 1,705 | 7.1× |
|
||||
| RTF | 1 | 0.230 | 0.234 | 1.0× |
|
||||
| RTF | 4 | 0.908 | 0.255 | 3.6× |
|
||||
| RTF | 10 | 2.186 | 0.292 | 7.5× |
|
||||
| Throughput (audio-s/wall-s) | 10 | 3.99 | 33.53 | 8.4× |
|
||||
|
||||
At concurrency 10, batching alone brings Qwen3-TTS RTF from 2.19 (slower than realtime) down to 0.29 (faster than realtime), and throughput from 4.0 to 33.5 audio-sec/wall-sec.
|
||||
|
||||
---
|
||||
|
||||
## CUDA Graph on the Critical Decode Path
|
||||
|
||||
### Why CUDA Graph helps here
|
||||
|
||||
In decode-heavy serving, repeatedly launching many small kernels from CPU can become a visible overhead. CUDA Graph reduces this overhead by capturing and replaying stable execution graphs.
|
||||
|
||||
In stage configs, this is represented by `enforce_eager: false` for stages where graph capture is desired (Thinker/Talker), while Code2Wav keeps eager mode depending on stage behavior.
|
||||
|
||||
### CUDA Graph results on top of batching
|
||||
|
||||
**Qwen3-Omni** (A100):
|
||||
|
||||
<table><tr>
|
||||
<td><img src="figures/omni/Mean_E2EL_ms_Batch_vs_Batch_CUDA_Graph.png" alt="Qwen3-Omni E2EL: Batch vs CUDA Graph" width="100%"/></td>
|
||||
<td><img src="figures/omni/Mean_AUDIO_TTFP_ms_Batch_vs_Batch_CUDA_Graph.png" alt="Qwen3-Omni TTFP: Batch vs CUDA Graph" width="100%"/></td>
|
||||
<td><img src="figures/omni/Mean_AUDIO_RTF_Batch_vs_Batch_CUDA_Graph.png" alt="Qwen3-Omni RTF: Batch vs CUDA Graph" width="100%"/></td>
|
||||
</tr></table>
|
||||
|
||||
| Metric | Concurrency | Batch | + CUDA Graph | Improvement |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| E2EL (ms) | 1 | 307,719 | 61,613 | 5.0× |
|
||||
| E2EL (ms) | 4 | 376,934 | 79,019 | 4.8× |
|
||||
| E2EL (ms) | 10 | 234,844 | 126,867 | 1.9× |
|
||||
| TTFP (ms) | 1 | 307,262 | 61,257 | 5.0× |
|
||||
| TTFP (ms) | 4 | 376,466 | 78,634 | 4.8× |
|
||||
| TTFP (ms) | 10 | 234,557 | 126,534 | 1.9× |
|
||||
| RTF | 1 | 1.51 | 0.32 | 4.7× |
|
||||
| RTF | 4 | 1.83 | 0.43 | 4.3× |
|
||||
| RTF | 10 | 2.28 | 0.90 | 2.5× |
|
||||
|
||||
For the larger Qwen3-Omni model (30B-A3B), CUDA Graph provides a significant improvement. At concurrency 1, E2EL drops from ~308 s to ~62 s; at concurrency 10, from ~235 s to ~127 s.
|
||||
|
||||
**Qwen3-TTS** (H200):
|
||||
|
||||
<table><tr>
|
||||
<td><img src="figures/tts/Mean_mean_e2e_ms_batch_vs_cuda_graph.png" alt="TTS E2EL: Batch vs +CG" width="100%"/></td>
|
||||
<td><img src="figures/tts/Mean_mean_ttfp_ms_batch_vs_cuda_graph.png" alt="TTS TTFP: Batch vs +CG" width="100%"/></td>
|
||||
<td><img src="figures/tts/Mean_mean_rtf_batch_vs_cuda_graph.png" alt="TTS RTF: Batch vs +CG" width="100%"/></td>
|
||||
</tr></table>
|
||||
|
||||
| Metric | Concurrency | Batch | + CUDA Graph | Improvement |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| E2EL (ms) | 1 | 1,339 | 733 | 1.8× |
|
||||
| E2EL (ms) | 4 | 1,471 | 987 | 1.5× |
|
||||
| E2EL (ms) | 10 | 1,705 | 1,197 | 1.4× |
|
||||
| RTF | 1 | 0.234 | 0.124 | 1.9× |
|
||||
| RTF | 10 | 0.292 | 0.203 | 1.4× |
|
||||
| Throughput (audio-s/wall-s) | 10 | 33.53 | 47.15 | 1.4× |
|
||||
|
||||
At concurrency 1, CUDA Graph reduces E2EL from 1,339 ms to 733 ms and RTF from 0.234 to 0.124 - nearly a 2x improvement. The benefit is consistent across all concurrency levels.
|
||||
|
||||
---
|
||||
|
||||
## Async Chunk and Streaming Output: Earlier Audio and Cross-Stage Overlap
|
||||
|
||||
### Why this step matters for first-packet latency
|
||||
|
||||
Two mechanisms work together to improve user-visible latency:
|
||||
|
||||
- **Streaming output**: audio streaming emits audio chunks as soon as they are decoded (lower **TTFP**). Without streaming, the client waits for larger buffers or end-of-sequence.
|
||||
- **Async chunk** is the main enabler for *earlier* audio: instead of handing off whole-request results between stages, each stage forwards **chunks** so the next stage can start as soon as the first chunk is ready. For Omni: Thinker -> Talker forwards hidden-state chunks; for both: Talker -> Code2Wav forwards codec chunks; Code2Wav decodes and emits packets incrementally. This **overlaps compute and communication** across stages and directly reduces time-to-first-audio-packet (TTFP) and end-to-end latency (E2EL).
|
||||
|
||||
So in practice: streaming output defines *how* bytes are sent to the client; async chunk defines *when* the pipeline can produce the first bytes.
|
||||
|
||||
**Dependency between the two:** Async chunk and audio streaming output are mutually dependent. Without async chunk, **audio streaming output cannot truly take effect**. Without audio streaming output, async chunk's **TTFP advantage is not fully realized**: the client would still wait for larger buffers or end-of-sequence instead of hearing the first packet as soon as it is ready. We therefore recommend enabling **both** on top of batching + CUDA Graph; the benchmarks in this post use both.
|
||||
|
||||
### Results: Batch + CUDA Graph vs. Batch + CUDA Graph + Async Chunk + Streaming Output
|
||||
|
||||
**Qwen3-Omni** (A100):
|
||||
|
||||
<table><tr>
|
||||
<td><img src="figures/omni/Mean_E2EL_ms_Batch_CUDA_Graph_vs_Async_Chunk.png" alt="Qwen3-Omni E2EL: CG vs Async Chunk" width="100%"/></td>
|
||||
<td><img src="figures/omni/Mean_AUDIO_TTFP_ms_Batch_CUDA_Graph_vs_Async_Chunk.png" alt="Qwen3-Omni TTFP: CG vs Async Chunk" width="100%"/></td>
|
||||
<td><img src="figures/omni/Mean_AUDIO_RTF_Batch_CUDA_Graph_vs_Async_Chunk.png" alt="Qwen3-Omni RTF: CG vs Async Chunk" width="100%"/></td>
|
||||
</tr></table>
|
||||
|
||||
| Metric | Concurrency | Batch + CG | + Async Chunk | Improvement |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| E2EL (ms) | 1 | 61,613 | 41,216 | 1.5× |
|
||||
| E2EL (ms) | 4 | 79,019 | 67,584 | 1.2× |
|
||||
| E2EL (ms) | 10 | 126,867 | 104,901 | 1.2× |
|
||||
| TTFP (ms) | 1 | 61,257 | 1,164 | 53× |
|
||||
| TTFP (ms) | 4 | 78,634 | 3,152 | 24.9× |
|
||||
| TTFP (ms) | 10 | 126,534 | 16,482 | 7.7× |
|
||||
| RTF | 1 | 0.32 | 0.21 | 1.5× |
|
||||
| RTF | 4 | 0.43 | 0.34 | 1.3× |
|
||||
| RTF | 10 | 0.90 | 0.74 | 1.2× |
|
||||
|
||||
Enabling both brings TTFP down sharply (concurrency 1: 61,257 ms -> 1,164 ms, **~98% reduction**; concurrency 4: 78,634 ms -> 3,152 ms, **~96% reduction**). E2EL and RTF also improve at every concurrency.
|
||||
|
||||
**Qwen3-TTS** (H200):
|
||||
|
||||
<table><tr>
|
||||
<td><img src="figures/tts/Mean_mean_e2e_ms_cuda_graph_vs_async_chunk.png" alt="Qwen3-TTS E2EL: CG vs Async Chunk" width="100%"/></td>
|
||||
<td><img src="figures/tts/Mean_mean_ttfp_ms_cuda_graph_vs_async_chunk.png" alt="Qwen3-TTS TTFP: CG vs Async Chunk" width="100%"/></td>
|
||||
<td><img src="figures/tts/Mean_mean_rtf_cuda_graph_vs_async_chunk.png" alt="Qwen3-TTS RTF: CG vs Async Chunk" width="100%"/></td>
|
||||
</tr></table>
|
||||
|
||||
| Metric | Concurrency | Batch + CG | + Async Chunk | Improvement |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| TTFP (ms) | 1 | 733 | **64** | **11.5×** |
|
||||
| TTFP (ms) | 4 | 987 | **119** | **8.3×** |
|
||||
| TTFP (ms) | 10 | 1,197 | **425** | **2.8×** |
|
||||
| E2EL (ms) | 1 | 733 | 941 | 0.8× |
|
||||
| E2EL (ms) | 10 | 1,197 | 1,767 | 0.7× |
|
||||
| RTF | 1 | 0.124 | 0.160 | 0.8× |
|
||||
| RTF | 10 | 0.203 | 0.314 | 0.6× |
|
||||
|
||||
The TTFP improvement is the headline result for both models. For Qwen3-TTS at concurrency 1, users hear the first audio in **64 ms** instead of 733 ms - an **11.5x reduction**. For Qwen3-Omni at concurrency 1, TTFP drops from 61 s to 1.2 s - a **53x reduction**.
|
||||
|
||||
### Why E2EL and RTF are higher with async chunk (TTS)
|
||||
|
||||
The table above shows that enabling async chunk + streaming *increases* E2EL and RTF for TTS compared to CUDA Graph alone. This is expected - the two configurations optimize for fundamentally different metrics:
|
||||
|
||||
- **CUDA Graph (no async chunk)** generates the entire audio end-to-end before returning. No chunking overhead, so total compute is minimized.
|
||||
- **Async Chunk + Streaming** splits the pipeline into incremental chunks, adding overhead from chunked transport, context overlap in Code2Wav (`codec_left_context_frames=25`), and smaller effective batch sizes per chunk.
|
||||
|
||||
**The tradeoff is intentional.** Async chunk trades ~30% higher total compute for **11x faster time-to-first-audio**. For interactive applications (voice assistants, chatbots), TTFP determines perceived responsiveness. For offline batch processing, CUDA Graph without async chunk is the better choice.
|
||||
|
||||
---
|
||||
|
||||
## TTS-Specific: Code Predictor Re-prefill + `torch.compile`
|
||||
|
||||
Qwen3-TTS has a **code predictor** - a small 5-layer transformer that generates residual codebook tokens (groups 1 through Q-1) autoregressively. Each AR step operates on very short sequences (2 to ~16 tokens).
|
||||
|
||||
The naive approach uses a KV cache for this small transformer, similar to the main Talker. But the KV cache machinery (block tables, slot mappings, paged attention) introduces significant overhead relative to the tiny model. Two optimizations replace that:
|
||||
|
||||
### Re-prefill (stateless forward, no KV cache)
|
||||
|
||||
Instead of maintaining a KV cache across steps, the code predictor **re-feeds the full growing sequence** at each AR step using `F.scaled_dot_product_attention`. With sequences of at most ~16 tokens through 5 layers, the O(T^2) attention cost is negligible - and removing the KV cache machinery (block table management, `set_forward_context`, slot mapping) saves far more time than it costs.
|
||||
|
||||
### `torch.compile` on the code predictor forward
|
||||
|
||||
The 5-layer transformer forward pass launches ~60 small CUDA kernels per step. `torch.compile(mode="default", dynamic=True)` fuses these into fewer kernels via Inductor:
|
||||
|
||||
```python
|
||||
self._compiled_model_fwd = torch.compile(
|
||||
self.model.forward,
|
||||
mode="default", # no Inductor CUDA graphs, avoids conflict with vLLM's CUDAGraphWrapper
|
||||
dynamic=True, # sequence length grows each step (2, 3, ..., num_groups+1)
|
||||
)
|
||||
```
|
||||
|
||||
`mode="default"` is used instead of `mode="reduce-overhead"` to avoid conflicts with vLLM's own CUDA graph capture on the main Talker model. `dynamic=True` handles the growing sequence length without recompilation.
|
||||
|
||||
These optimizations are always-on in the current codebase - all Qwen3-TTS benchmark results in this post include them.
|
||||
|
||||
---
|
||||
|
||||
## TTS-Specific: Dynamic Initial Chunk for Faster First Audio
|
||||
|
||||
In the async chunk pipeline, the standard `codec_chunk_frames` is 25 (each chunk = ~2 seconds of audio at 12 Hz). Waiting for 25 frames before forwarding the first chunk to Code2Wav adds unnecessary TTFP. The **initial codec chunk** optimization sends a smaller first chunk so Code2Wav can start decoding earlier.
|
||||
|
||||
**Dynamic initial chunk sizing (default behavior):**
|
||||
|
||||
Rather than using a fixed initial chunk size, vLLM-Omni dynamically selects it based on current server load. The initial chunk size is chosen from power-of-2 steps [2, 4, 8, 16] based on load factor (`active_requests / max_batch_size`):
|
||||
|
||||
| Server load | Initial chunk frames | Rationale |
|
||||
| --- | --- | --- |
|
||||
| Low (e.g. 1/10 active) | **2** (~167 ms of audio) | Minimize TTFP when there's headroom |
|
||||
| Medium (e.g. 5/10 active) | **4-8** | Balance TTFP vs decode efficiency |
|
||||
| High (e.g. 10/10 active) | **16** | Larger first chunk to amortize decode cost |
|
||||
|
||||
After the initial chunk, all subsequent chunks use the standard `codec_chunk_frames` (25) size.
|
||||
|
||||
**How it works in the pipeline:**
|
||||
|
||||
1. Talker generates codec tokens auto-regressively
|
||||
2. The stage input processor checks current load and picks an initial chunk size (e.g. **2 frames** at low load)
|
||||
3. After that many frames, the first chunk is forwarded to Code2Wav
|
||||
4. Code2Wav decodes this small chunk and emits the first audio packet
|
||||
5. Subsequent chunks use the standard 25-frame size for efficient batch decoding
|
||||
|
||||
**Per-request override:** Clients can also set a fixed initial chunk size via the API:
|
||||
|
||||
```json
|
||||
{"initial_codec_chunk_frames": 2}
|
||||
```
|
||||
|
||||
This overrides the dynamic calculation for that request.
|
||||
|
||||
**Config (server-side):**
|
||||
|
||||
```yaml
|
||||
runtime:
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
codec_streaming: true
|
||||
codec_chunk_frames: 25 # standard chunk size (~2s of audio)
|
||||
codec_left_context_frames: 25
|
||||
# initial chunk is computed dynamically by default
|
||||
# set initial_codec_chunk_frames: 2 to force a fixed value
|
||||
```
|
||||
|
||||
The 64 ms TTFP result reported above for Qwen3-TTS at concurrency 1 uses the dynamic initial chunk, which picks `initial_codec_chunk_frames=2` at low load. At higher concurrency the dynamic sizing increases the initial chunk to maintain decode efficiency.
|
||||
|
||||
---
|
||||
|
||||
## Live Demo: Streaming TTS over WebSocket
|
||||
|
||||
vLLM-Omni supports real-time streaming audio output for Qwen3-TTS over WebSocket ([PR #1719](https://github.com/vllm-project/vllm-omni/pull/1719)). With `stream_audio: true`, the server sends chunked PCM audio frames as they are generated, so clients can start playback before full sentence synthesis completes.
|
||||
|
||||
The WebSocket protocol uses `audio.start` / binary PCM chunks / `audio.done` framing per sentence:
|
||||
|
||||
```json
|
||||
// Client sends:
|
||||
{"type":"session.config","voice":"Vivian","response_format":"pcm","stream_audio":true}
|
||||
{"type":"input.text","text":"Hello world. This is a streaming demo."}
|
||||
{"type":"input.done"}
|
||||
|
||||
// Server streams back per sentence:
|
||||
{"type":"audio.start","sentence_index":0,"sentence_text":"Hello world.","format":"pcm","sample_rate":24000}
|
||||
<binary PCM chunk 1>
|
||||
<binary PCM chunk 2>
|
||||
...
|
||||
{"type":"audio.done","sentence_index":0,"total_bytes":96000,"error":false}
|
||||
{"type":"audio.start","sentence_index":1,"sentence_text":"This is a streaming demo.","format":"pcm","sample_rate":24000}
|
||||
<binary PCM chunk 1>
|
||||
...
|
||||
{"type":"audio.done","sentence_index":1,"total_bytes":72000,"error":false}
|
||||
{"type":"session.done","total_sentences":2}
|
||||
```
|
||||
|
||||
<iframe width="560" height="315" src="https://www.youtube.com/embed/O3IVniwwKNA" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>
|
||||
|
||||
---
|
||||
|
||||
## Deployment Playbook
|
||||
|
||||
### Qwen3-Omni
|
||||
|
||||
#### 1) Serve with the default 3-stage config
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
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.
|
||||
|
||||
#### 2) Enable async chunk
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
|
||||
```
|
||||
|
||||
#### 3) Key config knobs
|
||||
|
||||
```yaml
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0 # thinker
|
||||
runtime:
|
||||
max_batch_size: 64
|
||||
engine_args:
|
||||
enforce_eager: false
|
||||
max_num_batched_tokens: 32768
|
||||
custom_process_next_stage_input_func: >-
|
||||
vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk
|
||||
|
||||
- stage_id: 1 # talker
|
||||
runtime:
|
||||
max_batch_size: 64
|
||||
engine_args:
|
||||
enforce_eager: false
|
||||
max_num_batched_tokens: 32768
|
||||
custom_process_next_stage_input_func: >-
|
||||
vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk
|
||||
|
||||
- stage_id: 2 # code2wav
|
||||
runtime:
|
||||
max_batch_size: 64
|
||||
engine_args:
|
||||
enforce_eager: true
|
||||
max_num_batched_tokens: 51200
|
||||
```
|
||||
|
||||
#### Reproduce Qwen3-Omni benchmarks
|
||||
|
||||
```bash
|
||||
vllm bench serve \
|
||||
--dataset-name random \
|
||||
--port ${PORT} \
|
||||
--model ${MODEL_PATH} \
|
||||
--endpoint /v1/chat/completions \
|
||||
--backend openai-chat-omni \
|
||||
--max-concurrency ${MAX_CONCURRENCY} \
|
||||
--num-prompts ${NUM_PROMPTS} \
|
||||
--random-input-len 2500 \
|
||||
--ignore-eos \
|
||||
--percentile-metrics ttft,tpot,itl,e2el,audio_ttfp,audio_rtf \
|
||||
--random-output-len 900 \
|
||||
--extra_body '{"modalities": ["text","audio"]}'
|
||||
```
|
||||
|
||||
### Qwen3-TTS
|
||||
|
||||
#### 1) Serve with async chunk (recommended)
|
||||
|
||||
```bash
|
||||
vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--omni \
|
||||
--port 8000
|
||||
```
|
||||
|
||||
The default config (`qwen3_tts.yaml`) enables the full optimization stack:
|
||||
|
||||
- Batching with `max_batch_size: 10` on the Talker stage
|
||||
- CUDA Graph on the Talker (`enforce_eager: false`)
|
||||
- Async chunk with streaming transport
|
||||
|
||||
#### 2) Serve without async chunk (for comparison)
|
||||
|
||||
```bash
|
||||
vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--omni \
|
||||
--port 8000 \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml
|
||||
```
|
||||
|
||||
#### 3) Key config knobs
|
||||
|
||||
```yaml
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0 # Talker (AR decoder)
|
||||
runtime:
|
||||
max_batch_size: 10
|
||||
engine_args:
|
||||
enforce_eager: false
|
||||
max_num_batched_tokens: 512
|
||||
custom_process_next_stage_input_func: >-
|
||||
vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
|
||||
|
||||
- stage_id: 1 # Code2Wav (vocoder)
|
||||
runtime:
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
enforce_eager: true
|
||||
max_num_batched_tokens: 8192
|
||||
|
||||
runtime:
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
codec_streaming: true
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 25
|
||||
```
|
||||
|
||||
#### Reproduce Qwen3-TTS benchmarks
|
||||
|
||||
```bash
|
||||
GPU_DEVICE=0 \
|
||||
MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
NUM_PROMPTS=50 \
|
||||
CONCURRENCY="1 4 10" \
|
||||
bash benchmarks/qwen3-tts/vllm_omni/run_stacked_benchmark.sh
|
||||
```
|
||||
|
||||
This cycles through four configs (Baseline -> + Batch -> + CUDA Graph -> + Async Chunk + Streaming), benchmarks each at the specified concurrency levels, and generates all comparison figures automatically.
|
||||
@@ -1,39 +1,243 @@
|
||||
# Sleep Mode
|
||||
# Sleep Mode & ACK Protocol
|
||||
|
||||
vLLM-Omni’s **Sleep Mode** allows you to temporarily release most GPU memory used by a model—such as model weights and key-value (KV) caches (for autoregressive models)—**without stopping the server or unloading the Docker container**.
|
||||
vLLM-Omni’s **Sleep Mode** allows you to temporarily release most GPU memory used by a model—such as model weights and key-value (KV) caches—**without stopping the server or unloading the Docker container**.
|
||||
|
||||
This feature is inherited from [vLLM’s Sleep Mode](https://blog.vllm.ai/2025/10/26/sleep-mode.html), which provides zero-reload model switching for multi-model serving.
|
||||
|
||||
It is especially useful in **RLHF**, **training**, or **cost-saving scenarios**, where GPU resources must be freed between inference workloads.
|
||||
This feature is inherited from [vLLM’s Sleep Mode](https://blog.vllm.ai/2025/10/26/sleep-mode.html) and extended with the **Omni ACK Protocol** to support multi-stage pipelines and heterogeneous hardware backends (NVIDIA, AMD, Intel, Huawei). It is especially useful in **RLHF**, **dynamic model switching**, or **cost-saving scenarios**.
|
||||
|
||||
---
|
||||
|
||||
## Omni Model
|
||||
## 1. Feature Documentation
|
||||
|
||||
Omni model inherit the feature from vLLM' Sleep Mode
|
||||
### Overview
|
||||
Omni Sleep Mode provides a mechanism to "sleep" specific model stages. When a stage enters sleep, its physical VRAM is reclaimed by the system, while the process state is preserved for rapid "wake-up" without full re-initialization.
|
||||
|
||||
This means:
|
||||
### Sleep Levels
|
||||
We support two levels of hibernation to balance recovery speed and memory efficiency:
|
||||
|
||||
- Support both Level 1 and Level 2 sleep, allow to release and reset both model weights and KV Cache
|
||||
| Level | Name | Mechanism | Recovery Speed | Memory Freed |
|
||||
| :--- | :--- | :--- | :--- | :--- |
|
||||
| **Level 1** | **Weight Offloading** | Offloads weights to Host CPU RAM. | **Fast** (DMA) | Substantial |
|
||||
| **Level 2** | **Full De-mapping** | Physically releases memory pages via VRAM scavenging. | **Moderate** | **Maximum** (up to 95%+) |
|
||||
|
||||
## Diffusion Model Extension
|
||||
### Supported Platforms
|
||||
|
||||
We added Sleep Mode support for **diffusion models**, which previously lacked this functionality.
|
||||
In diffusion pipelines, this currently only offloads **model weight memory**, as these models typically do not use KV caches.
|
||||
Omni Sleep Mode is optimized for high-performance computing backends:
|
||||
|
||||
This means:
|
||||
* **NVIDIA**: Supported via Virtual Memory Management (VMM).
|
||||
* **AMD (ROCm)**: Fully supported with physical page de-mapping.
|
||||
* **Intel XPU**: Supported via Level Zero memory management.
|
||||
* **Huawei NPU**: Supported via Ascend memory scavenging.
|
||||
|
||||
- Diffusion models can now enter Level 1 sleep.
|
||||
- Pipeline states (e.g., noise schedulers, buffers) remain intact after waking.
|
||||
- Useful for releasing VRAM between image generation or training cycles.
|
||||
### Hardware Requirements
|
||||
* **Memory Considerations**: System RAM must be sufficient to hold offloaded weights during sleep.
|
||||
* **TP Support**: Tensor Parallel groups synchronize sleep/wake transitions across all workers.
|
||||
|
||||
---
|
||||
|
||||
## Enable sleep mode
|
||||
To enable sleep mode, set the `enable_sleep_mode` in `engine_args` to `True`
|
||||
|
||||
## 2. Usage Examples
|
||||
|
||||
### Python API Example
|
||||
You can programmatically control the lifecycle of stages using the `AsyncOmni` engine.
|
||||
|
||||
Example:
|
||||
```python
|
||||
omni = Omni(model=...,enable_sleep_mode=True)
|
||||
|
||||
import asyncio
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
|
||||
async def run_sleep_demo():
|
||||
# 1. initialization
|
||||
engine = AsyncOmni(
|
||||
model="ByteDance-Seed/BAGEL-7B-MoT",
|
||||
enable_sleep_mode=True
|
||||
)
|
||||
|
||||
# 2. sleep mode level2
|
||||
acks = await engine.sleep(stage_ids=[0], level=2)
|
||||
print(f"Freed {acks[0].freed_bytes / 1024**3:.2f} GiB on Stage 0")
|
||||
|
||||
# 3. wake up
|
||||
await engine.wake_up(stage_ids=[0])
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(run_sleep_demo())
|
||||
|
||||
```
|
||||
|
||||
### server command Example
|
||||
Start the server with sleep mode enabled:
|
||||
|
||||
The first method
|
||||
|
||||
```
|
||||
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--omni \
|
||||
--enable-sleep-mode \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.7
|
||||
|
||||
```
|
||||
|
||||
The second method
|
||||
|
||||
```
|
||||
|
||||
python3 -m vllm_omni.entrypoints.openai.api_server \
|
||||
--model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--omni \
|
||||
--enable-sleep-mode \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.7
|
||||
|
||||
```
|
||||
|
||||
|
||||
|
||||
|
||||
### Test Scenarios & Commands
|
||||
|
||||
#### Scenario 1: LLM Engine Sleep
|
||||
|
||||
Objective: Verify VRAM reclamation for Stage 0 (Thinker).
|
||||
|
||||
Trigger sleep (Level 1 or Level 2) via client:
|
||||
|
||||
```
|
||||
|
||||
curl -X POST http://localhost:8000/v1/omni/sleep \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"stage_ids": [0], "level": 2}'
|
||||
|
||||
```
|
||||
|
||||
Tip: Open a new terminal and run rocm-smi or nvidia-smi or to observe the immediate drop in VRAM usage.
|
||||
|
||||
|
||||
|
||||
#### Scenario 2: Diffusion Sleep
|
||||
Objective: Verify VRAM reclamation for Stage 1 (Diffusion).
|
||||
|
||||
Trigger sleep (Level 1 or Level 2) via client:
|
||||
|
||||
```
|
||||
|
||||
curl -X POST http://localhost:8000/v1/omni/sleep \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"stage_ids": [1], "level": 2}'
|
||||
|
||||
```
|
||||
|
||||
|
||||
|
||||
#### Scenario 3: Multi-Stage Coordinated Stress Test
|
||||
Objective: Test concurrent sleep and rapid wake-up across multiple stages.
|
||||
|
||||
Concurrent Sleep (Stage 0 & 1):
|
||||
|
||||
```
|
||||
|
||||
curl -X POST http://localhost:8000/v1/omni/sleep \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"stage_ids": [0, 1], "level": 2}'
|
||||
|
||||
```
|
||||
|
||||
|
||||
Rapid Wake-up:
|
||||
|
||||
```
|
||||
|
||||
curl -X POST http://localhost:8000/v1/omni/wakeup \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"stage_ids": [0, 1]}'
|
||||
|
||||
```
|
||||
|
||||
|
||||
#### Scenario 4: Full Lifecycle Memory Audit & Functional Integrity
|
||||
Objective: Audit the complete flow from Sleep to Wake-up followed by an Inference validation.
|
||||
|
||||
Check Initial State: Observe baseline VRAM usage.
|
||||
|
||||
Trigger Deep Sleep (Level 2):
|
||||
|
||||
```
|
||||
|
||||
curl -X POST http://localhost:8000/v1/omni/sleep \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"stage_ids": [0], "level": 2}'
|
||||
|
||||
```
|
||||
|
||||
Wake-up Model:
|
||||
|
||||
```
|
||||
|
||||
curl -X POST http://localhost:8000/v1/omni/wakeup \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"stage_ids": [0]}'
|
||||
|
||||
```
|
||||
|
||||
Verify Functional Integrity (Inference):
|
||||
Ensure the model still generates valid output after reloading weights.
|
||||
|
||||
```
|
||||
|
||||
curl -X POST http://localhost:8000/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"prompt": "A huge swimming pool, with many people swimming.",
|
||||
"model": "ByteDance-Seed/BAGEL-7B-MoT",
|
||||
"response_format": "b64_json",
|
||||
"extra_body": {"sampling_params": {"num_inference_steps": 4, "seed": 42}}
|
||||
}' > post.json
|
||||
|
||||
```
|
||||
|
||||
|
||||
|
||||
|
||||
## 3. API Reference
|
||||
|
||||
|
||||
### Methods
|
||||
|
||||
| Method | Arguments | Return Type | Description |
|
||||
| :--- | :--- | :--- | :--- |
|
||||
| **sleep** | `stage_ids: List[int], level: int` | `List[OmniACK]` | Triggers hibernation for specified stages. |
|
||||
| **wake_up** | `stage_ids: List[int]` | `List[OmniACK]` | Reloads weights and re-maps memory. |
|
||||
|
||||
|
||||
|
||||
### OmniACK Dataclass Fields
|
||||
|
||||
| Field | Type | Description |
|
||||
| :--- | :--- | :--- |
|
||||
| **task_id** | `str` | Unique identifier for the operation. |
|
||||
| **status** | `str` | `SUCCESS` or `ERROR`. |
|
||||
| **stage_id** | `int` | The ID of the stage that responded. |
|
||||
| **rank** | `int` | The rank ID within the Tensor Parallel group. |
|
||||
| **freed_bytes** | `int` | Actual amount of physical VRAM reclaimed. |
|
||||
| **metadata** | `dict` | Additional platform-specific metrics. |
|
||||
|
||||
Metadata Field Analysis
|
||||
The metadata field is a dynamic dictionary containing hardware-specific telemetry and audit data, primarily used for verifying memory reclamation on various backends (e.g., AMD ROCm, NVIDIA CUDA).
|
||||
|
||||
```
|
||||
"metadata": {
|
||||
"source": "Platform_AMD_Instinct_MI300X",
|
||||
"total_freed_gib": "78.57",
|
||||
"rank_residual_gib": "2.07"
|
||||
}
|
||||
```
|
||||
|
||||
#### Core Utility:
|
||||
**VRAM Reclamation Audit (total_freed_gib)**: Converts raw freed_bytes into human-readable GiB. It serves as the primary metric to verify that Level 2 sleep has successfully purged model weights from VRAM.
|
||||
|
||||
**Residual & Fragmentation Monitoring (rank_residual_gib)**: Reports the remaining VRAM footprint after memory de-mapping. A low residual value (e.g., 2.07 GiB) confirms a successful "clean" state, ensuring the device is ready for high-memory co-located tasks like training or diffusion pipelines.
|
||||
|
||||
**Backend Traceability (source)**: Identifies the underlying hardware driver or audit source. This is critical for debugging synchronization issues in multi-stage, distributed environments.
|
||||
|
||||
**Performance Analytics (Roadmap)**: Future updates will include latency_ms (context-switch overhead) and cuda_graph_recalled (graph engine status) to optimize performance in high-frequency sleep/wake scenarios.
|
||||
@@ -26,7 +26,7 @@ uv pip install vllm-omni
|
||||
|
||||
# Optional if want to run Qwen3 TTS
|
||||
uv pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
|
||||
uv pip install onnxruntime-rocm sox
|
||||
uv pip install onnxruntime-rocm
|
||||
```
|
||||
|
||||
# --8<-- [end:pre-built-wheels]
|
||||
|
||||
@@ -47,6 +47,7 @@ th {
|
||||
| `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` | ✅︎ | ✅︎ | | |
|
||||
| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -1,255 +0,0 @@
|
||||
# VLLM-Omni PR Reviewer
|
||||
|
||||
## Overview
|
||||
|
||||
The VLLM-Omni PR Reviewer is an automated code review bot powered by GLM-4.7 AI model. It helps maintain code quality by providing intelligent feedback on pull requests.
|
||||
|
||||
## Features
|
||||
|
||||
- **Intelligent Code Analysis**: Leverages GLM-4.7 for understanding code context and providing meaningful feedback
|
||||
- **Comprehensive Reviews**: Covers code quality, architecture, security, testing, and documentation
|
||||
- **Structured Output**: Provides well-formatted reviews with clear sections and actionable suggestions
|
||||
- **Rate Limiting**: Built-in cooldown mechanism to prevent excessive API usage
|
||||
- **Retry Logic**: Automatic retries with exponential backoff for transient API failures
|
||||
- **Defensive Parsing**: Robust validation of API responses to handle malformed data
|
||||
- **Cost Control**: Only repository members/collaborators/owners can trigger reviews
|
||||
|
||||
## How to Use
|
||||
|
||||
### Triggering a Review
|
||||
|
||||
To trigger an automated PR review, mention the bot in a PR comment:
|
||||
|
||||
```
|
||||
@vllm-omni-reviewer please review
|
||||
```
|
||||
|
||||
Or include in your PR description:
|
||||
|
||||
```
|
||||
@vllm-omni-reviewer
|
||||
```
|
||||
|
||||
The bot will automatically review your changes and post a detailed comment.
|
||||
|
||||
## What Gets Reviewed
|
||||
|
||||
- **vLLM Architecture Compatibility**: Ensures changes align with vLLM's design patterns
|
||||
- **Multi-modal Integration**: Reviews audio, vision, and text processing implementations
|
||||
- **Performance Implications**: Analyzes impact on inference latency and throughput
|
||||
- **Code Quality**: Checks Python best practices, type hints, and documentation
|
||||
- **Security Considerations**: Identifies potential security vulnerabilities
|
||||
- **Testing Coverage**: Recommends additional test cases when needed
|
||||
|
||||
## Review Output
|
||||
|
||||
The bot posts a structured review comment with:
|
||||
|
||||
- **Overview**: Brief summary of the PR's purpose
|
||||
- **Critical Issues (Must Fix)**: Blocking issues that need to be addressed
|
||||
- **Important Issues (Should Fix)**: Significant concerns that should be resolved
|
||||
- **Minor Issues & Suggestions**: Small improvements and optional suggestions
|
||||
- **Positive Aspects**: Highlights well-implemented features
|
||||
- **Performance Considerations**: Analysis of performance impact
|
||||
- **Testing Recommendations**: Suggestions for additional tests
|
||||
- **Overall Assessment**: Final recommendation (Approve/Request Changes/Needs Major Work)
|
||||
|
||||
## Rate Limiting and Cooldown
|
||||
|
||||
The bot includes a cooldown mechanism to prevent excessive API usage:
|
||||
|
||||
- **Default cooldown**: 5 minutes between reviews per PR
|
||||
- **Configurable**: Can be adjusted via `PR_REVIEWER_COOLDOWN_MINUTES` environment variable
|
||||
- **Smart detection**: Checks for previous bot comments before starting a review
|
||||
|
||||
If you trigger a review within the cooldown period, the bot will log a message and skip the review.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌─────────────────┐
|
||||
│ PR Comment │
|
||||
│ @vllm-omni- │
|
||||
│ reviewer │
|
||||
└────────┬────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────┐
|
||||
│ GitHub Actions Workflow │
|
||||
│ (.github/workflows/ │
|
||||
│ pr-reviewer.yml) │
|
||||
│ │
|
||||
│ - Python 3.11 │
|
||||
│ - requests==2.31.0 │
|
||||
│ - pyyaml==6.0.1 │
|
||||
└────────┬────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────┐
|
||||
│ PR Reviewer Script │
|
||||
│ (.github/scripts/ │
|
||||
│ pr_reviewer.py) │
|
||||
│ │
|
||||
│ 1. Check cooldown │
|
||||
│ 2. Fetch PR details & diff │
|
||||
│ 3. Build review prompt │
|
||||
│ 4. Call GLM-4.7 API │
|
||||
│ (with retry logic) │
|
||||
│ 5. Validate response │
|
||||
│ 6. Post review comment │
|
||||
└────────┬────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────┐
|
||||
│ GLM-4.7 API │
|
||||
│ (open.bigmodel.cn) │
|
||||
└─────────────────────────────────┘
|
||||
```
|
||||
|
||||
1. **GitHub Actions Workflow** (`.github/workflows/pr-reviewer.yml`): Triggers on @mention
|
||||
2. **Python Script** (`.github/scripts/pr_reviewer.py`): Fetches PR data and calls GLM-4.7 API
|
||||
3. **GLM-4.7 API**: Provides intelligent code analysis
|
||||
|
||||
## Testing
|
||||
|
||||
### Testing the PR Reviewer Bot
|
||||
|
||||
To test the PR reviewer bot before deploying to production:
|
||||
|
||||
1. **Create a test PR** - Make a small, safe change (e.g., documentation update)
|
||||
2. **Open the PR** - Create a pull request with a descriptive title
|
||||
3. **Trigger the review** - Comment `@vllm-omni-reviewer` on the PR
|
||||
4. **Monitor results** - Check the Actions tab for workflow execution logs
|
||||
|
||||
### Running Unit Tests
|
||||
|
||||
The bot includes comprehensive unit tests that can be run locally:
|
||||
|
||||
```bash
|
||||
# Run all tests
|
||||
pytest .github/tests/test_pr_reviewer.py -v
|
||||
|
||||
# Run specific test
|
||||
pytest .github/tests/test_pr_reviewer.py::TestCheckTrigger -v
|
||||
|
||||
# Run with coverage
|
||||
pytest .github/tests/test_pr_reviewer.py --cov=.github/scripts/pr_reviewer.py --cov-report=term-missing
|
||||
```
|
||||
|
||||
### What to Look For
|
||||
|
||||
When testing, verify that:
|
||||
- [ ] The workflow triggers on the `@vllm-omni-reviewer` comment
|
||||
- [ ] The cooldown mechanism works correctly
|
||||
- [ ] The GLM API call completes without errors (with retry if needed)
|
||||
- [ ] A review comment is posted to the PR
|
||||
- [ ] The review content is meaningful and well-structured
|
||||
- [ ] The cost is within the expected range (0.50-5 CNY)
|
||||
|
||||
### Safe Test Changes
|
||||
|
||||
For testing, consider making these types of safe changes:
|
||||
- Documentation updates (like adding this Testing section)
|
||||
- Comment improvements
|
||||
- README enhancements
|
||||
- Non-functional file additions
|
||||
|
||||
### Example Test PR
|
||||
|
||||
A good test PR might:
|
||||
- Update a documentation file
|
||||
- Add explanatory comments
|
||||
- Improve code formatting
|
||||
- Fix a minor typo
|
||||
|
||||
These changes are safe to merge if the test is successful and won't affect functionality.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Bot Doesn't Respond
|
||||
|
||||
1. **Check permissions** - Verify you have Owner/Member/Collaborator access
|
||||
2. **Check Actions tab** - Look for workflow execution and view logs
|
||||
3. **Check cooldown** - If another review was posted recently, wait for the cooldown period
|
||||
4. **Check API key** - Ensure `GLM_API_KEY` is configured in repository secrets
|
||||
|
||||
### API Errors
|
||||
|
||||
If the GLM API call fails:
|
||||
- Check the Actions tab for detailed error logs
|
||||
- Verify the `GLM_API_KEY` secret is correctly configured
|
||||
- Ensure sufficient API quota is available
|
||||
- The bot will automatically retry up to 3 times with exponential backoff
|
||||
|
||||
### Review Seems Truncated
|
||||
|
||||
If the review appears incomplete:
|
||||
- Large diffs may be truncated at 100,000 characters
|
||||
- Check the logs for truncation warnings
|
||||
- Consider breaking large PRs into smaller chunks
|
||||
|
||||
## Configuration
|
||||
|
||||
### Required Secrets
|
||||
|
||||
The following secret must be configured in the repository settings:
|
||||
|
||||
- `GLM_API_KEY` - Your GLM (BigModel) API key for accessing the GLM-4.7 API
|
||||
|
||||
To add the secret:
|
||||
1. Go to repository Settings → Secrets and variables → Actions
|
||||
2. Click "New repository secret"
|
||||
3. Name: `GLM_API_KEY`
|
||||
4. Value: Your GLM API key
|
||||
|
||||
### Optional Configuration
|
||||
|
||||
The following optional environment variables can be set in the workflow file:
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `GLM_API_URL` | `https://open.bigmodel.cn/api/paas/v4/chat/completions` | GLM API endpoint |
|
||||
| `GLM_MODEL` | `glm-4.7` | Model to use for reviews |
|
||||
| `PR_REVIEWER_COOLDOWN_MINUTES` | `5` | Cooldown period between reviews |
|
||||
| `PR_REVIEWER_MAX_RETRIES` | `3` | Maximum API retry attempts |
|
||||
| `PR_REVIEWER_RETRY_DELAY` | `1.0` | Base delay for retry backoff (seconds) |
|
||||
| `PR_REVIEWER_MAX_DIFF_SIZE` | `100000` | Maximum diff size before truncation |
|
||||
|
||||
### Workflow Customization
|
||||
|
||||
The workflow can be customized in `.github/workflows/pr-reviewer.yml`:
|
||||
- Change Python version (default: 3.11)
|
||||
- Adjust timeout value (default: 10 minutes)
|
||||
- Modify trigger conditions
|
||||
- Add additional dependencies
|
||||
|
||||
## Code Quality
|
||||
|
||||
The PR reviewer script follows vllm-omni coding standards:
|
||||
|
||||
- **Type hints**: All functions have complete type hints following mypy strict mode
|
||||
- **Logging**: Uses Python's logging module for structured logging
|
||||
- **Testing**: Comprehensive unit tests with pytest
|
||||
- **Pre-commit**: Script is checked by pre-commit hooks (flake8)
|
||||
|
||||
## Cost Estimate
|
||||
|
||||
| Component | Cost |
|
||||
|-----------|------|
|
||||
| GitHub Actions (public repo) | Free |
|
||||
| GLM API (glm-4.7) | ~0.50-5 CNY per PR (varies by size) |
|
||||
| Total (20 PRs/month) | ~10-100 CNY/month (~$2-15 USD) |
|
||||
|
||||
## Contributing
|
||||
|
||||
To improve the PR reviewer bot:
|
||||
|
||||
1. Edit `.github/scripts/pr_reviewer.py` for logic changes
|
||||
2. Edit `.github/workflows/pr-reviewer.yml` for workflow changes
|
||||
3. Add tests to `.github/tests/test_pr_reviewer.py`
|
||||
4. Run `pre-commit run --files .github/scripts/pr_reviewer.py` to check code quality
|
||||
5. Test thoroughly with a test PR before deploying to production
|
||||
|
||||
## License
|
||||
|
||||
This bot is part of the VLLM-Omni project and follows the same license terms.
|
||||
@@ -104,6 +104,8 @@ Content-Type: multipart/form-data
|
||||
| `guidance_scale` | float | model defaults | Classifier-free guidance scale (typically 0.0-20.0) |
|
||||
| `true_cfg_scale` | float | model defaults | True CFG scale (model-specific parameter, may be ignored if not supported) |
|
||||
| `seed` | integer | null | Random seed for reproducibility |
|
||||
| `reference_image` | string or array | null | Reference image for inpainting |
|
||||
| `mask_image` | string or array | null | Mask for inpainting (white areas will be inpainted) |
|
||||
|
||||
### Response Format
|
||||
|
||||
|
||||
@@ -15,28 +15,17 @@ Each server instance runs a single model (specified at startup via `vllm serve <
|
||||
```bash
|
||||
# Qwen3-TTS: CustomVoice model (predefined speakers)
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
|
||||
--deploy-config vllm_omni/deploy/qwen3_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
|
||||
# Fish Speech S2 Pro
|
||||
vllm-omni serve fishaudio/s2-pro \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.9
|
||||
vllm serve fishaudio/s2-pro --omni --port 8091
|
||||
|
||||
# Voxtral TTS
|
||||
vllm serve mistralai/Voxtral-4B-TTS-2603 \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
|
||||
```
|
||||
|
||||
### Generate Speech
|
||||
@@ -300,7 +289,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
```bash
|
||||
# Start server with VoiceDesign model first
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
|
||||
--deploy-config vllm_omni/deploy/qwen3_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
@@ -322,7 +311,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
```bash
|
||||
# Start server with Base model first
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
|
||||
--deploy-config vllm_omni/deploy/qwen3_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
@@ -517,15 +506,16 @@ for result in response.json()["results"]:
|
||||
|
||||
All items are fanned out to `generate()` concurrently. The engine's stage worker automatically batches them up to the configured `max_batch_size` and queues the rest — no client-side throttling needed.
|
||||
|
||||
For best throughput, use a batch-optimized stage config with `max_batch_size > 1`:
|
||||
For best throughput, set both stages' `max_num_seqs` to ≥4 via `--stage-overrides`:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml \
|
||||
--omni --port 8091 --trust-remote-code --enforce-eager
|
||||
--omni --port 8091 --trust-remote-code --enforce-eager \
|
||||
--stage-overrides '{"0":{"max_num_seqs":4,"gpu_memory_utilization":0.2},
|
||||
"1":{"max_num_seqs":4,"gpu_memory_utilization":0.2}}'
|
||||
```
|
||||
|
||||
The default `qwen3_tts.yaml` uses `max_batch_size: 1` (single request). The `qwen3_tts_batch.yaml` config sets `max_batch_size: 4` for ~4x throughput.
|
||||
The bundled `qwen3_tts.yaml` uses `max_num_seqs: 1` (single request) on both stages. Bumping to 4 yields roughly 4× throughput on the talker and lets stage 1 batch chunks across in-flight requests.
|
||||
|
||||
## Supported Models
|
||||
|
||||
@@ -617,7 +607,7 @@ Enable debug logging:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
|
||||
--deploy-config vllm_omni/deploy/qwen3_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
|
||||
|
Before Width: | Height: | Size: 134 KiB After Width: | Height: | Size: 226 KiB |
|
Before Width: | Height: | Size: 127 KiB After Width: | Height: | Size: 152 KiB |
@@ -283,3 +283,10 @@ Using Cache-DiT acceleration:
|
||||
|
||||
1. ✅ **Enable Cache-DiT** - Set `cache_backend="cache_dit"` to get 1.5x-3x speedup with optimized defaults
|
||||
2. ✅ **(Optional) Customize** - Adjust `cache_config` parameters for specific speed/quality trade-offs
|
||||
|
||||
---
|
||||
|
||||
## Additional Resources
|
||||
|
||||
- [Cache-DiT documentation](https://cache-dit.readthedocs.io/en/latest/)
|
||||
- [Cache-DiT API reference](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/)
|
||||
@@ -0,0 +1,92 @@
|
||||
# Frame Interpolation
|
||||
|
||||
## Overview
|
||||
|
||||
vLLM-Omni supports post-generation frame interpolation for supported video
|
||||
diffusion pipelines. This feature inserts synthesized intermediate frames
|
||||
between adjacent generated frames to improve temporal smoothness without
|
||||
rerunning the diffusion denoising loop.
|
||||
|
||||
Frame interpolation runs in the diffusion worker post-processing path instead
|
||||
of the API server encoding path. This allows the interpolation step to reuse
|
||||
the worker's current accelerator device and keeps the FastAPI event loop free
|
||||
from heavy synchronous PyTorch work.
|
||||
|
||||
For an input video with `N` generated frames and interpolation exponent `exp`,
|
||||
the output frame count is:
|
||||
|
||||
```text
|
||||
(N - 1) * 2**exp + 1
|
||||
```
|
||||
|
||||
The output FPS is multiplied by `2**exp` so the clip duration remains close to
|
||||
the original generated video.
|
||||
|
||||
## Supported Pipelines
|
||||
|
||||
Frame interpolation is currently supported for:
|
||||
|
||||
- `WanPipeline` (Wan2.2 text-to-video)
|
||||
- `WanImageToVideoPipeline`
|
||||
- `Wan22TI2VPipeline`
|
||||
|
||||
## Request Parameters
|
||||
|
||||
The video APIs `/v1/videos` and `/v1/videos/sync` accept:
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `enable_frame_interpolation` | bool | `false` | Enable post-generation frame interpolation |
|
||||
| `frame_interpolation_exp` | int | `1` | Interpolation exponent. `1=2x`, `2=4x`, etc. |
|
||||
| `frame_interpolation_scale` | float | `1.0` | RIFE inference scale |
|
||||
| `frame_interpolation_model_path` | str | `None` | Local directory or Hugging Face repo ID containing `flownet.pkl` |
|
||||
|
||||
## Execution Flow
|
||||
|
||||
For supported Wan2.2 pipelines, the execution order is:
|
||||
|
||||
1. Diffusion worker finishes denoising and decodes the raw video tensor.
|
||||
2. Worker-side model-specific post-processing runs.
|
||||
3. If frame interpolation is enabled, RIFE interpolates the decoded video
|
||||
tensor on the worker side and records a FPS multiplier in `custom_output`.
|
||||
4. The API server receives the already-interpolated video and only performs
|
||||
MP4 export.
|
||||
|
||||
This design keeps interpolation close to the generated tensor and avoids
|
||||
introducing another heavyweight GPU context in the API server process.
|
||||
|
||||
## Example
|
||||
|
||||
Start the server:
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091
|
||||
```
|
||||
|
||||
Run a sync request with interpolation enabled:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/videos/sync \
|
||||
-F "prompt=A dog running through a park" \
|
||||
-F "num_frames=81" \
|
||||
-F "width=832" \
|
||||
-F "height=480" \
|
||||
-F "fps=16" \
|
||||
-F "num_inference_steps=40" \
|
||||
-F "guidance_scale=1.0" \
|
||||
-F "guidance_scale_2=1.0" \
|
||||
-F "enable_frame_interpolation=true" \
|
||||
-F "frame_interpolation_exp=1" \
|
||||
-F "frame_interpolation_scale=1.0" \
|
||||
-F "seed=42" \
|
||||
-o sync_t2v_interpolated.mp4
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
- This is a post-processing feature. It does not modify the diffusion denoising
|
||||
schedule.
|
||||
- Higher interpolation exponents increase post-processing time and memory usage.
|
||||
- If the interpolation model weights are not available locally,
|
||||
`frame_interpolation_model_path` may point to a Hugging Face repo containing
|
||||
`flownet.pkl`.
|
||||
@@ -14,7 +14,7 @@ vLLM-Omni supports various advanced features for diffusion models:
|
||||
|
||||
- Acceleration: **cache methods**, **parallelism methods**, **startup optimizations**
|
||||
- Memory optimization: **cpu offloading**, **quantization**
|
||||
- Extensions: **LoRA inference**
|
||||
- Extensions: **LoRA inference**, **frame interpolation**
|
||||
- Execution modes: **step execution**
|
||||
|
||||
## Supported Features
|
||||
@@ -69,6 +69,7 @@ Extension methods add specialized capabilities to diffusion models beyond standa
|
||||
| Method | Description | Best For |
|
||||
|--------|-------------|----------|
|
||||
| **[LoRA Inference](diffusion/lora.md)** | Enables inference with Low-Rank Adaptation (LoRA) adapters weights | Reinforcement learning extensions |
|
||||
| **[Frame Interpolation](diffusion/frame_interpolation.md)** | Inserts intermediate video frames after generation for smoother motion | Video generation pipelines that need higher temporal smoothness |
|
||||
|
||||
|
||||
### Execution Modes
|
||||
@@ -108,13 +109,14 @@ The following tables show which models support each feature:
|
||||
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
@@ -138,16 +140,21 @@ The following tables show which models support each feature:
|
||||
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
|
||||
**Frame Interpolation Support**
|
||||
|
||||
- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
|
||||
- **Not supported**: Wan2.1-VACE, LTX-2, Helios, HunyuanVideo-1.5, DreamID-Omni
|
||||
|
||||
### AudioGen
|
||||
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
|
||||
|
||||
## Feature Compatibility
|
||||
@@ -258,6 +265,7 @@ Measured on NVIDIA H800:
|
||||
**Extensions:**
|
||||
|
||||
- **[LoRA Inference Guide](diffusion/lora.md)** - Low-Rank Adaptation for style customization and fine-tuning
|
||||
- **[Frame Interpolation Guide](diffusion/frame_interpolation.md)** - Worker-side post-generation video frame interpolation for smoother motion
|
||||
|
||||
**Execution Modes:**
|
||||
|
||||
|
||||
@@ -176,8 +176,6 @@ Example configuration for TP=2 on GPUs 0 and 1:
|
||||
|
||||
| Parameter | Value | Description |
|
||||
| :-------------------- | :------ | :------------------------------- |
|
||||
| `window_size` | `-1` | Window size (-1 means unlimited) |
|
||||
| `max_inflight` | `1` | Maximum inflight requests |
|
||||
| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB) |
|
||||
|
||||
## Using Mooncake Connector
|
||||
|
||||
@@ -61,10 +61,17 @@ Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
|
||||
- Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features.
|
||||
- Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform.
|
||||
|
||||
Stage wiring is configured in `vllm_omni/model_executor/stage_configs/cosyvoice3.yaml`.
|
||||
Pipeline topology lives in `vllm_omni/model_executor/models/cosyvoice3/pipeline.py`;
|
||||
runtime tunables (batch size, memory limits, sampling) live in
|
||||
`vllm_omni/deploy/cosyvoice3.yaml`. The deploy config auto-loads by
|
||||
HF `model_type` and defaults to `async_chunk: true` (shared-memory
|
||||
streaming). Pass `--no-async-chunk` on `vllm serve` to switch to the
|
||||
legacy sync path where stage 1 runs `text2flow` over the full
|
||||
speech-token sequence.
|
||||
|
||||
- Stage 0 emits latent speech tokens.
|
||||
- Stage 1 consumes them via `custom_process_input_func` and outputs audio.
|
||||
- Stage 1 consumes them via `sync_process_input_func` (sync mode) or the
|
||||
shared-memory connector (async-chunk mode) and outputs audio.
|
||||
|
||||
## Example materials
|
||||
|
||||
|
||||
@@ -38,7 +38,6 @@ Run a single sample for basic TTS:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft
|
||||
```
|
||||
@@ -47,7 +46,6 @@ Run batch samples for basic TTS:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft \
|
||||
--num-prompts {batch_size}
|
||||
@@ -65,7 +63,6 @@ Generate speech from text input:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft \
|
||||
--text "The weather is so nice today."
|
||||
@@ -77,7 +74,6 @@ Generate speech with explicit voice style instructions:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_instruct \
|
||||
--text "The weather is so nice today." \
|
||||
@@ -90,7 +86,6 @@ Generate speech using an audio reference for voice cloning:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_audio \
|
||||
--text "The weather is so nice today." \
|
||||
@@ -103,7 +98,6 @@ Generate speech from text containing natural voice descriptions:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_natural_instruction \
|
||||
--text "In a panting young male voice, he said: I can't run anymore, wait for me!"
|
||||
@@ -115,7 +109,6 @@ Transcribe audio to text:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_trancribing_sft \
|
||||
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
@@ -127,7 +120,6 @@ Understand and analyze audio content with text queries:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_understanding_sft \
|
||||
--text "Summarize the audio." \
|
||||
@@ -140,7 +132,6 @@ Audio understanding with reasoning chain:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_understanding_sft_with_thinking \
|
||||
--text "Summarize the audio." \
|
||||
@@ -153,7 +144,6 @@ Multi-turn dialogue with audio input and output:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type spoken_dialogue_sft_multiturn \
|
||||
--audio-path "./prompt_speech_zh_m.wav"
|
||||
@@ -167,7 +157,6 @@ Multi-turn dialogue converting speech to text:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type speech2text_dialogue_sft_multiturn
|
||||
```
|
||||
@@ -180,7 +169,6 @@ Multi-turn text-only dialogue:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type text_dialogue_sft_multiturn
|
||||
```
|
||||
|
||||
@@ -18,11 +18,11 @@ Please refer to the [stage configuration documentation](https://docs.vllm.ai/pro
|
||||
|
||||
### ROCm Dependencies
|
||||
|
||||
You will need to install these two dependencies `onnxruntime-rocm` and `sox`.
|
||||
You will need to install the dependency `onnxruntime-rocm`.
|
||||
|
||||
```
|
||||
pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
|
||||
pip install onnxruntime-rocm sox
|
||||
pip install onnxruntime-rocm
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
@@ -144,13 +144,13 @@ completes. This demonstrates that audio data is available progressively rather t
|
||||
|
||||
## Batched Decoding
|
||||
|
||||
The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_num_seqs > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
|
||||
The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, set `max_num_seqs > 1` on both stages via `--stage-overrides` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
|
||||
|
||||
```
|
||||
python end2end.py --query-type CustomVoice \
|
||||
--txt-prompts benchmark_prompts.txt \
|
||||
--batch-size 4 \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
|
||||
--stage-overrides '{"0":{"max_num_seqs":4,"gpu_memory_utilization":0.2},"1":{"max_num_seqs":4,"gpu_memory_utilization":0.2}}'
|
||||
```
|
||||
|
||||
**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_num_seqs >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_num_seqs`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
|
||||
|
||||
@@ -29,6 +29,22 @@ python text_to_audio.py \
|
||||
--output stable_audio_output.wav
|
||||
```
|
||||
|
||||
To reduce per-GPU memory for multi-GPU inference, launch with HSDP:
|
||||
|
||||
```bash
|
||||
python text_to_audio.py \
|
||||
--model stabilityai/stable-audio-open-1.0 \
|
||||
--prompt "The sound of a hammer hitting a wooden surface" \
|
||||
--negative-prompt "Low quality" \
|
||||
--seed 42 \
|
||||
--guidance-scale 7.0 \
|
||||
--audio-length 10.0 \
|
||||
--num-inference-steps 100 \
|
||||
--use-hsdp \
|
||||
--hsdp-shard-size 2 \
|
||||
--output stable_audio_output.wav
|
||||
```
|
||||
|
||||
Key arguments:
|
||||
|
||||
- `--prompt`: text description (string).
|
||||
@@ -37,6 +53,9 @@ Key arguments:
|
||||
- `--guidance-scale`: classifier-free guidance scale.
|
||||
- `--audio-length`: audio duration in seconds.
|
||||
- `--num-inference-steps`: diffusion sampling steps.(more steps = higher quality, slower).
|
||||
- `--use-hsdp`: enable HSDP weight sharding for the Stable Audio DiT.
|
||||
- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
|
||||
- `--hsdp-replicate-size`: number of HSDP replica groups.
|
||||
- `--output`: path to save the generated WAV file.
|
||||
|
||||
## Example materials
|
||||
|
||||
@@ -13,28 +13,24 @@ When `mistral_common` has `SpeechRequest` support, prompt token IDs are built vi
|
||||
```bash
|
||||
# Basic single-prompt with cheerful_female voice preset
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--write-audio --voice cheerful_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
|
||||
|
||||
# 32 replicate prompts with cheerful_female voice preset
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--num-prompts 32 --write-audio --voice cheerful_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
|
||||
|
||||
# Streaming with neutral_female voice preset
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--streaming --write-audio --voice neutral_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
|
||||
|
||||
# 32 prompts, 8 concurrent requests per wave, streaming with neutral_female voice
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
|
||||
@@ -42,7 +38,6 @@ python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
# Short debug prompt with reference audio
|
||||
# Note: Reference audio capability is not yet released.
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--write-audio \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "This is a test message." \
|
||||
@@ -57,7 +52,7 @@ python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
| `--text TEXT` | Text to synthesize (default: `"This is a test message."`) |
|
||||
| `--audio-path PATH` | Path to reference audio file for voice cloning |
|
||||
| `--output-dir DIR` | Directory to write output WAV files (default: `output_audio`) |
|
||||
| `--stage-configs-path PATH` | Path to stage configs YAML (currently it must be set for VoxtralTTS) |
|
||||
| `--deploy-config PATH` | Override the deploy config path. If unset, auto-loads `vllm_omni/deploy/voxtral_tts.yaml` from the HF `model_type`. |
|
||||
| `--num-prompts N` | Number of replicate prompts to run for measuring performance (default: 1) |
|
||||
| `--streaming` | Use streaming generation via `AsyncOmni` (default: blocking `Omni`) |
|
||||
| `--concurrency N` | Max concurrent requests per wave (must be used with `--streaming`, must evenly divide `--num-prompts`) |
|
||||
|
||||
@@ -31,9 +31,9 @@ dreamid_omni/
|
||||
```
|
||||
|
||||
### Run the Inference
|
||||
```
|
||||
```python
|
||||
python x_to_video_audio.py \
|
||||
--model /xx/dreamid_omni \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "Two people walking together and singing happily" \
|
||||
--image-path ./example0.png ./example1.png \
|
||||
--audio-path ./example0.wav ./example1.wav \
|
||||
@@ -43,11 +43,33 @@ python x_to_video_audio.py \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--output dreamid_omni.mp4
|
||||
--output out_dreamid_omni_twoip.mp4
|
||||
```
|
||||
In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
|
||||
The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
|
||||
|
||||
|
||||
You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
|
||||
|
||||
For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing.
|
||||
|
||||
```python
|
||||
# Example usage for oneip, ref media from the official repo DreamID-Omni
|
||||
python x_to_video_audio.py \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "<img1>: In the frame, a woman with black long hair is identified as <sub1>.\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: <sub1> is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: <sub1> tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n<sub1> maintains eye contact, swallows as if choosing her words, and says, <S>I keep telling myself I’m fine,but some nights it feels like I’m just performing calm.<E>" \
|
||||
--image-path 9.png \
|
||||
--audio-path 9.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--output out_dreamid_omni_oneip.mp4
|
||||
```
|
||||
|
||||
|
||||
Key arguments:
|
||||
- `--prompt`: text description (string).
|
||||
- `--model`: path to the model local directory.
|
||||
|
||||
@@ -41,15 +41,11 @@ Features:
|
||||
## Launch the Server
|
||||
|
||||
```bash
|
||||
vllm-omni serve fishaudio/s2-pro \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.9
|
||||
vllm serve fishaudio/s2-pro --omni --port 8091
|
||||
```
|
||||
|
||||
The deploy config is auto-loaded from `vllm_omni/deploy/fish_qwen3_omni.yaml`.
|
||||
|
||||
Or use the convenience script:
|
||||
|
||||
```bash
|
||||
|
||||
@@ -72,6 +72,9 @@ curl -X POST http://localhost:8091/v1/videos/sync \
|
||||
-F "guidance_scale_2=1.0" \
|
||||
-F "boundary_ratio=0.875" \
|
||||
-F "flow_shift=12.0" \
|
||||
-F "enable_frame_interpolation=true" \
|
||||
-F "frame_interpolation_exp=1" \
|
||||
-F "frame_interpolation_scale=1.0" \
|
||||
-F "seed=42" \
|
||||
-o sync_i2v_output.mp4
|
||||
```
|
||||
@@ -114,6 +117,9 @@ create_response=$(curl -s http://localhost:8091/v1/videos \
|
||||
-F "guidance_scale_2=1.0" \
|
||||
-F "boundary_ratio=0.875" \
|
||||
-F "flow_shift=12.0" \
|
||||
-F "enable_frame_interpolation=true" \
|
||||
-F "frame_interpolation_exp=1" \
|
||||
-F "frame_interpolation_scale=1.0" \
|
||||
-F "seed=42")
|
||||
|
||||
video_id=$(echo "$create_response" | jq -r '.id')
|
||||
@@ -172,9 +178,35 @@ curl -X POST http://localhost:8091/v1/videos \
|
||||
-F "guidance_scale_2=1.0" \
|
||||
-F "boundary_ratio=0.875" \
|
||||
-F "flow_shift=12.0" \
|
||||
-F "enable_frame_interpolation=true" \
|
||||
-F "frame_interpolation_exp=1" \
|
||||
-F "frame_interpolation_scale=1.0" \
|
||||
-F "seed=42"
|
||||
```
|
||||
|
||||
Frame interpolation is also available for supported Wan2.2 I2V requests. See
|
||||
[Frame Interpolation](../../diffusion/frame_interpolation.md) for worker-side
|
||||
execution details and feature constraints.
|
||||
|
||||
### Frame Interpolation Example
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/videos/sync \
|
||||
-F "prompt=A bear playing with yarn, smooth motion" \
|
||||
-F "input_reference=@/path/to/qwen-bear.png" \
|
||||
-F "width=832" \
|
||||
-F "height=480" \
|
||||
-F "num_frames=33" \
|
||||
-F "fps=16" \
|
||||
-F "num_inference_steps=40" \
|
||||
-F "guidance_scale=1.0" \
|
||||
-F "guidance_scale_2=1.0" \
|
||||
-F "enable_frame_interpolation=true" \
|
||||
-F "frame_interpolation_exp=1" \
|
||||
-F "frame_interpolation_scale=1.0" \
|
||||
-o sync_i2v_interpolated.mp4
|
||||
```
|
||||
|
||||
## Create Response Format
|
||||
|
||||
`POST /v1/videos` returns a job record, not inline base64 video data.
|
||||
|
||||
@@ -13,10 +13,10 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
|
||||
```bash
|
||||
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
|
||||
|
||||
vllm-omni serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \
|
||||
--served-model-name "MiMo-Audio-7B-Instruct" \
|
||||
--port 18091 --stage-configs-path ./vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--chat-template ./examples/online_serving/mimo_audio/chat_template.jinja
|
||||
vllm serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \
|
||||
--served-model-name "MiMo-Audio-7B-Instruct" \
|
||||
--port 18091 \
|
||||
--chat-template ./examples/online_serving/mimo_audio/chat_template.jinja
|
||||
```
|
||||
> ⚠️ **Important**
|
||||
> **MiMo-Audio is not compatible with the default chat template.**
|
||||
|
||||
@@ -18,12 +18,12 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
If you want to open async chunking for qwen3-omni, launch the server with command below
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
```
|
||||
|
||||
If you have custom stage configs file, launch the server with command below
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /path/to/deploy_config_file
|
||||
```
|
||||
|
||||
### Send Multi-modal Request
|
||||
@@ -187,7 +187,7 @@ The script supports the following arguments:
|
||||
- `--model`: Model name/path (default: Qwen/Qwen3-Omni-30B-A3B-Instruct)
|
||||
- `--server-port`: Port for vLLM server (default: 8091)
|
||||
- `--gradio-port`: Port for Gradio demo (default: 7861)
|
||||
- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
|
||||
- `--deploy-config`: Path to custom deploy config YAML file (optional)
|
||||
- `--server-host`: Host for vLLM server (default: 0.0.0.0)
|
||||
- `--gradio-ip`: IP for Gradio demo (default: 127.0.0.1)
|
||||
- `--share`: Share Gradio demo publicly (creates a public link)
|
||||
@@ -202,7 +202,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
|
||||
If you have custom stage configs file:
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /path/to/deploy_config_file
|
||||
```
|
||||
|
||||
**Step 2: Run the Gradio demo**
|
||||
|
||||