From ad3a162eb8a843c734b48f8462ecb63abbf4df9b Mon Sep 17 00:00:00 2001 From: rein yang <73573651+R2-Y@users.noreply.github.com> Date: Fri, 24 Jul 2026 17:20:39 +0800 Subject: [PATCH] [perf] Minicpm o4.5 optimization (#5228) Signed-off-by: R2-Y Co-authored-by: Sy03 <1370724210@qq.com> --- examples/online_serving/minicpmo/README.md | 120 +- .../minicpmo/streaming_chat_completion.py | 71 + recipes/OpenBMB/MiniCPM-o-4_5.md | 218 +- tests/benchmarks/test_serve_cli.py | 23 +- .../sched/test_omni_ar_scheduler_streaming.py | 22 + .../omni_connectors/test_adapter_and_flow.py | 17 + .../test_chunk_transfer_adapter.py | 52 + .../test_omni_connector_configs.py | 25 + .../test_orchestrator_stage_input_bridge.py | 73 + tests/engine/test_stage_engine_core_proc.py | 28 + .../online_serving/test_minicpmo_streaming.py | 29 + .../minicpmo_4_5/test_code2wav_batching.py | 431 ++++ .../models/minicpmo_4_5/test_hf_model_path.py | 74 - .../models/minicpmo_4_5/test_llm2tts.py | 62 +- .../models/minicpmo_4_5/test_pipeline.py | 103 +- .../minicpmo_4_5/test_talker_batching.py | 248 ++ .../test_minicpmo_4_5_async_chunk.py | 182 ++ .../test_minicpmo_4_5_omni.py | 1369 +---------- tests/test_config_factory.py | 32 +- tests/worker/test_gpu_ar_model_runner.py | 1 + tests/worker/test_native_duplex_hooks.py | 363 --- tests/worker/test_omni_gpu_model_runner.py | 27 + vllm_omni/config/stage_config.py | 4 + vllm_omni/core/sched/omni_ar_scheduler.py | 10 +- .../core/sched/omni_generation_scheduler.py | 24 +- vllm_omni/core/sched/omni_scheduler_mixin.py | 28 +- vllm_omni/data_entry_keys.py | 14 + vllm_omni/deploy/minicpmo_4_5.yaml | 57 +- vllm_omni/deploy/minicpmo_4_5_2gpu.yaml | 66 +- vllm_omni/deploy/minicpmo_4_5_3gpu.yaml | 61 +- vllm_omni/deploy/minicpmo_4_5_8x4090.yaml | 45 +- vllm_omni/deploy/minicpmo_4_5_batching.yaml | 88 + vllm_omni/deploy/minicpmo_4_5_duplex.yaml | 22 +- .../distributed/omni_connectors/adapter.py | 16 + .../chunk_transfer_adapter.py | 14 +- .../omni_connectors/utils/config.py | 28 + vllm_omni/engine/orchestrator.py | 14 +- vllm_omni/engine/stage_engine_core_proc.py | 7 + vllm_omni/engine/stage_init_utils.py | 10 + vllm_omni/engine/stage_pool.py | 2 + .../entrypoints/cli/benchmark/cli_args.py | 2 +- .../fullduplex/engine/intermediate.py | 17 +- .../models/minicpmo_4_5/batched_token2wav.py | 400 +++ .../minicpmo_4_5/minicpmo_4_5_code2wav.py | 547 +++++ .../models/minicpmo_4_5/minicpmo_4_5_omni.py | 249 +- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 2147 ++++------------- .../models/minicpmo_4_5/pipeline.py | 39 +- vllm_omni/model_executor/models/registry.py | 5 + .../minicpmo_4_5_omni.py | 290 ++- .../npu/worker/npu_ar_model_runner.py | 7 +- vllm_omni/utils/mm_outputs.py | 1 + vllm_omni/worker/gpu_ar_model_runner.py | 22 +- .../worker/gpu_generation_model_runner.py | 8 +- vllm_omni/worker/gpu_model_runner.py | 44 +- 54 files changed, 3861 insertions(+), 3997 deletions(-) create mode 100644 examples/online_serving/minicpmo/streaming_chat_completion.py create mode 100644 tests/engine/test_stage_engine_core_proc.py create mode 100644 tests/examples/online_serving/test_minicpmo_streaming.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py delete mode 100644 tests/model_executor/models/minicpmo_4_5/test_hf_model_path.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_talker_batching.py create mode 100644 tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py create mode 100644 vllm_omni/deploy/minicpmo_4_5_batching.yaml create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py diff --git a/examples/online_serving/minicpmo/README.md b/examples/online_serving/minicpmo/README.md index 11f6f8195..39bbef09b 100644 --- a/examples/online_serving/minicpmo/README.md +++ b/examples/online_serving/minicpmo/README.md @@ -24,18 +24,28 @@ including `librosa`. ## Start the backend server -Pick a deploy config that matches your GPU layout: +The deploy config auto-loads via `--omni`. +The default `vllm_omni/deploy/minicpmo_4_5.yaml` keeps all three stages on +logical device 0 with memory budgets of 65%, 15%, and 15%. For throughput, +`minicpmo_4_5_batching.yaml` gives the Thinker GPU 0 (90%) and colocates the +Talker (55%) and Code2Wav (35%) on GPU 1. Each stage admits at most four +concurrent sequences. -| config | GPUs | TP | Notes | -|---|---|---|---| -| `minicpmo_4_5.yaml` | 1 | 1 | Thinker and talker+t2w co-located on GPU0. | -| `minicpmo_4_5_2gpu.yaml` | 2 | 1 | Thinker on GPU0, talker+t2w on GPU1. | -| `minicpmo_4_5_3gpu.yaml` | 3 | 2 | Thinker 2-way TP on GPU0/1, talker+t2w share GPU2. | -| `minicpmo_4_5_8x4090.yaml` | 8 | 4 | Thinker 4-way TP on GPU0-3, talker+t2w on GPU4. | -| `minicpmo_4_5_3gpu_stage1_replicas.yaml` | 3 | 1 | Experimental validation profile: Thinker on GPU0, two talker+Token2wav replicas on GPU1/2. | -| `minicpmo_4_5_4gpu_stage1_replicas.yaml` | 4 | 1 | Experimental validation profile: Thinker on GPU0, three talker+Token2wav replicas on GPU1/2/3. | -| `minicpmo_4_5_8x4090_stage1_replicas.yaml` | 8 | 4 | Experimental validation profile: Thinker 4-way TP on GPU0-3, four talker+Token2wav replicas on GPU4-7. | -| `minicpmo_4_5_duplex.yaml` | 2 | 1 | Experimental native duplex profile. | +| deploy config | GPUs | Notes | +|---|---|---| +| `minicpmo_4_5.yaml` (default) | 1 | Memory-constrained compatibility layout. | +| `minicpmo_4_5_batching.yaml` | 2 | Recommended continuous-batching layout; Talker and Code2Wav share GPU 1. | +| `minicpmo_4_5_2gpu.yaml` | 2 | Backward-compatible alias for `minicpmo_4_5_batching.yaml`. | +| `minicpmo_4_5_3gpu.yaml` | 3 | One GPU per stage. | +| `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. | +| `minicpmo_4_5_duplex.yaml` | 1 | Experimental native full-duplex overlay. | + +The split pipeline preserves native-duplex epoch/turn identity, segment text, +turn completion, reference voice, and terminal-audio metadata through +Code2Wav. Focused CPU regressions cover this envelope; run the Realtime +scenario below for live barge-in validation on the target GPU. + +Default: ```bash vllm-omni serve openbmb/MiniCPM-o-4_5 \ @@ -49,13 +59,20 @@ For local ModelScope checkpoints, replace `openbmb/MiniCPM-o-4_5` with the checkpoint path. To start the experimental native duplex backend, use `vllm_omni/deploy/minicpmo_4_5_duplex.yaml`. -The `*_stage1_replicas.yaml` files exercise composite Stage1 replica routing -and failure recovery. They are validation profiles, not recommended production -entrypoints. +### Per-stage overrides -## Send chat requests +```bash +vllm serve openbmb/MiniCPM-o-4_5 --omni --trust-remote-code --port 8099 \ + --stage-overrides '{"0": {"gpu_memory_utilization": 0.65}}' +``` -From this directory, run the MiniCPM-specific curl or Python clients: +## Send multimodal requests + +```bash +cd examples/online_serving/minicpmo +``` + +### curl ```bash bash run_curl_multimodal_generation.sh text @@ -64,13 +81,40 @@ bash run_curl_multimodal_generation.sh use_audio '["text"]' python openai_chat_completion_client_for_multimodal_generation.py \ --query-type use_image \ - --host localhost \ + --port 8099 \ + --host localhost + +# Text-only (faster; no <|tts_bos|>) +python openai_chat_completion_client_for_multimodal_generation.py \ + --query-type text \ + --modalities text \ + --prompt "Briefly introduce yourself." +``` + +Streaming text + audio: + +```bash +python streaming_chat_completion.py \ + --base-url http://localhost:8099/v1 \ + --output minicpmo_stream.wav +``` + +The example prints text deltas immediately and joins the independently encoded +24 kHz WAV audio deltas into one valid WAV file. Add `--text-only` to skip +Talker and Code2Wav. + +Shared helpers also work if you pass MiniCPM defaults yourself: + +```bash +python ../openai_chat_completion_client_for_multimodal_generation.py \ + --model openbmb/MiniCPM-o-4_5 \ + --query-type text \ --port 8099 ``` -Speech output requires `chat_template_kwargs.use_tts_template=true`. Put that -field at the request root for curl; the OpenAI Python SDK can merge it from -`extra_body`. +Speech output no longer depends on a MiniCPM-specific default in the generic +serving layer. `chat_template_kwargs.use_tts_template=true` remains an +explicitly supported model option. ## Launch the Gradio demo @@ -86,6 +130,19 @@ python examples/online_serving/minicpmo/gradio_demo.py \ Open `http://:7862` in a browser. +## Daily-Omni accuracy + +Daily-Omni requires one A–D letter. Set +`chat_template_kwargs.enable_thinking=false` explicitly in `--extra-body`; +the generic benchmark CLI does not inject a MiniCPM-specific default. +Leaving reasoning enabled can exhaust the 256-token answer budget inside +`` and make first-letter extraction score reasoning text. + +For the established text benchmark, send +`--extra_body '{"modalities":["text"],"chat_template_kwargs":{"enable_thinking":false}}'`. +Requesting audio also benchmarks Talker and Code2Wav and changes the +assistant template, so it is not an apples-to-apples accuracy run. + ## Run the Realtime duplex CLI demo After the duplex backend is running, stream one WAV through the Realtime @@ -137,3 +194,26 @@ two-response WAV, its `--input-sha256`, and an - [Offline MiniCPM-o inference](../../offline_inference/minicpmo/) - [MiniCPM-o 4.5 recipe](../../../recipes/OpenBMB/MiniCPM-o-4_5.md) + +## Pipeline notes + +- Stage 1 performs request-owned AR continuous batching. Stage 2 keeps + request-owned Flow/HiFT caches and batches exact-shape-compatible chunks. +- Reference audio travels with the first codec chunk; Stage 2 owns its + temporary prompt WAV and evicts prompt features when the request finishes. +- Codec sampling reads the checkpoint `tts_config` (default deterministic + seed 42). Stage-1 YAML sampling parameters govern only the binary + continue/stop token exposed to vLLM. +- `StageRequestStats.batch_size` is request-scoped and does not report the + scheduler's execution batch. +- Stage 0 and Stage 1 use vLLM CUDA Graph capture. Stage 2 remains eager until + a dedicated exact-shape graph wrapper owns static I/O buffers and copies + request cache state outside capture. +- Co-locating all three stages minimizes hardware requirements but makes their + CUDA contexts contend for one GPU. Use the 8x4090 layout or a custom + multi-GPU deploy config when throughput is the primary goal. +- Output audio is base64 WAV in `message.audio.data` (24 kHz mono). +- Offline counterpart: + [`examples/offline_inference/minicpmo/`](../../offline_inference/minicpmo/) +- Recipe: + [`recipes/OpenBMB/MiniCPM-o-4_5.md`](../../../recipes/OpenBMB/MiniCPM-o-4_5.md) diff --git a/examples/online_serving/minicpmo/streaming_chat_completion.py b/examples/online_serving/minicpmo/streaming_chat_completion.py new file mode 100644 index 000000000..b4a4be2f0 --- /dev/null +++ b/examples/online_serving/minicpmo/streaming_chat_completion.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stream MiniCPM-o 4.5 text and audio from /v1/chat/completions.""" + +from __future__ import annotations + +import argparse +import base64 +import io +import wave +from pathlib import Path + +from openai import OpenAI + + +def append_wav_chunk(encoded: str, pcm_parts: list[bytes]) -> tuple[int, int, int]: + """Decode one base64 WAV delta and append its PCM frames.""" + with wave.open(io.BytesIO(base64.b64decode(encoded)), "rb") as wav: + channels = wav.getnchannels() + sample_width = wav.getsampwidth() + sample_rate = wav.getframerate() + pcm_parts.append(wav.readframes(wav.getnframes())) + return channels, sample_width, sample_rate + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--base-url", default="http://localhost:8099/v1") + parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") + parser.add_argument("--prompt", default="Say hello, then introduce vLLM in one sentence.") + parser.add_argument("--output", type=Path, default=Path("minicpmo_stream.wav")) + parser.add_argument("--text-only", action="store_true") + args = parser.parse_args() + + modalities = ["text"] if args.text_only else ["text", "audio"] + stream = OpenAI(base_url=args.base_url, api_key="EMPTY").chat.completions.create( + model=args.model, + messages=[{"role": "user", "content": args.prompt}], + modalities=modalities, + stream=True, + extra_body={"chat_template_kwargs": {"use_tts_template": not args.text_only}}, + ) + + pcm_parts: list[bytes] = [] + audio_format: tuple[int, int, int] | None = None + for chunk in stream: + for choice in chunk.choices: + content = getattr(choice.delta, "content", None) + if not content: + continue + if getattr(chunk, "modality", None) == "audio": + current_format = append_wav_chunk(content, pcm_parts) + if audio_format is not None and current_format != audio_format: + raise ValueError(f"Audio format changed: {audio_format} -> {current_format}") + audio_format = current_format + elif getattr(chunk, "modality", None) == "text": + print(content, end="", flush=True) + + print() + if pcm_parts and audio_format is not None: + channels, sample_width, sample_rate = audio_format + with wave.open(str(args.output), "wb") as wav: + wav.setnchannels(channels) + wav.setsampwidth(sample_width) + wav.setframerate(sample_rate) + wav.writeframes(b"".join(pcm_parts)) + print(f"Audio saved to {args.output} ({sample_rate} Hz)") + + +if __name__ == "__main__": + main() diff --git a/recipes/OpenBMB/MiniCPM-o-4_5.md b/recipes/OpenBMB/MiniCPM-o-4_5.md index 68e8aac9e..73d0fd5a5 100644 --- a/recipes/OpenBMB/MiniCPM-o-4_5.md +++ b/recipes/OpenBMB/MiniCPM-o-4_5.md @@ -17,22 +17,23 @@ Use this recipe as a known-good starting point for serving `openbmb/MiniCPM-o-4_5` on vLLM-Omni. MiniCPM-o 4.5 is the omni member -of the MiniCPM-o family — it pairs a multimodal-understanding thinker -LLM with a streaming `MiniCPMTTS + Token2Wav` talker so a single +of the MiniCPM-o family — it runs a multimodal thinker, a streaming +MiniCPMTTS codec talker, and a separate batched Code2Wav stage so a single `/v1/chat/completions` call can return text and 24 kHz speech in one -shot. The recipe covers the shipped GPU layouts (single / 2 / 3 / 8 GPUs): - the default co-locates both stages on one GPU, and the larger scale-out - layouts (2 / 3 / 8 GPUs) are selected via `--deploy-config`. +shot. The recommended batching deploy isolates the Thinker on GPU 0 and +co-locates Talker and Code2Wav on GPU 1; 1-GPU, 3-GPU, and 8x4090 layouts are +also provided. ## References - Default deploy configs (auto-loaded by HF `model_type=minicpmo` + `hf_config.version="4.5"`): - - Single-GPU layout (default): + - Default single-GPU compatibility layout (auto-loaded): [`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml) - - 2-GPU layout (thinker on GPU 0, talker on GPU 1): - [`vllm_omni/deploy/minicpmo_4_5_2gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_2gpu.yaml) - - 3-GPU layout (thinker TP=2): + - Recommended 2-GPU continuous-batching layout: + [`vllm_omni/deploy/minicpmo_4_5_batching.yaml`](../../vllm_omni/deploy/minicpmo_4_5_batching.yaml) + - 2-GPU and 3-GPU layouts: + [`vllm_omni/deploy/minicpmo_4_5_2gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_2gpu.yaml), [`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml) - 8x RTX 4090 layout: [`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml) @@ -42,7 +43,7 @@ shot. The recipe covers the shipped GPU layouts (single / 2 / 3 / 8 GPUs): [`examples/offline_inference/minicpmo/`](../../examples/offline_inference/minicpmo/) - Pipeline / talker source: [`vllm_omni/model_executor/models/minicpmo_4_5/`](../../vllm_omni/model_executor/models/minicpmo_4_5/) -- Stage-input processor (thinker → talker bridge): +- Stage-input processors (thinker → talker and talker → Code2Wav): [`vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py`](../../vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py) - Upstream model card: [`openbmb/MiniCPM-o-4_5`](https://huggingface.co/openbmb/MiniCPM-o-4_5) @@ -51,18 +52,25 @@ shot. The recipe covers the shipped GPU layouts (single / 2 / 3 / 8 GPUs): ## Hardware Support -Four GPU layouts ship with default deploy configs. Pick the layout that -matches your hardware and pass it via `--deploy-config`. The default -co-locates the thinker and the talker (`MiniCPMTTS + Token2Wav`, with its -in-process vocoder) on a single GPU; the multi-GPU layouts split the -talker onto its own GPU and scale the thinker out via TP. +Four hardware layouts ship with deploy configs. Every layout uses the +same strict three-stage topology. The Talker emits codec chunks only; +Code2Wav consumes them through a shared-memory async connector. -| Layout | Thinker | Talker + Token2Wav | Typical hardware | -| --- | --- | --- | --- | -| Single-GPU (default) | GPU 0 | GPU 0 | 1x A100/H100/H200 80GB | -| 2-GPU | GPU 0 | GPU 1 | 2x A100/H100/H200 80GB | -| 3-GPU (thinker TP=2) | GPU 0,1 (TP=2) | GPU 2 | 3x mid-tier GPUs | -| 8x RTX 4090 24GB | GPU 0–3 (TP=4) | GPU 4 | 8x RTX 4090 consumer | +| Layout | Thinker | Talker | Code2Wav | Typical hardware | +| --- | --- | --- | --- | --- | +| 1-GPU (default) | GPU 0 | GPU 0 | GPU 0 | 1x large-memory GPU | +| 2-GPU | GPU 0 | GPU 1 | GPU 1 | 2x large-memory GPU | +| 3-GPU | GPU 0 | GPU 1 | GPU 2 | 3x GPU | +| 8x RTX 4090 24GB | GPU 0–3 (TP=4) | GPU 4 | GPU 5 | 8x RTX 4090 consumer | + +### Migration from the fused deployment + +MiniCPM-o 4.5 now requires the three-stage topology: the Talker owns +request-local codec generation and Code2Wav owns waveform state and +reference-voice prompt features. `minicpmo_4_5.yaml` remains the stable +single-GPU entry point; `minicpmo_4_5_batching.yaml` is the recommended +two-GPU profile. The removed fused two-stage implementation is not retained as +a fallback because it would duplicate state machines and correctness paths. ## GPU @@ -70,11 +78,10 @@ talker onto its own GPU and scale the thinker out via TP. The default [`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml) -co-locates both stages on GPU 0: the thinker (`~80 %` memory) and the -talker + Token2Wav vocoder (`~10 %` memory, `max_num_seqs: 1`). This is -the recommended starting layout — works on a single 80GB-class GPU -(A100, H100, H200) as long as the thinker model weights and the talker's -in-process vocoder fit together. +co-locates Thinker, codec-only Talker, and Code2Wav on GPU 0. Their +`gpu_memory_utilization` budgets are 0.65, 0.15, and 0.15. This layout +minimizes the GPU count; use a large-memory accelerator and leave the +remaining 5% for runtime overhead. #### Environment @@ -95,6 +102,40 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \ The deploy config is auto-loaded by the model registry — no `--deploy-config` flag needed for this default single-GPU layout. +For the recommended two-GPU layout, add: + +```bash +--deploy-config vllm_omni/deploy/minicpmo_4_5_batching.yaml +``` + +#### Performance comparison + +Compare text-only and text+audio separately. Text-only isolates Thinker +generation; text+audio also schedules Talker and Code2Wav. The following full +Daily-Omni runs used the same two GPUs, 1197 samples, concurrency 10, and +identical `enable_thinking=false` / `use_tts_template=true` request settings. +The `origin/main` fused Talker ran eager because its graph capture copied an +unpinned CPU metadata tensor. + +| Metric | `origin/main` two-stage | Three-stage batching | +| --- | ---: | ---: | +| Accuracy | 64.83% | 64.83% | +| Throughput | 0.62 req/s | 1.97 req/s | +| Mean E2EL | 16.17 s | 5.07 s | +| Mean serving TTFT | 0.92 s | 1.28 s | +| Mean audio TTFP | 16.17 s | 3.24 s | +| Mean audio RTF | 5.97 | 2.11 | +| Stage 0 mean TPOT / ITL | 8.27 / 8.27 ms | 40.08 / 40.11 ms | +| Stage 0 median TPOT / ITL | 7.23 / 7.24 ms | 7.43 / 7.53 ms | + +The split pipeline improves throughput 3.19x and lowers audio TTFP by 80%. +Isolating the Thinker on GPU 0 also removes the prior single-GPU TPOT +regression: 40.08 ms is slightly better than the pre-rebase report (~44 ms). +Its median TPOT is effectively the same as main; the higher mean is queueing +tail latency because this profile bounds each stage to four sequences while +main's Thinker admits 16. Global TPOT/ITL remains zero when serving emits text +as one aggregated chunk, so the table reports Stage 0 metrics. + #### Verification **Quick smoke test (text-only output)**: @@ -109,10 +150,10 @@ curl http://localhost:8099/v1/chat/completions \ }' ``` -**Text + speech in one response** (the headline 4.5 feature). The TTS -path is gated by a Jinja flag on the chat template. Pass -`use_tts_template=true` via the **top-level** `chat_template_kwargs` -field (curl does not flatten nested `extra_body`): +**Text + speech in one response** (the headline 4.5 feature). The model +bridge conditions the Talker from the generated assistant span, so the +generic serving layer does not inject MiniCPM-specific template defaults. +`use_tts_template=true` remains supported when explicitly requested: ```bash curl http://localhost:8099/v1/chat/completions \ @@ -134,6 +175,17 @@ in another choice's `message.audio.data` (24 kHz mono, see Notes). With `modalities: ["text", "audio"]` you typically get two `choices` entries (one text, one audio). +**Streaming text + speech**: + +```bash +python examples/online_serving/minicpmo/streaming_chat_completion.py \ + --base-url http://localhost:8099/v1 \ + --output minicpmo_stream.wav +``` + +The client prints text deltas as they arrive and reconstructs one valid WAV +from the independently encoded audio deltas. + **Gradio demo (text + image + audio + video UI)**: ```bash @@ -150,68 +202,32 @@ speech output (TTS)"** checkbox on / off. #### Notes -- Memory budget: both stages share GPU 0 — the thinker at - `gpu_memory_utilization: 0.8`, the talker + Token2Wav vocoder at `0.1`. +- Memory budget: Thinker, Talker, and Code2Wav reserve 0.65, 0.15, and + 0.15 of GPU 0. The larger Thinker share protects its multimodal KV cache; + all three model processes still share one CUDA device. - `--trust-remote-code` is required — the HF repo ships a custom `MiniCPMO` config / model class. -- Stage 1 (talker) is hard-capped to `max_num_seqs: 1`: the talker - only consumes `runtime_additional_information[0]`, so any value > 1 - makes concurrent requests share request-0's audio. This is the same - cap baked into the deploy config. - -### 2 x GPU (talker on its own GPU) - -Use -[`vllm_omni/deploy/minicpmo_4_5_2gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_2gpu.yaml) -when you have two GPUs and want to give the talker + Token2Wav vocoder a -dedicated card instead of sharing GPU 0 with the thinker. The thinker -runs on GPU 0 (`~90 %` mem, TP=1) and the talker on GPU 1 (`~75 %` mem, -`max_num_seqs: 1`). This relieves the memory pressure of the default -single-GPU co-located layout and is the recommended step up when a -second 80GB-class card is available but full 3-way TP scale-out is not -needed. - -#### Command - -```bash -vllm serve openbmb/MiniCPM-o-4_5 --omni \ - --deploy-config vllm_omni/deploy/minicpmo_4_5_2gpu.yaml \ - --trust-remote-code \ - --host 0.0.0.0 --port 8099 -``` - -Verification and Notes mirror the single-GPU section; the only -difference is that the talker no longer competes with the thinker for -GPU 0 memory. - -### 3 x GPU (thinker TP=2) - -Use -[`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml) -when you have a third GPU available and want the thinker on 2-way -tensor parallel for higher throughput; the talker stays on its own -GPU (talker has its own in-process Token2Wav vocoder, so co-locating -it with the thinker risks OOM under load). - -#### Command - -```bash -vllm serve openbmb/MiniCPM-o-4_5 --omni \ - --deploy-config vllm_omni/deploy/minicpmo_4_5_3gpu.yaml \ - --trust-remote-code \ - --host 0.0.0.0 --port 8099 -``` - -Verification and Notes mirror the single-GPU section; thinker latency -roughly halves under load thanks to TP=2. +- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 remains + eager because its request-owned Flow/HiFT caches and variable chunk/cache + shapes are not yet exposed through a static exact-shape graph wrapper. +- All default stages use `max_num_seqs: 4` to reduce cross-process GPU + contention. Talker AR + state and Code2Wav caches are request-owned; Code2Wav batches only + exact-shape-compatible chunks and does not fall back to serial decode. +- `StageRequestStats.batch_size` is a request-scoped placeholder, not the + scheduler's execution batch. +- Single-GPU co-location trades throughput for hardware density: Stage 0/1 + CUDA Graph replay and eager Stage 2 vocoder kernels compete across three + CUDA contexts. Use the 8x4090 config or a custom multi-GPU mapping for + throughput-sensitive serving. ### 8 x RTX 4090 24GB (consumer-GPU layout) Use [`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml) on an 8x RTX 4090 host. Thinker uses 4-way TP across GPUs 0–3 -(`~85 %` mem each ≈ 20.4 GiB/card), talker + Token2Wav lives on GPU 4 -(`~90 %` mem). GPUs 5–7 are left free. +(`~85 %` mem each ≈ 20.4 GiB/card), Talker uses GPU 4, and Code2Wav +uses GPU 5. GPUs 6–7 are left free. #### Command @@ -233,8 +249,8 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \ ## Notes (applies to all layouts) -- **Talker dependency**: the `MiniCPM-o 4.5` talker calls - `from stepaudio2 import Token2wav` against the MiniCPM-o-flavored +- **Code2Wav dependency**: Stage 2 loads `Token2wav` from the + MiniCPM-o-flavored vocoder (PyPI package `stepaudio2-minicpmo` — NOT the upstream `stepfun-ai/Step-Audio2`, whose `Token2wav.__init__` signature rejects `n_timesteps`). Install via the published extra: @@ -247,13 +263,20 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \ missing dep raises `ImportError` at first request with the same install hint instead of silently emitting empty audio. -- **TTS trigger**: speech output requires - `chat_template_kwargs.use_tts_template=true` so the chat template - appends `<|tts_bos|>` before generation. Without it, Stage-1 talker - receives no TTS token span and returns silent audio (not text-only). - For **curl**, put `chat_template_kwargs` at the request root; nested - `extra_body.chat_template_kwargs` is ignored. The OpenAI Python SDK - may use `extra_body` because it flattens those fields into the root. +- **TTS conditioning**: the MiniCPM stage bridge can condition speech from + the generated assistant span without changing shared serving code. + `chat_template_kwargs.use_tts_template=true` remains supported when an + explicit `<|tts_bos|>` boundary is desired. For **curl**, put + `chat_template_kwargs` at the request root; the OpenAI Python SDK may use + `extra_body` because it flattens those fields into the root. + +- **Reference voice**: request audio is carried on the first codec chunk. + Code2Wav owns the temporary prompt WAV and prompt-feature cache, and removes + both when the stream ends. + +- **Talker sampling**: codec-token sampling reads the checkpoint `tts_config` + and defaults to deterministic seed 42. Stage-1 deploy sampling parameters + control only vLLM's binary continue/stop token. - **Output audio**: 24 kHz mono WAV inside the OpenAI-style `message.audio.data` (base64). The Gradio demo's WAV player decodes @@ -266,6 +289,9 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \ with this recipe's `--deploy-config` will be rejected at startup rather than silently misrouted. -- **Async chunking**: disabled in all three deploy configs - (`async_chunk: false`) — the talker batches a single full thinker - output, not chunks. +- **Async chunking**: enabled in all deploy configs. Talker sends + 25-code chunks with three-code left context to Code2Wav through + `SharedMemoryConnector`; terminal chunks flush held lookahead state. +- **Response choices**: text and audio are separate choices. SDK clients + should select the choice whose `message.audio.data` is populated rather + than assuming `choices[0]` contains audio. diff --git a/tests/benchmarks/test_serve_cli.py b/tests/benchmarks/test_serve_cli.py index 9ec30639c..7dd12d000 100644 --- a/tests/benchmarks/test_serve_cli.py +++ b/tests/benchmarks/test_serve_cli.py @@ -99,19 +99,21 @@ def test_update_omni_help_updates_upstream_actions() -> None: @pytest.mark.parametrize( - ("extra_body", "bot_task", "expected"), + ("backend", "extra_body", "bot_task", "expected"), [ - (None, "think", {"bot_task": "think"}), - ({"bot_task": "recaption"}, "think", {"bot_task": "recaption"}), - ({"custom": True}, None, {"custom": True}), + ("openai-chat-omni", None, "think", {}), + ("openai-image-edits-omni", None, "think", {"bot_task": "think"}), + ("openai-image-edits-omni", {"bot_task": "recaption"}, "think", {"bot_task": "recaption"}), + ("openai-image-edits-omni", {"custom": True}, None, {"custom": True}), ], ) def test_preprocess_serve_args_merges_bot_task_without_overriding_extra_body( + backend: str, extra_body: dict | None, bot_task: str | None, expected: dict, ) -> None: - args = argparse.Namespace(extra_body=extra_body, bot_task=bot_task) + args = argparse.Namespace(backend=backend, extra_body=extra_body, bot_task=bot_task) preprocess_serve_args(args) @@ -122,9 +124,15 @@ def test_preprocess_serve_args_merges_bot_task_without_overriding_extra_body( ("argv", "expected_extra_body", "expected_explicit"), [ ( - ["--print-stage", "--image-edits-bot-task", "recaption"], + [ + "--backend", + "openai-image-edits-omni", + "--print-stage", + "--image-edits-bot-task", + "recaption", + ], {"bot_task": "recaption"}, - {"print_stage", "bot_task"}, + {"backend", "print_stage", "bot_task"}, ), ( ["--extra-body", '{"bot_task":"vanilla"}'], @@ -140,6 +148,7 @@ def test_omni_args_parse_and_preprocess( ) -> None: parser = TrackingArgumentParser() parser.add_argument("--extra-body", type=json.loads, default=None) + parser.add_argument("--backend", default="openai-chat-omni") add_omni_args(parser) args = parser.parse_args(argv) diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index 41fb414cb..32284eeac 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -26,6 +26,7 @@ def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler: sched.num_waiting_for_streaming_input = 0 sched.log_stats = False sched.chunk_transfer_adapter = None + sched.skipped_waiting = set() return sched @@ -90,3 +91,24 @@ def test_stage0_streaming_update_keeps_all_computed_tokens_without_placeholder() assert session._output_token_ids == [] assert session.num_prompt_tokens == 8 assert sched._new_prompt_len_snapshot[session.request_id] == 2 + + +def test_explicit_streaming_payload_replaces_placeholder_prompt() -> None: + sched = _make_scheduler(stage_id=1) + sched.chunk_transfer_adapter = SimpleNamespace( + receives_chunks=False, + segment_finished_requests=set(), + ) + session = _make_request() + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + update = _make_update([10, 20]) + update.additional_information = { + "tts_token_ids": [10, 20], + "meta": {"replace_streaming_prompt": True}, + } + + sched._update_request_as_session(session, update) + + assert session.prompt_token_ids == [10, 20] + assert session.additional_information == update.additional_information + assert session.status == RequestStatus.WAITING diff --git a/tests/distributed/omni_connectors/test_adapter_and_flow.py b/tests/distributed/omni_connectors/test_adapter_and_flow.py index 92dc1d7c0..471d3108a 100644 --- a/tests/distributed/omni_connectors/test_adapter_and_flow.py +++ b/tests/distributed/omni_connectors/test_adapter_and_flow.py @@ -8,6 +8,7 @@ from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec, OmniTransferConfig from vllm_omni.distributed.omni_connectors.utils.initialization import get_connectors_config_for_stage +from vllm_omni.engine.stage_init_utils import get_stage_connector_spec pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -217,6 +218,22 @@ def test_get_connectors_for_stage(): assert stage_2_config["from_stage_1"]["spec"]["name"] == "C2" +def test_outgoing_only_async_stage_uses_sender_connector_spec(): + config = OmniTransferConfig( + connectors={ + ("1", "2"): ConnectorSpec( + name="SharedMemoryConnector", + extra={"codec_chunk_frames": 25}, + ) + } + ) + + spec = get_stage_connector_spec(config, stage_id=1, async_chunk=True) + + assert spec["name"] == "SharedMemoryConnector" + assert spec["extra"] == {"codec_chunk_frames": 25, "role": "sender"} + + def test_recv_with_missing_metadata(mocker: MockerFixture): """Test recv when queue payload is malformed (missing metadata).""" # Connector expects metadata but task doesn't have it diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 76d0cb550..1485ace60 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -13,6 +13,7 @@ from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.request import RequestStatus from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct +from vllm_omni.distributed.omni_connectors.adapter import construct_next_stage_streaming_input_prompt from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( OmniChunkTransferAdapter, @@ -43,6 +44,33 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None) ) +def test_streaming_payload_can_replace_placeholder_prompt(mocker: MockerFixture) -> None: + request = SimpleNamespace( + _all_token_ids=[0, 0, 7, 8], + _output_token_ids=[7, 8], + prompt_token_ids=[0, 0], + num_computed_tokens=4, + num_prompt_tokens=2, + update_block_hashes=mocker.Mock(), + ) + payload = { + "ids": {"prompt": [1, 2, 3]}, + "meta": { + "replace_streaming_prompt": True, + "next_stage_prompt_len": 7, + }, + } + + construct_next_stage_streaming_input_prompt(payload, request) + + assert request.prompt_token_ids == [0] * 7 + assert request._all_token_ids == [0] * 7 + assert request._output_token_ids == [] + assert request.num_computed_tokens == 0 + assert request.num_prompt_tokens == 7 + request.update_block_hashes.assert_called_once_with() + + @pytest.fixture def build_adapter(monkeypatch, mocker: MockerFixture): def _build( @@ -81,6 +109,10 @@ def build_adapter(monkeypatch, mocker: MockerFixture): worker_type=model_mode, max_num_seqs=max_num_seqs, active_stream_window=active_stream_window, + stage_connector_config={ + "name": "SharedMemoryConnector", + "extra": connector_extra or {}, + }, ) scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs) adapter = OmniChunkTransferAdapter( @@ -381,6 +413,26 @@ def test_load_poll_ar_request_additional_information_concats_tensors(build_adapt assert request.additional_information["meta"]["finished"].item() is True +def test_sender_only_adapter_does_not_park_or_clear_requests(build_adapter): + adapter, _ = build_adapter(stage_id=1, connector_extra={"role": "sender"}) + request = _req("req-1", RequestStatus.WAITING) + request.additional_information = {"tts_token_ids": torch.tensor([1])} + waiting_queue = DummyWaitingQueue([request]) + running_queue = [] + + adapter.load_async(request) + adapter.process_pending_chunks( + waiting_queue, + running_queue, + scheduler_requests={request.request_id: request}, + ) + + assert waiting_queue == [request] + assert request.status == RequestStatus.WAITING + assert request.additional_information["tts_token_ids"].item() == 1 + assert adapter._pending_load_reqs == deque() + + def test_process_and_restore_queues(build_adapter): adapter, _ = build_adapter(stage_id=1, max_num_seqs=8) waiting_req = _req("w1", RequestStatus.WAITING) diff --git a/tests/distributed/omni_connectors/test_omni_connector_configs.py b/tests/distributed/omni_connectors/test_omni_connector_configs.py index 82f13c5cd..6048de40a 100644 --- a/tests/distributed/omni_connectors/test_omni_connector_configs.py +++ b/tests/distributed/omni_connectors/test_omni_connector_configs.py @@ -2,10 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from pathlib import Path +from types import SimpleNamespace import pytest # Use the new import path for initialization utilities +from vllm_omni.distributed.omni_connectors.utils.config import ( + stage_receives_chunks, + stage_sends_async_output, +) from vllm_omni.distributed.omni_connectors.utils.initialization import load_omni_transfer_config pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -38,6 +43,26 @@ def get_config_files(): config_files = get_config_files() +@pytest.mark.parametrize( + ("role", "stage_id", "receives", "sends"), + [ + ("sender", 0, False, True), + ("receiver", 1, True, False), + (None, 0, True, False), + (None, None, True, True), + ], +) +def test_stage_chunk_direction_helpers(role, stage_id, receives, sends): + extra = {} if role is None else {"role": role} + model_config = SimpleNamespace( + stage_id=stage_id, + stage_connector_config={"extra": extra}, + ) + + assert stage_receives_chunks(model_config) is receives + assert stage_sends_async_output(model_config) is sends + + @pytest.mark.skipif(len(config_files) == 0, reason="No config files found or directory missing") @pytest.mark.parametrize("yaml_file", config_files, ids=lambda p: p.name) def test_load_qwen_yaml_configs(yaml_file): diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 3abb244eb..e74b9ea18 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -6,6 +6,7 @@ from __future__ import annotations import queue from types import SimpleNamespace from typing import Any +from unittest.mock import AsyncMock import janus import pytest @@ -91,6 +92,25 @@ class FakeInputProcessor: ) +class FakePrewarmPool: + stage_type = "llm" + + def __init__(self, role: str) -> None: + self.stage_vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + max_model_len=64, + stage_connector_config={"extra": {"role": role}}, + ) + ) + self.submitted: list[Any] = [] + + async def submit_initial(self, _request_id, _req_state, request, prompt_text=None): + self.submitted.append(request) + + def get_bound_replica_id(self, _request_id): + return 0 + + def _request_output(request_id: str) -> RequestOutput: completion = CompletionOutput( index=0, @@ -168,3 +188,56 @@ async def test_forward_text_prompt_uses_target_stage_input_processor() -> None: submitted_request = stage1.add_request_calls[0][0] assert submitted_request.prompt_token_ids == [101, 102] assert submitted_request.external_req_id == "req-text" + + +@pytest.mark.asyncio +async def test_async_prewarm_skips_outgoing_only_stage() -> None: + orchestrator = object.__new__(Orchestrator) + stage0 = FakePrewarmPool("sender") + stage1 = FakePrewarmPool("sender") + stage2 = FakePrewarmPool("receiver") + orchestrator.stage_pools = [stage0, stage1, stage2] + orchestrator._emit_tx_edge = lambda **_kwargs: None + req_state = OrchestratorRequestState( + request_id="req-prewarm", + prompt={"prompt_token_ids": [1, 2]}, + sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(3)], + final_stage_id=2, + ) + + await orchestrator._prewarm_async_chunk_stages( + "req-prewarm", + SimpleNamespace(prompt_token_ids=[1, 2], resumable=True), + req_state, + ) + + assert stage1.submitted == [] + assert len(stage2.submitted) == 1 + assert 1 not in req_state.stage_submit_ts + assert 2 in req_state.stage_submit_ts + + +@pytest.mark.asyncio +async def test_async_route_forwards_to_outgoing_only_stage() -> None: + orchestrator = object.__new__(Orchestrator) + orchestrator.async_chunk = True + orchestrator._pd_pair = None + orchestrator._cfg_tracker = SimpleNamespace( + is_companion=lambda _request_id: False, + has_companions=lambda _request_id: False, + ) + stage0 = SimpleNamespace(final_output=False) + stage1 = FakePrewarmPool("sender") + orchestrator.stage_pools = [stage0, stage1] + orchestrator._forward_to_next_stage = AsyncMock() + req_state = OrchestratorRequestState( + request_id="req-route", + sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(2)], + final_stage_id=1, + ) + req_state.stage_submit_ts[0] = 1.0 + output = SimpleNamespace(request_id="req-route", finished=True) + + await orchestrator._route_output(0, 0, output, req_state, None) + + orchestrator._forward_to_next_stage.assert_awaited_once() diff --git a/tests/engine/test_stage_engine_core_proc.py b/tests/engine/test_stage_engine_core_proc.py new file mode 100644 index 000000000..a7104b14d --- /dev/null +++ b/tests/engine/test_stage_engine_core_proc.py @@ -0,0 +1,28 @@ +from types import SimpleNamespace +from unittest.mock import patch + +from vllm.v1.engine.core import EngineCoreProc + +from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc + + +def test_preprocess_add_request_preserves_omni_fields(): + engine = StageEngineCoreProc.__new__(StageEngineCoreProc) + request = SimpleNamespace( + request_id="internal", + external_req_id="external", + additional_information={"conditioning": "payload"}, + ) + scheduler_request = SimpleNamespace() + + with patch.object( + EngineCoreProc, + "preprocess_add_request", + return_value=(scheduler_request, 3), + ): + result, current_wave = engine.preprocess_add_request(request) + + assert result is scheduler_request + assert current_wave == 3 + assert result.external_req_id == "external" + assert result.additional_information == {"conditioning": "payload"} diff --git a/tests/examples/online_serving/test_minicpmo_streaming.py b/tests/examples/online_serving/test_minicpmo_streaming.py new file mode 100644 index 000000000..90ba992eb --- /dev/null +++ b/tests/examples/online_serving/test_minicpmo_streaming.py @@ -0,0 +1,29 @@ +import base64 +import io +import runpy +import wave +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_append_wav_chunk_extracts_pcm_frames() -> None: + example = Path(__file__).parents[3] / "examples" / "online_serving" / "minicpmo" / "streaming_chat_completion.py" + append_wav_chunk = runpy.run_path(example)["append_wav_chunk"] + encoded_buffer = io.BytesIO() + with wave.open(encoded_buffer, "wb") as wav: + wav.setnchannels(1) + wav.setsampwidth(2) + wav.setframerate(24000) + wav.writeframes(b"\x01\x02\x03\x04") + + pcm_parts: list[bytes] = [] + audio_format = append_wav_chunk( + base64.b64encode(encoded_buffer.getvalue()).decode(), + pcm_parts, + ) + + assert audio_format == (1, 2, 24000) + assert pcm_parts == [b"\x01\x02\x03\x04"] diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py new file mode 100644 index 000000000..58f3d27b7 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py @@ -0,0 +1,431 @@ +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import ( + BatchedToken2Wav, +) +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import ( + MiniCPMO45Code2Wav, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeEncoder(nn.Module): + def __init__(self): + super().__init__() + self.calls: list[int] = [] + + def forward_chunk(self, xs, last_chunk=False, cnn_cache=None, att_cache=None): + batch, length, _ = xs.shape + self.calls.append(batch) + old_length = 0 if att_cache is None else att_cache.shape[3] + output = xs[:, : max(1, length - 1)] + cnn = xs[:, :1, :].transpose(1, 2).contiguous() + marker = xs[:, 0, 0].reshape(1, batch, 1, 1, 1) + att = marker.expand(1, batch, 1, old_length + output.shape[1], 1).clone() + return output, cnn, att + + +class _FakeBlock: + def __init__(self): + conv1 = SimpleNamespace(causal_padding=(1, 0)) + self.conv = SimpleNamespace( + in_channels=1, + out_channels=1, + block=[None, conv1], + ) + self.attn = SimpleNamespace(num_heads=1, head_dim=1) + + +class _FakeEstimator(nn.Module): + def __init__(self): + super().__init__() + self.blocks = [_FakeBlock()] + self.cfg_batches: list[int] = [] + self.speaker_order: list[list[float]] = [] + + def t_embedder(self, time): + return time[:, None] + + def blocks_forward_chunk( + self, + inputs, + time, + mask, + cnn_cache, + att_cache, + cnn_out, + att_out, + ): + del time, mask, cnn_cache, att_cache + self.cfg_batches.append(inputs.shape[0]) + self.speaker_order.append(inputs[:, 2, 0].tolist()) + marker = inputs[:, 1, 0] + cnn_out.copy_(marker.reshape(1, -1, 1, 1).expand_as(cnn_out)) + att_out.copy_(marker.reshape(1, -1, 1, 1, 1).expand_as(att_out)) + return inputs[:, 1:2] + + +class _FakeDecoder(nn.Module): + def __init__(self): + super().__init__() + self.estimator = _FakeEstimator() + self.inference_cfg_rate = 0.7 + self.register_buffer("rand_noise", torch.zeros(1, 1, 100), persistent=False) + + +class _FakeFlow(nn.Module): + def __init__(self): + super().__init__() + self.encoder = _FakeEncoder() + self.encoder_proj = nn.Identity() + self.decoder = _FakeDecoder() + self.spk_embed_affine_layer = nn.Identity() + + def input_embedding(self, tokens): + return tokens.to(torch.float32).unsqueeze(-1) + + +class _FakeHiFT(nn.Module): + def __init__(self): + super().__init__() + self.calls: list[int] = [] + + def forward(self, mel, source): + del source + self.calls.append(mel.shape[0]) + speech = mel[:, 0].repeat_interleave(3, dim=1) + generated_source = speech[:, None] + return speech, generated_source + + +class _FakeToken2Wav: + def __init__(self): + self.flow = _FakeFlow() + self.hift = _FakeHiFT() + self.float16 = False + self.n_timesteps = 2 + self.mel_cache_len = 1 + self.source_cache_len = 2 + self.speech_window = torch.hamming_window(4, periodic=False) + self.prompt_calls = 0 + + def _prepare_prompt(self, prompt_wav): + del prompt_wav + self.prompt_calls += 1 + return ( + torch.tensor([[5, 6]], dtype=torch.long), + torch.tensor([2], dtype=torch.int32), + torch.ones(1, 1), + torch.ones(1, 4, 1), + torch.tensor([4], dtype=torch.int32), + ) + + def stream(self, *args, **kwargs): + raise AssertionError("sequential stream fallback must never be called") + + def __call__(self, *args, **kwargs): + raise AssertionError("sequential __call__ fallback must never be called") + + +def _config(minimum: int = 1): + return SimpleNamespace( + model_config=SimpleNamespace( + model="/fake/model", + stage_connector_config={ + "extra": { + "code2wav_min_batch_size": minimum, + "prompt_cache_id": "shared", + "prompt_wav": "/fake/prompt.wav", + } + }, + ) + ) + + +def _model(): + token2wav = _FakeToken2Wav() + backend = BatchedToken2Wav(token2wav) + model = MiniCPMO45Code2Wav(vllm_config=_config()) + model.backend = backend + return model, token2wav + + +def _info( + request_id: str, + chunk_seq: int, + codes: list[int], + *, + last_chunk: bool = False, + cache_epoch: int = 0, +): + return { + "codes": {"audio": torch.tensor(codes, dtype=torch.long)}, + "meta": { + "request_id": request_id, + "chunk_seq": chunk_seq, + "cache_epoch": cache_epoch, + "last_chunk": last_chunk, + "prompt_cache_id": "shared", + }, + } + + +def _forward(model, infos, placeholder_counts=None, request_ids=None): + placeholder_counts = placeholder_counts or [1] * len(infos) + input_ids = torch.zeros(sum(placeholder_counts), dtype=torch.long) + return model( + input_ids=input_ids, + seq_token_counts=placeholder_counts, + runtime_additional_information=infos, + request_ids=request_ids, + ) + + +def test_adapter_runs_true_batch_cfg_and_splits_request_caches(): + token2wav = _FakeToken2Wav() + adapter = BatchedToken2Wav(token2wav) + prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") + states = adapter.setup_batch(prompt, 2) + audios, states = adapter.decode_batch( + torch.tensor([[10, 11], [20, 21]]), + prompt, + states, + last_chunk=False, + ) + + assert token2wav.prompt_calls == 1 + assert token2wav.flow.encoder.calls == [2, 2] + assert token2wav.flow.decoder.estimator.cfg_batches == [4, 4, 4, 4] + assert all(order == [1.0, 1.0, 0.0, 0.0] for order in token2wav.flow.decoder.estimator.speaker_order) + assert token2wav.hift.calls == [2] + assert len(audios) == 2 + cache0 = states[0].flow_cache["estimator_cnn_cache"] + cache1 = states[1].flow_cache["estimator_cnn_cache"] + assert cache0.data_ptr() != cache1.data_ptr() + assert cache0[0, 0, 0, 0, 0].item() == 10 + assert cache1[0, 0, 0, 0, 0].item() == 20 + + +def test_estimator_cache_stack_split_round_trip_preserves_cfg_rows(): + token2wav = _FakeToken2Wav() + adapter = BatchedToken2Wav(token2wav) + prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") + states = adapter.setup_batch(prompt, 2) + _, states = adapter.decode_batch( + torch.tensor([[10, 11], [20, 21]]), + prompt, + states, + last_chunk=False, + ) + + stacked = adapter._stack_flow_cache(states) + assert stacked["estimator_cnn_cache"].shape[2] == 4 + assert stacked["estimator_att_cache"].shape[2] == 4 + restored = adapter._split_flow_cache(stacked, 2) + for original, round_tripped in zip(states, restored, strict=True): + torch.testing.assert_close( + round_tripped["estimator_cnn_cache"], + original.flow_cache["estimator_cnn_cache"], + ) + torch.testing.assert_close( + round_tripped["estimator_att_cache"], + original.flow_cache["estimator_att_cache"], + ) + + +def test_model_preserves_output_slots_and_prefers_runtime_codes(): + model, token2wav = _model() + output = _forward( + model, + [_info("a", 0, [10, 11]), _info("b", 0, [20, 21])], + placeholder_counts=[3, 1], + ) + + audios = output.multimodal_outputs["model_outputs"] + assert len(audios) == 2 + assert len(output.multimodal_outputs["sr"]) == 2 + assert all(sr.item() == 24000 for sr in output.multimodal_outputs["sr"]) + assert all(audio.dtype == torch.float32 for audio in audios) + # Fake CFM uses two Euler steps whose deltas sum to one. Its conditional + # row is mu and its unconditional row is zero, so CFG produces 1.7 * mu. + torch.testing.assert_close(audios[0][0], torch.tensor(1.7 * 10)) + torch.testing.assert_close(audios[1][0], torch.tensor(1.7 * 20)) + assert token2wav.flow.encoder.calls[-1] == 2 + + +def test_mixed_final_exact_buckets_keep_order_and_release_only_final_states(): + model, _ = _model() + _forward( + model, + [_info(name, 0, [index + 1, index + 2]) for index, name in enumerate(("a", "b", "c", "d"))], + ) + output = _forward( + model, + [ + _info("a", 1, [11, 12]), + _info("c", 1, [31, 32, 33], last_chunk=True), + _info("b", 1, [21, 22]), + _info("d", 1, [41, 42, 43], last_chunk=True), + ], + ) + + audios = output.multimodal_outputs["model_outputs"] + window = torch.hamming_window(4, periodic=False) + overlap_scale = 1.7 * (window[0] + window[2]) + expected = torch.tensor([1, 3, 2, 4], dtype=torch.float32) * overlap_scale + actual = torch.stack([audio[0] for audio in audios]) + torch.testing.assert_close(actual, expected) + assert set(model._states) == {"a", "b"} + + +def test_empty_final_sentinel_emits_empty_and_releases_state_without_compute(): + model, token2wav = _model() + _forward(model, [_info("a", 0, [1, 2]), _info("b", 0, [3, 4])]) + hift_calls = list(token2wav.hift.calls) + output = _forward( + model, + [ + _info("a", 1, [], last_chunk=True), + _info("b", 1, [], last_chunk=True), + ], + ) + + assert [audio.numel() for audio in output.multimodal_outputs["model_outputs"]] == [0, 0] + assert model._states == {} + assert token2wav.hift.calls == hift_calls + + +def test_empty_final_ignores_generation_scheduler_placeholder_token(): + model, _ = _model() + _forward(model, [_info("a", 0, [1, 2]), _info("b", 0, [3, 4])]) + infos = [_info("a", 1, [], last_chunk=True), _info("b", 1, [], last_chunk=True)] + for info in infos: + info.pop("codes") + info["meta"]["code_flat_numel"] = 0 + + output = _forward(model, infos, placeholder_counts=[1, 1]) + + assert [audio.numel() for audio in output.multimodal_outputs["model_outputs"]] == [0, 0] + assert model._states == {} + + +@pytest.mark.parametrize( + ("info", "reason"), + [ + (_info("a", 0, [1, 2], cache_epoch=-1), "negative_stream_position"), + (_info("a", 0, [1, 2]), "stale_or_reordered_chunk"), + (_info("a", 2, [1, 2]), "stale_or_reordered_chunk"), + ], +) +def test_stale_epoch_and_reordered_chunks_are_rejected(info, reason): + model, _ = _model() + _forward(model, [_info("a", 0, [1, 2]), _info("b", 0, [3, 4])]) + + with pytest.raises(RuntimeError, match=reason): + _forward(model, [info, _info("b", 1, [3, 4])]) + + +def test_singleton_and_mixed_shape_buckets_use_same_batched_backend_without_fallback(): + model, token2wav = _model() + _forward(model, [_info("a", 0, [1, 2]), _info("b", 0, [3, 4])]) + output = _forward(model, [_info("a", 1, [5, 6]), _info("b", 1, [7, 8, 9])]) + + assert len(output.multimodal_outputs["model_outputs"]) == 2 + # Exact-shape buckets execute independently but both use the same vectorized + # adapter; there is no Token2wav.stream/__call__ fallback. + assert token2wav.hift.calls[-2:] == [1, 1] + + +def test_backend_failure_does_not_commit_any_request_state(monkeypatch): + model, _ = _model() + _forward( + model, + [_info(name, 0, [index + 1, index + 2]) for index, name in enumerate(("a", "b", "c", "d"))], + ) + before = dict(model._states) + original = model.backend.decode_batch + call_count = 0 + + def fail(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 2: + raise RuntimeError("injected failure") + return original(*args, **kwargs) + + monkeypatch.setattr(model.backend, "decode_batch", fail) + with pytest.raises(RuntimeError, match="injected failure"): + _forward( + model, + [ + _info("a", 1, [5, 6]), + _info("b", 1, [7, 8]), + _info("c", 1, [9, 10, 11]), + _info("d", 1, [12, 13, 14]), + ], + ) + assert call_count == 2 + assert model._states == before + + +def test_cleanup_and_profile_output_are_aligned(): + model, _ = _model() + _forward(model, [_info("a", 0, [1, 2]), _info("b", 0, [3, 4])]) + model.on_requests_finished(["a"]) + assert set(model._states) == {"b"} + + profile = model( + input_ids=torch.zeros(5, dtype=torch.long), + seq_token_counts=[2, 3], + ) + assert [audio.numel() for audio in profile.multimodal_outputs["model_outputs"]] == [0, 0] + assert set(model._states) == {"b"} + + +def test_cleanup_uses_generation_runner_internal_request_ids(): + model, _ = _model() + _forward( + model, + [_info("external-a", 0, [1, 2]), _info("external-b", 0, [3, 4])], + request_ids=["internal-a", "internal-b"], + ) + + model.on_requests_finished(["internal-a"]) + + assert set(model._states) == {"internal-b"} + + +def test_reference_voice_and_duplex_metadata_follow_request_lifecycle(): + model, _ = _model() + first = _info("voice-a", 0, [1, 2]) + first["codes"]["ref"] = torch.linspace(-0.1, 0.1, 160) + first["meta"].update( + ref_audio_sr=16000, + native_duplex_segment_text="hello", + duplex_turn_id=7, + duplex_epoch=3, + ) + first["meta"].pop("prompt_cache_id") + + output = _forward(model, [first]) + prompt_cache_id, prompt_wav = model._owned_prompt_wavs["voice-a"] + assert prompt_cache_id.startswith("voice-a:") + assert Path(prompt_wav).is_file() + assert bytes(output.multimodal_outputs["meta.llm_output_text_utf8"][0].tolist()).decode() == "hello" + assert output.multimodal_outputs["meta.duplex_turn_id"][0].item() == 7 + assert output.multimodal_outputs["meta.duplex_epoch"][0].item() == 3 + + final = _info("voice-a", 1, [3, 4], last_chunk=True) + final["meta"].pop("prompt_cache_id") + output = _forward(model, [final]) + + assert output.multimodal_outputs["meta.tts_is_last_chunk"][0].item() is True + assert "voice-a" not in model._owned_prompt_wavs + assert not Path(prompt_wav).exists() + assert (prompt_cache_id, prompt_wav) not in model.backend._prompt_features diff --git a/tests/model_executor/models/minicpmo_4_5/test_hf_model_path.py b/tests/model_executor/models/minicpmo_4_5/test_hf_model_path.py deleted file mode 100644 index 834b7b0e0..000000000 --- a/tests/model_executor/models/minicpmo_4_5/test_hf_model_path.py +++ /dev/null @@ -1,74 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Regression test for HF model-ID resolution and path caching in MiniCPM-o 4.5 TTS.""" - -from types import SimpleNamespace - -import pytest - -from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( - MiniCPMO45OmniTTSForConditionalGeneration, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class _FakeMiniCPMTTS: - def __init__(self, config, audio_tokenizer): - self.config = config - self.audio_tokenizer = audio_tokenizer - self.emb_text = object() - self.projector_semantic = object() - - def generate(self): - raise NotImplementedError - - -def test_hf_model_id_path_caching(mocker, tmp_path): - """Verify the resolved HF path is cached and reused for default ref audio.""" - resolved_path = str(tmp_path / "resolved_hf_model") - assets_dir = tmp_path / "resolved_hf_model" / "assets" - assets_dir.mkdir(parents=True, exist_ok=True) - ref_audio_file = assets_dir / "HT_ref_audio.wav" - ref_audio_file.write_bytes(b"dummy wav content") - - mock_download = mocker.patch( - "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts.download_weights_from_hf_specific", - return_value=resolved_path, - ) - - mocker.patch( - "transformers.dynamic_module_utils.get_class_from_dynamic_module", - return_value=_FakeMiniCPMTTS, - ) - - # Build fake vllm_config with HF model ID - model_id = "openbmb/MiniCPM-o-4_5" - tts_cfg = SimpleNamespace( - audio_bos_token_id=151687, - text_eos_token_id=151692, - num_audio_tokens=6562, - hidden_size=768, - normalize_projected_hidden=True, - top_p=0.8, - top_k=100, - repetition_penalty=1.02, - attn_implementation="sdpa", - ) - hf_config = SimpleNamespace(tts_config=tts_cfg) - vllm_config = SimpleNamespace( - model_config=SimpleNamespace( - model=model_id, - hf_config=hf_config, - ) - ) - - tts_model = MiniCPMO45OmniTTSForConditionalGeneration(vllm_config=vllm_config) - - tts_model._lazy_init_tts() - tts_model._lazy_init_tts() - - mock_download.assert_called_once_with(model_id, None, ["*"]) - assert tts_model._model_path == resolved_path - assert tts_model._assets_loaded is True - assert tts_model._resolve_prompt_wav_path(None, None) == (str(ref_audio_file), None) diff --git a/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py b/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py index 7ec8e8114..9c4b50967 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py +++ b/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py @@ -10,7 +10,7 @@ Covers ``vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.llm2t - structured model_intermediate_buffer carries the thinker ids and text - scheduler prompt tokens follow the selected TTS region or output tokens - MiniCPM-o 4.5 TTS region detection on 151703 / 151704 tokens - - No TTS markers present -> no ``tts_token_ids`` / ``tts_hidden_states`` keys + - plain chat without TTS markers conditions on the generated assistant span - prompt arg is normalized to a list and ``multi_modal_data`` is gated by ``requires_multimodal_data`` """ @@ -33,6 +33,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] _HIDDEN_DIM = 4 +@pytest.mark.skip(reason="Legacy fused Talker metadata is replaced by separate Talker and Code2Wav stage payloads.") def test_native_duplex_talker_emits_empty_text_metadata_to_clear_previous_segment() -> None: class _Talker: _ar_last_chunk_flags = [True] @@ -91,13 +92,15 @@ def _make_thinker_output( latent: torch.Tensor | None = None, hidden_states: torch.Tensor | None = None, multimodal_output: dict[str, object] | None = None, + request_multimodal_output: dict[str, object] | None = None, ): """Construct a minimal mock of a thinker engine output entry. The real ``llm2tts`` only reads a tight slice of fields: - - top-level: ``request_id``, ``prompt_token_ids``, ``outputs[0]`` - - per-output: ``multimodal_output`` (dict), ``hidden_states`` (opt), - ``text``, ``token_ids`` + - top-level: ``request_id``, ``prompt_token_ids``, optional + ``multimodal_output``, and ``outputs[0]`` + - per-output: fallback ``multimodal_output``, ``hidden_states`` (opt), + ``text``, and ``token_ids`` """ mm_output = dict(multimodal_output or {}) if latent is not None: @@ -109,11 +112,14 @@ def _make_thinker_output( ) if hidden_states is not None: output.hidden_states = hidden_states - return SimpleNamespace( + request_output = SimpleNamespace( request_id=request_id, prompt_token_ids=prompt_token_ids, outputs=[output], ) + if request_multimodal_output is not None: + request_output.multimodal_output = request_multimodal_output + return request_output class TestInputValidation: @@ -131,6 +137,19 @@ class TestInputValidation: class TestBasicShape: + def test_request_level_multimodal_output_feeds_orchestrator_bridge(self) -> None: + hidden = torch.zeros((2, _HIDDEN_DIM)) + thinker_output = _make_thinker_output( + prompt_token_ids=[10], + output_token_ids=[20], + request_multimodal_output={"latent": hidden}, + ) + + out = llm2tts([thinker_output], prompt=None) + + assert out[0]["prompt_token_ids"] == [0, 0, 0] + assert torch.equal(torch.tensor(out[0]["model_intermediate_buffer"]["hidden_states"]["tts"]), hidden[1:]) + def test_returns_one_entry_per_input(self) -> None: hidden = torch.zeros((3, _HIDDEN_DIM)) out = llm2tts( @@ -144,13 +163,13 @@ class TestBasicShape: ) assert len(out) == 2 - def test_talker_scheduler_prompt_uses_output_tokens(self) -> None: + def test_talker_scheduler_prompt_matches_condition_length(self) -> None: hidden = torch.zeros((2, _HIDDEN_DIM)) out = llm2tts( [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden)], prompt=None, ) - assert out[0]["prompt_token_ids"] == [20] + assert out[0]["prompt_token_ids"] == [0, 0, 0] assert "stream_output" not in out[0]["model_intermediate_buffer"] def test_model_intermediate_buffer_carries_thinker_outputs(self) -> None: @@ -174,6 +193,24 @@ class TestBasicShape: assert buffer["ids"]["output"] == out_ids assert buffer["llm_output_text"] == ["hello"] + def test_strict_talker_receives_canonical_nested_handoff(self) -> None: + hidden = torch.arange(5 * _HIDDEN_DIM, dtype=torch.float32).reshape(5, _HIDDEN_DIM) + result = llm2tts( + [ + _make_thinker_output( + prompt_token_ids=[10], + output_token_ids=[151703, 20, 21, 151704], + hidden_states=hidden, + ) + ], + prompt=None, + )[0] + + buffer = result["model_intermediate_buffer"] + assert result["prompt_token_ids"] == [0, 0, 0, 0] + assert buffer["ids"]["tts"] == [20, 21] + assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), hidden[2:4]) + def test_latent_in_multimodal_output_takes_precedence(self) -> None: # When both ``multimodal_output["latent"]`` and ``hidden_states`` are # present, the latent payload must win (this is the steady-state path @@ -241,13 +278,10 @@ class TestTtsRegionDetection: assert buffer["ids"]["tts"] == [30, 31] assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), hidden[3:5]) - def test_no_tts_markers_omits_slice_keys(self) -> None: - # If neither marker pair is present, the bridge should NOT populate - # ``tts_token_ids`` / ``tts_hidden_states`` — the talker should fall - # through to the dummy path. - buffer, _ = self._run([10, 11], [20, 21, 22]) - assert "tts" not in buffer["ids"] - assert "tts" not in buffer.get("hidden_states", {}) + def test_plain_chat_without_tts_markers_uses_assistant_span(self) -> None: + buffer, hidden = self._run([10, 11], [20, 21, 22]) + assert buffer["ids"]["tts"] == [20, 21, 22] + assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), hidden[2:5]) class TestPromptAndMultiModal: diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index b9c7938f0..1b8bd669f 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -5,8 +5,8 @@ Covers: - pipeline declared in the central registry - lazy loader returns the expected ``PipelineConfig`` - - 2-stage topology (thinker LLM_AR + talker LLM_AR with audio output) - - stage 1 routes through ``llm2tts`` custom input processor + - split 3-stage topology with no fused compatibility registration + - stage 1 routes through ``llm2tts`` and the async Code2Wav producer - ``hf_architectures`` covers both the shared ``MiniCPMO`` alias and the explicit 4.5 arch - ``hf_config_predicate`` selects MiniCPM-o 4.5 only and rejects 2.6 @@ -15,6 +15,7 @@ Covers: from __future__ import annotations +from pathlib import Path from types import SimpleNamespace import pytest @@ -23,20 +24,29 @@ from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.stage_config import ( PipelineConfig, StageExecutionType, + load_deploy_config, + merge_pipeline_deploy, ) +from vllm_omni.model_executor.models.registry import _OMNI_MODELS pytestmark = [pytest.mark.core_model, pytest.mark.cpu] _PIPELINE_KEY = "minicpmo_4_5" +_DEPLOY_DIR = Path(__file__).resolve().parents[4] / "vllm_omni" / "deploy" +_DEPLOY_LAYOUTS = { + "minicpmo_4_5.yaml": ["0", "0", "0"], + "minicpmo_4_5_batching.yaml": ["0", "1", "1"], + "minicpmo_4_5_2gpu.yaml": ["0", "1", "1"], + "minicpmo_4_5_3gpu.yaml": ["0", "1", "2"], + "minicpmo_4_5_8x4090.yaml": ["0,1,2,3", "4", "5"], +} class TestRegistryDeclaration: def test_declared_in_omni_pipelines(self) -> None: assert _PIPELINE_KEY in OMNI_PIPELINES - - def test_visible_in_central_registry(self) -> None: - assert _PIPELINE_KEY in OMNI_PIPELINES + assert "minicpmo_4_5_fused" not in OMNI_PIPELINES def test_lazy_load_returns_pipeline_config(self) -> None: pipeline = OMNI_PIPELINES[_PIPELINE_KEY] @@ -64,9 +74,9 @@ class TestPipelineTopology: def pipeline(self) -> PipelineConfig: return OMNI_PIPELINES[_PIPELINE_KEY] - def test_two_stages(self, pipeline: PipelineConfig) -> None: - assert len(pipeline.stages) == 2 - assert [s.stage_id for s in pipeline.stages] == [0, 1] + def test_three_stages(self, pipeline: PipelineConfig) -> None: + assert len(pipeline.stages) == 3 + assert [s.stage_id for s in pipeline.stages] == [0, 1, 2] def test_topology_validates(self, pipeline: PipelineConfig) -> None: # ``validate`` returns a list of structural errors; empty == valid. @@ -90,11 +100,15 @@ class TestPipelineTopology: assert talker.execution_type == StageExecutionType.LLM_AR # talker consumes thinker output assert talker.input_sources == (0,) - assert talker.final_output is True - assert talker.final_output_type == "audio" - assert talker.engine_output_type == "audio" + assert talker.final_output is False + assert talker.final_output_type is None + assert talker.engine_output_type == "latent" # scope KV cache / mrope sizing to talker sub-config assert talker.hf_config_name == "tts_config" + assert talker.custom_process_next_stage_input_func is None + assert talker.async_chunk_process_next_stage_input_func == ( + "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_async_chunk" + ) def test_talker_routes_through_llm2tts(self, pipeline: PipelineConfig) -> None: talker = pipeline.get_stage(1) @@ -106,6 +120,73 @@ class TestPipelineTopology: "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.llm2tts" ) + def test_code2wav_stage(self, pipeline: PipelineConfig) -> None: + code2wav = pipeline.get_stage(2) + assert code2wav is not None + assert code2wav.model_stage == "code2wav" + assert code2wav.execution_type == StageExecutionType.LLM_GENERATION + assert code2wav.input_sources == (1,) + assert code2wav.final_output is True + assert code2wav.final_output_type == "audio" + assert code2wav.engine_output_type == "audio" + assert code2wav.model_arch == "MiniCPMO45Code2Wav" + assert code2wav.sync_process_input_func is None + + +class TestDeployTopology: + @pytest.mark.parametrize(("filename", "devices"), _DEPLOY_LAYOUTS.items()) + def test_deploy_resolves_three_stage_pipeline(self, filename: str, devices: list[str]) -> None: + deploy = load_deploy_config(_DEPLOY_DIR / filename) + pipeline = OMNI_PIPELINES[deploy.pipeline] + stages = merge_pipeline_deploy(pipeline, deploy) + + assert deploy.pipeline == _PIPELINE_KEY + assert deploy.async_chunk is True + assert [stage.stage_id for stage in stages] == [0, 1, 2] + assert [stage.yaml_runtime["devices"] for stage in stages] == devices + assert "additional_config" not in stages[1].yaml_engine_args + assert "skip_mm_profiling" not in stages[0].yaml_engine_args + assert stages[1].yaml_engine_args["skip_mm_profiling"] is True + assert stages[2].yaml_engine_args["skip_mm_profiling"] is True + assert stages[1].yaml_extras["output_connectors"]["to_stage_2"] == "connector_of_shared_memory" + assert stages[2].yaml_extras["input_connectors"]["from_stage_1"] == "connector_of_shared_memory" + connector = deploy.connectors["connector_of_shared_memory"] + assert connector["name"] == "SharedMemoryConnector" + assert connector["extra"]["codec_chunk_frames"] == 25 + assert connector["extra"]["codec_left_context_frames"] == 3 + assert connector["extra"]["connector_get_max_wait_first_chunk"] == 3000 + assert connector["extra"]["connector_get_max_wait"] == 300 + if filename == "minicpmo_4_5.yaml": + assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [4, 4, 4] + assert [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] == [ + 0.65, + 0.15, + 0.15, + ] + elif filename in {"minicpmo_4_5_batching.yaml", "minicpmo_4_5_2gpu.yaml"}: + assert [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] == [ + 0.9, + 0.55, + 0.35, + ] + + def test_pipeline_exposes_no_full_payload_or_token_placeholder_hooks(self) -> None: + pipeline = OMNI_PIPELINES[_PIPELINE_KEY] + talker = pipeline.get_stage(1) + code2wav = pipeline.get_stage(2) + assert talker is not None + assert code2wav is not None + assert talker.custom_process_next_stage_input_func is None + assert code2wav.sync_process_input_func is None + + +def test_code2wav_model_is_lazily_registered() -> None: + assert _OMNI_MODELS["MiniCPMO45Code2Wav"] == ( + "minicpmo_4_5", + "minicpmo_4_5_code2wav", + "MiniCPMO45Code2Wav", + ) + class TestArchAliases: """``hf_architectures`` must cover both the shared and explicit names.""" diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py new file mode 100644 index 000000000..f0e9169b3 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Request-alignment tests for MiniCPM-o 4.5's native Talker.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, +) +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, + _max_audio_tokens, +) +from vllm_omni.utils.mm_outputs import to_payload_element + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeNativeTalker(nn.Module): + has_preprocess = True + + def __init__(self) -> None: + super().__init__() + self.forward_kwargs = None + + def forward(self, **kwargs): + self.forward_kwargs = kwargs + return torch.ones(2, 4) + + +def test_wrapper_always_delegates_talker_to_native_ar_path() -> None: + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + nn.Module.__init__(model) + model.model_stage = "tts" + model.talker = _FakeNativeTalker() + + output = model( + input_ids=torch.tensor([1, 2]), + positions=torch.arange(2), + model_intermediate_buffer=[{"request_id": "req"}], + ) + + assert output.shape == (2, 4) + assert model.talker.forward_kwargs["model_intermediate_buffer"][0]["request_id"] == "req" + + +def _make_talker() -> MiniCPMO45OmniTTSForConditionalGeneration: + talker = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + nn.Module.__init__(talker) + talker._num_audio_tokens = 8 + talker._batch_stop_logits = None + talker._request_generators = {} + talker._request_audio_states = {} + talker._deferred_cleanup_ids = set() + return talker + + +def _routed(output, index: int): + return to_payload_element( + output.multimodal_outputs, + index, + index, + index + 1, + seq_len=2, + scheduled_seq_len=2, + ) + + +@pytest.mark.parametrize( + ("condition_tokens", "expected"), + [(3, 64), (100, 1000), (1000, 2048)], +) +def test_audio_token_limit_scales_with_condition_length( + condition_tokens: int, + expected: int, +) -> None: + assert _max_audio_tokens(condition_tokens) == expected + + +def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> None: + talker = _make_talker() + seen: list[tuple[str, list[float], list[int]]] = [] + + def sample(hidden, history, request_id, step): + assert step == 0 + seen.append((request_id, hidden.reshape(-1).tolist(), history.tolist())) + return torch.tensor(2 if request_id == "req-a" else 3) + + mocker.patch.object(talker, "_sample_audio_code", side_effect=sample) + infos = [ + {"request_id": "req-a", "audio_codes": {"accumulated": torch.tensor([1])}}, + {"request_id": "req-b", "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}}, + ] + + output = talker.make_omni_output( + torch.tensor([[1.0, 0.0], [2.0, 0.0], [3.0, 0.0]]), + model_intermediate_buffer=infos, + request_token_spans=[(0, 2), (2, 3)], + ) + + assert seen == [ + ("req-a", [2.0, 0.0], [1]), + ("req-b", [3.0, 0.0], []), + ] + assert infos[0]["audio_codes"]["accumulated"].tolist() == [1, 2] + assert infos[1]["audio_codes"]["accumulated"].tolist() == [3] + assert set(output.multimodal_outputs) == {"codes", "meta"} + assert "model_outputs" not in output.multimodal_outputs + assert "sr" not in output.multimodal_outputs + assert _routed(output, 0)["codes"]["audio"].tolist() == [[2]] + assert _routed(output, 1)["codes"]["audio"].tolist() == [[3]] + assert _routed(output, 0)["meta"]["finished"].item() is False + assert talker.compute_logits(output.text_hidden_states).argmax(dim=-1).tolist() == [0, 0] + + +def test_incomplete_prefill_emits_no_code_and_does_not_advance_state(mocker) -> None: + talker = _make_talker() + sample = mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2)) + infos = [ + { + "request_id": "req-prefill", + "audio_state": {"step": 0}, + "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}, + }, + { + "request_id": "req-decode", + "audio_state": {"step": 4}, + "audio_codes": {"accumulated": torch.tensor([1])}, + }, + ] + + output = talker.make_omni_output( + torch.tensor([[1.0, 0.0], [2.0, 0.0], [3.0, 0.0]]), + model_intermediate_buffer=infos, + request_token_spans=[(0, 2), (2, 3)], + request_sample_eligible=[False, True], + ) + + sample.assert_called_once() + assert sample.call_args.args[2] == "req-decode" + assert infos[0]["audio_state"]["step"] == 0 + assert infos[0]["audio_codes"]["accumulated"].numel() == 0 + assert infos[1]["audio_state"]["step"] == 5 + assert _routed(output, 0)["codes"]["audio"].shape == (0, 1) + assert _routed(output, 1)["codes"]["audio"].tolist() == [[2]] + + +def test_eos_is_terminal_once_and_never_enters_codec_history(mocker) -> None: + talker = _make_talker() + sample = mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(7)) + info = { + "request_id": "req-stop", + "audio_state": {"step": 3}, + "audio_codes": {"accumulated": torch.tensor([4, 5])}, + } + + first = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + first_logits = talker.compute_logits(first.text_hidden_states) + second = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + + sample.assert_called_once() + assert info["audio_codes"]["accumulated"].tolist() == [4, 5] + assert first.multimodal_outputs["codes"]["audio"][0].shape == (0, 1) + assert first.multimodal_outputs["meta"]["finished"][0].item() is True + assert second.multimodal_outputs["meta"]["finished"][0].item() is False + assert first_logits.argmax(dim=-1).tolist() == [1] + assert talker.compute_logits(second.text_hidden_states).argmax(dim=-1).tolist() == [1] + + +def test_max_token_terminal_includes_only_new_codec_delta(mocker) -> None: + talker = _make_talker() + mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(3)) + info = { + "request_id": "req-limit", + "audio_state": {"step": 1, "max_tokens": 2}, + "audio_codes": {"accumulated": torch.tensor([4, 5])}, + } + + output = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + + assert info["audio_codes"]["accumulated"].tolist() == [4, 5, 3] + assert output.multimodal_outputs["codes"]["audio"][0].tolist() == [[3]] + assert output.multimodal_outputs["meta"]["finished"][0].item() is True + assert talker.compute_logits(output.text_hidden_states).argmax(dim=-1).tolist() == [1] + + +def test_request_local_state_survives_missing_runner_buffer_update(mocker) -> None: + talker = _make_talker() + mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(3)) + first_info = { + "request_id": "req-local-state", + "audio_state": {"step": 1, "max_tokens": 3}, + "audio_codes": {"accumulated": torch.tensor([4])}, + } + + talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[first_info], + request_token_spans=[(0, 1)], + ) + second = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[{"request_id": "req-local-state"}], + request_token_spans=[(0, 1)], + ) + + assert second.multimodal_outputs["meta"]["finished"][0].item() is True + assert talker._request_audio_states["req-local-state"]["step"] == 3 + + +def test_missing_conditioning_fails_clearly() -> None: + talker = _make_talker() + + with pytest.raises(ValueError, match="tts_token_ids and tts_hidden_states"): + talker.preprocess( + torch.tensor([0]), + None, + _omni_is_prefill=True, + request_id="req-invalid", + ) + + +def test_request_cleanup_evicts_ar_rng_and_decode_state() -> None: + talker = _make_talker() + talker._request_generators["req-done"] = torch.Generator() + talker._request_audio_states["req-done"] = {"step": 1} + + talker.on_requests_finished(["req-done"]) + talker._flush_deferred_cleanup() + + assert "req-done" not in talker._request_generators + assert "req-done" not in talker._request_audio_states diff --git a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py new file mode 100644 index 000000000..4f17b2c64 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections import defaultdict +from types import SimpleNamespace + +import pytest +import torch +from vllm.v1.request import RequestStatus + +from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import ( + tts2code2wav_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _manager(): + return SimpleNamespace( + connector=SimpleNamespace(config={"extra": {"codec_chunk_frames": 25, "codec_left_context_frames": 3}}), + code_prompt_token_ids=defaultdict(list), + request_payload={}, + put_req_chunk=defaultdict(int), + ) + + +def _request(external_id: str, internal_id: str | None = None): + request = SimpleNamespace( + external_req_id=external_id, + request_id=internal_id or external_id, + status=RequestStatus.RUNNING, + ) + request.is_finished = lambda: RequestStatus.is_finished(request.status) + return request + + +def _delta(*codes: int): + return { + "codes": {"audio": torch.tensor(codes, dtype=torch.long).reshape(-1, 1)}, + "meta": {"finished": torch.tensor(False)}, + } + + +def _codes(payload) -> list[int]: + assert payload.codes is not None + assert isinstance(payload.codes.audio, torch.Tensor) + assert payload.codes.audio.dtype == torch.long + assert payload.codes.audio.ndim == 1 + return payload.codes.audio.tolist() + + +@pytest.mark.parametrize(("count", "emitted"), [(24, False), (25, True), (26, True)]) +def test_first_chunk_threshold_is_25_generated_codes(count: int, emitted: bool) -> None: + manager = _manager() + payload = tts2code2wav_async_chunk( + transfer_manager=manager, + multimodal_output=_delta(*range(count)), + request=_request("req"), + is_finished=False, + ) + + assert (payload is not None) is emitted + if payload is not None: + assert _codes(payload) == [4218, 4218, 4218, *range(25)] + assert payload.meta.chunk_seq == 0 + assert payload.meta.code_flat_numel == 28 + + +def test_steady_chunk_has_three_code_overlap_and_25_new_codes() -> None: + manager = _manager() + request = _request("req") + + first = tts2code2wav_async_chunk(manager, _delta(*range(25)), request, False) + manager.put_req_chunk["req"] += 1 + steady = tts2code2wav_async_chunk(manager, _delta(*range(25, 50)), request, False) + + assert first is not None + assert steady is not None + assert _codes(steady) == [22, 23, 24, *range(25, 50)] + assert steady.meta.chunk_seq == 1 + + +def test_exact_boundary_final_flushes_held_lookahead() -> None: + manager = _manager() + request = _request("req") + + assert tts2code2wav_async_chunk(manager, _delta(*range(25)), request, False) is not None + manager.put_req_chunk["req"] += 1 + final = tts2code2wav_async_chunk(manager, None, request, True) + + assert final is not None + assert _codes(final) == [22, 23, 24] + assert final.meta.chunk_seq == 1 + assert final.meta.code_flat_numel == 3 + assert final.meta.last_chunk is True + assert final.meta.finished.item() is True + + +def test_short_final_flushes_silence_prefix_and_tail() -> None: + manager = _manager() + final = tts2code2wav_async_chunk(manager, _delta(*range(7)), _request("req"), True) + + assert final is not None + assert _codes(final) == [4218, 4218, 4218, *range(7)] + assert final.meta.last_chunk is True + assert final.meta.finished.item() is True + + +def test_first_chunk_forwards_reference_voice_and_duplex_identity() -> None: + manager = _manager() + request = _request("req") + request.additional_information = { + "codes": {"ref": [0.1, -0.1]}, + "meta": { + "ref_audio_sr": 16000, + "native_duplex_segment_text": "hello", + "segment_end": True, + "turn_end": True, + }, + "duplex": {"epoch": 3, "model_turn_id": 7}, + } + + payload = tts2code2wav_async_chunk(manager, _delta(*range(7)), request, True) + + assert payload is not None + assert payload.codes.ref.tolist() == pytest.approx([0.1, -0.1]) + assert payload.meta.ref_audio_sr == 16000 + assert payload.meta.native_duplex_segment_text == "hello" + assert payload.meta.duplex_epoch == 3 + assert payload.meta.duplex_turn_id == 7 + assert payload.meta.segment_end is True + assert payload.meta.turn_end is True + + +def test_empty_final_releases_wait_gate_once() -> None: + manager = _manager() + request = _request("req") + + final = tts2code2wav_async_chunk(manager, None, request, True) + duplicate = tts2code2wav_async_chunk(manager, None, request, True) + + assert final is not None + assert _codes(final) == [] + assert final.meta.chunk_seq == 0 + assert final.meta.request_id == "req" + assert final.meta.cache_epoch == 0 + assert final.meta.last_chunk is True + assert duplicate is None + + +def test_staggered_requests_keep_accumulators_isolated() -> None: + manager = _manager() + req_a = _request("a") + req_b = _request("b") + + assert tts2code2wav_async_chunk(manager, _delta(*range(24)), req_a, False) is None + out_b = tts2code2wav_async_chunk(manager, _delta(*range(100, 125)), req_b, False) + out_a = tts2code2wav_async_chunk(manager, _delta(24), req_a, False) + + assert out_b is not None + assert out_a is not None + assert _codes(out_b) == [4218, 4218, 4218, *range(100, 125)] + assert _codes(out_a) == [4218, 4218, 4218, *range(25)] + assert out_a.meta.request_id == "a" + assert out_b.meta.request_id == "b" + + +def test_cancel_drops_epoch_state_and_stale_request_cannot_publish() -> None: + manager = _manager() + stale = _request("req", "internal-0") + + assert tts2code2wav_async_chunk(manager, _delta(*range(10)), stale, False) is None + stale.status = RequestStatus.FINISHED_ABORTED + assert tts2code2wav_async_chunk(manager, None, stale, True) is None + assert tts2code2wav_async_chunk(manager, _delta(*range(25)), stale, False) is None + + replacement = _request("req", "internal-1") + payload = tts2code2wav_async_chunk(manager, _delta(*range(25)), replacement, False) + + assert payload is not None + assert payload.meta.cache_epoch == 1 + assert _codes(payload) == [4218, 4218, 4218, *range(25)] diff --git a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py index 453b3a595..c8381d881 100644 --- a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py +++ b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py @@ -1,75 +1,39 @@ -import base64 -import os -import struct from types import SimpleNamespace import pytest import torch -from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( - MiniCPMO45OmniTTSForConditionalGeneration, - _drain_native_duplex_emitted_text, - _queue_native_duplex_segment_text, - _TalkerTurnState, -) from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import ( _extract_first_audio_ref, - _native_duplex_segment_output_ids, llm2tts, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def test_native_duplex_transcript_waits_for_buffered_vocoder_audio(): - state = _TalkerTurnState(None, None, turn_id=1) - - _queue_native_duplex_segment_text(state, "好") - assert _drain_native_duplex_emitted_text(state, has_audio=False) == "" - - _queue_native_duplex_segment_text(state, "的，刚") - assert _drain_native_duplex_emitted_text(state, has_audio=True) == "好的，刚" - assert _drain_native_duplex_emitted_text(state, has_audio=True) == "" +def _output( + *, + prompt_ids: list[int], + output_ids: list[int], + latent: torch.Tensor, + multimodal_output: dict | None = None, + token_list: list[int] | None = None, +): + mm_output = dict(multimodal_output or {}) + mm_output["latent"] = latent + completion = SimpleNamespace( + token_ids=output_ids if token_list is None else token_list, + text="hello", + multimodal_output=mm_output, + ) + return SimpleNamespace( + request_id="req-1", + prompt_token_ids=prompt_ids, + outputs=[completion], + ) -def test_native_duplex_token2wav_cache_isolated_across_interleaved_sessions(): - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - model.audio_tokenizer = SimpleNamespace(stream_cache=None, hift_cache_dict={}) - observed = [] - - def stream_window(token_list, _prompt, *, last_chunk): - del token_list, last_chunk - marker = model.audio_tokenizer.stream_cache["marker"] - observed.append(marker) - prefix = marker[0] - index = int(marker[1:]) + 1 - model.audio_tokenizer.stream_cache = {"marker": f"{prefix}{index}"} - model.audio_tokenizer.hift_cache_dict = {"marker": f"{prefix}{index}"} - return torch.ones(1) - - model._t2w_stream_window = stream_window - a = _TalkerTurnState(None, None) - b = _TalkerTurnState(None, None) - a.stream_cache = {"marker": "A0"} - a.hift_cache_dict = {"marker": "A0"} - a.vocoder_initialized = True - b.stream_cache = {"marker": "B0"} - b.hift_cache_dict = {"marker": "B0"} - b.vocoder_initialized = True - - model._run_vocoder_window(a, [1], last_chunk=False) - model._run_vocoder_window(b, [1], last_chunk=False) - model._run_vocoder_window(a, [1], last_chunk=False) - model._run_vocoder_window(b, [1], last_chunk=True) - - assert observed == ["A0", "B0", "A1", "B1"] - assert a.stream_cache == {"marker": "A2"} - assert b.stream_cache == {"marker": "B2"} - assert model.audio_tokenizer.stream_cache is None - assert model.audio_tokenizer.hift_cache_dict == {} - - -def test_extract_first_audio_ref_accepts_dict_stereo_audio(): +def test_extract_first_audio_ref_accepts_dict_stereo_audio() -> None: ref = _extract_first_audio_ref( { "audio": { @@ -85,1263 +49,116 @@ def test_extract_first_audio_ref_accepts_dict_stereo_audio(): assert torch.allclose(waveform, torch.tensor([1.5, 3.5, 5.5])) -def test_llm2tts_carries_request_ref_audio_to_talker_payload(): - latent = torch.arange(20, dtype=torch.float16).reshape(5, 4) - output = SimpleNamespace( - token_ids=[11, 12, 9002], - text="hello", +def test_plain_chat_handoff_owns_talker_prompt_contract() -> None: + prompt_ids = [101, 102] + output_ids = [11, 12] + latent = torch.arange(16, dtype=torch.float32).reshape(4, 4) + + converted = llm2tts( + [_output(prompt_ids=prompt_ids, output_ids=output_ids, latent=latent)], + prompt=[{}], + )[0] + + info = converted["model_intermediate_buffer"] + assert info["ids"]["tts"] == output_ids + assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[2:4]) + assert converted["prompt_token_ids"] == [0, 0, 0, 0] + assert info["meta"]["replace_streaming_prompt"] is True + assert info["meta"]["next_stage_prompt_len"] == 4 + + +def test_llm2tts_carries_request_ref_audio() -> None: + latent = torch.arange(20, dtype=torch.float32).reshape(5, 4) + source = _output( + prompt_ids=[101, 9001], + output_ids=[11, 12, 9002], + latent=latent, multimodal_output={ - "latent": latent, "meta": { "tts_bos_token_id": 9001, "tts_eos_token_id": 9002, - }, + } }, ) - llm_output = SimpleNamespace( - request_id="req-1", - prompt_token_ids=[101, 9001], - outputs=[output], - ) ref_waveform = torch.tensor([0.1, 0.2, 0.3]) converted = llm2tts( - [llm_output], - prompt=[ - { - "multi_modal_data": { - "audio": (ref_waveform, 22050), - } - } - ], - ) + [source], + prompt=[{"multi_modal_data": {"audio": (ref_waveform, 22050)}}], + )[0] - assert "additional_information" not in converted[0] - info = converted[0]["model_intermediate_buffer"] + info = converted["model_intermediate_buffer"] assert info["codes"]["ref"] == ref_waveform.tolist() assert info["meta"]["ref_audio_sr"] == 22050 assert info["ids"]["tts"] == [11, 12] - assert torch.equal(torch.as_tensor(info["hidden_states"]["tts"]), latent[2:4].float()) - assert "tts_token_ids" not in info - assert "tts_hidden_states" not in info -def test_llm2tts_uses_duplex_prompt_token_ids_for_tts_boundary(): - latent = torch.arange(24, dtype=torch.float16).reshape(6, 4) - output = SimpleNamespace( - token_ids=[11, 12, 9002], - text="hello", +def test_native_duplex_speak_segment_reaches_split_talker() -> None: + prompt_ids = [101, 102] + output_ids = [9304, 21, 22, 9308] + latent = torch.arange(24, dtype=torch.float32).reshape(6, 4) + source = _output( + prompt_ids=prompt_ids, + output_ids=output_ids, + latent=latent, multimodal_output={ - "latent": latent, - "duplex_prompt_token_ids": [101, 102, 9001], - "meta": { - "tts_bos_token_id": 9001, - "tts_eos_token_id": 9002, - }, - }, - ) - llm_output = SimpleNamespace( - request_id="duplex-req", - prompt_token_ids=[0, 0, 0], - outputs=[output], - ) - - converted = llm2tts([llm_output], prompt=[{}]) - - assert "additional_information" not in converted[0] - info = converted[0]["model_intermediate_buffer"] - assert info["prompt_token_ids"] == [101, 102, 9001] - assert converted[0]["prompt_token_ids"] == [11, 12] - assert info["ids"]["tts"] == [11, 12] - assert torch.equal(torch.as_tensor(info["hidden_states"]["tts"]), latent[3:5].float()) - assert "tts_token_ids" not in info - assert "tts_hidden_states" not in info - - -def test_llm2tts_accepts_flat_tokenizer_derived_special_token_metadata(): - latent = torch.arange(24, dtype=torch.float16).reshape(6, 4) - output = SimpleNamespace( - token_ids=[11, 12, 9302], - text="hello", - multimodal_output={ - "latent": latent, - "duplex_prompt_token_ids": [101, 102, 9301], - "meta.tts_bos_token_id": torch.tensor([9301]), - "meta.tts_eos_token_id": torch.tensor([9302]), - "meta.listen_token_id": torch.tensor([9303]), - "meta.speak_token_id": torch.tensor([9304]), - "meta.chunk_eos_token_id": torch.tensor([9305]), - "meta.chunk_tts_eos_token_id": torch.tensor([9306]), - "meta.turn_eos_token_id": torch.tensor([9307]), - }, - ) - llm_output = SimpleNamespace( - request_id="duplex-flat-meta", - prompt_token_ids=[0, 0, 0], - outputs=[output], - ) - - converted = llm2tts([llm_output], prompt=[{}]) - - assert "additional_information" not in converted[0] - info = converted[0]["model_intermediate_buffer"] - assert converted[0]["prompt_token_ids"] == [11, 12] - assert info["ids"]["tts"] == [11, 12] - assert info["stream_output"] is True - assert info["native_duplex"] is True - assert torch.equal(torch.as_tensor(info["hidden_states"]["tts"]), latent[3:5].float()) - assert "tts_token_ids" not in info - assert "tts_hidden_states" not in info - - -def test_llm2tts_stops_at_tokenizer_derived_duplex_state_tokens(): - latent = torch.arange(32, dtype=torch.float16).reshape(8, 4) - output = SimpleNamespace( - token_ids=[11, 12, 9303, 13, 9302], - text="hello", - multimodal_output={ - "latent": latent, - "duplex_prompt_token_ids": [101, 102, 9301], + "duplex_prompt_token_ids": prompt_ids, "meta": { "tts_bos_token_id": 9301, "tts_eos_token_id": 9302, "listen_token_id": 9303, "speak_token_id": 9304, - "tts_pad_token_id": 9305, - "unit_token_id": 9306, - "unit_end_token_id": 9307, "chunk_eos_token_id": 9308, "chunk_tts_eos_token_id": 9309, "turn_eos_token_id": 9310, }, }, ) - llm_output = SimpleNamespace( - request_id="duplex-state-boundary", - prompt_token_ids=[0, 0, 0], - outputs=[output], - ) - - converted = llm2tts([llm_output], prompt=[{}]) - - info = converted[0]["model_intermediate_buffer"] - assert converted[0]["prompt_token_ids"] == [11, 12] - assert info["ids"]["tts"] == [11, 12] - assert torch.equal(torch.as_tensor(info["hidden_states"]["tts"]), latent[3:5].float()) - assert "tts_token_ids" not in info - assert "tts_hidden_states" not in info - - -def test_llm2tts_native_duplex_uses_speak_region_without_tts_bos_prompt(): - latent = torch.arange(36, dtype=torch.float16).reshape(9, 4) - output = SimpleNamespace( - token_ids=[9304, 11, 12, 9308, 13], - text="hello", - multimodal_output={ - "latent": latent, - "duplex_prompt_token_ids": [101, 102, 9306, 9306], - "meta": { - "tts_bos_token_id": 9301, - "tts_eos_token_id": 9302, - "listen_token_id": 9303, - "speak_token_id": 9304, - "tts_pad_token_id": 9305, - "unit_token_id": 9306, - "unit_end_token_id": 9307, - "chunk_eos_token_id": 9308, - "chunk_tts_eos_token_id": 9309, - "turn_eos_token_id": 9310, - }, - }, - ) - llm_output = SimpleNamespace( - request_id="duplex-speak-region", - prompt_token_ids=[0, 0, 0, 0], - outputs=[output], - ) - - converted = llm2tts([llm_output], prompt=[{}]) - - info = converted[0]["model_intermediate_buffer"] - assert converted[0]["prompt_token_ids"] == [11, 12] - assert info["ids"]["tts"] == [11, 12] - assert torch.equal(torch.as_tensor(info["hidden_states"]["tts"]), latent[5:7].float()) - assert "tts_token_ids" not in info - assert "tts_hidden_states" not in info - - -def test_llm2tts_native_duplex_uses_plain_text_region_without_speak_marker(): - latent = torch.arange(32, dtype=torch.float16).reshape(8, 4) - output = SimpleNamespace( - token_ids=[11, 12, 13, 9308], - text="hello", - multimodal_output={ - "latent": latent, - "duplex_prompt_token_ids": [101, 102, 9306, 9306], - "meta": { - "tts_bos_token_id": 9301, - "tts_eos_token_id": 9302, - "listen_token_id": 9303, - "speak_token_id": 9304, - "tts_pad_token_id": 9305, - "unit_token_id": 9306, - "unit_end_token_id": 9307, - "chunk_eos_token_id": 9308, - "chunk_tts_eos_token_id": 9309, - "turn_eos_token_id": 9310, - }, - }, - ) - llm_output = SimpleNamespace( - request_id="duplex-plain-text-region", - prompt_token_ids=[0, 0, 0, 0], - outputs=[output], - ) - - converted = llm2tts([llm_output], prompt=[{}]) - - info = converted[0]["model_intermediate_buffer"] - assert converted[0]["prompt_token_ids"] == [12, 13] - assert info["ids"]["tts"] == [12, 13] - assert torch.equal(torch.as_tensor(info["hidden_states"]["tts"]), latent[5:7].float()) - - -def test_llm2tts_native_duplex_conditions_on_turn_eos_and_midunit_speak(): - latent = torch.arange(36, dtype=torch.float16).reshape(9, 4) - # Official duplex includes mid-unit <|speak|> tokens AND the <|turn_eos|> - # token+hidden in the talker condition (its embedding is the trained - # stop signal); only chunk terminators bound the slice. - output = SimpleNamespace( - token_ids=[9304, 11, 9304, 12, 9310, 9308], - text="hello", - multimodal_output={ - "latent": latent, - "duplex_prompt_token_ids": [101, 102, 9306], - "meta": { - "tts_bos_token_id": 9301, - "tts_eos_token_id": 9302, - "listen_token_id": 9303, - "speak_token_id": 9304, - "tts_pad_token_id": 9305, - "unit_token_id": 9306, - "unit_end_token_id": 9307, - "chunk_eos_token_id": 9308, - "chunk_tts_eos_token_id": 9309, - "turn_eos_token_id": 9310, - }, - }, - ) - llm_output = SimpleNamespace( - request_id="duplex-turn-eos", - prompt_token_ids=[0, 0, 0], - outputs=[output], - ) - - converted = llm2tts([llm_output], prompt=[{}]) - - info = converted[0]["model_intermediate_buffer"] - assert info["ids"]["tts"] == [11, 9304, 12, 9310] - assert torch.equal(torch.as_tensor(info["hidden_states"]["tts"]), latent[4:8].float()) - assert info["meta"]["turn_eos_token_id"] == 9310 - - -def test_llm2tts_native_duplex_ignores_stale_tts_bos_inside_folded_prompt(): - latent = torch.arange(36, dtype=torch.float16).reshape(9, 4) - # A continuation unit: the resumable prompt has an EARLIER reply's - # <|tts_bos|> folded mid-prompt, and the model decided to LISTEN (no new - # speak text). The stale boundary must not re-slice already-spoken text. - output = SimpleNamespace( - token_ids=[9303, 9309], - text="", - multimodal_output={ - "latent": latent, - "duplex_prompt_token_ids": [101, 9301, 41, 42, 9306, 9306], - "meta": { - "tts_bos_token_id": 9301, - "tts_eos_token_id": 9302, - "listen_token_id": 9303, - "speak_token_id": 9304, - "tts_pad_token_id": 9305, - "unit_token_id": 9306, - "unit_end_token_id": 9307, - "chunk_eos_token_id": 9308, - "chunk_tts_eos_token_id": 9309, - "turn_eos_token_id": 9310, - }, - }, - ) - llm_output = SimpleNamespace( - request_id="duplex-stale-bos", - prompt_token_ids=[0, 0, 0, 0, 0, 0], - outputs=[output], - ) - - converted = llm2tts([llm_output], prompt=[{}]) - - assert converted == [] - - -def test_llm2tts_rejects_native_duplex_output_without_special_token_metadata(): - latent = torch.arange(16, dtype=torch.float16).reshape(4, 4) - output = SimpleNamespace( - token_ids=[11, 12], - text="hello", - multimodal_output={ - "latent": latent, - "duplex_prompt_token_ids": [101, 102, 9301], - }, - ) - llm_output = SimpleNamespace( - request_id="duplex-missing-meta", - prompt_token_ids=[0, 0, 0], - outputs=[output], - ) - - with pytest.raises(ValueError, match=r"<\|tts_bos\|>"): - llm2tts([llm_output], prompt=[{}]) - - -def test_minicpmo_talker_normalizes_list_handoff_payload(monkeypatch): - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - captured = {} - - monkeypatch.setattr(model, "_should_stream_output", lambda *args, **kwargs: False) - - def fake_generate_speech(tts_token_ids, tts_hidden_states, **kwargs): - captured["tts_token_ids"] = tts_token_ids - captured["tts_hidden_states"] = tts_hidden_states - return torch.tensor([0.0, 0.1], dtype=torch.float32) - - monkeypatch.setattr(model, "generate_speech", fake_generate_speech) - - waveform, mel = model.forward( - additional_information={ - "tts_token_ids": [11, 12], - "tts_hidden_states": [[0.1, 0.2], [0.3, 0.4]], - } - ) - - assert mel is None - assert torch.equal(waveform, torch.tensor([0.0, 0.1], dtype=torch.float32)) - assert torch.equal(captured["tts_token_ids"], torch.tensor([11, 12])) - assert captured["tts_hidden_states"].dtype == torch.float32 - assert captured["tts_hidden_states"].shape == (2, 2) - - -def test_minicpmo_talker_accepts_structured_omni_payload_handoff(monkeypatch): - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - captured = {} - - monkeypatch.setattr(model, "_should_stream_output", lambda *args, **kwargs: False) - - def fake_generate_speech(tts_token_ids, tts_hidden_states, **kwargs): - captured["tts_token_ids"] = tts_token_ids - captured["tts_hidden_states"] = tts_hidden_states - return torch.tensor([0.0, 0.1], dtype=torch.float32) - - monkeypatch.setattr(model, "generate_speech", fake_generate_speech) - - waveform, mel = model.forward( - additional_information={ - "ids": {"tts": [11, 12]}, - "hidden_states": {"tts": torch.tensor([[0.1, 0.2], [0.3, 0.4]])}, - } - ) - - assert mel is None - assert torch.equal(waveform, torch.tensor([0.0, 0.1], dtype=torch.float32)) - assert torch.equal(captured["tts_token_ids"], torch.tensor([11, 12])) - assert captured["tts_hidden_states"].dtype == torch.float32 - assert captured["tts_hidden_states"].shape == (2, 2) - - -def test_token2wav_prompt_cache_resets_when_reference_changes(): - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - model.audio_tokenizer = SimpleNamespace(cache="old") - - model._reset_token2wav_cache_if_needed("/tmp/ref_a.wav") - assert model.audio_tokenizer.cache is None - - model.audio_tokenizer.cache = "prepared" - model._reset_token2wav_cache_if_needed("/tmp/ref_a.wav") - assert model.audio_tokenizer.cache == "prepared" - - model._reset_token2wav_cache_if_needed("/tmp/ref_b.wav") - assert model.audio_tokenizer.cache is None - - -def test_request_ref_audio_prompt_wav_uses_content_cache(): - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - ref_audio = torch.tensor([0.0, 0.1, 0.2], dtype=torch.float32) - - path1 = model._write_ref_audio_prompt_wav(ref_audio, 16000) - path2 = model._write_ref_audio_prompt_wav(ref_audio.clone(), 16000) - - try: - assert path1 == path2 - assert os.path.exists(path1) - assert model._is_cached_ref_audio_prompt_wav(path1) - finally: - for path in set(getattr(model, "_ref_audio_prompt_files", {}).values()): - try: - os.unlink(path) - except OSError: - pass - - -def test_native_duplex_talker_uses_generate_chunk_continuation(monkeypatch): - class _FakeTTS: - audio_bos_token_id = 88 - - def __init__(self): - self.config = SimpleNamespace(num_audio_tokens=128) - self.model = SimpleNamespace(config=SimpleNamespace()) - self.emb_text = torch.nn.Embedding(256, 4) - self.calls = [] - self._responses = [ - (torch.tensor([[[10], [11], [12]]]), "kv1"), - (torch.tensor([[[20]]]), "kv2"), - ] - - def generate_chunk(self, **kwargs): - self.calls.append(kwargs) - return self._responses[len(self.calls) - 1] - - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - tts = _FakeTTS() - windows = [] - model.tts_obj = tts - model.audio_tokenizer = SimpleNamespace() - model._talker_turn_states = {} - model._talker_consumed_tokens = {} - - monkeypatch.setattr(model, "_lazy_init_tts", lambda: None) - monkeypatch.setattr( - model, "_build_tts_sampling_params", lambda: SimpleNamespace(temperature=0.8, repetition_penalty=1.05) - ) - monkeypatch.setattr(model, "_tts_runtime_config", lambda: SimpleNamespace(streaming_generator_chunk=2)) - monkeypatch.setattr(model, "_resolve_prompt_wav_path", lambda ref, sr: ("/tmp/ref.wav", None)) - monkeypatch.setattr(model, "_begin_turn_vocoder_cache", lambda prompt, **kwargs: None) - monkeypatch.setattr(model, "_t2w_pre_lookahead", lambda: 1) - monkeypatch.setattr( - model, - "_build_tts_condition_embeds", - lambda ids, hidden: torch.ones((ids.numel(), 4), dtype=torch.float32), - ) - - def fake_stream_window(token_list, prompt_wav_path, *, last_chunk): - windows.append((list(token_list), last_chunk)) - return torch.tensor([float(len(token_list))]) - - monkeypatch.setattr(model, "_t2w_stream_window", fake_stream_window) - - base_info = { - "request_id": "duplex-generate-chunk", - "meta": {"turn_eos_token_id": 9310}, - "ids": {"tts": [21]}, - "hidden_states": {"tts": [[0.1, 0.2, 0.3, 0.4]]}, - "codes": {"ref": None}, - "native_duplex": True, - } - list(model._create_native_duplex_stream_gen(base_info)) - - turn_end_info = { - **base_info, - "ids": {"tts": [21, 9310]}, - "hidden_states": {"tts": [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]]}, - } - list(model._create_native_duplex_stream_gen(turn_end_info)) - - assert len(tts.calls) == 2 - first, second = tts.calls - assert first["past_key_values"] is None - assert first["text_start_pos"] == 0 - assert first["min_new_tokens"] == 0 - assert first["max_new_token"] == 3 - assert torch.equal(first["eos_token"], torch.tensor([127])) - assert float(first["temperature"].item()) == pytest.approx(0.8) - assert first["repetition_penalty"] == pytest.approx(1.05) - - assert second["past_key_values"] == "kv1" - assert second["text_start_pos"] == 5 - assert second["min_new_tokens"] == 0 - assert second["max_new_token"] == 3 - assert "duplex-generate-chunk" not in model._talker_turn_states - assert windows[-1] == ([20], True) - - -def test_minicpmo_native_duplex_talker_turn_end_metadata_flushes_tail(monkeypatch): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( - MiniCPMO45OmniTTSForConditionalGeneration, - ) - - class _FakeEmbText: - weight = torch.ones((200, 4), dtype=torch.float32) - - def __call__(self, ids): - return torch.ones((ids.numel(), 4), dtype=torch.float32) - - class _FakeTTS: - audio_bos_token_id = 1 - config = SimpleNamespace(num_audio_tokens=128) - emb_text = _FakeEmbText() - model = SimpleNamespace(config=SimpleNamespace(rope_theta=10000.0)) - - def generate_chunk(self, **kwargs): - del kwargs - return torch.tensor([[[33]]]), None - - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - model.tts_obj = _FakeTTS() - model.audio_tokenizer = SimpleNamespace() - model._talker_turn_states = {} - model._talker_consumed_tokens = {} - - monkeypatch.setattr(model, "_lazy_init_tts", lambda: None) - monkeypatch.setattr( - model, "_build_tts_sampling_params", lambda: SimpleNamespace(temperature=0.8, repetition_penalty=1.05) - ) - monkeypatch.setattr(model, "_tts_runtime_config", lambda: SimpleNamespace(streaming_generator_chunk=25)) - monkeypatch.setattr(model, "_resolve_prompt_wav_path", lambda ref, sr: ("/tmp/ref.wav", None)) - monkeypatch.setattr(model, "_begin_turn_vocoder_cache", lambda prompt, **kwargs: None) - monkeypatch.setattr(model, "_t2w_pre_lookahead", lambda: 5) - monkeypatch.setattr( - model, - "_build_tts_condition_embeds", - lambda ids, hidden: torch.ones((ids.numel(), 4), dtype=torch.float32), - ) - - windows = [] - - def fake_stream_window(token_list, prompt_wav_path, *, last_chunk): - del prompt_wav_path - windows.append((list(token_list), last_chunk)) - return torch.tensor([float(len(token_list))]) - - monkeypatch.setattr(model, "_t2w_stream_window", fake_stream_window) - - info = { - "request_id": "duplex-metadata-tail-flush", - "meta": {"turn_eos_token_id": 9310}, - "ids": {"tts": [21]}, - "hidden_states": {"tts": [[0.1, 0.2, 0.3, 0.4]]}, - "codes": {"ref": None}, - "end_of_turn": True, - "native_duplex": True, - } - - list(model._create_native_duplex_stream_gen(info)) - - assert windows[-1][1] is True - assert "duplex-metadata-tail-flush" not in model._talker_turn_states - assert "duplex-metadata-tail-flush" not in model._talker_consumed_tokens - - -def test_minicpmo_native_duplex_terminal_only_new_turn_does_not_start_tts(monkeypatch): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( - MiniCPMO45OmniTTSForConditionalGeneration, - ) - - class _FakeEmbText: - weight = torch.ones((200, 4), dtype=torch.float32) - - def __call__(self, ids): - return torch.ones((ids.numel(), 4), dtype=torch.float32) - - class _FakeTTS: - audio_bos_token_id = 1 - config = SimpleNamespace(num_audio_tokens=128) - emb_text = _FakeEmbText() - model = SimpleNamespace(config=SimpleNamespace(rope_theta=10000.0)) - - def __init__(self): - self.calls = 0 - - def generate_chunk(self, **kwargs): - del kwargs - self.calls += 1 - return torch.tensor([[[33]]]), None - - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - tts = _FakeTTS() - model.tts_obj = tts - model.audio_tokenizer = SimpleNamespace() - model._talker_turn_states = {} - model._talker_consumed_tokens = {} - - monkeypatch.setattr(model, "_lazy_init_tts", lambda: None) - monkeypatch.setattr( - model, "_build_tts_sampling_params", lambda: SimpleNamespace(temperature=0.8, repetition_penalty=1.05) - ) - monkeypatch.setattr(model, "_tts_runtime_config", lambda: SimpleNamespace(streaming_generator_chunk=25)) - monkeypatch.setattr(model, "_resolve_prompt_wav_path", lambda ref, sr: ("/tmp/ref.wav", None)) - monkeypatch.setattr(model, "_begin_turn_vocoder_cache", lambda prompt, **kwargs: None) - monkeypatch.setattr(model, "_t2w_pre_lookahead", lambda: 5) - monkeypatch.setattr( - model, - "_build_tts_condition_embeds", - lambda ids, hidden: torch.ones((ids.numel(), 4), dtype=torch.float32), - ) - monkeypatch.setattr( - model, - "_t2w_stream_window", - lambda token_list, prompt_wav_path, *, last_chunk: torch.tensor([float(len(token_list))]), - ) - - outputs = list( - model._create_native_duplex_stream_gen( - { - "request_id": "duplex-terminal-only", - "duplex": {"epoch": 0, "turn_id": 1}, - "meta": {"turn_eos_token_id": 9310}, - "ids": {"tts": [9310]}, - "hidden_states": {"tts": [[0.1, 0.2, 0.3, 0.4]]}, - "codes": {"ref": None}, - "native_duplex": True, + context = SimpleNamespace( + bridge_states={ + "duplex": { + "epoch": 3, + "model_turn_id": 7, } - ) - ) - - assert tts.calls == 0 - assert len(outputs) == 1 - assert outputs[0][0].numel() == 0 - assert outputs[0][1] is True - assert model._talker_turn_states == {} - assert model._talker_consumed_tokens == {} - - -def test_minicpmo_native_duplex_talker_finished_request_closes_session_keyed_turn_state(): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( - MiniCPMO45OmniTTSForConditionalGeneration, - _TalkerTurnState, - ) - - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - model.audio_tokenizer = SimpleNamespace(stream_cache={"stale": True}, hift_cache_dict={"stale": True}) - model._stream_gens = {} - model._talker_turn_states = { - "sid-stage1": _TalkerTurnState(prompt_wav_path=None, temp_prompt_wav_path=None), - } - model._talker_consumed_tokens = {"sid-stage1": 12} - model._talker_request_keys = {"req-stage1": "sid-stage1"} - - model.on_requests_finished({"req-stage1"}) - - assert "sid-stage1" not in model._talker_turn_states - assert "sid-stage1" not in model._talker_consumed_tokens - assert "req-stage1" not in model._talker_request_keys - assert model.audio_tokenizer.stream_cache is None - assert model.audio_tokenizer.hift_cache_dict == {} - - -def test_minicpmo_native_duplex_talker_new_turn_reopens_session_keyed_state(monkeypatch): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( - MiniCPMO45OmniTTSForConditionalGeneration, - _TalkerTurnState, - ) - - class _FakeTTS: - audio_bos_token_id = 88 - - def __init__(self): - self.config = SimpleNamespace(num_audio_tokens=128) - self.model = SimpleNamespace(config=SimpleNamespace()) - self.emb_text = torch.nn.Embedding(256, 4) - self.calls = [] - - def generate_chunk(self, **kwargs): - self.calls.append(kwargs) - return torch.tensor([[[10]]]), "new-kv" - - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - tts = _FakeTTS() - model.tts_obj = tts - model.audio_tokenizer = SimpleNamespace(stream_cache={"old": True}, hift_cache_dict={"old": True}) - model._talker_turn_states = { - "sid-stage1": _TalkerTurnState( - prompt_wav_path=None, - temp_prompt_wav_path=None, - epoch=0, - turn_id=1, - ), - } - model._talker_turn_states["sid-stage1"].past_key_values = "old-kv" - model._talker_consumed_tokens = {"sid-stage1": 9} - model._talker_request_keys = {} - - monkeypatch.setattr(model, "_lazy_init_tts", lambda: None) - monkeypatch.setattr( - model, "_build_tts_sampling_params", lambda: SimpleNamespace(temperature=0.8, repetition_penalty=1.05) - ) - monkeypatch.setattr(model, "_tts_runtime_config", lambda: SimpleNamespace(streaming_generator_chunk=25)) - monkeypatch.setattr(model, "_resolve_prompt_wav_path", lambda ref, sr: ("/tmp/ref.wav", None)) - monkeypatch.setattr(model, "_begin_turn_vocoder_cache", lambda prompt, **kwargs: None) - monkeypatch.setattr(model, "_t2w_pre_lookahead", lambda: 5) - monkeypatch.setattr( - model, - "_build_tts_condition_embeds", - lambda ids, hidden: torch.ones((ids.numel(), 4), dtype=torch.float32), - ) - monkeypatch.setattr( - model, - "_t2w_stream_window", - lambda token_list, prompt_wav_path, *, last_chunk: torch.tensor([float(len(token_list))]), - ) - - info = { - "request_id": "sid-stage1", - "duplex": {"epoch": 0, "turn_id": 2}, - "meta": {"turn_eos_token_id": 9310}, - "ids": {"tts": [21]}, - "hidden_states": {"tts": [[0.1, 0.2, 0.3, 0.4]]}, - "codes": {"ref": None}, - "native_duplex": True, - } - - list(model._create_native_duplex_stream_gen(info)) - - assert tts.calls[0]["past_key_values"] is None - assert tts.calls[0]["text_start_pos"] == 0 - assert model._talker_consumed_tokens["sid-stage1"] == 1 - assert model._talker_turn_states["sid-stage1"].turn_id == 2 - - -def test_minicpmo_native_duplex_talker_segment_end_preserves_token2wav_stream_until_turn_end(monkeypatch): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( - MiniCPMO45OmniTTSForConditionalGeneration, - ) - - class _FakeEmbText: - weight = torch.ones((200, 4), dtype=torch.float32) - - def __call__(self, ids): - return torch.ones((ids.numel(), 4), dtype=torch.float32) - - class _FakeTTS: - audio_bos_token_id = 1 - config = SimpleNamespace(num_audio_tokens=128) - emb_text = _FakeEmbText() - model = SimpleNamespace(config=SimpleNamespace(rope_theta=10000.0)) - - def generate_chunk(self, **kwargs): - del kwargs - return torch.tensor([[[33]]]), "kv" - - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - model.tts_obj = _FakeTTS() - model.audio_tokenizer = SimpleNamespace() - model._talker_turn_states = {} - model._talker_consumed_tokens = {} - - monkeypatch.setattr(model, "_lazy_init_tts", lambda: None) - monkeypatch.setattr( - model, "_build_tts_sampling_params", lambda: SimpleNamespace(temperature=0.8, repetition_penalty=1.05) - ) - monkeypatch.setattr(model, "_tts_runtime_config", lambda: SimpleNamespace(streaming_generator_chunk=25)) - monkeypatch.setattr(model, "_resolve_prompt_wav_path", lambda ref, sr: ("/tmp/ref.wav", None)) - monkeypatch.setattr(model, "_begin_turn_vocoder_cache", lambda prompt, **kwargs: None) - monkeypatch.setattr(model, "_t2w_pre_lookahead", lambda: 5) - monkeypatch.setattr( - model, - "_build_tts_condition_embeds", - lambda ids, hidden: torch.ones((ids.numel(), 4), dtype=torch.float32), - ) - - windows = [] - - def fake_stream_window(token_list, prompt_wav_path, *, last_chunk): - del prompt_wav_path - windows.append((list(token_list), last_chunk)) - return torch.tensor([float(len(token_list))]) - - monkeypatch.setattr(model, "_t2w_stream_window", fake_stream_window) - - info = { - "request_id": "duplex-segment-tail-flush", - "meta": {"turn_eos_token_id": 9310, "segment_end": True}, - "ids": {"tts": [21]}, - "hidden_states": {"tts": [[0.1, 0.2, 0.3, 0.4]]}, - "codes": {"ref": None}, - "native_duplex": True, - } - - list(model._create_native_duplex_stream_gen(info)) - - state = model._talker_turn_states["duplex-segment-tail-flush"] - assert windows == [] - assert state.token2wav_buffer == [4218, 4218, 4218, 33] - - turn_end_info = { - **info, - "ids": {"tts": [21, 9310]}, - "hidden_states": { - "tts": [ - [0.1, 0.2, 0.3, 0.4], - [0.5, 0.6, 0.7, 0.8], - ] - }, - "meta": {"turn_eos_token_id": 9310, "turn_end": True}, - } - list(model._create_native_duplex_stream_gen(turn_end_info)) - - assert windows == [([4218, 4218, 4218, 33, 33], True)] - assert "duplex-segment-tail-flush" not in model._talker_turn_states - - -def test_tts_scheduler_eos_uses_tokenizer_im_end_when_config_has_no_eos(): - class _Tokenizer: - unk_token_id = 0 - - def convert_tokens_to_ids(self, token): - return {"<|im_end|>": 77}.get(token, self.unk_token_id) - - def encode(self, text, add_special_tokens=False): - del add_special_tokens - return [77] if text == "<|im_end|>" else [self.unk_token_id] - - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - model.config = SimpleNamespace(eos_token_id=None, vocab_size=100) - model.vllm_config = SimpleNamespace(model_config=SimpleNamespace()) - model._text_tokenizer = _Tokenizer() - model._ar_last_chunk_flags = [True] - - logits = model.compute_logits(torch.zeros(1, 4)) - - assert int(torch.argmax(logits[0]).item()) == 77 - - -def test_tts_streaming_non_final_chunk_does_not_mark_extra_rows_eos(): - class _Tokenizer: - unk_token_id = 0 - - def convert_tokens_to_ids(self, token): - return {"<|im_end|>": 77}.get(token, self.unk_token_id) - - def encode(self, text, add_special_tokens=False): - del add_special_tokens - return [77] if text == "<|im_end|>" else [self.unk_token_id] - - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - model.config = SimpleNamespace(eos_token_id=None, vocab_size=100) - model.vllm_config = SimpleNamespace(model_config=SimpleNamespace()) - model._text_tokenizer = _Tokenizer() - model._ar_last_chunk_flags = [False] - - logits = model.compute_logits(torch.zeros(3, 4)) - sampled = torch.argmax(logits, dim=-1).tolist() - - assert sampled == [1, 1, 1] - - -def test_native_duplex_turn_end_detection_is_not_segment_end_detection(): - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - turn_eos_id = 9310 - - assert model._native_duplex_input_ends_turn( - { - "native_duplex": True, - "ids": {"tts": [9304, 42, turn_eos_id]}, - "meta": {"turn_eos_token_id": turn_eos_id}, - } - ) - assert not model._native_duplex_input_ends_turn( - { - "native_duplex": True, - "ids": {"tts": [9304, 42, 9308]}, - "meta": { - "turn_eos_token_id": turn_eos_id, - "segment_end": True, - }, } ) + converted = llm2tts([source], prompt=[{}], _streaming_context=context)[0] -def _native_duplex_meta(): - return { - "tts_bos_token_id": 9301, - "tts_eos_token_id": 9302, - "listen_token_id": 9303, - "speak_token_id": 9304, - "tts_pad_token_id": 9305, - "unit_token_id": 9306, - "unit_end_token_id": 9307, - "chunk_eos_token_id": 9308, - "chunk_tts_eos_token_id": 9309, - "turn_eos_token_id": 9310, - } + info = converted["model_intermediate_buffer"] + assert info["native_duplex"] is True + assert info["ids"]["tts"] == [21, 22] + assert converted["prompt_token_ids"] == [0, 0, 0, 0] + assert info["meta"]["replace_streaming_prompt"] is True + assert info["meta"]["segment_end"] is True + assert info["duplex"]["epoch"] == 3 + assert info["duplex"]["turn_id"] == 7 -def _native_duplex_handoff(request_id, prompt_ids, cumulative_output_ids, *, text="hello"): - latent = torch.arange( - (len(prompt_ids) + len(cumulative_output_ids)) * 4, - dtype=torch.float16, - ).reshape(-1, 4) - output = SimpleNamespace( - token_ids=cumulative_output_ids, - text=text, - multimodal_output={ - "latent": latent, - "duplex_prompt_token_ids": list(prompt_ids), - "meta": _native_duplex_meta(), - }, - ) - return SimpleNamespace( - request_id=request_id, - prompt_token_ids=[0, 0, 0], - outputs=[output], +def test_native_duplex_requires_tokenizer_boundary_metadata() -> None: + latent = torch.zeros((3, 4)) + source = _output( + prompt_ids=[101], + output_ids=[21, 22], + latent=latent, + multimodal_output={"duplex_prompt_token_ids": [101]}, ) - -def _install_native_duplex_decoder(streaming_context, token_text): - def _decode(token_ids): - return "".join(token_text.get(int(token_id), "") for token_id in token_ids) - - streaming_context.source_token_decoder = _decode - return streaming_context + with pytest.raises(ValueError, match="tokenizer-derived.*metadata"): + llm2tts([source], prompt=[{}], _streaming_context=SimpleNamespace(bridge_states={})) -def test_native_duplex_segment_text_comes_from_token_delta_decode(): - streaming_context = _install_native_duplex_decoder( - SimpleNamespace(bridge_states={}), - { - 21: "今早", - 22: "吃东。", - 31: "西了吗？", - }, +def test_llm2tts_does_not_alias_live_thinker_token_list() -> None: + live_tokens = [11, 12] + latent = torch.zeros((3, 4)) + source = _output( + prompt_ids=[101], + output_ids=list(live_tokens), + latent=latent, + token_list=live_tokens, ) - ids1, text1 = _native_duplex_segment_output_ids( - [9304, 21, 22, 9308], - "今早吃东。", - streaming_context, - request_id="duplex-token-text", - ) - assert ids1 == [9304, 21, 22, 9308] - assert text1 == "今早吃东。" + converted = llm2tts([source], prompt=[{}])[0] + live_tokens.append(13) - ids2, text2 = _native_duplex_segment_output_ids( - [9304, 21, 22, 9308, 9304, 31, 9308], - "今早吃东。西了吗？如需陪聊，", - streaming_context, - request_id="duplex-token-text", - ) - assert ids2 == [9304, 31, 9308] - assert text2 == "西了吗？" - - -def test_native_duplex_segment_restart_does_not_reuse_character_cursor(): - streaming_context = _install_native_duplex_decoder( - SimpleNamespace(bridge_states={}), - { - 21: "你好", - 31: "你好", - 32: "，有什么想聊的吗？", - }, - ) - - _native_duplex_segment_output_ids( - [9304, 21, 9310], - "你好，第一轮尾巴", - streaming_context, - request_id="duplex-restart-text", - ) - - ids2, text2 = _native_duplex_segment_output_ids( - [9304, 31, 32, 9310], - "你好，第一轮尾巴你好，有什么想聊的吗？", - streaming_context, - request_id="duplex-restart-text", - ) - assert ids2 == [9304, 31, 32, 9310] - assert text2 == "你好，有什么想聊的吗？" - - -def test_native_duplex_segment_turn_id_change_keeps_cumulative_token_cursor(): - streaming_context = _install_native_duplex_decoder( - SimpleNamespace(bridge_states={"duplex": {"turn_id": 1}}), - { - 21: "你好", - 32: "，有什么想聊的吗？", - }, - ) - - ids1, text1 = _native_duplex_segment_output_ids( - [9304, 21], - "你好", - streaming_context, - request_id="duplex-turn-id-reset", - ) - assert ids1 == [9304, 21] - assert text1 == "你好" - - streaming_context.bridge_states["duplex"]["turn_id"] = 2 - ids2, text2 = _native_duplex_segment_output_ids( - [9304, 21, 32], - "你好，有什么想聊的吗？", - streaming_context, - request_id="duplex-turn-id-reset", - ) - assert ids2 == [32] - assert text2 == "，有什么想聊的吗？" - - -def test_llm2tts_native_duplex_hands_off_segment_deltas(): - """Each duplex handoff must carry only the current segment's tokens. - - The thinker's resumable request reports cumulative output ids while - earlier segments are already folded into the prompt, so re-forwarding - the full list misaligns the hidden states and grows the talker prompt - without bound. - """ - streaming_context = SimpleNamespace(bridge_states={}) - - # Segment 1: speak + two text tokens + chunk_eos. - seg1_output = [9304, 21, 22, 9308] - handoff1 = _native_duplex_handoff("duplex-delta", [101, 102], list(seg1_output)) - converted1 = llm2tts([handoff1], prompt=[{}], _streaming_context=streaming_context) - info1 = converted1[0]["model_intermediate_buffer"] - assert info1["ids"]["output"] == seg1_output - assert info1["llm_output_text"] == ["hello"] - assert info1["meta"]["native_duplex_segment_text"] == "hello" - assert "llm_output_text" in info1["meta"]["override_keys"] - assert ["meta", "native_duplex_segment_text"] in info1["meta"]["override_keys"] - assert converted1[0]["prompt_token_ids"] == [21, 22] - assert info1["meta"]["segment_end"] is True - - # Segment 2: prompt now contains the folded segment-1 tokens; the - # cumulative output list still carries segment 1 ahead of segment 2. - seg2_output = [9304, 31, 9308] - prompt2 = [101, 102, *seg1_output, 555] - handoff2 = _native_duplex_handoff("duplex-delta", prompt2, [*seg1_output, *seg2_output]) - converted2 = llm2tts([handoff2], prompt=[{}], _streaming_context=streaming_context) - info2 = converted2[0]["model_intermediate_buffer"] - assert info2["ids"]["output"] == seg2_output - # Cumulative text "hello" was fully delivered with segment 1; only the - # delta (empty here) rides along with segment 2. - assert info2["llm_output_text"] == [""] - assert info2["meta"]["native_duplex_segment_text"] == "" - assert converted2[0]["prompt_token_ids"] == [31] - assert info2["meta"]["segment_end"] is True - - -def test_llm2tts_native_duplex_marks_talker_handoff_as_data_plane(): - streaming_context = SimpleNamespace( - bridge_states={ - "duplex": { - "session_id": "sid-stage1-stop", - "incarnation": 2, - "epoch": 0, - "turn_id": 7, - "session_config": { - "extra_body": { - "duplex_stage_sampling_params": { - "1": {"stop_token_ids": [151645]}, - }, - }, - }, - }, - }, - ) - handoff = _native_duplex_handoff("duplex-stage1-stop", [101, 102], [9304, 21, 9308]) - - converted = llm2tts([handoff], prompt=[{}], _streaming_context=streaming_context) - - duplex = converted[0]["model_intermediate_buffer"]["duplex"] - assert duplex["data_plane"] is True - assert duplex["session_id"] == "sid-stage1-stop" - assert duplex["incarnation"] == 2 - assert duplex["epoch"] == 0 - assert duplex["turn_id"] == 7 - assert duplex["session_config"]["extra_body"]["duplex_stage_sampling_params"]["1"]["stop_token_ids"] == [151645] - - -def test_llm2tts_native_duplex_carries_runtime_ref_audio_to_talker_payload(): - ref_audio = base64.b64encode(struct.pack("<3f", 0.1, 0.2, 0.3)).decode("ascii") - streaming_context = SimpleNamespace( - bridge_states={ - "duplex": { - "session_id": "sid-runtime-ref", - "runtime_config": { - "ref_audio_data": ref_audio, - "ref_audio_format": "pcm_f32le", - "ref_audio_sample_rate_hz": 16000, - }, - }, - }, - ) - handoff = _native_duplex_handoff("duplex-runtime-ref", [101, 102], [9304, 21, 9308]) - - converted = llm2tts([handoff], prompt=[{}], _streaming_context=streaming_context) - - info = converted[0]["model_intermediate_buffer"] - assert torch.allclose(torch.tensor(info["codes"]["ref"]), torch.tensor([0.1, 0.2, 0.3])) - assert info["meta"]["ref_audio_sr"] == 16000 - assert info["duplex"]["runtime_config"]["ref_audio_data"] == ref_audio - - -def test_llm2tts_native_duplex_advances_model_turn_only_after_turn_eos(): - streaming_context = SimpleNamespace( - bridge_states={ - "duplex": { - "session_id": "sid-model-turn", - "epoch": 0, - "turn_id": 9, - "model_turn_id": 0, - }, - } - ) - - first = _native_duplex_handoff( - "duplex-model-turn", - [101], - [9304, 21, 9310], - text="first", - ) - converted_first = llm2tts([first], prompt=[{}], _streaming_context=streaming_context) - - first_duplex = converted_first[0]["model_intermediate_buffer"]["duplex"] - assert first_duplex["turn_id"] == 0 - assert streaming_context.bridge_states["duplex"]["model_turn_id"] == 1 - - second = _native_duplex_handoff( - "duplex-model-turn", - [101, 9304, 21, 9310, 555], - [9304, 21, 9310, 9304, 22, 9308], - text="second", - ) - converted_second = llm2tts([second], prompt=[{}], _streaming_context=streaming_context) - - second_duplex = converted_second[0]["model_intermediate_buffer"]["duplex"] - assert second_duplex["turn_id"] == 1 - assert streaming_context.bridge_states["duplex"]["model_turn_id"] == 1 - - -def test_llm2tts_native_duplex_accumulates_tts_condition_across_handoffs(): - """Every handoff must carry the FULL accumulated tts condition. - - The runner's resume-prefill path REPLACES the streaming buffer (only - in-place updates merge), so per-segment tts payloads were silently lost - for alternating segments; the talker then vocalized text it never saw. - Handing the complete history per handoff makes replacement lossless. - """ - streaming_context = SimpleNamespace(bridge_states={}) - - seg1_output = [9304, 21, 22, 9308] - handoff1 = _native_duplex_handoff("duplex-acc", [101, 102], list(seg1_output)) - converted1 = llm2tts([handoff1], prompt=[{}], _streaming_context=streaming_context) - info1 = converted1[0]["model_intermediate_buffer"] - assert info1["ids"]["tts"] == [21, 22] - assert len(info1["hidden_states"]["tts"]) == 2 - - seg2_output = [9304, 31, 9308] - prompt2 = [101, 102, *seg1_output, 555] - handoff2 = _native_duplex_handoff("duplex-acc", prompt2, [*seg1_output, *seg2_output]) - converted2 = llm2tts([handoff2], prompt=[{}], _streaming_context=streaming_context) - info2 = converted2[0]["model_intermediate_buffer"] - assert info2["ids"]["tts"] == [21, 22, 31] - assert len(info2["hidden_states"]["tts"]) == 3 - - -def test_llm2tts_native_duplex_preserves_text_cursor_across_turn_reset(): - streaming_context = _install_native_duplex_decoder( - SimpleNamespace(bridge_states={}), - { - 21: "hello", - 31: "hello", - }, - ) - - seg1_output = [9304, 21, 9310] - handoff1 = _native_duplex_handoff("duplex-turn-text", [101, 102], seg1_output, text="hello") - converted1 = llm2tts([handoff1], prompt=[{}], _streaming_context=streaming_context) - info1 = converted1[0]["model_intermediate_buffer"] - assert info1["meta"]["native_duplex_segment_text"] == "hello" - - seg2_output = [9304, 31, 9310] - handoff2 = _native_duplex_handoff( - "duplex-turn-text", - [101, 102, *seg1_output], - seg2_output, - text="hellohello", - ) - converted2 = llm2tts([handoff2], prompt=[{}], _streaming_context=streaming_context) - info2 = converted2[0]["model_intermediate_buffer"] - assert info2["ids"]["tts"] == [31, 9310] - assert info2["meta"]["native_duplex_segment_text"] == "hello" - - -def test_llm2tts_native_duplex_preserves_output_cursor_across_turn_reset(): - streaming_context = _install_native_duplex_decoder( - SimpleNamespace(bridge_states={}), - { - 21: "hello", - 31: "hello", - }, - ) - - seg1_output = [9304, 21, 9310] - handoff1 = _native_duplex_handoff("duplex-turn-output", [101, 102], seg1_output, text="hello") - converted1 = llm2tts([handoff1], prompt=[{}], _streaming_context=streaming_context) - info1 = converted1[0]["model_intermediate_buffer"] - assert info1["ids"]["output"] == seg1_output - assert info1["ids"]["tts"] == [21, 9310] - - seg2_output = [9304, 31, 9310] - handoff2 = _native_duplex_handoff( - "duplex-turn-output", - [101, 102, *seg1_output], - [*seg1_output, *seg2_output], - text="hellohello", - ) - converted2 = llm2tts([handoff2], prompt=[{}], _streaming_context=streaming_context) - info2 = converted2[0]["model_intermediate_buffer"] - assert info2["ids"]["output"] == seg2_output - assert info2["ids"]["tts"] == [31, 9310] - assert info2["meta"]["native_duplex_segment_text"] == "hello" - - -def test_llm2tts_native_duplex_resets_tts_condition_on_new_turn_without_resetting_output_cursor(): - streaming_context = _install_native_duplex_decoder( - SimpleNamespace(bridge_states={"duplex": {"turn_id": 0}}), - { - 21: "first", - 31: "second", - }, - ) - - seg1_output = [9304, 21, 9308] - handoff1 = _native_duplex_handoff("duplex-turn-condition", [101, 102], seg1_output, text="first") - converted1 = llm2tts([handoff1], prompt=[{}], _streaming_context=streaming_context) - info1 = converted1[0]["model_intermediate_buffer"] - assert info1["ids"]["output"] == seg1_output - assert info1["ids"]["tts"] == [21] - - streaming_context.bridge_states["duplex"]["turn_id"] = 1 - seg2_output = [9304, 31, 9308] - handoff2 = _native_duplex_handoff( - "duplex-turn-condition", - [101, 102, *seg1_output], - [*seg1_output, *seg2_output], - text="firstsecond", - ) - converted2 = llm2tts([handoff2], prompt=[{}], _streaming_context=streaming_context) - info2 = converted2[0]["model_intermediate_buffer"] - assert info2["ids"]["output"] == seg2_output - assert info2["ids"]["tts"] == [31] - assert info2["meta"]["native_duplex_segment_text"] == "second" - - -def test_llm2tts_never_aliases_thinker_token_list(): - """The talker prompt must never be the thinker's live token list object. - - CompletionOutput.token_ids can alias the upstream detokenizer's internal - list; forwarding that object lets the stage-1 streaming update extend the - list with itself, doubling the recorded output every segment. - """ - live_token_list = [9304, 21, 22, 9308] - handoff = _native_duplex_handoff("duplex-alias", [101, 102], live_token_list) - converted = llm2tts([handoff], prompt=[{}]) - scheduler_prompt = converted[0]["prompt_token_ids"] - assert scheduler_prompt is not live_token_list - scheduler_prompt.extend(scheduler_prompt) - assert live_token_list == [9304, 21, 22, 9308] + assert converted["model_intermediate_buffer"]["ids"]["output"] == [11, 12] diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 13a491896..10f423c84 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -1129,23 +1129,22 @@ class TestDeployConfigLoading: stages = merge_pipeline_deploy(pipeline, deploy) assert deploy.session_mode == "duplex" - assert deploy.async_chunk is False + assert deploy.async_chunk is True assert deploy.active_stream_window == 1 - assert [stage.session_mode for stage in stages] == ["duplex", "duplex"] - assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex"] - assert [stage.yaml_engine_args["async_scheduling"] for stage in stages] == [False, False] + assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"] + assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"] + assert [stage.yaml_engine_args["async_scheduling"] for stage in stages] == [False, False, False] assert all("Async" not in (stage.scheduler_cls or "") for stage in stages) assert deploy.stages[1].enforce_eager is False - assert deploy.stages[1].compilation_config == {"cudagraph_mode": "PIECEWISE"} - assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 4096 - assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [151645] + assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 2048 + assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [1] @pytest.mark.parametrize( - ("filename", "stage0_devices", "stage1_devices", "stage1_replicas"), + ("filename", "stage0_devices", "stage1_devices", "stage2_devices", "stage1_replicas"), [ - ("minicpmo_4_5_3gpu_stage1_replicas.yaml", "0", "1,2", 2), - ("minicpmo_4_5_4gpu_stage1_replicas.yaml", "0", "1,2,3", 3), - ("minicpmo_4_5_8x4090_stage1_replicas.yaml", "0,1,2,3", "4,5,6,7", 4), + ("minicpmo_4_5_3gpu_stage1_replicas.yaml", "0", "1,2", "0", 2), + ("minicpmo_4_5_4gpu_stage1_replicas.yaml", "0", "1,2,3", "0", 3), + ("minicpmo_4_5_8x4090_stage1_replicas.yaml", "0,1,2,3", "4,5,6,7", "5", 4), ], ) def test_load_minicpmo_replica_deploy_configs( @@ -1153,6 +1152,7 @@ class TestDeployConfigLoading: filename: str, stage0_devices: str, stage1_devices: str, + stage2_devices: str, stage1_replicas: int, ): deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / filename @@ -1164,11 +1164,15 @@ class TestDeployConfigLoading: stages = merge_pipeline_deploy(pipeline, deploy) assert deploy.session_mode == "duplex" - assert [stage.yaml_runtime["devices"] for stage in stages] == [stage0_devices, stage1_devices] + assert [stage.yaml_runtime["devices"] for stage in stages] == [ + stage0_devices, + stage1_devices, + stage2_devices, + ] assert stages[0].yaml_runtime["num_replicas"] == 1 assert stages[1].yaml_runtime["num_replicas"] == stage1_replicas - assert all(stage.yaml_engine_args["async_scheduling"] is False for stage in stages) - assert all("Async" not in (stage.scheduler_cls or "") for stage in stages) + assert all(stage.yaml_engine_args["async_scheduling"] is False for stage in stages[:2]) + assert all("Async" not in (stage.scheduler_cls or "") for stage in stages[:2]) def test_custom_voice_dir_is_pipeline_wide_engine_arg(self, tmp_path): deploy_path = tmp_path / "qwen3_tts_custom_voice.yaml" diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 90cf3afe1..83d270da4 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -663,6 +663,7 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk # inter_stage_outputs mirrors multimodal_outputs (PR #4792). assert output.inter_stage_outputs is not None assert output.multimodal_outputs is not None + assert output.inter_stage_outputs is output.multimodal_outputs assert torch.equal(output.inter_stage_outputs[0]["hidden"], output.multimodal_outputs[0]["hidden"]) assert torch.equal(output.inter_stage_outputs[1]["hidden"], output.multimodal_outputs[1]["hidden"]) assert torch.equal(output.multimodal_outputs[0]["hidden"], torch.tensor([[1.0, 10.0]])) diff --git a/tests/worker/test_native_duplex_hooks.py b/tests/worker/test_native_duplex_hooks.py index a5a81a1de..661656b7e 100644 --- a/tests/worker/test_native_duplex_hooks.py +++ b/tests/worker/test_native_duplex_hooks.py @@ -1,8 +1,6 @@ from __future__ import annotations import inspect -import os -import threading from types import SimpleNamespace import numpy as np @@ -544,220 +542,6 @@ def test_minicpmo_stage0_rejects_invalid_resolved_ref_audio(): ) -def test_minicpmo_tts_native_duplex_exports_segment_text_not_accumulated_condition_text(): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( - MiniCPMO45OmniForConditionalGeneration, - ) - - class _Talker: - _ar_last_chunk_flags = [False] - _ar_last_emitted_text = "你好，你有什莫想聊的吗？" - - def __call__(self, **kwargs): - del kwargs - return None, torch.zeros(8, dtype=torch.float32) - - model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) - model.model_stage = "tts" - model.config = SimpleNamespace(hidden_size=4) - model.talker = _Talker() - - output = model.forward( - input_ids=torch.zeros(1, dtype=torch.long), - positions=torch.zeros(1, dtype=torch.long), - runtime_additional_information=[ - { - "native_duplex": True, - "llm_output_text": ["你好，你有什莫想聊的吗？你好，你有什莫想聊的吗？"], - "meta": { - "native_duplex_segment_text": "你好，你有什莫想聊的吗？", - }, - } - ], - ) - - text_bytes = output.multimodal_outputs["meta.llm_output_text_utf8"].detach().cpu().tolist() - assert bytes(text_bytes).decode("utf-8") == "你好，你有什莫想聊的吗？" - assert int(output.multimodal_outputs["meta.audio_text_total_chars"].item()) == len("你好，你有什莫想聊的吗？") - - -def test_minicpmo_tts_native_duplex_exports_model_turn_end_metadata(): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( - MiniCPMO45OmniForConditionalGeneration, - ) - - class _Talker: - _ar_last_chunk_flags = [True] - _ar_turn_end_flags = [True] - - def __call__(self, **kwargs): - del kwargs - return None, torch.zeros(0, dtype=torch.float32) - - model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) - model.model_stage = "tts" - model.config = SimpleNamespace(hidden_size=4) - model.talker = _Talker() - - output = model.forward( - input_ids=torch.zeros(1, dtype=torch.long), - positions=torch.zeros(1, dtype=torch.long), - runtime_additional_information=[ - { - "native_duplex": True, - "meta": {"native_duplex_segment_text": ""}, - } - ], - ) - - assert int(output.multimodal_outputs["meta.turn_end"].item()) == 1 - - -def test_minicpmo_tts_routes_batched_requests_and_terminal_flags_independently(): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( - MiniCPMO45OmniForConditionalGeneration, - ) - - class _Talker: - def __init__(self): - self.calls = [] - self._ar_last_chunk_flags = [True] - self._ar_turn_end_flags = [False] - self._ar_last_emitted_text = "" - - def __call__(self, **kwargs): - info = kwargs["additional_information"] - request_id = info["global_request_id"][0] - self.calls.append(request_id) - is_last = request_id == "req-b" - self._ar_last_chunk_flags = [is_last] - self._ar_turn_end_flags = [is_last] - self._ar_last_emitted_text = request_id - value = 1.0 if request_id == "req-a" else 2.0 - return None, torch.full((2,), value, dtype=torch.float32) - - model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) - model.model_stage = "tts" - model.config = SimpleNamespace(hidden_size=4) - model.talker = _Talker() - - output = model.forward( - input_ids=torch.zeros(3, dtype=torch.long), - positions=torch.zeros(3, dtype=torch.long), - runtime_additional_information=[ - { - "global_request_id": ["req-a"], - "native_duplex": True, - "duplex": {"epoch": 3, "turn_id": 7}, - "meta": {"native_duplex_segment_text": "first"}, - }, - { - "global_request_id": ["req-idle"], - "native_duplex": True, - "duplex": {"epoch": 3, "turn_id": 9}, - "meta": {"native_duplex_segment_text": "idle"}, - }, - { - "global_request_id": ["req-b"], - "native_duplex": True, - "duplex": {"epoch": 4, "turn_id": 8}, - "meta": {"native_duplex_segment_text": "second"}, - }, - ], - request_token_spans=[(0, 2), (2, 2), (2, 3)], - ) - - assert model.talker.calls == ["req-a", "req-b"] - assert model.talker._ar_last_chunk_flags == [False, False, True] - assert model.talker._ar_turn_end_flags == [False, False, True] - assert output.multimodal_outputs["meta.req_id"] == ["req-a", "req-b"] - assert output.multimodal_outputs["meta.sparse_audio"] == ["1"] - assert [waveform.tolist() for waveform in output.multimodal_outputs["model_outputs"]] == [ - [1.0, 1.0], - [2.0, 2.0], - ] - assert [int(value.item()) for value in output.multimodal_outputs["meta.tts_is_last_chunk"]] == [0, 1] - assert [int(value.item()) for value in output.multimodal_outputs["meta.turn_end"]] == [0, 1] - assert [int(value.item()) for value in output.multimodal_outputs["meta.duplex_epoch"]] == [3, 4] - assert [int(value.item()) for value in output.multimodal_outputs["meta.duplex_turn_id"]] == [7, 8] - - -@pytest.mark.parametrize( - ("runtime_info", "request_token_spans", "expected_rows", "expected_ids"), - [ - ( - [ - { - "global_request_id": ["req-a"], - "native_duplex": True, - "duplex": {"epoch": 3, "turn_id": 7}, - "meta": {"native_duplex_segment_text": "first"}, - } - ], - [(0, 3)], - [3], - ["req-a"], - ), - ( - [ - { - "global_request_id": ["req-a"], - "native_duplex": True, - "duplex": {"epoch": 3, "turn_id": 7}, - "meta": {"native_duplex_segment_text": "first"}, - }, - { - "global_request_id": ["req-b"], - "native_duplex": True, - "duplex": {"epoch": 4, "turn_id": 8}, - "meta": {"native_duplex_segment_text": "second"}, - }, - ], - [(0, 2), (2, 3)], - [2, 1], - ["req-a", "req-b"], - ), - ], - ids=["single-request", "batched-requests"], -) -def test_minicpmo_tts_ignores_cuda_graph_padding_after_request_spans( - runtime_info, - request_token_spans, - expected_rows, - expected_ids, -): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( - MiniCPMO45OmniForConditionalGeneration, - ) - - class _Talker: - def __init__(self): - self.input_rows = [] - self._ar_last_chunk_flags = [True] - self._ar_turn_end_flags = [False] - self._ar_last_emitted_text = "" - - def __call__(self, **kwargs): - self.input_rows.append(int(kwargs["input_ids"].shape[0])) - return None, torch.ones(2, dtype=torch.float32) - - model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) - model.model_stage = "tts" - model.config = SimpleNamespace(hidden_size=4) - model.talker = _Talker() - - output = model.forward( - input_ids=torch.zeros(4, dtype=torch.long), - positions=torch.zeros(4, dtype=torch.long), - runtime_additional_information=runtime_info, - request_token_spans=request_token_spans, - ) - - assert model.talker.input_rows == expected_rows - assert output.text_hidden_states.shape == (4, 4) - assert output.multimodal_outputs["meta.req_id"] == expected_ids - - def test_minicpmo_stage0_special_token_ids_are_tokenizer_derived(): from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( MiniCPMO45Stage0DuplexRuntime, @@ -2427,150 +2211,3 @@ def test_minicpmo_stage0_session_context_includes_resolved_ref_audio(): assert state.context_token_ids == [1, 2, 3, 151683, 151683, 4, 5] assert len(state.context_embeds) == 6 - - -def test_minicpmo_tts_ref_audio_lru_evicts_matching_vocoder_base_cache(tmp_path): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( - MiniCPMO45OmniTTSForConditionalGeneration, - MiniCPMO45TTSRuntimeConfig, - ) - - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - model._runtime_config = MiniCPMO45TTSRuntimeConfig(ref_audio_file_cache_size=1) - model._t2w_base_caches = {} - model._token2wav_state_lock = threading.RLock() - first_path = model._write_ref_audio_prompt_wav(np.array([0.1, -0.1], dtype=np.float32), 16_000) - assert first_path is not None - model._t2w_base_caches[first_path] = (torch.tensor([1]), {"cache": torch.tensor([2])}) - - second_path = model._write_ref_audio_prompt_wav(np.array([0.2, -0.2], dtype=np.float32), 16_000) - - assert second_path is not None - assert second_path != first_path - assert first_path not in model._t2w_base_caches - assert not os.path.exists(first_path) - - -def test_minicpmo_tts_prompt_wav_path_does_not_default_to_model_asset(tmp_path): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( - MiniCPMO45OmniTTSForConditionalGeneration, - ) - - model_dir = tmp_path / "MiniCPM-o-4_5" - assets_dir = model_dir / "assets" - assets_dir.mkdir(parents=True) - (assets_dir / "HT_ref_audio.wav").write_bytes(b"placeholder") - - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - model.vllm_config = SimpleNamespace(model_config=SimpleNamespace(model=str(model_dir))) - model._write_ref_audio_prompt_wav = lambda ref_audio, ref_audio_sr: None - - assert model._resolve_prompt_wav_path(None, None) == (None, None) - - -def test_minicpmo_tts_stream_turn_id_prefers_duplex_identity(): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( - MiniCPMO45OmniTTSForConditionalGeneration, - ) - - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - - assert ( - model._stream_turn_id( - { - "duplex": {"turn_id": 0, "model_turn_id": 4}, - "meta": {"turn_id": 3}, - } - ) - == 4 - ) - assert ( - model._stream_turn_id( - { - "duplex": {"turn_id": 0}, - "meta": {"turn_id": 3}, - } - ) - == 0 - ) - assert model._stream_turn_id({"meta": {"turn_id": 3}, "turn_id": 2}) == 3 - assert model._stream_turn_id({"turn_id": 2}) == 2 - - -def test_minicpmo_tts_no_ref_audio_initializes_empty_vocoder_cache(): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( - MiniCPMO45OmniTTSForConditionalGeneration, - _TalkerTurnState, - ) - - class AudioTokenizer: - def set_stream_cache(self, prompt_wav_path): - raise AssertionError(f"unexpected reference-audio cache load: {prompt_wav_path!r}") - - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - model._t2w_base_caches = {} - model._token2wav_state_lock = threading.RLock() - model.audio_tokenizer = AudioTokenizer() - state = _TalkerTurnState(None, None) - - model._begin_turn_vocoder_cache(None, state=state) - - assert state.stream_cache is None - assert state.hift_cache_dict == {} - assert state.vocoder_initialized is True - assert model._t2w_base_caches[""] == (None, {}) - - -def test_minicpmo_tts_native_duplex_merges_and_left_pads_nonfinal_unit_audio(): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( - _native_duplex_unit_waveform, - ) - - waveform = _native_duplex_unit_waveform( - [torch.tensor([1.0, 2.0]), torch.tensor([3.0])], - turn_end=False, - target_samples=6, - ) - - assert waveform is not None - assert waveform.tolist() == [0.0, 0.0, 0.0, 1.0, 2.0, 3.0] - - -def test_minicpmo_tts_native_duplex_does_not_pad_final_unit_audio(): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( - _native_duplex_unit_waveform, - ) - - waveform = _native_duplex_unit_waveform( - [torch.tensor([1.0, 2.0]), torch.tensor([3.0])], - turn_end=True, - target_samples=6, - ) - - assert waveform is not None - assert waveform.tolist() == [1.0, 2.0, 3.0] - - -def test_minicpmo_tts_turn_cleanup_removes_deleted_ref_audio_vocoder_cache(tmp_path): - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( - MiniCPMO45OmniTTSForConditionalGeneration, - _TalkerTurnState, - ) - - prompt_path = tmp_path / "session-ref.wav" - prompt_path.write_bytes(b"wav") - model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) - model._talker_turn_states = { - "request": _TalkerTurnState(str(prompt_path), str(prompt_path)), - } - model._talker_consumed_tokens = {"request": 1} - model._talker_request_keys = {} - model._t2w_base_caches = { - str(prompt_path): (torch.tensor([1]), {"cache": torch.tensor([2])}), - } - model._token2wav_state_lock = threading.RLock() - model.audio_tokenizer = None - - assert model._close_turn_state("request") is True - assert str(prompt_path) not in model._t2w_base_caches - assert not prompt_path.exists() diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index cff36b9f6..701e3b589 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -440,6 +440,33 @@ def test_update_intermediate_buffer_accumulates(): assert torch.allclose(buf["b"], torch.tensor([2.0])) +def test_update_additional_information_deserializes_new_request_payload(): + from vllm_omni.engine.serialization import serialize_additional_information + + runner = _make_runner(req_ids=("r1",), hidden_size=4) + conditioning = { + "tts_token_ids": torch.tensor([1, 2]), + "tts_hidden_states": torch.ones(2, 4), + } + scheduler_output = SimpleNamespace( + scheduled_new_reqs=[ + SimpleNamespace( + req_id="r1", + additional_information=serialize_additional_information(conditioning), + ) + ], + scheduled_cached_reqs=SimpleNamespace(), + ) + + OmniGPUModelRunner._update_additional_information(runner, scheduler_output) + + assert torch.equal(runner.model_intermediate_buffer["r1"]["tts_token_ids"], conditioning["tts_token_ids"]) + assert torch.equal( + runner.model_intermediate_buffer["r1"]["tts_hidden_states"], + conditioning["tts_hidden_states"], + ) + + def test_update_intermediate_buffer_skips_empty_update(): """Validate that an empty update dict is a no-op.""" runner = _make_runner(req_ids=("r1",), hidden_size=4) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 5e16140b8..2a8efe824 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -921,6 +921,10 @@ def merge_pipeline_deploy( stage_type, worker_type = _resolve_execution_mode(ps.execution_type) input_proc, next_stage_proc = _select_processor_funcs(ps, deploy.async_chunk) engine_args = _build_engine_args(ps, ds, pipeline, deploy, next_stage_proc) + # Downstream stages may share a multimodal wrapper class without owning + # an encoder. Do not make vLLM profile dummy multimodal inputs for them. + if not ps.requires_multimodal_data: + engine_args.setdefault("skip_mm_profiling", True) sched_cls = _resolve_scheduler( ps.execution_type, engine_args.get("async_scheduling", True), diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index e88730353..93654f27f 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -700,9 +700,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): session.num_computed_tokens -= session.num_output_placeholders session.num_output_placeholders = 0 session.spec_token_ids = [] - if self.chunk_transfer_adapter: + stage_id = self.vllm_config.model_config.stage_id + if self.chunk_transfer_adapter and self.chunk_transfer_adapter.receives_chunks: self.chunk_transfer_adapter.requests_num_chunks_sent.pop(session.external_req_id, None) - if self.vllm_config.model_config.stage_id != 0: + if stage_id != 0: # Downstream async-chunk stages receive real payloads from the # connector. This update only resumes polling for the next segment. self.chunk_transfer_adapter.segment_finished_requests.discard(session.request_id) @@ -720,6 +721,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if self.log_stats: session.record_event(EngineCoreEventType.QUEUED) return + update_info = getattr(update, "additional_information", None) + update_meta = update_info.get("meta") if isinstance(update_info, dict) else None + if isinstance(update_meta, dict) and update_meta.get("replace_streaming_prompt") is True: + self._replace_streaming_session(session, update) + return super()._update_request_as_session(session, update) if hasattr(update, "model_intermediate_buffer"): session.model_intermediate_buffer = update.model_intermediate_buffer diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index c6f82ec23..3882fd42d 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -685,26 +685,4 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): Do not expend prompt id using update. """ - if self.chunk_transfer_adapter: - self.chunk_transfer_adapter.segment_finished_requests.discard(session.request_id) - session._output_token_ids.clear() - session._all_token_ids.clear() - new_prompt = update.prompt_token_ids or () - session._all_token_ids.extend(new_prompt) - session.num_computed_tokens = 0 - session.prompt_token_ids = update.prompt_token_ids or () - session.additional_information = update.additional_information or None - # Update block hashes for the new tokens. - session.update_block_hashes() - session.num_prompt_tokens = len(session.prompt_token_ids) - session.arrival_time = update.arrival_time - session.sampling_params = update.sampling_params - if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ: - self.num_waiting_for_streaming_input -= 1 - session.status = RequestStatus.WAITING - if session in self.skipped_waiting: - self.skipped_waiting.remove_requests((session,)) - self._enqueue_waiting_request(session) - - if self.log_stats: - session.record_event(EngineCoreEventType.QUEUED) + self._replace_streaming_session(session, update) diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 4eaa07cbb..3df47a357 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -7,8 +7,9 @@ from typing import Any from vllm.logger import init_logger from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.engine import EngineCoreEventType from vllm.v1.metrics.stats import SchedulerStats -from vllm.v1.request import RequestStatus +from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm_omni.core.sched.output import OmniChunkRecvHandle, OmniSchedulerOutput @@ -50,6 +51,31 @@ class OmniSchedulerMixin: if input_coordinator is not None: input_coordinator.free_finished_request(request_id) + def _replace_streaming_session(self, session: Request, update: StreamingUpdate) -> None: + """Replace a downstream stage's placeholder with its next payload.""" + adapter = getattr(self, "chunk_transfer_adapter", None) + if adapter is not None: + adapter.segment_finished_requests.discard(session.request_id) + session._output_token_ids.clear() + session._all_token_ids.clear() + new_prompt = update.prompt_token_ids or () + session._all_token_ids.extend(new_prompt) + session.num_computed_tokens = 0 + session.prompt_token_ids = new_prompt + session.additional_information = update.additional_information or None + session.update_block_hashes() + session.num_prompt_tokens = len(new_prompt) + session.arrival_time = update.arrival_time + session.sampling_params = update.sampling_params + if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + self.num_waiting_for_streaming_input -= 1 + session.status = RequestStatus.WAITING + if session in self.skipped_waiting: + self.skipped_waiting.remove_requests((session,)) + self._enqueue_waiting_request(session) + if self.log_stats: + session.record_event(EngineCoreEventType.QUEUED) + def _consume_pending_connector_output(self, model_mode: str) -> None: """Drain ``self._latest_omni_connector_output`` into the coordinator. diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py index eec3e5c33..db2cf15ea 100644 --- a/vllm_omni/data_entry_keys.py +++ b/vllm_omni/data_entry_keys.py @@ -61,12 +61,20 @@ class OmniPayloadMeta(TypedDict, total=False): finished: torch.Tensor is_segment_finished: torch.Tensor stream_finished: torch.Tensor + request_id: str + chunk_seq: int + cache_epoch: int + codec_chunk_frames: int + codec_left_context_frames: int + code_flat_numel: int + last_chunk: bool req_id: list[str] left_context_size: int right_holdback_size: int override_keys: list[tuple[str, str]] num_processed_tokens: int next_stage_prompt_len: int + replace_streaming_prompt: bool ar_width: int eol_token_id: int visual_token_start_id: int @@ -82,6 +90,7 @@ class OmniPayloadMeta(TypedDict, total=False): ref_context_request_id: str ref_context_included: bool talker_prefill_offset: int + omni_final_stage_id: int class OmniPayload(TypedDict, total=False): @@ -150,12 +159,17 @@ class MetaStruct(_StructBase): finished: torch.Tensor | None = None is_segment_finished: torch.Tensor | None = None stream_finished: torch.Tensor | None = None + request_id: str | None = None + chunk_seq: int | None = None + cache_epoch: int | None = None + last_chunk: bool | None = None req_id: list[str] | None = None left_context_size: int | None = None right_holdback_size: int | None = None override_keys: list[tuple[str, str]] | None = None num_processed_tokens: int | None = None next_stage_prompt_len: int | None = None + replace_streaming_prompt: bool | None = None ar_width: int | None = None eol_token_id: int | None = None visual_token_start_id: int | None = None diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index d191dddcc..da84667af 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -1,20 +1,24 @@ -# MiniCPM-o 4.5 default deploy: single-GPU layout. -# -# Both stages are co-located on GPU 0: -# Stage 0 (thinker): multimodal understanding, ~80% mem. -# Stage 1 (talker) : MiniCPMTTS + in-process Token2Wav vocoder, ~10% mem -# (talker LM plus the vocoder that emits the waveform). -# -# Other GPU layouts: minicpmo_4_5_3gpu.yaml, minicpmo_4_5_8x4090.yaml. -async_chunk: false +# MiniCPM-o 4.5 three-stage deploy on one large-memory GPU. +pipeline: minicpmo_4_5 +async_chunk: true trust_remote_code: true enable_prefix_caching: false +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 25 + codec_left_context_frames: 3 + stages: - stage_id: 0 - max_num_seqs: 16 - gpu_memory_utilization: 0.8 + max_num_seqs: 4 + gpu_memory_utilization: 0.65 enforce_eager: false tensor_parallel_size: 1 max_num_batched_tokens: 16384 @@ -31,21 +35,42 @@ stages: repetition_penalty: 1.1 - stage_id: 1 - # Talker only consumes runtime_additional_information[0]; any value > 1 - # makes concurrent requests share request-0's audio (mm_outputs fallback). - max_num_seqs: 1 - gpu_memory_utilization: 0.1 + max_num_seqs: 4 + gpu_memory_utilization: 0.15 enforce_eager: false max_num_batched_tokens: 8192 devices: "0" + output_connectors: + to_stage_2: connector_of_shared_memory + # Codec sampling comes from the checkpoint's tts_config. These parameters + # govern only vLLM's binary continue/stop token. default_sampling_params: temperature: 0.0 top_p: 1.0 top_k: -1 - max_tokens: 1 + min_tokens: 50 + max_tokens: 2048 + stop_token_ids: [1] seed: 42 detokenize: false + - stage_id: 2 + max_num_seqs: 4 + gpu_memory_utilization: 0.15 + enforce_eager: true + enable_chunked_prefill: false + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + input_connectors: + from_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + detokenize: true + platforms: npu: stages: diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index 7a9c400cc..eddfd0877 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml @@ -1,64 +1,2 @@ -# MiniCPM-o 4.5 deploy: 2-GPU layout (thinker on GPU 0, talker on GPU 1). -# -# Stage 0 (thinker): GPU 0 — multimodal understanding, ~90% mem (thinker -# gets a dedicated card, so it can take more of -# the GPU than the single-GPU co-located layout). -# Stage 1 (talker) : GPU 1 — MiniCPMTTS + Token2Wav, ~75% mem -# (covers both the talker LM and the in-process -# Token2Wav vocoder that emits the waveform). -# -# Use this when you have two GPUs and want the talker + Token2Wav vocoder -# on its own card instead of sharing GPU 0 with the thinker (the default -# single-GPU layout co-locates both and can OOM under load). -pipeline: minicpmo_4_5 -async_chunk: false - -trust_remote_code: true -enable_prefix_caching: false - -stages: - - stage_id: 0 - max_num_seqs: 16 - gpu_memory_utilization: 0.9 - enforce_eager: false - tensor_parallel_size: 1 - max_num_batched_tokens: 16384 - devices: "0" - limit_mm_per_prompt: - video: 1 - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 2048 - seed: 42 - detokenize: true - repetition_penalty: 1.1 - - - stage_id: 1 - # Talker only consumes runtime_additional_information[0]; any value > 1 - # makes concurrent requests share request-0's audio (mm_outputs fallback). - max_num_seqs: 1 - gpu_memory_utilization: 0.75 - enforce_eager: false - max_num_batched_tokens: 8192 - devices: "1" - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 1 - seed: 42 - detokenize: false - -platforms: - npu: - stages: - - stage_id: 0 - max_num_batched_tokens: 8192 - compilation_config: - cudagraph_mode: PIECEWISE - - stage_id: 1 - max_num_batched_tokens: 8192 - compilation_config: - cudagraph_mode: PIECEWISE +# Backward-compatible name for the canonical two-GPU batching profile. +base_config: minicpmo_4_5_batching.yaml diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index 69f77eb6d..9254d3c70 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml @@ -1,23 +1,31 @@ -# MiniCPM-o 4.5 deploy: 3-GPU layout (thinker TP=2 + talker on GPU 2). +# MiniCPM-o 4.5 three-stage deploy with one GPU per stage. # -# Stage 0 (thinker): GPU 0,1 — 2-way tensor parallel, ~70% mem each. -# Stage 1 (talker) : GPU 2 — MiniCPMTTS + Token2Wav, ~75% mem -# (covers both the talker LM and the in-process -# Token2Wav vocoder that emits the waveform). +# This replaces the former fused two-stage layout: Thinker, Talker, and +# Code2Wav run independently on GPUs 0, 1, and 2. pipeline: minicpmo_4_5 -async_chunk: false +async_chunk: true trust_remote_code: true enable_prefix_caching: false +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 25 + codec_left_context_frames: 3 + stages: - stage_id: 0 - max_num_seqs: 16 - gpu_memory_utilization: 0.7 + max_num_seqs: 4 + gpu_memory_utilization: 0.9 enforce_eager: false - tensor_parallel_size: 2 + tensor_parallel_size: 1 max_num_batched_tokens: 16384 - devices: "0,1" + devices: "0" limit_mm_per_prompt: video: 1 default_sampling_params: @@ -30,19 +38,40 @@ stages: repetition_penalty: 1.1 - stage_id: 1 - # Talker only consumes runtime_additional_information[0]; any value > 1 - # makes concurrent requests share request-0's audio (mm_outputs fallback). - max_num_seqs: 1 - gpu_memory_utilization: 0.75 + max_num_seqs: 4 + gpu_memory_utilization: 0.9 enforce_eager: false max_num_batched_tokens: 8192 - devices: "2" + devices: "1" + output_connectors: + to_stage_2: connector_of_shared_memory + # vLLM samples the binary continue/stop row; codec sampling comes from + # the checkpoint's tts_config and currently uses a deterministic seed. default_sampling_params: temperature: 0.0 top_p: 1.0 top_k: -1 - max_tokens: 1 + min_tokens: 50 + max_tokens: 2048 + stop_token_ids: [1] seed: 42 detokenize: false + - stage_id: 2 + max_num_seqs: 4 + gpu_memory_utilization: 0.9 + enforce_eager: true + enable_chunked_prefill: false + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "2" + input_connectors: + from_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + detokenize: true + platforms: {} diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml index dd249c70a..1b38059a3 100644 --- a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml @@ -4,10 +4,8 @@ # thinker weights take ~17.5 GB per card # on 4090, leaving ~3 GB for KV cache / # activations). -# Stage 1 (talker) : GPU 4 — MiniCPMTTS + Token2Wav, ~90% mem -# (covers both the talker LM and the -# in-process Token2Wav vocoder that emits -# the waveform). +# Stage 1 (talker) : GPU 4 — MiniCPMTTS codec generation, ~90% mem. +# Stage 2 (code2wav): GPU 5 — codec-to-waveform generation, ~90% mem. # # max_model_len is pinned to 4096 because the upstream config.json # advertises max_position_embeddings=40960 (tuned for H20 141GB); on a @@ -16,11 +14,21 @@ # this thinker model + vision/audio encoders + Whisper preallocation, # so 4096 is the empirically-verified safe ceiling. pipeline: minicpmo_4_5 -async_chunk: false +async_chunk: true trust_remote_code: true enable_prefix_caching: false +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 25 + codec_left_context_frames: 3 + stages: - stage_id: 0 max_num_seqs: 1 @@ -42,19 +50,40 @@ stages: repetition_penalty: 1.1 - stage_id: 1 - # Talker only consumes runtime_additional_information[0]; any value > 1 - # makes concurrent requests share request-0's audio (mm_outputs fallback). max_num_seqs: 1 gpu_memory_utilization: 0.9 enforce_eager: true max_num_batched_tokens: 8192 devices: "4" + output_connectors: + to_stage_2: connector_of_shared_memory + # vLLM samples the binary continue/stop row; codec sampling comes from + # the checkpoint's tts_config and currently uses a deterministic seed. default_sampling_params: temperature: 0.0 top_p: 1.0 top_k: -1 - max_tokens: 1 + min_tokens: 50 + max_tokens: 2048 + stop_token_ids: [1] seed: 42 detokenize: false + - stage_id: 2 + max_num_seqs: 1 + gpu_memory_utilization: 0.9 + enforce_eager: true + enable_chunked_prefill: false + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "5" + input_connectors: + from_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + detokenize: true + platforms: {} diff --git a/vllm_omni/deploy/minicpmo_4_5_batching.yaml b/vllm_omni/deploy/minicpmo_4_5_batching.yaml new file mode 100644 index 000000000..75693feb7 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_batching.yaml @@ -0,0 +1,88 @@ +# MiniCPM-o 4.5 two-GPU continuous-batching layout. +# +# Thinker owns GPU 0. Talker and Code2Wav share GPU 1, keeping text decoding +# isolated from codec/vocoder contention while preserving stage-2 batching. +pipeline: minicpmo_4_5 +async_chunk: true + +trust_remote_code: true +enable_prefix_caching: false + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 25 + codec_left_context_frames: 3 + +stages: + - stage_id: 0 + max_num_seqs: 4 + gpu_memory_utilization: 0.9 + enforce_eager: false + tensor_parallel_size: 1 + max_num_batched_tokens: 16384 + devices: "0" + limit_mm_per_prompt: + video: 1 + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 2048 + seed: 42 + detokenize: true + repetition_penalty: 1.1 + + - stage_id: 1 + max_num_seqs: 4 + gpu_memory_utilization: 0.55 + enforce_eager: false + max_num_batched_tokens: 8192 + devices: "1" + output_connectors: + to_stage_2: connector_of_shared_memory + # Codec sampling is model-owned by the checkpoint's tts_config (defaults: + # T=0.8, top-k=25, top-p=0.85, repetition penalty=1.05, seed=42). + # These parameters only govern vLLM's binary continue/stop token. + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + min_tokens: 50 + max_tokens: 2048 + stop_token_ids: [1] + seed: 42 + detokenize: false + + - stage_id: 2 + max_num_seqs: 4 + gpu_memory_utilization: 0.35 + enforce_eager: true + enable_chunked_prefill: false + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "1" + input_connectors: + from_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + detokenize: true + +platforms: + npu: + stages: + - stage_id: 0 + max_num_batched_tokens: 8192 + compilation_config: + cudagraph_mode: PIECEWISE + - stage_id: 1 + max_num_batched_tokens: 8192 + compilation_config: + cudagraph_mode: PIECEWISE diff --git a/vllm_omni/deploy/minicpmo_4_5_duplex.yaml b/vllm_omni/deploy/minicpmo_4_5_duplex.yaml index c227f2817..e58c515c7 100644 --- a/vllm_omni/deploy/minicpmo_4_5_duplex.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_duplex.yaml @@ -1,9 +1,9 @@ # MiniCPM-o 4.5 native duplex deployment. # -# The registered minicpmo_4_5 pipeline owns the fixed Stage0 -> Stage1 -# topology and llm2tts bridge. This overlay only changes session lifetime and -# the Stage1 token budget needed to drain streaming audio segments. -base_config: minicpmo_4_5.yaml +# The registered minicpmo_4_5 pipeline owns the fixed +# Thinker -> Talker -> Code2Wav topology. This overlay changes session +# lifetime and bounds all three stages to two live duplex sessions. +base_config: minicpmo_4_5_batching.yaml pipeline: minicpmo_4_5 session_mode: duplex active_stream_window: 1 @@ -27,10 +27,12 @@ stages: - stage_id: 1 max_num_seqs: 2 async_scheduling: false - # The streaming talker advances Python-owned per-request generators on - # every step, so only fixed-topology model regions are graph-capturable. - compilation_config: - cudagraph_mode: PIECEWISE + # The split Talker exposes a binary continue/stop vocabulary; native + # segment and turn boundaries travel as explicit output metadata. default_sampling_params: - max_tokens: 4096 - stop_token_ids: [151645] + max_tokens: 2048 + stop_token_ids: [1] + + - stage_id: 2 + max_num_seqs: 2 + async_scheduling: false diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index db253ebad..c124e8e43 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -229,6 +229,22 @@ def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], re extended prompt without discarding prior computed state. """ ids = payload_data.get("ids", {}) + meta = payload_data.get("meta", {}) + replace_prompt = isinstance(meta, dict) and meta.get("replace_streaming_prompt") is True + condition_len = meta.get("next_stage_prompt_len") if replace_prompt else None + if isinstance(condition_len, int) and condition_len > 0: + # Some downstream stages consume complete, independently conditioned + # segments instead of extending an existing KV prefix. The producer + # declares that transport behavior explicitly in payload metadata. + new_prompt = [0] * condition_len + request._output_token_ids.clear() + request._all_token_ids.clear() + request._all_token_ids.extend(new_prompt) + request.prompt_token_ids = new_prompt + request.num_computed_tokens = 0 + request.num_prompt_tokens = condition_len + request.update_block_hashes() + return prompt_token_ids = ids.get("prompt", None) if not prompt_token_ids: return diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 26be84012..ea3642c8d 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -13,7 +13,7 @@ from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_p from ..adapter import construct_next_stage_streaming_input_prompt from ..factory import OmniConnectorFactory -from ..utils.config import ConnectorSpec +from ..utils.config import ConnectorSpec, stage_receives_chunks from ..utils.logging import get_connector_logger from .base import OmniTransferAdapterBase @@ -56,6 +56,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self._active_window, ) self.connector = self.create_connector(model_config) + self.receives_chunks = stage_receives_chunks(model_config) super().__init__(model_config) self.model_mode = getattr(model_config, "worker_type", None) or "ar" # State specific to Chunk management @@ -135,7 +136,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): """ stage_id = self.connector.stage_id - if stage_id == 0: + if stage_id == 0 or not self.receives_chunks: return if not hasattr(request, "additional_information"): request.additional_information = None @@ -224,7 +225,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): payload_segment_finished = self._is_truthy_scalar(meta.get("is_segment_finished")) if self.model_mode == "ar": request.additional_information = payload_data - if chunk_id > 0 and request.resumable: + replace_prompt = meta.get("replace_streaming_prompt") is True + if getattr(request, "resumable", False) and (chunk_id > 0 or replace_prompt): # For new streaming input segment, we should update prompt from payload construct_next_stage_streaming_input_prompt(payload_data, request) @@ -464,6 +466,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): callers that don't track aborts may omit it to keep the prior (unguarded) behaviour. """ + if not self.receives_chunks: + return if self.connector.stage_id == 0: return @@ -628,6 +632,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): every parked request is restored unconditionally (the pre-purge behavior). """ + if not self.receives_chunks: + return if scheduler_requests is not None: self._purge_untracked_chunk_requests(self.waiting_for_chunk_waiting_requests, scheduler_requests) self._purge_untracked_chunk_requests(self.waiting_for_chunk_running_requests, scheduler_requests) @@ -659,6 +665,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): Add additional info for cached requests and clean up ready chunks from scheduler output. """ + if not self.receives_chunks: + return stage_id = self.connector.stage_id if stage_id == 0: diff --git a/vllm_omni/distributed/omni_connectors/utils/config.py b/vllm_omni/distributed/omni_connectors/utils/config.py index ee049904f..1a6d395c2 100644 --- a/vllm_omni/distributed/omni_connectors/utils/config.py +++ b/vllm_omni/distributed/omni_connectors/utils/config.py @@ -17,6 +17,34 @@ TRANSFER_ENGINE_CONNECTOR_NAMES = frozenset( ) +def get_stage_connector_role(model_config: Any) -> str | None: + """Return the configured stage connector direction, if explicit.""" + connector_config = getattr(model_config, "stage_connector_config", None) + if isinstance(connector_config, dict): + extra = connector_config.get("extra") + else: + extra = getattr(connector_config, "extra", None) + if isinstance(extra, dict): + role = extra.get("role") + return role if isinstance(role, str) else None + return None + + +def stage_receives_chunks(model_config: Any) -> bool: + """Whether connector chunks, rather than the orchestrator, feed a stage.""" + return get_stage_connector_role(model_config) != "sender" + + +def stage_sends_async_output(model_config: Any) -> bool: + """Whether async output should be partitioned for connector transport.""" + role = get_stage_connector_role(model_config) + if role is not None: + return role == "sender" + # Preserve legacy partitioning while keeping stage-0 orchestrator bridges + # on the normal RequestOutput path. + return getattr(model_config, "stage_id", None) != 0 + + @dataclass class ConnectorSpec: """Specification for a connector instance.""" diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index e10ce8845..e37111063 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -29,6 +29,7 @@ from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.membership_controller import MembershipController @@ -1243,7 +1244,7 @@ class Orchestrator: if ( (finished or (req_state.streaming.enabled and req_state.streaming.segment_finished)) and stage_id < req_state.final_stage_id - and not self.async_chunk + and (not self.async_chunk or not self._stage_receives_async_chunks(stage_id + 1)) and (not self._next_stage_already_submitted(stage_id, req_state) or req_state.streaming.enabled) ): if ( @@ -1291,6 +1292,12 @@ class Orchestrator: def _next_stage_already_submitted(self, stage_id: int, req_state: OrchestratorRequestState) -> bool: return (stage_id + 1) in req_state.stage_submit_ts + def _stage_receives_async_chunks(self, stage_id: int) -> bool: + """Whether a stage's connector supplies its runtime inputs.""" + pool = self.stage_pools[stage_id] + model_config = getattr(pool.stage_vllm_config, "model_config", None) + return stage_receives_chunks(model_config) + def _get_stage_input_processor(self, stage_id: int) -> Any: processor = self._stage_input_processors.get(stage_id) if processor is None: @@ -1963,6 +1970,11 @@ class Orchestrator: for next_stage_id in range(1, req_state.final_stage_id + 1): next_pool = self.stage_pools[next_stage_id] params = req_state.sampling_params_list[next_stage_id] + if not self._stage_receives_async_chunks(next_stage_id): + # Outgoing-only stages receive their first real input from the + # orchestrator. Pre-submitting a placeholder lets it race and + # execute before that conditioning payload arrives. + continue req_state.stage_submit_ts[next_stage_id] = _time.time() _t_submit_start = _time.perf_counter() diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index e3f75d9e3..c2989179c 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -51,6 +51,13 @@ class StageEngineCoreProc(EngineCoreProc): ``EngineCoreProc.run_engine_core()``. """ + def preprocess_add_request(self, request: OmniEngineCoreRequest) -> tuple[Any, int]: + """Preserve omni payloads when vLLM builds its scheduler request.""" + scheduler_request, current_wave = super().preprocess_add_request(request) + scheduler_request.additional_information = request.additional_information + scheduler_request.external_req_id = getattr(request, "external_req_id", request.request_id) + return scheduler_request, current_wave + @staticmethod def run_stage_core( *args: Any, diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index d741e9038..9861d7b38 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1063,6 +1063,16 @@ def get_stage_connector_spec( stage_connectors_cfg = get_stage_connector_config(omni_transfer_config, stage_id) for cfg in stage_connectors_cfg.values(): return dict(cfg.get("spec", {})) + + # A stage can be an async producer without consuming chunks itself. Keep + # its connector for save_async(), but mark it sender-only so the scheduler + # does not park normal orchestrator-provided inputs waiting for a chunk. + target_stage = str(stage_id) + for (from_stage, _to_stage), spec in getattr(omni_transfer_config, "connectors", {}).items(): + if from_stage == target_stage: + extra = dict(spec.extra or {}) + extra.setdefault("role", "sender") + return {"name": spec.name, "extra": extra} return {} diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 32a35ac8c..fe4ec5a78 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -669,6 +669,8 @@ class StagePool: num_tokens_out=num_tokens_out, stage_gen_time_ms=stage_gen_time_ms, batch_id=batch_id, + # This event summarizes one completed request. Execution batching + # happens inside the model runner and is not observable here. batch_size=1, replica_id=replica_id, rx_decode_time_ms=0.0, diff --git a/vllm_omni/entrypoints/cli/benchmark/cli_args.py b/vllm_omni/entrypoints/cli/benchmark/cli_args.py index 66ece628e..7f7671496 100644 --- a/vllm_omni/entrypoints/cli/benchmark/cli_args.py +++ b/vllm_omni/entrypoints/cli/benchmark/cli_args.py @@ -239,6 +239,6 @@ def preprocess_serve_args(args: argparse.Namespace) -> None: """Apply serving benchmark CLI transformations after parsing.""" extra_body = dict(getattr(args, "extra_body", None) or {}) bot_task = getattr(args, "bot_task", None) - if bot_task is not None: + if getattr(args, "backend", None) == "openai-image-edits-omni" and bot_task is not None: extra_body.setdefault("bot_task", bot_task) args.extra_body = extra_body diff --git a/vllm_omni/experimental/fullduplex/engine/intermediate.py b/vllm_omni/experimental/fullduplex/engine/intermediate.py index 07c6ab668..d85d8f712 100644 --- a/vllm_omni/experimental/fullduplex/engine/intermediate.py +++ b/vllm_omni/experimental/fullduplex/engine/intermediate.py @@ -65,7 +65,7 @@ def set_ref_audio(buffer: dict[str, object], waveform: object, sample_rate_hz: i def set_tts_handoff(buffer: dict[str, object], token_ids: object | None, hidden_states: object | None) -> None: - """Store the generic AR-stage token/hidden-state handoff consumed by TTS.""" + """Store the AR-to-TTS handoff used by the full-duplex stage bridge.""" if token_ids is not None: buffer.setdefault("ids", {})["tts"] = token_ids if hidden_states is not None: @@ -73,16 +73,15 @@ def set_tts_handoff(buffer: dict[str, object], token_ids: object | None, hidden_ def get_tts_handoff(info: dict[str, object]) -> tuple[object | None, object | None]: - """Read the generic AR-stage to TTS handoff, including legacy flat keys.""" + """Read the canonical handoff, including the legacy flat aliases.""" ids_info = info.get("ids") hidden_info = info.get("hidden_states") - tts_token_ids = ids_info.get("tts") if isinstance(ids_info, dict) else None - tts_hidden_states = hidden_info.get("tts") if isinstance(hidden_info, dict) else None - if tts_token_ids is None: - tts_token_ids = info.get("tts_token_ids") - if tts_hidden_states is None: - tts_hidden_states = info.get("tts_hidden_states") - return tts_token_ids, tts_hidden_states + token_ids = ids_info.get("tts") if isinstance(ids_info, dict) else None + hidden_states = hidden_info.get("tts") if isinstance(hidden_info, dict) else None + return ( + info.get("tts_token_ids") if token_ids is None else token_ids, + info.get("tts_hidden_states") if hidden_states is None else hidden_states, + ) def get_stream_request_key(info: dict[str, object]) -> str: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py new file mode 100644 index 000000000..34cc9a733 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py @@ -0,0 +1,400 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Strict, state-explicit batching for MiniCPM-o 4.5 Token2wav.""" + +from __future__ import annotations + +from contextlib import nullcontext +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +_SILENCE_TOKEN = 4218 + + +def tensor_signature(value: torch.Tensor) -> tuple[tuple[int, ...], str, str]: + return tuple(value.shape), str(value.dtype), value.device.type + + +def state_shape_signature(state: BatchedToken2WavState) -> tuple[Any, ...]: + flow = tuple((name, tensor_signature(state.flow_cache[name])) for name in sorted(state.flow_cache)) + hift = tuple((name, tensor_signature(state.hift_cache[name])) for name in sorted(state.hift_cache)) + return flow, hift + + +@dataclass(frozen=True) +class PromptFeatures: + speech_tokens: torch.Tensor + speaker_embedding: torch.Tensor + mels: torch.Tensor + + +@dataclass(frozen=True) +class BatchedToken2WavState: + flow_cache: dict[str, torch.Tensor] + hift_cache: dict[str, torch.Tensor] + + +class BatchedToken2Wav(nn.Module): + """Drive Token2wav's modules with dynamically-sized, request-owned caches. + + This class intentionally never calls ``Token2wav.stream`` or + ``Token2wav.__call__``. The upstream object is used only as a one-time + asset loader and prompt feature extractor. + """ + + def __init__(self, token2wav: Any): + super().__init__() + self._token2wav = token2wav + self.flow = token2wav.flow + self.hift = token2wav.hift + self.float16 = bool(token2wav.float16) + self.n_timesteps = int(token2wav.n_timesteps) + self.mel_cache_len = int(token2wav.mel_cache_len) + self.source_cache_len = int(token2wav.source_cache_len) + self.register_buffer( + "speech_window", + token2wav.speech_window.detach().clone(), + persistent=False, + ) + self._prompt_features: dict[tuple[str, str], PromptFeatures] = {} + + def prepare_prompt(self, prompt_cache_id: str, prompt_wav: str) -> PromptFeatures: + cache_key = (prompt_cache_id, prompt_wav) + cached = self._prompt_features.get(cache_key) + if cached is None: + # The generation runner may wrap model.forward in bf16 autocast, + # and vLLM constructs the model under a bf16 default dtype, while + # S3Tokenizer prompt extraction uses fp32 convolution weights. + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.float32) + with torch.amp.autocast("cuda", enabled=False): + values = self._token2wav._prepare_prompt(prompt_wav) + finally: + torch.set_default_dtype(previous_dtype) + cached = PromptFeatures( + speech_tokens=values[0], + speaker_embedding=values[2], + mels=values[3], + ) + self._prompt_features[cache_key] = cached + return cached + + def evict_prompt(self, prompt_cache_id: str, prompt_wav: str) -> None: + """Release request-owned prompt features after stream completion.""" + self._prompt_features.pop((prompt_cache_id, prompt_wav), None) + + @staticmethod + def _repeat_prompt(features: PromptFeatures, batch_size: int) -> tuple[torch.Tensor, ...]: + return ( + features.speech_tokens.expand(batch_size, -1), + features.speaker_embedding.expand(batch_size, -1), + features.mels.expand(batch_size, -1, -1), + ) + + def _autocast(self, device: torch.device): + if device.type != "cuda": + return nullcontext() + if not self.float16: + return torch.amp.autocast("cuda", enabled=False) + return torch.amp.autocast( + "cuda", + dtype=torch.float16, + ) + + def _encode_chunk( + self, + tokens: torch.Tensor, + *, + last_chunk: bool, + cnn_cache: torch.Tensor | None, + att_cache: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + embedded = self.flow.input_embedding(tokens) + hidden, new_cnn, new_att = self.flow.encoder.forward_chunk( + xs=embedded, + last_chunk=last_chunk, + cnn_cache=cnn_cache, + att_cache=att_cache, + ) + return self.flow.encoder_proj(hidden), new_cnn, new_att + + @staticmethod + def _estimator_buffers( + estimator: nn.Module, + x: torch.Tensor, + old_att: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + blocks = estimator.blocks + depth = len(blocks) + batch_size = int(x.shape[0]) + chunk_size = int(x.shape[2]) + old_att_len = int(old_att.shape[3]) if old_att is not None else 0 + block0 = blocks[0] + cnn_channels = int(block0.conv.in_channels + block0.conv.out_channels) + cnn_width = int(block0.conv.block[1].causal_padding[0]) + heads = int(block0.attn.num_heads) + att_width = int(block0.attn.head_dim * 2) + cnn = x.new_empty((depth, batch_size, cnn_channels, cnn_width)) + att = x.new_empty((depth, batch_size, heads, old_att_len + chunk_size, att_width)) + return cnn, att + + def _estimator_step( + self, + estimator: nn.Module, + *, + x: torch.Tensor, + mu: torch.Tensor, + time: torch.Tensor, + speakers: torch.Tensor, + cond: torch.Tensor, + cnn_cache: torch.Tensor | None, + att_cache: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + time_embedding = estimator.t_embedder(time).unsqueeze(1) + width = int(x.shape[-1]) + speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width) + estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1) + cnn_out, att_out = self._estimator_buffers(estimator, estimator_input, att_cache) + old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks) + old_att: Any = att_cache if att_cache is not None else [None] * len(estimator.blocks) + result = estimator.blocks_forward_chunk( + estimator_input, + time_embedding, + None, + old_cnn, + old_att, + cnn_out, + att_out, + ) + return result, cnn_out, att_out + + def _decode_cfm( + self, + mu: torch.Tensor, + speakers: torch.Tensor, + cond: torch.Tensor, + *, + cnn_cache: torch.Tensor | None, + att_cache: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + decoder = self.flow.decoder + estimator = decoder.estimator + batch_size = int(mu.shape[0]) + offset = int(att_cache.shape[4]) if att_cache is not None else 0 + end = offset + int(mu.shape[2]) + if end > int(decoder.rand_noise.shape[2]): + raise RuntimeError( + "MiniCPMO45Code2WavBatchError " + f'{{"reason":"noise_capacity","required":{end},' + f'"available":{int(decoder.rand_noise.shape[2])}}}' + ) + x = decoder.rand_noise[:, :, offset:end].expand(batch_size, -1, -1).clone() + timeline = torch.linspace( + 0, + 1, + self.n_timesteps + 1, + device=mu.device, + dtype=mu.dtype, + ) + timeline = 1 - torch.cos(timeline * 0.5 * torch.pi) + time = timeline[0].expand(batch_size) + mu_cfg = torch.cat((mu, torch.zeros_like(mu)), dim=0) + speakers_cfg = torch.cat((speakers, torch.zeros_like(speakers)), dim=0) + cond_cfg = torch.cat((cond, torch.zeros_like(cond)), dim=0) + next_cnn: list[torch.Tensor] = [] + next_att: list[torch.Tensor] = [] + dt = timeline[1] - timeline[0] + for step in range(self.n_timesteps): + old_cnn = cnn_cache[step] if cnn_cache is not None else None + old_att = att_cache[step] if att_cache is not None else None + estimate, step_cnn, step_att = self._estimator_step( + estimator, + x=torch.cat((x, x), dim=0), + mu=mu_cfg, + time=torch.cat((time, time), dim=0), + speakers=speakers_cfg, + cond=cond_cfg, + cnn_cache=old_cnn, + att_cache=old_att, + ) + conditional, unconditional = estimate.split(batch_size, dim=0) + velocity = (1.0 + decoder.inference_cfg_rate) * conditional - decoder.inference_cfg_rate * unconditional + x = x + dt * velocity + time = time + dt + if step + 1 < self.n_timesteps: + dt = timeline[step + 2] - time[0] + next_cnn.append(step_cnn) + next_att.append(step_att) + return x, torch.stack(next_cnn), torch.stack(next_att) + + @staticmethod + def _split_flow_cache(cache: dict[str, torch.Tensor], batch_size: int) -> list[dict[str, torch.Tensor]]: + result: list[dict[str, torch.Tensor]] = [] + for row in range(batch_size): + result.append( + { + "conformer_cnn_cache": cache["conformer_cnn_cache"][row : row + 1].detach().clone(), + "conformer_att_cache": cache["conformer_att_cache"][:, row : row + 1].detach().clone(), + "estimator_cnn_cache": torch.cat( + ( + cache["estimator_cnn_cache"][:, :, row : row + 1], + cache["estimator_cnn_cache"][:, :, batch_size + row : batch_size + row + 1], + ), + dim=2, + ).detach(), + "estimator_att_cache": torch.cat( + ( + cache["estimator_att_cache"][:, :, row : row + 1], + cache["estimator_att_cache"][:, :, batch_size + row : batch_size + row + 1], + ), + dim=2, + ).detach(), + } + ) + return result + + @staticmethod + def _stack_flow_cache(states: list[BatchedToken2WavState]) -> dict[str, torch.Tensor]: + flows = [state.flow_cache for state in states] + conditional_cnn = [flow["estimator_cnn_cache"][:, :, 0:1] for flow in flows] + unconditional_cnn = [flow["estimator_cnn_cache"][:, :, 1:2] for flow in flows] + conditional_att = [flow["estimator_att_cache"][:, :, 0:1] for flow in flows] + unconditional_att = [flow["estimator_att_cache"][:, :, 1:2] for flow in flows] + return { + "conformer_cnn_cache": torch.cat([flow["conformer_cnn_cache"] for flow in flows], dim=0), + "conformer_att_cache": torch.cat([flow["conformer_att_cache"] for flow in flows], dim=1), + "estimator_cnn_cache": torch.cat((*conditional_cnn, *unconditional_cnn), dim=2), + "estimator_att_cache": torch.cat((*conditional_att, *unconditional_att), dim=2), + } + + def setup_batch( + self, + features: PromptFeatures, + batch_size: int, + ) -> list[BatchedToken2WavState]: + prompt_tokens, speakers, prompt_mels = self._repeat_prompt(features, batch_size) + lookahead = prompt_tokens.new_full((batch_size, 3), _SILENCE_TOKEN) + with self._autocast(prompt_tokens.device): + hidden, conformer_cnn, conformer_att = self._encode_chunk( + torch.cat((prompt_tokens, lookahead), dim=1), + last_chunk=False, + cnn_cache=None, + att_cache=None, + ) + projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1)) + _, estimator_cnn, estimator_att = self._decode_cfm( + hidden.transpose(1, 2).contiguous(), + projected_speakers, + prompt_mels.transpose(1, 2).contiguous(), + cnn_cache=None, + att_cache=None, + ) + flow_cache = { + "conformer_cnn_cache": conformer_cnn, + "conformer_att_cache": conformer_att, + "estimator_cnn_cache": estimator_cnn, + "estimator_att_cache": estimator_att, + } + split = self._split_flow_cache(flow_cache, batch_size) + mel_channels = int(prompt_mels.shape[2]) + return [ + BatchedToken2WavState( + flow_cache=row, + hift_cache={ + "mel": prompt_mels.new_zeros((1, mel_channels, 0)), + "source": prompt_mels.new_zeros((1, 1, 0)), + "speech": prompt_mels.new_zeros((1, 0)), + }, + ) + for row in split + ] + + @staticmethod + def _fade_in_out( + speech: torch.Tensor, + previous: torch.Tensor, + window: torch.Tensor, + ) -> torch.Tensor: + overlap = int(window.shape[0] // 2) + result = speech.clone() + result[..., :overlap] = result[..., :overlap] * window[:overlap] + previous[..., -overlap:] * window[overlap:] + return result + + def decode_batch( + self, + tokens: torch.Tensor, + features: PromptFeatures, + states: list[BatchedToken2WavState], + *, + last_chunk: bool, + ) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]: + batch_size = int(tokens.shape[0]) + if batch_size != len(states): + raise ValueError(f"tokens batch {batch_size} != state batch {len(states)}") + flow_cache = self._stack_flow_cache(states) + speakers = features.speaker_embedding.expand(batch_size, -1) + with self._autocast(tokens.device): + hidden, conformer_cnn, conformer_att = self._encode_chunk( + tokens, + last_chunk=last_chunk, + cnn_cache=flow_cache["conformer_cnn_cache"], + att_cache=flow_cache["conformer_att_cache"], + ) + projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1)) + cond = torch.zeros_like(hidden).transpose(1, 2).contiguous() + chunk_mel, estimator_cnn, estimator_att = self._decode_cfm( + hidden.transpose(1, 2).contiguous(), + projected_speakers, + cond, + cnn_cache=flow_cache["estimator_cnn_cache"], + att_cache=flow_cache["estimator_att_cache"], + ) + + prompt_len = int(features.mels.shape[1]) + if estimator_att.shape[4] > prompt_len + 100: + estimator_att = torch.cat( + (estimator_att[..., :prompt_len, :], estimator_att[..., -100:, :]), + dim=4, + ) + if conformer_att.shape[3] > prompt_len + 100: + conformer_att = torch.cat( + (conformer_att[..., :prompt_len, :], conformer_att[..., -100:, :]), + dim=3, + ) + new_flow = self._split_flow_cache( + { + "conformer_cnn_cache": conformer_cnn, + "conformer_att_cache": conformer_att, + "estimator_cnn_cache": estimator_cnn, + "estimator_att_cache": estimator_att, + }, + batch_size, + ) + old_mel = torch.cat([state.hift_cache["mel"] for state in states], dim=0) + old_source = torch.cat([state.hift_cache["source"] for state in states], dim=0) + old_speech = torch.cat([state.hift_cache["speech"] for state in states], dim=0) + mel = torch.cat((old_mel, chunk_mel), dim=2) + speech, source = self.hift(mel, old_source) + if old_speech.shape[-1] > 0: + window = self.speech_window.to(device=speech.device, dtype=speech.dtype) + speech = self._fade_in_out(speech, old_speech, window) + next_hift = { + "mel": mel[..., -self.mel_cache_len :].detach(), + "source": source[..., -self.source_cache_len :].detach(), + "speech": speech[..., -self.source_cache_len :].detach(), + } + emitted = speech if last_chunk else speech[..., : -self.source_cache_len] + next_states = [ + BatchedToken2WavState( + flow_cache=new_flow[row], + hift_cache={name: value[row : row + 1].detach().clone() for name, value in next_hift.items()}, + ) + for row in range(batch_size) + ] + audios = [emitted[row].reshape(-1).to(dtype=torch.float32) for row in range(batch_size)] + return audios, next_states diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py new file mode 100644 index 000000000..93bd98e76 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -0,0 +1,547 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Strict batched codec-to-waveform stage for MiniCPM-o 4.5.""" + +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import soundfile as sf +import torch +import torch.nn as nn +from vllm.config import VllmConfig + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from .batched_token2wav import ( + BatchedToken2Wav, + BatchedToken2WavState, + state_shape_signature, +) + + +def _batch_error(reason: str, **details: Any) -> RuntimeError: + payload = {"reason": reason, **details} + return RuntimeError(f"MiniCPMO45Code2WavBatchError {json.dumps(payload, sort_keys=True)}") + + +def _scalar(value: Any, default: Any = None) -> Any: + if isinstance(value, torch.Tensor): + return value.reshape(-1)[0].item() if value.numel() else default + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + return _scalar(value[0], default) if value else default + return default if value is None else value + + +def _codec_tensor(value: Any, fallback: torch.Tensor) -> torch.Tensor: + if isinstance(value, torch.Tensor): + return value.reshape(-1).to(device=fallback.device, dtype=torch.long) + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + return torch.as_tensor(value, device=fallback.device, dtype=torch.long).reshape(-1) + return fallback.reshape(-1).to(dtype=torch.long) + + +@dataclass(frozen=True) +class _RequestState: + cache_epoch: int + chunk_seq: int + prompt_cache_id: str + prompt_wav: str + token2wav: BatchedToken2WavState + + +@dataclass(frozen=True) +class _WorkItem: + output_index: int + state_id: str + request_id: str + cache_epoch: int + chunk_seq: int + prompt_cache_id: str + prompt_wav: str + last_chunk: bool + tokens: torch.Tensor + previous: _RequestState | None + segment_text: str + duplex_turn_id: int | None + duplex_epoch: int | None + segment_end: bool + turn_end: bool + + +class MiniCPMO45Code2Wav(nn.Module): + """LLM_GENERATION model that admits only true exact-shape GPU batches.""" + + input_modalities = "audio" + have_multimodal_outputs = True + enable_update_additional_information = True + requires_raw_input_tokens = True + requires_request_ids = True + has_preprocess = False + has_postprocess = False + + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__() + del prefix + self.vllm_config = vllm_config + self.model_path = str(vllm_config.model_config.model) + self.backend: BatchedToken2Wav | None = None + self._states: dict[str, _RequestState] = {} + self._owned_prompt_wavs: dict[str, tuple[str, str]] = {} + extra = self._extra_config() + self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1)) + if self._min_batch_size < 1: + raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1") + self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio")) + self._default_prompt_wav = str( + extra.get( + "prompt_wav", + Path(self.model_path) / "assets" / "HT_ref_audio.wav", + ) + ) + + def _extra_config(self) -> dict[str, Any]: + model_config = getattr(self.vllm_config, "model_config", None) + connector = getattr(model_config, "stage_connector_config", None) + if isinstance(connector, Mapping): + extra = connector.get("extra", connector) + else: + extra = getattr(connector, "extra", None) + return dict(extra) if isinstance(extra, Mapping) else {} + + def _release_prompt(self, state_id: str) -> None: + owned = self._owned_prompt_wavs.pop(state_id, None) + if owned is None: + return + prompt_cache_id, prompt_wav = owned + if self.backend is not None: + self.backend.evict_prompt(prompt_cache_id, prompt_wav) + try: + os.unlink(prompt_wav) + except FileNotFoundError: + pass + + def _write_reference_prompt( + self, + state_id: str, + reference: object, + sample_rate: object, + ) -> tuple[str, str]: + waveform = torch.as_tensor(reference, dtype=torch.float32).reshape(-1).cpu() + if waveform.numel() == 0: + return self._default_prompt_id, self._default_prompt_wav + sample_rate_hz = int(_scalar(sample_rate, 24000)) + if sample_rate_hz <= 0: + raise _batch_error( + "invalid_ref_audio_sample_rate", + request_id=state_id, + sample_rate=sample_rate_hz, + ) + digest = hashlib.sha256() + digest.update(str(sample_rate_hz).encode("ascii")) + digest.update(waveform.numpy().tobytes()) + prompt_cache_id = f"{state_id}:{digest.hexdigest()}" + current = self._owned_prompt_wavs.get(state_id) + if current is not None and current[0] == prompt_cache_id and os.path.exists(current[1]): + return current + + self._release_prompt(state_id) + handle = tempfile.NamedTemporaryFile(prefix="minicpmo45-ref-", suffix=".wav", delete=False) + prompt_wav = handle.name + handle.close() + try: + sf.write(prompt_wav, waveform.numpy(), sample_rate_hz, format="WAV") + except Exception: + try: + os.unlink(prompt_wav) + except FileNotFoundError: + pass + raise + self._owned_prompt_wavs[state_id] = (prompt_cache_id, prompt_wav) + return prompt_cache_id, prompt_wav + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return torch.zeros((input_ids.numel(), 1), device=input_ids.device, dtype=torch.float32) + + def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None: + return None + + @staticmethod + def _split_segments(input_ids: torch.Tensor, counts: Any) -> list[torch.Tensor]: + flat = input_ids.reshape(-1) + if counts is None: + return [flat] + if not isinstance(counts, Sequence) or isinstance(counts, (str, bytes, bytearray)): + raise _batch_error("invalid_seq_token_counts", value_type=type(counts).__name__) + normalized = [int(value) for value in counts] + if any(value < 0 for value in normalized): + raise _batch_error("negative_seq_token_count", counts=normalized) + if sum(normalized) != int(flat.numel()): + raise _batch_error( + "seq_token_count_mismatch", + counts=normalized, + total=int(flat.numel()), + ) + return list(torch.split(flat, normalized)) + + def _parse_item( + self, + index: int, + state_id: str, + segment: torch.Tensor, + info: Mapping[str, Any], + ) -> _WorkItem: + meta = info.get("meta") + if not isinstance(meta, Mapping): + meta = info + request_id = str(_scalar(meta.get("request_id"), _scalar(info.get("request_id"), ""))) + if not request_id: + raise _batch_error("missing_request_id", output_index=index) + cache_epoch = int(_scalar(meta.get("cache_epoch"), 0)) + chunk_seq = int(_scalar(meta.get("chunk_seq"), 0)) + if cache_epoch < 0 or chunk_seq < 0: + raise _batch_error( + "negative_stream_position", + request_id=request_id, + cache_epoch=cache_epoch, + chunk_seq=chunk_seq, + ) + last_chunk = bool(_scalar(meta.get("last_chunk"), False)) + codes = info.get("codes") + audio = codes.get("audio") if isinstance(codes, Mapping) else None + reference = codes.get("ref") if isinstance(codes, Mapping) else None + tokens = _codec_tensor(audio, segment) + if last_chunk and int(_scalar(meta.get("code_flat_numel"), tokens.numel())) == 0: + # The generation scheduler reserves one placeholder token for an + # empty terminal chunk. The producer's explicit length is the + # authority, so do not decode that placeholder as codec data. + tokens = segment.new_empty(0, dtype=torch.long) + previous = self._states.get(state_id) + if previous is None: + if chunk_seq != 0: + raise _batch_error( + "missing_state_for_chunk", + request_id=request_id, + cache_epoch=cache_epoch, + chunk_seq=chunk_seq, + ) + elif cache_epoch < previous.cache_epoch: + raise _batch_error( + "stale_cache_epoch", + request_id=request_id, + expected=previous.cache_epoch, + actual=cache_epoch, + ) + elif cache_epoch > previous.cache_epoch: + if chunk_seq != 0: + raise _batch_error( + "new_epoch_requires_first_chunk", + request_id=request_id, + cache_epoch=cache_epoch, + chunk_seq=chunk_seq, + ) + self._release_prompt(state_id) + previous = None + elif chunk_seq != previous.chunk_seq + 1: + raise _batch_error( + "stale_or_reordered_chunk", + request_id=request_id, + expected=previous.chunk_seq + 1, + actual=chunk_seq, + ) + + explicit_prompt_id = _scalar(meta.get("prompt_cache_id")) + explicit_prompt_wav = _scalar(meta.get("prompt_wav")) + if previous is not None and explicit_prompt_id is None and explicit_prompt_wav is None: + prompt_cache_id = previous.prompt_cache_id + prompt_wav = previous.prompt_wav + elif reference is not None: + prompt_cache_id, prompt_wav = self._write_reference_prompt( + state_id, + reference, + meta.get("ref_audio_sr"), + ) + else: + prompt_cache_id = str(explicit_prompt_id or self._default_prompt_id) + prompt_wav = str(explicit_prompt_wav or self._default_prompt_wav) + + if previous is not None and prompt_cache_id != previous.prompt_cache_id: + raise _batch_error( + "prompt_changed_midstream", + request_id=request_id, + expected=previous.prompt_cache_id, + actual=prompt_cache_id, + ) + if previous is not None and prompt_wav != previous.prompt_wav: + raise _batch_error( + "prompt_changed_midstream", + request_id=request_id, + expected=previous.prompt_wav, + actual=prompt_wav, + ) + duplex_turn_id = _scalar(meta.get("duplex_turn_id")) + duplex_epoch = _scalar(meta.get("duplex_epoch")) + return _WorkItem( + output_index=index, + state_id=state_id, + request_id=request_id, + cache_epoch=cache_epoch, + chunk_seq=chunk_seq, + prompt_cache_id=prompt_cache_id, + prompt_wav=prompt_wav, + last_chunk=last_chunk, + tokens=tokens, + previous=previous, + segment_text=str(_scalar(meta.get("native_duplex_segment_text"), "")), + duplex_turn_id=int(duplex_turn_id) if duplex_turn_id is not None else None, + duplex_epoch=int(duplex_epoch) if duplex_epoch is not None else None, + segment_end=bool(_scalar(meta.get("segment_end"), False)), + turn_end=bool(_scalar(meta.get("turn_end"), False)), + ) + + @staticmethod + def _bucket_key(item: _WorkItem) -> tuple[Any, ...]: + cache_signature: Any + if item.previous is None: + cache_signature = ("uninitialized",) + else: + cache_signature = state_shape_signature(item.previous.token2wav) + return ( + item.prompt_cache_id, + item.prompt_wav, + int(item.tokens.numel()), + cache_signature, + item.last_chunk, + item.cache_epoch, + ) + + @torch.inference_mode() + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> OmniOutput: + del positions, intermediate_tensors, inputs_embeds + ids = input_ids if isinstance(input_ids, torch.Tensor) else torch.empty(0, dtype=torch.long) + segments = self._split_segments(ids, kwargs.get("seq_token_counts")) + empty = torch.empty(0, dtype=torch.float32, device=ids.device) + sample_rate = torch.tensor(24000, dtype=torch.int32) + if not runtime_additional_information: + count = len(segments) + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={ + "model_outputs": [empty for _ in range(count)], + "sr": [sample_rate for _ in range(count)], + }, + ) + if len(runtime_additional_information) != len(segments): + raise _batch_error( + "runtime_info_count_mismatch", + segments=len(segments), + runtime_infos=len(runtime_additional_information), + ) + if self.backend is None: + raise _batch_error("backend_not_loaded") + + state_ids = kwargs.get("request_ids") + if state_ids is None: + state_ids = [] + for index, info in enumerate(runtime_additional_information): + if not isinstance(info, Mapping): + state_ids.append(str(index)) + continue + meta = info.get("meta") + source = meta if isinstance(meta, Mapping) else info + state_ids.append(str(_scalar(source.get("request_id"), index))) + if len(state_ids) != len(segments): + raise _batch_error( + "request_id_count_mismatch", + segments=len(segments), + request_ids=len(state_ids), + ) + items: list[_WorkItem] = [] + for index, (state_id, segment, info) in enumerate( + zip(state_ids, segments, runtime_additional_information, strict=True) + ): + if not isinstance(info, Mapping): + raise _batch_error( + "invalid_runtime_info", + output_index=index, + value_type=type(info).__name__, + ) + items.append(self._parse_item(index, str(state_id), segment, info)) + state_ids = [item.state_id for item in items] + if len(state_ids) != len(set(state_ids)): + raise _batch_error("duplicate_request_in_forward", request_ids=state_ids) + outputs = [empty for _ in segments] + sentinels = [item for item in items if item.last_chunk and item.tokens.numel() == 0] + compute_items = [item for item in items if item.tokens.numel() > 0] + invalid_empty = [item.request_id for item in items if not item.last_chunk and item.tokens.numel() == 0] + if invalid_empty: + raise _batch_error("empty_nonfinal_chunk", request_ids=invalid_empty) + + buckets: dict[tuple[Any, ...], list[_WorkItem]] = {} + for item in compute_items: + buckets.setdefault(self._bucket_key(item), []).append(item) + undersized = [ + { + "size": len(bucket), + "request_ids": [item.request_id for item in bucket], + "codec_len": int(bucket[0].tokens.numel()), + } + for bucket in buckets.values() + if len(bucket) < self._min_batch_size + ] + if undersized: + raise _batch_error( + "exact_shape_bucket_below_minimum", + minimum=self._min_batch_size, + buckets=undersized, + ) + + pending: dict[str, _RequestState | None] = {item.state_id: None for item in sentinels} + for bucket in buckets.values(): + batch_size = len(bucket) + try: + features = self.backend.prepare_prompt( + bucket[0].prompt_cache_id, + bucket[0].prompt_wav, + ) + if bucket[0].previous is None: + states = self.backend.setup_batch(features, batch_size) + else: + states = [item.previous.token2wav for item in bucket if item.previous is not None] + tokens = torch.stack([item.tokens for item in bucket], dim=0) + audios, next_states = self.backend.decode_batch( + tokens, + features, + states, + last_chunk=bucket[0].last_chunk, + ) + except Exception as exc: + if isinstance(exc, RuntimeError) and str(exc).startswith("MiniCPMO45Code2WavBatchError "): + raise + raise _batch_error( + "backend_unsupported_or_failed", + request_ids=[item.request_id for item in bucket], + error_type=type(exc).__name__, + error=str(exc), + ) from exc + if len(audios) != batch_size or len(next_states) != batch_size: + raise _batch_error( + "backend_result_size_mismatch", + expected=batch_size, + audios=len(audios), + states=len(next_states), + ) + for item, audio, next_state in zip(bucket, audios, next_states, strict=True): + outputs[item.output_index] = audio.reshape(-1).to(dtype=torch.float32) + pending[item.state_id] = ( + None + if item.last_chunk + else _RequestState( + cache_epoch=item.cache_epoch, + chunk_seq=item.chunk_seq, + prompt_cache_id=item.prompt_cache_id, + prompt_wav=item.prompt_wav, + token2wav=next_state, + ) + ) + + for request_id, state in pending.items(): + if state is None: + self._states.pop(request_id, None) + self._release_prompt(request_id) + else: + self._states[request_id] = state + sample_rate_tensor = torch.as_tensor(sample_rate, dtype=torch.int32) + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={ + "model_outputs": outputs, + "sr": [sample_rate_tensor.clone() for _ in outputs], + "meta.llm_output_text_utf8": [ + torch.tensor(list(item.segment_text.encode("utf-8")), dtype=torch.uint8) for item in items + ], + "meta.tts_is_last_chunk": [torch.tensor(item.last_chunk, dtype=torch.bool) for item in items], + "meta.segment_end": [torch.tensor(item.segment_end, dtype=torch.bool) for item in items], + "meta.turn_end": [torch.tensor(item.turn_end, dtype=torch.bool) for item in items], + "meta.duplex_turn_id": [ + torch.tensor(item.duplex_turn_id, dtype=torch.int64) + if item.duplex_turn_id is not None + else torch.empty(0, dtype=torch.int64) + for item in items + ], + "meta.duplex_epoch": [ + torch.tensor(item.duplex_epoch, dtype=torch.int64) + if item.duplex_epoch is not None + else torch.empty(0, dtype=torch.int64) + for item in items + ], + }, + ) + + def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: + for request_id in finished_req_ids: + state_id = str(request_id) + self._states.pop(state_id, None) + self._release_prompt(state_id) + + def make_omni_output(self, model_outputs: Any, **_: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + if isinstance(model_outputs, tuple) and len(model_outputs) == len(OmniOutput._fields): + return OmniOutput(*model_outputs) + raise TypeError(f"MiniCPMO45Code2Wav expected OmniOutput, got {type(model_outputs).__name__}") + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + for _ in weights: + pass + if self.backend is not None: + return {name for name, _ in self.named_parameters()} + from stepaudio2.token2wav import Token2wav + + extra = self._extra_config() + prompt_path = Path(self._default_prompt_wav) + if not prompt_path.is_file(): + raise FileNotFoundError(f"MiniCPM-o Code2Wav prompt audio not found: {prompt_path}") + token2wav_path = Path(self.model_path) / "assets" / "token2wav" + if not token2wav_path.is_dir(): + raise FileNotFoundError(f"MiniCPM-o Code2Wav assets not found: {token2wav_path}") + use_float16 = bool(extra.get("token2wav_float16", False)) + previous_dtype = torch.get_default_dtype() + try: + # vLLM constructs bf16 models under a bf16 default-dtype context. + # Token2wav contains fp32-only S3Tokenizer/HiFT modules, so build + # its independent assets in their native precision. + torch.set_default_dtype(torch.float32) + token2wav = Token2wav( + str(token2wav_path), + float16=use_float16, + n_timesteps=int(extra.get("token2wav_n_timesteps", 10)), + ) + finally: + torch.set_default_dtype(previous_dtype) + self.backend = BatchedToken2Wav(token2wav) + # Token2wav loads flow.pt and hift.pt inside its constructor instead of + # from the parent MiniCPM checkpoint iterator. Report those registered + # parameters as initialized so vLLM's strict loader audit does not + # misclassify the independently loaded Stage-2 weights as missing. + return {name for name, _ in self.named_parameters()} diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index ba51c9cfc..c9083c261 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -34,7 +34,6 @@ from vllm.sequence import IntermediateTensors from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata -from vllm_omni.experimental.fullduplex.engine.intermediate import get_stream_request_key from vllm_omni.experimental.fullduplex.minicpmo45.policy import MiniCPMO45DuplexPolicy from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingRow from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( @@ -58,9 +57,11 @@ logger = init_logger(__name__) class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP, SupportsMRoPE): """MiniCPM-o 4.5 Omni model for conditional generation. - This model has two pipeline stages: - - llm: multimodal thinker and text generation - - tts: talker generation followed by its built-in Token2Wav vocoder + Three-stage pipeline: + - thinker (model_stage="llm"): image / video / audio encoders + 3D + resampler + the omni LLM that emits text + hidden states. + - talker (model_stage="tts"): native continuous MiniCPMTTS AR that emits + codec-token deltas for the separate Code2Wav stage. """ @classmethod @@ -106,7 +107,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp elif self.model_stage == "tts": self.thinker = None - # Initialize talker model (LLM generation) + # The Talker is always the runner-owned continuous codec producer. self.talker = init_vllm_registered_model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "talker"), @@ -125,12 +126,18 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp self.make_empty_intermediate_tensors = ( (self.thinker.make_empty_intermediate_tensors) if self.model_stage == "llm" and self.thinker is not None + else self.talker.make_empty_intermediate_tensors + if self.talker is not None else lambda: None ) self._language_model_names = ["model"] self.prefer_model_sampler = self.model_stage in {"llm", "tts"} - self.has_preprocess = self.model_stage == "llm" + # Both AR stages require model-specific embeddings. The Thinker uses + # preprocess for duplex audio, while the Talker converts the + # tts_token_ids/tts_hidden_states handoff into its conditioning + # embeddings and initializes request-local codec generation state. + self.has_preprocess = self.model_stage in {"llm", "tts"} @cached_property def sampler(self): @@ -299,6 +306,8 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp and sampler. This hook only turns the current duplex audio append into the prompt embeddings consumed by the normal runner forward. """ + if self.model_stage == "tts": + return self.talker.preprocess(input_ids=input_ids, input_embeds=input_embeds, **kwargs) if self.model_stage != "llm": embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids) return input_ids, embeds, {} @@ -474,174 +483,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp """vLLM V1 encoder profiling calls this; the inherited Protocol stub returns None.""" return self.get_multimodal_embeddings(**kwargs) - def _run_tts_request( - self, - *, - input_ids: torch.Tensor | None, - positions: torch.Tensor | None, - inputs_embeds: torch.Tensor | None, - talker_info: dict[str, Any], - device: torch.device, - ) -> tuple[dict[str, Any] | None, bool, bool]: - meta_info = talker_info.get("meta") if isinstance(talker_info.get("meta"), dict) else {} - if talker_info.get("native_duplex") is True: - # ids/hidden_states may be accumulated for the talker KV stream, - # but displayed transcript belongs to the current thinker segment. - tts_text = meta_info.get("native_duplex_segment_text", "") - else: - tts_text = talker_info.get("llm_output_text", "") - if isinstance(tts_text, list): - tts_text = ( - tts_text[-1] - if talker_info.get("native_duplex") is True and tts_text - else (tts_text[0] if tts_text else "") - ) - if not isinstance(tts_text, str): - tts_text = "" - - with torch.inference_mode(): - talker_result = self.talker( - input_ids=input_ids, - positions=positions, - inputs_embeds=inputs_embeds, - additional_information=talker_info, - ) - chunk_flags = getattr(self.talker, "_ar_last_chunk_flags", [True]) - chunk_is_last = bool(chunk_flags[-1]) if chunk_flags else True - turn_end_flags = getattr(self.talker, "_ar_turn_end_flags", [False]) - turn_ended = bool(turn_end_flags[-1]) if turn_end_flags else False - if talker_info.get("native_duplex") is True: - emitted_text = getattr(self.talker, "_ar_last_emitted_text", "") - tts_text = emitted_text if isinstance(emitted_text, str) else "" - - if not (isinstance(talker_result, tuple) and len(talker_result) == 2): - return None, chunk_is_last, turn_ended - - mel_spec, waveform = talker_result - mm_out: dict[str, Any] = {} - duplex_info = talker_info.get("duplex") if isinstance(talker_info.get("duplex"), dict) else {} - if talker_info.get("native_duplex") is True: - turn_id = duplex_info.get("turn_id") - if isinstance(turn_id, int): - mm_out["meta.duplex_turn_id"] = torch.tensor([turn_id], dtype=torch.int32, device=device) - epoch = duplex_info.get("epoch") - if isinstance(epoch, int): - mm_out["meta.duplex_epoch"] = torch.tensor([epoch], dtype=torch.int32, device=device) - mm_out["meta.tts_is_last_chunk"] = torch.tensor( - [int(chunk_is_last)], - dtype=torch.int32, - device=device, - ) - if talker_info.get("native_duplex") is True: - mm_out["meta.turn_end"] = torch.tensor( - [int(turn_ended)], - dtype=torch.int32, - device=device, - ) - if tts_text or talker_info.get("native_duplex") is True: - mm_out["meta.llm_output_text_utf8"] = torch.tensor( - list(tts_text.encode("utf-8")), - dtype=torch.uint8, - device=device, - ) - mm_out["meta.audio_text_total_chars"] = torch.tensor( - [len(tts_text)], - dtype=torch.int32, - device=device, - ) - if mel_spec is not None: - mm_out["mel_spec"] = [mel_spec] - if waveform is not None: - mm_out["model_outputs"] = [waveform] - elif mel_spec is not None: - mm_out["model_outputs"] = [mel_spec] - return mm_out, chunk_is_last, turn_ended - - @staticmethod - def _slice_tts_request_rows(value: torch.Tensor | None, start: int, end: int) -> torch.Tensor | None: - return value[start:end] if isinstance(value, torch.Tensor) else value - - def _run_batched_tts_requests( - self, - *, - input_ids: torch.Tensor | None, - positions: torch.Tensor | None, - inputs_embeds: torch.Tensor | None, - runtime_info: list[object], - request_token_spans: object, - num_tokens: int, - hidden_dim: int, - device: torch.device, - ) -> OmniOutput: - if not isinstance(request_token_spans, list) or len(request_token_spans) != len(runtime_info): - raise RuntimeError( - "MiniCPM-o 4.5 batched TTS requires one request_token_spans entry " - f"per request; got {request_token_spans!r} for {len(runtime_info)} requests" - ) - - request_ids: list[str] = [] - request_outputs: list[dict[str, Any]] = [] - row_chunk_flags = [True] * num_tokens - row_turn_end_flags = [False] * num_tokens - previous_end = 0 - for index, raw_info in enumerate(runtime_info): - if not isinstance(raw_info, dict): - raise RuntimeError(f"MiniCPM-o 4.5 TTS request metadata at index {index} is not a dict") - span = request_token_spans[index] - if not ( - isinstance(span, (list, tuple)) and len(span) == 2 and all(isinstance(value, int) for value in span) - ): - raise RuntimeError(f"Invalid MiniCPM-o 4.5 TTS token span at index {index}: {span!r}") - start, end = span - if start != previous_end or start < 0 or end < start or end > num_tokens: - raise RuntimeError( - f"Invalid MiniCPM-o 4.5 TTS token span at index {index}: " - f"expected start {previous_end} and 0 <= start <= end <= {num_tokens}, got {span!r}" - ) - previous_end = end - if start == end: - continue - mm_out, chunk_is_last, turn_ended = self._run_tts_request( - input_ids=self._slice_tts_request_rows(input_ids, start, end), - positions=self._slice_tts_request_rows(positions, start, end), - inputs_embeds=self._slice_tts_request_rows(inputs_embeds, start, end), - talker_info=raw_info, - device=device, - ) - row_chunk_flags[start:end] = [chunk_is_last] * (end - start) - row_turn_end_flags[start:end] = [turn_ended] * (end - start) - if mm_out is not None: - request_ids.append(get_stream_request_key(raw_info)) - request_outputs.append(mm_out) - # vLLM pads the flattened token tensor to a CUDA Graph capture size. - # The spans describe only scheduled request rows, so an uncovered tail - # is graph padding rather than a missing request. Contiguity checks - # above still reject gaps between real requests. - - self.talker._ar_last_chunk_flags = row_chunk_flags - self.talker._ar_turn_end_flags = row_turn_end_flags - dummy_hidden = torch.zeros(num_tokens, hidden_dim, device=device) - if not request_outputs: - return OmniOutput(text_hidden_states=dummy_hidden, multimodal_outputs=None) - - sparse_output: dict[str, Any] = { - "meta.req_id": request_ids, - "meta.sparse_audio": ["1"], - } - output_keys = set().union(*(output.keys() for output in request_outputs)) - for key in output_keys: - values: list[Any] = [] - for output in request_outputs: - value = output.get(key) - if isinstance(value, list) and len(value) == 1: - value = value[0] - if value is None: - dtype = torch.uint8 if key == "meta.llm_output_text_utf8" else torch.float32 - value = torch.empty(0, dtype=dtype, device=device) - values.append(value) - sparse_output[key] = values - return OmniOutput(text_hidden_states=dummy_hidden, multimodal_outputs=sparse_output) - def forward( self, input_ids: torch.Tensor, @@ -658,8 +499,9 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp Forward pass for MiniCPM-o Omni model. Workflow: - 1) LLM: multimodal thinker → hidden states and text tokens - 2) TTS: talker + Token2Wav → speech waveform + 1) Thinker (model_stage="llm"): Image / video / audio encoders + + 3D resampler + omni LLM → text + hidden states. + 2) Talker (model_stage="tts"): native MiniCPMTTS AR → codec deltas. """ if self.model_stage == "llm": # Normalize to batched inputs if caller provides 1D/2D unbatched tensors @@ -778,57 +620,24 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp multimodal_outputs=multimodal_outputs, ) - # Talker stage: runs TTS generation and its built-in Token2Wav vocoder. + # Talker stage: runner-owned native AR only. Waveform generation belongs + # to the separate Code2Wav stage. if self.model_stage == "tts": - if input_ids is not None: - num_tokens = input_ids.shape[0] - device = input_ids.device - elif inputs_embeds is not None: - num_tokens = inputs_embeds.shape[0] - device = inputs_embeds.device - else: - num_tokens = 1 - device = current_omni_platform.get_torch_device() - hidden_dim = self.config.hidden_size if hasattr(self.config, "hidden_size") else 2560 - - # Profile/dummy run: both input_ids and inputs_embeds are None. - # Note: SupportsMultiModal preprocessing converts input_ids to - # inputs_embeds, so input_ids=None alone does NOT indicate a dummy run. - if input_ids is None and inputs_embeds is None: - dummy_hidden = torch.zeros(num_tokens, hidden_dim, device=device) - return OmniOutput(text_hidden_states=dummy_hidden, multimodal_outputs=None) - - runtime_info = kwargs.get("runtime_additional_information") - request_token_spans = kwargs.get("request_token_spans") - if isinstance(runtime_info, list) and (len(runtime_info) > 1 or request_token_spans is not None): - return self._run_batched_tts_requests( - input_ids=input_ids, - positions=positions, - inputs_embeds=inputs_embeds, - runtime_info=runtime_info, - request_token_spans=request_token_spans, - num_tokens=num_tokens, - hidden_dim=hidden_dim, - device=device, - ) - talker_info = {} - if runtime_info and isinstance(runtime_info, list) and len(runtime_info) > 0: - talker_info = runtime_info[0] if isinstance(runtime_info[0], dict) else {} - dummy_hidden = torch.zeros(num_tokens, hidden_dim, device=device) - mm_out, _, _ = self._run_tts_request( + return self.talker( input_ids=input_ids, positions=positions, + intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, - talker_info=talker_info, - device=device, + **kwargs, ) - if mm_out is not None: - return OmniOutput(text_hidden_states=dummy_hidden, multimodal_outputs=mm_out) - - return OmniOutput(text_hidden_states=dummy_hidden, multimodal_outputs=None) raise ValueError(f"Unsupported model stage: {self.model_stage}") + def make_omni_output(self, model_outputs, **kwargs): + if self.model_stage != "tts": + return model_outputs + return self.talker.make_omni_output(model_outputs, **kwargs) + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput) -> torch.Tensor | None: # Handle OmniOutput type if isinstance(hidden_states, OmniOutput): @@ -1261,7 +1070,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp # MiniCPM-o checkpoint prefixes → stage mapping: # thinker: vpm, resampler, llm, apm, audio_projection_layer - # talker: tts (ConditionalChatTTS) + # talker: tts (native MiniCPMTTS AR codec producer) for k, v in weights: if k.startswith(("vpm.", "resampler.", "llm.", "apm.", "audio_projection_layer.")): thinker_weights.append((k, v)) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 73570ef74..09bbb2771 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -2,246 +2,129 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from: # https://huggingface.co/openbmb/MiniCPM-o-4_5/blob/main/modeling_minicpmo.py -"""MiniCPM-o 4.5 Talker + Token2Wav: MiniCPMTTS with hidden_text_merge condition. +"""MiniCPM-o 4.5 native autoregressive Talker. Pipeline: 1. Receive thinker hidden_states + full token IDs via additional_information 2. Extract tts_bos..tts_eos region 3. Build condition: emb_text(tokens) + projector_semantic(hidden) (hidden_text_merge) - 4. Run MiniCPMTTS.generate() -> discrete audio tokens - 5. Run Token2wav(tokens) -> waveform bytes -> numpy array + 4. Continuously generate request-aligned discrete audio-code deltas """ -import hashlib -import io -import logging -import os -import sys -import tempfile -import threading -import time -from collections import OrderedDict from collections.abc import Iterable -from dataclasses import dataclass -from importlib import import_module from typing import Any -import numpy as np -import soundfile as sf import torch import torch.nn as nn import torch.nn.functional as F +from transformers import LlamaConfig from vllm.config import VllmConfig from vllm.model_executor.models.interfaces import SupportsPP -from vllm.v1.outputs import SamplerOutput +from vllm.model_executor.models.llama import LlamaModel +from vllm.model_executor.models.utils import maybe_prefix +from vllm.v1.sample.sampler import Sampler -from vllm_omni.experimental.fullduplex.engine.intermediate import get_stream_request_key, get_tts_handoff -from vllm_omni.model_executor.model_loader.weight_utils import ( - download_weights_from_hf_specific, -) +from vllm_omni.experimental.fullduplex.engine.intermediate import get_tts_handoff +from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform -# The external vocoder hard-codes CUDA placement. Ascend uses the in-tree -# adapter, with the external package retained as a fallback for compatibility. -_stepaudio2_import_error: ImportError | None = None -if current_omni_platform.is_npu(): - try: - from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_token2wav import ( - MiniCPMO45Token2wav as _Token2wav, - ) - - _token2wav_backend = "step_audio2_core" - except ImportError as e: - _stepaudio2_import_error = e - try: - from stepaudio2 import Token2wav as _Token2wav - - _token2wav_backend = "stepaudio2_pkg" - except ImportError as fallback_error: - _Token2wav = None - _token2wav_backend = None - _stepaudio2_import_error = fallback_error -else: - try: - from stepaudio2 import Token2wav as _Token2wav - - _token2wav_backend = "stepaudio2_pkg" - except ImportError as e: - _Token2wav = None - _token2wav_backend = None - _stepaudio2_import_error = e - -_stepaudio2_available = _Token2wav is not None - -logger = logging.getLogger(__name__) +_REPETITION_WINDOW = 16 +_MIN_AUDIO_TOKENS = 64 +_MAX_AUDIO_TOKENS = 2048 +_AUDIO_TOKENS_PER_TEXT_TOKEN = 10 +# Codec-token sampling happens inside the model; vLLM sampling parameters +# only choose the Talker's binary continue/stop row. +_CODEC_SEED = 42 +_CODEC_TEMPERATURE = 0.8 +_CODEC_TOP_K = 25 +_CODEC_TOP_P = 0.85 +_CODEC_REPETITION_PENALTY = 1.05 +_CODEC_MIN_TOKENS = 50 -@dataclass(slots=True) -class MiniCPMO45TTSRuntimeConfig: - """Internal MiniCPM-o 4.5 Talker runtime defaults. +def _max_audio_tokens(condition_tokens: int) -> int: + """Bound codec generation with a conservative text-length estimate. - These values are deliberately not process environment knobs. If a value - needs to become user configurable, route it through stage/model config with - an explicit API contract instead of adding another ad hoc env var. + EOS is masked for the first 50 steps, so a direct ``text_tokens * 10`` + limit can terminate short responses before EOS is eligible. The 2048 + ceiling matches the checkpoint's native generation default and keeps the + sequence within the Talker's 4096-position context. """ - - token2wav_n_timesteps: int = 10 - tts_dtype: torch.dtype = torch.float32 - token2wav_autocast_dtype: torch.dtype | None = None - use_direct_token2wav: bool = True - ref_audio_file_cache_size: int = 16 - max_token_ratio: int = 32 - min_max_new_tokens: int = 256 - hard_max_new_tokens: int = 16384 - min_new_tokens: int = 50 - streaming_generator_chunk: int = 25 - streaming_vocoder_threshold: int = 2500 - streaming_vocoder_chunk: int = 50 - - -def _install_torchaudio_soundfile_shim() -> None: - """Monkey-patch torchaudio.load to use soundfile instead of the default - torchcodec backend, which requires libtorchcodec/ffmpeg shared libs that - may be missing on the deployment machine.""" - try: - import torchaudio - - if getattr(torchaudio, "_soundfile_shim_installed", False): - return - _orig_load = torchaudio.load - - def _patched_load(uri, *args, **kwargs): - try: - return _orig_load(uri, *args, **kwargs) - except Exception: - import numpy as _np - import soundfile as _sf - - data, sr = _sf.read(uri, dtype="float32", always_2d=True) - wav = torch.from_numpy(_np.ascontiguousarray(data.T)) - return wav, sr - - torchaudio.load = _patched_load - torchaudio._soundfile_shim_installed = True - logger.info("Installed torchaudio.load soundfile shim") - except Exception as _e: - logger.warning("Could not install torchaudio shim: %s", _e) - - -_install_torchaudio_soundfile_shim() - - -class _TalkerTurnState: - """Per-turn talker continuity for native duplex. - - The official duplex implementation runs ONE TTS stream per spoken turn: - each 1s unit appends its condition to the same KV (carried - past_key_values + text_start_pos) and the token2wav stream caches and - token buffer persist across units, reset only at turn end. Synthesizing - units as independent utterances resets prosody every second and garbles - the reply. - """ - - __slots__ = ( - "past_key_values", - "text_start_pos", - "token2wav_buffer", - "prompt_wav_path", - "temp_prompt_wav_path", - "epoch", - "turn_id", - "pending_text", - "stream_cache", - "hift_cache_dict", - "vocoder_initialized", + return max( + _MIN_AUDIO_TOKENS, + min(_MAX_AUDIO_TOKENS, condition_tokens * _AUDIO_TOKENS_PER_TEXT_TOKEN), ) - def __init__( - self, - prompt_wav_path, - temp_prompt_wav_path, - *, - epoch: int | None = None, - turn_id: int | None = None, - ): - self.past_key_values = None - self.text_start_pos = 0 - # Official seeds each turn's vocoder buffer with three silence - # tokens so the first synthesized window does not directly abut the - # reference-audio prompt cache (audible ref-voice bleed otherwise). - self.token2wav_buffer: list[int] = [_T2W_SILENCE_TOKEN] * 3 - self.prompt_wav_path = prompt_wav_path - self.temp_prompt_wav_path = temp_prompt_wav_path - self.epoch = epoch - self.turn_id = turn_id - self.pending_text = "" - self.stream_cache = None - self.hift_cache_dict: dict[str, Any] = {} - self.vocoder_initialized = False + +def _restore_weight_norm_weight(weight_g: torch.Tensor, weight_v: torch.Tensor) -> torch.Tensor: + """Materialize ``weight_norm(..., dim=0)`` checkpoint parameters.""" + row_norm = torch.linalg.vector_norm( + weight_v.float(), + dim=1, + keepdim=True, + ).clamp_min(torch.finfo(torch.float32).tiny) + return weight_v * (weight_g.to(dtype=weight_v.dtype) / row_norm.to(dtype=weight_v.dtype)) -def _queue_native_duplex_segment_text(state: _TalkerTurnState, text: object) -> None: - if isinstance(text, str) and text: - state.pending_text += text - - -def _drain_native_duplex_emitted_text(state: _TalkerTurnState, *, has_audio: bool) -> str: - if not has_audio: - return "" - text = state.pending_text - state.pending_text = "" - return text - - -_T2W_SILENCE_TOKEN = 4218 -_NATIVE_DUPLEX_UNIT_AUDIO_SAMPLES = 24_000 - - -def _native_duplex_unit_waveform( - waveforms: Iterable[torch.Tensor], +def _apply_repetition_penalty( + logits: torch.Tensor, + history: torch.Tensor, *, - turn_end: bool, - target_samples: int = _NATIVE_DUPLEX_UNIT_AUDIO_SAMPLES, -) -> torch.Tensor | None: - pieces = [torch.as_tensor(waveform).reshape(-1).cpu().contiguous() for waveform in waveforms] - pieces = [piece for piece in pieces if piece.numel() > 0] - if not pieces: - return None - waveform = pieces[0] if len(pieces) == 1 else torch.cat(pieces, dim=0) - if not turn_end and waveform.numel() < target_samples: - waveform = F.pad(waveform, (target_samples - waveform.numel(), 0)) - return waveform + penalty: float, + window_size: int, +) -> torch.Tensor: + """Match MiniCPMTTS' frequency-aware repetition penalty.""" + if penalty == 1.0 or history.numel() == 0: + return logits + recent = history.reshape(-1)[-window_size:].to(device=logits.device, dtype=torch.long) + frequencies = torch.bincount(recent, minlength=logits.shape[-1]).to(dtype=logits.dtype) + alpha = torch.pow(torch.as_tensor(penalty, device=logits.device, dtype=logits.dtype), frequencies) + return torch.where(logits < 0, logits * alpha, logits / alpha) -def _soundfile_patched_save(orig_save): - def _patched_save(uri, src, sample_rate, **kw): - kw.pop("backend", None) - if hasattr(uri, "write"): - sf.write(uri, src.cpu().numpy().T, sample_rate, format="WAV") - return - return orig_save(uri, src, sample_rate, backend="soundfile", **kw) - - return _patched_save +def _apply_top_k_top_p( + logits: torch.Tensor, + *, + top_k: int | None, + top_p: float | None, + min_tokens_to_keep: int = 3, +) -> torch.Tensor: + """Apply the same candidate floors as the upstream Transformers warpers.""" + filtered = logits.clone() + vocab_size = filtered.shape[-1] + # MiniCPM-o's gen_logits() appends TopPLogitsWarper before + # TopKLogitsWarper. The order is observable for fixed-seed sampling. + if top_p is not None and 0.0 < top_p < 1.0: + sorted_logits, sorted_indices = torch.sort(filtered, descending=False, dim=-1) + cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1) + remove = cumulative_probs <= (1.0 - float(top_p)) + remove[..., -min_tokens_to_keep:] = False + remove = remove.scatter(-1, sorted_indices, remove) + filtered.masked_fill_(remove, float("-inf")) + if top_k is not None and top_k > 0: + keep = min(vocab_size, max(int(top_k), min_tokens_to_keep)) + threshold = torch.topk(filtered, keep, dim=-1).values[..., -1, None] + filtered.masked_fill_(filtered < threshold, float("-inf")) + return filtered -def _torch_clone_recursive(obj): - if isinstance(obj, torch.Tensor): - return obj.clone() - if isinstance(obj, dict): - return {k: _torch_clone_recursive(v) for k, v in obj.items()} - if isinstance(obj, list): - return [_torch_clone_recursive(v) for v in obj] - if isinstance(obj, tuple): - return tuple(_torch_clone_recursive(v) for v in obj) - return obj +class _MiniCPMTTSProjector(nn.Module): + """Checkpoint-compatible hidden-state projector used by MiniCPMTTS.""" + + def __init__(self, input_size: int, hidden_size: int): + super().__init__() + self.linear1 = nn.Linear(input_size, hidden_size, bias=True) + self.relu = nn.ReLU() + self.linear2 = nn.Linear(hidden_size, hidden_size, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.linear2(self.relu(self.linear1(hidden_states))) class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): - """MiniCPM-o 4.5 Talker: MiniCPMTTS + Token2wav in a single forward pass.""" + """Runner-owned MiniCPM-o 4.5 Talker that emits codec tokens only.""" - # llm2tts hands the FULL accumulated condition per handoff (the runner's - # resume-prefill path REPLACES the streaming buffer, so runner-side - # accumulation is lossy); the per-turn state consumes it by cursor. + requires_request_sample_eligibility = True def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -250,1366 +133,360 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): config: MiniCPMOConfig = vllm_config.model_config.hf_config self.config = config self.vllm_config = vllm_config - self._runtime_config = MiniCPMO45TTSRuntimeConfig() - - self.tts = None - self.audio_tokenizer = None - self._assets_loaded = False - self._model_path: str | None = None - self._stream_gens: dict[str, Any] = {} - self._talker_turn_states: dict[str, _TalkerTurnState] = {} - # Consumed-cursor into the accumulated handoff condition for the - # currently open spoken turn. - self._talker_consumed_tokens: dict[str, int] = {} - self._talker_request_keys: dict[str, str] = {} - self._t2w_base_caches: dict[str, tuple[Any, Any]] = {} - self._token2wav_state_lock = threading.RLock() - self._ar_last_chunk_flags: list[bool] = [True] - self._ar_turn_end_flags: list[bool] = [False] - self._ar_last_emitted_text = "" - self._text_tokenizer = None + self._batch_stop_logits: torch.Tensor | None = None + self._request_generators: dict[str, torch.Generator] = {} + self._request_audio_states: dict[str, dict[str, Any]] = {} + self._deferred_cleanup_ids: set[str] = set() tts_config = getattr(config, "tts_config", None) + if tts_config is None and getattr(config, "model_type", None) == "minicpmtts": + tts_config = config if tts_config is not None: self._tts_config = tts_config - self._tts_bos_id = self._config_token_id(tts_config, "audio_bos_token_id") - self._text_eos_id = self._config_token_id(tts_config, "text_eos_token_id") + self._tts_bos_id = getattr(tts_config, "audio_bos_token_id", 151687) + self._text_eos_id = getattr(tts_config, "text_eos_token_id", 151692) self._num_audio_tokens = getattr(tts_config, "num_audio_tokens", 6562) self._hidden_size = getattr(tts_config, "hidden_size", 768) self._normalize = getattr(tts_config, "normalize_projected_hidden", True) + self._codec_seed = int(getattr(tts_config, "seed", _CODEC_SEED)) + self._codec_temperature = float(getattr(tts_config, "temperature", _CODEC_TEMPERATURE)) + self._codec_top_k = int(getattr(tts_config, "top_k", _CODEC_TOP_K)) + self._codec_top_p = float(getattr(tts_config, "top_p", _CODEC_TOP_P)) + self._codec_repetition_penalty = float(getattr(tts_config, "repetition_penalty", _CODEC_REPETITION_PENALTY)) + self._codec_min_tokens = int(getattr(tts_config, "min_new_tokens", _CODEC_MIN_TOKENS)) else: self._tts_config = None - def _tts_runtime_config(self) -> MiniCPMO45TTSRuntimeConfig: - cfg = getattr(self, "_runtime_config", None) - if cfg is None: - cfg = MiniCPMO45TTSRuntimeConfig() - self._runtime_config = cfg - return cfg + self.has_preprocess = True + self.has_postprocess = False + self.gpu_resident_buffer_keys: set[tuple[str, str]] = { + ("audio_codes", "current"), + ("audio_codes", "accumulated"), + } + self._init_native_talker(prefix) - @staticmethod - def _config_token_id(config: Any, attr: str) -> int: - value = getattr(config, attr, None) - if isinstance(value, (list, tuple)): - value = value[0] if value else None - if value is None: - raise ValueError(f"MiniCPM-o 4.5 TTS config missing required {attr}") - return int(value) - - def _get_text_tokenizer(self) -> Any: - tokenizer = getattr(self, "_text_tokenizer", None) - if tokenizer is not None: - return tokenizer - from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config - - tokenizer = cached_tokenizer_from_config(self.vllm_config.model_config) - self._text_tokenizer = tokenizer - return tokenizer - - def _tokenizer_token_id(self, token: str) -> int | None: - tokenizer = self._get_text_tokenizer() - unk_token_id = getattr(tokenizer, "unk_token_id", None) - convert = getattr(tokenizer, "convert_tokens_to_ids", None) - if callable(convert): - value = convert(token) - if isinstance(value, list): - value = value[0] if len(value) == 1 else None - try: - candidate = int(value) - except (TypeError, ValueError): - candidate = None - if candidate is not None and candidate >= 0 and candidate != unk_token_id: - return candidate - encode = getattr(tokenizer, "encode", None) - if callable(encode): - ids = list(encode(token, add_special_tokens=False)) - if len(ids) == 1: - candidate = int(ids[0]) - if candidate >= 0 and candidate != unk_token_id: - return candidate - return None - - def _scheduler_eos_token_id(self) -> int: - eos_raw = getattr(self.config, "eos_token_id", None) - if isinstance(eos_raw, (list, tuple)): - eos_raw = eos_raw[0] if eos_raw else None - if eos_raw is not None: - return int(eos_raw) - eos_id = self._tokenizer_token_id("<|im_end|>") - if eos_id is None: + def _init_native_talker(self, prefix: str) -> None: + if self._tts_config is None: + raise ValueError("MiniCPM-o continuous Talker requires tts_config") + cfg = self._tts_config + if int(getattr(cfg, "num_vq", 1)) != 1: raise ValueError( - "MiniCPM-o 4.5 TTS scheduler EOS requires config.eos_token_id or tokenizer-defined <|im_end|>" + "MiniCPM-o continuous Talker currently requires num_vq=1; " + f"checkpoint reports {getattr(cfg, 'num_vq', None)}" ) - return eos_id - - def _lazy_init_tts(self): - if self._assets_loaded or self._tts_config is None: - return - self._assets_loaded = True - try: - model_path = download_weights_from_hf_specific(self.vllm_config.model_config.model, None, ["*"]) - self._model_path = model_path - if model_path not in sys.path: - sys.path.insert(0, model_path) - from transformers import AutoImageProcessor - from transformers.dynamic_module_utils import get_class_from_dynamic_module - - # The remote processing module registers an image processor by - # string, which transformers>=5 rejects. The standalone talker does - # not use that registration, so ignore only the string form while - # importing MiniCPMTTS and restore the global method immediately. - original_register = AutoImageProcessor.register - AutoImageProcessor.register = ( # type: ignore[method-assign] - lambda key, *a, **k: None if isinstance(key, str) else original_register(key, *a, **k) - ) - try: - MiniCPMTTS = get_class_from_dynamic_module("modeling_minicpmo.MiniCPMTTS", model_path) - finally: - AutoImageProcessor.register = original_register # type: ignore[method-assign] - - prev_dtype = torch.get_default_dtype() - torch.set_default_dtype(torch.float32) - try: - for name, default in ( - ("top_p", 0.85), - ("top_k", 25), - ("repetition_penalty", 1.05), - ("temperature", 0.8), - ): - if not hasattr(self._tts_config, name): - setattr(self._tts_config, name, default) - self._tts_config.attn_implementation = "sdpa" - self.tts_obj = MiniCPMTTS(config=self._tts_config, audio_tokenizer=None) - finally: - torch.set_default_dtype(prev_dtype) - tts_module = import_module(self.tts_obj.__class__.__module__) - - def get_tts_module_attr(name: str): - return self.tts_obj.generate.__globals__.get(name) or getattr(tts_module, name, None) - - self._tts_sampling_params_cls = get_tts_module_attr("TTSSamplingParams") - self._tts_gen_logits = get_tts_module_attr("gen_logits") - self._tts_parametrize = get_tts_module_attr("P") - self._tts_streaming_generator_cls = get_tts_module_attr("TTSStreamingGenerator") - self.emb_text = self.tts_obj.emb_text - self.projector_semantic = self.tts_obj.projector_semantic - - token2wav_dir = os.path.join(model_path, "assets", "token2wav") - if os.path.isdir(token2wav_dir): - if not _stepaudio2_available: - raise ImportError( - "MiniCPM-o 4.5 token2wav stage requires the stepaudio2 package, " - "and all of its runtime dependencies." - ) from _stepaudio2_import_error - self._token2wav_n_timesteps = self._tts_runtime_config().token2wav_n_timesteps - prev_dtype2 = torch.get_default_dtype() - torch.set_default_dtype(torch.float32) - try: - self.audio_tokenizer = _Token2wav( - token2wav_dir, - float16=False, - n_timesteps=self._token2wav_n_timesteps, - ) - finally: - torch.set_default_dtype(prev_dtype2) - self.tts_obj.audio_tokenizer = self.audio_tokenizer - logger.info( - "Loaded Token2wav from %s (backend=%s, n_timesteps=%d)", - token2wav_dir, - _token2wav_backend, - self._token2wav_n_timesteps, - ) - except ImportError: - # Surface missing dependencies directly so users can act on them - # instead of getting a silent None waveform downstream. - raise - except Exception as e: - logger.error("Failed to init 4.5 TTS: %s", e, exc_info=True) - - def _build_tts_sampling_params(self): - params_cls = getattr(self, "_tts_sampling_params_cls", None) - if params_cls is None or not hasattr(self, "tts_obj"): - return None - - tts = self.tts_obj - - top_p = getattr(tts, "top_p", getattr(tts.config, "top_p", 0.85)) - top_k = getattr(tts, "top_k", getattr(tts.config, "top_k", 25)) - repetition_penalty = getattr( - tts, - "repetition_penalty", - getattr(tts.config, "repetition_penalty", 1.05), + llama_config = LlamaConfig( + vocab_size=32000, + hidden_size=int(cfg.hidden_size), + intermediate_size=int(cfg.intermediate_size), + num_hidden_layers=int(cfg.num_hidden_layers), + num_attention_heads=int(cfg.num_attention_heads), + num_key_value_heads=int(cfg.num_key_value_heads), + hidden_act=getattr(cfg, "hidden_act", "silu"), + max_position_embeddings=int(cfg.max_position_embeddings), + rms_norm_eps=float(getattr(cfg, "rms_norm_eps", 1e-6)), + tie_word_embeddings=False, ) - temperature = getattr(tts.config, "temperature", 0.8) - - return params_cls( - top_p=None if top_p is not None and top_p >= 1.0 else top_p, - top_k=None if top_k is not None and top_k <= 0 else top_k, - repetition_penalty=repetition_penalty, - temperature=temperature, + talker_config = self.vllm_config.with_hf_config(llama_config, architectures=["LlamaForCausalLM"]) + talker_config.model_config.hf_text_config = llama_config + self.tts_model = LlamaModel( + vllm_config=talker_config, + prefix=maybe_prefix(prefix, "tts_obj.model"), ) - - def _target_tts_dtype(self) -> torch.dtype: - return self._tts_runtime_config().tts_dtype - - def _token2wav_autocast_context(self): - dtype = self._tts_runtime_config().token2wav_autocast_dtype - device_type = str(current_omni_platform.device_type) - if dtype is None: - return ( - current_omni_platform.create_autocast_context( - device_type=device_type, - dtype=torch.float32, - enabled=False, - ), - "off", - ) - if dtype is torch.bfloat16: - return ( - current_omni_platform.create_autocast_context( - device_type=device_type, - dtype=torch.bfloat16, - enabled=True, - ), - "bf16", - ) - if dtype is torch.float16: - return ( - current_omni_platform.create_autocast_context( - device_type=device_type, - dtype=torch.float16, - enabled=True, - ), - "fp16", - ) - raise ValueError("MiniCPM-o 4.5 token2wav autocast only supports None, bfloat16, or float16") - - def _should_use_direct_token2wav(self) -> bool: - return self._tts_runtime_config().use_direct_token2wav - - def _should_stream_output(self, info: dict[str, Any] | None = None) -> bool: - if isinstance(info, dict): - for key in ("stream_output", "native_duplex"): - value = info.get(key) - if isinstance(value, bool): - return value - return False - - def _token2wav_prompt_cache_key(self, prompt_wav: str | None) -> str | None: - return os.path.abspath(prompt_wav) if prompt_wav else None - - def _reset_token2wav_cache_if_needed(self, prompt_wav: str | None) -> None: - token2wav = self.audio_tokenizer - if token2wav is None: - return - cache_key = self._token2wav_prompt_cache_key(prompt_wav) - if getattr(self, "_token2wav_prompt_cache_id", None) != cache_key: - token2wav.cache = None - self._token2wav_prompt_cache_id = cache_key - - def _normalize_ref_audio_tensor(self, ref_audio) -> np.ndarray | None: - if ref_audio is None: - return None - if isinstance(ref_audio, torch.Tensor): - waveform = ref_audio.detach().float().cpu().numpy() - else: - waveform = np.asarray(ref_audio, dtype=np.float32) - if waveform.ndim > 1: - if waveform.shape[0] <= 2 and waveform.shape[-1] > waveform.shape[0]: - waveform = waveform.mean(axis=0) - else: - waveform = waveform.mean(axis=-1) - waveform = np.asarray(waveform, dtype=np.float32).reshape(-1) - return waveform if waveform.size else None - - def _write_ref_audio_prompt_wav(self, ref_audio, ref_audio_sr: int | None) -> str | None: - waveform = self._normalize_ref_audio_tensor(ref_audio) - if waveform is None: - return None - sample_rate = int(ref_audio_sr or 24000) - cache_size = self._tts_runtime_config().ref_audio_file_cache_size - if cache_size > 0: - digest = hashlib.sha256() - digest.update(str(sample_rate).encode("ascii")) - digest.update(waveform.tobytes()) - cache_key = digest.hexdigest() - cache = getattr(self, "_ref_audio_prompt_files", None) - if cache is None: - cache = OrderedDict() - self._ref_audio_prompt_files = cache - cached_path = cache.get(cache_key) - if cached_path and os.path.exists(cached_path): - cache.move_to_end(cache_key) - return cached_path - - tmp_path = os.path.join(tempfile.gettempdir(), f"minicpmo45_ref_{cache_key[:24]}_{sample_rate}.wav") - if not os.path.exists(tmp_path): - sf.write(tmp_path, waveform, sample_rate, format="WAV") - cache[cache_key] = tmp_path - cache.move_to_end(cache_key) - while len(cache) > cache_size: - _, evicted_path = cache.popitem(last=False) - with self._token2wav_lock(): - self._t2w_base_caches.pop(evicted_path, None) - try: - os.unlink(evicted_path) - except OSError: - pass - return tmp_path - - tmp = tempfile.NamedTemporaryFile(prefix="minicpmo45_ref_", suffix=".wav", delete=False) - tmp_path = tmp.name - tmp.close() - sf.write(tmp_path, waveform, sample_rate, format="WAV") - return tmp_path - - def _is_cached_ref_audio_prompt_wav(self, prompt_wav: str | None) -> bool: - cache = getattr(self, "_ref_audio_prompt_files", None) - return bool(prompt_wav and cache and prompt_wav in cache.values()) - - def _run_token2wav_direct( - self, - generated_speech_tokens: list[int] | torch.Tensor, - prompt_wav: str | None, - ) -> tuple[torch.Tensor, int]: - """Run Token2wav without WAV encode/decode round-tripping. - - ``stepaudio2.Token2wav.__call__`` renders the GPU waveform to a WAV - BytesIO object and the vLLM-Omni adapter immediately decodes it back to - a numpy waveform. The engine already expects a float waveform tensor, - so keep the same flow/HIFT computation and return the waveform directly. - This also makes the configured ``n_timesteps`` apply to one-shot - inference; upstream currently hard-codes ``10`` in ``__call__``. - """ - token2wav = self.audio_tokenizer - if token2wav is None: - raise RuntimeError("Token2wav is not initialized") - required_attrs = ("_prepare_prompt", "flow", "hift") - if any(not hasattr(token2wav, attr) for attr in required_attrs): - raise RuntimeError("Token2wav direct path is incompatible with the installed stepaudio2 package") - - self._reset_token2wav_cache_if_needed(prompt_wav) - if token2wav.cache is None: - token2wav.cache = token2wav._prepare_prompt(prompt_wav) - prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens = token2wav.cache - - device = prompt_speech_tokens.device - if isinstance(generated_speech_tokens, torch.Tensor): - generated = generated_speech_tokens - if generated.ndim == 1: - generated = generated.unsqueeze(0) - elif generated.ndim == 3 and generated.shape[-1] == 1: - generated = generated.squeeze(-1) - generated = generated.to(device=device, dtype=torch.int32) - else: - generated = torch.tensor([generated_speech_tokens], dtype=torch.int32, device=device) - generated_lens = torch.tensor([generated.shape[1]], dtype=torch.int32, device=device) - mel = token2wav.flow.inference( - generated, - generated_lens, - prompt_speech_tokens, - prompt_speech_tokens_lens, - prompt_mels, - prompt_mels_lens, - spk_emb, - self._token2wav_n_timesteps, + self.emb_text = nn.Embedding(int(cfg.num_text_tokens), int(cfg.hidden_size)) + self.projector_semantic = _MiniCPMTTSProjector(int(cfg.llm_dim), int(cfg.hidden_size)) + self.emb_code = nn.ModuleList( + [nn.Embedding(int(cfg.num_audio_tokens), int(cfg.hidden_size)) for _ in range(int(cfg.num_vq))] ) - wav, _ = token2wav.hift(speech_feat=mel) - waveform = wav.squeeze(0).detach().float().reshape(-1).cpu().contiguous() - return waveform, 24000 + self.head_code = nn.ModuleList( + [nn.Linear(int(cfg.hidden_size), int(cfg.num_audio_tokens), bias=False) for _ in range(int(cfg.num_vq))] + ) + self.make_empty_intermediate_tensors = self.tts_model.make_empty_intermediate_tensors - def _build_tts_condition_embeds( + def _build_condition_embeddings( self, tts_token_ids: torch.Tensor, tts_hidden_states: torch.Tensor, ) -> torch.Tensor: - tts = self.tts_obj - device = tts.emb_text.weight.device - dtype = tts.emb_text.weight.dtype - llm_embeds = tts.emb_text(tts_token_ids.to(device)) - hidden_embeds = tts.projector_semantic(tts_hidden_states.to(device=device, dtype=dtype)) - if getattr(tts.config, "normalize_projected_hidden", False): - hidden_embeds = F.normalize(hidden_embeds, p=2, dim=-1) - return llm_embeds + hidden_embeds - - def _normalize_tts_handoff_tensors( - self, - tts_token_ids: Any, - tts_hidden_states: Any, - ) -> tuple[torch.Tensor, torch.Tensor]: - token_ids = torch.as_tensor(tts_token_ids, dtype=torch.long) - if token_ids.ndim == 0: - token_ids = token_ids.reshape(1) - elif token_ids.ndim > 1: - token_ids = token_ids.reshape(-1) - - hidden_states = torch.as_tensor(tts_hidden_states, dtype=torch.float32) - if hidden_states.ndim == 1: - hidden_states = hidden_states.unsqueeze(0) - elif hidden_states.ndim == 3 and hidden_states.shape[0] == 1: - hidden_states = hidden_states.squeeze(0) - elif hidden_states.ndim != 2: - hidden_states = hidden_states.reshape(token_ids.numel(), -1) - - if hidden_states.shape[0] != token_ids.numel(): + device = self.emb_text.weight.device + dtype = self.emb_text.weight.dtype + token_ids = tts_token_ids.to(device=device, dtype=torch.long).reshape(-1) + hidden = tts_hidden_states.to(device=device, dtype=dtype) + if hidden.shape[0] != token_ids.shape[0] and token_ids.shape[0] != 1: raise ValueError( - "MiniCPM-o 4.5 TTS handoff has mismatched token/hidden lengths: " - f"tokens={token_ids.numel()} hidden_rows={hidden_states.shape[0]}" + "MiniCPM-o Talker condition length mismatch: " + f"token_ids={token_ids.shape[0]} hidden_states={hidden.shape[0]}" ) - return token_ids.contiguous(), hidden_states.contiguous() - - def _resolve_prompt_wav_path(self, ref_audio, ref_audio_sr: int | None) -> tuple[str | None, str | None]: - temp_prompt_wav_path = self._write_ref_audio_prompt_wav(ref_audio, ref_audio_sr) - if temp_prompt_wav_path is not None: - return temp_prompt_wav_path, temp_prompt_wav_path - if (model_path := getattr(self, "_model_path", None)) is not None: - default_ref = os.path.join(model_path, "assets", "HT_ref_audio.wav") - if os.path.exists(default_ref): - return default_ref, None - return None, None - - def _max_tts_tokens_for_text(self, num_text: int) -> tuple[int, int]: - cfg = self._tts_runtime_config() - max_new_token = min( - cfg.hard_max_new_tokens, - max(cfg.min_max_new_tokens, num_text * cfg.max_token_ratio), - ) - return cfg.min_new_tokens, max_new_token - - def _stream_request_key(self, info: dict[str, Any]) -> str: - return get_stream_request_key(info) - - @staticmethod - def _coerce_request_key(value: Any) -> str | None: - if isinstance(value, (list, tuple)): - value = value[0] if value else None - if isinstance(value, bytes): - value = value.decode("utf-8", errors="replace") - if value is None: - return None - text = str(value) - return text if text else None - - @staticmethod - def _coerce_epoch(value: Any) -> int | None: - try: - return int(value) if value is not None else None - except (TypeError, ValueError): - return None - - _coerce_turn_id = _coerce_epoch - - def _stream_epoch(self, info: dict[str, Any]) -> int | None: - duplex = info.get("duplex") - if isinstance(duplex, dict): - epoch = self._coerce_epoch(duplex.get("epoch")) - if epoch is not None: - return epoch - meta = info.get("meta") - if isinstance(meta, dict): - epoch = self._coerce_epoch(meta.get("epoch")) - if epoch is not None: - return epoch - return self._coerce_epoch(info.get("epoch")) - - def _stream_turn_id(self, info: dict[str, Any]) -> int | None: - duplex = info.get("duplex") - if isinstance(duplex, dict): - turn_id = self._coerce_turn_id(duplex.get("model_turn_id")) - if turn_id is not None: - return turn_id - turn_id = self._coerce_turn_id(duplex.get("turn_id")) - if turn_id is not None: - return turn_id - meta = info.get("meta") - if isinstance(meta, dict): - turn_id = self._coerce_turn_id(meta.get("turn_id")) - if turn_id is not None: - return turn_id - return self._coerce_turn_id(info.get("turn_id")) - - def _remember_talker_request_key(self, info: dict[str, Any], key: str) -> None: - aliases = { - self._coerce_request_key(info.get("request_id")), - self._coerce_request_key(info.get("_omni_req_id")), - } - duplex = info.get("duplex") - if isinstance(duplex, dict): - aliases.add(self._coerce_request_key(duplex.get("request_id"))) - request_keys = getattr(self, "_talker_request_keys", None) - if request_keys is None: - request_keys = {} - self._talker_request_keys = request_keys - for alias in aliases: - if alias and alias != key: - request_keys[alias] = key - - def _empty_audio_chunk(self) -> torch.Tensor: - return torch.zeros((0,), dtype=torch.float32) - - @staticmethod - def _extract_tts_handoff(info: dict[str, Any]) -> tuple[Any, Any]: - return get_tts_handoff(info) - - def _native_duplex_input_ends_turn(self, info: dict[str, Any]) -> bool: - if info.get("native_duplex") is not True: - return False - meta = info.get("meta") - if not isinstance(meta, dict): - meta = {} - if bool(info.get("turn_end") or info.get("end_of_turn") or meta.get("turn_end") or meta.get("end_of_turn")): - return True - turn_eos_id = self._coerce_epoch(meta.get("turn_eos_token_id")) - if turn_eos_id is None: - return False - tts_token_ids, _ = self._extract_tts_handoff(info) - if isinstance(tts_token_ids, torch.Tensor): - return bool((tts_token_ids == turn_eos_id).any().item()) - if isinstance(tts_token_ids, np.ndarray): - return bool(np.any(tts_token_ids == turn_eos_id)) - if isinstance(tts_token_ids, (list, tuple)): - return turn_eos_id in tts_token_ids - return False - - def _t2w_pre_lookahead(self) -> int: - flow = getattr(self.audio_tokenizer, "flow", None) - try: - return int(getattr(flow, "pre_lookahead_len", 3) or 3) - except (TypeError, ValueError): - return 3 - - def _token2wav_lock(self) -> threading.RLock: - lock = getattr(self, "_token2wav_state_lock", None) - if lock is None: - lock = threading.RLock() - self._token2wav_state_lock = lock - return lock - - def _begin_turn_vocoder_cache( - self, - prompt_wav_path: str | None, - *, - state: _TalkerTurnState | None = None, - ) -> None: - """Restore a fresh per-turn clone of the ref-audio vocoder caches.""" - import torchaudio - - with self._token2wav_lock(): - cache_key = prompt_wav_path or "" - base = self._t2w_base_caches.get(cache_key) - if base is None: - if prompt_wav_path is None: - base = (None, {}) - else: - _orig_save = torchaudio.save - prev_dtype = torch.get_default_dtype() - torch.set_default_dtype(torch.float32) - try: - torchaudio.save = _soundfile_patched_save(_orig_save) - stream_cache, hift_cache_dict = self.audio_tokenizer.set_stream_cache(prompt_wav_path) - finally: - torch.set_default_dtype(prev_dtype) - torchaudio.save = _orig_save - base = ( - _torch_clone_recursive(stream_cache), - _torch_clone_recursive(hift_cache_dict), - ) - self._t2w_base_caches[cache_key] = base - stream_cache = _torch_clone_recursive(base[0]) - hift_cache_dict = _torch_clone_recursive(base[1]) - if state is None: - self.audio_tokenizer.stream_cache = stream_cache - self.audio_tokenizer.hift_cache_dict = hift_cache_dict - else: - state.stream_cache = stream_cache - state.hift_cache_dict = hift_cache_dict - state.vocoder_initialized = True - self.audio_tokenizer.stream_cache = None - self.audio_tokenizer.hift_cache_dict = {} - - def _t2w_stream_window(self, token_list: list[int], prompt_wav_path: str | None, *, last_chunk: bool): - import torchaudio - - _orig_save = torchaudio.save - prev_dtype = torch.get_default_dtype() - autocast_context, _ = self._token2wav_autocast_context() - torch.set_default_dtype(torch.float32) - try: - torchaudio.save = _soundfile_patched_save(_orig_save) - with autocast_context: - wav_np = self.audio_tokenizer.stream( - token_list, - prompt_wav_path, - last_chunk=bool(last_chunk), - return_waveform=True, - ) - finally: - torch.set_default_dtype(prev_dtype) - torchaudio.save = _orig_save - return torch.as_tensor(np.asarray(wav_np).reshape(-1), dtype=torch.float32).cpu().contiguous() - - def _run_vocoder_window( - self, - state: _TalkerTurnState, - token_list: list[int], - *, - last_chunk: bool, - ) -> torch.Tensor: - with self._token2wav_lock(): - if not state.vocoder_initialized: - self._begin_turn_vocoder_cache(state.prompt_wav_path, state=state) - self.audio_tokenizer.stream_cache = _torch_clone_recursive(state.stream_cache) - self.audio_tokenizer.hift_cache_dict = _torch_clone_recursive(state.hift_cache_dict) - try: - waveform = self._t2w_stream_window( - token_list, - state.prompt_wav_path, - last_chunk=last_chunk, - ) - state.stream_cache = _torch_clone_recursive(self.audio_tokenizer.stream_cache) - state.hift_cache_dict = _torch_clone_recursive(self.audio_tokenizer.hift_cache_dict) - return waveform - finally: - self.audio_tokenizer.stream_cache = None - self.audio_tokenizer.hift_cache_dict = {} - - def _native_duplex_vocode_tokens( - self, - state: _TalkerTurnState, - new_tokens: torch.Tensor, - *, - turn_end: bool, - force_flush: bool, - chunk_size: int, - ) -> list[torch.Tensor]: - """Run Token2wav with the same buffering policy as official duplex.""" - pre_lookahead = self._t2w_pre_lookahead() - token_list = new_tokens.reshape(-1).detach().cpu().tolist() - state.token2wav_buffer.extend(int(t) for t in token_list) - pieces: list[torch.Tensor] = [] - - if force_flush: - while len(state.token2wav_buffer) >= pre_lookahead + 5: - chunk_to_process = min(chunk_size + pre_lookahead, len(state.token2wav_buffer)) - window = state.token2wav_buffer[:chunk_to_process] - pieces.append(self._run_vocoder_window(state, window, last_chunk=False)) - state.token2wav_buffer = state.token2wav_buffer[min(chunk_size, chunk_to_process - pre_lookahead) :] - else: - while len(state.token2wav_buffer) >= chunk_size + pre_lookahead: - window = state.token2wav_buffer[: chunk_size + pre_lookahead] - pieces.append(self._run_vocoder_window(state, window, last_chunk=False)) - state.token2wav_buffer = state.token2wav_buffer[chunk_size:] - - if turn_end and state.token2wav_buffer: - pieces.append( - self._run_vocoder_window( - state, - list(state.token2wav_buffer), - last_chunk=True, - ) - ) - state.token2wav_buffer = [] - - return pieces - - def _close_turn_state( - self, - key: str, - *, - expected_epoch: int | None = None, - expected_turn_id: int | None = None, - ) -> bool: - state = self._talker_turn_states.get(key) - if ( - state is not None - and expected_epoch is not None - and state.epoch is not None - and state.epoch != expected_epoch - ): - return False - if ( - state is not None - and expected_turn_id is not None - and state.turn_id is not None - and state.turn_id != expected_turn_id - ): - return False - state = self._talker_turn_states.pop(key, None) - self._talker_consumed_tokens.pop(key, None) - request_keys = getattr(self, "_talker_request_keys", None) - if isinstance(request_keys, dict): - request_keys.pop(key, None) - for alias, mapped_key in list(request_keys.items()): - if mapped_key == key: - request_keys.pop(alias, None) - if state is None: - return True - if self.audio_tokenizer is not None: - with self._token2wav_lock(): - self.audio_tokenizer.stream_cache = None - self.audio_tokenizer.hift_cache_dict = {} - temp_path = state.temp_prompt_wav_path - if temp_path and not self._is_cached_ref_audio_prompt_wav(temp_path): - with self._token2wav_lock(): - self._t2w_base_caches.pop(temp_path, None) - try: - os.unlink(temp_path) - except OSError: - pass - return True - - @staticmethod - def _stream_identity_relation( - state: _TalkerTurnState | None, - *, - epoch: int | None, - turn_id: int | None, - ) -> str: - if state is None: - return "current" - if epoch is not None and state.epoch is not None: - if epoch < state.epoch: - return "stale" - if epoch > state.epoch: - return "newer" - if turn_id is not None and state.turn_id is not None: - if turn_id < state.turn_id: - return "stale" - if turn_id > state.turn_id: - return "newer" - return "current" - - def _warmup_duplex_vocoder(self) -> None: - """Pre-compile the two token2wav stream() modes used per turn. - - The first stream() call of each mode (mid-turn 25+pre_lookahead - window, and the variable-size last_chunk tail flush) costs ~20s of - one-time compilation; without this warmup the first spoken turn of - the first session stalls for both. Mirrors the official demo's - precompile step. - """ - if getattr(self, "_t2w_warmed", False): - return - self._t2w_warmed = True - try: - self._lazy_init_tts() - if self.audio_tokenizer is None: - return - prompt_wav_path, _ = self._resolve_prompt_wav_path(None, None) - if prompt_wav_path is None: - return - t0 = time.perf_counter() - chunk_size = self._tts_runtime_config().streaming_generator_chunk - pre_lookahead = self._t2w_pre_lookahead() - self._begin_turn_vocoder_cache(prompt_wav_path) - self._t2w_stream_window( - [_T2W_SILENCE_TOKEN] * (chunk_size + pre_lookahead), - prompt_wav_path, - last_chunk=False, - ) - self._begin_turn_vocoder_cache(prompt_wav_path) - self._t2w_stream_window( - [_T2W_SILENCE_TOKEN] * (chunk_size + pre_lookahead), - prompt_wav_path, - last_chunk=True, - ) - self.audio_tokenizer.stream_cache = None - self.audio_tokenizer.hift_cache_dict = {} - logger.info("4.5 Talker duplex vocoder warmup done in %.1fs", time.perf_counter() - t0) - except Exception: - logger.exception("4.5 Talker duplex vocoder warmup failed") - - def _create_native_duplex_stream_gen(self, info: dict[str, Any]): - """Per-segment generator over a persistent per-turn talker stream. - - Mirrors the official duplex talker: each unit calls - ``MiniCPMTTS.generate_chunk`` while carrying ``past_key_values`` and - ``text_start_pos`` across the spoken turn, then feeds the generated - audio tokens through the same per-turn Token2wav buffer/cache. - """ - key = self._stream_request_key(info) - self._remember_talker_request_key(info, key) - stream_epoch = self._stream_epoch(info) - stream_turn_id = self._stream_turn_id(info) - meta_info = info.get("meta") if isinstance(info.get("meta"), dict) else {} - codes_info = info.get("codes") if isinstance(info.get("codes"), dict) else {} - tts_token_ids, tts_hidden_states = self._extract_tts_handoff(info) - - self._lazy_init_tts() - if getattr(self, "tts_obj", None) is None or self.audio_tokenizer is None: - logger.warning("4.5 Talker duplex streaming: TTS runtime unavailable") - yield self._empty_audio_chunk(), True - return - - if isinstance(tts_token_ids, torch.Tensor): - ids_list = tts_token_ids.reshape(-1).tolist() - elif isinstance(tts_token_ids, list): - ids_list = [int(t) for t in tts_token_ids] - else: - ids_list = [] - - state = self._talker_turn_states.get(key) - identity_relation = self._stream_identity_relation( - state, - epoch=stream_epoch, - turn_id=stream_turn_id, - ) - if identity_relation == "stale": - logger.info( - "4.5 Talker duplex drop stale handoff: key=%s state_epoch=%s state_turn_id=%s " - "handoff_epoch=%s handoff_turn_id=%s", - key, - getattr(state, "epoch", None), - getattr(state, "turn_id", None), - stream_epoch, - stream_turn_id, - ) - yield self._empty_audio_chunk(), True - return - if identity_relation == "newer": - self._close_turn_state(key) - state = None - consumed = self._talker_consumed_tokens.get(key, 0) - if consumed > len(ids_list): - consumed = 0 - pending_ids = ids_list[consumed:] - - turn_eos_raw = meta_info.get("turn_eos_token_id") - try: - turn_eos_id = int(turn_eos_raw) if turn_eos_raw is not None else None - except (TypeError, ValueError): - turn_eos_id = None - explicit_turn_end = bool( - info.get("end_of_turn") or info.get("turn_end") or meta_info.get("end_of_turn") or meta_info.get("turn_end") - ) - turn_end = explicit_turn_end or (turn_eos_id is not None and turn_eos_id in pending_ids) - terminal_only_new_turn = ( - state is None - and turn_end - and turn_eos_id is not None - and bool(pending_ids) - and all(token_id == turn_eos_id for token_id in pending_ids) - ) - if terminal_only_new_turn: - # A model-owned empty turn can be [speak, turn_eos]. There is no - # spoken state to flush, so conditioning TTS on turn_eos alone - # would synthesize an unrelated audio-only response. - yield self._empty_audio_chunk(), True - return - if state is None and not pending_ids: - # No turn open and nothing new to speak: nothing to synthesize. - yield self._empty_audio_chunk(), True - return - - tts = self.tts_obj - if not hasattr(tts.model.config, "rope_theta"): - tts.model.config.rope_theta = 10000.0 - if not callable(getattr(tts, "generate_chunk", None)): - logger.warning("4.5 Talker duplex streaming: MiniCPMTTS.generate_chunk unavailable") - yield self._empty_audio_chunk(), True - return - sampling_params = self._build_tts_sampling_params() - if sampling_params is None: - logger.warning("4.5 Talker duplex streaming: sampling params unavailable") - yield self._empty_audio_chunk(), True - return - chunk_size = self._tts_runtime_config().streaming_generator_chunk - if chunk_size <= 0: - raise ValueError("MiniCPM-o 4.5 TTS streaming generator chunk must be positive") - - if state is None: - ref_audio = codes_info.get("ref", info.get("ref_audio")) - ref_audio_sr = meta_info.get("ref_audio_sr", info.get("ref_audio_sr")) - prompt_wav_path, temp_prompt_wav_path = self._resolve_prompt_wav_path(ref_audio, ref_audio_sr) - if prompt_wav_path is None: - logger.warning("4.5 Talker duplex streaming: no ref_audio prompt; skipping audio synthesis") - yield self._empty_audio_chunk(), True - return - state = _TalkerTurnState( - prompt_wav_path, - temp_prompt_wav_path, - epoch=stream_epoch, - turn_id=stream_turn_id, - ) - self._begin_turn_vocoder_cache(prompt_wav_path, state=state) - self._talker_turn_states[key] = state - - _queue_native_duplex_segment_text( - state, - meta_info.get("native_duplex_segment_text", ""), - ) - - if pending_ids: - pending_hidden = ( - tts_hidden_states[consumed:] - if isinstance(tts_hidden_states, list) - else torch.as_tensor(tts_hidden_states)[consumed:] - ) - cond_ids, cond_hidden = self._normalize_tts_handoff_tensors(pending_ids, pending_hidden) - condition = self._build_tts_condition_embeds(cond_ids, cond_hidden).unsqueeze(0) - else: - emb_dim = int(tts.emb_text.weight.shape[1]) - condition = tts.emb_text.weight.new_zeros((1, 0, emb_dim)) - audio_bos = tts.emb_text( - torch.tensor( - [tts.audio_bos_token_id], - dtype=torch.long, - device=tts.emb_text.weight.device, - ) - ).unsqueeze(0) - condition = torch.cat([condition, audio_bos], dim=1) - max_token_per_chunk = chunk_size + 1 - min_token_per_chunk = 0 if turn_end else max_token_per_chunk - force_flush = False - if state.text_start_pos == 0: - min_token_per_chunk = 0 - force_flush = True - eos_token = torch.tensor( - [tts.config.num_audio_tokens - 1], - dtype=torch.long, - device=tts.emb_text.weight.device, - ) - temperature = torch.tensor( - [float(sampling_params.temperature)], - dtype=torch.float, - device=tts.emb_text.weight.device, - ) - new_tokens, past_key_values = tts.generate_chunk( - inputs_embeds=condition, - temperature=temperature, - repetition_penalty=sampling_params.repetition_penalty, - eos_token=eos_token, - force_no_stop=False, - max_new_token=max_token_per_chunk, - min_new_tokens=min_token_per_chunk, - past_key_values=state.past_key_values, - logits_processors=None, - text_start_pos=state.text_start_pos, - ) - if turn_end: - state.past_key_values = None - state.text_start_pos = 0 - else: - state.past_key_values = past_key_values - state.text_start_pos += int(condition.shape[1]) + int(new_tokens.shape[1]) - waveforms = self._native_duplex_vocode_tokens( - state, - new_tokens, - turn_end=turn_end, - force_flush=force_flush, - chunk_size=chunk_size, - ) - self._talker_consumed_tokens[key] = len(ids_list) - unit_waveform = _native_duplex_unit_waveform(waveforms, turn_end=turn_end) - if unit_waveform is not None: - self._ar_last_emitted_text = _drain_native_duplex_emitted_text( - state, - has_audio=True, - ) - yield unit_waveform, False - if turn_end: - self._close_turn_state(key) - self._ar_last_emitted_text = "" - yield self._empty_audio_chunk(), True - return - self._ar_last_emitted_text = "" - yield self._empty_audio_chunk(), True - - def _create_stream_gen(self, info: dict[str, Any]): - """Yield waveform chunks from MiniCPM-o remote-code TTS streaming. - - This is the real vLLM streaming path: each yielded tensor is returned - through one scheduler step. The older streaming probe still concatenates - chunks inside generate_speech(), so it cannot improve API TTFA. - """ - if info.get("native_duplex") is True: - yield from self._create_native_duplex_stream_gen(info) - return - tts_token_ids, tts_hidden_states = self._extract_tts_handoff(info) - codes_info = info.get("codes") - meta_info = info.get("meta") - if not isinstance(codes_info, dict): - codes_info = {} - if not isinstance(meta_info, dict): - meta_info = {} - - ref_audio = codes_info.get("ref", info.get("ref_audio")) - ref_audio_sr = meta_info.get("ref_audio_sr", info.get("ref_audio_sr")) - - if tts_token_ids is None or tts_hidden_states is None: - logger.warning("4.5 Talker streaming: missing tts_token_ids or tts_hidden_states") - yield self._empty_audio_chunk(), True - return - tts_token_ids, tts_hidden_states = self._normalize_tts_handoff_tensors( - tts_token_ids, - tts_hidden_states, - ) - - self._lazy_init_tts() - if not hasattr(self, "tts_obj") or self.tts_obj is None: - logger.warning("4.5 Talker streaming: tts_obj not initialized") - yield self._empty_audio_chunk(), True - return - if self.audio_tokenizer is None: - logger.warning("4.5 Talker streaming: audio_tokenizer not initialized") - yield self._empty_audio_chunk(), True - return - - generator_cls = getattr(self, "_tts_streaming_generator_cls", None) - if generator_cls is None or self._tts_gen_logits is None: - logger.warning("4.5 Talker streaming: remote-code TTSStreamingGenerator unavailable") - waveform = self.generate_speech( - tts_token_ids, - tts_hidden_states, - ref_audio=ref_audio, - ref_audio_sr=ref_audio_sr, - ) - if waveform is None: - yield self._empty_audio_chunk(), True - else: - yield torch.as_tensor(waveform, dtype=torch.float32).reshape(-1).cpu().contiguous(), True - return - - tts = self.tts_obj - if not hasattr(tts.model.config, "rope_theta"): - tts.model.config.rope_theta = 10000.0 - - tts_embeds = self._build_tts_condition_embeds(tts_token_ids, tts_hidden_states) - num_text = int(tts_token_ids.shape[-1]) if tts_token_ids.ndim > 0 else 0 - min_new_token, max_new_token = self._max_tts_tokens_for_text(num_text) - sampling_params = self._build_tts_sampling_params() - if sampling_params is None: - logger.warning("4.5 Talker streaming: sampling params unavailable") - yield self._empty_audio_chunk(), True - return - - logits_warpers, logits_processors = self._tts_gen_logits( - num_code=tts.config.num_audio_tokens, - repetition_penalty=sampling_params.repetition_penalty, - top_p=sampling_params.top_p, - top_k=sampling_params.top_k, - ) - eos_token = torch.tensor([tts.config.num_audio_tokens - 1], dtype=torch.long, device=tts.emb_text.weight.device) - chunk_size = self._tts_runtime_config().streaming_generator_chunk - if chunk_size <= 0: - raise ValueError("MiniCPM-o 4.5 TTS streaming generator chunk must be positive") - - tts_streaming_generator = generator_cls( - model=tts, - temperature=sampling_params.temperature, - eos_token=eos_token, - chunk_size=chunk_size, - logits_processors=logits_processors, - logits_warpers=logits_warpers, - ) - - prompt_wav_path, temp_prompt_wav_path = self._resolve_prompt_wav_path(ref_audio, ref_audio_sr) - stream_cache = hift_cache_dict = None - import torchaudio - - _orig_save = torchaudio.save - - def _patched_save(uri, src, sample_rate, **kw): - kw.pop("backend", None) - if hasattr(uri, "write"): - sf.write(uri, src.cpu().numpy().T, sample_rate, format="WAV") - return - return _orig_save(uri, src, sample_rate, backend="soundfile", **kw) - - yielded_any = False - try: - torchaudio.save = _patched_save - prev_dtype = torch.get_default_dtype() - torch.set_default_dtype(torch.float32) - try: - stream_cache, hift_cache_dict = self.audio_tokenizer.set_stream_cache(prompt_wav_path) - finally: - torch.set_default_dtype(prev_dtype) - torchaudio.save = _orig_save - self.audio_tokenizer.stream_cache = stream_cache - self.audio_tokenizer.hift_cache_dict = hift_cache_dict - token_iter = tts_streaming_generator.generate_with_buffer( - condition=tts_embeds.unsqueeze(0), - text_finished=True, - max_new_token=max_new_token, - ) - while True: - try: - audio_token_chunk, is_last = next(token_iter) - except StopIteration: - break - if audio_token_chunk is None: - break - - token_list = audio_token_chunk.reshape(-1).detach().cpu().tolist() - if not token_list: - if is_last: - yield self._empty_audio_chunk(), True - yielded_any = True - break - continue - - autocast_context, _ = self._token2wav_autocast_context() - torchaudio.save = _patched_save - prev_dtype = torch.get_default_dtype() - torch.set_default_dtype(torch.float32) - try: - with autocast_context: - wav_np = self.audio_tokenizer.stream( - token_list, - prompt_wav_path, - last_chunk=bool(is_last), - return_waveform=True, - ) - finally: - torch.set_default_dtype(prev_dtype) - torchaudio.save = _orig_save - chunk = torch.as_tensor(np.asarray(wav_np).reshape(-1), dtype=torch.float32).cpu().contiguous() - yielded_any = True - yield chunk, bool(is_last) - if is_last: - break - finally: - torchaudio.save = _orig_save - self.audio_tokenizer.stream_cache = None - self.audio_tokenizer.hift_cache_dict = {} - if temp_prompt_wav_path and not self._is_cached_ref_audio_prompt_wav(temp_prompt_wav_path): - try: - os.unlink(temp_prompt_wav_path) - except OSError: - pass - - if not yielded_any: - yield self._empty_audio_chunk(), True - - def _move_tts_modules_to_device(self) -> torch.dtype: - device = current_omni_platform.get_torch_device() - target_dtype = torch.bfloat16 if current_omni_platform.is_npu() else self._target_tts_dtype() - if target_dtype is torch.float32: - self.tts_obj = self.tts_obj.to(device) - logger.info("Moved MiniCPM-o 4.5 TTS object to %s dtype=%s", device, target_dtype) - return target_dtype - - for module_name in ( - "emb_text", - "model", - "projector_spk", - "projector_semantic", - "emb_code", - "head_code", - ): - module = getattr(self.tts_obj, module_name, None) - if module is not None: - module.to(device=device, dtype=target_dtype) - logger.info("Moved MiniCPM-o 4.5 TTS AR modules to %s dtype=%s", device, target_dtype) - return target_dtype - - def generate_speech( - self, - tts_token_ids: torch.Tensor, - tts_hidden_states: torch.Tensor, - ref_audio=None, - ref_audio_sr: int | None = None, - ) -> torch.Tensor | np.ndarray | None: - """Run full 4.5 TTS pipeline using original MiniCPMTTS.generate.""" - self._lazy_init_tts() - if not hasattr(self, "tts_obj") or self.tts_obj is None: - logger.warning("generate_speech: tts_obj not initialized") - return None - - tts = self.tts_obj - device = tts.emb_text.weight.device - dtype = tts.emb_text.weight.dtype - - llm_embeds = tts.emb_text(tts_token_ids.to(device)) - hidden_embeds = tts.projector_semantic(tts_hidden_states.to(device=device, dtype=dtype)) - if getattr(tts.config, "normalize_projected_hidden", False): + text_embeds = self.emb_text(token_ids) + hidden_embeds = self.projector_semantic(hidden) + if self._normalize: hidden_embeds = F.normalize(hidden_embeds, p=2, dim=-1) - tts_embeds = llm_embeds + hidden_embeds + text_eos = self.emb_text(torch.tensor([self._text_eos_id], device=device, dtype=torch.long)) + audio_bos = self.emb_text(torch.tensor([self._tts_bos_id], device=device, dtype=torch.long)) + return torch.cat([text_embeds + hidden_embeds, text_eos, audio_bos], dim=0) - text_eos = tts.emb_text(torch.tensor([tts.config.text_eos_token_id], device=device, dtype=torch.long)) - audio_bos = tts.emb_text(torch.tensor([tts.audio_bos_token_id], device=device, dtype=torch.long)) - spk_embeds = torch.zeros(0, tts.config.hidden_size, device=device, dtype=tts_embeds.dtype) + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Build request-local prefill/decode embeddings for the vLLM runner.""" + del input_embeds + span_len = int(input_ids.shape[0]) + is_prefill = bool(info_dict.get("_omni_is_prefill", False)) + state = info_dict.get("audio_state") + first_call = not isinstance(state, dict) - inputs_embeds = torch.cat([spk_embeds, tts_embeds, text_eos, audio_bos], dim=0).unsqueeze(0) - - # Scale max_new_token with input text length. A fixed 2048-token floor - # can turn an EOS miss on a very short response into ~82s of audio and - # ~18s E2E latency. Keep a conservative short-text floor while bounding - # the tail. - num_text = int(tts_token_ids.shape[-1]) if tts_token_ids.ndim > 0 else 0 - min_new_token, max_new_token = self._max_tts_tokens_for_text(num_text) - - eos_token = torch.tensor([tts.config.num_audio_tokens - 1], dtype=torch.long, device=device) - sampling_params = self._build_tts_sampling_params() - generate_kwargs = { - "inputs_embeds": inputs_embeds, - "eos_token": eos_token, - "max_new_token": max_new_token, - "min_new_token": min_new_token, - "show_tqdm": False, - } - if sampling_params is not None: - generate_kwargs["sampling_params"] = sampling_params - - if self.audio_tokenizer is None: - logger.warning("No audio_tokenizer") - return None - - prompt_wav_path, temp_prompt_wav_path = self._resolve_prompt_wav_path(ref_audio, ref_audio_sr) - - try: - outputs = tts.generate(**generate_kwargs) - generated_tokens = outputs.new_ids.squeeze(-1) - - import torchaudio - - _orig_save = torchaudio.save - - def _patched_save(uri, src, sample_rate, **kw): - kw.pop("backend", None) - if hasattr(uri, "write"): - sf.write(uri, src.cpu().numpy().T, sample_rate, format="WAV") - return - return _orig_save(uri, src, sample_rate, backend="soundfile", **kw) - - torchaudio.save = _patched_save - prev_dtype = torch.get_default_dtype() - torch.set_default_dtype(torch.float32) - try: - autocast_context, token2wav_autocast = self._token2wav_autocast_context() - with autocast_context: - num_tokens = int(generated_tokens.shape[-1]) - - # For long outputs, the one-shot vocoder path - # (Token2wav.__call__ -> flow.inference) runs full O(N^2) self- - # attention over all audio tokens and OOMs on a 24GB card once - # N exceeds a few thousand (e.g. 4964 tokens needs ~3GiB for a - # single attention matmul). Switch to the chunked / streaming - # vocoder (set_stream_cache + stream) which truncates the flow - # attention caches to prompt_len + 100 steps on every chunk, - # keeping peak memory bounded regardless of total length. - STREAM_THRESHOLD = self._tts_runtime_config().streaming_vocoder_threshold # ~100s @ 25Hz - CHUNK_SIZE = self._tts_runtime_config().streaming_vocoder_chunk # ~2s per chunk - MIN_TAIL = 6 # must exceed flow.pre_lookahead_len (typically 3) - - if num_tokens <= STREAM_THRESHOLD: - if self._should_use_direct_token2wav() and token2wav_autocast == "off": - try: - waveform, _ = self._run_token2wav_direct(generated_tokens, prompt_wav_path) - except Exception as exc: - logger.warning( - "MiniCPM-o 4.5 direct Token2wav path failed; falling back to WAV path: %s", - exc, - exc_info=True, - ) - token_list = generated_tokens.squeeze(0).tolist() - self._reset_token2wav_cache_if_needed(prompt_wav_path) - wav_bytes = self.audio_tokenizer(token_list, prompt_wav_path) - waveform, _ = sf.read(io.BytesIO(wav_bytes)) - waveform = waveform.astype(np.float32) - else: - token_list = generated_tokens.squeeze(0).tolist() - self._reset_token2wav_cache_if_needed(prompt_wav_path) - wav_bytes = self.audio_tokenizer(token_list, prompt_wav_path) - waveform, _ = sf.read(io.BytesIO(wav_bytes)) - waveform = waveform.astype(np.float32) - else: - token_list = generated_tokens.squeeze(0).tolist() - # Build chunk boundaries, merging a too-small tail into the - # previous chunk so every chunk satisfies MIN_TAIL. - boundaries = [] - i = 0 - while i < num_tokens: - end = min(i + CHUNK_SIZE, num_tokens) - if 0 < num_tokens - end < MIN_TAIL: - end = num_tokens - boundaries.append((i, end)) - i = end - - logger.info( - "generate_speech: streaming vocoder, %d tokens -> %d chunks (chunk=%d)", - num_tokens, - len(boundaries), - CHUNK_SIZE, - ) - - stream_cache, hift_cache_dict = self.audio_tokenizer.set_stream_cache(prompt_wav_path) - self.audio_tokenizer.stream_cache = stream_cache - self.audio_tokenizer.hift_cache_dict = hift_cache_dict - - try: - pieces = [] - for idx, (s, e) in enumerate(boundaries): - is_last = idx == len(boundaries) - 1 - wav_np = self.audio_tokenizer.stream( - token_list[s:e], - prompt_wav_path, - last_chunk=is_last, - return_waveform=True, - ) - pieces.append(np.asarray(wav_np).reshape(-1)) - waveform = np.concatenate(pieces, axis=0).astype(np.float32) - finally: - # Free per-request streaming state so the next request starts clean - self.audio_tokenizer.stream_cache = None - self.audio_tokenizer.hift_cache_dict = {} - finally: - torch.set_default_dtype(prev_dtype) - torchaudio.save = _orig_save - - return waveform - finally: - if temp_prompt_wav_path and not self._is_cached_ref_audio_prompt_wav(temp_prompt_wav_path): - try: - os.unlink(temp_prompt_wav_path) - except OSError: - pass - - def _generate_tokens(self, inputs_embeds: torch.Tensor, max_new_token: int = 2048) -> torch.Tensor | None: - """Autoregressive generation of audio tokens using the TTS LlamaModel.""" - device = inputs_embeds.device - eos_token = self._num_audio_tokens - 1 - condition_length = inputs_embeds.shape[1] - num_vq = len(self.emb_code) - - new_tokens = torch.zeros(1, max_new_token, num_vq, device=device, dtype=torch.long) - past_key_values = None - finished = False - - for t in range(max_new_token): - if t == 0: - emb = inputs_embeds - position_ids = torch.arange(condition_length, device=device).unsqueeze(0) - else: - code_emb = [self.emb_code[q](new_tokens[:, t - 1 : t, q]) for q in range(num_vq)] - emb = torch.stack(code_emb, -1).sum(-1) - position_ids = torch.tensor([[condition_length + t - 1]], device=device) - - outputs = self.tts_model( - inputs_embeds=emb, - position_ids=position_ids, - past_key_values=past_key_values, - use_cache=True, + if is_prefill or first_call: + token_ids, hidden_states = get_tts_handoff(info_dict) + # Cross-process stage transport serializes CPU tensors as lists. + # Normalize both local tensor handoffs and transported payloads + # before validating/building the Talker condition. + if isinstance(token_ids, (list, tuple)): + token_ids = torch.as_tensor(token_ids, dtype=torch.long) + if isinstance(hidden_states, (list, tuple)): + hidden_states = torch.as_tensor(hidden_states, dtype=torch.float32) + if not isinstance(token_ids, torch.Tensor) or not isinstance(hidden_states, torch.Tensor): + available = sorted(key for key in info_dict if not key.startswith("_")) + raise ValueError( + "MiniCPM-o Talker requires tensor tts_token_ids and " + "tts_hidden_states conditioning; " + f"received token_ids={type(token_ids).__name__}, " + f"hidden_states={type(hidden_states).__name__}, " + f"available_keys={available}" + ) + if token_ids.numel() == 0 or hidden_states.numel() == 0: + raise ValueError("MiniCPM-o Talker conditioning must not be empty") + full_embeds = self._build_condition_embeddings(token_ids, hidden_states) + offset = int(info_dict.get("_omni_num_computed_tokens", 0)) + if offset == 0 and span_len > full_embeds.shape[0]: + # Async-chunk prewarms Stage 1 with placeholder token IDs. Two + # Thinker handoffs can arrive before the first one executes, + # so the scheduler prompt may already include the newer + # placeholder prefix while additional_information still + # carries the preceding condition. Materialize that prefix as + # the embeddings of its actual zero token IDs and keep the + # complete Talker condition tail-aligned (audio BOS remains + # the final prefill token). + prefix_len = span_len - full_embeds.shape[0] + placeholder_ids = torch.zeros( + prefix_len, + dtype=torch.long, + device=self.emb_text.weight.device, + ) + full_embeds = torch.cat([self.emb_text(placeholder_ids), full_embeds], dim=0) + embeds = full_embeds[offset : offset + span_len] + if embeds.shape[0] != span_len: + raise ValueError( + "MiniCPM-o Talker prefill span exceeds condition: " + f"request_id={info_dict.get('request_id')} offset={offset} " + f"span={span_len} condition={full_embeds.shape[0]} " + f"tts_ids={token_ids.shape[0]} tts_hidden={hidden_states.shape[0]} " + f"prompt_len={info_dict.get('_omni_prompt_len')}" + ) + max_tokens = _max_audio_tokens(int(token_ids.numel())) + state = { + "step": 0, + "max_tokens": max_tokens, + "finished": False, + } + request_id = str(info_dict.get("request_id", "0")) + request_states = getattr(self, "_request_audio_states", None) + if request_states is None: + request_states = {} + self._request_audio_states = request_states + request_states[request_id] = state + empty_codes = torch.empty(0, dtype=torch.long, device=embeds.device) + return ( + input_ids, + embeds, + { + "audio_state": state, + "audio_codes": { + "current": empty_codes, + "accumulated": empty_codes, + }, + }, ) - hidden = outputs.last_hidden_state - past_key_values = outputs.past_key_values - logits = torch.stack([self.head_code[q](hidden[:, -1]) for q in range(num_vq)], dim=-1) - logits = logits.float() / 0.8 + current = (info_dict.get("audio_codes", {}) or {}).get("current") + if not isinstance(current, torch.Tensor) or current.numel() != 1: + raise RuntimeError("MiniCPM-o Talker decode is missing the previous request-local audio code") + code = current.to(device=self.emb_code[0].weight.device, dtype=torch.long).reshape(1) + embeds = self.emb_code[0](code) + return input_ids, embeds, {} - if t < 50: - logits[:, eos_token, :] = -float("inf") + def _request_generator(self, request_id: str, device: torch.device) -> torch.Generator: + generator = self._request_generators.get(request_id) + if generator is None: + generator = torch.Generator(device=device) + generator.manual_seed(self._codec_seed) + self._request_generators[request_id] = generator + return generator - probs = F.softmax(logits, dim=1) - idx = torch.multinomial(probs.view(-1, probs.shape[1]), 1).view(1, num_vq) - new_tokens[:, t] = idx + def _sample_audio_code( + self, + hidden_state: torch.Tensor, + history: torch.Tensor, + request_id: str, + step: int, + ) -> torch.Tensor: + logits = self.head_code[0](hidden_state).float() / self._codec_temperature + eos_id = self._num_audio_tokens - 1 + logits = _apply_repetition_penalty( + logits, + history, + penalty=self._codec_repetition_penalty, + window_size=_REPETITION_WINDOW, + ) + if step < self._codec_min_tokens: + logits[..., eos_id] = float("-inf") + logits = _apply_top_k_top_p( + logits, + top_k=self._codec_top_k, + top_p=self._codec_top_p, + min_tokens_to_keep=3, + ) + probabilities = torch.softmax(logits, dim=-1) + return torch.multinomial( + probabilities, + num_samples=1, + generator=self._request_generator(request_id, probabilities.device), + ).reshape(()) - if (idx == eos_token).any(): - finished = True + def make_omni_output( + self, + model_outputs: torch.Tensor | OmniOutput, + **kwargs: Any, + ) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + hidden = model_outputs + infos = kwargs.get("model_intermediate_buffer") or [] + spans = kwargs.get("request_token_spans") + if spans is None or len(spans) != len(infos): + raise RuntimeError("MiniCPM-o continuous Talker requires one request_token_span per request") + sample_eligible = kwargs.get("request_sample_eligible") + if sample_eligible is None: + sample_eligible = [True] * len(infos) + if len(sample_eligible) != len(infos): + raise RuntimeError( + f"MiniCPM-o continuous Talker received {len(sample_eligible)} sampling flags for {len(infos)} requests" + ) + + stop_rows: list[torch.Tensor] = [] + codec_deltas: list[torch.Tensor] = [] + terminal_flags: list[torch.Tensor] = [] + empty_delta = hidden.new_empty((0, 1), dtype=torch.long) + for index, info in enumerate(infos): + if not isinstance(info, dict): + stop_rows.append(hidden.new_tensor([0.0, float("-inf")])) + codec_deltas.append(empty_delta) + terminal_flags.append(torch.tensor(False, dtype=torch.bool)) + continue + start, end = spans[index] + end = min(int(end), int(hidden.shape[0])) + if int(start) >= end: + stop_rows.append(hidden.new_tensor([0.0, float("-inf")])) + codec_deltas.append(empty_delta) + terminal_flags.append(torch.tensor(False, dtype=torch.bool)) + continue + request_id = str(info.get("request_id", index)) + request_states = getattr(self, "_request_audio_states", None) + if request_states is None: + request_states = {} + self._request_audio_states = request_states + state = request_states.get(request_id) + if not isinstance(state, dict): + state = dict(info.get("audio_state", {}) or {}) + request_states[request_id] = state + if state.get("finished"): + stop_rows.append(hidden.new_tensor([float("-inf"), 0.0])) + codec_deltas.append(empty_delta) + terminal_flags.append(torch.tensor(False, dtype=torch.bool)) + continue + if not sample_eligible[index]: + # vLLM computes a logit row for incomplete chunked prefills but + # discards its sampled token. Advancing codec/RNG state here + # would make output depend on prefill chunking and compaction. + stop_rows.append(hidden.new_tensor([0.0, float("-inf")])) + codec_deltas.append(empty_delta) + terminal_flags.append(torch.tensor(False, dtype=torch.bool)) + continue + codes = state.get("codes") + if not isinstance(codes, torch.Tensor): + codes = (info.get("audio_codes", {}) or {}).get("accumulated") + if not isinstance(codes, torch.Tensor): + codes = torch.empty(0, dtype=torch.long, device=hidden.device) + else: + codes = codes.to(device=hidden.device, dtype=torch.long).reshape(-1) + step = int(state.get("step", 0)) + sampled = self._sample_audio_code(hidden[end - 1 : end], codes, request_id, step) + is_eos = int(sampled.item()) == self._num_audio_tokens - 1 + state["step"] = int(state.get("step", 0)) + 1 + reached_limit = int(state["step"]) >= int(state.get("max_tokens", 2048)) + finished = is_eos or reached_limit + state["finished"] = finished + if not is_eos: + codes = torch.cat([codes[-(_REPETITION_WINDOW - 1) :], sampled.reshape(1)]) + delta = sampled.reshape(1, 1) + else: + delta = empty_delta + state["codes"] = codes + info["audio_state"] = state + info["audio_codes"] = { + "current": sampled.reshape(1), + "accumulated": codes, + } + codec_deltas.append(delta) + terminal_flags.append(torch.tensor(finished, dtype=torch.bool)) + stop_rows.append(hidden.new_tensor([float("-inf"), 0.0] if finished else [0.0, float("-inf")])) + + self._batch_stop_logits = torch.stack(stop_rows, dim=0) if stop_rows else hidden.new_empty((0, 2)) + # Lists are deliberate: the runner routes element i to request i, + # preserving compaction alignment while emitting only this step's code. + multimodal_outputs: dict[str, Any] = { + "codes": {"audio": codec_deltas}, + "meta": {"finished": terminal_flags}, + } + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs=multimodal_outputs, + ) + + def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: + self._deferred_cleanup_ids.update(str(req_id) for req_id in finished_req_ids) + + def _flush_deferred_cleanup(self) -> None: + request_audio_states = getattr(self, "_request_audio_states", {}) + for request_id in self._deferred_cleanup_ids: + self._request_generators.pop(request_id, None) + request_audio_states.pop(request_id, None) + self._deferred_cleanup_ids.clear() + + def _dummy_hidden_states( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor | None, + inputs_embeds: torch.Tensor | None, + ) -> torch.Tensor: + """Shape-correct zero tensor for vllm KV cache profiling. + + vllm's gpu_model_runner._dummy_run takes forward()'s return value as + ``hidden_states`` and does ``hidden_states[logit_indices_device]``; + returning None on the dummy path crashes with + ``TypeError: 'NoneType' object is not subscriptable``. + """ + for ref in (input_ids, positions, inputs_embeds): + if isinstance(ref, torch.Tensor): + num_tokens = int(ref.shape[0]) if ref.ndim >= 1 else 1 + device = ref.device break - - return new_tokens[:, : t + 1 if finished else t, :] + else: + num_tokens = 1 + device = current_omni_platform.get_torch_device() + hidden_size = int(getattr(self, "_hidden_size", 768) or 768) + return torch.zeros((num_tokens, hidden_size), device=device, dtype=torch.bfloat16) def forward( self, @@ -1617,163 +494,77 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): positions=None, intermediate_tensors=None, inputs_embeds=None, - additional_information=None, **kwargs, ): - if additional_information is None: - additional_information = {} - if not additional_information: - # Profile/dummy run: use it to pre-compile the per-turn vocoder - # stream modes so the first real spoken turn does not stall. - self._warmup_duplex_vocoder() - - tts_token_ids, tts_hidden_states = self._extract_tts_handoff(additional_information) - tts_text = additional_information.get("llm_output_text", [""]) - if isinstance(tts_text, list): - tts_text = tts_text[0] if tts_text else "" - codes_info = additional_information.get("codes") - meta_info = additional_information.get("meta") - if not isinstance(codes_info, dict): - codes_info = {} - if not isinstance(meta_info, dict): - meta_info = {} - ref_audio = codes_info.get("ref") - if ref_audio is None: - ref_audio = additional_information.get("ref_audio") - ref_audio_sr = meta_info.get("ref_audio_sr") - if ref_audio_sr is None: - ref_audio_sr = additional_information.get("ref_audio_sr") - - if tts_token_ids is None or tts_hidden_states is None: - logger.warning("4.5 Talker: missing tts_token_ids or tts_hidden_states") - self._ar_last_chunk_flags = [True] - self._ar_turn_end_flags = [False] - return None, None - tts_token_ids, tts_hidden_states = self._normalize_tts_handoff_tensors( - tts_token_ids, - tts_hidden_states, + self._flush_deferred_cleanup() + if input_ids is None and inputs_embeds is None: + return self._dummy_hidden_states(input_ids, positions, inputs_embeds) + return self.tts_model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, ) - if self._should_stream_output(additional_information): - request_key = self._stream_request_key(additional_information) - input_ends_turn = self._native_duplex_input_ends_turn(additional_information) - if request_key not in self._stream_gens: - self._stream_gens[request_key] = self._create_stream_gen(additional_information) - generator = self._stream_gens[request_key] - try: - waveform_chunk, is_last = next(generator) - except StopIteration: - self._stream_gens.pop(request_key, None) - waveform_chunk = self._empty_audio_chunk() - is_last = True - if is_last: - self._stream_gens.pop(request_key, None) - self._ar_last_chunk_flags = [bool(is_last)] - # A TTS generator also ends at ordinary chunk boundaries. Export - # turn_end only on the terminal output for a condition that - # actually contains the model's <|turn_eos|> decision. - self._ar_turn_end_flags = [bool(is_last and input_ends_turn)] - return None, waveform_chunk.reshape(-1).contiguous() - - self._ar_last_chunk_flags = [True] - self._ar_turn_end_flags = [False] - waveform = self.generate_speech( - tts_token_ids, - tts_hidden_states, - ref_audio=ref_audio, - ref_audio_sr=ref_audio_sr, - ) - if waveform is not None: - waveform_tensor = torch.as_tensor(waveform, dtype=torch.float32).detach() - if waveform_tensor.device.type != "cpu": - waveform_tensor = waveform_tensor.cpu() - return waveform_tensor.reshape(-1).contiguous(), None - return None, None - def compute_logits(self, hidden_states, *args, **kwargs): - device = hidden_states.device if isinstance(hidden_states, torch.Tensor) else torch.device("cuda") - if isinstance(hidden_states, torch.Tensor): - if hidden_states.ndim == 1: - num_rows = 1 - else: - num_rows = max(1, int(hidden_states.shape[0])) - else: - num_rows = 1 - eos_id = self._scheduler_eos_token_id() - vocab_size = max(int(getattr(self.config, "vocab_size", eos_id + 1) or (eos_id + 1)), eos_id + 1, 3) - safe_id = 1 if eos_id != 1 else 0 - logits = torch.full((num_rows, vocab_size), -1.0e9, dtype=torch.float32, device=device) - flags = self._ar_last_chunk_flags - default_is_last = bool(flags[-1]) if flags else True - for row in range(num_rows): - is_last = bool(flags[row]) if row < len(flags) else default_is_last - if is_last: - logits[row, eos_id] = 1.0e6 - else: - logits[row, safe_id] = 1.0e6 + if not isinstance(hidden_states, torch.Tensor): + return None + if self._batch_stop_logits is None: + return torch.zeros( + hidden_states.shape[0], + 2, + device=hidden_states.device, + dtype=torch.float32, + ) + logits = self._batch_stop_logits + self._batch_stop_logits = None return logits def sample(self, logits, sampling_metadata): - if logits is None or logits.numel() == 0: - return None - sampled = torch.argmax(logits, dim=-1).to(torch.int32) - return SamplerOutput(sampled_token_ids=sampled.unsqueeze(-1), logprobs_tensors=None) - - def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: - for req_id in finished_req_ids: - request_key = str(req_id) - mapped_key = getattr(self, "_talker_request_keys", {}).pop(request_key, None) - keys = {request_key} - if mapped_key: - keys.add(mapped_key) - for key in list(self._stream_gens): - if key in keys: - gen = self._stream_gens.pop(key, None) - if gen is not None: - try: - gen.close() - except Exception: - logger.exception("MiniCPM-o 4.5 failed to close stream gen for request %s", req_id) - for key in list(self._talker_turn_states): - if key in keys: - self._close_turn_state(key) - for key in list(self._talker_consumed_tokens): - if key in keys: - self._talker_consumed_tokens.pop(key, None) + return Sampler()(logits, sampling_metadata) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): - loaded = set() - tts_weights = {} - for k, v in weights: - if k.startswith("tts."): - local_name = k.replace("tts.", "", 1) - tts_weights[local_name] = v - loaded.add(f"tts_obj.{local_name}") + return self._load_native_weights(weights) - if tts_weights and self._tts_config is not None: - self._lazy_init_tts() - if hasattr(self, "tts_obj") and self.tts_obj is not None: - missing, unexpected = self.tts_obj.load_state_dict(tts_weights, strict=False) - if missing: - logger.warning("TTS missing keys (%d): %s", len(missing), missing[:5]) - if unexpected: - logger.warning("TTS unexpected keys (%d): %s", len(unexpected), unexpected[:5]) - tts_dtype = self._move_tts_modules_to_device() - if ( - not current_omni_platform.is_npu() - and self.audio_tokenizer is not None - and hasattr(self.audio_tokenizer, "to") - ): - self.audio_tokenizer.to("cuda") - self.emb_text = self.tts_obj.emb_text - self.projector_semantic = self.tts_obj.projector_semantic - logger.info( - "Loaded %d TTS weights, moved AR modules to %s dtype=%s", - len(tts_weights), - current_omni_platform.get_torch_device(), - tts_dtype, - ) + def _load_native_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loaded: set[str] = set() + backbone_weights: list[tuple[str, torch.Tensor]] = [] + direct_params = dict(self.named_parameters()) + head_g = head_v = None + for name, tensor in weights: + if not name.startswith("tts."): + continue + stripped = name[len("tts.") :] + if stripped.startswith("model."): + backbone_weights.append((stripped[len("model.") :], tensor)) + continue + if stripped == "head_code.0.parametrizations.weight.original0": + head_g = tensor + continue + if stripped == "head_code.0.parametrizations.weight.original1": + head_v = tensor + continue + target = stripped + parameter = direct_params.get(target) + if parameter is None: + continue + parameter.data.copy_(tensor.to(device=parameter.device, dtype=parameter.dtype)) + loaded.add(target) + + for name in self.tts_model.load_weights(backbone_weights): + loaded.add(f"tts_model.{name}") + + if head_g is None or head_v is None: + raise ValueError("MiniCPM-o checkpoint is missing weight-norm Talker head parameters") + restored = _restore_weight_norm_weight(head_g, head_v) + self.head_code[0].weight.data.copy_( + restored.to( + device=self.head_code[0].weight.device, + dtype=self.head_code[0].weight.dtype, + ) + ) + loaded.add("head_code.0.weight") return loaded def get_input_embeddings(self, input_ids, multimodal_embeddings=None, **kwargs): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index 57d8c4614..31d02c7cc 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -3,12 +3,11 @@ """MiniCPM-o 4.5 pipeline topology (frozen). Stage 0: Thinker — multimodal understanding + text generation. -Stage 1: Talker — MiniCPMTTS + Token2Wav, emits the final audio waveform. +Stage 1: Talker — MiniCPMTTS, emits codec tokens. +Stage 2: Code2Wav — codec tokens to the final audio waveform. -The thinker -> talker bridge passes the hidden states + token ids extracted -from the thinker output through ``minicpmo_4_5_omni.llm2tts``; the talker -runs MiniCPMTTS and the on-device Token2wav vocoder in the same process and -returns the waveform directly as the pipeline's final audio output. +The thinker -> talker bridge uses ``llm2tts``. The talker -> Code2Wav bridge +streams request-routed codec chunks. """ from vllm_omni.config.stage_config import ( @@ -56,28 +55,24 @@ MINICPMO_4_5_PIPELINE = PipelineConfig( StagePipelineConfig( stage_id=1, model_stage="tts", - # Stage 1 shares the top-level wrapper class - # (``MiniCPMO45OmniForConditionalGeneration``) inherited from - # ``model_arch`` above. The wrapper dispatches on ``model_stage`` - # and, for ``"tts"``, instantiates the standalone TTS submodule - # (``MiniCPMO45OmniTTSForConditionalGeneration``) internally. - # Routing through the wrapper is required so that the runner-side - # ``runtime_additional_information`` payload reaches the talker - # (the standalone TTS class only reads ``additional_information``, - # so wiring stage 1 directly to it would always trigger the dummy - # path) and so the resulting waveform is packaged as - # ``OmniOutput.multimodal_outputs["model_outputs"]`` instead of - # being returned as a bare tuple that the AR runner would mistake - # for hidden states. ``hf_config_name="tts_config"`` keeps KV - # cache / mrope sizing scoped to the talker sub-config. execution_type=StageExecutionType.LLM_AR, input_sources=(0,), + hf_config_name="tts_config", + engine_output_type="latent", + custom_process_input_func=f"{_PROC}.llm2tts", + async_chunk_process_next_stage_input_func=f"{_PROC}.tts2code2wav_async_chunk", + sampling_constraints={"detokenize": False}, + ), + StagePipelineConfig( + stage_id=2, + model_stage="code2wav", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(1,), final_output=True, final_output_type="audio", - hf_config_name="tts_config", engine_output_type="audio", - custom_process_input_func=f"{_PROC}.llm2tts", - sampling_constraints={"detokenize": False}, + model_arch="MiniCPMO45Code2Wav", + sampling_constraints={"detokenize": True}, ), ), ) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 6749d5cfc..6f0302499 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -360,6 +360,11 @@ _OMNI_MODELS = { "minicpmo_4_5_omni_tts", "MiniCPMO45OmniTTSForConditionalGeneration", ), + "MiniCPMO45Code2Wav": ( + "minicpmo_4_5", + "minicpmo_4_5_code2wav", + "MiniCPMO45Code2Wav", + ), "AuraQwen3VLForConditionalGeneration": ( "aura_omni", "qwen3_vl", diff --git a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py index 561c6ac8f..7d319de72 100644 --- a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py @@ -1,18 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Stage input processor for MiniCPM-o 4.5: Thinker (LLM) -> Talker (TTS). - -This bridge converts thinker hidden states and token ids into the talker -prompt payload. The talker runs Token2Wav internally; MiniCPM-o 4.5 does not -have an independent code2wav pipeline stage. -""" +"""MiniCPM-o 4.5 Thinker-to-Talker and Talker-to-Code2Wav bridges.""" import logging -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence +from typing import Any import torch from vllm.inputs import TextPrompt +from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct from vllm_omni.experimental.fullduplex.engine.intermediate import ( build_duplex_intermediate_buffer, set_ref_audio, @@ -21,6 +18,21 @@ from vllm_omni.experimental.fullduplex.engine.intermediate import ( from vllm_omni.inputs.data import OmniTokensPrompt logger = logging.getLogger(__name__) +_MINICPMO45_ASYNC_STATE = "_minicpmo45_async_codec_state" +_MINICPMO45_STREAM_RECORD = "_minicpmo45_async_stream_record" +_MINICPMO45_SILENCE_CODE = 4218 + + +class _MiniCPMO45MetaStruct(MetaStruct): + """Model-owned metadata for the split Talker-to-Code2Wav bridge.""" + + ref_audio_sr: int | None = None + native_duplex_segment_text: str | None = None + duplex_turn_id: int | None = None + duplex_epoch: int | None = None + segment_end: bool | None = None + turn_end: bool | None = None + tts_is_last_chunk: bool | None = None def _extract_first_audio_ref(multi_modal_data): @@ -111,6 +123,224 @@ def _coerce_int(value): return None +def _codec_config(transfer_manager: Any) -> tuple[int, int]: + connector = getattr(transfer_manager, "connector", None) + raw_config = getattr(connector, "config", {}) or {} + config = raw_config.get("extra", raw_config) if isinstance(raw_config, dict) else {} + config = config if isinstance(config, dict) else {} + chunk_frames = int(config.get("codec_chunk_frames", 25)) + left_context_frames = int(config.get("codec_left_context_frames", 3)) + if chunk_frames <= 0 or left_context_frames < 0: + raise ValueError( + "Invalid MiniCPM-o codec chunk config: " + f"codec_chunk_frames={chunk_frames}, " + f"codec_left_context_frames={left_context_frames}" + ) + return chunk_frames, left_context_frames + + +def _codec_scalars(value: Any) -> list[int]: + """Normalize one request-routed codec delta to CPU scalar token IDs.""" + if value is None: + return [] + if isinstance(value, torch.Tensor): + if value.numel() == 0: + return [] + return value.detach().to(device="cpu", dtype=torch.long).reshape(-1).tolist() + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + scalars: list[int] = [] + for item in value: + scalars.extend(_codec_scalars(item)) + return scalars + if isinstance(value, (int, bool)): + return [int(value)] + raise TypeError(f"Unsupported MiniCPM-o codec delta type: {type(value).__name__}") + + +def _extract_codec_delta(pooling_output: Any, request_id: str) -> list[int]: + if pooling_output is None: + return [] + if isinstance(pooling_output, Mapping): + meta = pooling_output.get("meta") + routed_id = meta.get("request_id") if isinstance(meta, Mapping) else None + routed_id = routed_id or pooling_output.get("request_id") + if routed_id is not None and str(routed_id) != request_id: + return [] + codes = pooling_output.get("codes") + audio = codes.get("audio") if isinstance(codes, Mapping) else None + return _codec_scalars(audio) + if isinstance(pooling_output, Sequence) and not isinstance( + pooling_output, + (str, bytes, bytearray), + ): + delta: list[int] = [] + for item in pooling_output: + values = _extract_codec_delta(item, request_id) if isinstance(item, Mapping) else _codec_scalars(item) + delta.extend(values) + return delta + return _codec_scalars(pooling_output) + + +def _drop_codec_state(transfer_manager: Any, request_id: str) -> None: + request_payload = getattr(transfer_manager, "request_payload", None) + if isinstance(request_payload, dict): + container = request_payload.get(request_id) + if isinstance(container, dict): + container.pop(_MINICPMO45_ASYNC_STATE, None) + if not container: + request_payload.pop(request_id, None) + else: + request_payload.pop(request_id, None) + code_accumulators = getattr(transfer_manager, "code_prompt_token_ids", None) + if hasattr(code_accumulators, "pop"): + code_accumulators.pop(request_id, None) + + +def _is_aborted(request: Any) -> bool: + status_name = getattr(getattr(request, "status", None), "name", "") + return any(marker in status_name for marker in ("ABORT", "CANCEL", "IGNORED", "ERROR")) + + +def tts2code2wav_async_chunk( + transfer_manager: Any, + multimodal_output: Any, + request: Any, + is_finished: bool = False, +) -> OmniPayloadStruct | None: + """Stream request-owned MiniCPM-o codec windows to Code2Wav.""" + external_id = getattr(request, "external_req_id", None) + internal_id = getattr(request, "request_id", None) + request_id = str(external_id if external_id is not None else internal_id) + internal_id = str(internal_id if internal_id is not None else request_id) + + request_payload = getattr(transfer_manager, "request_payload", None) + if request_payload is None: + request_payload = {} + transfer_manager.request_payload = request_payload + container = request_payload.get(request_id) + if not isinstance(container, dict): + container = {} + request_payload[request_id] = container + + record = container.get(_MINICPMO45_STREAM_RECORD) + if not isinstance(record, dict): + record = { + "internal_id": internal_id, + "cache_epoch": 0, + "retired_internal_ids": set(), + } + container[_MINICPMO45_STREAM_RECORD] = record + elif internal_id in record["retired_internal_ids"]: + return None + elif record["internal_id"] != internal_id: + record["retired_internal_ids"].add(record["internal_id"]) + record["internal_id"] = internal_id + record["cache_epoch"] = int(record["cache_epoch"]) + 1 + _drop_codec_state(transfer_manager, request_id) + + if _is_aborted(request): + record["retired_internal_ids"].add(internal_id) + _drop_codec_state(transfer_manager, request_id) + return None + + state = container.get(_MINICPMO45_ASYNC_STATE) + if not isinstance(state, dict): + request_info = getattr(request, "additional_information", None) + if not isinstance(request_info, Mapping): + request_info = {} + codes_info = request_info.get("codes") + if not isinstance(codes_info, Mapping): + codes_info = {} + meta_info = request_info.get("meta") + if not isinstance(meta_info, Mapping): + meta_info = {} + duplex_info = request_info.get("duplex") + if not isinstance(duplex_info, Mapping): + duplex_info = {} + state = { + "internal_id": internal_id, + "pending": [], + "left_context": [], + "codec_end": 0, + "ref_audio": codes_info.get("ref"), + "ref_audio_sr": meta_info.get("ref_audio_sr"), + "segment_text": meta_info.get("native_duplex_segment_text"), + "segment_end": bool(meta_info.get("segment_end", False)), + "turn_end": bool(meta_info.get("turn_end", False)), + "duplex_turn_id": duplex_info.get("model_turn_id", duplex_info.get("turn_id")), + "duplex_epoch": duplex_info.get("epoch"), + } + container[_MINICPMO45_ASYNC_STATE] = state + + pending = state["pending"] + pending.extend(_extract_codec_delta(multimodal_output, request_id)) + request_finished = getattr(request, "is_finished", None) + finished = bool(is_finished or (callable(request_finished) and request_finished())) + chunk_frames, left_context_frames = _codec_config(transfer_manager) + if not finished and len(pending) < chunk_frames: + return None + + new_token_count = len(pending) if finished else chunk_frames + new_codes = pending[:new_token_count] + del pending[:new_token_count] + codec_start = int(state["codec_end"]) + codec_end = codec_start + new_token_count + + if new_token_count: + if codec_start == 0: + context = [_MINICPMO45_SILENCE_CODE] * left_context_frames + else: + context = list(state["left_context"]) + output_codes = [*context, *new_codes] + history = [*state["left_context"], *new_codes] + state["left_context"] = history[-left_context_frames:] if left_context_frames else [] + elif finished and codec_start > 0 and state["left_context"]: + context = list(state["left_context"]) + output_codes = context + else: + context = [] + output_codes = [] + state["codec_end"] = codec_end + + last_chunk = bool(finished and not pending) + if last_chunk: + record["retired_internal_ids"].add(internal_id) + _drop_codec_state(transfer_manager, request_id) + + chunk_seq = int(getattr(transfer_manager, "put_req_chunk", {}).get(request_id, 0)) + finished_tensor = torch.tensor(last_chunk, dtype=torch.bool) + ref_audio = state.get("ref_audio") if codec_start == 0 else None + return OmniPayloadStruct( + codes=CodesStruct( + audio=torch.tensor(output_codes, dtype=torch.long), + ref=torch.as_tensor(ref_audio, dtype=torch.float32).reshape(-1) if ref_audio is not None else None, + ), + meta=_MiniCPMO45MetaStruct( + request_id=request_id, + chunk_seq=chunk_seq, + cache_epoch=int(record["cache_epoch"]), + code_flat_numel=len(output_codes), + codec_chunk_frames=new_token_count, + codec_left_context_frames=len(context), + left_context_size=len(context), + last_chunk=last_chunk, + stream_finished=finished_tensor, + finished=finished_tensor, + req_id=[request_id], + ref_audio_sr=_coerce_int(state.get("ref_audio_sr")), + native_duplex_segment_text=( + str(state["segment_text"]) if isinstance(state.get("segment_text"), str) else None + ), + duplex_turn_id=_coerce_int(state.get("duplex_turn_id")), + duplex_epoch=_coerce_int(state.get("duplex_epoch")), + segment_end=bool(state.get("segment_end", False)), + turn_end=bool(state.get("turn_end", False)), + tts_is_last_chunk=last_chunk, + ), + request_id=request_id, + ) + + def _special_token_ids_from_mm_output(mm_output): if not isinstance(mm_output, Mapping): return {} @@ -344,19 +574,7 @@ def llm2tts( requires_multimodal_data: bool = False, _streaming_context=None, ): - """Convert thinker stage output to talker stage input for MiniCPMO Omni. - - Extracts from thinker output: - - Full hidden states (prompt + generated) for speaker embedding extraction - - Prompt token IDs (for finding spk_bos/spk_eos positions) - - Generated token IDs (for decoding TTS text) - - The talker model will: - 1. Find <|spk_bos|>/<|spk_eos|> positions in prompt_token_ids - 2. Extract speaker embedding from hidden states at those positions - 3. Decode generated text and extract TTS content - 4. Run ConditionalChatTTS pipeline - """ + """Build Talker conditioning for ordinary and full-duplex streaming.""" if not source_outputs: raise ValueError("source_outputs cannot be empty") @@ -375,7 +593,14 @@ def llm2tts( for llm_output in llm_outputs: output = llm_output.outputs[0] - mm_output = output.multimodal_output if isinstance(output.multimodal_output, Mapping) else {} + request_mm_output = getattr(llm_output, "multimodal_output", None) + completion_mm_output = getattr(output, "multimodal_output", None) + if isinstance(request_mm_output, Mapping): + mm_output = request_mm_output + elif isinstance(completion_mm_output, Mapping): + mm_output = completion_mm_output + else: + mm_output = {} special_token_ids = _special_token_ids_from_mm_output(mm_output) prompt_token_ids = ( _coerce_token_id_list(mm_output.get("duplex_prompt_token_ids")) @@ -445,6 +670,11 @@ def llm2tts( for idx_t in range(search_start, len(full_token_ids)): if full_token_ids[idx_t] == tts_bos_id: tts_bos_idx = idx_t + 1 + if tts_bos_idx is None and not is_native_duplex_handoff and llm_output_ids: + # Audio routing is a model-stage concern, not an OpenAI serving + # default. Plain chat templates do not include <|tts_bos|>; in + # that case condition the Talker on the generated assistant span. + tts_bos_idx = prompt_token_ids_len tts_eos_idx = None if tts_bos_idx is not None: @@ -591,13 +821,21 @@ def llm2tts( continue set_tts_handoff(model_intermediate_buffer, handoff_ids, handoff_hidden) if native_turn_end_handoff: + model_intermediate_buffer.setdefault("meta", {})["turn_end"] = True _reset_native_tts_handoff(_streaming_context) - scheduler_prompt_token_ids = _build_tts_scheduler_prompt_token_ids( - tts_token_ids_slice, - llm_output_ids, - prompt_token_ids, - ) + if handoff_ids is not None and handoff_hidden is not None: + condition_length = max(len(handoff_ids), len(handoff_hidden)) + 2 + scheduler_prompt_token_ids = [0] * condition_length + handoff_meta = model_intermediate_buffer.setdefault("meta", {}) + handoff_meta["replace_streaming_prompt"] = True + handoff_meta["next_stage_prompt_len"] = condition_length + else: + scheduler_prompt_token_ids = _build_tts_scheduler_prompt_token_ids( + tts_token_ids_slice, + llm_output_ids, + prompt_token_ids, + ) tts_inputs.append( OmniTokensPrompt( prompt_token_ids=scheduler_prompt_token_ids, diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 0817d7bbf..904974829 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -43,6 +43,7 @@ from vllm_ascend.worker.model_runner_v1 import graph_capture from vllm_omni.data_entry_keys import flatten_payload from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager +from vllm_omni.distributed.omni_connectors.utils.config import stage_sends_async_output from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element @@ -1186,7 +1187,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): pooler_output.append(flatten_payload(payload)) pooler_output = pooler_output or [] - if self._async_chunk: + if self._async_chunk and stage_sends_async_output(self.model_config): pooler_inter, pooler_client = partition_payload_list(pooler_output) else: # Non-async-chunk ships the full payload to the next stage via @@ -1203,7 +1204,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): self.accumulate_full_payload_output(rid, pooler_inter[i], req_state) inter_stage_outputs = self._build_multimodal_outputs(pooler_inter) - multimodal_outputs = self._build_multimodal_outputs(pooler_client) + multimodal_outputs = ( + inter_stage_outputs if pooler_client is pooler_inter else self._build_multimodal_outputs(pooler_client) + ) model_runner_output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, diff --git a/vllm_omni/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py index 1dc4e37f7..af8f9ac14 100644 --- a/vllm_omni/utils/mm_outputs.py +++ b/vllm_omni/utils/mm_outputs.py @@ -32,6 +32,7 @@ _CLIENT_MM_META_KEYS: frozenset[str] = frozenset( "duplex_epoch", "duplex_turn_id", "llm_output_text_utf8", + "segment_end", "tts_is_last_chunk", "turn_end", } diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 3e1ea6ddb..ca7c2a21a 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -42,6 +42,10 @@ from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm_omni.data_entry_keys import flatten_payload from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager +from vllm_omni.distributed.omni_connectors.utils.config import ( + get_stage_connector_role, + stage_sends_async_output, +) from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner @@ -321,7 +325,10 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): "DyninOmniForConditionalGeneration", "IndexTTS2TalkerForConditionalGeneration", } - if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS: + if ( + getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS + or get_stage_connector_role(self.model_config) is not None + ): self.init_omni_connectors( model_config=self.model_config, kv_transfer_manager=self.kv_transfer_manager, @@ -1841,14 +1848,11 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): pooler_output.append(flatten_payload(payload)) pooler_output = pooler_output or [] - if self._async_chunk: + if self._async_chunk and stage_sends_async_output(self.model_config): pooler_inter, pooler_client = partition_payload_list(pooler_output) else: - # Non-async-chunk still ships the full payload to the next stage (via - # accumulate_full_payload_output and the inter_stage_outputs field); only - # client mm keys are split out when async_chunk is enabled. #4527 set this - # to (None, pooler_output), which skipped accumulation and starved the - # downstream stage (300s connector-input timeout / empty audio). (PR #4792) + # Connector-less stages expose the same payload through the + # orchestrator bridge; non-async stages preserve legacy behavior. pooler_inter, pooler_client = pooler_output, pooler_output if pooler_inter and self._should_accumulate_full_payload_output(): @@ -1860,7 +1864,9 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): with record_function_or_nullcontext("omni_output_builder:build_multimodal_outputs"): inter_stage_outputs = self._build_multimodal_outputs(pooler_inter) - multimodal_outputs = self._build_multimodal_outputs(pooler_client) + multimodal_outputs = ( + inter_stage_outputs if pooler_client is pooler_inter else self._build_multimodal_outputs(pooler_client) + ) with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): routed_experts_lists = None diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 6208f73e9..2c81643dd 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -36,6 +36,7 @@ from vllm.v1.worker.gpu_model_runner import ( from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs +from vllm_omni.distributed.omni_connectors.utils.config import get_stage_connector_role from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.utils.mm_outputs import partition_payload_list from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, _ensure_tensor_values @@ -68,7 +69,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin "DyninOmniForConditionalGeneration", "IndexTTS2S2MelDecoder", } - if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS: + if ( + getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS + or get_stage_connector_role(self.model_config) is not None + ): self.init_omni_connectors( model_config=self.model_config, ) @@ -317,6 +321,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin ) # [Omni] Pass token counts per request for code2wav output slicing model_kwargs["seq_token_counts"] = tokens + if getattr(self.model, "requires_request_ids", False): + model_kwargs["request_ids"] = list(req_ids) # Set cudagraph mode to none if calc_kv_scales is true. # KV scales calculation involves dynamic operations that are incompatible diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 0e14784ca..cb52a2b67 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1361,6 +1361,14 @@ class OmniGPUModelRunner(GPUModelRunner): if nstp is not None and len(nstp) == len(self.input_batch.req_ids): try: model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(nstp) + if getattr(self.model, "requires_request_sample_eligibility", False): + model_kwargs_extra["request_sample_eligible"] = [ + bool( + (req := self.requests.get(req_id)) is not None + and int(req.num_computed_tokens) + int(nstp[req_index]) >= int(req.num_tokens) + ) + for req_index, req_id in enumerate(self.input_batch.req_ids) + ] except Exception as e: # Visible on purpose: the fallback is the equal rows-per-request # split, which can re-introduce the cross-request corruption this @@ -1482,14 +1490,17 @@ class OmniGPUModelRunner(GPUModelRunner): if isinstance(model_buffer, dict) and model_buffer: self._update_intermediate_buffer(new_req.req_id, model_buffer) payload_info = getattr(new_req, "additional_information", None) - if isinstance(payload_info, dict): - self._update_intermediate_buffer(new_req.req_id, payload_info) + decoded_info = deserialize_additional_information(payload_info) + if decoded_info: + self._update_intermediate_buffer(new_req.req_id, decoded_info) if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"): cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {}) if isinstance(cached_infos, dict): for req_id, req_infos in cached_infos.items(): - self._update_intermediate_buffer(req_id, req_infos) + decoded_info = deserialize_additional_information(req_infos) + if decoded_info: + self._update_intermediate_buffer(req_id, decoded_info) def _maybe_attach_mimo_audio_req_infos( self, @@ -1661,9 +1672,8 @@ class OmniGPUModelRunner(GPUModelRunner): # (e.g., code2wav) so runtime_additional_information can be refreshed # from scheduler cached infos on every step. if hasattr(self.model, "has_preprocess") or hasattr(self.model, "enable_update_additional_information"): - if self.vllm_config.model_config.async_chunk: - self._update_additional_information(scheduler_output) - else: + self._update_additional_information(scheduler_output) + if not self.vllm_config.model_config.async_chunk: # In full-payload (non-async-chunk) mode, connector-delivered # stage payloads must override any earlier engine-level # additional_information written by the legacy @@ -1672,7 +1682,11 @@ class OmniGPUModelRunner(GPUModelRunner): self._sync_local_stage_payloads() if hasattr(self.model, "has_preprocess") and self.model.has_preprocess: - preprocess_device = input_ids.device if input_ids is not None else inputs_embeds.device + # Multimodal wrappers forward embeddings and therefore intentionally + # set ``input_ids`` to None. Model-specific preprocess hooks still + # need the scheduled token ids to build or replace those embeddings. + preprocess_input_ids = input_ids if input_ids is not None else self.input_ids.gpu[:num_input_tokens] + preprocess_device = preprocess_input_ids.device self._maybe_run_batch_preprocess(self.input_batch.req_ids, preprocess_device) # Overlay custom prompt_embeds per request for the prompt portion; @@ -1690,21 +1704,27 @@ class OmniGPUModelRunner(GPUModelRunner): req_ids_b = [item[0] for item in decode_batch_items] start_offsets_b = [item[1] for item in decode_batch_items] req_infos_b = [item[2] for item in decode_batch_items] - ids_b = torch.stack([input_ids[offset : offset + 1].reshape(-1)[0] for offset in start_offsets_b]) + ids_b = torch.stack( + [preprocess_input_ids[offset : offset + 1].reshape(-1)[0] for offset in start_offsets_b] + ) req_input_ids, req_embeds, last_talker_hidden, text_step, updates = batch_decode_preprocess( input_ids=ids_b, req_infos=req_infos_b, ) if inputs_embeds is None: inputs_embeds = torch.empty( - (input_ids.shape[0], req_embeds.shape[-1]), + (preprocess_input_ids.shape[0], req_embeds.shape[-1]), device=req_embeds.device, dtype=req_embeds.dtype, ) offsets_t = torch.tensor(start_offsets_b, device=req_embeds.device, dtype=torch.long) inputs_embeds.index_copy_(0, offsets_t, req_embeds) - input_ids.index_copy_(0, offsets_t, req_input_ids.reshape(-1).to(dtype=input_ids.dtype)) + preprocess_input_ids.index_copy_( + 0, + offsets_t, + req_input_ids.reshape(-1).to(dtype=preprocess_input_ids.dtype), + ) dst = slice(len(decode_req_ids), len(decode_req_ids) + len(req_ids_b)) self.talker_mtp_input_ids.gpu[dst].copy_(req_input_ids.reshape(-1)) @@ -1751,7 +1771,9 @@ class OmniGPUModelRunner(GPUModelRunner): embed_slice = inputs_embeds[s:e] if inputs_embeds is not None else None req_input_ids, req_embeds, update_dict = self.model.preprocess( - input_ids=preprocess_input_ids[s:e], input_embeds=embed_slice, **req_infos + input_ids=preprocess_input_ids[s:e], + input_embeds=embed_slice, + **req_infos, ) if inputs_embeds is None: inputs_embeds = torch.empty(