diff --git a/.buildkite/bootstrap-amd-omni.sh b/.buildkite/bootstrap-amd-omni.sh index a38b76220..9e7021493 100755 --- a/.buildkite/bootstrap-amd-omni.sh +++ b/.buildkite/bootstrap-amd-omni.sh @@ -90,10 +90,18 @@ upload_pipeline() { FAIL_FAST=$(fail_fast) cd .buildkite + + # Select test definition file: merge suite for main, ready suite for PRs + if [[ $BUILDKITE_BRANCH == "main" ]]; then + TEST_YAML="test-amd-merge.yml" + else + TEST_YAML="test-amd-ready.yaml" + fi + ( set -x # Output pipeline.yaml with all blank lines removed - minijinja-cli test-template.j2 test-amd.yaml \ + minijinja-cli test-template.j2 "$TEST_YAML" \ -D branch="$BUILDKITE_BRANCH" \ -D list_file_diff="$LIST_FILE_DIFF" \ -D run_all="$RUN_ALL" \ diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index f86b4b5d9..a06cf96bf 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -87,7 +87,11 @@ HF_CACHE="$(realpath ~)/huggingface" mkdir -p "${HF_CACHE}" HF_MOUNT="/root/.cache/huggingface" -commands=$@ +if [[ -n "${TEST_COMMAND:-}" ]]; then + commands="$TEST_COMMAND" +else + commands="$@" +fi echo "Commands:$commands" PARALLEL_JOB_COUNT=8 @@ -102,6 +106,7 @@ if [[ -z "$render_gid" ]]; then fi # check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs. +# TODO: @tjtanaa reenable to run VLLM_ROCM_USE_AITER=1 when AITER is shipped with prebuilt kernels. if [[ $commands == *"--shard-id="* ]]; then # assign job count as the number of shards used commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g') @@ -118,7 +123,7 @@ if [[ $commands == *"--shard-id="* ]]; then --rm \ -e MIOPEN_DEBUG_CONV_DIRECT=0 \ -e MIOPEN_DEBUG_CONV_GEMM=0 \ - -e VLLM_ROCM_USE_AITER=1 \ + -e VLLM_ROCM_USE_AITER=0 \ -e HIP_VISIBLE_DEVICES="${GPU}" \ -e HF_TOKEN \ -e AWS_ACCESS_KEY_ID \ @@ -153,7 +158,7 @@ else --rm \ -e MIOPEN_DEBUG_CONV_DIRECT=0 \ -e MIOPEN_DEBUG_CONV_GEMM=0 \ - -e VLLM_ROCM_USE_AITER=1 \ + -e VLLM_ROCM_USE_AITER=0 \ -e HF_TOKEN \ -e AWS_ACCESS_KEY_ID \ -e AWS_SECRET_ACCESS_KEY \ diff --git a/.buildkite/test-amd-merge.yml b/.buildkite/test-amd-merge.yml new file mode 100644 index 000000000..78d910b21 --- /dev/null +++ b/.buildkite/test-amd-merge.yml @@ -0,0 +1,197 @@ +steps: + +- label: "Simple Unit Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_ROCM_USE_AITER=0 + - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml" + +- label: "Diffusion Model Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model" + +- label: "Diffusion Images API LoRA E2E" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - timeout 20m pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py + +- label: "Diffusion Model CPU offloading Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - | + timeout 20m bash -c ' + set +e + pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py + EXIT1=\$? + pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py + EXIT2=\$? + exit \$((EXIT1 | EXIT2)) + ' + +## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274 +# - label: "Audio Generation Model Test" +# agent_pool: mi325_1 +# depends_on: amd-build +# mirror_hardwares: [amdproduction] +# grade: Blocking +# commands: +# - export GPU_ARCHS=gfx942 +# - export VLLM_LOGGING_LEVEL=DEBUG +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py + +- label: "Diffusion Cache Backend Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4" + +- label: "Diffusion Sequence Parallelism Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py + +# merge-only tests +- label: "Diffusion Tensor Parallelism Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py + +- label: "Diffusion GPU Worker Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py + +- label: "Benchmark & Engine Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - | + timeout 20m bash -c ' + set +e + pytest -s -v tests/benchmarks/test_serve_cli.py + EXIT1=\$? + pytest -s -v tests/engine/test_async_omni_engine_abort.py + EXIT2=\$? + exit \$((EXIT1 | EXIT2)) + ' + +- label: "Omni Model Test Qwen2-5-Omni" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py + +- label: "Omni Model Test Qwen3-Omni" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py + - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model" + +- label: "Qwen3-TTS E2E Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py + +- label: "Diffusion Image Edit Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py + +# split Bagel Model Test with H100 (Real Weights) into three tests +- label: "Bagel Text2Img Model Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_ROCM_USE_AITER_RMSNORM=0 + - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm" + +- label: "Bagel Img2Img Model Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_ROCM_USE_AITER_RMSNORM=0 + - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model" -k "rocm" + +- label: "Bagel Online Serving Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - export VLLM_IMAGE_FETCH_TIMEOUT=60 + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_ROCM_USE_AITER_RMSNORM=0 + - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model" -k "rocm" diff --git a/.buildkite/test-amd-ready.yaml b/.buildkite/test-amd-ready.yaml new file mode 100644 index 000000000..b09fe2ebf --- /dev/null +++ b/.buildkite/test-amd-ready.yaml @@ -0,0 +1,178 @@ +steps: + +- label: "Simple Unit Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_ROCM_USE_AITER=0 + - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml" + +- label: "Diffusion Model Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model" + +- label: "Diffusion Model CPU offloading Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - | + timeout 20m bash -c ' + set +e + pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py + EXIT1=\$? + pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py + EXIT2=\$? + exit \$((EXIT1 | EXIT2)) + ' + +## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274 +# - label: "Audio Generation Model Test" +# agent_pool: mi325_1 +# depends_on: amd-build +# mirror_hardwares: [amdproduction] +# grade: Blocking +# commands: +# - export GPU_ARCHS=gfx942 +# - export VLLM_LOGGING_LEVEL=DEBUG +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py + +- label: "Diffusion Cache Backend Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4" + +- label: "Diffusion Sequence Parallelism Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model + +- label: "Diffusion GPU Worker Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py + +- label: "Benchmark & Engine Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - | + timeout 30m bash -c ' + set +e + pytest -s -v tests/benchmarks/test_serve_cli.py + EXIT1=\$? + pytest -s -v tests/engine/test_async_omni_engine_abort.py + EXIT2=\$? + exit \$((EXIT1 | EXIT2)) + ' + +- label: "Omni Model Test Qwen2-5-Omni" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - timeout 17m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py + +- label: "Omni Model Test Qwen3-Omni" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py + - timeout 10m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model" + +- label: "Qwen3-TTS E2E Test" + agent_pool: mi325_2 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py + +- label: "Diffusion Image Edit Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py + +- label: "Bagel Text2Img Model Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_ROCM_USE_AITER_RMSNORM=0 + - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model" -k "rocm" + +- label: "Bagel Img2Img Model Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_ROCM_USE_AITER_RMSNORM=0 + - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model" -k "rocm" + +- label: "Bagel Online Serving Test" + agent_pool: mi325_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export GPU_ARCHS=gfx942 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - export VLLM_IMAGE_FETCH_TIMEOUT=60 + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_ROCM_USE_AITER_RMSNORM=0 + - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model" -k "rocm" diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index ae2e7f4c4..5b11e0f2d 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -340,3 +340,59 @@ steps: hostPath: path: /mnt/hf-cache type: DirectoryOrCreate + + # - label: "Voxtral TTS E2E Test" + # timeout_in_minutes: 20 + # depends_on: upload-merge-pipeline + # commands: + # - | + # timeout 20m bash -c ' + # export VLLM_LOGGING_LEVEL=DEBUG + # export VLLM_WORKER_MULTIPROC_METHOD=spawn + # pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model" + # ' + # agents: + # queue: "mithril-h100-pool" + # plugins: + # - kubernetes: + # podSpec: + # containers: + # - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + # resources: + # limits: + # nvidia.com/gpu: 1 + # volumeMounts: + # - name: devshm + # mountPath: /dev/shm + # - name: hf-cache + # mountPath: /root/.cache/huggingface + # env: + # - name: HF_HOME + # value: /root/.cache/huggingface + # nodeSelector: + # node.kubernetes.io/instance-type: gpu-h100-sxm + # volumes: + # - name: devshm + # emptyDir: + # medium: Memory + # - name: hf-cache + # hostPath: + # path: /mnt/hf-cache + # type: DirectoryOrCreate + + - label: "CUDA Unit Test" + timeout_in_minutes: 10 + depends_on: upload-merge-pipeline + commands: + - pytest -v -s -m 'core_model and cuda' + agents: + queue: "gpu_1_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 53a4f180c..6d5674bf5 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -11,6 +11,9 @@ steps: EXIT1=$$? pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model" EXIT2=$$? + # TODO: uncomment when mistralai/tts-model is public + # pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model" + # EXIT3=$$? exit $$((EXIT1 | EXIT2)) agents: queue: "mithril-h100-pool" diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index b922a868b..09bd85963 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -234,6 +234,45 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" + # TODO: uncomment when mistralai/tts-model is public + # - label: "Voxtral-TTS E2E Test" + # depends_on: upload-ready-pipeline + # commands: + # - | + # timeout 20m bash -c ' + # export VLLM_LOGGING_LEVEL=DEBUG + # export VLLM_WORKER_MULTIPROC_METHOD=spawn + # pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model" + # ' + # agents: + # queue: "mithril-h100-pool" + # plugins: + # - kubernetes: + # podSpec: + # containers: + # - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + # resources: + # limits: + # nvidia.com/gpu: 1 + # volumeMounts: + # - name: devshm + # mountPath: /dev/shm + # - name: hf-cache + # mountPath: /root/.cache/huggingface + # env: + # - name: HF_HOME + # value: /root/.cache/huggingface + # nodeSelector: + # node.kubernetes.io/instance-type: gpu-h100-sxm + # volumes: + # - name: devshm + # emptyDir: + # medium: Memory + # - name: hf-cache + # hostPath: + # path: /mnt/hf-cache + # type: DirectoryOrCreate + # - label: "Diffusion Image Edit Test with H100 (1 GPU)" # depends_on: upload-ready-pipeline # commands: diff --git a/.buildkite/test-template-amd-omni.j2 b/.buildkite/test-template-amd-omni.j2 index 291ed0a9a..4612f8ccd 100644 --- a/.buildkite/test-template-amd-omni.j2 +++ b/.buildkite/test-template-amd-omni.j2 @@ -40,9 +40,14 @@ {% else %} queue: amd_mi325_1 {% endif %} - command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh "(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }} ; {{ step.command or (step.commands | join(" && ")) | safe }}" +{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %} +{% set indented_cmd = cmd_body | replace("\n", "\n ") %} + command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh env: DOCKER_BUILDKIT: "1" + TEST_COMMAND: |- + (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }} + {{ indented_cmd | safe }} priority: 100 {% if step.grade and step.grade == "Blocking" %} soft_fail: false diff --git a/docs/.nav.yml b/docs/.nav.yml index 8042ffb70..b7d08e77e 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -58,6 +58,7 @@ nav: - Quantization: - Overview: user_guide/diffusion/quantization/overview.md - FP8: user_guide/diffusion/quantization/fp8.md + - Int8: user_guide/diffusion/quantization/int8.md - GGUF: user_guide/diffusion/quantization/gguf.md - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md diff --git a/docs/mkdocs/hooks/generate_api_readme.py b/docs/mkdocs/hooks/generate_api_readme.py index 9c344d3f1..5acd1891e 100644 --- a/docs/mkdocs/hooks/generate_api_readme.py +++ b/docs/mkdocs/hooks/generate_api_readme.py @@ -141,6 +141,7 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]: "vllm_omni.diffusion.models.qwen_image", "vllm_omni.entrypoints.async_diffusion", "vllm_omni.entrypoints.openai", + "vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts", ] if any(module_path.startswith(prefix) for prefix in excluded_prefixes): continue diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 87cdccef1..a6925b7a1 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -52,6 +52,7 @@ th { |`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | |`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | |`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` | +|`VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/tts-model` | ## List of Supported Models for NPU diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index 46b9154bb..17787e682 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -4,6 +4,7 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. - **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz. - **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz. +- **Voxtral TTS** (`mistralai/tts-model`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz. Each server instance runs a single model (specified at startup via `vllm serve --omni`). @@ -28,6 +29,14 @@ vllm-omni serve fishaudio/s2-pro \ --trust-remote-code \ --enforce-eager \ --gpu-memory-utilization 0.9 + +# Voxtral TTS +vllm serve mistralai/tts-model \ + --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager ``` ### Generate Speech @@ -347,6 +356,12 @@ curl -X POST http://localhost:8091/v1/audio/speech \ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech online serving example](../user_guide/examples/online_serving/fish_speech.md) for details. +### Voxtral TTS + +| Model | Description | +|-------|-------------| +| `mistralai/tts-model` | 3B AR + FlowMatching TTS. Supports text-to-speech with preset voices. | + ## Error Responses ### 400 Bad Request diff --git a/docs/user_guide/diffusion/quantization/int8.md b/docs/user_guide/diffusion/quantization/int8.md new file mode 100644 index 000000000..1e7853c3e --- /dev/null +++ b/docs/user_guide/diffusion/quantization/int8.md @@ -0,0 +1,75 @@ +# Int8 Quantization + +## Overview + +Int8 quantization converts BF16/FP16 weights to Int8 at model load time. No calibration or pre-quantized checkpoint needed. + +Depending on the model, either all layers can be quantized, or some sensitive layers should stay in BF16/FP16. See the [per-model table](#supported-models) for which case applies. + +## Configuration + +1. **Python API**: set `quantization="int8"`. To skip sensitive layers, use `quantization_config` with `ignored_layers`. + +```python +from vllm_omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +# All layers quantized +omni = Omni(model="", quantization="int8") + +# Skip sensitive layers +omni = Omni( + model="", + quantization_config={ + "method": "int8", + "ignored_layers": [""], + }, +) + +outputs = omni.generate( + "A cat sitting on a windowsill", + OmniDiffusionSamplingParams(num_inference_steps=50), +) +``` + +2. **CLI**: pass `--quantization int8` and optionally `--ignored-layers`. + +```bash +# All layers +python text_to_image.py --model --quantization int8 + +# Skip sensitive layers +python text_to_image.py --model --quantization int8 --ignored-layers "img_mlp" + +# Online serving +vllm serve --omni --quantization int8 +``` + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `method` | str | — | Quantization method (`"int8"`) | +| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16/FP16 | +| `activation_scheme` | str | `"dynamic"` | `"dynamic"` (no calibration) | + + +The available `ignored_layers` names depend on the model architecture (e.g., `to_qkv`, `to_out`, `img_mlp`, `txt_mlp`). Consult the transformer source for your target model. + +## Supported Models + +| Model | HF Models | Recommendation | `ignored_layers` | +|-------|-----------|---------------|------------------| +| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None | +| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | All layers | None | + +## Combining with Other Features + +Int8 quantization can be combined with cache acceleration: + +```python +omni = Omni( + model="", + quantization="int8", + cache_backend="tea_cache", + cache_config={"rel_l1_thresh": 0.2}, +) +``` diff --git a/docs/user_guide/diffusion/quantization/overview.md b/docs/user_guide/diffusion/quantization/overview.md index e4ce69677..a95afdbf4 100644 --- a/docs/user_guide/diffusion/quantization/overview.md +++ b/docs/user_guide/diffusion/quantization/overview.md @@ -7,12 +7,20 @@ vLLM-Omni supports quantization of DiT linear layers to reduce memory usage and | Method | Guide | |--------|-------| | FP8 | [FP8](fp8.md) | +| Int8 | [Int8](int8.md) | | GGUF | [GGUF](gguf.md) | -## Device Compatibility +## Device Compatibility for FP8 | GPU Generation | Example GPUs | FP8 Mode | |---------------|-------------------|----------| | Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware | Kernel selection is automatic. + +## Device Compatibility for Int8 + +| Device Type | Generation | Example | Int8 Mode | +|-------------|---------------|-------------------|----------| +| NVIDIA GPU | Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware | +| Ascend NPU | Atlas A2/Atlas A3 | Atlas 800T A2/Atlas 900 A3 | Full W8A8 with native hardware | diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index e9179b575..7ef470756 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/docs/user_guide/diffusion_acceleration.md @@ -16,7 +16,7 @@ Both methods can provide significant speedups (typically **1.5x-2.0x**) while ma vLLM-Omni also supports quantization methods: -3. **[FP8 Quantization](diffusion/quantization/overview.md)** - Reduces DiT linear layers from BF16 to FP8, providing ~1.28x speedup with minimal quality loss. Supports per-layer skip for sensitive layers. +3. **[Quantization](diffusion/quantization/overview.md)** - Reduces DiT linear layers from BF16 to FP8 or Int8, providing ~1.28x speedup with minimal quality loss. Supports per-layer skip for sensitive layers. vLLM-Omni also supports parallelism methods for diffusion models, including: @@ -46,6 +46,7 @@ vLLM-Omni also supports parallelism methods for diffusion models, including: | Method | Configuration | Description | Best For | |--------|--------------|-------------|----------| | **FP8** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup | +| **Int8** | `quantization="int8"` | Int8 W8A8 | Memory reduction, inference speedup | ## Supported Models @@ -81,11 +82,11 @@ The following table shows which models are currently supported by each accelerat ### Quantization -| Model | Model Identifier | FP8 | -|-------|------------------|:---:| -| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | -| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | -| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | +| Model | Model Identifier | FP8 | Int8 | +|-------|------------------|:---:|:---:| +| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | +| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | +| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ## Performance Benchmarks @@ -338,13 +339,30 @@ outputs = omni.generate( ) ``` +### Using Int8 Quantization + +```python +from vllm_omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +omni = Omni( + model="", + quantization="int8", +) + +outputs = omni.generate( + "A cat sitting on a windowsill", + OmniDiffusionSamplingParams(num_inference_steps=50), +) +``` + ## Documentation For detailed information on each acceleration method: - **[TeaCache Guide](diffusion/teacache.md)** - Complete TeaCache documentation, configuration options, and best practices - **[Cache-DiT Acceleration Guide](diffusion/cache_dit_acceleration.md)** - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters -- **[FP8 Quantization Guide](diffusion/quantization/overview.md)** - FP8 quantization for DiT models with per-layer control +- **[Quantization Guide](diffusion/quantization/overview.md)** - Quantization for DiT models with per-layer control - **[Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism)** - Guidance on how to enable TP for diffusion models. - **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration. - **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks. diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index e2463619a..ea17e6858 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -131,12 +131,10 @@ def parse_args() -> argparse.Namespace: "--quantization", type=str, default=None, - choices=["fp8", "gguf"], - help=( - "Quantization method for the transformer. " - "Options: 'fp8' (FP8 W8A8), 'gguf' (GGUF quantized weights). " - "Default: None (no quantization, uses BF16)." - ), + choices=["fp8", "int8", "gguf"], + help="Quantization method for the transformer. " + "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8), 'gguf' (GGUF quantized weights). " + "Default: None (no quantization, uses BF16).", ) parser.add_argument( "--gguf-model", diff --git a/examples/offline_inference/voxtral_tts/README.md b/examples/offline_inference/voxtral_tts/README.md new file mode 100644 index 000000000..7e0f84f2c --- /dev/null +++ b/examples/offline_inference/voxtral_tts/README.md @@ -0,0 +1,57 @@ +# Voxtral TTS Offline Inference + +`end2end.py` runs Voxtral TTS end-to-end offline inference using vLLM. It supports both blocking (`Omni`) and streaming (`AsyncOmni`) generation, batched prompts with configurable concurrency, and voice selection via preset name or reference audio file. + +When `mistral_common` has `SpeechRequest` support, prompt token IDs are built via `encode_speech_request`. Otherwise, the script falls back to manual token construction. + +## Usage Examples + +```bash +# Basic single-prompt with reference audio +python3 examples/offline_inference/voxtral_tts/end2end.py \ + --write-audio \ + --model mistralai/tts-model \ + --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?" \ + --audio-path path/to/reference_audio.wav + +# 32 replicate prompts with reference audio +python3 examples/offline_inference/voxtral_tts/end2end.py \ + --num-prompts 32 --write-audio \ + --model mistralai/tts-model \ + --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?" \ + --audio-path path/to/reference_audio.wav + +# Short debug prompt with reference audio +python3 examples/offline_inference/voxtral_tts/end2end.py \ + --write-audio \ + --model mistralai/tts-model \ + --text "This is a test message." \ + --audio-path path/to/reference_audio.wav + +# Streaming with neutral_female voice preset +python3 examples/offline_inference/voxtral_tts/end2end.py \ + --streaming --write-audio --voice neutral_female \ + --model mistralai/tts-model + +# 32 prompts, 8 concurrent requests per wave, streaming with casual_male voice +python3 examples/offline_inference/voxtral_tts/end2end.py \ + --num-prompts 32 --concurrency 8 --streaming --write-audio --voice casual_male \ + --model mistralai/tts-model +``` + +## Arguments + +| Argument | Description | +|---|---| +| `--model PATH` | HuggingFace repo ID or local directory path (default: `mistralai/tts-model`) | +| `--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 (auto-resolved from model if not set) | +| `--num-prompts N` | Number of replicate prompts to run for measuring performance (default: 1) | +| `--streaming` | Use streaming generation via `AsyncOmni` (default: blocking `Omni`) | +| `--concurrency N` | Max concurrent requests per wave (must be used with `--streaming`, must evenly divide `--num-prompts`) | +| `--voice NAME` | Voice preset to use instead of reference audio file (e.g., casual_female, casual_male, cheerful_female, neutral_female, neutral_male) | +| `--write-audio` | Write generated audio to WAV files | +| `--profiling-mode` | Enable profiling mode (reduces max tokens to 50) | +| `--log-stats` | Enable detailed statistics logging | diff --git a/examples/offline_inference/voxtral_tts/end2end.py b/examples/offline_inference/voxtral_tts/end2end.py new file mode 100644 index 000000000..569b5f6f3 --- /dev/null +++ b/examples/offline_inference/voxtral_tts/end2end.py @@ -0,0 +1,377 @@ +""" +This example shows how to use vLLM for running Voxtral TTS +""" + +import asyncio +import gc +import logging +import os +import time +import uuid +from argparse import Namespace +from pathlib import Path +from typing import Any + +import numpy as np +import soundfile as sf +import torch +from mistral_common.protocol.instruct.chunk import TextChunk + +try: + from mistral_common.protocol.speech.request import SpeechRequest + from mistral_common.tokens.tokenizers.mistral import MistralTokenizer +except ImportError: + raise ImportError( + "Could not import SpeechRequest or MistralTokenizer from mistral_common. " + "Please pull the latest mistral-common code and install it with: " + "pip install -e ." + ) +from vllm import SamplingParams +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni import AsyncOmni +from vllm_omni.entrypoints.omni import Omni + +logger = logging.getLogger(__name__) + + +# ---- Streaming version ---- +async def run_streaming(inputs, sampling_params_list, model_name, args, output_dir): + async_omni = AsyncOmni( + model=model_name, + stage_configs_path=args.stage_configs_path, + log_stats=args.log_stats, + ) + + # Normalize to a list so batch and single-input share the same code path + if isinstance(inputs, list): + inputs_list = inputs + else: + inputs_list = [inputs] + + total_audio_dur = 0.0 + total_gen_time = 0.0 + total_ttfa = 0.0 + ttfa_count = 0 + total_waits = 0 + total_chunks = 0 + results_lock = asyncio.Lock() + + async def _generate_one(batch_idx, single_input): + nonlocal total_audio_dur, total_gen_time, total_ttfa, ttfa_count + nonlocal total_waits, total_chunks + request_id = str(uuid.uuid4()) + all_audio_chunks = [] + chunk_arrival_times = [] + chunk_durations = [] + chunk_idx = 0 + gen_start = time.time() + ttfa = None + accumulated_sample = 0 + + async for stage_output in async_omni.generate(single_input, request_id, sampling_params_list): + mm_output = stage_output.multimodal_output + finished = stage_output.finished + if not mm_output or "audio" not in mm_output: + continue + + now = time.time() + if ttfa is None: + ttfa = now - gen_start + + audio_chunk = mm_output["audio"] + # audio_chunk is a 1-D tensor from decode_helper_async + if isinstance(audio_chunk, torch.Tensor): + if finished: + # Last chunk may return whole audio instead of chunk delta. + # Cut accumulated samples from previous chunks. + audio_numpy = audio_chunk[accumulated_sample:].float().detach().cpu().numpy() + else: + audio_numpy = audio_chunk.float().detach().cpu().numpy() + elif isinstance(audio_chunk, list): + # Audio_chunk list contain all previous chunks. Use chunk_idx to index + audio_numpy = audio_chunk[chunk_idx].float().detach().cpu().numpy() + else: + audio_numpy = audio_chunk + + accumulated_sample += len(audio_numpy) + all_audio_chunks.append(audio_numpy) + chunk_arrival_times.append(now) + chunk_durations.append(len(audio_numpy) / 24000) + chunk_idx += 1 + + gen_elapsed = time.time() - gen_start + + # Analyze wait/no-wait per chunk: + # A "wait" means the client has finished playing all previous audio + # before this chunk arrived, so there is a playback gap. + # buffer_time > 0: client still has audio to play (no wait) + # buffer_time <= 0: client ran out of audio and had to wait + chunk_labels = [] + accumulated_audio_dur = 0.0 + if chunk_arrival_times: + first_arrival = chunk_arrival_times[0] + for i in range(len(chunk_arrival_times)): + if i == 0: + chunk_labels.append("no_wait") + else: + playback_elapsed = chunk_arrival_times[i] - first_arrival + buffer_time = accumulated_audio_dur - playback_elapsed + if buffer_time <= 0: + chunk_labels.append("wait") + else: + chunk_labels.append("no_wait") + accumulated_audio_dur += chunk_durations[i] + + req_wait_count = sum(1 for chunk_label in chunk_labels if chunk_label == "wait") + req_wait_rate = req_wait_count / len(chunk_labels) if chunk_labels else 0.0 + + # Concatenate all chunks for this request's audio + if all_audio_chunks: + full_audio = np.concatenate(all_audio_chunks) + output_audio_dur = len(full_audio) / 24000 + output_path = os.path.join(output_dir, f"tts_output_{batch_idx}.wav") + if args.write_audio: + sf.write(output_path, full_audio, 24000) + print( + f"Request {batch_idx}: saved {len(full_audio)} samples ({output_audio_dur:.2f}s) to {output_path}" + ) + # Per-chunk wait details + for i, label in enumerate(chunk_labels): + dur_ms = chunk_durations[i] * 1000 + arrival_ms = (chunk_arrival_times[i] - first_arrival) * 1000 if i > 0 else 0.0 + print( + f" Request {batch_idx} chunk {i}: {label} | arrived={arrival_ms:.1f}ms | chunk_dur={dur_ms:.1f}ms" + ) + print( + f"Request {batch_idx}: TTFA={ttfa:.4f}s | " + f"Generation={gen_elapsed:.4f}s | " + f"Audio={output_audio_dur:.2f}s | " + f"RTF={output_audio_dur / gen_elapsed:.4f} | " + f"WaitRate={req_wait_rate:.2%} ({req_wait_count}/{len(chunk_labels)})" + ) + async with results_lock: + total_audio_dur += output_audio_dur + total_gen_time += gen_elapsed + total_waits += req_wait_count + total_chunks += len(chunk_labels) + if ttfa is not None: + total_ttfa += ttfa + ttfa_count += 1 + + # Launch requests in waves of `concurrency` size + concurrency = args.concurrency or len(inputs_list) + gen_start_all = time.time() + for wave_start in range(0, len(inputs_list), concurrency): + wave = inputs_list[wave_start : wave_start + concurrency] + wave_idx_offset = wave_start + print(f"\n--- Wave {wave_start // concurrency + 1} (requests {wave_start}-{wave_start + len(wave) - 1}) ---") + await asyncio.gather(*[_generate_one(wave_idx_offset + i, inp) for i, inp in enumerate(wave)]) + generation_time = time.time() - gen_start_all + avg_ttfa = total_ttfa / ttfa_count if ttfa_count else float("nan") + overall_wait_rate = total_waits / total_chunks if total_chunks else float("nan") + print( + f"\nAll requests: Generation={generation_time:.4f}s | " + f"TotalAudio={total_audio_dur:.2f}s | " + f"Concurrency={concurrency} | " + f"AvgTTFA={avg_ttfa:.4f}s | " + f"RTF(total)={total_audio_dur / generation_time:.4f} | " + f"RTF(per-request)={total_audio_dur / total_gen_time:.4f} | " + f"WaitRate={overall_wait_rate:.2%} ({total_waits}/{total_chunks})" + ) + + async_omni.shutdown() + torch.cuda.empty_cache() + gc.collect() + + +# ---- Non-streaming version ---- +def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir): + llm = Omni( + model=model_name, + log_stats=args.log_stats, + stage_configs_path=args.stage_configs_path, + ) + + if args.profiling_mode: + llm.start_profile() + + vllm_start = time.time() + outputs = llm.generate(inputs, sampling_params_list) + vllm_elapsed = time.time() - vllm_start + + if args.profiling_mode: + llm.stop_profile() + time.sleep(10) + + print(f"vLLM run time: {vllm_elapsed:.4f}s") + output_audio_dur = 0.0 + + for batch_idx, o in enumerate(outputs): + audio_tensor = torch.cat(o.request_output[0].outputs[0].multimodal_output["audio"]) + audio_array = audio_tensor.tolist() + output_audio_dur += float(len(audio_array)) / 24000 + if args.write_audio: + output_path = os.path.join(output_dir, f"tts_output_{batch_idx}.wav") + sf.write(output_path, audio_array, 24000) + print(f"Audio saved to {output_path}") + + print(f"Total audio duration: {output_audio_dur:.2f}s") + print(f"RTF: {output_audio_dur / vllm_elapsed:.4f}") + + del llm + torch.cuda.empty_cache() + gc.collect() + + +def parse_args() -> Namespace: + parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS") + parser.add_argument( + "--model", + type=str, + default="mistralai/tts-model", + help="Model name or path.", + ) + parser.add_argument( + "--text", + type=str, + default="This is a test message.", + help="Text to synthesize.", + ) + parser.add_argument( + "--audio-path", + type=str, + default=None, + help="Path to reference audio file for voice cloning.", + ) + parser.add_argument( + "--output-dir", + type=str, + default="output_audio", + help="Directory to write output wav files.", + ) + parser.add_argument( + "--stage-configs-path", + type=str, + default=None, + help="Path to stage configs YAML. Auto-resolved from model if not set.", + ) + parser.add_argument( + "--num-prompts", type=int, default=1, help="Number of replicate prompts to run for measuring performance" + ) + parser.add_argument( + "--profiling-mode", + default=False, + action="store_true", + help="Set max_num_tokens=2 to reduce profiling time", + ) + parser.add_argument( + "--log-stats", + action="store_true", + default=False, + help="Enable writing detailed statistics (default: disabled)", + ) + parser.add_argument( + "--write-audio", + action="store_true", + default=False, + help="Write audio output to files.", + ) + parser.add_argument( + "--streaming", + action="store_true", + default=False, + help="Use streaming generation via AsyncOmni instead of blocking Omni.", + ) + parser.add_argument( + "--concurrency", + type=int, + default=None, + help="Max concurrent requests per wave (default: all at once). Must evenly divide --num-prompts.", + ) + parser.add_argument( + "--voice", + type=str, + default=None, + help="Voice to use instead of audio file.", + ) + return parser.parse_args() + + +def compose_request( + model_name: str, + text_chunk: TextChunk, + audio_prompt_file: str, + args: Any, +) -> dict: + """Build the full TTS input dict (prompt_token_ids, multi_modal_data or additional_information).""" + inputs: dict[str, Any] = {} + if Path(model_name).is_dir(): + mistral_tokenizer = MistralTokenizer.from_file(str(Path(model_name) / "tekken.json")) + else: + mistral_tokenizer = MistralTokenizer.from_hf_hub(model_name) + instruct_tokenizer = mistral_tokenizer.instruct_tokenizer + + if args.voice is not None: + tokenized = instruct_tokenizer.encode_speech_request(SpeechRequest(input=text_chunk.text, voice=args.voice)) + inputs["additional_information"] = {"voice": [args.voice]} + inputs["prompt_token_ids"] = tokenized.tokens + else: + with open(audio_prompt_file, "rb") as f: + ref_audio_bytes = f.read() + tokenized = instruct_tokenizer.encode_speech_request( + SpeechRequest(input=text_chunk.text, ref_audio=ref_audio_bytes) + ) + audio = tokenized.audios[0] + inputs["multi_modal_data"] = {"audio": [(audio.audio_array, audio.sampling_rate)]} + inputs["prompt_token_ids"] = tokenized.tokens + + return inputs + + +def main(args: Any) -> None: + max_num_tokens = 50 if args.profiling_mode else 2500 + + model_name = args.model + output_dir = args.output_dir + + if args.voice is None and args.audio_path is None: + raise ValueError("Either --voice or --audio-path must be provided.") + + audio_prompt_file = args.audio_path + text_chunk = TextChunk(text=args.text) + + if args.write_audio: + os.makedirs(output_dir, exist_ok=True) + + inputs = compose_request(model_name, text_chunk, audio_prompt_file, args) + + sampling_params = SamplingParams( + max_tokens=max_num_tokens, + ) + sampling_params_list = [ + sampling_params, + sampling_params, + ] + + if args.num_prompts > 1: + inputs = [inputs] * args.num_prompts + + if args.concurrency is not None: + assert args.streaming, "--concurrency must be used with --streaming on AsyncOmni" + assert args.num_prompts % args.concurrency == 0, ( + f"--num-prompts ({args.num_prompts}) must be divisible by --concurrency ({args.concurrency})" + ) + + torch.cuda.empty_cache() + gc.collect() + + if args.streaming: + asyncio.run(run_streaming(inputs, sampling_params_list, model_name, args, output_dir)) + else: + run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir) + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/examples/online_serving/voxtral_tts/gradio_demo.py b/examples/online_serving/voxtral_tts/gradio_demo.py new file mode 100644 index 000000000..293f1c391 --- /dev/null +++ b/examples/online_serving/voxtral_tts/gradio_demo.py @@ -0,0 +1,255 @@ +import argparse +import io +import json +import logging +import socket +import uuid +from datetime import datetime +from pathlib import Path +from typing import Any + +import gradio as gr +import httpx +import numpy as np +import soundfile as sf + +logger = logging.getLogger() + +LOGFORMAT = "%(asctime)s - %(levelname)s - %(message)s" +TIMEFORMAT = "%Y-%m-%d %H:%M:%S" + +logging.basicConfig(level=logging.INFO, force=True, format=LOGFORMAT, datefmt=TIMEFORMAT) +logger.setLevel(logging.INFO) + + +_BASE_URL = f"{socket.gethostname()}:7860" + + +def run_inference( + voice_name: str, + text_prompt: str, + base_url: str, + model: str, +) -> tuple[int, np.ndarray]: + """Call /v1/audio/speech and return (sample_rate, audio_array).""" + text_prompt = text_prompt.strip() + if not text_prompt: + raise gr.Error("Please enter a text prompt.") + + payload: dict[str, Any] = { + "input": text_prompt, + "model": model, + "response_format": "wav", + "voice": voice_name, + } + + response = httpx.post( + f"{base_url}/audio/speech", + json=payload, + timeout=120.0, + ) + response.raise_for_status() + + audio_array, sr = sf.read(io.BytesIO(response.content), dtype="float32") + return sr, audio_array + + +def _save_example( + outputs_dir: Path, + voice_name: str, + text_prompt: str, + sr: int, + audio_array: np.ndarray, +) -> tuple[str, str]: + """ + Save inputs/outputs for sharing. + Returns (share_id, saved_audio_path) + """ + share_id = uuid.uuid4().hex + + # Save generated audio + saved_audio_path = outputs_dir / f"{share_id}_gen.wav" + sf.write(str(saved_audio_path), audio_array, sr) + + meta = { + "id": share_id, + "created_at": datetime.utcnow().isoformat(), + "voice_name": voice_name, + "text_prompt": text_prompt, + "generated_audio_path": str(saved_audio_path), + } + + with open(outputs_dir / f"{share_id}.json", "w") as f: + json.dump(meta, f, ensure_ascii=False) + + return share_id, str(saved_audio_path) + + +def _load_from_share( + outputs_dir: Path | None, + request: gr.Request, +) -> tuple[str | None, str, str | None, dict[str, Any], str]: + """ + Called on page load. If ?share_id=... is present, load stored example. + Returns: (voice_name, text_prompt, output_audio, submit_btn_update, share_link_text) + """ + if outputs_dir is None: + return "Neutral_Male", "", None, gr.update(interactive=False), "" + + share_id = None + if request and request.query_params: + share_id = request.query_params.get("share_id") + + if not share_id: + return "Neutral_Male", "", None, gr.update(interactive=False), "" + + meta_path = outputs_dir / f"{share_id}.json" + if not meta_path.exists(): + logger.warning("No stored example for share_id=%s", share_id) + return "Neutral_Male", "", None, gr.update(interactive=False), "" + + with open(meta_path) as f: + meta = json.load(f) + + voice = meta.get("voice_name", "Neutral_Male") + text_prompt = meta.get("text_prompt", "") + gen_path = meta.get("generated_audio_path") + + if _BASE_URL: + share_link = f"http://{_BASE_URL}?share_id={share_id}" + else: + share_link = f"Use this query on your current URL: ?share_id={share_id}" + + return voice, text_prompt, gen_path, gr.update(interactive=True), share_link + + +def main( + model: str, + output_dir: str | None = None, + debug: bool = False, +) -> None: + base_url = "http://localhost:8091/v1" + print(f"Using speech API at: {base_url}/audio/speech") + + outputs_dir: Path | None = None + if output_dir is not None: + outputs_dir = Path(output_dir) + outputs_dir.mkdir(parents=True, exist_ok=True) + + with gr.Blocks(title="Voxtral TTS", fill_height=True) as demo: + gr.Markdown("## Voxtral TTS") + + with gr.Row(): + with gr.Column(): + voice_name = gr.Dropdown( + ["Neutral_Male", "Neutral_Female", "Cheerful_Female", "Casual_Male", "Casual_Female"], + label="Voice", + value="Neutral_Male", + ) + text_prompt = gr.Textbox( + label="Text prompt", + placeholder="Enter the text you want to synthesize...", + lines=4, + ) + with gr.Row(): + reset_btn = gr.Button("Clear") + submit_btn = gr.Button("Generate audio", interactive=False) + + with gr.Column(): + output_audio = gr.Audio( + label="Generated audio", + show_download_button=True, + interactive=False, + autoplay=True, + type="filepath", + ) + share_link_box = gr.Textbox( + label="Shareable link", + interactive=False, + show_copy_button=True, + visible=outputs_dir is not None, + ) + + # --- UI logic: disable submit until text is non-empty --- + def _toggle_submit(text: str): + enabled = bool(text.strip()) + return gr.update(interactive=enabled) + + text_prompt.change( + fn=_toggle_submit, + inputs=[text_prompt], + outputs=submit_btn, + ) + + # --- Wiring inference + persistence to the button --- + def _on_submit(voice: str, text: str): + assert text.strip() != "" + sr, audio_array = run_inference(voice, text, base_url, model) + if outputs_dir is not None: + share_id, saved_audio_path = _save_example( + outputs_dir, + voice_name=voice, + text_prompt=text, + sr=sr, + audio_array=audio_array, + ) + share_link = f"{_BASE_URL}?share_id={share_id}" + return saved_audio_path, share_link + return (sr, audio_array), "" + + submit_btn.click( + fn=_on_submit, + inputs=[voice_name, text_prompt], + outputs=[output_audio, share_link_box], + ) + + # --- Clear everything and disable submit again --- + def _on_reset(): + return ( + "Neutral_Male", # voice_name + "", # text_prompt + None, # output_audio + gr.update(interactive=False), # submit_btn + "", # share_link_box + ) + + reset_btn.click( + fn=_on_reset, + inputs=[], + outputs=[voice_name, text_prompt, output_audio, submit_btn, share_link_box], + ) + + demo.load( + fn=lambda request: _load_from_share(outputs_dir, request), + inputs=[], + outputs=[voice_name, text_prompt, output_audio, submit_btn, share_link_box], + ) + + launch_kwargs: dict[str, Any] = { + "server_name": "0.0.0.0", + "share": True, + } + if outputs_dir is not None: + launch_kwargs["allowed_paths"] = [str(outputs_dir)] + demo.launch(**launch_kwargs) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Voxtral TTS Gradio Demo") + parser.add_argument("--model", type=str, required=True, help="Path to the model checkpoint") + parser.add_argument( + "--output-dir", + type=str, + default=None, + help="Directory to save generated audio and share links. " + "If not provided, save/share functionality is disabled.", + ) + parser.add_argument("--debug", action="store_true", default=False, help="Enable debug mode") + + args = parser.parse_args() + + main( + model=args.model, + output_dir=args.output_dir, + debug=args.debug, + ) diff --git a/tests/benchmarks/test_serve_cli.py b/tests/benchmarks/test_serve_cli.py index 8e9a3bfce..ee371330b 100644 --- a/tests/benchmarks/test_serve_cli.py +++ b/tests/benchmarks/test_serve_cli.py @@ -12,6 +12,8 @@ stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "q if current_omni_platform.is_xpu(): stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "xpu" / "qwen2_5_omni_ci.yaml")] +elif current_omni_platform.is_rocm(): + stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml")] # Create parameter combinations for model and stage config test_params = [ @@ -21,7 +23,7 @@ test_params = [ @pytest.mark.core_model @pytest.mark.benchmark -@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=3) +@hardware_test(res={"cuda": "L4", "xpu": "B60", "rocm": "MI325"}, num_cards=3) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_bench_serve_chat(omni_server): command = [ diff --git a/tests/conftest.py b/tests/conftest.py index 75cfaff98..d537ff531 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1236,8 +1236,6 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st port = params.port stage_config_path = params.stage_config_path if run_level == "advanced_model" and stage_config_path is not None: - # Dynamically detect stages from config to avoid KeyError - # for models with fewer stages (e.g., BAGEL has 2, Qwen3-Omni has 3) with open(stage_config_path, encoding="utf-8") as f: _cfg = yaml.safe_load(f) or {} _stage_ids = [s["stage_id"] for s in _cfg.get("stage_args", []) if "stage_id" in s] diff --git a/tests/diffusion/models/t5_encoder/test_t5_encoder_tp.py b/tests/diffusion/models/t5_encoder/test_t5_encoder_tp.py new file mode 100644 index 000000000..b36acc68a --- /dev/null +++ b/tests/diffusion/models/t5_encoder/test_t5_encoder_tp.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch +from transformers import T5Config +from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config + +from vllm_omni.diffusion.models.t5_encoder.t5_encoder import ( + T5EncoderModel, + T5SelfAttention, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_T5_MODULE = "vllm_omni.diffusion.models.t5_encoder.t5_encoder" + +SMALL_T5_CONFIG = dict( + d_model=64, + d_kv=8, + d_ff=128, + num_heads=8, + num_layers=2, + vocab_size=256, + relative_attention_num_buckets=32, + relative_attention_max_distance=128, + is_gated_act=True, + dense_act_fn="gelu_new", + layer_norm_epsilon=1e-6, + feed_forward_proj="gated-gelu", +) + + +@pytest.fixture(scope="module") +def t5_config() -> T5Config: + return T5Config(**SMALL_T5_CONFIG) + + +@pytest.fixture(scope="function", autouse=True) +def setup_tp_group(monkeypatch, mocker): + """Set up TP=2, rank=0, VllmConfig, and mock activation for all tests.""" + device_config = DeviceConfig(device="cpu") + + # TP world size + monkeypatch.setattr("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr( + "vllm.model_executor.layers.vocab_parallel_embedding.get_tensor_model_parallel_world_size", + lambda: 2, + ) + + monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + "vllm.model_executor.layers.vocab_parallel_embedding.get_tensor_model_parallel_rank", + lambda: 0, + ) + + # TP group + mock_tp_group = mocker.MagicMock() + mock_tp_group.world_size = 2 + mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_tp_group) + + monkeypatch.setattr(f"{_T5_MODULE}.get_act_fn", lambda _: torch.nn.GELU()) + + with set_current_vllm_config(VllmConfig(device_config=device_config)): + yield + + +class TestRelativePositionBiasTPSlicing: + """Verify compute_bias slices heads correctly per TP rank.""" + + def test_compute_bias_shape(self, t5_config): + attn = T5SelfAttention(t5_config, has_relative_attention_bias=True) + + seq_len = 6 + bias = attn.compute_bias(seq_len, seq_len, device=torch.device("cpu")) + + local_heads = t5_config.num_heads // 2 + assert bias.shape == (1, local_heads, seq_len, seq_len) + + def test_all_ranks_cover_all_heads(self, t5_config, monkeypatch): + seq_len = 4 + + biases = [] + ref_weight = None + for rank in range(2): + monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_rank", lambda r=rank: r) + attn = T5SelfAttention(t5_config, has_relative_attention_bias=True) + if rank > 0: + attn.relative_attention_bias.weight.data.copy_(ref_weight) + else: + ref_weight = attn.relative_attention_bias.weight.data.clone() + biases.append(attn.compute_bias(seq_len, seq_len, device=torch.device("cpu"))) + + full_bias = torch.cat(biases, dim=1) + assert full_bias.shape == (1, t5_config.num_heads, seq_len, seq_len) + + +class TestT5EncoderModelWeightLoading: + """Test weight loading at the top-level T5EncoderModel.""" + + def test_model_instantiation(self, t5_config): + model = T5EncoderModel(t5_config, prefix="text_encoder") + + assert model.config is t5_config + assert model.encoder is not None + assert len(model.encoder.block) == t5_config.num_layers + + def test_embedding_shape(self, t5_config): + model = T5EncoderModel(t5_config, prefix="text_encoder") + + assert model.shared.embedding_dim == t5_config.d_model + + def test_embed_input_ids(self, t5_config, monkeypatch): + # Verify method and output shape + model = T5EncoderModel(t5_config, prefix="text_encoder") + + # Mock all-reduce to be identity (no actual TP communication) + monkeypatch.setattr( + "vllm.model_executor.layers.vocab_parallel_embedding.tensor_model_parallel_all_reduce", + lambda x: x, + ) + + input_ids = torch.randint(0, t5_config.vocab_size, (2, 8)) + embeddings = model.embed_input_ids(input_ids) + + assert embeddings.shape == (2, 8, t5_config.d_model) + + def test_qkv_weights_loaded_through_blocks(self): + # Verify that HF-style separate Q/K/V weights can be loaded through + # the block hierarchy + config = T5Config(**{**SMALL_T5_CONFIG, "num_layers": 1}) + model = T5EncoderModel(config, prefix="text_encoder") + + inner_dim = config.num_heads * config.d_kv + prefix = "encoder.block.0.layer.0.SelfAttention." + + loaded = model.load_weights( + [ + (prefix + "q.weight", torch.randn(inner_dim, config.d_model)), + (prefix + "k.weight", torch.randn(inner_dim, config.d_model)), + (prefix + "v.weight", torch.randn(inner_dim, config.d_model)), + ] + ) + + assert len(loaded) > 0 + attn = model.encoder.block[0].layer[0].SelfAttention + expected_qkv_dim = 3 * (config.num_heads // 2) * config.d_kv + assert attn.qkv_proj.weight.shape == (expected_qkv_dim, config.d_model) + + +class TestTPConstraints: + """Verify that invalid TP configurations raise clear errors.""" + + def test_num_heads_not_divisible_by_tp(self): + config = T5Config(**{**SMALL_T5_CONFIG, "num_heads": 7}) + with pytest.raises(AssertionError, match=r"num_heads.*must be divisible by tp_size"): + T5SelfAttention(config) + + def test_num_heads_divisible_by_tp(self, t5_config): + attn = T5SelfAttention(t5_config) + assert attn.n_heads_per_partition == 4 diff --git a/tests/diffusion/quantization/test_int8_config.py b/tests/diffusion/quantization/test_int8_config.py new file mode 100644 index 000000000..9b5d67fcb --- /dev/null +++ b/tests/diffusion/quantization/test_int8_config.py @@ -0,0 +1,463 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for Int8 quantization config.""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch +from pytest_mock import MockerFixture +from torch.nn import Module, Parameter +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + +from vllm_omni.diffusion.quantization import ( + get_diffusion_quant_config, + get_vllm_quant_config_for_layers, +) +from vllm_omni.platforms import current_omni_platform + +npu_available = pytest.mark.skipif(not current_omni_platform.is_npu(), reason="NPU platform not available.") + +cuda_available = pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="GPU platform not available.") + + +def test_int8_config_creation(): + """Test that Int8 config can be created.""" + config = get_diffusion_quant_config("int8") + assert config is not None + assert config.get_name() == "int8" + + +def test_vllm_config_extraction(): + """Test that vLLM config can be extracted from diffusion config.""" + diff_config = get_diffusion_quant_config("int8") + vllm_config = get_vllm_quant_config_for_layers(diff_config) + assert vllm_config is not None + assert vllm_config.activation_scheme == "dynamic" + + +def test_none_quantization(): + """Test that None quantization returns None config.""" + config = get_diffusion_quant_config(None) + assert config is None + vllm_config = get_vllm_quant_config_for_layers(config) + assert vllm_config is None + + +def test_invalid_quantization(): + """Test that invalid quantization method raises error.""" + with pytest.raises(ValueError, match="Unknown quantization method"): + get_diffusion_quant_config("invalid_method") + + +def test_int8_config_with_custom_params(): + """Test Int8 config with custom parameters.""" + config = get_diffusion_quant_config( + "int8", + activation_scheme="dynamic", + ignored_layers=["proj_out"], + ) + assert config is not None + assert config.activation_scheme == "dynamic" + assert "proj_out" in config.ignored_layers + + +def test_supported_methods(): + """Test that supported methods list is correct.""" + from vllm_omni.diffusion.quantization import SUPPORTED_QUANTIZATION_METHODS + + assert "int8" in SUPPORTED_QUANTIZATION_METHODS + + +def test_quantization_integration(): + """Test end-to-end quantization flow through OmniDiffusionConfig.""" + from vllm_omni.diffusion.data import OmniDiffusionConfig + + # Test with quantization string only + config = OmniDiffusionConfig(model="test", quantization="int8") + assert config.quantization_config is not None + assert config.quantization_config.get_name() == "int8" + + # Test with quantization_config dict + config2 = OmniDiffusionConfig( + model="test", + quantization_config={"method": "int8", "activation_scheme": "dynamic"}, + ) + assert config2.quantization_config is not None + assert config2.quantization_config.get_name() == "int8" + assert config2.quantization_config.activation_scheme == "dynamic" + + # Test that vLLM config can be extracted + vllm_config = config.quantization_config.get_vllm_quant_config() + assert vllm_config is not None + + +def test_quantization_dict_not_mutated(): + """Test that passing a dict to quantization_config doesn't mutate it.""" + from vllm_omni.diffusion.data import OmniDiffusionConfig + + original_dict = {"method": "int8", "activation_scheme": "dynamic"} + dict_copy = original_dict.copy() + + OmniDiffusionConfig(model="test", quantization_config=original_dict) + + # Original dict should be unchanged + assert original_dict == dict_copy + + +def test_quantization_conflicting_methods_warning(caplog): + """Test warning when quantization and quantization_config['method'] conflict.""" + import logging + + from vllm_omni.diffusion.data import OmniDiffusionConfig + + with caplog.at_level(logging.WARNING): + config = OmniDiffusionConfig( + model="test", + quantization="int8", # This should be overridden + quantization_config={"method": "int8", "activation_scheme": "dynamic"}, + ) + # No warning when methods match + assert config.quantization_config is not None + + +def test_get_quant_method(mocker: MockerFixture): + """Test for get_quant_method method for GPU""" + from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod + + diff_config = get_diffusion_quant_config("int8") + vllm_config = get_vllm_quant_config_for_layers(diff_config) + + def _fake_init(self, quant_config): + pass + + layer = MagicMock(spec=LinearBase) + mocker.patch.object(Int8OnlineLinearMethod, "__init__", _fake_init) + + prefix = "test_layer" + + # Mock the platform to be GPU + with ( + patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=True), + patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=False), + ): + method = vllm_config.get_quant_method(layer, prefix) + assert isinstance(method, Int8OnlineLinearMethod) + + # Test skipping quantization for a layer + vllm_config.ignored_layers = [prefix] + method = vllm_config.get_quant_method(layer, prefix) + assert isinstance(method, UnquantizedLinearMethod) + + +def test_get_npu_quant_method(): + """Test for get_quant_method method for NPU""" + from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod + + diff_config = get_diffusion_quant_config("int8") + vllm_config = get_vllm_quant_config_for_layers(diff_config) + + layer = MagicMock(spec=LinearBase) + prefix = "test_layer" + + # Mock the platform to be NPU + with ( + patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=False), + patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=True), + ): + method = vllm_config.get_quant_method(layer, prefix) + assert isinstance(method, NPUInt8OnlineLinearMethod) + + # Test skipping quantization for a layer + vllm_config.ignored_layers = [prefix] + method = vllm_config.get_quant_method(layer, prefix) + assert isinstance(method, UnquantizedLinearMethod) + + +class TestInt8LinearMethod: + @pytest.fixture + def mock_quant_config(self, mocker): + return mocker.Mock() + + @pytest.fixture + def mock_kernel(self, mocker): + kernel = mocker.Mock() + kernel.process_weights_after_loading = mocker.Mock() + kernel.apply_weights = mocker.Mock(return_value=torch.randn(1, 10)) + return kernel + + @pytest.fixture + def patch_deps(self, mocker, mock_kernel): + # mock init_int8_linear_kernel + mocker.patch("vllm_omni.diffusion.quantization.int8.init_int8_linear_kernel", return_value=mock_kernel) + return mock_kernel + + def test_init(self, patch_deps, mock_quant_config): + # test for Int8LinearMethod init + from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod, init_int8_linear_kernel + + method = Int8LinearMethod(mock_quant_config) + + assert method.quant_config == mock_quant_config + init_int8_linear_kernel.assert_called_once_with( + is_channelwise=False, is_static_input_scheme=False, input_symmetric=True, module_name="Int8LinearMethod" + ) + assert method.int8_linear == patch_deps + + def test_process_weights_after_loading(self, patch_deps, mock_quant_config): + from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod + + method = Int8LinearMethod(mock_quant_config) + layer = Module() + + method.process_weights_after_loading(layer) + patch_deps.process_weights_after_loading.assert_called_once_with(layer) + + def test_apply(self, patch_deps, mock_quant_config): + from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod + + method = Int8LinearMethod(mock_quant_config) + layer = Module() + x = torch.randn(1, 128) + bias = torch.randn(128) + + output = method.apply(layer, x, bias) + + patch_deps.apply_weights.assert_called_once_with(layer, x, bias) + assert isinstance(output, torch.Tensor) + + +class TestInt8OnlineLinearMethod: + @pytest.fixture + def mock_quant_config(self, mocker): + return mocker.Mock() + + @pytest.fixture + def mock_deps(self, mocker): + # mock kernel + mock_kernel = mocker.Mock() + mock_kernel.layer_param_names = ("weight", "weight_scale", "input_scale", "input_zero_point", "azp_adj") + mocker.patch("vllm_omni.diffusion.quantization.int8.init_int8_linear_kernel", return_value=mock_kernel) + mocker.patch("vllm_omni.diffusion.quantization.int8.replace_parameter") + + # mock scaled_int8_quant return value + mock_qweight = torch.ones((128, 64), dtype=torch.int8) + mock_scale = torch.randn(128) + mock_quant = mocker.patch( + "vllm_omni.diffusion.quantization.int8.ops.scaled_int8_quant", return_value=(mock_qweight, mock_scale, None) + ) + return {"kernel": mock_kernel, "quant": mock_quant, "mock_qweight": mock_qweight, "mock_scale": mock_scale} + + def test_process_weights_after_loading(self, mock_deps, mock_quant_config): + from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod + + method = Int8OnlineLinearMethod(mock_quant_config) + layer = Module() + layer.weight = Parameter(torch.randn(128, 64)) + method.process_weights_after_loading(layer) + mock_deps["quant"].assert_called_once_with(layer.weight, scale=None) + + +@npu_available +class TestNPUInt8LinearMethod: + qweight_mock = torch.randn((128, 64)).to(dtype=torch.int8) + scale_mock = torch.randn(128) + out_mock = torch.randn((16, 128)) + + @pytest.fixture + def mock_torch_npu(self, mocker): + torch_npu = MagicMock() + + mocker.patch("vllm_omni.diffusion.quantization.int8.torch_npu", return_value=torch_npu) + mocker.patch( + "vllm_omni.diffusion.quantization.int8.torch_npu.npu_dynamic_quant", + return_value=(self.qweight_mock, self.scale_mock), + ) + mocker.patch("vllm_omni.diffusion.quantization.int8.torch_npu.npu_quant_matmul", return_value=self.out_mock) + return torch_npu + + @pytest.fixture + def mock_quant_config(self, mocker): + return mocker.Mock() + + @pytest.fixture + def mock_layer(self, mocker): + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(self.qweight_mock, requires_grad=False) + layer.weight_scale = torch.nn.Parameter(self.scale_mock, requires_grad=False) + return layer + + def test_npu_int8_process_weights_after_loading(self, mock_layer, mock_quant_config, mock_torch_npu): + from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod + + method = NPUInt8LinearMethod(mock_quant_config) + ori_weight_shape = mock_layer.weight.shape + + method.process_weights_after_loading(mock_layer) + + assert mock_layer.weight.shape == ori_weight_shape[::-1] + assert mock_layer.weight.is_contiguous() + + def test_npu_int8_apply(self, mock_layer, mock_quant_config, mock_torch_npu): + from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod + + method = NPUInt8LinearMethod(mock_quant_config) + x = torch.randn(1, 16, 64) + + output = method.apply(mock_layer, x) + assert output.shape == (1, 16, 128) + + def test_npu_int8_online_process_weights(self, mock_layer, mock_quant_config, mock_torch_npu): + from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod + + method = NPUInt8OnlineLinearMethod(mock_quant_config) + method.process_weights_after_loading(mock_layer) + + assert mock_layer.weight.shape == (64, 128) + assert torch.equal(mock_layer.weight_scale, self.scale_mock) + + +@pytest.fixture +def quant_config(): + """Shared quant config fixture for smoke tests.""" + from vllm_omni.diffusion.quantization.int8 import Int8Config + + return Int8Config( + is_checkpoint_int8_serialized=False, + activation_scheme="dynamic", + ) + + +@npu_available +class TestNPUInt8Smoke: + """Smoke tests using real torch_npu, only run on NPU.""" + + @pytest.fixture + def real_layer(self): + """Create a real linear layer with fp16 weights on NPU""" + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter( + torch.randn(128, 64, dtype=torch.float16, device="npu"), + requires_grad=False, + ) + layer.logical_widths = [128] + layer.input_size_per_partition = 64 + layer.output_size_per_partition = 128 + layer.orig_dtype = torch.float16 + return layer + + def test_real_npu_dynamic_quant_shape_contract(self, quant_config, real_layer): + """Smoke test: verify npu_dynamic_quant returns correct shapes.""" + import torch_npu + + # Call real torch_npu.npu_dynamic_quant + weight = real_layer.weight + qweight, scale = torch_npu.npu_dynamic_quant(weight) + + assert qweight.shape == weight.shape + assert qweight.dtype == torch.int8 + assert scale.shape == (weight.shape[0],) + + def test_real_npu_online_process_weights_after_loading(self, quant_config, real_layer): + """Smoke test: full process_weights_after_loading with real torch_npu.""" + from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod + + method = NPUInt8OnlineLinearMethod(quant_config) + + method.process_weights_after_loading(real_layer) + + assert real_layer.weight.shape == (64, 128) + assert real_layer.weight.dtype == torch.int8 + assert hasattr(real_layer, "weight_scale") + assert real_layer.weight_scale.shape == (128,) + + def test_real_npu_int8_apply_forward(self, quant_config): + """Smoke test: forward pass with real npu_quant_matmul.""" + import torch_npu + + from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod + + method = NPUInt8LinearMethod(quant_config) + + # Create layer with pre-processed weights on NPU + layer = torch.nn.Module() + weight_fp16 = torch.randn(128, 64, dtype=torch.float16, device="npu") + qweight, scale = torch_npu.npu_dynamic_quant(weight_fp16) + layer.weight = torch.nn.Parameter(qweight.t().contiguous(), requires_grad=False) + layer.weight_scale = torch.nn.Parameter(scale.squeeze(), requires_grad=False) + + # Forward pass on NPU + x = torch.randn(2, 16, 64, dtype=torch.float16, device="npu") + output = method.apply(layer, x) + + assert output.shape == (2, 16, 128) + assert output.dtype == torch.float16 + + +@cuda_available +class TestCudaInt8Smoke: + """Smoke tests using real CUDA kernels, only on CUDA""" + + @pytest.fixture + def real_layer(self): + """Create a real linear layer with fp16 weights on CUDA""" + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter( + torch.randn(128, 64, dtype=torch.float16, device="cuda"), + requires_grad=False, + ) + layer.logical_widths = [128] + layer.input_size_per_partition = 64 + layer.output_size_per_partition = 128 + layer.orig_dtype = torch.float16 + return layer + + def test_real_cuda_scaled_int8_quant_shape_contract(self, quant_config): + """Smoke test: verify scaled_int8_quant returns correct shapes.""" + from vllm import _custom_ops as ops + + weight = torch.randn(128, 64, dtype=torch.float16, device="cuda") + qweight, scale, _ = ops.scaled_int8_quant(weight, scale=None) + + assert qweight.shape == weight.shape + assert qweight.dtype == torch.int8 + assert scale.shape == (weight.shape[0], 1) + + def test_real_cuda_online_process_weights_after_loading(self, quant_config, real_layer): + """Smoke test: full process_weights_after_loading with real CUDA ops.""" + from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod + + method = Int8OnlineLinearMethod(quant_config) + + method.process_weights_after_loading(real_layer) + + assert real_layer.weight.shape == (64, 128) + assert real_layer.weight.dtype == torch.int8 + assert hasattr(real_layer, "weight_scale") + + def test_real_cuda_int8_apply_forward(self, quant_config): + """Smoke test: forward pass with real CUDA int8 kernel.""" + from vllm import _custom_ops as ops + + from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod + + method = Int8LinearMethod(quant_config) + + # Create layer with pre-processed weights + layer = torch.nn.Module() + weight_fp16 = torch.randn(128, 64, dtype=torch.float16, device="cuda") + qweight, scale, _ = ops.scaled_int8_quant(weight_fp16, scale=None) + layer.weight = torch.nn.Parameter(qweight.t(), requires_grad=False) + layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False) + + # Set required attributes for kernel + layer.input_scale = None + layer.input_zero_point = None + layer.azp_adj = None + + # Forward pass + x = torch.randn(2, 16, 64, dtype=torch.float16, device="cuda") + output = method.apply(layer, x) + + assert output.shape == (2, 16, 128) + assert output.dtype == torch.float16 diff --git a/tests/e2e/offline_inference/test_bagel_img2img.py b/tests/e2e/offline_inference/test_bagel_img2img.py index 8c734c6a2..c7df4f91b 100644 --- a/tests/e2e/offline_inference/test_bagel_img2img.py +++ b/tests/e2e/offline_inference/test_bagel_img2img.py @@ -5,7 +5,7 @@ End-to-end test for Bagel img2img generation. This test validates that the Bagel model generates images from an input image -and text prompt that match expected reference pixel values within a ±5 tolerance. +and text prompt that match expected reference pixel values within a ±10 tolerance. Equivalent to running: python3 examples/offline_inference/bagel/end2end.py \ @@ -25,6 +25,7 @@ from vllm.assets.image import ImageAsset from tests.conftest import modify_stage_config from tests.utils import hardware_test from vllm_omni.entrypoints.omni import Omni +from vllm_omni.platforms import current_omni_platform # Reference pixel data extracted from the known-good output image # Generated with seed=52, num_inference_steps=15, @@ -43,7 +44,21 @@ REFERENCE_PIXELS = [ {"position": (256, 256), "rgb": (181, 202, 222)}, ] -PIXEL_TOLERANCE = 5 +if current_omni_platform.is_rocm(): + REFERENCE_PIXELS = [ + {"position": (100, 100), "rgb": (156, 172, 215)}, + {"position": (400, 50), "rgb": (106, 144, 216)}, + {"position": (700, 100), "rgb": (118, 158, 231)}, + {"position": (150, 400), "rgb": (183, 23, 48)}, + {"position": (512, 336), "rgb": (218, 215, 191)}, + {"position": (700, 400), "rgb": (194, 14, 42)}, + {"position": (100, 600), "rgb": (105, 10, 16)}, + {"position": (400, 600), "rgb": (167, 33, 46)}, + {"position": (700, 600), "rgb": (102, 86, 92)}, + {"position": (256, 256), "rgb": (181, 201, 220)}, + ] + +PIXEL_TOLERANCE = 10 DEFAULT_PROMPT = "<|fim_middle|><|im_start|>Change the grass color to red<|im_end|>" @@ -191,7 +206,7 @@ def _resolve_stage_config(config_path: str, run_level: str) -> str: @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.diffusion -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) def test_bagel_img2img_shared_memory_connector(run_level): """Test Bagel img2img with shared memory connector.""" input_image = _load_input_image() diff --git a/tests/e2e/offline_inference/test_bagel_text2img.py b/tests/e2e/offline_inference/test_bagel_text2img.py index ed369aedd..505e12438 100644 --- a/tests/e2e/offline_inference/test_bagel_text2img.py +++ b/tests/e2e/offline_inference/test_bagel_text2img.py @@ -31,6 +31,7 @@ from PIL import Image from tests.conftest import modify_stage_config from tests.utils import hardware_test from vllm_omni.entrypoints.omni import Omni +from vllm_omni.platforms import current_omni_platform # Reference pixel data extracted from the known-good output image # Each entry contains (x, y) position and expected (R, G, B) values @@ -49,6 +50,20 @@ REFERENCE_PIXELS = [ {"position": (256, 256), "rgb": (171, 160, 153)}, ] +if current_omni_platform.is_rocm(): + REFERENCE_PIXELS = [ + {"position": (100, 100), "rgb": (123, 119, 100)}, + {"position": (400, 50), "rgb": (162, 161, 142)}, + {"position": (700, 100), "rgb": (171, 156, 127)}, + {"position": (150, 400), "rgb": (131, 128, 112)}, + {"position": (512, 512), "rgb": (134, 61, 59)}, + {"position": (700, 400), "rgb": (204, 107, 43)}, + {"position": (100, 700), "rgb": (201, 180, 165)}, + {"position": (400, 700), "rgb": (140, 108, 87)}, + {"position": (700, 700), "rgb": (247, 205, 145)}, + {"position": (256, 256), "rgb": (171, 160, 153)}, + ] + # Maximum allowed difference per color channel PIXEL_TOLERANCE = 5 @@ -181,7 +196,7 @@ def _resolve_stage_config(config_path: str, run_level: str) -> str: @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.diffusion -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) def test_bagel_text2img_shared_memory_connector(run_level): """Test Bagel text2img with shared memory connector.""" config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml") diff --git a/tests/e2e/online_serving/test_bagel_online.py b/tests/e2e/online_serving/test_bagel_online.py index a5e26db1e..62e79906c 100644 --- a/tests/e2e/online_serving/test_bagel_online.py +++ b/tests/e2e/online_serving/test_bagel_online.py @@ -97,7 +97,7 @@ def test_bagel_text2img_online(omni_server, openai_client) -> None: @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.diffusion -@hardware_test(res={"cuda": "H100"}) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_bagel_img2img_online(omni_server, openai_client) -> None: """Test Bagel img2img via OpenAI-compatible chat completions API.""" diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index ef4d40198..fcda20ba3 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -44,13 +44,9 @@ def get_chunk_config(): return path -# CI stage config for 2xH100-80G GPUs or AMD GPU MI325 -if current_omni_platform.is_rocm(): - # ROCm stage config optimized for MI325 GPU - stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")] -elif current_omni_platform.is_xpu(): +if current_omni_platform.is_xpu(): stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")] -else: +else: # MI325 GPU should share the same config as H100 stage_configs = [get_chunk_config()] # Create parameter combinations for model and stage config diff --git a/tests/e2e/online_serving/test_voxtral_tts.py b/tests/e2e/online_serving/test_voxtral_tts.py new file mode 100644 index 000000000..d4de073f8 --- /dev/null +++ b/tests/e2e/online_serving/test_voxtral_tts.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E Online tests for Voxtral TTS model with text input and audio output. + +These tests verify the /v1/audio/speech endpoint works correctly with +actual model inference, not mocks. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +from pathlib import Path + +import httpx +import pytest + +from tests.conftest import OmniServerParams +from tests.utils import hardware_test + +MODEL = "mistralai/tts-model" + +STAGE_CONFIG = str( + Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "voxtral_tts.yaml" +) +EXTRA_ARGS = ["--trust-remote-code", "--enforce-eager", "--disable-log-stats"] +TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG, server_args=EXTRA_ARGS)] + + +def make_speech_request( + host: str, + port: int, + text: str, + voice: str = "casual_female", + language: str = "English", + timeout: float = 120.0, +) -> httpx.Response: + """Make a request to the /v1/audio/speech endpoint.""" + url = f"http://{host}:{port}/v1/audio/speech" + payload = { + "input": text, + "voice": voice, + "language": language, + } + + with httpx.Client(timeout=timeout) as client: + return client.post(url, json=payload) + + +def verify_wav_audio(content: bytes) -> bool: + """Verify that content is valid WAV audio data.""" + # WAV files start with "RIFF" header + if len(content) < 44: # Minimum WAV header size + return False + return content[:4] == b"RIFF" and content[8:12] == b"WAVE" + + +# Minimum expected audio size for a short sentence (~1 second of 24kHz 16-bit mono WAV) +MIN_AUDIO_BYTES = 10000 + + +@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) +class TestVoxtralTTSFixedVoice: + """E2E tests for Voxtral TTS model.""" + + @pytest.mark.core_model + @pytest.mark.omni + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_speech_english_basic(self, omni_server) -> None: + """Test basic English TTS generation.""" + response = make_speech_request( + host=omni_server.host, + port=omni_server.port, + text="Hello, how are you?", + voice="casual_female", + language="English", + ) + + assert response.status_code == 200, f"Request failed: {response.text}" + assert response.headers.get("content-type") == "audio/wav" + assert verify_wav_audio(response.content), "Response is not valid WAV audio" + assert len(response.content) > MIN_AUDIO_BYTES, ( + f"Audio content too small ({len(response.content)} bytes), expected at least {MIN_AUDIO_BYTES} bytes" + ) + + @pytest.mark.advanced_model + @pytest.mark.omni + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_speech_different_voices(self, omni_server) -> None: + """Test TTS with different voice presets.""" + voices = ["casual_female", "neutral_male"] + for voice in voices: + response = make_speech_request( + host=omni_server.host, + port=omni_server.port, + text="Testing voice selection.", + voice=voice, + language="English", + ) + + assert response.status_code == 200, f"Request failed for voice {voice}: {response.text}" + assert verify_wav_audio(response.content), f"Invalid WAV for voice {voice}" + + @pytest.mark.advanced_model + @pytest.mark.omni + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_speech_invalid_voice_rejected(self, omni_server) -> None: + """Request with a non-existent voice should return an error response.""" + response = make_speech_request( + host=omni_server.host, + port=omni_server.port, + text="This voice does not exist.", + voice="nonexistent_voice_xyz", + language="English", + ) + + # vLLM returns ErrorResponse as JSON with HTTP 200; + # verify the body is a JSON error, not valid WAV audio. + data = response.json() + assert "error" in data or "message" in data, f"Expected error response for invalid voice, got: {data}" + assert not verify_wav_audio(response.content), "Should not return valid audio for invalid voice" + + @pytest.mark.advanced_model + @pytest.mark.omni + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_speech_binary_response_not_utf8_error(self, omni_server) -> None: + """ + Regression test: Verify binary audio is returned, not UTF-8 error. + + This test ensures the multimodal_output property correctly retrieves + audio from completion outputs, preventing the "TTS model did not + produce audio output" error. + """ + response = make_speech_request( + host=omni_server.host, + port=omni_server.port, + text="This should return binary audio, not a JSON error.", + voice="casual_female", + language="English", + ) + + # Should NOT be a JSON error response + assert response.status_code == 200, f"Request failed: {response.text}" + + # Verify it's binary audio, not JSON + try: + # If this succeeds and starts with {"error", it's a bug + text = response.content.decode("utf-8") + assert not text.startswith('{"error"'), f"Got error response instead of audio: {text}" + except UnicodeDecodeError: + # This is expected - binary audio can't be decoded as UTF-8 + pass + + assert verify_wav_audio(response.content), "Response is not valid WAV audio" + + +@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) +class TestVoxtralTTSAPIEndpoints: + """Test API endpoint functionality.""" + + @pytest.mark.advanced_model + @pytest.mark.omni + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_list_voices_endpoint(self, omni_server) -> None: + """Test the /v1/audio/voices endpoint returns available voices.""" + url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/voices" + + with httpx.Client(timeout=30.0) as client: + response = client.get(url) + + assert response.status_code == 200 + data = response.json() + assert "voices" in data + assert isinstance(data["voices"], list) + assert len(data["voices"]) > 0 + + @pytest.mark.advanced_model + @pytest.mark.omni + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_models_endpoint(self, omni_server) -> None: + """Test the /v1/models endpoint returns loaded model.""" + url = f"http://{omni_server.host}:{omni_server.port}/v1/models" + + with httpx.Client(timeout=30.0) as client: + response = client.get(url) + + assert response.status_code == 200 + data = response.json() + assert "data" in data + assert len(data["data"]) > 0 diff --git a/tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml b/tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml index 9fa5930c8..b95afd223 100644 --- a/tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml +++ b/tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml @@ -13,8 +13,8 @@ stage_args: model_arch: Qwen2_5OmniForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - max_model_len: 2400 - max_num_batched_tokens: 2400 + max_model_len: 16384 + max_num_batched_tokens: 16384 max_num_seqs: 1 gpu_memory_utilization: 0.8 skip_mm_profiling: true @@ -44,8 +44,8 @@ stage_args: model_arch: Qwen2_5OmniForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - max_model_len: 2400 - max_num_batched_tokens: 2400 + max_model_len: 16384 + max_num_batched_tokens: 16384 max_num_seqs: 1 gpu_memory_utilization: 0.8 skip_mm_profiling: true @@ -59,7 +59,7 @@ stage_args: temperature: 0.9 top_p: 0.8 top_k: 40 - max_tokens: 128 + max_tokens: 4096 seed: 42 detokenize: True repetition_penalty: 1.05 @@ -79,6 +79,8 @@ stage_args: trust_remote_code: true enable_prefix_caching: false engine_output_type: audio + max_num_batched_tokens: 4096 + max_model_len: 4096 engine_input_source: [1] final_output: true final_output_type: audio @@ -86,7 +88,7 @@ stage_args: temperature: 0.0 top_p: 1.0 top_k: -1 - max_tokens: 128 + max_tokens: 4096 seed: 42 detokenize: True repetition_penalty: 1.1 diff --git a/tests/model_executor/models/voxtral_tts/__init__.py b/tests/model_executor/models/voxtral_tts/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/model_executor/models/voxtral_tts/test_audio_tokenizer_parsing.py b/tests/model_executor/models/voxtral_tts/test_audio_tokenizer_parsing.py new file mode 100644 index 000000000..5a560d95a --- /dev/null +++ b/tests/model_executor/models/voxtral_tts/test_audio_tokenizer_parsing.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Tests for parse_batched_audio_input and apply_ctx_frames_cutting. + +These functions parse the [ctx_frames, context_length, ...tokens] wire format +produced by generator2tokenizer_async_chunk and cut leading context samples +from decoded audio arrays. +""" + +import pytest +import torch + +from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts import ( + apply_ctx_frames_cutting, + parse_batched_audio_input, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +NUM_CODEBOOKS = 37 + + +# ────────────────────────────────────────────────────────────────── +# parse_batched_audio_input tests +# ────────────────────────────────────────────────────────────────── + + +def _make_input_ids(*requests: tuple[int, int, torch.Tensor]) -> torch.Tensor: + """Build a flat input_ids tensor from (ctx_frames, context_length, tokens) tuples.""" + parts = [] + for ctx_frames, context_length, tokens in requests: + parts.append(torch.tensor([ctx_frames, context_length])) + parts.append(tokens.flatten()) + return torch.cat(parts) + + +def test_parse_single_request(): + """Single request is parsed correctly.""" + ctx_frames = 3 + context_length = 5 + total_frames = ctx_frames + context_length + tokens = torch.arange(total_frames * NUM_CODEBOOKS).reshape(total_frames, NUM_CODEBOOKS) + input_ids = _make_input_ids((ctx_frames, context_length, tokens)) + + all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS) + + assert len(all_audio_tokens) == 1 + assert len(all_ctx_frames) == 1 + assert all_ctx_frames[0] == ctx_frames + assert all_audio_tokens[0].shape == (total_frames, NUM_CODEBOOKS) + torch.testing.assert_close(all_audio_tokens[0], tokens) + + +def test_parse_multiple_requests(): + """Multiple requests in a batch are parsed correctly.""" + req1 = (2, 4, torch.ones((6, NUM_CODEBOOKS), dtype=torch.long)) + req2 = (0, 3, torch.full((3, NUM_CODEBOOKS), 7, dtype=torch.long)) + input_ids = _make_input_ids(req1, req2) + + all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS) + + assert len(all_audio_tokens) == 2 + assert all_ctx_frames == [2, 0] + assert all_audio_tokens[0].shape == (6, NUM_CODEBOOKS) + assert all_audio_tokens[1].shape == (3, NUM_CODEBOOKS) + assert torch.all(all_audio_tokens[0] == 1) + assert torch.all(all_audio_tokens[1] == 7) + + +def test_parse_zero_ctx_frames(): + """ctx_frames=0 means no context, just chunk data.""" + ctx_frames = 0 + context_length = 10 + tokens = torch.randint(0, 100, (context_length, NUM_CODEBOOKS)) + input_ids = _make_input_ids((ctx_frames, context_length, tokens)) + + all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS) + + assert all_ctx_frames[0] == 0 + assert all_audio_tokens[0].shape == (context_length, NUM_CODEBOOKS) + + +def test_parse_fails_on_misaligned_input(): + """Input not divisible by num_codebooks should raise AssertionError.""" + # 2 header elements + 5 data elements (not divisible by 37) + input_ids = torch.tensor([0, 1] + [0] * 5) + with pytest.raises(AssertionError, match="divisible by"): + parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS) + + +def test_parse_custom_num_codebooks(): + """Works with a non-default num_codebooks value.""" + cb = 4 + ctx_frames = 1 + context_length = 2 + total_frames = ctx_frames + context_length + tokens = torch.arange(total_frames * cb).reshape(total_frames, cb) + input_ids = _make_input_ids((ctx_frames, context_length, tokens)) + + all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=cb) + + assert all_audio_tokens[0].shape == (total_frames, cb) + torch.testing.assert_close(all_audio_tokens[0], tokens) + + +# ────────────────────────────────────────────────────────────────── +# apply_ctx_frames_cutting tests +# ────────────────────────────────────────────────────────────────── + + +def test_cut_removes_leading_samples(): + """Context frames are removed from the front of the audio array.""" + downsample_factor = 240 + ctx_frames = 5 + total_samples = 2400 # 10 frames * 240 + audio = torch.arange(total_samples, dtype=torch.float32) + + result = apply_ctx_frames_cutting([audio], [ctx_frames], downsample_factor) + + expected_cut = ctx_frames * downsample_factor # 1200 + assert len(result) == 1 + assert result[0].shape[0] == total_samples - expected_cut + torch.testing.assert_close(result[0], audio[expected_cut:]) + + +def test_cut_zero_ctx_frames_unchanged(): + """ctx_frames=0 leaves the audio unchanged.""" + downsample_factor = 240 + audio = torch.randn(2400) + + result = apply_ctx_frames_cutting([audio], [0], downsample_factor) + + assert result[0].shape == audio.shape + torch.testing.assert_close(result[0], audio) + + +def test_cut_multiple_requests(): + """Each request in the batch gets its own ctx_frames cut.""" + downsample_factor = 100 + audio1 = torch.arange(1000, dtype=torch.float32) # 10 frames + audio2 = torch.arange(500, dtype=torch.float32) # 5 frames + + result = apply_ctx_frames_cutting([audio1, audio2], [3, 0], downsample_factor) + + assert len(result) == 2 + assert result[0].shape[0] == 1000 - 3 * 100 # 700 + assert result[1].shape[0] == 500 # unchanged + torch.testing.assert_close(result[0], audio1[300:]) + torch.testing.assert_close(result[1], audio2) + + +def test_cut_all_frames_returns_empty(): + """Cutting all frames returns an empty tensor.""" + downsample_factor = 100 + ctx_frames = 5 + audio = torch.randn(500) # exactly 5 frames + + result = apply_ctx_frames_cutting([audio], [ctx_frames], downsample_factor) + + assert result[0].numel() == 0 diff --git a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py new file mode 100644 index 000000000..f0fba8a0d --- /dev/null +++ b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py @@ -0,0 +1,329 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Tests for CUDA Graph acoustic transformer wrapper numerical equivalence. + +Verifies that CUDA Graph-accelerated decoding produces results equivalent +to eager mode, with special attention to padding cases where zero-padding +may introduce small numerical differences due to the flow matching ODE. + +Architecture note: the real FlowMatchingAudioTransformer uses a multi-step +Euler ODE with CFG. The synthetic model here uses simple linear layers +to exercise the wrapper mechanism while keeping the test lightweight. +""" + +import importlib.util +import os + +import pytest +import torch +import torch.nn as nn + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.core_model, + pytest.mark.cuda, +] + +DEVICE = torch.device("cuda:0") +HIDDEN_DIM = 64 +N_ACOUSTIC_CODEBOOK = 7 +SEMANTIC_CODEBOOK_SIZE = 128 +ACOUSTIC_EMBEDDINGS_LEVELS = 1024 + +# Load CUDAGraphAcousticTransformerWrapper: try package import first, fall back to direct file load +try: + from vllm_omni.model_executor.models.voxtral_tts.cuda_graph_acoustic_transformer_wrapper import ( + CUDAGraphAcousticTransformerWrapper, + ) + from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation import ( + AudioSpecialTokens, + ) +except Exception: + _WRAPPER_PATH = os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + os.pardir, + os.pardir, + "vllm_omni", + "model_executor", + "models", + "voxtral_tts", + "cuda_graph_acoustic_transformer_wrapper.py", + ) + _spec = importlib.util.spec_from_file_location( + "cuda_graph_acoustic_transformer_wrapper", os.path.abspath(_WRAPPER_PATH) + ) + _mod = importlib.util.module_from_spec(_spec) + _spec.loader.exec_module(_mod) + CUDAGraphAcousticTransformerWrapper = _mod.CUDAGraphAcousticTransformerWrapper + + _AUDIO_GEN_PATH = os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + os.pardir, + os.pardir, + "vllm_omni", + "model_executor", + "models", + "voxtral_tts", + "voxtral_tts_audio_generation.py", + ) + _spec2 = importlib.util.spec_from_file_location("voxtral_tts_audio_generation", os.path.abspath(_AUDIO_GEN_PATH)) + _mod2 = importlib.util.module_from_spec(_spec2) + _spec2.loader.exec_module(_mod2) + AudioSpecialTokens = _mod2.AudioSpecialTokens + + +class SyntheticModelArgs: + """Mimics MultimodalAudioModelArgs interface.""" + + def __init__(self): + self.semantic_codebook_size = SEMANTIC_CODEBOOK_SIZE + self.n_acoustic_codebook = N_ACOUSTIC_CODEBOOK + + +class SyntheticAcousticTransformer(nn.Module): + """A small acoustic transformer mimicking FlowMatchingAudioTransformer's interface. + + Uses simple linear layers so the test stays lightweight while exercising + the full wrapper mechanism (semantic logit, flow matching ODE, quantization). + """ + + def __init__(self): + super().__init__() + self.model_args = SyntheticModelArgs() + self.acoustic_embeddings_levels = ACOUSTIC_EMBEDDINGS_LEVELS + + # semantic_codebook_output: hidden_dim -> padded_codebook_size + padded_semantic_size = len(AudioSpecialTokens) + SEMANTIC_CODEBOOK_SIZE + # Round up to multiple of 128 like the real model + padded_semantic_size = 128 * ((padded_semantic_size + 127) // 128) + self.semantic_codebook_output = nn.Linear(HIDDEN_DIM, padded_semantic_size, bias=False) + + # time_embedding: (B, 1) -> (B, dim) + self.time_embedding = nn.Linear(1, HIDDEN_DIM, bias=False) + + # _predict_velocity: takes x_t (B, C), llm_output (B, D), t_emb (B, D) -> (B, C) + self._velocity_proj = nn.Linear(N_ACOUSTIC_CODEBOOK + HIDDEN_DIM + HIDDEN_DIM, N_ACOUSTIC_CODEBOOK, bias=False) + + def _predict_velocity(self, x_t: torch.Tensor, llm_output: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor: + combined = torch.cat([x_t, llm_output, t_emb], dim=-1) + return self._velocity_proj(combined) + + +class SyntheticModel(nn.Module): + """Mimics VoxtralTTSAudioGenerationForConditionalGeneration interface.""" + + def __init__(self): + super().__init__() + self.acoustic_transformer = SyntheticAcousticTransformer() + end_audio_id = AudioSpecialTokens.id(AudioSpecialTokens.end_audio) + empty_audio_id = AudioSpecialTokens.id(AudioSpecialTokens.empty_audio) + self.end_audio_id = end_audio_id + self.empty_audio_id = empty_audio_id + + def compute_mm_logits( + self, + hidden_states: torch.Tensor, + mm_sampling_tensors=None, + ): + """Eager fallback path: replicate what the wrapper does.""" + at = self.acoustic_transformer + B = hidden_states.shape[0] + + empty_audio_id = self.empty_audio_id + end_audio_id = self.end_audio_id + semantic_mask_start = len(AudioSpecialTokens) + at.model_args.semantic_codebook_size + + # Semantic logits + semantic_logit = at.semantic_codebook_output(hidden_states).float() + semantic_logit[:, empty_audio_id] = -float("inf") + semantic_logit[:, semantic_mask_start:] = -float("inf") + semantic_code = semantic_logit.argmax(dim=-1, keepdim=True) + + # Flow matching Euler ODE + should_decode = semantic_code.squeeze(1) != end_audio_id + x = torch.randn(B, at.model_args.n_acoustic_codebook, device=hidden_states.device, dtype=hidden_states.dtype) + hidden_zero = torch.zeros_like(hidden_states) + timesteps = torch.linspace(0, 1, 16, device=hidden_states.device, dtype=hidden_states.dtype) + + for i in range(len(timesteps) - 1): + t = timesteps[i] + dt = timesteps[i + 1] - timesteps[i] + t_emb = at.time_embedding(t.view(-1, 1).repeat(B, 1)).to(hidden_states.dtype) + x_batched = torch.cat([x, x], dim=0) + llm_batched = torch.cat([hidden_states, hidden_zero], dim=0) + t_emb_batched = t_emb.repeat(2, 1) + v_all = at._predict_velocity(x_t=x_batched, llm_output=llm_batched, t_emb=t_emb_batched) + v_t, uncond_v_t = v_all[:B], v_all[B:] + cfg_alpha = 1.2 + v_t = cfg_alpha * v_t + (1 - cfg_alpha) * uncond_v_t + x = x + v_t * dt + + sampled = torch.clamp(x, -1, 1) + scaled_x = ((sampled + 1) / 2) * (at.acoustic_embeddings_levels - 1) + output_codes = scaled_x.round().long() + output_codes[~should_decode] = empty_audio_id + acoustic_codes = output_codes + len(AudioSpecialTokens) + + audio_codes = torch.cat([semantic_code, acoustic_codes], dim=1) + fake_eos = torch.where( + audio_codes[:, 0] == end_audio_id, + torch.tensor(1.0, dtype=hidden_states.dtype, device=hidden_states.device), + torch.tensor(0.0, dtype=hidden_states.dtype, device=hidden_states.device), + ) + + audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0)) + mm_tokens = {"audio": audio_list} + return fake_eos, mm_tokens + + +@pytest.fixture(scope="module") +def model(): + """Create a synthetic model on CUDA with fixed weights in bfloat16.""" + torch.manual_seed(42) + return SyntheticModel().to(device=DEVICE, dtype=torch.bfloat16).eval() + + +@pytest.fixture(scope="module") +def wrapper(model): + """Create a warmed-up CUDAGraphAcousticTransformerWrapper.""" + w = CUDAGraphAcousticTransformerWrapper( + model=model, + capture_sizes=[1, 2, 4, 8, 16, 32], + ) + w._warmup_and_capture( + device=DEVICE, + dtype=torch.bfloat16, + hidden_dim=HIDDEN_DIM, + ) + return w + + +def _random_hidden(batch_size, device=DEVICE, dtype=torch.bfloat16): + return torch.randn(batch_size, HIDDEN_DIM, device=device, dtype=dtype) + + +def _unpack_audio_codes(result): + """Unpack (fake_eos, {"audio": [list of tensors]}) into (fake_eos, audio_codes).""" + fake_eos, mm_tokens = result + audio_list = mm_tokens["audio"] + # Each element is (1, 1, 1+n_acoustic_codebook), concat along batch dim + audio_codes = torch.cat(audio_list, dim=0).squeeze(1) # (B, 1+n_acoustic) + return fake_eos, audio_codes + + +# ────────────────────────────────────────────────────────────────── +# 1. Exact-size inputs (no padding) → correct format and bounds +# ────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16, 32]) +def test_exact_size_output_format(model, wrapper, batch_size): + """Graph path returns correctly shaped and bounded outputs.""" + hidden = _random_hidden(batch_size) + with torch.no_grad(): + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + assert graph_eos.shape == (batch_size,) + assert graph_codes.shape == (batch_size, 1 + N_ACOUSTIC_CODEBOOK) + # fake_eos should be 0.0 or 1.0 + assert torch.all((graph_eos == 0.0) | (graph_eos == 1.0)) + # Audio codes should be non-negative + assert torch.all(graph_codes >= 0) + + +@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16, 32]) +def test_exact_size_deterministic(model, wrapper, batch_size): + """Same input + same RNG state produces identical CUDA graph output.""" + hidden = _random_hidden(batch_size) + with torch.no_grad(): + torch.manual_seed(42) + eos1, codes1 = _unpack_audio_codes(wrapper(hidden)) + torch.manual_seed(42) + eos2, codes2 = _unpack_audio_codes(wrapper(hidden)) + torch.testing.assert_close(eos1, eos2, atol=0, rtol=0) + torch.testing.assert_close(codes1, codes2, atol=0, rtol=0) + + +# ────────────────────────────────────────────────────────────────── +# 2. Padded inputs (zero-padding to nearest capture size) +# ────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("batch_size", [3, 5, 7, 10, 24]) +def test_padded_output_shape(model, wrapper, batch_size): + """Padded decode must return output trimmed to actual batch size.""" + hidden = _random_hidden(batch_size) + with torch.no_grad(): + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + assert graph_eos.shape == (batch_size,) + assert graph_codes.shape == (batch_size, 1 + N_ACOUSTIC_CODEBOOK) + + +@pytest.mark.parametrize("batch_size", [3, 5, 7, 10, 24]) +def test_padded_output_bounded(model, wrapper, batch_size): + """Padded output audio codes should be non-negative integers.""" + hidden = _random_hidden(batch_size) + with torch.no_grad(): + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + # fake_eos should be 0.0 or 1.0 + assert torch.all((graph_eos == 0.0) | (graph_eos == 1.0)) + # Audio codes should be non-negative + assert torch.all(graph_codes >= 0) + + +# ────────────────────────────────────────────────────────────────── +# 3. Batch size exceeds all capture sizes so fallback to eager +# ────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("batch_size", [33, 48, 64]) +def test_fallback_eager_exact_match(model, wrapper, batch_size): + """Cudagraph fallback to eager. Two eager runs should produce identical results.""" + hidden = _random_hidden(batch_size) + with torch.no_grad(): + torch.manual_seed(100) + eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden)) + torch.manual_seed(100) + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + torch.testing.assert_close(graph_eos, eager_eos, atol=0, rtol=0) + torch.testing.assert_close(graph_codes, eager_codes, atol=0, rtol=0) + + +# ────────────────────────────────────────────────────────────────── +# 4. Disabled wrapper so fallback to eager +# ────────────────────────────────────────────────────────────────── + + +def test_disabled_wrapper_matches_eager(model, wrapper): + """Cudagraph fallback to eager. Two eager runs should produce identical results.""" + hidden = _random_hidden(4) + wrapper.enabled = False + with torch.no_grad(): + torch.manual_seed(200) + eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden)) + torch.manual_seed(200) + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + wrapper.enabled = True + torch.testing.assert_close(graph_eos, eager_eos, atol=0, rtol=0) + torch.testing.assert_close(graph_codes, eager_codes, atol=0, rtol=0) + + +# ────────────────────────────────────────────────────────────────── +# 5. Determinism - same input + same RNG state on cudagraph runs +# ────────────────────────────────────────────────────────────────── + + +def test_deterministic_across_calls(model, wrapper): + """Same input + same RNG state. Two cudagraph runs should produce identical results.""" + hidden = _random_hidden(4) + with torch.no_grad(): + torch.manual_seed(300) + eos1, codes1 = _unpack_audio_codes(wrapper(hidden)) + torch.manual_seed(300) + eos2, codes2 = _unpack_audio_codes(wrapper(hidden)) + torch.testing.assert_close(eos1, eos2, atol=0, rtol=0) + torch.testing.assert_close(codes1, codes2, atol=0, rtol=0) diff --git a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py new file mode 100644 index 000000000..45e54eb69 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections import defaultdict +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.voxtral_tts import ( + generator2tokenizer, + generator2tokenizer_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# ---- Helpers for generator2tokenizer (non-streaming) ---- + + +def _make_stage(engine_outputs): + return SimpleNamespace(engine_outputs=engine_outputs) + + +def _make_output(audio_tensors): + """Build a single generator output with multimodal_output["audio"] = audio_tensors.""" + return SimpleNamespace( + outputs=[SimpleNamespace(multimodal_output={"audio": audio_tensors})], + ) + + +# ---- Tests for generator2tokenizer ---- + + +class TestGenerator2Tokenizer: + def test_single_output_single_chunk(self): + """Single output with one audio tensor produces correct prompt_token_ids.""" + audio = torch.tensor([10, 20, 30, 40]) + stage = _make_stage([_make_output([audio])]) + + result = generator2tokenizer([stage], engine_input_source=[0]) + + assert len(result) == 1 + assert result[0]["prompt_token_ids"] == [10, 20, 30, 40] + assert result[0]["multi_modal_data"] is None + + def test_2d_audio_tensors_flattened(self): + """2D audio tensors (e.g., multi-codebook frames) are flattened.""" + audio = torch.tensor([[1, 2, 3], [4, 5, 6]]) # (2, 3) + stage = _make_stage([_make_output([audio])]) + + result = generator2tokenizer([stage], engine_input_source=[0]) + + assert result[0]["prompt_token_ids"] == [1, 2, 3, 4, 5, 6] + + def test_empty_engine_input_source_raises(self): + """Empty engine_input_source raises ValueError.""" + stage = _make_stage([_make_output([torch.tensor([1])])]) + with pytest.raises(ValueError, match="cannot be empty"): + generator2tokenizer([stage], engine_input_source=[]) + + def test_invalid_stage_id_raises(self): + """Stage id beyond stage_list length raises IndexError.""" + stage = _make_stage([_make_output([torch.tensor([1])])]) + with pytest.raises(IndexError, match="Invalid stage_id"): + generator2tokenizer([stage], engine_input_source=[5]) + + def test_no_outputs_yet_raises(self): + """Stage with engine_outputs=None raises RuntimeError.""" + stage = _make_stage(engine_outputs=None) + with pytest.raises(RuntimeError, match="no outputs yet"): + generator2tokenizer([stage], engine_input_source=[0]) + + +# ---- Helpers for generator2tokenizer_async_chunk (streaming) ---- + + +def _req(external_req_id: str, *, finished: bool): + return SimpleNamespace( + external_req_id=external_req_id, + is_finished=lambda: finished, + ) + + +def _make_transfer_manager(codec_chunk_frames=25, codec_left_context_frames=25, codec_chunk_frames_at_begin=5): + return SimpleNamespace( + code_prompt_token_ids=defaultdict(list), + connector=SimpleNamespace( + config={ + "extra": { + "codec_chunk_frames": codec_chunk_frames, + "codec_left_context_frames": codec_left_context_frames, + "codec_chunk_frames_at_begin": codec_chunk_frames_at_begin, + } + } + ), + ) + + +def test_empty_chunk_when_not_finished(): + """Returns None when no audio data and not finished.""" + transfer_manager = _make_transfer_manager() + request = _req("rid-empty", finished=False) + + payload = generator2tokenizer_async_chunk( + transfer_manager=transfer_manager, + pooling_output={"audio": torch.zeros((0,))}, + request=request, + ) + + assert payload is None + + +def test_flush_tail_when_finished(): + """Emits remaining frames on finish with finished=True.""" + transfer_manager = _make_transfer_manager() + request_id = "rid-tail" + # Pre-populate with 24 frames of 8-element codes + transfer_manager.code_prompt_token_ids[request_id] = [[1, 2, 3, 4, 5, 6, 7, 8] for _ in range(24)] + request = _req(request_id, finished=True) + + payload = generator2tokenizer_async_chunk( + transfer_manager=transfer_manager, + pooling_output=None, # e.g. EOS step with no audio + request=request, + ) + + assert payload is not None + assert payload["finished"].item() is True + codes = payload["code_predictor_codes"] + # Format: [ctx_frames, context_length, ...flat_codes] + assert len(codes) >= 2 # At least ctx_frames + context_length header + ctx_frames = codes[0] + context_length = codes[1] + flat_codes = codes[2:] + # Total frames in window = ctx_frames + context_length + total_window_frames = ctx_frames + context_length + assert len(flat_codes) == total_window_frames * 8 + + +def test_eof_marker_when_finished_with_no_frames(): + """Emits EOF marker with empty codes when finished with no accumulated frames.""" + transfer_manager = _make_transfer_manager() + request = _req("rid-eof", finished=True) + + payload = generator2tokenizer_async_chunk( + transfer_manager=transfer_manager, + pooling_output=None, + request=request, + ) + + assert payload == { + "code_predictor_codes": [], + "finished": torch.tensor(True, dtype=torch.bool), + } + + +def test_normal_chunk_emission(): + """Emits chunks at codec_chunk_frames boundary.""" + transfer_manager = _make_transfer_manager( + codec_chunk_frames=25, + codec_left_context_frames=25, + codec_chunk_frames_at_begin=5, + ) + request_id = "rid-chunk" + request = _req(request_id, finished=False) + + # Feed 25 frames one by one + for i in range(25): + pooling_output = {"audio": torch.tensor([float(i)] * 4)} + payload = generator2tokenizer_async_chunk( + transfer_manager=transfer_manager, + pooling_output=pooling_output, + request=request, + ) + + # A chunk should be emitted + assert payload is not None + codes = payload["code_predictor_codes"] + ctx_frames = codes[0] + context_length = codes[1] + assert ctx_frames == 20 # 25 - 5(chunk_size_at_begin) + assert context_length == 5 # 25 - 20(chunk_size_at_begin * 4) + + +def test_small_initial_chunks(): + """Uses codec_chunk_frames_at_begin for initial frames when length <= chunk_size.""" + transfer_manager = _make_transfer_manager( + codec_chunk_frames=25, + codec_left_context_frames=25, + codec_chunk_frames_at_begin=5, + ) + request_id = "rid-begin" + request = _req(request_id, finished=False) + + # Feed 5 frames (should trigger emission because codec_chunk_frames_at_begin=5) + for i in range(5): + pooling_output = {"audio": torch.tensor([float(i + 1)] * 3)} + payload = generator2tokenizer_async_chunk( + transfer_manager=transfer_manager, + pooling_output=pooling_output, + request=request, + ) + + assert payload is not None + codes = payload["code_predictor_codes"] + ctx_frames = codes[0] + context_length = codes[1] + assert ctx_frames == 0 + assert context_length == 5 # chunk_size_at_begin + + +def test_no_emission_between_boundaries(): + """No chunk emitted when not at chunk boundary and not finished.""" + transfer_manager = _make_transfer_manager( + codec_chunk_frames=25, + codec_left_context_frames=25, + codec_chunk_frames_at_begin=5, + ) + request_id = "rid-mid" + request = _req(request_id, finished=False) + + # Feed 3 frames (3 % 5 != 0, should not emit) + for i in range(3): + pooling_output = {"audio": torch.tensor([float(i)] * 4)} + payload = generator2tokenizer_async_chunk( + transfer_manager=transfer_manager, + pooling_output=pooling_output, + request=request, + ) + + assert payload is None + + +def test_context_handling_format(): + """Verifies [ctx_frames, context_length, ...codes] format.""" + transfer_manager = _make_transfer_manager( + codec_chunk_frames=10, + codec_left_context_frames=5, + codec_chunk_frames_at_begin=5, + ) + request_id = "rid-ctx" + request = _req(request_id, finished=False) + + # Feed 5 frames to trigger a chunk (chunk_size_at_begin=5) + for i in range(5): + pooling_output = {"audio": torch.tensor([float(i + 10)] * 2)} # codebook_dim=2 + payload = generator2tokenizer_async_chunk( + transfer_manager=transfer_manager, + pooling_output=pooling_output, + request=request, + ) + + assert payload is not None + codes = payload["code_predictor_codes"] + # First two elements are ctx_frames and context_length + ctx_frames = codes[0] + context_length = codes[1] + flat_codes = codes[2:] + # The window has ctx_frames + context_length frames + assert isinstance(ctx_frames, int) + assert isinstance(context_length, int) + assert ctx_frames >= 0 + assert context_length > 0 + # flat_codes = total_window_frames * codebook_dim + total_window_frames = ctx_frames + context_length + assert len(flat_codes) == total_window_frames * 2 # codebook_dim=2 + + +def test_none_pooling_output_not_finished_returns_none(): + """None pooling_output when not finished returns None.""" + transfer_manager = _make_transfer_manager() + request = _req("rid-none", finished=False) + + payload = generator2tokenizer_async_chunk( + transfer_manager=transfer_manager, + pooling_output=None, + request=request, + ) + + assert payload is None diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index f4af0f49b..d12bd10cc 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -222,6 +222,7 @@ class StageConfigFactory: "qwen2_5_omni": "qwen2_5_omni", "bagel": "bagel", "qwen3_tts": "qwen3_tts", + "voxtral_tts": "voxtral_tts", "mimo_audio": "mimo_audio", "glm-image": "glm_image", "cosyvoice3": "cosyvoice3", diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux.py b/vllm_omni/diffusion/models/flux/pipeline_flux.py index 06af7c95d..b7aa0c7b4 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux.py +++ b/vllm_omni/diffusion/models/flux/pipeline_flux.py @@ -18,7 +18,7 @@ from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( ) from diffusers.utils.torch_utils import randn_tensor from torch import nn -from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast +from transformers import AutoConfig, CLIPTextModel, CLIPTokenizer, T5TokenizerFast from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig @@ -27,6 +27,7 @@ from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_g from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.flux import FluxTransformer2DModel +from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -148,7 +149,14 @@ class FluxPipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin): revision=None, prefix="transformer.", fall_back_to_pt=True, - ) + ), + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="text_encoder_2", + revision=None, + prefix="text_encoder_2.", + fall_back_to_pt=True, + ), ] self.device = get_local_device() @@ -162,9 +170,8 @@ class FluxPipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin): self.text_encoder = CLIPTextModel.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only ).to(self.device) - self.text_encoder_2 = T5EncoderModel.from_pretrained( - model, subfolder="text_encoder_2", local_files_only=local_files_only - ).to(self.device) + t5_config = AutoConfig.from_pretrained(model, subfolder="text_encoder_2", local_files_only=local_files_only) + self.text_encoder_2 = T5EncoderModel(t5_config, prefix="text_encoder_2").to(self.device) self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.device ) @@ -291,7 +298,7 @@ class FluxPipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin): f" {max_sequence_length} tokens: {removed_text}" ) - prompt_embeds = self.text_encoder_2(text_input_ids.to(self.device), output_hidden_states=False)[0] + prompt_embeds = self.text_encoder_2(text_input_ids.to(self.device))[0] dtype = self.text_encoder_2.dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device) diff --git a/vllm_omni/diffusion/models/t5_encoder/__init__.py b/vllm_omni/diffusion/models/t5_encoder/__init__.py new file mode 100644 index 000000000..f2ca52efa --- /dev/null +++ b/vllm_omni/diffusion/models/t5_encoder/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tensor-parallel T5 encoder model.""" + +from vllm_omni.diffusion.models.t5_encoder.t5_encoder import T5EncoderModel + +__all__ = ["T5EncoderModel"] diff --git a/vllm_omni/diffusion/models/t5_encoder/t5_encoder.py b/vllm_omni/diffusion/models/t5_encoder/t5_encoder.py new file mode 100644 index 000000000..7b0e842d0 --- /dev/null +++ b/vllm_omni/diffusion/models/t5_encoder/t5_encoder.py @@ -0,0 +1,389 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +import math +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import T5Config +from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size +from vllm.model_executor.layers.activation import get_act_fn +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + + +class T5SelfAttention(nn.Module): + def __init__( + self, + config: T5Config, + has_relative_attention_bias: bool = False, + prefix: str = "", + ): + super().__init__() + self.d_model = config.d_model + self.d_kv = config.d_kv + self.n_heads = config.num_heads + self.inner_dim = self.n_heads * self.d_kv + self.has_relative_attention_bias = has_relative_attention_bias + self.relative_attention_num_buckets = config.relative_attention_num_buckets + self.relative_attention_max_distance = config.relative_attention_max_distance + + tp_size = get_tensor_model_parallel_world_size() + assert self.n_heads % tp_size == 0, f"num_heads ({self.n_heads}) must be divisible by tp_size ({tp_size})" + self.n_heads_per_partition = self.n_heads // tp_size + + # Fused Q/K/V projection, sharded across heads. + # HF has separate q, k, v — we fuse into qkv_proj and use + # stacked_params_mapping in load_weights() to load each shard. + self.qkv_proj = QKVParallelLinear( + hidden_size=self.d_model, + head_size=self.d_kv, + total_num_heads=self.n_heads, + total_num_kv_heads=self.n_heads, # T5 uses MHA + bias=False, + ) + + # Output projection: all-reduce back to full d_model + # Named ``o`` to match HF's ``SelfAttention.o`` + self.o = RowParallelLinear( + self.inner_dim, + self.d_model, + bias=False, + input_is_parallel=True, + return_bias=False, + ) + + if has_relative_attention_bias: + # Store full embedding; slice heads per rank in forward + self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads) + + @staticmethod + def _relative_position_bucket( + relative_position: torch.Tensor, + bidirectional: bool = True, + num_buckets: int = 32, + max_distance: int = 128, + ) -> torch.Tensor: + relative_buckets = 0 + if bidirectional: + num_buckets //= 2 + relative_buckets += (relative_position > 0).to(torch.long) * num_buckets + relative_position = torch.abs(relative_position) + else: + relative_position = -torch.min(relative_position, torch.zeros_like(relative_position)) + max_exact = num_buckets // 2 + is_small = relative_position < max_exact + relative_position_if_large = max_exact + ( + torch.log(relative_position.float() / max_exact) + / math.log(max_distance / max_exact) + * (num_buckets - max_exact) + ).to(torch.long) + relative_position_if_large = torch.min( + relative_position_if_large, + torch.full_like(relative_position_if_large, num_buckets - 1), + ) + relative_buckets += torch.where(is_small, relative_position, relative_position_if_large) + return relative_buckets + + def compute_bias(self, query_length: int, key_length: int, device: torch.device) -> torch.Tensor: + """Compute relative position bias, returning only the local head shard.""" + context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] + memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] + relative_position = memory_position - context_position + relative_position_bucket = self._relative_position_bucket( + relative_position, + bidirectional=True, + num_buckets=self.relative_attention_num_buckets, + max_distance=self.relative_attention_max_distance, + ) + # values: (query_length, key_length, num_heads) + values = self.relative_attention_bias(relative_position_bucket) + # Slice to local heads for this TP rank + tp_rank = get_tensor_model_parallel_rank() + head_start = tp_rank * self.n_heads_per_partition + head_end = head_start + self.n_heads_per_partition + values = values[:, :, head_start:head_end] + # (1, local_heads, query_length, key_length) + values = values.permute(2, 0, 1).unsqueeze(0) + return values + + def forward( + self, + hidden_states: torch.Tensor, + mask: torch.Tensor | None = None, + position_bias: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, seq_length = hidden_states.shape[:2] + + # Fused QKV projection + qkv, _ = self.qkv_proj(hidden_states) + q_size = self.n_heads_per_partition * self.d_kv + kv_size = self.n_heads_per_partition * self.d_kv + query_states, key_states, value_states = qkv.split([q_size, kv_size, kv_size], dim=-1) + + # Reshape: (batch, seq, local_heads, d_kv) -> (batch, local_heads, seq, d_kv) + query_states = query_states.view(batch_size, seq_length, self.n_heads_per_partition, self.d_kv).transpose(1, 2) + key_states = key_states.view(batch_size, seq_length, self.n_heads_per_partition, self.d_kv).transpose(1, 2) + value_states = value_states.view(batch_size, seq_length, self.n_heads_per_partition, self.d_kv).transpose(1, 2) + + # Attention scores: (batch, local_heads, seq, seq) + scores = torch.matmul(query_states, key_states.transpose(3, 2)) + + if position_bias is None: + if self.has_relative_attention_bias: + position_bias = self.compute_bias(seq_length, seq_length, device=scores.device) + else: + position_bias = torch.zeros( + (1, self.n_heads_per_partition, seq_length, seq_length), + device=scores.device, + dtype=scores.dtype, + ) + if mask is not None: + position_bias = position_bias + mask + + scores += position_bias + + attn_weights = F.softmax(scores.float(), dim=-1).type_as(scores) + attn_output = torch.matmul(attn_weights, value_states) + + # (batch, local_heads, seq, d_kv) -> (batch, seq, local_heads * d_kv) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(batch_size, seq_length, -1) + + attn_output = self.o(attn_output) + + return attn_output, position_bias + + +class T5DenseGatedActDense(nn.Module): + def __init__(self, config: T5Config, prefix: str = ""): + super().__init__() + self.wi = MergedColumnParallelLinear( + config.d_model, + [config.d_ff, config.d_ff], + bias=False, + gather_output=False, + ) + self.wo = RowParallelLinear( + config.d_ff, + config.d_model, + bias=False, + input_is_parallel=True, + return_bias=False, + ) + self.act = get_act_fn(config.dense_act_fn) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # MergedColumnParallelLinear outputs concatenated [gate, up] + gate_up, _ = self.wi(hidden_states) + gate, up = gate_up.chunk(2, dim=-1) + hidden_states = self.act(gate) * up + hidden_states = self.wo(hidden_states) + return hidden_states + + +class T5DenseActDense(nn.Module): + def __init__(self, config: T5Config, prefix: str = ""): + super().__init__() + self.wi = ColumnParallelLinear( + config.d_model, + config.d_ff, + bias=False, + gather_output=False, + return_bias=False, + ) + self.wo = RowParallelLinear( + config.d_ff, + config.d_model, + bias=False, + input_is_parallel=True, + return_bias=False, + ) + self.act = get_act_fn(config.dense_act_fn) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.wi(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.wo(hidden_states) + return hidden_states + + +class T5LayerSelfAttention(nn.Module): + def __init__(self, config: T5Config, has_relative_attention_bias: bool = False, prefix: str = ""): + super().__init__() + self.SelfAttention = T5SelfAttention(config, has_relative_attention_bias, prefix=f"{prefix}.SelfAttention") + self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon) + + def forward( + self, + hidden_states: torch.Tensor, + mask: torch.Tensor | None = None, + position_bias: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + normed = self.layer_norm(hidden_states) + attn_output, position_bias = self.SelfAttention(normed, mask=mask, position_bias=position_bias) + hidden_states = hidden_states + attn_output + + # Clamp for fp16 stability (matching HF T5) + if hidden_states.dtype == torch.float16: + clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + return hidden_states, position_bias + + +class T5LayerFF(nn.Module): + def __init__(self, config: T5Config, prefix: str = ""): + super().__init__() + if config.is_gated_act: + self.DenseReluDense = T5DenseGatedActDense(config, prefix=f"{prefix}.DenseReluDense") + else: + self.DenseReluDense = T5DenseActDense(config, prefix=f"{prefix}.DenseReluDense") + self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + normed = self.layer_norm(hidden_states) + ff_output = self.DenseReluDense(normed) + hidden_states = hidden_states + ff_output + + if hidden_states.dtype == torch.float16: + clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + return hidden_states + + +class T5Block(nn.Module): + def __init__(self, config: T5Config, has_relative_attention_bias: bool = False, prefix: str = ""): + super().__init__() + self.layer = nn.ModuleList( + [ + T5LayerSelfAttention(config, has_relative_attention_bias, prefix=f"{prefix}.layer.0"), + T5LayerFF(config, prefix=f"{prefix}.layer.1"), + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + mask: torch.Tensor | None = None, + position_bias: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + hidden_states, position_bias = self.layer[0](hidden_states, mask=mask, position_bias=position_bias) + hidden_states = self.layer[1](hidden_states) + return hidden_states, position_bias + + +class T5Stack(nn.Module): + def __init__(self, config: T5Config, shared: nn.Embedding, prefix: str = ""): + super().__init__() + self.embed_tokens = shared + self.block = nn.ModuleList( + [ + T5Block(config, has_relative_attention_bias=(i == 0), prefix=f"{prefix}.block.{i}") + for i in range(config.num_layers) + ] + ) + self.final_layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon) + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + hidden_states = self.embed_tokens(input_ids) + + # Build attention mask: (batch, seq) -> (batch, 1, 1, seq) + if attention_mask is not None: + extended_mask = attention_mask[:, None, None, :].to(dtype=hidden_states.dtype) + extended_mask = (1.0 - extended_mask) * torch.finfo(hidden_states.dtype).min + else: + extended_mask = None + + position_bias = None + for block in self.block: + hidden_states, position_bias = block( + hidden_states, + mask=extended_mask, + position_bias=position_bias, + ) + + hidden_states = self.final_layer_norm(hidden_states) + return hidden_states + + +class T5EncoderModel(nn.Module): + """T5 encoder model applying upstream vLLM layers""" + + def __init__(self, config: T5Config, prefix: str = ""): + super().__init__() + self.config = config + self.shared = VocabParallelEmbedding(config.vocab_size, config.d_model) + self.encoder = T5Stack(config, self.shared, prefix=f"{prefix}.encoder") + + @property + def dtype(self) -> torch.dtype: + return next(self.parameters()).dtype + + @property + def device(self) -> torch.device: + return next(self.parameters()).device + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.shared(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, ...]: + hidden_states = self.encoder(input_ids, attention_mask=attention_mask) + return (hidden_states,) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + ("qkv_proj", "q", "q"), + ("qkv_proj", "k", "k"), + ("qkv_proj", "v", "v"), + ("wi", "wi_0", 0), + ("wi", "wi_1", 1), + ] + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + original_name = name + lookup_name = name + + for param_name, weight_name, shard_id in stacked_params_mapping: + if f".{weight_name}." not in name: + continue + lookup_name = name.replace(weight_name, param_name) + if lookup_name not in params_dict: + continue + param = params_dict[lookup_name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + loaded_params.add(original_name) + loaded_params.add(lookup_name) + + return loaded_params diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index b68b6f70d..b7367e9b3 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -279,6 +279,7 @@ class ZImageAttention(nn.Module): total_num_kv_heads=num_kv_heads, bias=False, quant_config=quant_config, + prefix="to_qkv", ) assert qk_norm is True @@ -297,6 +298,7 @@ class ZImageAttention(nn.Module): input_is_parallel=True, return_bias=False, quant_config=quant_config, + prefix="to_out", ) ] ) @@ -361,6 +363,7 @@ class FeedForward(nn.Module): dim: int, hidden_dim: int, quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() self.w13 = MergedColumnParallelLinear( @@ -369,6 +372,7 @@ class FeedForward(nn.Module): bias=False, return_bias=False, quant_config=quant_config, + prefix=prefix, ) self.act = SiluAndMul() self.w2 = RowParallelLinear( @@ -378,6 +382,7 @@ class FeedForward(nn.Module): input_is_parallel=True, return_bias=False, quant_config=quant_config, + prefix=prefix, ) def forward(self, x): @@ -409,9 +414,7 @@ class ZImageTransformerBlock(nn.Module): ) self.feed_forward = FeedForward( - dim=dim, - hidden_dim=int(dim / 3 * 8), - quant_config=quant_config, + dim=dim, hidden_dim=int(dim / 3 * 8), quant_config=quant_config, prefix="feed_forward" ) self.layer_id = layer_id diff --git a/vllm_omni/diffusion/quantization/__init__.py b/vllm_omni/diffusion/quantization/__init__.py index d297d51f1..eb4b8ea9f 100644 --- a/vllm_omni/diffusion/quantization/__init__.py +++ b/vllm_omni/diffusion/quantization/__init__.py @@ -29,6 +29,7 @@ from vllm.logger import init_logger from .base import DiffusionQuantizationConfig from .fp8 import DiffusionFp8Config from .gguf import DiffusionGgufConfig +from .int8 import DiffusionInt8Config if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( @@ -41,6 +42,7 @@ logger = init_logger(__name__) # To add a new method, create a new config class and register it here _QUANT_CONFIG_REGISTRY: dict[str, type[DiffusionQuantizationConfig]] = { "fp8": DiffusionFp8Config, + "int8": DiffusionInt8Config, "gguf": DiffusionGgufConfig, } @@ -110,6 +112,7 @@ def get_vllm_quant_config_for_layers( __all__ = [ "DiffusionQuantizationConfig", "DiffusionFp8Config", + "DiffusionInt8Config", "DiffusionGgufConfig", "get_diffusion_quant_config", "get_vllm_quant_config_for_layers", diff --git a/vllm_omni/diffusion/quantization/int8.py b/vllm_omni/diffusion/quantization/int8.py new file mode 100644 index 000000000..798038c9f --- /dev/null +++ b/vllm_omni/diffusion/quantization/int8.py @@ -0,0 +1,475 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""INT8 quantization config for diffusion transformers.""" + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, Optional + +import torch +from torch.nn import Module +from vllm import _custom_ops as ops +from vllm.logger import init_logger +from vllm.model_executor.kernels.linear import ( + init_int8_linear_kernel, +) +from vllm.model_executor.layers.linear import ( + LinearBase, + LinearMethodBase, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization import QuantizationMethods +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from vllm.model_executor.layers.quantization.fp8 import CopyNumelCounter, _copy_missing_attrs +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + is_layer_skipped, +) +from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight +from vllm.model_executor.parameter import ( + ChannelQuantScaleParameter, + ModelWeightParameter, +) +from vllm.model_executor.utils import replace_parameter + +from vllm_omni.platforms import current_omni_platform + +if current_omni_platform.is_npu(): + import torch_npu +else: + torch_npu = None + +from .base import DiffusionQuantizationConfig + +if TYPE_CHECKING: + from vllm.model_executor.models.utils import WeightsMapper + +# Dynamic quantization is supported first. +ACTIVATION_SCHEMES = ["dynamic"] + +logger = init_logger(__name__) + + +def create_weight_parameter( + output_size_per_partition: int, + input_size_per_partition: int, + weight_loader: Callable | None, + params_dtype: torch.dtype, +) -> torch.nn.Parameter: + """ + Create int8 weight parameter. + """ + return ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=params_dtype, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + + +class Int8Config(QuantizationConfig): + """ + Config class for Int8. + """ + + def __init__( + self, + is_checkpoint_int8_serialized: bool = False, + activation_scheme: str = "dynamic", + ignored_layers: list[str] | None = None, + ) -> None: + super().__init__() + + self.is_checkpoint_int8_serialized = is_checkpoint_int8_serialized + + if activation_scheme not in ACTIVATION_SCHEMES: + raise ValueError(f"Unsupported activation scheme {activation_scheme}") + self.activation_scheme = activation_scheme + self.ignored_layers = ignored_layers or [] + + @classmethod + def get_name(cls) -> QuantizationMethods: + return "int8" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.float16] + + @classmethod + def get_min_capability(cls) -> int: + # Have verified on A100 and H20, but not on oldest versions. + return 80 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return [] + + def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"): + if self.ignored_layers is not None: + self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers) + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "Int8Config": + quant_method = cls.get_from_keys(config, ["quant_method"]) + is_checkpoint_int8_serialized = "int8" in quant_method + activation_scheme = cls.get_from_keys_or(config, ["activation_scheme"], "dynamic") + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) + + if not ignored_layers: + ignored_layers = cls.get_from_keys_or(config, ["modules_to_not_convert"], None) + return cls( + is_checkpoint_int8_serialized=is_checkpoint_int8_serialized, + activation_scheme=activation_scheme, + ignored_layers=ignored_layers, + ) + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> Optional["QuantizeMethodBase"]: + if isinstance(layer, LinearBase): + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + if not self.is_checkpoint_int8_serialized: + if current_omni_platform.is_cuda(): + online_method = Int8OnlineLinearMethod(self) + elif current_omni_platform.is_npu(): + online_method = NPUInt8OnlineLinearMethod(self) + else: + raise NotImplementedError("The current platform is not supported int8 online quant.") + return online_method + else: + if current_omni_platform.is_cuda(): + offline_method = Int8LinearMethod(self) + elif current_omni_platform.is_npu(): + offline_method = NPUInt8LinearMethod(self) + else: + raise NotImplementedError("The current platform is not supported int8 offline quant.") + return offline_method + return None + + +class BaseInt8LinearMethod(LinearMethodBase): + """ + Linear method for Int8 + Supports loading Int8 checkpoints with static weight scale and dynamic activation scale. + + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: Int8Config): + self.quant_config = quant_config + self.out_dtype = torch.get_default_dtype() + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + + params_dtype = torch.int8 if self.quant_config.is_checkpoint_int8_serialized else params_dtype + weight = create_weight_parameter( + output_size_per_partition=output_size_per_partition, + input_size_per_partition=input_size_per_partition, + weight_loader=weight_loader, + params_dtype=params_dtype, + ) + layer.register_parameter("weight", weight) + + if self.quant_config.is_checkpoint_int8_serialized: + scale = ChannelQuantScaleParameter( + data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32), + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale", scale) + + def process_weights_after_loading(self, layer: Module) -> None: + raise NotImplementedError("No BaseInt8LinearMethod process_weights_after_loading implementation.") + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + raise NotImplementedError("No BaseInt8LinearMethod apply implementation.") + + +class LazyWeightMixin: + """ + Mixin for lazy weight loading with meta device. + weighs are created on meta device and materialized just-in-time during loadding. + """ + + uses_meta_device: bool = True + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + layer.weight_block_size = None + + # WEIGHT + def patched_weight_loader(param, loaded_weight, *args, **kwargs): + # track how many elements we have updated + if not hasattr(layer, "_loaded_numel"): + layer._loaded_numel = 0 + + # when the first `loaded_weight` is about to be + # loaded to `param`, materialize `param` just-in-time + weight = ModelWeightParameter( + data=torch.empty_like(layer.weight, device=layer._load_device), + input_dim=1, + output_dim=0, + weight_loader=patched_weight_loader, + ) + _copy_missing_attrs(layer.weight, weight) + layer.register_parameter("weight", weight) + del layer._load_device + + # refresh the reference to `param` to reflect just-in-time + # materialization + param = layer.weight + + # load the current weight chunk + copy_numel_counter = CopyNumelCounter() + with copy_numel_counter: + res = weight_loader(param, loaded_weight, *args, **kwargs) # type: ignore[misc] + layer._loaded_numel += copy_numel_counter.copied_numel + + # if we have loaded all of the elements, call + # process_weights_after_loading + target_loaded_numel = layer.weight.numel() + if layer._loaded_numel == target_loaded_numel: + self.process_weights_after_loading(layer) + + # Prevent the usual `process_weights_after_loading` call from doing + # anything + layer._already_called_process_weights_after_loading = True + + # Note that we keep `layer._loaded_numel` around just in case + # there is logic added to vllm in the future which calls a + # weight loader twice - we do not want to re-initialize in + # that case. + + return res + + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + # materialized just-in-time in `patched_weight_loader` + device="meta", + dtype=params_dtype, + ), + input_dim=1, + output_dim=0, + weight_loader=patched_weight_loader, + ) + # stash the correct device for `patched_weight_loader` + layer._load_device = torch.get_default_device() + layer.register_parameter("weight", weight) + + +class Int8LinearMethod(BaseInt8LinearMethod): + """ + Linear method for Int8 + Supports loading Int8 checkpoints. + + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: Int8Config): + super().__init__(quant_config) + + self.int8_linear = init_int8_linear_kernel( + is_channelwise=False, + is_static_input_scheme=False, + input_symmetric=True, + module_name=self.__class__.__name__, + ) + + def process_weights_after_loading(self, layer: Module) -> None: + self.int8_linear.process_weights_after_loading(layer) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.int8_linear.apply_weights(layer, x, bias) + + +class NPUInt8LinearMethod(BaseInt8LinearMethod): + """ + NPU Linear method for Int8 + Supports loading Int8 checkpoints. + + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: Int8Config): + super().__init__(quant_config) + + def process_weights_after_loading(self, layer: Module) -> None: + layer.weight.data = layer.weight.data.t().contiguous() + layer.weight_scale.data = layer.weight_scale.data.squeeze() + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + ori_shape = x.shape + ori_dtype = x.dtype + + x = x.reshape(-1, ori_shape[-1]) + quantized_x, pertoken_scale = torch_npu.npu_dynamic_quant(x) + + output = torch_npu.npu_quant_matmul( + quantized_x, + layer.weight, + layer.weight_scale, + bias=bias, + pertoken_scale=pertoken_scale, + output_dtype=ori_dtype, + ) + output = output.reshape(*ori_shape[:-1], -1) + return output + + +class Int8OnlineLinearMethod(LazyWeightMixin, Int8LinearMethod): + """ + Online version of Int8LinearMethod, loads the fp16/bf16 checkpoint + and quantized the weights during loading. + """ + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + if layer.weight.device == torch.device("meta"): + weight = ModelWeightParameter( + data=torch.empty_like(layer.weight, device=layer._load_device), + input_dim=1, + output_dim=0, + weight_loader=layer.weight.weight_loader, + ) + _copy_missing_attrs(layer.weight, weight) + layer.register_parameter("weight", weight) + initialize_single_dummy_weight(layer.weight) + + w_q_name, w_s_name, i_s_name, i_zp_name, azp_adj_name = self.int8_linear.layer_param_names + qweight, weight_scale, _ = ops.scaled_int8_quant(layer.weight, scale=None) + + # Update layer with new values. + replace_parameter(layer, w_q_name, torch.nn.Parameter(qweight.t().data, requires_grad=False)) + replace_parameter(layer, w_s_name, torch.nn.Parameter(weight_scale.data, requires_grad=False)) + + setattr(layer, i_s_name, None) + setattr(layer, i_zp_name, None) + setattr(layer, azp_adj_name, None) + + # Prevent duplicate processing (e.g., during weight reload) + layer._already_called_process_weights_after_loading = True + + +class NPUInt8OnlineLinearMethod(LazyWeightMixin, NPUInt8LinearMethod): + """ + NPU Online version of Int8LinearMethod, loads the fp16/bf16 checkpoint + and quantized the weights during loading. + """ + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + if layer.weight.device == torch.device("meta"): + weight = ModelWeightParameter( + data=torch.empty_like(layer.weight, device=layer._load_device), + input_dim=1, + output_dim=0, + weight_loader=layer.weight.weight_loader, + ) + _copy_missing_attrs(layer.weight, weight) + layer.register_parameter("weight", weight) + initialize_single_dummy_weight(layer.weight) + + weight = layer.weight + qweight, weight_scale = torch_npu.npu_dynamic_quant(weight) + + qweight = qweight.t().contiguous() + + # Update layer with new values. + replace_parameter(layer, "weight", qweight) + replace_parameter(layer, "weight_scale", weight_scale) + + # Prevent duplicate processing (e.g., during weight reload) + layer._already_called_process_weights_after_loading = True + + +class DiffusionInt8Config(DiffusionQuantizationConfig): + """ + Int8 quantization config optimized for diffusion transformers. + + Args: + activation_scheme: Activation quantization scheme. + - "dynamic": Per-token dynamic scaling (default, no calibration) + ignored_layers: List of layer name patterns to skip quantization. + """ + + quant_config_cls = Int8Config + + def __init__( + self, + activation_scheme: str = "dynamic", + ignored_layers: list[str] | None = None, + ): + self.activation_scheme = activation_scheme + self.ignored_layers = ignored_layers or [] + + # Create underlying vLLM Int8 config + self._vllm_config = Int8Config( + is_checkpoint_int8_serialized=False, # Online quantization + activation_scheme=activation_scheme, + ignored_layers=ignored_layers, + ) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index a1c9a05c8..fd1845c2c 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -22,16 +22,23 @@ def _register_omni_hf_configs() -> None: from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( Qwen3TTSConfig, ) + from vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts import ( + VoxtralTTSConfig, + ) except Exception as exc: # pragma: no cover - best-effort optional registration logger.warning("Skipping omni HF config registration due to import error: %s", exc) return - try: - AutoConfig.register("qwen3_tts", Qwen3TTSConfig) - AutoConfig.register("cosyvoice3", CosyVoice3Config) - except ValueError: - # Already registered elsewhere; ignore. - return + for model_type, config_cls in [ + ("qwen3_tts", Qwen3TTSConfig), + ("cosyvoice3", CosyVoice3Config), + ("mistral", VoxtralTTSConfig), + ]: + try: + AutoConfig.register(model_type, config_cls) + except ValueError: + # Already registered elsewhere; ignore. + pass def register_omni_models_to_vllm(): diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 50fb67597..56eb3b5cc 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -844,6 +844,7 @@ class AsyncOmniEngine: "enable_sleep_mode": kwargs.get("enable_sleep_mode", False), "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), + "quantization": kwargs.get("quantization", None), "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False), }, "final_output": True, diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index d58617e17..dc8c5a8ea 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -2178,11 +2178,17 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Convert images to base64 content image_contents: list[dict[str, Any]] = [] - for img in images: + flat_images = [] + for item in images: + if isinstance(item, list): + flat_images.extend(item) + else: + flat_images.append(item) + + for img in flat_images: with BytesIO() as buffer: img.save(buffer, format="PNG") - img_bytes = buffer.getvalue() - img_base64 = base64.b64encode(img_bytes).decode("utf-8") + img_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8") image_contents.append( { "type": "image_url", diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 6c2928e62..5ece2e750 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -30,7 +30,10 @@ from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) # TTS Configuration -_TTS_MODEL_STAGES: set[str] = {"qwen3_tts", "fish_speech_slow_ar"} +_VOXTRAL_TTS_MODEL_STAGES = {"audio_generation"} +_QWEN3_TTS_MODEL_STAGES = {"qwen3_tts"} +_FISH_TTS_MODEL_STAGES = {"fish_speech_slow_ar"} +_TTS_MODEL_STAGES: set[str] = _VOXTRAL_TTS_MODEL_STAGES | _QWEN3_TTS_MODEL_STAGES | _FISH_TTS_MODEL_STAGES _TTS_LANGUAGES: set[str] = { "Auto", "Chinese", @@ -145,6 +148,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ) self._fish_speech_tokenizer = None + # Determine TTS model type or None + self._tts_model_type = self._detect_tts_model_type() + # Cache TTS configuration values (computed once, reused per request) self._max_instructions_length = self._compute_max_instructions_length() @@ -203,6 +209,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return stage return None + def _detect_tts_model_type(self) -> str | None: + """Detect TTS model type from the stage's model_stage attribute.""" + if self._tts_stage is None: + return None + model_stage = getattr(self._tts_stage.engine_args, "model_stage", None) + if model_stage in _QWEN3_TTS_MODEL_STAGES: + return "qwen3_tts" + if model_stage in _VOXTRAL_TTS_MODEL_STAGES: + return "voxtral_tts" + if model_stage in _FISH_TTS_MODEL_STAGES: + return "fish_tts" + return None + def _compute_max_instructions_length(self) -> int: """Compute max instructions length with precedence: CLI > stage config > default. @@ -225,16 +244,22 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _load_supported_speakers(self) -> set[str]: """Load supported speakers (case-insensitive) from the model configuration.""" try: - talker_config = self.engine_client.model_config.hf_config.talker_config + if self._tts_model_type == "voxtral_tts": + config = self.engine_client.model_config.hf_config.audio_config + else: + # Default is qwen3_tts path + config = self.engine_client.model_config.hf_config.talker_config # Check for speakers in either spk_id or speaker_id for attr_name in ["spk_id", "speaker_id"]: - speakers_dict = getattr(talker_config, attr_name, None) + if isinstance(config, dict): + speakers_dict = config.get(attr_name) + else: + speakers_dict = getattr(config, attr_name, None) if speakers_dict and isinstance(speakers_dict, dict): - # Normalize to lowercase for case-insensitive matching return {speaker.lower() for speaker in speakers_dict.keys()} - logger.warning("No speakers found in talker_config (checked spk_id and speaker_id)") + logger.warning("No speakers found in config (checked spk_id and speaker_id)") except Exception as e: logger.warning(f"Could not load speakers from model config: {e}") @@ -497,6 +522,49 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _validate_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate TTS request parameters. Returns error message or None.""" + if self._tts_model_type == "voxtral_tts": + return self._validate_voxtral_tts_request(request) + return self._validate_qwen_tts_request(request) + + def _validate_ref_audio_format(self, ref_audio: str) -> str | None: + """Validate ref_audio is a supported URI format. Returns error or None.""" + if not ( + ref_audio.startswith(("http://", "https://")) + or ref_audio.startswith("data:") + or ref_audio.startswith("file://") + ): + return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)" + return None + + def _validate_voxtral_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: + """Validate Voxtral TTS request parameters. Returns error message or None.""" + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + + # Voxtral TTS requires either a preset voice or ref_audio for voice cloning. + if request.voice is None and request.ref_audio is None: + return "Either 'voice' (preset speaker) or 'ref_audio' (voice cloning) must be provided" + + if request.ref_audio is not None: + fmt_err = self._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + + if request.voice is not None: + request.voice = request.voice.lower() + if self.supported_speakers and request.voice not in self.supported_speakers: + return f"Invalid speaker '{request.voice}'. Supported: {', '.join(sorted(self.supported_speakers))}" + + if request.max_new_tokens is not None: + if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN: + return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}" + if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX: + return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}" + + return None + + def _validate_qwen_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: + """Validate Qwen TTS request parameters. Returns error message or None.""" # Infer Base task when ref_audio or ref_text is provided without explicit task_type. if request.task_type is None and (request.ref_audio is not None or request.ref_text is not None): request.task_type = "Base" @@ -530,13 +598,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if request.voice is None: if request.ref_audio is None: return "Base task requires 'ref_audio' for voice cloning" - # Validate ref_audio format (include file:// from upstream) - if not ( - request.ref_audio.startswith(("http://", "https://")) - or request.ref_audio.startswith("data:") - or request.ref_audio.startswith("file://") - ): - return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)" + fmt_err = self._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err # In-context voice cloning (default) requires non-empty ref_text. # x_vector_only_mode skips in-context and only uses speaker embedding. if not request.x_vector_only_mode: @@ -560,15 +624,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return ( f"Base task with built-in speaker '{request.voice}' requires 'ref_audio' for voice cloning" ) - # Validate ref_audio format for built-in speaker - if not ( - request.ref_audio.startswith(("http://", "https://")) - or request.ref_audio.startswith("data:") - or request.ref_audio.startswith("file://") - ): - return ( - "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)" - ) + fmt_err = self._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err # Validate cross-parameter dependencies if task_type != "Base": @@ -775,6 +833,40 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return params + # ---- Voxtral TTS helpers ---- + + async def _build_voxtral_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: + """Build Voxtral TTS engine prompt from shared TTS parameters.""" + from mistral_common.protocol.speech.request import SpeechRequest + + text = request.input + voice = request.voice + ref_audio = request.ref_audio + assert voice or ref_audio, "Either voice or ref_audio must be provided" + # Strip data URI prefix — mistral_common expects raw base64 + if ref_audio is not None and isinstance(ref_audio, str) and ref_audio.startswith("data:"): + _, _, ref_audio = ref_audio.partition(",") + if self._tts_tokenizer is None: + from vllm.tokenizers import cached_tokenizer_from_config + + mistral_tokenizer = cached_tokenizer_from_config(self.engine_client.model_config) + self._tts_tokenizer = mistral_tokenizer.instruct + if voice is not None: + tokens = self._tts_tokenizer.encode_speech_request(SpeechRequest(input=text, voice=voice)).tokens + return { + "prompt_token_ids": tokens, + "additional_information": {"voice": [voice]}, + } + else: + tokenized = self._tts_tokenizer.encode_speech_request(SpeechRequest(input=text, ref_audio=ref_audio)) + audio = tokenized.audios[0] + return { + "prompt_token_ids": tokenized.tokens, + "multi_modal_data": {"audio": [(audio.audio_array, audio.sampling_rate)]}, + } + + # ---- Fish Speech helpers ---- + def _build_fish_speech_prompt( self, request: OpenAICreateSpeechRequest, @@ -848,6 +940,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): "additional_information": additional_information, } + # ---- Common speech generation helpers ---- + async def _prepare_speech_generation( self, request: OpenAICreateSpeechRequest, @@ -871,25 +965,30 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if validation_error: raise ValueError(validation_error) - tts_params = self._build_tts_params(request) - if request.ref_audio is not None: - wav_list, sr = await self._resolve_ref_audio(request.ref_audio) - tts_params["ref_audio"] = [[wav_list, sr]] + if self._tts_model_type == "voxtral_tts": + prompt = await self._build_voxtral_prompt(request) + tts_params = {} + else: + tts_params = self._build_tts_params(request) + if request.ref_audio is not None: + wav_list, sr = await self._resolve_ref_audio(request.ref_audio) + tts_params["ref_audio"] = [[wav_list, sr]] - ph_len = self._estimate_prompt_len(tts_params) - prompt = {"prompt_token_ids": [1] * ph_len, "additional_information": tts_params} + ph_len = self._estimate_prompt_len(tts_params) + prompt = {"prompt_token_ids": [1] * ph_len, "additional_information": tts_params} else: tts_params = {} prompt = {"prompt": request.input} request_id = f"speech-{random_uuid()}" - model_type = ( - "fish_speech" - if self._is_fish_speech - else tts_params.get("task_type", ["unknown"])[0] - if self._is_tts - else "generic" - ) + if self._is_fish_speech: + model_type = "fish_speech" + elif self._tts_model_type == "voxtral_tts": + model_type = "voxtral_tts" + elif self._is_tts: + model_type = tts_params.get("task_type", ["unknown"])[0] + else: + model_type = "generic" logger.info( "TTS speech request %s: text=%r, model=%s", request_id, diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index f708b6ff9..b5ef92616 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -140,6 +140,18 @@ _OMNI_MODELS = { "fish_speech_dac_decoder", "FishSpeechDACDecoder", ), + ## Voxtral TTS + "VoxtralTTSForConditionalGeneration": ( + "voxtral_tts", + "voxtral_tts", + "VoxtralTTSForConditionalGeneration", + ), + "VoxtralTTSAudioGeneration": ( + "voxtral_tts", + "voxtral_tts_audio_generation", + "VoxtralTTSAudioGenerationForConditionalGeneration", + ), + "VoxtralTTSAudioTokenizer": ("voxtral_tts", "voxtral_tts_audio_tokenizer", "VoxtralTTSAudioTokenizer"), } diff --git a/vllm_omni/model_executor/models/voxtral_tts/configuration_voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/configuration_voxtral_tts.py new file mode 100644 index 000000000..d32a882e7 --- /dev/null +++ b/vllm_omni/model_executor/models/voxtral_tts/configuration_voxtral_tts.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from transformers import PretrainedConfig +from vllm.logger import init_logger +from vllm.transformers_utils.config import MistralConfigParser, register_config_parser + +logger = init_logger(__name__) + + +class VoxtralTTSConfig(PretrainedConfig): + """HuggingFace-style config for Voxtral TTS models.""" + + model_type = "voxtral_tts" + + def __init__( + self, + text_config: PretrainedConfig | dict | None = None, + audio_config: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + + if isinstance(text_config, PretrainedConfig): + self.text_config = text_config + elif isinstance(text_config, dict): + self.text_config = PretrainedConfig.from_dict(text_config) + else: + self.text_config = PretrainedConfig() + + self.audio_config = audio_config or {} + + def get_text_config(self, **kwargs: Any) -> PretrainedConfig: + return self.text_config + + +@register_config_parser("mistral") +class VoxtralTTSConfigParser(MistralConfigParser): + """Config parser that extends the base Mistral parser with TTS support. + + This only support voxtral_tts for now. + """ + + def _remap_mistral_audio_args(self, config_dict: dict) -> dict: + encoder_args = config_dict["multimodal"].pop("audio_model_args") + audio_tokenizer_args = config_dict["multimodal"].pop("audio_tokenizer_args", None) + audio_config = {} + if encoder_args is not None: + audio_config = { + "sampling_rate": encoder_args["audio_encoding_args"]["sampling_rate"], + "codec_args": audio_tokenizer_args, + "audio_model_args": encoder_args, + "speaker_id": audio_tokenizer_args.get("voice", {}), + } + return audio_config + + def parse( + self, + model: str | Path, + trust_remote_code: bool, + revision: str | None = None, + code_revision: str | None = None, + **kwargs: Any, + ) -> tuple[dict, PretrainedConfig]: + from vllm.transformers_utils.config import ( + _download_mistral_config_file, + ) + + config_dict = _download_mistral_config_file(model, revision) + + from vllm.transformers_utils.configs.mistral import ( + _remap_general_mistral_args, + _remap_mistral_quantization_args, + ) + + # Extract audio config before building text config + audio_config = {} + if (config_dict.get("multimodal") or {}).get("audio_model_args"): + audio_config = self._remap_mistral_audio_args(config_dict) + + # Build text_config from the top-level keys + non_text_keys = {"multimodal"} + text_config = {k: v for k, v in config_dict.items() if k not in non_text_keys} + text_config = _remap_general_mistral_args(text_config) + if text_config.get("quantization"): + text_config = _remap_mistral_quantization_args(text_config) + + # The text sub-model is a plain MistralForCausalLM + text_config.setdefault("architectures", ["MistralForCausalLM"]) + + config = VoxtralTTSConfig( + text_config=PretrainedConfig.from_dict(text_config), + audio_config=audio_config, + architectures=config_dict.get("architectures", ["VoxtralTTSForConditionalGeneration"]), + ) + + return config_dict, config diff --git a/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py b/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py new file mode 100644 index 000000000..ad7a2264b --- /dev/null +++ b/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py @@ -0,0 +1,256 @@ +""" +CUDA Graph wrapper for AcousticTransformer in VoxtralTTS. + +Currently it support FlowMatching implementation only. + +Captures the AcousticTransformer forward pass (semantic logit + +n-step Euler ODE with CFG) into CUDA graphs for fixed batch sizes, +eliminating kernel launch overhead on every decode step. +""" + +import torch +from torch.cuda import CUDAGraph +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation import ( + AudioSpecialTokens, +) + +logger = init_logger(__name__) + + +class CUDAGraphAcousticTransformerWrapper: + """ + CUDA Graph wrapper for the acoustic transformer. + + Replaces the vLLM Sampler path (which has Python-level branching and dynamic + tensor allocation) with a CUDA-graph-compatible path using torch.argmax + (equivalent to top_k=1 greedy sampling) and pre-allocated buffers. + """ + + DEFAULT_CAPTURE_SIZES = [1, 2, 4, 8, 16, 32] + + def __init__( + self, + model, + capture_sizes: list[int] | None = None, + ): + self.model = model + self.acoustic_transformer = model.acoustic_transformer + + self.capture_sizes = sorted(capture_sizes or self.DEFAULT_CAPTURE_SIZES) + + # Pre-compute constants from the acoustic transformer + self.empty_audio_token_id = AudioSpecialTokens.id(AudioSpecialTokens.empty_audio) + self.end_audio_token_id = AudioSpecialTokens.id(AudioSpecialTokens.end_audio) + self.semantic_mask_start = len(AudioSpecialTokens) + self.acoustic_transformer.model_args.semantic_codebook_size + self.n_acoustic_codebook = self.acoustic_transformer.model_args.n_acoustic_codebook + self.acoustic_embeddings_levels = self.acoustic_transformer.acoustic_embeddings_levels + + self.cfg_alpha = 1.2 + self.n_steps = 16 + + # Graph storage + self.graphs: dict[int, CUDAGraph] = {} + self.static_inputs: dict[int, torch.Tensor] = {} + self.static_noise: dict[int, torch.Tensor] = {} + self.static_fake_eos: dict[int, torch.Tensor] = {} + self.static_audio_codes: dict[int, torch.Tensor] = {} + + self.enabled = False + self._warmed_up = False + + def _warmup_and_capture(self, device: torch.device, dtype: torch.dtype, hidden_dim: int): + """Perform eager warmup and CUDA graph capture for all bucket sizes.""" + if self._warmed_up: + logger.warning("CUDAGraphAcousticTransformerWrapper already warmed up, skipping") + return + + logger.info( + "CUDAGraphAcousticTransformerWrapper: starting warmup and capture for sizes %s", + self.capture_sizes, + ) + + # Pre-create persistent buffers + self.timesteps = torch.linspace(0, 1, self.n_steps, device=device, dtype=dtype) + self.fake_eos_one = torch.tensor(1.0, dtype=dtype, device=device) + self.fake_eos_zero = torch.tensor(0.0, dtype=dtype, device=device) + + # Phase 1: Eager warmup for ALL capture sizes + for size in self.capture_sizes: + dummy = torch.zeros(size, hidden_dim, device=device, dtype=dtype) + with torch.no_grad(): + self._forward_cudagraph_compatible(dummy) + + torch.cuda.synchronize(device) + + # Phase 2: Capture graphs + for size in self.capture_sizes: + try: + self._capture_graph_for_size(size, device, dtype, hidden_dim) + logger.info(" Captured CUDA Graph for batch_size=%d", size) + except Exception: + logger.warning( + " Failed to capture CUDA Graph for batch_size=%d", + size, + exc_info=True, + ) + + self.enabled = True + self._warmed_up = True + logger.info( + "CUDAGraphAcousticTransformerWrapper warmup complete. Captured %d/%d graphs.", + len(self.graphs), + len(self.capture_sizes), + ) + + def _forward_cudagraph_compatible(self, hidden_states: torch.Tensor, noise: torch.Tensor | None = None): + """ + The actual computation captured by the CUDA graph. + + This replaces the full compute_mm_logits -> acoustic_transformer.forward() path + with a graph-compatible version: + - Uses argmax instead of vLLM Sampler (equivalent for top_k=1) + - Uses pre-created timesteps buffer instead of torch.linspace + - Uses pre-created scalar tensors for torch.where + - Calls _predict_velocity directly + - Uses a pre-allocated noise buffer to avoid baking random state + into the CUDA graph + """ + at = self.acoustic_transformer + B = hidden_states.shape[0] + + # --- Semantic logits via linear projection --- + semantic_logit = at.semantic_codebook_output(hidden_states).float() + semantic_logit[:, self.empty_audio_token_id] = -float("inf") + semantic_logit[:, self.semantic_mask_start :] = -float("inf") + + # argmax == top_k=1 greedy sampling + semantic_code = semantic_logit.argmax(dim=-1, keepdim=True) # (B, 1) + + # --- Flow matching: Euler ODE --- + should_decode = semantic_code.squeeze(1) != self.end_audio_token_id + + if noise is not None: + x = noise + else: + x = torch.randn(B, self.n_acoustic_codebook, device=hidden_states.device, dtype=hidden_states.dtype) + + # Pre-compute zero hidden states for unconditional CFG branch + hidden_states_zero = torch.zeros_like(hidden_states) + + timesteps = self.timesteps + for i in range(len(timesteps) - 1): + t = timesteps[i] + dt = timesteps[i + 1] - timesteps[i] + + # Batch conditional + unconditional velocity in a single forward pass + t_emb = at.time_embedding(t.view(-1, 1).repeat(B, 1)).to(hidden_states.dtype) + x_batched = torch.cat([x, x], dim=0) # (2B, C) + llm_batched = torch.cat([hidden_states, hidden_states_zero], dim=0) # (2B, D) + t_emb_batched = t_emb.repeat(2, 1) # (2B, D) + + v_all = at._predict_velocity(x_t=x_batched, llm_output=llm_batched, t_emb=t_emb_batched) + v_t, uncond_v_t = v_all[:B], v_all[B:] + + # CFG combination + v_t = self.cfg_alpha * v_t + (1 - self.cfg_alpha) * uncond_v_t + + x = x + v_t * dt + + # --- Quantize --- + sampled = torch.clamp(x, -1, 1) + scaled_x = ((sampled + 1) / 2) * (self.acoustic_embeddings_levels - 1) + output_codes = scaled_x.round().long() + output_codes[~should_decode] = self.empty_audio_token_id + acoustic_codes = output_codes + len(AudioSpecialTokens) + + # --- Combine semantic + acoustic --- + audio_codes = torch.cat([semantic_code, acoustic_codes], dim=1) # (B, 1 + n_acoustic) + + # --- Compute fake_eos --- + fake_eos = torch.where( + audio_codes[:, 0] == self.end_audio_token_id, + self.fake_eos_one, + self.fake_eos_zero, + ) + + return fake_eos, audio_codes + + def _capture_graph_for_size( + self, + size: int, + device: torch.device, + dtype: torch.dtype, + hidden_dim: int, + ): + """Capture a CUDA graph for a specific batch size.""" + static_input = torch.zeros(size, hidden_dim, device=device, dtype=dtype) + static_noise = torch.randn(size, self.n_acoustic_codebook, device=device, dtype=dtype) + + # Stabilizing eager run + with torch.no_grad(): + _ = self._forward_cudagraph_compatible(static_input, noise=static_noise) + + torch.cuda.synchronize(device) + + graph = CUDAGraph() + with torch.no_grad(): + with torch.cuda.graph(graph): + static_fake_eos, static_audio_codes = self._forward_cudagraph_compatible( + static_input, noise=static_noise + ) + + self.graphs[size] = graph + self.static_inputs[size] = static_input + self.static_noise[size] = static_noise + self.static_fake_eos[size] = static_fake_eos + self.static_audio_codes[size] = static_audio_codes + + def _get_padded_size(self, actual_size: int) -> int | None: + """Round up to the nearest captured bucket size.""" + for size in self.capture_sizes: + if actual_size <= size: + return size + return None + + def __call__( + self, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, dict[str, list[torch.Tensor]] | None]: + """ + Drop-in replacement for model.compute_mm_logits(). + + Falls back to eager execution if: + - CUDA graph is not enabled/warmed up + - Batch size exceeds largest captured size + """ + actual_size = hidden_states.shape[0] + + if not self.enabled or not self._warmed_up: + return self.model.compute_mm_logits(hidden_states) + + padded_size = self._get_padded_size(actual_size) + if padded_size is None or padded_size not in self.graphs: + return self.model.compute_mm_logits(hidden_states) + + # Zero static input, then copy actual data + self.static_inputs[padded_size].zero_() + self.static_inputs[padded_size][:actual_size] = hidden_states + + # Fill noise buffer with fresh random values before replay so the + # flow-matching ODE starts from different initial noise each time. + self.static_noise[padded_size].normal_() + + # Replay captured graph + self.graphs[padded_size].replay() + + # Clone and slice outputs for actual batch size + fake_eos = self.static_fake_eos[padded_size][:actual_size].clone() + audio_codes = self.static_audio_codes[padded_size][:actual_size].clone() + + # Package into expected format: (fake_eos, {"audio": [list of tensors]}) + audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0)) + mm_tokens = {"audio": audio_list} + + return fake_eos, mm_tokens diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py new file mode 100644 index 000000000..7e627135a --- /dev/null +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py @@ -0,0 +1,372 @@ +from collections.abc import Iterable +from functools import cached_property +from pathlib import Path + +import regex as re +import torch +import torch.nn as nn +from huggingface_hub import hf_hub_download +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.interfaces import SupportsMultiModal +from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.sequence import IntermediateTensors +from vllm.tokenizers import cached_tokenizer_from_config +from vllm.v1.outputs import SamplerOutput +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.sample.sampler import Sampler + +from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.voxtral_tts.cuda_graph_acoustic_transformer_wrapper import ( + CUDAGraphAcousticTransformerWrapper, +) +from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation import ( + VoxtralTTSDummyInputsBuilder, + VoxtralTTSMultiModalProcessor, + VoxtralTTSProcessingInfo, +) + +logger = init_logger(__name__) + + +def parse_batched_audio_input(input_ids: torch.Tensor, num_codebooks: int) -> tuple[list[torch.Tensor], list[int]]: + """Parse batched input_ids with [ctx_frames, context_length, ...tokens] format. + + Each request in the batch is laid out as: + [ctx_frames, context_length, ] + where the flat audio tokens have length (ctx_frames + context_length) * num_codebooks. + + Returns: + all_audio_tokens: list of (num_frames, num_codebooks) tensors per request. + all_ctx_frames: list of ctx_frames values per request. + """ + all_audio_tokens: list[torch.Tensor] = [] + all_ctx_frames: list[int] = [] + offset = 0 + while offset < input_ids.numel(): + ctx_frames = int(input_ids[offset].item()) + context_length = int(input_ids[offset + 1].item()) + offset += 2 + + req_input_ids_length = (ctx_frames + context_length) * num_codebooks + req_input_ids = input_ids[offset : offset + req_input_ids_length] + offset += req_input_ids_length + assert req_input_ids.numel() % num_codebooks == 0, ( + f"Number of elements must be divisible by {num_codebooks} but get {req_input_ids.numel()=}" + ) + num_of_audios = req_input_ids.numel() // num_codebooks + audio_tokens = req_input_ids.view(num_of_audios, num_codebooks) + all_audio_tokens.append(audio_tokens) + all_ctx_frames.append(ctx_frames) + return all_audio_tokens, all_ctx_frames + + +def apply_ctx_frames_cutting( + batch_audio_arrays: list[torch.Tensor], + all_ctx_frames: list[int], + downsample_factor: int, +) -> list[torch.Tensor]: + """Cut leading context samples from decoded audio arrays. + + Args: + batch_audio_arrays: raw decoded audio arrays, one per request. + all_ctx_frames: number of context frames per request. + downsample_factor: samples per frame in the audio tokenizer. + + Returns: + List of audio arrays with context samples removed. + """ + result: list[torch.Tensor] = [] + for audio_array, ctx_frames in zip(batch_audio_arrays, all_ctx_frames): + if ctx_frames > 0: + new_sample_cut = downsample_factor * ctx_frames + audio_array = audio_array[new_sample_cut:] + result.append(audio_array) + return result + + +@MULTIMODAL_REGISTRY.register_processor( + VoxtralTTSMultiModalProcessor, + info=VoxtralTTSProcessingInfo, + dummy_inputs=VoxtralTTSDummyInputsBuilder, +) +class VoxtralTTSForConditionalGeneration( + nn.Module, + SupportsMultiModal, + CustomProcessMixin, +): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.config = config + self.repo_id = vllm_config.model_config.model + self.is_hf_model = not Path(self.repo_id).is_dir() + self.model_stage = vllm_config.model_config.model_stage + if self.model_stage == "audio_generation": + self.has_preprocess = True + self.has_postprocess = True + self.requires_raw_input_tokens = True + self.set_custom_preprocess(self.tts_preprocess) + self.set_custom_postprocess(self.tts_postprocess) + self.audio_generation = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config, + prefix=maybe_prefix(prefix, "audio_generation"), + architectures=["VoxtralTTSAudioGeneration"], + ) + self.model = self.audio_generation + self.audio_tokenizer = None + self._cudagraph_acoustic_transformer = None + self._vllm_config = vllm_config + tokenizer = cached_tokenizer_from_config(vllm_config.model_config) + self._audio_token_id = tokenizer.instruct.audio_encoder.special_ids.audio + speaker_id = config.audio_config.get("speaker_id", None) + if speaker_id and self.is_hf_model: + self.voice_to_embedding = { + sid: torch.load(hf_hub_download(repo_id=self.repo_id, filename=f"voice_embedding/{sid}.pt")) + for sid in speaker_id + } + logger.info("Available voice embeddings: %s", list(self.voice_to_embedding.keys())) + else: + self.voice_to_embedding = {} + logger.warning("No speaker_id configured in audio_config.No voice embeddings will be available.") + elif self.model_stage == "audio_tokenizer": + self.requires_raw_input_tokens = True + self.audio_generation = None + self.audio_tokenizer = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config, + prefix=maybe_prefix(prefix, "audio_tokenizer"), + architectures=["VoxtralTTSAudioTokenizer"], + ) + self.model = self.audio_tokenizer + else: + raise ValueError("Invalid model stage") + + # -------------------- CUDA Graph for acoustic transformer -------------------- + def _enable_acoustic_transformer_cudagraph(self): + """Initialize and capture CUDA graphs for compute_mm_logits.""" + if self.model_stage != "audio_generation" or not hasattr(self, "_vllm_config"): + return + + model_cfg = getattr(self._vllm_config, "model_config", None) + if model_cfg and getattr(model_cfg, "enforce_eager", False): + logger.info("CUDA Graph for acoustic transformer not enabled: --enforce-eager is set") + return + + try: + acoustic_transformer = self.model.acoustic_transformer + device = next(acoustic_transformer.parameters()).device + dtype = next(acoustic_transformer.parameters()).dtype + hidden_dim = self.model.config.text_config.hidden_size + + wrapper = CUDAGraphAcousticTransformerWrapper(self.model) + wrapper._warmup_and_capture(device, dtype, hidden_dim) + self._cudagraph_acoustic_transformer = wrapper + logger.info("CUDA Graph for acoustic transformer enabled") + except Exception: + logger.warning("Failed to enable CUDA Graph for acoustic transformer", exc_info=True) + self._cudagraph_acoustic_transformer = None + + # -------------------- Device utilities -------------------- + @cached_property + def sampler(self): + if hasattr(self.model, "sampler"): + return self.model.sampler + return Sampler() + + def tts_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict | None): + self.post_process_idx = 0 + audio_tokens = info_dict.pop("audio", None) + if audio_tokens is not None: + kwargs = {"audio_tokens": audio_tokens.to(input_ids.device)} + multimodal_embeddings = self.model.embed_multimodal(**kwargs) + if input_ids[0] == self._audio_token_id: + input_embeds = multimodal_embeddings[0] + return input_ids, input_embeds, info_dict + voice = info_dict.pop("voice", None) + if voice is not None: + if isinstance(voice, list): + voice = voice[0] + multimodal_embeddings = self.voice_to_embedding[voice].to(input_ids.device).clone().detach() + is_multimodal = input_ids == self._audio_token_id + input_embeds = self.embed_input_ids( + input_ids=input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal + ) + return input_ids, input_embeds, info_dict + return input_ids, input_embeds, info_dict + + def tts_postprocess(self, hidden_states: torch.Tensor, multimodal_outputs: object, **info_dict: object | None): + update_dict = {} + if isinstance(multimodal_outputs, dict) and "audio" in multimodal_outputs: + assert self.post_process_idx < len(multimodal_outputs["audio"]), ( + f"Expect {self.post_process_idx=} < {len(multimodal_outputs['audio'])=}" + ) + update_dict["audio"] = multimodal_outputs["audio"][self.post_process_idx] + self.post_process_idx += 1 + return update_dict + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + is_multimodal=None, + ) -> torch.Tensor: + if self.model_stage == "audio_tokenizer": + return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.config.text_config.hidden_size) + return self.model.embed_input_ids( + input_ids=input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal + ) + + def embed_multimodal(self, **kwargs): + # Delegate to generation model for multimodal processing + return self.model.embed_multimodal(**kwargs) + + def last_index_of(self, list, value): + return len(list) - 1 - list[::-1].index(value) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors | OmniOutput: + if self.model_stage == "audio_generation": + if inputs_embeds is not None: + input_ids = None + hidden_states = self.model( + input_ids, + positions, + intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + return hidden_states + + if self.model_stage == "audio_tokenizer": + if (input_ids == 0).all(): + logger.info("audio_tokenizer: sample run with dummy input") + # TODO(chenyo): Move this to dummy_inputs creation + num_codebooks = self.audio_tokenizer.num_codebooks + audio_tokens = torch.randint(low=2, high=100, size=(116, num_codebooks), dtype=torch.int32) + audio_tokens[-1, :] = 0 + audio_tokens[-1, 0] = 1 + batch_audio_arrays = self.audio_tokenizer.decode_helper_batch_async([audio_tokens]) + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"audio": batch_audio_arrays}, + ) + all_audio_tokens, all_ctx_frames = parse_batched_audio_input( + input_ids, num_codebooks=self.audio_tokenizer.num_codebooks + ) + + # Batch decode all requests at once + batch_audio_arrays_raw = self.audio_tokenizer.decode_helper_batch_async(all_audio_tokens) + batch_audio_arrays = apply_ctx_frames_cutting( + batch_audio_arrays_raw, all_ctx_frames, self.audio_tokenizer.downsample_factor + ) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"audio": batch_audio_arrays}, + ) + + def make_omni_output( + self, model_outputs: torch.Tensor | OmniOutput | tuple, logits_index: int | None = None, **kwargs + ) -> OmniOutput: + if isinstance(model_outputs, torch.Tensor): + if self.model_stage == "audio_generation": + hidden_states = model_outputs + assert logits_index is not None + input_hidden_states = hidden_states[logits_index] + if self._cudagraph_acoustic_transformer is not None: + fake_eos, multimodal_outputs = self._cudagraph_acoustic_transformer(input_hidden_states) + else: + fake_eos, multimodal_outputs = self.model.compute_mm_logits(input_hidden_states) + hidden_states[logits_index, 0] = fake_eos + return OmniOutput( + text_hidden_states=hidden_states, + multimodal_outputs=multimodal_outputs, + ) + + raise ValueError(f"Unsupported {self.model_stage} for model_outputs type: {type(model_outputs)}") + if isinstance(model_outputs, OmniOutput): + if self.model_stage == "audio_tokenizer": + return model_outputs + + raise ValueError(f"Unsupported {self.model_stage} for model_outputs type: {type(model_outputs)}") + + raise ValueError(f"Unsupported model_outputs type: {type(model_outputs)}") + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput) -> torch.Tensor | None: + # Handle OmniOutput type + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if self.model_stage == "audio_generation": + text_logits = self.model.fake_logits_for_audio_tokens(fake_eos=hidden_states) + return text_logits + if self.model_stage == "audio_tokenizer": + raise ValueError("Invalid model stage") + + def sample( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> SamplerOutput | None: + return self.model.sample(logits, sampling_metadata) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loaded_weights = set() + remapping_rules = [ + (r"^audio_tokenizer\.(.*)$", r"\1"), # noqa: E501 + ( + r"^mm_audio_embeddings\.audio_codebook_embeddings\.embeddings\.(weight|bias)", + r"audio_token_embedding.embeddings.\1", + ), # noqa: E501 + (r"^mm_audio_embeddings\.tok_embeddings\.weight", r"tok_embeddings.weight"), # noqa: E501 + ] + + def llm_weights_generator(): + nonlocal loaded_weights + for name, w in weights: + is_audio_tokenizer = name.startswith( + "mm_audio_embeddings.audio_codebook_embeddings" + ) or name.startswith("audio_tokenizer.") + + if is_audio_tokenizer and self.audio_tokenizer is not None: + # Remap name only when loading audio_tokenizer (Stage-1). + # Yield the original name in Stage-0 so audio_generation.load_weights() + # can still filter by prefix using its own remapping rules. + remapped = name + for pattern, repl in remapping_rules: + if re.fullmatch(pattern, remapped): + remapped = re.sub(pattern, repl, remapped) + name = self.audio_tokenizer.load_weight((remapped, w)) + loaded_weights.add(f"audio_tokenizer.{name}") + continue + + yield (name, w) + + audio_generation_weights = list(llm_weights_generator()) + if self.audio_generation is not None: + for name in self.audio_generation.load_weights(audio_generation_weights): + loaded_weights.add(f"audio_generation.{name}") + + # If encoder weights were not in the checkpoint, mark them as + # "loaded" so the weight-validation does not fail. + # encode_waveforms() will raise at runtime if called without encoder weights. + if self.audio_tokenizer is not None and not self.audio_tokenizer._encoder_loaded: + for name, _ in self.audio_tokenizer.named_parameters(): + if name.startswith(self.audio_tokenizer._encoder_weight_prefixes): + loaded_weights.add(f"audio_tokenizer.{name}") + + # Capture CUDA graphs for compute_mm_logits after weights are loaded + self._enable_acoustic_transformer_cudagraph() + + return loaded_weights diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py new file mode 100644 index 000000000..9b45172bd --- /dev/null +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py @@ -0,0 +1,1083 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass, fields, is_dataclass +from enum import Enum +from functools import cached_property +from math import ceil +from typing import Any, Union, get_args, get_origin # Literal, cast, + +import numpy as np +import regex as re +import torch +import torch.nn as nn +import torch.nn.functional as F + +try: + from apex.normalization import FusedRMSNorm + + rms_norm = FusedRMSNorm +except ImportError: + from torch.nn import RMSNorm as RMSNorm + + rms_norm = RMSNorm + +from mistral_common.protocol.instruct.chunk import AudioChunk, RawAudio +from mistral_common.tokens.tokenizers.audio import Audio, AudioEncoder +from transformers import BatchFeature +from transformers.tokenization_utils_base import TextInput +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.interfaces import SupportsMultiModal +from vllm.model_executor.models.utils import ( + flatten_bn, + init_vllm_registered_model, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import ( + MultiModalDataDict, + MultiModalFieldConfig, + MultiModalKwargsItems, + NestedTensors, +) +from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems, MultiModalDataParser +from vllm.multimodal.processing import BaseDummyInputsBuilder, BaseMultiModalProcessor +from vllm.multimodal.processing.processor import ( + BaseProcessingInfo, + MultiModalProcessingInfo, + ProcessorInputs, + PromptReplacement, + PromptUpdate, + TimingContext, +) +from vllm.sequence import IntermediateTensors +from vllm.tokenizers import cached_tokenizer_from_config +from vllm.tokenizers.mistral import MistralTokenizer + +weight_norm = torch.nn.utils.parametrizations.weight_norm + + +logger = init_logger(__name__) + +SUPPORTED_LANGS = { + "en": "English", + "fr": "French", + "es": "Spanish", + "de": "German", + "ar": "Arabic", + "hi": "Hindi", + "it": "Italian", + "pt": "Portuguese", +} + + +# Audio +class AudioSpecialTokens(str, Enum): + """Special tokens predicted by audio codebook heads. + + These tokens are inserted by `audio_tokens_with_pattern`. They are not part of the text vocabulary. + We offset the output audio tokens from quantizer by `len(all_special_tokens)` to avoid conflicts with text tokens. + """ + + empty_audio = "[EMPTY_AUDIO]" + end_audio = "[END_AUDIO]" + + @staticmethod + def all_special_tokens() -> list["AudioSpecialTokens"]: + return [token for token in AudioSpecialTokens] + + @staticmethod + def id(token: "AudioSpecialTokens") -> int: + return AudioSpecialTokens.all_special_tokens().index(token) + + +@dataclass +class AcousticTransformerArgs: + input_dim: int + # Define some defaults + dim: int = 768 + n_layers: int = 3 + head_dim: int = 128 + hidden_dim: int = 2048 + n_heads: int = 6 + n_kv_heads: int = 2 + use_biases: bool = False + norm_eps: float = 1e-5 + sigma: float = 1e-5 # was 0.01 in beta version + + +@dataclass +class MultimodalAudioModelArgs: + # comma-separated list of codebook sizes + # The first token in a codebook should always be reserved to indicate + # absence. The codebook size should be inclusive of this. + semantic_codebook_size: int + acoustic_codebook_size: int + n_acoustic_codebook: int + acoustic_transformer_args: AcousticTransformerArgs + + @property + def codebook_sizes(self) -> list[int]: + return [ + self.semantic_codebook_size, + *[self.acoustic_codebook_size for _ in range(self.n_acoustic_codebook)], + ] + + def get_codebook_sizes( + self, pad_to_multiple: int | None = 128, include_special_tokens: bool = True + ) -> list[int]: # O + # TODO(@alexhliu): add one more argument to specify if we are computing the size of + # a single codebook merging all codebooks + # These codebook sizes should be exactly same as output of the audio quantizer. + # We will then add buffer for special tokens and pad to nearest multiple of 8 (for GPU efficiency) + + def _round_up_to_multiple_of_number(n: int, multiple: int) -> int: + return multiple * ((n + multiple - 1) // multiple) + + result_codebook_sizes = [] + for i, cb_size in enumerate(self.codebook_sizes): + # TODO(@alexhliu): special tokens resides only in first codebook + # so ideally we shouldn't need to add them to latter codebooks + if include_special_tokens: + cb_size += len(AudioSpecialTokens.all_special_tokens()) + if pad_to_multiple is not None: + cb_size = _round_up_to_multiple_of_number(cb_size, pad_to_multiple) + result_codebook_sizes.append(cb_size) + return result_codebook_sizes + + +def _repeat_interleave(t: torch.Tensor, repeats: int) -> torch.Tensor: + """ + Faster than + keys = torch.repeat_interleave( + keys, repeats, dim=2, output_size=xq.size(2) + ) + on gpu ? + """ + return t.unsqueeze(3).expand([-1, -1, -1, repeats, -1]).flatten(2, 3) + + +def repeat_kv(keys: torch.Tensor, values: torch.Tensor, repeats: int) -> tuple[torch.Tensor, torch.Tensor]: + if repeats > 1: + keys = _repeat_interleave(keys, repeats=repeats) + values = _repeat_interleave(values, repeats=repeats) + return keys, values + + +def from_nested_dict(cls, d): + """Recursively instantiate dataclasses from nested dicts.""" + if not is_dataclass(cls): + return d + + kwargs = {} + for f in fields(cls): + value = d.get(f.name, getattr(cls, f.name, None)) + field_type = f.type + + # Unwrap Optional / Union types + origin = get_origin(field_type) + if origin is Union: + args = get_args(field_type) + # take first non-None type + non_none_types = [a for a in args if a is not type(None)] + if len(non_none_types) == 1: + field_type = non_none_types[0] + + # Recurse if nested dataclass + if is_dataclass(field_type) and isinstance(value, dict): + value = from_nested_dict(field_type, value) + + kwargs[f.name] = value + + return cls(**kwargs) + + +class FeedForward(nn.Module): + def __init__( + self, + dim: int, + hidden_dim: int, + use_biases: bool, + ) -> None: + super().__init__() + + self.w1 = nn.Linear( + dim, + hidden_dim, + bias=False, + ) + self.w2 = nn.Linear( + hidden_dim, + dim, + bias=use_biases, + ) + self.w3 = nn.Linear( + dim, + hidden_dim, + bias=False, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +class BidirectionalAttention(nn.Module): + # TODO(@alexhliu): see if we can merge this with codec attention + """Attention layer (without any RoPE embeddings).""" + + def __init__( + self, + args: AcousticTransformerArgs, + layer_id: int, + ) -> None: + super().__init__() + self.args = args + + self.n_local_heads: int = args.n_heads + self.n_local_kv_heads: int = args.n_kv_heads + self.repeats = self.n_local_heads + self.layer_id = layer_id + + self.head_dim = args.head_dim + + # Note difference in how we set biases: + # - Mistral LM: only accepts bias=args.use_biases for wo, + # everything else is False + # - Whisper: wk always has bias=False, and bias=True for wq, wv and wo + self.wq = nn.Linear( + args.dim, + args.n_heads * args.head_dim, + bias=args.use_biases, + ) + self.wk = nn.Linear( + args.dim, + args.n_kv_heads * args.head_dim, + bias=False, + ) + self.wv = nn.Linear( + args.dim, + args.n_kv_heads * args.head_dim, + bias=args.use_biases, + ) + self.wo = nn.Linear( + (args.n_heads * args.head_dim), + args.dim, + bias=args.use_biases, + ) + + self.softmax_scale: float = self.args.head_dim**-0.5 + self.repeats = self.n_local_heads // self.n_local_kv_heads + + def _native_attention( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + ) -> torch.tensor: + scale = 1.0 / query.shape[-1] ** 0.5 + query = query * scale + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + attn = query @ key.transpose(-2, -1) + attn = attn.softmax(-1) + attn = attn @ value + return attn.transpose(1, 2).contiguous() + + def _forward_attention( # type: ignore + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + ) -> torch.Tensor: + key, value = repeat_kv(key, value, repeats=self.repeats) + bsz, seqlen, _, _ = query.shape + + output = self._native_attention(query, key, value) + + return output.view(bsz, seqlen, -1) + + def forward( + self, + x: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + if x.dim() == 2: + bsz, (seqlen, _) = 1, x.shape + else: + bsz, seqlen, _ = x.shape + + # compute xq + xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) + + xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim) + xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim) + xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim) + + output = self._forward_attention(query=xq, key=xk, value=xv, **kwargs) + output = output.view(bsz, seqlen, self.n_local_heads * self.head_dim) + return self.wo(output).squeeze(0) + + +class AcousticTransformerBlock(nn.Module): + def __init__( + self, + layer_id: int, + args: AcousticTransformerArgs, + ) -> None: + super().__init__() + self._layer_id = layer_id + self.n_heads = args.n_heads + self.dim = args.dim + self.attention = BidirectionalAttention( + args, + layer_id=layer_id, + ) + + self.feed_forward = FeedForward( + args.dim, + args.hidden_dim, + args.use_biases, + ) + + self.attention_norm = rms_norm( + args.dim, + eps=args.norm_eps, + ) + self.ffn_norm = rms_norm( + args.dim, + eps=args.norm_eps, + ) + self.args = args + + @property + def layer_id(self) -> int: + return self._layer_id + + def forward( + self, + x: torch.Tensor, + ) -> torch.Tensor: + r = self.attention.forward(self.attention_norm(x)) + h = x + r + r = self.feed_forward.forward(self.ffn_norm(h)) + out = h + r + return out + + +### Flow Matching Acoustic Transformer ### + + +class TimeEmbedding(nn.Module): + """Sinusoidal Embedding for encoding time""" + + def __init__(self, dim: int, theta: float = 10000.0) -> None: + super().__init__() + inv_freq = torch.exp(-math.log(theta) * torch.arange(dim // 2).float() / (dim // 2)) + # NOTE(@alexhliu): mistral codebase requires persistent = True for save/load + self.register_buffer("inv_freq", inv_freq, persistent=True) + + def forward(self, t: torch.Tensor) -> torch.Tensor: + emb = torch.einsum( + "bi, j -> bj", + t, + self.inv_freq, + ) + return torch.cat((emb.cos(), emb.sin()), dim=-1) + + +class FlowMatchingAudioTransformer(nn.Module): + def __init__( + self, + audio_model_args: dict, + ) -> None: + super().__init__() + if "codebook_sizes" in audio_model_args: + codebook_sizes = [int(c) for c in audio_model_args.pop("codebook_sizes").split(",")] + audio_model_args.update( + { + "semantic_codebook_size": codebook_sizes[0], + "acoustic_codebook_size": codebook_sizes[1], + "n_acoustic_codebook": len(codebook_sizes) - 1, + } + ) + self.model_args: MultimodalAudioModelArgs = from_nested_dict(MultimodalAudioModelArgs, audio_model_args) + # MultimodalAudioModelArgs + assert isinstance(self.model_args, MultimodalAudioModelArgs) + args = self.model_args.acoustic_transformer_args + # AcousticTransformerArgs + self.acoustic_transformer_args = args + assert isinstance(self.acoustic_transformer_args, AcousticTransformerArgs) + + # currently assuming always 1 semantic codebook + N acoustic codebooks + self.num_non_acoustic_embeddings = 1 + self.num_acoustic_codebooks = len(self.model_args.get_codebook_sizes()) - self.num_non_acoustic_embeddings + + # flow matching utils + self.sigma = args.sigma + + # codebook sizes + acoustic_codebook_sizes = self.model_args.get_codebook_sizes( + pad_to_multiple=None, include_special_tokens=False + )[1:] + assert len(set(acoustic_codebook_sizes)) == 1, "only 1 size for acoustic codebooks supported" + self.acoustic_embeddings_levels = acoustic_codebook_sizes[0] + self.acoustic_embeddings_dim = len(acoustic_codebook_sizes) + + self._init_audio_embeddings_layer() + self._init_output_layer() + self._init_layers() + + self._end_audio_token_id = AudioSpecialTokens.id(AudioSpecialTokens.end_audio) + self._empty_audio_token_id = AudioSpecialTokens.id(AudioSpecialTokens.empty_audio) + + # Flow matching constants + # TODO(chenyo): hardcoded, need to fix + self._acoustic_decode_iters = 16 + # TODO(chenyo): hardcoded, need to fix + self._cfg_alpha = 1.2 + self._noise_scale = 1.0 + self.register_buffer( + "_timesteps", + torch.linspace(0, 1, self._acoustic_decode_iters), + persistent=False, + ) + + def load_weight(self, weight: tuple[str, torch.Tensor]) -> str: + params_dict = dict(self.named_parameters()) + name, loaded_weight = weight + if name not in params_dict: + logger.warning(f"{name} not found in FlowMatchingAudioTransformer (UNUSED)") + return name + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + return name + + def _init_audio_embeddings_layer(self) -> None: + # time embedding for flow step + self.time_embedding = TimeEmbedding(self.acoustic_transformer_args.dim) + # override the input projection embedding + input_dim = self.acoustic_embeddings_dim + + self.input_projection = nn.Linear(input_dim, self.acoustic_transformer_args.dim, bias=False) + self.time_projection = nn.Linear( + self.acoustic_transformer_args.dim, + self.acoustic_transformer_args.dim, + bias=False, + ) + self.llm_projection = nn.Linear( + self.acoustic_transformer_args.input_dim, + self.acoustic_transformer_args.dim, + bias=False, + ) + + def _init_output_layer(self) -> None: + padded_codebook_sizes = self.model_args.get_codebook_sizes(pad_to_multiple=128) + self.semantic_codebook_output = nn.Linear( + self.acoustic_transformer_args.dim, + padded_codebook_sizes[0], + self.acoustic_transformer_args.use_biases, + ) + self.acoustic_codebook_output = nn.Linear( + in_features=self.acoustic_transformer_args.dim, + # we predict a float for every codebook + out_features=self.model_args.n_acoustic_codebook, + bias=False, + ) + + def _init_layers(self) -> None: + self.layers_ids: list[int] = list(range(self.acoustic_transformer_args.n_layers)) + self.layers = nn.ModuleDict() + for layer_id in self.layers_ids: + block = AcousticTransformerBlock( + layer_id=layer_id, + args=self.acoustic_transformer_args, + ) + self.layers[str(layer_id)] = block + + self.norm = rms_norm(self.acoustic_transformer_args.dim, self.acoustic_transformer_args.norm_eps) + + def forward_attention_layers(self, h: torch.Tensor) -> torch.Tensor: + for layer_id in self.layers_ids: + layer = self.layers[str(layer_id)] + h = layer(h) + return h + + def decode_one_frame( + self, + semantic_code: torch.Tensor, + llm_hidden: torch.Tensor, + ) -> torch.Tensor: + B = semantic_code.shape[0] + + # Skip decoding if codebook 0 is [END_AUDIO] token. + should_decode = semantic_code != self._end_audio_token_id + + # acoustic_codes starts from x_0 + x_0 = torch.randn(B, self.model_args.n_acoustic_codebook).to(dtype=llm_hidden.dtype, device=llm_hidden.device) + x_0 = self._noise_scale * x_0 + + timesteps = self._timesteps.to(dtype=llm_hidden.dtype) + llm_hidden_zero = torch.zeros_like(llm_hidden) + + # Euler integration with batched conditional + unconditional velocity + sampled = x_0 + for i in range(len(timesteps) - 1): + t = timesteps[i] + dt = timesteps[i + 1] - timesteps[i] + + t_emb = self.time_embedding(t.view(-1, 1).repeat(B, 1)).to(llm_hidden.dtype) + + # Batch cond + uncond into a single forward pass (2B batch) + x_batched = torch.cat([sampled, sampled], dim=0) + llm_batched = torch.cat([llm_hidden, llm_hidden_zero], dim=0) + t_emb_batched = torch.cat([t_emb, t_emb], dim=0) + + v_all = self._predict_velocity( + x_t=x_batched, + llm_output=llm_batched, + t_emb=t_emb_batched, + ) + v_t, uncond_v_t = v_all[:B], v_all[B:] + v_t = self._cfg_alpha * v_t + (1 - self._cfg_alpha) * uncond_v_t + + sampled = sampled + v_t * dt + + # quantize & mask end_of_audio + sampled = torch.clamp(sampled, -1, 1) # manually clip to [-1, 1] + scaled_x = ((sampled + 1) / 2) * (self.acoustic_embeddings_levels - 1) # scale to 0 ~ max level + output_codes = scaled_x.round().long() + output_codes[~should_decode] = self._empty_audio_token_id + return output_codes + len(AudioSpecialTokens) + + def _predict_velocity( + self, + x_t: torch.Tensor, # BxC + llm_output: torch.Tensor, # BxD + t_emb: torch.Tensor, # BxD + ) -> torch.Tensor: + x_t = x_t.to(llm_output.dtype) + + t_emb = self.time_projection(t_emb) + llm_output = self.llm_projection(llm_output) + + acoustic_and_semantic_embeddings = [ + self.input_projection(x_t.unsqueeze(1)), # Bx1xD + t_emb.unsqueeze(1), + llm_output.unsqueeze(1), + ] + acoustic_transformer_inputs = torch.concatenate(acoustic_and_semantic_embeddings, dim=1) + + # forward transformer + attn_output = self.forward_attention_layers(acoustic_transformer_inputs) + final_hidden = self.norm(attn_output) + final_hidden = final_hidden.view(-1, acoustic_transformer_inputs.shape[1], final_hidden.shape[-1]) + # predict v_t + v_t = self.acoustic_codebook_output(final_hidden[:, 0, :]) + + return v_t + + def forward( + self, + llm_hidden: torch.Tensor, + ) -> torch.Tensor: + # llm_hidden: BxD + semantic_logit = self.semantic_codebook_output(llm_hidden).float() + semantic_logit[:, self._empty_audio_token_id] = -float("inf") # 1 = eoa is allowed + semantic_logit[:, (len(AudioSpecialTokens) + self.model_args.semantic_codebook_size) :] = -float("inf") + + # semantic_logit: Bx1 + semantic_code = semantic_logit.argmax(dim=-1, keepdim=True) + + # acoustic codes, TODO(@chenyo): config sampling + acoustic_codes = self.decode_one_frame( + semantic_code.squeeze(1), + llm_hidden, + ) + + audio_codes = torch.concatenate( + [semantic_code, acoustic_codes], + dim=1, + ) + return audio_codes + + +class VoxtralTTSProcessorAdapter: + """ + Provide a HF-compatible interface + """ + + def __init__(self, tokenizer: MistralTokenizer) -> None: + super().__init__() + self.tokenizer = tokenizer + + @cached_property + def _audio_processor(self) -> AudioEncoder: + audio_encoder = self.tokenizer.instruct.audio_encoder + assert isinstance(audio_encoder, AudioEncoder) + return audio_encoder + + @cached_property + def audio_token_id(self) -> int: + return self._audio_processor.special_ids.audio + + @cached_property + def begin_audio_token_id(self) -> int: + return self._audio_processor.special_ids.begin_audio + + @cached_property + def sampling_rate(self) -> int: + return self._audio_processor.audio_config.sampling_rate + + @cached_property + def frame_rate(self) -> float: + return self._audio_processor.audio_config.frame_rate + + def get_num_audio_tokens( + self, + audio_length: int, + ) -> int: + # TODO(@chenyo): only +1 for TTS + return 1 + ceil(audio_length / (self.sampling_rate // self.frame_rate)) + + def __call__( + self, + text: TextInput | list[TextInput] | None = None, + audios: np.ndarray | list[np.ndarray] | None = None, + audio_tokens: np.ndarray | list[np.ndarray] | None = None, + **kwargs, + ) -> Mapping[str, NestedTensors]: + if text is None: + text = [] + if not isinstance(text, list): + text = [text] + if audios is None: + audios = [] + if not isinstance(audios, list): + audios = [audios] + if audio_tokens is None: + audio_tokens = [] + if not isinstance(audio_tokens, list): + # should be single audio tokens input? + audio_tokens = [audio_tokens] + + if audios: + # Allow dummy text, which is used for profiling as well as token inputs + if any(len(t) > 0 for t in text): + raise ValueError( + "You've passed text inputs instead of token inputs. " + "Make sure to process your input via `mistral_common`'s " + "tokenizer or pass a chat completion request. " + "For more info, see: " + "https://github.com/vllm-project/vllm/issues/8411." + ) + + audios_tokens = list[torch.Tensor]() + audios_processed = list[torch.Tensor]() + for audio in audios: + assert isinstance(audio, np.ndarray) + assert audio.ndim == 1 + num_audio_token = self.get_num_audio_tokens(len(audio)) + audio_tokens = [self.begin_audio_token_id] + [self.audio_token_id] * num_audio_token + audios_tokens.append(torch.tensor(audio_tokens)) + audios_processed.append(torch.tensor(audio)) + input_ids = torch.cat(audios_tokens)[None].expand(len(text), -1) + return BatchFeature( + { + "input_ids": input_ids, + "audio_arrays": audios_processed, + } + ) + elif audio_tokens: + # Allow dummy text, which is used for profiling as well as token inputs + if any(len(t) > 0 for t in text): + raise ValueError( + "You've passed text inputs instead of token inputs. " + "Make sure to process your input via `mistral_common`'s " + "tokenizer or pass a chat completion request. " + "For more info, see: " + "https://github.com/vllm-project/vllm/issues/8411." + ) + + text_tokens_for_audio = list[torch.Tensor]() + assert audios is not None + audio_tokens_pt = list[torch.Tensor]() + for audio_token_array in audio_tokens: + assert isinstance(audio_token_array, np.ndarray) + assert audio_token_array.ndim == 2 # (T, C) + num_audio_token = audio_token_array.shape[0] + text_token_ids = [self.begin_audio_token_id] + [self.audio_token_id] * num_audio_token + text_tokens_for_audio.append(torch.tensor(text_token_ids)) + audio_tokens_pt.append(torch.tensor(audio_token_array)) + input_ids = torch.cat(text_tokens_for_audio)[None].expand(len(text), -1) + + return BatchFeature( + { + "input_ids": input_ids, + "audio_tokens": audio_tokens_pt, + } + ) + else: + assert text is not None + return BatchFeature({"input_ids": torch.tensor(self.tokenizer(text).input_ids)}) + + +class VoxtralTTSProcessingInfo(BaseProcessingInfo): + def get_tokenizer(self) -> MistralTokenizer: + tokenizer = cached_tokenizer_from_config(self.ctx.model_config) + if not isinstance(tokenizer, MistralTokenizer): + raise ValueError("This model requires `--tokenizer-mode mistral`") + + return tokenizer + + def get_data_parser(self): + return MultiModalDataParser( + target_sr=self.get_hf_processor().sampling_rate, + expected_hidden_size=self._get_expected_hidden_size(), + ) + + def get_hf_processor(self, **kwargs: object) -> VoxtralTTSProcessorAdapter: + return VoxtralTTSProcessorAdapter(self.get_tokenizer()) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"audio": 1} + + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> Mapping[str, int]: + return {"audio": self.get_max_audio_tokens()} + + def get_max_audio_tokens(self) -> int: + return self.ctx.model_config.max_model_len + + def get_max_audio_array_len(self) -> int: + processor = self.get_hf_processor() + return self.get_max_audio_tokens() * int(processor.sampling_rate // processor.frame_rate) + + +class VoxtralTTSDummyInputsBuilder(BaseDummyInputsBuilder[VoxtralTTSProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return "" + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + # mm_options: Mapping[str, BaseDummyOptions] | None = None, + mm_options: Mapping[str, Any] | None = None, + ) -> MultiModalDataDict: + num_audios = mm_counts.get("audio", 0) + length = self.info.get_max_audio_array_len() + audio_overrides = mm_options.get("audio") if mm_options else None + return {"audio": self._get_dummy_audios(length=length, num_audios=num_audios, overrides=audio_overrides)} + + def get_dummy_processor_inputs( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, Any] | None = None, + ) -> ProcessorInputs: + dummy_mm_items = self.get_dummy_mm_data(seq_len, mm_counts, mm_options) + dummy_audios = dummy_mm_items.get("audio", []) + + audio_chunks: list[AudioChunk] = [] + for audio in dummy_audios: + audio_item = Audio( + audio_array=audio, + sampling_rate=self.info.get_hf_processor().sampling_rate, + format="wav", + ) + chunk = AudioChunk( + input_audio=RawAudio( + format="wav", + data=audio_item.to_base64(format="wav"), + ), + ) + audio_chunks.append(chunk) + + # hacked dummy tokens to make tts work + assert len(dummy_audios) == 1 + aud = [24] * (1 + int(math.ceil(len(dummy_audios[0]) / (24_000 // 12.5)))) + audio_tokens_all = [25, *aud] + dummy_tokens = [ + 1, + *audio_tokens_all, + 35, + 4380, + 1395, + 1261, + 2688, + 5117, + 1046, + 36, + 25, + ] + + parsed_mm_items = self.info.parse_mm_data(dummy_mm_items, validate=False) + + return ProcessorInputs(prompt=dummy_tokens, mm_data_items=parsed_mm_items) + + +class VoxtralTTSMultiModalProcessor(BaseMultiModalProcessor[VoxtralTTSProcessingInfo]): + def _get_mm_fields_config( + self, + hf_inputs: Mapping[str, NestedTensors], + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + return dict( + audio_arrays=MultiModalFieldConfig.batched("audio"), + audio_tokens=MultiModalFieldConfig.batched("audio"), + ) + # TODO(@alexhliu): + # set self.target_sr for resampling[] + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + audio_id = processor.audio_token_id + + def get_replacement(item_idx: int): + audios = mm_items.get_items("audio", AudioProcessorItems) + if isinstance(audios, AudioProcessorItems): + audio_len = audios.get_audio_length(item_idx) + audio_token_len = processor.get_num_audio_tokens(audio_len) + else: + raise ValueError(f"Unknown audio item type {type(audios)}") + return [audio_id] * audio_token_len + + return [ + PromptReplacement( + modality="audio", + target="", # Never match the prompt (see below note) + replacement=get_replacement, + ), + ] + + def _cached_apply_hf_processor( + self, + inputs: ProcessorInputs, + timing_ctx: TimingContext, + ) -> tuple[list[int], MultiModalProcessingInfo, bool]: + prompt_ids, mm_info, _ = super()._cached_apply_hf_processor(inputs, timing_ctx) + + # NOTE: The tokens are already inserted by the chat template + return prompt_ids, mm_info, True + + +@MULTIMODAL_REGISTRY.register_processor( + VoxtralTTSMultiModalProcessor, + info=VoxtralTTSProcessingInfo, + dummy_inputs=VoxtralTTSDummyInputsBuilder, +) +class VoxtralTTSAudioGenerationForConditionalGeneration(nn.Module, SupportsMultiModal): + supported_languages = SUPPORTED_LANGS + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.tokenizer = cached_tokenizer_from_config(vllm_config.model_config) + + config = vllm_config.model_config.hf_config + self.config = config + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), # Is this the backbone LLM? + ) + self.audio_tokenizer = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config, + prefix=maybe_prefix(prefix, "audio_tokenizer"), + architectures=["VoxtralTTSAudioTokenizer"], + ) + self.downsample_factor = self.audio_tokenizer.downsample_factor + self.acoustic_transformer = FlowMatchingAudioTransformer( + self.config.audio_config["audio_model_args"], + ) + + audio_encoder = self.tokenizer.instruct.audio_encoder + self.audio_tok_id = audio_encoder.audio_token + self.eos_tok_id = self.tokenizer.instruct.tokenizer.eos_id + self.vocab_size = config.text_config.vocab_size + + def get_language_model(self) -> torch.nn.Module: + return self.language_model + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + if intermediate_tensors is not None: + inputs_embeds = None + hidden_states = self.language_model.model( + input_ids, + positions, + intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + return hidden_states + + def embed_multimodal(self, **kwargs) -> list[torch.Tensor] | torch.Tensor | tuple[torch.Tensor, ...] | None: + audio_waveforms, audio_tokens = self._parse_and_validate_audio_arrays(**kwargs) + if (audio_waveforms is None) and (audio_tokens is None): + return None + + if audio_waveforms is not None and audio_tokens is None: + assert audio_tokens is None, f"{audio_tokens=} {audio_waveforms=}" + res = self.audio_tokenizer.encode_waveforms(audio_waveforms) + else: + res = self.audio_tokenizer.encode_tokens(audio_tokens) + return res + + def _parse_and_validate_audio_arrays( + self, **kwargs: object + ) -> tuple[list[torch.Tensor] | None, list[torch.Tensor] | None]: + audio_arrays = kwargs.pop("audio_arrays", None) + audio_tokens = kwargs.pop("audio_tokens", None) + if (audio_arrays is None) and (audio_tokens is None): + return None, None + + if audio_arrays is not None: + if not isinstance(audio_arrays, (torch.Tensor, list)): + raise ValueError(f"Incorrect type of images. Got type: {type(audio_arrays)}") + if isinstance(audio_arrays, torch.Tensor) and audio_arrays.dim() == 3: + audio_arrays = flatten_bn(audio_arrays) + if isinstance(audio_arrays, torch.Tensor): + audio_arrays = list(audio_arrays.unbind(0)) + + if audio_tokens is not None: + # audio tokens MIGHT come from previous decoding step: shape will be BxLxCB + + if isinstance(audio_tokens, list): + # concatenate along batch dim? is that ok/expected? + # audio_tokens = torch.concatenate( + # [a.transpose(1, 2) for a in audio_tokens], dim=0) + # audio_tokens = torch.concatenate(audio_tokens, dim=1) + + # these are of shape 1xLxCB + return None, [a.transpose(1, 2) for a in audio_tokens] + + if isinstance(audio_tokens, torch.Tensor): + if audio_tokens.dim() == 4: + # TODO(chenyo): Fix after rebase with main + audio_tokens = audio_tokens.squeeze(0).view(1, audio_tokens.size(-2), audio_tokens.size(-1)) + assert audio_tokens.dim() == 3, f"{audio_tokens.ndim=} {audio_tokens.shape=}" + audio_tokens = audio_tokens.transpose(1, 2) + audio_tokens = list(audio_tokens.unsqueeze(1).unbind(0)) + else: + raise NotImplementedError(f"Unsupported type for audio tokens: {type(audio_tokens)=} {audio_tokens=}") + + return audio_arrays, audio_tokens + + def fake_logits_for_audio_tokens( + self, + fake_eos: torch.Tensor, + ) -> torch.Tensor: + """ + creates fake logits with all-but-one -inf + to force decoding of audio tokens + """ + shape = (fake_eos.shape[0], self.vocab_size) + fake_logits = torch.full(shape, float("-inf"), device=fake_eos.device) + is_eos = fake_eos[:, 0].bool() + fake_logits[is_eos, self.eos_tok_id] = 1.0 + fake_logits[~is_eos, self.audio_tok_id] = 1.0 + return fake_logits + + # TODO(chenyo): Remove this + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + text_logits = self.language_model.compute_logits( + hidden_states, + ) + assert text_logits is not None + return text_logits + + def compute_mm_logits( + self, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + audio_codes = None + mm_tokens = None + audio_codes = self.acoustic_transformer( + llm_hidden=hidden_states, + ) + fake_eos = torch.where( + audio_codes[:, 0] == AudioSpecialTokens.id(AudioSpecialTokens.end_audio), + torch.tensor(1.0, dtype=torch.bfloat16), + torch.tensor(0.0, dtype=torch.bfloat16), + ) + # BxC -> Bx1xC since this is per-step + # Make it a list for vllm-omni processing + audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0)) + mm_tokens = {"audio": audio_list} + + return fake_eos, mm_tokens + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # fmt: on + remapping_rules = [ + (r"^acoustic_transformer\.(.*)$", r"\1"), # noqa: E501 + (r"^audio_tokenizer\.(.*)$", r"\1"), # noqa: E501 + ( + r"^mm_audio_embeddings\.audio_codebook_embeddings\.embeddings\.(weight|bias)", + r"audio_token_embedding.embeddings.\1", + ), # noqa: E501 + (r"^mm_audio_embeddings\.tok_embeddings\.weight", r"tok_embeddings.weight"), # noqa: E501 + ] + + loaded_weights = set() + + def llm_weights_generator(): + nonlocal loaded_weights + for name, w in weights: + is_audio_tokenizer = name.startswith( + "mm_audio_embeddings.audio_codebook_embeddings" + ) or name.startswith("audio_tokenizer.") + is_acoustic_transformer = name.startswith("acoustic_transformer.") + + for pattern, repl in remapping_rules: + if re.fullmatch(pattern, name): + name = re.sub(pattern, repl, name) + + if is_audio_tokenizer: + name = self.audio_tokenizer.load_weight((name, w)) + loaded_weights.add(f"audio_tokenizer.{name}") + continue + + if is_acoustic_transformer: + if self.acoustic_transformer is not None: + name = self.acoustic_transformer.load_weight((name, w)) + loaded_weights.add(f"acoustic_transformer.{name}") + continue + + yield (name, w) + + for name in self.language_model.load_weights(llm_weights_generator()): + loaded_weights.add(f"language_model.{name}") + + # If encoder weights were not in the checkpoint, mark them as + # "loaded" so the weight-validation does not fail. + # encode_waveforms() will raise at runtime if called without encoder weights. + if not self.audio_tokenizer._encoder_loaded: + for name, _ in self.audio_tokenizer.named_parameters(): + if name.startswith(self.audio_tokenizer._encoder_weight_prefixes): + loaded_weights.add(f"audio_tokenizer.{name}") + + return loaded_weights diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py new file mode 100644 index 000000000..4f488e2fc --- /dev/null +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py @@ -0,0 +1,1114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +from copy import deepcopy +from dataclasses import dataclass + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.config import VllmConfig + +from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation import ( + AudioSpecialTokens, + FeedForward, + MultimodalAudioModelArgs, + from_nested_dict, +) + +try: + from flash_attn import flash_attn_func + + HAS_FLASH_ATTN = True +except ImportError: + flash_attn_func = None + HAS_FLASH_ATTN = False + +try: + from apex.normalization import FusedRMSNorm + + rms_norm = FusedRMSNorm +except ImportError: + from torch.nn import RMSNorm as RMSNorm + + rms_norm = RMSNorm +from einops import rearrange +from vllm.logger import init_logger +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +logger = init_logger(__name__) + +if not HAS_FLASH_ATTN: + logger.warning( + "flash_attn is not installed. Falling back to PyTorch SDPA for " + "audio tokenizer attention. Install flash-attn for better performance." + ) + +weight_norm = torch.nn.utils.parametrizations.weight_norm + + +CODEC_NORM_EPS = 1e-2 + + +@dataclass +class AudioTokenizerArgs: + # audio setting + channels: int = 1 + sampling_rate: int = 24000 + pretransform_patch_size: int = 240 + patch_proj_kernel_size: int = 7 + + # quantizer setting + semantic_codebook_size: int = 8192 + semantic_dim: int = 256 + acoustic_codebook_size: int = 21 + acoustic_dim: int = 36 + + # architecture (general) + conv_weight_norm: bool = True + causal: bool = True + attn_sliding_window_size: int = 16 + half_attn_window_upon_downsampling: bool = True + dim: int = 1024 + hidden_dim: int = 4096 + head_dim: int = 128 + n_heads: int = 8 + n_kv_heads: int = 8 + qk_norm_eps: float = 1e-6 + qk_norm: bool = True + use_biases: bool = False + norm_eps: float = 1e-2 + layer_scale: bool = True + layer_scale_init: float | None = None + + # architecture (encoder) + encoder_transformer_lengths_str: str = "2,2,2,2" + encoder_convs_kernels_str: str = "4,4,4,3" + encoder_convs_strides_str: str = "2,2,2,1" + + # architecture (decoder) + decoder_transformer_lengths_str: str = "2,2,2,2" + decoder_convs_kernels_str: str = "3,4,4,4" + decoder_convs_strides_str: str = "1,2,2,2" + + def __post_init__(self) -> None: + assert ( + len(self.encoder_transformer_lengths) == len(self.encoder_convs_kernels) == len(self.encoder_convs_strides) + ) + assert ( + len(self.decoder_transformer_lengths) == len(self.decoder_convs_kernels) == len(self.decoder_convs_strides) + ) + + def __str2list__(self, input_str: str) -> tuple[int, ...]: + return tuple(int(i) for i in input_str.split(",")) + + @property + def encoder_transformer_lengths(self) -> tuple[int, ...]: + return self.__str2list__(self.encoder_transformer_lengths_str) + + @property + def encoder_convs_kernels(self) -> tuple[int, ...]: + return self.__str2list__(self.encoder_convs_kernels_str) + + @property + def encoder_convs_strides(self) -> tuple[int, ...]: + return self.__str2list__(self.encoder_convs_strides_str) + + @property + def decoder_transformer_lengths(self) -> tuple[int, ...]: + return self.__str2list__(self.decoder_transformer_lengths_str) + + @property + def decoder_convs_kernels(self) -> tuple[int, ...]: + return self.__str2list__(self.decoder_convs_kernels_str) + + @property + def decoder_convs_strides(self) -> tuple[int, ...]: + return self.__str2list__(self.decoder_convs_strides_str) + + @property + def frame_rate(self) -> float: + return self.sampling_rate / (self.pretransform_patch_size * math.prod(self.encoder_convs_strides)) + + +class SemanticCodebook(nn.Module): + """Euclidean distance-based codebook for semantic quantization.""" + + def __init__(self, codebook_size: int, codebook_dim: int) -> None: + super().__init__() + self.epsilon: float = 1e-5 + self.cluster_usage: torch.Tensor + self.register_buffer("cluster_usage", torch.ones(codebook_size)) + embedding = torch.zeros(codebook_size, codebook_dim) + self.embedding_sum: torch.Tensor + self.register_buffer("embedding_sum", embedding) + self.register_buffer("_embedding", None, persistent=False) + + @property + def embedding(self) -> torch.Tensor: + if self._embedding is None: + embedding = self.embedding_sum / self.cluster_usage.clamp(min=self.epsilon)[:, None] + self._embedding: torch.Tensor + self.register_buffer("_embedding", embedding, persistent=False) + return embedding + return self._embedding + + def encode(self, x: torch.Tensor) -> torch.Tensor: + assert x.dtype.is_floating_point, f"Input should be floats, got {x.dtype}" + B, D, T = x.shape + x = rearrange(x, "b d t -> b t d").view(B * T, D) + embedding = self.embedding.to(x.device) + distances = torch.cdist(x, embedding, p=2) # (B*T, V) + codes = distances.argmin(dim=-1).view(B, 1, T) + return codes + + def decode(self, codes: torch.Tensor) -> torch.Tensor: + assert not codes.dtype.is_floating_point, f"Codes should be integers, got {codes.dtype}" + assert codes.shape[1] == self.num_codebooks == 1 # only 1 semantic codebook for now + codes = codes.squeeze(1) # BxT + embedding = self.embedding.to(codes.device) + quantized = F.embedding(codes, embedding) + quantized = rearrange(quantized, "b t d -> b d t") + return quantized + + @property + def num_codebooks(self) -> int: + return 1 + + @property + def codebook_sizes(self) -> list[int]: + return [self.codebook_size] + + +class AcousticCodebook(nn.Module): + """Finite Scalar Quantization for acoustic codebooks.""" + + def __init__(self, codebook_size: int, codebook_dim: int) -> None: + super().__init__() + self.dim = codebook_dim + self.n_levels = codebook_size + self.num_codebooks = codebook_dim + + def _quantize(self, x: torch.Tensor, levels: torch.Tensor, ste: bool = True) -> torch.Tensor: + scaled_x = ((x + 1) / 2) * (levels - 1) + if ste: + quant_x = scaled_x + (scaled_x.round() - scaled_x).detach() + else: + quant_x = scaled_x.round() + return quant_x + + def _rescale(self, x: torch.Tensor, levels: int | torch.Tensor) -> torch.Tensor: + return (x * 2 / (levels - 1)) - 1 + + def _codes_from_quantized(self, quantized: torch.Tensor) -> torch.Tensor: + return quantized.long() + + def _quantized_from_codes(self, codes: torch.Tensor, levels: int) -> torch.Tensor: + return self._rescale(codes, levels) + + def encode(self, x: torch.Tensor) -> torch.Tensor: + assert x.dtype.is_floating_point, f"Input should be floats, got {x.dtype}" + x = torch.tanh(x) + levels = torch.ones_like(x) * self.n_levels + codes = self._codes_from_quantized(self._quantize(x, levels, ste=False)) + return codes + + def decode(self, codes: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor: + assert not codes.dtype.is_floating_point, f"Codes should be integers, got {codes.dtype}" + quantized = self._quantized_from_codes(codes, self.n_levels).to(dtype) + return quantized + + +class MistralAudioCodebook(nn.Module): + """Simplified audio codebook that combines semantic and acoustic quantization. + + This class replaces the complex SplitCodebook implementation with a simpler, + more direct approach that hard-codes the semantic and acoustic codebook types. + """ + + def __init__( + self, + audio_tokenizer_args: AudioTokenizerArgs, + ) -> None: + super().__init__() + + # Semantic codebook + self.semantic_codebook = SemanticCodebook( + codebook_size=audio_tokenizer_args.semantic_codebook_size, + codebook_dim=audio_tokenizer_args.semantic_dim, + ) + + # Acoustic codebook + self.acoustic_codebook = AcousticCodebook( + codebook_size=audio_tokenizer_args.acoustic_codebook_size, + codebook_dim=audio_tokenizer_args.acoustic_dim, + ) + + # Store dimensions + self.semantic_dim = audio_tokenizer_args.semantic_dim + self.acoustic_dim = audio_tokenizer_args.acoustic_dim + self.total_dim = self.semantic_dim + self.acoustic_dim + + @property + def num_codebooks(self) -> int: + """Total number of codebooks (semantic + acoustic).""" + return self.semantic_codebook.num_codebooks + self.acoustic_codebook.num_codebooks + + @property + def codebook_sizes(self) -> list[int]: + """List of sizes for all codebooks.""" + return self.semantic_codebook.codebook_sizes + self.acoustic_codebook.codebook_sizes + + def encode(self, x: torch.Tensor) -> torch.Tensor: + """Encode input tensor into discrete codes using both semantic and acoustic codebooks. + + Args: + x: Input tensor of shape [B, D, T] where D = semantic_dim + acoustic_dim + + Returns: + codes: Tensor of shape [B, K, T] where K = num_codebooks + """ + assert x.dtype.is_floating_point, f"Input should be floats, got {x.dtype}" + + # Split input into semantic and acoustic parts + semantic_part = x[:, : self.semantic_dim, :] + acoustic_part = x[:, self.semantic_dim :, :] + + # Quantize each part + semantic_codes = self.semantic_codebook.encode(semantic_part) + acoustic_codes = self.acoustic_codebook.encode(acoustic_part) + + # Combine codes along the codebook dimension + codes = torch.cat([semantic_codes, acoustic_codes], dim=1) + return codes + + def decode(self, codes: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Decode discrete codes back into continuous representations. + + Args: + codes: Input codes of shape [B, K, T] where K = num_codebooks + dtype: Target dtype for the output tensor + + Returns: + emb: Reconstructed tensor of shape [B, D, T] where D = semantic_dim + acoustic_dim + """ + assert not codes.dtype.is_floating_point, f"Codes should be integers, got {codes.dtype}" + + # Split codes into semantic and acoustic parts + semantic_codes = codes[:, : self.semantic_codebook.num_codebooks, :] + acoustic_codes = codes[:, self.semantic_codebook.num_codebooks :, :] + + # Decode each part + semantic_emb = self.semantic_codebook.decode(semantic_codes).to(dtype) # [B, semantic_dim, T] + acoustic_emb = self.acoustic_codebook.decode(acoustic_codes).to(dtype) # [B, acoustic_dim, T] + + # Combine embeddings along the dimension axis + emb = torch.cat([semantic_emb, acoustic_emb], dim=1) + return emb + + +### Encoder-decoder layers + + +def prepare_for_attention( + x: torch.Tensor, + time_last: bool = True, +) -> torch.Tensor: + if time_last: + x = rearrange(x, "b d t -> (b t) d") + else: + x = rearrange(x, "b t d -> (b t) d") + return x + + +def pad1d( + x: torch.Tensor, + paddings: tuple[int, int], + mode: str = "constant", + value: float = 0.0, +) -> torch.Tensor: + """ + Tiny wrapper around F.pad, just to allow for reflect padding + on small input. If this is the case, we insert extra 0 padding + to the right before the reflection happen. + """ + length = x.shape[-1] + padding_left, padding_right = paddings + assert padding_left >= 0 and padding_right >= 0, ( + padding_left, + padding_right, + ) + if mode == "reflect": + max_pad = max(padding_left, padding_right) + extra_pad = 0 + if length <= max_pad: + extra_pad = max_pad - length + 1 + x = F.pad(x, (0, extra_pad)) + padded = F.pad(x, paddings, mode, value) + end = padded.shape[-1] - extra_pad + return padded[..., :end] + else: + return F.pad(x, paddings, mode, value) + + +class CausalConv1d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + pad_mode: str = "reflect", + use_weight_norm: bool = True, + use_bias: bool = True, + ) -> None: + super().__init__() + conv = nn.Conv1d( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=0, + dilation=dilation, + bias=use_bias, + ) + self.conv = weight_norm(conv) if use_weight_norm else conv + self.pad_mode = pad_mode + self._stride = self.conv.stride[0] + self._effective_kernel_size = (kernel_size - 1) * self.conv.dilation[0] + 1 + self._padding_total = self._effective_kernel_size - self._stride + self.stride = self.conv.stride + + def forward(self, x: torch.Tensor) -> torch.Tensor: + n_frames = (x.shape[-1] - self._effective_kernel_size + self._padding_total) / self._stride + 1 + target_length = (math.ceil(n_frames) - 1) * self._stride + (self._effective_kernel_size - self._padding_total) + extra_padding = target_length - x.shape[-1] + x = pad1d(x, (self._padding_total, extra_padding), mode=self.pad_mode) + return self.conv(x) + + +class CausalConvTranspose1d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + groups: int = 1, + trim_ratio: float = 1.0, + use_weight_norm: bool = True, + use_bias: bool = True, + ) -> None: + super().__init__() + conv = nn.ConvTranspose1d( + in_channels, + out_channels, + kernel_size, + stride=stride, + groups=groups, + bias=use_bias, + ) + self.conv = weight_norm(conv) if use_weight_norm else conv + self.trim_ratio = trim_ratio + + def forward(self, x: torch.Tensor) -> torch.Tensor: + kernel_size = self.conv.kernel_size[0] + stride = self.conv.stride[0] + total_padding = kernel_size - stride + out = self.conv(x) + right_padding = math.ceil(total_padding * self.trim_ratio) + left_padding = total_padding - right_padding + return out[..., left_padding : out.shape[-1] - right_padding] + + +class MultiVocabEmbeddings(nn.Module): + def __init__( + self, + audio_model_args: dict, + embedding_dim: int, + ) -> None: + super().__init__() + self.model_args = from_nested_dict(MultimodalAudioModelArgs, audio_model_args) + # add one more for mask token (-1) + # also we only handle acoustic tokens here, + # so we skip the first codebook + self.codebook_sizes = [c for c in self.model_args.get_codebook_sizes(pad_to_multiple=None)] + self.offsets = torch.from_numpy(np.cumsum([0] + self.codebook_sizes[:-1])) + self.total_vocab_size = sum(self.codebook_sizes) + padded_size = 128 * ((self.total_vocab_size + 127) // 128) + self.embeddings = nn.Embedding( + padded_size, + embedding_dim, + ) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + # input_ids: BxCxL + self.offsets = self.offsets.to(input_ids.device) + input_ids = input_ids + self.offsets[torch.newaxis, :, torch.newaxis] + return self.embeddings(input_ids) + + +class Attention(nn.Module): + def __init__( + self, + args: AudioTokenizerArgs, + layer_id: int, + ) -> None: + super().__init__() + self.args = args + + self.n_local_heads: int = args.n_heads + self.n_local_kv_heads: int = args.n_kv_heads + self.repeats = self.n_local_heads // self.n_local_kv_heads + self.layer_id = layer_id + self.sliding_window = args.attn_sliding_window_size + + def get_alibi_slopes(n_heads: int) -> torch.Tensor: + def slopes_power_of_2(n: int) -> torch.Tensor: + # geometric sequence: 1, r, r^2, ..., r^(n-1) with r = 2^(-8/n) + r = 2.0 ** (-8.0 / n) + return torch.tensor([r**i for i in range(n)], dtype=torch.float32) + + if math.log2(n_heads).is_integer(): + slopes = slopes_power_of_2(n_heads) + else: + m = 2 ** math.floor(math.log2(n_heads)) # largest power of 2 < n_heads + slopes = torch.cat( + [ + slopes_power_of_2(m), + slopes_power_of_2(2 * m)[::2][: n_heads - m], + ] + ) + return slopes # shape [n_heads], values in (2^-8, 1] + + self.register_buffer( + "alibi_slopes", + get_alibi_slopes(self.n_local_heads), + persistent=False, + ) + + self.wq = nn.Linear( + args.dim, + args.n_heads * args.head_dim, + bias=False, + ) + self.wk = nn.Linear( + args.dim, + args.n_kv_heads * args.head_dim, + bias=False, + ) + self.wv = nn.Linear( + args.dim, + args.n_kv_heads * args.head_dim, + bias=False, + ) + self.wo = nn.Linear( + args.n_heads * args.head_dim, + args.dim, + bias=args.use_biases, + ) + + if args.qk_norm: + self.q_norm = rms_norm( + args.n_heads * args.head_dim, + eps=args.qk_norm_eps, + ) + self.k_norm = rms_norm( + args.n_kv_heads * args.head_dim, + eps=args.qk_norm_eps, + ) + + def _native_attention( + self, + xq: torch.Tensor, + xk: torch.Tensor, + xv: torch.Tensor, + ) -> torch.Tensor: + """SDPA fallback when flash_attn is not installed. + + Manually constructs the combined alibi + causal + sliding-window + bias and delegates to ``F.scaled_dot_product_attention``. + """ + B, S, H, D = xq.shape + Hkv = xk.shape[2] + + # (B, S, H, D) -> (B, H, S, D) + q = xq.transpose(1, 2) + k = xk.transpose(1, 2) + v = xv.transpose(1, 2) + + # Expand KV heads for GQA + if H != Hkv: + repeats = H // Hkv + k = k.repeat_interleave(repeats, dim=1) + v = v.repeat_interleave(repeats, dim=1) + + # Build attention bias: alibi + causal + sliding window + positions = torch.arange(S, device=xq.device) + # rel_pos[i, j] = j - i + rel_pos = positions.unsqueeze(0) - positions.unsqueeze(1) # (S, S) + + # Alibi bias: slope * (j - i), shape (H, S, S) + alibi_slopes = self.alibi_slopes.to(dtype=xq.dtype, device=xq.device) + attn_bias = alibi_slopes.view(H, 1, 1) * rel_pos.unsqueeze(0).to(xq.dtype) + + # Causal mask: mask out j > i (future positions) + if self.args.causal: + attn_bias = attn_bias.masked_fill(rel_pos.unsqueeze(0) > 0, float("-inf")) + + # Sliding window mask + window_left = self.sliding_window + window_right = 0 if self.args.causal else self.sliding_window + outside_window = (rel_pos < -window_left) | (rel_pos > window_right) + attn_bias = attn_bias.masked_fill(outside_window.unsqueeze(0), float("-inf")) + + output = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_bias.unsqueeze(0)) + # (B, H, S, D) -> (B, S, H, D) + return output.transpose(1, 2) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.dim() == 2: + bsz, (seqlen, _) = 1, x.shape + else: + bsz, seqlen, _ = x.shape + xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) + if self.args.qk_norm: + xq = self.q_norm(xq) + xk = self.k_norm(xk) + xq = xq.view(bsz, seqlen, self.n_local_heads, self.args.head_dim) + xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.args.head_dim) + xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.args.head_dim) + + if HAS_FLASH_ATTN: + alibi_slopes = self.alibi_slopes.to(torch.float32) + output = flash_attn_func( + xq, + xk, + xv, + causal=self.args.causal, + window_size=( + self.sliding_window, + 0 if self.args.causal else self.sliding_window, + ), + alibi_slopes=alibi_slopes, + ) + else: + output = self._native_attention(xq, xk, xv) + + output = output.view(bsz, seqlen, self.n_local_heads * self.args.head_dim) + return self.wo(output).squeeze(0) + + +class TransformerBlock(nn.Module): + def __init__( + self, + layer_id: int, + args: AudioTokenizerArgs, + ) -> None: + super().__init__() + self._layer_id = layer_id + self.n_heads = args.n_heads + self.dim = args.dim + self.attention = Attention(args, layer_id=layer_id) + + self.feed_forward = FeedForward( + dim=args.dim, + hidden_dim=args.hidden_dim, + use_biases=args.use_biases, + ) + self.attention_norm = rms_norm( + args.dim, + eps=args.norm_eps, + ) + self.ffn_norm = rms_norm( + args.dim, + eps=args.norm_eps, + ) + self.post_attention_norm: nn.Module | None + self.post_attention_norm = None + self.post_ffn_norm: nn.Module | None + self.post_ffn_norm = None + self.args = args + + self.layer_scale = args.layer_scale + if self.layer_scale: + if args.layer_scale_init is None: + if layer_id < 18: + init_scale = 0.1 + elif layer_id <= 24: + init_scale = 1e-5 + else: + init_scale = 1e-6 + else: + init_scale = args.layer_scale_init + self.attention_scale = nn.Parameter(torch.full((args.dim,), init_scale, requires_grad=True)) + self.ffn_scale = nn.Parameter(torch.full((args.dim,), init_scale, requires_grad=True)) + + @property + def layer_id(self) -> int: + return self._layer_id + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + r = self.attention.forward(self.attention_norm(x)) + if self.post_attention_norm is not None: + r = self.post_attention_norm(r) + if self.layer_scale: + r = self.attention_scale * r + h = x + r + r = self.feed_forward.forward(self.ffn_norm(h)) + if self.post_ffn_norm is not None: + r = self.post_ffn_norm(r) + if self.layer_scale: + r = self.ffn_scale * r + out = h + r + return out + + +class Transformer(nn.Module): + def __init__( + self, + args: AudioTokenizerArgs, + n_layers: int, + ) -> None: + super().__init__() + + self.args = args + self.n_layers = n_layers + + self.layers_ids: list[int] = list(range(n_layers)) + + self.layers = torch.nn.ModuleDict() + for layer_id in self.layers_ids: + block = TransformerBlock(layer_id=layer_id, args=args) + self.layers[str(layer_id)] = block + + assert len(self.layers) == len(self.layers_ids), ( + len(self.layers), + len(self.layers_ids), + ) + + def device(self) -> torch.device: + return next(self.parameters()).device + + def dtype(self) -> torch.dtype: + return next(self.parameters()).dtype + + def forward( + self, + input_ids: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + assert len(self.layers_ids) == len(self.layers), ( + len(self.layers_ids), + len(self.layers), + ) + h = input_ids + for layer_id in self.layers_ids: + layer = self.layers[str(layer_id)] + assert layer.layer_id == layer_id, (layer.layer_id, layer_id) + h = layer( + h, + ) + + return h + + +class VoxtralTTSAudioTokenizer(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config = vllm_config.model_config.hf_config + self.config = config + args = from_nested_dict(AudioTokenizerArgs, self.config.audio_config["codec_args"]) + self.args = args + + if not args.causal: + # causal mask is hard-coded in forward function + raise NotImplementedError + + self.patch_size = args.pretransform_patch_size + self.latent_dim = args.semantic_dim + args.acoustic_dim + + self.input_proj = CausalConv1d( + args.pretransform_patch_size * args.channels, + args.dim, + kernel_size=args.patch_proj_kernel_size, + use_weight_norm=args.conv_weight_norm, + use_bias=False, + ) + + ### Encoder + encoder_blocks: list[nn.Module] = [] + encoder_transformer_lengths = args.encoder_transformer_lengths + encoder_convs_kernels = args.encoder_convs_kernels + encoder_convs_strides = args.encoder_convs_strides + + assert encoder_transformer_lengths is not None + assert encoder_convs_kernels is not None + assert encoder_convs_strides is not None + + # window_size might change at different layer + cur_window_size = args.attn_sliding_window_size + + for idx, n_layers in enumerate(encoder_transformer_lengths): + # transformer + layer_args = deepcopy(args) + layer_args.attn_sliding_window_size = cur_window_size + assert layer_args.qk_norm, "qk_norm must be True for decoder" + encoder_transformer = Transformer( + args=layer_args, + n_layers=n_layers, + ) + encoder_blocks.append(encoder_transformer) + # projection + is_last_layer = idx == len(encoder_transformer_lengths) - 1 + proj_output_dim = self.latent_dim if is_last_layer else args.dim + if ( + (encoder_convs_kernels[idx] != 1) + or (encoder_convs_strides[idx] != 1) + or is_last_layer # (args.dim != proj_output_dim) + ): + encoder_blocks.append( + CausalConv1d( + args.dim, + proj_output_dim, + kernel_size=encoder_convs_kernels[idx], + stride=encoder_convs_strides[idx], + pad_mode="replicate", + use_bias=False, + ) + ) + if args.half_attn_window_upon_downsampling and (encoder_convs_strides[idx] > 1): + assert encoder_convs_strides[idx] == 2, "only supporting 2x downsampling" + cur_window_size = cur_window_size // 2 + assert cur_window_size >= 2 + + self.encoder_blocks = nn.ModuleList(encoder_blocks) + + ### Audio token lookup table for LLM + # placed here assuming we will never need to return tokens + # from encoder to client + self.audio_token_embedding = MultiVocabEmbeddings( + audio_model_args=self.config.audio_config["audio_model_args"], + embedding_dim=self.config.text_config.hidden_size, + ) + + ### Decoder + decoder_blocks: list[nn.Module] = [] + decoder_convs_kernels = args.decoder_convs_kernels + decoder_convs_strides = args.decoder_convs_strides + decoder_transformer_lengths = args.decoder_transformer_lengths + assert decoder_convs_kernels is not None + assert decoder_convs_strides is not None + assert decoder_transformer_lengths is not None + # first projection layer is necessary + decoder_blocks.append( + CausalConv1d( + self.latent_dim, + args.dim, + kernel_size=decoder_convs_kernels[0], + stride=decoder_convs_strides[0], + pad_mode="replicate", + use_bias=False, + ) + ) + if args.half_attn_window_upon_downsampling and (decoder_convs_strides[0] > 1): + assert decoder_convs_strides[0] == 2, "only supporting 2x upsampling" + cur_window_size = cur_window_size * 2 + + for idx, n_layers in enumerate(decoder_transformer_lengths): + # transformer + layer_args = deepcopy(args) + layer_args.attn_sliding_window_size = cur_window_size + decoder_transformer = Transformer( + args=layer_args, + n_layers=n_layers, + ) + decoder_blocks.append(decoder_transformer) + # projection + if (idx + 1 != len(decoder_transformer_lengths)) and ( + (decoder_convs_kernels[idx + 1] != 1) or (decoder_convs_strides[idx + 1] != 1) + ): + decoder_blocks.append( + CausalConvTranspose1d( + args.dim, + args.dim, + kernel_size=decoder_convs_kernels[idx + 1], + stride=decoder_convs_strides[idx + 1], + use_bias=False, + ) + ) + if args.half_attn_window_upon_downsampling and (decoder_convs_strides[idx + 1] > 1): + assert decoder_convs_strides[idx + 1] == 2, "only supporting 2x upsampling" + cur_window_size = cur_window_size * 2 + + self.decoder_blocks = nn.ModuleList(decoder_blocks) + + self.quantizer = MistralAudioCodebook(args) + + self.output_proj = CausalConv1d( + args.dim, + args.pretransform_patch_size, + kernel_size=args.patch_proj_kernel_size, + use_weight_norm=args.conv_weight_norm, + use_bias=False, + ) + + scale_factor = math.prod(encoder_convs_strides) + assert scale_factor == math.prod(decoder_convs_strides) + self._frame_rate = args.sampling_rate / (self.patch_size * scale_factor) + self._sampling_rate = args.sampling_rate + self._channels = args.channels + if self._channels != 1: + raise NotImplementedError + + # Encoder weight prefixes — these may be absent in open-source + # checkpoints that only ship decoder / quantizer / embedding weights. + self._encoder_weight_prefixes = ("input_proj.", "encoder_blocks.") + self._encoder_loaded = False + + @property + def channels(self) -> int: + return self._channels + + @property + def frame_rate(self) -> float: + return self._frame_rate + + @property + def sampling_rate(self) -> int: + return self._sampling_rate + + @property + def downsample_factor(self) -> int: + assert self._sampling_rate % self._frame_rate == 0 + return int(self._sampling_rate / self._frame_rate) + + @property + def num_codebooks(self) -> int: + """Active number of codebooks used by the quantizer.""" + return self.quantizer.num_codebooks + + @property + def codebook_sizes(self) -> list[int]: + """List of size of each codebook""" + return self.quantizer.codebook_sizes + + def load_weight(self, weight: tuple[str, torch.Tensor]) -> str: + params_dict = dict(self.named_parameters()) + name, loaded_weight = weight + if name not in params_dict: + if name == "quantizer.semantic_codebook.cluster_usage": + setattr(self.quantizer.semantic_codebook, "cluster_usage", loaded_weight) + elif name == "quantizer.semantic_codebook.embedding_sum": + setattr(self.quantizer.semantic_codebook, "embedding_sum", loaded_weight) + else: + raise KeyError(f"Weight {name} not found in model parameters.") + else: + param = params_dict[name] + + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + if name.startswith(self._encoder_weight_prefixes): + self._encoder_loaded = True + return name + + def _forward_encoder(self, x: torch.Tensor) -> torch.Tensor: + assert x.dim() == 3 + emb = rearrange( + x, + "b c (t h) -> b (c h) t", + h=self.patch_size, + ) + + emb = self.input_proj(emb) + emb = rearrange(emb, "b d t -> b t d").contiguous() + + for _, block in enumerate(self.encoder_blocks): + if type(block) is CausalConv1d: + emb = rearrange(emb, "b t d -> b d t") + emb = block(emb) + emb = rearrange(emb, "b d t -> b t d") + else: + bsz, _, _ = emb.shape + emb = prepare_for_attention(emb, False) # (b * t, d) + emb = block(emb) # (b * t, d) + emb = rearrange(emb, "(b t) d -> b t d", b=bsz) # (b, t, d) + + emb = rearrange(emb, "b t d -> b d t") + return emb + + def _tokenize_audio(self, x: torch.Tensor) -> torch.Tensor: + """Encode the given input tensor to quantized representation. + + Args: + x (torch.Tensor): Float tensor of shape [B, C, T] + n_q (int): number of codebooks, default each dimension is a codebook + + Returns: + codes (torch.Tensor): an int tensor of shape [B, K, T] + with K the number of codebooks used and T the timestep. + """ + if x.shape[-1] % self.patch_size != 0: + pad_length = self.patch_size - (x.shape[-1] % self.patch_size) + x = F.pad(x, (0, pad_length), mode="constant", value=0) + with torch.autocast(dtype=torch.bfloat16, device_type="cuda"): + # bf16 to use alibi bias in flash attn + emb = self._forward_encoder(x) # (b, d, t) + codes = self.quantizer.encode(emb) # (b, k, t) + + return codes + + def encode_tokens(self, x: list[torch.Tensor]) -> list[torch.Tensor]: + audio_embeddings = [] + + for audio_code in x: + # expected input shape: BxCBxL + # BxCBxLxD (B=1 for now w/o batching) + multi_codebook_embedding = self.audio_token_embedding(audio_code) + # BxLxD, TODO(@alexhliu): add assertion to make sure + # input_embedding_concat_type == sum + multi_codebook_embedding = multi_codebook_embedding.sum(dim=1) + # LxD + multi_codebook_embedding = multi_codebook_embedding.squeeze(0) + audio_embeddings.append(multi_codebook_embedding) + return audio_embeddings + + def encode_waveforms(self, x: list[torch.Tensor]) -> list[torch.Tensor]: + if not self._encoder_loaded: + raise RuntimeError( + "encode_waveforms requires encoder weights which are not available in the open-source checkpoint." + ) + audio_codes = [] + for waveform in x: + # TODO(@alexhliu): we can batch this for better performance + assert waveform.dim() == 1 + if torch.min(waveform) < -1.0: + logger.warning("Min value of input waveform signal is %s", torch.min(waveform)) + if torch.max(waveform) > 1.0: + logger.warning("Max value of input waveform signal is %s", torch.max(waveform)) + audio_code = self._tokenize_audio(waveform.unsqueeze(0).unsqueeze(0)) # BxVxL, B==1 + audio_code = audio_code + len( + AudioSpecialTokens.all_special_tokens() + ) # TODO(@alexhliu): do this properly somewhere else + # there is always EOA after waveform + B, V, _ = audio_code.shape + eoa = torch.zeros([B, V, 1], dtype=audio_code.dtype, device=audio_code.device) + eoa[:, 0, :] = 1 + audio_code = torch.cat([audio_code, eoa], dim=-1) + audio_codes.append(audio_code) + return self.encode_tokens(audio_codes) + + def _forward_decoder(self, emb: torch.Tensor) -> torch.Tensor: + emb = rearrange(emb, "b d t -> b t d").contiguous() + + for idx, block in enumerate(self.decoder_blocks): + if type(block) in [CausalConvTranspose1d, CausalConv1d]: + emb = rearrange(emb, "b t d -> b d t") + emb = block(emb) + emb = rearrange(emb, "b d t -> b t d") + else: + # This part changed for implementing batch inference + emb = block(emb) # (b, t, d) - supports batched attention via flash_attn + + emb = rearrange(emb, "b t d -> b d t") + emb = self.output_proj(emb) + + out = rearrange( + emb, + "b (c h) t -> b c (t h)", + h=self.patch_size, + ) + return out + + def decode(self, codes: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Decode the given codes to a reconstructed representation. + + Args: + codes (torch.Tensor): Int tensor of shape [B, K, T] + + Returns: + out (torch.Tensor): Float tensor of shape [B, C, T], + the reconstructed audio. + """ + emb = self.quantizer.decode(codes, dtype) # (b, k, t) + return self._forward_decoder(emb) + + def decode_helper_batch_async(self, codes_list: list[torch.Tensor]) -> list[torch.Tensor]: + """Batch decode a list of code tensors to reconstructed audio. + + Args: + codes_list: List of Int tensors, each of shape [T_i, K] where T_i + varies per request and K is codebooks (e.g. 37). + + Returns: + List of Float 1-D tensors of reconstructed audio waveform samples. + """ + chunk_size = 375 # TODO(chenyos): Hardcode. Fix it + + # Pre-process: find EOA and unshift tokens + processed = [] + for codes in codes_list: + eoa_mask = codes[:, 0] == 1 + eoa_indices = eoa_mask.nonzero(as_tuple=False) + cutting_point = eoa_indices[0].item() if len(eoa_indices) > 0 else len(codes) + audio_tokens = codes[:cutting_point] - 2 + processed.append(audio_tokens) + + # Separate empty and non-empty + results: list[torch.Tensor | None] = [None] * len(processed) + non_empty: list[tuple[int, torch.Tensor]] = [] + for idx, tokens in enumerate(processed): + if len(tokens) == 0: + results[idx] = torch.tensor([], dtype=torch.float32) + else: + non_empty.append((idx, tokens)) + + if not non_empty: + return results + + # Split all requests into chunks and collect + all_chunks: list[torch.Tensor] = [] + chunk_lengths: list[int] = [] + chunk_map: list[tuple[int, list[int]]] = [] + for orig_idx, tokens in non_empty: + req_chunk_indices = [] + for i in range(0, len(tokens), chunk_size): + chunk = tokens[i : i + chunk_size] + req_chunk_indices.append(len(all_chunks)) + chunk_lengths.append(len(chunk)) + all_chunks.append(chunk) + chunk_map.append((orig_idx, req_chunk_indices)) + + # Pad chunks to max length and batch decode + max_chunk_len = max(chunk_lengths) + K = all_chunks[0].shape[1] + padded = torch.zeros( + len(all_chunks), + max_chunk_len, + K, + dtype=all_chunks[0].dtype, + device=all_chunks[0].device, + ) + for i, chunk in enumerate(all_chunks): + padded[i, : len(chunk)] = chunk + + audio_codes = padded.to(device=torch.device("cuda")) # [B, T, K] + audio_values = self.decode(audio_codes.transpose(1, 2), dtype=torch.bfloat16) # [B, 1, T_out] + audio_values = audio_values.detach().cpu().float().squeeze(1) # [B, T_out] + if torch.min(audio_values) < -1.0: + logger.warning("Min value of decoded waveform signal is %s", torch.min(audio_values)) + if torch.max(audio_values) > 1.0: + logger.warning("Max value of decoded waveform signal is %s", torch.max(audio_values)) + + # Trim padding and reassemble per request + for orig_idx, chunk_indices in chunk_map: + audio_parts = [] + for ci in chunk_indices: + expected_samples = chunk_lengths[ci] * self.downsample_factor + audio_parts.append(audio_values[ci, :expected_samples]) + results[orig_idx] = torch.cat(audio_parts, dim=0) if len(audio_parts) > 1 else audio_parts[0] + + return results diff --git a/vllm_omni/model_executor/stage_configs/voxtral_tts.yaml b/vllm_omni/model_executor/stage_configs/voxtral_tts.yaml new file mode 100644 index 000000000..bb791e080 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/voxtral_tts.yaml @@ -0,0 +1,105 @@ +async_chunk: true +stage_args: + - stage_id: 0 + stage_type: llm # Use llm stage type to launch OmniLLM + runtime: + process: true # Run this stage in a separate process + devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device) + max_batch_size: 32 + engine_args: + model_stage: audio_generation + model_arch: VoxtralTTSForConditionalGeneration + worker_type: ar + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.8 + enforce_eager: false + trust_remote_code: true + async_scheduling: true + engine_output_type: latent + enable_prefix_caching: false + tokenizer_mode: mistral + config_format: mistral + load_format: mistral + skip_mm_profiling: true + enable_chunked_prefill: false + max_model_len: 4096 + custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.voxtral_tts.generator2tokenizer_async_chunk + output_connectors: + to_stage_1: connector_of_shared_memory + is_comprehension: true + final_output: false + 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 + stage_type: llm # Use llm stage type to launch OmniLLM + runtime: + process: true + devices: "0" + max_batch_size: 32 + engine_args: + model_stage: audio_tokenizer + model_arch: VoxtralTTSForConditionalGeneration + worker_type: generation + worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + async_scheduling: false + gpu_memory_utilization: 0.1 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + skip_mm_profiling: true + engine_output_type: audio + tokenizer_mode: mistral + config_format: mistral + load_format: mistral + max_num_batched_tokens: 65536 + max_model_len: 65536 + engine_input_source: [0] + is_comprehension: false + 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.9 + top_p: 0.8 + top_k: 40 + max_tokens: 2048 + seed: 42 + detokenize: True + repetition_penalty: 1.05 + +# 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 + + 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_chunk_frames_at_begin: 5 + codec_left_context_frames: 25 + + edges: + - from: 0 # language_model → acoustic_transformer: trigger only after receiving full input (-1) + to: 1 + window_size: -1 diff --git a/vllm_omni/model_executor/stage_configs/voxtral_tts_1773881731.yaml b/vllm_omni/model_executor/stage_configs/voxtral_tts_1773881731.yaml new file mode 100644 index 000000000..d004402b2 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/voxtral_tts_1773881731.yaml @@ -0,0 +1,46 @@ +async_chunk: true +stage_args: +- stage_id: 0 + stage_type: llm + runtime: {process: true, devices: '0', max_batch_size: 32} + engine_args: {model_stage: audio_generation, model_arch: VoxtralTTSForConditionalGeneration, + worker_type: ar, worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker, scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler, + gpu_memory_utilization: 0.8, enforce_eager: false, trust_remote_code: true, async_scheduling: true, + engine_output_type: latent, enable_prefix_caching: false, tokenizer_mode: mistral, + config_format: mistral, skip_mm_profiling: true, enable_chunked_prefill: false, + max_model_len: 4096, custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.voxtral_tts.generator2tokenizer_async_chunk} + output_connectors: {to_stage_1: connector_of_shared_memory} + is_comprehension: true + final_output: false + 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 + stage_type: llm + runtime: {process: true, devices: '0', max_batch_size: 32} + engine_args: {model_stage: audio_tokenizer, model_arch: VoxtralTTSForConditionalGeneration, + worker_type: generation, worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker, + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler, + async_scheduling: false, gpu_memory_utilization: 0.1, enforce_eager: true, trust_remote_code: true, + enable_prefix_caching: false, skip_mm_profiling: true, engine_output_type: audio, + tokenizer_mode: mistral, config_format: mistral, max_num_batched_tokens: 65536, + max_model_len: 65536} + engine_input_source: [0] + is_comprehension: false + 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.9, top_p: 0.8, top_k: 40, max_tokens: 2048, + seed: 42, detokenize: true, repetition_penalty: 1.05} +runtime: + enabled: true + defaults: {window_size: -1, max_inflight: 1} + 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_chunk_frames_at_begin: 5, codec_left_context_frames: 25} + edges: + - {from: 0, to: 1, window_size: -1} diff --git a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py new file mode 100644 index 000000000..2e71a080b --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py @@ -0,0 +1,110 @@ +from typing import Any + +import torch +from vllm.inputs import TextPrompt +from vllm.logger import init_logger + +from vllm_omni.inputs.data import OmniTokensPrompt + +logger = init_logger(__name__) + + +def generator2tokenizer( + stage_list, + engine_input_source, + prompt: OmniTokensPrompt | TextPrompt = None, + requires_multimodal_data: bool = False, +): + if not engine_input_source: + raise ValueError("engine_input_source cannot be empty") + source_stage_id = engine_input_source[0] + if source_stage_id >= len(stage_list): + raise IndexError(f"Invalid stage_id: {source_stage_id}") + if stage_list[source_stage_id].engine_outputs is None: + raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") + + generator_outputs = stage_list[source_stage_id].engine_outputs + tokenizer_inputs = [] + for generator_output in generator_outputs: + output = generator_output.outputs[0] + audio_tokens = torch.cat(output.multimodal_output["audio"], dim=-1).flatten().detach().cpu().tolist() + tokenizer_inputs.append( + OmniTokensPrompt( + prompt_token_ids=audio_tokens, + multi_modal_data=None, + mm_processor_kwargs=None, + ) + ) + + return tokenizer_inputs + + +def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None: + audio = pooling_output.get("audio") + if not isinstance(audio, torch.Tensor) or audio.numel() == 0: + return None + return audio.flatten() + + +def generator2tokenizer_async_chunk( + transfer_manager: Any, + pooling_output: dict[str, Any], + request: Any, + is_finished: bool = False, +) -> dict[str, Any] | None: + request_id = request.external_req_id + finished = bool(is_finished or request.is_finished()) + + if isinstance(pooling_output, dict): + frame = _extract_last_frame(pooling_output) + if frame is not None: + codec_codes = frame.cpu().tolist() + transfer_manager.code_prompt_token_ids[request_id].append(codec_codes) + elif not finished: + # Some steps may not produce pooling_output. Only flush on finish. + return None + + connector = getattr(transfer_manager, "connector", None) + raw_cfg = getattr(connector, "config", {}) or {} + cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} + chunk_size = int(cfg.get("codec_chunk_frames", 25)) + chunk_size_at_begin = int(cfg.get("codec_chunk_frames_at_begin", 5)) + left_context_size = int(cfg.get("codec_left_context_frames", 25)) + if chunk_size <= 0 or left_context_size < 0: + raise ValueError( + f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, " + f"codec_left_context_frames={left_context_size}" + ) + length = len(transfer_manager.code_prompt_token_ids[request_id]) + + # Avoid emitting empty chunks during normal streaming. If the request is + # finished and nothing was produced, emit an EOF marker. + if length <= 0: + if finished: + return { + "code_predictor_codes": [], + "finished": torch.tensor(True, dtype=torch.bool), + } + return None + + # Use a small chunk size at begin + if length <= chunk_size: + chunk_size = chunk_size_at_begin + + chunk_length = length % chunk_size + + if chunk_length != 0 and not finished: + return None + + context_length = chunk_length if chunk_length != 0 else chunk_size + end_index = min(length, left_context_size + context_length) + ctx_frames = max(0, int(end_index - context_length)) + window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:] + + # Pack context + chunk into codebook-major flat codes for adapter. + code_predictor_codes = torch.tensor(window_frames).reshape(-1).tolist() + + return { + "code_predictor_codes": [int(ctx_frames)] + [int(context_length)] + code_predictor_codes, + "finished": torch.tensor(finished, dtype=torch.bool), + } diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index ecab7af05..697c39d24 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -576,10 +576,15 @@ class GPUARModelRunner(OmniGPUModelRunner): if sub_dict: mm_cpu[k] = sub_dict elif isinstance(v, list): - element = v[0] - if isinstance(element, torch.Tensor): - element = element.detach().to("cpu").contiguous() - mm_cpu[k] = element + if len(v) == 0: + continue + cpu_list = [] + for elem in v: + if isinstance(elem, torch.Tensor): + cpu_list.append(elem.detach().to("cpu").contiguous()) + else: + cpu_list.append(elem) + mm_cpu[k] = cpu_list except Exception as e: logger.error(f"Error in merge multimodal outputs: {e}") @@ -598,6 +603,12 @@ class GPUARModelRunner(OmniGPUModelRunner): mm_payload[k] = v[start:end].contiguous() elif isinstance(v, dict): mm_payload[k] = {sk: sv[start:end].contiguous() for sk, sv in v.items()} + elif isinstance(v, list): + element = v[idx] if idx < len(v) else v[0] + # Clone tensors to avoid cross-request aliasing + if isinstance(element, torch.Tensor): + element = element.clone() + mm_payload[k] = element elif isinstance(v, torch.Tensor): # List-derived tensor payloads are request-invariant; clone to # avoid accidental cross-request aliasing on downstream mutation. diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index bea2f9a22..c8928a16c 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -92,9 +92,11 @@ class OmniGPUModelRunner(GPUModelRunner): # 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. + self.has_talker_mtp = False talker_mtp = getattr(self.model, "talker_mtp", None) if talker_mtp is not None: self.talker_mtp = talker_mtp # type: ignore[assignment] + self.has_talker_mtp = True cudagraph_mode = self.compilation_config.cudagraph_mode assert cudagraph_mode is not None # Only wrap talker_mtp in CUDAGraphWrapper for Omni models that @@ -805,7 +807,7 @@ class OmniGPUModelRunner(GPUModelRunner): slot_mapping=slot_mappings, ), ): - if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"): + if getattr(self.model, "talker", None) is not None and self.has_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] @@ -1049,7 +1051,9 @@ class OmniGPUModelRunner(GPUModelRunner): s, e = start_offset, start_offset + sched_tokens # only consider to store data into update dict. hidden_states_slice = hidden_states[s:e] - update_dict = self.model.postprocess(hidden_states_slice, **req_infos) + update_dict = self.model.postprocess( + hidden_states_slice, multimodal_outputs=multimodal_outputs, **req_infos + ) self._update_intermediate_buffer(req_id, update_dict) except Exception as e: logger.error( @@ -1276,7 +1280,7 @@ class OmniGPUModelRunner(GPUModelRunner): dtype=req_embeds.dtype, ) - if hasattr(self.model, "talker_mtp") and span_len == 1: + if self.has_talker_mtp and span_len == 1: last_talker_hidden, text_step = update_dict.pop("mtp_inputs") decode_slice = slice(len(decode_req_ids), len(decode_req_ids) + 1) self.talker_mtp_input_ids.gpu[decode_slice].copy_(req_input_ids) @@ -1295,7 +1299,7 @@ class OmniGPUModelRunner(GPUModelRunner): input_ids[s : s + seg_len] = req_input_ids # run talker mtp decode - if hasattr(self.model, "talker_mtp"): + if self.has_talker_mtp: self._talker_mtp_forward(decode_req_ids, inputs_embeds) return ( @@ -1362,7 +1366,7 @@ class OmniGPUModelRunner(GPUModelRunner): **model_kwargs_extra, ) 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) + model_output = self.model.make_omni_output(model_output, **model_kwargs, **model_kwargs_extra) # Cache model output so later sample_tokens can consume multimodal results. self._omni_last_model_output = model_output return model_output