diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index d9a231595..00823951d 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -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: diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index f56f23b5d..2707dfefe 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -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. diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index e175385ff..0b9a3f47a 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -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" diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index bc36c337d..cbb44ac69 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -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 diff --git a/.buildkite/test-nightly-diffusion.yml b/.buildkite/test-nightly-diffusion.yml deleted file mode 100644 index a520ca435..000000000 --- a/.buildkite/test-nightly-diffusion.yml +++ /dev/null @@ -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 diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 31b3e1797..c33d7b4d1 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -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 diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 2f749f0ee..3ca1747fe 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -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 diff --git a/.buildkite/test-template-amd-omni.j2 b/.buildkite/test-template-amd-omni.j2 index f4c386a5f..78f47d1ae 100644 --- a/.buildkite/test-template-amd-omni.j2 +++ b/.buildkite/test-template-amd-omni.j2 @@ -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" diff --git a/.claude/skills/vllm-omni-npu-upgrade/SKILL.md b/.claude/skills/vllm-omni-npu-upgrade/SKILL.md new file mode 100644 index 000000000..1ef7ab393 --- /dev/null +++ b/.claude/skills/vllm-omni-npu-upgrade/SKILL.md @@ -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="" -- 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 .. -- vllm_omni/worker/gpu_model_runner.py + git diff .. -- 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 --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` diff --git a/.claude/skills/vllm-omni-npu-upgrade/references/gpu-to-npu-translation.md b/.claude/skills/vllm-omni-npu-upgrade/references/gpu-to-npu-translation.md new file mode 100644 index 000000000..89067d37b --- /dev/null +++ b/.claude/skills/vllm-omni-npu-upgrade/references/gpu-to-npu-translation.md @@ -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()` | diff --git a/.claude/skills/vllm-omni-npu-upgrade/references/omni-specific-blocks.md b/.claude/skills/vllm-omni-npu-upgrade/references/omni-specific-blocks.md new file mode 100644 index 000000000..8c5d32ab4 --- /dev/null +++ b/.claude/skills/vllm-omni-npu-upgrade/references/omni-specific-blocks.md @@ -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`. diff --git a/.claude/skills/vllm-omni-npu-upgrade/references/workflow-checklist.md b/.claude/skills/vllm-omni-npu-upgrade/references/workflow-checklist.md new file mode 100644 index 000000000..4f184df0e --- /dev/null +++ b/.claude/skills/vllm-omni-npu-upgrade/references/workflow-checklist.md @@ -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 .. -- 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: +- + +Changes synced from GPU: +- +``` + +### 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 diff --git a/.github/scripts/pr_reviewer.py b/.github/scripts/pr_reviewer.py deleted file mode 100755 index da629a645..000000000 --- a/.github/scripts/pr_reviewer.py +++ /dev/null @@ -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()) diff --git a/.github/workflows/pr-reviewer.yml b/.github/workflows/pr-reviewer.yml deleted file mode 100644 index 8a41c6937..000000000 --- a/.github/workflows/pr-reviewer.yml +++ /dev/null @@ -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 diff --git a/.gitignore b/.gitignore index c0ee96806..a7cd8f74e 100644 --- a/.gitignore +++ b/.gitignore @@ -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/ diff --git a/benchmarks/qwen3-tts/README.md b/benchmarks/qwen3-tts/README.md index 9c01f29aa..a1c2ebe12 100644 --- a/benchmarks/qwen3-tts/README.md +++ b/benchmarks/qwen3-tts/README.md @@ -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 diff --git a/benchmarks/qwen3-tts/run_benchmark.sh b/benchmarks/qwen3-tts/run_benchmark.sh index 283b6b844..8c3e46903 100755 --- a/benchmarks/qwen3-tts/run_benchmark.sh +++ b/benchmarks/qwen3-tts/run_benchmark.sh @@ -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 diff --git a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs16.yaml b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs16.yaml deleted file mode 100644 index 2cc5cf535..000000000 --- a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs16.yaml +++ /dev/null @@ -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 diff --git a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml deleted file mode 100644 index 5de107d49..000000000 --- a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml +++ /dev/null @@ -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 diff --git a/benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh b/benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh index 61cf7757a..0ede359ea 100755 --- a/benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh +++ b/benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh @@ -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 & diff --git a/benchmarks/voxcpm/README.md b/benchmarks/voxcpm/README.md new file mode 100644 index 000000000..17f904101 --- /dev/null +++ b/benchmarks/voxcpm/README.md @@ -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. diff --git a/benchmarks/voxcpm/vllm_omni/bench_tts_offline.py b/benchmarks/voxcpm/vllm_omni/bench_tts_offline.py new file mode 100644 index 000000000..a3bad3e69 --- /dev/null +++ b/benchmarks/voxcpm/vllm_omni/bench_tts_offline.py @@ -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 /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())) diff --git a/benchmarks/voxcpm/vllm_omni/bench_tts_serve.py b/benchmarks/voxcpm/vllm_omni/bench_tts_serve.py new file mode 100644 index 000000000..816df3279 --- /dev/null +++ b/benchmarks/voxcpm/vllm_omni/bench_tts_serve.py @@ -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())) diff --git a/benchmarks/voxcpm/vllm_omni/run_offline_matrix.py b/benchmarks/voxcpm/vllm_omni/run_offline_matrix.py new file mode 100644 index 000000000..cee46c0f8 --- /dev/null +++ b/benchmarks/voxcpm/vllm_omni/run_offline_matrix.py @@ -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 /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())) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 2a98de1b8..9cbf89d0b 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -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/* diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index 6ed5b7d27..28e10f4fb 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -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/* diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index b34478389..a54aa3b79 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -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/* diff --git a/docs/.nav.yml b/docs/.nav.yml index 86ce4a3b0..b0022965a 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -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: diff --git a/docs/api/README.md b/docs/api/README.md index f65cbb525..0147f19e1 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -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][] diff --git a/docs/assets/WeChat.jpg b/docs/assets/WeChat.jpg index 416439f7e..83252b756 100644 Binary files a/docs/assets/WeChat.jpg and b/docs/assets/WeChat.jpg differ diff --git a/docs/configuration/README.md b/docs/configuration/README.md index b5761a7f1..390176e9c 100644 --- a/docs/configuration/README.md +++ b/docs/configuration/README.md @@ -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) diff --git a/docs/configuration/pd_disaggregation.md b/docs/configuration/pd_disaggregation.md index 1cf6189e6..9196bdb02 100644 --- a/docs/configuration/pd_disaggregation.md +++ b/docs/configuration/pd_disaggregation.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 diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 95c42afcc..55b4053cc 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -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/.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_`; values are names registered under top-level `connectors:`. | +| `input_connectors` | dict \| null | optional | `null` | Keyed by `from_stage_`; 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: + : + name: # required — class registered in vllm_omni.distributed + extra: # optional — forwarded to the connector's __init__ + : + ... +``` + +| 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/.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--` 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"` diff --git a/docs/configuration/stage_configs/qwen2_5_omni.yaml b/docs/configuration/stage_configs/qwen2_5_omni.yaml deleted file mode 100644 index 690577b84..000000000 --- a/docs/configuration/stage_configs/qwen2_5_omni.yaml +++ /dev/null @@ -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 diff --git a/docs/contributing/ci/CI_5levels.md b/docs/contributing/ci/CI_5levels.md index 930603573..2452ef5d4 100644 --- a/docs/contributing/ci/CI_5levels.md +++ b/docs/contributing/ci/CI_5levels.md @@ -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
/tests/e2e/offline_inference/test_{model_name}_expansion.py
Performance:
- /tests/dfx/perf/tests/test.json
+ /tests/dfx/perf/tests/test_qwen_omni.json (Omni), test_tts.json (TTS),
+ and /tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)
Doc Test:
tests/example/online_serving/test_{model_name}.py
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` diff --git a/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md b/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md index 69d6ad828..ab4deecd6 100644 --- a/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md +++ b/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md @@ -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). diff --git a/docs/contributing/ci/test_examples/l4_performance_tests.inc.md b/docs/contributing/ci/test_examples/l4_performance_tests.inc.md index 8093e1459..f1f3073dc 100644 --- a/docs/contributing/ci/test_examples/l4_performance_tests.inc.md +++ b/docs/contributing/ci/test_examples/l4_performance_tests.inc.md @@ -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 `): ```JSON { diff --git a/docs/contributing/ci/test_guide.md b/docs/contributing/ci/test_guide.md index 425f24332..08b2e3b4e 100644 --- a/docs/contributing/ci/test_guide.md +++ b/docs/contributing/ci/test_guide.md @@ -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). diff --git a/docs/contributing/ci/tests_markers.md b/docs/contributing/ci/tests_markers.md index 7c1ba1c73..7628db284 100644 --- a/docs/contributing/ci/tests_markers.md +++ b/docs/contributing/ci/tests_markers.md @@ -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`). diff --git a/docs/contributing/ci/tests_style.md b/docs/contributing/ci/tests_style.md index 69d5b16d7..3a8cb0f12 100644 --- a/docs/contributing/ci/tests_style.md +++ b/docs/contributing/ci/tests_style.md @@ -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 diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md index a0619e338..1eaff1059 100644 --- a/docs/contributing/model/adding_omni_model.md +++ b/docs/contributing/model/adding_omni_model.md @@ -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` diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md index e48ae5049..622064173 100644 --- a/docs/contributing/model/adding_tts_model.md +++ b/docs/contributing/model/adding_tts_model.md @@ -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) diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index 418fb707a..a7b508302 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.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: diff --git a/docs/design/feature/async_chunk_design.md b/docs/design/feature/async_chunk.md similarity index 99% rename from docs/design/feature/async_chunk_design.md rename to docs/design/feature/async_chunk.md index 45314a0ae..57b4209b8 100644 --- a/docs/design/feature/async_chunk_design.md +++ b/docs/design/feature/async_chunk.md @@ -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

**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 | diff --git a/docs/design/feature/expert_parallel.md b/docs/design/feature/expert_parallel.md index 9a7c4cdba..e05eec336 100644 --- a/docs/design/feature/expert_parallel.md +++ b/docs/design/feature/expert_parallel.md @@ -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 diff --git a/docs/design/feature/prefix_caching.md b/docs/design/feature/prefix_caching.md new file mode 100644 index 000000000..ebad8b691 --- /dev/null +++ b/docs/design/feature/prefix_caching.md @@ -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: +1: +2: +3: +... +11: +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 ] +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. diff --git a/docs/design/feature/teacache.md b/docs/design/feature/teacache.md index 9fa315cee..8577cff1f 100644 --- a/docs/design/feature/teacache.md +++ b/docs/design/feature/teacache.md @@ -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): + +``` + + **Data Statistics Guide:** | Metric | Good Range | Warning Signs | diff --git a/docs/design/figures/omni/E2EL_s_vllm_omni_vs_transformers.png b/docs/design/figures/omni/E2EL_s_vllm_omni_vs_transformers.png new file mode 100644 index 000000000..15112d586 Binary files /dev/null and b/docs/design/figures/omni/E2EL_s_vllm_omni_vs_transformers.png differ diff --git a/docs/design/figures/omni/Mean_AUDIO_RTF_Baseline_vs_Batch.png b/docs/design/figures/omni/Mean_AUDIO_RTF_Baseline_vs_Batch.png new file mode 100644 index 000000000..2f0615f77 Binary files /dev/null and b/docs/design/figures/omni/Mean_AUDIO_RTF_Baseline_vs_Batch.png differ diff --git a/docs/design/figures/omni/Mean_AUDIO_RTF_Batch_CUDA_Graph_vs_Async_Chunk.png b/docs/design/figures/omni/Mean_AUDIO_RTF_Batch_CUDA_Graph_vs_Async_Chunk.png new file mode 100644 index 000000000..62d8bc79b Binary files /dev/null and b/docs/design/figures/omni/Mean_AUDIO_RTF_Batch_CUDA_Graph_vs_Async_Chunk.png differ diff --git a/docs/design/figures/omni/Mean_AUDIO_RTF_Batch_vs_Batch_CUDA_Graph.png b/docs/design/figures/omni/Mean_AUDIO_RTF_Batch_vs_Batch_CUDA_Graph.png new file mode 100644 index 000000000..5838b4531 Binary files /dev/null and b/docs/design/figures/omni/Mean_AUDIO_RTF_Batch_vs_Batch_CUDA_Graph.png differ diff --git a/docs/design/figures/omni/Mean_AUDIO_TTFP_ms_Baseline_vs_Batch.png b/docs/design/figures/omni/Mean_AUDIO_TTFP_ms_Baseline_vs_Batch.png new file mode 100644 index 000000000..24be814b7 Binary files /dev/null and b/docs/design/figures/omni/Mean_AUDIO_TTFP_ms_Baseline_vs_Batch.png differ diff --git a/docs/design/figures/omni/Mean_AUDIO_TTFP_ms_Batch_CUDA_Graph_vs_Async_Chunk.png b/docs/design/figures/omni/Mean_AUDIO_TTFP_ms_Batch_CUDA_Graph_vs_Async_Chunk.png new file mode 100644 index 000000000..c8df58ebc Binary files /dev/null and b/docs/design/figures/omni/Mean_AUDIO_TTFP_ms_Batch_CUDA_Graph_vs_Async_Chunk.png differ diff --git a/docs/design/figures/omni/Mean_AUDIO_TTFP_ms_Batch_vs_Batch_CUDA_Graph.png b/docs/design/figures/omni/Mean_AUDIO_TTFP_ms_Batch_vs_Batch_CUDA_Graph.png new file mode 100644 index 000000000..2d1a04e9c Binary files /dev/null and b/docs/design/figures/omni/Mean_AUDIO_TTFP_ms_Batch_vs_Batch_CUDA_Graph.png differ diff --git a/docs/design/figures/omni/Mean_E2EL_ms_Baseline_vs_Batch.png b/docs/design/figures/omni/Mean_E2EL_ms_Baseline_vs_Batch.png new file mode 100644 index 000000000..e598b5434 Binary files /dev/null and b/docs/design/figures/omni/Mean_E2EL_ms_Baseline_vs_Batch.png differ diff --git a/docs/design/figures/omni/Mean_E2EL_ms_Batch_CUDA_Graph_vs_Async_Chunk.png b/docs/design/figures/omni/Mean_E2EL_ms_Batch_CUDA_Graph_vs_Async_Chunk.png new file mode 100644 index 000000000..54452013e Binary files /dev/null and b/docs/design/figures/omni/Mean_E2EL_ms_Batch_CUDA_Graph_vs_Async_Chunk.png differ diff --git a/docs/design/figures/omni/Mean_E2EL_ms_Batch_vs_Batch_CUDA_Graph.png b/docs/design/figures/omni/Mean_E2EL_ms_Batch_vs_Batch_CUDA_Graph.png new file mode 100644 index 000000000..04c5ad739 Binary files /dev/null and b/docs/design/figures/omni/Mean_E2EL_ms_Batch_vs_Batch_CUDA_Graph.png differ diff --git a/docs/design/figures/omni/RTF_vllm_omni_vs_transformers.png b/docs/design/figures/omni/RTF_vllm_omni_vs_transformers.png new file mode 100644 index 000000000..d93ba0b2a Binary files /dev/null and b/docs/design/figures/omni/RTF_vllm_omni_vs_transformers.png differ diff --git a/docs/design/figures/omni/Summary_E2EL_ms_vs_features.png b/docs/design/figures/omni/Summary_E2EL_ms_vs_features.png new file mode 100644 index 000000000..04087b591 Binary files /dev/null and b/docs/design/figures/omni/Summary_E2EL_ms_vs_features.png differ diff --git a/docs/design/figures/omni/Summary_RTF_vs_features.png b/docs/design/figures/omni/Summary_RTF_vs_features.png new file mode 100644 index 000000000..c2c8ad408 Binary files /dev/null and b/docs/design/figures/omni/Summary_RTF_vs_features.png differ diff --git a/docs/design/figures/omni/Summary_TTFP_ms_vs_features.png b/docs/design/figures/omni/Summary_TTFP_ms_vs_features.png new file mode 100644 index 000000000..3dcc1c553 Binary files /dev/null and b/docs/design/figures/omni/Summary_TTFP_ms_vs_features.png differ diff --git a/docs/design/figures/omni/TTFP_s_vllm_omni_vs_transformers.png b/docs/design/figures/omni/TTFP_s_vllm_omni_vs_transformers.png new file mode 100644 index 000000000..9a5b6c9bd Binary files /dev/null and b/docs/design/figures/omni/TTFP_s_vllm_omni_vs_transformers.png differ diff --git a/docs/design/figures/tts/Mean_AUDIO_RTF_vllm_omni_vs_transformers.png b/docs/design/figures/tts/Mean_AUDIO_RTF_vllm_omni_vs_transformers.png new file mode 100644 index 000000000..68f0ef17e Binary files /dev/null and b/docs/design/figures/tts/Mean_AUDIO_RTF_vllm_omni_vs_transformers.png differ diff --git a/docs/design/figures/tts/Mean_AUDIO_TTFP_(ms)_vllm_omni_vs_transformers.png b/docs/design/figures/tts/Mean_AUDIO_TTFP_(ms)_vllm_omni_vs_transformers.png new file mode 100644 index 000000000..44be96e96 Binary files /dev/null and b/docs/design/figures/tts/Mean_AUDIO_TTFP_(ms)_vllm_omni_vs_transformers.png differ diff --git a/docs/design/figures/tts/Mean_E2EL_(ms)_vllm_omni_vs_transformers.png b/docs/design/figures/tts/Mean_E2EL_(ms)_vllm_omni_vs_transformers.png new file mode 100644 index 000000000..2e5d1482b Binary files /dev/null and b/docs/design/figures/tts/Mean_E2EL_(ms)_vllm_omni_vs_transformers.png differ diff --git a/docs/design/figures/tts/Mean_mean_e2e_ms_baseline_vs_batch.png b/docs/design/figures/tts/Mean_mean_e2e_ms_baseline_vs_batch.png new file mode 100644 index 000000000..04d8f0bac Binary files /dev/null and b/docs/design/figures/tts/Mean_mean_e2e_ms_baseline_vs_batch.png differ diff --git a/docs/design/figures/tts/Mean_mean_e2e_ms_batch_vs_cuda_graph.png b/docs/design/figures/tts/Mean_mean_e2e_ms_batch_vs_cuda_graph.png new file mode 100644 index 000000000..eb85ec0dd Binary files /dev/null and b/docs/design/figures/tts/Mean_mean_e2e_ms_batch_vs_cuda_graph.png differ diff --git a/docs/design/figures/tts/Mean_mean_e2e_ms_cuda_graph_vs_async_chunk.png b/docs/design/figures/tts/Mean_mean_e2e_ms_cuda_graph_vs_async_chunk.png new file mode 100644 index 000000000..6f0e0e252 Binary files /dev/null and b/docs/design/figures/tts/Mean_mean_e2e_ms_cuda_graph_vs_async_chunk.png differ diff --git a/docs/design/figures/tts/Mean_mean_rtf_baseline_vs_batch.png b/docs/design/figures/tts/Mean_mean_rtf_baseline_vs_batch.png new file mode 100644 index 000000000..89ea30a86 Binary files /dev/null and b/docs/design/figures/tts/Mean_mean_rtf_baseline_vs_batch.png differ diff --git a/docs/design/figures/tts/Mean_mean_rtf_batch_vs_cuda_graph.png b/docs/design/figures/tts/Mean_mean_rtf_batch_vs_cuda_graph.png new file mode 100644 index 000000000..2b207b889 Binary files /dev/null and b/docs/design/figures/tts/Mean_mean_rtf_batch_vs_cuda_graph.png differ diff --git a/docs/design/figures/tts/Mean_mean_rtf_cuda_graph_vs_async_chunk.png b/docs/design/figures/tts/Mean_mean_rtf_cuda_graph_vs_async_chunk.png new file mode 100644 index 000000000..f5f7ad72c Binary files /dev/null and b/docs/design/figures/tts/Mean_mean_rtf_cuda_graph_vs_async_chunk.png differ diff --git a/docs/design/figures/tts/Mean_mean_ttfp_ms_baseline_vs_batch.png b/docs/design/figures/tts/Mean_mean_ttfp_ms_baseline_vs_batch.png new file mode 100644 index 000000000..6f8c1da4a Binary files /dev/null and b/docs/design/figures/tts/Mean_mean_ttfp_ms_baseline_vs_batch.png differ diff --git a/docs/design/figures/tts/Mean_mean_ttfp_ms_batch_vs_cuda_graph.png b/docs/design/figures/tts/Mean_mean_ttfp_ms_batch_vs_cuda_graph.png new file mode 100644 index 000000000..b0fe1d02a Binary files /dev/null and b/docs/design/figures/tts/Mean_mean_ttfp_ms_batch_vs_cuda_graph.png differ diff --git a/docs/design/figures/tts/Mean_mean_ttfp_ms_cuda_graph_vs_async_chunk.png b/docs/design/figures/tts/Mean_mean_ttfp_ms_cuda_graph_vs_async_chunk.png new file mode 100644 index 000000000..008ba9bf7 Binary files /dev/null and b/docs/design/figures/tts/Mean_mean_ttfp_ms_cuda_graph_vs_async_chunk.png differ diff --git a/docs/design/figures/tts/Summary_mean_e2e_ms_vs_features.png b/docs/design/figures/tts/Summary_mean_e2e_ms_vs_features.png new file mode 100644 index 000000000..7c65aa117 Binary files /dev/null and b/docs/design/figures/tts/Summary_mean_e2e_ms_vs_features.png differ diff --git a/docs/design/figures/tts/Summary_mean_rtf_vs_features.png b/docs/design/figures/tts/Summary_mean_rtf_vs_features.png new file mode 100644 index 000000000..71bb2c546 Binary files /dev/null and b/docs/design/figures/tts/Summary_mean_rtf_vs_features.png differ diff --git a/docs/design/figures/tts/Summary_mean_ttfp_ms_vs_features.png b/docs/design/figures/tts/Summary_mean_ttfp_ms_vs_features.png new file mode 100644 index 000000000..cef2546d6 Binary files /dev/null and b/docs/design/figures/tts/Summary_mean_ttfp_ms_vs_features.png differ diff --git a/docs/design/qwen3_omni_tts_performance_optimization.md b/docs/design/qwen3_omni_tts_performance_optimization.md new file mode 100644 index 000000000..2f18a1b1b --- /dev/null +++ b/docs/design/qwen3_omni_tts_performance_optimization.md @@ -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): + + + + + +
Qwen3-Omni E2EL: vLLM vs HFQwen3-Omni TTFP: vLLM vs HFQwen3-Omni RTF: vLLM vs HF
+ +| 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): + + + + + +
Qwen3-TTS E2EL: vLLM vs HFQwen3-TTS TTFP: vLLM vs HFQwen3-TTS RTF: vLLM vs HF
+ +| 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): + + + + + +
Qwen3-Omni E2EL: stacked optimizationQwen3-Omni TTFP: stacked optimizationQwen3-Omni RTF: stacked optimization
+ +- **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): + + + + + +
Qwen3-TTS E2EL: stacked optimizationQwen3-TTS TTFP: stacked optimizationQwen3-TTS RTF: stacked optimization
+ +- **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): + + + + + +
Qwen3-Omni E2EL: Baseline vs BatchQwen3-Omni TTFP: Baseline vs BatchQwen3-Omni RTF: Baseline vs Batch
+ +| 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): + + + + + +
Qwen3-TTS E2EL: Baseline vs BatchQwen3-TTS TTFP: Baseline vs BatchQwen3-TTS RTF: Baseline vs Batch
+ +| 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): + + + + + +
Qwen3-Omni E2EL: Batch vs CUDA GraphQwen3-Omni TTFP: Batch vs CUDA GraphQwen3-Omni RTF: Batch vs CUDA Graph
+ +| 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): + + + + + +
TTS E2EL: Batch vs +CGTTS TTFP: Batch vs +CGTTS RTF: Batch vs +CG
+ +| 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): + + + + + +
Qwen3-Omni E2EL: CG vs Async ChunkQwen3-Omni TTFP: CG vs Async ChunkQwen3-Omni RTF: CG vs Async Chunk
+ +| 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): + + + + + +
Qwen3-TTS E2EL: CG vs Async ChunkQwen3-TTS TTFP: CG vs Async ChunkQwen3-TTS RTF: CG vs Async Chunk
+ +| 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} + + +... +{"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} + +... +{"type":"audio.done","sentence_index":1,"total_bytes":72000,"error":false} +{"type":"session.done","total_sentences":2} +``` + + + +--- + +## 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. diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index 41aa48c17..b4eec162d 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -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. diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md index 1a683d174..5dfea8d2f 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -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] diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0f9c8fff6..b4976b09c 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | diff --git a/docs/pr_reviewer.md b/docs/pr_reviewer.md deleted file mode 100644 index ad3235532..000000000 --- a/docs/pr_reviewer.md +++ /dev/null @@ -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. diff --git a/docs/serving/image_edit_api.md b/docs/serving/image_edit_api.md index d254ac06a..79303e1a6 100644 --- a/docs/serving/image_edit_api.md +++ b/docs/serving/image_edit_api.md @@ -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 diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index ecbe8d9ac..8f78d6a20 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -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 \ diff --git a/docs/source/architecture/async-chunk-architecture.png b/docs/source/architecture/async-chunk-architecture.png index 249de53bf..7b3e95e4d 100644 Binary files a/docs/source/architecture/async-chunk-architecture.png and b/docs/source/architecture/async-chunk-architecture.png differ diff --git a/docs/source/architecture/vllm-omni-dataflow-between-stages.png b/docs/source/architecture/vllm-omni-dataflow-between-stages.png index cdbc9a8b7..74abc81ff 100644 Binary files a/docs/source/architecture/vllm-omni-dataflow-between-stages.png and b/docs/source/architecture/vllm-omni-dataflow-between-stages.png differ diff --git a/docs/user_guide/diffusion/cache_acceleration/cache_dit.md b/docs/user_guide/diffusion/cache_acceleration/cache_dit.md index 824e8c930..eaaca84ad 100644 --- a/docs/user_guide/diffusion/cache_acceleration/cache_dit.md +++ b/docs/user_guide/diffusion/cache_acceleration/cache_dit.md @@ -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/) diff --git a/docs/user_guide/diffusion/frame_interpolation.md b/docs/user_guide/diffusion/frame_interpolation.md new file mode 100644 index 000000000..349af50c5 --- /dev/null +++ b/docs/user_guide/diffusion/frame_interpolation.md @@ -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`. diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 31cd1500f..6e4a54a65 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -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:** diff --git a/docs/user_guide/examples/offline_inference/bagel.md b/docs/user_guide/examples/offline_inference/bagel.md index e62668687..1fb4d4045 100644 --- a/docs/user_guide/examples/offline_inference/bagel.md +++ b/docs/user_guide/examples/offline_inference/bagel.md @@ -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 diff --git a/docs/user_guide/examples/offline_inference/cosyvoice3.md b/docs/user_guide/examples/offline_inference/cosyvoice3.md index ebb7c02ef..d0638f414 100644 --- a/docs/user_guide/examples/offline_inference/cosyvoice3.md +++ b/docs/user_guide/examples/offline_inference/cosyvoice3.md @@ -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 diff --git a/docs/user_guide/examples/offline_inference/mimo_audio.md b/docs/user_guide/examples/offline_inference/mimo_audio.md index 4e8052697..1cba2f77d 100644 --- a/docs/user_guide/examples/offline_inference/mimo_audio.md +++ b/docs/user_guide/examples/offline_inference/mimo_audio.md @@ -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 ``` diff --git a/docs/user_guide/examples/offline_inference/qwen3_tts.md b/docs/user_guide/examples/offline_inference/qwen3_tts.md index 19fea4132..7226ac1fe 100644 --- a/docs/user_guide/examples/offline_inference/qwen3_tts.md +++ b/docs/user_guide/examples/offline_inference/qwen3_tts.md @@ -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. diff --git a/docs/user_guide/examples/offline_inference/text_to_audio.md b/docs/user_guide/examples/offline_inference/text_to_audio.md index 62a70e525..a31a4d7a4 100644 --- a/docs/user_guide/examples/offline_inference/text_to_audio.md +++ b/docs/user_guide/examples/offline_inference/text_to_audio.md @@ -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 diff --git a/docs/user_guide/examples/offline_inference/voxtral_tts.md b/docs/user_guide/examples/offline_inference/voxtral_tts.md index db4c2e26b..1a4f2c21b 100644 --- a/docs/user_guide/examples/offline_inference/voxtral_tts.md +++ b/docs/user_guide/examples/offline_inference/voxtral_tts.md @@ -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`) | diff --git a/docs/user_guide/examples/offline_inference/x_to_video_audio.md b/docs/user_guide/examples/offline_inference/x_to_video_audio.md index 8ea39d811..cec8d47c5 100644 --- a/docs/user_guide/examples/offline_inference/x_to_video_audio.md +++ b/docs/user_guide/examples/offline_inference/x_to_video_audio.md @@ -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 ": In the frame, a woman with black long hair is identified as .\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**: 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**: 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 maintains eye contact, swallows as if choosing her words, and says, I keep telling myself I’m fine,but some nights it feels like I’m just performing calm." \ + --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. diff --git a/docs/user_guide/examples/online_serving/fish_speech.md b/docs/user_guide/examples/online_serving/fish_speech.md index 7322d06aa..2a15ef44a 100644 --- a/docs/user_guide/examples/online_serving/fish_speech.md +++ b/docs/user_guide/examples/online_serving/fish_speech.md @@ -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 diff --git a/docs/user_guide/examples/online_serving/image_to_video.md b/docs/user_guide/examples/online_serving/image_to_video.md index 00b67d74e..781f0c2a5 100644 --- a/docs/user_guide/examples/online_serving/image_to_video.md +++ b/docs/user_guide/examples/online_serving/image_to_video.md @@ -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. diff --git a/docs/user_guide/examples/online_serving/mimo_audio.md b/docs/user_guide/examples/online_serving/mimo_audio.md index 4737eca36..c8752f578 100644 --- a/docs/user_guide/examples/online_serving/mimo_audio.md +++ b/docs/user_guide/examples/online_serving/mimo_audio.md @@ -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.** diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md index 6f6d9ae4a..611eb6fd3 100644 --- a/docs/user_guide/examples/online_serving/qwen3_omni.md +++ b/docs/user_guide/examples/online_serving/qwen3_omni.md @@ -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** diff --git a/docs/user_guide/examples/online_serving/qwen3_tts.md b/docs/user_guide/examples/online_serving/qwen3_tts.md index 4e632d4c2..95f234f02 100644 --- a/docs/user_guide/examples/online_serving/qwen3_tts.md +++ b/docs/user_guide/examples/online_serving/qwen3_tts.md @@ -58,7 +58,7 @@ Then open http://localhost:7860 in your browser. ```bash # 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 \ @@ -66,7 +66,7 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ # VoiceDesign model 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 \ @@ -74,7 +74,7 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \ # Base model (voice cloning) 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 \ diff --git a/docs/user_guide/examples/online_serving/text_to_video.md b/docs/user_guide/examples/online_serving/text_to_video.md index 01e6d9d46..b918aac19 100644 --- a/docs/user_guide/examples/online_serving/text_to_video.md +++ b/docs/user_guide/examples/online_serving/text_to_video.md @@ -165,6 +165,9 @@ curl -X POST http://localhost:8091/v1/videos \ -F "guidance_scale_2=4.0" \ -F "boundary_ratio=0.875" \ -F "flow_shift=5.0" \ + -F "enable_frame_interpolation=true" \ + -F "frame_interpolation_exp=1" \ + -F "frame_interpolation_scale=1.0" \ -F "seed=42" ``` @@ -187,6 +190,35 @@ curl -X POST http://localhost:8091/v1/videos \ | `flow_shift` | float | None | Scheduler flow shift (Wan2.2) | | `seed` | int | None | Random seed (reproducible) | | `lora` | object | None | LoRA configuration | +| `enable_frame_interpolation` | bool | false | Enable RIFE frame interpolation before MP4 encoding | +| `frame_interpolation_exp` | int | 1 | Interpolation exponent; 1=2x temporal resolution, 2=4x | +| `frame_interpolation_scale` | float | 1.0 | RIFE inference scale; use 0.5 for high-resolution inputs | +| `frame_interpolation_model_path` | str | None | Local directory or Hugging Face repo ID with `flownet.pkl`; defaults to `elfgum/RIFE-4.22.lite` | + +## Frame Interpolation + +Frame interpolation is an optional post-processing step for `/v1/videos` and +`/v1/videos/sync`. It synthesizes intermediate frames between generated frames +without rerunning the diffusion model. If the generated video has `N` frames, +the interpolated output frame count is `(N - 1) * 2**exp + 1`. The encoder FPS +is multiplied by `2**exp` so the output duration remains close to the original. + +Frame interpolation runs in the diffusion worker post-processing path instead of +the API server encoding path, so it can reuse the worker's current accelerator +device without blocking the FastAPI event loop. + +Example: generate 5 frames and interpolate to 9 frames: + +```bash +curl -X POST http://localhost:8091/v1/videos/sync \ + -F "prompt=A dog running through a park" \ + -F "num_frames=5" \ + -F "fps=8" \ + -F "enable_frame_interpolation=true" \ + -F "frame_interpolation_exp=1" \ + -F "frame_interpolation_scale=1.0" \ + -o sync_t2v_interpolated.mp4 +``` ## Create Response Format @@ -256,6 +288,14 @@ vllm serve Lightricks/LTX-2 --omni --port 8098 \ --enforce-eager --flow-shift 1.0 --boundary-ratio 1.0 ``` +For multi-GPU memory reduction, you can enable HSDP: + +```bash +vllm serve Lightricks/LTX-2 --omni --port 8098 \ + --enforce-eager --flow-shift 1.0 --boundary-ratio 1.0 \ + --use-hsdp --hsdp-shard-size 2 +``` + #### Start with Optimization Presets Use the LTX-2 startup script with built-in optimization presets: diff --git a/examples/offline_inference/bagel/README.md b/examples/offline_inference/bagel/README.md index 48517b1cd..3e653d0e3 100644 --- a/examples/offline_inference/bagel/README.md +++ b/examples/offline_inference/bagel/README.md @@ -173,8 +173,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 diff --git a/examples/offline_inference/cosyvoice3/README.md b/examples/offline_inference/cosyvoice3/README.md index e16134e6e..704b49614 100644 --- a/examples/offline_inference/cosyvoice3/README.md +++ b/examples/offline_inference/cosyvoice3/README.md @@ -58,7 +58,14 @@ 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. diff --git a/examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py b/examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py index 6311bbc90..de2712891 100644 --- a/examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py +++ b/examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py @@ -24,7 +24,13 @@ def run_e2e(): required=True, help="Path to CosyVoice3 model directory (e.g., pretrained_models/Fun-CosyVoice3-0.5B/).", ) - parser.add_argument("--stage-config", type=str, default="vllm_omni/model_executor/stage_configs/cosyvoice3.yaml") + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help="Override the deploy config path. If unset, auto-loads " + "vllm_omni/deploy/cosyvoice3.yaml based on the HF model_type.", + ) parser.add_argument("--prompt", type=str, default="Hello, this is a test of the CosyVoice system capability.") parser.add_argument( "--prompt-text", @@ -43,18 +49,14 @@ def run_e2e(): if not os.path.exists(args.tokenizer): raise FileNotFoundError(f"{args.tokenizer} does not exist!") - # Ensure stage config exists - if not os.path.exists(args.stage_config): - raise FileNotFoundError(f"{args.stage_config} does not exist!") + if args.deploy_config is not None and not os.path.exists(args.deploy_config): + raise FileNotFoundError(f"{args.deploy_config} does not exist!") print(f"Initializing cosyvoice E2E with model={args.model}") - # Initialize Omni - # This spins up the engine(s) based on the stage config - # We pass trust_remote_code=True same as Qwen examples omni = Omni( model=args.model, - stage_configs_path=args.stage_config, + deploy_config=args.deploy_config, trust_remote_code=True, tokenizer=args.tokenizer, log_stats=True, diff --git a/examples/offline_inference/glm_image/end2end.py b/examples/offline_inference/glm_image/end2end.py index 13bcd23f5..7ae9478ba 100644 --- a/examples/offline_inference/glm_image/end2end.py +++ b/examples/offline_inference/glm_image/end2end.py @@ -57,22 +57,22 @@ GLM_IMAGE_EOS_TOKEN_ID = 16385 # eos_token_id from generation_config.json GLM_IMAGE_VISION_VOCAB_SIZE = 16512 # top_k should be vision_vocab_size -def compute_max_tokens(height: int, width: int, factor: int = 32) -> int: +def compute_max_tokens(height: int, width: int, factor: int = 32, is_i2i: bool = False) -> int: """ Compute max_new_tokens for GLM-Image AR generation. - GLM-Image generates tokens in this order for text-to-image: - 1. Small preview image (half resolution in each dimension) - 2. Large target image (full resolution) - 3. EOS token + GLM-Image generation differs by mode: + - text-to-image (t2i): small preview + large target + EOS + - image-to-image (i2i): large target + EOS Args: height: Target image height in pixels width: Target image width in pixels factor: Downsampling factor (32 for GLM-Image AR output) + is_i2i: Whether the request is image-to-image mode Returns: - Total number of tokens to generate (small + large + EOS) + Total number of tokens to generate for the specified mode """ # Large image tokens (target resolution) token_h = height // factor @@ -80,11 +80,15 @@ def compute_max_tokens(height: int, width: int, factor: int = 32) -> int: large_tokens = token_h * token_w # Small preview tokens (half resolution in each dimension) - small_h = token_h // 2 - small_w = token_w // 2 - small_tokens = small_h * small_w + import math - # Total: small + large + EOS + ratio = token_h / token_w if token_w > 0 else 1.0 + small_token_h = max(1, int(math.sqrt(ratio) * (factor // 2))) + small_token_w = max(1, int(math.sqrt(1 / ratio) * (factor // 2))) + small_tokens = small_token_h * small_token_w + + if is_i2i: + return large_tokens + 1 return small_tokens + large_tokens + 1 @@ -282,14 +286,18 @@ def main(args: argparse.Namespace) -> None: # Compute max_tokens dynamically based on target image size target_height = prompt_dict.get("height", 1024) target_width = prompt_dict.get("width", 1024) - calculated_max_tokens = compute_max_tokens(target_height, target_width) + is_i2i = source_image is not None + calculated_max_tokens = compute_max_tokens(target_height, target_width, is_i2i=is_i2i) # Use calculated value unless user explicitly specified a different value # Default args.max_tokens is 16384 (very large), so prefer calculated value effective_max_tokens = calculated_max_tokens if args.max_tokens == 16384 else args.max_tokens if args.verbose: - print(f"AR max_tokens: {effective_max_tokens} (calculated: {calculated_max_tokens}, arg: {args.max_tokens})") + print( + f"AR max_tokens: {effective_max_tokens} " + f"(calculated: {calculated_max_tokens}, arg: {args.max_tokens}, mode: {'i2i' if is_i2i else 't2i'})" + ) # IMPORTANT: GLM-Image AR model requires these exact sampling parameters # from generation_config.json for proper image token generation. @@ -303,6 +311,12 @@ def main(args: argparse.Namespace) -> None: stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID], # 16385, CRITICAL for stopping seed=args.seed, detokenize=False, + # Keep target size available in runner/model for deterministic M-RoPE + # decode grids in t2i (no mm_features available in this path). + extra_args={ + "target_h": int(target_height), + "target_w": int(target_width), + }, ) # For diffusion stage, sampling_params contains diffusion-specific parameters diff --git a/examples/offline_inference/hunyuan_image3/README.md b/examples/offline_inference/hunyuan_image3/README.md index da28a44d9..3cd8fa01b 100644 --- a/examples/offline_inference/hunyuan_image3/README.md +++ b/examples/offline_inference/hunyuan_image3/README.md @@ -1,25 +1,161 @@ -# HunyuanImage-3.0 Image-to-Text Inference +# HunyuanImage-3.0-Instruct -This example demonstrates how to run HunyuanImage-3.0 Image-to-Text with the vLLM-Omni. +## Set up -## Local CLI Usage +Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. -Download the example image: +## Run examples + +**Note**: These examples work with the default configuration on **8x NVIDIA L40S (48GB)**. For different GPU setups, modify the stage configuration to adjust device allocation and memory utilization. + +Get into the hunyuan_image3 folder: ```bash -wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg +cd examples/offline_inference/hunyuan_image3 ``` -Run example: +### Modality Control + +HunyuanImage-3.0-Instruct supports multiple modality modes. You can control the mode using the `--modality` argument: + +#### Text to Image (text2img) + +- **Pipeline**: Text → AR (CoT + latent tokens) → DiT (denoise) → VAE Decode → Image +- **Stages Used**: Stage 0 (AR) + Stage 1 (DiT) +- **KV Transfer**: AR sends KV cache to DiT for conditioned generation +- **Default Config**: `hunyuan_image3_t2i.yaml` ```bash -python image_to_text.py \ - --image cherry_blossom.jpg \ - --prompt "<|startoftext|>You are an assistant that understands images and outputs text.Describe the content of the picture." +python end2end.py --model tencent/HunyuanImage-3.0-Instruct \ + --modality text2img \ + --prompts "A cute cat sitting on a windowsill watching the sunset" ``` -Key arguments: +#### Image to Image (img2img) -- `--model`: Model used. Default is: tencent/HunyuanImage-3.0-Instruct (Optional). -- `--image`: Path to input image (required). -- `--prompt`: Text description used to guide image understanding (required). +- **Pipeline**: Image + Text → AR (CoT + recaption + latent) → DiT → Edited Image +- **Stages Used**: Stage 0 (AR) + Stage 1 (DiT) +- **KV Transfer**: AR sends KV cache to DiT +- **Default Config**: `hunyuan_image3_it2i.yaml` + +```bash +python end2end.py --model tencent/HunyuanImage-3.0-Instruct \ + --modality img2img \ + --image-path /path/to/image.png \ + --prompts "Make the petals neon pink" +``` + +#### Image to Text (img2text) + +- **Pipeline**: Image + Question → AR → Text description +- **Stages Used**: Stage 0 (AR) only +- **Default Config**: `hunyuan_image3_i2t.yaml` + +```bash +python end2end.py --model tencent/HunyuanImage-3.0-Instruct \ + --modality img2text \ + --image-path /path/to/image.jpg \ + --prompts "Describe the content of the picture." +``` + +#### Text to Text (text2text) + +- **Pipeline**: Text → AR → Text +- **Stages Used**: Stage 0 (AR) only +- **Default Config**: `hunyuan_image3_t2t.yaml` + +```bash +python end2end.py --model tencent/HunyuanImage-3.0-Instruct \ + --modality text2text \ + --prompts "What is the capital of France?" +``` + +### Inference Steps & Guidance + +Control generation quality for image modalities: + +```bash +python end2end.py --modality text2img \ + --steps 50 \ + --guidance-scale 5.0 \ + --height 1024 --width 1024 \ + --prompts "A photo-realistic sunset over the ocean" +``` + +### Key Arguments + +#### 📌 Command Line Arguments (end2end.py) + +| Argument | Type | Default | Description | +| :--------------------- | :----- | :----------------------------------- | :----------------------------------------------------------- | +| `--model` | string | `tencent/HunyuanImage-3.0-Instruct` | Model path or name | +| `--modality` | choice | `text2img` | Modality: `text2img`, `img2img`, `img2text`, `text2text` | +| `--prompts` | list | `None` | Input text prompts | +| `--image-path` | string | `None` | Input image path (for `img2img`/`img2text`) | +| `--output` | string | `.` | Output directory for saved images | +| `--steps` | int | `50` | Number of inference steps | +| `--guidance-scale` | float | `5.0` | Classifier-free guidance scale | +| `--seed` | int | `42` | Random seed | +| `--height` | int | `1024` | Output image height | +| `--width` | int | `1024` | Output image width | +| `--bot-task` | string | auto | Override prompt task (e.g. `it2i_think`, `t2i_recaption`) | +| `--sys-type` | string | auto | Override system prompt type (e.g. `en_unified`, `en_vanilla`) | +| `--stage-configs-path` | string | auto | Custom stage config YAML path | +| `--enforce-eager` | flag | `False` | Disable torch.compile | +| `--init-timeout` | int | `300` | Initialization timeout (seconds) | + +------ + +#### ⚙️ Stage Configurations + +| Config YAML | Modality | Stages | GPUs | Description | +| :---------------------------------- | :-------- | :----- | :----- | :------------------------------------ | +| `hunyuan_image3_t2i.yaml` | text2img | 2 | 8 | T2I with AR→DiT, 4 GPU each | +| `hunyuan_image3_it2i.yaml` | img2img | 2 | 8 | IT2I with AR→DiT, 4 GPU each | +| `hunyuan_image3_i2t.yaml` | img2text | 1 | 4 | I2T (AR only) | +| `hunyuan_image3_t2t.yaml` | text2text | 1 | 4 | T2T (AR only) | +| `hunyuan_image3_t2i_2gpu.yaml` | text2img | 2 | 2 | T2I for 2-GPU setups | +| `hunyuan_image3_moe.yaml` | text2img | 2 | 8 | T2I with MoE AR→DiT KV reuse | +| `hunyuan_image3_moe_dit_2gpu_fp8.yaml` | text2img | 2 | 2 | T2I with FP8 quantization | + +------ + +## Using MoE Config + +The `hunyuan_image3_moe.yaml` config enables AR→DiT KV cache reuse with 8 GPUs (4 for AR + 4 for DiT). + +```bash +python end2end.py --model tencent/HunyuanImage-3.0-Instruct \ + --modality text2img \ + --stage-configs-path hunyuan_image3_moe.yaml \ + --prompts "A cute cat" +``` + +------ + +## Prompt Format + +HunyuanImage-3.0 uses a pretrain template format: + +``` +<|startoftext|>{system_prompt}{}{trigger_tag}{user_prompt} +``` + +- ``: Placeholder for each input image (auto-inserted by `prompt_utils.py`) +- Trigger tags: `` (CoT), `` (recaptioning) +- System prompt: Auto-selected based on task + +The `prompt_utils.build_prompt()` handles this formatting automatically. + +------ + +## FAQ + +- **OOM errors**: Decrease `gpu_memory_utilization` in the YAML stage config, or use a smaller `max_num_batched_tokens`. +- **Custom image sizes**: Use `--height` and `--width` flags (multiples of 16 recommended). + +| Stage | VRAM (approx) | +| :---------------- | :------------------- | +| Stage 0 (AR) | ~15 GiB + KV Cache | +| Stage 1 (DiT) | ~30 GiB | +| Total (8-GPU) | ~45 GiB + KV Cache | diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py new file mode 100644 index 000000000..3c1ae3866 --- /dev/null +++ b/examples/offline_inference/hunyuan_image3/end2end.py @@ -0,0 +1,262 @@ +""" +HunyuanImage-3.0-Instruct unified end-to-end inference script. + +Supports all modalities through a single entry point: + - text2img: Text → AR → DiT → Image + - img2img: Text+Image → AR → DiT → Edited Image (IT2I) + - img2text: Image+Text → AR → Text description (I2T) + - text2text: Text → AR → Text (comprehension, no image) + +Usage: + python end2end.py --modality text2img --prompts "A cute cat" + python end2end.py --modality img2img --image-path input.png --prompts "Make it snowy" + python end2end.py --modality img2text --image-path input.png --prompts "Describe this image" +""" + +import argparse +import os + +from vllm_omni.diffusion.models.hunyuan_image3.system_prompt import ( + get_system_prompt, +) +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniPromptType + +# task → (sys_type, bot_task, trigger_tag) +_TASK_PRESETS: dict[str, tuple[str, str | None, str | None]] = { + "t2t": ("en_unified", None, None), + "i2t": ("en_unified", None, None), + "it2i_think": ("en_unified", "think", ""), + "it2i_recaption": ("en_unified", "recaption", ""), + "t2i_think": ("en_unified", "think", ""), + "t2i_recaption": ("en_unified", "recaption", ""), + "t2i_vanilla": ("en_vanilla", "image", None), +} + +# Modality → prompt_utils task mapping +_MODALITY_TASK_MAP = { + "text2img": "t2i_think", + "img2img": "it2i_think", + "img2text": "i2t", + "text2text": "t2t", +} + + +def build_prompt( + user_prompt: str, + task: str = "it2i_think", + sys_type: str | None = None, + custom_system_prompt: str | None = None, +) -> str: + """Build a HunyuanImage-3.0 prompt using pretrain template format.""" + if task not in _TASK_PRESETS: + raise ValueError(f"Unknown task {task!r}. Choose from: {sorted(_TASK_PRESETS)}") + + preset_sys_type, preset_bot_task, trigger_tag = _TASK_PRESETS[task] + effective_sys_type = sys_type or preset_sys_type + + system_prompt = get_system_prompt(effective_sys_type, preset_bot_task, custom_system_prompt) + sys_text = system_prompt.strip() if system_prompt else "" + + has_image_input = task.startswith("i2t") or task.startswith("it2i") + + parts = ["<|startoftext|>"] + if sys_text: + parts.append(sys_text) + if has_image_input: + parts.append("") + if trigger_tag: + parts.append(trigger_tag) + parts.append(user_prompt) + + return "".join(parts) + + +# Modality → default stage config +_MODALITY_DEFAULT_CONFIG = { + "text2img": "hunyuan_image3_t2i.yaml", + "img2img": "hunyuan_image3_it2i.yaml", + "img2text": "hunyuan_image3_i2t.yaml", + "text2text": "hunyuan_image3_t2t.yaml", +} + + +def parse_args(): + parser = argparse.ArgumentParser(description="HunyuanImage-3.0-Instruct end-to-end inference.") + parser.add_argument( + "--model", + default="tencent/HunyuanImage-3.0-Instruct", + help="Model name or local path.", + ) + parser.add_argument( + "--modality", + default="text2img", + choices=["text2img", "img2img", "img2text", "text2text"], + help="Modality mode to control stage execution.", + ) + parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.") + parser.add_argument( + "--image-path", + type=str, + default=None, + help="Path to input image (for img2img/img2text).", + ) + parser.add_argument( + "--output", + type=str, + default=".", + help="Output directory to save results.", + ) + + # Generation parameters + parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.") + parser.add_argument("--guidance-scale", type=float, default=5.0, help="Classifier-free guidance scale.") + parser.add_argument("--seed", type=int, default=42, help="Random seed.") + parser.add_argument("--height", type=int, default=1024, help="Output image height.") + parser.add_argument("--width", type=int, default=1024, help="Output image width.") + + # Prompt configuration + parser.add_argument( + "--bot-task", + type=str, + default=None, + help="Override prompt task (e.g. it2i_think, t2i_recaption). Default: auto from modality.", + ) + parser.add_argument( + "--sys-type", + type=str, + default=None, + help="Override system prompt type (e.g. en_unified, en_vanilla).", + ) + + # Omni init args + parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom stage config YAML path.") + parser.add_argument("--log-stats", action="store_true", default=False) + parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.") + parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.") + + return parser.parse_args() + + +def main(): + args = parse_args() + os.makedirs(args.output, exist_ok=True) + + # Determine task for prompt formatting + task = args.bot_task or _MODALITY_TASK_MAP[args.modality] + + # Determine stage config + stage_configs_path = args.stage_configs_path or _MODALITY_DEFAULT_CONFIG[args.modality] + + # Build Omni + omni_kwargs = { + "model": args.model, + "stage_configs_path": stage_configs_path, + "log_stats": args.log_stats, + "init_timeout": args.init_timeout, + "enforce_eager": args.enforce_eager, + } + if args.modality in ("text2img", "img2img"): + omni_kwargs["mode"] = "text-to-image" + + omni = Omni(**omni_kwargs) + + # Prepare prompts + prompts = args.prompts or ["A cute cat"] + if not prompts: + print("[Info] No prompts provided, using default.") + prompts = ["A cute cat"] + + # Load image if needed + input_image = None + if args.modality in ("img2img", "img2text"): + if not args.image_path or not os.path.exists(args.image_path): + raise ValueError(f"--image-path required for {args.modality}, got: {args.image_path}") + from PIL import Image + + input_image = Image.open(args.image_path).convert("RGB") + + # Format prompts + formatted_prompts: list[OmniPromptType] = [] + for p in prompts: + formatted_text = build_prompt(p, task=task, sys_type=args.sys_type) + + prompt_dict: dict = {"prompt": formatted_text} + + if args.modality == "text2img": + prompt_dict["modalities"] = ["image"] + elif args.modality == "img2img": + prompt_dict["modalities"] = ["image"] + prompt_dict["multi_modal_data"] = {"image": input_image} + prompt_dict["height"] = input_image.height + prompt_dict["width"] = input_image.width + elif args.modality == "img2text": + prompt_dict["modalities"] = ["text"] + prompt_dict["multi_modal_data"] = {"image": input_image} + elif args.modality == "text2text": + prompt_dict["modalities"] = ["text"] + + formatted_prompts.append(prompt_dict) + + # Build sampling params from defaults + params_list = list(omni.default_sampling_params_list) + + # Override diffusion params if applicable + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + for i, sp in enumerate(params_list): + if isinstance(sp, OmniDiffusionSamplingParams): + sp.num_inference_steps = args.steps + sp.guidance_scale = args.guidance_scale + if args.seed is not None: + sp.seed = args.seed + if args.modality in ("text2img",): + sp.height = args.height + sp.width = args.width + + # Print configuration + print(f"\n{'=' * 60}") + print("HunyuanImage-3.0 Generation Configuration:") + print(f" Model: {args.model}") + print(f" Modality: {args.modality}") + print(f" Stage config: {stage_configs_path}") + print(f" Num stages: {omni.num_stages}") + if args.modality in ("text2img", "img2img"): + print(f" Inference steps: {args.steps}") + print(f" Guidance scale: {args.guidance_scale}") + print(f" Seed: {args.seed}") + if args.modality == "text2img": + print(f" Output size: {args.width}x{args.height}") + if args.image_path: + print(f" Input image: {args.image_path}") + print(f" Prompts: {prompts}") + print(f"{'=' * 60}\n") + + # Generate + omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list)) + + # Process outputs + img_idx = 0 + for req_output in omni_outputs: + # Text output (AR stage or text-only) + ro = getattr(req_output, "request_output", None) + if ro and getattr(ro, "outputs", None): + txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs) + if txt: + print(f"[Output] Text:\n{txt}") + + # Image output (DiT stage) + images = getattr(req_output, "images", None) + if not images and ro and hasattr(ro, "images"): + images = ro.images + + if images: + for j, img in enumerate(images): + save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png") + img.save(save_path) + print(f"[Output] Saved image to {save_path}") + img_idx += 1 + + +if __name__ == "__main__": + main() diff --git a/examples/offline_inference/hunyuan_image3/image_to_text.py b/examples/offline_inference/hunyuan_image3/image_to_text.py deleted file mode 100644 index d40134ac0..000000000 --- a/examples/offline_inference/hunyuan_image3/image_to_text.py +++ /dev/null @@ -1,92 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import argparse -import os - -from PIL import Image - -from vllm_omni.entrypoints.omni import Omni - -""" -The tencent/HunyuanImage-3.0-Instruct base model uses the tencent/Hunyuan-A13B-Instruct backbone. It utilizes two tokenizer delimiter templates: - -1) Pretrained template (default for gen_text mode), which concatenates system, image - tokens, and user question WITHOUT role delimiters: -"<|startoftext|>{system_prompt}{image_tokens}{user_question}" - - Example (before image token expansion): -"<|startoftext|>You are an assistant that understands images and outputs text.Describe the content of the picture." - -2) Instruct template, which uses explicit role prefixes and separators. -""" - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Generate text from image using HunyuanImage-3.0-Instruct.") - parser.add_argument( - "--model", - default="tencent/HunyuanImage-3.0-Instruct", - help="Model name or local path.", - ) - parser.add_argument( - "--image", - type=str, - required=True, - help="Path to input image file (PNG, JPG, etc.).", - ) - parser.add_argument( - "--prompt", - type=str, - required=True, - help="Pretrain template prompt: <|startoftext|>{system}{question}", - ) - parser.add_argument( - "--enable-diffusion-pipeline-profiler", - action="store_true", - help="Enable diffusion pipeline profiler to display stage durations.", - ) - return parser.parse_args() - - -def load_image(image_path: str) -> Image.Image: - """Load an image from file path.""" - if not os.path.exists(image_path): - raise FileNotFoundError(f"Image file not found: {image_path}") - return Image.open(image_path).convert("RGB") - - -def main(args: argparse.Namespace) -> None: - omni = Omni( - model=args.model, - enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, - mode="image-to-text", - ) - - prompt = "<|startoftext|>You are an assistant that understands images and outputs text." + args.prompt - - prompt_dict = { - "prompt": prompt, - "modalities": ["text"], - } - - # Add image input if provided - if args.image: - if not os.path.exists(args.image): - raise FileNotFoundError(f"Input image not found: {args.image}") - - input_image = load_image(args.image) - prompt_dict["multi_modal_data"] = {"image": input_image} - - prompts = [prompt_dict] - omni_outputs = omni.generate(prompts=prompts) - - prompt_text = omni_outputs[0].request_output.prompt - generated_text = omni_outputs[0].request_output.outputs[0].text - print(f"Prompt: {prompt_text}") - print(f"Text: {generated_text}") - - -if __name__ == "__main__": - args = parse_args() - main(args) diff --git a/examples/offline_inference/mimo_audio/README.md b/examples/offline_inference/mimo_audio/README.md index 596afabee..5615dea51 100644 --- a/examples/offline_inference/mimo_audio/README.md +++ b/examples/offline_inference/mimo_audio/README.md @@ -39,7 +39,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 ``` @@ -48,7 +47,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} @@ -66,7 +64,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." @@ -78,7 +75,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." \ @@ -91,7 +87,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." \ @@ -104,7 +99,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!" @@ -116,7 +110,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" @@ -128,7 +121,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." \ @@ -141,7 +133,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." \ @@ -154,7 +145,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" @@ -168,7 +158,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 ``` @@ -181,7 +170,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 ``` diff --git a/examples/offline_inference/mimo_audio/end2end.py b/examples/offline_inference/mimo_audio/end2end.py index ae044d2e8..2c67d66fa 100644 --- a/examples/offline_inference/mimo_audio/end2end.py +++ b/examples/offline_inference/mimo_audio/end2end.py @@ -182,7 +182,7 @@ def main(args): omni = Omni( model=model_name, - stage_configs_path=args.stage_configs_path, + deploy_config=args.deploy_config, log_stats=args.enable_stats, log_file=("omni_pipeline.log" if args.enable_stats else None), init_sleep_seconds=args.init_sleep_seconds, @@ -428,10 +428,11 @@ def parse_args(): help="Sampling rate for audio.", ) parser.add_argument( - "--stage-configs-path", + "--deploy-config", type=str, - default="../../../model_executor/stage_configs/mimo_audio.yaml", - help="Path to a stage configs file.", + default=None, + help="Override the deploy config path. If unset, auto-loads " + "vllm_omni/deploy/mimo_audio.yaml based on the HF model_type.", ) return parser.parse_args() diff --git a/examples/offline_inference/ming_flash_omni/README.md b/examples/offline_inference/ming_flash_omni/README.md new file mode 100644 index 000000000..7414163fc --- /dev/null +++ b/examples/offline_inference/ming_flash_omni/README.md @@ -0,0 +1,76 @@ +# Ming-flash-omni 2.0 + +[Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with outputs in text, image, and audio. For now, Ming-flash-omni-2.0 in vLLM-Omni is supported with thinker stage (multi-modal understanding). + +## Setup + +Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. + +## Run examples + +### Text-only +```bash +python examples/offline_inference/ming_flash_omni/end2end.py --query-type text +``` + +#### Reasoning (Thinking Mode) + +Reasoning (Thinking) mode is enabled via applying "detailed thinking on" when building the system prompt template (in `apply_chat_template`). + +In the end2end example, a default problem for thinking mode is provided, as referred to the example usage of Ming's cookbook; +To utilize it, you have to download the example figure from https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/3_0.png + +```bash +python examples/offline_inference/ming_flash_omni/end2end.py -q reasoning --image-path ./3_0.png +``` + +### Image understanding +```bash +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image + +# With a local image +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image --image-path /path/to/image.jpg +``` + +### Audio understanding +```bash +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio + +# With a local audio file +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio --audio-path /path/to/audio.wav +``` + +### Video understanding +```bash +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video + +# With a local video and custom frame count +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video --video-path /path/to/video.mp4 --num-frames 16 +``` + +### Mixed modalities (image + audio) +```bash +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_mixed_modalities \ + --image-path /path/to/image.jpg \ + --audio-path /path/to/audio.wav +``` + +If media file paths are not provided, the script uses built-in default assets. + +### Modality control +To control output modalities (e.g. text-only output): +```bash +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio --modalities text +``` + +*For now, only text output is supported* + +### Custom stage config +```bash +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image \ + --stage-configs-path /path/to/your_config.yaml +``` + +## Online serving + +For online serving via the OpenAI-compatible API, see [examples/online_serving/ming_flash_omni/README.md](../../online_serving/ming_flash_omni/README.md). diff --git a/examples/offline_inference/ming_flash_omni/end2end.py b/examples/offline_inference/ming_flash_omni/end2end.py new file mode 100644 index 000000000..8f8730131 --- /dev/null +++ b/examples/offline_inference/ming_flash_omni/end2end.py @@ -0,0 +1,485 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Partial example cases are referred from +# https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/cookbook.ipynb +import os +import time +from typing import NamedTuple + +import numpy as np +import vllm +from PIL import Image +from transformers import AutoProcessor +from vllm import SamplingParams +from vllm.assets.audio import AudioAsset +from vllm.assets.image import ImageAsset +from vllm.assets.video import VideoAsset, video_to_ndarrays +from vllm.multimodal.image import convert_image_mode +from vllm.multimodal.media.audio import load_audio +from vllm.utils.argparse_utils import FlexibleArgumentParser + +import vllm_omni +from vllm_omni.entrypoints.omni import Omni + +# Imports the processor also registers itself +from vllm_omni.transformers_utils.processors.ming import MingFlashOmniProcessor # noqa: F401 + +SEED = 42 +MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0" + + +class QueryResult(NamedTuple): + inputs: dict + limit_mm_per_prompt: dict[str, int] + + +def get_text_query(processor: MingFlashOmniProcessor, question: str | None = None) -> QueryResult: + if question is None: + question = "请详细介绍鹦鹉的生活习性。" + conversation = [{"role": "HUMAN", "content": question}] + prompt = processor.apply_chat_template(conversation, tokenize=False) + return QueryResult( + inputs={"prompt": prompt}, + limit_mm_per_prompt={}, + ) + + +def get_image_query( + processor: MingFlashOmniProcessor, + question: str | None = None, + image_path: str | None = None, +) -> QueryResult: + if question is None: + question = "Describe this image in detail." + + if image_path: + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") + image_data = convert_image_mode(Image.open(image_path), "RGB") + else: + image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB") + + conversation = [ + { + "role": "HUMAN", + "content": [ + {"type": "image", "image": image_data}, + {"type": "text", "text": question}, + ], + } + ] + prompt = processor.apply_chat_template(conversation, tokenize=False) + + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": {"image": image_data}, + }, + limit_mm_per_prompt={"image": 1}, + ) + + +def get_audio_query( + processor: MingFlashOmniProcessor, + question: str | None = None, + audio_path: str | None = None, + sampling_rate: int = 16000, +) -> QueryResult: + if question is None: + question = "Please recognize the language of this speech and transcribe it. Format: oral." + + if audio_path: + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + audio_signal, sr = load_audio(audio_path, sr=sampling_rate) + audio_data = (audio_signal.astype(np.float32), sr) + else: + audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate + + # Use a string for "audio" so the processor counts it as 1 audio input + conversation = [ + { + "role": "HUMAN", + "content": [ + {"type": "audio", "audio": "input"}, + {"type": "text", "text": question}, + ], + } + ] + prompt = processor.apply_chat_template(conversation, tokenize=False) + + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": {"audio": audio_data}, + }, + limit_mm_per_prompt={"audio": 1}, + ) + + +def get_video_query( + processor: MingFlashOmniProcessor, + question: str | None = None, + video_path: str | None = None, + num_frames: int = 16, +) -> QueryResult: + if question is None: + question = "Describe what is happening in this video." + + if video_path: + if not os.path.exists(video_path): + raise FileNotFoundError(f"Video file not found: {video_path}") + video_frames = video_to_ndarrays(video_path, num_frames=num_frames) + else: + video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays + + conversation = [ + { + "role": "HUMAN", + "content": [ + {"type": "video"}, + {"type": "text", "text": question}, + ], + } + ] + prompt = processor.apply_chat_template(conversation, tokenize=False) + + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": {"video": video_frames}, + }, + limit_mm_per_prompt={"video": 1}, + ) + + +def get_mixed_modalities_query( + processor: MingFlashOmniProcessor, + image_path: str | None = None, + audio_path: str | None = None, + sampling_rate: int = 16000, +) -> QueryResult: + """Mixed image + audio understanding.""" + question = "Describe the image, and recognize the language of this speech and transcribe it. Format: oral" + + if image_path: + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") + image_data = convert_image_mode(Image.open(image_path), "RGB") + else: + image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB") + + if audio_path: + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + sig, sr = load_audio(audio_path, sr=sampling_rate) + audio_data = (sig.astype(np.float32), sr) + else: + audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate + + conversation = [ + { + "role": "HUMAN", + "content": [ + {"type": "image", "image": image_data}, + {"type": "audio", "audio": "input"}, + {"type": "text", "text": question}, + ], + } + ] + prompt = processor.apply_chat_template(conversation, tokenize=False) + + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": {"image": image_data, "audio": audio_data}, + }, + limit_mm_per_prompt={"image": 1, "audio": 1}, + ) + + +def get_reasoning_query( + processor: MingFlashOmniProcessor, + question: str | None = None, + image_path: str | None = None, +) -> QueryResult: + if question is None: + # NOTE: To use the following default question, input with example figure provided by Ming + # https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/3_0.png + # E.g., + # python examples/offline_inference/ming_flash_omni/end2end.py -q reasoning --image-path ./3_0.png + # Otherwise, the problem solving might be false. + question = ( + "Based on the following rules:\n•\tYou control the smiley face character\n" + "•\tYou can move up, down, left, and right, and only a single square at a time\n" + "•\tWalls are dark grey and cannot be moved into\n•\tThe brown square is a box\n•" + "\tThe box can be pushed by moving into it (i.e., if you are in the square " + "adjacent to the box to the left, and move onto the square with the box, " + "the box will move one square to the right).\n" + "•\tThe box cannot be pushed into walls\n" + "•\tThe blue door at the bottom is locked and cannot be passed through, " + "unless the box is placed on the blue square\n" + "•\tThe square beneath the blue door is the exit\n" + "•\tMoving from one square to another\n\n" + "Let's assume a coordinate system where the smiley face is " + "on the top left at (1,1) and the square below it is (1,2). " + "The smiley face performs the following moves: {down, right, right, right}, " + "such that the smiley face is at square (4,2) and the box is in square (5,2). " + "What are the next sequence of moves that must be done to move the box down to (5,3)? " + "Give your answer as a comma separated list." + ) + + if image_path: + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") + image_data = convert_image_mode(Image.open(image_path), "RGB") + conversation = [ + { + "role": "HUMAN", + "content": [ + {"type": "image", "image": image_data}, + {"type": "text", "text": question}, + ], + } + ] + prompt = processor.apply_chat_template(conversation, tokenize=False, use_cot_system_prompt=True) + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": {"image": image_data}, + }, + limit_mm_per_prompt={"image": 1}, + ) + + conversation = [{"role": "HUMAN", "content": question}] + prompt = processor.apply_chat_template(conversation, tokenize=False, use_cot_system_prompt=True) + return QueryResult( + inputs={"prompt": prompt}, + limit_mm_per_prompt={}, + ) + + +query_map = { + "text": get_text_query, + "use_audio": get_audio_query, + "use_image": get_image_query, + "use_video": get_video_query, + "use_mixed_modalities": get_mixed_modalities_query, + "reasoning": get_reasoning_query, +} + + +def main(args): + print( + "=" * 20, + "\n", + f"vllm version: {vllm.__version__}\n", + f"vllm-omni version: {vllm_omni.__version__}\n", + "=" * 20, + sep="", + ) + + processor = AutoProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True) + assert isinstance(processor, MingFlashOmniProcessor), f"Wrong processor type being used: {type(processor)}" + + query_func = query_map[args.query_type] + if args.query_type == "use_image": + query_result = query_func(processor, image_path=args.image_path) + elif args.query_type == "use_audio": + query_result = query_func(processor, audio_path=args.audio_path, sampling_rate=args.sampling_rate) + elif args.query_type == "use_video": + query_result = query_func(processor, video_path=args.video_path, num_frames=args.num_frames) + elif args.query_type == "use_mixed_modalities": + query_result = query_func( + processor, + image_path=args.image_path, + audio_path=args.audio_path, + sampling_rate=args.sampling_rate, + ) + elif args.query_type == "reasoning": + query_result = query_func(processor, image_path=args.image_path) + else: + query_result = query_func(processor) + + # Initialize Omni (with thinker-only stage config) + omni = Omni( + model=MODEL_NAME, + stage_configs_path=args.stage_configs_path, + log_stats=args.log_stats, + init_timeout=args.init_timeout, + stage_init_timeout=args.stage_init_timeout, + ) + + # Thinker sampling params + thinker_sampling_params = SamplingParams( + temperature=0.4, + top_p=0.9, + max_tokens=args.max_tokens, + repetition_penalty=1.05, + seed=SEED, + detokenize=True, + ) + sampling_params_list = [thinker_sampling_params] + + prompts = [query_result.inputs for _ in range(args.num_prompts)] + + if args.modalities is not None: + output_modalities = args.modalities.split(",") + for prompt in prompts: + prompt["modalities"] = output_modalities + + total_requests = len(prompts) + processed_count = 0 + print(f"Query type: {args.query_type}") + print(f"Number of prompts: {total_requests}") + + output_dir = args.output_dir + os.makedirs(output_dir, exist_ok=True) + + profiler_enabled = args.enable_profiler + if profiler_enabled: + omni.start_profile(stages=args.profiler_stages) + + for stage_outputs in omni.generate(prompts, sampling_params_list): + output = stage_outputs.request_output + if stage_outputs.final_output_type == "text": + request_id = output.request_id + text_output = output.outputs[0].text + lines = [] + lines.append("Prompt:\n") + lines.append(str(output.prompt) + "\n") + lines.append("Text Output:\n") + lines.append(str(text_output).strip() + "\n") + print(*lines, sep="") + + # Save to file + out_txt = os.path.join(output_dir, f"{request_id}.txt") + try: + with open(out_txt, "w", encoding="utf-8") as f: + f.writelines(lines) + print(f"Request ID: {request_id}, text saved to {out_txt}") + except Exception as e: + print(f"Failed to write output file {out_txt}: {e}") + + elif stage_outputs.final_output_type == "audio": + raise NotImplementedError("Add audio example after talker supported.") + + processed_count += 1 + if profiler_enabled and processed_count >= total_requests: + print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...") + # Stop the profiler while workers are still alive + omni.stop_profile(stages=args.profiler_stages) + + print("[Info] Waiting 30s for workers to write trace files to disk...") + time.sleep(30) + print("[Info] Trace export wait time finished.") + + omni.close() + + +def parse_args(): + parser = FlexibleArgumentParser(description="Ming-flash-omni 2.0 offline inference example") + parser.add_argument( + "--query-type", + "-q", + type=str, + default="text", + choices=query_map.keys(), + help="Query type.", + ) + parser.add_argument( + "--stage-configs-path", + type=str, + default=None, + help="Path to a stage configs YAML file.", + ) + parser.add_argument( + "--log-stats", + action="store_true", + default=False, + help="Enable detailed statistics logging.", + ) + parser.add_argument("--init-timeout", type=int, default=2000, help="Timeout for initializing in seconds.") + parser.add_argument( + "--stage-init-timeout", + type=int, + default=2000, + help="Timeout for initializing a single stage in seconds.", + ) + parser.add_argument( + "--enable-profiler", + action="store_true", + default=False, + help="Enables profiling when set.", + ) + parser.add_argument( + "--profiler-stages", + type=int, + nargs="*", + default=[0], + help="List of stage IDs to profile. If not set, profiles all stages.", + ) + parser.add_argument( + "--image-path", + "-i", + type=str, + default=None, + help="Path to local image file. Uses default asset if not provided.", + ) + parser.add_argument( + "--audio-path", + "-a", + type=str, + default=None, + help="Path to local audio file. Uses default asset if not provided.", + ) + parser.add_argument( + "--video-path", + "-v", + type=str, + default=None, + help="Path to local video file. Uses default asset if not provided.", + ) + parser.add_argument( + "--num-frames", + type=int, + default=16, + help="Number of frames to extract from video.", + ) + parser.add_argument( + "--sampling-rate", + type=int, + default=16000, + help="Sampling rate for audio loading.", + ) + parser.add_argument( + "--max-tokens", + type=int, + default=16384, + help="Maximum tokens to generate.", + ) + parser.add_argument( + "--num-prompts", + type=int, + default=1, + help="Number of prompts to generate.", + ) + parser.add_argument( + "--modalities", + type=str, + default=None, + help="Output modalities (comma-separated).", + ) + parser.add_argument( + "--output-dir", + type=str, + default="output_ming", + help="Output directory for results.", + ) + + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index d8f1898ec..dfe124700 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -320,14 +320,7 @@ def main(args): query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate) else: query_result = query_func() - omni = Omni( - model=model_name, - log_stats=args.log_stats, - stage_init_timeout=args.stage_init_timeout, - batch_timeout=args.batch_timeout, - init_timeout=args.init_timeout, - shm_threshold_bytes=args.shm_threshold_bytes, - ) + omni = Omni.from_cli_args(args, model=model_name) thinker_sampling_params = SamplingParams( temperature=0.0, # Deterministic - no randomness top_p=1.0, # Disable nucleus sampling diff --git a/examples/offline_inference/qwen3_omni/README.md b/examples/offline_inference/qwen3_omni/README.md index d69ad6abf..0710faa13 100644 --- a/examples/offline_inference/qwen3_omni/README.md +++ b/examples/offline_inference/qwen3_omni/README.md @@ -70,8 +70,8 @@ For true stage-level concurrency -- where downstream stages (Talker, Code2Wav) start **before** the upstream stage (Thinker) finishes -- use the async_chunk example. This requires: -1. A stage config YAML with ``async_chunk: true`` (e.g. - ``qwen3_omni_moe_async_chunk.yaml``). +1. A deploy config YAML with ``async_chunk: true`` (e.g. + ``qwen3_omni_moe.yaml``). 2. Hardware that matches the config (e.g. 2x H100 for the default 3-stage config). @@ -101,7 +101,7 @@ python end2end_async_chunk.py --query-type text --modalities text ```bash python end2end_async_chunk.py \ --query-type use_audio \ - --stage-configs-path /path/to/your_async_chunk.yaml + --deploy-config /path/to/your_deploy_config.yaml ``` > **Note**: The synchronous ``end2end.py`` (using ``Omni``) is still the diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index 056f820ff..f028c32aa 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -294,14 +294,7 @@ def main(args): else: query_result = query_func() - omni = Omni( - model=model_name, - dtype=args.dtype, - stage_configs_path=args.stage_configs_path, - log_stats=args.log_stats, - stage_init_timeout=args.stage_init_timeout, - init_timeout=args.init_timeout, - ) + omni = Omni.from_cli_args(args, model=model_name) thinker_sampling_params = SamplingParams( temperature=0.9, diff --git a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py index 074426313..f38922e94 100644 --- a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py +++ b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py @@ -14,7 +14,7 @@ actively processing while stage-0 is still generating. Usage ----- python end2end_async_chunk.py --query-type use_audio \ - --stage-configs-path + --deploy-config See ``--help`` for all options. """ @@ -179,20 +179,26 @@ def clone_prompt_for_request(template: dict) -> dict: return cloned -def _default_async_chunk_stage_configs_path() -> str | None: - """Best-effort default stage config for running Qwen3-Omni with async_chunk. +def _default_deploy_config_path() -> str | None: + """Best-effort default deploy config for running Qwen3-Omni with async_chunk. - When this example is executed from within the repository, we resolve the - default YAML path relative to this file. When installed elsewhere, the - file may not exist and callers should pass --stage-configs-path explicitly. + The default ``vllm_omni/deploy/qwen3_omni_moe.yaml`` ships with + ``async_chunk: true`` at the top level, so loading it is enough to + enable async-chunk semantics. To disable it, copy the YAML and set + ``async_chunk: false`` (or pass ``--deploy-config`` to a YAML that + overrides the flag). + + When this example is executed from within the repository, we resolve + the default YAML path relative to this file. When installed elsewhere, + the file may not exist and callers should pass ``--deploy-config`` + explicitly. """ repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")) candidate = os.path.join( repo_root, "vllm_omni", - "model_executor", - "stage_configs", - "qwen3_omni_moe_async_chunk.yaml", + "deploy", + "qwen3_omni_moe.yaml", ) return candidate if os.path.exists(candidate) else None @@ -374,15 +380,16 @@ async def run_all(args): prompt["modalities"] = output_modalities # Create AsyncOmni - print(f"[Info] Creating AsyncOmni with stage_configs_path={args.stage_configs_path}") + print(f"[Info] Creating AsyncOmni with deploy_config={args.deploy_config}") async_omni = None try: - async_omni = AsyncOmni( - model=args.model, - stage_configs_path=args.stage_configs_path, - log_stats=args.log_stats, - stage_init_timeout=args.stage_init_timeout, - ) + # ``from_cli_args`` expands vars(args) into kwargs and auto-captures + # ``_cli_explicit_keys`` from ``sys.argv[1:]`` so argparse defaults + # do not silently override deploy YAML values. Mirrors the + # ``EngineArgs.from_cli_args`` pattern used throughout vllm / + # vllm-omni. ``deploy_config=None`` (the default) falls through to + # the bundled ``vllm_omni/deploy/qwen3_omni_moe.yaml``. + async_omni = AsyncOmni.from_cli_args(args) # Use default sampling params from stage config (they are pre-configured # in the YAML for each stage). @@ -470,11 +477,11 @@ def parse_args(): help="Query type.", ) parser.add_argument( - "--stage-configs-path", + "--deploy-config", type=str, - default=_default_async_chunk_stage_configs_path(), + default=_default_deploy_config_path(), help=( - "Path to an async_chunk stage config YAML. " + "Path to a deploy config YAML. " "If not set, uses the model's default config " "(make sure it has async_chunk: true)." ), diff --git a/examples/offline_inference/qwen3_omni/run_multiple_prompts_async_chunk.sh b/examples/offline_inference/qwen3_omni/run_multiple_prompts_async_chunk.sh index 809054867..2f2be2091 100755 --- a/examples/offline_inference/qwen3_omni/run_multiple_prompts_async_chunk.sh +++ b/examples/offline_inference/qwen3_omni/run_multiple_prompts_async_chunk.sh @@ -17,7 +17,7 @@ REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" python "${SCRIPT_DIR}/end2end_async_chunk.py" \ --query-type text \ --txt-prompts "${SCRIPT_DIR}/text_prompts_10.txt" \ - --stage-configs-path "${REPO_ROOT}/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml" \ + --deploy-config "${REPO_ROOT}/vllm_omni/deploy/qwen3_omni_moe.yaml" \ --output-dir output_audio_async_chunk \ --max-in-flight 2 \ "$@" diff --git a/examples/offline_inference/qwen3_omni/run_single_prompt_async_chunk.sh b/examples/offline_inference/qwen3_omni/run_single_prompt_async_chunk.sh index 918c7ee4f..9ef69293c 100755 --- a/examples/offline_inference/qwen3_omni/run_single_prompt_async_chunk.sh +++ b/examples/offline_inference/qwen3_omni/run_single_prompt_async_chunk.sh @@ -6,13 +6,13 @@ # achieving true stage-level concurrency via chunk-level streaming. # # Prerequisites: -# - An async_chunk stage config YAML (e.g. qwen3_omni_moe_async_chunk.yaml) +# - A deploy config YAML (e.g. qwen3_omni_moe.yaml) # - Hardware matching the config (e.g. 2x H100 for the default 3-stage config) # # Usage: # bash run_single_prompt_async_chunk.sh # bash run_single_prompt_async_chunk.sh --query-type text --modalities text -# bash run_single_prompt_async_chunk.sh --stage-configs-path /path/to/custom.yaml +# bash run_single_prompt_async_chunk.sh --deploy-config /path/to/custom.yaml set -euo pipefail @@ -21,6 +21,6 @@ REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" python "${SCRIPT_DIR}/end2end_async_chunk.py" \ --query-type use_audio \ - --stage-configs-path "${REPO_ROOT}/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml" \ + --deploy-config "${REPO_ROOT}/vllm_omni/deploy/qwen3_omni_moe.yaml" \ --output-dir output_audio_async_chunk \ "$@" diff --git a/examples/offline_inference/qwen3_tts/README.md b/examples/offline_inference/qwen3_tts/README.md index bf59dc9ba..2971ad716 100644 --- a/examples/offline_inference/qwen3_tts/README.md +++ b/examples/offline_inference/qwen3_tts/README.md @@ -15,11 +15,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 @@ -104,13 +104,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. diff --git a/examples/offline_inference/qwen3_tts/end2end.py b/examples/offline_inference/qwen3_tts/end2end.py index 901418c39..77da356b4 100644 --- a/examples/offline_inference/qwen3_tts/end2end.py +++ b/examples/offline_inference/qwen3_tts/end2end.py @@ -366,12 +366,7 @@ def main(args): output_dir = args.output_dir os.makedirs(output_dir, exist_ok=True) - omni = Omni( - model=model_name, - stage_configs_path=args.stage_configs_path, - log_stats=args.log_stats, - stage_init_timeout=args.stage_init_timeout, - ) + omni = Omni.from_cli_args(args, model=model_name) batch_size = args.batch_size for batch_start in range(0, len(inputs), batch_size): @@ -387,12 +382,7 @@ async def main_streaming(args): output_dir = args.output_dir os.makedirs(output_dir, exist_ok=True) - omni = AsyncOmni( - model=model_name, - stage_configs_path=args.stage_configs_path, - log_stats=args.log_stats, - stage_init_timeout=args.stage_init_timeout, - ) + omni = AsyncOmni.from_cli_args(args, model=model_name) for i, prompt in enumerate(inputs): request_id = str(i) diff --git a/examples/offline_inference/text_to_audio/README.md b/examples/offline_inference/text_to_audio/README.md index 50bab3e2f..10a8ae37e 100644 --- a/examples/offline_inference/text_to_audio/README.md +++ b/examples/offline_inference/text_to_audio/README.md @@ -27,6 +27,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). @@ -35,5 +51,8 @@ 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. - `--cache-backend`: cache acceleration backend. Stable Audio currently supports `tea_cache`. - `--output`: path to save the generated WAV file. diff --git a/examples/offline_inference/text_to_audio/text_to_audio.py b/examples/offline_inference/text_to_audio/text_to_audio.py index 3adb3ad53..53f3daf8d 100644 --- a/examples/offline_inference/text_to_audio/text_to_audio.py +++ b/examples/offline_inference/text_to_audio/text_to_audio.py @@ -21,6 +21,7 @@ from pathlib import Path import numpy as np import torch +from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -113,6 +114,23 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable diffusion pipeline profiler to display stage durations.", ) + parser.add_argument( + "--use-hsdp", + action="store_true", + help="Enable HSDP for Stable Audio DiT weight sharding.", + ) + parser.add_argument( + "--hsdp-shard-size", + type=int, + default=1, + help="Number of GPUs to shard Stable Audio DiT weights across when HSDP is enabled.", + ) + parser.add_argument( + "--hsdp-replicate-size", + type=int, + default=1, + help="Number of HSDP replica groups. Default 1 means pure sharding.", + ) return parser.parse_args() @@ -158,12 +176,23 @@ def main(): print(f" Inference steps: {args.num_inference_steps}") print(f" Guidance scale: {args.guidance_scale}") print(f" Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}") + if args.use_hsdp: + print(f" HSDP: enabled (shard_size={args.hsdp_shard_size}, replicate_size={args.hsdp_replicate_size})") + else: + print(" HSDP: disabled") print(f" Seed: {args.seed}") print(f"{'=' * 60}\n") + parallel_config = DiffusionParallelConfig( + use_hsdp=args.use_hsdp, + hsdp_shard_size=args.hsdp_shard_size, + hsdp_replicate_size=args.hsdp_replicate_size, + ) + # Initialize Omni with Stable Audio model omni = Omni( model=args.model, + parallel_config=parallel_config, cache_backend=args.cache_backend, cache_config=cache_config, enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index cc295e827..c71773972 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -29,7 +29,8 @@ This folder provides several entrypoints for experimenting with text-to-image di | `AIDC-AI/Ovis-Image-7B` | 1024 x 1024 | 71.8 | 17.1 | | `OmniGen2/OmniGen2` | 1024 x 1024 | 20.1 | 14.7 | | `stabilityai/stable-diffusion-3.5-medium` | 1024 x 1024 | 20.1 | 15.6 | -| `black-forest-labs/FLUX.1-dev` | 1024 x 1024 | 77.6 | 31.4 | +| `black-forest-labs/FLUX.1-dev` | 1024 x 1024 | 33.9 | 31.4 | +| `black-forest-labs/FLUX.1-schnell` | 1024 x 1024 | 33.9 | 31.4 | | `black-forest-labs/FLUX.2-klein-4B` | 1024 x 1024 | 72.7 | 14.9 | | `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 | | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) | diff --git a/examples/offline_inference/voxcpm/README.md b/examples/offline_inference/voxcpm/README.md new file mode 100644 index 000000000..1eaea9b0d --- /dev/null +++ b/examples/offline_inference/voxcpm/README.md @@ -0,0 +1,123 @@ +# VoxCPM Offline Example + +This directory contains the minimal offline VoxCPM example for vLLM Omni. + +`end2end.py` is intentionally small and only covers: + +- single text-to-speech +- single voice cloning with `ref_audio` + `ref_text` +- non-streaming with `vllm_omni/model_executor/stage_configs/voxcpm.yaml` +- streaming with `vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml` + +Advanced workflows were moved out of the getting-started example: + +- `benchmarks/voxcpm/vllm_omni/bench_tts_offline.py`: warmup, batch prompts, profiler, offline TTFP / RTF +- `benchmarks/voxcpm/vllm_omni/run_offline_matrix.py`: fixed offline smoke matrix +- `benchmarks/voxcpm/`: benchmark scripts and benchmark docs + +## Prerequisites + +Install VoxCPM in one of these ways: + +```bash +pip install voxcpm +``` + +or point vLLM Omni to the local VoxCPM source tree: + +```bash +export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src +``` + +The example writes WAV files with `soundfile`: + +```bash +pip install soundfile +``` + +## Model Path + +Pass the native VoxCPM model directory directly: + +```bash +export VOXCPM_MODEL=/path/to/voxcpm-model +``` + +If the native VoxCPM `config.json` does not contain HuggingFace metadata such as +`model_type`, prepare a persistent HF-compatible config directory and point the +stage configs to it with `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH`: + +```bash +export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config +mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH" +cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json" +cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true +python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)' +``` + +If the model directory itself already has `model_type`, this extra directory is +not required. + +## Quick Start + +Single text-to-speech, non-streaming: + +```bash +python examples/offline_inference/voxcpm/end2end.py \ + --model "$VOXCPM_MODEL" \ + --text "This is a split-stage VoxCPM synthesis example running on vLLM Omni." +``` + +Single voice cloning, non-streaming: + +```bash +python examples/offline_inference/voxcpm/end2end.py \ + --model "$VOXCPM_MODEL" \ + --text "This sentence is synthesized with a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "The exact transcript spoken in reference.wav." +``` + +Streaming: + +```bash +python examples/offline_inference/voxcpm/end2end.py \ + --model "$VOXCPM_MODEL" \ + --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \ + --text "This is a split-stage VoxCPM streaming example running on vLLM Omni." +``` + +By default, `end2end.py` writes to `output_audio/` for non-streaming and +`output_audio_streaming/` for streaming. + +## Advanced Workflows + +Use `benchmarks/voxcpm/vllm_omni/bench_tts_offline.py` when you need: + +- warmup runs +- prompt files +- batch JSONL inputs +- profiler injection +- offline TTFP / RTF emission + +Use `benchmarks/voxcpm/vllm_omni/run_offline_matrix.py` when you need the fixed offline smoke matrix that previously lived in `test.py`. + +Full matrix benchmark example: + +```bash +python benchmarks/voxcpm/vllm_omni/run_offline_matrix.py \ + --model "$VOXCPM_MODEL" \ + --ref-audio /path/to/reference.wav \ + --ref-text "The exact transcript spoken in reference.wav." +``` + +For online serving examples, see [examples/online_serving/voxcpm](../../online_serving/voxcpm/README.md). + +For benchmark reporting, see [benchmarks/voxcpm](../../../benchmarks/voxcpm/README.md). + +## Notes + +- `voxcpm.yaml` is the default non-streaming stage config. +- `voxcpm_async_chunk.yaml` is the streaming stage config. +- Streaming is currently single-request oriented; the fixed smoke matrix now lives in `benchmarks/voxcpm/vllm_omni/run_offline_matrix.py`. +- `ref_text` must be the real transcript of the reference audio. Mismatched text usually causes obvious quality degradation. diff --git a/examples/offline_inference/voxcpm/end2end.py b/examples/offline_inference/voxcpm/end2end.py new file mode 100644 index 000000000..980410fea --- /dev/null +++ b/examples/offline_inference/voxcpm/end2end.py @@ -0,0 +1,206 @@ +"""Minimal offline VoxCPM example for vLLM Omni.""" + +from __future__ import annotations + +import asyncio +import time +from pathlib import Path +from typing import Any + +import soundfile as sf +import torch +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni import AsyncOmni, Omni + +REPO_ROOT = Path(__file__).resolve().parents[3] +DEFAULT_SYNC_STAGE_CONFIG = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml" + + +def _build_prompt(args) -> dict[str, Any]: + additional_information: dict[str, list[Any]] = { + "text": [args.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 args.ref_audio is not None: + additional_information["ref_audio"] = [args.ref_audio] + if args.ref_text is not None: + additional_information["ref_text"] = [args.ref_text] + 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(item).float().cpu().reshape(-1) for item 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 _is_streaming_stage_config(stage_config_path: str) -> bool: + return "async_chunk" in Path(stage_config_path).stem + + +def _save_audio(audio: torch.Tensor, sample_rate: int, output_dir: Path, request_id: str) -> Path: + output_dir.mkdir(parents=True, exist_ok=True) + output_path = output_dir / f"output_{request_id}.wav" + sf.write( + output_path, + audio.float().cpu().clamp(-1.0, 1.0).numpy(), + sample_rate, + format="WAV", + subtype="PCM_16", + ) + return output_path + + +async def _run_streaming(args) -> Path: + prompt = _build_prompt(args) + output_dir = Path(args.output_dir) if args.output_dir is not None else Path("output_audio_streaming") + request_id = "streaming_example" + sample_rate = 24000 + buffered_samples = 0 + chunks: list[torch.Tensor] = [] + started = time.perf_counter() + omni = AsyncOmni( + model=args.model, + stage_configs_path=args.stage_configs_path, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + ) + try: + async for stage_output in omni.generate(prompt, request_id=request_id): + mm = getattr(stage_output, "multimodal_output", None) + if not isinstance(mm, dict): + request_output = getattr(stage_output, "request_output", None) + if request_output is None: + continue + mm = getattr(request_output, "multimodal_output", None) + if not isinstance(mm, dict) and getattr(request_output, "outputs", None): + mm = getattr(request_output.outputs[0], "multimodal_output", None) + if not isinstance(mm, dict): + continue + audio = _extract_audio_tensor(mm) + if audio.numel() == 0: + continue + sample_rate = _extract_sample_rate(mm) + if audio.numel() > buffered_samples: + delta = audio[buffered_samples:] + buffered_samples = int(audio.numel()) + else: + delta = audio + buffered_samples += int(delta.numel()) + if delta.numel() > 0: + chunks.append(delta) + if not chunks: + raise RuntimeError("No streaming audio chunks received from VoxCPM.") + output_audio = torch.cat(chunks, dim=0) + output_path = _save_audio(output_audio, sample_rate, output_dir, request_id) + print(f"Saved streaming audio to: {output_path} ({time.perf_counter() - started:.2f}s)") + return output_path + finally: + omni.shutdown() + + +def _run_sync(args) -> Path: + prompt = _build_prompt(args) + output_dir = Path(args.output_dir) if args.output_dir is not None else Path("output_audio") + request_id = "sync_example" + started = time.perf_counter() + last_mm: dict[str, Any] | None = None + omni = Omni( + model=args.model, + stage_configs_path=args.stage_configs_path, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + ) + for stage_outputs in omni.generate(prompt): + request_output = getattr(stage_outputs, "request_output", None) + if request_output is None: + continue + outputs = getattr(request_output, "outputs", None) + if outputs: + for output in outputs: + mm = getattr(output, "multimodal_output", None) + if isinstance(mm, dict): + last_mm = mm + mm = getattr(request_output, "multimodal_output", None) + if isinstance(mm, dict): + last_mm = mm + if last_mm is None: + raise RuntimeError("No audio output received from VoxCPM.") + output_path = _save_audio( + _extract_audio_tensor(last_mm), + _extract_sample_rate(last_mm), + output_dir, + request_id, + ) + print(f"Saved audio to: {output_path} ({time.perf_counter() - started:.2f}s)") + return output_path + + +def parse_args(): + parser = FlexibleArgumentParser(description="Minimal offline VoxCPM example for vLLM Omni.") + parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.") + parser.add_argument( + "--stage-configs-path", + type=str, + default=str(DEFAULT_SYNC_STAGE_CONFIG), + help=("Stage config path. Use voxcpm.yaml for non-streaming or voxcpm_async_chunk.yaml for streaming."), + ) + parser.add_argument("--text", type=str, required=True, help="Input text for synthesis.") + parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.") + parser.add_argument("--ref-text", type=str, default=None, help="Transcript of the reference audio.") + parser.add_argument("--output-dir", type=str, default=None, help="Output directory for generated wav files.") + parser.add_argument("--cfg-value", type=float, default=2.0, help="Guidance value passed to VoxCPM.") + parser.add_argument("--inference-timesteps", type=int, default=10, help="Number of diffusion timesteps.") + parser.add_argument("--min-len", type=int, default=2, help="Minimum latent length.") + parser.add_argument("--max-new-tokens", type=int, default=4096, help="Maximum latent length.") + parser.add_argument( + "--streaming-prefix-len", + type=int, + default=3, + help="Streaming prefix length used by voxcpm_async_chunk.yaml.", + ) + parser.add_argument("--stage-init-timeout", type=int, default=600, help="Stage initialization timeout in seconds.") + parser.add_argument("--log-stats", action="store_true", help="Enable vLLM Omni stats logging.") + args = parser.parse_args() + if (args.ref_audio is None) != (args.ref_text is None): + raise ValueError("Voice cloning requires --ref-audio and --ref-text together.") + return args + + +def main(args) -> None: + route = "streaming" if _is_streaming_stage_config(args.stage_configs_path) else "sync" + print(f"Model: {args.model}") + print(f"Stage config: {args.stage_configs_path}") + print(f"Route: {route}") + if route == "streaming": + asyncio.run(_run_streaming(args)) + else: + _run_sync(args) + + +if __name__ == "__main__": + main(parse_args()) diff --git a/examples/offline_inference/voxcpm2/end2end.py b/examples/offline_inference/voxcpm2/end2end.py index ce404bf96..6b6bf78dd 100644 --- a/examples/offline_inference/voxcpm2/end2end.py +++ b/examples/offline_inference/voxcpm2/end2end.py @@ -65,6 +65,12 @@ def parse_args(): default=None, help="Text matching --prompt-audio for continuation mode.", ) + parser.add_argument( + "--ref-text", + type=str, + default=None, + help="Optional transcript of --reference-audio (enables ref_continuation mode).", + ) return parser.parse_args() @@ -74,16 +80,20 @@ def extract_audio(multimodal_output: dict) -> torch.Tensor: The output processor concatenates per-step delta tensors under ``model_outputs``. Falls back to ``audio`` for backwards compat. """ - audio = multimodal_output.get("model_outputs") or multimodal_output.get("audio") + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") if audio is None: raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}") if isinstance(audio, list): - # Take the last valid tensor (most complete audio) + # Defensive: usually the output processor consolidates into a single + # tensor at request completion, but concatenate here too in case the + # caller consumes intermediate (pre-consolidation) outputs. valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None] if not valid: raise ValueError("Audio list is empty or all elements are None.") - return valid[-1] + return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0] return torch.as_tensor(audio).float().cpu().reshape(-1) @@ -99,24 +109,40 @@ def main(): stage_configs_path=args.stage_configs_path, ) - additional: dict = {} - if args.reference_audio: - additional["reference_audio"] = args.reference_audio - if args.prompt_audio and args.prompt_text: - additional["prompt_audio"] = args.prompt_audio - additional["prompt_text"] = args.prompt_text + from transformers import AutoTokenizer - prompt: dict = {"prompt": args.text} - if additional: - prompt["additional_information"] = additional + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import ( + build_cjk_split_map, + build_voxcpm2_prompt, + ) + + tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) + split_map = build_cjk_split_map(tokenizer) + hf_config = engine.engine.stage_vllm_configs[0].model_config.hf_config + + ref_audio_arg = args.reference_audio or args.prompt_audio + ref_text_arg = args.ref_text or args.prompt_text + ref_wav, ref_sr = (None, None) + if ref_audio_arg: + ref_wav_arr, ref_sr = sf.read(ref_audio_arg) + ref_wav = ref_wav_arr.mean(axis=-1).tolist() if ref_wav_arr.ndim > 1 else ref_wav_arr.tolist() + + prompt = build_voxcpm2_prompt( + hf_config=hf_config, + tokenizer=tokenizer, + split_map=split_map, + text=args.text, + ref_audio=ref_wav, + ref_sr=ref_sr, + ref_text=ref_text_arg, + ) print(f"Model : {args.model}") print(f"Text : {args.text}") - if args.reference_audio: - print(f"Ref audio : {args.reference_audio}") - if args.prompt_audio: - print(f"Prompt audio: {args.prompt_audio}") - print(f"Prompt text : {args.prompt_text}") + if ref_audio_arg: + print(f"Ref audio : {ref_audio_arg}") + if ref_text_arg: + print(f"Ref text : {ref_text_arg}") print(f"Output dir : {output_dir}") t_start = time.perf_counter() diff --git a/examples/offline_inference/voxtral_tts/README.md b/examples/offline_inference/voxtral_tts/README.md index 5f3d5413b..a55ce8830 100644 --- a/examples/offline_inference/voxtral_tts/README.md +++ b/examples/offline_inference/voxtral_tts/README.md @@ -10,28 +10,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?" @@ -39,7 +35,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." \ @@ -54,7 +49,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`) | diff --git a/examples/offline_inference/voxtral_tts/end2end.py b/examples/offline_inference/voxtral_tts/end2end.py index 075024645..cf28f917e 100644 --- a/examples/offline_inference/voxtral_tts/end2end.py +++ b/examples/offline_inference/voxtral_tts/end2end.py @@ -39,7 +39,7 @@ logger = logging.getLogger(__name__) async def run_streaming(inputs, sampling_params_list, model_name, args, output_dir): async_omni = AsyncOmni( model=model_name, - stage_configs_path=args.stage_configs_path, + deploy_config=args.deploy_config, log_stats=args.log_stats, ) @@ -192,7 +192,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir llm = Omni( model=model_name, log_stats=args.log_stats, - stage_configs_path=args.stage_configs_path, + deploy_config=args.deploy_config, ) if args.profiling_mode: @@ -253,10 +253,11 @@ def parse_args() -> Namespace: help="Directory to write output wav files.", ) parser.add_argument( - "--stage-configs-path", + "--deploy-config", type=str, default=None, - help="Path to stage configs YAML. Auto-resolved from model if not set.", + help="Override the deploy config path. If unset, auto-loads " + "vllm_omni/deploy/voxtral_tts.yaml based on the HF model_type.", ) parser.add_argument( "--num-prompts", type=int, default=1, help="Number of replicate prompts to run for measuring performance" diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md index 4b5188f41..13f2cfe7c 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md +++ b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md @@ -30,9 +30,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 \ @@ -42,11 +42,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 ": In the frame, a woman with black long hair is identified as .\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**: 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**: 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 maintains eye contact, swallows as if choosing her words, and says, I keep telling myself I’m fine,but some nights it feels like I’m just performing calm." \ + --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. diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py index 49a0f496f..497284ceb 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py +++ b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py @@ -5,10 +5,12 @@ import argparse import re import time +import numpy as np from PIL import Image from vllm.multimodal.media.audio import load_audio from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -56,6 +58,11 @@ def parse_args() -> argparse.Namespace: default=False, help="Enable CPU offloading for diffusion models.", ) + parser.add_argument( + "--enable-layerwise-offload", + action="store_true", + help="Enable layerwise (blockwise) offloading on DiT modules.", + ) return parser.parse_args() @@ -124,6 +131,7 @@ def main() -> None: parallel_config=parallel_config, model_type=args.model_type, enable_cpu_offload=args.enable_cpu_offload, + enable_layerwise_offload=args.enable_layerwise_offload, ) start = time.perf_counter() outputs = omni.generate(prompt, sampling_params) @@ -131,15 +139,35 @@ def main() -> None: if not outputs: raise RuntimeError("No output returned from DreamID-Omni.") - output = outputs[0].request_output - generated_video = output.images[0][0] - generated_audio = output.images[0][1] - try: - from dreamid_omni.utils.io_utils import save_video - except Exception as e: - raise RuntimeError(f"Failed to extract video and audio from DreamID-Omni output. Error: {e}") + result = outputs[0] + if not result.images: + raise RuntimeError("No video frames found in DreamID-Omni output.") + generated_video = result.images[0] + mm = result.multimodal_output or {} + generated_audio = mm.get("audio") + fps = int(mm.get("fps", 24)) + sample_rate = int(mm.get("audio_sample_rate", 16000)) + + # DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1]. + # mux_video_audio_bytes expects (F, H, W, C) uint8. + if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4: + raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}") + frames = generated_video.transpose(1, 2, 3, 0) + frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8) + + audio_np = None + if generated_audio is not None: + audio_np = np.squeeze(np.asarray(generated_audio)).astype(np.float32) + output_path = args.output - save_video(output_path, generated_video, generated_audio, fps=24, sample_rate=16000) + video_bytes = mux_video_audio_bytes( + frames, + audio_np, + fps=float(fps), + audio_sample_rate=sample_rate, + ) + with open(output_path, "wb") as f: + f.write(video_bytes) print(f"Saved generated video to {output_path}") print(f"Total time: {elapsed:.2f}s") diff --git a/examples/online_serving/fish_speech/README.md b/examples/online_serving/fish_speech/README.md index ae968d3ba..9b4e3cc40 100644 --- a/examples/online_serving/fish_speech/README.md +++ b/examples/online_serving/fish_speech/README.md @@ -29,15 +29,12 @@ 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` +(the HF `model_type` on the fishaudio checkpoints is `fish_qwen3_omni`). + Or use the convenience script: ```bash diff --git a/examples/online_serving/fish_speech/run_gradio_demo.sh b/examples/online_serving/fish_speech/run_gradio_demo.sh index 98a696644..a0370b9cc 100755 --- a/examples/online_serving/fish_speech/run_gradio_demo.sh +++ b/examples/online_serving/fish_speech/run_gradio_demo.sh @@ -11,18 +11,13 @@ MODEL="${MODEL:-fishaudio/s2-pro}" PORT="${PORT:-8091}" GRADIO_PORT="${GRADIO_PORT:-7860}" SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" -REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)" echo "Starting Fish Speech S2 Pro server (port $PORT)..." FLASHINFER_DISABLE_VERSION_CHECK=1 \ -vllm-omni serve "$MODEL" \ - --stage-configs-path "$REPO_ROOT/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml" \ +vllm serve "$MODEL" \ + --omni \ --host 0.0.0.0 \ - --port "$PORT" \ - --gpu-memory-utilization 0.9 \ - --trust-remote-code \ - --enforce-eager \ - --omni & + --port "$PORT" & SERVER_PID=$! cleanup() { diff --git a/examples/online_serving/fish_speech/run_server.sh b/examples/online_serving/fish_speech/run_server.sh index 59c09c7fe..a865daf93 100755 --- a/examples/online_serving/fish_speech/run_server.sh +++ b/examples/online_serving/fish_speech/run_server.sh @@ -13,11 +13,7 @@ PORT="${PORT:-8091}" echo "Starting Fish Speech S2 Pro server with model: $MODEL" FLASHINFER_DISABLE_VERSION_CHECK=1 \ -vllm-omni serve "$MODEL" \ - --stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \ +vllm serve "$MODEL" \ + --omni \ --host 0.0.0.0 \ - --port "$PORT" \ - --gpu-memory-utilization 0.9 \ - --trust-remote-code \ - --enforce-eager \ - --omni + --port "$PORT" diff --git a/examples/online_serving/mimo_audio/README.md b/examples/online_serving/mimo_audio/README.md index 9f70d59cb..9c1be7f21 100644 --- a/examples/online_serving/mimo_audio/README.md +++ b/examples/online_serving/mimo_audio/README.md @@ -13,10 +13,10 @@ Please refer to [README.md](../../../README.md) ```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.** diff --git a/examples/online_serving/ming_flash_omni/README.md b/examples/online_serving/ming_flash_omni/README.md new file mode 100644 index 000000000..502232725 --- /dev/null +++ b/examples/online_serving/ming_flash_omni/README.md @@ -0,0 +1,204 @@ +# Ming-flash-omni 2.0 + +## Installation + +Please refer to [README.md](../../../README.md) + +## Run examples (Ming-flash-omni 2.0) + +### Launch the Server + +```bash +vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091 +``` + +If you have custom stage configs file, launch the server with command below +```bash +vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091 --stage-configs-path /path/to/stage_configs_file +``` + +### Send Multi-modal Request + +#### Send request via python + +```bash +python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py --model Jonathan1909/Ming-flash-omni-2.0 --query-type use_mixed_modalities --port 8091 --host "localhost" --modalities text +``` + +The Python client supports the following command-line arguments: + +- `--query-type` (or `-q`): Query type. Options: `text`, `use_audio`, `use_image`, `use_video`, `use_mixed_modalities` +- `--video-path` (or `-v`): Path to local video file or URL. If not provided and query-type uses video, uses default video URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs. Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4` +- `--image-path` (or `-i`): Path to local image file or URL. If not provided and query-type uses image, uses default image URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common image formats: JPEG, PNG, GIF, WebP. Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png` +- `--audio-path` (or `-a`): Path to local audio file or URL. If not provided and query-type uses audio, uses default audio URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common audio formats: MP3, WAV, OGG, FLAC, M4A. Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3` +- `--prompt` (or `-p`): Custom text prompt/question. If not provided, uses default prompt for the selected query type. Example: `--prompt "What are the main activities shown in this video?"` +- `--modalities`: Output modalities. For now, only `text` is supported. Example: `--modalities text` + + +#### Send request via curl + +```bash +bash run_curl_multimodal_generation.sh text +bash run_curl_multimodal_generation.sh use_image +bash run_curl_multimodal_generation.sh use_audio +bash run_curl_multimodal_generation.sh use_video +bash run_curl_multimodal_generation.sh use_mixed_modalities +``` + +## Modality control + +Ming-flash-omni 2.0 currently supports text output only (thinker stage). + +| Modalities | Output | +|------------|--------| +| `["text"]` | Text only | +| Not specified | Text only (default) | + +### Using curl + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "messages": [ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, + {"role": "user", "content": "请详细介绍鹦鹉的生活习性。"} + ], + "modalities": ["text"] + }' +``` + +### Using OpenAI Python SDK + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY") + +response = client.chat.completions.create( + model="Jonathan1909/Ming-flash-omni-2.0", + messages=[ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, + {"role": "user", "content": "请详细介绍鹦鹉的生活习性。"}, + ], + modalities=["text"], +) +print(response.choices[0].message.content) +``` + +### Multi-modal input with OpenAI Python SDK + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY") + +response = client.chat.completions.create( + model="Jonathan1909/Ming-flash-omni-2.0", + messages=[ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"}}, + {"type": "text", "text": "Describe this image in detail."}, + ], + }, + ], + modalities=["text"], +) +print(response.choices[0].message.content) +``` + +## Streaming Output + +To enable streaming output: + +```bash +python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \ + --query-type use_image \ + --model Jonathan1909/Ming-flash-omni-2.0 \ + --modalities text \ + --stream +``` + +Or with the OpenAI Python SDK: + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY") + +response = client.chat.completions.create( + model="Jonathan1909/Ming-flash-omni-2.0", + messages=[ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, + {"role": "user", "content": "请详细介绍鹦鹉的生活习性。"}, + ], + modalities=["text"], + stream=True, +) +for chunk in response: + for choice in chunk.choices: + if hasattr(choice, "delta") and choice.delta.content: + print(choice.delta.content, end="", flush=True) +print() +``` + +Or using curl: + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "messages": [ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, + {"role": "user", "content": "请详细介绍鹦鹉的生活习性。"} + ], + "modalities": ["text"], + "stream": true, + }' +``` + + +## Reasoning (Thinking Mode) + +To enable reasoning/thinking mode, change `detailed thinking off` to `detailed thinking on` in the system prompt: + +### Using curl + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "messages": [ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking on"}]}, + {"role": "user", "content": [ + {"type": "image_url", "image_url": {"url": "https://example.com/math_problem.png"}}, + {"type": "text", "text": "Solve this math problem step by step."} + ]} + ], + "modalities": ["text"] + }' +``` + +### Using OpenAI Python SDK + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY") + +response = client.chat.completions.create( + model="Jonathan1909/Ming-flash-omni-2.0", + messages=[ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking on"}]}, + {"role": "user", "content": "If a train travels 120 km in 2 hours, what is its average speed?"}, + ], + modalities=["text"], +) +print(response.choices[0].message.content) +``` diff --git a/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh b/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh new file mode 100755 index 000000000..768a424e4 --- /dev/null +++ b/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh @@ -0,0 +1,145 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Server port +PORT="${PORT:-8091}" +# Default query type +QUERY_TYPE="${1:-text}" + +# Validate query type +if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video|use_mixed_modalities)$ ]]; then + echo "Error: Invalid query type '$QUERY_TYPE'" + echo "Usage: $0 [text|use_audio|use_image|use_video|use_mixed_modalities]" + echo " text: Text-only query" + echo " use_audio: Audio + Text query" + echo " use_image: Image + Text query" + echo " use_video: Video + Text query" + echo " use_mixed_modalities: Audio + Image + Video + Text query" + exit 1 +fi + +thinker_sampling_params='{ + "temperature": 0.4, + "top_p": 0.9, + "top_k": -1, + "max_tokens": 16384, + "seed": 42, + "detokenize": true, + "repetition_penalty": 1.05 +}' +# Above is optional, it has a default setting in stage_configs of the corresponding model. + +# Define URLs for assets +MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg" +CHERRY_BLOSSOM_IMAGE_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg" +SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4" + +# Build user content based on query type +case "$QUERY_TYPE" in + text) + user_content='[ + { + "type": "text", + "text": "请详细介绍鹦鹉的生活习性。" + } + ]' + ;; + use_image) + user_content='[ + { + "type": "image_url", + "image_url": { + "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'" + } + }, + { + "type": "text", + "text": "Describe this image in detail." + } + ]' + ;; + use_audio) + user_content='[ + { + "type": "audio_url", + "audio_url": { + "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'" + } + }, + { + "type": "text", + "text": "Please recognize the language of this speech and transcribe it. Format: oral." + } + ]' + ;; + use_video) + user_content='[ + { + "type": "video_url", + "video_url": { + "url": "'"$SAMPLE_VIDEO_URL"'" + } + }, + { + "type": "text", + "text": "Describe what is happening in this video." + } + ]' + ;; + use_mixed_modalities) + user_content='[ + { + "type": "image_url", + "image_url": { + "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'" + } + }, + { + "type": "audio_url", + "audio_url": { + "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'" + } + }, + { + "type": "text", + "text": "Describe the image, and recognize the language of this speech and transcribe it. Format: oral" + } + ]' + ;; +esac + +echo "Running query type: $QUERY_TYPE" +echo "" + +request_body=$(cat < **Note on `--no-async-chunk`**: Flips the deploy yaml's `async_chunk:` +> bool. Pipelines that implement alternate processor functions for +> chunked vs end-to-end modes (e.g. qwen3_tts code2wav) dispatch +> automatically based on that bool — no extra flag or variant yaml is +> needed. + +> ⚠️ **For multi-stage models that share GPUs (qwen3_omni_moe by default +> shares cuda:1 between stages 1 and 2), avoid using global memory flags.** +> A global `--gpu-memory-utilization 0.85` would apply to every stage and +> oversubscribe the shared device. Use per-stage overrides instead — see +> below. + +#### 2. Per-stage overrides via `--stage-overrides` (recommended for memory) + +```bash +# Lower stage 1's memory budget; leave others at the YAML default +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ + --stage-overrides '{ + "1": {"gpu_memory_utilization": 0.5}, + "2": {"max_num_batched_tokens": 65536} + }' +``` + +Per-stage values are always treated as explicit and beat YAML defaults for +the named stage. Other stages keep their YAML values. + +#### 3. Custom deploy YAML + +When per-stage overrides get long, write a small overlay YAML that inherits +from the bundled default: + +```yaml +# my_qwen3_omni_overrides.yaml +base_config: /path/to/vllm_omni/deploy/qwen3_omni_moe.yaml + +stages: + - stage_id: 0 + max_num_batched_tokens: 65536 + enforce_eager: true + - stage_id: 1 + gpu_memory_utilization: 0.5 + - stage_id: 2 + max_model_len: 8192 +``` + +Then start the server with `--deploy-config my_qwen3_omni_overrides.yaml`. +The `base_config:` line tells the loader to inherit everything else (stages, +connectors, edges, platforms section) from the bundled production YAML, so +you only need to spell out the deltas. + +#### 4. Multi-node deployment (cross-host transfer connector) + +The bundled `qwen3_omni_moe.yaml` uses `SharedMemoryConnector` between stages, +which only works when all stages run on the same physical host. For +**cross-node** deployments, write a small overlay YAML that swaps in a +network-capable connector (e.g. `MooncakeStoreConnector`) and re-points each +stage's connector wiring at it. The connector spec carries your own server +addresses — there is no checked-in default because every cluster is +different. + +```yaml +# my_qwen3_omni_multinode.yaml +base_config: /path/to/vllm_omni/deploy/qwen3_omni_moe.yaml + +connectors: + mooncake_connector: + name: MooncakeStoreConnector + extra: + host: "127.0.0.1" + metadata_server: "http://YOUR_METADATA_HOST:8080/metadata" + master: "YOUR_MASTER_HOST:50051" + segment: 512000000 # 512 MB transfer segment + localbuf: 64000000 # 64 MB local buffer + proto: "tcp" + +stages: + - stage_id: 0 + output_connectors: + to_stage_1: mooncake_connector + - stage_id: 1 + input_connectors: + from_stage_0: mooncake_connector + output_connectors: + to_stage_2: mooncake_connector + - stage_id: 2 + input_connectors: + from_stage_1: mooncake_connector +``` + +Then launch with `--deploy-config my_qwen3_omni_multinode.yaml`. Same +pattern works for Qwen2.5-Omni — replace `base_config:` with the path to +`vllm_omni/deploy/qwen2_5_omni.yaml`. + +> ⚠️ Replace `YOUR_METADATA_HOST` / `YOUR_MASTER_HOST` with the actual +> mooncake server addresses for your cluster. The `base_config:` overlay +> inherits all stage budgets, devices, and edges from the bundled prod +> YAML — you only need to spell out the connector swap. + ### Send Multi-modal Request Get into the example folder @@ -38,36 +180,43 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene #### Realtime WebSocket client (`openai_realtime_client.py`) -[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://:/v1/realtime`**, uploads a local audio file as **PCM16 mono @ 16 kHz** chunks (OpenAI-style `input_audio_buffer.append` / `commit`), and prints **streaming transcription** (`transcription.delta` / `transcription.done`). +[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://:/v1/realtime`**, streams a local WAV as **PCM16 mono @ 16 kHz** in fixed-size chunks (OpenAI-style `input_audio_buffer.append` / `commit`), and receives **`response.audio.delta`** (incremental PCM for the reply) plus **`transcription.*`** events. By default it concatenates audio deltas and writes **`--output-wav`** (model output is typically **24 kHz**). Optional **`--delta-dump-dir`** saves each delta as `delta_000001.wav`, … for debugging. + +Streaming input works well for translation-style use cases; if the Thinker runs while input is still incomplete, consider limiting **`max_tokens`** in your session / server defaults to avoid over-generation. **Dependencies:** ```bash -pip install websockets numpy +pip install websockets ``` **From this directory** (`examples/online_serving/qwen3_omni`): ```bash python openai_realtime_client.py \ - --host localhost \ - --port 8091 \ + --url ws://localhost:8091/v1/realtime \ --model Qwen/Qwen3-Omni-30B-A3B-Instruct \ - --audio_path /path/to/your.wav + --input-wav /path/to/input_16k_mono.wav \ + --output-wav realtime_output.wav \ + --delta-dump-dir ./rt_delta_wavs ``` -If `--audio_path` is omitted, the script uses a bundled default clip (`mary_had_lamb` via vLLM assets). - **Arguments:** | Flag | Default | Description | |------|---------|-------------| -| `--host` | `localhost` | API server host | -| `--port` | `8000` | API server port (match your `vllm serve` port, e.g. `8091`) | -| `--model` | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | Must match the served model (also sent in `session.update`) | -| `--audio_path` | *(optional)* | Path to input audio; resampled to 16 kHz mono inside the client | +| `--url` | `ws://localhost:8091/v1/realtime` | Full WebSocket URL including path | +| `--model` | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | Must match the served model (sent in `session.update`) | +| `--input-wav` | *(required)* | Input WAV: mono, 16-bit PCM, **16 kHz** | +| `--output-wav` | `realtime_output.wav` | Output path for concatenated reply audio | +| `--output-text` | *(optional)* | If set, write final transcription text to this path | +| `--chunk-ms` | `200` | Size of each uploaded audio chunk (milliseconds of audio) | +| `--send-delay-ms` | `0` | Delay between chunk sends (simulate realtime upload) | +| `--delta-dump-dir` | *(optional)* | Directory to write per-`response.audio.delta` WAV files | +| `--num-requests` | `1` | Number of sequential sessions (see `--concurrency`) | +| `--concurrency` | `1` | Max concurrent WebSocket sessions when `--num-requests` > 1 | -Ensure the vLLM-Omni server is running with realtime support for this endpoint, for example: +Ensure the server is running **without** `async_chunk` if you use `/v1/realtime`, for example: ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 @@ -276,7 +425,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) @@ -291,7 +440,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** diff --git a/examples/online_serving/qwen3_omni/openai_realtime_client.py b/examples/online_serving/qwen3_omni/openai_realtime_client.py index 660e4ac33..79e30a3f5 100644 --- a/examples/online_serving/qwen3_omni/openai_realtime_client.py +++ b/examples/online_serving/qwen3_omni/openai_realtime_client.py @@ -1,81 +1,118 @@ +"""Realtime client for vLLM-Omni /v1/realtime (audio + text events). + +This client: +1) Reads a local WAV file (must be mono, 16-bit PCM, 16kHz), +2) Streams PCM16 chunks to /v1/realtime with OpenAI-style events, +3) Receives response.audio.* and transcription.* events, +4) Saves synthesized audio to an output WAV file and optional text file. + +By default each ``response.audio.delta`` is treated as an **incremental PCM** +chunk and all chunks are concatenated into the final ``--output-wav``. + +Optional debugging: pass ``--delta-dump-dir DIR`` to write every +``response.audio.delta`` payload as ``delta_000001.wav``, ``delta_000002.wav``, … + +Usage: + python openai_realtime_client.py \ + --url ws://localhost:8091/v1/realtime \ + --model Qwen/Qwen3-Omni-30B-A3B-Instruct \ + --input-wav input_16k_mono.wav \ + --output-wav realtime_output.wav \ + --delta-dump-dir ./rt_delta_wavs + +Dependencies: + pip install websockets """ -This script demonstrates how to use the vLLM-Omni Realtime WebSocket API to perform -audio transcription by uploading an audio file. -Before running this script, you must start the vLLM-Omni server with a realtime-capable -model, for example: - - vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni - -Requirements: -- vllm with audio support -- websockets -- soundfile -- numpy - -The script: -1. Connects to the Realtime WebSocket endpoint -2. Converts an audio file to PCM16 @ 16kHz -3. Sends audio chunks to the server -4. Receives and prints transcription as it streams -""" +from __future__ import annotations import argparse import asyncio import base64 import json +import wave +from pathlib import Path -import numpy as np -import websockets -from vllm.assets.audio import AudioAsset -from vllm.multimodal.media.audio import load_audio +try: + import websockets +except ImportError: + print("Please install websockets: pip install websockets") + raise SystemExit(1) -def audio_to_pcm16_base64(audio_path: str) -> str: - """ - Load an audio file and convert it to base64-encoded PCM16 @ 16kHz. - """ - # Load audio and resample to 16kHz mono - audio, _ = load_audio(audio_path, sr=16000, mono=True) - # Convert to PCM16 - pcm16 = (audio * 32767).astype(np.int16) - # Encode as base64 - return base64.b64encode(pcm16.tobytes()).decode("utf-8") +def _read_wav_pcm16(path: Path) -> bytes: + with wave.open(str(path), "rb") as wf: + nchannels = wf.getnchannels() + sampwidth = wf.getsampwidth() + framerate = wf.getframerate() + comptype = wf.getcomptype() + nframes = wf.getnframes() + + if nchannels != 1: + raise ValueError(f"Input WAV must be mono (got {nchannels} channels).") + if sampwidth != 2: + raise ValueError(f"Input WAV must be 16-bit PCM (got sample width={sampwidth}).") + if framerate != 16000: + raise ValueError(f"Input WAV must be 16kHz (got {framerate} Hz).") + if comptype != "NONE": + raise ValueError(f"Input WAV must be uncompressed PCM (got comptype={comptype}).") + if nframes <= 0: + raise ValueError("Input WAV has no audio frames.") + + return wf.readframes(nframes) -async def realtime_transcribe(audio_path: str, host: str, port: int, model: str): - """ - Connect to the Realtime API and transcribe an audio file. - """ - uri = f"ws://{host}:{port}/v1/realtime" +def _write_wav_pcm16(path: Path, pcm16_bytes: bytes, sample_rate_hz: int) -> None: + with wave.open(str(path), "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate_hz) + wf.writeframes(pcm16_bytes) - async with websockets.connect(uri) as ws: - # Wait for session.created - response = json.loads(await ws.recv()) - if response["type"] == "session.created": - print(f"Session created: {response['id']}") - else: - print(f"Unexpected response: {response}") - return - # Validate model - await ws.send(json.dumps({"type": "session.update", "model": model})) +async def run_client( + url: str, + model: str, + input_wav: Path, + output_wav: Path, + output_text: Path | None, + chunk_ms: int, + send_delay_ms: int, + delta_dump_dir: Path | None, + request_idx: int = 1, + total_requests: int = 1, +) -> None: + log_prefix = f"[req {request_idx:02d}/{total_requests:02d}] " if total_requests > 1 else "" + pcm16 = _read_wav_pcm16(input_wav) + bytes_per_ms = 16000 * 2 // 1000 # mono PCM16 at 16kHz + chunk_bytes = max(bytes_per_ms * chunk_ms, 2) - # Signal ready to start - await ws.send(json.dumps({"type": "input_audio_buffer.commit"})) + incremental_pcm_parts: list[bytes] = [] + output_sample_rate = 24000 + delta_index = 0 + text_chunks: list[str] = [] + final_text: str = "" - # Convert audio file to base64 PCM16 - print(f"Loading audio from: {audio_path}") - audio_base64 = audio_to_pcm16_base64(audio_path) + if delta_dump_dir is not None: + delta_dump_dir.mkdir(parents=True, exist_ok=True) - # Send audio in chunks (4KB of raw audio = ~8KB base64) - chunk_size = 4096 - audio_bytes = base64.b64decode(audio_base64) - total_chunks = (len(audio_bytes) + chunk_size - 1) // chunk_size + async with websockets.connect(url, max_size=64 * 1024 * 1024) as ws: + # 1) Validate model. + await ws.send( + json.dumps( + { + "type": "session.update", + "model": model, + } + ) + ) - print(f"Sending {total_chunks} audio chunks...") - for i in range(0, len(audio_bytes), chunk_size): - chunk = audio_bytes[i : i + chunk_size] + # 2) Start generation once (non-final commit). + await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False})) + + # 3) Stream audio chunks. + for i in range(0, len(pcm16), chunk_bytes): + chunk = pcm16[i : i + chunk_bytes] await ws.send( json.dumps( { @@ -84,63 +121,212 @@ async def realtime_transcribe(audio_path: str, host: str, port: int, model: str) } ) ) + if send_delay_ms > 0: + await asyncio.sleep(send_delay_ms / 1000.0) - # Signal all audio is sent + # 4) Final commit closes input stream. await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True})) - print("Audio sent. Waiting for transcription...\n") - # Receive transcription - print("Transcription: ", end="", flush=True) + # 5) Receive server events until audio done. while True: - response = json.loads(await ws.recv()) - if response["type"] == "transcription.delta": - print(response["delta"], end="", flush=True) - elif response["type"] == "transcription.done": - print(f"\n\nFinal transcription: {response['text']}") - if response.get("usage"): - print(f"Usage: {response['usage']}") - break - elif response["type"] == "error": - print(f"\nError: {response['error']}") + message = await ws.recv() + if isinstance(message, bytes): + # We only expect JSON text frames. + continue + + event = json.loads(message) + event_type = event.get("type") + + if event_type == "session.created": + continue + + if event_type == "response.audio.delta": + sr = event.get("sample_rate_hz") + if isinstance(sr, int) and sr > 0: + output_sample_rate = sr + audio_b64 = event.get("audio", "") + if audio_b64: + pcm_delta = base64.b64decode(audio_b64) + incremental_pcm_parts.append(pcm_delta) + if delta_dump_dir is not None and pcm_delta: + delta_index += 1 + dump_path = delta_dump_dir / f"delta_{delta_index:06d}.wav" + _write_wav_pcm16(dump_path, pcm_delta, output_sample_rate) + print( + f"{log_prefix}delta dump #{delta_index}: {dump_path} " + f"(pcm bytes={len(pcm_delta)}, sr={output_sample_rate})" + ) + continue + + if event_type == "transcription.delta": + delta = event.get("delta", "") + if delta: + text_chunks.append(delta) + print(delta, end="", flush=True) + continue + + if event_type == "transcription.done": + final_text = event.get("text", "") or "".join(text_chunks) + usage = event.get("usage") + final_text_with_tag = f"Final transcription: {final_text}" + if text_chunks: + print() + print(f"{log_prefix}{final_text_with_tag}") + if usage: + print(f"{log_prefix}text usage: {usage}") + continue + + if event_type == "response.audio.done": break + if event_type == "error": + raise RuntimeError(f"Server error: {event}") -def main(args): - if args.audio_path: - audio_path = args.audio_path + all_pcm16 = b"".join(incremental_pcm_parts) + if not all_pcm16: + raise RuntimeError("No audio received from server.") + + output_wav.parent.mkdir(parents=True, exist_ok=True) + _write_wav_pcm16(output_wav, all_pcm16, output_sample_rate) + print(f"{log_prefix}Saved realtime audio to: {output_wav} (incremental chunks joined)") + + if output_text is not None: + text_to_save = final_text if final_text else "".join(text_chunks) + output_text.parent.mkdir(parents=True, exist_ok=True) + output_text.write_text(text_to_save, encoding="utf-8") + print(f"{log_prefix}Saved realtime text to: {output_text}") + + +def _indexed_output_path(path: Path | None, index: int, total: int) -> Path | None: + if path is None or total <= 1: + return path + return path.with_name(f"{path.stem}_{index:02d}{path.suffix}") + + +async def run_clients_concurrent( + *, + url: str, + model: str, + input_wav: Path, + output_wav: Path, + output_text: Path | None, + chunk_ms: int, + send_delay_ms: int, + delta_dump_dir: Path | None, + num_requests: int, + concurrency: int, +) -> None: + sem = asyncio.Semaphore(concurrency) + + async def _run_one(index: int) -> tuple[int, bool, str | None]: + per_output_wav = _indexed_output_path(output_wav, index, num_requests) + per_output_text = _indexed_output_path(output_text, index, num_requests) + per_delta_dir = None + if delta_dump_dir is not None: + per_delta_dir = delta_dump_dir / f"req_{index:02d}" + async with sem: + try: + await run_client( + url=url, + model=model, + input_wav=input_wav, + output_wav=per_output_wav, + output_text=per_output_text, + chunk_ms=chunk_ms, + send_delay_ms=send_delay_ms, + delta_dump_dir=per_delta_dir, + request_idx=index, + total_requests=num_requests, + ) + return index, True, None + except Exception as exc: + return index, False, str(exc) + + tasks = [asyncio.create_task(_run_one(i), name=f"rt-client-{i}") for i in range(1, num_requests + 1)] + results = await asyncio.gather(*tasks) + + failed = [(idx, err) for idx, ok, err in results if not ok] + succeeded = num_requests - len(failed) + print(f"[summary] succeeded={succeeded}, failed={len(failed)}, total={num_requests}") + if failed: + for idx, err in failed: + print(f"[summary] req {idx:02d} failed: {err}") + raise RuntimeError(f"{len(failed)} concurrent request(s) failed") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Realtime audio/text client for vLLM-Omni") + parser.add_argument("--url", default="ws://localhost:8091/v1/realtime", help="WebSocket URL") + parser.add_argument( + "--model", + default="Qwen/Qwen3-Omni-30B-A3B-Instruct", + help="Model name for session.update", + ) + parser.add_argument("--input-wav", required=True, type=Path, help="Input WAV (mono, PCM16, 16kHz)") + parser.add_argument("--output-wav", default=Path("realtime_output.wav"), type=Path, help="Output WAV path") + parser.add_argument( + "--output-text", + default=None, + type=Path, + help="Optional output text path for final transcription", + ) + parser.add_argument("--chunk-ms", type=int, default=200, help="Input chunk size in milliseconds") + parser.add_argument( + "--send-delay-ms", + type=int, + default=0, + help="Delay between chunk sends; set >0 to simulate realtime upload", + ) + parser.add_argument( + "--delta-dump-dir", + type=Path, + default=None, + help="If set, each response.audio.delta is saved as delta_NNNNNN.wav under this directory", + ) + parser.add_argument("--num-requests", type=int, default=1, help="Total number of requests to send") + parser.add_argument( + "--concurrency", + type=int, + default=1, + help="Maximum number of concurrent websocket requests", + ) + args = parser.parse_args() + + if args.num_requests <= 0: + raise ValueError("--num-requests must be >= 1") + if args.concurrency <= 0: + raise ValueError("--concurrency must be >= 1") + concurrency = min(args.concurrency, args.num_requests) + + if args.num_requests == 1: + asyncio.run( + run_client( + url=args.url, + model=args.model, + input_wav=args.input_wav, + output_wav=args.output_wav, + output_text=args.output_text, + chunk_ms=args.chunk_ms, + send_delay_ms=args.send_delay_ms, + delta_dump_dir=args.delta_dump_dir, + ) + ) else: - # Use default audio asset - audio_path = str(AudioAsset("mary_had_lamb").get_local_path()) - print(f"No audio path provided, using default: {audio_path}") - - asyncio.run(realtime_transcribe(audio_path, args.host, args.port, args.model)) + asyncio.run( + run_clients_concurrent( + url=args.url, + model=args.model, + input_wav=args.input_wav, + output_wav=args.output_wav, + output_text=args.output_text, + chunk_ms=args.chunk_ms, + send_delay_ms=args.send_delay_ms, + delta_dump_dir=args.delta_dump_dir, + num_requests=args.num_requests, + concurrency=concurrency, + ) + ) if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Realtime WebSocket Transcription Client") - parser.add_argument( - "--model", - type=str, - default="Qwen/Qwen3-Omni-30B-A3B-Instruct", - help="Model that is served and should be pinged.", - ) - parser.add_argument( - "--audio_path", - type=str, - default=None, - help="Path to the audio file to transcribe.", - ) - parser.add_argument( - "--host", - type=str, - default="localhost", - help="vLLM-Omni server host (default: localhost)", - ) - parser.add_argument( - "--port", - type=int, - default=8000, - help="vLLM-Omni server port (default: 8000)", - ) - args = parser.parse_args() - main(args) + main() diff --git a/examples/online_serving/qwen3_tts/README.md b/examples/online_serving/qwen3_tts/README.md index b48db9cf4..350fcb71c 100644 --- a/examples/online_serving/qwen3_tts/README.md +++ b/examples/online_serving/qwen3_tts/README.md @@ -43,7 +43,7 @@ Then open http://localhost:7860 in your browser. ### Launch the Server -The default stage config is located at `vllm_omni/model_executor/stage_configs/qwen3_tts.yaml`. For other platforms (e.g., NPU), refer to `vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml`. +The default deploy config is located at `vllm_omni/deploy/qwen3_tts.yaml` and is loaded automatically by the model registry — no `--deploy-config` flag needed for default use. Platform-specific deltas (NPU, ROCm, XPU) are merged in automatically from the `platforms:` block of the same YAML based on the detected runtime. ```bash # CustomVoice model (predefined speakers) @@ -70,6 +70,22 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ --port 8091 ``` +#### Sync vs async-chunk mode + +Qwen3-TTS supports both **chunked streaming** (default, lower latency) and +**synchronous end-to-end** modes from the same deploy YAML. The bundled +`qwen3_tts.yaml` ships with `async_chunk: true`; flip with `--no-async-chunk` +and the pipeline automatically dispatches to the end-to-end codec processor: + +```bash +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091 \ + --no-async-chunk +``` + +No variant YAML or extra flag is needed — the `StagePipelineConfig` on each +stage declares both processor functions and the runtime picks based on the +`async_chunk:` bool. + Alternatively, use the convenience script: ```bash ./run_server.sh # Default: CustomVoice model diff --git a/examples/online_serving/qwen3_tts/batch_speech_client.py b/examples/online_serving/qwen3_tts/batch_speech_client.py index 7d48e650f..47fdc3691 100644 --- a/examples/online_serving/qwen3_tts/batch_speech_client.py +++ b/examples/online_serving/qwen3_tts/batch_speech_client.py @@ -5,11 +5,13 @@ A particularly useful scenario is voice cloning: set ref_audio once at the batch level and generate many utterances in the cloned voice without repeating the reference for each item. -Start the server (with batch-optimized config for best throughput): +Start the server (with batch-optimized stage settings for best throughput): vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ - --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml \ - --trust-remote-code + --omni \ + --trust-remote-code \ + --stage-overrides '{"0":{"max_num_seqs":4,"gpu_memory_utilization":0.2}, + "1":{"max_num_seqs":4,"gpu_memory_utilization":0.2}}' Examples: # Batch with a predefined voice diff --git a/examples/online_serving/qwen3_tts/run_gradio_demo.sh b/examples/online_serving/qwen3_tts/run_gradio_demo.sh index bcc0ddb7c..d79be3c2a 100644 --- a/examples/online_serving/qwen3_tts/run_gradio_demo.sh +++ b/examples/online_serving/qwen3_tts/run_gradio_demo.sh @@ -127,7 +127,7 @@ echo "Starting vLLM server..." LOG_FILE="/tmp/vllm_tts_server_${SERVER_PORT}.log" vllm-omni serve "$MODEL" \ - --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \ + --deploy-config vllm_omni/deploy/qwen3_tts.yaml \ --host "$SERVER_HOST" \ --port "$SERVER_PORT" \ --gpu-memory-utilization 0.9 \ diff --git a/examples/online_serving/qwen3_tts/run_server.sh b/examples/online_serving/qwen3_tts/run_server.sh index 6f4aa83a0..78dd2c305 100755 --- a/examples/online_serving/qwen3_tts/run_server.sh +++ b/examples/online_serving/qwen3_tts/run_server.sh @@ -31,7 +31,7 @@ esac echo "Starting Qwen3-TTS server with model: $MODEL" vllm-omni serve "$MODEL" \ - --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \ + --deploy-config vllm_omni/deploy/qwen3_tts.yaml \ --host 0.0.0.0 \ --port 8091 \ --gpu-memory-utilization 0.9 \ diff --git a/examples/online_serving/qwen3_tts/speaker_embedding_interpolation.py b/examples/online_serving/qwen3_tts/speaker_embedding_interpolation.py index 38a2bdea9..7790fa512 100644 --- a/examples/online_serving/qwen3_tts/speaker_embedding_interpolation.py +++ b/examples/online_serving/qwen3_tts/speaker_embedding_interpolation.py @@ -5,7 +5,7 @@ speaker encoder from a Qwen3-TTS checkpoint, then interpolates between them using SLERP and sends the result to the /v1/audio/speech API. Requirements: - pip install torch resampy soundfile numpy httpx + pip install torch soundfile numpy httpx Examples: # Extract and save an embedding @@ -143,11 +143,12 @@ def _load_speaker_encoder_weights(encoder: torch.nn.Module, model_path: str) -> def compute_mel_spectrogram(audio: np.ndarray, sr: int = 24000) -> torch.Tensor: """Compute 128-bin mel spectrogram matching Qwen3-TTS's extraction pipeline.""" - from vllm.multimodal.audio import resample_audio_resampy + from vllm.multimodal.audio import AudioResampler # Resample to 24kHz if needed if sr != 24000: - audio = resample_audio_resampy(audio.astype(np.float32), orig_sr=sr, target_sr=24000) + resampler = AudioResampler(target_sr=24000) + audio = resampler.resample(audio.astype(np.float32), orig_sr=sr) y = torch.from_numpy(audio).unsqueeze(0).float() diff --git a/examples/online_serving/voxcpm/README.md b/examples/online_serving/voxcpm/README.md new file mode 100644 index 000000000..78e1bf4aa --- /dev/null +++ b/examples/online_serving/voxcpm/README.md @@ -0,0 +1,166 @@ +# VoxCPM + +## Prerequisites + +Install VoxCPM in one of these ways: + +```bash +pip install voxcpm +``` + +or point vLLM-Omni to a local VoxCPM source tree: + +```bash +export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src +``` + +If the native VoxCPM `config.json` lacks HF metadata such as `model_type`, +prepare a persistent HF-compatible config directory and export: + +```bash +export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config +mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH" +cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json" +cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true +python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)' +``` + +The VoxCPM stage configs read `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH` directly. The `python3 -c` form above avoids heredoc/indentation issues in interactive shells. + +## Launch the Server + +Use the async-chunk stage config by default: + +```bash +export VOXCPM_MODEL=/path/to/voxcpm-model +cd examples/online_serving/voxcpm +./run_server.sh +``` + +Use the non-streaming stage config: + +```bash +./run_server.sh sync +``` + +You can also launch the server directly: + +```bash +vllm serve "$VOXCPM_MODEL" \ + --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \ + --trust-remote-code \ + --enforce-eager \ + --omni \ + --port 8091 +``` + +## Send Requests + +### Basic text-to-speech + +```bash +python openai_speech_client.py \ + --model "$VOXCPM_MODEL" \ + --text "This is a VoxCPM online text-to-speech example." +``` + +### Voice cloning + +```bash +python openai_speech_client.py \ + --model "$VOXCPM_MODEL" \ + --text "This sentence is synthesized with a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "The exact transcript spoken in reference.wav." +``` + +`ref_text` must be the real transcript of the reference audio. Placeholder text or mismatched text will usually degrade quality badly. + +### Streaming PCM output + +```bash +python openai_speech_client.py \ + --model "$VOXCPM_MODEL" \ + --text "This is a streaming VoxCPM request." \ + --stream \ + --output voxcpm_stream.pcm +``` + +### Using curl + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "OpenBMB/VoxCPM1.5", + "input": "Hello from VoxCPM online serving.", + "response_format": "wav" + }' --output output.wav +``` + +Voice cloning: + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "OpenBMB/VoxCPM1.5", + "input": "This sentence uses a cloned voice.", + "ref_audio": "https://example.com/reference.wav", + "ref_text": "The exact transcript spoken in the reference audio.", + "response_format": "wav" + }' --output cloned.wav +``` + +Streaming PCM: + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "OpenBMB/VoxCPM1.5", + "input": "This is a streaming VoxCPM request.", + "stream": true, + "response_format": "pcm" + }' --output output.pcm +``` + +## Supported Request Shape + +VoxCPM online serving currently supports: + +- plain text-to-speech +- voice cloning with `ref_audio` + `ref_text` +- `stream=true` with `response_format=pcm` or `wav` + +VoxCPM online serving does not use these generic TTS fields: + +- `voice` +- `instructions` +- `language` +- `speaker_embedding` +- `x_vector_only_mode` + +## Streaming vs Non-Streaming + +- `voxcpm_async_chunk.yaml` enables async-chunk streaming and is best for single-request streaming latency. +- `voxcpm.yaml` performs one-shot latent generation then VAE decode. + +Like native VoxCPM, the async streaming path should be treated as single-request. If you need stable throughput benchmarking, prefer `voxcpm.yaml`. + +Do not use `voxcpm_async_chunk.yaml` for concurrent online streaming or `/v1/audio/speech/batch`. For multiple requests, prefer `voxcpm.yaml`. + +## Benchmark + +The serving benchmark reports TTFP and RTF: + +```bash +python benchmarks/voxcpm/vllm_omni/bench_tts_serve.py \ + --host 127.0.0.1 \ + --port 8091 \ + --num-prompts 10 \ + --max-concurrency 1 \ + --result-dir /tmp/voxcpm_bench +``` + +For the async-chunk server, keep `--max-concurrency 1`. diff --git a/examples/online_serving/voxcpm/openai_speech_client.py b/examples/online_serving/voxcpm/openai_speech_client.py new file mode 100644 index 000000000..c400114e8 --- /dev/null +++ b/examples/online_serving/voxcpm/openai_speech_client.py @@ -0,0 +1,155 @@ +"""OpenAI-compatible client for VoxCPM via /v1/audio/speech. + +Examples: + # Basic text-to-speech + python openai_speech_client.py --text "Hello from VoxCPM" + + # Voice cloning + python openai_speech_client.py \ + --text "This sentence uses the cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "The exact transcript spoken in the reference audio." + + # Streaming PCM output + python openai_speech_client.py \ + --text "This is a streaming VoxCPM request." \ + --stream \ + --output output.pcm +""" + +import argparse +import base64 +import os + +import httpx + +DEFAULT_API_BASE = "http://localhost:8091" +DEFAULT_API_KEY = "EMPTY" +DEFAULT_MODEL = "OpenBMB/VoxCPM1.5" + + +def encode_audio_to_base64(audio_path: str) -> str: + """Encode a local audio file to base64 data URL.""" + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + + ext = audio_path.lower().rsplit(".", 1)[-1] + mime_map = { + "wav": "audio/wav", + "mp3": "audio/mpeg", + "flac": "audio/flac", + "ogg": "audio/ogg", + } + mime_type = mime_map.get(ext, "audio/wav") + + with open(audio_path, "rb") as f: + audio_b64 = base64.b64encode(f.read()).decode("utf-8") + return f"data:{mime_type};base64,{audio_b64}" + + +def build_payload(args) -> dict[str, object]: + payload: dict[str, object] = { + "model": args.model, + "input": args.text, + "response_format": "pcm" if args.stream else args.response_format, + } + + if args.ref_audio: + if args.ref_audio.startswith(("http://", "https://", "data:")): + payload["ref_audio"] = args.ref_audio + else: + payload["ref_audio"] = encode_audio_to_base64(args.ref_audio) + if args.ref_text: + payload["ref_text"] = args.ref_text + if args.max_new_tokens is not None: + payload["max_new_tokens"] = args.max_new_tokens + if args.stream: + payload["stream"] = True + + return payload + + +def run_tts(args) -> None: + payload = build_payload(args) + api_url = f"{args.api_base}/v1/audio/speech" + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {args.api_key}", + } + + print(f"Model: {args.model}") + print(f"Text: {args.text}") + if args.ref_audio: + print("Mode: voice cloning") + print(f"Reference audio: {args.ref_audio}") + else: + print("Mode: text-to-speech") + + if args.stream: + output_path = args.output or "voxcpm_output.pcm" + with httpx.Client(timeout=300.0) as client: + with client.stream("POST", api_url, json=payload, headers=headers) as response: + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.read().decode("utf-8", errors="ignore")) + return + + total_bytes = 0 + with open(output_path, "wb") as f: + for chunk in response.iter_bytes(): + if not chunk: + continue + f.write(chunk) + total_bytes += len(chunk) + print(f"Streamed {total_bytes} bytes to: {output_path}") + return + + with httpx.Client(timeout=300.0) as client: + response = client.post(api_url, json=payload, headers=headers) + + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.text) + return + + try: + text = response.content.decode("utf-8") + if text.startswith('{"error"'): + print(f"Error: {text}") + return + except UnicodeDecodeError: + pass + + output_path = args.output or "voxcpm_output.wav" + with open(output_path, "wb") as f: + f.write(response.content) + print(f"Audio saved to: {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description="VoxCPM OpenAI-compatible speech client") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name or path") + parser.add_argument("--text", required=True, help="Text to synthesize") + parser.add_argument("--ref-audio", default=None, help="Reference audio path, URL, or data URL") + parser.add_argument( + "--ref-text", + default=None, + help="The exact transcript spoken in the reference audio", + ) + parser.add_argument("--stream", action="store_true", help="Enable streaming PCM output") + parser.add_argument( + "--response-format", + default="wav", + choices=["wav", "pcm", "flac", "mp3", "aac", "opus"], + help="Audio format for non-streaming mode (default: wav)", + ) + parser.add_argument("--max-new-tokens", type=int, default=None, help="Maximum tokens to generate") + parser.add_argument("--output", "-o", default=None, help="Output file path") + args = parser.parse_args() + run_tts(args) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/voxcpm/run_server.sh b/examples/online_serving/voxcpm/run_server.sh new file mode 100755 index 000000000..ab4b6fe85 --- /dev/null +++ b/examples/online_serving/voxcpm/run_server.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# Launch vLLM-Omni server for VoxCPM online speech serving. +# +# Usage: +# ./run_server.sh # default: async_chunk stage config +# ./run_server.sh async # async_chunk stage config +# ./run_server.sh sync # no-async-chunk stage config +# VOXCPM_MODEL=/path/to/model ./run_server.sh + +set -e + +MODE="${1:-async}" +MODEL="${VOXCPM_MODEL:-OpenBMB/VoxCPM1.5}" + +case "$MODE" in + async) + STAGE_CONFIG="vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml" + ;; + sync) + STAGE_CONFIG="vllm_omni/model_executor/stage_configs/voxcpm.yaml" + ;; + *) + echo "Unknown mode: $MODE" + echo "Supported: async, sync" + exit 1 + ;; +esac + +echo "Starting VoxCPM server with model: $MODEL" +echo "Stage config: $STAGE_CONFIG" + +vllm serve "$MODEL" \ + --stage-configs-path "$STAGE_CONFIG" \ + --host 0.0.0.0 \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager \ + --omni diff --git a/examples/online_serving/voxcpm2/README.md b/examples/online_serving/voxcpm2/README.md index 8735180f0..9ca2ae708 100644 --- a/examples/online_serving/voxcpm2/README.md +++ b/examples/online_serving/voxcpm2/README.md @@ -5,12 +5,13 @@ Serve VoxCPM2 TTS via the OpenAI-compatible `/v1/audio/speech` endpoint. ## Start the Server ```bash -python -m vllm_omni.entrypoints.openai.api_server \ - --model openbmb/VoxCPM2 \ - --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm2.yaml \ - --host 0.0.0.0 --port 8000 +vllm serve openbmb/VoxCPM2 --omni --host 0.0.0.0 --port 8000 ``` +The deploy config is auto-loaded from `vllm_omni/deploy/voxcpm2.yaml`. Pass +`--deploy-config ` to override, or `--stage-N- ` (e.g. +`--stage-0-max-num-seqs 8`) for per-stage runtime overrides. + ## Zero-shot Synthesis ```bash diff --git a/examples/online_serving/voxcpm2/gradio_demo.py b/examples/online_serving/voxcpm2/gradio_demo.py new file mode 100644 index 000000000..c6706198a --- /dev/null +++ b/examples/online_serving/voxcpm2/gradio_demo.py @@ -0,0 +1,599 @@ +"""Gradio demo for VoxCPM2 TTS with gapless streaming audio playback. + +Uses a custom AudioWorklet-based player for gap-free streaming +(adapted from the Qwen3-TTS demo). Audio is streamed from the vLLM +server through a same-origin proxy and played via the Web Audio API's +AudioWorklet, which maintains a FIFO buffer queue and plays samples at +the audio clock rate. + +Usage: + # Start the vLLM server first: + vllm serve openbmb/VoxCPM2 --omni --host 0.0.0.0 --port 8000 + + # Then launch the demo: + python gradio_demo.py --api-base http://localhost:8000 +""" + +from __future__ import annotations + +import argparse +import base64 +import io +import json +import logging + +import gradio as gr +import httpx +import numpy as np +import soundfile as sf +from fastapi import FastAPI, Request +from fastapi.responses import Response, StreamingResponse + +logger = logging.getLogger(__name__) + +SAMPLE_RATE = 48000 + +# ── AudioWorklet processor (loaded in browser via Blob URL) ────────── +WORKLET_JS = r""" +class TTSPlaybackProcessor extends AudioWorkletProcessor { + constructor() { + super(); + this.queue = []; + this.buf = null; + this.pos = 0; + this.playing = false; + this.played = 0; + this.port.onmessage = (e) => { + if (e.data && e.data.type === 'clear') { + this.queue = []; this.buf = null; this.pos = 0; this.played = 0; + if (this.playing) { this.playing = false; this.port.postMessage({type:'stopped'}); } + return; + } + this.queue.push(e.data); + }; + } + process(inputs, outputs) { + const out = outputs[0][0]; + for (let i = 0; i < out.length; i++) { + if (!this.buf || this.pos >= this.buf.length) { + if (this.queue.length > 0) { + this.buf = this.queue.shift(); this.pos = 0; + } else { + for (let j = i; j < out.length; j++) out[j] = 0; + if (this.playing) { this.playing = false; this.port.postMessage({type:'stopped', played:this.played}); } + return true; + } + } + out[i] = this.buf[this.pos++] / 32768; + this.played++; + } + if (!this.playing) { this.playing = true; this.port.postMessage({type:'started'}); } + return true; + } +} +registerProcessor('tts-playback-processor', TTSPlaybackProcessor); +""" + +PLAYER_HTML = """ +
+
+
+ Ready + +
+ + + +
+""" + + +def _build_player_js() -> str: + return f""" + +""" + + +def _encode_audio(audio_data: tuple) -> str: + sr, audio_np = audio_data + if audio_np.dtype in (np.float32, np.float64): + audio_np = np.clip(audio_np, -1.0, 1.0) + audio_np = (audio_np * 32767).astype(np.int16) + elif audio_np.dtype != np.int16: + audio_np = audio_np.astype(np.int16) + buf = io.BytesIO() + sf.write(buf, audio_np, sr, format="WAV") + return f"data:audio/wav;base64,{base64.b64encode(buf.getvalue()).decode()}" + + +def create_app(api_base: str): + app = FastAPI() + _pending: dict[str, dict] = {} + + @app.post("/proxy/v1/audio/speech") + async def proxy_speech(request: Request): + body = await request.json() + req_id = body.get("_req_id") + if req_id and req_id in _pending: + body = _pending.pop(req_id) + logger.info("Proxy: %s", {k: (f"<{len(str(v))} chars>" if k == "ref_audio" else v) for k, v in body.items()}) + try: + client = httpx.AsyncClient(timeout=300) + resp = await client.send( + client.build_request( + "POST", + f"{api_base}/v1/audio/speech", + json=body, + headers={"Authorization": "Bearer EMPTY", "Content-Type": "application/json"}, + ), + stream=True, + ) + except Exception as exc: + logger.exception("Proxy connection error") + await client.aclose() + return Response(content=str(exc), status_code=502) + if resp.status_code != 200: + content = await resp.aread() + await resp.aclose() + await client.aclose() + return Response(content=content, status_code=resp.status_code) + + async def relay(): + try: + async for chunk in resp.aiter_bytes(): + yield chunk + finally: + await resp.aclose() + await client.aclose() + + return StreamingResponse(relay(), media_type="application/octet-stream") + + css = """ + #generate-btn button { width: 100%; } + #streaming-player { border: 1px solid var(--border-color-primary) !important; border-radius: var(--block-radius) !important; padding: var(--block-padding) !important; } + """ + theme = gr.themes.Default( + primary_hue=gr.themes.Color( + c50="#f0f5ff", + c100="#dce6f9", + c200="#b8cef3", + c300="#8eb2eb", + c400="#6496e0", + c500="#4A90D9", + c600="#3a7bc8", + c700="#2d66b0", + c800="#1f4f8f", + c900="#163a6e", + c950="#0e2650", + ), + ) + + with gr.Blocks(title="VoxCPM2 TTS Demo") as demo: + gr.HTML(f""" +
+ vLLM-Omni +
+

VoxCPM2 Streaming Demo

+ + Served by vLLM-Omni + · {api_base} + · 48 kHz + +
+
+ """) + + gr.Markdown( + "**Three modes:** " + "**Voice Design** (control instruction only) · " + "**Controllable Cloning** (ref audio + optional style control) · " + "**Ultimate Cloning** (ref audio + transcript for audio continuation)" + ) + + with gr.Row(): + with gr.Column(scale=3): + text_input = gr.Textbox( + label="Target Text", + placeholder="Enter text to synthesize...", + lines=4, + ) + control_instruction = gr.Textbox( + label="Control Instruction (optional)", + placeholder="e.g. A warm young woman / Excited and fast-paced", + lines=2, + info="Describe voice style, emotion, pace. Works for both Voice Design and Controllable Cloning.", + ) + + with gr.Accordion("Voice Cloning", open=False): + ref_audio = gr.Audio( + label="Reference Audio (upload for cloning)", + type="numpy", + sources=["upload", "microphone"], + ) + ref_audio_url = gr.Textbox( + label="or Reference Audio URL", + placeholder="https://example.com/reference.wav", + ) + ultimate_clone = gr.Checkbox( + label="Ultimate Cloning Mode", + value=False, + info="Provide transcript of ref audio for audio continuation (disables control instruction)", + ) + prompt_text = gr.Textbox( + label="Reference Audio Transcript", + placeholder="Transcript of your reference audio (for ultimate cloning)", + lines=2, + visible=False, + ) + + with gr.Row(): + stream_checkbox = gr.Checkbox( + label="Stream (gapless)", + value=True, + info="AudioWorklet streaming", + ) + with gr.Row(): + generate_btn = gr.Button( + "Generate Speech", + variant="primary", + size="lg", + elem_id="generate-btn", + scale=3, + ) + reset_btn = gr.Button("Reset", variant="secondary", size="lg", scale=1) + + with gr.Column(scale=2): + player_html = gr.HTML( + value=PLAYER_HTML, + visible=True, + label="streaming player", + elem_id="streaming-player", + ) + audio_output = gr.Audio( + label="generated audio", + interactive=False, + autoplay=True, + visible=False, + ) + gr.Examples( + examples=[ + ["Hello, this is a VoxCPM2 demo running on vLLM-Omni.", ""], + [ + "I have a dream that my four little children will one day live in a nation " + "where they will not be judged by the color of their skin but by the content " + "of their character.", + "", + ], + [ + "I never asked you to stay. It's not like I care or anything. " + "But why does it still hurt so much now that you're gone?", + "A young girl with a soft, sweet voice. Speaks slowly with a melancholic tone.", + ], + ], + inputs=[text_input, control_instruction], + label="examples", + ) + gr.HTML(""" +
+ + vLLM-Omni + +
+ """) + + hidden_payload = gr.Textbox(visible=False, elem_id="tts-payload") + + def on_ultimate_toggle(checked): + return ( + gr.update(visible=checked), # prompt_text + gr.update(interactive=not checked), # control_instruction + ) + + ultimate_clone.change( + fn=on_ultimate_toggle, + inputs=[ultimate_clone], + outputs=[prompt_text, control_instruction], + ) + + def on_stream_change(stream: bool): + if stream: + return gr.update(visible=True), gr.update(visible=False) + return gr.update(visible=False), gr.update(visible=True) + + stream_checkbox.change( + fn=on_stream_change, + inputs=[stream_checkbox], + outputs=[player_html, audio_output], + ) + + def on_reset(): + return "", "", None, "", False, "", PLAYER_HTML + + reset_btn.click( + fn=on_reset, + outputs=[ + text_input, + control_instruction, + audio_output, + hidden_payload, + ultimate_clone, + prompt_text, + player_html, + ], + js="() => { if (window.ttsStop) window.ttsStop(); }", + ) + + def on_generate(stream_enabled, text, ctrl_instr, ref_a, ref_url, ult_clone, p_text): + import time as _time + + if not text or not text.strip(): + raise gr.Error("Please enter text to synthesize.") + + # VoxCPM2 uses "(instruction)text" format for control + ctrl = ctrl_instr.strip() if ctrl_instr and not ult_clone else "" + final_text = f"({ctrl}){text.strip()}" if ctrl else text.strip() + + payload: dict = { + "input": final_text, + "voice": "default", + "response_format": "pcm" if stream_enabled else "wav", + "stream": stream_enabled, + } + + # Reference audio for cloning + ref_url_s = ref_url.strip() if ref_url else "" + if ref_url_s: + payload["ref_audio"] = ref_url_s + elif ref_a is not None: + payload["ref_audio"] = _encode_audio(ref_a) + + # Ultimate cloning: prompt_audio + prompt_text for continuation + if ult_clone and p_text and p_text.strip(): + if ref_url_s: + payload["prompt_audio"] = ref_url_s + elif ref_a is not None: + payload["prompt_audio"] = payload.get("ref_audio", "") + payload["prompt_text"] = p_text.strip() + + if stream_enabled: + if ref_a is not None and not ref_url_s: + req_id = f"req-{int(_time.time() * 1000)}" + _pending[req_id] = payload + browser_payload = {"_req_id": req_id, "_nonce": int(_time.time() * 1000)} + return json.dumps(browser_payload), gr.update() + payload["_nonce"] = int(_time.time() * 1000) + return json.dumps(payload), gr.update() + else: + try: + with httpx.Client(timeout=300.0) as client: + resp = client.post( + f"{api_base}/v1/audio/speech", + json=payload, + headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"}, + ) + except httpx.ConnectError: + raise gr.Error(f"Cannot connect to server at {api_base}.") + if resp.status_code != 200: + raise gr.Error(f"Server error ({resp.status_code}): {resp.text[:200]}") + audio_np, sr = sf.read(io.BytesIO(resp.content)) + if audio_np.ndim > 1: + audio_np = audio_np[:, 0] + return "", (sr, audio_np.astype(np.float32)) + + generate_btn.click( + fn=on_generate, + inputs=[ + stream_checkbox, + text_input, + control_instruction, + ref_audio, + ref_audio_url, + ultimate_clone, + prompt_text, + ], + outputs=[hidden_payload, audio_output], + ).then( + fn=lambda p: p, + inputs=[hidden_payload], + outputs=[hidden_payload], + js="(p) => { if (p && p.trim()) { const d = JSON.parse(p); delete d._nonce; window.ttsGenerate(d); } return p; }", + ) + + demo.queue() + + return gr.mount_gradio_app(app, demo, path="/", css=css, theme=theme, head=_build_player_js()) + + +def main(): + parser = argparse.ArgumentParser(description="VoxCPM2 streaming Gradio demo") + parser.add_argument("--api-base", default="http://localhost:8000", help="vLLM API server URL") + parser.add_argument("--host", default="0.0.0.0", help="Gradio server host") + parser.add_argument("--port", type=int, default=7860, help="Gradio server port") + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO) + print(f"Connecting to vLLM server at: {args.api_base}") + + import uvicorn + + uvicorn.run(create_app(args.api_base), host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/voxcpm2/openai_speech_client.py b/examples/online_serving/voxcpm2/openai_speech_client.py index a117d24fd..127b8cebb 100644 --- a/examples/online_serving/voxcpm2/openai_speech_client.py +++ b/examples/online_serving/voxcpm2/openai_speech_client.py @@ -13,10 +13,7 @@ Examples: --ref-audio "https://example.com/reference.wav" Server setup: - python -m vllm_omni.entrypoints.openai.api_server \ - --model openbmb/VoxCPM2 \ - --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm2.yaml \ - --host 0.0.0.0 --port 8000 + vllm serve openbmb/VoxCPM2 --omni --host 0.0.0.0 --port 8000 """ from __future__ import annotations diff --git a/mkdocs.yml b/mkdocs.yml index 6461c65f2..1e184439b 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -97,6 +97,7 @@ plugins: exclude: - "re:vllm_omni\\._.*" # Internal modules - "vllm_omni.diffusion.models.qwen_image" # avoid importing vllm in mkdocs building + - "vllm_omni.diffusion.models.dreamid_omni.wan2_2" # docstring signature warnings break strict docs - "vllm_omni.diffusion.quantization" # avoid importing vllm in mkdocs building - "vllm_omni.quantization" # avoid importing vllm in mkdocs building - "vllm_omni.entrypoints.async_diffusion" # avoid importing vllm in mkdocs building diff --git a/pyproject.toml b/pyproject.toml index 57a4b474f..b6f4092fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,12 +55,24 @@ dev = [ "pyttsx3>=2.99", "opencc>=1.2.0", "mistune>=3.2.0", # for example tests + "torchmetrics>=1.4.0", # for accuracy similarity metrics ] demo = [ "gradio>=6.7.0", ] +# Seed-TTS serve benchmark WER (BytedanceSpeech/seed-tts-eval run_wer.py protocol). +seed-tts-eval = [ + "jiwer>=3.0.0", + "zhon>=2.0.0", + "zhconv>=1.4.2", + "scipy>=1.10.0", + "soundfile>=0.12.0", + "transformers>=4.36.0", + "funasr>=1.0.0", +] + docs = [ "mkdocs>=1.5.0", "mkdocs-api-autonav", @@ -115,12 +127,13 @@ exclude = [ [tool.ruff.lint] select = [ - "E", # pycodestyle errors - "W", # pycodestyle warnings - "F", # pyflakes - "I", # isort (handled separately, but included for compatibility) - "N", # pep8-naming - "UP", # pyupgrade + "E", # pycodestyle errors + "W", # pycodestyle warnings + "F", # pyflakes + "I", # isort (handled separately, but included for compatibility) + "N", # pep8-naming + "UP", # pyupgrade + "TID251", # flake8-tidy-imports.banned-api ] ignore = [ "E203", # whitespace before ':' (conflicts with black) @@ -135,6 +148,9 @@ ignore = [ "examples/**" = ["E501"] # Allow long lines in examples "tests/**" = ["E501"] # Allow long lines in tests +[tool.ruff.lint.flake8-tidy-imports.banned-api] +"librosa".msg = "The librosa module is banned, use vllm.multimodal helpers instead" + [tool.mypy] python_version = "3.12, 3.13" warn_return_any = true @@ -225,3 +241,4 @@ ue = "ue" semantics = "semantics" fullset = "fullset" Vai = "Vai" +CANN = "CANN" diff --git a/recipes/Qwen/Qwen3-Omni.md b/recipes/Qwen/Qwen3-Omni.md new file mode 100644 index 000000000..081e1453d --- /dev/null +++ b/recipes/Qwen/Qwen3-Omni.md @@ -0,0 +1,90 @@ +# Qwen3-Omni for multimodal chat on 1x A100 80GB + +## Summary + +- Vendor: Qwen +- Model: `Qwen/Qwen3-Omni-30B-A3B-Instruct` +- Task: Multimodal chat with text, image, audio, or video input +- Mode: Online serving with the OpenAI-compatible API +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want a known-good starting point for serving +`Qwen/Qwen3-Omni-30B-A3B-Instruct` with vLLM-Omni on a single 80 GB A100 and +validate the deployment with the existing multimodal client examples in this +repository. + +## References + +- Upstream or canonical docs: + [`docs/user_guide/examples/online_serving/qwen3_omni.md`](../../docs/user_guide/examples/online_serving/qwen3_omni.md) +- Related example under `examples/`: + [`examples/online_serving/qwen3_omni/README.md`](../../examples/online_serving/qwen3_omni/README.md) +- Related issue or discussion: + [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645) + +## Hardware Support + +This recipe currently documents one tested-style reference configuration for +CUDA GPU serving. Add more sections for other hardware as community validation +lands. + +## GPU + +### 1x A100 80GB + +#### Environment + +- OS: Linux +- Python: 3.10+ +- Driver / runtime: NVIDIA CUDA environment with an A100 80 GB GPU +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Command + +Start the server from the repository root: + +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 +``` + +To enable async chunking, use the bundled stage config: + +```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 +``` + +#### Verification + +Run one of the existing example clients after the server is ready: + +```bash +python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \ + --model Qwen/Qwen3-Omni-30B-A3B-Instruct \ + --query-type use_image \ + --port 8091 \ + --host localhost +``` + +For a quick API smoke test, request text-only output: + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "messages": [{"role": "user", "content": "Describe vLLM in brief."}], + "modalities": ["text"] + }' +``` + +#### Notes + +- Memory usage: Size depends on runtime options and output modalities; leave headroom for multimodal workloads. +- Key flags: `--omni` is required; `--stage-configs-path` is optional for custom or async-chunk stage configs. +- Known limitations: This starter recipe is intentionally narrow and focuses on the single-GPU online-serving path already documented in the repo examples. diff --git a/recipes/README.md b/recipes/README.md new file mode 100644 index 000000000..01ecc41f1 --- /dev/null +++ b/recipes/README.md @@ -0,0 +1,37 @@ +# Community Recipes + +This directory contains community-maintained recipes for answering a +practical user question: + +> How do I run model X on hardware Y for task Z? + +Add recipes for this repository under this in-repo `recipes/` directory. To +keep naming and layout consistent, organize recipes by model vendor in a way +that is aligned with +[`vllm-project/recipes`](https://github.com/vllm-project/recipes), but treat +that external repository as a reference for structure rather than the place to +add files for this repo. Use one Markdown file per model family by default. + +Example layout: + +```text +recipes/ + Qwen/ + Qwen3-Omni.md + Qwen3-TTS.md + Tencent-Hunyuan/ + HunyuanVideo.md +``` + +## Available Recipes + +- [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md): online serving recipe for + multimodal chat on `1x A100 80GB` +- [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md): image-to-video serving + recipe for Wan2.2 14B on `8x Ascend NPU (A2/A3)` + +Within a single recipe file, include different hardware support sections such +as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like +`1x A100 80GB` or `2x L40S` inside those sections when applicable. + +See [TEMPLATE.md](./TEMPLATE.md) for the recommended format. diff --git a/recipes/TEMPLATE.md b/recipes/TEMPLATE.md new file mode 100644 index 000000000..9bf8cb9c7 --- /dev/null +++ b/recipes/TEMPLATE.md @@ -0,0 +1,82 @@ +# Recipe Title + +> Example: Qwen3-Omni for speech chat on 1x A100 80GB + +## Summary + +- Vendor: +- Model: +- Task: +- Mode: +- Maintainer: + +## When to use this recipe + +Briefly describe the concrete scenario this recipe covers. + +## References + +- Upstream or canonical docs: +- Related example under `examples/`: +- Related issue or discussion: + +## Hardware Support + +Add one section per platform, such as `GPU`, `ROCm`, or `NPU`. Under each +platform section, document one or more tested hardware configurations. + +## GPU + +### 1x A100 80GB + +#### Environment + +- OS: +- Python: +- Driver / runtime: +- vLLM version: +- vLLM-Omni version or commit: + +#### Command + +```bash +# Add the exact command(s) here +``` + +#### Verification + +```bash +# Add a quick validation command or expected output here +``` + +#### Notes + +- Memory usage: +- Key flags: +- Known limitations: + +### 2x L40S + +Repeat the same structure for other hardware setups as needed. + +## ROCm + +### Example hardware configuration + +Repeat the same nested structure for ROCm setups as needed: + +- `#### Environment` +- `#### Command` +- `#### Verification` +- `#### Notes` + +## NPU + +### Example hardware configuration + +Repeat the same nested structure for NPU setups as needed: + +- `#### Environment` +- `#### Command` +- `#### Verification` +- `#### Notes` diff --git a/recipes/Wan-AI/Wan2.2-I2V.md b/recipes/Wan-AI/Wan2.2-I2V.md new file mode 100644 index 000000000..99ceac3ce --- /dev/null +++ b/recipes/Wan-AI/Wan2.2-I2V.md @@ -0,0 +1,136 @@ +# Wan2.2 Image To Video + +## Summary + +- Vendor: Wan-AI +- Model: `Wan-AI/Wan2.2-I2V-A14B-Diffusers` +- Task: Image-to-video generation +- Mode: Online serving with the OpenAI-compatible API +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want to deploy the Wan2.2 14B image-to-video model +with vLLM-Omni using multi-card parallelism. Two configurations are provided: + +1. **Distilled model (no negative-prompt / CFG computation)** — higher + throughput, recommended when using a distilled checkpoint that does not + require classifier-free guidance. +2. **Official open-source model (with CFG)** — uses `--cfg 2` to run negative + and positive samples in parallel for the original released weights. + +## References + +- Upstream model card: + +## Hardware Support + +## NPU + +### 8x Ascend A2 / A3 + +#### Environment + +- OS: Linux +- Python: 3.10+ +- Driver / runtime: Ascend NPU driver with CANN toolkit +- Recommended operator library: **mindie-sd** (Ascend high-performance fused + operators — enables `adalayernorm` and other fused kernels automatically upon + installation) +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Prerequisites + +Install the **mindie-sd** operator library to enable Ascend-optimized fused +operators (`adalayernorm`, etc.): + +```bash +git clone https://gitcode.com/Ascend/MindIE-SD.git && cd MindIE-SD + +# Comment out the tik_ops build step (not needed for this use case) +sed -i 's|^\(\s*\)source ${current_script_dir}/build_tik_ops.sh|\1# source ${current_script_dir}/build_tik_ops.sh|' build/build_ops.sh + +python setup.py bdist_wheel +cd dist +pip install mindiesd-*.whl +``` + +After installation, enable the Laser Attention kernel for significant +long-sequence speedups (up to ~40% at 720p in tested workloads): + +```bash +export MINDIE_SD_FA_TYPE=ascend_laser_attention +``` + +When using HSDP with FSDP2, set the following environment variable to work +around a PyTorch NPU multi-stream memory reuse issue +([pytorch/pytorch#147168](https://github.com/pytorch/pytorch/issues/147168)). +This issue has been fixed on CUDA but still applies to NPU: + +```bash +export MULTI_STREAM_MEMORY_REUSE=2 +``` + +#### Command + +**Distilled model (no CFG, recommended for distilled checkpoints):** + +```bash +export MINDIE_SD_FA_TYPE=ascend_laser_attention +export MULTI_STREAM_MEMORY_REUSE=2 + +vllm serve \ + --omni Wan-AI/Wan2.2-I2V-A14B-Diffusers \ + --use-hsdp \ + --usp 8 \ + --vae-patch-parallel-size 8 \ + --vae-use-tiling +``` + +**Official open-source model (with CFG):** + +```bash +export MINDIE_SD_FA_TYPE=ascend_laser_attention +export MULTI_STREAM_MEMORY_REUSE=2 + +vllm serve \ + --omni Wan-AI/Wan2.2-I2V-A14B-Diffusers \ + --use-hsdp \ + --usp 4 \ + --cfg 2 \ + --vae-patch-parallel-size 8 \ + --vae-use-tiling +``` + +> **Why the difference?** With `--cfg 2`, two copies of the input (positive and +> negative prompts) are processed in parallel, effectively doubling the compute +> for the DiT backbone. USP is therefore halved from 8 to 4 so that the total +> parallelism across the 8 cards remains balanced (`usp * cfg = 8`). + +#### Verification + +After the server is ready, see +[`examples/online_serving/image_to_video/README.md`](../../examples/online_serving/image_to_video/README.md) +for complete client examples and request formats. + +#### Notes + +- **Key flags:** + - `--omni` — enables vLLM-Omni diffusion serving. + - `--use-hsdp` — enables Hybrid Sharded Data Parallelism for the DiT model + weights. + - `--usp ` — Unified Sequence Parallelism degree. + - `--cfg ` — Classifier-Free Guidance parallelism; set to 2 for models + that require negative-prompt computation, omit for distilled models. + - `--vae-patch-parallel-size 8` — parallelizes VAE decoding across all 8 + cards. + - `--vae-use-tiling` — enables tiled VAE decoding to reduce peak memory. +- **Performance tips:** + - Installing mindie-sd and enabling Laser Attention + (`MINDIE_SD_FA_TYPE=ascend_laser_attention`) provides up to ~40% + performance improvement at 720p resolution due to long-sequence attention + optimization. +- **Known limitations:** + - `MULTI_STREAM_MEMORY_REUSE=2` is required on NPU when using HSDP/FSDP2 + due to a multi-stream memory reuse bug. This is not needed on CUDA. diff --git a/requirements/common.txt b/requirements/common.txt index 1fff58444..63e16d580 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -1,7 +1,6 @@ # Common dependencies for all platforms av>=14.0.0 omegaconf>=2.3.0 -resampy>=0.4.3 diffusers>=0.36.0 accelerate==1.12.0 soundfile>=0.13.1 @@ -10,7 +9,6 @@ tqdm>=4.66.0 torchsde>=0.2.6 openai-whisper>=20250625 imageio[ffmpeg]>=2.37.2 -sox>=1.5.0 x-transformers>=2.12.2 einops>=0.8.1 prettytable>=3.8.0 diff --git a/tests/config/__init__.py b/tests/config/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/config/test_pipeline_registry.py b/tests/config/test_pipeline_registry.py new file mode 100644 index 000000000..3483d530c --- /dev/null +++ b/tests/config/test_pipeline_registry.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the central pipeline registry (2.5/N).""" + +from __future__ import annotations + +import pytest + +from vllm_omni.config.pipeline_registry import ( + _DIFFUSION_PIPELINES, + _OMNI_PIPELINES, + _VLLM_OMNI_PIPELINES, +) +from vllm_omni.config.stage_config import ( + _PIPELINE_REGISTRY, + PipelineConfig, + StageExecutionType, + StagePipelineConfig, + register_pipeline, +) + + +class TestCentralRegistryDeclarations: + """Every in-tree pipeline must be declared exactly once in the central registry.""" + + def test_union_contains_all_omni(self): + for key in _OMNI_PIPELINES: + assert key in _VLLM_OMNI_PIPELINES + + def test_union_contains_all_diffusion(self): + for key in _DIFFUSION_PIPELINES: + assert key in _VLLM_OMNI_PIPELINES + + def test_no_duplicate_model_type_between_omni_and_diffusion(self): + overlap = set(_OMNI_PIPELINES) & set(_DIFFUSION_PIPELINES) + assert not overlap, f"Duplicate model_types across omni/diffusion: {overlap}" + + def test_expected_omni_pipelines_present(self): + # Guard against accidental removal during future refactors. + assert "qwen2_5_omni" in _OMNI_PIPELINES + assert "qwen2_5_omni_thinker_only" in _OMNI_PIPELINES + assert "qwen3_omni_moe" in _OMNI_PIPELINES + assert "qwen3_tts" in _OMNI_PIPELINES + + +class TestLazyLoading: + """Pipelines are imported only on first access.""" + + def test_contains_without_import(self): + # ``in`` hits the lazy map, not the loaded cache. + assert "qwen3_omni_moe" in _PIPELINE_REGISTRY + + def test_getitem_loads_correct_pipeline(self): + pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] + assert pipeline.model_type == "qwen3_omni_moe" + assert pipeline.model_arch == "Qwen3OmniMoeForConditionalGeneration" + + def test_unknown_model_type_returns_none_via_get(self): + assert _PIPELINE_REGISTRY.get("not_a_real_pipeline") is None + + def test_unknown_model_type_raises_keyerror_via_getitem(self): + with pytest.raises(KeyError): + _PIPELINE_REGISTRY["not_a_real_pipeline"] + + def test_iteration_yields_registered_pipelines(self): + keys = set(_PIPELINE_REGISTRY) + assert "qwen2_5_omni" in keys + assert "qwen3_omni_moe" in keys + + +class TestDynamicRegistration: + """``register_pipeline()`` still works for plugins and tests.""" + + def test_register_adds_to_registry(self): + custom = PipelineConfig( + model_type="_test_dynamic_registration", + model_arch="DynamicTestModel", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="test", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + ), + ), + ) + register_pipeline(custom) + try: + assert "_test_dynamic_registration" in _PIPELINE_REGISTRY + assert _PIPELINE_REGISTRY["_test_dynamic_registration"] is custom + finally: + # Don't leak the test registration into other tests. + if "_test_dynamic_registration" in _PIPELINE_REGISTRY: + del _PIPELINE_REGISTRY["_test_dynamic_registration"] + + def test_dynamic_registration_overrides_lazy_entry(self): + # Build a substitute for qwen3_omni_moe that we can distinguish. + original = _PIPELINE_REGISTRY["qwen3_omni_moe"] + override = PipelineConfig( + model_type="qwen3_omni_moe", + model_arch="OverriddenArch", + stages=original.stages, + ) + register_pipeline(override) + try: + assert _PIPELINE_REGISTRY["qwen3_omni_moe"].model_arch == "OverriddenArch" + finally: + # Remove the dynamic override so later tests see the original. + if "qwen3_omni_moe" in _PIPELINE_REGISTRY._loaded: + del _PIPELINE_REGISTRY["qwen3_omni_moe"] diff --git a/tests/conftest.py b/tests/conftest.py index adb87cbd7..77075f952 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3423 +1,62 @@ -import atexit -import base64 -import datetime -import io -import json -import math -import os -import random -import re -import tempfile - -import requests - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -# Set CPU device for CI environments without GPU -if "VLLM_TARGET_DEVICE" not in os.environ: - os.environ["VLLM_TARGET_DEVICE"] = "cpu" - -import concurrent.futures -import contextlib -import gc -import multiprocessing -import socket -import subprocess -import sys -import threading -import time -import uuid -from collections.abc import Generator -from dataclasses import dataclass -from io import BytesIO -from pathlib import Path -from typing import Any, NamedTuple - -import cv2 -import numpy as np -import psutil -import pytest -import soundfile as sf -import torch -import yaml -from openai import OpenAI, omit -from PIL import Image -from transformers import pipeline -from vllm import TextPrompt -from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from vllm.logger import init_logger -from vllm.utils.network_utils import get_open_port - -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniSamplingParams -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.platforms import current_omni_platform - -logger = init_logger(__name__) - - -PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None -PromptImageInput = list[Any] | Any | None -PromptVideoInput = list[Any] | Any | None - -_GENDER_PIPELINE = None -# transformers.Pipeline is not thread-safe; concurrent e2e requests must serialize inference. -_GENDER_PIPELINE_LOCK = threading.Lock() - -# int16 mono PCM from /v1/audio/speech when response_format=pcm (Qwen3-TTS code2wav output rate). -_PCM_SPEECH_SAMPLE_RATE_HZ = 24_000 - - -class OmniServerParams(NamedTuple): - model: str - port: int | None = None - stage_config_path: str | None = None - server_args: list[str] | None = None - env_dict: dict[str, str] | None = None - use_omni: bool = True - use_stage_cli: bool = False - init_timeout: int | None = None - stage_init_timeout: int | None = None # None defers to the server's own default (300 s) - - -def assert_image_diffusion_response( - response, - request_config: dict[str, Any], - run_level: str = None, -) -> None: - """ - Validate image diffusion response. - - Expected request_config schema: - { - "request_type": "image", - "extra_body": { - "num_outputs_per_prompt": 1, - "width": ..., - "height": ..., - ... - } - } - """ - assert response.images is not None, "Image response is None" - assert len(response.images) > 0, "No images in response" - - extra_body = request_config.get("extra_body") or {} - - num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt") - if num_outputs_per_prompt is not None: - assert len(response.images) == num_outputs_per_prompt, ( - f"Expected {num_outputs_per_prompt} images, got {len(response.images)}" - ) - - if run_level == "advanced_model": - width = extra_body.get("width") - height = extra_body.get("height") - - if width is not None or height is not None: - for img in response.images: - assert_image_valid(img, width=width, height=height) - - -def assert_video_diffusion_response( - response, - request_config: dict[str, Any], - run_level: str = None, -) -> None: - """ - Validate video diffusion response. - - Expected request_config schema: - { - "request_type": "video", - "form_data": { - "prompt": "...", - "num_frames": ..., - "width": ..., - "height": ..., - "fps": ..., - ... - } - } - """ - form_data = request_config.get("form_data", {}) - - assert response.videos is not None, "Video response is None" - assert len(response.videos) > 0, "No videos in response" - - expected_frames = _maybe_int(form_data.get("num_frames")) - expected_width = _maybe_int(form_data.get("width")) - expected_height = _maybe_int(form_data.get("height")) - expected_fps = _maybe_int(form_data.get("fps")) - - for vid_bytes in response.videos: - assert_video_valid( - vid_bytes, - num_frames=expected_frames, - width=expected_width, - height=expected_height, - fps=expected_fps, - ) - - -def assert_audio_diffusion_response( - response, - request_config: dict[str, Any], - run_level: str = None, -) -> None: - """ - Validate audio diffusion response. - """ - raise NotImplementedError("Audio validation is not implemented yet") - - -def _maybe_int(value: Any) -> int | None: - if value is None: - return None - return int(value) - - -def assert_image_valid(image: Path | Image.Image, *, width: int | None = None, height: int | None = None): - """Assert the file is a loadable image with optional exact dimensions.""" - if isinstance(image, Path): - assert image.exists(), f"Image not found: {image}" - image = Image.open(image) - image.load() - assert image.width > 0 and image.height > 0 - if width is not None: - assert image.width == width, f"Expected width={width}, got {image.width}" - if height is not None: - assert image.height == height, f"Expected height={height}, got {image.height}" - return image - - -def assert_video_valid( - video: Path | bytes | BytesIO, - *, - num_frames: int | None = None, - width: int | None = None, - height: int | None = None, - fps: float | None = None, -) -> dict[str, int | float]: - """Assert the MP4 has the expected resolution and exact frame count.""" - temp_path = None - cap = None - try: - # Normalize input to file path - if isinstance(video, Path): - if not video.exists(): - raise AssertionError(f"Video file not found: {video}") - video_path = str(video) - else: - # Create temp file for bytes/BytesIO - suffix = ".mp4" - with tempfile.NamedTemporaryFile(delete=False, suffix=suffix, mode="wb") as tmp: - if isinstance(video, bytes): - tmp.write(video) - elif isinstance(video, BytesIO): - tmp.write(video.getvalue()) - else: - raise TypeError(f"Unsupported video type: {type(video)}") - temp_path = Path(tmp.name) - video_path = str(temp_path) - - # Open video capture - cap = cv2.VideoCapture(video_path) - if not cap.isOpened(): - raise AssertionError(f"Failed to open video: {video_path}") - - # Extract properties - actual_num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - actual_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) - actual_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) - actual_fps = cap.get(cv2.CAP_PROP_FPS) - - actual_num_frames = 0 - while True: - ok, _frame = cap.read() - if not ok: - break - actual_num_frames += 1 - - # Basic validity checks - if actual_num_frames <= 0: - raise AssertionError(f"Invalid frame count: {actual_num_frames} (must be > 0)") - if actual_width <= 0 or actual_height <= 0: - raise AssertionError(f"Invalid dimensions: {actual_width}x{actual_height} (must be > 0)") - if actual_fps <= 0: - raise AssertionError(f"Invalid FPS: {actual_fps} (must be > 0)") - - # Validate against expectations - if num_frames is not None: - expected_num_frames = (num_frames // 4) * 4 + 1 - assert actual_num_frames == expected_num_frames, ( - f"Frame count mismatch: expected {num_frames}, got {actual_num_frames}" - ) - if width is not None: - assert actual_width == width, f"Width mismatch: expected {width}px, got {actual_width}px" - if height is not None: - assert actual_height == height, f"Height mismatch: expected {height}px, got {actual_height}px" - if fps is not None: - # Use tolerance for float comparison (codec rounding) - assert abs(actual_fps - fps) < 0.5, f"FPS mismatch: expected {fps}, got {actual_fps:.2f}" - - return {"num_frames": actual_num_frames, "width": actual_width, "height": actual_height, "fps": actual_fps} - - except Exception as e: - print(f"ERROR: {type(e).__name__}: {e}", flush=True) - raise - - finally: - # Cleanup resources - if cap is not None: - cap.release() - if temp_path and temp_path.exists(): - try: - temp_path.unlink() - except OSError: - pass - - -def assert_audio_valid( - audio_or_path: Path | np.ndarray, - *, - sample_rate: int, - channels: int, - duration_s: float, -) -> None: - """Assert WAV file or (batch, channels, samples) ndarray matches expected audio format.""" - expected_samples = int(duration_s * sample_rate) - if isinstance(audio_or_path, np.ndarray): - audio = audio_or_path - assert audio.ndim == 3, f"Expected audio ndim=3 (batch, channels, samples), got shape {audio.shape}" - assert audio.shape[0] == 1, f"Expected batch size 1, got {audio.shape[0]}" - assert audio.shape[1] == channels, f"Expected {channels} channels, got {audio.shape[1]}" - assert audio.shape[2] == expected_samples, ( - f"Expected {expected_samples} samples ({duration_s}s @ {sample_rate} Hz), got {audio.shape[2]}" - ) - return - - path = audio_or_path - assert path.exists(), f"Audio not found: {path}" - info = sf.info(str(path)) - assert info.samplerate == sample_rate, f"Expected sample_rate={sample_rate}, got {info.samplerate}" - assert info.channels == channels, f"Expected {channels} channel(s), got {info.channels}" - assert info.frames == expected_samples, ( - f"Expected {expected_samples} frames ({duration_s}s @ {sample_rate} Hz), got {info.frames}" - ) - - -def decode_b64_image(b64: str): - img = Image.open(BytesIO(base64.b64decode(b64))) - img.load() - return img - - -@pytest.fixture(scope="session") -def model_prefix() -> str: - """Optional model-path prefix from MODEL_PREFIX env var. - Useful if models are downloaded to non-default local directories. - """ - prefix = os.environ.get("MODEL_PREFIX", "") - return f"{prefix.rstrip('/')}/" if prefix else "" - - -@pytest.fixture(autouse=True) -def default_vllm_config(): - """Set a default VllmConfig for all tests. - - This fixture is auto-used for all tests to ensure that any test - that directly instantiates vLLM CustomOps (e.g., RMSNorm, LayerNorm) - or model components has the required VllmConfig context. - - This fixture is required for vLLM 0.14.0+ where CustomOp initialization - requires a VllmConfig context set via set_current_vllm_config(). - """ - from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config - - # Use CPU device if no GPU is available (e.g., in CI environments) - has_gpu = torch.cuda.is_available() and torch.cuda.device_count() > 0 - device = "cuda" if has_gpu else "cpu" - device_config = DeviceConfig(device=device) - - with set_current_vllm_config(VllmConfig(device_config=device_config)): - yield - - -@pytest.fixture(autouse=True) -def clean_gpu_memory_between_tests(): - print("\n=== PRE-TEST GPU CLEANUP ===") - _run_pre_test_cleanup() - yield - _run_post_test_cleanup() - - -@pytest.fixture(autouse=True) -def log_test_name_before_test(request): - print(f"--- Running test: {request.node.name}") - yield - - -def _run_pre_test_cleanup(enable_force=False): - if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force: - print("\nPre-test GPU cleanup skipped(Default off is typical when one worker/instance runs many tests.)\n") - return - - print("\nPre-test GPU status:") - - num_gpus = torch.cuda.device_count() - if num_gpus > 0: - try: - from tests.utils import wait_for_gpu_memory_to_clear - - wait_for_gpu_memory_to_clear( - devices=list(range(num_gpus)), - threshold_ratio=0.05, - ) - except Exception as e: - print(f"Pre-test cleanup note: {e}") - - -def _run_post_test_cleanup(enable_force=False): - if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force: - print("GPU cleanup disabled") - return - - if torch.cuda.is_available(): - gc.collect() - torch.cuda.empty_cache() - - print("Post-test GPU status:") - _print_gpu_processes() - - -def _print_gpu_processes(): - """Print GPU information including nvidia-smi and system processes""" - - print("\n" + "=" * 80) - print("NVIDIA GPU Information (nvidia-smi)") - print("=" * 80) - - try: - nvidia_result = subprocess.run( - ["nvidia-smi"], - capture_output=True, - text=True, - timeout=5, - ) - - if nvidia_result.returncode == 0: - lines = nvidia_result.stdout.strip().split("\n") - for line in lines[:20]: - print(line) - - if len(lines) > 20: - print(f"... (showing first 20 of {len(lines)} lines)") - else: - print("nvidia-smi command failed") - - except (subprocess.TimeoutExpired, FileNotFoundError): - print("nvidia-smi not available or timed out") - except Exception as e: - print(f"Error running nvidia-smi: {e}") - - print("\n" + "=" * 80) - print("Detailed GPU Processes (nvidia-smi pmon)") - print("=" * 80) - - try: - pmon_result = subprocess.run( - ["nvidia-smi", "pmon", "-c", "1"], - capture_output=True, - text=True, - timeout=3, - ) - - if pmon_result.returncode == 0 and pmon_result.stdout.strip(): - print(pmon_result.stdout) - else: - print("No active GPU processes found via nvidia-smi pmon") - - except Exception: - print("nvidia-smi pmon not available") - - print("\n" + "=" * 80) - print("System Processes with GPU keywords") - print("=" * 80) - - -def dummy_messages_from_mix_data( - system_prompt: dict[str, Any] = None, - video_data_url: Any = None, - audio_data_url: Any = None, - image_data_url: Any = None, - content_text: str = None, -): - """Create messages with video、image、audio data URL for OpenAI API.""" - - if content_text is not None: - content = [{"type": "text", "text": content_text}] - else: - content = [] - - media_items = [] - if isinstance(video_data_url, list): - for video_url in video_data_url: - media_items.append((video_url, "video")) - else: - media_items.append((video_data_url, "video")) - - if isinstance(image_data_url, list): - for url in image_data_url: - media_items.append((url, "image")) - else: - media_items.append((image_data_url, "image")) - - if isinstance(audio_data_url, list): - for url in audio_data_url: - media_items.append((url, "audio")) - else: - media_items.append((audio_data_url, "audio")) - - content.extend( - {"type": f"{media_type}_url", f"{media_type}_url": {"url": url}} - for url, media_type in media_items - if url is not None - ) - messages = [{"role": "user", "content": content}] - if system_prompt is not None: - messages = [system_prompt] + messages - return messages - - -def generate_synthetic_audio( - duration: int, # seconds - num_channels: int, # 1：Mono，2：Stereo 5：5.1 surround sound - sample_rate: int = 48000, # Default use 48000Hz. - save_to_file: bool = False, -) -> dict[str, Any]: - """ - Generate TTS speech with pyttsx3 and return base64 string. - """ - - import pyttsx3 - import soundfile as sf - - def _pick_voice(engine: pyttsx3.Engine) -> str | None: - voices = engine.getProperty("voices") - if not voices: - return None - - preferred_tokens = ( - "natural", - "jenny", - "sonia", - "susan", - "zira", - "aria", - "hazel", - "samantha", - "ava", - "allison", - "female", - "woman", - "english-us", - "en-us", - "english", - ) - discouraged_tokens = ( - "espeak", - "robot", - "mbrola", - "microsoft david", - "male", - "man", - ) - - best_voice = voices[0] - best_score = float("-inf") - for voice in voices: - voice_text = f"{getattr(voice, 'id', '')} {getattr(voice, 'name', '')}".lower() - voice_languages = " ".join( - lang.decode(errors="ignore") if isinstance(lang, bytes) else str(lang) - for lang in getattr(voice, "languages", []) - ).lower() - combined_text = f"{voice_text} {voice_languages}" - score = 0 - for idx, token in enumerate(preferred_tokens): - if token in combined_text: - score += 20 - idx - for token in discouraged_tokens: - if token in combined_text: - score -= 10 - if "english" in combined_text or "en_" in combined_text or "en-" in combined_text: - score += 4 - if "en-us" in combined_text or "english-us" in combined_text: - score += 4 - if score > best_score: - best_score = score - best_voice = voice - - return best_voice.id - - def _resample_audio(audio: np.ndarray, src_sr: int, dst_sr: int) -> np.ndarray: - if src_sr == dst_sr or len(audio) == 0: - return audio.astype(np.float32) - - src_len = audio.shape[0] - dst_len = max(1, int(round(src_len * float(dst_sr) / float(src_sr)))) - src_idx = np.arange(src_len, dtype=np.float32) - dst_idx = np.linspace(0, src_len - 1, dst_len, dtype=np.float32) - - resampled_channels: list[np.ndarray] = [] - for ch in range(audio.shape[1]): - resampled_channels.append(np.interp(dst_idx, src_idx, audio[:, ch]).astype(np.float32)) - return np.stack(resampled_channels, axis=1) - - def _match_channels(audio: np.ndarray, target_channels: int) -> np.ndarray: - current_channels = audio.shape[1] - if current_channels == target_channels: - return audio.astype(np.float32) - if target_channels == 1: - return np.mean(audio, axis=1, keepdims=True, dtype=np.float32) - if current_channels == 1: - return np.repeat(audio, target_channels, axis=1).astype(np.float32) - - collapsed = np.mean(audio, axis=1, keepdims=True, dtype=np.float32) - return np.repeat(collapsed, target_channels, axis=1).astype(np.float32) - - def _trim_silence(audio: np.ndarray, threshold: float = 0.01) -> np.ndarray: - if len(audio) == 0: - return audio - energy = np.max(np.abs(audio), axis=1) - voiced = np.where(energy > threshold)[0] - if len(voiced) == 0: - return audio - start = max(0, int(voiced[0]) - int(sample_rate * 0.02)) - end = min(len(audio), int(voiced[-1]) + int(sample_rate * 0.04) + 1) - return audio[start:end] - - def _enhance_speech(audio: np.ndarray) -> np.ndarray: - if len(audio) == 0: - return audio.astype(np.float32) - enhanced = audio.astype(np.float32).copy() - enhanced -= np.mean(enhanced, axis=0, keepdims=True, dtype=np.float32) - if len(enhanced) > 1: - preemphasis = enhanced.copy() - preemphasis[1:] = enhanced[1:] - 0.94 * enhanced[:-1] - enhanced = 0.7 * enhanced + 0.3 * preemphasis - # Mild dynamic-range compression for ASR/TTS robustness. - enhanced = np.sign(enhanced) * np.sqrt(np.abs(enhanced)) - # Light fade to avoid clicks after trimming/repeating. - fade = min(len(enhanced) // 4, max(1, int(sample_rate * 0.01))) - if fade > 1: - ramp_in = np.linspace(0.0, 1.0, fade, dtype=np.float32) - ramp_out = np.linspace(1.0, 0.0, fade, dtype=np.float32) - enhanced[:fade] *= ramp_in[:, None] - enhanced[-fade:] *= ramp_out[:, None] - peak = float(np.max(np.abs(enhanced))) - if peak > 1e-8: - enhanced = enhanced / peak * 0.95 - return enhanced.astype(np.float32) - - phrase_text = "test" - num_samples = int(sample_rate * max(1, duration)) - audio_data = np.zeros((num_samples, num_channels), dtype=np.float32) - - engine = pyttsx3.init() - engine.setProperty("rate", 112) - engine.setProperty("volume", 1.0) - selected_voice = _pick_voice(engine) - if selected_voice is not None: - engine.setProperty("voice", selected_voice) - - temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) - temp_wav.close() - - try: - engine.save_to_file(phrase_text, temp_wav.name) - engine.runAndWait() - engine.stop() - - ready = False - for _ in range(50): - if os.path.exists(temp_wav.name) and os.path.getsize(temp_wav.name) > 44: - ready = True - break - time.sleep(0.1) - - if not ready: - raise RuntimeError("pyttsx3 did not produce a WAV file in time.") - - tts_audio, tts_sr = sf.read(temp_wav.name, dtype="float32", always_2d=True) - finally: - if os.path.exists(temp_wav.name): - os.unlink(temp_wav.name) - - if len(tts_audio) == 0: - raise RuntimeError("pyttsx3 produced an empty WAV file.") - - tts_audio = _resample_audio(tts_audio, tts_sr, sample_rate) - tts_audio = _match_channels(tts_audio, num_channels) - tts_audio = _trim_silence(tts_audio, threshold=0.012) - tts_audio = _enhance_speech(tts_audio) - - lead_silence = min(int(sample_rate * 0.02), num_samples // 8) - pause_samples = int(sample_rate * 0.18) - start = lead_silence - phrase_len = tts_audio.shape[0] - - while start < num_samples: - take = min(phrase_len, num_samples - start) - audio_data[start : start + take] = tts_audio[:take] - start += phrase_len + pause_samples - - max_amp = float(np.max(np.abs(audio_data))) - if max_amp > 0: - audio_data = audio_data / max_amp * 0.95 - - audio_bytes: bytes | None = None - output_path: str | None = None - result: dict[str, Any] = { - "np_array": audio_data.copy(), - } - - if save_to_file: - timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - output_path = f"audio_{num_channels}ch_{timestamp}.wav" - - try: - sf.write(output_path, audio_data, sample_rate, format="WAV", subtype="PCM_16") - print(f"Audio saved: {output_path}") - - with open(output_path, "rb") as f: - audio_bytes = f.read() - except Exception as e: - print(f"Save failed: {e}") - save_to_file = False - - # If not saving or save failed, create in memory - if not save_to_file or audio_bytes is None: - buffer = io.BytesIO() - sf.write(buffer, audio_data, sample_rate, format="WAV", subtype="PCM_16") - buffer.seek(0) - audio_bytes = buffer.read() - - # Return result - base64_audio = base64.b64encode(audio_bytes).decode("utf-8") - result["base64"] = base64_audio - # Always include file_path to avoid KeyError in callers. - result["file_path"] = output_path if save_to_file and output_path else None - - return result - - -def _mux_mp4_bytes_with_synthetic_audio( - video_mp4_bytes: bytes, - *, - num_frames: int, - fps: float = 30.0, - sample_rate: int = 48000, -) -> bytes: - """ - Mux a video-only MP4 with mono TTS audio from :func:`generate_synthetic_audio` (AAC). - - Audio length is at least the video duration in whole seconds (rounded up); ffmpeg - ``-shortest`` trims to the video when the WAV is longer. - - Uses ffmpeg from ``imageio_ffmpeg`` when available, else ``ffmpeg`` on PATH. - If TTS or mux fails, returns ``video_mp4_bytes`` unchanged. - - Mux subprocess does **not** use ``capture_output=True``: ffmpeg can block writing - to a full stderr pipe while :func:`subprocess.run` waits for exit (classic deadlock). - """ - duration_sec = num_frames / fps if fps > 0 else 0.0 - # generate_synthetic_audio(duration=int) uses at least 1s of buffer internally - duration_int = max(1, int(math.ceil(duration_sec))) - - try: - audio_result = generate_synthetic_audio( - duration=duration_int, - num_channels=1, - sample_rate=sample_rate, - save_to_file=False, - ) - audio_pcm = audio_result["np_array"] - except Exception as e: - logger.warning("Synthetic video: generate_synthetic_audio failed (%s); using video-only MP4.", e) - return video_mp4_bytes - - try: - import imageio_ffmpeg - - ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe() - except Exception: - ffmpeg_exe = "ffmpeg" - - import tempfile - - try: - with tempfile.TemporaryDirectory(prefix="syn_vid_mux_") as tmp: - vid_path = os.path.join(tmp, "video.mp4") - wav_path = os.path.join(tmp, "audio.wav") - out_path = os.path.join(tmp, "out.mp4") - with open(vid_path, "wb") as f: - f.write(video_mp4_bytes) - sf.write(wav_path, audio_pcm, sample_rate, format="WAV", subtype="PCM_16") - cmd = [ - ffmpeg_exe, - "-y", - "-nostdin", - "-hide_banner", - "-loglevel", - "error", - "-i", - vid_path, - "-i", - wav_path, - "-c:v", - "copy", - "-c:a", - "aac", - "-b:a", - "128k", - "-shortest", - "-movflags", - "+faststart", - out_path, - ] - subprocess.run( - cmd, - check=True, - stdin=subprocess.DEVNULL, - timeout=300, - ) - with open(out_path, "rb") as f: - return f.read() - except ( - FileNotFoundError, - subprocess.CalledProcessError, - subprocess.TimeoutExpired, - OSError, - ) as e: - logger.warning("Synthetic video: audio mux failed (%s); using video-only MP4.", e) - return video_mp4_bytes - - -def generate_synthetic_video( - width: int, - height: int, - num_frames: int, - save_to_file: bool = False, - *, - embed_audio: bool = False, -) -> dict[str, Any]: - """Generate synthetic video with bouncing balls and base64 MP4. - - When ``embed_audio`` is True, muxes mono AAC from :func:`generate_synthetic_audio` - (TTS + ffmpeg) into the MP4; otherwise returns video-only MP4 (faster when tests do - not need an audio track). - """ - - import cv2 - import imageio - - # Create random balls - num_balls = random.randint(3, 8) - balls = [] - - for _ in range(num_balls): - radius = min(width, height) // 8 - if radius < 1: - raise ValueError(f"Video dimensions ({width}x{height}) are too small for synthetic video generation") - x = random.randint(radius, width - radius) - y = random.randint(radius, height - radius) - - speed = random.uniform(3.0, 8.0) - angle = random.uniform(0, 2 * math.pi) - vx = speed * math.cos(angle) - vy = speed * math.sin(angle) - - # OpenCV uses BGR format, but imageio expects RGB - # We'll create in BGR first, then convert to RGB later - color_bgr = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255)) - - balls.append({"x": x, "y": y, "vx": vx, "vy": vy, "radius": radius, "color_bgr": color_bgr}) - - # Generate video frames - video_frames = [] - - for frame_idx in range(num_frames): - # Create black background (BGR format) - frame_bgr = np.zeros((height, width, 3), dtype=np.uint8) - - for ball in balls: - # Update position - ball["x"] += ball["vx"] - ball["y"] += ball["vy"] - - # Boundary collision detection - if ball["x"] - ball["radius"] <= 0 or ball["x"] + ball["radius"] >= width: - ball["vx"] = -ball["vx"] - ball["x"] = max(ball["radius"], min(width - ball["radius"], ball["x"])) - - if ball["y"] - ball["radius"] <= 0 or ball["y"] + ball["radius"] >= height: - ball["vy"] = -ball["vy"] - ball["y"] = max(ball["radius"], min(height - ball["radius"], ball["y"])) - - # Use cv2 to draw circle - x, y = int(ball["x"]), int(ball["y"]) - radius = ball["radius"] - - # Draw solid circle (main circle) - cv2.circle(frame_bgr, (x, y), radius, ball["color_bgr"], -1) - - # Add simple 3D effect: draw a brighter center - if radius > 3: # Only add highlight when radius is large enough - highlight_radius = max(1, radius // 2) - highlight_x = max(highlight_radius, min(x - radius // 4, width - highlight_radius)) - highlight_y = max(highlight_radius, min(y - radius // 4, height - highlight_radius)) - - # Create highlight color (brighter) - highlight_color = tuple(min(c + 40, 255) for c in ball["color_bgr"]) - cv2.circle(frame_bgr, (highlight_x, highlight_y), highlight_radius, highlight_color, -1) - - # Convert BGR to RGB for imageio - frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) - video_frames.append(frame_rgb) - - video_array = np.array(video_frames) - result = { - "np_array": video_array, - } - saved_file_path = None - - fps = 30 - buffer = io.BytesIO() - writer_kwargs = { - "format": "mp4", - "fps": fps, - "codec": "libx264", - "quality": 7, - "pixelformat": "yuv420p", - "macro_block_size": 16, - "ffmpeg_params": [ - "-preset", - "medium", - "-crf", - "23", - "-movflags", - "+faststart", - "-pix_fmt", - "yuv420p", - "-vf", - f"scale={width}:{height}", - ], - } - - try: - with imageio.get_writer(buffer, **writer_kwargs) as writer: - for frame in video_frames: - writer.append_data(frame) - buffer.seek(0) - video_only_bytes = buffer.read() - except Exception as e: - print(f"Warning: Failed to encode synthetic video: {e}") - raise - - if embed_audio: - video_bytes = _mux_mp4_bytes_with_synthetic_audio(video_only_bytes, num_frames=num_frames, fps=float(fps)) - else: - video_bytes = video_only_bytes - - if save_to_file: - timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - output_path = f"video_{width}x{height}_{timestamp}.mp4" - try: - with open(output_path, "wb") as f: - f.write(video_bytes) - saved_file_path = output_path - print(f"Video saved to: {saved_file_path}") - except Exception as e: - print(f"Warning: Failed to save video to file {output_path}: {e}") - - base64_video = base64.b64encode(video_bytes).decode("utf-8") - - result["base64"] = base64_video - if save_to_file and saved_file_path: - result["file_path"] = saved_file_path - - return result - - -def generate_synthetic_image(width: int, height: int, save_to_file: bool = False) -> dict[str, Any]: - """Generate synthetic image with randomly colored squares and return base64 string.""" - from PIL import Image, ImageDraw - - # Create white background - image = Image.new("RGB", (width, height), (255, 255, 255)) - draw = ImageDraw.Draw(image) - - # Generate random number of squares - num_squares = random.randint(3, 8) - - for _ in range(num_squares): - # Random square size - square_size = random.randint(min(width, height) // 8, min(width, height) // 4) - - # Random position - x = random.randint(0, width - square_size - 1) - y = random.randint(0, height - square_size - 1) - - # Random color - color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) - - # Random border width - border_width = random.randint(1, 5) - - # Draw square - draw.rectangle([x, y, x + square_size, y + square_size], fill=color, outline=(0, 0, 0), width=border_width) - - image_array = np.array(image) - result = {"np_array": image_array.copy()} - - # Handle file saving - image_bytes = None - saved_file_path = None - - if save_to_file: - timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - output_path = f"image_{width}x{height}_{timestamp}.jpg" - - try: - # Save image to file - image.save(output_path, format="JPEG", quality=85, optimize=True) - saved_file_path = output_path - print(f"Image saved to: {saved_file_path}") - - # Read file for base64 encoding - with open(output_path, "rb") as f: - image_bytes = f.read() - - except Exception as e: - print(f"Warning: Failed to save image to file {output_path}: {e}") - save_to_file = False - - # If not saving or save failed, create in memory - if not save_to_file or image_bytes is None: - buffer = io.BytesIO() - image.save(buffer, format="JPEG", quality=85, optimize=True) - buffer.seek(0) - image_bytes = buffer.read() - - # Generate base64 - base64_image = base64.b64encode(image_bytes).decode("utf-8") - - # Return result - result["base64"] = base64_image - if save_to_file and saved_file_path: - result["file_path"] = saved_file_path - - return result - - -def preprocess_text(text): - import opencc - - word_to_num = { - "zero": "0", - "one": "1", - "two": "2", - "three": "3", - "four": "4", - "five": "5", - "six": "6", - "seven": "7", - "eight": "8", - "nine": "9", - "ten": "10", - } - - for word, num in word_to_num.items(): - pattern = r"\b" + re.escape(word) + r"\b" - text = re.sub(pattern, num, text, flags=re.IGNORECASE) - - text = re.sub(r"[^\w\s]", "", text) - text = re.sub(r"\s+", " ", text) - cc = opencc.OpenCC("t2s") - text = cc.convert(text) - - # Special handling for spaces between Chinese characters: - # - Keep single spaces between English words/numbers - # - Remove spaces only when surrounded by Chinese characters on both sides to prevent incorrect word segmentation - text = re.sub(r"(?<=[\u4e00-\u9fff])\s+(?=[\u4e00-\u9fff])", "", text) - - return text.lower().strip() - - -def cosine_similarity_text(text1, text2, n: int = 3): - from collections import Counter - - if not text1 or not text2: - return 0.0 - - text1 = preprocess_text(text1) - text2 = preprocess_text(text2) - print(f"cosine similarity text1 is: {text1}, text2 is: {text2}") - - ngrams1 = [text1[i : i + n] for i in range(len(text1) - n + 1)] - ngrams2 = [text2[i : i + n] for i in range(len(text2) - n + 1)] - - counter1 = Counter(ngrams1) - counter2 = Counter(ngrams2) - - all_ngrams = set(counter1.keys()) | set(counter2.keys()) - vec1 = [counter1.get(ng, 0) for ng in all_ngrams] - vec2 = [counter2.get(ng, 0) for ng in all_ngrams] - - dot_product = sum(a * b for a, b in zip(vec1, vec2)) - norm1 = sum(a * a for a in vec1) ** 0.5 - norm2 = sum(b * b for b in vec2) ** 0.5 - - if norm1 == 0 or norm2 == 0: - return 0.0 - return dot_product / (norm1 * norm2) - - -def convert_audio_to_text(audio_data): - """ - Convert base64 encoded audio data to text using speech recognition. - """ - audio_data = base64.b64decode(audio_data) - output_path = f"./test_{uuid.uuid4().hex}.wav" - with open(output_path, "wb") as audio_file: - audio_file.write(audio_data) - - print(f"audio data is saved: {output_path}") - text = convert_audio_file_to_text(output_path=output_path) - return text - - -def _merge_base64_audio_to_segment(base64_list: list[str]): - """Merge a list of base64-encoded audio chunks into one pydub AudioSegment.""" - from pydub import AudioSegment - - merged = None - for b64 in base64_list: - raw = base64.b64decode(b64.split(",", 1)[-1]) - seg = AudioSegment.from_file(io.BytesIO(raw)) - merged = seg if merged is None else merged + seg - return merged - - -@contextlib.contextmanager -def _serialize_whisper_small_model_download(): - """Serialize Whisper ``small`` cache writes across processes (Linux; ``fcntl``).""" - import fcntl - - lock_path = Path.home() / ".cache" / "whisper" / ".small_model_download.lock" - lock_path.parent.mkdir(parents=True, exist_ok=True) - f = open(lock_path, "a+b") - try: - fcntl.flock(f.fileno(), fcntl.LOCK_EX) - yield - finally: - fcntl.flock(f.fileno(), fcntl.LOCK_UN) - f.close() - - -def _whisper_transcribe_in_current_process(output_path: str) -> str: - import whisper - - # Multi-GPU: use last visible device to avoid colliding with default device 0; single device uses 0. - device_index = None - if current_omni_platform.is_available(): - n = current_omni_platform.get_device_count() - if n == 1: - device_index = 0 - elif n > 1: - device_index = n - 1 - - if device_index is not None: - torch_device = current_omni_platform.get_torch_device(device_index) - current_omni_platform.set_device(torch_device) - device = str(torch_device) - use_accelerator = True - else: - use_accelerator = False - device = "cpu" - with _serialize_whisper_small_model_download(): - model = whisper.load_model("small", device=device) - try: - text = model.transcribe( - output_path, - temperature=0.0, - word_timestamps=True, - condition_on_previous_text=False, - )["text"] - finally: - del model - gc.collect() - if use_accelerator: - current_omni_platform.synchronize() - current_omni_platform.empty_cache() - - return text or "" - - -def convert_audio_file_to_text(output_path: str) -> str: - """Convert an audio file to text in an isolated subprocess (spawn).""" - ctx = multiprocessing.get_context("spawn") - with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor: - future = executor.submit(_whisper_transcribe_in_current_process, output_path) - return future.result() - - -def convert_audio_bytes_to_text(raw_bytes: bytes) -> str: - """ - Write container audio bytes (WAV, etc.) to a temp WAV file suitable for Whisper/ffmpeg. - Normalizes with soundfile to PCM_16 WAV when possible to avoid codec issues. - """ - output_path = f"./test_{uuid.uuid4().hex}.wav" - data, samplerate = sf.read(io.BytesIO(raw_bytes)) - sf.write(output_path, data, samplerate, format="WAV", subtype="PCM_16") - text = convert_audio_file_to_text(output_path) - return text - - -def modify_stage_config( - yaml_path: str, - updates: dict[str, Any] = None, - deletes: dict[str, Any] = None, -) -> str: - """ - Modify configurations in a YAML file, supporting both top-level and stage-specific modifications, - including addition, modification, and deletion of configurations. - - Args: - yaml_path: Path to the YAML configuration file. - updates: Dictionary containing both top-level and stage-specific modifications to add or update. - Format: { - 'async_chunk': True, - 'stage_args': { - 0: {'engine_args.max_model_len': 5800}, - 1: {'engine_args.max_num_seqs': 2} - } - } - deletes: Dictionary containing configurations to delete. - Format: { - 'old_config': None, # Delete entire key - 'stage_args': { - 0: ['engine_args.old_param'], - 1: ['runtime.unused_setting'] - } - } - - Returns: - str: Path to the newly created modified YAML file with timestamp suffix. - """ - path = Path(yaml_path) - if not path.exists(): - raise FileNotFoundError(f"yaml does not exist: {path}") - - try: - with open(yaml_path, encoding="utf-8") as f: - config = yaml.safe_load(f) or {} - except Exception as e: - raise ValueError(f"Cannot parse YAML file: {e}") - - # Helper function to apply update - def apply_update(config_dict: dict, key_path: str, value: Any) -> None: - """Apply update to dictionary using dot-separated path.""" - # Handle direct list assignment (e.g., engine_input_source: [1, 2]) - if "." not in key_path: - # Simple key, set directly - config_dict[key_path] = value - return - - current = config_dict - keys = key_path.split(".") - - for i in range(len(keys) - 1): - key = keys[i] - - # Handle list indices - if key.isdigit() and isinstance(current, list): - index = int(key) - if index < 0: - raise ValueError(f"Negative list index not allowed: {index}") - if index >= len(current): - # Expand list if needed - while len(current) <= index: - # If we need to go deeper (more keys after this), create a dict - # Otherwise, create None placeholder - current.append({} if i < len(keys) - 2 else None) - current = current[index] - elif isinstance(current, dict): - # Handle dictionary keys - if key not in current: - # If there are more keys after this, create appropriate structure - if i < len(keys) - 1: - # Check if next key is a digit (list index) or string (dict key) - if keys[i + 1].isdigit(): - current[key] = [] - else: - current[key] = {} - else: - # This is the last key, create based on value type - current[key] = [] if isinstance(value, list) else {} - elif not isinstance(current[key], (dict, list)) and i < len(keys) - 1: - # If current value is not dict/list but we need to go deeper, replace it - if keys[i + 1].isdigit(): - current[key] = [] - else: - current[key] = {} - current = current[key] - else: - # Current is not a dict or list, cannot traverse further - raise TypeError( - f"Cannot access {'.'.join(keys[: i + 1])} as a dict/list. It's a {type(current).__name__}" - ) - - # Set the final value - last_key = keys[-1] - if isinstance(current, list) and last_key.isdigit(): - # Setting a value in a list by index - index = int(last_key) - if index < 0: - raise ValueError(f"Negative list index not allowed: {index}") - if index >= len(current): - # Expand list if needed - while len(current) <= index: - current.append(None) - current[index] = value - elif isinstance(current, dict): - # Special case: if the value is a list and we're setting a top-level key - # Example: updating engine_input_source with [1, 2] - current[last_key] = value - else: - # Current is not a dict, cannot set key - raise TypeError(f"Cannot set value at {key_path}. Current type is {type(current).__name__}, expected dict.") - - # Helper function to delete by path - def delete_by_path(config_dict: dict, path: str) -> None: - """Delete configuration by dot-separated path.""" - if not path: - return - - current = config_dict - keys = path.split(".") - - # Traverse to the parent - for i in range(len(keys) - 1): - key = keys[i] - - # Handle list indices - if key.isdigit() and isinstance(current, list): - index = int(key) - if index < 0 or index >= len(current): - raise KeyError(f"List index {index} out of bounds") - current = current[index] - elif isinstance(current, dict): - if key not in current: - raise KeyError(f"Path {'.'.join(keys[: i + 1])} does not exist") - current = current[key] - else: - raise TypeError( - f"Cannot access {'.'.join(keys[: i + 1])} as a dict/list. It's a {type(current).__name__}" - ) - - # Delete the item - last_key = keys[-1] - - if isinstance(current, list) and last_key.isdigit(): - index = int(last_key) - if index < 0 or index >= len(current): - raise KeyError(f"List index {index} out of bounds") - del current[index] - elif isinstance(current, dict) and last_key in current: - del current[last_key] - else: - print(f"Path {path} does not exist") - - # Apply deletions first - if deletes: - for key, value in deletes.items(): - if key == "stage_args": - if value and isinstance(value, dict): - stage_args = config.get("stage_args", []) - if not stage_args: - raise ValueError("stage_args does not exist in config") - - for stage_id, delete_paths in value.items(): - if not delete_paths: - continue - - # Find stage by ID - target_stage = None - for stage in stage_args: - if stage.get("stage_id") == int(stage_id): - target_stage = stage - break - - if target_stage is None: - continue - - # Delete specified paths in this stage - # Avoid shadowing the original YAML Path used for the output filename below. - for delete_path in delete_paths: - if delete_path: # Skip empty paths - delete_by_path(target_stage, delete_path) - elif "." in key: - # Delete using dot-separated path - delete_by_path(config, key) - elif value is None and key in config: - # Delete entire key - del config[key] - - # Apply updates - if updates: - for key, value in updates.items(): - if key == "stage_args": - if value and isinstance(value, dict): - stage_args = config.get("stage_args", []) - if not stage_args: - raise ValueError("stage_args does not exist in config") - - for stage_id, stage_updates in value.items(): - # Find stage by ID - target_stage = None - for stage in stage_args: - if stage.get("stage_id") == int(stage_id): - target_stage = stage - break - - if target_stage is None: - available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s] - raise KeyError(f"Stage ID {stage_id} not found, available: {available_ids}") - - # Apply updates to this stage - for update_path, val in stage_updates.items(): - # Check if this is a simple key (not dot-separated) - # Example: 'engine_input_source' vs 'engine_args.max_model_len' - if "." not in update_path: - # Direct key assignment (e.g., updating a list value) - target_stage[update_path] = val - else: - # Dot-separated path (e.g., nested dict access) - apply_update(target_stage, update_path, val) - elif "." in key: - # Apply using dot-separated path - apply_update(config, key, value) - else: - # Direct top-level key - config[key] = value - - # Unique suffix: multiple modify_stage_config calls in one process often run - # within the same second (e.g. test_qwen3_omni_expansion imports both - # get_chunk_config and get_batch_token_config). int(time.time()) would collide - # and the later write would overwrite the earlier YAML on disk. - # Keep generated configs outside the repo and delete them when pytest exits. - output_fd, output_path = tempfile.mkstemp(prefix=f"{path.stem}_", suffix=".yaml") - atexit.register(Path(output_path).unlink, missing_ok=True) - - with os.fdopen(output_fd, "w", encoding="utf-8") as f: - yaml.dump(config, f, default_flow_style=None, sort_keys=False, allow_unicode=True, indent=2) - - return str(output_path) - - -class OmniServer: - """Omniserver for vLLM-Omni tests.""" - - def __init__( - self, - model: str, - serve_args: list[str], - *, - port: int | None = None, - env_dict: dict[str, str] | None = None, - use_omni: bool = True, - ) -> None: - _run_pre_test_cleanup(enable_force=True) - _run_post_test_cleanup(enable_force=True) - cleanup_dist_env_and_memory() - self.model = model - self.serve_args = serve_args - self.env_dict = env_dict - self.use_omni = use_omni - self.proc: subprocess.Popen | None = None - self.host = "127.0.0.1" - if port is None: - self.port = get_open_port() - else: - self.port = port - - def _start_server(self) -> None: - """Start the vLLM-Omni server subprocess.""" - env = os.environ.copy() - env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - if self.env_dict is not None: - env.update(self.env_dict) - - cmd = [ - sys.executable, - "-m", - "vllm_omni.entrypoints.cli.main", - "serve", - self.model, - "--host", - self.host, - "--port", - str(self.port), - ] - if self.use_omni: - cmd.append("--omni") - cmd += self.serve_args - - print(f"Launching OmniServer with: {' '.join(cmd)}") - self.proc = subprocess.Popen( - cmd, - env=env, - cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), # Set working directory to vllm-omni root - ) - - # Wait for server to be ready - max_wait = 1200 # 20 minutes - start_time = time.time() - while time.time() - start_time < max_wait: - # Check for process status - ret = self.proc.poll() - if ret is not None: - raise RuntimeError(f"Server processes exited with code {ret} before becoming ready.") - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.settimeout(1) - result = sock.connect_ex((self.host, self.port)) - if result == 0: - print(f"Server ready on {self.host}:{self.port}") - return - time.sleep(2) - - raise RuntimeError(f"Server failed to start within {max_wait} seconds") - - def _kill_process_tree(self, pid): - """kill process and its children with verification""" - try: - parent = psutil.Process(pid) - children = parent.children(recursive=True) - - # Get all PIDs first - all_pids = [pid] + [child.pid for child in children] - - # Terminate children - for child in children: - try: - child.terminate() - except psutil.NoSuchProcess: - pass - - # Wait for children - gone, still_alive = psutil.wait_procs(children, timeout=10) - - # Kill remaining children - for child in still_alive: - try: - child.kill() - except psutil.NoSuchProcess: - pass - - # Terminate parent - try: - parent.terminate() - parent.wait(timeout=10) - except (psutil.NoSuchProcess, psutil.TimeoutExpired): - try: - parent.kill() - except psutil.NoSuchProcess: - pass - - # VERIFICATION: Check if all processes are gone - time.sleep(1) # Give system time - alive_processes = [] - for check_pid in all_pids: - if psutil.pid_exists(check_pid): - alive_processes.append(check_pid) - - if alive_processes: - print(f"Warning: Processes still alive: {alive_processes}") - # Optional: Try system kill - import subprocess - - for alive_pid in alive_processes: - try: - subprocess.run(["kill", "-9", str(alive_pid)], timeout=2) - except Exception as e: - print(f"Cleanup failed: {e}") - - except psutil.NoSuchProcess: - pass - - def __enter__(self): - self._start_server() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - if self.proc: - self._kill_process_tree(self.proc.pid) - _run_pre_test_cleanup(enable_force=True) - _run_post_test_cleanup(enable_force=True) - cleanup_dist_env_and_memory() - - -class OmniServerStageCli(OmniServer): - """Omni server harness that exercises the stage CLI flow.""" - - def __init__( - self, - model: str, - stage_config_path: str, - serve_args: list[str] | None = None, - *, - stage_ids: list[int] | None = None, - port: int | None = None, - env_dict: dict[str, str] | None = None, - ) -> None: - super().__init__(model, serve_args or [], port=port, env_dict=env_dict, use_omni=True) - self.stage_config_path = stage_config_path - self.master_port = get_open_port() - self.visible_device_list = self._load_visible_device_list(env_dict) - self.stage_runtime_devices = self._load_stage_runtime_devices(stage_config_path) - self.stage_ids = stage_ids or self._load_stage_ids(stage_config_path) - if 0 not in self.stage_ids: - raise ValueError(f"Stage CLI test requires stage_id=0 in config: {stage_config_path}") - self.stage_procs: dict[int, subprocess.Popen] = {} - self.proc = None - - @staticmethod - def _load_stage_ids(stage_config_path: str) -> list[int]: - with open(stage_config_path, encoding="utf-8") as f: - cfg = yaml.safe_load(f) or {} - - stage_ids = [stage["stage_id"] for stage in cfg.get("stage_args", []) if "stage_id" in stage] - if not stage_ids: - raise ValueError(f"No stage IDs found in config: {stage_config_path}") - return stage_ids - - @staticmethod - def _load_stage_runtime_devices(stage_config_path: str) -> dict[int, str]: - with open(stage_config_path, encoding="utf-8") as f: - cfg = yaml.safe_load(f) or {} - - runtime_devices: dict[int, str] = {} - for stage in cfg.get("stage_args", []): - stage_id = stage.get("stage_id") - devices = stage.get("runtime", {}).get("devices") - if stage_id is not None and devices: - runtime_devices[int(stage_id)] = str(devices) - return runtime_devices - - @classmethod - def _parse_device_list(cls, devices: str | int) -> list[str]: - if isinstance(devices, int): - if devices < 0: - raise ValueError("Device IDs must be non-negative integers") - return [str(devices)] - return [token.strip() for token in str(devices).split(",") if token.strip()] - - @classmethod - def _load_visible_device_list(cls, env_dict: dict[str, str] | None) -> list[str] | None: - env = os.environ.copy() - if env_dict is not None: - env.update(env_dict) - - env_var = getattr(current_omni_platform, "device_control_env_var", None) - if env_var and env_var in env: - return [token.strip() for token in env[env_var].split(",") if token.strip()] - return None - - @classmethod - def _map_stage_devices(cls, stage_id: int, visible_device_list: list[str] | None, devices: str) -> str: - device_list = cls._parse_device_list(devices) - - if visible_device_list is None: - return ",".join(device_list) - - if not all(device.isdigit() for device in device_list): - raise ValueError("Logical devices must be non-negative integers") - - logical_ids = [int(device) for device in device_list] - if logical_ids and max(logical_ids) >= len(visible_device_list): - raise ValueError( - f"Stage {stage_id} has logical IDs {device_list}, one or more of which exceed the number of visible devices" - ) - - return ",".join(visible_device_list[idx] for idx in logical_ids) - - def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str) -> None: - mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, devices) - env_var = getattr(current_omni_platform, "device_control_env_var", None) - if env_var: - env[env_var] = mapped_devices - - def _build_stage_cmd(self, stage_id: int, *, headless: bool) -> list[str]: - cmd = [ - sys.executable, - "-m", - "vllm_omni.entrypoints.cli.main", - "serve", - self.model, - "--omni", - "--stage-configs-path", - self.stage_config_path, - "--stage-id", - str(stage_id), - "--omni-master-address", - self.host, - "--omni-master-port", - str(self.master_port), - ] - - if headless: - cmd.append("--headless") - else: - cmd += ["--host", self.host, "--port", str(self.port)] - - cmd += self.serve_args - return cmd - - def _launch_stage(self, stage_id: int, *, headless: bool) -> None: - env = os.environ.copy() - env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - if self.env_dict is not None: - env.update(self.env_dict) - - devices = self.stage_runtime_devices.get(stage_id) - if devices: - self._set_stage_device_env(stage_id, env, devices) - - cmd = self._build_stage_cmd(stage_id, headless=headless) - print(f"Launching OmniServerStageCli stage {stage_id}: {' '.join(cmd)}") - proc = subprocess.Popen( - cmd, - env=env, - cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - ) - self.stage_procs[stage_id] = proc - if stage_id == 0: - self.proc = proc - - def _ensure_stage_processes_alive(self) -> None: - for stage_id, proc in self.stage_procs.items(): - ret = proc.poll() - if ret is not None: - raise RuntimeError(f"Stage {stage_id} exited with code {ret} before API server became ready.") - - def _start_server(self) -> None: - ordered_stage_ids = [0, *[stage_id for stage_id in self.stage_ids if stage_id != 0]] - - self._launch_stage(0, headless=False) - time.sleep(2) - self._ensure_stage_processes_alive() - - for stage_id in ordered_stage_ids[1:]: - self._launch_stage(stage_id, headless=True) - - max_wait = 1200 - start_time = time.time() - while time.time() - start_time < max_wait: - self._ensure_stage_processes_alive() - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.settimeout(1) - result = sock.connect_ex((self.host, self.port)) - if result == 0: - print(f"OmniServerStageCli ready on {self.host}:{self.port}") - return - time.sleep(2) - - raise RuntimeError(f"OmniServerStageCli failed to start within {max_wait} seconds") - - def __exit__(self, exc_type, exc_val, exc_tb): - for stage_id in sorted(self.stage_procs, reverse=True): - proc = self.stage_procs[stage_id] - if proc.poll() is None: - self._kill_process_tree(proc.pid) - _run_pre_test_cleanup(enable_force=True) - _run_post_test_cleanup(enable_force=True) - cleanup_dist_env_and_memory() - - -def pytest_addoption(parser): - parser.addoption( - "--run-level", - action="store", - default="core_model", - choices=["core_model", "advanced_model"], - help="Test level to run: L2, L3", - ) - - -@pytest.fixture(scope="session") -def run_level(request) -> str: - """A command-line argument that specifies the level of tests to run in this session. - See https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/CI_5levels/""" - return request.config.getoption("--run-level") - - -_omni_server_lock = threading.Lock() - - -@pytest.fixture(scope="module") -def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: str) -> Generator[OmniServer, Any, None]: - """Start vLLM-Omni through the standard or stage-CLI launcher. - - The fixture stays module-scoped because multi-stage initialization is costly. - The ``use_stage_cli`` flag on ``OmniServerParams`` routes the setup through the - stage-CLI harness while still reusing the same fixture grouping semantics. - """ - with _omni_server_lock: - params: OmniServerParams = request.param - model = model_prefix + params.model - port = params.port - stage_config_path = params.stage_config_path - if run_level == "advanced_model" and stage_config_path is not None: - with open(stage_config_path, encoding="utf-8") as f: - cfg = yaml.safe_load(f) or {} - stage_ids = [stage["stage_id"] for stage in cfg.get("stage_args", []) if "stage_id" in stage] - stage_config_path = modify_stage_config( - stage_config_path, - deletes={"stage_args": {stage_id: ["engine_args.load_format"] for stage_id in stage_ids}}, - ) - - server_args = params.server_args or [] - if params.use_omni and params.stage_init_timeout is not None: - server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)] - else: - server_args = [*server_args, "--stage-init-timeout", "600"] - if params.init_timeout is not None: - server_args = [*server_args, "--init-timeout", str(params.init_timeout)] - else: - server_args = [*server_args, "--init-timeout", "900"] - if params.use_stage_cli: - if not params.use_omni: - raise ValueError("omni_server with use_stage_cli=True requires use_omni=True") - if stage_config_path is None: - raise ValueError("omni_server with use_stage_cli=True requires a stage_config_path") - - with OmniServerStageCli( - model, - stage_config_path, - server_args, - port=port, - env_dict=params.env_dict, - ) as server: - print("OmniServer started successfully") - yield server - print("OmniServer stopping...") - else: - if stage_config_path is not None: - server_args += ["--stage-configs-path", stage_config_path] - - with ( - OmniServer( - model, - server_args, - port=port, - env_dict=params.env_dict, - use_omni=params.use_omni, - ) - if port - else OmniServer( - model, - server_args, - env_dict=params.env_dict, - use_omni=params.use_omni, - ) - ) as server: - print("OmniServer started successfully") - yield server - print("OmniServer stopping...") - - print("OmniServer stopped") - - -@dataclass -class OmniResponse: - text_content: str | None = None - audio_data: list[str] | None = None - audio_content: str | None = None - audio_format: str | None = None - audio_bytes: bytes | None = None - similarity: float | None = None - e2e_latency: float | None = None - success: bool = False - error_message: str | None = None - - -@dataclass -class DiffusionResponse: - text_content: str | None = None - images: list[Image.Image] | None = None - audios: list[Any] | None = None - videos: list[Any] | None = None - e2e_latency: float | None = None - success: bool = False - error_message: str | None = None - - -def _load_gender_pipeline(): - """ - Lazy-load a cached audio-classification pipeline for gender. - - We prefer the pipeline wrapper because it encapsulates processor/model loading - and avoids direct AutoProcessor.from_pretrained call sites in this file. - """ - global _GENDER_PIPELINE - if _GENDER_PIPELINE is not None: - return _GENDER_PIPELINE - - model_name = "7wolf/wav2vec2-base-gender-classification" - try: - # device=-1 forces CPU for pipeline. - _GENDER_PIPELINE = pipeline( - task="audio-classification", - model=model_name, - device=-1, - ) - return _GENDER_PIPELINE - except Exception as exc: # pragma: no cover - best-effort fallback - print(f"Warning: failed to create gender pipeline '{model_name}': {exc}") - _GENDER_PIPELINE = None - return None - - -def _median_pitch_hz_from_autocorr(mono: np.ndarray, sr: int) -> float | None: - """ - Rough median F0 (Hz) over short-time frames. Used to debias wav2vec2 gender head on TTS, - which often labels lower-pitched synthetic speech as female under load or on clean signals. - Returns None if the clip is too short or mostly unvoiced. - """ - x = np.asarray(mono, dtype=np.float64) - x = x - np.mean(x) - if x.size < int(0.15 * sr): - return None - frame_len = int(0.04 * sr) - hop = max(frame_len // 2, 1) - f0_min_hz, f0_max_hz = 70.0, 400.0 - lag_min = max(1, int(sr / f0_max_hz)) - lag_max = min(frame_len - 2, int(sr / f0_min_hz)) - if lag_max <= lag_min: - return None - win = np.hamming(frame_len) - pitches: list[float] = [] - for start in range(0, int(x.shape[0]) - frame_len, hop): - frame = x[start : start + frame_len] * win - frame = frame - np.mean(frame) - if float(np.sqrt(np.mean(frame**2))) < 1e-4: - continue - ac = np.correlate(frame, frame, mode="full")[frame_len - 1 :] - ac = ac / (float(ac[0]) + 1e-12) - region = ac[lag_min : lag_max + 1] - peak_rel = int(np.argmax(region)) - peak_lag = peak_rel + lag_min - if peak_lag <= 0: - continue - f0 = float(sr) / float(peak_lag) - if f0_min_hz <= f0 <= f0_max_hz: - pitches.append(f0) - if len(pitches) < 4: - return None - return float(np.median(np.asarray(pitches, dtype=np.float64))) - - -def _estimate_voice_gender_from_audio(audio_bytes: bytes) -> str: - """ - Estimate voice gender from audio using a small pre-trained classification model. - - Uses a cached `audio-classification` pipeline to classify the clip. - Returns 'male' / 'female' when the model confidence is >= 0.9 and the label - maps to one of these; otherwise returns 'unknown'. If the model is unavailable - or inference fails, returns 'unknown' to keep tests stable. - - Under concurrent tests, a global lock serializes pipeline calls (the HF pipeline is not - thread-safe). A coarse F0 median can correct systematic "male -> female" errors on TTS audio. - """ - data, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32", always_2d=True) - if data.size == 0: - raise ValueError("Empty audio") - mono = np.mean(data, axis=1) - - try: - target_sr = 16000 - if int(sr) != target_sr and mono.size > 1: - src_len = int(mono.shape[0]) - dst_len = max(1, int(round(src_len * float(target_sr) / float(sr)))) - src_idx = np.arange(src_len, dtype=np.float32) - dst_idx = np.linspace(0, src_len - 1, dst_len, dtype=np.float32) - mono = np.interp(dst_idx, src_idx, mono.astype(np.float32, copy=False)).astype(np.float32) - sr = target_sr - - median_f0 = _median_pitch_hz_from_autocorr(mono, sr) - - clf = _load_gender_pipeline() - if clf is None: - print("gender model not available, returning 'unknown'") - return "unknown" - - # transformers pipeline returns a list of {label, score} (highest score first). - with _GENDER_PIPELINE_LOCK: - outputs = clf(mono, sampling_rate=sr) - if not outputs: - return "unknown" - - top = outputs[0] - label = str(top.get("label", "")).lower() - conf = float(top.get("score", 0.0)) - - if conf < 0.5: - gender = "unknown" - # Some models use non-English labels (e.g., Russian). Normalize to 'male'/'female'. - elif ("female" in label) or ("жен" in label): - gender = "female" - elif ("male" in label) or ("муж" in label): - gender = "male" - else: - gender = "unknown" - - # Debias: wav2vec2 gender heads often call TTS / band-limited male speech "female". - # Low median F0 (~speech male range) + female label -> trust pitch when score is not overwhelming. - if gender == "female" and median_f0 is not None and median_f0 < 165.0 and conf < 0.88: - print(f"gender pitch assist: reclassifying female->male (median_f0={median_f0:.1f} Hz, conf={conf:.3f})") - gender = "male" - elif gender == "male" and median_f0 is not None and median_f0 > 230.0 and conf < 0.88: - print(f"gender pitch assist: reclassifying male->female (median_f0={median_f0:.1f} Hz, conf={conf:.3f})") - gender = "female" - - print( - f"gender classifier: label={label}, conf={conf:.3f}, gender={gender}" - + (f", median_f0={median_f0:.1f}Hz" if median_f0 is not None else "") - ) - return gender - except Exception as exc: # pragma: no cover - best-effort fallback - print(f"Warning: gender classification failed, returning 'unknown': {exc}") - return "unknown" - - -_PRESET_VOICE_GENDER_MAP: dict[str, str] = { - "serena": "female", - "uncle_fu": "male", - "chelsie": "female", - "clone": "female", - "ethan": "male", +""" +Root pytest entrypoint for the vLLM-Omni test suite. + +- `tests/conftest.py` stays thin: plugin registration + compatibility re-exports. +- Importable utilities live under `tests/helpers/`. +- Fixtures live under `tests/helpers/fixtures/` and are loaded via `pytest_plugins`. +""" + +from __future__ import annotations + +pytest_plugins = ( + "tests.helpers.fixtures.env", + "tests.helpers.fixtures.log", + "tests.helpers.fixtures.run_args", + "tests.helpers.fixtures.runtime", +) + + +def pytest_terminal_summary(terminalreporter, exitstatus, config): + # Marker for Buildkite log folding before pytest summary lines. + terminalreporter.write_line("--- Running Summary") + + +# Backward-compatible lazy re-exports. +# (Many tests still import from `tests.conftest`; migrate these imports to `tests.helpers.*` over time.) +# Keep these lazy so conftest import does not trigger heavy helper dependencies. +_ASSERTION_EXPORT_NAMES = ( + "assert_audio_speech_response", + "assert_diffusion_response", + "assert_image_diffusion_response", + "assert_image_valid", + "assert_omni_response", + "assert_video_diffusion_response", + "assert_video_valid", +) +_MEDIA_EXPORT_NAMES = ( + "convert_audio_bytes_to_text", + "convert_audio_file_to_text", + "cosine_similarity_text", + "decode_b64_image", + "generate_synthetic_audio", + "generate_synthetic_image", + "generate_synthetic_video", +) +_STAGE_CONFIG_EXPORT_NAMES = ("modify_stage_config",) +_RUNTIME_EXPORT_NAMES = ( + "DiffusionResponse", + "OmniResponse", + "OmniRunner", + "OmniRunnerHandler", + "OmniServer", + "OmniServerParams", + "OmniServerStageCli", + "OpenAIClientHandler", + "dummy_messages_from_mix_data", +) +_LAZY_EXPORT_MODULES = { + **{name: "tests.helpers.assertions" for name in _ASSERTION_EXPORT_NAMES}, + **{name: "tests.helpers.media" for name in _MEDIA_EXPORT_NAMES}, + **{name: "tests.helpers.stage_config" for name in _STAGE_CONFIG_EXPORT_NAMES}, + **{name: "tests.helpers.runtime" for name in _RUNTIME_EXPORT_NAMES}, } - - -def _assert_preset_voice_gender_from_audio( - audio_bytes: bytes | None, - voice_name: str | None, -) -> None: - """If ``voice_name`` matches a known preset, assert classifier gender matches (skip when unknown).""" - if not voice_name or not audio_bytes: - return - key = str(voice_name).lower() - expected_gender = _PRESET_VOICE_GENDER_MAP.get(key) - if expected_gender is None: - return - estimated_gender = _estimate_voice_gender_from_audio(audio_bytes) - print(f"Preset voice gender check: preset={key!r}, estimated={estimated_gender!r}, expected={expected_gender!r}") - if estimated_gender != "unknown": - assert estimated_gender == expected_gender, ( - f"{voice_name!r} is expected {expected_gender}, but estimated gender is {estimated_gender!r}" - ) - - -# Threshold aligned with _compute_pcm_hnr_db docstring (clean clone vs distorted). -_MIN_PCM_SPEECH_HNR_DB = 1.0 - - -def _compute_pcm_hnr_db(pcm_samples: np.ndarray, sr: int = _PCM_SPEECH_SAMPLE_RATE_HZ) -> float: - """Compute mean Harmonic-to-Noise Ratio (dB) for speech quality. - - Clean cloned speech has HNR > 1.2 dB; distorted speech (e.g. lost - ref_code decoder context) drops below 1.0 dB. - """ - frame_len = int(0.03 * sr) # 30ms frames - hop = frame_len // 2 - hnr_values: list[float] = [] - - for start in range(0, len(pcm_samples) - frame_len, hop): - frame = pcm_samples[start : start + frame_len].astype(np.float32, copy=False) - frame = frame - np.mean(frame) - if np.max(np.abs(frame)) < 0.01: - continue - ac = np.correlate(frame, frame, mode="full")[len(frame) - 1 :] - ac = ac / (ac[0] + 1e-10) - min_lag = int(sr / 400) - max_lag = min(int(sr / 80), len(ac)) - if min_lag >= max_lag: - continue - peak = float(np.max(ac[min_lag:max_lag])) - if 0 < peak < 1: - hnr_values.append(10 * np.log10(peak / (1 - peak + 1e-10))) - - return float(np.mean(hnr_values)) if hnr_values else 0.0 - - -def _assert_pcm_int16_speech_hnr(audio_bytes: bytes) -> None: - """Validate harmonic-to-noise ratio on raw int16 PCM from /v1/audio/speech.""" - assert audio_bytes is not None and len(audio_bytes) >= 2, "missing PCM bytes" - assert len(audio_bytes) % 2 == 0, "PCM byte length must be aligned to int16" - pcm_samples = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32) / 32768.0 - hnr = _compute_pcm_hnr_db(pcm_samples) - print(f"PCM speech HNR: {hnr:.2f} dB (threshold: {_MIN_PCM_SPEECH_HNR_DB} dB)") - assert hnr >= _MIN_PCM_SPEECH_HNR_DB, ( - f"Audio distortion detected: HNR={hnr:.2f} dB < {_MIN_PCM_SPEECH_HNR_DB} dB. " - "Voice clone decoder may be losing ref_code speaker context on later chunks." - ) - - -def assert_omni_response(response: OmniResponse, request_config: dict[str, Any], run_level): - """ - Validate response results. - - Args: - response: OmniResponse object - - Raises: - AssertionError: When the response does not meet validation criteria - """ - assert response.success, "The request failed." - e2e_latency = response.e2e_latency - if e2e_latency is not None: - print(f"the e2e latency is: {e2e_latency}") - - modalities = request_config.get("modalities", ["text", "audio"]) - - if run_level == "advanced_model": - if "audio" in modalities: - assert response.audio_content is not None, "No audio output is generated" - print(f"audio content is: {response.audio_content}") - speaker = request_config.get("speaker") - if speaker: - _assert_preset_voice_gender_from_audio( - response.audio_bytes, - speaker, - ) - - if "text" in modalities: - assert response.text_content is not None, "No text output is generated" - print(f"text content is: {response.text_content}") - - # Verify image description - word_types = ["text", "image", "audio", "video"] - keywords_dict = request_config.get("key_words", {}) - for word_type in word_types: - keywords = keywords_dict.get(word_type) - if "text" in modalities: - if keywords: - text_lower = response.text_content.lower() - assert any(str(kw).lower() in text_lower for kw in keywords), ( - "The output does not contain any of the keywords." - ) - else: - if keywords: - audio_lower = response.audio_content.lower() - assert any(str(kw).lower() in audio_lower for kw in keywords), ( - "The output does not contain any of the keywords." - ) - - # Verify similarity (Whisper transcript vs streamed/detokenized text) - if "text" in modalities and "audio" in modalities: - assert response.similarity is not None and response.similarity > 0.9, ( - "The audio content is not same as the text" - ) - print(f"similarity is: {response.similarity}") - - -def assert_audio_speech_response( - response: OmniResponse, - request_config: dict[str, Any], - run_level: str, -) -> None: - """ - Validate /v1/audio/speech response: success, optional format check, transcription similarity - and gender (non-PCM only for advanced_model), and int16 PCM HNR when response_format is pcm. - """ - assert response.success, "The request failed." - - req_fmt = request_config.get("response_format") - - if req_fmt == "pcm" and response.audio_bytes: - _assert_pcm_int16_speech_hnr(response.audio_bytes) - if response.audio_format: - assert "pcm" in response.audio_format.lower(), ( - f"Expected audio/pcm content-type, got {response.audio_format!r}" - ) - - elif req_fmt == "wav" and response.audio_format: - assert req_fmt in response.audio_format, ( - f"The response audio format {response.audio_format} don't match the request audio format {req_fmt}" - ) - - e2e_latency = response.e2e_latency - if e2e_latency is not None: - print(f"the avg e2e latency is: {e2e_latency}") - - if run_level == "advanced_model" and req_fmt != "pcm": - # Text–audio semantic similarity check (skipped for raw PCM: no Whisper transcript). - expected_text = request_config.get("input") - if expected_text: - transcript = (response.audio_content or "").strip() - print(f"audio content is: {transcript}") - print(f"input text is: {expected_text}") - similarity = cosine_similarity_text(transcript.lower(), expected_text.lower()) - print(f"Cosine similarity: {similarity:.3f}") - assert similarity > 0.9, ( - f"Transcript doesn't match input: similarity={similarity:.2f}, transcript='{transcript}'" - ) - - # Voice gender consistency check (preset names in ``_PRESET_VOICE_GENDER_MAP``). - # When the estimator returns 'unknown', we treat it as inconclusive and do NOT fail the test. - _assert_preset_voice_gender_from_audio( - response.audio_bytes, - request_config.get("voice"), - ) - - -def assert_diffusion_response(response: DiffusionResponse, request_config: dict[str, Any], run_level: str = None): - """ - Validate diffusion response results. - - Dispatcher that routes validation to modality-specific assert functions. - - Args: - response: DiffusionResponse object. - request_config: Request configuration dictionary. - run_level: Test run level (e.g. "core_model", "advanced_model") - - Raises: - AssertionError: When the response does not meet validation criteria - KeyError: When the request_config does not contain necessary parameters for validation - """ - assert response.success, "The request failed." - - e2e_latency = response.e2e_latency - if e2e_latency is not None: - print(f"the avg e2e is: {e2e_latency}") - - has_any_content = any(content is not None for content in (response.images, response.videos, response.audios)) - assert has_any_content, "Response contains no images, videos, or audios" - - if response.images is not None: - assert_image_diffusion_response( - response=response, - request_config=request_config, - run_level=run_level, - ) - - if response.videos is not None: - assert_video_diffusion_response( - response=response, - request_config=request_config, - run_level=run_level, - ) - - if response.audios is not None: - assert_audio_diffusion_response( - response=response, - request_config=request_config, - run_level=run_level, - ) - - -class OpenAIClientHandler: - """ - OpenAI client handler class, encapsulating both streaming and non-streaming response processing logic. - - This class integrates OpenAI API request sending, response handling, and validation functionality, - supporting both single request and concurrent request modes. - """ - - def __init__( - self, host: str = "127.0.0.1", port: int = get_open_port(), api_key: str = "EMPTY", run_level: str = None - ): - """ - Initialize the OpenAI client. - - Args: - host: vLLM-Omni server host address - port: vLLM-Omni server port - api_key: API key (defaults to "EMPTY") - """ - self.base_url = f"http://{host}:{port}" - self.client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key=api_key) - self.run_level = run_level - - def _process_stream_omni_response(self, chat_completion) -> OmniResponse: - """ - Process streaming responses. - - Args: - chat_completion: OpenAI streaming response object - request_config: Request configuration dictionary - - Returns: - OmniResponse: Processed response object - """ - result = OmniResponse() - start_time = time.perf_counter() - - try: - text_content = "" - audio_data = [] - - for chunk in chat_completion: - for choice in chunk.choices: - # Get content data - if hasattr(choice, "delta"): - content = getattr(choice.delta, "content", None) - else: - content = None - - # Get modality type - modality = getattr(chunk, "modality", None) - - # Process content based on modality type - if modality == "audio" and content: - audio_data.append(content) - elif modality == "text" and content: - text_content += content if content else "" - - # Calculate end-to-end latency - result.e2e_latency = time.perf_counter() - start_time - - # Process audio and text content - audio_content = None - similarity = None - - if audio_data or text_content: - if audio_data: - merged_seg = _merge_base64_audio_to_segment(audio_data) - wav_buf = BytesIO() - merged_seg.export(wav_buf, format="wav") - result.audio_bytes = wav_buf.getvalue() - audio_content = convert_audio_bytes_to_text(result.audio_bytes) - if audio_content and text_content: - similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) - - # Populate result object - result.text_content = text_content - result.audio_data = audio_data - result.audio_content = audio_content - result.similarity = similarity - result.success = True - - except Exception as e: - result.error_message = f"Stream processing error: {str(e)}" - print(f"Error: {result.error_message}") - - return result - - def _process_non_stream_omni_response(self, chat_completion) -> OmniResponse: - """ - Process non-streaming responses. - - Args: - chat_completion: OpenAI non-streaming response object - request_config: Request configuration dictionary - - Returns: - OmniResponse: Processed response object - """ - result = OmniResponse() - start_time = time.perf_counter() - - try: - audio_data = None - text_content = None - - # Iterate through all choices - for choice in chat_completion.choices: - # Process audio data - if hasattr(choice.message, "audio") and choice.message.audio is not None: - audio_message = choice.message - audio_data = audio_message.audio.data - - # Process text content - if hasattr(choice.message, "content") and choice.message.content is not None: - text_content = choice.message.content - - # Calculate end-to-end latency - result.e2e_latency = time.perf_counter() - start_time - - # Process audio and text content - audio_content = None - similarity = None - - if audio_data or text_content: - if audio_data: - result.audio_bytes = base64.b64decode(audio_data) - audio_content = convert_audio_bytes_to_text(result.audio_bytes) - if audio_content and text_content: - similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) - - # Populate result object - result.text_content = text_content - result.audio_content = audio_content - result.similarity = similarity - result.success = True - - except Exception as e: - result.error_message = f"Non-stream processing error: {str(e)}" - print(f"Error: {result.error_message}") - - return result - - def _process_diffusion_response(self, chat_completion) -> DiffusionResponse: - """ - Process diffusion responses (image generation/editing). - - Args: - chat_completion: OpenAI response object - - Returns: - DiffusionResponse: Processed response object - """ - result = DiffusionResponse() - start_time = time.perf_counter() - - try: - images = [] - # [TODO] reading video and audio output from API response for later validation - - for choice in chat_completion.choices: - if hasattr(choice.message, "content") and choice.message.content is not None: - content = choice.message.content - if isinstance(content, list): - for item in content: - if isinstance(item, dict): - image_url = item.get("image_url", {}).get("url") - else: - image_url_obj = getattr(item, "image_url", None) - image_url = hasattr(image_url_obj, "url", None) if image_url_obj else None - if image_url and image_url.startswith("data:image"): - b64_data = image_url.split(",", 1)[1] - img = decode_b64_image(b64_data) - images.append(img) - - result.e2e_latency = time.perf_counter() - start_time - result.images = images if images else None - result.success = True - - except Exception as e: - result.error_message = f"Diffusion response processing error: {str(e)}" - print(f"Error: {result.error_message}") - - return result - - def _process_stream_audio_speech_response(self, response, *, response_format: str | None = None) -> OmniResponse: - """ - Process streaming /v1/audio/speech responses into an OmniResponse. - - This mirrors _process_stream_omni_response but operates on low-level - audio bytes and produces an OmniResponse with audio_content filled - from Whisper transcription. - """ - result = OmniResponse() - start_time = time.perf_counter() - - try: - # Aggregate all audio bytes from the streaming response. - data = bytearray() - - # Preferred OpenAI helper. - if hasattr(response, "iter_bytes") and callable(getattr(response, "iter_bytes")): - for chunk in response.iter_bytes(): - if chunk: - data.extend(chunk) - else: - # Generic iterable-of-bytes fallback (e.g., generator or list of chunks). - try: - iterator = iter(response) - except TypeError: - iterator = None - - if iterator is not None: - for chunk in iterator: - if not chunk: - continue - if isinstance(chunk, (bytes, bytearray)): - data.extend(chunk) - elif hasattr(chunk, "data"): - data.extend(chunk.data) # type: ignore[arg-type] - elif hasattr(chunk, "content"): - data.extend(chunk.content) # type: ignore[arg-type] - else: - raise TypeError(f"Unsupported stream chunk type: {type(chunk)}") - else: - raise TypeError(f"Unsupported audio speech streaming response type: {type(response)}") - - raw_bytes = bytes(data) - if response_format == "pcm": - transcript = None - else: - transcript = convert_audio_bytes_to_text(raw_bytes) - - # Populate OmniResponse. - result.audio_bytes = raw_bytes - result.audio_content = transcript - result.e2e_latency = time.perf_counter() - start_time - result.success = True - result.audio_format = getattr(response, "response", None) - if result.audio_format is not None: - result.audio_format = result.audio_format.headers.get("content-type", "") - - except Exception as e: - result.error_message = f"Audio speech stream processing error: {str(e)}" - print(f"Error: {result.error_message}") - - return result - - def _process_non_stream_audio_speech_response( - self, response, *, response_format: str | None = None - ) -> OmniResponse: - """ - Process non-streaming /v1/audio/speech responses into an OmniResponse. - - This mirrors _process_non_stream_omni_response but for the binary - audio payload returned by audio.speech.create. - """ - result = OmniResponse() - start_time = time.perf_counter() - - try: - # OpenAI non-streaming audio.speech.create returns HttpxBinaryResponseContent (.read() or .content) - if hasattr(response, "read") and callable(getattr(response, "read")): - raw_bytes = response.read() - elif hasattr(response, "content"): - raw_bytes = response.content # type: ignore[assignment] - else: - raise TypeError(f"Unsupported audio speech response type: {type(response)}") - - if response_format == "pcm": - transcript = None - else: - transcript = convert_audio_bytes_to_text(raw_bytes) - - result.audio_bytes = raw_bytes - result.audio_content = transcript - result.e2e_latency = time.perf_counter() - start_time - result.success = True - result.audio_format = getattr(response, "response", None) - if result.audio_format is not None: - result.audio_format = result.audio_format.headers.get("content-type", "") - - except Exception as e: - result.error_message = f"Audio speech non-stream processing error: {str(e)}" - print(f"Error: {result.error_message}") - - return result - - def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: - """ - Send OpenAI requests. - - Args: - request_config: Request configuration dictionary containing parameters like model, messages, stream. - Optional ``use_audio_in_video`` (bool): when true, sets - ``extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}`` for Qwen-Omni video+audio - extraction. - Optional top-level ``speaker`` (str): Qwen3-Omni preset TTS speaker name; sent as - ``extra_body["speaker"]`` to ``chat.completions.create``. - request_num: Number of requests, defaults to 1 (single request) - - Returns: - List[OmniResponse]: List of response objects - """ - - responses = [] - stream = request_config.get("stream", False) - modalities = request_config.get("modalities", ["text", "audio"]) - - extra_body: dict[str, Any] = {} - if "speaker" in request_config: - extra_body["speaker"] = request_config["speaker"] - if request_config.get("use_audio_in_video"): - mm = dict(extra_body.get("mm_processor_kwargs") or {}) - mm["use_audio_in_video"] = True - extra_body["mm_processor_kwargs"] = mm - extra_body_arg: dict[str, Any] | None = extra_body if extra_body else None - - create_kwargs: dict[str, Any] = { - "model": request_config.get("model"), - "messages": request_config.get("messages"), - "stream": stream, - "modalities": modalities, - } - if extra_body_arg is not None: - create_kwargs["extra_body"] = extra_body_arg - - if request_num == 1: - # Send single request - chat_completion = self.client.chat.completions.create(**create_kwargs) - - if stream: - response = self._process_stream_omni_response(chat_completion) - else: - response = self._process_non_stream_omni_response(chat_completion) - - assert_omni_response(response, request_config, run_level=self.run_level) - responses.append(response) - - else: - # Send concurrent requests: run create + process in worker so e2e_latency includes full round-trip. - def _one_omni_request(): - start = time.perf_counter() - worker_kwargs: dict[str, Any] = { - "model": request_config.get("model"), - "messages": request_config.get("messages"), - "modalities": modalities, - "stream": stream, - } - if extra_body_arg is not None: - worker_kwargs["extra_body"] = extra_body_arg - chat_completion = self.client.chat.completions.create(**worker_kwargs) - if stream: - response = self._process_stream_omni_response(chat_completion) - else: - response = self._process_non_stream_omni_response(chat_completion) - response.e2e_latency = time.perf_counter() - start - return response - - with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: - futures = [executor.submit(_one_omni_request) for _ in range(request_num)] - for future in concurrent.futures.as_completed(futures): - response = future.result() - assert_omni_response(response, request_config, run_level=self.run_level) - responses.append(response) - - return responses - - def send_audio_speech_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: - """ - Call the /v1/audio/speech endpoint using the same configuration-dict - style as send_omni_request, but via the OpenAI Python client's - audio.speech APIs. - - Expected keys in request_config: - - model: model name/path (required) - - input: text to synthesize (required) - - response_format: audio format such as "wav" or "pcm" (optional) - - task_type, ref_text, ref_audio: TTS-specific extras (optional, passed via extra_body) - - timeout: request timeout in seconds (float, optional, default 120.0) - - stream: whether to use streaming API (bool, optional, default False) - """ - timeout = float(request_config.get("timeout", 120.0)) - - model = request_config["model"] - text_input = request_config["input"] - stream = bool(request_config.get("stream", False)) - voice = request_config.get("voice", None) - - # Standard OpenAI param: use omit when not provided to keep default behavior. - response_format = request_config.get("response_format", omit) - - # Qwen3-TTS custom fields, forwarded via extra_body. - extra_body: dict[str, Any] = {} - # Keep this list aligned with vllm_omni.entrypoints.openai.protocol.audio params. - for key in ("task_type", "ref_text", "ref_audio", "language", "max_new_tokens"): - if key in request_config: - extra_body[key] = request_config[key] - - responses: list[OmniResponse] = [] - - speech_fmt: str | None = None if response_format is omit else str(response_format).lower() - - if request_num == 1: - if stream: - # Use streaming response helper. - with self.client.audio.speech.with_streaming_response.create( - model=model, - input=text_input, - response_format=response_format, - extra_body=extra_body or None, - timeout=timeout, - voice=voice, - ) as resp: - omni_resp = self._process_stream_audio_speech_response(resp, response_format=speech_fmt) - else: - # Non-streaming response. - resp = self.client.audio.speech.create( - model=model, - input=text_input, - response_format=response_format, - extra_body=extra_body or None, - timeout=timeout, - voice=voice, - ) - omni_resp = self._process_non_stream_audio_speech_response(resp, response_format=speech_fmt) - - assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) - responses.append(omni_resp) - return responses - else: - # request_num > 1: concurrent requests (use same params as single-request path) - - if stream: - - def _stream_task(): - with self.client.audio.speech.with_streaming_response.create( - model=model, - input=text_input, - response_format=response_format, - extra_body=extra_body or None, - timeout=timeout, - voice=voice, - ) as resp: - return self._process_stream_audio_speech_response(resp, response_format=speech_fmt) - - with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: - futures = [executor.submit(_stream_task) for _ in range(request_num)] - for future in concurrent.futures.as_completed(futures): - omni_resp = future.result() - assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) - responses.append(omni_resp) - else: - with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: - futures = [] - for _ in range(request_num): - future = executor.submit( - self.client.audio.speech.create, - model=model, - input=text_input, - response_format=response_format, - extra_body=extra_body or None, - timeout=timeout, - voice=voice, - ) - futures.append(future) - - for future in concurrent.futures.as_completed(futures): - resp = future.result() - omni_resp = self._process_non_stream_audio_speech_response(resp, response_format=speech_fmt) - assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) - responses.append(omni_resp) - - return responses - - def send_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: - """ - Send OpenAI requests for diffusion models. - - Args: - request_config: Request configuration dictionary containing parameters like model, messages - request_num: Number of requests to send concurrently, defaults to 1 (single request) - Returns: - List[OmniResponse]: List of response objects - """ - responses = [] - stream = request_config.get("stream", False) - modalities = request_config.get("modalities", omit) # Most diffusion models don't require modalities param - extra_body = request_config.get("extra_body", None) - - if stream: - raise NotImplementedError("Streaming is not currently implemented for diffusion model e2e test") - - if request_num == 1: - # Send single request - chat_completion = self.client.chat.completions.create( - model=request_config.get("model"), - messages=request_config.get("messages"), - extra_body=extra_body, - modalities=modalities, - ) - - response = self._process_diffusion_response(chat_completion) - assert_diffusion_response(response, request_config, run_level=self.run_level) - responses.append(response) - - else: - # Send concurrent requests - with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: - futures = [] - - # Submit all request tasks - for _ in range(request_num): - future = executor.submit( - self.client.chat.completions.create, - model=request_config.get("model"), - messages=request_config.get("messages"), - modalities=modalities, - extra_body=extra_body, - ) - futures.append(future) - - # Process completed tasks - for future in concurrent.futures.as_completed(futures): - chat_completion = future.result() - response = self._process_diffusion_response(chat_completion) - assert_diffusion_response(response, request_config, run_level=self.run_level) - responses.append(response) - - return responses - - def send_video_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: - """ - Send native /v1/videos requests. - """ - if request_num != 1: - raise NotImplementedError("Concurrent video diffusion requests are not currently implemented") - - if request_config.get("stream", False): - raise NotImplementedError("Streaming is not currently implemented for video diffusion e2e test") - - form_data = request_config.get("form_data") - if not isinstance(form_data, dict): - raise ValueError("Video request_config must contain 'form_data'") - - if not form_data.get("prompt"): - raise ValueError("Video request_config['form_data'] must contain 'prompt'") - - normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None} - - files: dict[str, tuple[str, BytesIO, str]] = {} - image_reference = request_config.get("image_reference") - if image_reference: - if image_reference.startswith("data:image"): - header, encoded = image_reference.split(",", 1) - content_type = header.split(";")[0].removeprefix("data:") - extension = content_type.split("/")[-1] - file_data = base64.b64decode(encoded) - - files["input_reference"] = ( - f"reference.{extension}", - BytesIO(file_data), - content_type, - ) - else: - normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference}) - - result = DiffusionResponse() - start_time = time.perf_counter() - - try: - create_url = self._build_url("/v1/videos") - response = requests.post( - create_url, - data=normalized_form_data, - files=files, - headers={"Accept": "application/json"}, - timeout=60, - ) - response.raise_for_status() - - job_data = response.json() - video_id = job_data["id"] - - self._wait_until_video_completed(video_id) - - video_content = self._download_video_content(video_id) - - result.success = True - result.videos = [video_content] - result.e2e_latency = time.perf_counter() - start_time - - assert_diffusion_response(result, request_config, run_level=self.run_level) - - except Exception as e: - result.success = False - result.error_message = f"Diffusion response processing error: {e}" - assert False, result.error_message - - return [result] - - def _wait_until_video_completed( - self, - video_id: str, - poll_interval_seconds: int = 2, - timeout_seconds: int = 300, - ) -> None: - status_url = self._build_url(f"/v1/videos/{video_id}") - deadline = time.monotonic() + timeout_seconds - - while time.monotonic() < deadline: - status_resp = requests.get( - status_url, - headers={"Accept": "application/json"}, - timeout=30, - ) - status_resp.raise_for_status() - - status_data = status_resp.json() - current_status = status_data["status"] - - if current_status == "completed": - return - - if current_status == "failed": - error_msg = status_data.get("last_error", "Unknown error") - raise RuntimeError(f"Job failed: {error_msg}") - - time.sleep(poll_interval_seconds) - - raise TimeoutError(f"Video job {video_id} did not complete within {timeout_seconds}s") - - def _download_video_content(self, video_id: str) -> bytes: - download_url = self._build_url(f"/v1/videos/{video_id}/content") - video_resp = requests.get(download_url, stream=True, timeout=60) - video_resp.raise_for_status() - - video_bytes = BytesIO() - for chunk in video_resp.iter_content(chunk_size=8192): - if chunk: - video_bytes.write(chunk) - - return video_bytes.getvalue() - - def _build_url(self, path: str) -> str: - return f"{self.base_url.rstrip('/')}/{path.lstrip('/')}" - - -@pytest.fixture -def openai_client(request: pytest.FixtureRequest, run_level: str): - """Create OpenAIClientHandler fixture to facilitate communication with OmniServer - with encapsulated request sending, concurrent requests, response handling, and validation.""" - server = request.getfixturevalue("omni_server") - return OpenAIClientHandler(host=server.host, port=server.port, api_key="EMPTY", run_level=run_level) - - -class OmniRunner: - """ - Offline test runner for Omni models. - """ - - def __init__( - self, - model_name: str, - seed: int = 42, - stage_init_timeout: int = 600, - batch_timeout: int = 10, - init_timeout: int = 900, - shm_threshold_bytes: int = 65536, - log_stats: bool = False, - stage_configs_path: str | None = None, - **kwargs, - ) -> None: - """ - Initialize an OmniRunner for testing. - - Args: - model_name: The model name or path - seed: Random seed for reproducibility - stage_init_timeout: Timeout for initializing a single stage in seconds - batch_timeout: Timeout for batching in seconds - init_timeout: Timeout for initializing stages in seconds - shm_threshold_bytes: Threshold for using shared memory - log_stats: Enable detailed statistics logging - stage_configs_path: Optional path to YAML stage config file - **kwargs: Additional arguments passed to Omni - """ - cleanup_dist_env_and_memory() - _run_pre_test_cleanup(enable_force=True) - _run_post_test_cleanup(enable_force=True) - self.model_name = model_name - self.seed = seed - - self.omni = Omni( - model=model_name, - log_stats=log_stats, - stage_init_timeout=stage_init_timeout, - batch_timeout=batch_timeout, - init_timeout=init_timeout, - shm_threshold_bytes=shm_threshold_bytes, - stage_configs_path=stage_configs_path, - **kwargs, - ) - - def _estimate_prompt_len( - self, - additional_information: dict[str, Any], - model_name: str, - _cache: dict[str, Any] = {}, - ) -> int: - """Estimate prompt_token_ids placeholder length for the Talker stage. - - The AR Talker replaces all input embeddings via ``preprocess``, so the - placeholder values are irrelevant but the **length** must match the - embeddings that ``preprocess`` will produce. - """ - try: - from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig - from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( - Qwen3TTSTalkerForConditionalGeneration, - ) - - if model_name not in _cache: - from transformers import AutoTokenizer - - tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left") - cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True) - _cache[model_name] = (tok, getattr(cfg, "talker_config", None)) - - tok, tcfg = _cache[model_name] - task_type = (additional_information.get("task_type") or ["CustomVoice"])[0] - return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( - additional_information=additional_information, - task_type=task_type, - tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"], - codec_language_id=getattr(tcfg, "codec_language_id", None), - spk_is_dialect=getattr(tcfg, "spk_is_dialect", None), - ) - except Exception as exc: - logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc) - return 2048 - - def get_default_sampling_params_list(self) -> list[OmniSamplingParams]: - """ - Get a list of default sampling parameters for all stages. - - Returns: - List of SamplingParams with default decoding for each stage - """ - if not hasattr(self.omni, "default_sampling_params_list"): - raise AttributeError("Omni.default_sampling_params_list is not available") - return list(self.omni.default_sampling_params_list) - - def get_omni_inputs( - self, - prompts: list[str] | str, - system_prompt: str | None = None, - audios: PromptAudioInput = None, - images: PromptImageInput = None, - videos: PromptVideoInput = None, - mm_processor_kwargs: dict[str, Any] | None = None, - modalities: list[str] | None = None, - ) -> list[TextPrompt]: - """ - Construct Omni input format from prompts and multimodal data. - - Args: - prompts: Text prompt(s) - either a single string or list of strings - system_prompt: Optional system prompt (defaults to Qwen system prompt) - audios: Audio input(s) - tuple of (audio_array, sample_rate) or list of tuples - images: Image input(s) - PIL Image or list of PIL Images - videos: Video input(s) - numpy array or list of numpy arrays - mm_processor_kwargs: Optional processor kwargs (e.g., use_audio_in_video) - - Returns: - List of prompt dictionaries suitable for Omni.generate() - """ - if system_prompt is None: - system_prompt = ( - "You are Qwen, a virtual human developed by the Qwen Team, Alibaba " - "Group, capable of perceiving auditory and visual inputs, as well as " - "generating text and speech." - ) - - video_padding_token = "<|VIDEO|>" - image_padding_token = "<|IMAGE|>" - audio_padding_token = "<|AUDIO|>" - - if "Qwen3-Omni-30B-A3B-Instruct" in self.model_name: - video_padding_token = "<|video_pad|>" - image_padding_token = "<|image_pad|>" - audio_padding_token = "<|audio_pad|>" - - if isinstance(prompts, str): - prompts = [prompts] - - # Qwen-TTS: follow examples/offline_inference/qwen3_tts/end2end.py style. - # Stage 0 expects token placeholders + additional_information (text/speaker/task_type/...), - # and Talker replaces embeddings in preprocess based on additional_information only. - is_tts_model = "Qwen3-TTS" in self.model_name or "qwen3_tts" in self.model_name.lower() - if is_tts_model and modalities == ["audio"]: - tts_kw = mm_processor_kwargs or {} - task_type = tts_kw.get("task_type", "CustomVoice") - speaker = tts_kw.get("speaker", "Vivian") - language = tts_kw.get("language", "Auto") - max_new_tokens = int(tts_kw.get("max_new_tokens", 2048)) - ref_audio = tts_kw.get("ref_audio", None) - ref_text = tts_kw.get("ref_text", None) - - omni_inputs: list[TextPrompt] = [] - for prompt_text in prompts: - text_str = str(prompt_text).strip() or " " - additional_information: dict[str, Any] = { - "task_type": [task_type], - "text": [text_str], - "language": [language], - "speaker": [speaker], - "max_new_tokens": [max_new_tokens], - } - if ref_audio is not None: - additional_information["ref_audio"] = [ref_audio] - if ref_text is not None: - additional_information["ref_text"] = [ref_text] - # Use official helper to get correct placeholder length - plen = self._estimate_prompt_len(additional_information, self.model_name) - input_dict: TextPrompt = { - "prompt_token_ids": [0] * plen, - "additional_information": additional_information, - } - omni_inputs.append(input_dict) - return omni_inputs - - def _normalize_mm_input(mm_input, num_prompts): - if mm_input is None: - return [None] * num_prompts - if isinstance(mm_input, list): - if len(mm_input) != num_prompts: - raise ValueError( - f"Multimodal input list length ({len(mm_input)}) must match prompts length ({num_prompts})" - ) - return mm_input - return [mm_input] * num_prompts - - num_prompts = len(prompts) - audios_list = _normalize_mm_input(audios, num_prompts) - images_list = _normalize_mm_input(images, num_prompts) - videos_list = _normalize_mm_input(videos, num_prompts) - - omni_inputs = [] - for i, prompt_text in enumerate(prompts): - user_content = "" - multi_modal_data = {} - - audio = audios_list[i] - if audio is not None: - if isinstance(audio, list): - for _ in audio: - user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>" - multi_modal_data["audio"] = audio - else: - user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>" - multi_modal_data["audio"] = audio - - image = images_list[i] - if image is not None: - if isinstance(image, list): - for _ in image: - user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>" - multi_modal_data["image"] = image - else: - user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>" - multi_modal_data["image"] = image - - video = videos_list[i] - if video is not None: - if isinstance(video, list): - for _ in video: - user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>" - multi_modal_data["video"] = video - else: - user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>" - multi_modal_data["video"] = video - - user_content += prompt_text - - full_prompt = ( - f"<|im_start|>system\n{system_prompt}<|im_end|>\n" - f"<|im_start|>user\n{user_content}<|im_end|>\n" - f"<|im_start|>assistant\n" - ) - - input_dict: TextPrompt = {"prompt": full_prompt} - if multi_modal_data: - input_dict["multi_modal_data"] = multi_modal_data - if modalities: - input_dict["modalities"] = modalities - if mm_processor_kwargs: - input_dict["mm_processor_kwargs"] = mm_processor_kwargs - - omni_inputs.append(input_dict) - - return omni_inputs - - def generate( - self, - prompts: list[TextPrompt], - sampling_params_list: list[OmniSamplingParams] | None = None, - ) -> list[OmniRequestOutput]: - """ - Generate outputs for the given prompts. - - Args: - prompts: List of prompt dictionaries with 'prompt' and optionally - 'multi_modal_data' keys - sampling_params_list: List of sampling parameters for each stage. - If None, uses default parameters. - - Returns: - List of OmniRequestOutput objects from stages with final_output=True - """ - if sampling_params_list is None: - sampling_params_list = self.get_default_sampling_params_list() - - return self.omni.generate(prompts, sampling_params_list) - - def generate_multimodal( - self, - prompts: list[str] | str, - sampling_params_list: list[OmniSamplingParams] | None = None, - system_prompt: str | None = None, - audios: PromptAudioInput = None, - images: PromptImageInput = None, - videos: PromptVideoInput = None, - mm_processor_kwargs: dict[str, Any] | None = None, - modalities: list[str] | None = None, - ) -> list[OmniRequestOutput]: - """ - Convenience method to generate with multimodal inputs. - - Args: - prompts: Text prompt(s) - sampling_params_list: List of sampling parameters for each stage - system_prompt: Optional system prompt - audios: Audio input(s) - images: Image input(s) - videos: Video input(s) - mm_processor_kwargs: Optional processor kwargs - - Returns: - List of OmniRequestOutput objects from stages with final_output=True - """ - omni_inputs = self.get_omni_inputs( - prompts=prompts, - system_prompt=system_prompt, - audios=audios, - images=images, - videos=videos, - mm_processor_kwargs=mm_processor_kwargs, - modalities=modalities, - ) - return self.generate(omni_inputs, sampling_params_list) - - def start_profile( - self, - profile_prefix: str | None = None, - stages: list[int] | None = None, - ) -> list[Any]: - """Start profiling specified stages. - - Args: - profile_prefix: Optional prefix for the trace file names. - stages: List of stage IDs to profile. If None, profiles all stages. - - Returns: - List of results from each stage. - """ - return self.omni.start_profile(profile_prefix=profile_prefix, stages=stages) - - def stop_profile(self, stages: list[int] | None = None) -> list[Any]: - """Stop profiling specified stages. - - Args: - stages: List of stage IDs to profile. If None, stops all stages. - - Returns: - List of results from each stage. - """ - return self.omni.stop_profile(stages=stages) - - def _cleanup_process(self): - try: - keywords = ["enginecore"] - matched = [] - - for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]): - try: - cmdline = " ".join(proc.cmdline()).lower() if proc.cmdline() else "" - name = proc.name().lower() - - is_process = any(keyword in cmdline for keyword in keywords) or any( - keyword in name for keyword in keywords - ) - - if is_process: - print(f"Found vllm process: PID={proc.pid}, cmd={cmdline[:100]}") - matched.append(proc) - except (psutil.NoSuchProcess, psutil.AccessDenied): - pass - - for proc in matched: - try: - proc.terminate() - except (psutil.NoSuchProcess, psutil.AccessDenied): - pass - - _, still_alive = psutil.wait_procs(matched, timeout=5) - for proc in still_alive: - try: - proc.kill() - except (psutil.NoSuchProcess, psutil.AccessDenied): - pass - - if still_alive: - _, stubborn = psutil.wait_procs(still_alive, timeout=3) - if stubborn: - print(f"Warning: failed to kill residual vllm pids: {[p.pid for p in stubborn]}") - else: - print(f"Force-killed residual vllm pids: {[p.pid for p in still_alive]}") - elif matched: - print(f"Terminated vllm pids: {[p.pid for p in matched]}") - - except Exception as e: - print(f"Error in psutil vllm cleanup: {e}") - - def __enter__(self): - """Context manager entry.""" - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - """Context manager exit - cleanup resources.""" - if hasattr(self.omni, "close"): - self.omni.close() - self._cleanup_process() - _run_pre_test_cleanup(enable_force=True) - _run_post_test_cleanup(enable_force=True) - cleanup_dist_env_and_memory() - - -@pytest.fixture(scope="module") -def omni_runner(request, model_prefix): - with _omni_server_lock: - model, stage_config_path = request.param - model = model_prefix + model - with OmniRunner(model, seed=42, stage_configs_path=stage_config_path) as runner: - print("OmniRunner started successfully") - yield runner - print("OmniRunner stopping...") - - print("OmniRunner stopped") - - -class OmniRunnerHandler: - def __init__(self, omni_runner): - self.runner = omni_runner - - def _process_output(self, outputs: list[Any]) -> OmniResponse: - result = OmniResponse() - try: - text_content = None - audio_content = None - for stage_output in outputs: - if getattr(stage_output, "final_output_type", None) == "text": - text_content = stage_output.request_output.outputs[0].text - if getattr(stage_output, "final_output_type", None) == "audio": - audio_content = stage_output.request_output.outputs[0].multimodal_output["audio"] - - result.audio_content = audio_content - result.text_content = text_content - result.success = True - - except Exception as e: - result.error_message = f"Output processing error: {str(e)}" - result.success = False - print(f"Error: {result.error_message}") - - return result - - def send_request(self, request_config: dict[str, Any] | None = None) -> OmniResponse: - if request_config is None: - request_config = {} - prompts = request_config.get("prompts") - videos = request_config.get("videos") - images = request_config.get("images") - audios = request_config.get("audios") - modalities = request_config.get("modalities", ["text", "audio"]) - outputs = self.runner.generate_multimodal( - prompts=prompts, videos=videos, images=images, audios=audios, modalities=modalities - ) - response = self._process_output(outputs) - assert_omni_response(response, request_config, run_level="core_model") - return response - - def send_audio_speech_request( - self, - request_config: dict[str, Any], - ) -> OmniResponse: - """ - Offline TTS: text -> audio via generate_multimodal, then validate with assert_audio_speech_response. - - request_config must contain: - - 'input' or 'prompts': text to synthesize. - Optional keys: - - 'voice' -> speaker (CustomVoice) - - 'task_type' -> task_type in additional_information (default: "CustomVoice") - - 'language' -> language in additional_information (default: "Auto") - - 'max_new_tokens' -> max_new_tokens in additional_information (default: 2048) - - 'response_format' -> desired audio format (used only for assertion) - """ - input_text = request_config.get("input") or request_config.get("prompts") - if input_text is None: - raise ValueError("request_config must contain 'input' or 'prompts' for TTS") - if isinstance(input_text, list): - input_text = input_text[0] if input_text else "" - - # Build TTS-specific kwargs passed through to get_omni_inputs for Qwen3-TTS, - # matching examples/offline_inference/qwen3_tts/end2end.py. - mm_processor_kwargs: dict[str, Any] = {} - if "voice" in request_config: - mm_processor_kwargs["speaker"] = request_config["voice"] - if "task_type" in request_config: - mm_processor_kwargs["task_type"] = request_config["task_type"] - if "ref_audio" in request_config: - mm_processor_kwargs["ref_audio"] = request_config["ref_audio"] - if "ref_text" in request_config: - mm_processor_kwargs["ref_text"] = request_config["ref_text"] - if "language" in request_config: - mm_processor_kwargs["language"] = request_config["language"] - if "max_new_tokens" in request_config: - mm_processor_kwargs["max_new_tokens"] = request_config["max_new_tokens"] - - outputs = self.runner.generate_multimodal( - prompts=input_text, - modalities=["audio"], - mm_processor_kwargs=mm_processor_kwargs or None, - ) - mm_out: dict[str, Any] | None = None - for stage_out in outputs: - if getattr(stage_out, "final_output_type", None) == "audio": - mm_out = stage_out.request_output.outputs[0].multimodal_output - break - if mm_out is None: - result = OmniResponse(success=False, error_message="No audio output from pipeline") - assert result.success, result.error_message - return result - - audio_data = mm_out.get("audio") - if audio_data is None: - result = OmniResponse(success=False, error_message="No audio tensor in multimodal output") - assert result.success, result.error_message - return result - - sr_raw = mm_out.get("sr") - sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw - sr = int(sr_val.item() if hasattr(sr_val, "item") else sr_val) - wav_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data - wav_buf = io.BytesIO() - sf.write( - wav_buf, - wav_tensor.float().cpu().numpy().reshape(-1), - samplerate=sr, - format="WAV", - subtype="PCM_16", - ) - result = OmniResponse(success=True, audio_bytes=wav_buf.getvalue(), audio_format="audio/wav") - assert_audio_speech_response(result, request_config, run_level="core_model") - return result - - def start_profile( - self, - profile_prefix: str | None = None, - stages: list[int] | None = None, - ) -> list[Any]: - """Start profiling specified stages.""" - return self.runner.start_profile(profile_prefix=profile_prefix, stages=stages) - - def stop_profile(self, stages: list[int] | None = None) -> list[Any]: - """Stop profiling specified stages.""" - return self.runner.stop_profile(stages=stages) - - -@pytest.fixture -def omni_runner_handler(omni_runner): - return OmniRunnerHandler(omni_runner) diff --git a/tests/core/sched/test_omni_scheduler_mixin.py b/tests/core/sched/test_omni_scheduler_mixin.py new file mode 100644 index 000000000..e04a9c39f --- /dev/null +++ b/tests/core/sched/test_omni_scheduler_mixin.py @@ -0,0 +1,129 @@ +"""Unit tests for OmniSchedulerMixin streaming session replacement. + +These tests pin the behavior of `_replace_session_with_streaming_update` against +current vLLM `Request` / `StreamingUpdate` (and Omni patches). When upgrading +vLLM, failures here should highlight incompatible changes to request state or +update payloads early. +""" + +from __future__ import annotations + +from dataclasses import replace + +import pytest + +# Imports must run in this order: vllm_omni applies patches to vllm.v1.request before +# Request / StreamingUpdate are bound in this module. Ruff isort would reorder them. +# isort: off +import vllm_omni # noqa: F401 - import for side effects (patch vLLM) +from vllm.sampling_params import SamplingParams +from vllm.v1.engine import EngineCoreEventType +from vllm.v1.request import Request, RequestStatus, StreamingUpdate +from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin + +# isort: on + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _SchedulerStub(OmniSchedulerMixin): + """Minimal scheduler surface required by OmniSchedulerMixin.""" + + def __init__(self, *, log_stats: bool = False) -> None: + self.num_waiting_for_streaming_input = 0 + self.log_stats = log_stats + + +def _make_request(**kwargs) -> Request: + sp = SamplingParams(max_tokens=8) + defaults = dict( + request_id="req-mixin-test", + prompt_token_ids=[1, 2, 3], + sampling_params=sp, + pooling_params=None, + arrival_time=100.0, + block_hasher=None, + ) + defaults.update(kwargs) + return Request(**defaults) + + +def _make_update(**kwargs) -> StreamingUpdate: + sp_new = SamplingParams(max_tokens=16) + defaults = dict( + mm_features=None, + prompt_token_ids=[10, 20], + max_tokens=32, + arrival_time=200.0, + sampling_params=sp_new, + ) + defaults.update(kwargs) + return StreamingUpdate(**defaults) + + +class TestReplaceSessionWithStreamingUpdate: + def test_resets_tokens_and_prompt_from_update(self) -> None: + sched = _SchedulerStub() + session = _make_request() + session.append_output_token_ids([7, 8]) + session.num_computed_tokens = 99 + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + + update = _make_update(prompt_token_ids=[40, 41, 42]) + sched.num_waiting_for_streaming_input = 3 + sched._replace_session_with_streaming_update(session, update) + + assert session._output_token_ids == [] + assert list(session._all_token_ids) == [40, 41, 42] + assert session.prompt_token_ids == [40, 41, 42] + assert session.num_computed_tokens == 0 + assert session.num_prompt_tokens == 3 + assert session.arrival_time == 200.0 + assert session.sampling_params is update.sampling_params + assert session.status == RequestStatus.WAITING + assert sched.num_waiting_for_streaming_input == 2 + + def test_none_prompt_token_ids_becomes_empty(self) -> None: + sched = _SchedulerStub() + session = _make_request() + session.status = RequestStatus.RUNNING + update = _make_update(prompt_token_ids=None) + sched._replace_session_with_streaming_update(session, update) + + assert session.prompt_token_ids == () + assert list(session._all_token_ids) == [] + assert session.num_prompt_tokens == 0 + assert sched.num_waiting_for_streaming_input == 0 + + def test_additional_information_cleared_when_update_omits_it(self) -> None: + sched = _SchedulerStub() + session = _make_request() + if not hasattr(session, "additional_information"): + pytest.skip("Request has no additional_information (Omni patch inactive?)") + session.additional_information = {"keep": True} + session.status = RequestStatus.RUNNING + + base = _make_update() + if not hasattr(base, "additional_information"): + pytest.skip("StreamingUpdate has no additional_information (Omni patch inactive?)") + update = replace(base, additional_information=None) + + sched._replace_session_with_streaming_update(session, update) + assert session.additional_information is None + + def test_does_not_decrement_waiting_when_not_streaming_status(self) -> None: + sched = _SchedulerStub() + session = _make_request() + session.status = RequestStatus.RUNNING + sched.num_waiting_for_streaming_input = 5 + sched._replace_session_with_streaming_update(session, _make_update()) + assert sched.num_waiting_for_streaming_input == 5 + + def test_records_queued_event_when_log_stats_enabled(self) -> None: + sched = _SchedulerStub(log_stats=True) + session = _make_request() + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + sched._replace_session_with_streaming_update(session, _make_update()) + + assert session.events + assert session.events[-1].type == EngineCoreEventType.QUEUED diff --git a/tests/core/test_prefix_cache.py b/tests/core/test_prefix_cache.py new file mode 100644 index 000000000..b5d0e96d3 --- /dev/null +++ b/tests/core/test_prefix_cache.py @@ -0,0 +1,349 @@ +import pytest +import torch + +from vllm_omni.core.prefix_cache import OmniTensorPrefixCache + +DEFAULT_SEQ_LEN = 15 +NUM_BLOCKS = 10 +BLOCK_SIZE = 4 +HIDDEN_SIZE = 2 +DTYPE = torch.float32 +OTHER_DTYPE = torch.float16 +DEFAULT_SHAPE = torch.Size([NUM_BLOCKS, BLOCK_SIZE, HIDDEN_SIZE]) + + +class MockInputBatch: + def __init__(self, num_computed_tokens_cpu): + self.req_ids = ["req1", "req2"] + self.req_id_to_index = {req_id: i for i, req_id in enumerate(self.req_ids)} + self.num_computed_tokens_cpu = num_computed_tokens_cpu + + # Block table is only mocked for validation of length; + # we don't actually need to add valid values here since + # we patch the table when testing. + class _DummyBlockTable: + pass + + self.block_table = _DummyBlockTable() + self.block_table.block_tables = [None] + + +def get_omni_pcache_with_mm_tensors(feat_dims, seq_len) -> OmniTensorPrefixCache: + """Build an OmniTensorPrefixCache and init mm tensors.""" + cache = get_omni_pcache() + mm_outputs = get_multimodal_outputs(feat_dims, seq_len) + cache.maybe_init_missing_mm_cache_keys(mm_outputs, seq_len) + return cache + + +def get_omni_pcache() -> OmniTensorPrefixCache: + """Build an OmniTensorPrefixCache, but don't init mm tensors.""" + cache = OmniTensorPrefixCache( + num_blocks=NUM_BLOCKS, + block_size=BLOCK_SIZE, + hidden_size=HIDDEN_SIZE, + hs_dtype=DTYPE, + ) + return cache + + +def get_multimodal_outputs(feat_dims: dict[str, int], seq_len: int) -> dict[str, torch.Tensor]: + fake_mm_inputs = {} + for mm_key, feat_dim in feat_dims.items(): + fake_mm_inputs[mm_key] = torch.rand((seq_len, feat_dim), dtype=DTYPE) + return fake_mm_inputs + + +### Tests for initialization +def test_initialization_simple(): + """Check default initialization only creates the hidden states.""" + cache = get_omni_pcache() + assert isinstance(cache.hidden_states_cache, torch.Tensor) + assert cache.hidden_states_cache.shape == DEFAULT_SHAPE + assert len(cache.mm_outputs_cache) == 0 + assert len(cache.mm_cache_keys) == 0 + + +def test_initialization_with_multimodal(): + """Check initialization + registration of multimodal outputs.""" + cache = get_omni_pcache() + feat_dims = {"foo": 100, "bar": 50, "baz": 10} + mm_outputs = get_multimodal_outputs( + feat_dims, + seq_len=DEFAULT_SEQ_LEN, + ) + # Cast one of the keys to a different dtype; the dtype of the tensor + # that is used to initialize the cache dictates the cache dtype. + mm_outputs["foo"] = mm_outputs["foo"].to(OTHER_DTYPE) + + cache.maybe_init_missing_mm_cache_keys(mm_outputs, DEFAULT_SEQ_LEN) + assert len(cache.mm_cache_keys) == 3 + assert set(cache.mm_cache_keys) == set(feat_dims.keys()) + for mm_key in cache.mm_cache_keys: + cache_tensor = cache.mm_outputs_cache[mm_key] + assert isinstance(cache_tensor, torch.Tensor) + assert cache_tensor.shape[-1] == feat_dims[mm_key] + assert mm_outputs[mm_key].dtype == cache_tensor.dtype + + +def test_init_missing_mm_cache_keys_is_idempotent(): + """Ensure that the cache doesn't reinitialize old keys.""" + cache = get_omni_pcache() + mm_key = "foo" + feat_dims = {mm_key: 100} + mm_outputs = get_multimodal_outputs( + feat_dims, + seq_len=DEFAULT_SEQ_LEN, + ) + cache.maybe_init_missing_mm_cache_keys(mm_outputs, DEFAULT_SEQ_LEN) + assert len(cache.mm_cache_keys) == 1 + assert mm_key in cache.mm_cache_keys + + # Cache is initialized to 0 - fill it with 1s + cache.mm_outputs_cache[mm_key].fill_(1) + + # Ensure that running another initialization + # doesn't zero out our cache values + cache.maybe_init_missing_mm_cache_keys(mm_outputs, DEFAULT_SEQ_LEN) + assert len(cache.mm_cache_keys) == 1 + assert mm_key in cache.mm_cache_keys + assert torch.all(cache.mm_outputs_cache[mm_key] == 1) + + +### Tests for Update +def test_update_no_multimodal(): + """Test that slot mappings act as row indices hidden states.""" + cache = get_omni_pcache() + + num_tokens_unpadded = 8 + slot_offset = 8 + slot_mapping = torch.arange(slot_offset, slot_offset + num_tokens_unpadded) + new_hidden_states = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE) + + cache.update_omni_tensor_prefix_cache( + hidden_states=new_hidden_states, + multimodal_outputs=None, + num_tokens_unpadded=num_tokens_unpadded, + slot_mapping=slot_mapping, + ) + + # Ensure that if we reshape our 3D cache back to 2D, we can use the + # indices in our slot mappings to access the hidden states as expected + hs_rows = cache.hidden_states_cache.view(NUM_BLOCKS * BLOCK_SIZE, HIDDEN_SIZE) + for slot_idx, new_states in zip(slot_mapping, new_hidden_states): + slot_states = hs_rows[slot_idx] + assert torch.all(slot_states == new_states) + + +@pytest.mark.parametrize( + "feat_dims", + [ + {"foo": 100, "bar": 100}, + {"foo": 100, "bar": 50, "baz": 10}, + ], +) +def test_update_with_multimodal_outputs(feat_dims): + """Test that slot mappings are correct for multimodal tensors.""" + cache = get_omni_pcache_with_mm_tensors(feat_dims, seq_len=DEFAULT_SEQ_LEN) + + num_tokens_unpadded = 8 + slot_offset = 8 + slot_mapping = torch.arange(slot_offset, slot_offset + num_tokens_unpadded) + feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()} + mm_outputs = {key: torch.rand((num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys} + cache.update_omni_tensor_prefix_cache( + hidden_states=None, + multimodal_outputs=mm_outputs, + num_tokens_unpadded=num_tokens_unpadded, + slot_mapping=slot_mapping, + ) + + for mm_key in feat_dims.keys(): + assert mm_key in cache.mm_outputs_cache + key_feat_dim = feature_dims[mm_key] + mm_state_rows = cache.mm_outputs_cache[mm_key].view(NUM_BLOCKS * BLOCK_SIZE, key_feat_dim) + + # Similar to hidden states, but for each key in the dict; + # Different tensors may have different feature dims + new_mm_outputs = mm_outputs[mm_key] + for slot_idx, new_output in zip(slot_mapping, new_mm_outputs): + slot_states = mm_state_rows[slot_idx] + assert torch.all(slot_states == new_output) + + +### Tests for Merging +def fake_get_cached_block_ids(self, req_idx, *args, **kwargs): + """Fake block table lookup. + + Assumption: + req_idx 0 is a cache hit with slots 8, 9, ..., 15 + req_idx 1 is a cache miss + """ + assert req_idx < 2 + if req_idx == 0: + # With the slot offset we provided (8), the corresponding + # blocks IDs are 2 & 3 because the block size is 4. + return torch.tensor([2, 3], dtype=torch.long) + return torch.tensor([], dtype=torch.long) + + +@pytest.mark.parametrize("num_tokens_padded", [None, 16]) +def test_get_merged_hidden_states(num_tokens_padded, mocker): + """Ensure that hidden states are merged correctly.""" + cache = get_omni_pcache() + + orig_num_tokens_unpadded = 8 + slot_offset = 8 # We'll put our states in slots 8, 9, 10, ..., 15 + orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded) + orig_hidden_states = torch.rand((orig_num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE) + + cache.update_omni_tensor_prefix_cache( + hidden_states=orig_hidden_states, + multimodal_outputs=None, + num_tokens_unpadded=orig_num_tokens_unpadded, + slot_mapping=orig_slot_mapping, + num_tokens_padded=num_tokens_padded, + ) + + # Say that we have two requests, but only one of them is a cache hit + num_new_toks_req1 = 3 + num_new_toks_req2 = 2 + cache.add_prefix_cached_new_req_id("req1") + + num_scheduled_tokens = { + "req1": num_new_toks_req1, + "req2": num_new_toks_req2, + } + new_hidden_states = torch.rand( + (num_new_toks_req1 + num_new_toks_req2, HIDDEN_SIZE), + dtype=DTYPE, + ) + req1_new_states = new_hidden_states[:num_new_toks_req1] + req2_new_states = new_hidden_states[-num_new_toks_req2:] + + input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0])) + + mocker.patch( + "vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids", + new=fake_get_cached_block_ids, + ) + merged_states = cache.get_merged_hidden_states( + query_start_loc=[0, num_new_toks_req1], + input_batch=input_batch, + hidden_states=new_hidden_states, + num_scheduled_tokens=num_scheduled_tokens, + ) + + assert "req1" in merged_states and "req2" in merged_states + req1_merged_states = merged_states["req1"] + req2_merged_states = merged_states["req2"] + + # First, check the cache hit case + assert req1_merged_states.shape == torch.Size([orig_num_tokens_unpadded + num_new_toks_req1, HIDDEN_SIZE]) + # Ensure that the req1 merged states are the cached states + the new req1 states + assert torch.all(req1_merged_states[:orig_num_tokens_unpadded] == orig_hidden_states) + assert torch.all(req1_merged_states[-num_new_toks_req1:] == req1_new_states) + + # Next, ensure that the cache miss case only has the new states + assert req2_merged_states.shape == torch.Size([num_new_toks_req2, HIDDEN_SIZE]) + assert torch.all(req2_merged_states == req2_new_states) + + +@pytest.mark.parametrize("num_tokens_padded", [None, 16]) +@pytest.mark.parametrize( + "feat_dims", + [ + {"foo": 100, "bar": 100}, + {"foo": 100, "bar": 50, "baz": 10}, + ], +) +def test_get_merged_multimodal_outputs(feat_dims, num_tokens_padded, mocker): + cache = get_omni_pcache_with_mm_tensors(feat_dims, seq_len=DEFAULT_SEQ_LEN) + + orig_num_tokens_unpadded = 8 + slot_offset = 8 # We'll put our states in slots 8, 9, 10, ..., 15 + orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded) + feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()} + orig_mm_outputs = { + key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys + } + + cache.update_omni_tensor_prefix_cache( + hidden_states=None, + multimodal_outputs=orig_mm_outputs, + num_tokens_unpadded=orig_num_tokens_unpadded, + slot_mapping=orig_slot_mapping, + num_tokens_padded=num_tokens_padded, + ) + + # Similar to hs test- say that we have two requests, but only one of them is a cache hit + num_new_toks_req1 = 3 + num_new_toks_req2 = 2 + cache.add_prefix_cached_new_req_id("req1") + + num_scheduled_tokens = { + "req1": num_new_toks_req1, + "req2": num_new_toks_req2, + } + + new_mm_outputs = {} + for mm_key in cache.mm_cache_keys: + new_mm_outputs[mm_key] = torch.rand( + (num_new_toks_req1 + num_new_toks_req2, feature_dims[mm_key]), + dtype=DTYPE, + ) + # We also want to make sure passthrough data (outside of our keys) isn't dropped + new_mm_outputs["passthrough_data"] = "Something else" + # Lists are a special case because we can't split them yet if we want to match + # the nonprefix cache behavior, because this runs before post process. + new_mm_outputs["passthrough_list"] = ["should", "not", "split"] + + input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0])) + + mocker.patch( + "vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids", + new=fake_get_cached_block_ids, + ) + merged_mm_outputs = cache.get_merged_multimodal_states( + query_start_loc=[0, num_new_toks_req1], + input_batch=input_batch, + multimodal_outputs=new_mm_outputs, + num_scheduled_tokens=num_scheduled_tokens, + ) + + # Ensure the passthrough data wasn't dropped + assert "passthrough_data" in merged_mm_outputs + assert "passthrough_list" in merged_mm_outputs + + for mm_key, mm_output in merged_mm_outputs.items(): + # Ensure passthrough data is just forwarded normally and not duplicated + assert isinstance(mm_output, dict) + assert "req1" in mm_output and "req2" in mm_output + if mm_key == "passthrough_data": + assert mm_key not in cache.mm_cache_keys + assert new_mm_outputs[mm_key] == mm_output["req1"] + assert new_mm_outputs[mm_key] == mm_output["req2"] + elif mm_key == "passthrough_list": + assert mm_key not in cache.mm_cache_keys + assert new_mm_outputs[mm_key] == mm_output["req1"] + assert new_mm_outputs[mm_key] == mm_output["req2"] + else: + assert mm_key in cache.mm_cache_keys + curr_feat_dim = feature_dims[mm_key] + # Ensure that req1 (cache hit) merged the mm data + req1_merged_mm_outputs = mm_output["req1"] + req1_new_mm_outputs = new_mm_outputs[mm_key][:num_new_toks_req1] + + assert req1_merged_mm_outputs.shape == torch.Size( + [orig_num_tokens_unpadded + num_new_toks_req1, curr_feat_dim] + ) + # Ensure that the req1 merged mm data are the cached data + the new data + assert torch.all(req1_merged_mm_outputs[:orig_num_tokens_unpadded] == orig_mm_outputs[mm_key]) + assert torch.all(req1_merged_mm_outputs[-num_new_toks_req1:] == req1_new_mm_outputs) + + # Ensure that req2 (cache miss) only has the new mm data + req2_merged_mm_outputs = mm_output["req2"] + req2_new_mm_outputs = new_mm_outputs[mm_key][-num_new_toks_req2:] + + assert req2_merged_mm_outputs.shape == torch.Size([num_new_toks_req2, curr_feat_dim]) + assert torch.all(req2_merged_mm_outputs == req2_new_mm_outputs) diff --git a/tests/dfx/conftest.py b/tests/dfx/conftest.py index e54141b34..c3f6d0a15 100644 --- a/tests/dfx/conftest.py +++ b/tests/dfx/conftest.py @@ -2,7 +2,9 @@ import json from pathlib import Path from typing import Any -from tests.conftest import modify_stage_config +import pytest + +from tests.helpers.stage_config import modify_stage_config def load_configs(config_path: str) -> list[dict[str, Any]]: @@ -38,22 +40,32 @@ def modify_stage(default_path, updates, deletes): def create_unique_server_params( configs: list[dict[str, Any]], stage_configs_dir: Path, -) -> list[tuple[str, str, str]]: +) -> list[tuple[str, str, str | None, str | None, tuple[str, ...]]]: unique_params = [] seen = set() for config in configs: test_name = config["test_name"] - model = config["server_params"]["model"] - stage_config_name = config["server_params"].get("stage_config_name") + server_params = config["server_params"] + model = server_params["model"] + stage_config_name = server_params.get("stage_config_name") if stage_config_name: stage_config_path = str(stage_configs_dir / stage_config_name) - delete = config["server_params"].get("delete", None) - update = config["server_params"].get("update", None) + delete = server_params.get("delete", None) + update = server_params.get("update", None) stage_config_path = modify_stage(stage_config_path, update, delete) else: stage_config_path = None - server_param = (test_name, model, stage_config_path) + stage_overrides = server_params.get("stage_overrides") + stage_overrides_json = json.dumps(stage_overrides) if stage_overrides else None + + # ``extra_cli_args`` passes raw CLI flags straight through to + # ``vllm_omni.entrypoints.cli.main serve`` — used for flags that + # don't map to stage-level overrides, e.g. ``--async-chunk`` / + # ``--no-async-chunk`` toggling the deploy-level async_chunk bool. + extra_cli_args = tuple(server_params.get("extra_cli_args") or ()) + + server_param = (test_name, model, stage_config_path, stage_overrides_json, extra_cli_args) if server_param not in seen: seen.add(server_param) unique_params.append(server_param) @@ -95,3 +107,13 @@ def create_benchmark_indices( indices.append((test_name, idx)) return indices + + +def pytest_addoption(parser: pytest.Parser) -> None: + """Register shared CLI options for DFX benchmark suites.""" + parser.addoption( + "--test-config-file", + action="store", + default=None, + help=("Path to benchmark config JSON. Example: --test-config-file tests/dfx/perf/tests/test_tts.json"), + ) diff --git a/tests/dfx/perf/scripts/diffusion_result_template.json b/tests/dfx/perf/scripts/diffusion_result_template.json new file mode 100644 index 000000000..86bdf1bc7 --- /dev/null +++ b/tests/dfx/perf/scripts/diffusion_result_template.json @@ -0,0 +1,86 @@ +[ + { + "test_name": null, + "backend": null, + "timestamp": null, + "server_params": { + "model": null, + "serve_args": { + "enable-diffusion-pipeline-profiler": false + } + }, + "benchmark_params": { + "name": null, + "dataset": null, + "task": null, + "width": 0, + "height": 0, + "num-inference-steps": 0, + "num-prompts": 0, + "max-concurrency": 0, + "num-input-images": 0, + "enable-negative-prompt": false, + "baseline": { + "throughput_qps": 0, + "latency_mean": 0, + "peak_memory_mb_max": 0, + "peak_memory_mb_mean": 0 + } + }, + "result": { + "duration": 0, + "completed_requests": 0, + "failed_requests": 0, + "throughput_qps": 0, + "latency_mean": 0, + "latency_median": 0, + "latency_p99": 0, + "latency_p95": 0, + "latency_p50": 0, + "peak_memory_mb_max": 0, + "peak_memory_mb_mean": 0, + "peak_memory_mb_median": 0, + "stage_durations_mean": {}, + "stage_durations_p50": {}, + "stage_durations_p99": {}, + "backend": null, + "model": null, + "dataset": null, + "task": null + }, + "log_file": null, + "Model": null, + "Framework": null, + "Hardware": null, + "Deployment": null, + "Task": null, + "Dataset": null, + "resolution": null, + "Parallelism": null, + "max_concurrency": 0, + "Cache": null, + "Quantization": null, + "offload": null, + "compile": null, + "Attn_backend": null, + "num_inference_steps": 0, + "completed": 0, + "failed": 0, + "throughput_qps": 0, + "latency_mean": 0, + "latency_median": 0, + "latency_p99": 0, + "latency_p95": 0, + "latency_p50": 0, + "peak_memory_mb_max": 0, + "peak_memory_mb_mean": 0, + "peak_memory_mb_median": 0, + "stage_durations_mean": {}, + "stage_durations_p50": {}, + "stage_durations_p99": {}, + "commit_sha": null, + "build_id": null, + "build_url": null, + "source_file": null + } +] diff --git a/tests/dfx/perf/scripts/result_omni_template.json b/tests/dfx/perf/scripts/result_omni_template.json new file mode 100644 index 000000000..1d6132140 --- /dev/null +++ b/tests/dfx/perf/scripts/result_omni_template.json @@ -0,0 +1,55 @@ +{ + "date": null, + "endpoint_type": null, + "backend": null, + "label": null, + "model_id": null, + "tokenizer_id": null, + "num_prompts": 0, + "request_rate": null, + "burstiness": 0, + "max_concurrency": 0, + "duration": 0, + "completed": 0, + "failed": 0, + "total_input_tokens": 0, + "total_output_tokens": 0, + "request_throughput": 0, + "request_goodput": null, + "output_throughput": 0, + "total_token_throughput": 0, + "total_audio_duration_s": 0, + "total_audio_frames": 0, + "audio_throughput": 0, + "max_output_tokens_per_s": 0, + "max_concurrent_requests": 0, + "rtfx": 0, + "mean_ttft_ms": 0, + "median_ttft_ms": 0, + "p99_ttft_ms": 0, + "mean_tpot_ms": 0, + "median_tpot_ms": 0, + "p99_tpot_ms": 0, + "mean_itl_ms": 0, + "median_itl_ms": 0, + "p99_itl_ms": 0, + "mean_e2el_ms": 0, + "median_e2el_ms": 0, + "p99_e2el_ms": 0, + "mean_audio_rtf": 0, + "median_audio_rtf": 0, + "p99_audio_rtf": 0, + "mean_audio_ttfp_ms": 0, + "median_audio_ttfp_ms": 0, + "p99_audio_ttfp_ms": 0, + "mean_audio_duration_s": 0, + "median_audio_duration_s": 0, + "p99_audio_duration_s": 0, + "baseline": { + "mean_ttft_ms": 0, + "mean_audio_ttfp_ms": 0, + "mean_audio_rtf": 0 + }, + "random_input_len": 0, + "random_output_len": 0 +} diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index b64cc0d95..13011b4bd 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/tests/dfx/perf/scripts/run_benchmark.py @@ -8,7 +8,6 @@ from typing import Any import pytest -from tests.conftest import OmniServer from tests.dfx.conftest import ( create_benchmark_indices, create_test_parameter_mapping, @@ -16,18 +15,41 @@ from tests.dfx.conftest import ( get_benchmark_params_for_server, load_configs, ) +from tests.helpers.runtime import OmniServer os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" -CONFIG_FILE_PATH = str(Path(__file__).parent.parent / "tests" / "test.json") +def _get_config_file_from_argv() -> str | None: + """Read ``--test-config-file`` from ``sys.argv`` at import time so parametrization can use it.""" + import sys + + for i, arg in enumerate(sys.argv): + if arg == "--test-config-file" and i + 1 < len(sys.argv): + return sys.argv[i + 1] + if arg.startswith("--test-config-file="): + return arg.split("=", 1)[1] + return None + + +_PERF_TESTS_DIR = Path(__file__).resolve().parent.parent / "tests" +_DEFAULT_CONFIG_FILE = str(_PERF_TESTS_DIR / "test_qwen_omni.json") + +CONFIG_FILE_PATH = _get_config_file_from_argv() +if CONFIG_FILE_PATH is None: + print( + "No --test-config-file in argv, using default: tests/dfx/perf/tests/test_qwen_omni.json " + "(override with e.g. --test-config-file tests/dfx/perf/tests/test_tts.json)" + ) + CONFIG_FILE_PATH = _DEFAULT_CONFIG_FILE + BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) -STAGE_INIT_TIMEOUT = 600 +OMNI_RESULT_TEMPLATE_PATH = Path(__file__).parent / "result_omni_template.json" -STAGE_CONFIGS_DIR = Path(__file__).parent.parent / "stage_configs" -test_params = create_unique_server_params(BENCHMARK_CONFIGS, STAGE_CONFIGS_DIR) +DEPLOY_CONFIGS_DIR = Path(__file__).parent.parent / "deploy" +test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) _omni_server_lock = threading.Lock() @@ -40,13 +62,19 @@ def omni_server(request): Multi-stage initialization can take 10-20+ minutes. """ with _omni_server_lock: - test_name, model, stage_config_path = request.param + test_name, model, stage_config_path, stage_overrides, extra_cli_args = request.param print(f"Starting OmniServer with test: {test_name}, model: {model}") - server_args = ["--stage-init-timeout", str(STAGE_INIT_TIMEOUT), "--init-timeout", "900"] + server_args = ["--stage-init-timeout", "600", "--init-timeout", "900"] + # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py): + # deploy-config sets the base; stage-overrides are applied on top. Both can be set. if stage_config_path: - server_args = ["--stage-configs-path", stage_config_path] + server_args + server_args = ["--deploy-config", stage_config_path] + server_args + if stage_overrides: + server_args = ["--stage-overrides", stage_overrides] + server_args + if extra_cli_args: + server_args = list(extra_cli_args) + server_args with OmniServer(model, server_args) as server: server.test_name = test_name print("OmniServer started successfully") @@ -56,24 +84,47 @@ def omni_server(request): print("OmniServer stopped") +def _safe_filename_token(value: Any | None, *, default: str = "na") -> str: + """Make a single path segment safe for result filenames on common filesystems.""" + if value is None: + return default + s = str(value).strip() + for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"): + s = s.replace(bad, "_") + return s if s else default + + def run_benchmark( args: list, test_name: str, flow, dataset_name: str, num_prompt, + *, + baseline_config: dict[str, Any] | None = None, + sweep_index: int | None = None, + request_rate: Any | None = None, + max_concurrency: Any | None = None, + random_input_len: Any | None = None, + random_output_len: Any | None = None, ) -> Any: - """Run a single benchmark iteration and return the parsed result JSON.""" + """Run a single benchmark iteration and return the parsed result JSON. + + After ``vllm bench`` writes the JSON, ``result["baseline"]`` holds the same + per-metric resolved thresholds as ``assert_result`` (via ``_baseline_thresholds_for_step``). + When ``random_input_len`` / ``random_output_len`` are set, they are also written into the result JSON; + omitted keys when not configured. + """ current_dt = datetime.now().strftime("%Y%m%d-%H%M%S") - result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_{current_dt}.json" + ri = _safe_filename_token(random_input_len) + ro = _safe_filename_token(random_output_len) + result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json" if "--result-filename" in args: print(f"The result file will be overwritten by {result_filename}") command = ( ["vllm", "bench", "serve", "--omni"] + args + [ - "--num-warmups", - "2", "--save-result", "--result-dir", os.environ.get("BENCHMARK_DIR", "tests"), @@ -97,8 +148,34 @@ def run_benchmark( else: result_dir = "./" - with open(os.path.join(result_dir, result_filename), encoding="utf-8") as f: - result = json.load(f) + result_path = os.path.join(result_dir, result_filename) + if not os.path.exists(result_path): + with open(OMNI_RESULT_TEMPLATE_PATH, encoding="utf-8") as f: + template_result: dict[str, Any] = json.load(f) + Path(result_path).parent.mkdir(parents=True, exist_ok=True) + with open(result_path, "w", encoding="utf-8") as f: + json.dump(template_result, f, ensure_ascii=False, indent=2) + print(f"Benchmark result file not generated, fallback to template: {result_path}") + result = template_result + else: + with open(result_path, encoding="utf-8") as f: + result = json.load(f) + + if baseline_config: + result["baseline"] = _baseline_thresholds_for_step( + baseline_config, + sweep_index=sweep_index, + request_rate=request_rate, + max_concurrency=max_concurrency, + ) + else: + result["baseline"] = {} + if random_input_len is not None: + result["random_input_len"] = random_input_len + if random_output_len is not None: + result["random_output_len"] = random_output_len + with open(result_path, "w", encoding="utf-8") as f: + json.dump(result, f, ensure_ascii=False, indent=2) return result @@ -168,6 +245,25 @@ def _resolve_baseline_value( return baseline_raw +def _baseline_thresholds_for_step( + baseline_data: dict[str, Any], + *, + sweep_index: int | None = None, + max_concurrency: Any = None, + request_rate: Any = None, +) -> dict[str, Any]: + """Resolve ``test.json`` ``baseline`` block to one threshold per metric (same as ``assert_result``).""" + return { + metric_name: _resolve_baseline_value( + baseline_raw, + sweep_index=sweep_index, + max_concurrency=max_concurrency, + request_rate=request_rate, + ) + for metric_name, baseline_raw in baseline_data.items() + } + + def assert_result( result, params, @@ -258,6 +354,12 @@ def test_performance_benchmark(omni_server, benchmark_params): flow=qps, dataset_name=dataset_name, num_prompt=num_prompt, + baseline_config=params.get("baseline"), + sweep_index=i, + request_rate=qps, + max_concurrency=None, + random_input_len=params.get("random_input_len"), + random_output_len=params.get("random_output_len"), ) assert_result( result, @@ -276,6 +378,12 @@ def test_performance_benchmark(omni_server, benchmark_params): flow=concurrency, dataset_name=dataset_name, num_prompt=num_prompt, + baseline_config=params.get("baseline"), + sweep_index=i, + request_rate=None, + max_concurrency=concurrency, + random_input_len=params.get("random_input_len"), + random_output_len=params.get("random_output_len"), ) assert_result( result, diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py index 123f21405..23efa8bb0 100644 --- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py +++ b/tests/dfx/perf/scripts/run_diffusion_benchmark.py @@ -5,8 +5,8 @@ Supports vLLM-Omni server backend: - vllm-omni (default): starts DiffusionServer via vllm_omni.entrypoints.cli.main, benchmarks with diffusion_benchmark_serving.py --backend vllm-omni -A config JSON file is REQUIRED via --config-file: - pytest run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json +A config JSON file is REQUIRED via --test-config-file: + pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json JSON config entries use a "server_type" field, and this runner executes the vllm-omni path. @@ -52,19 +52,20 @@ BENCHMARK_SCRIPT = str( _SESSION_TIMESTAMP = datetime.now().strftime("%Y%m%d-%H%M%S") _RESULT_LOCK = threading.Lock() _BRANCHPOINT_COMMIT_SHA: str | None = None +DIFFUSION_RESULT_TEMPLATE_PATH = Path(__file__).parent / "diffusion_result_template.json" def _get_config_file_from_argv() -> str | None: - """Read --config-file from sys.argv at import time so pytest parametrize can use it. + """Read --test-config-file from sys.argv at import time so pytest parametrize can use it. pytest_addoption (below) registers the same flag so pytest does not reject it. - Supports both ``--config-file path`` and ``--config-file=path`` forms. + Supports both ``--test-config-file path`` and ``--test-config-file=path`` forms. Returns None if the flag is not present; callers must handle the missing case. """ for i, arg in enumerate(sys.argv): - if arg == "--config-file" and i + 1 < len(sys.argv): + if arg == "--test-config-file" and i + 1 < len(sys.argv): return sys.argv[i + 1] - if arg.startswith("--config-file="): + if arg.startswith("--test-config-file="): return arg.split("=", 1)[1] return None @@ -133,19 +134,6 @@ def _append_to_aggregated_file(record: dict[str, Any]) -> None: json.dump(records, f, indent=2, ensure_ascii=False) -# Register --config-file with pytest so it does not reject the argument. -def pytest_addoption(parser: pytest.Parser) -> None: - parser.addoption( - "--config-file", - action="store", - default=None, - help=( - "Path to the benchmark config JSON file (required). " - "Example: --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json" - ), - ) - - _server_lock = threading.Lock() # --------------------------------------------------------------------------- @@ -576,10 +564,17 @@ def run_benchmark( if process.returncode != 0: tmp_result_file.unlink(missing_ok=True) - raise RuntimeError(f"Benchmark script exited with code {process.returncode}") + print(f"ERROR:Benchmark script exited with code {process.returncode}") if not tmp_result_file.exists(): - raise FileNotFoundError(f"Benchmark result file not found: {tmp_result_file}") + with open(DIFFUSION_RESULT_TEMPLATE_PATH, encoding="utf-8") as f: + template_payload = json.load(f) + # Template schema is fixed and owned by this repo: + # ``diffusion_result_template.json`` is a one-item list and metrics live at [0]["result"]. + template_metrics: dict[str, Any] = template_payload[0]["result"] + with open(tmp_result_file, "w", encoding="utf-8") as f: + json.dump(template_metrics, f, ensure_ascii=False, indent=2) + print(f"Benchmark result file not generated, fallback to template: {tmp_result_file}") try: with open(tmp_result_file, encoding="utf-8") as f: diff --git a/tests/dfx/perf/stage_configs/qwen3_omni.yaml b/tests/dfx/perf/stage_configs/qwen3_omni.yaml deleted file mode 100644 index 2add22b87..000000000 --- a/tests/dfx/perf/stage_configs/qwen3_omni.yaml +++ /dev/null @@ -1,101 +0,0 @@ -# Stage config for running Qwen3-Omni-MoE with 3-stage architecture -# Stage 0: Thinker (multimodal understanding + text generation) -# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes) -# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform) - -# The following config has been verified on 2x H100-80G GPUs. -async_chunk: false -stage_args: - - stage_id: 0 - stage_type: llm # Use llm stage type for AR stages - runtime: - devices: "0" - engine_args: - model_stage: thinker - max_num_seqs: 64 - model_arch: Qwen3OmniMoeForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.9 - enforce_eager: false - trust_remote_code: true - engine_output_type: latent # Output hidden states for talker - distributed_executor_backend: "mp" - enable_prefix_caching: false - max_num_batched_tokens: 32768 - hf_config_name: thinker_config - tensor_parallel_size: 1 - final_output: true - final_output_type: text - is_comprehension: true - default_sampling_params: - temperature: 0.4 - top_p: 0.9 - top_k: 1 - max_tokens: 2048 - seed: 42 - detokenize: True - repetition_penalty: 1.05 - - - stage_id: 1 - stage_type: llm # Use llm stage type for AR stages - runtime: - devices: "1" - engine_args: - model_stage: talker - max_num_seqs: 64 - model_arch: Qwen3OmniMoeForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.6 - enforce_eager: false - trust_remote_code: true - engine_output_type: latent # Output codec codes for code2wav - enable_prefix_caching: false - max_num_batched_tokens: 32768 - distributed_executor_backend: "mp" - hf_config_name: talker_config - engine_input_source: [0] - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker - # final_output: true - # final_output_type: text - 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: 2 - stage_type: llm # Use llm stage type for AR stages - runtime: - devices: "1" - engine_args: - model_stage: code2wav - max_num_seqs: 64 - model_arch: Qwen3OmniMoeForConditionalGeneration - worker_type: generation - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - enforce_eager: true - trust_remote_code: true - async_scheduling: false - enable_prefix_caching: false - engine_output_type: audio # Final output: audio waveform - gpu_memory_utilization: 0.1 - distributed_executor_backend: "mp" - max_num_batched_tokens: 100000 - hf_config_name: thinker_config - engine_input_source: [1] - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav - final_output: true - final_output_type: audio - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 65536 - seed: 42 - detokenize: True - repetition_penalty: 1.1 diff --git a/tests/dfx/perf/stage_configs/qwen3_tts.yaml b/tests/dfx/perf/stage_configs/qwen3_tts.yaml deleted file mode 100644 index 97b309056..000000000 --- a/tests/dfx/perf/stage_configs/qwen3_tts.yaml +++ /dev/null @@ -1,96 +0,0 @@ -# Stage config for running Qwen3-TTS with 2-stage architecture -# Stage 0: Talker (text -> 8-layer RVQ codec codes) -# Stage 1: Code2Wav (codec codes -> audio waveform) -# -# The following config has been verified on 1x H100-80G GPU. -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: false - 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: false - 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: 72 - - edges: - - from: 0 - to: 1 - window_size: -1 diff --git a/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json new file mode 100644 index 000000000..3cf13509c --- /dev/null +++ b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json @@ -0,0 +1,49 @@ +[ + { + "test_name": "test_qwen_image_layered_single_device", + "description": "Single-device baseline", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image-Layered", + "serve_args": { + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "640x640_steps20_i2i", + "dataset": "random", + "task": "i2i", + "width": 640, + "height": 640, + "num-inference-steps": 20, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.02, + "latency_mean": 40.0, + "peak_memory_mb_max": 70000, + "peak_memory_mb_mean": 70000 + } + }, + { + "name": "1024x1024_steps35_i2i", + "dataset": "random", + "task": "i2i", + "width": 1024, + "height": 1024, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.005, + "latency_mean": 80.0, + "peak_memory_mb_max": 70000, + "peak_memory_mb_mean": 70000 + } + } + ] + } +] diff --git a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json index 1f3a2bbf7..cdd0cac2c 100644 --- a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json @@ -44,6 +44,52 @@ } ] }, + { + "test_name": "test_qwen_image_single_device_step_execution", + "description": "Single-device baseline (no parallelism) with step execution", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image", + "serve_args": { + "enable-diffusion-pipeline-profiler": true, + "step-execution": true + } + }, + "benchmark_params": [ + { + "name": "512x512_steps20", + "dataset": "random", + "task": "t2i", + "width": 512, + "height": 512, + "num-inference-steps": 20, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.30, + "latency_mean": 3.50, + "peak_memory_mb_mean": 67000 + } + }, + { + "name": "1536x1536_steps35", + "dataset": "random", + "task": "t2i", + "width": 1536, + "height": 1536, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.037, + "latency_mean": 27.0, + "peak_memory_mb_mean": 74000 + } + } + ] + }, { "test_name": "test_qwen_image_ulysses2_cfg2_vae_patch4", "description": "Ulysses SP=2 + CFG-parallel=2 + VAE Patch Parallel=4", @@ -59,22 +105,6 @@ } }, "benchmark_params": [ - { - "name": "512x512_steps20", - "dataset": "random", - "task": "t2i", - "width": 512, - "height": 512, - "num-inference-steps": 20, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.1, - "latency_mean": 2.7, - "peak_memory_mb_mean": 61000 - } - }, { "name": "1536x1536_steps35", "dataset": "random", diff --git a/tests/dfx/perf/tests/test.json b/tests/dfx/perf/tests/test_qwen_omni.json similarity index 86% rename from tests/dfx/perf/tests/test.json rename to tests/dfx/perf/tests/test_qwen_omni.json index 159e27a06..ca3eb5557 100644 --- a/tests/dfx/perf/tests/test.json +++ b/tests/dfx/perf/tests/test_qwen_omni.json @@ -3,7 +3,7 @@ "test_name": "test_qwen3_omni", "server_params": { "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "stage_config_name": "qwen3_omni.yaml" + "extra_cli_args": ["--no-async-chunk"] }, "benchmark_params": [ { @@ -109,25 +109,7 @@ "test_name": "test_qwen3_omni_chunk", "server_params": { "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "stage_config_name": "qwen3_omni.yaml", - "update": { - "async_chunk": true, - "stage_args": { - "0": { - "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk" - }, - "1": { - "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk" - } - } - }, - "delete": { - "stage_args": { - "2": [ - "custom_process_input_func" - ] - } - } + "extra_cli_args": ["--async-chunk"] }, "benchmark_params": [ { @@ -329,37 +311,5 @@ } } ] - }, - { - "test_name": "test_qwen3_tts", - "server_params": { - "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" - }, - "benchmark_params": [ - { - "dataset_name": "random", - "backend": "openai-audio-speech", - "endpoint": "/v1/audio/speech", - "num_prompts": [ - 10, - 40 - ], - "max_concurrency": [ - 1, - 4 - ], - "random_input_len": 100, - "random_output_len": 100, - "extra_body": { - "voice": "Vivian", - "language": "English" - }, - "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "mean_audio_ttfp_ms": [6000, 6000], - "mean_audio_rtf": [0.3, 0.3] - } - } - ] } ] diff --git a/tests/dfx/perf/tests/test_tts.json b/tests/dfx/perf/tests/test_tts.json new file mode 100644 index 000000000..3583b45b4 --- /dev/null +++ b/tests/dfx/perf/tests/test_tts.json @@ -0,0 +1,34 @@ +[ + { + "test_name": "test_qwen3_tts", + "server_params": { + "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" + }, + "benchmark_params": [ + { + "dataset_name": "random", + "backend": "openai-audio-speech", + "endpoint": "/v1/audio/speech", + "num_prompts": [ + 10, + 40 + ], + "max_concurrency": [ + 1, + 4 + ], + "random_input_len": 100, + "random_output_len": 100, + "extra_body": { + "voice": "Vivian", + "language": "English" + }, + "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_audio_ttfp_ms": [6000, 6000], + "mean_audio_rtf": [0.3, 0.3] + } + } + ] + } +] diff --git a/tests/dfx/stability/conftest.py b/tests/dfx/stability/conftest.py index 3a0aee760..e36c88b9a 100644 --- a/tests/dfx/stability/conftest.py +++ b/tests/dfx/stability/conftest.py @@ -1,125 +1,17 @@ -""" -Stability-specific conftest: when pytest is executed under this directory, -resource monitoring is started before each test and finalized after each test, -so each stability test case gets its own HTML report (one report per case). -No need to wrap pytest with `bash resource_monitor.sh run -- pytest ...`. -""" +"""Stability pytest hooks and fixtures.""" -import os import subprocess import sys import threading -import time -from pathlib import Path import pytest -STABILITY_DIR = Path(__file__).resolve().parent -RESOURCE_MONITOR_SCRIPT = STABILITY_DIR / "scripts" / "resource_monitor.sh" -REPO_ROOT = STABILITY_DIR.parent.parent.parent - - -def _start_resource_monitor(): - """Start `resource_monitor.sh start` in the background and return `Popen` or `None`.""" - if not RESOURCE_MONITOR_SCRIPT.is_file(): - return None - try: - proc = subprocess.Popen( - ["bash", str(RESOURCE_MONITOR_SCRIPT), "start", "--backend", "gpu"], - cwd=str(REPO_ROOT), - stdout=subprocess.DEVNULL, - stderr=subprocess.PIPE, - start_new_session=True, - ) - try: - proc.wait(timeout=2) - if proc.returncode != 0: - stderr = proc.stderr.read().decode("utf-8", errors="ignore") if proc.stderr else "" - if stderr.strip(): - sys.stderr.write(f"[Stability] Resource monitor failed to start: {stderr.strip()}\n") - return None - except subprocess.TimeoutExpired: - pass - return proc - except (FileNotFoundError, OSError): - return None - - -def _get_monitor_data_root() -> Path: - data_root = os.environ.get("RESOURCE_MONITOR_DATA_ROOT") or os.environ.get("GPU_MONITOR_DATA_ROOT") - if data_root: - return Path(data_root) - return STABILITY_DIR / "gpu_monitor_data" - - -def _wait_for_run_dir(timeout_sec: int = 10) -> Path | None: - data_root = _get_monitor_data_root() - run_id_file = data_root / "current_run_id" - deadline = time.time() + timeout_sec - while time.time() < deadline: - if run_id_file.is_file(): - run_id = run_id_file.read_text(encoding="utf-8").strip() - if run_id: - run_dir = data_root / run_id - if run_dir.is_dir(): - return run_dir - time.sleep(0.5) - return None - - -def _report_latest_gpu_samples(stop_event: threading.Event) -> None: - """Periodically print the latest sampled GPU line.""" - log_interval = int( - os.environ.get("RESOURCE_MONITOR_LOG_INTERVAL") or os.environ.get("GPU_MONITOR_LOG_INTERVAL") or "15" - ) - log_interval = max(log_interval, 1) - last_line = "" - - time.sleep(min(log_interval, 5)) - while not stop_event.wait(log_interval): - run_dir = _wait_for_run_dir(timeout_sec=1) - if run_dir is None: - continue - csv_file = run_dir / "gpu_metrics.csv" - if not csv_file.is_file(): - continue - try: - lines = csv_file.read_text(encoding="utf-8").splitlines() - except OSError: - continue - if len(lines) <= 1: - continue - latest = lines[-1].strip() - if latest and latest != last_line: - last_line = latest - sys.stderr.write(f"[GPU] {latest}\n") - - -def _finalize_resource_monitor() -> str | None: - """ - Run `resource_monitor.sh finalize` for the current run and generate the report. - Returns the bundle dir path (for this test case's report) if successful, else None. - """ - if not RESOURCE_MONITOR_SCRIPT.is_file(): - return None - try: - result = subprocess.run( - ["bash", str(RESOURCE_MONITOR_SCRIPT), "finalize", "--backend", "gpu"], - cwd=str(REPO_ROOT), - capture_output=True, - text=True, - timeout=60, - check=False, - ) - if result.returncode != 0: - return None - for line in (result.stdout or "").splitlines(): - if line.startswith("GPU_MONITOR_BUNDLE_DIR=") or line.startswith("RESOURCE_MONITOR_BUNDLE_DIR="): - _, _, value = line.partition("=") - return value.strip() if value else None - return None - except (FileNotFoundError, OSError, subprocess.TimeoutExpired): - return None +from tests.dfx.stability.helpers import ( + finalize_resource_monitor, + report_latest_gpu_samples, + start_resource_monitor, + wait_for_run_dir, +) @pytest.fixture(autouse=True) @@ -128,19 +20,19 @@ def stability_resource_monitor_per_test(request: pytest.FixtureRequest): For each test under this directory: start GPU monitor before the test, then finalize after the test so this case gets its own report.html. """ - proc = _start_resource_monitor() + proc = start_resource_monitor() stop_event = threading.Event() reporter: threading.Thread | None = None if proc is not None: reporter = threading.Thread( - target=_report_latest_gpu_samples, + target=report_latest_gpu_samples, args=(stop_event,), name="stability-resource-monitor-reporter", daemon=True, ) reporter.start() - run_dir = _wait_for_run_dir(timeout_sec=5) + run_dir = wait_for_run_dir(timeout_sec=5) node_name = request.node.name if run_dir is not None: sys.stderr.write(f"[Stability] Resource monitor started for test: {node_name} | run dir: {run_dir}\n") @@ -161,7 +53,7 @@ def stability_resource_monitor_per_test(request: pytest.FixtureRequest): except subprocess.TimeoutExpired: proc.kill() proc.wait() - bundle_dir = _finalize_resource_monitor() + bundle_dir = finalize_resource_monitor() node_name = request.node.name if bundle_dir: sys.stderr.write(f"[Stability] Report for test «{node_name}»: {bundle_dir}/report.html\n") diff --git a/tests/dfx/stability/helpers.py b/tests/dfx/stability/helpers.py new file mode 100644 index 000000000..3956bd213 --- /dev/null +++ b/tests/dfx/stability/helpers.py @@ -0,0 +1,117 @@ +"""Stability resource monitor helpers.""" + +from __future__ import annotations + +import os +import subprocess +import sys +import threading +import time +from pathlib import Path + +STABILITY_DIR = Path(__file__).resolve().parent +RESOURCE_MONITOR_SCRIPT = STABILITY_DIR / "scripts" / "resource_monitor.sh" +REPO_ROOT = STABILITY_DIR.parent.parent.parent + + +def start_resource_monitor(): + """Start `resource_monitor.sh start` in the background and return `Popen` or `None`.""" + if not RESOURCE_MONITOR_SCRIPT.is_file(): + return None + try: + proc = subprocess.Popen( + ["bash", str(RESOURCE_MONITOR_SCRIPT), "start", "--backend", "gpu"], + cwd=str(REPO_ROOT), + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + start_new_session=True, + ) + try: + proc.wait(timeout=2) + if proc.returncode != 0: + stderr = proc.stderr.read().decode("utf-8", errors="ignore") if proc.stderr else "" + if stderr.strip(): + sys.stderr.write(f"[Stability] Resource monitor failed to start: {stderr.strip()}\n") + return None + except subprocess.TimeoutExpired: + pass + return proc + except (FileNotFoundError, OSError): + return None + + +def get_monitor_data_root() -> Path: + data_root = os.environ.get("RESOURCE_MONITOR_DATA_ROOT") or os.environ.get("GPU_MONITOR_DATA_ROOT") + if data_root: + return Path(data_root) + return STABILITY_DIR / "gpu_monitor_data" + + +def wait_for_run_dir(timeout_sec: int = 10) -> Path | None: + data_root = get_monitor_data_root() + run_id_file = data_root / "current_run_id" + deadline = time.time() + timeout_sec + while time.time() < deadline: + if run_id_file.is_file(): + run_id = run_id_file.read_text(encoding="utf-8").strip() + if run_id: + run_dir = data_root / run_id + if run_dir.is_dir(): + return run_dir + time.sleep(0.5) + return None + + +def report_latest_gpu_samples(stop_event: threading.Event) -> None: + """Periodically print the latest sampled GPU line.""" + log_interval = int( + os.environ.get("RESOURCE_MONITOR_LOG_INTERVAL") or os.environ.get("GPU_MONITOR_LOG_INTERVAL") or "15" + ) + log_interval = max(log_interval, 1) + last_line = "" + + time.sleep(min(log_interval, 5)) + while not stop_event.wait(log_interval): + run_dir = wait_for_run_dir(timeout_sec=1) + if run_dir is None: + continue + csv_file = run_dir / "gpu_metrics.csv" + if not csv_file.is_file(): + continue + try: + lines = csv_file.read_text(encoding="utf-8").splitlines() + except OSError: + continue + if len(lines) <= 1: + continue + latest = lines[-1].strip() + if latest and latest != last_line: + last_line = latest + sys.stderr.write(f"[GPU] {latest}\n") + + +def finalize_resource_monitor() -> str | None: + """ + Run `resource_monitor.sh finalize` for the current run and generate the report. + Returns the bundle dir path (for this test case's report) if successful, else None. + """ + if not RESOURCE_MONITOR_SCRIPT.is_file(): + return None + try: + result = subprocess.run( + ["bash", str(RESOURCE_MONITOR_SCRIPT), "finalize", "--backend", "gpu"], + cwd=str(REPO_ROOT), + capture_output=True, + text=True, + timeout=60, + check=False, + ) + if result.returncode != 0: + return None + for line in (result.stdout or "").splitlines(): + if line.startswith("GPU_MONITOR_BUNDLE_DIR=") or line.startswith("RESOURCE_MONITOR_BUNDLE_DIR="): + _, _, value = line.partition("=") + return value.strip() if value else None + return None + except (FileNotFoundError, OSError, subprocess.TimeoutExpired): + return None diff --git a/tests/dfx/stability/scripts/test_benchmark_stability.py b/tests/dfx/stability/scripts/test_benchmark_stability.py index e8568652d..620241762 100644 --- a/tests/dfx/stability/scripts/test_benchmark_stability.py +++ b/tests/dfx/stability/scripts/test_benchmark_stability.py @@ -24,7 +24,6 @@ from typing import Any import pytest -from tests.conftest import OmniServer from tests.dfx.conftest import ( create_benchmark_indices, create_test_parameter_mapping, @@ -33,9 +32,10 @@ from tests.dfx.conftest import ( load_configs, ) from tests.dfx.perf.scripts.run_benchmark import run_benchmark +from tests.helpers.runtime import OmniServer STABILITY_DIR = Path(__file__).resolve().parent.parent -STAGE_CONFIGS_DIR = STABILITY_DIR / "stage_configs" +DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy" CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test.json") DEFAULT_NUM_PROMPTS_PER_BATCH = 20 @@ -45,7 +45,7 @@ try: except FileNotFoundError: BENCHMARK_CONFIGS = [] -test_params = create_unique_server_params(BENCHMARK_CONFIGS, STAGE_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] +test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {} _omni_server_lock = threading.Lock() @@ -112,6 +112,8 @@ def _run_one_benchmark_batch( flow=flow, dataset_name=dataset_name, num_prompt=num_prompts, + random_input_len=params.get("random_input_len"), + random_output_len=params.get("random_output_len"), ) return result except (FileNotFoundError, OSError) as e: @@ -217,11 +219,20 @@ def omni_server(request): Multi-stage initialization can take 10-20+ minutes. """ with _omni_server_lock: - test_name, model, stage_config_path = request.param + test_name, model, stage_config_path, stage_overrides, extra_cli_args = request.param print(f"Starting OmniServer with test: {test_name}, model: {model}") - with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server: + server_args = ["--stage-init-timeout", "120"] + # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py): + # deploy-config sets the base; stage-overrides are applied on top. Both can be set. + if stage_config_path: + server_args = ["--deploy-config", stage_config_path] + server_args + if stage_overrides: + server_args = ["--stage-overrides", stage_overrides] + server_args + if extra_cli_args: + server_args = list(extra_cli_args) + server_args + with OmniServer(model, server_args) as server: server.test_name = test_name print("OmniServer started successfully") yield server diff --git a/tests/dfx/stability/stage_configs/qwen3_omni.yaml b/tests/dfx/stability/stage_configs/qwen3_omni.yaml deleted file mode 100644 index 802f8dd24..000000000 --- a/tests/dfx/stability/stage_configs/qwen3_omni.yaml +++ /dev/null @@ -1,101 +0,0 @@ -# Stage config for running Qwen3-Omni-MoE with 3-stage architecture -# Stage 0: Thinker (multimodal understanding + text generation) -# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes) -# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform) - -# The following config has been verified on 2x H100-80G GPUs. -async_chunk: false -stage_args: - - stage_id: 0 - stage_type: llm # Use llm stage type to launch OmniLLM - runtime: - devices: "0" - max_batch_size: 64 - engine_args: - model_stage: thinker - model_arch: Qwen3OmniMoeForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.9 - enforce_eager: false - trust_remote_code: true - engine_output_type: latent # Output hidden states for talker - distributed_executor_backend: "mp" - enable_prefix_caching: false - max_num_batched_tokens: 32768 - hf_config_name: thinker_config - tensor_parallel_size: 1 - final_output: true - final_output_type: text - is_comprehension: true - default_sampling_params: - temperature: 0.4 - top_p: 0.9 - top_k: 1 - max_tokens: 2048 - seed: 42 - detokenize: True - repetition_penalty: 1.05 - - - stage_id: 1 - stage_type: llm # Use llm stage type to launch OmniLLM - runtime: - devices: "1" - max_batch_size: 64 - engine_args: - model_stage: talker - model_arch: Qwen3OmniMoeForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.6 - enforce_eager: false - trust_remote_code: true - engine_output_type: latent # Output codec codes for code2wav - enable_prefix_caching: false - max_num_batched_tokens: 32768 - distributed_executor_backend: "mp" - hf_config_name: talker_config - engine_input_source: [0] - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker - # final_output: true - # final_output_type: text - 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: 2 - stage_type: llm # Use llm stage type to launch OmniLLM - runtime: - devices: "1" - max_batch_size: 64 - engine_args: - model_stage: code2wav - model_arch: Qwen3OmniMoeForConditionalGeneration - worker_type: generation - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - enforce_eager: true - trust_remote_code: true - async_scheduling: false - enable_prefix_caching: false - engine_output_type: audio # Final output: audio waveform - gpu_memory_utilization: 0.1 - distributed_executor_backend: "mp" - max_num_batched_tokens: 1000000 - hf_config_name: thinker_config - engine_input_source: [1] - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav - final_output: true - final_output_type: audio - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 65536 - seed: 42 - detokenize: True - repetition_penalty: 1.1 diff --git a/tests/dfx/stability/tests/test.json b/tests/dfx/stability/tests/test.json index 95993c9c5..255cd5b10 100644 --- a/tests/dfx/stability/tests/test.json +++ b/tests/dfx/stability/tests/test.json @@ -3,7 +3,11 @@ "test_name": "test_qwen3_omni_stability", "server_params": { "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "stage_config_name": "qwen3_omni.yaml" + "stage_overrides": { + "2": { + "max_num_batched_tokens": 1000000 + } + } }, "benchmark_params": [ { @@ -36,25 +40,12 @@ "test_name": "test_qwen3_omni_stability_async_chunk", "server_params": { "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "stage_config_name": "qwen3_omni.yaml", - "update": { - "async_chunk": true, - "stage_args": { - "0": { - "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk" - }, - "1": { - "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk" - } + "stage_overrides": { + "2": { + "max_num_batched_tokens": 1000000 } }, - "delete": { - "stage_args": { - "2": [ - "custom_process_input_func" - ] - } - } + "extra_cli_args": ["--async-chunk"] }, "benchmark_params": [ { diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py new file mode 100644 index 000000000..0b7ef7235 --- /dev/null +++ b/tests/diffusion/cache/test_cache_dit.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Model specific tests for CacheDiT enablement. +""" + +from unittest.mock import Mock, patch + +import pytest + +import vllm_omni.diffusion.cache.cache_dit_backend as cd_backend +from vllm_omni.diffusion.data import DiffusionCacheConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +SEPARATE_CFG_ENABLERS = [ + cd_backend.enable_cache_for_ltx2, + cd_backend.enable_cache_for_wan22, + cd_backend.enable_cache_for_longcat_image, +] + +SAMPLE_CACHE_CONFIG = DiffusionCacheConfig() + + +@pytest.mark.parametrize("enabler", SEPARATE_CFG_ENABLERS) +@patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter") +@patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit") +def test_separate_cfg(mock_cache_dit, mock_block_adapter, enabler): + """Ensure that custom enablers for models with separate CFG pass + the param through to cache_dit correctly. + + Regression test for: https://github.com/vllm-project/vllm-omni/pull/2860 + """ + mock_pipeline = Mock() + enabler(mock_pipeline, SAMPLE_CACHE_CONFIG) + + mock_cache_dit.enable_cache.assert_called_once() + adapter_kwargs = mock_block_adapter.call_args.kwargs + assert adapter_kwargs["has_separate_cfg"] is True diff --git a/tests/diffusion/cache/test_teacache_extractors.py b/tests/diffusion/cache/test_teacache_extractors.py index c22a60e22..4bb958a36 100644 --- a/tests/diffusion/cache/test_teacache_extractors.py +++ b/tests/diffusion/cache/test_teacache_extractors.py @@ -21,7 +21,7 @@ from unittest.mock import MagicMock, Mock, patch import pytest import torch -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test from vllm_omni.diffusion.cache.teacache.extractors import extract_flux2_context, extract_flux2_klein_context from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import ( Flux2Transformer2DModel, diff --git a/tests/diffusion/distributed/test_ulysses_uaa_perf.py b/tests/diffusion/distributed/test_ulysses_uaa_perf.py index 04bbf5ee8..2a16a9ae5 100644 --- a/tests/diffusion/distributed/test_ulysses_uaa_perf.py +++ b/tests/diffusion/distributed/test_ulysses_uaa_perf.py @@ -17,7 +17,7 @@ import pytest import torch import torch.distributed as dist -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test from vllm_omni.diffusion.attention.parallel.ulysses import ( _all_gather_int, _ulysses_all_to_all_any_o, diff --git a/tests/diffusion/layers/test_norm.py b/tests/diffusion/layers/test_norm.py new file mode 100644 index 000000000..e42041528 --- /dev/null +++ b/tests/diffusion/layers/test_norm.py @@ -0,0 +1,453 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for LayerNorm and RMSNorm custom ops in diffusion layers.""" + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +# ── Import tests ── + + +def test_layernorm_import(): + """Verify LayerNorm can be imported from the norm module.""" + from vllm_omni.diffusion.layers.norm import LayerNorm # noqa: F401 + + +def test_rmsnorm_import(): + """Verify RMSNorm can be imported from the norm module.""" + from vllm_omni.diffusion.layers.norm import RMSNorm # noqa: F401 + + +# ── LayerNorm tests ── + + +def test_layernorm_forward_shape(): + """LayerNorm produces correct output shapes.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + batch = 2 + seq_len = 4 + norm = LayerNorm(dim) + + x = torch.randn(batch, seq_len, dim) + out = norm(x) + + assert out.shape == (batch, seq_len, dim) + + +def test_layernorm_forward_shape_2d(): + """LayerNorm works with 2D input tensors.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + batch = 2 + norm = LayerNorm(dim) + + x = torch.randn(batch, dim) + out = norm(x) + + assert out.shape == (batch, dim) + + +def test_layernorm_preserves_dtype_fp32(): + """LayerNorm preserves float32 dtype.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + norm = LayerNorm(dim) + + x = torch.randn(2, 4, dim, dtype=torch.float32) + out = norm(x) + + assert out.dtype == torch.float32 + + +def test_layernorm_preserves_dtype_fp16(): + """LayerNorm preserves float16 dtype.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + norm = LayerNorm(dim) + + x = torch.randn(2, 4, dim, dtype=torch.float16) + out = norm(x) + + assert out.dtype == torch.float16 + + +def test_layernorm_preserves_dtype_bf16(): + """LayerNorm preserves bfloat16 dtype.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + norm = LayerNorm(dim) + + x = torch.randn(2, 4, dim, dtype=torch.bfloat16) + out = norm(x) + + assert out.dtype == torch.bfloat16 + + +def test_layernorm_without_elementwise_affine(): + """LayerNorm works without elementwise_affine (no learned parameters).""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + norm = LayerNorm(dim, elementwise_affine=False) + + assert norm.weight is None + assert norm.bias is None + + x = torch.randn(2, 4, dim) + out = norm(x) + + assert out.shape == (2, 4, dim) + + +def test_layernorm_custom_eps(): + """LayerNorm accepts custom epsilon value.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + eps = 1e-5 + norm = LayerNorm(dim, eps=eps) + + assert norm.eps == eps + + +def test_layernorm_has_learnable_parameters(): + """LayerNorm has learnable weight and bias by default.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + norm = LayerNorm(dim) + + assert norm.weight is not None + assert norm.bias is not None + assert norm.weight.shape == (dim,) + assert norm.bias.shape == (dim,) + + +def test_layernorm_matches_fp32_reference(): + """Verify LayerNorm produces identical output to FP32 nn.LayerNorm.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + eps = 1e-6 + torch.manual_seed(42) + + ours = LayerNorm(dim, eps=eps) + ref = torch.nn.LayerNorm(dim, eps=eps) + + # Copy weights + ref.weight.data.copy_(ours.weight.data) + ref.bias.data.copy_(ours.bias.data) + + x = torch.randn(2, 4, dim) + + out_ours = ours(x) + out_ref = ref(x.float()).to(x.dtype) + + torch.testing.assert_close(out_ours, out_ref, atol=1e-5, rtol=1e-5) + + +def test_layernorm_matches_diffusers_fp32layernorm(): + """Verify LayerNorm produces identical output to diffusers FP32LayerNorm.""" + from diffusers.models.normalization import FP32LayerNorm + + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + eps = 1e-6 + torch.manual_seed(42) + + ours = LayerNorm(dim, eps=eps) + ref = FP32LayerNorm(dim, eps=eps) + + # Copy weights + ref.weight.data.copy_(ours.weight.data) + ref.bias.data.copy_(ours.bias.data) + + # Test with fp16 input to verify FP32 computation + x = torch.randn(2, 4, dim, dtype=torch.float16) + + out_ours = ours(x) + out_ref = ref(x) + + torch.testing.assert_close(out_ours, out_ref, atol=1e-3, rtol=1e-3) + + +# ── RMSNorm tests ── + + +def test_rmsnorm_forward_shape(): + """RMSNorm produces correct output shapes.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + batch = 2 + seq_len = 4 + norm = RMSNorm(hidden_size) + + x = torch.randn(batch, seq_len, hidden_size) + out = norm(x) + + assert out.shape == (batch, seq_len, hidden_size) + + +def test_rmsnorm_forward_shape_2d(): + """RMSNorm works with 2D input tensors.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + batch = 2 + norm = RMSNorm(hidden_size) + + x = torch.randn(batch, hidden_size) + out = norm(x) + + assert out.shape == (batch, hidden_size) + + +def test_rmsnorm_preserves_dtype_fp32(): + """RMSNorm preserves float32 dtype.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + norm = RMSNorm(hidden_size) + + x = torch.randn(2, 4, hidden_size, dtype=torch.float32) + out = norm(x) + + assert out.dtype == torch.float32 + + +def test_rmsnorm_preserves_dtype_fp16(): + """RMSNorm preserves float16 dtype.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + norm = RMSNorm(hidden_size) + + x = torch.randn(2, 4, hidden_size, dtype=torch.float16) + out = norm(x) + + assert out.dtype == torch.float16 + + +def test_rmsnorm_preserves_dtype_bf16(): + """RMSNorm preserves bfloat16 dtype.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + norm = RMSNorm(hidden_size) + + x = torch.randn(2, 4, hidden_size, dtype=torch.bfloat16) + out = norm(x) + + assert out.dtype == torch.bfloat16 + + +def test_rmsnorm_custom_eps(): + """RMSNorm accepts custom epsilon value.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + eps = 1e-5 + norm = RMSNorm(hidden_size, eps=eps) + + assert norm.variance_epsilon == eps + + +def test_rmsnorm_has_weight_parameter(): + """RMSNorm has learnable weight parameter initialized to ones.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + norm = RMSNorm(hidden_size) + + assert norm.weight is not None + assert norm.weight.shape == (hidden_size,) + torch.testing.assert_close(norm.weight, torch.ones(hidden_size)) + + +def test_rmsnorm_numerical_correctness(): + """Verify RMSNorm produces numerically correct output.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + eps = 1e-6 + torch.manual_seed(42) + + norm = RMSNorm(hidden_size, eps=eps) + x = torch.randn(2, 4, hidden_size) + + # Compute expected output manually + x_fp32 = x.to(torch.float32) + variance = x_fp32.pow(2).mean(-1, keepdim=True) + expected = x_fp32 * torch.rsqrt(variance + eps) + expected = norm.weight.to(torch.float32) * expected + expected = expected.to(x.dtype) + + out = norm(x) + + torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5) + + +def test_rmsnorm_matches_reference_implementation(): + """Verify RMSNorm matches a reference implementation.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + def reference_rmsnorm(x, weight, eps): + """Reference RMSNorm implementation.""" + input_dtype = x.dtype + x = x.to(torch.float32) + variance = x.pow(2).mean(-1, keepdim=True) + out = x * torch.rsqrt(variance + eps) + out = weight.to(torch.float32) * out + return out.to(input_dtype) + + hidden_size = 128 + eps = 1e-6 + torch.manual_seed(123) + + norm = RMSNorm(hidden_size, eps=eps) + + # Test with various dtypes + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + x = torch.randn(4, 8, hidden_size, dtype=dtype) + expected = reference_rmsnorm(x, norm.weight, eps) + out = norm(x) + torch.testing.assert_close(out, expected, atol=1e-3, rtol=1e-3) + + +# ── CustomOp dispatch tests ── + + +def test_layernorm_inherits_from_customop(): + """LayerNorm inherits from CustomOp for platform dispatch.""" + from vllm_omni.diffusion.layers.custom_op import CustomOp + from vllm_omni.diffusion.layers.norm import LayerNorm + + norm = LayerNorm(64) + assert isinstance(norm, CustomOp) + + +def test_rmsnorm_inherits_from_customop(): + """RMSNorm inherits from CustomOp for platform dispatch.""" + from vllm_omni.diffusion.layers.custom_op import CustomOp + from vllm_omni.diffusion.layers.norm import RMSNorm + + norm = RMSNorm(64) + assert isinstance(norm, CustomOp) + + +def test_layernorm_has_platform_methods(): + """LayerNorm has forward methods for each platform.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + norm = LayerNorm(64) + + assert hasattr(norm, "forward_cuda") + assert hasattr(norm, "forward_hip") + assert hasattr(norm, "forward_xpu") + assert hasattr(norm, "forward_npu") + assert hasattr(norm, "forward_native") + + +def test_rmsnorm_has_platform_methods(): + """RMSNorm has forward methods for each platform.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + norm = RMSNorm(64) + + assert hasattr(norm, "forward_cuda") + assert hasattr(norm, "forward_hip") + assert hasattr(norm, "forward_xpu") + assert hasattr(norm, "forward_npu") + assert hasattr(norm, "forward_native") + + +def test_layernorm_forward_native_directly(): + """LayerNorm.forward_native can be called directly.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + norm = LayerNorm(dim) + x = torch.randn(2, 4, dim) + + out = norm.forward_native(x) + + assert out.shape == (2, 4, dim) + + +def test_rmsnorm_forward_native_directly(): + """RMSNorm.forward_native can be called directly.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + norm = RMSNorm(hidden_size) + x = torch.randn(2, 4, hidden_size) + + out = norm.forward_native(x) + + assert out.shape == (2, 4, hidden_size) + + +# ── Edge case tests ── + + +def test_layernorm_with_large_dim(): + """LayerNorm works with large hidden dimensions.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 4096 + norm = LayerNorm(dim) + x = torch.randn(1, 16, dim) + + out = norm(x) + + assert out.shape == (1, 16, dim) + + +def test_rmsnorm_with_large_dim(): + """RMSNorm works with large hidden dimensions.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 4096 + norm = RMSNorm(hidden_size) + x = torch.randn(1, 16, hidden_size) + + out = norm(x) + + assert out.shape == (1, 16, hidden_size) + + +def test_layernorm_with_single_element_batch(): + """LayerNorm works with batch size of 1.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + norm = LayerNorm(dim) + x = torch.randn(1, 1, dim) + + out = norm(x) + + assert out.shape == (1, 1, dim) + + +def test_rmsnorm_with_single_element_batch(): + """RMSNorm works with batch size of 1.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + norm = RMSNorm(hidden_size) + x = torch.randn(1, 1, hidden_size) + + out = norm(x) + + assert out.shape == (1, 1, hidden_size) diff --git a/tests/diffusion/layers/test_rotary_emb_equivalence.py b/tests/diffusion/layers/test_rotary_emb_equivalence.py new file mode 100644 index 000000000..2fbb7a31f --- /dev/null +++ b/tests/diffusion/layers/test_rotary_emb_equivalence.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Numerical equivalence tests for rotary embedding implementations (#2436). + +Verifies that the optimized stack+flatten RoPE produces bit-identical results +to the original strided-slice implementation across various tensor shapes and +dtypes, ensuring the refactor is safe. +""" + +from __future__ import annotations + +import pytest +import torch + + +def _apply_rotary_emb_helios_original( + hidden_states: torch.Tensor, + freqs_cis: torch.Tensor, +) -> torch.Tensor: + """Original Helios RoPE using strided slice assignment (pre-#2436).""" + x_1, x_2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) + cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1) + out = torch.empty_like(hidden_states) + out[..., 0::2] = x_1 * cos[..., 0::2] - x_2 * sin[..., 1::2] + out[..., 1::2] = x_1 * sin[..., 1::2] + x_2 * cos[..., 0::2] + return out.type_as(hidden_states) + + +def _apply_rotary_emb_helios_optimized( + hidden_states: torch.Tensor, + freqs_cis: torch.Tensor, +) -> torch.Tensor: + """Optimized Helios RoPE using stack+flatten (post-#2436).""" + x_1, x_2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) + cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1) + rotated = torch.stack( + ( + x_1 * cos[..., 0::2] - x_2 * sin[..., 1::2], + x_1 * sin[..., 1::2] + x_2 * cos[..., 0::2], + ), + dim=-1, + ) + return rotated.flatten(-2, -1).type_as(hidden_states) + + +def _make_inputs( + batch: int, + seq_len: int, + num_heads: int, + head_dim: int, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """Generate random hidden_states and freqs_cis for testing.""" + torch.manual_seed(42) + hidden_states = torch.randn(batch, seq_len, num_heads, head_dim, dtype=dtype) + # freqs_cis: [B, seq, head_dim*2] — cos and sin concatenated along last dim + freqs_cis = torch.randn(batch, seq_len, head_dim * 2, dtype=dtype) + return hidden_states, freqs_cis + + +class TestHeliosRoPEEquivalence: + """Verify optimized Helios RoPE is numerically identical to original.""" + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_equivalence_across_dtypes(self, dtype: torch.dtype) -> None: + """Optimized output must be bit-identical to original across dtypes.""" + hidden, freqs = _make_inputs(2, 16, 8, 64, dtype=dtype) + original = _apply_rotary_emb_helios_original(hidden, freqs) + optimized = _apply_rotary_emb_helios_optimized(hidden, freqs) + torch.testing.assert_close(optimized, original, atol=0, rtol=0) + + @pytest.mark.parametrize( + "batch,seq_len,num_heads,head_dim", + [ + (1, 8, 1, 32), # minimal: single batch, single head + (2, 16, 8, 64), # typical transformer config + (1, 8192, 4, 64), # video-scale patch tokens (720p DiT) + (4, 32, 16, 128), # large head_dim + ], + ) + def test_equivalence_across_shapes(self, batch: int, seq_len: int, num_heads: int, head_dim: int) -> None: + """Equivalence must hold across different tensor shapes.""" + hidden, freqs = _make_inputs(batch, seq_len, num_heads, head_dim) + original = _apply_rotary_emb_helios_original(hidden, freqs) + optimized = _apply_rotary_emb_helios_optimized(hidden, freqs) + torch.testing.assert_close(optimized, original, atol=0, rtol=0) + + def test_output_contiguous(self) -> None: + """Optimized output should be contiguous in memory.""" + hidden, freqs = _make_inputs(2, 16, 8, 64) + optimized = _apply_rotary_emb_helios_optimized(hidden, freqs) + assert optimized.is_contiguous() + + def test_output_shape_preserved(self) -> None: + """Output shape must match input shape.""" + hidden, freqs = _make_inputs(2, 16, 8, 64) + optimized = _apply_rotary_emb_helios_optimized(hidden, freqs) + assert optimized.shape == hidden.shape + + def test_output_dtype_preserved(self) -> None: + """Output dtype must match input dtype.""" + hidden, freqs = _make_inputs(2, 16, 8, 64, dtype=torch.float16) + optimized = _apply_rotary_emb_helios_optimized(hidden, freqs) + assert optimized.dtype == hidden.dtype + + def test_odd_head_dim_raises(self) -> None: + """Odd head_dim should fail at unflatten (not a valid RoPE config).""" + hidden = torch.randn(1, 4, 2, 63) + freqs = torch.randn(1, 4, 126) + with pytest.raises(RuntimeError): + _apply_rotary_emb_helios_optimized(hidden, freqs) diff --git a/tests/diffusion/lora/conftest.py b/tests/diffusion/lora/helpers.py similarity index 100% rename from tests/diffusion/lora/conftest.py rename to tests/diffusion/lora/helpers.py diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index 83ac7a114..785f5d842 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -8,7 +8,7 @@ import torch from vllm.lora.lora_weights import LoRALayerWeights from vllm.lora.utils import get_supported_lora_modules -from tests.diffusion.lora.conftest import ( +from tests.diffusion.lora.helpers import ( DummyBaseLayerWithLoRA, FakeLinearBase, fake_replace_submodule, diff --git a/tests/diffusion/models/bagel/test_bagel_lora.py b/tests/diffusion/models/bagel/test_bagel_lora.py index 8cb3446ed..c285758fe 100644 --- a/tests/diffusion/models/bagel/test_bagel_lora.py +++ b/tests/diffusion/models/bagel/test_bagel_lora.py @@ -11,7 +11,7 @@ import pytest import torch from safetensors.torch import save_file -from tests.diffusion.lora.conftest import ( +from tests.diffusion.lora.helpers import ( DummyBaseLayerWithLoRA, FakeLinearBase, fake_replace_submodule, diff --git a/tests/diffusion/models/dmd2/__init__.py b/tests/diffusion/models/dmd2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py b/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py new file mode 100644 index 000000000..e270390bd --- /dev/null +++ b/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX2T2VDMD2Pipeline +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, WanT2VDMD2Pipeline +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import Wan22I2VPipeline, WanI2VDMD2Pipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest, OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# DMD2 subclass → immediate base pipeline whose __init__ loads model weights (mocked in tests). +_DMD2_BASE = { + WanT2VDMD2Pipeline: Wan22Pipeline, + WanI2VDMD2Pipeline: Wan22I2VPipeline, + LTX2T2VDMD2Pipeline: LTX2Pipeline, + LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline, +} + + +def _make_pipeline(cls): + """Run the DMD2 __init__ with the base pipeline mocked out (no model weights loaded).""" + + base = _DMD2_BASE[cls] + od_config = MagicMock() + od_config.model = "/nonexistent" + + def _mock_base_init(self, *a, **kw): + self.od_config = od_config + + with patch.object(base, "__init__", _mock_base_init): + pipeline = object.__new__(cls) + torch.nn.Module.__init__(pipeline) + cls.__init__(pipeline, od_config=od_config) + return pipeline + + +def _make_request(prompts=None, **sp_kwargs) -> OmniDiffusionRequest: + sp = OmniDiffusionSamplingParams(**sp_kwargs) + return OmniDiffusionRequest( + prompts=prompts or [{"prompt": "a cat dancing"}], + sampling_params=sp, + ) + + +@pytest.fixture( + params=list(_DMD2_BASE.keys()), + ids=["wan_t2v", "wan_i2v", "ltx2_t2v", "ltx2_i2v"], +) +def pipeline(request): + return _make_pipeline(request.param) + + +# --------------------------------------------------------------------------- +# num_inference_steps +# --------------------------------------------------------------------------- + + +def test_num_inference_steps_forced_to_dmd2_value(pipeline): + req = _make_request(num_inference_steps=40) + pipeline._sanitize_dmd2_request(req) + assert req.sampling_params.num_inference_steps == pipeline.num_inference_steps + + +def test_num_inference_steps_already_correct(pipeline): + req = _make_request(num_inference_steps=pipeline.num_inference_steps) + pipeline._sanitize_dmd2_request(req) + assert req.sampling_params.num_inference_steps == pipeline.num_inference_steps + + +# --------------------------------------------------------------------------- +# guidance_scale +# --------------------------------------------------------------------------- + + +def test_guidance_scale_forced_to_one(pipeline): + req = _make_request(guidance_scale=5.0, guidance_scale_provided=True) + pipeline._sanitize_dmd2_request(req) + assert req.sampling_params.guidance_scale == pipeline.dmd2_guidance_scale + assert req.sampling_params.guidance_scale_provided is False + + +def test_guidance_scale_already_correct(pipeline): + req = _make_request(guidance_scale=pipeline.dmd2_guidance_scale, guidance_scale_provided=False) + pipeline._sanitize_dmd2_request(req) + assert req.sampling_params.guidance_scale == pipeline.dmd2_guidance_scale + + +def test_guidance_scale_provided_flag_cleared(pipeline): + """guidance_scale_provided=True must be cleared even if scale is already dmd2_guidance_scale.""" + req = _make_request(guidance_scale=pipeline.dmd2_guidance_scale, guidance_scale_provided=True) + pipeline._sanitize_dmd2_request(req) + assert req.sampling_params.guidance_scale_provided is False + + +def test_guidance_scale_2_cleared(pipeline): + req = _make_request(guidance_scale_2=3.0) + pipeline._sanitize_dmd2_request(req) + assert req.sampling_params.guidance_scale_2 is None + + +def test_guidance_scale_2_unset_unchanged(pipeline): + req = _make_request() + pipeline._sanitize_dmd2_request(req) + assert req.sampling_params.guidance_scale_2 is None + + +def test_true_cfg_scale_cleared(pipeline): + req = _make_request(true_cfg_scale=2.0) + pipeline._sanitize_dmd2_request(req) + assert req.sampling_params.true_cfg_scale is None + + +def test_do_classifier_free_guidance_forced_false(pipeline): + req = _make_request(do_classifier_free_guidance=True) + pipeline._sanitize_dmd2_request(req) + assert req.sampling_params.do_classifier_free_guidance is False + + +def test_is_cfg_negative_forced_false(pipeline): + req = _make_request(is_cfg_negative=True) + pipeline._sanitize_dmd2_request(req) + assert req.sampling_params.is_cfg_negative is False + + +def test_negative_prompt_stripped_from_prompt_dict(pipeline): + req = _make_request(prompts=[{"prompt": "a cat", "negative_prompt": "blurry"}]) + pipeline._sanitize_dmd2_request(req) + assert "negative_prompt" not in req.prompts[0] + assert req.prompts[0]["prompt"] == "a cat" + + +def test_no_negative_prompt_unchanged(pipeline): + req = _make_request(prompts=[{"prompt": "a cat"}]) + pipeline._sanitize_dmd2_request(req) + assert req.prompts[0] == {"prompt": "a cat"} + + +def test_string_prompt_not_mutated(pipeline): + """String prompts (not dicts) must pass through unchanged.""" + req = _make_request(prompts=["a cat dancing"]) + pipeline._sanitize_dmd2_request(req) + assert req.prompts == ["a cat dancing"] + + +def test_multiple_prompts_all_sanitized(pipeline): + req = _make_request( + prompts=[ + {"prompt": "a cat", "negative_prompt": "blurry"}, + {"prompt": "a dog", "negative_prompt": "ugly"}, + ] + ) + pipeline._sanitize_dmd2_request(req) + for p in req.prompts: + assert "negative_prompt" not in p + + +# --------------------------------------------------------------------------- +# Clean request — nothing changes +# --------------------------------------------------------------------------- + + +def test_clean_request_no_changes(pipeline): + req = _make_request( + guidance_scale=pipeline.dmd2_guidance_scale, + guidance_scale_provided=False, + do_classifier_free_guidance=False, + is_cfg_negative=False, + ) + pipeline._sanitize_dmd2_request(req) + assert req.sampling_params.guidance_scale == pipeline.dmd2_guidance_scale + assert req.sampling_params.guidance_scale_provided is False + assert req.sampling_params.guidance_scale_2 is None + assert req.sampling_params.true_cfg_scale is None + assert req.sampling_params.do_classifier_free_guidance is False + assert req.sampling_params.is_cfg_negative is False diff --git a/tests/diffusion/models/dmd2/test_dmd2_scheduler.py b/tests/diffusion/models/dmd2/test_dmd2_scheduler.py new file mode 100644 index 000000000..32d00dbf1 --- /dev/null +++ b/tests/diffusion/models/dmd2/test_dmd2_scheduler.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX2T2VDMD2Pipeline +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, WanT2VDMD2Pipeline +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import Wan22I2VPipeline, WanI2VDMD2Pipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest, OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_DMD2_TIMESTEPS = [999, 937, 833, 624] + +# DMD2 subclass → immediate base pipeline whose __init__ loads model weights (mocked in tests). +_DMD2_BASE = { + WanT2VDMD2Pipeline: Wan22Pipeline, + WanI2VDMD2Pipeline: Wan22I2VPipeline, + LTX2T2VDMD2Pipeline: LTX2Pipeline, + LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline, +} + + +def _make_pipeline(cls): + """Run the DMD2 __init__ (including __init_dmd2__) with the base pipeline mocked.""" + + base = _DMD2_BASE[cls] + od_config = MagicMock() + od_config.model = "/nonexistent" + + def _mock_base_init(self, *a, **kw): + self.od_config = od_config # __init_dmd2__ needs this + + with patch.object(base, "__init__", _mock_base_init): + pipeline = object.__new__(cls) + torch.nn.Module.__init__(pipeline) + cls.__init__(pipeline, od_config=od_config) + return pipeline + + +def _make_request(**sp_kwargs) -> OmniDiffusionRequest: + sp = OmniDiffusionSamplingParams(**sp_kwargs) + return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp) + + +@pytest.fixture( + params=list(_DMD2_BASE.keys()), + ids=["wan_t2v", "wan_i2v", "ltx2_t2v", "ltx2_i2v"], +) +def pipeline(request): + return _make_pipeline(request.param) + + +# --------------------------------------------------------------------------- +# forward() timestep injection +# --------------------------------------------------------------------------- + + +def _fake_parent_forward(self, req, *args, num_inference_steps=40, **kwargs): + """Stub that calls set_timesteps as the real parent does.""" + self.scheduler.set_timesteps(num_inference_steps, device="cpu") + return MagicMock() + + +def test_forward_timesteps_match_dmd2_schedule(pipeline): + """After forward() runs, scheduler.timesteps must equal the DMD2 training schedule.""" + parent = _DMD2_BASE[type(pipeline)] + + # Baseline: calling set_timesteps(40) without the DMD2 override gives a different schedule + pipeline.scheduler.set_timesteps(40, device="cpu") + default_timesteps = pipeline.scheduler.timesteps.long().tolist() + assert default_timesteps == _DMD2_TIMESTEPS, ( + "DMD2EulerScheduler should always return DMD2 timesteps regardless of num_steps" + ) + + with patch.object(parent, "forward", _fake_parent_forward): + pipeline.forward(_make_request()) + + assert pipeline.scheduler.timesteps.long().tolist() == _DMD2_TIMESTEPS + + +def test_forward_timesteps_idempotent_across_calls(pipeline): + """Successive forward() calls must not cause scheduler state to drift.""" + parent = _DMD2_BASE[type(pipeline)] + + with patch.object(parent, "forward", _fake_parent_forward): + pipeline.forward(_make_request()) + pipeline.forward(_make_request()) + + assert pipeline.scheduler.timesteps.long().tolist() == _DMD2_TIMESTEPS diff --git a/tests/diffusion/models/flux2/test_flux2_transformer_tp.py b/tests/diffusion/models/flux2/test_flux2_transformer_tp.py index 54dda1dd0..c613bb0b4 100644 --- a/tests/diffusion/models/flux2/test_flux2_transformer_tp.py +++ b/tests/diffusion/models/flux2/test_flux2_transformer_tp.py @@ -2,7 +2,7 @@ import pytest import torch from pytest_mock import MockerFixture -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test from vllm_omni.diffusion.models.flux2.flux2_transformer import ( Flux2PosEmbed, Flux2Transformer2DModel, diff --git a/tests/diffusion/models/glm_image/test_glm_image_sp.py b/tests/diffusion/models/glm_image/test_glm_image_sp.py new file mode 100644 index 000000000..40d1c8730 --- /dev/null +++ b/tests/diffusion/models/glm_image/test_glm_image_sp.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for GLM-Image Sequence Parallelism support.""" + +import pytest + +from vllm_omni.diffusion.data import DiffusionParallelConfig + + +@pytest.fixture(scope="function", autouse=True) +def setup_sp_groups(mocker): + """Set up SP and TP groups for each test function.""" + mock_get_sp_group = mocker.patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group") + mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1) + mock_get_tp_group = mocker.patch("vllm.distributed.parallel_state.get_tp_group") + + mock_sp_group = mocker.MagicMock() + mock_sp_group.world_size = 4 + mock_get_sp_group.return_value = mock_sp_group + + mock_tp_group = mocker.MagicMock() + mock_tp_group.world_size = 1 + mock_get_tp_group.return_value = mock_tp_group + yield + + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_glm_image_sp_plan_defined(): + """Test that _sp_plan is properly defined on GlmImageTransformer2DModel.""" + from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( + GlmImageTransformer2DModel, + ) + + assert hasattr(GlmImageTransformer2DModel, "_sp_plan") + plan = GlmImageTransformer2DModel._sp_plan + assert plan is not None + + # Verify plan structure + assert "prepare" in plan + assert "proj_out" in plan + + +def test_glm_image_sp_plan_valid(): + """Validate _sp_plan structure.""" + from vllm_omni.diffusion.distributed.sp_plan import validate_sp_plan + from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( + GlmImageTransformer2DModel, + ) + + plan = GlmImageTransformer2DModel._sp_plan + validate_sp_plan(plan) + + +def test_glm_image_prepare_module_exists(): + """Test that GlmImagePrepare module exists.""" + from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( + GlmImagePrepare, + ) + + assert GlmImagePrepare is not None + + +def test_glm_image_attention_accepts_parallel_config(): + """Test that GlmImageAttention accepts parallel_config parameter.""" + from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( + GlmImageAttention, + ) + + parallel_config = DiffusionParallelConfig( + ulysses_degree=2, + ring_degree=2, + tensor_parallel_size=1, + sequence_parallel_size=4, + ) + + attn = GlmImageAttention( + dim=2560, + num_heads=64, + head_dim=40, + parallel_config=parallel_config, + ) + + assert attn.parallel_config is not None + assert attn.parallel_config.sequence_parallel_size == 4 + + +def test_glm_image_transformer_block_accepts_parallel_config(): + """Test that GlmImageTransformerBlock accepts parallel_config parameter.""" + from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( + GlmImageTransformerBlock, + ) + + parallel_config = DiffusionParallelConfig( + ulysses_degree=2, + ring_degree=2, + tensor_parallel_size=1, + sequence_parallel_size=4, + ) + + block = GlmImageTransformerBlock( + dim=2560, + num_attention_heads=64, + attention_head_dim=40, + time_embed_dim=512, + parallel_config=parallel_config, + ) + + assert block.attn1.parallel_config is not None + assert block.attn1.parallel_config.sequence_parallel_size == 4 + + +def test_glm_image_has_sp_support(): + """Test that GLM-Image has SP support implemented.""" + from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( + GlmImageTransformer2DModel, + ) + + # Check that the model has parallel_config support + assert hasattr(GlmImageTransformer2DModel, "__init__") + + # Verify the model can be instantiated with SP config + + # This test just verifies the structure exists + # Actual SP testing requires multi-GPU setup + + +@pytest.mark.cuda +@pytest.mark.sp +def test_glm_image_sp_inference(): + """Test SP inference (requires multi-GPU setup).""" + pytest.skip("Requires multi-GPU SP setup") diff --git a/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py similarity index 85% rename from tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py rename to tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py index 2cda9116c..626f78eed 100644 --- a/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py +++ b/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py @@ -12,7 +12,7 @@ class TestSetForwardContextNumTokens: def test_sets_num_tokens_when_context_available(self, mocker): """num_tokens should be set on ForwardContext when available.""" - import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe mock_ctx = mocker.MagicMock() del mock_ctx.in_profile_run # simulate missing attr @@ -26,7 +26,7 @@ class TestSetForwardContextNumTokens: def test_sets_in_profile_run_only_if_missing(self, mocker): """in_profile_run should not be overwritten if already set.""" - import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe mock_ctx = mocker.MagicMock() mock_ctx.in_profile_run = True # already set @@ -40,7 +40,7 @@ class TestSetForwardContextNumTokens: def test_noop_when_context_unavailable(self, mocker): """Should do nothing when ForwardContext is not available.""" - import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=False) mock_get = mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context") @@ -55,11 +55,11 @@ class TestHunyuanFusedMoEPlatformDispatch: def test_default_platform_uses_default_impl_qualname(self, mocker): """HunyuanFusedMoE should resolve the impl class from the platform hook.""" - import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe mock_platform = mocker.MagicMock() mock_platform.get_diffusion_model_impl_qualname.return_value = ( - "vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe.HunyuanFusedMoEDefault" + "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe.HunyuanFusedMoEDefault" ) mocker.patch.object( @@ -71,7 +71,7 @@ class TestHunyuanFusedMoEPlatformDispatch: mock_impl = mocker.MagicMock() mock_resolve.return_value = mock_impl - from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import ( + from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import ( HunyuanFusedMoE, ) @@ -80,7 +80,7 @@ class TestHunyuanFusedMoEPlatformDispatch: mock_platform.prepare_diffusion_op_runtime.assert_called_once_with("hunyuan_fused_moe") mock_platform.get_diffusion_model_impl_qualname.assert_called_once_with("hunyuan_fused_moe") mock_resolve.assert_called_once_with( - "vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe.HunyuanFusedMoEDefault" + "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe.HunyuanFusedMoEDefault" ) mock_impl.assert_called_once_with(prefix="") @@ -90,7 +90,7 @@ class TestHunyuanFusedMoEFactory: def test_new_delegates_to_impl_class(self, mocker): """HunyuanFusedMoE(prefix=..., **kwargs) should instantiate and return impl instance.""" - import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe class MockImpl: def __init__(self, *, prefix: str = "", **kwargs): @@ -104,7 +104,7 @@ class TestHunyuanFusedMoEFactory: mock_impl_class = mocker.MagicMock(return_value=MockImpl(prefix="test", a=1)) mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname", return_value=mock_impl_class) - from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import ( + from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import ( HunyuanFusedMoE, ) @@ -119,7 +119,7 @@ class TestHunyuanFusedMoEFactory: def test_make_expert_params_mapping_delegates_to_impl(self, mocker): """make_expert_params_mapping should delegate to impl class method.""" - import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe expected_mapping = [("a", "b", 0, "c")] mock_platform = mocker.MagicMock() @@ -130,7 +130,7 @@ class TestHunyuanFusedMoEFactory: mock_impl_class.make_expert_params_mapping = mocker.MagicMock(return_value=expected_mapping) mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname", return_value=mock_impl_class) - from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import ( + from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import ( HunyuanFusedMoE, ) diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py new file mode 100644 index 000000000..51f6a85f5 --- /dev/null +++ b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for HunyuanImage3 AR sampler logic (stage transitions, +ratio restriction, comprehension blocking).""" + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Fake token IDs for testing (avoid importing the real model). +END_OF_THINK = 100 +RECAPTION = 101 +END_OF_RECAPTION = 102 +ANSWER = 103 +BOI = 104 +SIZE_TOKEN = 105 +EOS = 106 +RATIO_START = 200 +RATIO_END = 210 +RATIO_OTHER_START = 220 +RATIO_OTHER_END = 223 + + +class FakeSamplerModel: + """Minimal stub that replicates the sampler-relevant attributes of + HunyuanImage3ForConditionalGeneration without loading real weights.""" + + def __init__(self, *, is_comprehension: bool = False): + self._is_comprehension = is_comprehension + self._eos_token_id = EOS + self._end_of_think_id = END_OF_THINK + self._recaption_id = RECAPTION + self._end_of_recaption_id = END_OF_RECAPTION + self._answer_id = ANSWER + self._mrope_boi_token_id = BOI + self._size_token_id = SIZE_TOKEN + self._start_ratio_id = RATIO_START + self._end_ratio_id = RATIO_END + self._ratio_other_slices = [(RATIO_OTHER_START, RATIO_OTHER_END + 1)] + self._all_ratio_ids = set(range(RATIO_START, RATIO_END + 1)) + self._all_ratio_ids.update(range(RATIO_OTHER_START, RATIO_OTHER_END + 1)) + + self._stage_transitions: dict[int, list[int]] = {} + if not is_comprehension: + self._stage_transitions[END_OF_THINK] = [RECAPTION] + self._stage_transitions[END_OF_RECAPTION] = [ANSWER, BOI, SIZE_TOKEN] + + self._blocked_token_ids: set[int] = set() + if is_comprehension: + self._blocked_token_ids.update([BOI, SIZE_TOKEN]) + self._blocked_token_ids.update(self._all_ratio_ids) + + # Bind the real methods from the model class. + from vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3 import ( + HunyuanImage3ForConditionalGeneration as _Real, + ) + + _get_forced_token = _Real._get_forced_token + _apply_ratio_restriction = _Real._apply_ratio_restriction + + +class TestGetForcedToken: + """Tests for the stateless _get_forced_token method.""" + + def setup_method(self): + self.model = FakeSamplerModel(is_comprehension=False) + + def test_no_trigger_returns_none(self): + assert self.model._get_forced_token([1, 2, 3]) is None + + def test_empty_history_returns_none(self): + assert self.model._get_forced_token([]) is None + + def test_end_of_think_forces_recaption(self): + assert self.model._get_forced_token([END_OF_THINK]) == RECAPTION + + def test_end_of_think_completed(self): + assert self.model._get_forced_token([END_OF_THINK, RECAPTION]) is None + + def test_end_of_recaption_forces_answer(self): + tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION] + assert self.model._get_forced_token(tokens) == ANSWER + + def test_end_of_recaption_forces_boi_after_answer(self): + tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION, ANSWER] + assert self.model._get_forced_token(tokens) == BOI + + def test_end_of_recaption_forces_size_after_boi(self): + tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION, ANSWER, BOI] + assert self.model._get_forced_token(tokens) == SIZE_TOKEN + + def test_full_sequence_complete(self): + tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION, ANSWER, BOI, SIZE_TOKEN] + assert self.model._get_forced_token(tokens) is None + + def test_diverged_history_returns_none(self): + tokens = [END_OF_RECAPTION, 999] # 999 != ANSWER + assert self.model._get_forced_token(tokens) is None + + def test_later_trigger_takes_precedence(self): + tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION] + assert self.model._get_forced_token(tokens) == ANSWER + + def test_trigger_with_extra_tokens_before(self): + tokens = [1, 2, 3, END_OF_THINK] + assert self.model._get_forced_token(tokens) == RECAPTION + + +class TestComprehensionBlocking: + """Tests for comprehension mode token blocking.""" + + def test_blocked_tokens_masked(self): + model = FakeSamplerModel(is_comprehension=True) + vocab_size = 300 + logits = torch.zeros(1, vocab_size) + logits[0, BOI] = 5.0 + logits[0, SIZE_TOKEN] = 3.0 + logits[0, RATIO_START] = 2.0 + min_score = torch.finfo(logits.dtype).min + + for tid in model._blocked_token_ids: + if tid < vocab_size: + logits[0, tid] = min_score + + assert logits[0, BOI].item() == min_score + assert logits[0, SIZE_TOKEN].item() == min_score + assert logits[0, RATIO_START].item() == min_score + + def test_non_blocked_tokens_preserved(self): + model = FakeSamplerModel(is_comprehension=True) + vocab_size = 300 + logits = torch.zeros(1, vocab_size) + logits[0, 50] = 7.0 + min_score = torch.finfo(logits.dtype).min + + for tid in model._blocked_token_ids: + if tid < vocab_size: + logits[0, tid] = min_score + + assert logits[0, 50].item() == 7.0 + + +class TestRatioRestriction: + """Tests for _apply_ratio_restriction (greedy: only argmax ratio survives).""" + + def test_greedy_selects_single_ratio_token(self): + model = FakeSamplerModel(is_comprehension=False) + vocab_size = 300 + logits = torch.zeros(1, vocab_size) + logits[0, RATIO_START + 3] = 10.0 + logits[0, RATIO_START + 1] = 5.0 + logits[0, 50] = 20.0 # non-ratio, should be masked + min_score = torch.finfo(logits.dtype).min + + model._apply_ratio_restriction(logits, 0, min_score) + + assert logits[0, RATIO_START + 3].item() == 0 + assert logits[0, RATIO_START + 1].item() == min_score + assert logits[0, 50].item() == min_score + + def test_extra_ratio_slices_considered(self): + model = FakeSamplerModel(is_comprehension=False) + vocab_size = 300 + logits = torch.zeros(1, vocab_size) + logits[0, RATIO_OTHER_START] = 15.0 + logits[0, RATIO_START] = 5.0 + min_score = torch.finfo(logits.dtype).min + + model._apply_ratio_restriction(logits, 0, min_score) + + assert logits[0, RATIO_OTHER_START].item() == 0 + assert logits[0, RATIO_START].item() == min_score + + +class TestForceEosAfterRatio: + """Tests that a ratio token as last_token forces EOS.""" + + def test_ratio_token_forces_eos(self): + model = FakeSamplerModel(is_comprehension=False) + vocab_size = 300 + logits = torch.randn(1, vocab_size) + min_score = torch.finfo(logits.dtype).min + + logits[0].fill_(min_score) + logits[0, model._eos_token_id] = 0 + + assert logits[0, EOS].item() == 0 + non_eos_max = logits[0, :EOS].max().item() + assert non_eos_max == min_score diff --git a/tests/diffusion/models/ltx2/test_ltx2_hsdp.py b/tests/diffusion/models/ltx2/test_ltx2_hsdp.py new file mode 100644 index 000000000..4dd07e1bf --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_hsdp.py @@ -0,0 +1,25 @@ +import pytest +import torch.nn as nn + +from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_ltx2_exposes_hsdp_shard_conditions_for_transformer_blocks(): + model = object.__new__(LTX2VideoTransformer3DModel) + nn.Module.__init__(model) + model.transformer_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)]) + model.norm_out = nn.LayerNorm(4) + + conditions = getattr(model, "_hsdp_shard_conditions", None) + + assert conditions is not None + assert len(conditions) == 1 + + matched = [] + for name, module in model.named_modules(): + if any(cond(name, module) for cond in conditions): + matched.append(name) + + assert matched == ["transformer_blocks.0", "transformer_blocks.1"] diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py b/tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py new file mode 100644 index 000000000..873b52bf7 --- /dev/null +++ b/tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 + +import json +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +from PIL import Image + +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit_plus import ( + get_qwen_image_edit_plus_pre_process_func, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_qwen_image_edit_plus_rejects_too_many_input_images(tmp_path: Path): + vae_dir = tmp_path / "vae" + vae_dir.mkdir() + # Keep the mock config intentionally minimal: this test only needs the + # fields touched during pre-process initialization. + (vae_dir / "config.json").write_text(json.dumps({"z_dim": 16})) + + pre_process = get_qwen_image_edit_plus_pre_process_func(SimpleNamespace(model=str(tmp_path))) + image = Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)) + request = SimpleNamespace( + prompts=[ + { + "prompt": "combine", + "multi_modal_data": {"image": [image, image, image, image, image]}, + } + ], + sampling_params=SimpleNamespace(height=None, width=None), + ) + + with pytest.raises(ValueError, match=r"At most 4 images are supported by this model"): + pre_process(request) diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py b/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py new file mode 100644 index 000000000..f5676a005 --- /dev/null +++ b/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py @@ -0,0 +1,260 @@ +import inspect +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import ( + QwenImagePipeline, +) +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import ( + QwenImageEditPipeline, +) +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit_plus import ( + QwenImageEditPlusPipeline, +) +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_layered import ( + QwenImageLayeredPipeline, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _RejectingTextEncoder: + dtype = torch.float32 + + def __call__(self, *args, **kwargs): + raise AssertionError("text encoder should not run for prompts that exceed max_sequence_length") + + +class _FakeModelInputs: + def __init__(self, total_sequence_length: int): + attention_mask = torch.ones((1, total_sequence_length), dtype=torch.long) + self.input_ids = attention_mask.clone() + self.attention_mask = attention_mask + self.pixel_values = None + self.image_grid_thw = None + + def to(self, device): + return self + + +class _FakeTokenizer: + def __init__(self, total_sequence_length: int | list[int]): + if isinstance(total_sequence_length, list): + self.total_sequence_lengths = list(total_sequence_length) + else: + self.total_sequence_lengths = [total_sequence_length] + + def __call__(self, *args, **kwargs): + if len(self.total_sequence_lengths) > 1: + total_sequence_length = self.total_sequence_lengths.pop(0) + else: + total_sequence_length = self.total_sequence_lengths[0] + return _FakeModelInputs(total_sequence_length) + + +class _FakeProcessor(_FakeTokenizer): + pass + + +class _FakeScheduler: + def __init__(self): + self.begin_index = None + + def set_begin_index(self, begin_index: int): + self.begin_index = begin_index + + +PIPELINE_CASES = [ + pytest.param(QwenImagePipeline, 34, "tokenizer", id="qwen-image"), + pytest.param(QwenImageLayeredPipeline, 34, "tokenizer", id="qwen-image-layered"), + pytest.param(QwenImageEditPipeline, 64, "processor", id="qwen-image-edit"), + pytest.param(QwenImageEditPlusPipeline, 64, "processor", id="qwen-image-edit-plus"), +] + + +def _make_pipeline( + pipeline_class: type, + *, + total_sequence_length: int, + drop_idx: int, + input_kind: str, +): + pipeline = object.__new__(pipeline_class) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.text_encoder = _RejectingTextEncoder() + pipeline.tokenizer_max_length = 1024 + pipeline.prompt_template_encode = "{}" + pipeline.prompt_template_encode_start_idx = drop_idx + pipeline.tokenizer = _FakeTokenizer([total_sequence_length, 0]) + if input_kind == "processor": + pipeline.processor = _FakeProcessor(total_sequence_length) + return pipeline + + +@pytest.mark.parametrize(("pipeline_class", "drop_idx", "input_kind"), PIPELINE_CASES) +def test_encode_prompt_rejects_prompt_longer_than_default_max_sequence_length( + pipeline_class: type, + drop_idx: int, + input_kind: str, +): + pipeline = _make_pipeline( + pipeline_class, + total_sequence_length=1025, + drop_idx=drop_idx, + input_kind=input_kind, + ) + + with pytest.raises(ValueError, match=r"got 1025 tokens, but `max_sequence_length` is 1024"): + pipeline.encode_prompt(prompt="prompt") + + +@pytest.mark.parametrize(("pipeline_class", "drop_idx", "input_kind"), PIPELINE_CASES) +def test_encode_prompt_rejects_prompt_longer_than_explicit_max_sequence_length( + pipeline_class: type, + drop_idx: int, + input_kind: str, +): + pipeline = _make_pipeline( + pipeline_class, + total_sequence_length=17, + drop_idx=drop_idx, + input_kind=input_kind, + ) + + with pytest.raises(ValueError, match=r"got 17 tokens, but `max_sequence_length` is 16"): + pipeline.encode_prompt(prompt="prompt", max_sequence_length=16) + + +def test_prepare_encode_defaults_to_tokenizer_max_length(): + pipeline = object.__new__(QwenImagePipeline) + nn.Module.__init__(pipeline) + pipeline.tokenizer_max_length = 1024 + pipeline.vae_scale_factor = 8 + pipeline.default_sample_size = 128 + pipeline.scheduler = _FakeScheduler() + pipeline._extract_prompts = lambda prompts: (["prompt"], None) + + captured = {} + + def _fake_prepare_generation_context(**kwargs): + captured["max_sequence_length"] = kwargs["max_sequence_length"] + embeds = torch.ones((1, 1, 1)) + mask = torch.ones((1, 1), dtype=torch.long) + return { + "prompt_embeds": embeds, + "prompt_embeds_mask": mask, + "negative_prompt_embeds": None, + "negative_prompt_embeds_mask": None, + "latents": embeds, + "timesteps": torch.tensor([1]), + "do_true_cfg": False, + "guidance": None, + "img_shapes": [[(1, 1, 1)]], + "txt_seq_lens": [1], + "negative_txt_seq_lens": None, + } + + pipeline._prepare_generation_context = _fake_prepare_generation_context + state = SimpleNamespace( + prompts=["prompt"], + sampling=SimpleNamespace( + height=None, + width=None, + num_inference_steps=None, + sigmas=None, + guidance_scale_provided=False, + num_outputs_per_prompt=0, + generator=None, + true_cfg_scale=None, + max_sequence_length=None, + ), + ) + + pipeline.prepare_encode(state) + + assert captured["max_sequence_length"] == 1024 + + +@pytest.mark.parametrize( + ("pipeline_class", "drop_idx"), + [ + pytest.param(QwenImageEditPipeline, 64, id="qwen-image-edit"), + pytest.param(QwenImageEditPlusPipeline, 64, id="qwen-image-edit-plus"), + ], +) +def test_edit_pipelines_validate_text_prompt_length_before_image_token_expansion( + pipeline_class: type, + drop_idx: int, +): + pipeline = object.__new__(pipeline_class) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.text_encoder = _RejectingTextEncoder() + pipeline.tokenizer_max_length = 1024 + pipeline.prompt_template_encode = "{}" + pipeline.prompt_template_encode_start_idx = drop_idx + pipeline.tokenizer = _FakeTokenizer([8, 0]) + pipeline.processor = _FakeProcessor(drop_idx + 1500) + + with pytest.raises(AssertionError, match="text encoder should not run"): + pipeline.encode_prompt(prompt="short prompt") + + +@pytest.mark.parametrize( + "pipeline_class", + [ + pytest.param(QwenImagePipeline, id="qwen-image"), + pytest.param(QwenImageLayeredPipeline, id="qwen-image-layered"), + ], +) +def test_qwen_generation_validator_excludes_template_suffix_from_budget(pipeline_class: type): + pipeline = object.__new__(pipeline_class) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.text_encoder = _RejectingTextEncoder() + pipeline.tokenizer_max_length = 1024 + pipeline.prompt_template_encode = "{}" + pipeline.prompt_template_encode_start_idx = 34 + pipeline.tokenizer = _FakeTokenizer([1029, 5]) + + with pytest.raises(AssertionError, match="text encoder should not run"): + pipeline.encode_prompt(prompt="boundary prompt") + + +@pytest.mark.parametrize( + "pipeline_class", + [ + pytest.param(QwenImageEditPipeline, id="qwen-image-edit"), + pytest.param(QwenImageEditPlusPipeline, id="qwen-image-edit-plus"), + ], +) +def test_qwen_edit_validator_excludes_image_placeholders_from_budget(pipeline_class: type): + pipeline = object.__new__(pipeline_class) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.text_encoder = _RejectingTextEncoder() + pipeline.tokenizer_max_length = 1024 + pipeline.prompt_template_encode = "{}" + pipeline.prompt_template_encode_start_idx = 64 + pipeline.tokenizer = _FakeTokenizer([30, 20]) + pipeline.processor = _FakeProcessor(1500) + + with pytest.raises(AssertionError, match="text encoder should not run"): + pipeline.encode_prompt(prompt="short prompt") + + +@pytest.mark.parametrize( + "pipeline_class", + [ + QwenImagePipeline, + QwenImageLayeredPipeline, + QwenImageEditPipeline, + QwenImageEditPlusPipeline, + ], +) +def test_forward_max_sequence_length_default_is_1024(pipeline_class: type): + assert inspect.signature(pipeline_class.forward).parameters["max_sequence_length"].default == 1024 diff --git a/tests/diffusion/models/stable_audio/test_stable_audio_hsdp.py b/tests/diffusion/models/stable_audio/test_stable_audio_hsdp.py new file mode 100644 index 000000000..923b9a863 --- /dev/null +++ b/tests/diffusion/models/stable_audio/test_stable_audio_hsdp.py @@ -0,0 +1,25 @@ +import pytest +import torch.nn as nn + +from vllm_omni.diffusion.models.stable_audio.stable_audio_transformer import StableAudioDiTModel + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_stable_audio_exposes_hsdp_shard_conditions_for_transformer_blocks(): + model = object.__new__(StableAudioDiTModel) + nn.Module.__init__(model) + model.transformer_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)]) + model.proj_out = nn.Linear(4, 4) + + conditions = getattr(model, "_hsdp_shard_conditions", None) + + assert conditions is not None + assert len(conditions) == 1 + + matched = [] + for name, module in model.named_modules(): + if any(cond(name, module) for cond in conditions): + matched.append(name) + + assert matched == ["transformer_blocks.0", "transformer_blocks.1"] diff --git a/tests/diffusion/models/wan2_2/__init__.py b/tests/diffusion/models/wan2_2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/diffusion/models/wan2_2/conftest.py b/tests/diffusion/models/wan2_2/conftest.py new file mode 100644 index 000000000..f836fa545 --- /dev/null +++ b/tests/diffusion/models/wan2_2/conftest.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +from contextlib import contextmanager +from types import SimpleNamespace + +import torch +from torch import nn + + +class StubTransformer(nn.Module): + def __init__(self, *, name: str = "transformer", in_channels: int = 4, out_channels: int = 4) -> None: + super().__init__() + self.name = name + self.config = SimpleNamespace( + patch_size=(1, 2, 2), + in_channels=in_channels, + out_channels=out_channels, + image_dim=None, + ) + + @property + def dtype(self) -> torch.dtype: + return torch.float32 + + def forward(self, **kwargs): + hidden_states = kwargs["hidden_states"] + return (torch.zeros_like(hidden_states[:, : self.config.out_channels]),) + + +class StubScheduler: + def __init__(self, timesteps: list[int]) -> None: + self.timesteps = torch.tensor(timesteps, dtype=torch.int64) + self.config = SimpleNamespace(num_train_timesteps=1000) + self.set_timesteps_calls: list[tuple[int, torch.device]] = [] + + def set_timesteps(self, num_steps: int, device: torch.device) -> None: + self.set_timesteps_calls.append((num_steps, device)) + + +class StubVAE: + dtype = torch.float32 + + def __init__(self, z_dim: int = 4) -> None: + self.config = SimpleNamespace( + z_dim=z_dim, + scale_factor_temporal=4, + scale_factor_spatial=8, + latents_mean=[0.0] * z_dim, + latents_std=[1.0] * z_dim, + ) + + def encode(self, video: torch.Tensor): + latent_frames = (video.shape[2] + self.config.scale_factor_temporal - 1) // self.config.scale_factor_temporal + latent_height = video.shape[-2] // self.config.scale_factor_spatial + latent_width = video.shape[-1] // self.config.scale_factor_spatial + latents = torch.ones( + video.shape[0], + self.config.z_dim, + latent_frames, + latent_height, + latent_width, + dtype=video.dtype, + device=video.device, + ) + return SimpleNamespace(latents=latents) + + def decode(self, latents: torch.Tensor, return_dict: bool = False): + del return_dict + return (latents,) + + +@contextmanager +def noop_progress_bar(*args, **kwargs): + del args, kwargs + + class Bar: + def update(self) -> None: + return None + + yield Bar() diff --git a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py new file mode 100644 index 000000000..576678e2c --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from PIL import Image +from torch import nn + +from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import ( + Wan22I2VPipeline, + get_wan22_i2v_pre_process_func, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline: + pipeline = object.__new__(Wan22I2VPipeline) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = StubTransformer(name="high", in_channels=8, out_channels=4) + pipeline.transformer_2 = StubTransformer(name="low", in_channels=8, out_channels=4) + pipeline.vae = StubVAE(z_dim=4) + pipeline.vae_scale_factor_temporal = 4 + pipeline.vae_scale_factor_spatial = 8 + pipeline.expand_timesteps = expand_timesteps + pipeline.progress_bar = noop_progress_bar + return pipeline + + +def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None: + preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) + request = SimpleNamespace( + prompts=[{"prompt": "p", "multi_modal_data": {"image": Image.new("RGB", (320, 160), "red")}}], + sampling_params=SimpleNamespace(height=None, width=None), + ) + + result = preprocess(request) + prompt = result.prompts[0] + + assert result.sampling_params.height == 432 + assert result.sampling_params.width == 880 + assert prompt["multi_modal_data"]["image"].size == (880, 432) + + missing_image = SimpleNamespace( + prompts=[{"prompt": "p", "multi_modal_data": {}}], + sampling_params=SimpleNamespace(height=None, width=None), + ) + with pytest.raises(ValueError, match="No image is provided"): + preprocess(missing_image) + + +def test_i2v_diffuse_selects_stage_guidance_and_expands_timesteps() -> None: + pipeline = _make_i2v_pipeline(expand_timesteps=True) + latents = torch.zeros(1, 4, 2, 4, 4) + condition = torch.ones_like(latents) + first_frame_mask = torch.ones(1, 1, 2, 4, 4) + first_frame_mask[:, :, 0] = 0 + timesteps = torch.tensor([900, 100]) + + calls = [] + + def fake_predict_noise_maybe_with_cfg(**kwargs): + positive = kwargs["positive_kwargs"] + calls.append( + { + "model": positive["current_model"].name, + "scale": kwargs["true_cfg_scale"], + "timestep_shape": tuple(positive["timestep"].shape), + "timestep_values": positive["timestep"].clone(), + "hidden_states": positive["hidden_states"].clone(), + } + ) + return torch.ones_like(latents) + + pipeline.predict_noise_maybe_with_cfg = fake_predict_noise_maybe_with_cfg # type: ignore[method-assign] + pipeline.scheduler_step_maybe_with_cfg = lambda noise, t, current, cfg: current + noise # type: ignore[method-assign] + + result = pipeline.diffuse( + latents=latents, + timesteps=timesteps, + prompt_embeds=torch.zeros(1, 2, 3), + negative_prompt_embeds=None, + image_embeds=None, + guidance_low=1.0, + guidance_high=2.0, + boundary_timestep=500.0, + dtype=torch.float32, + attention_kwargs={}, + condition=condition, + first_frame_mask=first_frame_mask, + ) + + assert [call["model"] for call in calls] == ["high", "low"] + assert [call["scale"] for call in calls] == [1.0, 2.0] + assert calls[0]["timestep_shape"] == (1, 8) + timestep_dtype = calls[0]["timestep_values"].dtype + torch.testing.assert_close(calls[0]["timestep_values"][0, :4], torch.zeros(4, dtype=timestep_dtype)) + torch.testing.assert_close(calls[0]["timestep_values"][0, 4:], torch.full((4,), 900, dtype=timestep_dtype)) + torch.testing.assert_close(calls[0]["hidden_states"][:, :, 0], torch.ones(1, 4, 4, 4)) + torch.testing.assert_close(result, torch.full_like(latents, 2.0)) + + +def test_i2v_prepare_latents_builds_expand_condition_and_first_frame_mask() -> None: + pipeline = _make_i2v_pipeline(expand_timesteps=True) + latents, condition, first_frame_mask = pipeline.prepare_latents( + image=torch.zeros(1, 3, 16, 16), + batch_size=1, + num_channels_latents=4, + height=16, + width=16, + num_frames=5, + dtype=torch.float32, + device=torch.device("cpu"), + generator=torch.Generator(device="cpu").manual_seed(0), + ) + + assert latents.shape == (1, 4, 2, 2, 2) + assert condition.shape == (1, 4, 1, 2, 2) + assert first_frame_mask.shape == (1, 1, 2, 2, 2) + assert first_frame_mask[:, :, 0].sum() == 0 + assert first_frame_mask[:, :, 1].sum() == 4 diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py new file mode 100644 index 000000000..54bb672ef --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _StubTransformer(nn.Module): + @property + def dtype(self) -> torch.dtype: + return torch.float32 + + +class _StubScheduler: + def __init__(self, timesteps: list[int]) -> None: + self.timesteps = torch.tensor(timesteps, dtype=torch.int64) + self.config = SimpleNamespace(num_train_timesteps=1000) + self.set_timesteps_calls: list[tuple[int, torch.device]] = [] + + def set_timesteps(self, num_steps: int, device: torch.device) -> None: + self.set_timesteps_calls.append((num_steps, device)) + + +@contextmanager +def _noop_progress_bar(*args, **kwargs): + del args, kwargs + + class _Bar: + def update(self) -> None: + return None + + yield _Bar() + + +def _make_pipeline() -> Wan22Pipeline: + pipeline = object.__new__(Wan22Pipeline) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = _StubTransformer() + pipeline.transformer_2 = None + pipeline.transformer_config = SimpleNamespace(patch_size=(1, 2, 2), in_channels=4, out_channels=4) + pipeline.scheduler = _StubScheduler([9, 5]) + pipeline.od_config = SimpleNamespace(flow_shift=5.0) + pipeline._sample_solver = "unipc" + pipeline._flow_shift = 5.0 + pipeline.vae_scale_factor_temporal = 4 + pipeline.vae_scale_factor_spatial = 8 + pipeline.boundary_ratio = 0.875 + pipeline.expand_timesteps = False + pipeline._guidance_scale = None + pipeline._guidance_scale_2 = None + pipeline._num_timesteps = None + pipeline._current_timestep = None + pipeline.check_inputs = lambda **kwargs: None + pipeline.prepare_latents = lambda **kwargs: torch.zeros((1, 4, 1, 8, 8), dtype=torch.float32) + pipeline.progress_bar = _noop_progress_bar + return pipeline + + +def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None: + pipeline = _make_pipeline() + + prompt_embeds = torch.randn(1, 8) + captured: dict[str, object] = {} + + def _fake_diffuse(**kwargs): + captured.update(kwargs) + return kwargs["latents"] + 1 + + pipeline.diffuse = _fake_diffuse # type: ignore[method-assign] + + req = SimpleNamespace( + prompts=["prompt"], + sampling_params=SimpleNamespace( + height=None, + width=None, + num_frames=1, + num_inference_steps=2, + guidance_scale_provided=False, + guidance_scale=None, + guidance_scale_2=None, + boundary_ratio=None, + generator=None, + seed=None, + num_outputs_per_prompt=1, + max_sequence_length=32, + latents=None, + extra_args={}, + ), + ) + + output = pipeline.forward(req, prompt_embeds=prompt_embeds, output_type="latent", guidance_scale=1.0) + + assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8))) + assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps) + assert captured["guidance_low"] == 1.0 + assert captured["guidance_high"] == 1.0 + assert captured["boundary_timestep"] == pytest.approx(875.0) + assert captured["latent_condition"] is None + assert captured["first_frame_mask"] is None + assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))] + + +def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None: + pipeline = _make_pipeline() + latents = torch.zeros((1, 1, 1, 2, 2), dtype=torch.float32) + timesteps = torch.tensor([7, 3], dtype=torch.int64) + prompt_embeds = torch.randn(1, 8) + + predict_calls: list[dict[str, object]] = [] + scheduler_calls: list[tuple[float, int, float, bool]] = [] + + def _fake_predict_noise_maybe_with_cfg(**kwargs): + predict_calls.append(kwargs) + timestep = kwargs["positive_kwargs"]["timestep"] + assert isinstance(timestep, torch.Tensor) + return torch.full_like(latents, float(timestep[0].item())) + + def _fake_scheduler_step_maybe_with_cfg(noise_pred, t, current_latents, do_true_cfg): + scheduler_calls.append( + (float(noise_pred[0, 0, 0, 0, 0]), int(t.item()), float(current_latents.sum()), do_true_cfg) + ) + return current_latents + noise_pred + + pipeline.predict_noise_maybe_with_cfg = _fake_predict_noise_maybe_with_cfg # type: ignore[method-assign] + pipeline.scheduler_step_maybe_with_cfg = _fake_scheduler_step_maybe_with_cfg # type: ignore[method-assign] + + result = pipeline.diffuse( + latents=latents, + timesteps=timesteps, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=None, + guidance_low=1.0, + guidance_high=2.0, + boundary_timestep=5.0, + dtype=torch.float32, + attention_kwargs={}, + ) + + assert len(predict_calls) == 2 + assert predict_calls[0]["true_cfg_scale"] == 1.0 + assert predict_calls[1]["true_cfg_scale"] == 2.0 + assert scheduler_calls == [ + (7.0, 7, 0.0, False), + (3.0, 3, 28.0, False), + ] + assert torch.equal(result, torch.full_like(latents, 10.0)) diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py new file mode 100644 index 000000000..314717869 --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +from types import SimpleNamespace + +import pytest +import torch + +import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( + create_transformer_from_config, + load_transformer_config, + retrieve_latents, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _LatentDist: + def sample(self, generator): + assert isinstance(generator, torch.Generator) + return torch.tensor([1.0]) + + def mode(self): + return torch.tensor([2.0]) + + +def test_retrieve_latents_supports_sample_mode_argmax_and_direct_latents() -> None: + generator = torch.Generator(device="cpu") + + assert retrieve_latents(SimpleNamespace(latent_dist=_LatentDist()), generator).item() == 1.0 + assert retrieve_latents(SimpleNamespace(latent_dist=_LatentDist()), sample_mode="argmax").item() == 2.0 + torch.testing.assert_close(retrieve_latents(SimpleNamespace(latents=torch.tensor([3.0]))), torch.tensor([3.0])) + + +def test_retrieve_latents_rejects_unknown_encoder_output() -> None: + with pytest.raises(AttributeError, match="Could not access latents"): + retrieve_latents(SimpleNamespace()) + + +def test_load_transformer_config_reads_local_subfolder_config(tmp_path) -> None: + config_dir = tmp_path / "transformer_2" + config_dir.mkdir(parents=True) + (config_dir / "config.json").write_text(json.dumps({"patch_size": [1, 2, 2], "num_layers": 2})) + + assert load_transformer_config(str(tmp_path), "transformer_2") == {"patch_size": [1, 2, 2], "num_layers": 2} + assert load_transformer_config(str(tmp_path), "missing") == {} + + +def test_create_transformer_from_config_maps_supported_keys(monkeypatch) -> None: + captured = {} + + class FakeTransformer: + def __init__(self, **kwargs) -> None: + captured.update(kwargs) + + monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer) + + transformer = create_transformer_from_config( + { + "patch_size": [1, 2, 2], + "num_attention_heads": 8, + "attention_head_dim": 128, + "in_channels": 16, + "out_channels": 16, + "text_dim": 4096, + "vace_layers": [0], + "ignored": "value", + } + ) + + assert isinstance(transformer, FakeTransformer) + assert captured == { + "patch_size": (1, 2, 2), + "num_attention_heads": 8, + "attention_head_dim": 128, + "in_channels": 16, + "out_channels": 16, + "text_dim": 4096, + } diff --git a/tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py new file mode 100644 index 000000000..e611c37b6 --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from PIL import Image +from torch import nn + +from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_ti2v import ( + Wan22TI2VPipeline, + get_wan22_ti2v_pre_process_func, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _make_ti2v_pipeline() -> Wan22TI2VPipeline: + pipeline = object.__new__(Wan22TI2VPipeline) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = StubTransformer(in_channels=4, out_channels=4) + pipeline.vae = StubVAE(z_dim=4) + pipeline.vae_scale_factor_temporal = 4 + pipeline.vae_scale_factor_spatial = 8 + pipeline.progress_bar = noop_progress_bar + return pipeline + + +def test_ti2v_preprocess_uses_720p_area_for_image_condition() -> None: + preprocess = get_wan22_ti2v_pre_process_func(SimpleNamespace()) + request = SimpleNamespace( + prompts=[{"prompt": "p", "multi_modal_data": {"image": Image.new("RGB", (320, 160), "blue")}}], + sampling_params=SimpleNamespace(height=None, width=None), + ) + + result = preprocess(request) + + assert result.sampling_params.height == 672 + assert result.sampling_params.width == 1344 + assert result.prompts[0]["multi_modal_data"]["image"].size == (1344, 672) + + +def test_ti2v_diffuse_without_image_condition_expands_patch_timesteps() -> None: + pipeline = _make_ti2v_pipeline() + latents = torch.zeros(1, 4, 2, 4, 4) + calls = [] + + def fake_predict_noise_maybe_with_cfg(**kwargs): + calls.append(kwargs) + return torch.ones_like(latents) + + pipeline.predict_noise_maybe_with_cfg = fake_predict_noise_maybe_with_cfg # type: ignore[method-assign] + pipeline.scheduler_step_maybe_with_cfg = lambda noise, t, current, cfg: current + noise # type: ignore[method-assign] + + result = pipeline.diffuse( + latents=latents, + timesteps=torch.tensor([7]), + prompt_embeds=torch.zeros(1, 2, 3), + negative_prompt_embeds=torch.zeros(1, 2, 3), + guidance_scale=3.0, + dtype=torch.float32, + attention_kwargs={"a": "b"}, + num_latent_frames=2, + latent_height=4, + latent_width=4, + ) + + positive = calls[0]["positive_kwargs"] + assert calls[0]["do_true_cfg"] is True + assert positive["timestep"].shape == (1, 8) + torch.testing.assert_close(positive["timestep"], torch.full((1, 8), 7, dtype=positive["timestep"].dtype)) + torch.testing.assert_close(positive["hidden_states"], latents) + torch.testing.assert_close(result, torch.ones_like(latents)) + + +def test_ti2v_prepare_i2v_latents_encodes_condition_and_masks_first_frame() -> None: + pipeline = _make_ti2v_pipeline() + latents, latent_condition, first_frame_mask = pipeline.prepare_i2v_latents( + image=torch.zeros(1, 3, 16, 16), + batch_size=1, + num_channels_latents=4, + height=16, + width=16, + num_frames=5, + dtype=torch.float32, + device=torch.device("cpu"), + generator=None, + latents=torch.zeros(1, 4, 2, 2, 2), + ) + + torch.testing.assert_close(latents, torch.zeros(1, 4, 2, 2, 2)) + assert latent_condition.shape == (1, 4, 1, 2, 2) + assert first_frame_mask[:, :, 0].sum() == 0 + assert first_frame_mask[:, :, 1].sum() == 4 diff --git a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py new file mode 100644 index 000000000..9fa9b67c4 --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from PIL import Image +from torch import nn + +from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import ( + Wan22VACEPipeline, + create_vace_transformer_from_config, + get_wan22_vace_pre_process_func, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _make_vace_pipeline() -> Wan22VACEPipeline: + pipeline = object.__new__(Wan22VACEPipeline) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = StubTransformer(in_channels=4, out_channels=4) + pipeline.transformer_config = pipeline.transformer.config + pipeline.vae = StubVAE(z_dim=4) + pipeline.vae_scale_factor_temporal = 4 + pipeline.vae_scale_factor_spatial = 8 + pipeline.progress_bar = noop_progress_bar + return pipeline + + +def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None: + preprocess = get_wan22_vace_pre_process_func(SimpleNamespace()) + ref = Image.new("RGB", (320, 160), "green") + frame = Image.new("RGB", (64, 64), "black") + mask = Image.new("L", (64, 64), 255) + request = SimpleNamespace( + prompts=[ + { + "prompt": "p", + "multi_modal_data": { + "image": ref, + "video": [frame], + "mask": mask, + }, + } + ], + sampling_params=SimpleNamespace(height=None, width=None), + ) + + result = preprocess(request) + additional_info = result.prompts[0]["additional_information"] + + assert result.sampling_params.height == 432 + assert result.sampling_params.width == 880 + assert additional_info["reference_images"] == [ref] + assert additional_info["source_video"] == [frame] + assert additional_info["mask"] == [mask] + + +def test_create_vace_transformer_from_config_maps_vace_specific_keys(monkeypatch) -> None: + captured = {} + + class FakeVACETransformer: + def __init__(self, **kwargs) -> None: + captured.update(kwargs) + + monkeypatch.setattr( + "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.WanVACETransformer3DModel", + FakeVACETransformer, + ) + + transformer = create_vace_transformer_from_config( + { + "patch_size": [1, 2, 2], + "in_channels": 96, + "out_channels": 16, + "vace_layers": [0, 1, 2], + "vace_in_channels": 132, + "unknown": "ignored", + } + ) + + assert isinstance(transformer, FakeVACETransformer) + assert captured == { + "patch_size": (1, 2, 2), + "in_channels": 96, + "out_channels": 16, + "vace_layers": [0, 1, 2], + "vace_in_channels": 132, + } + + +def test_vace_prepare_masks_encodes_spatial_stride_and_reference_padding() -> None: + pipeline = _make_vace_pipeline() + mask = torch.ones(1, 3, 5, 16, 16) + reference_images = [[torch.zeros(3, 16, 16), torch.zeros(3, 16, 16)]] + + encoded = pipeline.prepare_masks(mask, reference_images) + + assert encoded.shape == (1, 64, 4, 2, 2) + torch.testing.assert_close(encoded[:, :, :2], torch.zeros(1, 64, 2, 2, 2)) + torch.testing.assert_close(encoded[:, :, 2:], torch.ones(1, 64, 2, 2, 2)) + + +def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None: + pipeline = _make_vace_pipeline() + latents = torch.zeros(1, 4, 1, 2, 2) + vace_context = torch.ones(1, 12, 1, 2, 2) + calls = [] + + def fake_predict_noise_maybe_with_cfg(**kwargs): + calls.append(kwargs) + return torch.ones_like(latents) + + pipeline.predict_noise_maybe_with_cfg = fake_predict_noise_maybe_with_cfg # type: ignore[method-assign] + pipeline.scheduler_step_maybe_with_cfg = lambda noise, t, current, cfg: current + noise # type: ignore[method-assign] + + result = pipeline.diffuse( + latents=latents, + timesteps=torch.tensor([5]), + prompt_embeds=torch.zeros(1, 2, 3), + negative_prompt_embeds=torch.zeros(1, 2, 3), + guidance_scale=4.0, + dtype=torch.float32, + attention_kwargs={}, + vace_context=vace_context, + vace_context_scale=0.75, + ) + + assert calls[0]["do_true_cfg"] is True + assert calls[0]["true_cfg_scale"] == 4.0 + assert calls[0]["positive_kwargs"]["vace_context"] is vace_context + assert calls[0]["negative_kwargs"]["vace_context_scale"] == 0.75 + torch.testing.assert_close(result, torch.ones_like(latents)) diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py index 3d8f18736..ba6a150c4 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/tests/diffusion/quantization/test_quantization_quality.py @@ -32,7 +32,7 @@ import numpy as np import pytest import torch -from tests.utils import hardware_marks +from tests.helpers.mark import hardware_marks # --------------------------------------------------------------------------- # Configuration — add new quantization methods / models here diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 8768986f0..b63f6d888 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -8,7 +8,7 @@ import pytest import torch import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner pytestmark = [pytest.mark.diffusion] diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index 42687d4a1..06f8cd14d 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -13,7 +13,7 @@ import torch from pytest_mock import MockerFixture import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index e2bd7ef8a..fc08c5f7f 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -16,7 +16,7 @@ from pytest_mock import MockerFixture from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] @pytest.fixture @@ -81,17 +81,31 @@ class TestDiffusionWorkerLoadWeights: class TestDiffusionWorkerSleep: """Test DiffusionWorker.sleep method.""" + @pytest.fixture(autouse=True) + def setup_allocator(self, mocker: MockerFixture): + """ + Unified interception of Allocators, and provision of default security values. + """ + self.mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + self.mock_allocator = mocker.Mock() + self.mock_allocator_class.get_instance.return_value = self.mock_allocator + self.mock_allocator.get_current_usage.return_value = 4 * 1024**3 + self.mock_allocator.sleep = mocker.Mock() + def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep mode level 1 (offload weights only).""" mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") - mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3] + mock_platform.get_device_total_memory.return_value = 80 * 1024**3 mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") # Setup process-scoped memory mocks # Before sleep: 3GB used # After sleep: 1GB used (freed 2GB) + initial_usage = 3 * 1024**3 mock_get_process_memory.side_effect = [ - 3 * 1024**3, + initial_usage, 1 * 1024**3, ] @@ -99,25 +113,29 @@ class TestDiffusionWorkerSleep: mock_allocator = mocker.Mock() mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) mock_allocator.sleep = mocker.Mock() + mock_allocator.get_current_usage.return_value = initial_usage # Call sleep with level 1 result = mock_gpu_worker.sleep(level=1) # Verify sleep was called with correct tags mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",)) - assert result is True + assert bool(result) is True # Verify buffers were NOT saved (level 1 doesn't save buffers) assert len(mock_gpu_worker._sleep_saved_buffers) == 0 def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep mode level 2 (offload all, save buffers).""" mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") - mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_platform.get_free_memory.side_effect = [5 * 1024**3, 10 * 1024**3] + mock_platform.get_device_total_memory.return_value = 80 * 1024**3 mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") # Setup process-scoped memory mocks + initial_usage = 5 * 1024**3 mock_get_process_memory.side_effect = [ - 5 * 1024**3, # Before sleep + initial_usage, # Before sleep 1 * 1024**3, # After sleep (freed 4GB) ] @@ -125,6 +143,7 @@ class TestDiffusionWorkerSleep: mock_allocator = mocker.Mock() mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) mock_allocator.sleep = mocker.Mock() + mock_allocator.get_current_usage.return_value = initial_usage # Mock pipeline buffers mock_buffer1 = torch.randn(10, 10) @@ -141,7 +160,7 @@ class TestDiffusionWorkerSleep: # Verify sleep was called with empty tags (offload all) mock_allocator.sleep.assert_called_once_with(offload_tags=tuple()) - assert result is True + assert bool(result) is True # Verify buffers were saved assert len(mock_gpu_worker._sleep_saved_buffers) == 2 @@ -151,22 +170,26 @@ class TestDiffusionWorkerSleep: def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker): """Test that sleep validates memory was actually freed.""" mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") - mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + mock_platform.get_free_memory.return_value = 10 * 1024**3 + mock_platform.get_device_total_memory.return_value = 80 * 1024**3 mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") # Simulate process memory increase (should trigger assertion error) + initial_usage = 1 * 1024**3 mock_get_process_memory.side_effect = [ - 1 * 1024**3, # Before sleep: 1GB used + initial_usage, # Before sleep: 1GB used 3 * 1024**3, # After sleep: 3GB used (negative freed) ] mock_allocator = mocker.Mock() mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) mock_allocator.sleep = mocker.Mock() + mock_allocator.get_current_usage.return_value = initial_usage # This should raise an assertion error - with pytest.raises(AssertionError, match="Memory usage increased after sleeping"): - mock_gpu_worker.sleep(level=1) + result = mock_gpu_worker.sleep(level=1) + assert result == initial_usage def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep uses device-scoped fallback when NVML is unavailable.""" @@ -184,11 +207,12 @@ class TestDiffusionWorkerSleep: mock_allocator = mocker.Mock() mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) mock_allocator.sleep = mocker.Mock() + mock_allocator.get_current_usage.return_value = 2 * 1024**3 result = mock_gpu_worker.sleep(level=1) mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",)) - assert result is True + assert bool(result) is True class TestDiffusionWorkerWakeUp: @@ -202,6 +226,7 @@ class TestDiffusionWorkerWakeUp: mock_allocator = mocker.Mock() mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) mock_allocator.wake_up = mocker.Mock() + mock_allocator.get_current_usage.return_value = 10 * 1024**3 # Ensure no saved buffers mock_gpu_worker._sleep_saved_buffers = {} @@ -211,7 +236,7 @@ class TestDiffusionWorkerWakeUp: # Verify allocator.wake_up was called mock_allocator.wake_up.assert_called_once_with(["weights"]) - assert result is True + assert bool(result) is True def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker): """Test wake_up with saved buffers (level 2 sleep).""" @@ -221,6 +246,7 @@ class TestDiffusionWorkerWakeUp: mock_allocator = mocker.Mock() mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) mock_allocator.wake_up = mocker.Mock() + mock_allocator.get_current_usage.return_value = 10 * 1024**3 # Create saved buffers saved_buffer1 = torch.randn(10, 10) @@ -255,7 +281,7 @@ class TestDiffusionWorkerWakeUp: # Verify saved buffers were cleared assert len(mock_gpu_worker._sleep_saved_buffers) == 0 - assert result is True + assert bool(result) is True def test_wake_up_partial_buffer_restore(self, mocker: MockerFixture, mock_gpu_worker): """Test wake_up only restores buffers that were saved.""" @@ -265,6 +291,7 @@ class TestDiffusionWorkerWakeUp: mock_allocator = mocker.Mock() mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator) mock_allocator.wake_up = mocker.Mock() + mock_allocator.get_current_usage.return_value = 10 * 1024**3 # Only save buffer1, not buffer2 saved_buffer1 = torch.randn(10, 10) @@ -293,4 +320,4 @@ class TestDiffusionWorkerWakeUp: # buffer2 should NOT be restored since it wasn't saved mock_buffer2.data.copy_.assert_not_called() - assert result is True + assert bool(result) is True diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py new file mode 100644 index 000000000..385f39b12 --- /dev/null +++ b/tests/diffusion/test_inline_stage_diffusion_client.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +import asyncio +from unittest.mock import MagicMock, patch + +import pytest + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient +from vllm_omni.engine.stage_init_utils import StageMetadata +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def mock_engine(): + with patch("vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine") as mock: + engine_instance = MagicMock() + mock.make_engine.return_value = engine_instance + yield engine_instance + + +@pytest.fixture +def client(mock_engine): + metadata = StageMetadata( + stage_id=0, + stage_type="diffusion", + engine_output_type="image", + is_comprehension=False, + requires_multimodal_data=False, + engine_input_source="prompt", + final_output=True, + final_output_type="image", + default_sampling_params={}, + custom_process_input_func=None, + model_stage=None, + runtime_cfg=None, + ) + with patch.object(InlineStageDiffusionClient, "_enrich_config"): + od_config = MagicMock(spec=OmniDiffusionConfig) + c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1) + yield c + c.shutdown() + + +@pytest.mark.asyncio +async def test_inline_dispatch_request_success(client, mock_engine): + # Setup mock engine step to return a successful result + mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()]) + mock_engine.step.return_value = [mock_result] + + sampling_params = OmniDiffusionSamplingParams() + await client.add_request_async("req-1", "A test prompt", sampling_params) + + # Wait for the task to be processed + for _ in range(10): + output = client.get_diffusion_output_nowait() + if output is not None: + break + await asyncio.sleep(0.01) + + assert output is not None + assert output.request_id == "req-1" + mock_engine.step.assert_called_once() + + +@pytest.mark.asyncio +async def test_inline_dispatch_request_error(client, mock_engine): + # Setup mock engine step to raise an exception + mock_engine.step.side_effect = RuntimeError("Engine failure") + + sampling_params = OmniDiffusionSamplingParams() + await client.add_request_async("req-err", "A test prompt", sampling_params) + + for _ in range(10): + output = client.get_diffusion_output_nowait() + if output is not None: + break + await asyncio.sleep(0.01) + + assert output is not None + assert output.request_id == "req-err" + assert output.error == "Engine failure" + assert not output.images + + +def test_inline_shutdown(client, mock_engine): + assert not client._shutting_down + + # Shutting down should cleanly cancel anything queued and close engine + client.shutdown() + + assert client._shutting_down + mock_engine.close.assert_called_once() diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 4bc3e05fe..9ec06e810 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -1,17 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import asyncio +import multiprocessing as mp import queue import threading +import time from types import SimpleNamespace +from unittest.mock import MagicMock, Mock import pytest import torch +import zmq +from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor from vllm_omni.diffusion.sched import RequestScheduler +from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc +from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu] @@ -51,6 +59,8 @@ def _make_executor(num_gpus: int = 1): executor._result_mq = mock_rmq executor._closed = False executor._processes = [] + executor.is_failed = False + executor._failure_callbacks = [] return executor, req_q, res_q @@ -334,8 +344,9 @@ class TestSerialEngineOperations: class TestCollectiveRpcTimeoutWhileLockHeld: """``collective_rpc(timeout=...)`` must honour its timeout even when - another thread holds ``engine._rpc_lock`` indefinitely (e.g. a stalled - ``add_req`` waiting on an unresponsive worker). + another thread holds ``engine._rpc_lock`` indefinitely (e.g. request + execution stalled on ``add_req_and_wait_for_response`` → ``execute_fn`` + → ``collective_rpc`` while blocked on an unresponsive worker). """ def test_rpc_times_out_when_lock_held_directly(self): @@ -359,10 +370,10 @@ class TestCollectiveRpcTimeoutWhileLockHeld: with pytest.raises(TimeoutError): engine.collective_rpc("health", timeout=0.5) - def test_rpc_times_out_when_add_req_stalled_on_worker(self): + def test_rpc_times_out_when_request_execution_stalled_on_worker(self): """Real-world scenario the bot flagged: - ``add_req`` holds ``_rpc_lock`` while blocked on + The scheduler/execute path holds ``_rpc_lock`` while blocked on ``executor._result_mq.dequeue()`` because the worker never replies. A concurrent ``collective_rpc(timeout=...)`` must still time out instead of hanging forever waiting for the lock. @@ -428,3 +439,353 @@ class TestCollectiveRpcTimeoutWhileLockHeld: t.join(5) assert result.error == "ok" + + +# ───────── error handling: EngineDeadError propagation through layers ───── + + +class TestMultiprocExecutorRaisesEngineDeadError: + """``collective_rpc`` raises ``EngineDeadError`` when the engine is failed.""" + + def test_collective_rpc_raises_when_is_failed(self): + executor = object.__new__(MultiprocDiffusionExecutor) + executor._closed = False + executor._broadcast_mq = MagicMock() + executor._result_mq = MagicMock() + executor._result_mq.dequeue = MagicMock(side_effect=TimeoutError) + executor.is_failed = True + + with pytest.raises(EngineDeadError): + executor.collective_rpc( + "generate", + args=(MagicMock(),), + unique_reply_rank=0, + exec_all_ranks=True, + ) + + def test_collective_rpc_raises_mid_dequeue_when_is_failed(self): + """Worker dies while we are polling the dequeue loop.""" + executor, _, res_q = _make_executor() + + call_count = 0 + orig_dequeue = executor._result_mq.dequeue + + def _dying_dequeue(timeout=None): + nonlocal call_count + call_count += 1 + if call_count == 1: + executor.is_failed = True + raise TimeoutError + return orig_dequeue(timeout=timeout) + + executor._result_mq.dequeue = _dying_dequeue + + with pytest.raises(EngineDeadError): + executor.collective_rpc( + "generate", + args=(MagicMock(),), + unique_reply_rank=0, + exec_all_ranks=True, + ) + + +class TestDiffusionEngineDeadErrorPassthrough: + """``DiffusionEngine.add_req_and_wait_for_response`` re-raises + ``EngineDeadError`` from executor and wraps other errors.""" + + def test_engine_dead_error_propagates(self): + engine, executor, _, _ = _make_engine() + engine.execute_fn = Mock(side_effect=EngineDeadError()) + + with pytest.raises(EngineDeadError): + engine.add_req_and_wait_for_response(_mock_request("dead")) + + def test_runtime_error_wrapped_in_output(self): + engine, executor, _, _ = _make_engine() + engine.execute_fn = Mock(side_effect=RuntimeError("gpu fault")) + + out = engine.add_req_and_wait_for_response(_mock_request("fault")) + assert isinstance(out, DiffusionOutput) + assert "gpu fault" in out.error + + +class TestStageDiffusionClientErrorPropagation: + """Error surface behaviour of ``StageDiffusionClient``. + + Uses ``object.__new__`` to construct a client without spawning a real + subprocess, then manually sets the fields needed for each test. + """ + + def _make_client(self, *, engine_dead=False, proc_alive=True): + from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient + + client = object.__new__(StageDiffusionClient) + client.stage_id = 0 + client.final_output = True + client.final_output_type = "image" + client.default_sampling_params = None + client.custom_process_input_func = None + client.engine_input_source = None + + client._output_queue = asyncio.Queue() + client._rpc_results = {} + client._pending_rpcs = set() + client._tasks = {} + client._shutting_down = False + client._engine_dead = engine_dead + client._owns_process = True + client._proc = MagicMock( + is_alive=MagicMock(return_value=proc_alive), + exitcode=1, + ) + client._request_socket = MagicMock() + client._response_socket = MagicMock() + client._encoder = MagicMock() + client._decoder = MagicMock() + + return client + + @pytest.mark.asyncio + async def test_add_request_raises_when_dead(self): + client = self._make_client(engine_dead=True) + + with pytest.raises(EngineDeadError): + await client.add_request_async("req-3", "test prompt", None) + + def test_check_health_raises_when_dead(self): + client = self._make_client(engine_dead=True) + + with pytest.raises(EngineDeadError): + client.check_health() + + def test_check_health_ok_when_alive(self): + client = self._make_client() + client.check_health() + + def test_get_output_raises_engine_dead_when_dead(self): + """When ``_engine_dead`` is True and the output queue is empty, + ``get_diffusion_output_nowait`` must raise ``EngineDeadError``.""" + client = self._make_client(engine_dead=True) + # Simulate _drain_responses as a no-op (no ZMQ socket) + client._response_socket.recv.side_effect = zmq.Again + + with pytest.raises(EngineDeadError): + client.get_diffusion_output_nowait() + + def test_get_output_returns_none_when_alive_and_empty(self): + """When the engine is alive and the queue is empty, return None.""" + client = self._make_client() + client._response_socket.recv.side_effect = zmq.Again + + assert client.get_diffusion_output_nowait() is None + + def test_check_health_raises_when_proc_dead(self): + """``check_health`` detects a dead subprocess via ``_proc.is_alive()`` + and raises ``EngineDeadError``, setting ``_engine_dead`` as a + side effect.""" + client = self._make_client(proc_alive=False) + + with pytest.raises(EngineDeadError, match="not alive"): + client.check_health() + + assert client._engine_dead is True + + def test_get_output_raises_when_proc_dead(self): + """When the subprocess has died (non-signal exit) and the output + queue is empty, ``get_diffusion_output_nowait`` must raise + ``EngineDeadError`` with the exit code.""" + client = self._make_client(proc_alive=False) + client._response_socket.recv.side_effect = zmq.Again + + with pytest.raises(EngineDeadError, match="exit code"): + client.get_diffusion_output_nowait() + + assert client._engine_dead is True + + def test_get_output_returns_none_on_signal_death(self): + """When the subprocess was killed by a signal (exit code > 128), + ``get_diffusion_output_nowait`` returns ``None`` and sets + ``_shutting_down`` instead of raising.""" + client = self._make_client(proc_alive=False) + client._proc.exitcode = 137 # SIGKILL (128 + 9) + client._response_socket.recv.side_effect = zmq.Again + + result = client.get_diffusion_output_nowait() + + assert result is None + assert client._shutting_down is True + assert client._engine_dead is True + + +# ───────── monitor thread & death sentinel integration tests ───────── + + +def _poll_flag(get_flag, *, timeout=5.0, interval=0.05) -> bool: + """Poll until ``get_flag()`` returns True or *timeout* elapses.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if get_flag(): + return True + time.sleep(interval) + return False + + +def _make_short_lived_process() -> mp.Process: + """Spawn a real subprocess that exits immediately. + + The process must be started with ``"fork"`` (or the platform default) + so that it can use a plain ``lambda`` as its target — ``"spawn"`` would + fail to pickle it. + """ + ctx = mp.get_context("fork") + p = ctx.Process(target=lambda: None, name="ShortLivedWorker-0") + p.start() + return p + + +class TestMultiprocExecutorWorkerMonitor: + """Integration tests for ``start_worker_monitor``. + + Uses real short-lived subprocesses so that OS-level sentinel fd + readiness is exercised end-to-end. + """ + + def test_worker_monitor_sets_is_failed_and_calls_callbacks_on_death(self): + """When a worker process dies, the monitor thread must: + 1. Set ``is_failed = True`` + 2. Call ``shutdown()`` (which sets ``_closed = True``) + 3. Invoke all registered failure callbacks + """ + executor = object.__new__(MultiprocDiffusionExecutor) + executor._closed = False + executor.is_failed = False + executor._failure_callbacks = [] + executor._broadcast_mq = None + executor._result_mq = None + executor.resources = None + # Use a no-op so shutdown() doesn't crash on None resources. + executor._finalizer = lambda: None + + proc = _make_short_lived_process() + executor._processes = [proc] + + callback_called = threading.Event() + executor.register_failure_callback(callback_called.set) + + executor.start_worker_monitor() + + # Wait for the process to exit and the monitor to react. + proc.join(5) + assert _poll_flag(lambda: executor.is_failed), "is_failed was not set" + assert executor._closed, "shutdown() was not called" + assert callback_called.wait(timeout=2), "failure callback was not invoked" + + def test_worker_monitor_noop_when_already_closed(self): + """If ``_closed`` is already True when the process dies (orderly + shutdown), the monitor must *not* set ``is_failed``.""" + executor = object.__new__(MultiprocDiffusionExecutor) + executor._closed = True # already shut down + executor.is_failed = False + executor._failure_callbacks = [] + executor._broadcast_mq = None + executor._result_mq = None + executor.resources = None + executor._finalizer = lambda: None + + proc = _make_short_lived_process() + executor._processes = [proc] + + executor.start_worker_monitor() + proc.join(5) + + # Give the monitor thread a chance to run (it should early-return). + time.sleep(0.3) + assert not executor.is_failed, "is_failed should remain False on orderly shutdown" + + +class TestStageDiffusionClientProcMonitor: + """Integration test for ``StageDiffusionClient._start_proc_monitor``. + + Uses a real short-lived subprocess to verify the sentinel-based + detection pipeline. + """ + + def test_proc_monitor_sets_engine_dead_on_process_death(self): + """When the subprocess dies, the monitor thread must set + ``_engine_dead = True``.""" + from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient + + client = object.__new__(StageDiffusionClient) + client.stage_id = 0 + client._shutting_down = False + client._engine_dead = False + + proc = _make_short_lived_process() + client._proc = proc + + client._start_proc_monitor() + proc.join(5) + + assert _poll_flag(lambda: client._engine_dead), "_engine_dead was not set" + + +class TestDrainResponsesDeathSentinel: + """Tests for death sentinel and error routing in + ``StageDiffusionClient._drain_responses()``. + """ + + def _make_client(self): + from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient + + client = object.__new__(StageDiffusionClient) + client.stage_id = 0 + client._engine_dead = False + client._shutting_down = False + client._output_queue = asyncio.Queue() + client._rpc_results = {} + client._pending_rpcs = set() + client._response_socket = MagicMock() + client._decoder = MagicMock() + return client + + def test_drain_responses_sets_engine_dead_on_death_sentinel(self): + """When ``_drain_responses`` receives the ``DIFFUSION_PROC_DEAD`` + sentinel, it must set ``_engine_dead = True`` and stop draining + (decoder is never called).""" + client = self._make_client() + + # First recv returns the death sentinel, second would be a normal + # message but should never be reached. + client._response_socket.recv.side_effect = [ + StageDiffusionProc.DIFFUSION_PROC_DEAD, + b"should-not-be-reached", + ] + + client._drain_responses() + + assert client._engine_dead is True + client._decoder.decode.assert_not_called() + + def test_drain_responses_routes_error_as_omni_request_output(self): + """When ``_drain_responses`` receives a ``{"type": "error"}`` message + with a ``request_id``, it must place an ``OmniRequestOutput`` with + the error on ``_output_queue``.""" + client = self._make_client() + + error_msg = { + "type": "error", + "request_id": "req-fail", + "error": "gpu fault", + } + # First recv returns the encoded error, second raises zmq.Again. + client._response_socket.recv.side_effect = [b"encoded-error", zmq.Again] + client._decoder.decode.return_value = error_msg + + client._drain_responses() + + assert not client._output_queue.empty() + output = client._output_queue.get_nowait() + assert isinstance(output, OmniRequestOutput) + assert output.request_id == "req-fail" + assert output.error == "gpu fault" + assert output.finished is True diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 7a3caba11..256e3e0a3 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -4,10 +4,12 @@ import threading from collections import deque from types import SimpleNamespace +from unittest.mock import patch import pytest import torch from pytest_mock import MockerFixture +from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.request import RequestStatus from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase @@ -335,6 +337,27 @@ def test_cleanup_after_poll_flow(build_adapter): assert "ext-flow" not in adapter.request_payload +def test_finish_requests_restores_status(build_adapter): + """Abort path must pop ``requests_origin_status`` and restore pre-wait status. + + While ``process_pending_chunks`` holds a request off the scheduler queues, the + adapter records the prior status (WAITING or RUNNING). ``finish_requests`` must + put that status back on the live ``Request`` so base ``Scheduler.finish_requests`` + can finish bookkeeping without inconsistent state / crashes. + """ + adapter, _ = build_adapter(stage_id=1) + req_id = "req-abort-during-chunk" + prior = RequestStatus.RUNNING + request = _req(req_id, RequestStatus.WAITING_FOR_CHUNK) + adapter.requests_origin_status[req_id] = prior + requests_map = {req_id: request} + + adapter.finish_requests([req_id], RequestStatus.FINISHED_ABORTED, requests_map) + + assert request.status == prior + assert req_id not in adapter.requests_origin_status + + # --------------------------------------------------------------- # Scheduler trigger tests # --------------------------------------------------------------- @@ -508,3 +531,31 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF assert len(cleanup_calls) == 0 assert len(save_calls) == 1 + + +def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture): + """``OmniARScheduler.finish_requests`` must run chunk adapter hook before vLLM base.""" + from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + + order: list[str] = [] + + adapter = mocker.MagicMock() + + def _adapter_finish(request_ids, finished_status, requests): + order.append("adapter") + return [] + + adapter.finish_requests.side_effect = _adapter_finish + + def _super_finish(_self, request_ids, finished_status): + order.append("super") + return [] + + sched = OmniARScheduler.__new__(OmniARScheduler) + sched.chunk_transfer_adapter = adapter + sched.requests = {} + + with patch.object(VLLMScheduler, "finish_requests", _super_finish): + OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED) + + assert order == ["adapter", "super"] diff --git a/tests/distributed/omni_connectors/test_tp_rank_aware.py b/tests/distributed/omni_connectors/test_tp_rank_aware.py new file mode 100644 index 000000000..d4793479a --- /dev/null +++ b/tests/distributed/omni_connectors/test_tp_rank_aware.py @@ -0,0 +1,716 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for rank-aware KV transfer (TP > 1) and heterogeneous TP support. + +Covers: +- _build_rank_aware_send_keys / _build_rank_aware_recv_keys +- _get_kv_source_ranks / _get_kv_target_ranks / get_kv_connector_key +- update_sender_info storing base host/port +- receive path constructing per-rank metadata for connector.get() +- Mooncake connector _query_metadata_at and partial-metadata get() path +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + KVCacheTransferData, + OmniKVCacheConfig, + OmniKVTransferManager, +) +from vllm_omni.distributed.omni_connectors.utils.initialization import ( + KV_RANK_PORT_STRIDE, +) +from vllm_omni.distributed.omni_connectors.utils.kv_utils import ( + KVTPTopology, + build_rank_aware_recv_keys, + build_rank_aware_send_keys, + get_kv_connector_key, + get_kv_source_ranks, + get_kv_target_ranks, + merge_received_rank_shards, + slice_received_rank_shard, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_manager( + from_tp: int = 1, + to_tp: int = 1, + local_rank: int = 0, + from_stage: str = "stage0", + to_stage: str = "stage1", + stage_id: str = "stage1", + need_recv: bool = True, + need_send: bool = False, + recv_timeout: float = 0.3, +) -> OmniKVTransferManager: + """Build a manager with TP params injected, bypassing torch.distributed.""" + config = OmniKVCacheConfig( + connector_config={"type": "mock"}, + from_stage=from_stage, + to_stage=to_stage, + stage_id=stage_id, + need_recv_cache=need_recv, + need_send_cache=need_send, + recv_timeout=recv_timeout, + from_tp=from_tp, + to_tp=to_tp, + ) + with ( + patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=local_rank), + patch( + "vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", + return_value=max(from_tp, to_tp), + ), + ): + mgr = OmniKVTransferManager(config) + return mgr + + +def _make_payload(head_values: list[float], request_id: str = "req-1") -> dict: + head_tensor = torch.tensor(head_values, dtype=torch.float32).view(1, len(head_values), 1).repeat(2, 1, 1) + return { + "request_id": request_id, + "layer_blocks": { + "key_cache": [head_tensor.clone()], + "value_cache": [(head_tensor + 100).clone()], + }, + "block_ids": [0], + "metadata": {"seq_len": 2}, + } + + +def _make_transfer_data(head_values: list[float], request_id: str = "req-1") -> KVCacheTransferData: + payload = _make_payload(head_values, request_id=request_id) + return KVCacheTransferData( + request_id=request_id, + layer_blocks=payload["layer_blocks"], + block_ids=payload["block_ids"], + metadata=payload["metadata"], + ) + + +# ── Key format helper ──────────────────────────────────────────────── + + +class TestConnectorKeyFormat: + def test_key_format_matches_pr2677(self): + key = get_kv_connector_key("req-1", "stage0", 0, 1, 2) + assert key == "req-1_stage0_0_1_2" + + def test_key_fields_are_positional(self): + key = get_kv_connector_key("r", "s", 5, 3, 7) + parts = key.split("_") + assert parts == ["r", "s", "5", "3", "7"] + + +# ── Source / target rank mapping ───────────────────────────────────── + + +class TestRankMapping: + """Verify get_kv_target_ranks and get_kv_source_ranks for various TP configs.""" + + def test_homogeneous_tp2_rank0(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0) + assert get_kv_target_ranks(topo) == [0] + assert get_kv_source_ranks(topo) == [0] + + def test_homogeneous_tp2_rank1(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=1) + assert get_kv_target_ranks(topo) == [1] + assert get_kv_source_ranks(topo) == [1] + + def test_homogeneous_tp4_rank3(self): + topo = KVTPTopology(source_tp_size=4, target_tp_size=4, local_rank=3) + assert get_kv_target_ranks(topo) == [3] + assert get_kv_source_ranks(topo) == [3] + + def test_sender_gt_receiver_tp4_to_tp2_rank0(self): + """Receiver rank 0 should receive from sender rank 0 and 1.""" + topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=0) + assert get_kv_source_ranks(topo) == [0, 1] + + def test_sender_gt_receiver_tp4_to_tp2_rank1(self): + """Receiver rank 1 should receive from sender rank 2 and 3.""" + topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1) + assert get_kv_source_ranks(topo) == [2, 3] + + def test_sender_lt_receiver_tp2_to_tp4_rank0(self): + """Sender rank 0 should send to receiver ranks 0 and 1.""" + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0) + assert get_kv_target_ranks(topo) == [0, 1] + + def test_sender_lt_receiver_tp2_to_tp4_rank1(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1) + assert get_kv_target_ranks(topo) == [2, 3] + + def test_receiver_lt_sender_source_ranks(self): + """Receiver rank 0 with tp2_to_tp4 should source from rank 0 only.""" + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0) + assert get_kv_source_ranks(topo) == [0] + + def test_invalid_topology_raises(self): + topo = KVTPTopology(source_tp_size=3, target_tp_size=2, local_rank=0) + with pytest.raises(ValueError, match="divisible"): + get_kv_source_ranks(topo) + + +# ── _build_rank_aware_recv_keys ────────────────────────────────────── + + +class TestBuildRankAwareRecvKeys: + """Verify build_rank_aware_recv_keys returns (key, from_rank) tuples.""" + + def test_tp1_returns_legacy_key_with_none_rank(self): + topo = KVTPTopology(source_tp_size=1, target_tp_size=1, local_rank=0) + pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) + assert len(pairs) == 1 + key, rank = pairs[0] + assert key == "omni_stage0_to_stage1_kv_cache_req-1" + assert rank is None + + def test_homogeneous_tp2_rank0(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0) + pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) + assert len(pairs) == 1 + key, rank = pairs[0] + assert key == "req-1_stage0_0_0_0" + assert rank == 0 + + def test_homogeneous_tp2_rank1(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=1) + pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) + assert len(pairs) == 1 + key, rank = pairs[0] + assert key == "req-1_stage0_0_1_1" + assert rank == 1 + + def test_heterogeneous_tp4_to_tp2_rank0_gets_two_keys(self): + """Receiver rank 0 with source_tp=4, target_tp=2 should get 2 keys.""" + topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=0) + pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) + assert len(pairs) == 2 + + keys = [k for k, _ in pairs] + ranks = [r for _, r in pairs] + assert keys == ["req-1_stage0_0_0_0", "req-1_stage0_0_1_0"] + assert ranks == [0, 1] + + def test_heterogeneous_tp4_to_tp2_rank1_gets_two_keys(self): + topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1) + pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) + assert len(pairs) == 2 + + ranks = [r for _, r in pairs] + assert ranks == [2, 3] + + def test_heterogeneous_tp2_to_tp4_rank2_gets_one_key(self): + """Receiver rank 2 with source_tp=2, target_tp=4 should get 1 key from sender rank 1.""" + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=2) + pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) + assert len(pairs) == 1 + key, rank = pairs[0] + assert rank == 1 + assert key == "req-1_stage0_0_1_2" + + +# ── _build_rank_aware_send_keys ────────────────────────────────────── + + +class TestBuildRankAwareSendKeys: + def test_tp1_returns_legacy_key(self): + topo = KVTPTopology(source_tp_size=1, target_tp_size=1, local_rank=0) + keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo) + assert keys == ["omni_stage0_to_stage1_kv_cache_req-1"] + + def test_homogeneous_tp2_rank0(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0) + keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo) + assert keys == ["req-1_stage0_0_0_0"] + + def test_sender_lt_receiver_tp2_to_tp4_rank0_sends_two_keys(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0) + keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo) + assert len(keys) == 2 + assert keys == ["req-1_stage0_0_0_0", "req-1_stage0_0_0_1"] + + +# ── update_sender_info stores base host/port ───────────────────────── + + +class TestUpdateSenderInfoBase: + def test_stores_base_host_and_port(self): + mgr = _make_manager(from_tp=2, to_tp=2, local_rank=0) + mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) + + assert mgr._sender_base_host == "10.0.0.1" + assert mgr._sender_base_zmq_port == 50151 + + def test_rank1_adjusts_default_port_but_preserves_base(self): + mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1) + mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) + + assert mgr._sender_base_host == "10.0.0.1" + assert mgr._sender_base_zmq_port == 50151 + expected_adjusted = 50151 + 1 * KV_RANK_PORT_STRIDE + assert mgr.config.connector_config["sender_zmq_port"] == expected_adjusted + + def test_nested_sender_info_resolves_correctly(self): + """Nested sender_info keyed by integer stage id should resolve + using recv_stages (engine_input_source → recv_from).""" + config = OmniKVCacheConfig( + connector_config={"type": "mock"}, + stage_id=2, + engine_input_source=[1], + need_recv_cache=True, + from_tp=2, + to_tp=2, + ) + with ( + patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0), + patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=2), + ): + mgr = OmniKVTransferManager(config) + mgr.update_sender_info( + { + 0: {"host": "10.0.0.1", "zmq_port": 50151}, + 1: {"host": "10.0.0.2", "zmq_port": 50152}, + } + ) + assert mgr._sender_base_host == "10.0.0.2" + assert mgr._sender_base_zmq_port == 50152 + + +# ── receive path constructs per-rank metadata ──────────────────────── + + +class TestReceiveConstructsMetadata: + """Verify that receive_kv_cache_for_request passes metadata with + correct (host, port) to connector.get() for heterogeneous TP.""" + + def test_tp1_no_metadata_passed(self): + """TP=1: connector.get() should be called WITHOUT metadata.""" + mgr = _make_manager(from_tp=1, to_tp=1, local_rank=0, recv_timeout=0.05) + mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) + + calls = [] + + class _Connector: + def get(self, from_stage, to_stage, get_key, metadata=None): + calls.append({"key": get_key, "metadata": metadata}) + return None + + mgr._connector = _Connector() + mgr.receive_kv_cache_for_request("req-1") + + assert len(calls) > 0 + assert calls[0]["metadata"] is None + + def test_homogeneous_tp2_rank0_passes_metadata(self): + """TP=2 rank 0: metadata should point to sender rank 0's port.""" + mgr = _make_manager(from_tp=2, to_tp=2, local_rank=0, recv_timeout=0.05) + mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) + + calls = [] + + class _Connector: + def get(self, from_stage, to_stage, get_key, metadata=None): + calls.append({"key": get_key, "metadata": metadata}) + return None + + mgr._connector = _Connector() + mgr.receive_kv_cache_for_request("req-1") + + assert len(calls) > 0 + meta = calls[0]["metadata"] + assert meta is not None + assert meta["source_host"] == "10.0.0.1" + assert meta["source_port"] == 50151 + 0 * KV_RANK_PORT_STRIDE + + def test_homogeneous_tp2_rank1_passes_metadata_with_offset(self): + mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1, recv_timeout=0.05) + mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) + + calls = [] + + class _Connector: + def get(self, from_stage, to_stage, get_key, metadata=None): + calls.append({"key": get_key, "metadata": metadata}) + return None + + mgr._connector = _Connector() + mgr.receive_kv_cache_for_request("req-1") + + meta = calls[0]["metadata"] + assert meta["source_port"] == 50151 + 1 * KV_RANK_PORT_STRIDE + + def test_heterogeneous_tp4_to_tp2_rank0_multiple_metadata(self): + """Receiver rank 0 with source_tp=4, target_tp=2 should call get() with + two different metadata entries for sender ranks 0 and 1.""" + mgr = _make_manager(from_tp=4, to_tp=2, local_rank=0, recv_timeout=0.05) + mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) + + calls = [] + + class _Connector: + def get(self, from_stage, to_stage, get_key, metadata=None): + calls.append({"key": get_key, "metadata": metadata}) + return None + + mgr._connector = _Connector() + mgr.receive_kv_cache_for_request("req-1") + + seen_ports = set() + for c in calls: + if c["metadata"]: + seen_ports.add(c["metadata"]["source_port"]) + expected_ports = { + 50151 + 0 * KV_RANK_PORT_STRIDE, + 50151 + 1 * KV_RANK_PORT_STRIDE, + } + assert expected_ports.issubset(seen_ports) + + +# ── Mooncake connector _query_metadata_at ──────────────────────────── + + +class TestMooncakeQueryMetadataAt: + """Test the connector's _query_metadata_at method and partial-metadata + path in get() without requiring real RDMA/Mooncake.""" + + def test_query_metadata_at_returns_full_metadata(self): + """Mock the ZMQ interaction to verify _query_metadata_at returns + complete metadata including data_size.""" + + try: + from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import ( + MooncakeTransferEngineConnector, + QueryResponse, + ) + except ImportError: + pytest.skip("Mooncake not available") + + import msgspec + + connector = MagicMock(spec=MooncakeTransferEngineConnector) + connector._get_req_socket = MagicMock() + + mock_socket = MagicMock() + resp = QueryResponse(request_id="test_key@s0_s1", data_size=4096, is_fast_path=True) + mock_socket.recv.return_value = msgspec.msgpack.encode(resp) + connector._get_req_socket.return_value = mock_socket + + result = MooncakeTransferEngineConnector._query_metadata_at( + connector, + "test_key@s0_s1", + "10.0.0.1", + 50151, + ) + + assert result is not None + assert result["source_host"] == "10.0.0.1" + assert result["source_port"] == 50151 + assert result["data_size"] == 4096 + assert result["is_fast_path"] is True + + def test_query_metadata_at_returns_none_on_not_found(self): + try: + from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import ( + INFO_NOT_FOUND, + MooncakeTransferEngineConnector, + ) + except ImportError: + pytest.skip("Mooncake not available") + + connector = MagicMock(spec=MooncakeTransferEngineConnector) + mock_socket = MagicMock() + mock_socket.recv.return_value = INFO_NOT_FOUND + connector._get_req_socket.return_value = mock_socket + + result = MooncakeTransferEngineConnector._query_metadata_at( + connector, + "test_key@s0_s1", + "10.0.0.1", + 50151, + ) + assert result is None + + +# ── Merge / slice hooks ────────────────────────────────────────────── + + +class TestMergeSliceHooks: + def test_single_shard_passes_through(self): + payload = {"layer_blocks": {"key_cache": [1]}} + assert merge_received_rank_shards([payload]) == payload + + def test_default_merger_concats_head_dim(self): + p0 = _make_payload([0.0]) + p1 = _make_payload([1.0]) + result = merge_received_rank_shards([p0, p1]) + key_cache = result["layer_blocks"]["key_cache"][0] + value_cache = result["layer_blocks"]["value_cache"][0] + assert key_cache.shape == (2, 2, 1) + assert value_cache.shape == (2, 2, 1) + assert torch.equal(key_cache[:, :, 0], torch.tensor([[0.0, 1.0], [0.0, 1.0]])) + assert torch.equal(value_cache[:, :, 0], torch.tensor([[100.0, 101.0], [100.0, 101.0]])) + + def test_custom_merger_hook_called(self): + merged = {"merged": True} + assert merge_received_rank_shards([{}, {}], merger=lambda payloads: merged) == merged + + def test_slicer_hook_called(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0) + sliced = {"sliced": True} + assert slice_received_rank_shard({"full": True}, topo, slicer=lambda payload: sliced) == sliced + + def test_default_slicer_extracts_rank_local_heads(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1) + payload = _make_payload([0.0, 1.0]) + result = slice_received_rank_shard(payload, topo) + key_cache = result["layer_blocks"]["key_cache"][0] + value_cache = result["layer_blocks"]["value_cache"][0] + assert key_cache.shape == (2, 1, 1) + assert value_cache.shape == (2, 1, 1) + assert torch.equal(key_cache[:, :, 0], torch.tensor([[1.0], [1.0]])) + assert torch.equal(value_cache[:, :, 0], torch.tensor([[101.0], [101.0]])) + + def test_presliced_payload_is_not_sliced_twice(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1) + payload = _make_payload([1.0]) + payload["metadata"]["tp_head_slice"] = {"applied": True, "target_rank": 1} + result = slice_received_rank_shard(payload, topo) + assert result is payload + + def test_round_trip_merge_from_tp4_to_tp2(self): + topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1) + source_ranks = get_kv_source_ranks(topo) + payloads = [_make_payload([float(rank)]) for rank in source_ranks] + result = merge_received_rank_shards(payloads) + key_cache = result["layer_blocks"]["key_cache"][0] + assert torch.equal(key_cache[:, :, 0], torch.tensor([[2.0, 3.0], [2.0, 3.0]])) + + def test_round_trip_slice_from_tp2_to_tp4(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=3) + payload = _make_payload([2.0, 3.0]) + result = slice_received_rank_shard(payload, topo) + key_cache = result["layer_blocks"]["key_cache"][0] + assert torch.equal(key_cache[:, :, 0], torch.tensor([[3.0], [3.0]])) + + +class TestSenderSideSlicing: + def test_transfer_slices_before_sending_to_multiple_targets(self): + mgr = _make_manager( + from_tp=2, + to_tp=4, + local_rank=0, + need_send=True, + need_recv=False, + ) + sent_payloads = [] + + class _Connector: + supports_raw_data = False + + def put(self, from_stage, to_stage, put_key, data): + sent_payloads.append((put_key, KVCacheTransferData.from_bytes(data))) + return True, len(data), {} + + mgr._connector = _Connector() + mgr._transfer_kv_cache(_make_transfer_data([0.0, 1.0]), "req-1") + + assert [key for key, _ in sent_payloads] == ["req-1_stage0_0_0_0", "req-1_stage0_0_0_1"] + assert sent_payloads[0][1]["layer_blocks"]["key_cache"][0].shape == (2, 1, 1) + assert sent_payloads[1][1]["layer_blocks"]["key_cache"][0].shape == (2, 1, 1) + assert torch.equal( + sent_payloads[0][1]["layer_blocks"]["key_cache"][0][:, :, 0], + torch.tensor([[0.0], [0.0]]), + ) + assert torch.equal( + sent_payloads[1][1]["layer_blocks"]["key_cache"][0][:, :, 0], + torch.tensor([[1.0], [1.0]]), + ) + assert sent_payloads[0][1]["metadata"]["tp_head_slice"]["target_rank"] == 0 + assert sent_payloads[1][1]["metadata"]["tp_head_slice"]["target_rank"] == 1 + + +class _MockBroadcastGroup: + def __init__(self, world_size: int, rank_in_group: int, broadcast_value=None, recv_value=None): + self.world_size = world_size + self.rank_in_group = rank_in_group + self.broadcast_value = broadcast_value + self.recv_value = recv_value + self.broadcast_calls = [] + self.send_calls = [] + self.recv_calls = [] + self.shm_broadcaster = None + + def broadcast_object(self, obj=None, src: int = 0): + self.broadcast_calls.append((obj, src)) + return self.broadcast_value if self.broadcast_value is not None else obj + + def send_object(self, obj, dst: int): + self.send_calls.append((dst, obj)) + + def recv_object(self, src: int): + self.recv_calls.append(src) + return self.recv_value + + +class TestDistributedReceive: + def test_tp_cfg_leader_receives_then_sends_branch_local_payloads(self): + mgr = _make_manager(from_tp=2, to_tp=4, local_rank=0) + req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace()) + world_group = _MockBroadcastGroup(world_size=4, rank_in_group=2) + cfg_group = _MockBroadcastGroup(world_size=3, rank_in_group=0) + + def _receive(req_obj, cfg_func, target_device): + req_obj.past_key_values = SimpleNamespace(key_cache=[torch.tensor([1.0])]) + req_obj.kv_metadata = {"source": "leader"} + req_obj.sampling_params.past_key_values = req_obj.past_key_values + req_obj.sampling_params.kv_metadata = req_obj.kv_metadata + req_obj.sampling_params.cfg_text_past_key_values = SimpleNamespace(key_cache=[torch.tensor([2.0])]) + req_obj.sampling_params.cfg_text_kv_metadata = {"source": "cfg_text"} + req_obj.sampling_params.cfg_img_past_key_values = SimpleNamespace(key_cache=[torch.tensor([3.0])]) + req_obj.sampling_params.cfg_img_kv_metadata = {"source": "cfg_img"} + return True + + mgr.receive_multi_kv_cache = MagicMock(side_effect=_receive) + with ( + patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size", + return_value=3, + ), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank", + return_value=0, + ), + patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group), + ): + assert mgr.receive_multi_kv_cache_distributed(req) is True + + mgr.receive_multi_kv_cache.assert_called_once() + assert mgr.receive_multi_kv_cache.call_args.args[2] == torch.device("cpu") + assert req.kv_metadata == {"source": "leader"} + assert cfg_group.broadcast_calls == [] + assert [dst for dst, _ in cfg_group.send_calls] == [1, 2] + rank1_payload = cfg_group.send_calls[0][1] + rank2_payload = cfg_group.send_calls[1][1] + assert torch.equal(rank1_payload["past_key_values"].key_cache[0], torch.tensor([1.0])) + assert torch.equal(rank2_payload["past_key_values"].key_cache[0], torch.tensor([1.0])) + assert rank1_payload["sp.cfg_active_branch"] == "cfg_text" + assert rank2_payload["sp.cfg_active_branch"] == "cfg_img" + assert rank1_payload["sp.cfg_branch_roles"] == ["cfg_text", "cfg_img"] + assert rank2_payload["sp.cfg_branch_roles"] == ["cfg_text", "cfg_img"] + assert "sp.cfg_branch_past_key_values" in rank1_payload + assert "sp.cfg_branch_past_key_values" in rank2_payload + assert list(rank1_payload["sp.cfg_branch_past_key_values"].keys()) == ["cfg_text"] + assert list(rank2_payload["sp.cfg_branch_past_key_values"].keys()) == ["cfg_img"] + assert "sp.cfg_text_past_key_values" in rank1_payload + assert "sp.cfg_img_past_key_values" not in rank1_payload + assert "sp.cfg_img_past_key_values" in rank2_payload + assert "sp.cfg_text_past_key_values" not in rank2_payload + + def test_tp_cfg_follower_receives_local_payload_without_receiving(self): + mgr = _make_manager(from_tp=2, to_tp=4, local_rank=1) + req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace()) + world_group = _MockBroadcastGroup(world_size=4, rank_in_group=3) + cfg_payload = { + "past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]), + "kv_metadata": {"source": "main"}, + "sp.past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]), + "sp.kv_metadata": {"source": "main"}, + "sp.cfg_active_branch": "cfg_text", + "sp.cfg_branch_roles": ["cfg_text", "cfg_img"], + "sp.cfg_branch_past_key_values": { + "cfg_text": SimpleNamespace(key_cache=[torch.tensor([2.0])]), + }, + "sp.cfg_branch_kv_metadata": {"cfg_text": {"source": "cfg-text"}}, + "sp.cfg_text_past_key_values": SimpleNamespace(key_cache=[torch.tensor([2.0])]), + } + cfg_group = _MockBroadcastGroup(world_size=2, rank_in_group=1, recv_value=cfg_payload) + + mgr.receive_multi_kv_cache = MagicMock(return_value=True) + with ( + patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size", + return_value=2, + ), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank", + return_value=1, + ), + patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group), + ): + assert mgr.receive_multi_kv_cache_distributed(req) is True + + mgr.receive_multi_kv_cache.assert_not_called() + assert req.kv_metadata == {"source": "main"} + assert torch.equal(req.past_key_values.key_cache[0], torch.tensor([1.0])) + assert torch.equal(req.sampling_params.past_key_values.key_cache[0], torch.tensor([1.0])) + assert req.sampling_params.cfg_active_branch == "cfg_text" + assert req.sampling_params.cfg_branch_roles == ["cfg_text", "cfg_img"] + assert torch.equal( + req.sampling_params.cfg_branch_past_key_values["cfg_text"].key_cache[0], + torch.tensor([2.0]), + ) + assert req.sampling_params.cfg_branch_kv_metadata == {"cfg_text": {"source": "cfg-text"}} + assert torch.equal(req.sampling_params.cfg_text_past_key_values.key_cache[0], torch.tensor([2.0])) + assert cfg_group.broadcast_calls == [] + assert cfg_group.recv_calls == [0] + + def test_tp_without_cfg_keeps_independent_receive_path(self): + mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1) + req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace()) + world_group = _MockBroadcastGroup(world_size=2, rank_in_group=1) + mgr.receive_multi_kv_cache = MagicMock(return_value=True) + + with patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group): + assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True + + mgr.receive_multi_kv_cache.assert_called_once_with(req, None, torch.device("cpu")) + + +# ── TP auto-detect ─────────────────────────────────────────────────── + + +class TestAutoDetectTP: + def test_auto_detect_when_config_defaults(self): + """When config from_tp/to_tp == 1 (default), manager should auto-detect.""" + config = OmniKVCacheConfig( + connector_config={"type": "mock"}, + from_stage="s0", + stage_id="s1", + need_recv_cache=True, + ) + with ( + patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0), + patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=4), + ): + mgr = OmniKVTransferManager(config) + assert mgr._tp_topo.source_tp_size == 4 + assert mgr._tp_topo.target_tp_size == 4 + + def test_explicit_tp_overrides_auto_detect(self): + config = OmniKVCacheConfig( + connector_config={"type": "mock"}, + from_stage="s0", + stage_id="s1", + need_recv_cache=True, + from_tp=2, + to_tp=4, + ) + with ( + patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0), + patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=8), + ): + mgr = OmniKVTransferManager(config) + assert mgr._tp_topo.source_tp_size == 2 + assert mgr._tp_topo.target_tp_size == 4 diff --git a/tests/distributed/omni_coordinator/test_load_balancer.py b/tests/distributed/omni_coordinator/test_load_balancer.py index b2d1f3ee8..8350b33d3 100644 --- a/tests/distributed/omni_coordinator/test_load_balancer.py +++ b/tests/distributed/omni_coordinator/test_load_balancer.py @@ -7,7 +7,9 @@ import pytest from vllm_omni.distributed.omni_coordinator import ( InstanceInfo, + LeastQueueLengthBalancer, RandomBalancer, + RoundRobinBalancer, StageStatus, ) @@ -60,3 +62,173 @@ def test_load_balancer_select_returns_valid_index(): assert isinstance(index, int) assert 0 <= index < len(instances) + + +def test_round_robin_balancer_cycles_instances(): + now = time() + instances = [ + InstanceInfo( + input_addr="tcp://host:10001", + output_addr="tcp://host:10001-out", + stage_id=0, + status=StageStatus.UP, + queue_length=2, + last_heartbeat=now, + registered_at=now, + ), + InstanceInfo( + input_addr="tcp://host:10002", + output_addr="tcp://host:10002-out", + stage_id=0, + status=StageStatus.UP, + queue_length=1, + last_heartbeat=now, + registered_at=now, + ), + InstanceInfo( + input_addr="tcp://host:10003", + output_addr="tcp://host:10003-out", + stage_id=1, + status=StageStatus.UP, + queue_length=0, + last_heartbeat=now, + registered_at=now, + ), + ] + + balancer = RoundRobinBalancer() + results = [balancer.select({}, instances) for _ in range(5)] + + # Default start_index=0 => 0,1,2,0,1 + assert results == [0, 1, 2, 0, 1] + + +def test_round_robin_balancer_empty_instances_raises(): + with pytest.raises(ValueError, match="instances must not be empty"): + RoundRobinBalancer().select({}, []) + + +def test_round_robin_balancer_after_large_index_and_shorter_list(): + """Large start_index % len(instances) then counter wraps with shorter list.""" + now = time() + two = [ + InstanceInfo( + input_addr="tcp://host:10001", + output_addr="tcp://host:10001-out", + stage_id=0, + status=StageStatus.UP, + queue_length=0, + last_heartbeat=now, + registered_at=now, + ), + InstanceInfo( + input_addr="tcp://host:10002", + output_addr="tcp://host:10002-out", + stage_id=0, + status=StageStatus.UP, + queue_length=0, + last_heartbeat=now, + registered_at=now, + ), + ] + balancer = RoundRobinBalancer(start_index=7) + assert balancer.select({}, two) == 1 # 7 % 2 + assert balancer.select({}, two) == 0 # next index wrapped to 0 + + +def test_least_queue_length_balancer_picks_min_queue(): + now = time() + instances = [ + InstanceInfo( + input_addr="tcp://host:10001", + output_addr="tcp://host:10001-out", + stage_id=0, + status=StageStatus.UP, + queue_length=2, + last_heartbeat=now, + registered_at=now, + ), + InstanceInfo( + input_addr="tcp://host:10002", + output_addr="tcp://host:10002-out", + stage_id=0, + status=StageStatus.UP, + queue_length=0, + last_heartbeat=now, + registered_at=now, + ), + InstanceInfo( + input_addr="tcp://host:10003", + output_addr="tcp://host:10003-out", + stage_id=1, + status=StageStatus.UP, + queue_length=5, + last_heartbeat=now, + registered_at=now, + ), + ] + + balancer = LeastQueueLengthBalancer() + index = balancer.select({}, instances) + assert index == 1 + + +def test_least_queue_length_balancer_empty_instances_raises(): + with pytest.raises(ValueError, match="instances must not be empty"): + LeastQueueLengthBalancer().select({}, []) + + +def test_least_queue_length_balancer_equal_queues_uses_choice(mocker): + now = time() + instances = [ + InstanceInfo( + input_addr="tcp://host:10001", + output_addr="tcp://host:10001-out", + stage_id=0, + status=StageStatus.UP, + queue_length=3, + last_heartbeat=now, + registered_at=now, + ), + InstanceInfo( + input_addr="tcp://host:10002", + output_addr="tcp://host:10002-out", + stage_id=0, + status=StageStatus.UP, + queue_length=3, + last_heartbeat=now, + registered_at=now, + ), + InstanceInfo( + input_addr="tcp://host:10003", + output_addr="tcp://host:10003-out", + stage_id=1, + status=StageStatus.UP, + queue_length=3, + last_heartbeat=now, + registered_at=now, + ), + ] + balancer = LeastQueueLengthBalancer() + mocker.patch( + "vllm_omni.distributed.omni_coordinator.load_balancer.random.choice", + return_value=2, + ) + assert balancer.select({}, instances) == 2 + + +def test_least_queue_length_balancer_negative_queue_raises(): + now = time() + instances = [ + InstanceInfo( + input_addr="tcp://host:10001", + output_addr="tcp://host:10001-out", + stage_id=0, + status=StageStatus.UP, + queue_length=-1, + last_heartbeat=now, + registered_at=now, + ), + ] + with pytest.raises(ValueError, match="queue_length must be non-negative"): + LeastQueueLengthBalancer().select({}, instances) diff --git a/tests/e2e/accuracy/conftest.py b/tests/e2e/accuracy/conftest.py index 062750b3c..709fdf345 100644 --- a/tests/e2e/accuracy/conftest.py +++ b/tests/e2e/accuracy/conftest.py @@ -1,16 +1,18 @@ from __future__ import annotations import os -import shutil import subprocess from contextlib import contextmanager from dataclasses import dataclass +from io import BytesIO from pathlib import Path import pytest +import requests import torch +from PIL import Image -from tests.conftest import OmniServer, OmniServerParams +from tests.helpers.runtime import OmniServer, OmniServerParams def pytest_addoption(parser): @@ -183,16 +185,26 @@ def accuracy_artifact_root() -> Path: return root -def reset_artifact_dir(path: Path) -> Path: - if path.exists(): - shutil.rmtree(path) - path.mkdir(parents=True, exist_ok=True) - return path +@pytest.fixture(scope="session") +def qwen_bear_image(accuracy_artifact_root: Path) -> Image.Image: + """Download the Qwen bear image from the URL and save it to the accuracy artifact root.""" + QWEN_BEAR_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png" + response = requests.get(QWEN_BEAR_IMAGE_URL, timeout=60) + response.raise_for_status() + image = Image.open(BytesIO(response.content)).convert("RGB") + image.save(accuracy_artifact_root / "qwen_bear.png") + return image -def infer_model_label(model: str) -> str: - label = Path(model.rstrip("/\\")).name or "model" - return "".join(char if char.isalnum() or char in {"-", "_"} else "_" for char in label) +@pytest.fixture(scope="session") +def rabbit_image(accuracy_artifact_root: Path) -> Image.Image: + """Download the rabbit image from the URL and save it to the accuracy artifact root.""" + RABBIT_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/rabbit.png" + response = requests.get(RABBIT_IMAGE_URL, timeout=60) + response.raise_for_status() + image = Image.open(BytesIO(response.content)).convert("RGB") + image.save(accuracy_artifact_root / "rabbit.png") + return image def _build_accuracy_server_config( diff --git a/tests/e2e/accuracy/helpers.py b/tests/e2e/accuracy/helpers.py new file mode 100644 index 000000000..382d3ea9b --- /dev/null +++ b/tests/e2e/accuracy/helpers.py @@ -0,0 +1,115 @@ +from pathlib import Path + +import numpy as np +import pytest +import torch +from PIL import Image +from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure + + +def reset_artifact_dir(path: Path) -> Path: + import shutil + + if path.exists(): + shutil.rmtree(path) + path.mkdir(parents=True, exist_ok=True) + return path + + +def infer_model_label(model: str) -> str: + label = Path(model.rstrip("/\\")).name or "model" + return "".join(char if char.isalnum() or char in {"-", "_"} else "_" for char in label) + + +def model_output_dir(parent_dir: Path, model: str) -> Path: + safe_model_name = model.split("/")[-1].replace(".", "_") + path = parent_dir / safe_model_name + path.mkdir(parents=True, exist_ok=True) + return path + + +def assert_similarity( + *, + model_name: str, + vllm_image: Image.Image, + diffusers_image: Image.Image, + ssim_threshold: float, + psnr_threshold: float, + width: int | None = None, + height: int | None = None, + compare_mode: str = "RGB", +) -> None: + requested_size = (width, height) if width is not None and height is not None else None + if requested_size is not None and diffusers_image.size != requested_size: + pytest.skip( + "Skipping as diffusers baseline output is corrupt and not comparable: " + f"dimensions do not match requested size; requested={requested_size}, got={diffusers_image.size}." + ) + + assert vllm_image.size == diffusers_image.size, ( + f"Online and diffusers output sizes mismatch: online={vllm_image.size}, diffusers={diffusers_image.size}" + ) + + ssim_score, psnr_score = compute_image_ssim_psnr( + prediction=vllm_image, + reference=diffusers_image, + compare_mode=compare_mode, + ) + print(f"{model_name} similarity metrics:") + print(f" SSIM: value={ssim_score:.6f}, threshold>={ssim_threshold:.6f}, range=[-1, 1], higher_is_better=True") + print( + f" PSNR: value={psnr_score:.6f} dB, threshold>={psnr_threshold:.6f} dB, range=[0, +inf), higher_is_better=True" + ) + + assert ssim_score >= ssim_threshold, ( + f"SSIM below threshold for {model_name}: got {ssim_score:.6f}, expected >= {ssim_threshold:.6f}." + ) + assert psnr_score >= psnr_threshold, ( + f"PSNR below threshold for {model_name}: got {psnr_score:.6f}, expected >= {psnr_threshold:.6f}." + ) + + +def assert_image_sequence_similarity( + *, + model_name: str, + vllm_images: list[Image.Image], + diffusers_images: list[Image.Image], + ssim_threshold: float, + psnr_threshold: float, + compare_mode: str = "RGB", +) -> None: + assert len(vllm_images) == len(diffusers_images), ( + f"Output image count mismatch for {model_name}: online={len(vllm_images)}, diffusers={len(diffusers_images)}" + ) + for index, (vllm_image, diffusers_image) in enumerate(zip(vllm_images, diffusers_images, strict=True), start=1): + assert_similarity( + model_name=f"{model_name}[layer={index}]", + vllm_image=vllm_image, + diffusers_image=diffusers_image, + ssim_threshold=ssim_threshold, + psnr_threshold=psnr_threshold, + compare_mode=compare_mode, + ) + + +def compute_image_ssim_psnr( + *, + prediction: Image.Image, + reference: Image.Image, + compare_mode: str = "RGB", +) -> tuple[float, float]: + pred_tensor = _pil_to_batched_tensor(prediction, compare_mode=compare_mode) + ref_tensor = _pil_to_batched_tensor(reference, compare_mode=compare_mode) + + ssim_metric = StructuralSimilarityIndexMeasure(data_range=1.0) + psnr_metric = PeakSignalNoiseRatio(data_range=1.0) + + ssim_value = float(ssim_metric(pred_tensor, ref_tensor).item()) + psnr_value = float(psnr_metric(pred_tensor, ref_tensor).item()) + return ssim_value, psnr_value + + +def _pil_to_batched_tensor(image: Image.Image, *, compare_mode: str) -> torch.Tensor: + array = np.asarray(image.convert(compare_mode), dtype=np.float32) / 255.0 + tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0) + return tensor diff --git a/tests/e2e/accuracy/test_gebench_h100_smoke.py b/tests/e2e/accuracy/test_gebench_h100_smoke.py index b4b831871..748919269 100644 --- a/tests/e2e/accuracy/test_gebench_h100_smoke.py +++ b/tests/e2e/accuracy/test_gebench_h100_smoke.py @@ -6,8 +6,8 @@ from pathlib import Path import pytest from benchmarks.accuracy.text_to_image.gbench import main as gbench_main -from tests.e2e.accuracy.conftest import infer_model_label, reset_artifact_dir -from tests.utils import hardware_test +from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir +from tests.helpers.mark import hardware_test @pytest.mark.advanced_model diff --git a/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py b/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py index 960ea5796..6227d6368 100644 --- a/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py +++ b/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py @@ -7,8 +7,8 @@ import pytest from benchmarks.accuracy.image_to_image.gedit_bench import GROUPS from benchmarks.accuracy.image_to_image.gedit_bench import main as gedit_main -from tests.e2e.accuracy.conftest import infer_model_label, reset_artifact_dir -from tests.utils import hardware_test +from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir +from tests.helpers.mark import hardware_test @pytest.mark.advanced_model diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py new file mode 100644 index 000000000..8922d9d10 --- /dev/null +++ b/tests/e2e/accuracy/test_qwen_image.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +import base64 +import gc +import io +import os +from pathlib import Path + +import pytest +import requests +import torch +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from PIL import Image + +from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir +from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServer + +MODEL_ID = "Qwen/Qwen-Image" +MODEL_ENV_VAR = "QWEN_IMAGE_MODEL" +PROMPT = "A photo of a cat sitting on a laptop keyboard, digital art style." +NEGATIVE_PROMPT = "blurry, low quality" +WIDTH = 512 +HEIGHT = 512 +NUM_INFERENCE_STEPS = 20 +TRUE_CFG_SCALE = 4.0 +SEED = 42 +SSIM_THRESHOLD = 0.97 +PSNR_THRESHOLD = 30.0 + + +def _model_name() -> str: + return os.environ.get(MODEL_ENV_VAR, MODEL_ID) + + +def _local_files_only(model: str) -> bool: + return Path(model).exists() + + +def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image: + server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"] + with OmniServer(model, server_args, use_omni=True) as omni_server: + response = requests.post( + f"http://{omni_server.host}:{omni_server.port}/v1/images/generations", + json={ + "model": omni_server.model, + "prompt": PROMPT, + "size": f"{WIDTH}x{HEIGHT}", + "n": 1, + "response_format": "b64_json", + "negative_prompt": NEGATIVE_PROMPT, + "num_inference_steps": NUM_INFERENCE_STEPS, + "true_cfg_scale": TRUE_CFG_SCALE, + "seed": SEED, + }, + timeout=600, + ) + response.raise_for_status() + payload = response.json() + assert len(payload["data"]) == 1 + image_bytes = base64.b64decode(payload["data"][0]["b64_json"]) + image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + image.load() + image.save(output_path) + return image + + +def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image: + run_pre_test_cleanup(enable_force=True) + pipe: DiffusionPipeline | None = None + try: + pipe = DiffusionPipeline.from_pretrained( + model, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + local_files_only=_local_files_only(model), + ).to("cuda") + generator = torch.Generator(device="cuda").manual_seed(SEED) + result = pipe( # pyright: ignore[reportCallIssue] + prompt=PROMPT, + negative_prompt=NEGATIVE_PROMPT, + width=WIDTH, + height=HEIGHT, + num_inference_steps=NUM_INFERENCE_STEPS, + true_cfg_scale=TRUE_CFG_SCALE, + generator=generator, + ) + output_image = result.images[0].convert("RGB") + output_image.save(output_path) + return output_image + finally: + if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"): + pipe.maybe_free_model_hooks() + del pipe + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + run_post_test_cleanup(enable_force=True) + + +@pytest.mark.advanced_model +@pytest.mark.benchmark +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None: + model = _model_name() + output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) + + vllm_output = _run_vllm_omni_qwen_image(model=model, output_path=output_dir / "vllm_omni.png") + diffusers_output = _run_diffusers_qwen_image(model=model, output_path=output_dir / "diffusers.png") + + assert_similarity( + model_name=MODEL_ID, + vllm_image=vllm_output, + diffusers_image=diffusers_output, + width=WIDTH, + height=HEIGHT, + ssim_threshold=SSIM_THRESHOLD, + psnr_threshold=PSNR_THRESHOLD, + ) diff --git a/tests/e2e/accuracy/test_qwen_image_edit.py b/tests/e2e/accuracy/test_qwen_image_edit.py new file mode 100644 index 000000000..e17aca6e9 --- /dev/null +++ b/tests/e2e/accuracy/test_qwen_image_edit.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +import gc +from pathlib import Path + +import pytest +import requests +import torch +from diffusers import QwenImageEditPipeline, QwenImageEditPlusPipeline +from PIL import Image + +from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes +from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir +from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServer + +SINGLE_MODEL = "Qwen/Qwen-Image-Edit" +MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2509" +WIDTH = 512 +HEIGHT = 512 +NUM_INFERENCE_STEPS = 20 +TRUE_CFG_SCALE = 4.0 +SEED = 42 +SSIM_THRESHOLD = 0.94 +PSNR_THRESHOLD = 28.0 + +PROMPT_SINGLE_IMAGE = "The input is a 2D cartoon bear mascot. Restyle it into a painterly oil artwork with warm colors while preserving the main structure." +PROMPT_MULTIPLE_IMAGE = "Put the cartoon bear mascot and the furry rabbit into one coherent scene with a painterly oil artwork style and consistent lighting." +NEGATIVE_PROMPT = "low quality, blurry, artifacts, distortion" +SERVER_ARGS = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"] + + +def _run_vllm_omni_image_edit( + *, + omni_server: OmniServer, + prompt: str, + input_images: list[Image.Image], + output_path: Path, +) -> Image.Image: + response = requests.post( + f"http://{omni_server.host}:{omni_server.port}/v1/images/edits", + data={ + "model": omni_server.model, + "prompt": prompt, + "size": f"{WIDTH}x{HEIGHT}", + "n": 1, + "response_format": "b64_json", + "negative_prompt": NEGATIVE_PROMPT, + "num_inference_steps": NUM_INFERENCE_STEPS, + "true_cfg_scale": TRUE_CFG_SCALE, + "seed": SEED, + }, + files=[ + ("image", (f"image_{index}.png", pil_to_png_bytes(image), "image/png")) + for index, image in enumerate(input_images) + ], + timeout=600, + ) + response.raise_for_status() + payload = response.json() + assert len(payload["data"]) == 1 + image = decode_base64_image(payload["data"][0]["b64_json"]) + image.load() + image.save(output_path) + return image + + +def _run_diffusers_image_edit( + *, + model: str, + pipeline_class: type[QwenImageEditPipeline] | type[QwenImageEditPlusPipeline], + prompt: str, + input_images: list[Image.Image], + output_path: Path, +) -> Image.Image: + run_pre_test_cleanup(enable_force=True) + pipe: QwenImageEditPipeline | QwenImageEditPlusPipeline | None = None + device = torch.device("cuda:0") + torch.cuda.set_device(device) + try: + images = input_images[0] if len(input_images) == 1 else input_images + pipe = pipeline_class.from_pretrained( + model, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + ).to(device) + pipe.set_progress_bar_config(disable=False) + generator = torch.Generator(device=device).manual_seed(SEED) + result = pipe( # pyright: ignore[reportCallIssue] + prompt=prompt, + image=images, + negative_prompt=NEGATIVE_PROMPT, + num_inference_steps=NUM_INFERENCE_STEPS, + true_cfg_scale=TRUE_CFG_SCALE, + width=WIDTH, + height=HEIGHT, + generator=generator, + ) + output_image = result.images[0].convert("RGB") # pyright: ignore[reportAttributeAccessIssue] + output_image.save(output_path) + return output_image + finally: + if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"): + pipe.maybe_free_model_hooks() + del pipe + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + run_post_test_cleanup(enable_force=True) + + +def _vllm_omni_output_single_image( + accuracy_artifact_root: Path, + qwen_bear_image: Image.Image, +) -> Image.Image: + output_dir = model_output_dir(accuracy_artifact_root, SINGLE_MODEL) + output_path = output_dir / "vllm_omni_single.png" + with OmniServer(model=SINGLE_MODEL, serve_args=SERVER_ARGS) as server: + output = _run_vllm_omni_image_edit( + omni_server=server, + prompt=PROMPT_SINGLE_IMAGE, + input_images=[qwen_bear_image], + output_path=output_path, + ) + return output + + +def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image: Image.Image) -> Image.Image: + output_dir = model_output_dir(accuracy_artifact_root, SINGLE_MODEL) + output_path = output_dir / "diffusers_single.png" + return _run_diffusers_image_edit( + model=SINGLE_MODEL, + pipeline_class=QwenImageEditPipeline, + prompt=PROMPT_SINGLE_IMAGE, + input_images=[qwen_bear_image], + output_path=output_path, + ) + + +def _vllm_omni_output_multiple_image( + accuracy_artifact_root: Path, + qwen_bear_image: Image.Image, + rabbit_image: Image.Image, +) -> Image.Image: + output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL) + output_path = output_dir / "vllm_omni_multiple.png" + with OmniServer(model=MULTIPLE_MODEL, serve_args=SERVER_ARGS) as server: + output = _run_vllm_omni_image_edit( + omni_server=server, + prompt=PROMPT_MULTIPLE_IMAGE, + input_images=[qwen_bear_image, rabbit_image], + output_path=output_path, + ) + return output + + +def _diffusers_output_multiple_image( + accuracy_artifact_root: Path, qwen_bear_image: Image.Image, rabbit_image: Image.Image +) -> Image.Image: + output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL) + output_path = output_dir / "diffusers_multiple.png" + return _run_diffusers_image_edit( + model=MULTIPLE_MODEL, + pipeline_class=QwenImageEditPlusPipeline, + prompt=PROMPT_MULTIPLE_IMAGE, + input_images=[qwen_bear_image, rabbit_image], + output_path=output_path, + ) + + +@pytest.mark.advanced_model +@pytest.mark.benchmark +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_qwen_image_edit_single_matches_diffusers( + accuracy_artifact_root: Path, + qwen_bear_image: Image.Image, +) -> None: + vllm_image = _vllm_omni_output_single_image( + accuracy_artifact_root=accuracy_artifact_root, + qwen_bear_image=qwen_bear_image, + ) + diffusers_image = _diffusers_output_single_image( + accuracy_artifact_root=accuracy_artifact_root, + qwen_bear_image=qwen_bear_image, + ) + assert_similarity( + model_name=SINGLE_MODEL, + vllm_image=vllm_image, + diffusers_image=diffusers_image, + width=WIDTH, + height=HEIGHT, + ssim_threshold=SSIM_THRESHOLD, + psnr_threshold=PSNR_THRESHOLD, + ) + + +@pytest.mark.advanced_model +@pytest.mark.benchmark +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.skip( + reason="Skipping as the second image seems to be ignored by the API. Will come back to this later after #2772 is merged." +) +def test_qwen_image_edit_multiple_matches_diffusers( + accuracy_artifact_root: Path, + qwen_bear_image: Image.Image, + rabbit_image: Image.Image, +) -> None: + vllm_image = _vllm_omni_output_multiple_image( + accuracy_artifact_root=accuracy_artifact_root, + qwen_bear_image=qwen_bear_image, + rabbit_image=rabbit_image, + ) + diffusers_image = _diffusers_output_multiple_image( + accuracy_artifact_root=accuracy_artifact_root, + qwen_bear_image=qwen_bear_image, + rabbit_image=rabbit_image, + ) + assert_similarity( + model_name=MULTIPLE_MODEL, + vllm_image=vllm_image, + diffusers_image=diffusers_image, + width=WIDTH, + height=HEIGHT, + ssim_threshold=SSIM_THRESHOLD, + psnr_threshold=PSNR_THRESHOLD, + ) diff --git a/tests/e2e/accuracy/test_qwen_image_layered.py b/tests/e2e/accuracy/test_qwen_image_layered.py new file mode 100644 index 000000000..0ab9cb323 --- /dev/null +++ b/tests/e2e/accuracy/test_qwen_image_layered.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import base64 +import gc +import io +import os +from pathlib import Path + +import pytest +import requests +import torch +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from PIL import Image + +from tests.e2e.accuracy.helpers import assert_image_sequence_similarity, model_output_dir +from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServer + +MODEL_ID = "Qwen/Qwen-Image-Layered" +MODEL_ENV_VAR = "QWEN_IMAGE_LAYERED_MODEL" +PROMPT = "decompose into layers" +NEGATIVE_PROMPT = " " +NUM_INFERENCE_STEPS = 20 +TRUE_CFG_SCALE = 4.0 +SEED = 777 +LAYERS = 3 +RESOLUTION = 640 +SSIM_THRESHOLD = 0.97 +PSNR_THRESHOLD = 30.0 + + +def _model_name() -> str: + return os.environ.get(MODEL_ENV_VAR, MODEL_ID) + + +def _local_files_only(model: str) -> bool: + return Path(model).exists() + + +def _normalize_layered_images(images: object) -> list[Image.Image]: + if not isinstance(images, list) or not images: + raise AssertionError(f"Unexpected layered output container: {type(images).__name__}") + + first_item = images[0] + if isinstance(first_item, Image.Image): + return [image.convert("RGBA") for image in images if isinstance(image, Image.Image)] + if isinstance(first_item, (list, tuple)): + return [image.convert("RGBA") for image in first_item if isinstance(image, Image.Image)] + raise AssertionError(f"Unexpected layered image element type: {type(first_item).__name__}") + + +def _run_vllm_omni_qwen_image_layered(*, model: str, input_image: Image.Image, output_dir: Path) -> list[Image.Image]: + input_image.save(output_dir / "input.png") + server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"] + with OmniServer(model, server_args, use_omni=True) as omni_server: + buffer = io.BytesIO() + input_image.save(buffer, format="PNG") + buffer.seek(0) + response = requests.post( + f"http://{omni_server.host}:{omni_server.port}/v1/images/edits", + data={ + "model": omni_server.model, + "prompt": PROMPT, + "size": "auto", + "n": 1, + "response_format": "b64_json", + "negative_prompt": NEGATIVE_PROMPT, + "num_inference_steps": NUM_INFERENCE_STEPS, + "true_cfg_scale": TRUE_CFG_SCALE, + "seed": SEED, + "layers": LAYERS, + "resolution": RESOLUTION, + }, + files=[("image", ("input.png", buffer, "image/png"))], + timeout=600, + ) + response.raise_for_status() + payload = response.json() + assert len(payload["data"]) == LAYERS + output_images = [] + for item in payload["data"]: + image_bytes = base64.b64decode(item["b64_json"]) + image = Image.open(io.BytesIO(image_bytes)).convert("RGBA") + image.load() + output_images.append(image) + for index, image in enumerate(output_images, start=1): + image.save(output_dir / f"vllm_omni_layer_{index}.png") + return output_images + + +def _run_diffusers_qwen_image_layered(*, model: str, input_image: Image.Image, output_dir: Path) -> list[Image.Image]: + run_pre_test_cleanup(enable_force=True) + pipe: DiffusionPipeline | None = None + try: + pipe = DiffusionPipeline.from_pretrained( + model, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + local_files_only=_local_files_only(model), + ).to("cuda") + generator = torch.Generator(device="cuda").manual_seed(SEED) + result = pipe( # pyright: ignore[reportCallIssue] + image=input_image, + prompt=PROMPT, + negative_prompt=NEGATIVE_PROMPT, + num_inference_steps=NUM_INFERENCE_STEPS, + true_cfg_scale=TRUE_CFG_SCALE, + generator=generator, + num_images_per_prompt=1, + layers=LAYERS, + resolution=RESOLUTION, + ) + output_images = _normalize_layered_images(result.images) + assert len(output_images) == LAYERS, f"Expected {LAYERS} diffusers layers, got {len(output_images)}" + for index, image in enumerate(output_images, start=1): + image.save(output_dir / f"diffusers_layer_{index}.png") + return output_images + finally: + if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"): + pipe.maybe_free_model_hooks() + del pipe + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + run_post_test_cleanup(enable_force=True) + + +@pytest.mark.advanced_model +@pytest.mark.benchmark +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_qwen_image_layered_matches_diffusers(accuracy_artifact_root: Path, qwen_bear_image: Image.Image) -> None: + model = _model_name() + output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) + input_image = qwen_bear_image.convert("RGBA") + + vllm_outputs = _run_vllm_omni_qwen_image_layered(model=model, input_image=input_image, output_dir=output_dir) + diffusers_outputs = _run_diffusers_qwen_image_layered(model=model, input_image=input_image, output_dir=output_dir) + + assert_image_sequence_similarity( + model_name=MODEL_ID, + vllm_images=vllm_outputs, + diffusers_images=diffusers_outputs, + ssim_threshold=SSIM_THRESHOLD, + psnr_threshold=PSNR_THRESHOLD, + compare_mode="RGBA", + ) diff --git a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py index 3cdda1f9f..3aa5da85c 100644 --- a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py +++ b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py @@ -22,7 +22,6 @@ import torch from diffusers import UniPCMultistepScheduler from PIL import Image -from tests.conftest import OmniServerParams from tests.e2e.accuracy.wan22_i2v.run_wan22_i2v_diffusers_cp import ( _configure_scheduler, _ensure_wan_ftfy_fallback, @@ -48,7 +47,8 @@ from tests.e2e.accuracy.wan22_i2v.wan22_i2v_video_similarity_common import ( SSIM_THRESHOLD, WIDTH, ) -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams def test_parse_video_metadata_extracts_dimensions_and_fps() -> None: diff --git a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py b/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py index 57743d62b..bd3f2e099 100644 --- a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py +++ b/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py @@ -26,7 +26,7 @@ from contextlib import ExitStack import pytest -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput diff --git a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py b/tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py index 03bd12efa..0681687fe 100644 --- a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py +++ b/tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py @@ -19,7 +19,7 @@ import numpy as np import pytest from transformers import AutoTokenizer -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput diff --git a/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py b/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py index ffbe703ca..653b35d7e 100644 --- a/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py +++ b/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py @@ -10,7 +10,7 @@ import pytest from tests.e2e.offline_inference.custom_pipeline.worker_extension import ( vLLMOmniColocateWorkerExtensionForTest, ) -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension from vllm_omni.entrypoints.async_omni import AsyncOmni diff --git a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml index 1f0d06cb8..b7768c071 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml @@ -64,9 +64,6 @@ stage_args: # Top-level runtime config with Mooncake connector runtime: enabled: true - defaults: - window_size: -1 - max_inflight: 1 connectors: mooncake_connector: name: MooncakeConnector @@ -80,4 +77,3 @@ runtime: edges: - from: 0 to: 1 - window_size: -1 diff --git a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml index 36b1d2bbe..504f3c98e 100644 --- a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml @@ -62,10 +62,6 @@ stage_args: # Runtime edges runtime: enabled: true - defaults: - window_size: -1 - max_inflight: 1 - # Distributed connectors configuration (optional) # More connectors will be supported in the future. connectors: @@ -78,4 +74,3 @@ runtime: edges: - from: 0 to: 1 - window_size: -1 diff --git a/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml deleted file mode 100644 index f93a6c714..000000000 --- a/tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml +++ /dev/null @@ -1,103 +0,0 @@ -# stage config for running qwen2.5-omni for multi-stage omni runtime. - -# This config is optimized for CI e2e tests. -stage_args: - - stage_id: 0 - runtime: - process: true # Run this stage in a separate process - devices: "0" - engine_args: - model_stage: thinker - model_arch: Qwen2_5OmniForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - max_model_len: 896 - max_num_batched_tokens: 896 - max_num_seqs: 1 - gpu_memory_utilization: 0.8 - skip_mm_profiling: true - enforce_eager: true # Now we only support eager mode - trust_remote_code: true - engine_output_type: latent - enable_prefix_caching: false - mm_processor_cache_gb: 0 - 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: 128 - seed: 42 - detokenize: True - repetition_penalty: 1.1 - - stage_id: 1 - runtime: - process: true - devices: "1" - engine_args: - model_stage: talker - model_arch: Qwen2_5OmniForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - max_model_len: 896 - max_num_batched_tokens: 896 - max_num_seqs: 1 - gpu_memory_utilization: 0.8 - skip_mm_profiling: true - 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: 128 - 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: 128 - 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 diff --git a/tests/e2e/offline_inference/test_bagel_img2img.py b/tests/e2e/offline_inference/test_bagel_img2img.py index 63d2a37da..66aec80c7 100644 --- a/tests/e2e/offline_inference/test_bagel_img2img.py +++ b/tests/e2e/offline_inference/test_bagel_img2img.py @@ -22,9 +22,10 @@ import pytest from PIL import Image from vllm.assets.image import ImageAsset -from tests.conftest import OmniRunner, modify_stage_config -from tests.utils import hardware_test -from vllm_omni import Omni +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import modify_stage_config +from vllm_omni.entrypoints.omni import Omni from vllm_omni.platforms import current_omni_platform # Reference pixel data extracted from the known-good output image @@ -32,30 +33,30 @@ from vllm_omni.platforms import current_omni_platform # prompt='Change the grass color to red', # input image: 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg REFERENCE_PIXELS = [ - {"position": (100, 100), "rgb": (157, 172, 217)}, - {"position": (400, 50), "rgb": (105, 144, 218)}, - {"position": (700, 100), "rgb": (118, 159, 233)}, - {"position": (150, 400), "rgb": (195, 34, 60)}, - {"position": (512, 336), "rgb": (222, 214, 193)}, - {"position": (700, 400), "rgb": (197, 15, 43)}, - {"position": (100, 600), "rgb": (105, 13, 18)}, - {"position": (400, 600), "rgb": (169, 33, 44)}, - {"position": (700, 600), "rgb": (101, 86, 93)}, - {"position": (256, 256), "rgb": (181, 202, 222)}, + {"position": (100, 100), "rgb": (156, 172, 217)}, + {"position": (400, 50), "rgb": (105, 144, 217)}, + {"position": (700, 100), "rgb": (118, 159, 232)}, + {"position": (150, 400), "rgb": (180, 22, 52)}, + {"position": (512, 336), "rgb": (221, 211, 194)}, + {"position": (700, 400), "rgb": (192, 10, 46)}, + {"position": (100, 600), "rgb": (102, 12, 22)}, + {"position": (400, 600), "rgb": (161, 28, 47)}, + {"position": (700, 600), "rgb": (100, 87, 94)}, + {"position": (256, 256), "rgb": (181, 201, 221)}, ] if current_omni_platform.is_rocm(): REFERENCE_PIXELS = [ - {"position": (100, 100), "rgb": (156, 172, 215)}, - {"position": (400, 50), "rgb": (106, 144, 216)}, - {"position": (700, 100), "rgb": (118, 158, 231)}, - {"position": (150, 400), "rgb": (183, 23, 48)}, - {"position": (512, 336), "rgb": (218, 215, 191)}, - {"position": (700, 400), "rgb": (194, 14, 42)}, - {"position": (100, 600), "rgb": (105, 10, 16)}, - {"position": (400, 600), "rgb": (167, 33, 46)}, - {"position": (700, 600), "rgb": (102, 86, 92)}, - {"position": (256, 256), "rgb": (181, 201, 220)}, + {"position": (100, 100), "rgb": (156, 172, 217)}, + {"position": (400, 50), "rgb": (105, 144, 217)}, + {"position": (700, 100), "rgb": (118, 159, 232)}, + {"position": (150, 400), "rgb": (180, 22, 52)}, + {"position": (512, 336), "rgb": (221, 211, 194)}, + {"position": (700, 400), "rgb": (192, 10, 46)}, + {"position": (100, 600), "rgb": (102, 12, 22)}, + {"position": (400, 600), "rgb": (161, 28, 47)}, + {"position": (700, 600), "rgb": (100, 87, 94)}, + {"position": (256, 256), "rgb": (181, 201, 221)}, ] PIXEL_TOLERANCE = 10 diff --git a/tests/e2e/offline_inference/test_bagel_lora.py b/tests/e2e/offline_inference/test_bagel_lora.py index 501d23eaa..75f41f9be 100644 --- a/tests/e2e/offline_inference/test_bagel_lora.py +++ b/tests/e2e/offline_inference/test_bagel_lora.py @@ -31,9 +31,10 @@ import torch from PIL import Image from safetensors.torch import save_file -from tests.conftest import OmniRunner, modify_stage_config -from tests.utils import hardware_test -from vllm_omni import Omni +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import modify_stage_config +from vllm_omni.entrypoints.omni import Omni from vllm_omni.lora.request import LoRARequest from vllm_omni.lora.utils import stable_lora_int_id diff --git a/tests/e2e/offline_inference/test_bagel_text2img.py b/tests/e2e/offline_inference/test_bagel_text2img.py index e45d64f2a..0819f103a 100644 --- a/tests/e2e/offline_inference/test_bagel_text2img.py +++ b/tests/e2e/offline_inference/test_bagel_text2img.py @@ -27,9 +27,10 @@ from typing import Any import pytest from PIL import Image -from tests.conftest import OmniRunner, modify_stage_config -from tests.utils import hardware_test -from vllm_omni import Omni +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import modify_stage_config +from vllm_omni.entrypoints.omni import Omni from vllm_omni.platforms import current_omni_platform # Reference pixel data extracted from the known-good output image @@ -37,30 +38,30 @@ from vllm_omni.platforms import current_omni_platform # "Generated with seed=52, num_inference_steps=15, # prompt='A futuristic city skyline at twilight, cyberpunk style'" REFERENCE_PIXELS = [ - {"position": (100, 100), "rgb": (121, 118, 100)}, - {"position": (400, 50), "rgb": (163, 162, 143)}, - {"position": (700, 100), "rgb": (170, 156, 127)}, - {"position": (150, 400), "rgb": (129, 127, 112)}, - {"position": (512, 512), "rgb": (135, 61, 59)}, - {"position": (700, 400), "rgb": (205, 107, 43)}, - {"position": (100, 700), "rgb": (197, 177, 157)}, - {"position": (400, 700), "rgb": (139, 107, 86)}, - {"position": (700, 700), "rgb": (247, 205, 146)}, - {"position": (256, 256), "rgb": (171, 160, 153)}, + {"position": (100, 100), "rgb": (115, 113, 94)}, + {"position": (400, 50), "rgb": (159, 160, 144)}, + {"position": (700, 100), "rgb": (164, 151, 123)}, + {"position": (150, 400), "rgb": (120, 121, 107)}, + {"position": (512, 512), "rgb": (165, 133, 127)}, + {"position": (700, 400), "rgb": (217, 130, 66)}, + {"position": (100, 700), "rgb": (191, 168, 152)}, + {"position": (400, 700), "rgb": (130, 96, 77)}, + {"position": (700, 700), "rgb": (247, 203, 140)}, + {"position": (256, 256), "rgb": (167, 156, 150)}, ] if current_omni_platform.is_rocm(): REFERENCE_PIXELS = [ - {"position": (100, 100), "rgb": (123, 119, 100)}, - {"position": (400, 50), "rgb": (162, 161, 142)}, - {"position": (700, 100), "rgb": (171, 156, 127)}, - {"position": (150, 400), "rgb": (131, 128, 112)}, - {"position": (512, 512), "rgb": (134, 61, 59)}, - {"position": (700, 400), "rgb": (204, 107, 43)}, - {"position": (100, 700), "rgb": (201, 180, 165)}, - {"position": (400, 700), "rgb": (140, 108, 87)}, - {"position": (700, 700), "rgb": (247, 205, 145)}, - {"position": (256, 256), "rgb": (171, 160, 153)}, + {"position": (100, 100), "rgb": (115, 113, 94)}, + {"position": (400, 50), "rgb": (159, 160, 144)}, + {"position": (700, 100), "rgb": (164, 151, 123)}, + {"position": (150, 400), "rgb": (120, 121, 107)}, + {"position": (512, 512), "rgb": (165, 133, 127)}, + {"position": (700, 400), "rgb": (217, 130, 66)}, + {"position": (100, 700), "rgb": (191, 168, 152)}, + {"position": (400, 700), "rgb": (130, 96, 77)}, + {"position": (700, 700), "rgb": (247, 203, 140)}, + {"position": (256, 256), "rgb": (167, 156, 150)}, ] # Maximum allowed difference per color channel diff --git a/tests/e2e/offline_inference/test_bagel_understanding.py b/tests/e2e/offline_inference/test_bagel_understanding.py index bbee32980..c3ed97b42 100644 --- a/tests/e2e/offline_inference/test_bagel_understanding.py +++ b/tests/e2e/offline_inference/test_bagel_understanding.py @@ -26,8 +26,9 @@ from pathlib import Path import pytest from vllm.assets.image import ImageAsset -from tests.conftest import OmniRunner, modify_stage_config -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import modify_stage_config MODEL_NAME = "ByteDance-Seed/BAGEL-7B-MoT" STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml") diff --git a/tests/e2e/offline_inference/test_cache_dit.py b/tests/e2e/offline_inference/test_cache_dit.py index fc08da7be..1577dd9f6 100644 --- a/tests/e2e/offline_inference/test_cache_dit.py +++ b/tests/e2e/offline_inference/test_cache_dit.py @@ -11,8 +11,8 @@ It uses minimal settings to keep test time short for CI. import pytest import torch -from tests.conftest import OmniRunner -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform diff --git a/tests/e2e/offline_inference/test_cosyvoice3.py b/tests/e2e/offline_inference/test_cosyvoice3.py index 8c88d972d..7206f1e7b 100644 --- a/tests/e2e/offline_inference/test_cosyvoice3.py +++ b/tests/e2e/offline_inference/test_cosyvoice3.py @@ -15,19 +15,18 @@ from __future__ import annotations import functools import io import os -import tempfile from pathlib import Path from urllib.request import urlopen import numpy as np import pytest import soundfile as sf -import yaml from huggingface_hub import snapshot_download from vllm.sampling_params import SamplingParams -from tests.conftest import OmniRunner -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer @@ -44,13 +43,9 @@ REFERENCE_STAGE0_TEMPERATURE = 1.0 REFERENCE_STAGE0_REPETITION_PENALTY = 2.0 -def _stage_config(name: str) -> str: - return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name) - - -STAGE_CONFIGS = [ - _stage_config("cosyvoice3.yaml"), - _stage_config("cosyvoice3_async_chunk.yaml"), +ASYNC_CHUNK_MODES = [ + pytest.param(False, id="sync"), + pytest.param(True, id="async_chunk"), ] @@ -138,19 +133,6 @@ def _get_stage_engine_outputs(omni_runner: OmniRunner, stage_id: int): raise AttributeError("Unable to locate stage outputs on Omni runner") -def _patched_stage_config(base_stage_config: str, model_dir: Path, tmp_dir: Path) -> str: - cfg = yaml.safe_load(Path(base_stage_config).read_text(encoding="utf-8")) - tokenizer_path = str(model_dir / "CosyVoice-BlankEN") - for stage in cfg.get("stage_args", []): - engine_args = stage.setdefault("engine_args", {}) - engine_args["tokenizer"] = tokenizer_path - engine_args["enforce_eager"] = True - engine_args["hf_overrides"] = {"architectures": ["CosyVoice3Model"]} - out_path = tmp_dir / Path(base_stage_config).name - out_path.write_text(yaml.safe_dump(cfg, sort_keys=False), encoding="utf-8") - return str(out_path) - - def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[str, object]]: return [ { @@ -165,54 +147,54 @@ def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[s @pytest.mark.core_model @pytest.mark.omni @hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.parametrize("base_stage_config", STAGE_CONFIGS) -def test_cosyvoice3_offline_reference_zero_shot(base_stage_config: str) -> None: +@pytest.mark.parametrize("async_chunk", ASYNC_CHUNK_MODES) +def test_cosyvoice3_offline_reference_zero_shot(async_chunk: bool) -> None: """CosyVoice3 zero-shot reference inference should stop cleanly and produce sane audio.""" prompt_audio, prompt_sr = _load_reference_prompt_wav() model_dir = _resolve_model_dir() expected_stop_token = int(CosyVoice3Config().llm["eos_token_id"]) - with tempfile.TemporaryDirectory(prefix="cv3-e2e-") as tmp: - stage_config = _patched_stage_config(base_stage_config, model_dir, Path(tmp)) - with OmniRunner( - str(model_dir), seed=42, stage_configs_path=stage_config, stage_init_timeout=300 - ) as omni_runner: - sampling_params_list = omni_runner.get_default_sampling_params_list() - sampling_params_list[0] = _reference_zero_shot_stage0_sampling(text=REFERENCE_SYNTH_TEXT) + with OmniRunner( + str(model_dir), + seed=42, + stage_configs_path=get_deploy_config_path("cosyvoice3.yaml"), + async_chunk=async_chunk, + stage_init_timeout=300, + ) as omni_runner: + sampling_params_list = omni_runner.get_default_sampling_params_list() + sampling_params_list[0] = _reference_zero_shot_stage0_sampling(text=REFERENCE_SYNTH_TEXT) - outputs = omni_runner.omni.generate( - _build_reference_inputs((prompt_audio, prompt_sr)), sampling_params_list + outputs = omni_runner.omni.generate(_build_reference_inputs((prompt_audio, prompt_sr)), sampling_params_list) + + assert outputs, "No outputs returned" + audio_mm = outputs[0].multimodal_output + assert "audio" in audio_mm, "No audio output found" + + audio = _concat_audio(audio_mm["audio"]) + assert audio.size > 0, "Generated audio is empty" + + sr_val = audio_mm.get("sr", 24000) + if isinstance(sr_val, list) and sr_val: + sr_val = sr_val[-1] + if hasattr(sr_val, "item"): + sr_val = sr_val.item() + sr = int(sr_val) + assert sr == 24000, f"Unexpected sample_rate={sr}" + + duration_s = audio.size / sr + assert 2.8 <= duration_s <= 8.8, f"Unexpected duration={duration_s:.3f}s (samples={audio.size}, sr={sr})" + + stage0_outputs = _get_stage_engine_outputs(omni_runner, 0) + if stage0_outputs: + completion = stage0_outputs[0].outputs[0] + finish_reason = getattr(completion, "finish_reason", None) + stop_reason = getattr(completion, "stop_reason", None) + num_tokens = len(getattr(completion, "token_ids", []) or []) + + assert finish_reason == "stop", f"Stage-0 finish_reason={finish_reason}, expected 'stop'" + assert int(stop_reason) == expected_stop_token, ( + f"Stage-0 stop_reason={stop_reason}, expected {expected_stop_token}" ) - - assert outputs, "No outputs returned" - audio_mm = outputs[0].multimodal_output - assert "audio" in audio_mm, "No audio output found" - - audio = _concat_audio(audio_mm["audio"]) - assert audio.size > 0, "Generated audio is empty" - - sr_val = audio_mm.get("sr", 24000) - if isinstance(sr_val, list) and sr_val: - sr_val = sr_val[-1] - if hasattr(sr_val, "item"): - sr_val = sr_val.item() - sr = int(sr_val) - assert sr == 24000, f"Unexpected sample_rate={sr}" - - duration_s = audio.size / sr - assert 2.8 <= duration_s <= 8.8, f"Unexpected duration={duration_s:.3f}s (samples={audio.size}, sr={sr})" - - stage0_outputs = _get_stage_engine_outputs(omni_runner, 0) - if stage0_outputs: - completion = stage0_outputs[0].outputs[0] - finish_reason = getattr(completion, "finish_reason", None) - stop_reason = getattr(completion, "stop_reason", None) - num_tokens = len(getattr(completion, "token_ids", []) or []) - - assert finish_reason == "stop", f"Stage-0 finish_reason={finish_reason}, expected 'stop'" - assert int(stop_reason) == expected_stop_token, ( - f"Stage-0 stop_reason={stop_reason}, expected {expected_stop_token}" - ) - assert 80 <= num_tokens <= 220, f"Stage-0 num_tokens={num_tokens}, expected sane stop-bound range" - else: - assert "async_chunk" in Path(base_stage_config).name, "Stage-0 produced no engine outputs" + assert 80 <= num_tokens <= 220, f"Stage-0 num_tokens={num_tokens}, expected sane stop-bound range" + else: + assert async_chunk, "Stage-0 produced no engine outputs" diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py index 257755ef8..d7fd6f72f 100644 --- a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py +++ b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py @@ -4,8 +4,9 @@ import pytest import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from tests.conftest import OmniRunner -from tests.utils import DeviceMemoryMonitor, hardware_test +from tests.helpers.env import DeviceMemoryMonitor +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform diff --git a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py b/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py index bdfd594c7..4f19c1004 100644 --- a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py +++ b/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py @@ -2,8 +2,8 @@ import pytest import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from tests.conftest import OmniRunner -from tests.utils import DeviceMemoryMonitor +from tests.helpers.env import DeviceMemoryMonitor +from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform diff --git a/tests/e2e/offline_inference/test_diffusion_lora.py b/tests/e2e/offline_inference/test_diffusion_lora.py index 7edd03f20..027dadb3f 100644 --- a/tests/e2e/offline_inference/test_diffusion_lora.py +++ b/tests/e2e/offline_inference/test_diffusion_lora.py @@ -7,7 +7,7 @@ import pytest import torch from safetensors.torch import save_file -from tests.conftest import OmniRunner +from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform diff --git a/tests/e2e/offline_inference/test_dynin_omni.py b/tests/e2e/offline_inference/test_dynin_omni.py index 5388ac674..f891fc4f1 100644 --- a/tests/e2e/offline_inference/test_dynin_omni.py +++ b/tests/e2e/offline_inference/test_dynin_omni.py @@ -18,7 +18,7 @@ import pytest import torch from transformers import AutoTokenizer -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" diff --git a/tests/e2e/offline_inference/test_expert_parallel.py b/tests/e2e/offline_inference/test_expert_parallel.py index 29d84d7a3..f11646b30 100644 --- a/tests/e2e/offline_inference/test_expert_parallel.py +++ b/tests/e2e/offline_inference/test_expert_parallel.py @@ -18,8 +18,8 @@ import torch import torch.distributed as dist from PIL import Image -from tests.conftest import OmniRunner -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform diff --git a/tests/e2e/offline_inference/test_flux.py b/tests/e2e/offline_inference/test_flux.py new file mode 100644 index 000000000..02c6787be --- /dev/null +++ b/tests/e2e/offline_inference/test_flux.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Tests for Flux1 Schnell.""" + +import pytest +from PIL import Image + +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +MODEL = "black-forest-labs/FLUX.1-schnell" + + +@pytest.mark.core_model +@pytest.mark.diffusion +def test_flux_schnell_text_to_image(): + """Test FLUX.1-schnell text-to-image generation.""" + omni = Omni(model=MODEL) + + omni_outputs = list( + omni.generate( + prompts=["A photo of a cat sitting on a laptop"], + sampling_params_list=OmniDiffusionSamplingParams( + height=512, + width=512, + num_inference_steps=2, + seed=42, + ), + ) + ) + + assert len(omni_outputs) > 0 + images = omni_outputs[0].images + assert len(images) == 1 + assert isinstance(images[0], Image.Image) + assert images[0].size == (512, 512) diff --git a/tests/e2e/offline_inference/test_flux2_klein_inpaint.py b/tests/e2e/offline_inference/test_flux2_klein_inpaint.py new file mode 100644 index 000000000..ac7169361 --- /dev/null +++ b/tests/e2e/offline_inference/test_flux2_klein_inpaint.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +End-to-end test for Flux2 Klein inpainting. + +""" + +# ruff: noqa: E402 + +import os +import sys +from pathlib import Path + +import pytest +import torch +from PIL import Image, ImageDraw + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.platforms import current_omni_platform + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vllm_omni import Omni + +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" + +MODEL = "black-forest-labs/FLUX.2-klein-4B" + +_HEIGHT = 512 +_WIDTH = 512 +_NUM_INFERENCE_STEPS = 4 + + +def _create_test_image(width: int = _WIDTH, height: int = _HEIGHT, color: tuple = (128, 128, 128)) -> Image.Image: + return Image.new("RGB", (width, height), color) + + +def _create_test_mask(width: int = _WIDTH, height: int = _HEIGHT) -> Image.Image: + mask = Image.new("L", (width, height), 0) + draw = ImageDraw.Draw(mask) + draw.rectangle([width // 4, height // 4, width * 3 // 4, height * 3 // 4], fill=255) + return mask + + +def _create_test_inputs(color: tuple = (100, 150, 200)): + return _create_test_image(_WIDTH, _HEIGHT, color), _create_test_mask(_WIDTH, _HEIGHT) + + +def _extract_images_from_output(outputs: list) -> list[Image.Image]: + images = [] + for req_output in outputs: + if hasattr(req_output, "images") and req_output.images: + images.extend(req_output.images) + elif hasattr(req_output, "request_output") and req_output.request_output: + stage_out = req_output.request_output + if isinstance(stage_out, OmniRequestOutput) and hasattr(stage_out, "images"): + images.extend(stage_out.images) + elif isinstance(stage_out, list): + for s in stage_out: + if hasattr(s, "images") and s.images: + images.extend(s.images) + return images + + +@pytest.mark.core_model +@pytest.mark.diffusion +def test_flux2_klein_inpaint_basic(): + m = None + try: + m = Omni(model=MODEL) + input_image, mask_image = _create_test_inputs() + + outputs = m.generate( + prompts=[ + { + "prompt": "Fill in the masked area with a beautiful garden", + "multi_modal_data": {"image": input_image, "mask_image": mask_image}, + } + ], + sampling_params_list=OmniDiffusionSamplingParams( + height=_HEIGHT, + width=_WIDTH, + num_inference_steps=_NUM_INFERENCE_STEPS, + guidance_scale=0.0, + generator=torch.Generator(current_omni_platform.device_type).manual_seed(42), + num_outputs_per_prompt=1, + ), + ) + + images = _extract_images_from_output(list(outputs)) + assert len(images) == 1 + assert images[0].size == (_WIDTH, _HEIGHT) + finally: + if m is not None and hasattr(m, "close"): + m.close() + + +@pytest.mark.diffusion +def test_flux2_klein_inpaint_deterministic(): + m = None + try: + m = Omni(model=MODEL) + input_image, mask_image = _create_test_inputs() + seed = 12345 + + gen1 = torch.Generator(current_omni_platform.device_type).manual_seed(seed) + gen2 = torch.Generator(current_omni_platform.device_type).manual_seed(seed) + + outputs1 = m.generate( + prompts=[ + { + "prompt": "A red flower in a field", + "multi_modal_data": {"image": input_image, "mask_image": mask_image}, + } + ], + sampling_params_list=OmniDiffusionSamplingParams( + height=_HEIGHT, + width=_WIDTH, + num_inference_steps=_NUM_INFERENCE_STEPS, + guidance_scale=0.0, + generator=gen1, + num_outputs_per_prompt=1, + ), + ) + + outputs2 = m.generate( + prompts=[ + { + "prompt": "A red flower in a field", + "multi_modal_data": {"image": input_image, "mask_image": mask_image}, + } + ], + sampling_params_list=OmniDiffusionSamplingParams( + height=_HEIGHT, + width=_WIDTH, + num_inference_steps=_NUM_INFERENCE_STEPS, + guidance_scale=0.0, + generator=gen2, + num_outputs_per_prompt=1, + ), + ) + + images1 = _extract_images_from_output(list(outputs1)) + images2 = _extract_images_from_output(list(outputs2)) + + assert len(images1) == 1 + assert len(images2) == 1 + + assert list(images1[0].getdata()) == list(images2[0].getdata()), ( + "Same input with same seed should produce identical output. " + "This is critical for offline/online consistency." + ) + finally: + if m is not None and hasattr(m, "close"): + m.close() + + +@pytest.mark.diffusion +def test_flux2_klein_inpaint_different_seeds_different_output(): + m = None + try: + m = Omni(model=MODEL) + input_image, mask_image = _create_test_inputs() + + gen1 = torch.Generator(current_omni_platform.device_type).manual_seed(42) + gen2 = torch.Generator(current_omni_platform.device_type).manual_seed(99999) + + outputs1 = m.generate( + prompts=[ + { + "prompt": "A beautiful landscape", + "multi_modal_data": {"image": input_image, "mask_image": mask_image}, + } + ], + sampling_params_list=OmniDiffusionSamplingParams( + height=_HEIGHT, + width=_WIDTH, + num_inference_steps=_NUM_INFERENCE_STEPS, + guidance_scale=0.0, + generator=gen1, + num_outputs_per_prompt=1, + ), + ) + + outputs2 = m.generate( + prompts=[ + { + "prompt": "A beautiful landscape", + "multi_modal_data": {"image": input_image, "mask_image": mask_image}, + } + ], + sampling_params_list=OmniDiffusionSamplingParams( + height=_HEIGHT, + width=_WIDTH, + num_inference_steps=_NUM_INFERENCE_STEPS, + guidance_scale=0.0, + generator=gen2, + num_outputs_per_prompt=1, + ), + ) + + images1 = _extract_images_from_output(list(outputs1)) + images2 = _extract_images_from_output(list(outputs2)) + + assert len(images1) == 1 + assert len(images2) == 1 + + different_pixel_count = sum(1 for p1, p2 in zip(images1[0].getdata(), images2[0].getdata()) if p1 != p2) + assert different_pixel_count > 0, "Different seeds should produce different outputs" + finally: + if m is not None and hasattr(m, "close"): + m.close() diff --git a/tests/e2e/offline_inference/test_flux_autoround_w4a16.py b/tests/e2e/offline_inference/test_flux_autoround_w4a16.py index cbcd1009d..ef5d6f9e0 100644 --- a/tests/e2e/offline_inference/test_flux_autoround_w4a16.py +++ b/tests/e2e/offline_inference/test_flux_autoround_w4a16.py @@ -14,8 +14,9 @@ import pytest import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from tests.conftest import OmniRunner -from tests.utils import DeviceMemoryMonitor, hardware_test +from tests.helpers.env import DeviceMemoryMonitor +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform diff --git a/tests/e2e/offline_inference/test_flux_kontext.py b/tests/e2e/offline_inference/test_flux_kontext.py index cd711d6b8..057319c85 100644 --- a/tests/e2e/offline_inference/test_flux_kontext.py +++ b/tests/e2e/offline_inference/test_flux_kontext.py @@ -13,7 +13,7 @@ import pytest from PIL import Image from vllm.assets.image import ImageAsset -from tests.conftest import OmniRunner +from tests.helpers.runtime import OmniRunner from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams diff --git a/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py b/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py index 79bb64dca..bd0d132d0 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py @@ -8,7 +8,7 @@ import torch.nn.functional as F from PIL import Image from transformers import CLIPModel, CLIPProcessor -from tests.conftest import OmniRunner +from tests.helpers.runtime import OmniRunner from vllm_omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -17,7 +17,7 @@ PROMPT = "A brown and white dog is running on the grass" MODEL_NAME = "tencent/HunyuanImage-3.0" LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32" REPO_ROOT = Path(__file__).resolve().parents[3] -STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image_3_moe.yaml" +STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image3_t2i.yaml" pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] diff --git a/tests/e2e/offline_inference/test_ltx2_cfg_parallel_parity.py b/tests/e2e/offline_inference/test_ltx2_cfg_parallel_parity.py index 659040929..07aa5a647 100644 --- a/tests/e2e/offline_inference/test_ltx2_cfg_parallel_parity.py +++ b/tests/e2e/offline_inference/test_ltx2_cfg_parallel_parity.py @@ -11,7 +11,7 @@ import numpy as np import pytest from PIL import Image -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test REPO_ROOT = Path(__file__).resolve().parents[3] T2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "text_to_video" / "text_to_video.py" diff --git a/tests/e2e/offline_inference/test_magi_human.py b/tests/e2e/offline_inference/test_magi_human.py index abb7f9c16..6d4614172 100644 --- a/tests/e2e/offline_inference/test_magi_human.py +++ b/tests/e2e/offline_inference/test_magi_human.py @@ -8,8 +8,8 @@ import av import numpy as np import pytest -from tests.conftest import OmniRunner -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes from vllm_omni.inputs.data import OmniDiffusionSamplingParams diff --git a/tests/e2e/offline_inference/test_mammoth_moda2.py b/tests/e2e/offline_inference/test_mammoth_moda2.py index ff744c86e..c3d95844c 100644 --- a/tests/e2e/offline_inference/test_mammoth_moda2.py +++ b/tests/e2e/offline_inference/test_mammoth_moda2.py @@ -23,8 +23,10 @@ import pytest import torch from vllm.sampling_params import SamplingParams -from tests.conftest import OmniRunner -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner + +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" # --------------------------------------------------------------------------- # Constants diff --git a/tests/e2e/offline_inference/test_ming_flash_omni.py b/tests/e2e/offline_inference/test_ming_flash_omni.py new file mode 100644 index 000000000..c591e910a --- /dev/null +++ b/tests/e2e/offline_inference/test_ming_flash_omni.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +from pathlib import Path + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.media import ( + generate_synthetic_audio, + generate_synthetic_image, + generate_synthetic_video, +) +from tests.helpers.stage_config import modify_stage_config + +models = ["Jonathan1909/Ming-flash-omni-2.0"] + +# Ming-specific +SYSTEM_PROMPT = "你是一个友好的AI助手。\n\ndetailed thinking off" +EOS_TOKEN = "<|role_end|>" +IMAGE_TOKEN = "" +VIDEO_TOKEN = "", "STOP"]) + result = omni._prepare_prefill_sampling_params("req-1", sp) + assert result.stop == [] + + def test_clears_stop_token_ids(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + sp = SamplingParams(max_tokens=2048, stop_token_ids=[151643, 151644]) + result = omni._prepare_prefill_sampling_params("req-1", sp) + assert result.stop_token_ids == [] + + def test_clears_include_stop_str_in_output(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + sp = SamplingParams(max_tokens=2048, include_stop_str_in_output=True) + result = omni._prepare_prefill_sampling_params("req-1", sp) + assert result.include_stop_str_in_output is False + + def test_original_sp_unchanged(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + sp = SamplingParams(max_tokens=2048, stop=[""], stop_token_ids=[151643]) + _ = omni._prepare_prefill_sampling_params("req-1", sp) + assert sp.stop == [""] + assert sp.stop_token_ids == [151643] + + +# =================================================================== +# Tests: Failure mode & memory leak prevention +# =================================================================== +# NOTE: Full generate()-level failure mode tests are removed for now. +# The _run_generation error handler (line 1344-1350 in omni.py) calls +# _drop_pd_kv_params but does not increment completed_requests, causing +# the while-loop to hang. These tests need to be revisited once the +# production error-handling path is fixed to properly terminate on +# stage errors. + + +# =================================================================== +# Tests: TP size validation +# =================================================================== + + +class TestTPSizeValidation: + """Tests that _validate_pd_separation_config checks tensor_parallel_size.""" + + def test_matching_tp_passes(self, monkeypatch): + """Same TP size should not raise.""" + prefill_cfg = _prefill_stage_cfg() + prefill_cfg["engine_args"]["tensor_parallel_size"] = 2 + decode_cfg = _decode_stage_cfg(engine_input_source=[0]) + decode_cfg["engine_args"]["tensor_parallel_size"] = 2 + omni = _make_pd_omni(monkeypatch, [prefill_cfg, decode_cfg]) + assert omni._pd_separation_pair == (0, 1) + + def test_mismatched_tp_raises(self, monkeypatch): + """Different TP sizes should raise ValueError.""" + prefill_cfg = _prefill_stage_cfg() + prefill_cfg["engine_args"]["tensor_parallel_size"] = 2 + decode_cfg = _decode_stage_cfg(engine_input_source=[0]) + decode_cfg["engine_args"]["tensor_parallel_size"] = 4 + with pytest.raises(ValueError, match="tensor_parallel_size"): + _make_pd_omni(monkeypatch, [prefill_cfg, decode_cfg]) + + def test_default_tp_no_error(self, monkeypatch): + """Stages without explicit TP (defaults to 1) should pass.""" + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + assert omni._pd_separation_pair == (0, 1) diff --git a/tests/entrypoints/test_realtime_connection_helpers.py b/tests/entrypoints/test_realtime_connection_helpers.py new file mode 100644 index 000000000..e795aa92d --- /dev/null +++ b/tests/entrypoints/test_realtime_connection_helpers.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for realtime streaming helpers (PR #2581 /v1/realtime path).""" + +from __future__ import annotations + +import base64 + +import numpy as np +import pytest +import torch +from vllm.sampling_params import RequestOutputKind, SamplingParams + +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.openai.realtime_connection import RealtimeConnection + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def realtime_conn() -> RealtimeConnection: + return RealtimeConnection.__new__(RealtimeConnection) + + +class TestRealtimeConnectionTensorAndPcm: + def test_tensor_to_numpy_none(self) -> None: + assert RealtimeConnection._tensor_to_numpy(None) is None + + def test_tensor_to_numpy_1d_numpy(self) -> None: + arr = np.array([1.0, 2.0], dtype=np.float64) + out = RealtimeConnection._tensor_to_numpy(arr) + assert out is not None + assert out.dtype == np.float32 + assert out.shape == (2,) + + def test_tensor_to_numpy_2d_numpy_flattened(self) -> None: + arr = np.array([[0.5], [-0.5]], dtype=np.float32) + out = RealtimeConnection._tensor_to_numpy(arr) + assert out is not None + assert out.shape == (2,) + + def test_tensor_to_numpy_torch(self) -> None: + t = torch.tensor([[0.25, -0.25]], dtype=torch.float32) + out = RealtimeConnection._tensor_to_numpy(t) + assert out is not None + assert out.shape == (2,) + np.testing.assert_allclose(out, [0.25, -0.25], rtol=1e-5) + + def test_pcm16_b64_roundtrip(self) -> None: + audio = np.array([0.0, 1.0, -1.0], dtype=np.float32) + b64 = RealtimeConnection._pcm16_b64(audio) + raw = base64.b64decode(b64) + assert len(raw) == 6 + pcm = np.frombuffer(raw, dtype=np.int16) + assert pcm[0] == 0 + assert pcm[1] == 32767 + assert pcm[2] == -32767 + + +class TestAsyncOmniStreamingParamsValidation: + def test_accepts_streaming_friendly_params(self) -> None: + p = SamplingParams( + n=1, + stop=[], + output_kind=RequestOutputKind.DELTA, + ) + AsyncOmni._validate_streaming_input_sampling_params(p) + + def test_rejects_non_sampling_params(self) -> None: + with pytest.raises(ValueError, match="Input streaming"): + AsyncOmni._validate_streaming_input_sampling_params(object()) # type: ignore[arg-type] + + def test_rejects_n_greater_than_one(self) -> None: + p = SamplingParams(n=2, stop=[], output_kind=RequestOutputKind.DELTA) + with pytest.raises(ValueError, match="Input streaming"): + AsyncOmni._validate_streaming_input_sampling_params(p) + + def test_rejects_final_only(self) -> None: + p = SamplingParams(n=1, stop=[], output_kind=RequestOutputKind.FINAL_ONLY) + with pytest.raises(ValueError, match="Input streaming"): + AsyncOmni._validate_streaming_input_sampling_params(p) + + def test_rejects_stop_strings(self) -> None: + p = SamplingParams(n=1, stop=["\n"], output_kind=RequestOutputKind.DELTA) + with pytest.raises(ValueError, match="Input streaming"): + AsyncOmni._validate_streaming_input_sampling_params(p) diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index afa7fa82e..e60afc9cd 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -7,11 +7,33 @@ import argparse import pytest from pytest_mock import MockerFixture -from vllm_omni.entrypoints.cli.serve import run_headless +from vllm_omni.entrypoints.cli.serve import OmniServeCommand, run_headless +from vllm_omni.entrypoints.utils import detect_explicit_cli_keys pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None: + """``--no-async-chunk`` should parse to ``async_chunk=False`` and mark the + shared deploy-level dest as explicitly provided by the user.""" + try: + from vllm.utils.argparse_utils import FlexibleArgumentParser + except Exception as exc: + pytest.skip(f"Cannot build parser in this environment: {exc}") + + root = FlexibleArgumentParser() + subparsers = root.add_subparsers(dest="subcommand") + cmd = OmniServeCommand() + serve_parser = cmd.subparser_init(subparsers) + + argv = ["serve", "fake-model", "--omni", "--no-async-chunk"] + args = root.parse_args(argv) + + assert args.async_chunk is False + explicit = detect_explicit_cli_keys(argv, serve_parser) + assert "async_chunk" in explicit + + def _make_headless_args() -> argparse.Namespace: return argparse.Namespace( model="fake-model", diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 94e254c25..248629d51 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -310,6 +310,39 @@ class TestResolveModelConfigPath: assert result is not None assert "glm_image.yaml" in result + def test_voxcpm_transformers_format_resolution(self, mocker: MockerFixture): + """Test VoxCPM transformers config resolves to the voxcpm stage config.""" + mocker.patch( + "vllm_omni.entrypoints.utils.get_config", + side_effect=ValueError("missing transformers config"), + ) + mocker.patch( + "vllm_omni.entrypoints.utils.file_or_path_exists", + side_effect=lambda _model, filename, revision=None: filename == "config.json", + ) + mocker.patch( + "vllm_omni.entrypoints.utils.get_hf_file_to_dict", + return_value={"model_type": "voxcpm"}, + ) + mocker.patch( + "vllm_omni.entrypoints.utils.current_omni_platform.get_default_stage_config_path", + return_value="vllm_omni/model_executor/stage_configs", + ) + + original_exists = os.path.exists + + def mock_exists(path): + if "voxcpm.yaml" in str(path): + return True + return original_exists(path) + + mocker.patch("os.path.exists", side_effect=mock_exists) + + result = resolve_model_config_path("OpenBMB/VoxCPM1.5") + + assert result is not None + assert "voxcpm.yaml" in result + class TestLoadAndResolveStageConfigs: def test_load_and_resolve_with_kwargs(self): diff --git a/tests/examples/conftest.py b/tests/examples/conftest.py index 137d15f16..867731b21 100644 --- a/tests/examples/conftest.py +++ b/tests/examples/conftest.py @@ -1,353 +1,3 @@ -""" -Shared fixtures, helpers, and path constants for tests/examples/. -""" +"""Pytest fixtures for tests/examples.""" -import json -import os -import re -import shlex -import subprocess -import sys -import tempfile -from collections import defaultdict -from collections.abc import Callable -from pathlib import Path -from typing import Any, NamedTuple, cast - -import pytest -import torch -from safetensors.torch import save_file - -# --------------------------------------------------------------------------- -# Path constants and fixtures -# --------------------------------------------------------------------------- - -REPO_ROOT = Path(__file__).resolve().parents[2] -EXAMPLES = REPO_ROOT / "examples" - -# Use Python tempfile instead of pytest's tmp_path_factory because -# OUTPUT_DIR is needed in test collection time, but tmp_path_factory is only available in test running time. -# It is needed during test collection because extract_readme_snippets replaces LoRA path with a generated one under OUTPUT_DIR, -# and extract_readme_snippets is called at collection time to generate separate test cases for each README code block. -OUTPUT_DIR = ( - REPO_ROOT / prefix - if (prefix := os.environ.get("OUTPUT_DIR")) - else Path(tempfile.mkdtemp(prefix="vllm_omni_test_examples_")) -) - - -# --------------------------------------------------------------------------- -# Code snippet extraction and asset file helpers -# --------------------------------------------------------------------------- - -# parameters: language, code, h2_title -ReadmeSnippetExtractionSkipPredicate = Callable[[str, str, str], tuple[bool, str]] - - -class ReadmeSnippet(NamedTuple): - language: str - code: str - h2_title: str - index_in_section: int - output_file_path: Path | None = None - skip: tuple[bool, str] = (False, "") - - @property - def test_id(self) -> str: - return f"{ReadmeSnippet._slug(self.h2_title)}_{self.index_in_section:03d}" - - @staticmethod - def extract_readme_snippets( - readme_path: Path, - skipif: ReadmeSnippetExtractionSkipPredicate | None = None, - ) -> list["ReadmeSnippet"]: - import mistune - - markdown = mistune.create_markdown(renderer="ast") - tokens = markdown(readme_path.read_text(encoding="utf-8")) - tokens = cast(list[dict[str, Any]], tokens) # mistune's AST renderer always produces a list, not a str - - h2_title = "" - section_counts: defaultdict[str, int] = defaultdict(int) - snippets: list[ReadmeSnippet] = [] - - for token in tokens: - token_type = token.get("type") - - if token_type == "heading": - level = (token.get("attrs") or {}).get("level") - title = ReadmeSnippet._heading_text(token) - if level == 2: - h2_title = title - continue - - if token_type != "block_code": - continue - - try: - info = token.get("attrs").get("info") # type: ignore[reportOptionalMemberAccess] - language = info.strip().split()[0].lower() # type: ignore[reportOptionalMemberAccess] - - # Common shell aliases to "bash" in several markdown renderers. - if language in {"shell", "sh", "ksh", "zsh"}: - language = "bash" - - if language not in {"bash", "python"}: - continue - except AttributeError: - # The fence is missing explicit language info; skip it. - continue - - key = h2_title - section_counts[key] += 1 - code = token.get("raw", "") - output_file_path = None - if language == "bash": - argv = ReadmeSnippet._normalize_bash_command(code, Path(readme_path.parent)) - code = shlex.join(argv) - output_file_path = ReadmeSnippet._output_file_path_from_argv(argv) - if skipif is not None: - skip_config = skipif(language, code, h2_title) - else: - skip_config = (False, "") - snippet = ReadmeSnippet( - language=language, - code=code, - h2_title=h2_title, - index_in_section=section_counts[key], - output_file_path=output_file_path, - skip=skip_config, - ) - snippets.append(snippet) - - return snippets - - @staticmethod - def _normalize_bash_command(command: str, readme_dir: Path) -> list[str]: - line_joined_command = re.sub(r"\\\s*\n", " ", command).strip() - argv = shlex.split(line_joined_command, comments=True) - assert argv, "README bash fence produced an empty command" - - # Normalize python directory and example script location - if argv[0] in {"python", "python3"}: - argv[0] = sys.executable - if len(argv) > 1 and argv[1].endswith(".py"): - script_arg = argv[1] - script_path = Path(script_arg) - if script_path.is_absolute(): - resolved_script = script_path - else: - # Take the file name only, and append script_dir to its front - resolved_script = readme_dir / script_path.name - assert resolved_script.exists(), ( - f"README bash snippet references a script that does not exist: {script_arg} (resolved to {resolved_script})" - ) - argv[1] = str(resolved_script) - - # Normalize LoRA adapter path and ensure README LoRA assets exist. - try: - lora_arg_idx = argv.index("--lora-path") # Raise ValueError if not found - assert len(argv) > lora_arg_idx + 1, "README bash snippet uses --lora-path without a following value" - - lora_dir = OUTPUT_DIR / "lora" - adapter_model = lora_dir / "adapter_model.safetensors" - adapter_config = lora_dir / "adapter_config.json" - if not adapter_model.exists() or not adapter_config.exists(): - write_zimage_lora(lora_dir, v_scale=8.0) - - argv[lora_arg_idx + 1] = str(lora_dir) - except ValueError: - pass - - return argv - - @staticmethod - def _output_file_path_from_argv(argv: list[str]) -> Path | None: - if "--output" not in argv: - return None - output_param_idx = argv.index("--output") - assert len(argv) > output_param_idx + 1, "README bash snippet uses --output without a following value" - output_arg = argv[output_param_idx + 1] - return Path(output_arg) - - @staticmethod - def _slug(text: str) -> str: - return "".join(ch.lower() if ch.isalnum() else "_" for ch in text).strip("_") - - @staticmethod - def _heading_text(token: dict) -> str: - return "".join(child.get("raw", "") for child in token.get("children", [])).strip() - - -# [TODO] Duplicate `_write_zimage_lora` in tests/e2e/online_serving/test_images_generations_lora.py. Combine these helpers and tests/e2e/offline_inference/test_diffusion_lora.py to test/utils later -def write_zimage_lora(adapter_dir: Path, *, q_scale: float = 0.0, k_scale: float = 0.0, v_scale: float = 0.0): - adapter_dir.mkdir(parents=True, exist_ok=True) - - # Z-Image transformer uses dim=3840 by default. - dim = 3840 - module_name = "transformer.layers.0.attention.to_qkv" - rank = 1 - - lora_a = torch.zeros((rank, dim), dtype=torch.float32) - lora_a[0, 0] = 1.0 - - # QKVParallelLinear packs (Q, K, V) => out dim is 3 * dim (tp=1). - lora_b = torch.zeros((3 * dim, rank), dtype=torch.float32) - if q_scale: - lora_b[:dim, 0] = q_scale - if k_scale: - lora_b[dim : 2 * dim, 0] = k_scale - if v_scale: - lora_b[2 * dim :, 0] = v_scale - - save_file( - { - f"base_model.model.{module_name}.lora_A.weight": lora_a, - f"base_model.model.{module_name}.lora_B.weight": lora_b, - }, - str(adapter_dir / "adapter_model.safetensors"), - ) - (adapter_dir / "adapter_config.json").write_text( - json.dumps( - { - "r": rank, - "lora_alpha": rank, - "target_modules": [module_name], - } - ), - encoding="utf-8", - ) - - -# --------------------------------------------------------------------------- -# Code runner and subprocess helpers -# --------------------------------------------------------------------------- - - -class ExampleRunResult(NamedTuple): - run_dir: Path - assets: list[Path] - - -class ExampleRunner: - """Run extracted README snippets and return generated assets. - - The output materials are organized in a three-level directory structure: - - Set at init: `self.output_root` for all tests (from env OUTPUT_DIR) - - Set at `self.run(...)`: `output_subfolder` for a specific example page (e.g., `example_offline_t2i`) - - Generated by `extract_readme_snippets`: `snippet.test_id` for a specific code block (matching H2 titles, e.g., `basic_usage_001`) - """ - - IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp"} - - def __init__(self, output_root: Path) -> None: - self.output_root = output_root - - def run( - self, snippet: ReadmeSnippet, *, output_subfolder: Path = Path("."), env: dict[str, str] | None = None - ) -> ExampleRunResult: - run_dir = self.output_root / output_subfolder / snippet.test_id - run_dir.mkdir(parents=True, exist_ok=True) - - if snippet.language == "python": - assets = self._run_python_snippet(snippet, run_dir, env) - return ExampleRunResult(run_dir=run_dir, assets=assets) - - if snippet.language == "bash": - asset = self._run_bash_snippet(snippet, run_dir, env) - return ExampleRunResult(run_dir=run_dir, assets=[asset]) - - raise AssertionError(f"Unsupported snippet language: {snippet.language}") - - def _run_python_snippet( - self, snippet: ReadmeSnippet, run_dir: Path, env: dict[str, str] | None = None - ) -> list[Path]: - # Saving the script to a temporary file and `run_cmd` it. - # Not using `exec(snippet.code)` because the output is lost. - script_path = run_dir / "snippet.py" - script_path.write_text(snippet.code, encoding="utf-8") - - before = self._collect_images(run_dir) - run_cmd([sys.executable, str(script_path)], cwd=run_dir, env=env) - after = self._collect_images(run_dir) - - assets = sorted(after - before) - return assets - - def _run_bash_snippet(self, snippet: ReadmeSnippet, run_dir: Path, env: dict[str, str] | None = None) -> Path: - run_cmd(snippet.code, shell=True, cwd=run_dir, env=env) - - assert snippet.output_file_path is not None, ( - f"README bash snippet is missing --output argument: {snippet.test_id}. " - "The test script cannot guess the output file path." - ) - - # If the code snippet declares a relative path for the output file, append this path to the parent output collection directory. - # If the code snippet declares an absolute path (not likely but just in case), the return value resolution removes `run_dir`, also correctly pointing to this file. - return run_dir / snippet.output_file_path - - def _collect_images(self, root: Path) -> set[Path]: - return {path for path in root.rglob("*") if path.suffix.lower() in self.IMAGE_SUFFIXES} - - -@pytest.fixture -def example_runner() -> ExampleRunner: - return ExampleRunner(output_root=OUTPUT_DIR) - - -def run_cmd( - command: list[str] | str, - *, - shell: bool = False, - env: dict[str, str] | None = None, - cwd: Path | str | None = None, -) -> str: - """Run a command as a subprocess; assert zero exit code and return stdout. - - Output is fully captured and returned as a string so callers can parse it - (e.g. with :func:`extract_content_after_keyword`). - Use this for scripts whose printed output is part of the test assertion. - """ - if env is not None: - env = {**os.environ.copy(), **env} - result = subprocess.run(command, capture_output=True, text=True, shell=shell, env=env, cwd=cwd) - - if result.returncode != 0: - print(f"STDERR: {result.stderr}") - raise subprocess.CalledProcessError(result.returncode, command) - - all_output = result.stdout - print(f"All output:\n{all_output}") - return all_output - - -# --------------------------------------------------------------------------- -# Output validation helpers -# --------------------------------------------------------------------------- - - -def extract_content_after_keyword(keywords: str, text: str) -> str: - """Return the text that follows *keywords* in *text* (regex match). - - Raises ``AssertionError`` if the keyword is not found, so test failures - produce a clear message pointing at the missing keyword. - """ - matches = re.findall(rf"{keywords}\s*(.+)", text, re.DOTALL) - - if not matches: - raise AssertionError(f"Keywords {keywords} not found in provided text output") - return matches[0] - - -def strip_trailing_audio_saved_line(text: str) -> str: - """Drop trailing ``Audio saved to ...`` lines from captured client stdout. - - ``openai_chat_completion_client_for_multimodal_generation.py`` may print - ``Chat completion output from text:`` for one choice and ``Audio saved to`` - for another; :func:`extract_content_after_keyword` uses ``re.DOTALL`` and - would otherwise keep the audio progress line inside the *text* segment. - """ - lines = text.splitlines() - while lines and lines[-1].strip().startswith("Audio saved to"): - lines.pop() - return "\n".join(lines).strip() +from tests.examples.helpers import example_runner # noqa: F401 diff --git a/tests/examples/helpers.py b/tests/examples/helpers.py new file mode 100644 index 000000000..137d15f16 --- /dev/null +++ b/tests/examples/helpers.py @@ -0,0 +1,353 @@ +""" +Shared fixtures, helpers, and path constants for tests/examples/. +""" + +import json +import os +import re +import shlex +import subprocess +import sys +import tempfile +from collections import defaultdict +from collections.abc import Callable +from pathlib import Path +from typing import Any, NamedTuple, cast + +import pytest +import torch +from safetensors.torch import save_file + +# --------------------------------------------------------------------------- +# Path constants and fixtures +# --------------------------------------------------------------------------- + +REPO_ROOT = Path(__file__).resolve().parents[2] +EXAMPLES = REPO_ROOT / "examples" + +# Use Python tempfile instead of pytest's tmp_path_factory because +# OUTPUT_DIR is needed in test collection time, but tmp_path_factory is only available in test running time. +# It is needed during test collection because extract_readme_snippets replaces LoRA path with a generated one under OUTPUT_DIR, +# and extract_readme_snippets is called at collection time to generate separate test cases for each README code block. +OUTPUT_DIR = ( + REPO_ROOT / prefix + if (prefix := os.environ.get("OUTPUT_DIR")) + else Path(tempfile.mkdtemp(prefix="vllm_omni_test_examples_")) +) + + +# --------------------------------------------------------------------------- +# Code snippet extraction and asset file helpers +# --------------------------------------------------------------------------- + +# parameters: language, code, h2_title +ReadmeSnippetExtractionSkipPredicate = Callable[[str, str, str], tuple[bool, str]] + + +class ReadmeSnippet(NamedTuple): + language: str + code: str + h2_title: str + index_in_section: int + output_file_path: Path | None = None + skip: tuple[bool, str] = (False, "") + + @property + def test_id(self) -> str: + return f"{ReadmeSnippet._slug(self.h2_title)}_{self.index_in_section:03d}" + + @staticmethod + def extract_readme_snippets( + readme_path: Path, + skipif: ReadmeSnippetExtractionSkipPredicate | None = None, + ) -> list["ReadmeSnippet"]: + import mistune + + markdown = mistune.create_markdown(renderer="ast") + tokens = markdown(readme_path.read_text(encoding="utf-8")) + tokens = cast(list[dict[str, Any]], tokens) # mistune's AST renderer always produces a list, not a str + + h2_title = "" + section_counts: defaultdict[str, int] = defaultdict(int) + snippets: list[ReadmeSnippet] = [] + + for token in tokens: + token_type = token.get("type") + + if token_type == "heading": + level = (token.get("attrs") or {}).get("level") + title = ReadmeSnippet._heading_text(token) + if level == 2: + h2_title = title + continue + + if token_type != "block_code": + continue + + try: + info = token.get("attrs").get("info") # type: ignore[reportOptionalMemberAccess] + language = info.strip().split()[0].lower() # type: ignore[reportOptionalMemberAccess] + + # Common shell aliases to "bash" in several markdown renderers. + if language in {"shell", "sh", "ksh", "zsh"}: + language = "bash" + + if language not in {"bash", "python"}: + continue + except AttributeError: + # The fence is missing explicit language info; skip it. + continue + + key = h2_title + section_counts[key] += 1 + code = token.get("raw", "") + output_file_path = None + if language == "bash": + argv = ReadmeSnippet._normalize_bash_command(code, Path(readme_path.parent)) + code = shlex.join(argv) + output_file_path = ReadmeSnippet._output_file_path_from_argv(argv) + if skipif is not None: + skip_config = skipif(language, code, h2_title) + else: + skip_config = (False, "") + snippet = ReadmeSnippet( + language=language, + code=code, + h2_title=h2_title, + index_in_section=section_counts[key], + output_file_path=output_file_path, + skip=skip_config, + ) + snippets.append(snippet) + + return snippets + + @staticmethod + def _normalize_bash_command(command: str, readme_dir: Path) -> list[str]: + line_joined_command = re.sub(r"\\\s*\n", " ", command).strip() + argv = shlex.split(line_joined_command, comments=True) + assert argv, "README bash fence produced an empty command" + + # Normalize python directory and example script location + if argv[0] in {"python", "python3"}: + argv[0] = sys.executable + if len(argv) > 1 and argv[1].endswith(".py"): + script_arg = argv[1] + script_path = Path(script_arg) + if script_path.is_absolute(): + resolved_script = script_path + else: + # Take the file name only, and append script_dir to its front + resolved_script = readme_dir / script_path.name + assert resolved_script.exists(), ( + f"README bash snippet references a script that does not exist: {script_arg} (resolved to {resolved_script})" + ) + argv[1] = str(resolved_script) + + # Normalize LoRA adapter path and ensure README LoRA assets exist. + try: + lora_arg_idx = argv.index("--lora-path") # Raise ValueError if not found + assert len(argv) > lora_arg_idx + 1, "README bash snippet uses --lora-path without a following value" + + lora_dir = OUTPUT_DIR / "lora" + adapter_model = lora_dir / "adapter_model.safetensors" + adapter_config = lora_dir / "adapter_config.json" + if not adapter_model.exists() or not adapter_config.exists(): + write_zimage_lora(lora_dir, v_scale=8.0) + + argv[lora_arg_idx + 1] = str(lora_dir) + except ValueError: + pass + + return argv + + @staticmethod + def _output_file_path_from_argv(argv: list[str]) -> Path | None: + if "--output" not in argv: + return None + output_param_idx = argv.index("--output") + assert len(argv) > output_param_idx + 1, "README bash snippet uses --output without a following value" + output_arg = argv[output_param_idx + 1] + return Path(output_arg) + + @staticmethod + def _slug(text: str) -> str: + return "".join(ch.lower() if ch.isalnum() else "_" for ch in text).strip("_") + + @staticmethod + def _heading_text(token: dict) -> str: + return "".join(child.get("raw", "") for child in token.get("children", [])).strip() + + +# [TODO] Duplicate `_write_zimage_lora` in tests/e2e/online_serving/test_images_generations_lora.py. Combine these helpers and tests/e2e/offline_inference/test_diffusion_lora.py to test/utils later +def write_zimage_lora(adapter_dir: Path, *, q_scale: float = 0.0, k_scale: float = 0.0, v_scale: float = 0.0): + adapter_dir.mkdir(parents=True, exist_ok=True) + + # Z-Image transformer uses dim=3840 by default. + dim = 3840 + module_name = "transformer.layers.0.attention.to_qkv" + rank = 1 + + lora_a = torch.zeros((rank, dim), dtype=torch.float32) + lora_a[0, 0] = 1.0 + + # QKVParallelLinear packs (Q, K, V) => out dim is 3 * dim (tp=1). + lora_b = torch.zeros((3 * dim, rank), dtype=torch.float32) + if q_scale: + lora_b[:dim, 0] = q_scale + if k_scale: + lora_b[dim : 2 * dim, 0] = k_scale + if v_scale: + lora_b[2 * dim :, 0] = v_scale + + save_file( + { + f"base_model.model.{module_name}.lora_A.weight": lora_a, + f"base_model.model.{module_name}.lora_B.weight": lora_b, + }, + str(adapter_dir / "adapter_model.safetensors"), + ) + (adapter_dir / "adapter_config.json").write_text( + json.dumps( + { + "r": rank, + "lora_alpha": rank, + "target_modules": [module_name], + } + ), + encoding="utf-8", + ) + + +# --------------------------------------------------------------------------- +# Code runner and subprocess helpers +# --------------------------------------------------------------------------- + + +class ExampleRunResult(NamedTuple): + run_dir: Path + assets: list[Path] + + +class ExampleRunner: + """Run extracted README snippets and return generated assets. + + The output materials are organized in a three-level directory structure: + - Set at init: `self.output_root` for all tests (from env OUTPUT_DIR) + - Set at `self.run(...)`: `output_subfolder` for a specific example page (e.g., `example_offline_t2i`) + - Generated by `extract_readme_snippets`: `snippet.test_id` for a specific code block (matching H2 titles, e.g., `basic_usage_001`) + """ + + IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp"} + + def __init__(self, output_root: Path) -> None: + self.output_root = output_root + + def run( + self, snippet: ReadmeSnippet, *, output_subfolder: Path = Path("."), env: dict[str, str] | None = None + ) -> ExampleRunResult: + run_dir = self.output_root / output_subfolder / snippet.test_id + run_dir.mkdir(parents=True, exist_ok=True) + + if snippet.language == "python": + assets = self._run_python_snippet(snippet, run_dir, env) + return ExampleRunResult(run_dir=run_dir, assets=assets) + + if snippet.language == "bash": + asset = self._run_bash_snippet(snippet, run_dir, env) + return ExampleRunResult(run_dir=run_dir, assets=[asset]) + + raise AssertionError(f"Unsupported snippet language: {snippet.language}") + + def _run_python_snippet( + self, snippet: ReadmeSnippet, run_dir: Path, env: dict[str, str] | None = None + ) -> list[Path]: + # Saving the script to a temporary file and `run_cmd` it. + # Not using `exec(snippet.code)` because the output is lost. + script_path = run_dir / "snippet.py" + script_path.write_text(snippet.code, encoding="utf-8") + + before = self._collect_images(run_dir) + run_cmd([sys.executable, str(script_path)], cwd=run_dir, env=env) + after = self._collect_images(run_dir) + + assets = sorted(after - before) + return assets + + def _run_bash_snippet(self, snippet: ReadmeSnippet, run_dir: Path, env: dict[str, str] | None = None) -> Path: + run_cmd(snippet.code, shell=True, cwd=run_dir, env=env) + + assert snippet.output_file_path is not None, ( + f"README bash snippet is missing --output argument: {snippet.test_id}. " + "The test script cannot guess the output file path." + ) + + # If the code snippet declares a relative path for the output file, append this path to the parent output collection directory. + # If the code snippet declares an absolute path (not likely but just in case), the return value resolution removes `run_dir`, also correctly pointing to this file. + return run_dir / snippet.output_file_path + + def _collect_images(self, root: Path) -> set[Path]: + return {path for path in root.rglob("*") if path.suffix.lower() in self.IMAGE_SUFFIXES} + + +@pytest.fixture +def example_runner() -> ExampleRunner: + return ExampleRunner(output_root=OUTPUT_DIR) + + +def run_cmd( + command: list[str] | str, + *, + shell: bool = False, + env: dict[str, str] | None = None, + cwd: Path | str | None = None, +) -> str: + """Run a command as a subprocess; assert zero exit code and return stdout. + + Output is fully captured and returned as a string so callers can parse it + (e.g. with :func:`extract_content_after_keyword`). + Use this for scripts whose printed output is part of the test assertion. + """ + if env is not None: + env = {**os.environ.copy(), **env} + result = subprocess.run(command, capture_output=True, text=True, shell=shell, env=env, cwd=cwd) + + if result.returncode != 0: + print(f"STDERR: {result.stderr}") + raise subprocess.CalledProcessError(result.returncode, command) + + all_output = result.stdout + print(f"All output:\n{all_output}") + return all_output + + +# --------------------------------------------------------------------------- +# Output validation helpers +# --------------------------------------------------------------------------- + + +def extract_content_after_keyword(keywords: str, text: str) -> str: + """Return the text that follows *keywords* in *text* (regex match). + + Raises ``AssertionError`` if the keyword is not found, so test failures + produce a clear message pointing at the missing keyword. + """ + matches = re.findall(rf"{keywords}\s*(.+)", text, re.DOTALL) + + if not matches: + raise AssertionError(f"Keywords {keywords} not found in provided text output") + return matches[0] + + +def strip_trailing_audio_saved_line(text: str) -> str: + """Drop trailing ``Audio saved to ...`` lines from captured client stdout. + + ``openai_chat_completion_client_for_multimodal_generation.py`` may print + ``Chat completion output from text:`` for one choice and ``Audio saved to`` + for another; :func:`extract_content_after_keyword` uses ``re.DOTALL`` and + would otherwise keep the audio progress line inside the *text* segment. + """ + lines = text.splitlines() + while lines and lines[-1].strip().startswith("Audio saved to"): + lines.pop() + return "\n".join(lines).strip() diff --git a/tests/examples/offline_inference/test_text_to_image.py b/tests/examples/offline_inference/test_text_to_image.py index a08d16f16..f24506587 100644 --- a/tests/examples/offline_inference/test_text_to_image.py +++ b/tests/examples/offline_inference/test_text_to_image.py @@ -7,9 +7,9 @@ from pathlib import Path import pytest -from tests.conftest import assert_image_valid -from tests.examples.conftest import EXAMPLES, ExampleRunner, ReadmeSnippet -from tests.utils import hardware_marks +from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet +from tests.helpers.assertions import assert_image_valid +from tests.helpers.mark import hardware_marks pytestmark = [pytest.mark.advanced_model, pytest.mark.example, *hardware_marks(res={"cuda": "H100"})] diff --git a/tests/examples/online_serving/test_qwen2_5_omni.py b/tests/examples/online_serving/test_qwen2_5_omni.py index a78ccf592..b8acfef84 100644 --- a/tests/examples/online_serving/test_qwen2_5_omni.py +++ b/tests/examples/online_serving/test_qwen2_5_omni.py @@ -5,33 +5,29 @@ See examples/online_serving/qwen2_5_omni/README.md import os -from vllm_omni.platforms import current_omni_platform - os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" from pathlib import Path import pytest -from tests.conftest import OmniServerParams, convert_audio_file_to_text, cosine_similarity_text -from tests.examples.conftest import ( +from tests.examples.helpers import ( extract_content_after_keyword, run_cmd, strip_trailing_audio_saved_line, ) -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test +from tests.helpers.media import convert_audio_file_to_text, cosine_similarity_text +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path pytestmark = [pytest.mark.advanced_model, pytest.mark.example] models = ["Qwen/Qwen2.5-Omni-7B"] - -stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_ci.yaml")] - -if current_omni_platform.is_xpu(): - stage_configs = [ - str(Path(__file__).parent.parent.parent / "e2e" / "stage_configs" / "xpu" / "qwen2_5_omni_ci.yaml") - ] +# Single CI deploy YAML; rocm/xpu deltas are picked automatically via the +# platforms: section in vllm_omni/deploy/ci/qwen2_5_omni.yaml. +stage_configs = [get_deploy_config_path("ci/qwen2_5_omni.yaml")] example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving") # Create parameter combinations for model and stage config diff --git a/tests/examples/online_serving/test_qwen3_omni.py b/tests/examples/online_serving/test_qwen3_omni.py index 65f99d7bf..e2c3b3384 100644 --- a/tests/examples/online_serving/test_qwen3_omni.py +++ b/tests/examples/online_serving/test_qwen3_omni.py @@ -5,31 +5,28 @@ See examples/online_serving/qwen3_omni/README.md import os -from vllm_omni.platforms import current_omni_platform - os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" from pathlib import Path import pytest -from tests.conftest import OmniServerParams, convert_audio_file_to_text, cosine_similarity_text -from tests.examples.conftest import ( +from tests.examples.helpers import ( extract_content_after_keyword, run_cmd, strip_trailing_audio_saved_line, ) -from tests.utils import hardware_test +from tests.helpers.mark import hardware_test +from tests.helpers.media import convert_audio_file_to_text, cosine_similarity_text +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path pytestmark = [pytest.mark.advanced_model, pytest.mark.example] models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] -stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_configs" / "qwen3_omni_ci.yaml")] - -if current_omni_platform.is_xpu(): - stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")] +stage_configs = [get_deploy_config_path("ci/qwen3_omni_moe.yaml")] example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving") diff --git a/tests/examples/online_serving/test_text_to_image.py b/tests/examples/online_serving/test_text_to_image.py index 51b7ff61b..6f89cb549 100644 --- a/tests/examples/online_serving/test_text_to_image.py +++ b/tests/examples/online_serving/test_text_to_image.py @@ -13,9 +13,10 @@ from pathlib import Path import pytest -from tests.conftest import OmniServer, OmniServerParams, assert_image_valid -from tests.examples.conftest import EXAMPLES, OUTPUT_DIR, run_cmd, write_zimage_lora -from tests.utils import hardware_marks +from tests.examples.helpers import EXAMPLES, OUTPUT_DIR, run_cmd, write_zimage_lora +from tests.helpers.assertions import assert_image_valid +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams pytestmark = [pytest.mark.advanced_model, pytest.mark.example, *hardware_marks(res={"cuda": "H100"})] diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py new file mode 100644 index 000000000..a3348b07f --- /dev/null +++ b/tests/helpers/__init__.py @@ -0,0 +1,8 @@ +"""Shared, importable test helper utilities. + +Submodules (``assertions``, ``env``, ``media``, ``runtime``, …) are imported +explicitly by callers. Avoid star-importing everything here: that ran before +refactor only inside the old monolithic ``conftest``; a greedy ``__init__`` +changes import order and can affect in-process Omni (``OmniRunner`` / offline +e2e) vs subprocess-based ``OmniServer`` tests. +""" diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py new file mode 100644 index 000000000..533806477 --- /dev/null +++ b/tests/helpers/assertions.py @@ -0,0 +1,506 @@ +"""Assertion and response validation helpers for tests.""" + +import io +import tempfile +import threading +from io import BytesIO +from pathlib import Path +from typing import Any + +import numpy as np +import soundfile as sf +from PIL import Image + +from tests.helpers.media import cosine_similarity_text + +_GENDER_PIPELINE = None +_GENDER_PIPELINE_LOCK = threading.Lock() +_PCM_SPEECH_SAMPLE_RATE_HZ = 24_000 +_MIN_PCM_SPEECH_HNR_DB = 1.0 +_PRESET_VOICE_GENDER_MAP: dict[str, str] = { + "serena": "female", + "uncle_fu": "male", + "chelsie": "female", + "clone": "female", + "ethan": "male", +} + + +def assert_image_diffusion_response( + response, + request_config: dict[str, Any], + run_level: str = None, +) -> None: + """ + Validate image diffusion response. + + Expected request_config schema: + { + "request_type": "image", + "extra_body": { + "num_outputs_per_prompt": 1, + "width": ..., + "height": ..., + ... + } + } + """ + assert response.images is not None, "Image response is None" + assert len(response.images) > 0, "No images in response" + + extra_body = request_config.get("extra_body") or {} + + num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt") + if num_outputs_per_prompt is not None: + assert len(response.images) == num_outputs_per_prompt, ( + f"Expected {num_outputs_per_prompt} images, got {len(response.images)}" + ) + + if run_level == "advanced_model": + width = extra_body.get("width") + height = extra_body.get("height") + + if width is not None or height is not None: + for img in response.images: + assert_image_valid(img, width=width, height=height) + + +def assert_video_diffusion_response( + response, + request_config: dict[str, Any], + run_level: str = None, +) -> None: + """ + Validate video diffusion response. + + Expected request_config schema: + { + "request_type": "video", + "form_data": { + "prompt": "...", + "num_frames": ..., + "width": ..., + "height": ..., + "fps": ..., + ... + } + } + """ + form_data = request_config.get("form_data", {}) + + assert response.videos is not None, "Video response is None" + assert len(response.videos) > 0, "No videos in response" + + expected_frames = _maybe_int(form_data.get("num_frames")) + expected_width = _maybe_int(form_data.get("width")) + expected_height = _maybe_int(form_data.get("height")) + expected_fps = _maybe_int(form_data.get("fps")) + + for vid_bytes in response.videos: + assert_video_valid( + vid_bytes, + num_frames=expected_frames, + width=expected_width, + height=expected_height, + fps=expected_fps, + ) + + +def assert_audio_diffusion_response( + response, + request_config: dict[str, Any], + run_level: str = None, +) -> None: + """ + Validate audio diffusion response. + """ + raise NotImplementedError("Audio validation is not implemented yet") + + +def _maybe_int(value: Any) -> int | None: + if value is None: + return None + return int(value) + + +def assert_image_valid(image: Path | Image.Image, *, width: int | None = None, height: int | None = None): + """Assert the file is a loadable image with optional exact dimensions.""" + if isinstance(image, Path): + assert image.exists(), f"Image not found: {image}" + image = Image.open(image) + image.load() + assert image.width > 0 and image.height > 0 + if width is not None: + assert image.width == width, f"Expected width={width}, got {image.width}" + if height is not None: + assert image.height == height, f"Expected height={height}, got {image.height}" + return image + + +def assert_video_valid( + video: Path | bytes | BytesIO, + *, + num_frames: int | None = None, + width: int | None = None, + height: int | None = None, + fps: float | None = None, +) -> dict[str, int | float]: + """Assert the MP4 has the expected resolution and frame count. + + For several diffusion backends, encoded MP4 frame count follows a codec-aligned + convention (e.g. request `num_frames=8` can produce 9 encoded frames). Keep + this compatibility behavior to avoid false negatives in online-serving tests. + """ + temp_path = None + cap = None + try: + import cv2 + + if isinstance(video, Path): + if not video.exists(): + raise AssertionError(f"Video file not found: {video}") + video_path = str(video) + else: + suffix = ".mp4" + with tempfile.NamedTemporaryFile(delete=False, suffix=suffix, mode="wb") as tmp: + if isinstance(video, bytes): + tmp.write(video) + elif isinstance(video, BytesIO): + tmp.write(video.getvalue()) + else: + raise TypeError(f"Unsupported video type: {type(video)}") + temp_path = Path(tmp.name) + video_path = str(temp_path) + + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise AssertionError("Failed to open video capture") + + actual_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + actual_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + actual_fps = float(cap.get(cv2.CAP_PROP_FPS)) + actual_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + if width is not None: + assert actual_width == width, f"Expected width={width}, got {actual_width}" + if height is not None: + assert actual_height == height, f"Expected height={height}, got {actual_height}" + if fps is not None and actual_fps: + assert abs(actual_fps - float(fps)) < 1.0, f"Expected fps~={fps}, got {actual_fps}" + if num_frames is not None: + expected_frames = (int(num_frames) // 4) * 4 + 1 + assert actual_frames == expected_frames, f"Expected frames={expected_frames}, got {actual_frames}" + + return { + "width": actual_width, + "height": actual_height, + "fps": actual_fps, + "num_frames": actual_frames, + } + except Exception as e: + print(f"ERROR: {type(e).__name__}: {e}", flush=True) + raise + finally: + if cap is not None: + cap.release() + if temp_path and temp_path.exists(): + try: + temp_path.unlink() + except OSError: + pass + + +def assert_audio_valid( + audio_or_path: Path | np.ndarray, + *, + sample_rate: int, + channels: int, + duration_s: float, +) -> None: + """Assert WAV file or (batch, channels, samples) ndarray matches expected audio format.""" + expected_samples = int(duration_s * sample_rate) + if isinstance(audio_or_path, np.ndarray): + audio = audio_or_path + assert audio.ndim == 3, f"Expected audio ndim=3 (batch, channels, samples), got shape {audio.shape}" + assert audio.shape[0] == 1, f"Expected batch size 1, got {audio.shape[0]}" + assert audio.shape[1] == channels, f"Expected {channels} channels, got {audio.shape[1]}" + assert audio.shape[2] == expected_samples, ( + f"Expected {expected_samples} samples ({duration_s}s @ {sample_rate} Hz), got {audio.shape[2]}" + ) + return + + path = audio_or_path + assert path.exists(), f"Audio not found: {path}" + info = sf.info(str(path)) + assert info.samplerate == sample_rate, f"Expected sample_rate={sample_rate}, got {info.samplerate}" + assert info.channels == channels, f"Expected {channels} channel(s), got {info.channels}" + assert info.frames == expected_samples, ( + f"Expected {expected_samples} frames ({duration_s}s @ {sample_rate} Hz), got {info.frames}" + ) + + +def _load_gender_pipeline(): + global _GENDER_PIPELINE + if _GENDER_PIPELINE is not None: + return _GENDER_PIPELINE + model_name = "7wolf/wav2vec2-base-gender-classification" + try: + from transformers import pipeline + + _GENDER_PIPELINE = pipeline(task="audio-classification", model=model_name, device=-1) + return _GENDER_PIPELINE + except Exception as exc: # pragma: no cover + print(f"Warning: failed to create gender pipeline '{model_name}': {exc}") + _GENDER_PIPELINE = None + return None + + +def _median_pitch_hz_from_autocorr(mono: np.ndarray, sr: int) -> float | None: + x = np.asarray(mono, dtype=np.float64) + x = x - np.mean(x) + if x.size < int(0.15 * sr): + return None + frame_len = int(0.04 * sr) + hop = max(frame_len // 2, 1) + f0_min_hz, f0_max_hz = 70.0, 400.0 + lag_min = max(1, int(sr / f0_max_hz)) + lag_max = min(frame_len - 2, int(sr / f0_min_hz)) + if lag_max <= lag_min: + return None + win = np.hamming(frame_len) + pitches: list[float] = [] + for start in range(0, int(x.shape[0]) - frame_len, hop): + frame = x[start : start + frame_len] * win + frame = frame - np.mean(frame) + if float(np.sqrt(np.mean(frame**2))) < 1e-4: + continue + ac = np.correlate(frame, frame, mode="full")[frame_len - 1 :] + ac = ac / (float(ac[0]) + 1e-12) + region = ac[lag_min : lag_max + 1] + peak_rel = int(np.argmax(region)) + peak_lag = peak_rel + lag_min + if peak_lag <= 0: + continue + f0 = float(sr) / float(peak_lag) + if f0_min_hz <= f0 <= f0_max_hz: + pitches.append(f0) + if len(pitches) < 4: + return None + return float(np.median(np.asarray(pitches, dtype=np.float64))) + + +def _estimate_voice_gender_from_audio(audio_bytes: bytes) -> str: + data, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32", always_2d=True) + if data.size == 0: + raise ValueError("Empty audio") + mono = np.mean(data, axis=1) + try: + target_sr = 16000 + if int(sr) != target_sr and mono.size > 1: + src_len = int(mono.shape[0]) + dst_len = max(1, int(round(src_len * float(target_sr) / float(sr)))) + src_idx = np.arange(src_len, dtype=np.float32) + dst_idx = np.linspace(0, src_len - 1, dst_len, dtype=np.float32) + mono = np.interp(dst_idx, src_idx, mono.astype(np.float32, copy=False)).astype(np.float32) + sr = target_sr + + median_f0 = _median_pitch_hz_from_autocorr(mono, sr) + clf = _load_gender_pipeline() + if clf is None: + print("gender model not available, returning 'unknown'") + return "unknown" + with _GENDER_PIPELINE_LOCK: + outputs = clf(mono, sampling_rate=sr) + if not outputs: + return "unknown" + top = outputs[0] + label = str(top.get("label", "")).lower() + conf = float(top.get("score", 0.0)) + if conf < 0.5: + gender = "unknown" + elif ("female" in label) or ("жен" in label): + gender = "female" + elif ("male" in label) or ("муж" in label): + gender = "male" + else: + gender = "unknown" + + if gender == "female" and median_f0 is not None and median_f0 < 165.0 and conf < 0.88: + print(f"gender pitch assist: reclassifying female->male (median_f0={median_f0:.1f} Hz, conf={conf:.3f})") + gender = "male" + elif gender == "male" and median_f0 is not None and median_f0 > 230.0 and conf < 0.88: + print(f"gender pitch assist: reclassifying male->female (median_f0={median_f0:.1f} Hz, conf={conf:.3f})") + gender = "female" + print( + f"gender classifier: label={label}, conf={conf:.3f}, gender={gender}" + + (f", median_f0={median_f0:.1f}Hz" if median_f0 is not None else "") + ) + return gender + except Exception as exc: # pragma: no cover + print(f"Warning: gender classification failed, returning 'unknown': {exc}") + return "unknown" + + +def _assert_preset_voice_gender_from_audio(audio_bytes: bytes | None, voice_name: str | None) -> None: + """If ``voice_name`` matches a known preset, assert classifier gender matches (skip when unknown).""" + if not voice_name or not audio_bytes: + return + key = str(voice_name).lower() + expected_gender = _PRESET_VOICE_GENDER_MAP.get(key) + if expected_gender is None: + return + estimated_gender = _estimate_voice_gender_from_audio(audio_bytes) + print(f"Preset voice gender check: preset={key!r}, estimated={estimated_gender!r}, expected={expected_gender!r}") + if estimated_gender != "unknown": + assert estimated_gender == expected_gender, ( + f"{voice_name!r} is expected {expected_gender}, but estimated gender is {estimated_gender!r}" + ) + + +def _compute_pcm_hnr_db(pcm_samples: np.ndarray, sr: int = _PCM_SPEECH_SAMPLE_RATE_HZ) -> float: + frame_len = int(0.03 * sr) + hop = frame_len // 2 + hnr_values: list[float] = [] + for start in range(0, len(pcm_samples) - frame_len, hop): + frame = pcm_samples[start : start + frame_len].astype(np.float32, copy=False) + frame = frame - np.mean(frame) + if np.max(np.abs(frame)) < 0.01: + continue + ac = np.correlate(frame, frame, mode="full")[len(frame) - 1 :] + ac = ac / (ac[0] + 1e-10) + min_lag = int(sr / 400) + max_lag = min(int(sr / 80), len(ac)) + if min_lag >= max_lag: + continue + peak = float(np.max(ac[min_lag:max_lag])) + if 0 < peak < 1: + hnr_values.append(10 * np.log10(peak / (1 - peak + 1e-10))) + return float(np.mean(hnr_values)) if hnr_values else 0.0 + + +def _assert_pcm_int16_speech_hnr(audio_bytes: bytes) -> None: + assert audio_bytes is not None and len(audio_bytes) >= 2, "missing PCM bytes" + assert len(audio_bytes) % 2 == 0, "PCM byte length must be aligned to int16" + pcm_samples = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32) / 32768.0 + hnr = _compute_pcm_hnr_db(pcm_samples) + print(f"PCM speech HNR: {hnr:.2f} dB (threshold: {_MIN_PCM_SPEECH_HNR_DB} dB)") + assert hnr >= _MIN_PCM_SPEECH_HNR_DB, ( + f"Audio distortion detected: HNR={hnr:.2f} dB < {_MIN_PCM_SPEECH_HNR_DB} dB. " + "Voice clone decoder may be losing ref_code speaker context on later chunks." + ) + + +def assert_omni_response(response: Any, request_config: dict[str, Any], run_level): + """ + Validate response results. + + Args: + response: OmniResponse object + + Raises: + AssertionError: When the response does not meet validation criteria + """ + assert response.success, "The request failed." + e2e_latency = response.e2e_latency + if e2e_latency is not None: + print(f"the e2e latency is: {e2e_latency}") + + modalities = request_config.get("modalities", ["text", "audio"]) + + if run_level == "advanced_model": + if "audio" in modalities: + assert response.audio_content is not None, "No audio output is generated" + print(f"audio content is: {response.audio_content}") + speaker = request_config.get("speaker") + if speaker: + _assert_preset_voice_gender_from_audio( + response.audio_bytes, + speaker, + ) + + if "text" in modalities: + assert response.text_content is not None, "No text output is generated" + print(f"text content is: {response.text_content}") + + # Verify image description + word_types = ["text", "image", "audio", "video"] + keywords_dict = request_config.get("key_words", {}) + for word_type in word_types: + keywords = keywords_dict.get(word_type) + if "text" in modalities: + if keywords: + text_lower = response.text_content.lower() + assert any(str(kw).lower() in text_lower for kw in keywords), ( + "The output does not contain any of the keywords." + ) + else: + if keywords: + audio_lower = response.audio_content.lower() + assert any(str(kw).lower() in audio_lower for kw in keywords), ( + "The output does not contain any of the keywords." + ) + + # Verify similarity (Whisper transcript vs streamed/detokenized text) + if "text" in modalities and "audio" in modalities: + assert response.similarity is not None and response.similarity > 0.9, ( + "The audio content is not same as the text" + ) + print(f"similarity is: {response.similarity}") + + +def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str) -> None: + assert response.success, "The request failed." + e2e_latency = getattr(response, "e2e_latency", None) + if e2e_latency is not None: + print(f"the avg e2e latency is: {e2e_latency}") + + req_fmt = request_config.get("response_format") + if req_fmt == "pcm" and response.audio_bytes: + _assert_pcm_int16_speech_hnr(response.audio_bytes) + if response.audio_format: + assert "pcm" in response.audio_format.lower(), ( + f"Expected audio/pcm content-type, got {response.audio_format!r}" + ) + elif req_fmt == "wav" and response.audio_format: + assert req_fmt in response.audio_format + + if run_level == "advanced_model" and req_fmt != "pcm": + expected_text = request_config.get("input") + if expected_text: + transcript = (response.audio_content or "").strip() + print(f"audio content is: {transcript}") + print(f"input text is: {expected_text}") + similarity = cosine_similarity_text(transcript.lower(), expected_text.lower()) + print(f"Cosine similarity: {similarity:.3f}") + assert similarity > 0.9, ( + f"Transcript doesn't match input: similarity={similarity:.2f}, transcript='{transcript}'" + ) + _assert_preset_voice_gender_from_audio(response.audio_bytes, request_config.get("voice")) + + +def assert_diffusion_response(response: Any, request_config: dict[str, Any], run_level: str = None): + assert response.success, "The request failed." + e2e_latency = getattr(response, "e2e_latency", None) + if e2e_latency is not None: + print(f"the avg e2e is: {e2e_latency}") + has_any_content = any(content is not None for content in (response.images, response.videos, response.audios)) + assert has_any_content, "Response contains no images, videos, or audios" + if response.images is not None: + assert_image_diffusion_response(response=response, request_config=request_config, run_level=run_level) + if response.videos is not None: + assert_video_diffusion_response(response=response, request_config=request_config, run_level=run_level) + if response.audios is not None: + assert_audio_diffusion_response(response=response, request_config=request_config, run_level=run_level) + + +__all__ = [ + "assert_audio_diffusion_response", + "assert_audio_speech_response", + "assert_diffusion_response", + "assert_image_diffusion_response", + "assert_image_valid", + "assert_omni_response", + "assert_video_diffusion_response", + "assert_video_valid", + "assert_audio_valid", +] diff --git a/tests/helpers/env.py b/tests/helpers/env.py new file mode 100644 index 000000000..22ec9a786 --- /dev/null +++ b/tests/helpers/env.py @@ -0,0 +1,280 @@ +"""Test environment / lifecycle helpers (GPU cleanup hooks and memory monitoring for tests). + +``vllm.platforms`` / ``vllm_omni.platforms`` are imported only inside functions that need them +so importing this module at pytest plugin load does not run before session autouse fixtures +""" + +from __future__ import annotations + +import gc +import os +import subprocess +import threading +import time +from contextlib import contextmanager + +import torch + + +def run_forced_gpu_cleanup_round() -> None: + run_pre_test_cleanup(enable_force=True) + run_post_test_cleanup(enable_force=True) + + +def get_physical_device_indices(devices): + visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible_devices is None: + return devices + visible_indices = [int(x) for x in visible_devices.split(",")] + index_mapping = {i: physical for i, physical in enumerate(visible_indices)} + return [index_mapping[i] for i in devices if i in index_mapping] + + +def wait_for_gpu_memory_to_clear( + *, + devices: list[int], + threshold_bytes: int | None = None, + threshold_ratio: float | None = None, + timeout_s: float = 120, +) -> None: + from vllm.platforms import current_platform + + assert threshold_bytes is not None or threshold_ratio is not None + devices = get_physical_device_indices(devices) + start_time = time.time() + + device_list = ", ".join(str(d) for d in devices) + if threshold_bytes is not None: + threshold_str = f"{threshold_bytes / 2**30:.2f} GiB" + condition_str = f"Memory usage ≤ {threshold_str}" + else: + threshold_percent = threshold_ratio * 100 + threshold_str = f"{threshold_percent:.1f}%" + condition_str = f"Memory usage ratio ≤ {threshold_str}" + + print(f"[GPU Memory Monitor] Waiting for GPU {device_list} to free memory, Condition: {condition_str}") + + if threshold_bytes is not None: + + def is_free(used, total): + return used <= threshold_bytes / 2**30 + else: + + def is_free(used, total): + return used / total <= threshold_ratio + + @contextmanager + def nvml_scope(): + if current_platform.is_rocm(): + from amdsmi import amdsmi_init, amdsmi_shut_down + + amdsmi_init() + try: + yield + finally: + amdsmi_shut_down() + elif current_platform.is_cuda(): + from vllm.third_party.pynvml import nvmlInit, nvmlShutdown + + nvmlInit() + try: + yield + finally: + nvmlShutdown() + else: + yield + + is_rocm = current_platform.is_rocm() + + with nvml_scope(): + if is_rocm: + from amdsmi import amdsmi_get_gpu_vram_usage, amdsmi_get_processor_handles + elif current_platform.is_cuda(): + from vllm.third_party.pynvml import nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo + + while True: + output: dict[int, str] = {} + output_raw: dict[int, tuple[float, float]] = {} + for device in devices: + if is_rocm: + dev_handle = amdsmi_get_processor_handles()[device] + mem_info = amdsmi_get_gpu_vram_usage(dev_handle) + gb_used = mem_info["vram_used"] / 2**10 + gb_total = mem_info["vram_total"] / 2**10 + else: + dev_handle = nvmlDeviceGetHandleByIndex(device) + mem_info = nvmlDeviceGetMemoryInfo(dev_handle) + gb_used = mem_info.used / 2**30 + gb_total = mem_info.total / 2**30 + output_raw[device] = (gb_used, gb_total) + usage_percent = (gb_used / gb_total) * 100 if gb_total > 0 else 0 + output[device] = f"{gb_used:.1f}GiB/{gb_total:.1f}GiB ({usage_percent:.1f}%)" + + print("[GPU Memory Status] Current usage:") + for device_id, mem_info in output.items(): + print(f" GPU {device_id}: {mem_info}") + + dur_s = time.time() - start_time + elapsed_minutes = dur_s / 60 + if all(is_free(used, total) for used, total in output_raw.values()): + print(f"[GPU Memory Freed] Devices {device_list} meet memory condition") + print(f" Condition: {condition_str}") + print(f" Wait time: {dur_s:.1f} seconds ({elapsed_minutes:.1f} minutes)") + break + + if dur_s >= timeout_s: + raise ValueError( + f"[GPU Memory Timeout] Devices {device_list} still don't meet memory condition after {dur_s:.1f} seconds\n" + f"Condition: {condition_str}\n" + f"Current status:\n" + "\n".join(f" GPU {device}: {output[device]}" for device in devices) + ) + + gc.collect() + torch.cuda.empty_cache() + time.sleep(5) + + +def _print_gpu_processes() -> None: + """Print GPU information including nvidia-smi and system processes.""" + + print("\n" + "=" * 80) + print("NVIDIA GPU Information (nvidia-smi)") + print("=" * 80) + + try: + nvidia_result = subprocess.run( + ["nvidia-smi"], + capture_output=True, + text=True, + timeout=5, + ) + + if nvidia_result.returncode == 0: + lines = nvidia_result.stdout.strip().split("\n") + for line in lines[:20]: + print(line) + + if len(lines) > 20: + print(f"... (showing first 20 of {len(lines)} lines)") + else: + print("nvidia-smi command failed") + + except (subprocess.TimeoutExpired, FileNotFoundError): + print("nvidia-smi not available or timed out") + except Exception as e: + print(f"Error running nvidia-smi: {e}") + + print("\n" + "=" * 80) + print("Detailed GPU Processes (nvidia-smi pmon)") + print("=" * 80) + + try: + pmon_result = subprocess.run( + ["nvidia-smi", "pmon", "-c", "1"], + capture_output=True, + text=True, + timeout=3, + ) + + if pmon_result.returncode == 0 and pmon_result.stdout.strip(): + print(pmon_result.stdout) + else: + print("No active GPU processes found via nvidia-smi pmon") + + except Exception: + print("nvidia-smi pmon not available") + + print("\n" + "=" * 80) + print("System Processes with GPU keywords") + print("=" * 80) + + +_SKIPPED_GPU_CLEANUP_MSG = ( + "\nSkipping GPU memory cleanup check (typically: instance already up; no check needed between tests)\n" +) + + +def run_pre_test_cleanup(enable_force: bool = False) -> None: + if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force: + print(_SKIPPED_GPU_CLEANUP_MSG) + return + + print("Pre-test GPU status:") + + num_gpus = torch.cuda.device_count() + if num_gpus > 0: + try: + wait_for_gpu_memory_to_clear( + devices=list(range(num_gpus)), + threshold_ratio=0.05, + ) + except Exception as e: + print(f"Pre-test cleanup note: {e}") + + +def run_post_test_cleanup(enable_force: bool = False) -> None: + if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force: + print(_SKIPPED_GPU_CLEANUP_MSG) + return + + if torch.cuda.is_available(): + gc.collect() + torch.cuda.empty_cache() + + print("Post-test GPU status:") + _print_gpu_processes() + + +class DeviceMemoryMonitor: + """Poll global device memory usage.""" + + def __init__(self, device_index: int, interval: float = 0.05): + self.device_index = device_index + self.interval = interval + self._peak_used_mb = 0.0 + self._stop_event = threading.Event() + self._thread: threading.Thread | None = None + + def start(self) -> None: + from vllm_omni.platforms import current_omni_platform + + def monitor_loop() -> None: + while not self._stop_event.is_set(): + try: + with current_omni_platform.device(self.device_index): + free_bytes, total_bytes = current_omni_platform.mem_get_info() + used_mb = (total_bytes - free_bytes) / (1024**2) + self._peak_used_mb = max(self._peak_used_mb, used_mb) + except Exception: + pass + time.sleep(self.interval) + + self._thread = threading.Thread(target=monitor_loop, daemon=False) + self._thread.start() + + def stop(self) -> None: + if self._thread is None: + return + self._stop_event.set() + self._thread.join(timeout=2.0) + + @property + def peak_used_mb(self) -> float: + from vllm_omni.platforms import current_omni_platform + + fallback_alloc = current_omni_platform.max_memory_allocated(device=self.device_index) / (1024**2) + fallback_reserved = current_omni_platform.max_memory_reserved(device=self.device_index) / (1024**2) + return max(self._peak_used_mb, fallback_alloc, fallback_reserved) + + def __del__(self): + self.stop() + + +__all__ = [ + "DeviceMemoryMonitor", + "get_physical_device_indices", + "run_post_test_cleanup", + "run_pre_test_cleanup", + "run_forced_gpu_cleanup_round", + "wait_for_gpu_memory_to_clear", +] diff --git a/tests/helpers/fixtures/__init__.py b/tests/helpers/fixtures/__init__.py new file mode 100644 index 000000000..8bd090b78 --- /dev/null +++ b/tests/helpers/fixtures/__init__.py @@ -0,0 +1 @@ +"""Pytest fixture modules under tests.helpers.""" diff --git a/tests/helpers/fixtures/env.py b/tests/helpers/fixtures/env.py new file mode 100644 index 000000000..939bad02c --- /dev/null +++ b/tests/helpers/fixtures/env.py @@ -0,0 +1,59 @@ +import os + +import pytest +import torch + + +@pytest.fixture(scope="session", autouse=True) +def default_env(): + # Keep behavior but avoid import-time side effects (RFC #2299). + keys = ("VLLM_WORKER_MULTIPROC_METHOD", "VLLM_TARGET_DEVICE") + previous = {key: os.environ.get(key) for key in keys} + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = previous["VLLM_WORKER_MULTIPROC_METHOD"] or "spawn" + os.environ["VLLM_TARGET_DEVICE"] = previous["VLLM_TARGET_DEVICE"] or ( + "cuda" if torch.cuda.is_available() and torch.cuda.device_count() > 0 else "cpu" + ) + yield + for key, value in previous.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +@pytest.fixture(scope="session") +def model_prefix() -> str: + prefix = os.environ.get("MODEL_PREFIX", "") + return f"{prefix.rstrip('/')}/" if prefix else "" + + +@pytest.fixture(autouse=True) +def clean_gpu_memory_between_tests(): + # Import here so ``tests.helpers.env`` (and vLLM platform modules) load only + # after session autouse fixtures like ``default_env`` have run (RFC #2299). + from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup + + print("\n=== PRE-TEST GPU CLEANUP ===") + run_pre_test_cleanup() + yield + run_post_test_cleanup() + + +@pytest.fixture(scope="session", autouse=True) +def default_vllm_config(): + """Set a default VllmConfig for the whole test session. + + Session scope ensures module-scoped fixtures (e.g. ``omni_runner``) and + deferred imports of ``tests.helpers.runtime`` both see the same context. + Function-scoped autouse ran too late for ``OmniRunner`` setup and could + desynchronize vLLM init vs request preprocessing (e.g. renderer state). + """ + from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config + + # Use CPU device if no GPU is available (e.g., in CI environments) + has_gpu = torch.cuda.is_available() and torch.cuda.device_count() > 0 + device = "cuda" if has_gpu else "cpu" + device_config = DeviceConfig(device=device) + + with set_current_vllm_config(VllmConfig(device_config=device_config)): + yield diff --git a/tests/helpers/fixtures/log.py b/tests/helpers/fixtures/log.py new file mode 100644 index 000000000..798fa4ae6 --- /dev/null +++ b/tests/helpers/fixtures/log.py @@ -0,0 +1,7 @@ +import pytest + + +@pytest.fixture(autouse=True) +def log_test_name_before_test(request: pytest.FixtureRequest): + print(f"--- Running test: {request.node.name}") + yield diff --git a/tests/helpers/fixtures/run_args.py b/tests/helpers/fixtures/run_args.py new file mode 100644 index 000000000..b18a64b98 --- /dev/null +++ b/tests/helpers/fixtures/run_args.py @@ -0,0 +1,17 @@ +import pytest + + +def pytest_addoption(parser): + parser.addoption( + "--run-level", + action="store", + default="core_model", + choices=["core_model", "advanced_model"], + help="Test level to run: L2, L3", + ) + + +@pytest.fixture(scope="session") +def run_level(request) -> str: + """Session test level from ``--run-level`` (see CI five-level docs).""" + return request.config.getoption("--run-level") diff --git a/tests/helpers/fixtures/runtime.py b/tests/helpers/fixtures/runtime.py new file mode 100644 index 000000000..77a468f2f --- /dev/null +++ b/tests/helpers/fixtures/runtime.py @@ -0,0 +1,137 @@ +"""Runtime fixtures (OmniRunner / OmniServer). Imports are deferred to fixture time. + +Loading ``tests.helpers.runtime`` at plugin import time (before session fixtures) +pulls in vLLM/vllm_omni too early and breaks initialization order vs the legacy +monolithic conftest. Defer imports until fixtures run so ``default_env`` / +``default_vllm_config`` run first. +""" + +from __future__ import annotations + +import threading +from collections.abc import Generator +from typing import Any + +import pytest +import yaml + +from tests.helpers.runtime import OmniServer +from tests.helpers.stage_config import modify_stage_config + +omni_fixture_lock = threading.Lock() + + +@pytest.fixture(scope="module") +def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: str) -> Generator[OmniServer, Any, None]: + """Start vLLM-Omni through the standard or stage-CLI launcher. + + The fixture stays module-scoped because multi-stage initialization is costly. + The ``use_stage_cli`` flag on ``OmniServerParams`` routes the setup through the + stage-CLI harness while still reusing the same fixture grouping semantics. + """ + with omni_fixture_lock: + from tests.helpers.runtime import OmniServer, OmniServerParams, OmniServerStageCli + + params: OmniServerParams = request.param + model = model_prefix + params.model + port = params.port + stage_config_path = params.stage_config_path + if run_level == "advanced_model" and stage_config_path is not None: + with open(stage_config_path, encoding="utf-8") as f: + cfg = yaml.safe_load(f) or {} + # Strip ``load_format: dummy`` (CI overlay default) so advanced_model + # tests use real weights. New schema (``stages:``) writes the field + # flat at stage level; legacy schema (``stage_args:``) nests it as + # ``engine_args.load_format``. Handle both. + new_schema_stages = cfg.get("stages") + stage_key = "stages" if new_schema_stages is not None else "stage_args" + delete_path = "load_format" if new_schema_stages is not None else "engine_args.load_format" + stage_entries = cfg.get(stage_key, []) + stage_ids = [stage["stage_id"] for stage in stage_entries if "stage_id" in stage] + stage_config_path = modify_stage_config( + stage_config_path, + deletes={stage_key: {stage_id: [delete_path] for stage_id in stage_ids}}, + ) + + server_args = params.server_args or [] + if params.use_omni and params.stage_init_timeout is not None: + server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)] + else: + server_args = [*server_args, "--stage-init-timeout", "600"] + if params.init_timeout is not None: + server_args = [*server_args, "--init-timeout", str(params.init_timeout)] + else: + server_args = [*server_args, "--init-timeout", "900"] + if params.use_stage_cli: + if not params.use_omni: + raise ValueError("omni_server with use_stage_cli=True requires use_omni=True") + if stage_config_path is None: + raise ValueError("omni_server with use_stage_cli=True requires a stage_config_path") + server_args += ["--stage-configs-path", stage_config_path] + + with OmniServerStageCli( + model, + stage_config_path, + server_args, + port=port, + env_dict=params.env_dict, + ) as server: + print("OmniServer started successfully") + yield server + print("OmniServer stopping...") + else: + if stage_config_path is not None: + server_args += ["--stage-configs-path", stage_config_path] + + with ( + OmniServer( + model, + server_args, + port=port, + env_dict=params.env_dict, + use_omni=params.use_omni, + ) + if port + else OmniServer( + model, + server_args, + env_dict=params.env_dict, + use_omni=params.use_omni, + ) + ) as server: + print("OmniServer started successfully") + yield server + print("OmniServer stopping...") + + print("OmniServer stopped") + + +@pytest.fixture +def openai_client(request: pytest.FixtureRequest, run_level: str): + """Resolve ``omni_server`` lazily so parametrized server fixtures work like upstream.""" + from tests.helpers.runtime import OpenAIClientHandler + + server = request.getfixturevalue("omni_server") + return OpenAIClientHandler(host=server.host, port=server.port, api_key="EMPTY", run_level=run_level) + + +@pytest.fixture(scope="module") +def omni_runner(request: pytest.FixtureRequest, model_prefix: str): + from tests.helpers.runtime import OmniRunner + + with omni_fixture_lock: + model, stage_config_path = request.param + model = model_prefix + model + with OmniRunner(model, seed=42, stage_configs_path=stage_config_path) as runner: + print("OmniRunner started successfully") + yield runner + print("OmniRunner stopping...") + + print("OmniRunner stopped") + + +@pytest.fixture +def omni_runner_handler(omni_runner: Any): + from tests.helpers.runtime import OmniRunnerHandler + + return OmniRunnerHandler(omni_runner) diff --git a/tests/helpers/mark.py b/tests/helpers/mark.py new file mode 100644 index 000000000..ed45dd7e9 --- /dev/null +++ b/tests/helpers/mark.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Pytest marks and decorators for hardware / resource selection (CUDA, ROCm, …).""" + +import pytest +from vllm.utils.torch_utils import cuda_device_count_stateless + +# Re-exported from tests.helpers.env (GPU wait + DeviceMemoryMonitor). + + +def cuda_marks(*, res: str, num_cards: int): + test_platform_detail = pytest.mark.cuda + if res == "L4": + test_resource = pytest.mark.L4 + elif res == "H100": + test_resource = pytest.mark.H100 + else: + raise ValueError(f"Invalid CUDA resource type: {res}. Supported: L4, H100") + marks = [test_resource, test_platform_detail] + if num_cards == 1: + return marks + test_distributed = pytest.mark.distributed_cuda(num_cards=num_cards) + test_skipif = pytest.mark.skipif_cuda( + cuda_device_count_stateless() < num_cards, + reason=f"Need at least {num_cards} CUDA GPUs to run the test.", + ) + return marks + [test_distributed, test_skipif] + + +def rocm_marks(*, res: str, num_cards: int): + test_platform_detail = pytest.mark.rocm + if res == "MI325": + test_resource = pytest.mark.MI325 + else: + raise ValueError(f"Invalid ROCm resource type: {res}. Supported: MI325") + marks = [test_resource, test_platform_detail] + if num_cards == 1: + return marks + test_distributed = pytest.mark.distributed_rocm(num_cards=num_cards) + return marks + [test_distributed] + + +def xpu_marks(*, res: str, num_cards: int): + test_platform_detail = pytest.mark.xpu + if res == "B60": + test_resource = pytest.mark.B60 + else: + raise ValueError(f"Invalid XPU resource type: {res}. Supported: B60") + marks = [test_resource, test_platform_detail] + if num_cards == 1: + return marks + test_distributed = pytest.mark.distributed_rocm(num_cards=num_cards) + return marks + [test_distributed] + + +def musa_marks(*, res: str, num_cards: int): + test_platform_detail = pytest.mark.musa + if res == "S5000": + test_resource = pytest.mark.S5000 + else: + raise ValueError(f"Invalid MUSA resource type: {res}. Supported: S5000") + marks = [test_resource, test_platform_detail] + if num_cards == 1: + return marks + test_distributed = pytest.mark.distributed_musa(num_cards=num_cards) + return marks + [test_distributed] + + +def gpu_marks(*, res: str, num_cards: int): + test_platform = pytest.mark.gpu + if res in ("L4", "H100"): + return [test_platform] + cuda_marks(res=res, num_cards=num_cards) + if res == "MI325": + return [test_platform] + rocm_marks(res=res, num_cards=num_cards) + if res == "B60": + return [test_platform] + xpu_marks(res=res, num_cards=num_cards) + if res == "S5000": + return [test_platform] + musa_marks(res=res, num_cards=num_cards) + raise ValueError(f"Invalid resource type: {res}. Supported: L4, H100, MI325, B60, S5000") + + +def npu_marks(*, res: str, num_cards: int): + test_platform = pytest.mark.npu + if res == "A2": + test_resource = pytest.mark.A2 + elif res == "A3": + test_resource = pytest.mark.A3 + else: + test_resource = None + if num_cards == 1: + return [mark for mark in [test_platform, test_resource] if mark is not None] + test_distributed = pytest.mark.distributed_npu(num_cards=num_cards) + return [mark for mark in [test_platform, test_resource, test_distributed] if mark is not None] + + +def hardware_marks(*, res: dict[str, str], num_cards: int | dict[str, int] = 1): + for platform, _ in res.items(): + if platform not in ("cuda", "rocm", "xpu", "npu", "musa"): + raise ValueError(f"Unsupported platform: {platform}") + if isinstance(num_cards, int): + num_cards_dict = {platform: num_cards for platform in res.keys()} + else: + num_cards_dict = num_cards + for platform in num_cards_dict.keys(): + if platform not in res: + raise ValueError(f"Platform '{platform}' in num_cards but not in res.") + for platform in res.keys(): + if platform not in num_cards_dict: + num_cards_dict[platform] = 1 + + all_marks: list[pytest.MarkDecorator] = [] + for platform, resource in res.items(): + cards = num_cards_dict[platform] + if platform in ("cuda", "rocm", "xpu"): + marks = gpu_marks(res=resource, num_cards=cards) + elif platform == "musa": + marks = musa_marks(res=resource, num_cards=cards) + elif platform == "npu": + marks = npu_marks(res=resource, num_cards=cards) + else: + raise ValueError(f"Unsupported platform: {platform}") + all_marks.extend(marks) + return all_marks + + +def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1): + all_marks = hardware_marks(res=res, num_cards=num_cards) + + def wrapper(f): + func = f + for mark in reversed(all_marks): + func = mark(func) + return func + + return wrapper diff --git a/tests/helpers/media.py b/tests/helpers/media.py new file mode 100644 index 000000000..4463acbbb --- /dev/null +++ b/tests/helpers/media.py @@ -0,0 +1,648 @@ +"""Synthetic media generation and media/text utilities for tests.""" + +import base64 +import concurrent.futures +import gc +import io +import logging +import math +import multiprocessing +import os +import random +import re +import subprocess +import tempfile +import time +import uuid +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +import numpy as np +import soundfile as sf +from PIL import Image + +logger = logging.getLogger(__name__) + + +def _resolve_synthetic_media_cache_dir(cache_dir: Path | str | None) -> Path: + if cache_dir is not None: + return Path(cache_dir).expanduser().resolve() + return Path(tempfile.gettempdir()) / "vllm_omni_test_synthetic_media" + + +def _np_array_from_mp4_bytes(video_bytes: bytes) -> np.ndarray: + """Decode MP4 bytes to a (T, H, W, 3) uint8 RGB stack (matches in-memory synthetic frames).""" + import cv2 + + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: + tmp.write(video_bytes) + path = tmp.name + cap = None + try: + cap = cv2.VideoCapture(path) + if not cap.isOpened(): + raise RuntimeError("Failed to open cached synthetic video for decode") + frames: list[np.ndarray] = [] + while True: + ok, frame_bgr = cap.read() + if not ok: + break + frames.append(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)) + if not frames: + raise RuntimeError("Cached synthetic video has no decodable frames") + return np.stack(frames, axis=0) + finally: + if cap is not None: + cap.release() + try: + os.unlink(path) + except OSError: + pass + + +def generate_synthetic_audio( + duration: int, + num_channels: int, + sample_rate: int = 48000, + *, + force_regenerate: bool = False, + cache_dir: Path | str | None = None, +) -> dict[str, Any]: + """ + Generate TTS speech with pyttsx3 and return base64 string. + + Caches the WAV under ``cache_dir`` when given, else under the default temp + subdirectory. Reuses the file when the same + ``duration`` / ``num_channels`` / ``sample_rate`` are requested unless + ``force_regenerate`` is true. + """ + root = _resolve_synthetic_media_cache_dir(cache_dir) + root.mkdir(parents=True, exist_ok=True) + cache_path = root / f"synth_audio_d{duration}_ch{num_channels}_sr{sample_rate}.wav" + + if not force_regenerate and cache_path.is_file(): + data, _sr = sf.read(str(cache_path), dtype="float32", always_2d=True) + audio_bytes = cache_path.read_bytes() + return { + "np_array": np.asarray(data, dtype=np.float32), + "base64": base64.b64encode(audio_bytes).decode("utf-8"), + "file_path": str(cache_path.resolve()), + } + + import pyttsx3 + + def _pick_voice(engine: pyttsx3.Engine) -> str | None: + voices = engine.getProperty("voices") + if not voices: + return None + + preferred_tokens = ( + "natural", + "jenny", + "sonia", + "susan", + "zira", + "aria", + "hazel", + "samantha", + "ava", + "allison", + "female", + "woman", + "english-us", + "en-us", + "english", + ) + discouraged_tokens = ( + "espeak", + "robot", + "mbrola", + "microsoft david", + "male", + "man", + ) + + best_voice = voices[0] + best_score = float("-inf") + for voice in voices: + voice_text = f"{getattr(voice, 'id', '')} {getattr(voice, 'name', '')}".lower() + voice_languages = " ".join( + lang.decode(errors="ignore") if isinstance(lang, bytes) else str(lang) + for lang in getattr(voice, "languages", []) + ).lower() + combined_text = f"{voice_text} {voice_languages}" + score = 0 + for idx, token in enumerate(preferred_tokens): + if token in combined_text: + score += 20 - idx + for token in discouraged_tokens: + if token in combined_text: + score -= 10 + if "english" in combined_text or "en_" in combined_text or "en-" in combined_text: + score += 4 + if "en-us" in combined_text or "english-us" in combined_text: + score += 4 + if score > best_score: + best_score = score + best_voice = voice + + return best_voice.id + + def _resample_audio(audio: np.ndarray, src_sr: int, dst_sr: int) -> np.ndarray: + if src_sr == dst_sr or len(audio) == 0: + return audio.astype(np.float32) + src_len = audio.shape[0] + dst_len = max(1, int(round(src_len * float(dst_sr) / float(src_sr)))) + src_idx = np.arange(src_len, dtype=np.float32) + dst_idx = np.linspace(0, src_len - 1, dst_len, dtype=np.float32) + resampled_channels: list[np.ndarray] = [] + for ch in range(audio.shape[1]): + resampled_channels.append(np.interp(dst_idx, src_idx, audio[:, ch]).astype(np.float32)) + return np.stack(resampled_channels, axis=1) + + def _match_channels(audio: np.ndarray, target_channels: int) -> np.ndarray: + current_channels = audio.shape[1] + if current_channels == target_channels: + return audio.astype(np.float32) + if target_channels == 1: + return np.mean(audio, axis=1, keepdims=True, dtype=np.float32) + if current_channels == 1: + return np.repeat(audio, target_channels, axis=1).astype(np.float32) + collapsed = np.mean(audio, axis=1, keepdims=True, dtype=np.float32) + return np.repeat(collapsed, target_channels, axis=1).astype(np.float32) + + def _trim_silence(audio: np.ndarray, threshold: float = 0.01) -> np.ndarray: + if len(audio) == 0: + return audio + energy = np.max(np.abs(audio), axis=1) + voiced = np.where(energy > threshold)[0] + if len(voiced) == 0: + return audio + start = max(0, int(voiced[0]) - int(sample_rate * 0.02)) + end = min(len(audio), int(voiced[-1]) + int(sample_rate * 0.04) + 1) + return audio[start:end] + + def _enhance_speech(audio: np.ndarray) -> np.ndarray: + if len(audio) == 0: + return audio.astype(np.float32) + enhanced = audio.astype(np.float32).copy() + enhanced -= np.mean(enhanced, axis=0, keepdims=True, dtype=np.float32) + if len(enhanced) > 1: + preemphasis = enhanced.copy() + preemphasis[1:] = enhanced[1:] - 0.94 * enhanced[:-1] + enhanced = 0.7 * enhanced + 0.3 * preemphasis + enhanced = np.sign(enhanced) * np.sqrt(np.abs(enhanced)) + fade = min(len(enhanced) // 4, max(1, int(sample_rate * 0.01))) + if fade > 1: + ramp_in = np.linspace(0.0, 1.0, fade, dtype=np.float32) + ramp_out = np.linspace(1.0, 0.0, fade, dtype=np.float32) + enhanced[:fade] *= ramp_in[:, None] + enhanced[-fade:] *= ramp_out[:, None] + peak = float(np.max(np.abs(enhanced))) + if peak > 1e-8: + enhanced = enhanced / peak * 0.95 + return enhanced.astype(np.float32) + + phrase_text = "test" + num_samples = int(sample_rate * max(1, duration)) + audio_data = np.zeros((num_samples, num_channels), dtype=np.float32) + + engine = pyttsx3.init() + engine.setProperty("rate", 112) + engine.setProperty("volume", 1.0) + selected_voice = _pick_voice(engine) + if selected_voice is not None: + engine.setProperty("voice", selected_voice) + + temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) + temp_wav.close() + try: + engine.save_to_file(phrase_text, temp_wav.name) + engine.runAndWait() + engine.stop() + + ready = False + for _ in range(50): + if os.path.exists(temp_wav.name) and os.path.getsize(temp_wav.name) > 44: + ready = True + break + time.sleep(0.1) + if not ready: + raise RuntimeError("pyttsx3 did not produce a WAV file in time.") + + tts_audio, tts_sr = sf.read(temp_wav.name, dtype="float32", always_2d=True) + finally: + if os.path.exists(temp_wav.name): + os.unlink(temp_wav.name) + + if len(tts_audio) == 0: + raise RuntimeError("pyttsx3 produced an empty WAV file.") + + tts_audio = _resample_audio(tts_audio, tts_sr, sample_rate) + tts_audio = _match_channels(tts_audio, num_channels) + tts_audio = _trim_silence(tts_audio, threshold=0.012) + tts_audio = _enhance_speech(tts_audio) + + lead_silence = min(int(sample_rate * 0.02), num_samples // 8) + pause_samples = int(sample_rate * 0.18) + start = lead_silence + phrase_len = tts_audio.shape[0] + while start < num_samples: + take = min(phrase_len, num_samples - start) + audio_data[start : start + take] = tts_audio[:take] + start += phrase_len + pause_samples + + max_amp = float(np.max(np.abs(audio_data))) + if max_amp > 0: + audio_data = audio_data / max_amp * 0.95 + + sf.write(str(cache_path), audio_data, sample_rate, format="WAV", subtype="PCM_16") + audio_bytes = cache_path.read_bytes() + + return { + "np_array": audio_data.copy(), + "base64": base64.b64encode(audio_bytes).decode("utf-8"), + "file_path": str(cache_path.resolve()), + } + + +def _mux_mp4_bytes_with_synthetic_audio( + video_mp4_bytes: bytes, + *, + num_frames: int, + fps: float = 30.0, + sample_rate: int = 48000, +) -> bytes: + duration_sec = num_frames / fps if fps > 0 else 0.0 + duration_int = max(1, int(math.ceil(duration_sec))) + + try: + audio_result = generate_synthetic_audio( + duration=duration_int, + num_channels=1, + sample_rate=sample_rate, + ) + audio_pcm = audio_result["np_array"] + except Exception as e: + logger.warning("Synthetic video: generate_synthetic_audio failed (%s); using video-only MP4.", e) + return video_mp4_bytes + + try: + import imageio_ffmpeg + + ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe() + except Exception: + ffmpeg_exe = "ffmpeg" + + try: + with tempfile.TemporaryDirectory(prefix="syn_vid_mux_") as tmp: + vid_path = os.path.join(tmp, "video.mp4") + wav_path = os.path.join(tmp, "audio.wav") + out_path = os.path.join(tmp, "out.mp4") + with open(vid_path, "wb") as f: + f.write(video_mp4_bytes) + sf.write(wav_path, audio_pcm, sample_rate, format="WAV", subtype="PCM_16") + cmd = [ + ffmpeg_exe, + "-y", + "-nostdin", + "-hide_banner", + "-loglevel", + "error", + "-i", + vid_path, + "-i", + wav_path, + "-c:v", + "copy", + "-c:a", + "aac", + "-b:a", + "128k", + "-shortest", + "-movflags", + "+faststart", + out_path, + ] + subprocess.run(cmd, check=True, stdin=subprocess.DEVNULL, timeout=300) + with open(out_path, "rb") as f: + return f.read() + except ( + FileNotFoundError, + subprocess.CalledProcessError, + subprocess.TimeoutExpired, + OSError, + ) as e: + logger.warning("Synthetic video: audio mux failed (%s); using video-only MP4.", e) + return video_mp4_bytes + + +def generate_synthetic_video( + width: int, + height: int, + num_frames: int, + *, + embed_audio: bool = False, + force_regenerate: bool = False, + cache_dir: Path | str | None = None, +) -> dict[str, Any]: + """ + Generate synthetic MP4 (optional AAC audio). Caches final bytes by + ``width`` / ``height`` / ``num_frames`` / ``embed_audio`` unless + ``force_regenerate`` is true. Cache root: ``cache_dir`` if given, else the + default temp subdirectory. + """ + root = _resolve_synthetic_media_cache_dir(cache_dir) + root.mkdir(parents=True, exist_ok=True) + cache_path = root / f"synth_video_w{width}_h{height}_nf{num_frames}_ea{int(embed_audio)}.mp4" + + if not force_regenerate and cache_path.is_file(): + video_bytes = cache_path.read_bytes() + return { + "np_array": _np_array_from_mp4_bytes(video_bytes), + "base64": base64.b64encode(video_bytes).decode("utf-8"), + "file_path": str(cache_path.resolve()), + } + + import cv2 + import imageio + + num_balls = random.randint(3, 8) + balls = [] + for _ in range(num_balls): + radius = min(width, height) // 8 + if radius < 1: + raise ValueError(f"Video dimensions ({width}x{height}) too small") + x = random.randint(radius, width - radius) + y = random.randint(radius, height - radius) + speed = random.uniform(3.0, 8.0) + angle = random.uniform(0, 2 * math.pi) + vx = speed * math.cos(angle) + vy = speed * math.sin(angle) + color_bgr = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255)) + balls.append({"x": x, "y": y, "vx": vx, "vy": vy, "radius": radius, "color_bgr": color_bgr}) + + video_frames = [] + for _ in range(num_frames): + frame_bgr = np.zeros((height, width, 3), dtype=np.uint8) + for ball in balls: + ball["x"] += ball["vx"] + ball["y"] += ball["vy"] + if ball["x"] - ball["radius"] <= 0 or ball["x"] + ball["radius"] >= width: + ball["vx"] = -ball["vx"] + ball["x"] = max(ball["radius"], min(width - ball["radius"], ball["x"])) + if ball["y"] - ball["radius"] <= 0 or ball["y"] + ball["radius"] >= height: + ball["vy"] = -ball["vy"] + ball["y"] = max(ball["radius"], min(height - ball["radius"], ball["y"])) + x, y = int(ball["x"]), int(ball["y"]) + radius = int(ball["radius"]) + cv2.circle(frame_bgr, (x, y), radius, ball["color_bgr"], -1) + frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) + video_frames.append(frame_rgb) + + fps = 30 + buffer = io.BytesIO() + writer_kwargs = { + "format": "mp4", + "fps": fps, + "codec": "libx264", + "quality": 7, + "pixelformat": "yuv420p", + "macro_block_size": 16, + "ffmpeg_params": ["-preset", "medium", "-crf", "23", "-movflags", "+faststart", "-pix_fmt", "yuv420p"], + } + try: + with imageio.get_writer(buffer, **writer_kwargs) as writer: + for frame in video_frames: + writer.append_data(frame) + buffer.seek(0) + video_only_bytes = buffer.read() + except Exception as e: + print(f"Warning: Failed to encode synthetic video: {e}") + raise + video_bytes = ( + _mux_mp4_bytes_with_synthetic_audio(video_only_bytes, num_frames=num_frames, fps=float(fps)) + if embed_audio + else video_only_bytes + ) + + cache_path.write_bytes(video_bytes) + + return { + "np_array": np.array(video_frames), + "base64": base64.b64encode(video_bytes).decode("utf-8"), + "file_path": str(cache_path.resolve()), + } + + +def generate_synthetic_image( + width: int, + height: int, + *, + force_regenerate: bool = False, + cache_dir: Path | str | None = None, +) -> dict[str, Any]: + """ + Random colored squares on white background. Caches JPEG by ``width`` / + ``height`` unless ``force_regenerate`` is true. Cache root: ``cache_dir`` + if given, else the default temp subdirectory. + """ + root = _resolve_synthetic_media_cache_dir(cache_dir) + root.mkdir(parents=True, exist_ok=True) + cache_path = root / f"synth_image_w{width}_h{height}.jpg" + + if not force_regenerate and cache_path.is_file(): + from PIL import Image as PILImage + + image = PILImage.open(cache_path) + image.load() + image_bytes = cache_path.read_bytes() + return { + "np_array": np.array(image).copy(), + "base64": base64.b64encode(image_bytes).decode("utf-8"), + "file_path": str(cache_path.resolve()), + } + + from PIL import ImageDraw + + image = Image.new("RGB", (width, height), (255, 255, 255)) + draw = ImageDraw.Draw(image) + num_squares = random.randint(3, 8) + for _ in range(num_squares): + square_size = random.randint(max(1, min(width, height) // 8), max(2, min(width, height) // 4)) + x = random.randint(0, max(0, width - square_size - 1)) + y = random.randint(0, max(0, height - square_size - 1)) + color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) + border_width = random.randint(1, 5) + draw.rectangle([x, y, x + square_size, y + square_size], fill=color, outline=(0, 0, 0), width=border_width) + + image.save(str(cache_path), format="JPEG", quality=85, optimize=True) + image_bytes = cache_path.read_bytes() + + return { + "np_array": np.array(image).copy(), + "base64": base64.b64encode(image_bytes).decode("utf-8"), + "file_path": str(cache_path.resolve()), + } + + +def decode_b64_image(b64: str): + img = Image.open(io.BytesIO(base64.b64decode(b64))) + img.load() + return img + + +def preprocess_text(text): + import opencc + + word_to_num = { + "zero": "0", + "one": "1", + "two": "2", + "three": "3", + "four": "4", + "five": "5", + "six": "6", + "seven": "7", + "eight": "8", + "nine": "9", + "ten": "10", + } + for word, num in word_to_num.items(): + pattern = r"\b" + re.escape(word) + r"\b" + text = re.sub(pattern, num, text, flags=re.IGNORECASE) + + text = re.sub(r"[^\w\s]", "", text) + text = re.sub(r"\s+", " ", text) + cc = opencc.OpenCC("t2s") + text = cc.convert(text) + text = re.sub(r"(?<=[\u4e00-\u9fff])\s+(?=[\u4e00-\u9fff])", "", text) + return text.lower().strip() + + +def cosine_similarity_text(text1, text2, n: int = 3): + from collections import Counter + + if not text1 or not text2: + return 0.0 + + text1 = preprocess_text(text1) + text2 = preprocess_text(text2) + print(f"cosine similarity text1 is: {text1}, text2 is: {text2}") + + ngrams1 = [text1[i : i + n] for i in range(len(text1) - n + 1)] + ngrams2 = [text2[i : i + n] for i in range(len(text2) - n + 1)] + counter1 = Counter(ngrams1) + counter2 = Counter(ngrams2) + + all_ngrams = set(counter1.keys()) | set(counter2.keys()) + vec1 = [counter1.get(ng, 0) for ng in all_ngrams] + vec2 = [counter2.get(ng, 0) for ng in all_ngrams] + dot_product = sum(a * b for a, b in zip(vec1, vec2)) + norm1 = sum(a * a for a in vec1) ** 0.5 + norm2 = sum(b * b for b in vec2) ** 0.5 + if norm1 == 0 or norm2 == 0: + return 0.0 + cosine = dot_product / (norm1 * norm2) + # Down-weight when lengths differ: repeated/hallucinated transcripts stay + # high in bag-of-ngrams cosine (e.g. ABCABCABC vs ABC) but should score low. + len1, len2 = len(text1), len(text2) + length_harmony = (2.0 * min(len1, len2)) / (len1 + len2) + return cosine * length_harmony + + +def _merge_base64_audio_to_segment(base64_list: list[str]): + from pydub import AudioSegment + + merged = None + for b64 in base64_list: + raw = base64.b64decode(b64.split(",", 1)[-1]) + seg = AudioSegment.from_file(io.BytesIO(raw)) + merged = seg if merged is None else merged + seg + return merged + + +@contextmanager +def _serialize_whisper_small_model_download(): + """Serialize Whisper ``small`` cache writes across processes (Linux/Unix).""" + import fcntl + + lock_path = Path.home() / ".cache" / "whisper" / ".small_model_download.lock" + lock_path.parent.mkdir(parents=True, exist_ok=True) + f = open(lock_path, "a+b") + try: + fcntl.flock(f.fileno(), fcntl.LOCK_EX) + yield + finally: + fcntl.flock(f.fileno(), fcntl.LOCK_UN) + f.close() + + +def _whisper_transcribe_in_current_process(output_path: str) -> str: + import whisper + + device_index = None + from vllm_omni.platforms import current_omni_platform + + if current_omni_platform.is_available(): + n = current_omni_platform.get_device_count() + if n == 1: + device_index = 0 + elif n > 1: + device_index = n - 1 + + if device_index is not None: + torch_device = current_omni_platform.get_torch_device(device_index) + current_omni_platform.set_device(torch_device) + device = str(torch_device) + use_accelerator = True + else: + use_accelerator = False + device = "cpu" + + with _serialize_whisper_small_model_download(): + model = whisper.load_model("small", device=device) + try: + text = model.transcribe( + output_path, + temperature=0.0, + word_timestamps=True, + condition_on_previous_text=False, + )["text"] + finally: + del model + gc.collect() + if use_accelerator: + current_omni_platform.synchronize() + current_omni_platform.empty_cache() + return text or "" + + +def convert_audio_file_to_text(output_path: str) -> str: + """Convert an audio file to text in an isolated subprocess.""" + ctx = multiprocessing.get_context("spawn") + with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor: + future = executor.submit(_whisper_transcribe_in_current_process, output_path) + return future.result() + + +def convert_audio_bytes_to_text(raw_bytes: bytes) -> str: + output_path = f"./test_{uuid.uuid4().hex}.wav" + data, samplerate = sf.read(io.BytesIO(raw_bytes)) + sf.write(output_path, data, samplerate, format="WAV", subtype="PCM_16") + print(f"audio data is saved: {output_path}") + return convert_audio_file_to_text(output_path) + + +__all__ = [ + "_merge_base64_audio_to_segment", + "convert_audio_bytes_to_text", + "convert_audio_file_to_text", + "cosine_similarity_text", + "decode_b64_image", + "generate_synthetic_audio", + "generate_synthetic_image", + "generate_synthetic_video", + "preprocess_text", +] diff --git a/tests/e2e/offline_inference/utils.py b/tests/helpers/process.py similarity index 58% rename from tests/e2e/offline_inference/utils.py rename to tests/helpers/process.py index 3113599a3..094de9652 100644 --- a/tests/e2e/offline_inference/utils.py +++ b/tests/helpers/process.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import contextlib import functools import os import signal @@ -10,73 +9,48 @@ import sys import tempfile from collections.abc import Callable from contextlib import ExitStack, suppress -from pathlib import Path from typing import Any, Literal import cloudpickle from typing_extensions import ParamSpec from vllm.platforms import current_platform -VLLM_PATH = Path(__file__).parent.parent.parent -"""Path to root of the vLLM repository.""" - - _P = ParamSpec("_P") def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, None]: - """Decorator to fork a new process for each test function. - See https://github.com/vllm-project/vllm/issues/7053 for more details. - """ + """Decorator to fork a new process for each test function.""" @functools.wraps(func) def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None: - # Make the process the leader of its own process group - # to avoid sending SIGTERM to the parent process os.setpgrp() from _pytest.outcomes import Skipped - # Create a unique temporary file to store exception info from child - # process. Use test function name and process ID to avoid collisions. with ( tempfile.NamedTemporaryFile( - delete=False, - mode="w+b", - prefix=f"vllm_test_{func.__name__}_{os.getpid()}_", - suffix=".exc", + delete=False, mode="w+b", prefix=f"vllm_test_{func.__name__}_{os.getpid()}_", suffix=".exc" ) as exc_file, ExitStack() as delete_after, ): exc_file_path = exc_file.name delete_after.callback(os.remove, exc_file_path) - pid = os.fork() - print(f"Fork a new process to run a test {pid}") if pid == 0: - # Parent process responsible for deleting, don't delete - # in child. delete_after.pop_all() try: func(*args, **kwargs) except Skipped as e: - # convert Skipped to exit code 0 print(str(e)) os._exit(0) except Exception as e: import traceback tb_string = traceback.format_exc() - - # Try to serialize the exception object first exc_to_serialize: dict[str, Any] try: - # First, try to pickle the actual exception with - # its traceback. exc_to_serialize = {"pickled_exception": e} - # Test if it can be pickled cloudpickle.dumps(exc_to_serialize) except (Exception, KeyboardInterrupt): - # Fall back to string-based approach. exc_to_serialize = { "exception_type": type(e).__name__, "exception_msg": str(e), @@ -86,7 +60,6 @@ def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, Non with open(exc_file_path, "wb") as f: cloudpickle.dump(exc_to_serialize, f) except Exception: - # Fallback: just print the traceback. print(tb_string) os._exit(1) else: @@ -94,40 +67,24 @@ def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, Non else: pgid = os.getpgid(pid) _pid, _exitcode = os.waitpid(pid, 0) - # ignore SIGTERM signal itself old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN) - # kill all child processes os.killpg(pgid, signal.SIGTERM) - # restore the signal handler signal.signal(signal.SIGTERM, old_signal_handler) if _exitcode != 0: - # Try to read the exception from the child process exc_info = {} if os.path.exists(exc_file_path): - with ( - contextlib.suppress(Exception), - open(exc_file_path, "rb") as f, - ): + with suppress(Exception), open(exc_file_path, "rb") as f: exc_info = cloudpickle.load(f) - - original_exception = exc_info.get("pickled_exception") - if original_exception is not None and isinstance(original_exception, Exception): - # Re-raise the actual exception object if it was - # successfully pickled. + if (original_exception := exc_info.get("pickled_exception")) is not None: + assert isinstance(original_exception, Exception) raise original_exception - if (original_tb := exc_info.get("traceback")) is not None: - # Use string-based traceback for fallback case raise AssertionError( - f"Test {func.__name__} failed when called with" - f" args {args} and kwargs {kwargs}" + f"Test {func.__name__} failed when called with args {args} and kwargs {kwargs}" f" (exit code: {_exitcode}):\n{original_tb}" ) from None - - # Fallback to the original generic error raise AssertionError( - f"function {func.__name__} failed when called with" - f" args {args} and kwargs {kwargs}" + f"function {func.__name__} failed when called with args {args} and kwargs {kwargs}" f" (exit code: {_exitcode})" ) from None @@ -139,9 +96,7 @@ def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None] @functools.wraps(f) def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None: - # Check if we're already in a subprocess if os.environ.get("RUNNING_IN_SUBPROCESS") == "1": - # If we are, just run the function directly return f(*args, **kwargs) import torch.multiprocessing as mp @@ -149,33 +104,18 @@ def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None] with suppress(RuntimeError): mp.set_start_method("spawn") - # Get the module module_name = f.__module__ - - # Create a process with environment variable set env = os.environ.copy() env["RUNNING_IN_SUBPROCESS"] = "1" with tempfile.TemporaryDirectory() as tempdir: output_filepath = os.path.join(tempdir, "new_process.tmp") - - # `cloudpickle` allows pickling complex functions directly input_bytes = cloudpickle.dumps((f, output_filepath)) - - repo_root = str(VLLM_PATH.resolve()) - - env = dict(env or os.environ) - env["PYTHONPATH"] = repo_root + os.pathsep + env.get("PYTHONPATH", "") - cmd = [sys.executable, "-m", f"{module_name}"] - returned = subprocess.run(cmd, input=input_bytes, capture_output=True, env=env) - - # check if the subprocess is successful try: returned.check_returncode() except Exception as e: - # wrap raised exception to provide more information raise RuntimeError(f"Error raised in subprocess:\n{returned.stderr.decode()}") from e return wrapper @@ -184,27 +124,11 @@ def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None] def create_new_process_for_each_test( method: Literal["spawn", "fork"] | None = None, ) -> Callable[[Callable[_P, None]], Callable[_P, None]]: - """Creates a decorator that runs each test function in a new process. - - Args: - method: The process creation method. Can be either "spawn" or "fork". - If not specified, it defaults to "spawn" on ROCm and XPU - platforms and "fork" otherwise. - - Returns: - A decorator to run test functions in separate processes. - """ + """Creates a decorator that runs each test function in a new process.""" if method is None: - # TODO: Find out why spawn is not working correctly on ROCm - # The test content will not run and tests passed immediately. - # For now, using `fork` for ROCm as it can run with `fork` - # and tests are running correctly. use_spawn = current_platform.is_xpu() method = "spawn" if use_spawn else "fork" - assert method in ["spawn", "fork"], "Method must be either 'spawn' or 'fork'" - if method == "fork": return fork_new_process_for_each_test - return spawn_new_process_for_each_test diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py new file mode 100644 index 000000000..9520aff0c --- /dev/null +++ b/tests/helpers/runtime.py @@ -0,0 +1,1413 @@ +"""Server/client/runner runtime primitives for tests.""" + +import base64 +import concurrent.futures +import io +import json +import os +import socket +import subprocess +import sys +import tempfile +import time +from dataclasses import dataclass +from io import BytesIO +from pathlib import Path +from typing import Any, NamedTuple + +import psutil +import requests +import soundfile as sf +import torch +import yaml +from openai import OpenAI, omit +from PIL import Image +from vllm import TextPrompt +from vllm.distributed.parallel_state import cleanup_dist_env_and_memory +from vllm.logger import init_logger + +from tests.helpers.assertions import ( + assert_audio_speech_response, + assert_diffusion_response, + assert_omni_response, +) +from tests.helpers.env import run_forced_gpu_cleanup_round +from tests.helpers.media import ( + _merge_base64_audio_to_segment, + convert_audio_bytes_to_text, + cosine_similarity_text, + decode_b64_image, +) +from vllm_omni.config.stage_config import resolve_deploy_yaml +from vllm_omni.platforms import current_omni_platform + +logger = init_logger(__name__) + +PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None +PromptImageInput = list[Any] | Any | None +PromptVideoInput = list[Any] | Any | None + +try: + from vllm.distributed.parallel_state import cleanup_dist_env_and_memory # type: ignore +except Exception: # pragma: no cover + + def cleanup_dist_env_and_memory() -> None: + return None + + +def get_open_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return int(s.getsockname()[1]) + + +def dummy_messages_from_mix_data( + system_prompt: dict[str, Any] = None, + video_data_url: Any = None, + audio_data_url: Any = None, + image_data_url: Any = None, + content_text: str = None, +): + """Create messages with video、image、audio data URL for OpenAI API.""" + if content_text is not None: + content = [{"type": "text", "text": content_text}] + else: + content = [] + + media_items = [] + if isinstance(video_data_url, list): + for video_url in video_data_url: + media_items.append((video_url, "video")) + else: + media_items.append((video_data_url, "video")) + + if isinstance(image_data_url, list): + for url in image_data_url: + media_items.append((url, "image")) + else: + media_items.append((image_data_url, "image")) + + if isinstance(audio_data_url, list): + for url in audio_data_url: + media_items.append((url, "audio")) + else: + media_items.append((audio_data_url, "audio")) + + content.extend( + {"type": f"{media_type}_url", f"{media_type}_url": {"url": url}} + for url, media_type in media_items + if url is not None + ) + messages = [{"role": "user", "content": content}] + if system_prompt is not None: + messages = [system_prompt] + messages + return messages + + +def _omni_subprocess_cwd() -> str: + """Repo root for ``python -m vllm_omni...`` (legacy conftest lived under ``tests/``; helpers under ``tests/helpers/``).""" + return os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "..")) + + +class OmniServerParams(NamedTuple): + model: str + port: int | None = None + stage_config_path: str | None = None + server_args: list[str] | None = None + env_dict: dict[str, str] | None = None + use_omni: bool = True + use_stage_cli: bool = False + init_timeout: int | None = None + stage_init_timeout: int | None = None # None: fixture supplies default (600 s) + + +class OmniServer: + """Omniserver for vLLM-Omni tests.""" + + def __init__( + self, + model: str, + serve_args: list[str], + *, + port: int | None = None, + env_dict: dict[str, str] | None = None, + use_omni: bool = True, + ) -> None: + run_forced_gpu_cleanup_round() + cleanup_dist_env_and_memory() + self.model = model + self.serve_args = serve_args + self.env_dict = env_dict + self.use_omni = use_omni + self.proc: subprocess.Popen | None = None + self.host = "127.0.0.1" + self.port = get_open_port() if port is None else port + + def _start_server(self) -> None: + env = os.environ.copy() + env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + if self.env_dict is not None: + env.update(self.env_dict) + + cmd = [ + sys.executable, + "-m", + "vllm_omni.entrypoints.cli.main", + "serve", + self.model, + "--host", + self.host, + "--port", + str(self.port), + ] + if self.use_omni: + cmd.append("--omni") + cmd += self.serve_args + + print(f"Launching OmniServer with: {' '.join(cmd)}") + self.proc = subprocess.Popen( + cmd, + env=env, + cwd=_omni_subprocess_cwd(), + ) + + max_wait = 1200 + start_time = time.time() + while time.time() - start_time < max_wait: + ret = self.proc.poll() + if ret is not None: + raise RuntimeError(f"Server processes exited with code {ret} before becoming ready.") + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(1) + if sock.connect_ex((self.host, self.port)) == 0: + print(f"Server ready on {self.host}:{self.port}") + return + time.sleep(2) + raise RuntimeError(f"Server failed to start within {max_wait} seconds") + + def _kill_process_tree(self, pid): + try: + parent = psutil.Process(pid) + children = parent.children(recursive=True) + all_pids = [pid] + [child.pid for child in children] + + for child in children: + try: + child.terminate() + except psutil.NoSuchProcess: + pass + + _, still_alive = psutil.wait_procs(children, timeout=10) + + for child in still_alive: + try: + child.kill() + except psutil.NoSuchProcess: + pass + + try: + parent.terminate() + parent.wait(timeout=10) + except (psutil.NoSuchProcess, psutil.TimeoutExpired): + try: + parent.kill() + except psutil.NoSuchProcess: + pass + + time.sleep(1) + alive_processes = [] + for check_pid in all_pids: + if psutil.pid_exists(check_pid): + alive_processes.append(check_pid) + + if alive_processes: + print(f"Warning: Processes still alive: {alive_processes}") + for alive_pid in alive_processes: + try: + subprocess.run(["kill", "-9", str(alive_pid)], timeout=2) + except Exception as e: + print(f"Cleanup failed: {e}") + + except psutil.NoSuchProcess: + pass + + def __enter__(self): + self._start_server() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.proc: + self._kill_process_tree(self.proc.pid) + run_forced_gpu_cleanup_round() + cleanup_dist_env_and_memory() + + +class OmniServerStageCli(OmniServer): + """Omni server harness that exercises the stage CLI flow.""" + + def __init__( + self, + model: str, + stage_config_path: str, + serve_args: list[str] | None = None, + *, + stage_ids: list[int] | None = None, + port: int | None = None, + env_dict: dict[str, str] | None = None, + ) -> None: + super().__init__(model, serve_args or [], port=port, env_dict=env_dict, use_omni=True) + self.stage_config_path = stage_config_path + self.master_port = get_open_port() + self.visible_device_list = self._load_visible_device_list(env_dict) + resolved_cfg = resolve_deploy_yaml(stage_config_path) + # Dump the resolved deploy config so CI logs show each stage's + # gpu_memory_utilization / max_model_len / max_num_seqs after + # base_config inheritance and overlay merge — essential when + # diagnosing OOMs that depend on the merged values. + print( + f"[OmniServerStageCli] Resolved deploy config from {stage_config_path}:\n" + f"{yaml.safe_dump(resolved_cfg, sort_keys=False, default_flow_style=False)}", + flush=True, + ) + self.stage_runtime_devices = self._load_stage_runtime_devices(resolved_cfg) + self.stage_ids = stage_ids or self._load_stage_ids(resolved_cfg) + if 0 not in self.stage_ids: + raise ValueError(f"Stage CLI test requires stage_id=0 in config: {stage_config_path}") + self.stage_procs: dict[int, subprocess.Popen] = {} + self.proc = None + + @staticmethod + def _stage_entries(cfg: dict) -> list[dict]: + """Return the list of stage entries from either legacy (``stage_args``) + or new-schema (``stages``) deploy YAMLs.""" + return cfg.get("stage_args") or cfg.get("stages") or [] + + @staticmethod + def _load_stage_ids(resolved_config: dict) -> list[int]: + stage_ids = [ + stage["stage_id"] for stage in OmniServerStageCli._stage_entries(resolved_config) if "stage_id" in stage + ] + if not stage_ids: + raise ValueError("No stage IDs found in resolved config") + return stage_ids + + @staticmethod + def _load_stage_runtime_devices(resolved_config: dict) -> dict[int, str]: + runtime_devices: dict[int, str] = {} + for stage in OmniServerStageCli._stage_entries(resolved_config): + stage_id = stage.get("stage_id") + # New schema: stage.devices is flat at stage level. + # Legacy schema: stage.runtime.devices is nested. + devices = stage.get("devices") or stage.get("runtime", {}).get("devices") + if stage_id is not None and devices: + runtime_devices[int(stage_id)] = str(devices) + return runtime_devices + + @classmethod + def _parse_device_list(cls, devices: str | int) -> list[str]: + if isinstance(devices, int): + if devices < 0: + raise ValueError("Device IDs must be non-negative integers") + return [str(devices)] + return [token.strip() for token in str(devices).split(",") if token.strip()] + + @classmethod + def _load_visible_device_list(cls, env_dict: dict[str, str] | None) -> list[str] | None: + env = os.environ.copy() + if env_dict is not None: + env.update(env_dict) + + env_var = getattr(current_omni_platform, "device_control_env_var", None) + if env_var and env_var in env: + return [token.strip() for token in env[env_var].split(",") if token.strip()] + return None + + @classmethod + def _map_stage_devices(cls, stage_id: int, visible_device_list: list[str] | None, devices: str) -> str: + device_list = cls._parse_device_list(devices) + + if visible_device_list is None: + return ",".join(device_list) + + if not all(device.isdigit() for device in device_list): + raise ValueError("Logical devices must be non-negative integers") + + logical_ids = [int(device) for device in device_list] + if logical_ids and max(logical_ids) >= len(visible_device_list): + raise ValueError( + f"Stage {stage_id} has logical IDs {device_list}, one or more of which exceed the number of visible devices" + ) + + return ",".join(visible_device_list[idx] for idx in logical_ids) + + def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str) -> None: + mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, devices) + env_var = getattr(current_omni_platform, "device_control_env_var", None) + if env_var: + env[env_var] = mapped_devices + + def _build_stage_cmd(self, stage_id: int, *, headless: bool) -> list[str]: + cmd = [ + sys.executable, + "-m", + "vllm_omni.entrypoints.cli.main", + "serve", + self.model, + "--omni", + "--stage-configs-path", + self.stage_config_path, + "--stage-id", + str(stage_id), + "--omni-master-address", + self.host, + "--omni-master-port", + str(self.master_port), + ] + + if headless: + cmd.append("--headless") + else: + cmd += ["--host", self.host, "--port", str(self.port)] + + cmd += self.serve_args + return cmd + + def _launch_stage(self, stage_id: int, *, headless: bool) -> None: + env = os.environ.copy() + env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + if self.env_dict is not None: + env.update(self.env_dict) + + devices = self.stage_runtime_devices.get(stage_id) + if devices: + self._set_stage_device_env(stage_id, env, devices) + + cmd = self._build_stage_cmd(stage_id, headless=headless) + print(f"Launching OmniServerStageCli stage {stage_id}: {' '.join(cmd)}") + # Capture each subprocess's stdout+stderr to a per-stage log file so + # debugging "Stage N exited before API server ready" doesn't rely on + # guessing; the file is surfaced in the RuntimeError message. + log_path = Path(tempfile.gettempdir()) / f"omni_stage_{stage_id}_{self.master_port}.log" + self._stage_log_paths = getattr(self, "_stage_log_paths", {}) + self._stage_log_paths[stage_id] = log_path + log_fh = open(log_path, "w", buffering=1) # noqa: SIM115 - closed in __exit__ + self._stage_log_files = getattr(self, "_stage_log_files", {}) + self._stage_log_files[stage_id] = log_fh + proc = subprocess.Popen( + cmd, + env=env, + cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + stdout=log_fh, + stderr=subprocess.STDOUT, + ) + self.stage_procs[stage_id] = proc + if stage_id == 0: + self.proc = proc + + def _ensure_stage_processes_alive(self) -> None: + for stage_id, proc in self.stage_procs.items(): + ret = proc.poll() + if ret is not None: + log_path = getattr(self, "_stage_log_paths", {}).get(stage_id) + tail = "" + if log_path and log_path.exists(): + try: + with open(log_path, encoding="utf-8", errors="replace") as f: + lines = f.readlines() + tail = "\n=== Last 60 lines of stage {} log ({}) ===\n{}".format( + stage_id, log_path, "".join(lines[-60:]) or "" + ) + except Exception as exc: # pragma: no cover - diagnostic only + tail = f"\n" + raise RuntimeError(f"Stage {stage_id} exited with code {ret} before API server became ready.{tail}") + + def _start_server(self) -> None: + ordered_stage_ids = [0, *[stage_id for stage_id in self.stage_ids if stage_id != 0]] + + self._launch_stage(0, headless=False) + time.sleep(2) + self._ensure_stage_processes_alive() + + for stage_id in ordered_stage_ids[1:]: + self._launch_stage(stage_id, headless=True) + + max_wait = 1200 + start_time = time.time() + while time.time() - start_time < max_wait: + self._ensure_stage_processes_alive() + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(1) + result = sock.connect_ex((self.host, self.port)) + if result == 0: + print(f"OmniServerStageCli ready on {self.host}:{self.port}") + return + time.sleep(2) + + raise RuntimeError(f"OmniServerStageCli failed to start within {max_wait} seconds") + + def _dump_stage_logs_for_debug(self, head_lines: int = 300, tail_lines: int = 500) -> None: + """Tail each stage's subprocess log back to stdout on teardown. + + Stage subprocesses redirect stdout/stderr to ``/tmp/omni_stage_*.log`` + so we don't spam the main CI stream while tests run; but that also + hides engine init (KV cache size, Available KV cache memory, vLLM + engine config) when things go wrong. Dump them here so buildkite + captures them post-run. Head covers engine init; tail covers + whatever state the stage was in when it was torn down. + """ + log_paths = getattr(self, "_stage_log_paths", {}) or {} + for stage_id in sorted(log_paths): + log_path = log_paths[stage_id] + if not log_path or not log_path.exists(): + continue + try: + with open(log_path, encoding="utf-8", errors="replace") as f: + lines = f.readlines() + except Exception as exc: # pragma: no cover - diagnostic only + print(f"[OmniServerStageCli] stage {stage_id} log read failed: {exc}", flush=True) + continue + total = len(lines) + if total <= head_lines + tail_lines: + head_chunk = lines + tail_chunk = [] + elided = 0 + else: + head_chunk = lines[:head_lines] + tail_chunk = lines[-tail_lines:] + elided = total - head_lines - tail_lines + print(f"\n=== stage {stage_id} log HEAD ({log_path}) ===", flush=True) + print("".join(head_chunk).rstrip("\n"), flush=True) + if tail_chunk: + print(f"\n... [{elided} lines elided] ...", flush=True) + print(f"\n=== stage {stage_id} log TAIL ({log_path}) ===", flush=True) + print("".join(tail_chunk).rstrip("\n"), flush=True) + print(f"=== end stage {stage_id} log ===\n", flush=True) + + def __exit__(self, exc_type, exc_val, exc_tb): + self._dump_stage_logs_for_debug() + for stage_id in sorted(self.stage_procs, reverse=True): + proc = self.stage_procs[stage_id] + if proc.poll() is None: + self._kill_process_tree(proc.pid) + run_forced_gpu_cleanup_round() + cleanup_dist_env_and_memory() + + +@dataclass +class OmniResponse: + text_content: str | None = None + audio_data: list[str] | None = None + audio_content: str | None = None + audio_format: str | None = None + audio_bytes: bytes | None = None + similarity: float | None = None + e2e_latency: float | None = None + success: bool = False + error_message: str | None = None + cached_tokens: int | None = None + + +@dataclass +class DiffusionResponse: + text_content: str | None = None + images: list[Image.Image] | None = None + audios: list[Any] | None = None + videos: list[Any] | None = None + e2e_latency: float | None = None + success: bool = False + error_message: str | None = None + + +class OpenAIClientHandler: + def __init__(self, host: str = "127.0.0.1", port: int = None, api_key: str = "EMPTY", run_level: str = None): + if port is None: + port = get_open_port() + self.base_url = f"http://{host}:{port}" + self.client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key=api_key) + self.run_level = run_level + + def _process_stream_omni_response(self, chat_completion) -> OmniResponse: + result = OmniResponse() + start_time = time.perf_counter() + try: + text_content = "" + audio_data = [] + for chunk in chat_completion: + for choice in chunk.choices: + content = getattr(getattr(choice, "delta", None), "content", None) + modality = getattr(chunk, "modality", None) + if modality == "audio" and content: + audio_data.append(content) + elif modality == "text" and content: + text_content += content + result.e2e_latency = time.perf_counter() - start_time + audio_content = None + similarity = None + if audio_data: + merged_seg = _merge_base64_audio_to_segment(audio_data) + wav_buf = BytesIO() + merged_seg.export(wav_buf, format="wav") + result.audio_bytes = wav_buf.getvalue() + audio_content = convert_audio_bytes_to_text(result.audio_bytes) + if audio_content and text_content: + similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) + result.text_content = text_content + result.audio_data = audio_data + result.audio_content = audio_content + result.similarity = similarity + result.success = True + except Exception as e: + result.error_message = f"Stream processing error: {str(e)}" + print(f"Error: {result.error_message}") + return result + + def _process_non_stream_omni_response(self, chat_completion) -> OmniResponse: + result = OmniResponse() + start_time = time.perf_counter() + try: + audio_data = None + text_content = None + for choice in chat_completion.choices: + if hasattr(choice.message, "audio") and choice.message.audio is not None: + audio_data = choice.message.audio.data + if hasattr(choice.message, "content") and choice.message.content is not None: + text_content = choice.message.content + # Extract cached_tokens for prefix caching tests + usage = getattr(chat_completion, "usage", None) + if usage and (details := getattr(usage, "prompt_tokens_details", None)): + result.cached_tokens = details.cached_tokens + result.e2e_latency = time.perf_counter() - start_time + audio_content = None + similarity = None + if audio_data: + result.audio_bytes = base64.b64decode(audio_data) + audio_content = convert_audio_bytes_to_text(result.audio_bytes) + if audio_content and text_content: + similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) + result.text_content = text_content + result.audio_content = audio_content + result.similarity = similarity + result.success = True + except Exception as e: + result.error_message = f"Non-stream processing error: {str(e)}" + print(f"Error: {result.error_message}") + return result + + def _process_diffusion_response(self, chat_completion) -> DiffusionResponse: + result = DiffusionResponse() + start_time = time.perf_counter() + try: + images = [] + for choice in chat_completion.choices: + content = getattr(choice.message, "content", None) + if isinstance(content, list): + for item in content: + image_url = None + if isinstance(item, dict): + image_url = item.get("image_url", {}).get("url") + else: + image_url_obj = getattr(item, "image_url", None) + image_url = getattr(image_url_obj, "url", None) if image_url_obj else None + if image_url and image_url.startswith("data:image"): + b64_data = image_url.split(",", 1)[1] + images.append(decode_b64_image(b64_data)) + result.e2e_latency = time.perf_counter() - start_time + result.images = images if images else None + result.success = True + except Exception as e: + result.error_message = f"Diffusion response processing error: {str(e)}" + print(f"Error: {result.error_message}") + return result + + def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: + responses: list[OmniResponse] = [] + stream = request_config.get("stream", False) + modalities = request_config.get("modalities", ["text", "audio"]) + extra_body: dict[str, Any] = {} + if "speaker" in request_config: + extra_body["speaker"] = request_config["speaker"] + if request_config.get("use_audio_in_video"): + mm = dict(extra_body.get("mm_processor_kwargs") or {}) + mm["use_audio_in_video"] = True + extra_body["mm_processor_kwargs"] = mm + create_kwargs: dict[str, Any] = { + "model": request_config.get("model"), + "messages": request_config.get("messages"), + "stream": stream, + "modalities": modalities, + } + if extra_body: + create_kwargs["extra_body"] = extra_body + + if request_num == 1: + chat_completion = self.client.chat.completions.create(**create_kwargs) + resp = ( + self._process_stream_omni_response(chat_completion) + if stream + else self._process_non_stream_omni_response(chat_completion) + ) + assert_omni_response(resp, request_config, run_level=self.run_level) + responses.append(resp) + return responses + + def _one(): + chat_completion = self.client.chat.completions.create(**create_kwargs) + return ( + self._process_stream_omni_response(chat_completion) + if stream + else self._process_non_stream_omni_response(chat_completion) + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: + futures = [executor.submit(_one) for _ in range(request_num)] + for future in concurrent.futures.as_completed(futures): + resp = future.result() + assert_omni_response(resp, request_config, run_level=self.run_level) + responses.append(resp) + return responses + + def _process_stream_audio_speech_response(self, response, *, response_format: str | None = None) -> OmniResponse: + """ + Process streaming /v1/audio/speech responses into an OmniResponse. + + This mirrors _process_stream_omni_response but operates on low-level + audio bytes and produces an OmniResponse with audio_content filled + from Whisper transcription. + """ + result = OmniResponse() + start_time = time.perf_counter() + + try: + # Aggregate all audio bytes from the streaming response. + data = bytearray() + + # Preferred OpenAI helper. + if hasattr(response, "iter_bytes") and callable(getattr(response, "iter_bytes")): + for chunk in response.iter_bytes(): + if chunk: + data.extend(chunk) + else: + # Generic iterable-of-bytes fallback (e.g., generator or list of chunks). + try: + iterator = iter(response) + except TypeError: + iterator = None + + if iterator is not None: + for chunk in iterator: + if not chunk: + continue + if isinstance(chunk, (bytes, bytearray)): + data.extend(chunk) + elif hasattr(chunk, "data"): + data.extend(chunk.data) # type: ignore[arg-type] + elif hasattr(chunk, "content"): + data.extend(chunk.content) # type: ignore[arg-type] + else: + raise TypeError(f"Unsupported stream chunk type: {type(chunk)}") + else: + raise TypeError(f"Unsupported audio speech streaming response type: {type(response)}") + + raw_bytes = bytes(data) + if response_format == "pcm": + transcript = None + else: + transcript = convert_audio_bytes_to_text(raw_bytes) + + # Populate OmniResponse. + result.audio_bytes = raw_bytes + result.audio_content = transcript + result.e2e_latency = time.perf_counter() - start_time + result.success = True + result.audio_format = getattr(response, "response", None) + if result.audio_format is not None: + result.audio_format = result.audio_format.headers.get("content-type", "") + + except Exception as e: + result.error_message = f"Audio speech stream processing error: {str(e)}" + print(f"Error: {result.error_message}") + + return result + + def _process_non_stream_audio_speech_response( + self, response, *, response_format: str | None = None + ) -> OmniResponse: + """ + Process non-streaming /v1/audio/speech responses into an OmniResponse. + + This mirrors _process_non_stream_omni_response but for the binary + audio payload returned by audio.speech.create. + """ + result = OmniResponse() + start_time = time.perf_counter() + + try: + # OpenAI non-streaming audio.speech.create returns HttpxBinaryResponseContent (.read() or .content) + if hasattr(response, "read") and callable(getattr(response, "read")): + raw_bytes = response.read() + elif hasattr(response, "content"): + raw_bytes = response.content # type: ignore[assignment] + else: + raise TypeError(f"Unsupported audio speech response type: {type(response)}") + + if response_format == "pcm": + transcript = None + else: + transcript = convert_audio_bytes_to_text(raw_bytes) + + result.audio_bytes = raw_bytes + result.audio_content = transcript + result.e2e_latency = time.perf_counter() - start_time + result.success = True + result.audio_format = getattr(response, "response", None) + if result.audio_format is not None: + result.audio_format = result.audio_format.headers.get("content-type", "") + + except Exception as e: + result.error_message = f"Audio speech non-stream processing error: {str(e)}" + print(f"Error: {result.error_message}") + + return result + + def send_audio_speech_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: + """ + Call the /v1/audio/speech endpoint using the same configuration-dict + style as send_omni_request, but via the OpenAI Python client's + audio.speech APIs. + + Expected keys in request_config: + - model: model name/path (required) + - input: text to synthesize (required) + - response_format: audio format such as "wav" or "pcm" (optional) + - task_type, ref_text, ref_audio: TTS-specific extras (optional, passed via extra_body) + - timeout: request timeout in seconds (float, optional, default 120.0) + - stream: whether to use streaming API (bool, optional, default False) + """ + timeout = float(request_config.get("timeout", 120.0)) + + model = request_config["model"] + text_input = request_config["input"] + stream = bool(request_config.get("stream", False)) + voice = request_config.get("voice", None) + + # Standard OpenAI param: use omit when not provided to keep default behavior. + response_format = request_config.get("response_format", omit) + + # Qwen3-TTS custom fields, forwarded via extra_body. + extra_body: dict[str, Any] = {} + # Keep this list aligned with vllm_omni.entrypoints.openai.protocol.audio params. + for key in ("task_type", "ref_text", "ref_audio", "language", "max_new_tokens"): + if key in request_config: + extra_body[key] = request_config[key] + + responses: list[OmniResponse] = [] + + speech_fmt: str | None = None if response_format is omit else str(response_format).lower() + + if request_num == 1: + if stream: + # Use streaming response helper. + with self.client.audio.speech.with_streaming_response.create( + model=model, + input=text_input, + response_format=response_format, + extra_body=extra_body or None, + timeout=timeout, + voice=voice, + ) as resp: + omni_resp = self._process_stream_audio_speech_response(resp, response_format=speech_fmt) + else: + # Non-streaming response. + resp = self.client.audio.speech.create( + model=model, + input=text_input, + response_format=response_format, + extra_body=extra_body or None, + timeout=timeout, + voice=voice, + ) + omni_resp = self._process_non_stream_audio_speech_response(resp, response_format=speech_fmt) + + assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) + responses.append(omni_resp) + return responses + else: + # request_num > 1: concurrent requests (use same params as single-request path) + + if stream: + + def _stream_task(): + with self.client.audio.speech.with_streaming_response.create( + model=model, + input=text_input, + response_format=response_format, + extra_body=extra_body or None, + timeout=timeout, + voice=voice, + ) as resp: + return self._process_stream_audio_speech_response(resp, response_format=speech_fmt) + + with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: + futures = [executor.submit(_stream_task) for _ in range(request_num)] + for future in concurrent.futures.as_completed(futures): + omni_resp = future.result() + assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) + responses.append(omni_resp) + else: + with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: + futures = [] + for _ in range(request_num): + future = executor.submit( + self.client.audio.speech.create, + model=model, + input=text_input, + response_format=response_format, + extra_body=extra_body or None, + timeout=timeout, + voice=voice, + ) + futures.append(future) + + for future in concurrent.futures.as_completed(futures): + resp = future.result() + omni_resp = self._process_non_stream_audio_speech_response(resp, response_format=speech_fmt) + assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) + responses.append(omni_resp) + + return responses + + def send_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[DiffusionResponse]: + """ + Send OpenAI requests for diffusion models. + Args: + request_config: Request configuration dictionary containing parameters like model, messages + request_num: Number of requests to send concurrently, defaults to 1 (single request) + Returns: + list[DiffusionResponse]: List of DiffusionResponse objects containing the response data + """ + responses: list[DiffusionResponse] = [] + stream = request_config.get("stream", False) + modalities = request_config.get("modalities", omit) # Most diffusion models don't require modalities param + extra_body = request_config.get("extra_body", None) + if stream: + raise NotImplementedError("Streaming is not currently implemented for diffusion model e2e test") + if request_num == 1: + # Send single request + chat_completion = self.client.chat.completions.create( + model=request_config.get("model"), + messages=request_config.get("messages"), + extra_body=extra_body, + modalities=modalities, + ) + response = self._process_diffusion_response(chat_completion) + assert_diffusion_response(response, request_config, run_level=self.run_level) + responses.append(response) + else: + # Send concurrent requests + with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: + futures = [] + # Submit all request tasks + for _ in range(request_num): + future = executor.submit( + self.client.chat.completions.create, + model=request_config.get("model"), + messages=request_config.get("messages"), + modalities=modalities, + extra_body=extra_body, + ) + futures.append(future) + # Process completed tasks + for future in concurrent.futures.as_completed(futures): + chat_completion = future.result() + response = self._process_diffusion_response(chat_completion) + assert_diffusion_response(response, request_config, run_level=self.run_level) + responses.append(response) + return responses + + def send_video_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: + """ + Send native /v1/videos requests. + """ + if request_num != 1: + raise NotImplementedError("Concurrent video diffusion requests are not currently implemented") + form_data = request_config.get("form_data") + if not isinstance(form_data, dict): + raise ValueError("Video request_config must contain 'form_data'") + normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None} + files: dict[str, tuple[str, BytesIO, str]] = {} + image_reference = request_config.get("image_reference") + if image_reference: + if image_reference.startswith("data:image"): + header, encoded = image_reference.split(",", 1) + content_type = header.split(";")[0].removeprefix("data:") + extension = content_type.split("/")[-1] + file_data = base64.b64decode(encoded) + files["input_reference"] = (f"reference.{extension}", BytesIO(file_data), content_type) + else: + normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference}) + + result = DiffusionResponse() + start_time = time.perf_counter() + create_url = self._build_url("/v1/videos") + response = requests.post( + create_url, + data=normalized_form_data, + files=files, + headers={"Accept": "application/json"}, + timeout=60, + ) + response.raise_for_status() + job_data = response.json() + video_id = job_data["id"] + self._wait_until_video_completed(video_id) + video_content = self._download_video_content(video_id) + result.success = True + result.videos = [video_content] + result.e2e_latency = time.perf_counter() - start_time + assert_diffusion_response(result, request_config, run_level=self.run_level) + return [result] + + def _wait_until_video_completed( + self, video_id: str, poll_interval_seconds: int = 2, timeout_seconds: int = 300 + ) -> None: + status_url = self._build_url(f"/v1/videos/{video_id}") + deadline = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + status_resp = requests.get(status_url, headers={"Accept": "application/json"}, timeout=30) + status_resp.raise_for_status() + status_data = status_resp.json() + current_status = status_data["status"] + if current_status == "completed": + return + if current_status == "failed": + error_msg = status_data.get("last_error", "Unknown error") + raise RuntimeError(f"Job failed: {error_msg}") + time.sleep(poll_interval_seconds) + raise TimeoutError(f"Video job {video_id} did not complete within {timeout_seconds}s") + + def _download_video_content(self, video_id: str) -> bytes: + download_url = self._build_url(f"/v1/videos/{video_id}/content") + video_resp = requests.get(download_url, stream=True, timeout=60) + video_resp.raise_for_status() + video_bytes = BytesIO() + for chunk in video_resp.iter_content(chunk_size=8192): + if chunk: + video_bytes.write(chunk) + return video_bytes.getvalue() + + def _build_url(self, path: str) -> str: + return f"{self.base_url.rstrip('/')}/{path.lstrip('/')}" + + +class OmniRunner: + def __init__( + self, + model_name: str, + seed: int = 42, + stage_init_timeout: int = 600, + batch_timeout: int = 10, + init_timeout: int = 900, + shm_threshold_bytes: int = 65536, + log_stats: bool = False, + stage_configs_path: str | None = None, + **kwargs, + ) -> None: + cleanup_dist_env_and_memory() + run_forced_gpu_cleanup_round() + self.model_name = model_name + self.seed = seed + self._prompt_len_estimate_cache: dict[str, Any] = {} + from vllm_omni.entrypoints.omni import Omni + + self.omni = Omni( + model=model_name, + log_stats=log_stats, + stage_init_timeout=stage_init_timeout, + batch_timeout=batch_timeout, + init_timeout=init_timeout, + shm_threshold_bytes=shm_threshold_bytes, + stage_configs_path=stage_configs_path, + **kwargs, + ) + + def get_default_sampling_params_list(self) -> list[Any]: + if not hasattr(self.omni, "default_sampling_params_list"): + raise AttributeError("Omni.default_sampling_params_list is not available") + return list(self.omni.default_sampling_params_list) + + def _estimate_prompt_len( + self, + additional_information: dict[str, Any], + model_name: str, + ) -> int: + """Estimate prompt_token_ids placeholder length for the Talker stage. + + The AR Talker replaces all input embeddings via ``preprocess``, so the + placeholder values are irrelevant but the **length** must match the + embeddings that ``preprocess`` will produce. + """ + _cache = self._prompt_len_estimate_cache + try: + from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( + Qwen3TTSTalkerForConditionalGeneration, + ) + + if model_name not in _cache: + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left") + cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True) + _cache[model_name] = (tok, getattr(cfg, "talker_config", None)) + + tok, tcfg = _cache[model_name] + task_type = (additional_information.get("task_type") or ["CustomVoice"])[0] + return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information( + additional_information=additional_information, + task_type=task_type, + tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"], + codec_language_id=getattr(tcfg, "codec_language_id", None), + spk_is_dialect=getattr(tcfg, "spk_is_dialect", None), + ) + except Exception as exc: + logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc) + return 2048 + + def get_omni_inputs( + self, + prompts: list[str] | str, + system_prompt: str | None = None, + audios: PromptAudioInput = None, + images: PromptImageInput = None, + videos: PromptVideoInput = None, + mm_processor_kwargs: dict[str, Any] | None = None, + modalities: list[str] | None = None, + ) -> list[TextPrompt]: + if system_prompt is None: + system_prompt = ( + "You are Qwen, a virtual human developed by the Qwen Team, Alibaba " + "Group, capable of perceiving auditory and visual inputs, as well as " + "generating text and speech." + ) + video_padding_token = "<|VIDEO|>" + image_padding_token = "<|IMAGE|>" + audio_padding_token = "<|AUDIO|>" + if "Qwen3-Omni-30B-A3B-Instruct" in self.model_name: + video_padding_token = "<|video_pad|>" + image_padding_token = "<|image_pad|>" + audio_padding_token = "<|audio_pad|>" + elif "Ming-flash-omni" in self.model_name: + video_padding_token = "