[zhangj1an]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Latest Status [8 May ]

Will mark as ready for review by this week after finish testing the last 3 models, also write a more detailed PR description

Purpose

Closes #3131.

Currently, there are 6 TTS/omni models' audio encoder that chose their own attention instead of relying on our existing flash attention backend in vllm_omni/diffusion/attention/backends/utils/fa.py. Also the order in fa is a bit messy as pointed out by a code comment in #3302.

This PR updated the priority order for attention in fa.py to be as follows:

1. use vllm.vllm_flash_attn. This method will run FA3 on Hopper, and FA2 on other GPU architecture.
2. prebuilt FA3 wheel (fa3_fwd_interface)
3. source build FA3 (flash_attn_interface),
4. source build FA2 (flash_attn).

For each model, we followed the Ming-flash-omni-2.0 refactor in #3146, and route attention selection via fa.py dispatch. We mainly updated function signatures to match with fa.py. The following models are affected: mimo_audio, voxtral_tts, ming_flash_omni, qwen2_5_omni, qwen3_omni and qwen3_tts. We tested each model and show the accuracy stays the same on this branch.

Also, voxtral tts hit the same #3327 error. It was using a method (flash_attn_func) available in ROCm and MUSA but not on CUDA, so we added the same fix by adding an if-else branch that uses flash_attn_varlen_func (available on CUDA) when architecture is CUDA.

Test Plan

We firstly ran pytests for each of the 6 models, and then tested their accuracy.

Env configs

All runs were executed against the top of this branch, on 2 RTX PRO 6000 Blackwell (sm_120, 96 GB)

# Run with flash-attn UNINSTALLED so fa.py exercises vllm.vllm_flash_attn
pip uninstall -y flash-attn flash_attn_interface fa3_fwd_interface || true

python -c "
from vllm_omni.diffusion.attention.backends.utils.fa import HAS_FLASH_ATTN, flash_attn_varlen_func
print('HAS_FLASH_ATTN=', HAS_FLASH_ATTN, 'fn=', flash_attn_varlen_func.__module__)
"
# Expected: HAS_FLASH_ATTN=True, fn=vllm.vllm_flash_attn.flash_attn_interface

# Env required for the runs to reach the model under test
export VLLM_USE_FLASHINFER_SAMPLER=0       # flashinfer JIT path mis-detects Blackwell
export TORCH_CUDA_ARCH_LIST="12.0+PTX"     # so JIT compiles target sm_120
export LD_LIBRARY_PATH="/usr/local/lib/python3.12/dist-packages/nvidia/cu13/lib:$LD_LIBRARY_PATH"
                                            # libnvrtc-builtins.so.13.0 lives in the nvidia/cu13 wheel
                                            # (driver is CUDA 13, system toolkit is 12.8)

e2e commands to run each model

mimo_audio

Default model: XiaomiMiMo/MiMo-Audio-7B-Instruct.

export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
OUT="$FA_TEST_ROOT/migrated/mimo_audio"; mkdir -p "$OUT"
python3 -u examples/offline_inference/mimo_audio/end2end.py \
  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
  --query-type tts_sft --output-dir "$OUT/tts_sft" 2>&1 | tee "$OUT/tts_sft.log"
python3 -u examples/offline_inference/mimo_audio/end2end.py \
  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
  --query-type audio_understanding_sft \
  --text "Summarize the audio." \
  --audio-path examples/offline_inference/mimo_audio/spoken_dialogue_assistant_turn_1.wav \
  --output-dir "$OUT/audio_understanding" 2>&1 | tee "$OUT/audio_understanding.log"

voxtral_tts

Default model: mistralai/Voxtral-4B-TTS-2603.

OUT="$FA_TEST_ROOT/migrated/voxtral_tts"; mkdir -p "$OUT"
python3 examples/offline_inference/voxtral_tts/end2end.py \
  --write-audio --voice cheerful_female \
  --model mistralai/Voxtral-4B-TTS-2603 \
  --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?" \
  --output-dir "$OUT/blocking_single" 2>&1 | tee "$OUT/blocking_single.log"
python3 examples/offline_inference/voxtral_tts/end2end.py \
  --streaming --write-audio --voice neutral_female \
  --model mistralai/Voxtral-4B-TTS-2603 \
  --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?" \
  --output-dir "$OUT/streaming_single" 2>&1 | tee "$OUT/streaming_single.log"

