anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Merge branch 'main' into dev/vllm-align

Browse Source
This commit is contained in:
tzhouam committed 2026-03-20 01:38:43 +00:00
2 parents 367ce6328c 6ac7ba692d
commit 3c52432100
55 files changed
+7222 -100
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.buildkite/bootstrap-amd-omni.sh
+9 -1
View File
@@ -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" \
.buildkite/scripts/hardware_ci/run-amd-test.sh
+8 -3
View File
@@ -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 \
.buildkite/test-amd-merge.yml
+197
View File
@@ -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"
.buildkite/test-amd-ready.yaml
+178
View File
@@ -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"
.buildkite/test-merge.yml
+56
View File
@@ -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"
.buildkite/test-nightly.yml
+3
View File
@@ -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"
.buildkite/test-ready.yml
+39
View File
@@ -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:
.buildkite/test-template-amd-omni.j2
+6 -1
View File
@@ -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
docs/.nav.yml
+1
View File
@@ -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
docs/mkdocs/hooks/generate_api_readme.py
+1
View File
@@ -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
docs/models/supported_models.md
+1
View File
@@ -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
docs/serving/speech_api.md
+15
View File
@@ -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 <model> --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
docs/user_guide/diffusion/quantization/int8.md
+75
View File
@@ -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="<your-model>", quantization="int8")
# Skip sensitive layers
omni = Omni(
model="<your-model>",
quantization_config={
"method": "int8",
"ignored_layers": ["<layer-name>"],
},
)
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 <your-model> --quantization int8
# Skip sensitive layers
python text_to_image.py --model <your-model> --quantization int8 --ignored-layers "img_mlp"
# Online serving
vllm serve <your-model> --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="<your-model>",
quantization="int8",
cache_backend="tea_cache",
cache_config={"rel_l1_thresh": 0.2},
)
```
docs/user_guide/diffusion/quantization/overview.md
+9 -1
View File
@@ -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 |
docs/user_guide/diffusion_acceleration.md
+25 -7
View File
@@ -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="<your-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.
examples/offline_inference/text_to_image/text_to_image.py
+4 -6
View File
@@ -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",
examples/offline_inference/voxtral_tts/README.md
+57
View File
@@ -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 |
examples/offline_inference/voxtral_tts/end2end.py
+377
View File
@@ -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)
examples/online_serving/voxtral_tts/gradio_demo.py
+255
View File
@@ -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,
)
tests/benchmarks/test_serve_cli.py
+3 -1
View File
@@ -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 = [
tests/conftest.py
-2
View File
@@ -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]
tests/diffusion/models/t5_encoder/test_t5_encoder_tp.py
+161
View File
@@ -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
tests/diffusion/quantization/test_int8_config.py
+463
View File
@@ -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
tests/e2e/offline_inference/test_bagel_img2img.py
+18 -3
View File
@@ -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()
tests/e2e/offline_inference/test_bagel_text2img.py
+16 -1
View File
@@ -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")
tests/e2e/online_serving/test_bagel_online.py
+1 -1
View File
@@ -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."""
tests/e2e/online_serving/test_qwen3_omni.py
+2 -6
View File
@@ -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
tests/e2e/online_serving/test_voxtral_tts.py
+192
View File
@@ -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
tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml
+8 -6
View File
@@ -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
tests/model_executor/models/voxtral_tts/__init__.py
View File
Whitespace-only changes.
tests/model_executor/models/voxtral_tts/test_audio_tokenizer_parsing.py
+161
View File
@@ -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
tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
+329
View File
@@ -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)
tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py
+281
View File
@@ -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
vllm_omni/config/stage_config.py
+1
View File
@@ -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",
vllm_omni/diffusion/models/flux/pipeline_flux.py
+13 -6
View File
@@ -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)
vllm_omni/diffusion/models/t5_encoder/__init__.py
+7
View File
@@ -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"]
vllm_omni/diffusion/models/t5_encoder/t5_encoder.py
+389
View File
@@ -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
vllm_omni/diffusion/models/z_image/z_image_transformer.py
+6 -3
View File
@@ -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
vllm_omni/diffusion/quantization/__init__.py
+3
View File
@@ -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",
vllm_omni/diffusion/quantization/int8.py
+475
View File
@@ -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,
)
vllm_omni/engine/arg_utils.py
+13 -6
View File
@@ -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():
vllm_omni/engine/async_omni_engine.py
+1
View File
@@ -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,
vllm_omni/entrypoints/openai/serving_chat.py
+9 -3
View File
@@ -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",
vllm_omni/entrypoints/openai/serving_speech.py
+133 -34
View File
@@ -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,
vllm_omni/model_executor/models/registry.py
+12
View File
@@ -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"),
}
vllm_omni/model_executor/models/voxtral_tts/configuration_voxtral_tts.py
+99
View File
@@ -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
vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py
+256
View File
@@ -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
vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py
+372
View File
@@ -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, <flat audio tokens>]
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
vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py
+1083
View File
File diff suppressed because it is too large. Load diff
vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py
+1114
View File
File diff suppressed because it is too large. Load diff
vllm_omni/model_executor/stage_configs/voxtral_tts.yaml
+105
View File
@@ -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
vllm_omni/model_executor/stage_configs/voxtral_tts_1773881731.yaml
+46
View File
@@ -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}
vllm_omni/model_executor/stage_input_processors/voxtral_tts.py
+110
View File
@@ -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),
}
vllm_omni/worker/gpu_ar_model_runner.py
+15 -4
View File
@@ -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.
vllm_omni/worker/gpu_model_runner.py
+9 -5
View File
@@ -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