diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 6105513c3ee..3dbc65f4031 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -40,7 +40,10 @@ ENV VLLM_OMNI_TARGET_DEVICE=xpu # Keep the upstream oneCCL, torch and Triton packages paired. The final # vLLM XPU wheel requires triton==3.8.0+xpu; installing the old triton-xpu # overlay would replace that runtime with an incompatible version. +# The kernels egg_info hook imports kernels_data. Install that runtime before +# building Omni without isolation in the shared environment. RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install "kernels-data>=0.14.0.dev1" && \ uv pip install --no-build-isolation --excludes <(printf "torch\ntriton-xpu\n") \ ".[dev]" diff --git a/docs/design/feature/prefix_caching.md b/docs/design/feature/prefix_caching.md index 1c9009f0cb0..9a4d97ad8a8 100644 --- a/docs/design/feature/prefix_caching.md +++ b/docs/design/feature/prefix_caching.md @@ -290,7 +290,8 @@ Hit spans come from adapter-classified `STARTED` events for new requests: `EXTENDED`, so the manager does not mark a hit: those rows were delivered in earlier steps and re-emitting them would duplicate output. A `delivered_upto` span for this case is Phase 2. -- Preemption + reschedule: vLLM resets `num_computed_tokens` to 0 on +- Preemption + reschedule: for stages with `recompute_preemption="allow"` + (the default), vLLM resets `num_computed_tokens` to 0 on preemption and re-runs prefix matching on resume, so the resumed request can come back with a fresh hit. With the V1 model runner it arrives through `scheduled_cached_reqs` (id in `resumed_req_ids`, `new_block_ids` @@ -303,6 +304,22 @@ Hit spans come from adapter-classified `STARTED` events for new requests: hidden states should be sized so preemption does not occur while prefix caching is on. Same as before this refactor; tracked for Phase 2. +Stages whose decode state cannot be reconstructed from token IDs declare +`recompute_preemption="fail"` in their AR pipeline topology. Capacity +preemption then finishes the request with one terminal ERROR and releases its +request/KV state instead of requeueing it, including when a KV-delivery +connector asks to discard stale output. Explicit `reset_prefix_cache` retains +its existing resume behavior, subject to `supports_running_prefix_cache_reset`; +its reset scope, rather than the shared `drop_stale_output` flag, identifies +that exception. Fail-mode preemption also notifies the auxiliary-output +connector before removing the request, so the worker receives its finished +notification and subsequent metadata assembly cannot reference a freed request. + +Fish Omni retains the default `allow` policy: its MRv2 model state saves the +generated codebooks before a slot is reused and restores full frame rows on +resume. Missing codebooks in the token stream alone do not imply that a stage +cannot replay; model-state replay support must also be considered. + The adapter is the sole owner of started-request observations. Its single `translate_step` entry point captures an immutable `PrefixCacheStep`. `EXTENDED` is represented by `step.extended_req_ids`, with token counts in diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index f5700a318f9..d7c738ce75f 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -978,7 +978,22 @@ See the [Gepard section of the online TTS hub](../user_guide/examples/online_ser | Model | Description | | ------- | ------------- | -| `k2-fsa/OmniVoice` | Pure-diffusion TTS. Supports voice cloning via `ref_audio` (with optional `ref_text`); no built-in voice presets. | +| `k2-fsa/OmniVoice` | OmniVoice generates mono speech at 24 kHz. It supports text-to-speech without a reference clip, voice cloning with `ref_audio` and optional `ref_text`, and voice design. It has no built-in voice presets. | + +For voice cloning, send `ref_audio`. If `ref_text` is missing or blank, the default deploy config leaves Whisper unloaded until the first such request, then loads `openai/whisper-large-v3-turbo` on the worker device and transcribes the clip. Set `additional_config.omnivoice_asr.load_asr_on_startup` to `true` to load it during worker startup. Set `additional_config.omnivoice_asr.asr_device` to `cuda:0` or `cpu` to override the worker device. Sending `ref_text` avoids ASR and uses the supplied transcript. + +OmniVoice prepares the reference clip before it builds the voice cloning input. It performs these steps: + +- It converts the clip to mono, resamples it to 24 kHz, and measures its original root mean square (RMS) level. +- When the original RMS is greater than zero and less than `0.1`, it raises the waveform level to `0.1` and keeps the original RMS for output volume control. +- When `ref_text` is missing, it trims a clip longer than 20 seconds at a detected silence. Supplied text skips this special long-audio trim. +- It removes middle silence and trims silence at both edges whether the transcript is automatic or supplied by the caller. +- It aligns the sample count down to the audio tokenizer hop size. +- When `ref_text` is missing, it sends the prepared waveform to Whisper. +- It adds a final period for English text or a Chinese full stop for Chinese text when the reference text has no ending punctuation. +- It sends the same prepared waveform to the audio tokenizer. + +OmniVoice processes voice-cloning audio after decoding. It removes middle gaps of at least 500 milliseconds. It trims edge silence while keeping 100 milliseconds at each edge. It scales cloned output using the original reference RMS when that value is below `0.1`, and leaves the level unchanged when the value is at least `0.1`. It fades the first and last 100 milliseconds, then adds 100 milliseconds of silence at each edge. Text-only output returns the decoder tensor without this voice-cloning postprocessing. OmniVoice uses packed variable-length attention for batched generator execution. The attention operator accepts FP16 and BF16 inputs, so its query, key, and value tensors are evaluated in BF16 even when the stage is configured with `dtype: float32`; the attention output is converted back to the model's diff --git a/docs/user_guide/examples/offline_inference/text_to_speech.md b/docs/user_guide/examples/offline_inference/text_to_speech.md index cee7c4a321e..940a6b05ca4 100644 --- a/docs/user_guide/examples/offline_inference/text_to_speech.md +++ b/docs/user_guide/examples/offline_inference/text_to_speech.md @@ -36,7 +36,7 @@ python examples/offline_inference/text_to_speech//end2end.py \ --ref-text "Transcript of the reference audio." ``` -`--ref-audio` and `--ref-text` are optional (text-only synthesis works without them) and must be provided together for voice cloning. The exotic scripts — Qwen3-TTS, Voxtral TTS, CosyVoice3 — accept additional model-specific flags documented in their per-model section below. +For most models, `--ref-audio` and `--ref-text` are optional for text-only synthesis but must be provided together for voice cloning. OmniVoice also accepts `--ref-audio` without `--ref-text` and transcribes the reference lazily with Whisper. Qwen3-TTS, Voxtral TTS, and CosyVoice3 accept additional model-specific flags documented in their sections below. --- @@ -252,7 +252,7 @@ Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 ## OmniVoice -Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design). +OmniVoice creates speech in more than 600 languages. It supports text-to-speech without a reference clip, voice cloning, and voice design. ### Prerequisites @@ -260,7 +260,7 @@ Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / c huggingface-cli download k2-fsa/OmniVoice ``` -Voice cloning requires `transformers>=5.3.0`. Auto and design modes work with `transformers>=4.57.0`. +Voice cloning requires `transformers>=5.3.0`. Automatic voice and voice design use `transformers>=4.57.0`. ### Quick start (auto voice) @@ -272,6 +272,8 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ ### Voice cloning +With an explicit transcript, voice cloning does not load Whisper: + ```bash python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --model k2-fsa/OmniVoice \ @@ -280,6 +282,30 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --ref-text "This is the reference transcription." ``` +If `--ref-text` is omitted, OmniVoice loads `openai/whisper-large-v3-turbo` on the worker device and transcribes the reference audio after it prepares the waveform. The first request downloads and loads Whisper, so it takes longer and uses more device memory. Set `additional_config.omnivoice_asr.load_asr_on_startup` to `true` to load Whisper during worker startup, or set `additional_config.omnivoice_asr.asr_device` to use another GPU or the CPU. + +```bash +python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ + --model k2-fsa/OmniVoice \ + --text "hello" \ + --ref-audio trump_ref.wav +``` + +### Reference audio processing + +OmniVoice prepares the reference clip before it transcribes the clip or builds the voice cloning input. It performs these steps: + +- It converts the clip to mono and resamples it to 24 kHz. +- It measures the original root mean square (RMS) level. When the RMS is greater than zero and below `0.1`, it raises the waveform level to `0.1` and keeps the original RMS for output volume control. +- When `--ref-text` is omitted, it trims a clip longer than 20 seconds at a detected silence. +- It removes middle silence and trims silence at both edges whether the transcript is automatic or supplied by the caller. +- It aligns the sample count down to the audio tokenizer hop size. +- When `--ref-text` is omitted, it sends the prepared waveform to Whisper. +- It adds a final period for English text or a Chinese full stop for Chinese text when the reference text has no ending punctuation. +- It sends the same prepared waveform to the audio tokenizer and uses the resulting reference tokens to estimate the target duration. + +OmniVoice processes voice cloning audio after decoding. It removes middle gaps of at least 500 milliseconds and trims edge silence while keeping 100 milliseconds at each edge. It uses the original reference RMS when the RMS is below `0.1`, and it leaves the level unchanged when the RMS is at least `0.1`. It fades the first and last 100 milliseconds, then adds 100 milliseconds of silence at both edges. Text only output returns the decoder tensor without this processing. + ### Voice design ```bash @@ -302,6 +328,11 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ - Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking. - Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz. +- By default, the first request with reference audio and no `--ref-text` downloads + and loads `openai/whisper-large-v3-turbo`, so it has extra latency and GPU + memory use. Set `additional_config.omnivoice_asr.load_asr_on_startup` to `true` + to move that load to worker startup. +- Generated audio remains mono at 24 kHz. --- diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index 6e3e9354302..80ee5dd66f3 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -448,7 +448,7 @@ bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh ## OmniVoice -Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline. +OmniVoice creates speech in more than 600 languages. It supports text-to-speech without a reference clip, voice cloning, and voice design. ### Prerequisites @@ -456,7 +456,7 @@ Zero-shot multilingual TTS (600+ languages). Online serving currently exposes ** huggingface-cli download k2-fsa/OmniVoice ``` -Voice cloning (offline) needs `transformers>=5.3.0`; auto voice works with `transformers>=4.57.0`. +Voice cloning needs `transformers>=5.3.0`. When `ref_text` is omitted, OmniVoice loads `openai/whisper-large-v3-turbo` on the worker device to transcribe the reference audio. Set `additional_config.omnivoice_asr.load_asr_on_startup` to `true` to load Whisper during worker startup, or set `additional_config.omnivoice_asr.asr_device` to use another GPU or the CPU. ### Launch @@ -472,13 +472,35 @@ vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code cd examples/online_serving/text_to_speech/omnivoice python speech_client.py --text "Hello, how are you?" python speech_client.py --text "Bonjour, comment allez-vous?" --language French +python speech_client.py --text "hello" --ref-audio /path/to/ref_audio.wav ``` -The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--output`. +The client supports `--api-base`, `--model`, `--text`, `--response-format`, +`--language`, `--voice`, `--ref-audio`, `--ref-text`, `--instructions`, +`--seed`, and `--output`. ### Notes -- Voice cloning and voice design require offline inference; see the [offline OmniVoice section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#omnivoice). +- The explicit `--ref-text` path is faster because it does not load Whisper. +- By default, the first request with reference audio and no reference text loads + `openai/whisper-large-v3-turbo` on the worker device, which adds latency and + device memory use. Set `additional_config.omnivoice_asr.load_asr_on_startup` to + `true` to move that load to worker startup. Generated audio is mono at 24 kHz. + +### Reference audio and output processing + +When a request includes `ref_audio`, OmniVoice prepares the reference clip before it transcribes the clip or builds the voice cloning input. It performs these steps: + +- It converts the clip to mono and resamples it to 24 kHz. +- It measures the original root mean square (RMS) level. When the RMS is greater than zero and below `0.1`, it raises the waveform level to `0.1` and keeps the original RMS for output volume control. +- When `ref_text` is omitted, it trims a clip longer than 20 seconds at a detected silence. +- It removes middle silence and trims silence at both edges whether the transcript is automatic or supplied by the caller. +- It aligns the sample count down to the audio tokenizer hop size. +- When `ref_text` is omitted, it sends the prepared waveform to Whisper. +- It adds a final period for English text or a Chinese full stop for Chinese text when the reference text has no ending punctuation. +- It sends the same prepared waveform to the audio tokenizer. + +After decoding, OmniVoice processes voice cloning output by removing middle gaps of at least 500 milliseconds and trimming edge silence while keeping 100 milliseconds at each edge. It uses the original reference RMS when the RMS is below `0.1`, and it leaves the level unchanged when the RMS is at least `0.1`. It fades the first and last 100 milliseconds, then adds 100 milliseconds of silence at both edges. Text only output returns the decoder tensor without this processing. --- diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index 1495a40068e..235f23c7f4c 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -44,7 +44,7 @@ python examples/offline_inference/text_to_speech//end2end.py \ --ref-text "Transcript of the reference audio." ``` -`--ref-audio` and `--ref-text` are optional (text-only synthesis works without them) and must be provided together for voice cloning. The exotic scripts — Qwen3-TTS, Voxtral TTS, CosyVoice3 — accept additional model-specific flags documented in their per-model section below. Qwen3-TTS in particular uses its own argparse surface (`--query-type`, `--audio-path`, etc.) and does not follow the common shape; see its section. +For most models, `--ref-audio` and `--ref-text` are optional for text-only synthesis but must be provided together for voice cloning. OmniVoice also accepts `--ref-audio` without `--ref-text` and transcribes the reference lazily with Whisper. Qwen3-TTS, Voxtral TTS, and CosyVoice3 accept additional model-specific flags documented in their sections below. Qwen3-TTS uses its own argparse surface with options such as `--query-type` and `--audio-path`, so it does not follow the common command shape. --- @@ -407,7 +407,7 @@ python examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py \ ## OmniVoice -Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design). +OmniVoice creates speech in more than 600 languages. It supports text-to-speech without a reference clip, voice cloning, and voice design. ### Prerequisites @@ -415,7 +415,7 @@ Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / c huggingface-cli download k2-fsa/OmniVoice ``` -Voice cloning requires `transformers>=5.3.0`. Auto and design modes work with `transformers>=4.57.0`. +Voice cloning requires `transformers>=5.3.0`. Automatic voice and voice design use `transformers>=4.57.0`. ### Quick start (auto voice) @@ -427,6 +427,8 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ ### Voice cloning +With an explicit transcript, voice cloning does not load Whisper: + ```bash python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --model k2-fsa/OmniVoice \ @@ -435,6 +437,30 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --ref-text "This is the reference transcription." ``` +Omit `--ref-text` to use automatic transcription. The default OmniVoice deploy config loads Whisper on the worker device when the first request needs it. Set `load_asr_on_startup: true` in `additional_config.omnivoice_asr` to load it during worker startup. Set `asr_device: "cuda:0"` or `asr_device: "cpu"` in the same block to override the worker device: + +```bash +python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ + --model k2-fsa/OmniVoice \ + --text "hello" \ + --ref-audio trump_ref.wav +``` + +### Reference audio processing + +OmniVoice prepares the reference clip before it transcribes the clip or builds the voice cloning input. It performs these steps: + +- It converts the clip to mono and resamples it to 24 kHz. +- It measures the original root mean square (RMS) level. When the RMS is greater than zero and below `0.1`, it raises the waveform level to `0.1` and keeps the original RMS for output volume control. +- When `--ref-text` is omitted, it trims a clip longer than 20 seconds at a detected silence. +- It removes middle silence and trims silence at both edges whether the transcript is automatic or supplied by the caller. +- It aligns the sample count down to the audio tokenizer hop size. +- When `--ref-text` is omitted, it sends the prepared waveform to Whisper. +- It adds a final period for English text or a Chinese full stop for Chinese text when the reference text has no ending punctuation. +- It sends the same prepared waveform to the audio tokenizer. + +After decoding, OmniVoice applies silence removal, RMS scaling, fades, and padding only to voice-cloning output. Text-only output returns the decoder tensor unchanged. Output audio is mono at 24 kHz. + ### Voice design ```bash @@ -466,6 +492,11 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ - Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking. - Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz. +- The default deploy config uses `openai/whisper-large-v3-turbo` for lazy ASR on + the worker device. Set `load_asr_on_startup: true` in `additional_config.omnivoice_asr` + to load it during worker startup. Set `asr_device` in the same block to use a + different GPU or the CPU. Supplying `--ref-text` avoids ASR. +- Generated audio remains mono at 24 kHz. --- diff --git a/examples/offline_inference/text_to_speech/omnivoice/end2end.py b/examples/offline_inference/text_to_speech/omnivoice/end2end.py index a54ae863d4f..bf040d14c96 100644 --- a/examples/offline_inference/text_to_speech/omnivoice/end2end.py +++ b/examples/offline_inference/text_to_speech/omnivoice/end2end.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """End-to-end OmniVoice TTS inference via vLLM-Omni. Supports: - Auto voice mode: text only → generated speech - Voice cloning mode: text + reference audio → cloned voice speech +- Automatic reference transcription: reference audio without reference text Usage: # Auto voice @@ -13,6 +14,11 @@ # Voice cloning python end2end.py --model k2-fsa/OmniVoice --text "Hello" \ --ref-audio ref.wav --ref-text "reference transcription" + + # Automatic reference transcription (lazy by default; set + # load_asr_on_startup: true to preload Whisper at worker startup) + python end2end.py --model k2-fsa/OmniVoice --text "hello" \ + --ref-audio trump_ref.wav """ import argparse @@ -113,7 +119,8 @@ def run_e2e(): audio_signal, sr = load_audio(args.ref_audio, sr=None) multi_modal_data["audio"] = (audio_signal.astype(np.float32), sr) - mm_processor_kwargs["ref_text"] = args.ref_text or "" + if args.ref_text is not None: + mm_processor_kwargs["ref_text"] = args.ref_text mm_processor_kwargs["sample_rate"] = sr if args.lang: diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 6a8a2468662..53725cb740f 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -726,7 +726,7 @@ Then open in your browser. ## OmniVoice -Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline. +OmniVoice creates speech in more than 600 languages. It supports text-to-speech without a reference clip, voice cloning, and voice design. ### Prerequisites @@ -734,7 +734,7 @@ Zero-shot multilingual TTS (600+ languages). Online serving currently exposes ** huggingface-cli download k2-fsa/OmniVoice ``` -Voice cloning (offline) needs `transformers>=5.3.0`; auto voice works with `transformers>=4.57.0`. +Voice cloning needs `transformers>=5.3.0`. The default OmniVoice deploy config loads `openai/whisper-large-v3-turbo` lazily on the worker device when `ref_text` is missing. Set `load_asr_on_startup: true` in `additional_config.omnivoice_asr` to load it during worker startup. Set `asr_device: "cuda:0"` or `asr_device: "cpu"` in the same block to override the worker device. ### Launch @@ -759,6 +759,11 @@ python speech_client.py \ --ref-audio /path/to/ref_audio.wav \ --ref-text "Bonjour, comment allez-vous?" +# Voice cloning with automatic reference transcription +python speech_client.py \ +--text "hello" \ +--ref-audio /path/to/ref_audio.wav + # Style instruction (voice design-style control) python speech_client.py \ --text "Bonjour, comment allez-vous?" \ @@ -771,6 +776,23 @@ python speech_client.py --text "Hello, how are you?" --seed 42 The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--voice`, `--ref-audio`, `--ref-text`, `--instructions`, `--seed`, and `--output`. +### Reference audio and output processing + +The explicit `--ref-text` path avoids Whisper. With the default config, missing `ref_text` loads `openai/whisper-large-v3-turbo` on the worker device during the first request that needs transcription. Set `load_asr_on_startup: true` in `additional_config.omnivoice_asr` to move that load to worker startup. Set `asr_device` in the same block to choose another GPU or the CPU. + +OmniVoice prepares the reference clip before it transcribes the clip or builds the voice cloning input. It performs these steps: + +- It converts the clip to mono and resamples it to 24 kHz. +- It measures the original root mean square (RMS) level. When the RMS is greater than zero and below `0.1`, it raises the waveform level to `0.1` and keeps the original RMS for output volume control. +- When `--ref-text` is omitted, it trims a clip longer than 20 seconds at a detected silence. Supplied text skips this special long-audio trim. +- It removes middle silence and trims silence at both edges whether the transcript is automatic or supplied by the caller. +- It aligns the sample count down to the audio tokenizer hop size. +- When `--ref-text` is omitted, it sends the prepared waveform to Whisper. +- It adds a final period for English text or a Chinese full stop for Chinese text when the reference text has no ending punctuation. +- It sends the same prepared waveform to the audio tokenizer. + +After decoding, OmniVoice applies silence removal, RMS scaling, fades, and padding only to voice-cloning output. Text-only output returns the decoder tensor unchanged. Output audio is mono at 24 kHz. + ## Qwen3-TTS Three model variants exposed via separate checkpoints: diff --git a/examples/online_serving/text_to_speech/omnivoice/speech_client.py b/examples/online_serving/text_to_speech/omnivoice/speech_client.py index 9c9b8128ea6..fa9c56c2b21 100644 --- a/examples/online_serving/text_to_speech/omnivoice/speech_client.py +++ b/examples/online_serving/text_to_speech/omnivoice/speech_client.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Client for OmniVoice TTS via /v1/audio/speech endpoint. Examples: @@ -9,6 +12,9 @@ # Use a specific uploaded/supported voice python speech_client.py --text "Hello" --voice my_uploaded_voice + + # Clone a reference voice and let the server transcribe the reference audio + python speech_client.py --text "hello" --ref-audio /path/to/ref_audio.wav """ import argparse diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index bcf6b8c54fd..69d897ac982 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -533,7 +533,7 @@ def test_minicpmo_npu_additional_config_reaches_engine_args(monkeypatch, deploy_ monkeypatch.setattr(current_omni_platform, "device_name", "npu") monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None) stage = _from_pipeline_key("minicpmo_4_5", deploy_config_path=deploy_name).stage_by_id(2) - expected = {"code2wav_enable_npu_graph": True, "code2wav_max_npu_graphs": 32} + expected: dict[str, object] = {"code2wav_enable_npu_graph": True, "code2wav_max_npu_graphs": 32} if deploy_name == "minicpmo_4_5": # Single-card deploy: batch sizes round up to a capture bucket so the # 32-entry pool covers the shape space. @@ -1152,6 +1152,7 @@ def test_sub_config_fields_match_structured_scopes(): "media_io_kwargs", "final_output", "supports_running_prefix_cache_reset", + "recompute_preemption", "active_stream_window", "session_mode", "duplex_max_sessions", @@ -2511,3 +2512,28 @@ def test_mrv2_duplex_session_validation(): ps_no_duplex = replace(ps, supports_duplex_mrv2=False) with pytest.raises(ValueError, match="supports session_mode 'turn' only"): validate_native_mrv2_session(deploy_duplex, ps_no_duplex, "v2") + + +@pytest.mark.parametrize("execution_type", list(StageExecutionType)) +@pytest.mark.parametrize("mode", ["allow", "fail"]) +def test_recompute_preemption_projection_is_ar_only(execution_type, mode): + pipeline = PipelineConfig( + model_type="test_preemption", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="test", + execution_type=execution_type, + recompute_preemption=mode, + final_output=True, + ), + ), + ) + deploy = DeployConfig(stages=[StageDeployConfig(stage_id=0)]) + legacy = merge_pipeline_deploy(pipeline, deploy)[0] + structured = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0) + expected = mode if execution_type is StageExecutionType.LLM_AR else "allow" + assert legacy.yaml_engine_args.get("recompute_preemption", "allow") == expected + assert structured.model_config.recompute_preemption == expected + if execution_type is not StageExecutionType.LLM_AR: + assert "recompute_preemption" not in legacy.yaml_engine_args diff --git a/tests/core/sched/test_fast_allocate.py b/tests/core/sched/test_fast_allocate.py index 6ef5bc9861f..a36b6e9afc2 100644 --- a/tests/core/sched/test_fast_allocate.py +++ b/tests/core/sched/test_fast_allocate.py @@ -4,6 +4,7 @@ import pytest import torch +import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers from vllm.sampling_params import SamplingParams from vllm.v1.core.kv_cache_manager import KVCacheManager from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec @@ -17,6 +18,10 @@ def _manager(num_blocks=40, enable_caching=False, groups=1): + # vLLM lazily registers built-in KV specs. Another test may already have + # registered an out-of-tree spec, so the registry can be non-empty before + # FullAttentionSpec has a manager. Initialize the native mappings first. + native_kv_managers.register_all_kvcache_specs(None) spec = FullAttentionSpec(block_size=BLOCK, num_kv_heads=1, head_size=8, dtype=torch.float16) config = KVCacheConfig( num_blocks=num_blocks, @@ -70,6 +75,69 @@ def test_fast_path_matches_full_allocation_through_block_boundaries(skip_zeroing assert fast.allocate_slots(r_fast, 1, reserved_blocks=reserve) is None +def test_foreign_registry_still_builds_full_attention_manager(): + """A prior out-of-tree spec must not hide FullAttentionSpec from ``_manager``. + + ``KVCacheSpecRegistry._ensure_registered`` returns once the registry is + non-empty, which is the weekly ``No manager registered for KVCacheSpec + <...FullAttentionSpec>`` failure. ``_manager`` calls + ``register_all_kvcache_specs`` anyway, so the same construction succeeds. + The process registry is restored afterwards. + """ + from dataclasses import dataclass + + import vllm.v1.kv_cache_spec_registry as kv_spec_registry + from vllm.v1.kv_cache_interface import KVCacheSpec + + @dataclass(frozen=True) + class _ForeignKVCacheSpec(KVCacheSpec): + """Occupies the registry without registering a native attention spec.""" + + class _ForeignKVManager: + """Never constructed. Its registration only makes the registry non-empty.""" + + spec_registry = kv_spec_registry._REGISTRY_KVCACHESPEC_LIST + role_registry = kv_spec_registry._REGISTRY_ROLE_MANAGERS + saved_specs = dict(spec_registry) + saved_roles = dict(role_registry) + spec = FullAttentionSpec(block_size=BLOCK, num_kv_heads=1, head_size=8, dtype=torch.float16) + config = KVCacheConfig( + num_blocks=40, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(["layer0"], spec)], + ) + + def construct_without_native_registration(): + return KVCacheManager( + config, + max_model_len=256, + scheduler_block_size=BLOCK, + hash_block_size=BLOCK, + enable_caching=False, + ) + + try: + spec_registry.clear() + kv_spec_registry.KVCacheSpecRegistry.register(_ForeignKVCacheSpec, _ForeignKVManager) + kv_spec_registry.KVCacheSpecRegistry._ensure_registered(None) + assert FullAttentionSpec not in spec_registry + + with pytest.raises( + AssertionError, + match=r"No manager registered for KVCacheSpec ", + ): + construct_without_native_registration() + + manager = _manager() + assert manager.coordinator.single_type_managers + assert FullAttentionSpec in spec_registry + finally: + spec_registry.clear() + spec_registry.update(saved_specs) + role_registry.clear() + role_registry.update(saved_roles) + + def test_fast_path_only_installs_where_it_is_exact(): assert not install_decode_allocate_fast_path(_manager(enable_caching=True)) assert not install_decode_allocate_fast_path(_manager(groups=2)) diff --git a/tests/core/sched/test_omni_recompute_preemption.py b/tests/core/sched/test_omni_recompute_preemption.py new file mode 100644 index 00000000000..097ae26b21e --- /dev/null +++ b/tests/core/sched/test_omni_recompute_preemption.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Capacity preemption through the production scheduler and real KV allocator. + +Only model configuration, runner output and connector I/O are synthetic. No +schedule, finish, allocation or terminal-output method is replaced. +""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from vllm.config import CacheConfig, SchedulerConfig, VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1 +from vllm.sampling_params import SamplingParams +from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs +from vllm.v1.engine import FinishReason +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec +from vllm.v1.outputs import ModelRunnerOutput +from vllm.v1.request import Request, RequestStatus +from vllm.v1.structured_output import StructuredOutputManager + +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _scheduler(mode="fail", requires_kv_delivery=False, auxiliary=False, block_size=16, num_blocks=2): + config = VllmConfig() + # Avoid downloading model weights/config; all scheduler/cache objects are real. + config.model_config = SimpleNamespace( + multimodal_config=None, + uses_mrope=False, + is_encoder_decoder=False, + max_model_len=64, + runner_type="generate", + is_diffusion=False, + supports_multimodal_inputs=False, + return_sampling_mask=False, + recompute_preemption=mode, + ) + config.scheduler_config = SchedulerConfig( + max_model_len=64, + max_num_seqs=2, + max_num_batched_tokens=64, + is_encoder_decoder=False, + ) + config.aux_output_config.enable_return_routed_experts = auxiliary + config.cache_config = CacheConfig(block_size=block_size, enable_prefix_caching=False) + # By default: one null block and one usable block, enough only for prefill. + config.cache_config.num_gpu_blocks = num_blocks + # Match EngineCore initialization before constructing the scheduler. Custom + # specs registered during test collection can suppress lazy built-in registration. + register_all_kvcache_specs(config) + kv_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec(block_size=block_size, num_kv_heads=1, head_size=8, dtype=torch.float16), + ) + ], + ) + scheduler = OmniARScheduler(config, kv_config, Mock(spec=StructuredOutputManager), block_size) + if requires_kv_delivery: + connector = Mock(spec=KVConnectorBase_V1) + connector.requires_kv_delivery = True + connector.has_pending_block_frees.return_value = False + connector.get_num_new_matched_tokens.return_value = (0, False) + connector.request_finished.return_value = (False, None) + connector.build_connector_meta.return_value = None + connector.take_events.return_value = () + connector.get_kv_connector_stats.return_value = None + scheduler.connector = connector + scheduler.requires_kv_delivery = connector.requires_kv_delivery + return scheduler + + +def _output(request=None): + ids = [] if request is None else [request.request_id] + return ModelRunnerOutput( + req_ids=ids, + req_id_to_index={rid: i for i, rid in enumerate(ids)}, + sampled_token_ids=[] if request is None else [[1]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + +def _prefill(scheduler, prompt_len=16): + request = Request("req-preempt", list(range(prompt_len)), SamplingParams(max_tokens=32), None, client_index=3) + scheduler.add_request(request) + scheduled = scheduler.schedule() + assert scheduled.num_scheduled_tokens == {request.request_id: prompt_len} + scheduler.update_from_output(scheduled, _output(request)) + assert request.status == RequestStatus.RUNNING + assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == 0 + return request + + +@pytest.mark.parametrize("requires_kv_delivery", [False, True]) +@pytest.mark.parametrize("auxiliary", [False, True]) +def test_capacity_preemption_emits_one_error_and_frees_state(requires_kv_delivery, auxiliary): + scheduler = _scheduler(requires_kv_delivery=requires_kv_delivery, auxiliary=auxiliary) + request = _prefill(scheduler) + scheduled = scheduler.schedule() + outputs = scheduler.update_from_output(scheduled, _output()) + + assert request.status == RequestStatus.FINISHED_ERROR + assert request.request_id not in scheduler.requests + assert not scheduler.running and not scheduler.waiting and not scheduler.kv_holding_waiting + assert request not in scheduler.deferred_waiting + assert request not in scheduler._inflight_prefills + assert not any(scheduler.kv_cache_manager.get_block_ids(request.request_id)) + assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == 1 + assert scheduled.finished_req_ids == {request.request_id} + assert scheduled.num_scheduled_tokens == {} + [terminal] = outputs[3].outputs + assert terminal.finish_reason == FinishReason.ERROR + assert terminal.new_token_ids == [] + assert "cannot replay" in terminal.stop_reason and "truncated" in terminal.stop_reason + assert outputs[3].finished_requests == {request.request_id} + if requires_kv_delivery: + scheduler.connector.request_finished.assert_called_once() + + if auxiliary: + assert request.request_id not in scheduler.aux_output_connector._sent_hash_counts + assert scheduled.aux_output_connector_metadata.finished_requests == (request.request_id,) + + # A late repeat and subsequent empty engine step must not emit ERROR/ABORT twice. + scheduler._preempt_request(request, 1.0) + next_step = scheduler.schedule() + if auxiliary: + assert next_step.aux_output_connector_metadata.finished_requests == () + assert scheduler.update_from_output(next_step, _output()) == {} + + +@pytest.mark.parametrize("requires_kv_delivery", [False, True]) +def test_allow_mode_capacity_preemption_requeues(requires_kv_delivery): + scheduler = _scheduler("allow", requires_kv_delivery) + request = _prefill(scheduler) + scheduled = scheduler.schedule() + assert request.status == RequestStatus.PREEMPTED + assert request in scheduler.waiting + assert request.num_computed_tokens == 0 + assert request.drop_stale_output is requires_kv_delivery + assert scheduler.update_from_output(scheduled, _output()) == {} + + +@pytest.mark.parametrize("requires_kv_delivery", [False, True]) +def test_explicit_prefix_cache_reset_preserves_resume(requires_kv_delivery): + scheduler = _scheduler(requires_kv_delivery=requires_kv_delivery) + # The prompt plus its first sampled token must still fit for reset replay. + request = _prefill(scheduler, prompt_len=15) + assert scheduler.reset_prefix_cache(reset_running_requests=True) + assert request.status == RequestStatus.PREEMPTED + assert request in scheduler.waiting + assert request.drop_stale_output + assert not scheduler._resetting_prefix_cache + assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == 1 + # The resumed prefill fills the only block, then capacity must fail again: + # the reset exception cannot leak into a subsequent scheduler iteration. + scheduled = scheduler.schedule() + scheduler.update_from_output(scheduled, _output(request)) + outputs = scheduler.update_from_output(scheduler.schedule(), _output()) + [terminal] = outputs[3].outputs + assert terminal.finish_reason == FinishReason.ERROR + + +def test_prefix_cache_reset_scope_is_cleared_on_error(monkeypatch): + scheduler = _scheduler() + _prefill(scheduler) + monkeypatch.setattr(scheduler.kv_cache_manager, "reset_prefix_cache", lambda: False) + with pytest.raises(RuntimeError, match="Failed to reset KV cache"): + scheduler.reset_prefix_cache(reset_running_requests=True) + assert not scheduler._resetting_prefix_cache diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 7c11112c8e8..8f03eace75b 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/tests/diffusion/attention/test_attention_config.py @@ -1271,9 +1271,9 @@ class TestDense4DMaskRouting: """ @staticmethod - def _backend(*, supports_4d_mask: bool): + def _backend(*, supports_4d_mask: bool, piecewise: bool = False): class _Backend: - supports_piecewise_spans = False + supports_piecewise_spans = piecewise @classmethod def supports_dense_4d_mask(cls, attention_spec=None): @@ -1296,6 +1296,7 @@ def _make_attention( monkeypatch, *, supports_4d_mask: bool = False, + piecewise: bool = False, backend_explicit: bool = False, paged_kv: bool = False, with_sdpa: bool = True, @@ -1307,7 +1308,7 @@ def _make_attention( attention._scheduler_paged_kv = False attention.role = "self" attention.attn_spec = None - attention.attn_backend = self._backend(supports_4d_mask=supports_4d_mask) + attention.attn_backend = self._backend(supports_4d_mask=supports_4d_mask, piecewise=piecewise) attention.backend_explicit = backend_explicit attention.paged_kv_cache_role = "primary" if paged_kv else None backend_forward = Mock(return_value=torch.ones(1)) @@ -1400,6 +1401,44 @@ def test_dense_4d_mask_capable_backend_is_not_diverted(self, monkeypatch): backend_forward.assert_called_once() sdpa_forward.assert_not_called() + @pytest.mark.parametrize("backend_explicit", [False, True]) + def test_piecewise_spans_reach_backend_despite_dense_4d_mask(self, monkeypatch, backend_explicit): + # HunyuanImage3 sends a dense 4D mask together with full_attn_spans; a piecewise + # backend honors the spans and never reads the mask, so the call must not be + # diverted to SDPA or rejected for an explicit FLASH_ATTN selection. + attention, backend_forward, sdpa_forward = self._make_attention( + monkeypatch, piecewise=True, backend_explicit=backend_explicit + ) + query, key, value = self._tensors() + metadata = AttentionMetadata(attn_mask=self._dense_4d_mask(), full_attn_spans=[[(1, 3)], [(1, 3)]]) + + output = attention._run_local_attention(query, key, value, metadata) + + assert output is backend_forward.return_value + backend_forward.assert_called_once() + sdpa_forward.assert_not_called() + + def test_piecewise_backend_without_spans_still_diverts_dense_4d_mask(self, monkeypatch): + attention, backend_forward, sdpa_forward = self._make_attention(monkeypatch, piecewise=True) + query, key, value = self._tensors() + metadata = AttentionMetadata(attn_mask=self._dense_4d_mask()) + + output = attention._run_local_attention(query, key, value, metadata) + + assert output is sdpa_forward.return_value + backend_forward.assert_not_called() + + def test_spans_on_non_piecewise_backend_still_fail_closed(self, monkeypatch): + attention, backend_forward, sdpa_forward = self._make_attention(monkeypatch, backend_explicit=True) + query, key, value = self._tensors() + metadata = AttentionMetadata(attn_mask=self._dense_4d_mask(), full_attn_spans=[[(1, 3)], [(1, 3)]]) + + with pytest.raises(NotImplementedError, match="does not support a dense 4D attn_mask"): + attention._run_local_attention(query, key, value, metadata) + + backend_forward.assert_not_called() + sdpa_forward.assert_not_called() + @pytest.mark.parametrize("is_npu", [False, True]) def test_flash_attention_capability_matches_platform(self, monkeypatch, is_npu): from vllm_omni.diffusion.attention.backends import flash_attn as flash_attn_mod diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py index 96b4616ae5e..1da5894ccc0 100644 --- a/tests/diffusion/attention/test_flash_attn.py +++ b/tests/diffusion/attention/test_flash_attn.py @@ -1373,6 +1373,138 @@ def fake_attention_forward(query, key, value, **kwargs): torch.testing.assert_close(out, torch.ones(1, 8, 2, 4)) +# --- Test group E: NPU piecewise block-causal attention ---------------------- +# +# ``forward_fa_npu`` must honour ``full_attn_spans`` the way ``forward_cuda`` +# does. Before this was wired up, span metadata was silently dropped and the +# whole sequence ran bidirectional attention, which is wrong for block-causal +# models (e.g. Qwen-Image-2.1 prefill) and poisoned the cached prefix K/V. + + +def _npu_fusion_attention_reference( + query, + key, + value, + *, + atten_mask=None, + head_num, + input_layout="BSND", + scale=1.0, + keep_prob=1.0, + sparse_mode=0, + inner_precise=0, +): + """Pure-torch stand-in for ``npu_fusion_attention`` (BSND, no dropout). + + Implements the two sparse modes the piecewise path relies on: 0 (no mask, + full attention) and 3 (right-down causal, bottom-right aligned). The shared + compressed mask is validated for shape and convention, but the alignment + itself comes from ``sparse_mode`` -- exactly as the real kernel documents. + """ + del inner_precise + assert input_layout == "BSND" + assert query.shape[2] == head_num == key.shape[2] + assert keep_prob == 1.0 + query_length, key_length = query.shape[1], key.shape[1] + if sparse_mode == 0: + assert atten_mask is None, "full segments must not carry a mask" + keep = torch.ones(query_length, key_length, dtype=torch.bool) + elif sparse_mode == 3: + assert atten_mask is not None and atten_mask.shape == (2048, 2048) + assert atten_mask.dtype == torch.bool + # Ascend convention: True = masked out, so the mask is the strict + # upper triangle. + strict_upper = torch.triu(torch.ones(2048, 2048, dtype=torch.bool), diagonal=1) + assert torch.equal(atten_mask, strict_upper) + assert query_length <= key_length, "right-down causal requires Sq <= Skv" + query_positions = torch.arange(query_length).unsqueeze(1) + key_positions = torch.arange(key_length).unsqueeze(0) + keep = key_positions <= query_positions + (key_length - query_length) + else: + raise AssertionError(f"unexpected sparse_mode {sparse_mode}") + + scores = torch.einsum("bqhd,bkhd->bhqk", query.float(), key.float()) * scale + scores = scores.masked_fill(~keep, float("-inf")) + out = torch.einsum("bhqk,bkhd->bqhd", torch.softmax(scores, dim=-1), value.float()) + return out.to(query.dtype), None, None, None, 0, 0, 0 + + +def _block_causal_reference(query, key, value, image_ids, scale): + """Dense ``(q >= kv) or same_image_block`` reference for the joint sequence.""" + seq_len = query.shape[1] + query_positions = torch.arange(seq_len).unsqueeze(1) + key_positions = torch.arange(seq_len).unsqueeze(0) + same_block = (image_ids.view(-1, 1) == image_ids.view(1, -1)) & (image_ids.view(-1, 1) >= 0) + keep = (query_positions >= key_positions) | same_block + scores = torch.einsum("bqhd,bkhd->bhqk", query.float(), key.float()) * scale + scores = scores.masked_fill(~keep, float("-inf")) + return torch.einsum("bhqk,bkhd->bqhd", torch.softmax(scores, dim=-1), value.float()) + + +@pytest.mark.cpu +@pytest.mark.parametrize("causal", [False, True]) +def test_npu_piecewise_reports_its_own_execution_path(npu_contract, causal): + _, query, _, _ = npu_contract + impl = _npu_impl(causal=causal) + metadata = AttentionMetadata(full_attn_spans=[[(3, 5)]]) + _assert_platform_path(impl, "npu", query, metadata, "npu_piecewise") + + +@pytest.mark.cpu +def test_npu_causal_without_spans_stays_unverified(npu_contract): + _, query, _, _ = npu_contract + impl = _npu_impl(causal=True) + _assert_platform_path(impl, "npu", query, None, "npu_unverified") + + +@pytest.mark.cpu +@pytest.mark.parametrize("causal", [False, True]) +def test_npu_piecewise_runs_one_fusion_attention_per_segment(monkeypatch, causal): + fusion_attention = Mock(side_effect=lambda q, k, v, **kwargs: (torch.zeros_like(q), None, None, None, 0, 0, 0)) + fake_torch_npu = _fake_torch_npu(monkeypatch, npu_fusion_attention=fusion_attention) + impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=causal) + query = torch.randn(1, 5, 2, 4) + # Two text tokens, then one condition-image block covering [2, 5). + metadata = AttentionMetadata(full_attn_spans=[[(2, 5)]]) + + impl.forward_fa_npu(query, query, query, metadata) + + assert fake_torch_npu.npu_fusion_attention.call_count == 2 + causal_call, full_call = fake_torch_npu.npu_fusion_attention.call_args_list + # Causal text segment: K truncated to its own end, compressed mask, mode 3. + assert causal_call.args[0].shape == (1, 2, 2, 4) + assert causal_call.args[1].shape == (1, 2, 2, 4) + assert causal_call.kwargs["sparse_mode"] == 3 + assert causal_call.kwargs["atten_mask"].shape == (2048, 2048) + # Full image segment: K spans the whole prefix, no mask materialized. + assert full_call.args[0].shape == (1, 3, 2, 4) + assert full_call.args[1].shape == (1, 5, 2, 4) + assert full_call.kwargs["sparse_mode"] == 0 + assert full_call.kwargs["atten_mask"] is None + + +@pytest.mark.cpu +@pytest.mark.parametrize("causal", [False, True]) +def test_npu_piecewise_matches_block_causal_reference(monkeypatch, causal): + """Segments reproduce the dense block-causal pattern end to end.""" + torch.manual_seed(0) + _fake_torch_npu(monkeypatch, npu_fusion_attention=_npu_fusion_attention_reference) + impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=causal) + # text(0-1) | image A(2-4) | text(5) | image B(6-8) + image_ids = torch.tensor([-1, -1, 0, 0, 0, -1, 1, 1, 1]) + metadata = AttentionMetadata(full_attn_spans=[[(2, 5), (6, 9)]]) + seq_len = image_ids.shape[0] + query = torch.randn(1, seq_len, 2, 4) + key = torch.randn(1, seq_len, 2, 4) + value = torch.randn(1, seq_len, 2, 4) + + out = impl.forward_fa_npu(query, key, value, metadata) + expected = _block_causal_reference(query, key, value, image_ids, impl.softmax_scale) + + assert out.shape == query.shape + torch.testing.assert_close(out, expected, rtol=1e-5, atol=1e-6) + + if __name__ == "__main__": print("Running FlashAttention Padding Tests...") print("=" * 60) diff --git a/tests/diffusion/distributed/test_wan_attention_sdpa_alignment.py b/tests/diffusion/distributed/test_wan_attention_sdpa_alignment.py new file mode 100644 index 00000000000..69e11709093 --- /dev/null +++ b/tests/diffusion/distributed/test_wan_attention_sdpa_alignment.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CPU checks for Wan VAE attention sequence strides. + +Cosmos3-Nano decodes with the Wan 2.2 VAE (spatial compression 16, mid-block +channels 1024) and the diffusers tile defaults (256px minimum, 192px stride). +A corner tile whose latent is 1x1 makes ``WanAttentionBlock`` hand SDPA a query +with sequence stride 1. The memory-efficient CUDA kernel rejects that with +``query is not correctly aligned (strideM)``. These tests do not run that +kernel; they check the stride condition and that realignment preserves values. +""" + +from __future__ import annotations + +import warnings +from collections.abc import Iterator +from contextlib import contextmanager + +import pytest +import torch +import torch.nn.functional as F +from torch import nn +from torch.nn.attention import SDPBackend, sdpa_kernel + +from vllm_omni.diffusion.distributed.autoencoders.wan_sdpa_alignment import ( + align_sdpa_qkv, + install_wan_attention_sdpa_alignment, + sdpa_sequence_stride_alignment, + sequence_stride_is_aligned, + split_wan_attention_qkv, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# nvidia/Cosmos3-Nano ``vae/config.json`` plus diffusers' Wan tile defaults. +_COSMOS3_SPATIAL_COMPRESSION = 16 +_COSMOS3_MID_CHANNELS = 256 * 4 +_TILE_LATENT_MIN = 256 // _COSMOS3_SPATIAL_COMPRESSION +_TILE_LATENT_STRIDE = 192 // _COSMOS3_SPATIAL_COMPRESSION +# tests/dfx/stability/tests/test_cosmos3.json samples inclusive integer ranges. +_STABILITY_HEIGHTS = range(720, 1025) +_STABILITY_WIDTHS = range(1024, 1281) + + +def _tile_lengths(length: int) -> list[int]: + """Latent extents produced by ``DistributedAutoencoderKLWan.tile_split``.""" + return [min(_TILE_LATENT_MIN, length - start) for start in range(0, length, _TILE_LATENT_STRIDE)] + + +def _stability_corner_tiles() -> set[tuple[int, int]]: + tiles: set[tuple[int, int]] = set() + for height in _STABILITY_HEIGHTS: + latent_h = height // _COSMOS3_SPATIAL_COMPRESSION + for width in _STABILITY_WIDTHS: + latent_w = width // _COSMOS3_SPATIAL_COMPRESSION + for tile_h in _tile_lengths(latent_h): + for tile_w in _tile_lengths(latent_w): + tiles.add((tile_h, tile_w)) + return tiles + + +def _qkv_projection(batch_time: int, channels: int, height: int, width: int, dtype: torch.dtype) -> torch.Tensor: + generator = torch.Generator().manual_seed(0) + return torch.randn(batch_time, channels * 3, height, width, dtype=dtype, generator=generator) + + +@pytest.mark.parametrize("channels_last", [False, True]) +def test_channels_last_projection_has_the_same_stride_rule(channels_last: bool) -> None: + projection = _qkv_projection(1, _COSMOS3_MID_CHANNELS, 1, 1, torch.bfloat16) + if channels_last: + projection = projection.contiguous(memory_format=torch.channels_last) + query, key, value = split_wan_attention_qkv(projection) + assert query.stride(-2) == 1 + aligned_query, aligned_key, aligned_value = align_sdpa_qkv(query, key, value) + assert sequence_stride_is_aligned(aligned_query) + assert torch.equal(aligned_query, query) + assert torch.equal(aligned_key, key) + assert torch.equal(aligned_value, value) + + wide = _qkv_projection(1, _COSMOS3_MID_CHANNELS, 4, 9, torch.bfloat16) + if channels_last: + wide = wide.contiguous(memory_format=torch.channels_last) + wide_query, _, _ = split_wan_attention_qkv(wide) + realigned, _, _ = align_sdpa_qkv(wide_query, wide_query, wide_query) + assert realigned is wide_query + + +def test_stability_mix_includes_one_by_one_corner_tile() -> None: + tiles = _stability_corner_tiles() + assert (1, 1) in tiles + # Every other edge remainder in this grid keeps a legal sequence stride. + for tile_h, tile_w in tiles: + projection = _qkv_projection(1, _COSMOS3_MID_CHANNELS, tile_h, tile_w, torch.float32) + query, _, _ = split_wan_attention_qkv(projection) + aligned = sequence_stride_is_aligned(query) + assert aligned is (tile_h != 1 or tile_w != 1) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize( + ("height", "width"), + [(1, 1), (1, 16), (16, 1), (16, 16), (9, 4), (4, 8)], +) +def test_align_sdpa_qkv_fixes_only_misaligned_queries(dtype: torch.dtype, height: int, width: int) -> None: + projection = _qkv_projection(1, _COSMOS3_MID_CHANNELS, height, width, dtype) + query, key, value = split_wan_attention_qkv(projection) + aligned_query, aligned_key, aligned_value = align_sdpa_qkv(query, key, value) + + assert sequence_stride_is_aligned(aligned_query) + assert sequence_stride_is_aligned(aligned_key) + assert sequence_stride_is_aligned(aligned_value) + assert torch.equal(aligned_query, query) + assert torch.equal(aligned_key, key) + assert torch.equal(aligned_value, value) + + if height == 1 and width == 1: + alignment = sdpa_sequence_stride_alignment(query) + assert query.stride(-2) == 1 + assert query.stride(-2) % alignment != 0 + assert aligned_query is not query + assert aligned_query.stride(-2) % alignment == 0 + else: + assert aligned_query is query + assert aligned_key is key + assert aligned_value is value + + +class WanAttentionBlock(nn.Module): + """Stand-in with the diffusers block's attribute names. Not a numeric reference.""" + + def __init__(self, dim: int, dtype: torch.dtype) -> None: + super().__init__() + self.norm = nn.Identity() + self.to_qkv = nn.Conv2d(dim, dim * 3, 1).to(dtype=dtype) + self.proj = nn.Conv2d(dim, dim, 1).to(dtype=dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + raise AssertionError("unpatched Wan attention forward was called") + + +class _Root(nn.Module): + def __init__(self, block: WanAttentionBlock) -> None: + super().__init__() + self.block = block + + +@contextmanager +def _math_sdpa() -> Iterator[None]: + """Compare attention values with the math kernel. + + ROCm's default SDPA is not bitwise identical for a misaligned 1x1 query + and the repacked copy. These checks verify that realignment preserves + values, not that a vendor kernel is layout-invariant. + """ + with sdpa_kernel(SDPBackend.MATH): + yield + + +def _reference_attention(block: nn.Module, hidden: torch.Tensor) -> torch.Tensor: + """Diffusers' attention math on the raw, possibly misaligned QKV views.""" + identity = hidden + batch_size, channels, time, height, width = hidden.size() + features = hidden.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, height, width) + features = block.norm(features) + query, key, value = split_wan_attention_qkv(block.to_qkv(features)) + features = F.scaled_dot_product_attention(query, key, value) + features = features.squeeze(1).permute(0, 2, 1).reshape(batch_size * time, channels, height, width) + features = block.proj(features) + features = features.view(batch_size, time, channels, height, width).permute(0, 2, 1, 3, 4) + return features + identity + + +def _assert_matches_unaligned_math(block: nn.Module, hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + with _math_sdpa(): + actual = block(hidden) + expected = _reference_attention(block, hidden) + assert torch.equal(actual, expected) + return actual, expected + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize(("height", "width"), [(1, 1), (4, 9), (16, 16)]) +def test_aligned_forward_matches_unaligned_math(dtype: torch.dtype, height: int, width: int) -> None: + torch.manual_seed(0) + block = WanAttentionBlock(32, dtype) + root = _Root(block) + install_wan_attention_sdpa_alignment(root) + hidden = torch.randn(1, 32, 1, height, width, dtype=dtype) + + actual, expected = _assert_matches_unaligned_math(block, hidden) + assert actual.stride() == expected.stride() + # A second install must leave the patched forward in place. + forward = block.forward + install_wan_attention_sdpa_alignment(root) + assert block.forward is forward + + +def test_math_reference_rejects_value_changing_alignment(monkeypatch) -> None: + """The exact comparator must reject a realignment that changes values.""" + from vllm_omni.diffusion.distributed.autoencoders import wan_sdpa_alignment + + original_align = wan_sdpa_alignment.align_sdpa_qkv + + def corrupt_value(query, key, value): + aligned_query, aligned_key, aligned_value = original_align(query, key, value) + return aligned_query, aligned_key, aligned_value + 1 + + monkeypatch.setattr(wan_sdpa_alignment, "align_sdpa_qkv", corrupt_value) + torch.manual_seed(0) + block = WanAttentionBlock(32, torch.float32) + root = _Root(block) + install_wan_attention_sdpa_alignment(root) + hidden = torch.randn(1, 32, 1, 1, 1) + + with pytest.raises(AssertionError): + _assert_matches_unaligned_math(block, hidden) + + +@pytest.mark.parametrize("alignment_first", [False, True]) +def test_alignment_preserves_spatial_shard_gather(monkeypatch, alignment_first: bool) -> None: + """Both installation orders must retain global attention before resharding.""" + from diffusers.models.autoencoders.autoencoder_kl_wan import WanAttentionBlock as DiffusersWanAttentionBlock + + from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard + + torch.manual_seed(0) + block = DiffusersWanAttentionBlock(32) + root = _Root(block) + local = torch.randn(1, 32, 1, 2, 2) + remote = torch.randn_like(local) + full = torch.cat((local, remote), dim=-1) + expected = block(full)[..., :2] + calls = [] + + def gather(x, *, group, dim): + calls.append(dim) + return torch.cat((x, remote), dim=dim) + + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 2)) + monkeypatch.setattr(wan_spatial_shard, "all_gather_along_dim", gather) + if alignment_first: + install_wan_attention_sdpa_alignment(root) + wan_spatial_shard._patch_attention_block(block, None, "width") + sharded_forward = block.forward + install_wan_attention_sdpa_alignment(root) + assert block.forward is sharded_forward + actual = block(local) + assert calls == [-1] + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("installed_version", ["0.40.0", "0.41.0", "0.40.0.dev0"]) +def test_alignment_warns_on_unverified_diffusers_version(monkeypatch, installed_version: str) -> None: + """A dependency upgrade must not silently retain the copied forward.""" + from importlib import metadata + + original_version = metadata.version + monkeypatch.setattr( + metadata, "version", lambda name: installed_version if name == "diffusers" else original_version(name) + ) + root = _Root(WanAttentionBlock(32, torch.float32)) + with warnings.catch_warnings(record=True) as recorded: + warnings.simplefilter("always") + install_wan_attention_sdpa_alignment(root) + forward = root.block.forward + install_wan_attention_sdpa_alignment(root) + assert root.block.forward is forward + if installed_version == "0.40.0": + assert not recorded + else: + assert len(recorded) == 1 + assert "validated against diffusers 0.40.0" in str(recorded[0].message) + assert installed_version in str(recorded[0].message) + hidden = torch.randn(1, 32, 1, 1, 1) + _assert_matches_unaligned_math(root.block, hidden) diff --git a/tests/diffusion/distributed/test_wan_attention_sdpa_alignment_cuda.py b/tests/diffusion/distributed/test_wan_attention_sdpa_alignment_cuda.py new file mode 100644 index 00000000000..08d3b423b06 --- /dev/null +++ b/tests/diffusion/distributed/test_wan_attention_sdpa_alignment_cuda.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Real Wan attention regression for the CUDA memory-efficient SDPA backend. + +No checkpoint is needed. The installed patch must run under the memory-efficient +backend and match the math backend. A torch build that still rejects an +unpatched 1x1 query records that strideM error and continues; a build that +accepts the query does not fail this case. This is kernel coverage, not the +Cosmos3-Nano multi-rank tiled-decode acceptance run. +""" + +import re + +import pytest +import torch +from diffusers.models.autoencoders.autoencoder_kl_wan import WanAttentionBlock +from torch.nn.attention import SDPBackend, sdpa_kernel + +from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.distributed.autoencoders.wan_sdpa_alignment import install_wan_attention_sdpa_alignment + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize(("height", "width"), [(1, 1), (4, 9)]) +def test_real_wan_attention_memory_efficient_sdpa(dtype: torch.dtype, height: int, width: int) -> None: + """Exercise the actual kernel, including an already-aligned tile control.""" + torch.manual_seed(0) + # Cosmos3-Nano's Wan VAE mid-block has 1024 channels. + block = WanAttentionBlock(1024).to(device="cuda", dtype=dtype).eval() + hidden = torch.randn(1, 1024, 1, height, width, device="cuda", dtype=dtype) + with torch.inference_mode(): + with sdpa_kernel(SDPBackend.MATH): + expected = block(hidden) + if (height, width) == (1, 1): + # Current memory-efficient SDPA rejects this query. A later torch + # that accepts it should not fail the regression; the patched + # comparison below is what guards the fix. + with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION): + try: + block(hidden) + except RuntimeError as exc: + if re.search(r"not correctly aligned.*strideM", str(exc)) is None: + raise + else: + with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION): + unpatched = block(hidden) + install_wan_attention_sdpa_alignment(block) + with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION): + actual = block(hidden) + tolerance = 2e-2 if dtype == torch.bfloat16 else 2e-3 + torch.testing.assert_close(actual, expected, rtol=tolerance, atol=tolerance) + if (height, width) != (1, 1): + torch.testing.assert_close(actual, unpatched, rtol=0, atol=0) diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index 487671a7350..b77c8ced316 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -366,6 +366,46 @@ def test_move_moe_lora_buffers_keeps_all_stacks_and_views_on_device(): ] +def test_lora_manager_homes_lora_buffers_to_compute_device_without_dlo(monkeypatch): + # Non-DLO deployments used to leave `_resident_lora_device` unset, so + # buffers created while the base weights were offloaded stayed on CPU + # after wake, and rank-growth reinit re-stranded them (#8817). + import vllm_omni.diffusion.lora.manager as manager_mod + + class _LoRALayer(DummyBaseLayerWithLoRA): + lora_buffer_device = None + + def _set_diffusion_lora_buffer_device(self, device): + self.lora_buffer_device = device + self._diffusion_lora_buffer_device = device + + monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _LoRALayer) + monkeypatch.setattr( + manager_mod, + "from_layer_diffusion", + lambda *, layer, **_kwargs: _LoRALayer(layer), + ) + + pipeline = torch.nn.Module() + pipeline.transformer = torch.nn.Module() + pipeline.transformer.foo = _FakeLinearBase() + compute = torch.device("meta") + manager = DiffusionLoRAManager( + pipeline=pipeline, + device=compute, + dtype=torch.bfloat16, + ) + + manager._replace_layers_with_lora(type("_PH", (), {"r": 1})()) + + layer = pipeline.transformer.foo + assert layer.lora_buffer_device == compute + # The recorded device is what create_lora_weights consults when a later + # rank increase re-allocates the buffers, so the second exposure from + # #8817 is pinned by this attribute as well. + assert layer._diffusion_lora_buffer_device == compute + + def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch): import vllm_omni.diffusion.lora.manager as manager_mod diff --git a/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py b/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py index 026e4e9907a..1851b0d1a46 100644 --- a/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py +++ b/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py @@ -24,7 +24,6 @@ pytestmark = [ pytest.mark.core_model, pytest.mark.advanced_model, - pytest.mark.cuda, pytest.mark.usefixtures("mock_tp1"), *hardware_marks(res={"cuda": "L4"}, num_cards=1), ] @@ -128,6 +127,7 @@ def test_production_rope_uses_fused_qk_norm_rope(monkeypatch, force_native): .to(torch.bfloat16) .eval() ) + initialize_block_weights(block) hidden = torch.randn(2, seq, DIM, device="cuda", dtype=torch.bfloat16) mask = torch.ones(2, seq, dtype=torch.bool, device="cuda") mask[0, -37:] = False diff --git a/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py b/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py index 9a530bf9006..c5503f15556 100644 --- a/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py +++ b/tests/diffusion/models/minimax_h3/test_fasth3_checkpoint.py @@ -401,6 +401,35 @@ def test_native_local_model_root_keeps_partition_layout(tmp_path): assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder") +def test_native_local_model_root_with_modular_index_keeps_partition_layout(tmp_path): + # The official MiniMaxAI/MiniMax-H3 repository ships modular_model_index.json at its root. + root = tmp_path / "MiniMax-H3" + (root / "FL2VA" / "text_encoder").mkdir(parents=True) + (root / "FL2VA" / "text_encoder" / "config.json").write_text("{}", encoding="utf-8") + (root / "FL2VA" / "model_index.json").write_text("{}", encoding="utf-8") + (root / "Ref2VA").mkdir() + (root / "modular_model_index.json").write_text( + json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}), + encoding="utf-8", + ) + + assert is_minimax_h3_modular(str(root)) is False + assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder") + assert resolve_minimax_h3_diffusion_model_path(str(root), None, "fl2va") == str(root / "FL2VA") + + +def test_empty_partition_dirs_do_not_override_modular_marker(tmp_path): + root = tmp_path / "FastH3" + (root / "FL2VA").mkdir(parents=True) + (root / "Ref2VA").mkdir() + (root / "modular_model_index.json").write_text( + json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}), + encoding="utf-8", + ) + + assert is_minimax_h3_modular(str(root)) is True + + def test_diffusion_resolver_returns_local_modular_root_without_partition_files(tmp_path): root = tmp_path / "FastH3" root.mkdir() diff --git a/tests/diffusion/models/omnivoice/test_prompt_tokenization.py b/tests/diffusion/models/omnivoice/test_prompt_tokenization.py index b8fae42414a..46d0d08bef8 100644 --- a/tests/diffusion/models/omnivoice/test_prompt_tokenization.py +++ b/tests/diffusion/models/omnivoice/test_prompt_tokenization.py @@ -28,23 +28,26 @@ def test_prepare_request_tokenizes_text_once_with_reference_audio(audio_mode): pipeline.tokenizer = SimpleNamespace( encode=Mock(side_effect=[SimpleNamespace(ids=[1, 2]), SimpleNamespace(ids=[5, 6, 7])]) ) - pipeline.audio_tokenizer = Mock() + pipeline.audio_tokenizer = SimpleNamespace(config=SimpleNamespace(sample_rate=24000, hop_length=240)) ref_tokens = torch.full((8, 4), 8, dtype=torch.long) pipeline._encode_ref_audio = Mock(return_value=ref_tokens) pipeline._speaker_cache = Mock() - pipeline._speaker_cache.get.return_value = {"ref_audio_tokens": ref_tokens} if audio_mode == "cached" else None + pipeline._speaker_cache.get.return_value = ( + {"ref_audio_tokens": ref_tokens, "reference_rms": 0.2} if audio_mode == "cached" else None + ) prompt = { "text": "hello world", "mm_processor_kwargs": {"ref_text": "reference"}, } if audio_mode != "none": - prompt["multi_modal_data"] = {"audio": [(np.zeros(2400, dtype=np.float32), 24000)]} + prompt["multi_modal_data"] = {"audio": [(np.full(2400, 0.2, dtype=np.float32), 24000)]} prompt["voice_name"] = "speaker" prepared = pipeline._prepare_request_input(prompt, {"seed": 42}) assert pipeline.tokenizer.encode.call_count == 2 # One style and one text encoding. - assert pipeline.tokenizer.encode.call_args.args == ("<|text_start|>reference hello world<|text_end|>",) + expected_ref_text = "reference" if audio_mode == "none" else "reference." + assert pipeline.tokenizer.encode.call_args.args == (f"<|text_start|>{expected_ref_text} hello world<|text_end|>",) torch.testing.assert_close(prepared.input_ids[:5], torch.tensor([1, 2, 5, 6, 7]).unsqueeze(1).repeat(1, 8)) assert prepared.target_len == 2 assert prepared.seed == 42 @@ -54,7 +57,7 @@ def test_prepare_request_tokenizes_text_once_with_reference_audio(audio_mode): if audio_mode == "uncached": pipeline._encode_ref_audio.assert_called_once() audio, sr = pipeline._encode_ref_audio.call_args.args - assert tuple(audio.shape) == (2400,) + assert tuple(audio.shape) == (1, 2400) assert sr == 24000 pipeline._speaker_cache.put.assert_called_once() else: diff --git a/tests/diffusion/models/qwen_image_21/test_vae.py b/tests/diffusion/models/qwen_image_21/test_vae.py index a9c1b39c03b..f50fe7c8de5 100644 --- a/tests/diffusion/models/qwen_image_21/test_vae.py +++ b/tests/diffusion/models/qwen_image_21/test_vae.py @@ -123,7 +123,8 @@ def test_distributed_tile_hooks_match_tiled_decode(): tiled = vae.tiled_decode(z, return_dict=False)[0] tasks, grid_spec = DistributedAutoencoderKLQwenImage21.tile_split(distributed, z) - # 16x16 latent, 8x8 latent tiles, stride 6 -> range(0, 16, 6) = 3 stops/axis. + # 16x16 latent, 8x8 latent tiles, stride 6 -> _tile_starts spreads the + # 8-cell span evenly: [0, 4, 8], so still 3 stops per axis. assert len(tasks) == 9 assert grid_spec.grid_shape == (3, 3) assert grid_spec.split_dims == (3, 4) @@ -137,6 +138,58 @@ def test_distributed_tile_hooks_match_tiled_decode(): torch.testing.assert_close(merged, tiled) +def test_trailing_tile_keeps_the_full_decode_window(): + """Regression: a clipped trailing tile leaves a visible seam (vllm-omni#8850). + + Tiles are placed at a fixed stride, so a size that is not a multiple of the + stride used to clip the trailing tile at the image border (4 latent cells + instead of 8 here). A clipped tile decodes with less context than its + neighbours, which leaves a seam that the blend does not hide. The layout must + snap the trailing tile back so every tile keeps the full decode window. + """ + torch.manual_seed(0) + vae = _make_vae(DistributedAutoencoderKLQwenImage21) + vae.enable_tiling( + tile_sample_min_height=32, + tile_sample_min_width=32, + tile_sample_stride_height=24, + tile_sample_stride_width=24, + ) + tile_latent = 32 // vae.spatial_compression_ratio # 8 latent cells + + # 16 and 22 are not multiples of the latent stride (6), so the trailing tile + # is clipped under the old layout. + for height, width in ((16, 22), (22, 16), (16, 16)): + z = torch.randn(1, 4, 1, height, width) + tasks, _ = DistributedAutoencoderKLQwenImage21.tile_split(vae, z) + windows = {(task.tensor[0].shape[-2], task.tensor[0].shape[-1]) for task in tasks} + assert windows == {(tile_latent, tile_latent)}, ( + f"latent {height}x{width}: decoded tiles with a clipped window {windows}; " + f"every tile must keep the full {tile_latent}x{tile_latent} window" + ) + + +def test_tile_starts_spread_evenly_and_keep_full_windows(): + """The layout invariant behind vllm-omni#8850, over every plausible axis size. + + Every start must leave room for the full window, the trailing tile must end + on the border (so nothing is clipped), the spacing must never exceed the + configured stride, and the spacing must be as even as integer cells allow. + """ + stride, tile_min = 24, 32 + for size in range(1, 257): + starts = AutoencoderKLQwenImage21._tile_starts(size, tile_min, stride) + assert starts[0] == 0 + assert starts == sorted(set(starts)), f"size {size}: starts not strictly increasing" + if size <= tile_min: + assert starts == [0], f"size {size}: a single tile should cover the axis" + continue + assert starts[-1] == size - tile_min, f"size {size}: trailing tile must end on the border" + gaps = [starts[i + 1] - starts[i] for i in range(len(starts) - 1)] + assert all(1 <= gap <= stride for gap in gaps), f"size {size}: gaps {gaps}" + assert max(gaps) - min(gaps) <= 1, f"size {size}: spacing must be even, got {gaps}" + + def test_tiled_encode_output_shape_and_closeness(): """tiled_encode is engaged for image-conditioned generation when tiling is on.""" torch.manual_seed(0) diff --git a/tests/e2e/offline_inference/test_omnivoice_expansion.py b/tests/e2e/offline_inference/test_omnivoice_expansion.py index dd65cbdccc6..d507d354d76 100644 --- a/tests/e2e/offline_inference/test_omnivoice_expansion.py +++ b/tests/e2e/offline_inference/test_omnivoice_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E offline tests for OmniVoice TTS model with text input and audio output. @@ -12,13 +12,18 @@ import numpy as np import pytest +import soundfile as sf +from vllm.multimodal.media.audio import load_audio from tests.helpers.mark import hardware_test +from tests.helpers.media import convert_audio_file_to_text, get_asset_path from tests.helpers.runtime import OmniRunner from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.inputs.data import OmniDiffusionSamplingParams MODEL = "k2-fsa/OmniVoice" DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml") +_OMNIVOICE_REF_AUDIO_SEED = 102 # OmniRunner tuple: model, deploy config path, extra Omni kwargs. _OMNI_RUNNER_PARAM = ( @@ -45,8 +50,6 @@ def test_omnivoice_text_to_audio(omni_runner: OmniRunner) -> None: Input Modal: text Output Modal: audio """ - from vllm_omni.inputs.data import OmniDiffusionSamplingParams - prompts = {"prompt": "Hello, this is a test for text to audio."} sampling_params_list = [OmniDiffusionSamplingParams()] @@ -56,12 +59,7 @@ def test_omnivoice_text_to_audio(omni_runner: OmniRunner) -> None: # Check final output has audio final_output = outputs[-1] - ro = final_output - assert ro is not None, "No request_output" - - mm = getattr(ro, "multimodal_output", None) - if not mm and ro.outputs: - mm = getattr(ro.outputs[0], "multimodal_output", None) + mm = final_output.multimodal_output assert mm is not None, "No multimodal_output" assert "audio" in mm, f"No 'audio' key in multimodal_output: {mm.keys()}" @@ -77,3 +75,48 @@ def test_omnivoice_text_to_audio(omni_runner: OmniRunner) -> None: assert rms > 0.01, f"Audio RMS too low ({rms:.4f}), likely silence" print(f"Generated audio: {len(audio_np) / 24000:.2f}s, rms={rms:.4f}") + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_omnivoice_ref_audio_without_ref_text(omni_runner: OmniRunner, tmp_path) -> None: + """Reference audio without a transcript is transcribed before cloning.""" + ref_audio_path = get_asset_path("qwen3_tts/clone_2.wav") + audio_signal, sample_rate = load_audio(str(ref_audio_path), sr=None) + prompts = { + "prompt": "hello", + "multi_modal_data": { + "audio": (audio_signal.astype(np.float32), sample_rate), + }, + "mm_processor_kwargs": {"sample_rate": sample_rate}, + } + + def generate_audio(): + sampling_params_list = [OmniDiffusionSamplingParams(extra_args={"seed": _OMNIVOICE_REF_AUDIO_SEED})] + outputs = list(omni_runner.omni.generate(prompts, sampling_params_list=sampling_params_list)) + assert outputs, "No outputs generated" + final_output = outputs[-1] + mm = final_output.multimodal_output + assert mm is not None, "No multimodal_output" + + audio = mm.get("audio") + assert audio is not None, "No audio output" + audio_np = audio if isinstance(audio, np.ndarray) else audio.cpu().numpy().squeeze() + return np.asarray(audio_np).squeeze(), int(mm.get("sr", 24000)) + + audio_np, output_sample_rate = generate_audio() + assert audio_np.size > 1, "Audio output is empty" + assert np.isfinite(audio_np).all(), "Audio contains non-finite samples" + assert np.unique(audio_np).size > 1, "Audio has no sample variation" + rms = np.sqrt(np.mean(audio_np**2)) + assert rms > 0.01, f"Audio RMS too low ({rms:.4f}), likely white noise or silence" + assert output_sample_rate == 24000 + + repeated_audio, repeated_sample_rate = generate_audio() + assert repeated_sample_rate == output_sample_rate + assert repeated_audio.shape == audio_np.shape + np.testing.assert_allclose(repeated_audio, audio_np, rtol=1e-5, atol=1e-4) + + output_path = tmp_path / "omnivoice_ref_audio_only.wav" + sf.write(str(output_path), audio_np, output_sample_rate) + transcript = convert_audio_file_to_text(str(output_path), language="en").lower() + assert "hello" in transcript, f"Generated audio transcript does not contain 'hello': {transcript!r}" diff --git a/tests/e2e/online_serving/test_omnivoice_expansion.py b/tests/e2e/online_serving/test_omnivoice_expansion.py index c66750b7dce..b4b2c0eb899 100644 --- a/tests/e2e/online_serving/test_omnivoice_expansion.py +++ b/tests/e2e/online_serving/test_omnivoice_expansion.py @@ -14,7 +14,9 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +import numpy as np import pytest +import soundfile as sf from openai import BadRequestError as OpenAIBadRequestError from tests.helpers.mark import hardware_test @@ -60,6 +62,7 @@ ] # Lower this in ``request_config`` via ``min_audio_bytes`` if a run produces legitimately short WAVs. _DEFAULT_MIN_AUDIO_BYTES = 5000 +_OMNIVOICE_REF_AUDIO_SEED = 102 REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True) @@ -184,6 +187,8 @@ def test_speech_auto_voice_seed_deterministic(self, omni_server, online_client) r1 = online_client.send_audio_speech_request(cfg)[0] r2 = online_client.send_audio_speech_request(cfg)[0] + assert r1.audio_bytes is not None + assert r2.audio_bytes is not None assert r1.audio_bytes == r2.audio_bytes @hardware_test(res={"cuda": "L4"}, num_cards=1) @@ -210,16 +215,31 @@ class TestOmniVoiceVoiceCloning: @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_voice_clone_ref_audio_only(self, omni_server, online_client) -> None: - """Test voice cloning with ref_audio only (x_vector mode).""" + """Test automatic reference transcription with ref_audio only.""" request_config = { "model": omni_server.model, - "input": get_prompt("text"), + "input": "hello", "ref_audio": REF_AUDIO_URL, "response_format": "wav", + "seed": _OMNIVOICE_REF_AUDIO_SEED, "timeout": 180.0, "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES, + "transcript_language": "en", } - online_client.send_audio_speech_request(request_config) + response = online_client.send_audio_speech_request(request_config)[0] + assert response.audio_bytes is not None + audio, sample_rate = sf.read(io.BytesIO(response.audio_bytes), dtype="float32") + assert sample_rate == 24000 + assert np.isfinite(audio).all() + assert np.unique(audio).size > 1 + assert np.sqrt(np.mean(audio**2)) > 0.01 + + repeated_response = online_client.send_audio_speech_request(request_config)[0] + assert repeated_response.audio_bytes is not None + repeated_audio, repeated_sample_rate = sf.read(io.BytesIO(repeated_response.audio_bytes), dtype="float32") + assert repeated_sample_rate == sample_rate + assert repeated_audio.shape == audio.shape + np.testing.assert_allclose(repeated_audio, audio, rtol=1e-5, atol=1e-4) @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_voice_clone_ref_audio_and_text(self, omni_server, online_client) -> None: diff --git a/tests/engine/duplex/test_session_runner.py b/tests/engine/duplex/test_session_runner.py index 07710626c98..61da03f2182 100644 --- a/tests/engine/duplex/test_session_runner.py +++ b/tests/engine/duplex/test_session_runner.py @@ -376,6 +376,12 @@ def listen_output(request_id: str) -> SimpleNamespace: ) +def _is_silence_unit(submission: DuplexStageSubmission) -> bool: + """A continuation unit: the only append whose audio is all zeros.""" + payload = submission.prompt["model_intermediate_buffer"]["duplex"]["payload"] + return set(base64.b64decode(payload["audio"])) == {0} + + def types(events: Sequence[DuplexEvent]) -> list[str]: return [event.type for event in events] @@ -1915,6 +1921,81 @@ async def test_speech_after_silence_continuation_is_not_the_first_audio_window() await close_harness(h) +@pytest.mark.asyncio +async def test_answer_continues_when_the_segment_end_arrives_before_the_answer() -> None: + """A continuation planned for the model turn is re-aimed at the answer that opened (#7729).""" + h = await open_harness() + try: + for _ in range(4): + await h.run(append_audio()) + request_id = h.stage0_request_id() + + # A TTS segment ends with no text or audio before any response exists. + # This schedules a continuation, which waits for up to one chunk period. + h.deliver(tts_output(request_id, samples=0, text="", tts_is_last_chunk=True)) + await asyncio.sleep(0.1) + + # During that wait the last real units arrive and the answer starts. + h.submit(append_audio()) + await asyncio.sleep(0.1) + h.deliver(tts_output(request_id, samples=24000, text="no")) + await asyncio.sleep(0.1) + h.submit(append_audio()) + h.submit(commands.Commit()) + await h.settle(timeout_s=5.0) + + # More of the answer. Its segment only ends after another unit is processed. + for samples, text in ((48000, "no fi"), (72000, "no fight, bu"), (96000, "no fight, but yes,")): + await h.deliver_and_settle(tts_output(request_id, samples=samples, text=text)) + await h.settle(idle_s=2.5, timeout_s=8.0) + + assert h.session.active_response_id is not None + # 6 real units, then a silence unit that keeps the answer going. + assert len(h.port.submissions) == 7 + assert _is_silence_unit(h.port.submissions[6]) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_a_dropped_continuation_is_not_re_aimed_across_a_barge_in() -> None: + """The re-aim in #7729 stays inside the epoch that planned the unit. + + The next answer must open while the old continuation is still waiting out + its chunk period, so this drives the runner with ``submit``/``deliver`` and + short sleeps. ``run``/``settle`` would wait for that continuation to drop + first, and then a fix that re-arms with the current epoch passes too. + """ + h = await open_harness() + try: + for _ in range(2): + await h.run(append_audio()) + request_id = h.stage0_request_id() + + # Plan a continuation for the model turn of epoch 0. It waits for up to + # one chunk period (1 s) before it submits. + h.deliver(tts_output(request_id, samples=0, text="", tts_is_last_chunk=True)) + await asyncio.sleep(0.1) + + # The user barges in while it waits, and the next turn opens its own answer. + h.submit(commands.BargeIn()) + await asyncio.sleep(0.1) + assert h.session.epoch == 1 + h.submit(append_audio()) + await asyncio.sleep(0.1) + assert any(task.get_name() == "duplex-continue" and not task.done() for task in h.runner._background_tasks), ( + "the old continuation should still be waiting when the new answer opens" + ) + h.deliver(tts_output(h.stage0_request_id(), samples=24000, text="hi", epoch=1)) + await h.settle(idle_s=2.5, timeout_s=8.0) + + assert h.session.epoch == 1 + assert h.session.active_response_id is not None + assert [s for s in h.port.submissions if _is_silence_unit(s)] == [] + finally: + await close_harness(h) + + @pytest.mark.asyncio async def test_stage_failure_fails_the_active_response() -> None: h = await open_harness() diff --git a/tests/engine/test_duplex_omni_engine.py b/tests/engine/test_duplex_omni_engine.py index 18211f94b8c..ddd4197ec77 100644 --- a/tests/engine/test_duplex_omni_engine.py +++ b/tests/engine/test_duplex_omni_engine.py @@ -91,6 +91,27 @@ async def test_open_session_runs_the_correlated_rpc_with_the_typed_config() -> N assert engine.rpc_client.block_on_submit_flags == [True] +@pytest.mark.asyncio +async def test_open_uses_the_cold_start_budget_and_controls_keep_the_short_one() -> None: + # #9072: the first open on a cold engine also initializes per-session runner + # state, which can exceed the 10 s control budget. Only open gets the longer + # budget; control ops on a live session keep the short one. + from vllm_omni.engine.duplex_omni_engine import _DEFAULT_CONTROL_TIMEOUT_S, _DEFAULT_OPEN_TIMEOUT_S + + assert _DEFAULT_OPEN_TIMEOUT_S > _DEFAULT_CONTROL_TIMEOUT_S + output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512) + + open_engine = _engine( + _ok("open", capabilities=DuplexCapabilities(), public_session={"id": "sid"}, lease_generation=0) + ) + await open_engine.open_session_async("sid", DuplexSessionConfig(), output_buffer=output_buffer) + assert open_engine.rpc_client.calls[-1][2] == _DEFAULT_OPEN_TIMEOUT_S + + control_engine = _engine(_ok("close")) + await control_engine.close_session_async("sid") + assert control_engine.rpc_client.calls[-1][2] == _DEFAULT_CONTROL_TIMEOUT_S + + @pytest.mark.asyncio async def test_close_resume_and_touch_build_their_messages_without_incarnation() -> None: engine = _engine(_ok("close")) diff --git a/tests/engine/test_omni_request_output.py b/tests/engine/test_omni_request_output.py index 0179a221bdd..e3c3c696d1f 100644 --- a/tests/engine/test_omni_request_output.py +++ b/tests/engine/test_omni_request_output.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for OmniRequestOutput class.""" import numpy as np @@ -8,10 +8,32 @@ from vllm.outputs import CompletionOutput, RequestOutput from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.outputs.mm_outputs import MultimodalCompletionOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.mark.parametrize("payload_location", ["request", "completion", "missing"]) +def test_multimodal_output_audio_accessor(payload_location): + audio = {"audio": [0.1, -0.1], "sr": 24000} + output = OmniRequestOutput(request_id="audio-accessor") + if payload_location == "request": + output = OmniRequestOutput.from_diffusion(request_id="audio-accessor", images=[], multimodal_output=audio) + elif payload_location == "completion": + output.outputs = [ + MultimodalCompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=None, + logprobs=None, + multimodal_output=audio, + ) + ] + + assert output.multimodal_output == ({} if payload_location == "missing" else audio) + + # ── helpers ──────────────────────────────────────────────────────────────── diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 83c45b57bc3..6a063965b0f 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -473,12 +473,12 @@ def __init__( if 0 not in self.stage_ids: raise ValueError(f"Stage CLI test requires stage_id=0 in config: {stage_config_path}") self.stage_replica_counts = load_stage_replica_counts(resolved_cfg) - self.stage_procs: dict[tuple[int, int], subprocess.Popen] = {} - self._stage_log_paths: dict[tuple[int, int], Path] = {} - self._stage_log_files: dict[tuple[int, int], Any] = {} + self.stage_procs: dict[int, subprocess.Popen] = {} + self._stage_log_paths: dict[int, Path] = {} + self._stage_log_files: dict[int, Any] = {} self.proc = None - def _build_stage_cmd(self, stage_id: int, *, headless: bool, replica_id: int = 0) -> list[str]: + def _build_stage_cmd(self, stage_id: int, *, headless: bool) -> list[str]: cmd = [ sys.executable, "-m", @@ -494,10 +494,15 @@ def _build_stage_cmd(self, stage_id: int, *, headless: bool, replica_id: int = 0 self.host, "--omni-master-port", str(self.master_port), - "--replica-id", - str(replica_id), ] + # One process owns all local replicas of a stage: serve ignores + # ``--replica-id``, and only ``--omni-dp-size-local`` splits the stage + # ``devices`` pool so each replica gets its own card. + num_local_replicas = self.stage_replica_counts.get(stage_id, 1) + if num_local_replicas > 1: + cmd += ["--omni-dp-size-local", str(num_local_replicas)] + if headless: cmd.append("--headless") else: @@ -506,21 +511,20 @@ def _build_stage_cmd(self, stage_id: int, *, headless: bool, replica_id: int = 0 cmd += self.serve_args return cmd - def _launch_stage(self, stage_id: int, *, headless: bool, replica_id: int = 0) -> None: + def _launch_stage(self, stage_id: int, *, headless: bool) -> None: env = os.environ.copy() if self.env_dict is not None: env.update(self.env_dict) - cmd = self._build_stage_cmd(stage_id, headless=headless, replica_id=replica_id) - print(f"Launching OmniServerStageCli stage {stage_id} replica {replica_id}: {' '.join(cmd)}") + cmd = self._build_stage_cmd(stage_id, headless=headless) + print(f"Launching OmniServerStageCli stage {stage_id}: {' '.join(cmd)}") # Capture each subprocess's stdout+stderr to a per-stage log file so # debugging "Stage N exited before API server ready" doesn't rely on # guessing; the file is surfaced in the RuntimeError message. - log_path = Path(tempfile.gettempdir()) / f"omni_stage_{stage_id}_replica_{replica_id}_{self.master_port}.log" - stage_key = (stage_id, replica_id) - self._stage_log_paths[stage_key] = log_path + log_path = Path(tempfile.gettempdir()) / f"omni_stage_{stage_id}_{self.master_port}.log" + self._stage_log_paths[stage_id] = log_path log_fh = open(log_path, "w", buffering=1) # noqa: SIM115 - closed in ``__exit__`` - self._stage_log_files[stage_key] = log_fh + self._stage_log_files[stage_id] = log_fh proc = subprocess.Popen( cmd, env=env, @@ -531,40 +535,37 @@ def _launch_stage(self, stage_id: int, *, headless: bool, replica_id: int = 0) - stdout=log_fh, stderr=subprocess.STDOUT, ) - self.stage_procs[stage_key] = proc - if stage_id == 0 and replica_id == 0: + self.stage_procs[stage_id] = proc + if stage_id == 0: self.proc = proc def _ensure_stage_processes_alive(self) -> None: - for (stage_id, replica_id), proc in self.stage_procs.items(): + for stage_id, proc in self.stage_procs.items(): ret = proc.poll() if ret is not None: - log_path = self._stage_log_paths.get((stage_id, replica_id)) + log_path = self._stage_log_paths.get(stage_id) tail = "" if log_path and log_path.exists(): try: with open(log_path, encoding="utf-8", errors="replace") as f: lines = f.readlines() - tail = "\n=== Last 60 lines of stage {} replica {} log ({}) ===\n{}".format( - stage_id, replica_id, log_path, "".join(lines[-60:]) or "" + tail = "\n=== Last 60 lines of stage {} log ({}) ===\n{}".format( + stage_id, log_path, "".join(lines[-60:]) or "" ) except Exception as exc: # pragma: no cover - diagnostic only tail = f"\n" - raise RuntimeError( - f"Stage {stage_id} replica {replica_id} exited with code {ret} before API server became ready.{tail}" - ) + raise RuntimeError(f"Stage {stage_id} exited with code {ret} before API server became ready.{tail}") def _start_server(self) -> None: startup_t0 = time.perf_counter() ordered_stage_ids = [0, *[stage_id for stage_id in self.stage_ids if stage_id != 0]] - self._launch_stage(0, headless=False, replica_id=0) + self._launch_stage(0, headless=False) time.sleep(2) self._ensure_stage_processes_alive() for stage_id in ordered_stage_ids[1:]: - for replica_id in range(self.stage_replica_counts.get(stage_id, 1)): - self._launch_stage(stage_id, headless=True, replica_id=replica_id) + self._launch_stage(stage_id, headless=True) max_wait = self.startup_timeout # System clock corrections must not shorten or extend startup waits. @@ -598,15 +599,15 @@ def _dump_stage_logs_for_debug(self, head_lines: int = 300, tail_lines: int = 50 whatever state the stage was in when it was torn down. """ log_paths = self._stage_log_paths - for stage_id, replica_id in sorted(log_paths): - log_path = log_paths[(stage_id, replica_id)] + for stage_id in sorted(log_paths): + log_path = log_paths[stage_id] if not log_path or not log_path.exists(): continue try: with open(log_path, encoding="utf-8", errors="replace") as f: lines = f.readlines() except Exception as exc: # pragma: no cover - diagnostic only - print(f"[OmniServerStageCli] stage {stage_id} replica {replica_id} log read failed: {exc}", flush=True) + print(f"[OmniServerStageCli] stage {stage_id} log read failed: {exc}", flush=True) continue total = len(lines) if total <= head_lines + tail_lines: @@ -617,18 +618,18 @@ def _dump_stage_logs_for_debug(self, head_lines: int = 300, tail_lines: int = 50 head_chunk = lines[:head_lines] tail_chunk = lines[-tail_lines:] elided = total - head_lines - tail_lines - print(f"\n=== stage {stage_id} replica {replica_id} log HEAD ({log_path}) ===", flush=True) + print(f"\n=== stage {stage_id} log HEAD ({log_path}) ===", flush=True) print("".join(head_chunk).rstrip("\n"), flush=True) if tail_chunk: print(f"\n... [{elided} lines elided] ...", flush=True) - print(f"\n=== stage {stage_id} replica {replica_id} log TAIL ({log_path}) ===", flush=True) + print(f"\n=== stage {stage_id} log TAIL ({log_path}) ===", flush=True) print("".join(tail_chunk).rstrip("\n"), flush=True) - print(f"=== end stage {stage_id} replica {replica_id} log ===\n", flush=True) + print(f"=== end stage {stage_id} log ===\n", flush=True) def __exit__(self, exc_type, exc_val, exc_tb): self._dump_stage_logs_for_debug() - for stage_key in sorted(self.stage_procs, reverse=True): - proc = self.stage_procs[stage_key] + for stage_id in sorted(self.stage_procs, reverse=True): + proc = self.stage_procs[stage_id] if proc.poll() is None: self._kill_process_tree(proc.pid) @@ -636,15 +637,15 @@ def __exit__(self, exc_type, exc_val, exc_tb): keep_logs = os.environ.get("VLLM_OMNI_KEEP_LOG", "").lower() in ("1", "true", "yes") log_files = self._stage_log_files log_paths = self._stage_log_paths - for stage_key in set(log_files) | set(log_paths): - log_fh = log_files.get(stage_key) + for stage_id in set(log_files) | set(log_paths): + log_fh = log_files.get(stage_id) if log_fh is not None: try: log_fh.close() except Exception: pass if not keep_logs: - log_path = log_paths.get(stage_key) + log_path = log_paths.get(stage_id) if log_path is not None: try: Path(log_path).unlink(missing_ok=True) diff --git a/tests/helpers/tests/test_runtime_stage_cli.py b/tests/helpers/tests/test_runtime_stage_cli.py index 712a87e0936..c4fed6c1e68 100644 --- a/tests/helpers/tests/test_runtime_stage_cli.py +++ b/tests/helpers/tests/test_runtime_stage_cli.py @@ -9,6 +9,7 @@ from contextlib import nullcontext from pathlib import Path from types import SimpleNamespace +from unittest.mock import MagicMock import pytest @@ -37,14 +38,46 @@ def _write_stage_config(tmp_path: Path) -> str: def test_stage_cli_builds_headless_replica_cmd(tmp_path: Path) -> None: server = OmniServerStageCli("fake-model", _write_stage_config(tmp_path), ["--disable-log-stats"]) - cmd = server._build_stage_cmd(1, headless=True, replica_id=2) + cmd = server._build_stage_cmd(1, headless=True) assert "--headless" in cmd assert cmd[cmd.index("--stage-id") + 1] == "1" - assert cmd[cmd.index("--replica-id") + 1] == "2" + # ``--replica-id`` is ignored by serve; one process must own all local + # replicas so ``devices: "1,2,3"`` is split one card per replica. + assert "--replica-id" not in cmd + assert cmd[cmd.index("--omni-dp-size-local") + 1] == "3" assert cmd[cmd.index("--deploy-config") + 1] == server.stage_config_path +def test_stage_cli_omits_dp_size_local_for_single_replica_stage(tmp_path: Path) -> None: + server = OmniServerStageCli("fake-model", _write_stage_config(tmp_path), []) + + cmd = server._build_stage_cmd(0, headless=False) + + assert "--omni-dp-size-local" not in cmd + assert "--replica-id" not in cmd + + +def test_stage_cli_launches_one_process_per_stage(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + server = OmniServerStageCli("fake-model", _write_stage_config(tmp_path), []) + launched: list[tuple[int, bool]] = [] + + def _fake_launch(stage_id: int, *, headless: bool) -> None: + launched.append((stage_id, headless)) + + sock = MagicMock() + sock.__enter__.return_value = sock + sock.connect_ex.return_value = 0 + monkeypatch.setattr(server, "_launch_stage", _fake_launch) + monkeypatch.setattr(server, "_ensure_stage_processes_alive", lambda: None) + monkeypatch.setattr("tests.helpers.runtime.time.sleep", lambda *_args, **_kwargs: None) + monkeypatch.setattr(runtime.socket, "socket", lambda *args, **kwargs: sock) + + server._start_server() + + assert launched == [(0, False), (1, True)] + + def test_stage_cli_loads_stage_ids_and_replica_counts(tmp_path: Path) -> None: server = OmniServerStageCli("fake-model", _write_stage_config(tmp_path), []) @@ -102,16 +135,15 @@ def test_stage_cli_enter_rolls_back_launched_stages_on_startup_failure( killed: list[int] = [] monkeypatch.setattr(server, "_kill_process_tree", lambda pid: killed.append(pid)) - def _fake_launch(stage_id: int, *, headless: bool, replica_id: int = 0) -> None: - stage_key = (stage_id, replica_id) - server.stage_procs[stage_key] = proc # type: ignore[assignment] - server._stage_log_paths[stage_key] = log_path - server._stage_log_files[stage_key] = log_fh - if stage_id == 0 and replica_id == 0: + def _fake_launch(stage_id: int, *, headless: bool) -> None: + server.stage_procs[stage_id] = proc # type: ignore[assignment] + server._stage_log_paths[stage_id] = log_path + server._stage_log_files[stage_id] = log_fh + if stage_id == 0: server.proc = proc # type: ignore[assignment] def _fail_alive() -> None: - raise RuntimeError("Stage 0 replica 0 exited with code 1") + raise RuntimeError("Stage 0 exited with code 1") monkeypatch.setattr(server, "_launch_stage", _fake_launch) monkeypatch.setattr(server, "_ensure_stage_processes_alive", _fail_alive) diff --git a/tests/helpers/tests/test_runtime_startup_timeout.py b/tests/helpers/tests/test_runtime_startup_timeout.py index 9432a2bd113..85bfc0b20c7 100644 --- a/tests/helpers/tests/test_runtime_startup_timeout.py +++ b/tests/helpers/tests/test_runtime_startup_timeout.py @@ -29,8 +29,8 @@ def startup_server(request, tmp_path, monkeypatch): config.write_text("stages:\n - stage_id: 0\n - stage_id: 1\n") server = runtime.OmniServerStageCli("fake-model", str(config), []) - def launch(stage_id, *, headless, replica_id=0): - server.stage_procs[(stage_id, replica_id)] = proc + def launch(stage_id, *, headless): + server.stage_procs[stage_id] = proc monkeypatch.setattr(server, "_launch_stage", launch) return server diff --git a/tests/model_executor/models/fish_omni/test_fish_omni_model_state.py b/tests/model_executor/models/fish_omni/test_fish_omni_model_state.py index b6e88750480..b5ff5bd3585 100644 --- a/tests/model_executor/models/fish_omni/test_fish_omni_model_state.py +++ b/tests/model_executor/models/fish_omni/test_fish_omni_model_state.py @@ -213,3 +213,51 @@ def test_max_top_k_is_a_power_of_two_from_the_ras_top_k_to_the_codebook(value, o else: with pytest.raises(ValueError, match="power of two"): max_top_k(config, 4096) + + +def test_pipeline_capacity_preemption_restores_frames_after_slot_reuse(make_state): + from vllm.sampling_params import SamplingParams + from vllm.v1.request import Request, RequestStatus + + from tests.core.sched.test_omni_recompute_preemption import _output, _scheduler + from vllm_omni.model_executor.models.fish_omni.pipeline import FISH_OMNI_PIPELINE + + stage = FISH_OMNI_PIPELINE.get_stage(0) + scheduler = _scheduler(stage.recompute_preemption, block_size=4, num_blocks=3) + scheduler.use_v2_model_runner = True + state = make_state() + # Hold one of two usable blocks: prefill fits, first decode cannot extend. + blocker = Request("blocker", [1] * 4, SamplingParams(max_tokens=8), None) + assert scheduler.kv_cache_manager.allocate_slots(blocker, 4) is not None + request = Request("a", [1] * 4, SamplingParams(max_tokens=8), None) + scheduler.add_request(request) + prefill = scheduler.schedule() + state.add_request(1, prefill.scheduled_new_reqs[0]) + codes = torch.tensor([[5, 6, 7]], dtype=torch.int16) + state.buffers.slot_history[1, :1] = codes + state.buffers.slot_generated[1] = 1 + scheduler.update_from_output(prefill, _output(request)) + + preempt = scheduler.schedule() + assert request.status == RequestStatus.PREEMPTED + assert preempt.preempted_req_ids == {"a"} + assert scheduler.update_from_output(preempt, _output()) == {} + state.on_request_preempted("a", 1) + # Reuse the former slot for another request before resuming in a new slot. + state.add_request(1, _request("other")) + state.buffers.slot_history[1].fill_(99) + state.buffers.slot_generated[1] = 0 + scheduler.kv_cache_manager.free(blocker) + + resumed = scheduler.schedule() + assert request.status == RequestStatus.RUNNING + assert resumed.num_scheduled_tokens == {"a": 5} + state.add_request(2, resumed.scheduled_new_reqs[0]) + assert np.array_equal(state.prompts[2][:4], PROMPT.numpy()) + assert state.prompts[2][4:, 0].tolist() == list(request.output_token_ids) + assert np.array_equal(state.prompts[2][4:, 1:], codes.numpy()) + assert torch.equal(state._pending_history[2], codes) + assert state._pending_slots[2][1][1] == 1 + assert "a" not in state._suspended + outputs = scheduler.update_from_output(resumed, _output(request)) + assert all(output.finish_reason is None for batch in outputs.values() for output in batch.outputs) diff --git a/tests/model_executor/models/omnivoice/test_asr.py b/tests/model_executor/models/omnivoice/test_asr.py new file mode 100644 index 00000000000..c61e218730e --- /dev/null +++ b/tests/model_executor/models/omnivoice/test_asr.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from vllm_omni.diffusion.models.omnivoice import pipeline_omnivoice + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeASRModel: + def __init__(self): + self.to_calls = [] + + def to(self, device): + self.to_calls.append(device) + return self + + +class _FakeASRPipeline: + def __init__(self, result): + self.result = result + self.inputs = [] + self.model = _FakeASRModel() + self.device = torch.device("cpu") + + def __call__(self, audio_input): + self.inputs.append(audio_input) + return self.result + + +@pytest.fixture +def pipeline() -> pipeline_omnivoice.OmniVoicePipeline: + model = pipeline_omnivoice.OmniVoicePipeline.__new__(pipeline_omnivoice.OmniVoicePipeline) + torch.nn.Module.__init__(model) + model.device = torch.device("cpu") + model._asr_model_name = pipeline_omnivoice._ASR_MODEL_NAME + model._asr_device = "cpu" + model._asr_pipeline = None + return model + + +def _install_fake_asr(monkeypatch, result=None): + asr = _FakeASRPipeline(result={"text": " hello there "} if result is None else result) + load_calls = [] + + def load_pipeline(*args, **kwargs): + load_calls.append((args, kwargs)) + return asr + + monkeypatch.setattr(pipeline_omnivoice, "hf_pipeline", load_pipeline) + return asr, load_calls + + +def test_missing_reference_text_loads_asr_and_strips_transcript(pipeline, monkeypatch): + asr, load_calls = _install_fake_asr(monkeypatch) + + transcript = pipeline._resolve_ref_text((np.zeros((1, 4), dtype=np.float32), 16000), None) + + assert transcript == "hello there" + assert len(load_calls) == 1 + assert asr.inputs[0]["sampling_rate"] == 16000 + assert asr.inputs[0]["array"].shape == (4,) + + +def test_second_missing_reference_text_request_reuses_asr(pipeline, monkeypatch): + asr, load_calls = _install_fake_asr(monkeypatch) + ref_audio = (np.zeros(4, dtype=np.float32), 16000) + + assert pipeline._resolve_ref_text(ref_audio, None) == "hello there" + assert pipeline._resolve_ref_text(ref_audio, " ") == "hello there" + + assert len(load_calls) == 1 + assert len(asr.inputs) == 2 + + +def test_explicit_reference_text_bypasses_asr(pipeline, monkeypatch): + _install_fake_asr(monkeypatch) + ref_text = " caller supplied transcript " + + assert pipeline._resolve_ref_text((np.zeros(4), 16000), ref_text) == ref_text + assert pipeline._asr_pipeline is None + + +def test_text_only_input_bypasses_asr(pipeline, monkeypatch): + _install_fake_asr(monkeypatch) + + assert pipeline._resolve_ref_text(None, None) is None + assert pipeline._asr_pipeline is None + + +@pytest.mark.parametrize("ref_text", [None, "", " ", "\t\n"]) +def test_empty_reference_text_uses_asr(pipeline, monkeypatch, ref_text): + _install_fake_asr(monkeypatch) + + assert pipeline._resolve_ref_text((np.zeros(4), 16000), ref_text) == "hello there" + + +def test_empty_asr_result_is_rejected(pipeline, monkeypatch): + _install_fake_asr(monkeypatch, result={"text": " \t"}) + + with pytest.raises(ValueError, match="empty reference transcription"): + pipeline._resolve_ref_text((np.zeros(4), 16000), None) + + +def test_malformed_asr_result_is_rejected(pipeline, monkeypatch): + _install_fake_asr(monkeypatch, result={"chunks": []}) + + with pytest.raises(RuntimeError, match="without a 'text' field"): + pipeline._resolve_ref_text((np.zeros(4), 16000), None) + + +def test_asr_load_failure_includes_checkpoint_and_device(pipeline, monkeypatch): + def load_pipeline(*args, **kwargs): + raise OSError("download failed") + + monkeypatch.setattr(pipeline_omnivoice, "hf_pipeline", load_pipeline) + + with pytest.raises(RuntimeError) as exc_info: + pipeline._load_asr_pipeline() + + message = str(exc_info.value) + assert pipeline_omnivoice._ASR_MODEL_NAME in message + assert "cpu" in message + assert "download failed" in message + + +def test_asr_device_placement_failure_is_reported_and_can_be_retried(pipeline, monkeypatch): + asr, load_calls = _install_fake_asr(monkeypatch) + pipeline._asr_device = "cuda:1" + + def move_model(device): + raise RuntimeError(f"cannot move to {device}") + + monkeypatch.setattr(asr.model, "to", move_model) + + with pytest.raises(RuntimeError) as exc_info: + pipeline._load_asr_pipeline() + + message = str(exc_info.value) + assert pipeline_omnivoice._ASR_MODEL_NAME in message + assert "cuda:1" in message + assert "cannot move" in message + assert pipeline._asr_pipeline is None + + monkeypatch.setattr(asr.model, "to", lambda device: asr.model) + assert pipeline._load_asr_pipeline() is asr + assert len(load_calls) == 2 + assert asr.device == torch.device("cuda:1") + + +@pytest.mark.parametrize( + ("additional_config", "expected"), + [ + ({}, (False, pipeline_omnivoice._ASR_MODEL_NAME, None)), + ( + {"omnivoice_asr": {"load_asr_on_startup": False, "asr_device": None}}, + (False, pipeline_omnivoice._ASR_MODEL_NAME, None), + ), + ( + { + "omnivoice_asr": { + "load_asr_on_startup": True, + "asr_model_name": "local/whisper", + "asr_device": "cuda:1", + } + }, + (True, "local/whisper", "cuda:1"), + ), + ], +) +def test_asr_config_defaults_and_values(additional_config, expected): + assert pipeline_omnivoice._parse_asr_config(additional_config) == expected + + +@pytest.mark.parametrize( + "additional_config", + [ + {"omnivoice_asr": {"load_asr_on_startup": "true"}}, + {"omnivoice_asr": {"asr_model_name": " "}}, + {"omnivoice_asr": {"asr_device": ""}}, + {"omnivoice_asr": []}, + {"omnivoice_asr": None}, + ], +) +def test_asr_config_rejects_invalid_values(additional_config): + with pytest.raises((TypeError, ValueError)): + pipeline_omnivoice._parse_asr_config(additional_config) + + +def test_asr_loader_uses_configured_model_and_device(pipeline, monkeypatch): + asr, load_calls = _install_fake_asr(monkeypatch) + pipeline._asr_model_name = "local/whisper" + pipeline._asr_device = "cpu" + + pipeline._load_asr_pipeline() + + assert asr is pipeline._asr_pipeline + assert asr.model.to_calls == [torch.device("cpu")] + assert asr.device == torch.device("cpu") + assert load_calls == [ + ( + ("automatic-speech-recognition",), + {"model": "local/whisper", "dtype": torch.float32, "device": "cpu"}, + ) + ] + + +def test_eager_asr_is_loaded_during_initialization(monkeypatch): + model = pipeline_omnivoice.OmniVoicePipeline.__new__(pipeline_omnivoice.OmniVoicePipeline) + torch.nn.Module.__init__(model) + load_calls = [] + + def load_asr(self): + load_calls.append((self._asr_model_name, self._asr_device)) + + monkeypatch.setattr(pipeline_omnivoice.OmniVoicePipeline, "_load_asr_pipeline", load_asr) + + model._initialize_asr( + { + "omnivoice_asr": { + "load_asr_on_startup": True, + "asr_model_name": "local/whisper", + "asr_device": "cpu", + } + } + ) + + assert load_calls == [("local/whisper", "cpu")] + + +def test_lazy_asr_is_not_loaded_during_initialization(monkeypatch): + model = pipeline_omnivoice.OmniVoicePipeline.__new__(pipeline_omnivoice.OmniVoicePipeline) + torch.nn.Module.__init__(model) + load_calls = [] + + def load_asr(self): + load_calls.append(True) + + monkeypatch.setattr(pipeline_omnivoice.OmniVoicePipeline, "_load_asr_pipeline", load_asr) + + model._initialize_asr({"omnivoice_asr": {"load_asr_on_startup": False}}) + + assert load_calls == [] + + +def test_lazy_asr_uses_worker_device_on_first_request(pipeline, monkeypatch): + asr, load_calls = _install_fake_asr(monkeypatch) + pipeline._asr_device = None + + assert pipeline._resolve_ref_text((np.zeros(4, dtype=np.float32), 16000), None) == "hello there" + assert asr is pipeline._asr_pipeline + assert load_calls[0][1]["device"] == torch.device("cpu") + + +@pytest.mark.parametrize( + "waveform", + [ + np.arange(4, dtype=np.float32)[None, :], + torch.arange(4, dtype=torch.float32)[None, :], + ], +) +def test_asr_input_accepts_numpy_and_torch_waveforms_without_mutating(waveform, pipeline, monkeypatch): + asr, _ = _install_fake_asr(monkeypatch) + original = waveform.clone() if isinstance(waveform, torch.Tensor) else waveform.copy() + + pipeline._transcribe_ref_audio((waveform, 22050)) + + asr_input = asr.inputs[0] + assert asr_input["array"].shape == (4,) + np.testing.assert_array_equal(asr_input["array"], np.arange(4, dtype=np.float32)) + asr_input["array"][0] = 100 + if isinstance(waveform, torch.Tensor): + assert torch.equal(waveform, original) + else: + np.testing.assert_array_equal(waveform, original) + + +@pytest.mark.parametrize( + ("device", "dtype"), + [ + (torch.device("cpu"), torch.float32), + (torch.device("cuda:0"), torch.float16), + ("xpu:0", torch.float16), + ], +) +def test_loader_receives_checkpoint_device_and_dtype(pipeline, monkeypatch, device, dtype): + asr, load_calls = _install_fake_asr(monkeypatch) + pipeline._asr_device = str(device) + + pipeline._load_asr_pipeline() + + assert load_calls[0][1]["model"] == "openai/whisper-large-v3-turbo" + assert load_calls[0][1]["device"] == str(device) + assert load_calls[0][1]["dtype"] is dtype + assert asr.model.to_calls == [torch.device(device)] + assert asr.device == torch.device(device) + + +def test_transcript_is_combined_before_target_text(pipeline, monkeypatch): + _install_fake_asr(monkeypatch) + ref_audio = (np.zeros(4, dtype=np.float32), 16000) + + ref_text = pipeline._resolve_ref_text(ref_audio, None) + + assert pipeline_omnivoice._combine_text("hello", ref_text) == "hello there hello" + + +def test_voice_clone_duration_uses_resolved_text_and_reference_token_count(pipeline): + calls = [] + + class DurationEstimator: + def estimate_duration(self, text, ref_text, ref_duration): + calls.append((text, ref_text, ref_duration)) + return 17.8 + + pipeline.duration_estimator = DurationEstimator() + ref_audio_tokens = torch.zeros(8, 73, dtype=torch.long) + + assert pipeline._estimate_target_len("target", "transcribed reference", ref_audio_tokens) == 17 + assert calls == [("target", "transcribed reference", 73)] + + +def test_text_only_duration_keeps_fallback_inputs(pipeline): + calls = [] + + class DurationEstimator: + def estimate_duration(self, text, ref_text, ref_duration): + calls.append((text, ref_text, ref_duration)) + return 4 + + pipeline.duration_estimator = DurationEstimator() + + assert pipeline._estimate_target_len("target", None, None) == 4 + assert calls == [("target", "Nice to meet you.", 25)] diff --git a/tests/model_executor/models/omnivoice/test_audio_processing.py b/tests/model_executor/models/omnivoice/test_audio_processing.py new file mode 100644 index 00000000000..39909940255 --- /dev/null +++ b/tests/model_executor/models/omnivoice/test_audio_processing.py @@ -0,0 +1,609 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from collections import OrderedDict +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import soundfile as sf +import torch + +from vllm_omni.diffusion.models.omnivoice import audio as audio_utils +from vllm_omni.diffusion.models.omnivoice import pipeline_omnivoice +from vllm_omni.diffusion.models.omnivoice.audio import ( + PreparedReferenceAudio, + add_reference_punctuation, + postprocess_generated_audio, + prepare_reference_audio, + remove_silence, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.diffusion.worker.utils import StepRequestState +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_SAMPLE_RATE = 24000 +_HOP_LENGTH = 960 + + +def _prepare(waveform, sample_rate=_SAMPLE_RATE): + return prepare_reference_audio( + waveform, + sample_rate, + target_sample_rate=_SAMPLE_RATE, + hop_length=_HOP_LENGTH, + trim_long=False, + ) + + +def test_reference_audio_normalization_and_hop_alignment(): + waveform = np.full(1010, 0.05, dtype=np.float32) + + prepared = _prepare(waveform) + + assert prepared.waveform.shape == (1, _HOP_LENGTH) + assert prepared.original_rms == pytest.approx(0.05, abs=1e-6) + assert np.sqrt(np.mean(prepared.waveform**2)) == pytest.approx(0.1, abs=1e-4) + + +def test_reference_audio_normalization_clips_peaky_quiet_waveform(): + waveform = np.zeros(7056, dtype=np.float32) + waveform[waveform.size // 2] = 1.0 + + prepared = _prepare(waveform) + + assert prepared.original_rms == pytest.approx(0.01190476, abs=1e-8) + assert np.max(np.abs(prepared.waveform)) == 1.0 + + +def test_reference_audio_at_or_above_target_rms_is_not_rescaled(): + prepared = _prepare(np.full(960, 0.2, dtype=np.float32)) + + assert prepared.original_rms == pytest.approx(0.2, abs=1e-6) + assert np.sqrt(np.mean(prepared.waveform**2)) == pytest.approx(0.2, abs=1e-4) + + +def test_reference_audio_resamples_and_downmixes(): + waveform = np.full((2, 800), 0.1, dtype=np.float32) + + prepared = _prepare(waveform, sample_rate=16000) + + assert prepared.waveform.shape == (1, 960) + assert prepared.sample_rate == _SAMPLE_RATE + + +@pytest.mark.parametrize("waveform", [np.full(960, 0.1, dtype=np.float32), torch.full((960,), 0.1)]) +def test_reference_audio_does_not_mutate_input(waveform): + original = waveform.clone() if isinstance(waveform, torch.Tensor) else waveform.copy() + + _prepare(waveform) + + if isinstance(waveform, torch.Tensor): + assert torch.equal(waveform, original) + else: + np.testing.assert_array_equal(waveform, original) + + +def test_reference_audio_rejects_empty_after_silence_removal(): + with pytest.raises(ValueError, match="empty after silence removal"): + _prepare(np.zeros(960, dtype=np.float32)) + + +def test_reference_audio_rejects_channels_that_cancel_when_downmixed(): + waveform = np.stack([np.full(960, 0.1), np.full(960, -0.1)]) + + with pytest.raises(ValueError, match="empty after silence removal"): + _prepare(waveform) + + +def test_reference_audio_matches_official_asset_preparation(): + asset_path = Path(__file__).resolve().parents[4] / "tests/assets/qwen3_tts/clone_2.wav" + waveform, sample_rate = sf.read(asset_path, always_2d=False) + + prepared = prepare_reference_audio( + waveform, + sample_rate, + target_sample_rate=_SAMPLE_RATE, + hop_length=_HOP_LENGTH, + trim_long=True, + ) + + assert prepared.waveform.shape == (1, 179520) + assert prepared.original_rms == pytest.approx(0.07468347, abs=1e-5) + + +@pytest.mark.parametrize( + ("text", "expected"), + [ + ("hello", "hello."), + ("你好", "你好。"), + ("hello!", "hello!"), + ("hello,", "hello,"), + ("hello]", "hello]"), + ("", ""), + ], +) +def test_reference_punctuation(text, expected): + assert add_reference_punctuation(text) == expected + + +def test_remove_silence_uses_reference_thresholds(): + waveform = np.concatenate( + [ + np.zeros(2400, dtype=np.float32), + np.full(12000, 0.2, dtype=np.float32), + np.zeros(9600, dtype=np.float32), + np.full(12000, 0.2, dtype=np.float32), + np.zeros(7200, dtype=np.float32), + ] + )[np.newaxis, :] + + processed = remove_silence( + waveform, + _SAMPLE_RATE, + middle_silence_ms=200, + leading_silence_ms=100, + trailing_silence_ms=200, + ) + + assert processed.shape[-1] < waveform.shape[-1] + assert processed.shape[-1] > 24000 + + +@pytest.mark.parametrize( + ("gap_ms", "expected_removed_ms"), + [(199, 0), (200, 0), (500, 100)], +) +def test_remove_silence_preserves_middle_gap_boundary(gap_ms, expected_removed_ms): + active = np.full(_SAMPLE_RATE, 0.2, dtype=np.float32) + gap = np.zeros(round(gap_ms * _SAMPLE_RATE / 1000), dtype=np.float32) + waveform = np.concatenate([active, gap, active])[np.newaxis, :] + + processed = remove_silence( + waveform, + _SAMPLE_RATE, + middle_silence_ms=200, + leading_silence_ms=0, + trailing_silence_ms=0, + ) + + removed_ms = round(1000 * (waveform.shape[-1] - processed.shape[-1]) / _SAMPLE_RATE) + assert removed_ms == expected_removed_ms + + +@pytest.mark.parametrize( + ("gap_level", "is_silent"), + [ + (10 ** (-50 / 20) * 0.99, True), + (10 ** (-50 / 20), True), + (10 ** (-50 / 20) * 1.01, False), + ], +) +def test_remove_silence_uses_reference_threshold(gap_level, is_silent): + active = np.full(_SAMPLE_RATE, 0.2, dtype=np.float32) + gap = np.full(round(250 * _SAMPLE_RATE / 1000), gap_level, dtype=np.float32) + waveform = np.concatenate([active, gap, active])[np.newaxis, :] + + silent_ranges = audio_utils._detect_silence_ranges( + waveform, + _SAMPLE_RATE, + min_silence_ms=200, + silence_threshold_db=-50, + seek_step_ms=10, + ) + + assert bool(silent_ranges) is is_silent + + +@pytest.mark.parametrize("duration_s", [19.9, 20.0, 20.1]) +def test_trim_long_audio_only_runs_above_twenty_seconds(duration_s): + waveform = np.full((1, round(duration_s * _SAMPLE_RATE)), 0.2, dtype=np.float32) + + trimmed = audio_utils.trim_long_audio(waveform, _SAMPLE_RATE) + + if duration_s <= 20.0: + assert trimmed.shape == waveform.shape + else: + assert trimmed.shape[-1] == 15 * _SAMPLE_RATE + + +def test_generated_audio_scales_with_reference_rms_and_pads(): + audio = np.full((1, 10000), 0.2, dtype=np.float32) + + processed = postprocess_generated_audio( + audio, + sample_rate=_SAMPLE_RATE, + reference_rms=0.05, + ) + + # NumPy keeps the source sample count and uses sample-based padding. + assert processed.shape[-1] == 14800 + assert processed[0, 5000] == pytest.approx(0.1, abs=1e-5) + assert processed[0, 0] == 0.0 + assert processed[0, -1] == 0.0 + + +def test_generated_audio_keeps_reference_rms_at_or_above_target(): + processed = postprocess_generated_audio( + np.full((1, 10000), 0.2, dtype=np.float32), + sample_rate=_SAMPLE_RATE, + reference_rms=0.2, + ) + + assert processed[0, 5000] == pytest.approx(0.2, abs=3e-5) + + +def test_generated_audio_uses_output_silence_threshold(monkeypatch): + calls = [] + + def fake_remove_silence(audio, sample_rate, **kwargs): + calls.append((sample_rate, kwargs)) + return audio + + monkeypatch.setattr(audio_utils, "remove_silence", fake_remove_silence) + audio_utils.postprocess_generated_audio( + np.ones((1, 1000), dtype=np.float32), + sample_rate=_SAMPLE_RATE, + reference_rms=0.05, + ) + + assert calls == [ + ( + _SAMPLE_RATE, + { + "middle_silence_ms": 500, + "leading_silence_ms": 100, + "trailing_silence_ms": 100, + }, + ) + ] + + +def test_generated_zero_audio_remains_finite(): + processed = postprocess_generated_audio( + np.zeros((1, 1000), dtype=np.float32), + sample_rate=_SAMPLE_RATE, + reference_rms=0.05, + ) + + assert np.isfinite(processed).all() + + +class _FakeASR: + def __init__(self, text: str): + self.text = text + self.inputs: list[dict[str, object]] = [] + + def __call__(self, audio_input): + self.inputs.append(audio_input) + return {"text": self.text} + + +class _FakeGenerator: + def __init__(self, num_codebooks: int): + self.num_codebooks = num_codebooks + + def __call__(self, **kwargs): + target_len = sum(kwargs["target_lens"]) + return torch.zeros((1, self.num_codebooks, target_len), dtype=torch.long) + + +class _FakeDecoder: + def __call__(self, tokens): + return torch.full((1, 1, 10000), 0.2, dtype=torch.float32) + + +class _FakeDurationEstimator: + def __init__(self): + self.calls = [] + + def estimate_duration(self, text, ref_text, ref_audio_tokens): + self.calls.append((text, ref_text, ref_audio_tokens)) + return 4 + + +def _build_fake_pipeline(monkeypatch, prepared_waveform, *, asr_text="reference transcript"): + model = pipeline_omnivoice.OmniVoicePipeline.__new__(pipeline_omnivoice.OmniVoicePipeline) + torch.nn.Module.__init__(model) + model.device = torch.device("cpu") + model.config = SimpleNamespace(num_audio_codebook=2, audio_mask_id=-1) + model.audio_tokenizer = SimpleNamespace(config=SimpleNamespace(sample_rate=_SAMPLE_RATE, hop_length=_HOP_LENGTH)) + model.tokenizer = SimpleNamespace(encode=lambda text: SimpleNamespace(ids=[1, 2])) + model.generator = _FakeGenerator(num_codebooks=2) + model.decoder = _FakeDecoder() + model.duration_estimator = _FakeDurationEstimator() + model._asr_pipeline = _FakeASR(asr_text) + model.num_step = 1 + model.guidance_scale = 1.0 + model.t_shift = 1.0 + model.layer_penalty_factor = 0.0 + model.position_temperature = 1.0 + model.class_temperature = 1.0 + model.sample_rate = _SAMPLE_RATE + model._inline_reference_cache = OrderedDict() + + prepare_calls = [] + encoded_calls = [] + + def fake_prepare(waveform, sample_rate, **kwargs): + prepare_calls.append((waveform, sample_rate, kwargs)) + return PreparedReferenceAudio(prepared_waveform, _SAMPLE_RATE, 0.07) + + def fake_encode(self, audio_signal, sample_rate): + encoded_calls.append((audio_signal.detach().clone(), sample_rate)) + return torch.zeros((2, 3), dtype=torch.long) + + monkeypatch.setattr(pipeline_omnivoice, "prepare_reference_audio", fake_prepare) + monkeypatch.setattr(pipeline_omnivoice.OmniVoicePipeline, "_encode_ref_audio", fake_encode) + return model, prepare_calls, encoded_calls + + +def _request(*prompts): + return DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt=prompt, + sampling_params=OmniDiffusionSamplingParams(), + request_id=str(index), + ) + for index, prompt in enumerate(prompts) + ], + ) + + +def test_pipeline_uses_one_prepared_waveform_for_asr_and_tokenizer(monkeypatch): + prepared = np.arange(_HOP_LENGTH, dtype=np.float32)[np.newaxis, :] + model, prepare_calls, encoded_calls = _build_fake_pipeline(monkeypatch, prepared) + + result = model.forward( + _request( + { + "prompt": "hello", + "multi_modal_data": {"audio": (np.ones(1000, dtype=np.float32), 16000)}, + } + ) + ) + + assert result[0].error is None + assert len(prepare_calls) == 1 + assert prepare_calls[0][1] == 16000 + assert prepare_calls[0][2]["trim_long"] is True + assert prepare_calls[0][2]["target_sample_rate"] == _SAMPLE_RATE + assert prepare_calls[0][2]["hop_length"] == _HOP_LENGTH + assert len(model._asr_pipeline.inputs) == 1 + np.testing.assert_array_equal(model._asr_pipeline.inputs[0]["array"], prepared[0]) + assert model._asr_pipeline.inputs[0]["sampling_rate"] == _SAMPLE_RATE + assert len(encoded_calls) == 1 + torch.testing.assert_close(encoded_calls[0][0], torch.from_numpy(prepared)) + assert encoded_calls[0][1] == _SAMPLE_RATE + assert model.duration_estimator.calls == [("hello", "reference transcript.", 3)] + + +def test_pipeline_explicit_reference_text_prepares_audio_without_asr(monkeypatch): + prepared = np.ones((1, _HOP_LENGTH), dtype=np.float32) + model, prepare_calls, encoded_calls = _build_fake_pipeline(monkeypatch, prepared) + + result = model.forward( + _request( + { + "prompt": "hello", + "ref_audio": (np.ones(1000, dtype=np.float32), 16000), + "ref_text": "caller supplied transcript", + } + ) + ) + + assert result[0].error is None + assert model._asr_pipeline.inputs == [] + assert prepare_calls[0][2]["trim_long"] is False + assert len(encoded_calls) == 1 + assert model.duration_estimator.calls == [("hello", "caller supplied transcript.", 3)] + + +def test_pipeline_text_only_skips_reference_audio_processing(monkeypatch): + model, prepare_calls, encoded_calls = _build_fake_pipeline( + monkeypatch, + np.ones((1, _HOP_LENGTH), dtype=np.float32), + ) + + result = model.forward(_request("hello")) + + assert result[0].error is None + torch.testing.assert_close(result[0].output, torch.full((1, 1, 10000), 0.2)) + assert prepare_calls == [] + assert encoded_calls == [] + assert model._asr_pipeline.inputs == [] + assert model.duration_estimator.calls == [("hello", "Nice to meet you.", 25)] + + +def test_pipeline_text_only_bypasses_audio_postprocessing(monkeypatch): + model, _, _ = _build_fake_pipeline( + monkeypatch, + np.ones((1, _HOP_LENGTH), dtype=np.float32), + ) + + def fail_postprocess(*args, **kwargs): + raise AssertionError("text-only output must bypass voice-cloning postprocessing") + + monkeypatch.setattr(pipeline_omnivoice, "postprocess_generated_audio", fail_postprocess) + + result = model.forward(_request("hello")) + + assert result[0].error is None + + +def test_pipeline_reuses_inline_reference_cache(monkeypatch): + prepared = np.arange(_HOP_LENGTH, dtype=np.float32)[np.newaxis, :] + model, prepare_calls, encoded_calls = _build_fake_pipeline(monkeypatch, prepared) + request = _request( + { + "prompt": "hello", + "ref_audio": (np.ones(1000, dtype=np.float32), 16000), + } + ) + + first = model.forward(request) + second = model.forward(request) + + assert first[0].error is None + assert second[0].error is None + assert len(prepare_calls) == 2 + assert len(model._asr_pipeline.inputs) == 1 + assert len(encoded_calls) == 1 + + +def test_pipeline_keeps_explicit_and_asr_cache_entries_separate(monkeypatch): + prepared = np.arange(_HOP_LENGTH, dtype=np.float32)[np.newaxis, :] + model, prepare_calls, encoded_calls = _build_fake_pipeline(monkeypatch, prepared) + explicit_request = _request( + { + "prompt": "hello", + "ref_audio": (np.ones(1000, dtype=np.float32), 16000), + "ref_text": "caller supplied transcript", + } + ) + missing_request = _request( + { + "prompt": "hello", + "ref_audio": (np.ones(1000, dtype=np.float32), 16000), + } + ) + + model.forward(explicit_request) + model.forward(missing_request) + model.forward(missing_request) + + assert len(prepare_calls) == 3 + assert len(encoded_calls) == 2 + assert len(model._asr_pipeline.inputs) == 1 + assert {key[0] for key in model._inline_reference_cache} == {"explicit", "asr"} + + +def test_named_cache_separates_explicit_and_asr_preparation(monkeypatch): + prepared = np.arange(_HOP_LENGTH, dtype=np.float32)[np.newaxis, :] + model, prepare_calls, encoded_calls = _build_fake_pipeline(monkeypatch, prepared) + model._speaker_cache = SpeakerEmbeddingCache(max_bytes=1024 * 1024) + explicit_request = _request( + { + "prompt": "hello", + "ref_audio": (np.ones(1000, dtype=np.float32), 16000), + "ref_text": "caller supplied transcript", + "voice_name": "alice", + "voice_created_at": 1, + } + ) + missing_request = _request( + { + "prompt": "hello", + "ref_audio": (np.ones(1000, dtype=np.float32), 16000), + "voice_name": "alice", + "voice_created_at": 1, + } + ) + + model.forward(explicit_request) + model.forward(missing_request) + model.forward(missing_request) + + assert model._speaker_cache.stats()["entries"] == 2 + assert len(prepare_calls) == 2 + assert len(encoded_calls) == 2 + assert len(model._asr_pipeline.inputs) == 1 + + +def test_inline_reference_cache_eviction_is_bounded(monkeypatch): + model, _, _ = _build_fake_pipeline( + monkeypatch, + np.ones((1, _HOP_LENGTH), dtype=np.float32), + ) + + for index in range(pipeline_omnivoice._INLINE_CACHE_MAX_ENTRIES + 1): + model._put_inline_cache(("asr", index), {"index": index}) + + assert len(model._inline_reference_cache) == pipeline_omnivoice._INLINE_CACHE_MAX_ENTRIES + assert ("asr", 0) not in model._inline_reference_cache + assert ("asr", pipeline_omnivoice._INLINE_CACHE_MAX_ENTRIES) in model._inline_reference_cache + + +def test_inline_reference_cache_key_preserves_original_rms(): + waveform = np.ones((1, _HOP_LENGTH), dtype=np.float32) + quiet = PreparedReferenceAudio(waveform, _SAMPLE_RATE, 0.05) + loud = PreparedReferenceAudio(waveform, _SAMPLE_RATE, 0.08) + + assert pipeline_omnivoice.OmniVoicePipeline._inline_cache_key( + quiet, "asr" + ) != pipeline_omnivoice.OmniVoicePipeline._inline_cache_key(loud, "asr") + + +@pytest.mark.parametrize("cache_kind", ["inline", "named"]) +def test_batch_keeps_reference_rms_per_request_and_on_cache_hits(monkeypatch, cache_kind): + model, _, encoded_calls = _build_fake_pipeline(monkeypatch, np.ones((1, _HOP_LENGTH), dtype=np.float32)) + model._speaker_cache = SpeakerEmbeddingCache(max_bytes=1024 * 1024) + + def prepare(waveform, sample_rate, **kwargs): + return PreparedReferenceAudio(waveform[np.newaxis, :], sample_rate, float(waveform[0])) + + monkeypatch.setattr(pipeline_omnivoice, "prepare_reference_audio", prepare) + prompts = ["hello", {"prompt": ""}] + for index, rms in enumerate((0.03, 0.08)): + prompt = { + "prompt": "hello", + "ref_audio": (np.full(_HOP_LENGTH, rms, dtype=np.float32), _SAMPLE_RATE), + } + if cache_kind == "named": + prompt["voice_name"] = f"speaker-{index}" + prompts.append(prompt) + + for _ in range(2): + outputs = model.forward(_request(*prompts)) + assert len(outputs) == 4 + torch.testing.assert_close(outputs[0].output, torch.full((1, 1, 10000), 0.2)) + assert outputs[1].error == "Empty text prompt" + for output, rms in zip(outputs[2:], (0.03, 0.08)): + assert output.error is None + assert output.output.shape == (1, 1, 14800) + assert output.output[0, 0, 5000].item() == pytest.approx(0.2 * rms / 0.1, abs=1e-5) + assert len(encoded_calls) == 2 + assert len(model._asr_pipeline.inputs) == 2 + + +@pytest.mark.parametrize("clone", [False, True]) +def test_step_execution_preserves_reference_postprocessing(monkeypatch, clone): + model, prepare_calls, _ = _build_fake_pipeline(monkeypatch, np.ones((1, _HOP_LENGTH), dtype=np.float32)) + prompt: str | dict = "hello" + if clone: + prompt = { + "prompt": "hello", + "ref_audio": (np.ones(_HOP_LENGTH, dtype=np.float32), _SAMPLE_RATE), + } + else: + + def fail_postprocess(*args, **kwargs): + raise AssertionError("text-only output must bypass voice-cloning postprocessing") + + monkeypatch.setattr(pipeline_omnivoice, "postprocess_generated_audio", fail_postprocess) + + state = StepRequestState( + request_id="step-request", + prompt=prompt, + sampling=OmniDiffusionSamplingParams(num_inference_steps=1), + ) + model.prepare_encode(state) + output = model.post_decode(state) + + assert output.error is None + if clone: + assert state.extra["reference_rms"] == 0.07 + assert output.output.shape == (1, 1, 14800) + assert output.output[0, 0, 5000].item() == pytest.approx(0.14, abs=1e-5) + else: + assert state.extra["reference_rms"] is None + assert prepare_calls == [] + torch.testing.assert_close(output.output, torch.full((1, 1, 10000), 0.2)) diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_preempt_replay.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_preempt_replay.py new file mode 100644 index 00000000000..301f41bfca1 --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_preempt_replay.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Why the Qwen3-Omni talker declares ``recompute_preemption="fail"``. + +A talker decode step feeds ``summed_embeddings(layer0..7) + text_step`` into +PagedAttention, but only the layer-0 codec id is kept in ``output_token_ids``. +When ``_preempt_request()`` zeroes ``num_computed_tokens`` the replay is +scheduled as a prefill of ``prompt + every frame already generated``; the talker +rebuilds only the prompt part, and the runner's +``seg_len = min(span_len, req_embeds.shape[0])`` leaves the remaining rows of +its ``torch.empty`` buffer uninitialised (vllm-project/vllm-omni#6179). + +Unlike VoxCPM2, the layer-0 code stays inside the codec vocabulary, so the +replay raises nothing: the damage is silent. The declaration itself is pinned +in ``tests/model_executor/models/test_recompute_preemption_declarations.py``; +this file pins the arithmetic that makes it necessary. +""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +HIDDEN = 8 +PROMPT_LEN = 12 # rows _thinker_to_talker_prefill builds +GEN_FRAMES = 5 # frames already delivered before preemption + + +def test_replay_span_exceeds_the_embeddings_the_talker_can_rebuild(): + """Mirror of OmniGPUModelRunner._preprocess's copy.""" + span_len = PROMPT_LEN + GEN_FRAMES # what _preempt_request() re-schedules + rebuilt = torch.zeros(PROMPT_LEN, HIDDEN) # prompt-only, as prefill builds it + + inputs_embeds = torch.empty(span_len, HIDDEN) + seg_len = min(span_len, rebuilt.shape[0]) + inputs_embeds[0:seg_len] = rebuilt[:seg_len] + + assert seg_len == PROMPT_LEN + assert inputs_embeds.shape[0] - seg_len == GEN_FRAMES + + +def test_only_layer0_of_a_talker_frame_reaches_the_token_stream(): + """The shortfall above is structural: 7 of 8 codes never enter the stream.""" + num_code_groups = 8 # talker_config.num_code_groups + assert num_code_groups - 1 == 7 diff --git a/tests/model_executor/models/test_recompute_preemption_declarations.py b/tests/model_executor/models/test_recompute_preemption_declarations.py new file mode 100644 index 00000000000..ffc4d8f5224 --- /dev/null +++ b/tests/model_executor/models/test_recompute_preemption_declarations.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Which stages declare that recompute preemption cannot replay them. + +A stage qualifies when its decode step feeds PagedAttention something the token +stream does not carry, so `_preempt_request()`'s prompt replay rebuilds a +different KV than the one that produced the audio already delivered +(vllm-project/vllm-omni#6179): + +- a sampled RVQ codebook folded into `inputs_embeds` (Qwen3-Omni talker, + VoxCPM2, audio8_tts, fish_speech, qwen3_tts, moss_tts), +- a per-frame temporal hidden plus codebook (personaplex), +- 31 side-channel FSQ codes per frame (Gepard). + +Stages that embed the sampled token alone -- every thinker, Qwen2.5-Omni's +talker, which keeps all codebook layers in the token stream -- keep upstream's +recompute mode. Fish Omni also keeps recompute: its MRv2 model state saves and +restores the full generated codebooks across preemption. + +This pins the declaration, not the failure mode: a stage dropped from `_FAIL` +silently goes back to splicing or corrupting audio under memory pressure. +""" + +from __future__ import annotations + +import pytest + +from vllm_omni.model_executor.models.audio8_tts.pipeline import AUDIO8_TTS_PIPELINE +from vllm_omni.model_executor.models.fish_omni.pipeline import FISH_OMNI_PIPELINE +from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE +from vllm_omni.model_executor.models.gepard.pipeline import GEPARD_PIPELINE +from vllm_omni.model_executor.models.moss_tts.pipeline import ( + MOSS_TTS_LOCAL_PIPELINE, + MOSS_TTS_PIPELINE, + MOSS_TTS_REALTIME_PIPELINE, +) +from vllm_omni.model_executor.models.personaplex.pipeline import PERSONAPLEX_PIPELINE +from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import QWEN2_5_OMNI_PIPELINE +from vllm_omni.model_executor.models.qwen3_omni.pipeline import QWEN3_OMNI_PIPELINE +from vllm_omni.model_executor.models.qwen3_tts.pipeline import ( + QWEN3_TTS_FUSED_PIPELINE, + QWEN3_TTS_PIPELINE, +) +from vllm_omni.model_executor.models.voxcpm2.pipeline import VOXCPM2_PIPELINE + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# (pipeline, stage that cannot replay) +_FAIL = [ + (QWEN3_OMNI_PIPELINE, "talker"), + (VOXCPM2_PIPELINE, "latent_generator"), + (GEPARD_PIPELINE, "gepard"), + (AUDIO8_TTS_PIPELINE, "audio8_tts_slow_ar"), + (FISH_SPEECH_PIPELINE, "fish_speech_slow_ar"), + (QWEN3_TTS_PIPELINE, "qwen3_tts"), + (QWEN3_TTS_FUSED_PIPELINE, "qwen3_tts"), + (PERSONAPLEX_PIPELINE, "personaplex"), + (MOSS_TTS_PIPELINE, "moss_tts"), + (MOSS_TTS_REALTIME_PIPELINE, "moss_tts"), + (MOSS_TTS_LOCAL_PIPELINE, "moss_tts_local"), +] + +_ALLOW = [ + # Fish Omni MRv2 retains all generated codebooks across slot reuse. + (FISH_OMNI_PIPELINE, "fish_omni"), + # Text AR: the token stream is the state. + (QWEN3_OMNI_PIPELINE, "thinker"), + (QWEN2_5_OMNI_PIPELINE, "thinker"), + (QWEN2_5_OMNI_PIPELINE, "talker"), +] + + +def _stage(pipeline, model_stage): + (stage,) = [s for s in pipeline.stages if s.model_stage == model_stage] + return stage + + +@pytest.mark.parametrize( + ("pipeline", "model_stage"), + _FAIL, + ids=[f"{p.model_type}:{m}" for p, m in _FAIL], +) +def test_unreplayable_stage_fails_instead_of_replaying(pipeline, model_stage): + stage = _stage(pipeline, model_stage) + assert stage.recompute_preemption == "fail", ( + f"{pipeline.model_type}/{model_stage} feeds decode embeddings that its " + "token stream cannot rebuild; a recompute preemption would splice or " + "corrupt its audio." + ) + + +@pytest.mark.parametrize( + ("pipeline", "model_stage"), + _ALLOW, + ids=[f"{p.model_type}:{m}" for p, m in _ALLOW], +) +def test_replayable_stage_keeps_upstream_recompute(pipeline, model_stage): + assert _stage(pipeline, model_stage).recompute_preemption == "allow" + + +def test_default_is_allow_so_untouched_pipelines_are_unaffected(): + from vllm_omni.config.stage_config import StagePipelineConfig + + assert StagePipelineConfig(stage_id=0, model_stage="x").recompute_preemption == "allow" diff --git a/tests/platforms/test_rocm_groupnorm_patch.py b/tests/platforms/test_rocm_groupnorm_patch.py index c95157648f1..c49a561e826 100644 --- a/tests/platforms/test_rocm_groupnorm_patch.py +++ b/tests/platforms/test_rocm_groupnorm_patch.py @@ -20,6 +20,14 @@ patch_groupnorm._registry_mod.initialize_model = patch_groupnorm._original_initialize_model +@pytest.fixture(autouse=True) +def aiter_available(monkeypatch): + # Exercise replacement on CPU so a missing Hunyuan guard cannot pass silently. + aiter_ops = types.ModuleType("vllm._aiter_ops") + setattr(aiter_ops, "is_aiter_found_and_supported", lambda: True) + monkeypatch.setitem(sys.modules, "vllm._aiter_ops", aiter_ops) + + def test_hunyuan_image3_keeps_pytorch_groupnorm(monkeypatch): model = types.SimpleNamespace(vae=nn.Sequential(nn.GroupNorm(4, 8))) monkeypatch.setattr(patch_groupnorm, "_original_initialize_model", lambda _: model) @@ -46,10 +54,6 @@ def replace_groupnorm(_vae): return True monkeypatch.setattr(patch_groupnorm, "_replace_groupnorm_with_aiter", replace_groupnorm) - aiter_ops = types.ModuleType("vllm._aiter_ops") - setattr(aiter_ops, "is_aiter_found_and_supported", lambda: True) - monkeypatch.setitem(sys.modules, "vllm._aiter_ops", aiter_ops) - od_config = types.SimpleNamespace(model_class_name="OtherDiffusionModel") assert patch_groupnorm._patched_initialize_model(od_config) is model assert replaced diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 3c43b53a5da..b04fbac9df1 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from dataclasses import MISSING, field -from typing import Any +from typing import Any, Literal from pydantic import ConfigDict, TypeAdapter from transformers import PretrainedConfig @@ -131,6 +131,7 @@ class OmniModelConfig(ModelConfig): session_mode: str = "turn" retains_state_across_chunks: bool = False supports_running_prefix_cache_reset: bool = True + recompute_preemption: Literal["allow", "fail"] = "allow" use_v2_model_runner: bool = False supports_native_mrv2_data_plane: bool = False # Stage-1 active stream slots; 0 keeps legacy chunk-level round-robin. diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 88dde4321ea..114fd02f6b9 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -493,6 +493,7 @@ class OmniStageModelConfig(_TrackExplicitConfigFields): media_io_kwargs: dict[str, Any] | None = None final_output: bool = False supports_running_prefix_cache_reset: bool = True + recompute_preemption: Literal["allow", "fail"] = "allow" active_stream_window: int = Field(default=0, ge=0) session_mode: str = "turn" duplex_max_sessions: int = Field(default=1, ge=1) @@ -1394,6 +1395,7 @@ def _validate_stage_engine_override_ownership( "model_stage", "retains_state_across_chunks", "supports_running_prefix_cache_reset", + "recompute_preemption", "scheduler_cls", "stage_connector_spec", "worker_type", @@ -2063,6 +2065,10 @@ def _build_model_config( kwargs["final_output"] = topology.final_output if not topology.supports_running_prefix_cache_reset: kwargs["supports_running_prefix_cache_reset"] = False + # AR-only capability: diffusion stages have no KV-recompute preemption and + # reject unknown config fields. + if topology.execution_type is StageExecutionType.LLM_AR and topology.recompute_preemption != "allow": + kwargs["recompute_preemption"] = topology.recompute_preemption if "custom_voice_dir" not in kwargs and deploy.custom_voice_dir is not None: kwargs["custom_voice_dir"] = _copy_value(deploy.custom_voice_dir) stage_runner = resolve_stage_model_runner(deploy, stage_deploy) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 5d790777636..424afd2dafc 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -356,6 +356,14 @@ class StagePipelineConfig: retains_state_across_chunks: bool = False # Some stateful audio stages cannot roll back already consumed frames. supports_running_prefix_cache_reset: bool = True + # Recompute preemption (upstream's only mode) frees KV and replays the + # prompt. A native-AR stage whose decode step feeds more than the sampled + # token into the model -- side-channel codes, per-request hidden state, + # vocoder progress -- cannot rebuild that KV from the token stream. Under + # ``"fail"`` the scheduler reclaims the request's resources and emits one + # terminal error instead of requeuing it, so the caller never receives the + # spliced audio the replay would produce. + recompute_preemption: Literal["allow", "fail"] = "allow" sampling_constraints: dict[str, Any] = field(default_factory=dict) custom_process_input_func: str | None = None custom_process_next_stage_input_func: str | None = None @@ -1360,6 +1368,8 @@ def _build_engine_args( engine_args["requires_full_payload_input"] = ps.requires_full_payload_input if not ps.supports_running_prefix_cache_reset: engine_args["supports_running_prefix_cache_reset"] = False + if ps.execution_type is StageExecutionType.LLM_AR and ps.recompute_preemption != "allow": + engine_args["recompute_preemption"] = ps.recompute_preemption return engine_args diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 6a60dfcbeee..46b34379674 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -197,6 +197,10 @@ def _init_omni_io_scheduling_state(self) -> None: # (deploy ``model_runner``, optionally per stage). model_config = self.vllm_config.model_config self.use_v2_model_runner = bool(getattr(model_config, "use_v2_model_runner", False)) + # Native-AR stages whose decode step feeds more than the sampled token + # into the model declare "fail": a recompute preemption replays the + # prompt from the token stream, which cannot rebuild that state. + self.recompute_preemption_fails = getattr(model_config, "recompute_preemption", "allow") == "fail" self._native_data_plane = uses_native_mrv2_data_plane( model_config, use_v2_model_runner=self.use_v2_model_runner, @@ -552,6 +556,70 @@ def _finish_input_timeout_requests(self, request_ids: set[str]) -> None: ) self.finish_requests(request_ids, RequestStatus.FINISHED_ERROR) + def reset_prefix_cache(self, reset_running_requests: bool = False, reset_connector: bool = False) -> bool: + # Upstream uses drop_stale_output for both reset and KV-delivery + # capacity preemption. Only the explicit reset operation may resume. + self._resetting_prefix_cache = True + try: + return super().reset_prefix_cache(reset_running_requests, reset_connector) + finally: + self._resetting_prefix_cache = False + + def _preempt_request( + self, + request: Request, + timestamp: float, + drop_stale_output: bool = False, + ) -> None: + """Reclaim a preempted request instead of replaying it when the stage cannot. + + Upstream's only preemption mode is recompute: free the KV, zero + ``num_computed_tokens`` and replay the prompt. That replay is lossy for + a stage whose decode step feeds more than the sampled token into the + model -- side-channel codes, per-request hidden state, vocoder progress + -- because none of it is in the token stream. Requeueing such a request + would restart it at frame 0 and splice two unrelated halves together. + + Under ``recompute_preemption="fail"`` the request is finished here with + one terminal error instead, so the caller learns the output is + truncated rather than receiving audio that silently restarts. + """ + if not getattr(self, "recompute_preemption_fails", False) or getattr(self, "_resetting_prefix_cache", False): + return super()._preempt_request(request, timestamp, drop_stale_output=drop_stale_output) + if request.is_finished(): + return + reason = ( + "Request was preempted and this stage cannot replay its decode state; " + "the output generated so far is truncated. Raise the stage's " + "gpu_memory_utilization or lower max_num_seqs to avoid preemption." + ) + logger.warning( + "Preempting request %s after %d computed token(s): %s", + request.request_id, + request.num_computed_tokens, + reason, + ) + pending = getattr(self, "_pending_input_timeout_outputs", None) + if pending is None: + pending = self._pending_input_timeout_outputs = {} + pending[request.request_id] = ( + request.client_index, + OmniEngineCoreOutput( + request_id=request.request_id, + new_token_ids=[], + finish_reason=FinishReason.ERROR, + stop_reason=reason, + ), + ) + request.stop_reason = reason + # Preserve upstream preemption's auxiliary-output lifecycle: finish_requests + # removes the request before this step builds connector metadata. + if self.aux_output_connector is not None: + self.aux_output_connector.request_finished(request) + # The caller already popped this request from ``running``; freeing here + # is what upstream's requeue path would otherwise do later. + self.finish_requests([request.request_id], RequestStatus.FINISHED_ERROR) + def _log_failed_chunk_sends(self) -> None: """Surface chunks the sender gave up on (R1.2 of #4855). diff --git a/vllm_omni/deploy/omnivoice.yaml b/vllm_omni/deploy/omnivoice.yaml index df9e6a5daf0..0028529a42d 100644 --- a/vllm_omni/deploy/omnivoice.yaml +++ b/vllm_omni/deploy/omnivoice.yaml @@ -11,5 +11,12 @@ stages: trust_remote_code: true distributed_executor_backend: "mp" dtype: "float32" + additional_config: + omnivoice_asr: + # false defers Whisper until ref_audio is sent without ref_text; + # it does not disable automatic transcription. Set true to load at startup. + load_asr_on_startup: false + asr_model_name: "openai/whisper-large-v3-turbo" + asr_device: null max_num_seqs: 8 request_batch_max_wait_ms: 10 diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index 853bc4b7229..1fea772ddd1 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -381,12 +381,23 @@ def _resolve_npu_execution_path( context.parallel_strategy is not ParallelStrategy.NONE or context.outer_boundaries or context.paged_kv - or self.causal or extra.get("kv_cache_dtype") not in (None, "float", "auto") - or (attn_metadata is not None and attn_metadata.full_attn_spans is not None) ): return result + # Match forward_fa_npu: spans take precedence over the global causal flag. + if attn_metadata is not None and attn_metadata.full_attn_spans is not None: + return replace( + result, + path="npu_piecewise", + support=CapabilityResult.unmigrated( + "NPU piecewise routing only; kernel and compilation validation require Ascend" + ), + ) + + if self.causal: + return result + path = "npu_masked" if attn_metadata is not None and attn_metadata.attn_mask is not None else "npu_dense" if extra.get("npu_attn_varlen", False): if self._resolve_packed_seq_npu(query, key, extra) is not None: @@ -1110,6 +1121,38 @@ def forward_fa_quant_npu( # Execution errors propagate. Never retry with another precision. return runtime(query, key, value, **kwargs) + @staticmethod + def _npu_fusion_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + atten_mask: torch.Tensor | None, + sparse_mode: int, + softmax_scale: float, + inner_precise: int = 0, + ) -> torch.Tensor: + """Run ``torch_npu.npu_fusion_attention`` on BSND tensors. + + ``atten_mask`` follows Ascend's convention (True = masked out). When it + is ``None`` the sparse mode has no effect and full attention is + computed, so no mask needs to be materialized. + """ + import torch_npu + + return torch_npu.npu_fusion_attention( + query.contiguous(), + key.contiguous(), + value.contiguous(), + head_num=query.shape[2], + input_layout="BSND", + atten_mask=atten_mask, + scale=float(softmax_scale), + keep_prob=1.0, + inner_precise=inner_precise, + sparse_mode=sparse_mode, + )[0] + def forward_fa_npu( self, query: torch.Tensor, @@ -1117,10 +1160,14 @@ def forward_fa_npu( value: torch.Tensor, attn_metadata: AttentionMetadata | None = None, ) -> torch.Tensor: + # Piecewise block-causal attention takes precedence, mirroring + # ``forward_cuda``: the caller emitted ``full_attn_spans`` instead of a + # mask, so the generic mask handling below cannot express it. + if attn_metadata is not None and attn_metadata.full_attn_spans is not None: + return self._forward_piecewise_npu(query, key, value, attn_metadata) + attention_mask = attn_metadata.attn_mask if attn_metadata else None if self.causal: - import torch_npu - if attention_mask is None: # The compressed upper-triangular mask with sparse_mode=3 # applies FlashAttention's bottom-right causal alignment when @@ -1151,18 +1198,15 @@ def forward_fa_npu( # the Sq/Skv relationship. inner_precise = 2 - return torch_npu.npu_fusion_attention( - query.contiguous(), - key.contiguous(), - value.contiguous(), - head_num=query.shape[2], - input_layout="BSND", + return self._npu_fusion_attention( + query, + key, + value, atten_mask=npu_attention_mask, - scale=float(self.softmax_scale), - keep_prob=1.0, - inner_precise=inner_precise, sparse_mode=sparse_mode, - )[0] + softmax_scale=self.softmax_scale, + inner_precise=inner_precise, + ) from mindiesd import attention_forward @@ -1213,6 +1257,73 @@ def forward_fa_npu( layout=layout, ) + def _npu_piecewise_attn_func( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + causal: bool, + softmax_scale: float, + ) -> torch.Tensor: + """``attn_func`` for ``piecewise_attn`` on Ascend. + + Both segment kinds run on ``npu_fusion_attention``: + - causal segments (text runs) carry the shared compressed + (2048, 2048) mask with ``sparse_mode=3`` (rightDownCausal). + ``piecewise_attn`` gives causal segments ``kv_end == q_end`` in + global coordinates, so they always satisfy ``Sq <= Skv``; + bottom-right alignment reproduces FlashAttention's + ``attn(Q[:, s:e], K[:, :e], causal=True)`` contract for both + ``Sq < Skv`` and ``Sq == Skv``. + - full segments (condition/target image blocks) pass no mask — + ``sparse_mode`` is ignored when ``atten_mask`` is ``None`` — so + they are plain bidirectional attention over the truncated K/V + prefix. + + Inputs follow the model-side ``(B, S, H, D)`` convention == BSND. + """ + if causal: + atten_mask = _get_npu_compressed_causal_mask(query.device) + sparse_mode = 3 + else: + atten_mask = None + sparse_mode = 0 + return self._npu_fusion_attention( + query, + key, + value, + atten_mask=atten_mask, + sparse_mode=sparse_mode, + softmax_scale=softmax_scale, + ) + + def _forward_piecewise_npu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + """Piecewise block-causal attention on Ascend, mirroring ``forward_cuda``. + + Same segment plan as the CUDA path (``piecewise_attn``); only the + per-segment kernel differs (``npu_fusion_attention`` instead of + ``flash_attn_func``). + """ + assert attn_metadata.full_attn_spans is not None + self._warn_fa_deterministic_non_dense("piecewise-npu") + logger.debug("Using piecewise NPU attention for mixed causal/full mask") + return piecewise_attn( + query, + key, + value, + attn_metadata.full_attn_spans, + self.softmax_scale, + self._npu_piecewise_attn_func, + query_ranges=attn_metadata.query_ranges, + ) + def _resolve_packed_seq_npu( self, query: torch.Tensor, diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 9e4de74c5c8..3838f147745 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -771,7 +771,19 @@ def _run_local_attention(self, query, key, value, attn_metadata): raise NotImplementedError(blocked) return self.sdpa_fallback.forward(query, key, value, attn_metadata) dense_mask = attn_metadata.attn_mask if attn_metadata is not None else None - if dense_mask is not None and dense_mask.ndim == 4 and self.attn_backend is not None: + # Piecewise backends honor full_attn_spans ahead of attn_mask and never read the + # dense mask (e.g. HunyuanImage3, which sends both), so there is nothing to divert. + spans_take_precedence = ( + attn_metadata is not None + and attn_metadata.full_attn_spans is not None + and getattr(self.attn_backend, "supports_piecewise_spans", False) + ) + if ( + dense_mask is not None + and dense_mask.ndim == 4 + and self.attn_backend is not None + and not spans_take_precedence + ): # A dense 4D mask (for example a padded DiT batch) is a semantic # requirement, not a backend hint. Backends that only implement a # 2D key-validity mask must not reach their kernel wrapper, which diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py index 4ee0a8a4b08..0c258dac378 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py @@ -40,24 +40,29 @@ def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio - tile_sample_stride_height = self.tile_sample_stride_height - tile_sample_stride_width = self.tile_sample_stride_width if self.config.patch_size is not None: sample_height = sample_height // self.config.patch_size sample_width = sample_width // self.config.patch_size - tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size - tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size - blend_height = self.tile_sample_min_height // self.config.patch_size - tile_sample_stride_height - blend_width = self.tile_sample_min_width // self.config.patch_size - tile_sample_stride_width - else: - blend_height = self.tile_sample_min_height - tile_sample_stride_height - blend_width = self.tile_sample_min_width - tile_sample_stride_width - - # Split z into overlapping tiles and decode them separately. - # The tiles have an overlap to avoid seams between tiles. + + # Split z into overlapping tiles and decode them separately. The starts + # are spread evenly so every tile keeps the full decode window (see + # AutoencoderKLQwenImage21._tile_starts, vllm-omni#8850). + starts_height = self._tile_starts(height, tile_latent_min_height, tile_latent_stride_height) + starts_width = self._tile_starts(width, tile_latent_min_width, tile_latent_stride_width) + # Decoded cells contributed by one latent cell, before unpatchify. + cells_per_latent = self.spatial_compression_ratio // (self.config.patch_size or 1) + overlap_height = [ + (starts_height[n] + tile_latent_min_height - starts_height[n + 1]) * cells_per_latent + for n in range(len(starts_height) - 1) + ] + overlap_width = [ + (starts_width[n] + tile_latent_min_width - starts_width[n + 1]) * cells_per_latent + for n in range(len(starts_width) - 1) + ] + tiletask_list: list[TileTask] = [] - for i in range(0, height, tile_latent_stride_height): - for j in range(0, width, tile_latent_stride_width): + for i_index, i in enumerate(starts_height): + for j_index, j in enumerate(starts_width): time_list = [] for k in range(num_frames): tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width] @@ -65,7 +70,7 @@ def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: tiletask_list.append( TileTask( len(tiletask_list), - (i // tile_latent_stride_height, j // tile_latent_stride_width), + (i_index, j_index), time_list, workload=time_list[0].shape[3] * time_list[0].shape[4], ) @@ -73,14 +78,15 @@ def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: tile_spec = { "sample_height": sample_height, "sample_width": sample_width, - "blend_height": blend_height, - "blend_width": blend_width, - "tile_sample_stride_height": tile_sample_stride_height, - "tile_sample_stride_width": tile_sample_stride_width, + "starts_height": starts_height, + "starts_width": starts_width, + "cells_per_latent": cells_per_latent, + "overlap_height": overlap_height, + "overlap_width": overlap_width, } grid_spec = GridSpec( split_dims=(3, 4), - grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), + grid_shape=(len(starts_height), len(starts_width)), tile_spec=tile_spec, output_dtype=self.dtype, ) @@ -103,24 +109,30 @@ def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid grid_h, grid_w = grid_spec.grid_shape self.clear_cache() + starts_height = grid_spec.tile_spec["starts_height"] + starts_width = grid_spec.tile_spec["starts_width"] + cells_per_latent = grid_spec.tile_spec["cells_per_latent"] + overlap_height = grid_spec.tile_spec["overlap_height"] + overlap_width = grid_spec.tile_spec["overlap_width"] + result_rows = [] for i in range(grid_h): result_row = [] for j in range(grid_w): tile = coord_tensor_map[(i, j)] if i > 0: - tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"]) + tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, overlap_height[i - 1]) if j > 0: - tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"]) - result_row.append( - tile[ - :, - :, - :, - : grid_spec.tile_spec["tile_sample_stride_height"], - : grid_spec.tile_spec["tile_sample_stride_width"], - ] + tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, overlap_width[j - 1]) + # A tile contributes up to the start of the next one; the trailing + # tile contributes the rest and is trimmed by the final crop. + keep_h = ( + (starts_height[i + 1] - starts_height[i]) * cells_per_latent if i + 1 < grid_h else tile.shape[-2] + ) + keep_w = ( + (starts_width[j + 1] - starts_width[j]) * cells_per_latent if j + 1 < grid_w else tile.shape[-1] ) + result_row.append(tile[:, :, :, :keep_h, :keep_w]) result_rows.append(torch.cat(result_row, dim=-1)) dec = torch.cat(result_rows, dim=3)[ :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"] diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py index bd4ca81d0dd..eb28d09a035 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -19,6 +19,9 @@ GridSpec, TileTask, ) +from vllm_omni.diffusion.distributed.autoencoders.wan_sdpa_alignment import ( + install_wan_attention_sdpa_alignment, +) from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import ( can_encode_frames, decode_frames, @@ -49,6 +52,10 @@ def _execution_context(self): enabled=True, ) + def _install_attention_alignment(self) -> None: + """Realign 1x1 Wan attention inputs before the first encode or decode.""" + install_wan_attention_sdpa_alignment(self) + def encode(self, x: torch.Tensor, return_dict: bool = True): with self._execution_context(): return super().encode(x, return_dict=return_dict) @@ -56,6 +63,7 @@ def encode(self, x: torch.Tensor, return_dict: bool = True): def _encode(self, x: torch.Tensor): # Keep the parent's tiling dispatch and patchified tile coordinates. # Only the untiled, supported inference schedule uses preallocation. + self._install_attention_alignment() if not is_encoder_installed(self): return super()._encode(x) if can_encode_frames(self, x): @@ -76,6 +84,7 @@ def decode(self, z: torch.Tensor, return_dict: bool = True): def _decode(self, z: torch.Tensor, return_dict: bool = True): """Decode one latent batch, preserving the parent Diffusers contract.""" + self._install_attention_alignment() _, _, _, height, width = z.shape tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio @@ -191,6 +200,7 @@ def _discard_decoded_chunk(_chunk: torch.Tensor) -> None: def _decode_with_chunks(self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer) -> None: """Run the temporal streaming primitive without API or hook handling.""" + self._install_attention_alignment() distributed_enabled = False if hasattr(self, "is_distributed_enabled"): distributed_enabled = self.is_distributed_enabled() diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_sdpa_alignment.py b/vllm_omni/diffusion/distributed/autoencoders/wan_sdpa_alignment.py new file mode 100644 index 00000000000..3bee3eca171 --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/wan_sdpa_alignment.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Align Wan VAE attention inputs for memory-efficient SDPA. + +Diffusers' ``WanAttentionBlock`` builds Q, K, and V by chunking one QKV +projection along the last dimension. For a 1x1 spatial activation that chunk +is a size-1 sequence whose stride stays 1: PyTorch treats the tensor as +contiguous and ``contiguous()`` does not rewrite unit-dimension strides. The +CUDA memory-efficient kernel then rejects it with +``query is not correctly aligned (strideM)``. + +Distributed Wan tiling (256px tiles, stride 192, spatial compression 16) +produces that 1x1 latent at the corner of some resolutions, including ones +sampled by the Cosmos3 stability mix. Repacking only those tensors gives the +kernel a canonical sequence stride and leaves every already-aligned tile on +its original storage. +""" + +from __future__ import annotations + +import warnings +from importlib import metadata +from types import MethodType + +import torch +import torch.nn.functional as F +from torch import nn + +# 128-bit vector width used by the memory-efficient SDPA strideM check. +_SDPA_ALIGNMENT_BYTES = 16 +_ROOT_ATTR = "_vllm_omni_wan_attention_sdpa_aligned" + + +def sdpa_sequence_stride_alignment(tensor: torch.Tensor) -> int: + """Element count the sequence stride must be divisible by for ``tensor``'s dtype.""" + element_size = tensor.element_size() + if element_size <= 0 or element_size > _SDPA_ALIGNMENT_BYTES: + return 1 + return _SDPA_ALIGNMENT_BYTES // element_size + + +def sequence_stride_is_aligned(tensor: torch.Tensor) -> bool: + """Whether ``tensor`` satisfies the memory-efficient SDPA sequence-stride check. + + The kernel reads ``stride(-2)`` (sequence) and requires the last dimension + to be contiguous. Size-1 sequence dimensions are included: ``is_contiguous`` + ignores them, which is how a 1x1 query keeps sequence stride 1. + """ + if tensor.ndim < 2 or tensor.stride(-1) != 1: + return False + alignment = sdpa_sequence_stride_alignment(tensor) + return tensor.stride(-2) % alignment == 0 + + +def _repack_contiguous(tensor: torch.Tensor) -> torch.Tensor: + """Copy ``tensor`` into a buffer with canonical strides, including size-1 dims.""" + repacked = torch.empty(tensor.shape, dtype=tensor.dtype, device=tensor.device) + repacked.copy_(tensor) + return repacked + + +def align_sdpa_qkv( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return Q, K, V views whose sequence strides the memory-efficient kernel accepts. + + Already-aligned inputs are returned unchanged, so normal tiles keep their + storage and values. Misaligned inputs are copied into a canonical buffer; + the values are unchanged. + """ + if sequence_stride_is_aligned(query) and sequence_stride_is_aligned(key) and sequence_stride_is_aligned(value): + return query, key, value + return _repack_contiguous(query), _repack_contiguous(key), _repack_contiguous(value) + + +def split_wan_attention_qkv(qkv: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Match diffusers 0.40 ``WanAttentionBlock`` QKV views from ``[B*T, 3C, H, W]``.""" + batch_time, channels_times_3, _, _ = qkv.shape + fused = qkv.reshape(batch_time, 1, channels_times_3, -1).permute(0, 1, 3, 2).contiguous() + query, key, value = fused.chunk(3, dim=-1) + return query, key, value + + +def aligned_wan_attention_forward(self: nn.Module, x: torch.Tensor) -> torch.Tensor: + """``WanAttentionBlock.forward`` with SDPA inputs realigned when required. + + The arithmetic matches diffusers 0.40.0. The only extra work is + :func:`align_sdpa_qkv`, which is a no-op for every tile whose sequence + stride is already legal. + """ + identity = x + batch_size, channels, time, height, width = x.size() + + hidden = x.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, height, width) + hidden = self.norm(hidden) + query, key, value = align_sdpa_qkv(*split_wan_attention_qkv(self.to_qkv(hidden))) + hidden = F.scaled_dot_product_attention(query, key, value) + + hidden = hidden.squeeze(1).permute(0, 2, 1).reshape(batch_size * time, channels, height, width) + hidden = self.proj(hidden) + hidden = hidden.view(batch_size, time, channels, height, width) + hidden = hidden.permute(0, 2, 1, 3, 4) + return hidden + identity + + +def install_wan_attention_sdpa_alignment(root: nn.Module) -> None: + """Point each ``WanAttentionBlock`` under ``root`` at the aligned forward. + + Idempotent. Installed at decode/encode time so the Wan fast path, which + rejects a custom forward at install time, still sees the diffusers method + while the model is loading. Existing spatial-shard wrappers are preserved: + they gather the full spatial extent before attention, rather than tiles. + Warn on unverified diffusers versions because the forward is copied from + 0.40.0; keep the alignment fix active rather than restoring the tile crash. + """ + if getattr(root, _ROOT_ATTR, False): + return + diffusers_version = metadata.version("diffusers") + if diffusers_version != "0.40.0": + warnings.warn( + "Wan attention SDPA alignment was validated against diffusers 0.40.0, " + f"but {diffusers_version} is installed. The copied WanAttentionBlock.forward " + "must be checked against this version for compatibility.", + RuntimeWarning, + stacklevel=2, + ) + for module in root.modules(): + if module.__class__.__name__ != "WanAttentionBlock": + continue + # LingBot installs this wrapper before the first decode. Replacing it + # would drop the all-gather/reshard and silently attend to local shards. + if getattr(module, "_vllm_omni_spatial_shard_attention", False): + continue + module.forward = MethodType(aligned_wan_attention_forward, module) + setattr(root, _ROOT_ATTR, True) diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 934f976d10c..b3022aeb528 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -116,18 +116,13 @@ def __init__( ): raise ValueError("max_num_batched_tokens must be a positive integer") self._max_num_batched_tokens = max_num_batched_tokens - # Imported here: vllm_omni.diffusion.data imports this package while it - # is still initializing, and the offloader package imports that module. - from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload_strategy - - od_config = getattr(pipeline, "od_config", None) - # DLO owns the base-weight lifecycle. Keep request-switchable LoRA - # sidecars resident instead of rebuilding DLO host shards per request. - self._resident_lora_device = ( - device - if od_config is not None and resolve_offload_strategy(od_config) is OffloadStrategy.DISTRIBUTED_LAYER_WISE - else None - ) + # LoRA sidecars are allocated outside the base-weight lifecycle, so no + # sleep/offload path knows to move them. Keep them resident on the + # execution device under every offload strategy: buffers created or + # re-created while the base weights sit on CPU (the standard RL + # weight-sync moment) would otherwise stay there after wake, and + # apply() would matmul device activations against host buffers. + self._resident_lora_device = device # Cache supported/expected module suffixes once, before any layer # replacement happens. After LoRA layers are injected, the original @@ -629,12 +624,10 @@ def _matches_target(module_name: str) -> bool: ) if lora_layer is not module and isinstance(lora_layer, BaseLayerWithLoRA): - if self._resident_lora_device is not None: - set_buffer_device = getattr(lora_layer, "_set_diffusion_lora_buffer_device", None) - if not callable(set_buffer_device): - raise RuntimeError( - f"{type(lora_layer).__name__} cannot keep dynamic LoRA buffers resident for DLO" - ) + # Layers without the re-homing hook keep the legacy + # placement (buffers follow the base weight device). + set_buffer_device = getattr(lora_layer, "_set_diffusion_lora_buffer_device", None) + if callable(set_buffer_device): set_buffer_device(self._resident_lora_device) replace_submodule(component, module_name, lora_layer) self._lora_modules[full_module_name] = lora_layer diff --git a/vllm_omni/diffusion/models/omnivoice/audio.py b/vllm_omni/diffusion/models/omnivoice/audio.py new file mode 100644 index 00000000000..b20a3f93c6a --- /dev/null +++ b/vllm_omni/diffusion/models/omnivoice/audio.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Audio preparation and output processing for OmniVoice.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import torch +import torchaudio + +_END_PUNCTUATION = { + ";", + ":", + ",", + ".", + "!", + "?", + "…", + ")", + "]", + "}", + '"', + "'", + "“", + "”", + "‘", + "’", + ";", + ":", + ",", + "。", + "!", + "?", + "、", + ")", + "】", +} + + +@dataclass(frozen=True) +class PreparedReferenceAudio: + """Reference waveform shared by ASR and the audio tokenizer.""" + + waveform: np.ndarray + sample_rate: int + original_rms: float + + +def _segment_length_ms(num_samples: int, sample_rate: int) -> int: + """Return the rounded millisecond length used by the reference algorithm.""" + return round(1000 * num_samples / sample_rate) + + +def _frame_index(milliseconds: int, sample_rate: int, num_samples: int) -> int: + """Map a millisecond position to a bounded sample index.""" + return min(num_samples, max(0, int(milliseconds * sample_rate / 1000.0))) + + +def _slice_milliseconds(audio: np.ndarray, start_ms: int, end_ms: int | None, sample_rate: int) -> np.ndarray: + num_samples = audio.shape[-1] + start = _frame_index(start_ms, sample_rate, num_samples) + end = num_samples if end_ms is None else _frame_index(end_ms, sample_rate, num_samples) + return audio[:, start:end] + + +def _rms(audio: np.ndarray) -> float: + if audio.size == 0: + return 0.0 + squared = np.square(np.asarray(audio, dtype=np.float32), dtype=np.float32) + return float(np.sqrt(np.mean(squared, dtype=np.float32))) + + +def _dbfs(audio: np.ndarray) -> float: + rms = _rms(audio) + if rms == 0.0: + return -float("inf") + return 20.0 * float(np.log10(rms)) + + +def _window_rms(audio: np.ndarray, starts_ms: list[int], window_ms: int, sample_rate: int) -> np.ndarray: + starts = np.fromiter( + (_frame_index(start, sample_rate, audio.shape[-1]) for start in starts_ms), + dtype=np.int64, + ) + ends = np.fromiter( + (_frame_index(start + window_ms, sample_rate, audio.shape[-1]) for start in starts_ms), + dtype=np.int64, + ) + squared = np.square(audio, dtype=np.float32) + cumulative = np.concatenate( + [np.zeros((audio.shape[0], 1), dtype=np.float64), np.cumsum(squared, axis=-1, dtype=np.float64)], + axis=-1, + ) + sums = cumulative[:, ends] - cumulative[:, starts] + counts = ends - starts + rms = np.zeros(len(starts_ms), dtype=np.float32) + valid = counts > 0 + if np.any(valid): + means = sums[:, valid].sum(axis=0) / (counts[valid] * audio.shape[0]) + rms[valid] = np.sqrt(means).astype(np.float32) + return rms + + +def _detect_silence_ranges( + audio: np.ndarray, + sample_rate: int, + *, + min_silence_ms: int, + silence_threshold_db: float, + seek_step_ms: int, +) -> list[list[int]]: + """Find silent ranges with the same window and range rules as pydub.""" + segment_length_ms = _segment_length_ms(audio.shape[-1], sample_rate) + if segment_length_ms < min_silence_ms: + return [] + + threshold = 10.0 ** (silence_threshold_db / 20.0) + last_start = segment_length_ms - min_silence_ms + silence_starts = list(range(0, last_start + 1, seek_step_ms)) + if last_start % seek_step_ms: + silence_starts.append(last_start) + + window_rms = _window_rms(audio, silence_starts, min_silence_ms, sample_rate) + silent_starts = [start for start, rms in zip(silence_starts, window_rms) if rms <= threshold] + if not silent_starts: + return [] + + silent_ranges: list[list[int]] = [] + previous = silent_starts[0] + current_start = previous + for start in silent_starts[1:]: + continuous = start == previous + seek_step_ms + has_gap = start > previous + min_silence_ms + if not continuous and has_gap: + silent_ranges.append([current_start, previous + min_silence_ms]) + current_start = start + previous = start + silent_ranges.append([current_start, previous + min_silence_ms]) + return silent_ranges + + +def _detect_nonsilent_ranges( + audio: np.ndarray, + sample_rate: int, + *, + min_silence_ms: int, + silence_threshold_db: float, + seek_step_ms: int, +) -> list[list[int]]: + segment_length_ms = _segment_length_ms(audio.shape[-1], sample_rate) + silent_ranges = _detect_silence_ranges( + audio, + sample_rate, + min_silence_ms=min_silence_ms, + silence_threshold_db=silence_threshold_db, + seek_step_ms=seek_step_ms, + ) + if not silent_ranges: + return [[0, segment_length_ms]] + if silent_ranges[0] == [0, segment_length_ms]: + return [] + + nonsilent_ranges: list[list[int]] = [] + previous_end = 0 + for start, end in silent_ranges: + nonsilent_ranges.append([previous_end, start]) + previous_end = end + if previous_end != segment_length_ms: + nonsilent_ranges.append([previous_end, segment_length_ms]) + if nonsilent_ranges and nonsilent_ranges[0] == [0, 0]: + nonsilent_ranges.pop(0) + return nonsilent_ranges + + +def _split_on_silence( + audio: np.ndarray, + sample_rate: int, + *, + min_silence_ms: int, + silence_threshold_db: float, + keep_silence_ms: int, + seek_step_ms: int, +) -> np.ndarray: + ranges = _detect_nonsilent_ranges( + audio, + sample_rate, + min_silence_ms=min_silence_ms, + silence_threshold_db=silence_threshold_db, + seek_step_ms=seek_step_ms, + ) + output_ranges = [[start - keep_silence_ms, end + keep_silence_ms] for start, end in ranges] + for previous, current in zip(output_ranges, output_ranges[1:]): + if current[0] < previous[1]: + split = (previous[1] + current[0]) // 2 + previous[1] = split + current[0] = split + + segment_length_ms = _segment_length_ms(audio.shape[-1], sample_rate) + segments = [ + _slice_milliseconds(audio, max(start, 0), min(end, segment_length_ms), sample_rate) + for start, end in output_ranges + ] + if not segments: + return np.empty((audio.shape[0], 0), dtype=np.float32) + return np.concatenate(segments, axis=-1) + + +def _detect_leading_silence_ms(audio: np.ndarray, sample_rate: int, *, threshold_db: float, chunk_ms: int = 10) -> int: + segment_length_ms = _segment_length_ms(audio.shape[-1], sample_rate) + trim_ms = 0 + while trim_ms < segment_length_ms: + chunk = _slice_milliseconds(audio, trim_ms, trim_ms + chunk_ms, sample_rate) + if _dbfs(chunk) >= threshold_db: + break + trim_ms += chunk_ms + return min(trim_ms, segment_length_ms) + + +def remove_silence( + audio: np.ndarray, + sample_rate: int, + *, + middle_silence_ms: int, + leading_silence_ms: int, + trailing_silence_ms: int, +) -> np.ndarray: + """Remove long gaps and trim edge silence using official thresholds.""" + wave = np.asarray(audio, dtype=np.float32) + if wave.ndim == 1: + wave = wave[np.newaxis, :] + if wave.ndim != 2: + raise ValueError(f"OmniVoice audio must be 1D or 2D, got {wave.ndim} dimensions.") + if middle_silence_ms > 0: + wave = _split_on_silence( + wave, + sample_rate, + min_silence_ms=middle_silence_ms, + silence_threshold_db=-50, + keep_silence_ms=middle_silence_ms, + seek_step_ms=10, + ) + + start_idx = _detect_leading_silence_ms(wave, sample_rate, threshold_db=-50) + start_idx = max(0, start_idx - leading_silence_ms) + wave = _slice_milliseconds(wave, start_idx, None, sample_rate) + + wave = wave[:, ::-1] + start_idx = _detect_leading_silence_ms(wave, sample_rate, threshold_db=-50) + start_idx = max(0, start_idx - trailing_silence_ms) + wave = _slice_milliseconds(wave, start_idx, None, sample_rate) + wave = wave[:, ::-1] + + return np.ascontiguousarray(wave, dtype=np.float32) + + +def trim_long_audio( + audio: np.ndarray, + sample_rate: int, + *, + max_duration_s: float = 15.0, + min_duration_s: float = 3.0, + trim_threshold_s: float = 20.0, +) -> np.ndarray: + """Trim long reference audio at a suitable silence boundary.""" + if audio.shape[-1] / sample_rate <= trim_threshold_s: + return audio + + non_silent_ranges = _detect_nonsilent_ranges( + np.asarray(audio, dtype=np.float32), + sample_rate=sample_rate, + min_silence_ms=100, + silence_threshold_db=-40, + seek_step_ms=10, + ) + if not non_silent_ranges: + return audio + + max_ms = int(max_duration_s * 1000) + min_ms = int(min_duration_s * 1000) + best_split = 0 + for start, end in non_silent_ranges: + if start > best_split and start <= max_ms: + best_split = start + if end > max_ms: + break + + if best_split < min_ms: + best_split = min(max_ms, _segment_length_ms(audio.shape[-1], sample_rate)) + return _slice_milliseconds(np.asarray(audio, dtype=np.float32), 0, best_split, sample_rate) + + +def prepare_reference_audio( + waveform: np.ndarray | torch.Tensor, + sample_rate: int, + *, + target_sample_rate: int, + hop_length: int, + trim_long: bool, +) -> PreparedReferenceAudio: + """Prepare reference audio before ASR and audio-tokenizer encoding.""" + if isinstance(waveform, torch.Tensor): + waveform = waveform.detach().cpu().numpy() + waveform = np.array(waveform, dtype=np.float32, copy=True) + if waveform.ndim == 1: + waveform = waveform[np.newaxis, :] + elif waveform.ndim != 2: + raise ValueError(f"OmniVoice reference audio must be 1D or 2D, got {waveform.ndim} dimensions.") + if waveform.shape[0] > 1: + waveform = waveform.mean(axis=0, keepdims=True) + if not waveform.size or not np.any(waveform): + raise ValueError("Reference audio is empty after silence removal.") + + if sample_rate != target_sample_rate: + waveform = torchaudio.functional.resample( + torch.from_numpy(waveform), + orig_freq=sample_rate, + new_freq=target_sample_rate, + ).numpy() + + original_rms = float(np.sqrt(np.mean(waveform**2))) if waveform.size else 0.0 + if 0 < original_rms < 0.1: + waveform = waveform * (0.1 / original_rms) + waveform = np.clip(waveform, -1.0, 1.0) + + if trim_long: + waveform = trim_long_audio(waveform, target_sample_rate) + waveform = remove_silence( + waveform, + target_sample_rate, + middle_silence_ms=200, + leading_silence_ms=100, + trailing_silence_ms=200, + ) + if waveform.shape[-1] == 0: + raise ValueError("Reference audio is empty after silence removal.") + + remainder = waveform.shape[-1] % hop_length + if remainder: + waveform = waveform[:, :-remainder] + if waveform.shape[-1] == 0: + raise ValueError("Reference audio is shorter than one audio-tokenizer hop after preprocessing.") + + return PreparedReferenceAudio( + waveform=np.ascontiguousarray(waveform, dtype=np.float32), + sample_rate=target_sample_rate, + original_rms=original_rms, + ) + + +def add_reference_punctuation(text: str) -> str: + """Add an English or Chinese sentence terminator when one is missing.""" + text = text.strip() + if not text: + return text + if text[-1] in _END_PUNCTUATION: + return text + if any("\u4e00" <= char <= "\u9fff" for char in text): + return f"{text}。" + return f"{text}." + + +def postprocess_generated_audio( + audio: np.ndarray, + *, + sample_rate: int, + reference_rms: float, +) -> np.ndarray: + """Apply official OmniVoice output silence, volume, fade, and padding rules.""" + audio = np.array(audio, dtype=np.float32, copy=True) + if audio.ndim == 1: + audio = audio[np.newaxis, :] + if audio.ndim != 2: + raise ValueError(f"OmniVoice generated audio must be 1D or 2D, got {audio.ndim} dimensions.") + + audio = remove_silence( + audio, + sample_rate, + middle_silence_ms=500, + leading_silence_ms=100, + trailing_silence_ms=100, + ) + + if audio.shape[-1] == 0: + return audio + + if reference_rms < 0.1: + audio = audio * (reference_rms / 0.1) + + fade_samples = int(0.1 * sample_rate) + if fade_samples > 0: + fade_length = min(fade_samples, audio.shape[-1] // 2) + if fade_length > 0: + fade_in = np.linspace(0, 1, fade_length, dtype=np.float32)[np.newaxis, :] + fade_out = np.linspace(1, 0, fade_length, dtype=np.float32)[np.newaxis, :] + audio[:, :fade_length] *= fade_in + audio[:, -fade_length:] *= fade_out + + pad_samples = int(0.1 * sample_rate) + if pad_samples > 0: + padding = np.zeros((audio.shape[0], pad_samples), dtype=np.float32) + audio = np.concatenate([padding, audio, padding], axis=-1) + return np.ascontiguousarray(audio, dtype=np.float32) diff --git a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py index bad1c6154e7..d94da8b070d 100644 --- a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py +++ b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py @@ -11,17 +11,20 @@ from __future__ import annotations +import hashlib import json import math import os import random import re -from collections.abc import Iterable, Sequence +from collections import OrderedDict +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from typing import Any, ClassVar import numpy as np import torch +import torchaudio from tokenizers import Tokenizer as HFTokenizer from torch import nn from vllm.logger import init_logger @@ -29,6 +32,11 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.models.interface import SupportAudioOutput +from vllm_omni.diffusion.models.omnivoice.audio import ( + add_reference_punctuation, + postprocess_generated_audio, + prepare_reference_audio, +) from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.diffusion.worker.utils import StepRequestState @@ -48,10 +56,42 @@ except ImportError: HiggsAudioV2TokenizerModel = None -import torchaudio +try: + from transformers import pipeline as hf_pipeline +except ImportError: + hf_pipeline = None logger = init_logger(__name__) +_ASR_MODEL_NAME = "openai/whisper-large-v3-turbo" +_INLINE_CACHE_MAX_ENTRIES = 8 + + +def _parse_asr_config(additional_config: Mapping[str, object] | None) -> tuple[bool, str, str | None]: + """Return validated OmniVoice ASR settings from the model config.""" + if additional_config is None: + additional_config = {} + if not isinstance(additional_config, Mapping): + raise TypeError(f"additional_config must be a mapping or None, got {type(additional_config)!r}") + + raw_config = additional_config.get("omnivoice_asr", {}) + if not isinstance(raw_config, Mapping): + raise TypeError(f"additional_config['omnivoice_asr'] must be a mapping, got {type(raw_config)!r}") + + load_asr_on_startup = raw_config.get("load_asr_on_startup", False) + if not isinstance(load_asr_on_startup, bool): + raise TypeError("additional_config['omnivoice_asr']['load_asr_on_startup'] must be a bool") + + model_name = raw_config.get("asr_model_name", _ASR_MODEL_NAME) + if not isinstance(model_name, str) or not model_name.strip(): + raise ValueError("additional_config['omnivoice_asr']['asr_model_name'] must be a non-empty string") + + asr_device = raw_config.get("asr_device") + if asr_device is not None and (not isinstance(asr_device, str) or not asr_device.strip()): + raise ValueError("additional_config['omnivoice_asr']['asr_device'] must be a non-empty string") + + return load_asr_on_startup, model_name.strip(), asr_device.strip() if asr_device is not None else None + @dataclass class _PreparedOmniVoiceRequest: @@ -60,6 +100,7 @@ class _PreparedOmniVoiceRequest: cond_len: int target_len: int seed: int | None + reference_rms: float | None = None def get_omnivoice_post_process_func(od_config: OmniDiffusionConfig): @@ -159,6 +200,8 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): self.od_config = od_config self.device = get_local_device() self.model_path = od_config.model + self._initialize_asr(getattr(od_config, "additional_config", None)) + self._inline_reference_cache: OrderedDict[tuple[object, ...], dict[str, object]] = OrderedDict() # Resolve model path (HF hub ID → local cache) if not os.path.isdir(self.model_path): @@ -206,6 +249,101 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): self.class_temperature = self.config.class_temperature self.sample_rate = self.config.sample_rate + def _initialize_asr(self, additional_config: Mapping[str, object] | None) -> None: + load_asr_on_startup, self._asr_model_name, self._asr_device = _parse_asr_config(additional_config) + self._asr_pipeline = None + if load_asr_on_startup: + self._load_asr_pipeline() + + def _load_asr_pipeline(self): + """Load the configured reference-audio ASR pipeline.""" + if self._asr_pipeline is not None: + return self._asr_pipeline + asr_device = self._asr_device if self._asr_device is not None else self.device + if hf_pipeline is None: + raise RuntimeError( + "OmniVoice automatic transcription requires the Hugging Face " + f"ASR pipeline ({self._asr_model_name!r}) on device {asr_device}." + ) + + asr_dtype = torch.float16 if str(asr_device).lower().startswith(("cuda", "xpu")) else torch.float32 + logger.info( + "Loading OmniVoice ASR model %s on %s", + self._asr_model_name, + asr_device, + ) + try: + asr_pipeline = hf_pipeline( + "automatic-speech-recognition", + model=self._asr_model_name, + dtype=asr_dtype, + device=asr_device, + ) + # Transformers keeps a newly loaded pipeline model on CPU when a + # torch.distributed process group is already initialized. vLLM + # workers always have one, so explicitly honor the configured + # device after construction, as other auxiliary models do here. + target_device = torch.device(asr_device) + asr_pipeline.model = asr_pipeline.model.to(target_device) + asr_pipeline.device = target_device + except Exception as exc: + raise RuntimeError( + f"Failed to load OmniVoice ASR model {self._asr_model_name!r} on device {asr_device}: {exc}" + ) from exc + self._asr_pipeline = asr_pipeline + logger.info("OmniVoice ASR model loaded on %s", asr_device) + return self._asr_pipeline + + @torch.inference_mode() + def _transcribe_ref_audio(self, ref_audio) -> str: + """Transcribe a reference waveform with the lazily loaded ASR model.""" + waveform, sr = ref_audio + if isinstance(waveform, torch.Tensor): + waveform = waveform.detach().cpu().numpy() + waveform = np.squeeze(np.array(waveform, copy=True)) + result = self._load_asr_pipeline()( + { + "array": waveform, + "sampling_rate": int(sr), + } + ) + if not isinstance(result, dict) or "text" not in result: + raise RuntimeError("OmniVoice ASR returned a malformed result without a 'text' field.") + transcript = result["text"] + if not isinstance(transcript, str): + raise RuntimeError("OmniVoice ASR returned a malformed result: 'text' must be a string.") + transcript = transcript.strip() + if not transcript: + raise ValueError("OmniVoice ASR returned an empty reference transcription.") + return transcript + + def _resolve_ref_text(self, ref_audio, ref_text: str | None) -> str | None: + """Resolve missing reference text only when reference audio is present.""" + if ref_audio is not None and (ref_text is None or not ref_text.strip()): + logger.debug("Automatically transcribing OmniVoice reference audio") + return self._transcribe_ref_audio(ref_audio) + return ref_text + + def _estimate_target_len( + self, + text: str, + ref_text: str | None, + ref_audio_tokens: torch.Tensor | None, + ) -> int: + """Estimate target audio tokens using resolved voice-clone conditioning.""" + if ref_audio_tokens is None or not ref_text: + ref_text = "Nice to meet you." + num_ref_audio_tokens = 25 + else: + num_ref_audio_tokens = ref_audio_tokens.size(-1) + + target_len = self.duration_estimator.estimate_duration( + text, + ref_text, + num_ref_audio_tokens, + ) + return max(1, int(target_len)) + def _encode_ref_audio(self, audio_signal: torch.Tensor, sr: int) -> torch.Tensor: """Encode reference audio to 8-codebook tokens for voice cloning.""" if self.audio_tokenizer is None: @@ -226,6 +364,30 @@ def _encode_ref_audio(self, audio_signal: torch.Tensor, sr: int) -> torch.Tensor tokens = tokens.squeeze(0) # [8, T_ref] return tokens + @staticmethod + def _inline_cache_key(prepared_audio, preparation_mode: str) -> tuple[object, ...]: + waveform = np.ascontiguousarray(prepared_audio.waveform, dtype=np.float32) + digest = hashlib.sha256(waveform.tobytes()).digest() + return ( + preparation_mode, + digest, + tuple(waveform.shape), + int(prepared_audio.sample_rate), + float(prepared_audio.original_rms), + ) + + def _get_inline_cache(self, key: tuple[object, ...]) -> dict[str, object] | None: + cached = self._inline_reference_cache.pop(key, None) + if cached is not None: + self._inline_reference_cache[key] = cached + return cached + + def _put_inline_cache(self, key: tuple[object, ...], artifacts: dict[str, object]) -> None: + self._inline_reference_cache[key] = artifacts + self._inline_reference_cache.move_to_end(key) + while len(self._inline_reference_cache) > _INLINE_CACHE_MAX_ENTRIES: + self._inline_reference_cache.popitem(last=False) + def _prepare_request_input( self, prompt: Any, @@ -271,6 +433,10 @@ def _prepare_request_input( lang = mm_kwargs.get("lang") if instruct is None: instruct = mm_kwargs.get("instruct") + + if isinstance(ref_text, str): + ref_text = ref_text.strip() or None + if not text: return DiffusionOutput(error="Empty text prompt") lang = lang or "None" @@ -280,49 +446,119 @@ def _prepare_request_input( if not text: return DiffusionOutput(error="Empty text prompt") - target_len = self.duration_estimator.estimate_duration(text, "Nice to meet you.", 25) - target_len = max(1, int(target_len)) - - style_text = f"<|denoise|><|lang_start|>{lang}<|lang_end|><|instruct_start|>{instruct}<|instruct_end|>" - full_text = _combine_text(ref_text=ref_text, text=text) - wrapped_text = f"<|text_start|>{full_text}<|text_end|>" - style_tokens = self.tokenizer.encode(style_text).ids - text_tokens = _tokenize_with_nonverbal_tags(wrapped_text, self.tokenizer) - encoding_ids = style_tokens + text_tokens - text_tokens_tensor = torch.tensor(encoding_ids, dtype=torch.long, device=self.device) - text_len = text_tokens_tensor.shape[0] + device = self.device + num_cb = self.config.num_audio_codebook + mask_id = self.config.audio_mask_id ref_audio_tokens = None + reference_rms = None if ref_audio is not None: + needs_asr = ref_text is None if self.audio_tokenizer is None: raise RuntimeError( "Voice cloning requires transformers>=5.3.0. Try: uv pip install 'transformers>=5.3.0'" ) - cache_key = None + preparation_mode = "asr" if needs_asr else "explicit" + _cache_key = None if voice_name: - cache_key = self._speaker_cache.make_cache_key( + _cache_key = self._speaker_cache.make_cache_key( voice_name, - model_type="omnivoice", + model_type=f"omnivoice_{preparation_mode}", created_at=int(prompt.get("voice_created_at") or 0), ) - cached = self._speaker_cache.get(cache_key) - if cached is not None: - ref_audio_tokens = cached["ref_audio_tokens"].to(self.device) - cache_key = None + cached = self._speaker_cache.get(_cache_key) + if cached is not None and (not needs_asr or cached.get("ref_text")): + ref_audio_tokens = cached["ref_audio_tokens"].to(device) + reference_rms = cached["reference_rms"] + if needs_asr: + ref_text = cached["ref_text"] + _cache_key = None # hit → don't store again logger.debug("Speaker cache HIT for OmniVoice speaker '%s'", voice_name) if ref_audio_tokens is None: - audio_signal, sr = ref_audio - if isinstance(audio_signal, np.ndarray): - audio_signal = torch.from_numpy(audio_signal).float() - ref_audio_tokens = self._encode_ref_audio(audio_signal, int(sr)).to(self.device) - if cache_key is not None: - self._speaker_cache.put(cache_key, {"ref_audio_tokens": ref_audio_tokens.cpu()}) + audio_signal, sample_rate = ref_audio + try: + prepared_audio = prepare_reference_audio( + audio_signal, + int(sample_rate), + target_sample_rate=self.audio_tokenizer.config.sample_rate, + hop_length=self.audio_tokenizer.config.hop_length, + trim_long=needs_asr, + ) + except (RuntimeError, ValueError) as exc: + return DiffusionOutput(error=str(exc)) + reference_rms = prepared_audio.original_rms + reference_duration = prepared_audio.waveform.shape[-1] / prepared_audio.sample_rate + if reference_duration > 20.0: + logger.warning( + "OmniVoice reference audio is %.1fs long (>20s); this may increase memory use " + "and reduce cloning quality.", + reference_duration, + ) + + inline_cache_key = None + if not voice_name: + inline_cache_key = self._inline_cache_key(prepared_audio, preparation_mode) + cached = self._get_inline_cache(inline_cache_key) + if cached is not None: + ref_audio_tokens = cached["ref_audio_tokens"].to(device) + reference_rms = cached["reference_rms"] + if needs_asr: + ref_text = cached["ref_text"] + logger.debug("Inline OmniVoice reference cache HIT") + + if ref_audio_tokens is None and needs_asr: + try: + ref_text = self._resolve_ref_text( + (prepared_audio.waveform, prepared_audio.sample_rate), + ref_text, + ) + except (RuntimeError, ValueError) as exc: + return DiffusionOutput(error=str(exc)) + if ref_audio_tokens is None: + ref_audio_tokens = self._encode_ref_audio( + torch.from_numpy(prepared_audio.waveform), + prepared_audio.sample_rate, + ).to(device) + + # Store named and inline entries for the matching preparation mode. + if _cache_key is not None: + self._speaker_cache.put( + _cache_key, + { + "ref_audio_tokens": ref_audio_tokens.cpu(), + "ref_text": ref_text if needs_asr else None, + "reference_rms": reference_rms, + }, + ) logger.debug("Speaker cache STORE for OmniVoice speaker '%s'", voice_name) + elif not voice_name: + self._put_inline_cache( + inline_cache_key, + { + "ref_audio_tokens": ref_audio_tokens.cpu(), + "ref_text": ref_text if needs_asr else None, + "reference_rms": reference_rms, + }, + ) + logger.debug("Inline OmniVoice reference cache STORE") + + if ref_text: + ref_text = add_reference_punctuation(ref_text) + + target_len = self._estimate_target_len(text, ref_text, ref_audio_tokens) + + # Build text prompt with control tokens + style_text = f"<|denoise|><|lang_start|>{lang}<|lang_end|><|instruct_start|>{instruct}<|instruct_end|>" + full_text = _combine_text(ref_text=ref_text, text=text) + wrapped_text = f"<|text_start|>{full_text}<|text_end|>" + style_tokens = self.tokenizer.encode(style_text).ids + text_tokens = _tokenize_with_nonverbal_tags(wrapped_text, self.tokenizer) + encoding_ids = style_tokens + text_tokens + text_tokens = torch.tensor(encoding_ids, dtype=torch.long, device=device) + text_len = text_tokens.shape[0] - num_cb = self.config.num_audio_codebook - mask_id = self.config.audio_mask_id - text_ids = text_tokens_tensor.unsqueeze(0).repeat(num_cb, 1) + text_ids = text_tokens.unsqueeze(0).repeat(num_cb, 1) target_ids = torch.full((num_cb, target_len), mask_id, dtype=torch.long, device=self.device) cond_ids = ( torch.cat([text_ids, ref_audio_tokens, target_ids], dim=1) @@ -343,6 +579,7 @@ def _prepare_request_input( cond_len=cond_len, target_len=target_len, seed=seed, + reference_rms=reference_rms, ) def _collate_request_inputs( @@ -431,6 +668,7 @@ def prepare_encode(self, state: StepRequestState) -> StepRequestState: state.extra["target_len"] = target_len state.extra["audio_mask"] = audio_mask state.extra["tokens"] = tokens + state.extra["reference_rms"] = prepared_request.reference_rms return state def denoise_step(self, input_batch: InputBatch, *, states: Sequence[StepRequestState] | None = None, **kwargs: Any): @@ -540,11 +778,23 @@ def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kw state.latents = noise_pred state.step_index += 1 + def _postprocess_audio(self, audio: torch.Tensor, reference_rms: float) -> torch.Tensor: + audio_np = audio.squeeze(0).detach().cpu().float().numpy() + audio_np = postprocess_generated_audio( + audio_np, + sample_rate=self.sample_rate, + reference_rms=reference_rms, + ) + return torch.from_numpy(audio_np).unsqueeze(0) + def post_decode(self, state: StepRequestState, **kwargs: Any): tokens = state.extra["tokens"] if tokens.dim() == 2: tokens = tokens.unsqueeze(0) audio = self.decoder(tokens) + reference_rms = state.extra["reference_rms"] + if reference_rms is not None: + audio = self._postprocess_audio(audio, reference_rms) return DiffusionOutput(output=audio) @torch.inference_mode() @@ -596,6 +846,9 @@ def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: for i, target_len in enumerate(batch_target_len): request_tokens = tokens[:, :, target_offset : target_offset + target_len] audio = self.decoder(request_tokens) + reference_rms = prepared_requests[i].reference_rms + if reference_rms is not None: + audio = self._postprocess_audio(audio, reference_rms) outputs[prepared_indices[i]] = DiffusionOutput(output=audio) target_offset += target_len return outputs diff --git a/vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py b/vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py index 3eceeb93b98..3d4d0f31cc5 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py +++ b/vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py @@ -19,6 +19,7 @@ import functools +import math from contextlib import contextmanager import torch @@ -1264,7 +1265,8 @@ def __init__( self.tile_sample_min_width = 512 # The minimal distance between two spatial tiles (75% of the tile - # size, i.e. 25% overlap) + # size, i.e. 25% overlap). At decode time it is an upper bound on the + # spacing: the actual starts are spread evenly, see _tile_starts. self.tile_sample_stride_height = 384 self.tile_sample_stride_width = 384 @@ -1429,18 +1431,23 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | t def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + # Every tile overlap is at least one cell (see _tile_starts), and an + # empty range would leave b untouched, so the window never divides by zero. for y in range(blend_extent): - b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( - y / blend_extent - ) + # Raised-cosine window: its slope vanishes at both ends, so the + # cross-fade does not leave a luma step where it meets a single + # tile. A linear ramp does, because its slope jumps there + # (vllm-omni#8850). + weight = 0.5 - 0.5 * math.cos(math.pi * (y + 0.5) / blend_extent) + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - weight) + b[:, :, :, y, :] * weight return b def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) for x in range(blend_extent): - b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( - x / blend_extent - ) + # Raised-cosine window, see blend_v. + weight = 0.5 - 0.5 * math.cos(math.pi * (x + 0.5) / blend_extent) + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - weight) + b[:, :, :, :, x] * weight return b @_with_cudnn_deterministic @@ -1562,6 +1569,33 @@ def _tiled_decode_adaptive(self, z: torch.Tensor, return_dict: bool = True) -> D self.tile_sample_stride_width, ) = saved + @staticmethod + def _tile_starts(size: int, tile_min: int, stride: int) -> list[int]: + r"""Start offset of every tile that covers a ``size``-cell axis. + + The starts are spread evenly over ``[0, size - tile_min]``: every tile + keeps the full decode window, and each pair of neighbours overlaps by + roughly the same amount. ``stride`` is therefore an upper bound on the + spacing, not the exact spacing. + + Two layouts are wrong here (vllm-omni#8850). Stepping by ``stride`` until + it runs out of axis clips the trailing tile at the border; a tile + narrower than ``tile_min`` decodes with less context than its neighbours + and leaves a seam the blend cannot hide. Snapping only the trailing tile + fixes that but leaves one unusually wide cross-fade and can make the last + blend read a region that was already blended. + """ + if size <= tile_min: + return [0] + span = size - tile_min + count = math.ceil(span / stride) + 1 # fewest tiles whose spacing stays <= stride + base, extra = divmod(span, count - 1) + starts, offset = [0], 0 + for n in range(count - 1): + offset += base + (1 if n < extra else 0) + starts.append(offset) + return starts + @_with_cudnn_deterministic def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: r""" @@ -1585,25 +1619,39 @@ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutp tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio - tile_sample_stride_height = self.tile_sample_stride_height - tile_sample_stride_width = self.tile_sample_stride_width if self.config.patch_size is not None: sample_height = sample_height // self.config.patch_size sample_width = sample_width // self.config.patch_size - tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size - tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size - blend_height = self.tile_sample_min_height // self.config.patch_size - tile_sample_stride_height - blend_width = self.tile_sample_min_width // self.config.patch_size - tile_sample_stride_width - else: - blend_height = self.tile_sample_min_height - tile_sample_stride_height - blend_width = self.tile_sample_min_width - tile_sample_stride_width - # Split z into overlapping tiles and decode them separately. - # The tiles have an overlap to avoid seams between tiles. + # Split z into overlapping tiles and decode them separately. The starts + # are spread evenly so every tile keeps the full decode window and the + # blends stay similar in width (see _tile_starts). + starts_height = self._tile_starts(height, tile_latent_min_height, tile_latent_stride_height) + starts_width = self._tile_starts(width, tile_latent_min_width, tile_latent_stride_width) + # Decoded cells contributed by one latent cell, before unpatchify. + cells_per_latent = self.spatial_compression_ratio // (self.config.patch_size or 1) + # A tile contributes up to the start of the next one; the trailing tile + # contributes the rest and is trimmed by the final crop. + keep_height = [ + (starts_height[n + 1] - starts_height[n]) * cells_per_latent for n in range(len(starts_height) - 1) + ] + keep_width = [(starts_width[n + 1] - starts_width[n]) * cells_per_latent for n in range(len(starts_width) - 1)] + # Neighbours are cross-faded over their real overlap. Snapping the + # trailing tile makes that overlap wider than the nominal blend extent, + # and blending a fixed extent would mix non-corresponding columns. + overlap_height = [ + (starts_height[n] + tile_latent_min_height - starts_height[n + 1]) * cells_per_latent + for n in range(len(starts_height) - 1) + ] + overlap_width = [ + (starts_width[n] + tile_latent_min_width - starts_width[n + 1]) * cells_per_latent + for n in range(len(starts_width) - 1) + ] + rows = [] - for i in range(0, height, tile_latent_stride_height): + for i in starts_height: row = [] - for j in range(0, width, tile_latent_stride_width): + for j in starts_width: self.clear_cache() time = [] for k in range(num_frames): @@ -1622,10 +1670,12 @@ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutp # blend the above tile and the left tile # to the current tile and add the current tile to the result row if i > 0: - tile = self.blend_v(rows[i - 1][j], tile, blend_height) + tile = self.blend_v(rows[i - 1][j], tile, overlap_height[i - 1]) if j > 0: - tile = self.blend_h(row[j - 1], tile, blend_width) - result_row.append(tile[:, :, :, :tile_sample_stride_height, :tile_sample_stride_width]) + tile = self.blend_h(row[j - 1], tile, overlap_width[j - 1]) + keep_h = keep_height[i] if i < len(keep_height) else tile.shape[-2] + keep_w = keep_width[j] if j < len(keep_width) else tile.shape[-1] + result_row.append(tile[:, :, :, :keep_h, :keep_w]) result_rows.append(torch.cat(result_row, dim=-1)) dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width] diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 19c31d3381f..cfd02b2c608 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -7,7 +7,7 @@ import tempfile from collections.abc import Callable from dataclasses import dataclass, field, fields -from typing import Any, cast +from typing import Any, Literal, cast from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs from vllm.logger import init_logger @@ -209,6 +209,7 @@ class OmniEngineArgs(EngineArgs): session_mode: str = "turn" retains_state_across_chunks: bool = False supports_running_prefix_cache_reset: bool = True + recompute_preemption: Literal["allow", "fail"] = "allow" use_v2_model_runner: bool = False supports_native_mrv2_data_plane: bool = False # WS-A: Stage-1 active stream slots. 0 = legacy preempt-everything. @@ -446,6 +447,7 @@ def create_model_config(self) -> OmniModelConfig: session_mode=self.session_mode, retains_state_across_chunks=self.retains_state_across_chunks, supports_running_prefix_cache_reset=self.supports_running_prefix_cache_reset, + recompute_preemption=self.recompute_preemption, use_v2_model_runner=self.use_v2_model_runner, supports_native_mrv2_data_plane=self.supports_native_mrv2_data_plane, active_stream_window=self.active_stream_window, diff --git a/vllm_omni/engine/duplex/session/model_channel.py b/vllm_omni/engine/duplex/session/model_channel.py index 26bab17f73b..b6fa73f3c40 100644 --- a/vllm_omni/engine/duplex/session/model_channel.py +++ b/vllm_omni/engine/duplex/session/model_channel.py @@ -1771,3 +1771,12 @@ async def maybe_continue_response( if scheduled and not self._ctx.run.closing: model_state.continuation_owner_id = owner_id model_state.continuation_units = count + 1 + return + if not response_owned and session.active_response_id is not None: + # The unit was planned for the model turn, and the answer opened while it + # waited: it is stale against its old owner, and the answer's own segment + # end only arrives after one more unit is processed. Re-aim the unit at the + # answer instead of dropping it, which would leave the answer half spoken + # for a client that has stopped sending audio. The re-aimed call re-checks + # the session state, so a closed session or a newer epoch still drops it. + await self.maybe_continue_response(expected_epoch=expected_epoch) diff --git a/vllm_omni/engine/duplex_omni_engine.py b/vllm_omni/engine/duplex_omni_engine.py index 22c02cef9cd..14f7bcf0b00 100644 --- a/vllm_omni/engine/duplex_omni_engine.py +++ b/vllm_omni/engine/duplex_omni_engine.py @@ -47,6 +47,12 @@ logger = init_logger(__name__) _DEFAULT_CONTROL_TIMEOUT_S = 10.0 +# Opening a session also initializes per-session runner state (session runner, +# talker admission, first stage reservation), so the first open on a cold engine +# can take tens of seconds while later opens are fast. Keep the short control +# budget for control ops on a live session and give open its own budget, or a +# cold start fails the first open with a spurious timeout (#9072). +_DEFAULT_OPEN_TIMEOUT_S = 60.0 _COMMAND_PUT_TIMEOUT_S = 30.0 @@ -173,7 +179,7 @@ def _open_session( session_config: DuplexSessionConfig, *, output_buffer: DuplexOutputBuffer, - timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + timeout: float | None = _DEFAULT_OPEN_TIMEOUT_S, ) -> DuplexControlResultMessage: control_id = uuid.uuid4().hex return self._execute_control( @@ -195,7 +201,7 @@ async def open_session_async( session_config: DuplexSessionConfig, *, output_buffer: DuplexOutputBuffer, - timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + timeout: float | None = _DEFAULT_OPEN_TIMEOUT_S, ) -> DuplexControlResultMessage: loop = asyncio.get_running_loop() return await loop.run_in_executor( diff --git a/vllm_omni/entrypoints/duplex_omni.py b/vllm_omni/entrypoints/duplex_omni.py index 508b630015a..1e297aaff1e 100644 --- a/vllm_omni/entrypoints/duplex_omni.py +++ b/vllm_omni/entrypoints/duplex_omni.py @@ -50,6 +50,9 @@ async def consume(): logger = init_logger(__name__) _DEFAULT_CONTROL_TIMEOUT_S = 10.0 +# A cold engine can spend tens of seconds initializing per-session state before +# the first open lands; control ops on a live session keep the short budget. +_DEFAULT_OPEN_TIMEOUT_S = 60.0 class DuplexSessionHandle: @@ -375,7 +378,7 @@ async def open_session( self, config: DuplexSessionConfig | Mapping[str, object] | None = None, *, - timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + timeout: float | None = _DEFAULT_OPEN_TIMEOUT_S, ) -> DuplexSessionHandle: """Open an engine-resident session and return its handle. diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 54b373c64ed..57dd0429f41 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -2583,8 +2583,7 @@ async def _create_diffusion_speech( # to allow tests to catch invalid voices generally request.voice = self._get_normalized_voice(request.voice) - # Assume that this will follow the same adapter pattern - # once all RFC is implemented + assert self._adapter is not None validation_error = validate_tts_request(self._adapter, request) if validation_error is not None: raise ValueError(validation_error) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/omnivoice.py b/vllm_omni/entrypoints/openai/tts_adapters/omnivoice.py index b2baa72898e..8766204bd96 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/omnivoice.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/omnivoice.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """OmniVoice serving adapter (AR engine_client path). OmniVoice can also be served through the pure-diffusion engine via @@ -41,7 +42,7 @@ async def build( if request.ref_audio: wav, sr, _ = await server._resolve_ref_audio(request.ref_audio) prompt["ref_audio"] = (np.asarray(wav, dtype=np.float32), sr) - if request.ref_text: + if request.ref_text and request.ref_text.strip(): prompt["ref_text"] = request.ref_text if request.voice: voice_lower = request.voice.lower() diff --git a/vllm_omni/model_executor/models/audio8_tts/pipeline.py b/vllm_omni/model_executor/models/audio8_tts/pipeline.py index 836ab71d98b..43c6b36e9ef 100644 --- a/vllm_omni/model_executor/models/audio8_tts/pipeline.py +++ b/vllm_omni/model_executor/models/audio8_tts/pipeline.py @@ -34,6 +34,10 @@ "detokenize": False, "stop_token_ids": [AUDIO8_TTS_EOS_TOKEN_ID], }, + # talker_mtp folds the sampled codebook embeddings into this + # step's inputs_embeds; a replay rebuilds only the text part and + # pads the rest (vllm-project/vllm-omni#6179). + recompute_preemption="fail", ), StagePipelineConfig( stage_id=1, diff --git a/vllm_omni/model_executor/models/fish_omni/pipeline.py b/vllm_omni/model_executor/models/fish_omni/pipeline.py index 511ace54a03..9440806d498 100644 --- a/vllm_omni/model_executor/models/fish_omni/pipeline.py +++ b/vllm_omni/model_executor/models/fish_omni/pipeline.py @@ -33,6 +33,8 @@ # Drama 3 end of speech (config eos_token_id). "stop_token_ids": [248046], }, + # MRv2 saves generated codebooks on preemption and restores the full + # frame rows on resume, so retain the default recompute policy. ), StagePipelineConfig( stage_id=1, diff --git a/vllm_omni/model_executor/models/fish_speech/pipeline.py b/vllm_omni/model_executor/models/fish_speech/pipeline.py index 1cb113094c2..ed92919614e 100644 --- a/vllm_omni/model_executor/models/fish_speech/pipeline.py +++ b/vllm_omni/model_executor/models/fish_speech/pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Fish Speech S2 Pro pipeline topology (frozen). Stage 0: slow_ar — text → RVQ codec tokens (LLM autoregressive). @@ -36,6 +36,9 @@ # <|im_end|> — stop when the model emits end-of-turn. "stop_token_ids": [151645], }, + # talker_mtp adds the sampled codebook embeddings to this step's + # inputs_embeds; a replay pads the unbuilt rows (vllm-project/vllm-omni#6179). + recompute_preemption="fail", ), StagePipelineConfig( stage_id=1, diff --git a/vllm_omni/model_executor/models/gepard/pipeline.py b/vllm_omni/model_executor/models/gepard/pipeline.py index e6ce4600dea..18e9a7d7a1e 100644 --- a/vllm_omni/model_executor/models/gepard/pipeline.py +++ b/vllm_omni/model_executor/models/gepard/pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Gepard-1.0 pipeline topology. Single-stage AR TTS. The Qwen3.5 backbone runs under paged attention, the 32 @@ -40,6 +40,11 @@ "detokenize": False, "stop_token_ids": [_GEPARD_STOP_TOKEN_ID], }, + # A frame's 31 side-channel FSQ codes never enter the token stream, + # so a replay cannot rebuild the KV behind audio already delivered + # (#6179; the preemption warning at gepard_talker.py only reports + # the damage after it happens). + recompute_preemption="fail", ), ), ) diff --git a/vllm_omni/model_executor/models/minimax_h3/checkpoint.py b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py index b39af5cf860..4d4735d6792 100644 --- a/vllm_omni/model_executor/models/minimax_h3/checkpoint.py +++ b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py @@ -10,6 +10,7 @@ ) _FASTH3_MODULAR_MODEL_PREFIX = "FastVideo/" +_NATIVE_PARTITIONS = ("FL2VA", "Ref2VA") MINIMAX_H3_ENCODER_DOWNLOAD_PATTERNS = { partition: [ @@ -26,8 +27,11 @@ def is_minimax_h3_modular(model: str, revision: str | None = None) -> bool: path = Path(model) if path.is_dir(): # Native MiniMax-H3 snapshots can advertise the Diffusers modular - # class in ``model_index.json`` while retaining the legacy FL2VA / - # Ref2VA layout. FastH3 modular releases carry this marker file. + # class and even ship ``modular_model_index.json`` while retaining the + # legacy FL2VA / Ref2VA layout, so the partitions take precedence over + # the marker files that FastH3 modular releases carry. + if any((path / partition / "model_index.json").is_file() for partition in _NATIVE_PARTITIONS): + return False return (path / "modular_model_index.json").is_file() or (path / "fastvideo_inference.json").is_file() # The FastVideo release is the modular checkpoint handled by this # resolver. Do not classify the native MiniMax-H3 Hub repository from diff --git a/vllm_omni/model_executor/models/moss_tts/pipeline.py b/vllm_omni/model_executor/models/moss_tts/pipeline.py index 31053a92b83..ff4f998c24d 100644 --- a/vllm_omni/model_executor/models/moss_tts/pipeline.py +++ b/vllm_omni/model_executor/models/moss_tts/pipeline.py @@ -41,6 +41,9 @@ sampling_constraints={ "detokenize": False, }, + # The local depth transformer folds each sampled code frame into this + # step's inputs_embeds inside talker_mtp (vllm-project/vllm-omni#6179). + recompute_preemption="fail", ), StagePipelineConfig( stage_id=1, @@ -72,6 +75,9 @@ engine_output_type="latent", async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_raw_async_chunk"), sampling_constraints={"detokenize": False}, + # The local depth transformer folds each sampled code frame into this + # step's inputs_embeds inside talker_mtp (vllm-project/vllm-omni#6179). + recompute_preemption="fail", ), StagePipelineConfig( stage_id=1, @@ -111,6 +117,9 @@ # result; codec audio output remains incremental. "output_kind": RequestOutputKind.FINAL_ONLY, }, + # The local depth transformer folds each sampled code frame into this + # step's inputs_embeds inside talker_mtp (vllm-project/vllm-omni#6179). + recompute_preemption="fail", ), StagePipelineConfig( stage_id=1, diff --git a/vllm_omni/model_executor/models/personaplex/pipeline.py b/vllm_omni/model_executor/models/personaplex/pipeline.py index fe02b1fc42e..0e4d95ac694 100644 --- a/vllm_omni/model_executor/models/personaplex/pipeline.py +++ b/vllm_omni/model_executor/models/personaplex/pipeline.py @@ -50,6 +50,10 @@ sampling_constraints={ "detokenize": False, }, + # Each frame's base embedding carries the previous frame's temporal + # hidden plus the text step; talker_mtp adds the agent codebook + # (vllm-project/vllm-omni#6179). + recompute_preemption="fail", ), StagePipelineConfig( stage_id=1, diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py index a32462d3293..03537852acb 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py @@ -71,6 +71,12 @@ "stop_token_ids": [2150], }, requires_full_payload_input=True, + # A decode step feeds the code predictor's summed embedding plus the + # thinker's text_step into PagedAttention, but only the layer-0 + # codec id enters the token stream. Replaying the prompt after a + # preemption cannot rebuild layers 1..7 or last_talker_hidden, so + # fail the request instead of resuming it from frame 0. + recompute_preemption="fail", ), StagePipelineConfig( stage_id=2, diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index e1644f24215..a47b064e444 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py @@ -49,6 +49,9 @@ "detokenize": False, "stop_token_ids": [2150], }, + # talker_mtp sums the residual codebook embeddings and adds + # text_step; neither is in the token stream (vllm-project/vllm-omni#6179). + recompute_preemption="fail", ), StagePipelineConfig( stage_id=1, @@ -101,6 +104,9 @@ "detokenize": False, "stop_token_ids": [2150], }, + # talker_mtp sums the residual codebook embeddings and adds + # text_step; neither is in the token stream (vllm-project/vllm-omni#6179). + recompute_preemption="fail", ), ), ) diff --git a/vllm_omni/model_executor/models/voxcpm2/pipeline.py b/vllm_omni/model_executor/models/voxcpm2/pipeline.py index c92cf61ba92..b7c475ab7fe 100644 --- a/vllm_omni/model_executor/models/voxcpm2/pipeline.py +++ b/vllm_omni/model_executor/models/voxcpm2/pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """VoxCPM2 pipeline topology (frozen). Single-stage AR TTS: text → speech waveform in one pass. @@ -32,6 +32,13 @@ "detokenize": False, "stop_token_ids": [1], }, + # ``curr_embed_for_next`` / ``prev_feat_embed`` / CFM-VAE progress + # are per-request side state; replaying the prompt embeds audio ids + # through the text embedding table and faults (#6179). + # Decode feeds PagedAttention embeddings this stage cannot rebuild + # from ``output_token_ids``; a recompute preemption would splice or + # corrupt the audio (vllm-project/vllm-omni#6179). + recompute_preemption="fail", ), ), )