qwen2_5_omni

Default model: Qwen/Qwen2.5-Omni-7B.

OUT="$FA_TEST_ROOT/migrated/qwen2_5_omni"; mkdir -p "$OUT"
python examples/offline_inference/qwen2_5_omni/end2end.py \
  --output-wav "$OUT/use_audio" --query-type use_audio 2>&1 | tee "$OUT/use_audio.log"
python examples/offline_inference/qwen2_5_omni/end2end.py \
  --output-wav "$OUT/use_mixed_modalities" --query-type use_mixed_modalities 2>&1 | tee "$OUT/mix.log"

qwen3_omni

Default model: Qwen/Qwen3-Omni-30B-A3B-Instruct.

OUT="$FA_TEST_ROOT/migrated/qwen3_omni"; mkdir -p "$OUT"
python examples/offline_inference/qwen3_omni/end2end.py \
  --output-wav "$OUT/use_audio" --query-type use_audio 2>&1 | tee "$OUT/use_audio.log"
python examples/offline_inference/qwen3_omni/end2end_async_chunk.py \
  --output-wav "$OUT/async_chunk" --query-type use_audio 2>&1 | tee "$OUT/async_chunk.log"

ming_flash_omni

Default model: Jonathan1909/Ming-flash-omni-2.0.

OUT="$FA_TEST_ROOT/migrated/ming_flash_omni"; mkdir -p "$OUT"
python examples/offline_inference/ming_flash_omni/end2end.py \
  --deploy-config vllm_omni/deploy/ming_flash_omni_thinker_only.yaml \
  --query-type use_audio 2>&1 | tee "$OUT/thinker_use_audio.log"
python examples/offline_inference/ming_flash_omni/end2end.py \
  --query-type use_audio 2>&1 | tee "$OUT/omni_speech_use_audio.log"

qwen3_tts

Default model: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice (and -VoiceDesign / -Base variants).

OUT="$FA_TEST_ROOT/migrated/qwen3_tts"; mkdir -p "$OUT"
cd examples/offline_inference/qwen3_tts
python end2end.py --query-type CustomVoice  --output-dir "$OUT/CustomVoice"  2>&1 | tee "$OUT/CustomVoice.log"
python end2end.py --query-type VoiceDesign  --output-dir "$OUT/VoiceDesign"  2>&1 | tee "$OUT/VoiceDesign.log"
python end2end.py --query-type Base         --output-dir "$OUT/Base"         2>&1 | tee "$OUT/Base.log"

Test Result

All pytests for the 6 models still pass.

Details

qwen3_tts — Qwen/Qwen3-TTS-12Hz-0.6B-Base and -1.7B-CustomVoice

$ pytest tests/e2e/offline_inference/test_qwen3_tts_base.py -v
tests/e2e/offline_inference/test_qwen3_tts_base.py::test_text_to_audio_001[no_cuda_graph] PASSED [100%]
================== 1 passed, 20 warnings in 131.07s (0:02:11) ==================

$ pytest tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -v
tests/e2e/offline_inference/test_qwen3_tts_customvoice.py::test_text_to_audio_001[no_cuda_graph] PASSED [100%]
======================= 1 passed, 22 warnings in 48.32s ========================

voxtral_tts — mistralai/Voxtral-4B-TTS-2603

$ pytest tests/e2e/offline_inference/test_voxtral_tts.py -v
tests/e2e/offline_inference/test_voxtral_tts.py::test_voxtral_tts_offline_basic PASSED [ 50%]
tests/e2e/offline_inference/test_voxtral_tts.py::test_voxtral_tts_offline_streaming PASSED [100%]
================== 2 passed, 38 warnings in 115.46s (0:01:55) ==================

