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

[perf] Minicpm o4.5 optimization (#5228)

Browse Source 
Signed-off-by: R2-Y <ruiruyang2@gmail.com>
Co-authored-by: Sy03 <1370724210@qq.com>
This commit is contained in:
rein yangandSy03 authored and GitHub committed 2026-07-24 17:20:39 +08:00
1 parent acf941da4d
commit ad3a162eb8
54 files changed
+3861 -3997
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

examples/online_serving/minicpmo/README.md
+100 -20
View File
@@ -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://<host>: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
`<think>` 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)
examples/online_serving/minicpmo/streaming_chat_completion.py
+71
View File
@@ -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()
recipes/OpenBMB/MiniCPM-o-4_5.md
+122 -96
View File
@@ -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.
tests/benchmarks/test_serve_cli.py
+16 -7
View File
@@ -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)
tests/core/sched/test_omni_ar_scheduler_streaming.py
+22
View File
@@ -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
tests/distributed/omni_connectors/test_adapter_and_flow.py
+17
View File
@@ -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
tests/distributed/omni_connectors/test_chunk_transfer_adapter.py
+52
View File
@@ -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)
tests/distributed/omni_connectors/test_omni_connector_configs.py
+25
View File
@@ -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):
tests/engine/test_orchestrator_stage_input_bridge.py
+73
View File
@@ -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()
tests/engine/test_stage_engine_core_proc.py
+28
View File
@@ -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"}
tests/examples/online_serving/test_minicpmo_streaming.py
+29
View File
@@ -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"]
tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
+431
View File
@@ -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
tests/model_executor/models/minicpmo_4_5/test_hf_model_path.py
-74
View File
@@ -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)
tests/model_executor/models/minicpmo_4_5/test_llm2tts.py
+48 -14
View File
@@ -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:
tests/model_executor/models/minicpmo_4_5/test_pipeline.py
+92 -11
View File
@@ -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."""
tests/model_executor/models/minicpmo_4_5/test_talker_batching.py
+248
View File
@@ -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
tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py
+182
View File
@@ -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)]
tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py
+93 -1276
View File
File diff suppressed because it is too large. Load diff
tests/test_config_factory.py
+18 -14
View File
@@ -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"
tests/worker/test_gpu_ar_model_runner.py
+1
View File
@@ -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]]))
tests/worker/test_native_duplex_hooks.py
-363
View File
@@ -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()
tests/worker/test_omni_gpu_model_runner.py
+27
View File
@@ -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)
vllm_omni/config/stage_config.py
+4
View File
@@ -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),
vllm_omni/core/sched/omni_ar_scheduler.py
+8 -2
View File
@@ -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
vllm_omni/core/sched/omni_generation_scheduler.py
+1 -23
View File
@@ -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)
vllm_omni/core/sched/omni_scheduler_mixin.py
+27 -1
View File
@@ -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.
vllm_omni/data_entry_keys.py
+14
View File
@@ -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
vllm_omni/deploy/minicpmo_4_5.yaml
+41 -16
View File
@@ -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:
vllm_omni/deploy/minicpmo_4_5_2gpu.yaml
+2 -64
View File
@@ -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
vllm_omni/deploy/minicpmo_4_5_3gpu.yaml
+45 -16
View File
@@ -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: {}
vllm_omni/deploy/minicpmo_4_5_8x4090.yaml
+37 -8
View File
@@ -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: {}
vllm_omni/deploy/minicpmo_4_5_batching.yaml
+88
View File
@@ -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
vllm_omni/deploy/minicpmo_4_5_duplex.yaml
+12 -10
View File
@@ -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
vllm_omni/distributed/omni_connectors/adapter.py
+16
View File
@@ -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
vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py
+11 -3
View File
@@ -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:
vllm_omni/distributed/omni_connectors/utils/config.py
+28
View File
@@ -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."""
vllm_omni/engine/orchestrator.py
+13 -1
View File
@@ -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()
vllm_omni/engine/stage_engine_core_proc.py
+7
View File
@@ -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,
vllm_omni/engine/stage_init_utils.py
+10
View File
@@ -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 {}
vllm_omni/engine/stage_pool.py
+2
View File
@@ -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,
vllm_omni/entrypoints/cli/benchmark/cli_args.py
+1 -1
View File
@@ -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
vllm_omni/experimental/fullduplex/engine/intermediate.py
+8 -9
View File
@@ -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:
vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py
+400
View File
@@ -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
vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py
+547
View File
@@ -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()}
vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py
+29 -220
View File
@@ -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))
vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py
+469 -1678
View File
File diff suppressed because it is too large. Load diff
vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py
+17 -22
View File
@@ -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},
),
),
)
vllm_omni/model_executor/models/registry.py
+5
View File
@@ -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",
vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
+264 -26
View File
@@ -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,
vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+5 -2
View File
@@ -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,
vllm_omni/utils/mm_outputs.py
+1
View File
@@ -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",
}
vllm_omni/worker/gpu_ar_model_runner.py
+14 -8
View File
@@ -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
vllm_omni/worker/gpu_generation_model_runner.py
+7 -1
View File
@@ -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
vllm_omni/worker/gpu_model_runner.py
+33 -11
View File
@@ -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(