$ pytest tests/e2e/online_serving/test_voxtral_tts.py -v
tests/e2e/online_serving/test_voxtral_tts.py::TestVoxtralTTSFixedVoice::test_speech_english_basic[omni_server0]               PASSED [ 14%]
tests/e2e/online_serving/test_voxtral_tts.py::TestVoxtralTTSFixedVoice::test_speech_english_streaming[omni_server0]           PASSED [ 28%]
tests/e2e/online_serving/test_voxtral_tts.py::TestVoxtralTTSFixedVoice::test_speech_different_voices[omni_server0]            PASSED [ 42%]
tests/e2e/online_serving/test_voxtral_tts.py::TestVoxtralTTSFixedVoice::test_speech_invalid_voice_rejected[omni_server0]      PASSED [ 57%]
tests/e2e/online_serving/test_voxtral_tts.py::TestVoxtralTTSFixedVoice::test_speech_binary_response_not_utf8_error[omni_server0] PASSED [ 71%]
tests/e2e/online_serving/test_voxtral_tts.py::TestVoxtralTTSAPIEndpoints::test_list_voices_endpoint[omni_server0]              PASSED [ 85%]
tests/e2e/online_serving/test_voxtral_tts.py::TestVoxtralTTSAPIEndpoints::test_models_endpoint[omni_server0]                   PASSED [100%]
================== 7 passed, 16 warnings in 85.18s (0:01:25) ===================

mimo_audio — XiaomiMiMo/MiMo-Audio-7B-Instruct

$ pytest tests/e2e/online_serving/test_mimo_audio.py -v
tests/e2e/online_serving/test_mimo_audio.py::test_text_to_text_001[omni_server0] PASSED [100%]
============= 1 passed, 1 skipped, 16 warnings in 94.36s (0:01:34) =============

qwen2_5_omni — Qwen/Qwen2.5-Omni-7B

still pending

ming_flash_omni — Jonathan1909/Ming-flash-omni-2.0

still pending

qwen3_omni — Qwen/Qwen3-Omni-30B-A3B-Instruct

still pending

The end 2 end test for each model still works, and we attach the generated audio files.

Model	config	Input prompt	Main Branch's Output	This PR's Output	Whisper-medium transcript	WER / CER
qwen3_tts	Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice, --query-type CustomVoice	其实我真的有发现，我是一个特别善于观察别人情绪的人。 (examples/offline_inference/qwen3_tts/end2end.py:get_customvoice_query)	qwen3_tts_main.wav	qwen3_tts_branch.wav	其实我真的有发现我是一个特别善于观察别人情绪的人	CER 0.0000 ✅
voxtral_tts	mistralai/Voxtral-4B-TTS-2603, --voice cheerful_female, blocking path	That eerie silence after the first storm was just the calm before another round of chaos, wasn't it? (examples/offline_inference/voxtral_tts/end2end.py)	voxtral_tts_main.wav	voxtral_branch.wav	That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?	WER 0.0000 ✅
mimo_audio	XiaomiMiMo/MiMo-Audio-7B-Instruct, --query-type tts_sft --text "The weather is so nice today."	Please convert this text to speech: The weather is so nice today.	mimo_main.wav	mimo_branch.wav	The weather is S O O D	WER 0.6667, the main branch model also does not work as reported in #3452.
qwen2_5_omni	Qwen/Qwen2.5-Omni-7B, --query-type use_audio	n/a — deferred	n/a	n/a	n/a	n/a — CI deploy YAML hard-codes 3 GPUs (stage 0→cuda:0, 1→cuda:1, 2→cuda:2); to be run on rental hardware per §8.7
ming_flash_omni	Jonathan1909/Ming-flash-omni-2.0, --query-type use_audio	n/a — deferred	n/a	n/a	n/a	n/a — weights ~64 GB, will not fit on 80 GB system disk alongside other test models
qwen3_omni	Qwen/Qwen3-Omni-30B-A3B-Instruct, --query-type use_audio	n/a — deferred	n/a	n/a	n/a	n/a — 30B params (~60 GB bf16), multi-GPU per deploy YAML

Details

For mimo-audio, the main branch does not work,

The accuracy for each model compared to the main branch remains the same.

Details

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
• (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[hsliuustc0106]

Draft PR - ready for full review when WIP status removed.

[hsliuustc0106]

please pay attention to the accuracy

[BeatSeat]

Any updates here?

[zhangj1an]

Any updates here?

I will try finish this PR by tomorrow!

[dtrifiro]

We should also remove fa3_fwd_config from requirements.txt

[lishunyang12]

Since this pr has been released over 3 months, I will close it if there is no further update.

[zhangj1an]

I will not be updating this PR anymore so will close it. I am happy for future contributors to continue from there