Merge branch 'main' into qwen-image-2.1
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
This commit is contained in:
51 files changed
+2189
-70
No files matched your search
@@ -64,3 +64,21 @@ steps:
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=1
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py::test_text_to_text_audio_001 -m "full_model and rocm and MI325 and cards_2" --run-level "full_model"
|
||||
|
||||
- group: ":speaker: CosyVoice3 MI300 Nightly"
|
||||
steps:
|
||||
# Exercise all sync/async and Chinese/English cases with the platform's
|
||||
# default attention backend. On ROCm this retains end-to-end AITER
|
||||
# coverage without putting its cold JIT build on the READY critical path.
|
||||
- label: "CosyVoice3-TTS E2E Test"
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: NonBlocking
|
||||
timeout_in_minutes: 90
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: 134
|
||||
limit: 1
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"
|
||||
@@ -328,7 +328,10 @@ steps:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
|
||||
|
||||
- label: "CosyVoice3-TTS E2E Test"
|
||||
# Keep a fast model/server integration check in the blocking lane. The
|
||||
# full default-backend suite remains in the nightly pipeline because a
|
||||
# cold ROCm AITER MHA build can add more than 25 minutes to this job.
|
||||
- label: "CosyVoice3-TTS E2E Smoke (SDPA)"
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
@@ -340,7 +343,8 @@ steps:
|
||||
- exit_status: 134
|
||||
limit: 1
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"
|
||||
- export DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA
|
||||
- pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py::test_voice_clone_zh_002 -m "slow" --run-level "core_model"
|
||||
|
||||
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
|
||||
# Need to debug before reneable numerical changes across large PRs
|
||||
|
||||
@@ -255,6 +255,8 @@ source_file_dependencies:
|
||||
- *minicpmo_4_5_duplex
|
||||
- tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
|
||||
- tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
|
||||
- tests/e2e/online_serving/test_minicpmo_4_5_window.py
|
||||
- tests/e2e/online_serving/helpers/minicpmo_window_e2e.py
|
||||
- tests/e2e/online_serving/test_duplex_client_live.py
|
||||
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
|
||||
- tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
|
||||
|
||||
@@ -130,7 +130,7 @@ steps:
|
||||
source_file_dependencies: omni_minicpmo_4_5_duplex_function
|
||||
timeout_in_minutes: 50
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py tests/e2e/online_serving/test_minicpmo_4_5_window.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Diffusion · Qwen Image 2.1 Test"
|
||||
|
||||
@@ -88,6 +88,7 @@ depends on the installed kernels and model path.
|
||||
| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; defaults to on for Qwen3-TTS and off for other pipelines | Pipeline default computed from `pipeline_config.model_type` at engine initialization; env override resolved at `Orchestrator` construction and separately when the serving-side final-output drain starts | An explicit env value wins; otherwise both consumers use the engine's pipeline default. Values are stripped and case-normalized; any unrecognized value selects the legacy poll loop. Set before server startup. | Experimental |
|
||||
| `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control |
|
||||
| `VLLM_OMNI_ORCH_MONITOR_PATH` | Filesystem path; default `<current-working-directory>/vllm_omni_orch_monitor_<timestamp>.json` | Orchestrator monitor enabled by `--enable-orch-monitor`; read when the monitor is created | Environment-only path override. Parent directories are created; write errors are logged. | Diagnostic |
|
||||
| `VLLM_OMNI_SPEAKER_REGISTRATION_POLICY` | `overwrite` or `immutable`; default `overwrite` | Speech server; read when speaker storage initializes | Environment-only setting. `immutable` rejects re-registering an existing uploaded voice name until it is deleted; any other value raises `ValueError` at startup. | Experimental |
|
||||
| `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` | Float seconds; default `600` | Synchronous Videos API; read when the API server module imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental |
|
||||
| `VLLM_VIDEO_ASYNC_CHUNK` | `on` or `off`; default `on` | Streaming video output; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `on`. | Experimental |
|
||||
| `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast` or `slow`; default `fast` | Streaming video audio deltas; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `fast`. | Experimental |
|
||||
|
||||
@@ -180,17 +180,52 @@ cache_dit.enable_cache(
|
||||
|
||||
> **Note:** For single transformer with multiple block lists, `refresh_context` works the same as standard models.
|
||||
|
||||
### Example 3: Sequential-CFG Diffusion Pipeline (MammothModa2)
|
||||
|
||||
MammothModa2's DiT stage runs on the diffusion runner, which enables the
|
||||
configured backend at startup and transfers ownership to the pipeline through
|
||||
the request-scoped protocol (`adopt_cache_dit_backend` / `is_cache_dit_enabled`);
|
||||
`forward()` then reconciles per-request state through
|
||||
`RequestScopedCacheDiTRuntime` (MiniMax H3 precedent).
|
||||
|
||||
**Key difference:** The enabler caches only the repeated main-layer stack and
|
||||
marks the pipeline as requiring paired CFG forwards, because cache-dit's
|
||||
separate-CFG accounting distinguishes cond/uncond passes purely by forward
|
||||
parity:
|
||||
|
||||
```python
|
||||
pipeline._cache_dit_requires_paired_cfg = True
|
||||
block_adapter = BlockAdapter(
|
||||
transformer=pipeline.gen_transformer,
|
||||
blocks=[pipeline.gen_transformer.layers], # Refiners stay outside
|
||||
forward_pattern=[ForwardPattern.Pattern_3],
|
||||
has_separate_cfg=True,
|
||||
check_forward_pattern=True,
|
||||
)
|
||||
```
|
||||
|
||||
Per request, the pipeline calls `runtime.prepare(spec)` for CFG requests
|
||||
(refreshing the step context) and `runtime.prepare(None)` for `guidance=1.0`
|
||||
requests (disabling hooks — their single forward per step cannot be expressed
|
||||
by the parity accounting).
|
||||
|
||||
### Registering Custom Implementations
|
||||
|
||||
After writing your custom enabler, register it in `CUSTOM_DIT_ENABLERS` in
|
||||
After writing your custom enabler, register it in
|
||||
`CUSTOM_DIT_ENABLERS` via `register_custom_dit_enablers()` in
|
||||
`vllm_omni/diffusion/cache/cachedit/model_specific.py`:
|
||||
|
||||
```python
|
||||
CUSTOM_DIT_ENABLERS = {
|
||||
"Wan22Pipeline": enable_cache_for_wan22,
|
||||
"LongCatImagePipeline": enable_cache_for_longcat_image,
|
||||
"YourCustomPipeline": enable_cache_for_your_model, # Add here
|
||||
}
|
||||
CUSTOM_DIT_ENABLERS.update(
|
||||
{
|
||||
"Wan22Pipeline": enable_cache_for_wan22,
|
||||
"Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3,
|
||||
"Krea2Pipeline": enable_cache_for_krea2,
|
||||
"Magi2Pipeline": enable_cache_for_magi2,
|
||||
"MammothModa2DiTPipeline": enable_cache_for_mammothmoda2,
|
||||
"YourCustomPipeline": enable_cache_for_your_model, # Add here
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
---
|
||||
@@ -274,6 +309,7 @@ Complete examples in the codebase:
|
||||
| **Wan2.2** | `cachedit.model_specific::enable_cache_for_wan22` | Single or dual-transformer | Auto-detects mode based on transformer_2 presence |
|
||||
| **LongCat** | `cachedit.config::CacheDiTAdapterConfig` | Declarative block adapter | Two block lists in one transformer |
|
||||
| **BAGEL** | `cachedit.model_specific::BagelCachedAdapter` | Custom cached adapter | Complex architecture |
|
||||
| **MammothModa2** | `cachedit.model_specific::enable_cache_for_mammothmoda2` | Sub-stack + paired CFG | Generation-runner pipeline; lifecycle owned by `pipeline_mammothmoda2_dit.py` |
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -289,6 +289,19 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
|
||||
|
||||
Fields `ref_text` and `speaker_description` are omitted when not provided at upload time.
|
||||
|
||||
**Naming rules:**
|
||||
|
||||
- Names that collide with one of the model's built-in or precomputed voices are rejected (400). Voice
|
||||
files already on disk under such a name are ignored at startup with a warning, so an upload can never
|
||||
shadow a built-in voice.
|
||||
- Re-uploading an existing uploaded name overwrites it in place (the previous audio file is
|
||||
deleted). Set `VLLM_OMNI_SPEAKER_REGISTRATION_POLICY=immutable` on the server to reject duplicates instead,
|
||||
requiring an explicit `DELETE /v1/audio/voices/{name}` before re-registering — useful when the
|
||||
endpoint is reachable by multiple writers and silent overwrites are a risk.
|
||||
- The voice registry has no per-user ownership: any client that can reach the endpoint can
|
||||
overwrite (default policy) or delete any uploaded voice. For multi-tenant deployments, add
|
||||
authentication at a proxy and namespace voice names per user (e.g. `{user}.{name}`).
|
||||
|
||||
**Usage Example:**
|
||||
|
||||
```bash
|
||||
@@ -577,6 +590,7 @@ by `GET /v1/audio/voices`. Valid precomputed voices can be used in
|
||||
| ---------- | --------- | ------------- |
|
||||
| `SPEAKER_SAMPLES_DIR` | `~/.cache/vllm-omni/speakers` | Directory for persisted uploaded speakers (`.safetensors` files). |
|
||||
| `SPEAKER_MAX_UPLOADED` | `1000` | Maximum number of uploaded speakers kept on disk. Upload requests past the cap return 400. |
|
||||
| `VLLM_OMNI_SPEAKER_REGISTRATION_POLICY` | `overwrite` | `immutable` rejects re-uploading an existing uploaded name (400) until it is deleted; any other value fails startup. |
|
||||
|
||||
The in-memory LRU has a fixed 512 MiB byte budget.
|
||||
|
||||
|
||||
@@ -72,6 +72,39 @@ omni = Omni(
|
||||
|
||||
## Example Script
|
||||
|
||||
### Multi-Stage Models (MammothModa2)
|
||||
|
||||
MammothModa2 is a multi-stage AR→DiT pipeline; the diffusion runner owns the
|
||||
Cache-DiT lifecycle of its DiT stage. Configure it on the **DiT stage entry**
|
||||
of the deploy YAML so the acceleration only applies to the denoising stage:
|
||||
|
||||
```yaml
|
||||
# deploy YAML (see vllm_omni/deploy/mammoth_moda2.yaml)
|
||||
stages:
|
||||
- stage_id: 1
|
||||
# ... other DiT stage settings ...
|
||||
cache_backend: cache_dit
|
||||
cache_config:
|
||||
Fn_compute_blocks: 1
|
||||
Bn_compute_blocks: 0
|
||||
max_warmup_steps: 4
|
||||
residual_diff_threshold: 0.24
|
||||
max_continuous_cached_steps: 3
|
||||
enable_cache_dit_summary: true # log skip-ratio stats per request
|
||||
```
|
||||
|
||||
Behavior notes:
|
||||
|
||||
- Only the repeated main-layer stack is cached; the Q-Former refiners always run.
|
||||
- Sequential-CFG parity requires the unconditional pass on every denoise step,
|
||||
so the `cfg_range` skip optimization is not used: outside the interval CFG is
|
||||
neutralized with `scale=1.0` instead of skipping the uncond forward.
|
||||
- Requests with `text_guidance_scale = 1.0` run with cache hooks disabled.
|
||||
- `residual_diff_threshold` defaults to `0.24` (measured ~1.7-2.2x at 50 steps,
|
||||
PSNR 27-34 dB vs uncached). For 20-step generation — especially with a
|
||||
partial `cfg_range` — use `0.12` for noticeably closer parity at a modest
|
||||
speed cost.
|
||||
|
||||
### Offline Inference
|
||||
|
||||
Use the example script under `examples/offline_inference/text_to_image`:
|
||||
|
||||
@@ -130,7 +130,7 @@ The following tables show which models support each feature:
|
||||
| **Krea 2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **MammothModa2(T2I)** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **OmniGen2** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
@@ -155,7 +155,8 @@ The following tables show which models support each feature:
|
||||
> 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
|
||||
> 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
|
||||
> 6. BAGEL step execution supports image generation with `bagel.yaml`, `bagel_think.yaml`, and `bagel_single_stage.yaml`; two-stage Thinker execution and explicit single-stage text output remain on their existing complete-request paths. Image requests require `num_inference_steps >= 2`. BAGEL step execution cannot currently be combined with sequence parallelism or a diffusion cache backend; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
|
||||
> 7. Qwen-Image-2.1 (`QwenImage21Pipeline`) handles both text-to-image and image-conditioned generation (up to 4 condition images) and supports LoRA through `QwenImageLoraLoaderMixin`. Sequence parallelism is Ulysses-only (ring attention is not supported), and cache acceleration (TeaCache / Cache-DiT) is unavailable; its built-in KV cache for the text/condition prefix is a separate feature from those cache backends.
|
||||
> 7. MammothModa2 runs its DiT stage on the diffusion runner (`StageExecutionType.DIFFUSION`); Cache-DiT is enabled through the standard diffusion-stage knobs on the stage entry of the deploy YAML (`cache_backend: cache_dit`, plus optional `cache_config` / `enable_cache_dit_summary`). The runner installs the backend at startup and the pipeline adopts it per request. Only the repeated main-layer stack is cached; requests with `text_guidance_scale = 1.0` bypass the cache hooks.
|
||||
> 8. Qwen-Image-2.1 (`QwenImage21Pipeline`) handles both text-to-image and image-conditioned generation (up to 4 condition images) and supports LoRA through `QwenImageLoraLoaderMixin`. Sequence parallelism is Ulysses-only (ring attention is not supported), and cache acceleration (TeaCache / Cache-DiT) is unavailable; its built-in KV cache for the text/condition prefix is a separate feature from those cache backends.
|
||||
|
||||
### VideoGen
|
||||
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech>.
|
||||
|
||||
|
||||
vLLM-Omni supports several autoregressive TTS models. They share a common
|
||||
CLI shape (`--text`, `--ref-audio`, `--ref-text`, `--output-dir`) and live
|
||||
together in this hub. Each model has its own subdirectory containing a
|
||||
@@ -15,7 +14,7 @@ list of supported architectures across all modalities, see
|
||||
## Supported Models
|
||||
|
||||
| Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
|
||||
|---|---|---|---|---|---|---|
|
||||
| --- | --- | --- | --- | --- | --- | --- |
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | — | continuation (`--ref-audio` + `--ref-text`) | 48 kHz |
|
||||
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ (`async_chunk: true` default) | — | 22.05 kHz |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ (`--streaming`) | — | 44.1 kHz |
|
||||
@@ -45,13 +44,15 @@ python examples/offline_inference/text_to_speech/<model>/end2end.py \
|
||||
Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
```bash
|
||||
pip install voxcpm
|
||||
# or, for a local source checkout:
|
||||
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/voxcpm
|
||||
pip install -e /path/to/voxcpm
|
||||
```
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
|
||||
--model openbmb/VoxCPM2 \
|
||||
@@ -59,7 +60,9 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
|
||||
Pass a reference audio for isolated cloning, or both `--ref-audio` + `--ref-text` for prompt continuation:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
|
||||
--text "Hello, this is a voice clone demo." \
|
||||
@@ -68,6 +71,7 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 48 kHz mono WAV.
|
||||
- Deploy config: `vllm_omni/deploy/voxcpm2.yaml` (default).
|
||||
|
||||
@@ -78,12 +82,14 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
|
||||
2-stage TTS pipeline (`talker` + `code2wav`) at 22.05 kHz.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
```bash
|
||||
uv pip install -e .
|
||||
# Includes soundfile, onnxruntime, x-transformers, einops via requirements.
|
||||
```
|
||||
|
||||
Download the model snapshot:
|
||||
|
||||
```python
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512',
|
||||
@@ -91,15 +97,18 @@ snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512',
|
||||
```
|
||||
|
||||
If your downloaded checkpoint lacks `config.json`, add it:
|
||||
|
||||
```json
|
||||
{
|
||||
"model_type": "cosyvoice3",
|
||||
"architectures": ["CosyVoice3Model"]
|
||||
}
|
||||
```
|
||||
|
||||
This is required because `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers the class mapping; the loader still reads `model_type` from `config.json` to select the class.
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
|
||||
--model pretrained_models/Fun-CosyVoice3-0.5B \
|
||||
@@ -107,7 +116,9 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
|
||||
Pass a reference audio. Note that CosyVoice3's `--prompt-text` is a system-style prompt for the GPT stage, not a reference transcript:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
|
||||
--model pretrained_models/Fun-CosyVoice3-0.5B \
|
||||
@@ -117,6 +128,7 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
|
||||
- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
|
||||
- `async_chunk: true` is the default; pass `--no-async-chunk` to switch to the legacy synchronous path.
|
||||
@@ -129,12 +141,14 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
|
||||
No extra packages are required; the DAC codec is vendored in vLLM-Omni.
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
|
||||
--text "Hello, this is a test of the Fish Speech text to speech system."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
|
||||
--text "Hello, this is a cloned voice." \
|
||||
@@ -143,14 +157,17 @@ python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
|
||||
```
|
||||
|
||||
### Streaming
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
|
||||
--text "Hello, this is a streaming test." \
|
||||
--streaming
|
||||
```
|
||||
|
||||
Streaming requires `async_chunk: true` in the deploy config.
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 44.1 kHz mono WAV.
|
||||
- DAC codec weights (`codec.pth`) are loaded lazily from the model directory.
|
||||
|
||||
@@ -161,6 +178,7 @@ Streaming requires `async_chunk: true` in the deploy config.
|
||||
2-stage TTS pipeline (AR + DiT flow-matching) at 24 kHz. Every request requires reference audio and its transcript for zero-shot voice cloning.
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
|
||||
--model zai-org/GLM-TTS \
|
||||
@@ -171,12 +189,14 @@ python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
|
||||
```
|
||||
|
||||
### Architecture
|
||||
```
|
||||
|
||||
```text
|
||||
Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz)
|
||||
(Llama-based) (32k vocab) (Flow Matching)
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- `--ref-audio` and `--ref-text` are **required** together; GLM-TTS does not support text-only synthesis.
|
||||
- Reference audio should be 3-10 seconds.
|
||||
- First run may be slow due to lazy loading of WhisperVQ tokenizer and CampPlus ONNX speaker embedder.
|
||||
@@ -190,12 +210,15 @@ Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24
|
||||
Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design).
|
||||
|
||||
### Prerequisites
|
||||
|
||||
```bash
|
||||
huggingface-cli download k2-fsa/OmniVoice
|
||||
```
|
||||
|
||||
Voice cloning requires `transformers>=5.3.0`. Auto and design modes work with `transformers>=4.57.0`.
|
||||
|
||||
### Quick start (auto voice)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
@@ -203,6 +226,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
@@ -212,6 +236,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
```
|
||||
|
||||
### Voice design
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
@@ -220,6 +245,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
```
|
||||
|
||||
### Language hint
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
@@ -228,6 +254,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking.
|
||||
- Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz.
|
||||
|
||||
@@ -238,13 +265,16 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
3-task-variant TTS with 24 kHz output. Has its own argparse surface (this script does not follow the common `--text` / `--ref-audio` shape).
|
||||
|
||||
### Prerequisites
|
||||
|
||||
For ROCm builds, replace `onnxruntime` with `onnxruntime-rocm`:
|
||||
|
||||
```bash
|
||||
pip uninstall onnxruntime
|
||||
pip install onnxruntime-rocm
|
||||
```
|
||||
|
||||
### Task variants
|
||||
|
||||
- `CustomVoice`: predefined speaker (speaker ID) with optional style instruction.
|
||||
- `VoiceDesign`: text + descriptive instruction designs a new voice.
|
||||
- `Base`: voice cloning from reference audio + transcript.
|
||||
@@ -264,16 +294,20 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty
|
||||
```
|
||||
|
||||
### Streaming
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
--query-type CustomVoice \
|
||||
--streaming \
|
||||
--output-dir /tmp/out_stream
|
||||
```
|
||||
|
||||
Streaming requires `async_chunk: true` in the deploy config.
|
||||
|
||||
### Batched decoding
|
||||
|
||||
The Code2Wav stage supports batched decoding through the SpeechTokenizer. Configure both stages with `max_num_seqs > 1` via `--stage-overrides` and pass multiple prompts via `--txt-prompts`:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
--query-type CustomVoice \
|
||||
@@ -281,9 +315,11 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
--batch-size 4 \
|
||||
--stage-overrides '{"0":{"max_num_seqs":4,"gpu_memory_utilization":0.2},"1":{"max_num_seqs":4,"gpu_memory_utilization":0.2}}'
|
||||
```
|
||||
|
||||
`--batch-size` must match a CUDA-graph capture size (1, 2, 4, 8, 16…).
|
||||
|
||||
### Notes
|
||||
|
||||
- Run `--help` for the full argument surface.
|
||||
- See `qwen3_tts/end2end.py` for the prompt-length-estimation logic the Talker uses.
|
||||
|
||||
@@ -294,12 +330,15 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
Voxtral-4B-TTS (Mistral). Has its own argparse surface; uses voice presets and the `mistral_common` `SpeechRequest` protocol.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Latest `mistral_common` with `SpeechRequest` support:
|
||||
|
||||
```bash
|
||||
pip install -e /path/to/mistral-common # or upgrade from PyPI when available
|
||||
```
|
||||
|
||||
### Quick start (voice preset)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
|
||||
--write-audio --voice cheerful_female \
|
||||
@@ -308,6 +347,7 @@ python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
|
||||
```
|
||||
|
||||
### Voice cloning (capability gated upstream)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
|
||||
--write-audio \
|
||||
@@ -317,15 +357,18 @@ python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
|
||||
```
|
||||
|
||||
### Streaming + concurrency
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
|
||||
--num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "..."
|
||||
```
|
||||
|
||||
Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
|
||||
|
||||
### Notes
|
||||
|
||||
- `--num-prompts N` replicates the prompt for performance measurement.
|
||||
- `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
|
||||
- Run `--help` for the full argument surface.
|
||||
|
||||
@@ -158,7 +158,11 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument(
|
||||
"--enable-cache-dit-summary",
|
||||
action="store_true",
|
||||
help="Enable cache-dit summary logging after diffusion forward passes.",
|
||||
default=None,
|
||||
help=(
|
||||
"Enable cache-dit summary logging after diffusion forward passes. "
|
||||
"Default: unset (defer to the deploy YAML's enable_cache_dit_summary)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ulysses-degree",
|
||||
|
||||
@@ -563,7 +563,7 @@ Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → F
|
||||
```bash
|
||||
pip install voxcpm
|
||||
# or, for a local source checkout:
|
||||
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/voxcpm
|
||||
pip install -e /path/to/voxcpm
|
||||
```
|
||||
|
||||
### Quick start
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Offline VoxCPM2 inference example (native AR pipeline).
|
||||
|
||||
Uses the single-stage native AR config (voxcpm2.yaml).
|
||||
Requires the `voxcpm` package or VLLM_OMNI_VOXCPM_CODE_PATH env var.
|
||||
Requires the `voxcpm` package (pip install voxcpm).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -67,11 +67,31 @@ def test_ready_diffusion_cpu_suite_is_sharded() -> None:
|
||||
assert "--shard-id=$$BUILDKITE_PARALLEL_JOB" in pytest_command
|
||||
|
||||
|
||||
def test_cosyvoice_gpu_abort_gets_one_fresh_job_retry() -> None:
|
||||
step = _find_step("CosyVoice3-TTS E2E Test", AMD_READY_PIPELINE)
|
||||
def test_cosyvoice_ready_smoke_uses_sdpa() -> None:
|
||||
step = _find_step("CosyVoice3-TTS E2E Smoke (SDPA)", AMD_READY_PIPELINE)
|
||||
|
||||
assert step["grade"] == "Blocking"
|
||||
assert step["retry"] == {"automatic": [{"exit_status": 134, "limit": 1}]}
|
||||
assert "export DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA" in step["commands"]
|
||||
|
||||
pytest_command = next(command for command in step["commands"] if "pytest" in command)
|
||||
assert "tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py::test_voice_clone_zh_002" in pytest_command
|
||||
|
||||
|
||||
def test_cosyvoice_full_default_backend_suite_runs_nightly() -> None:
|
||||
step = _find_step("CosyVoice3-TTS E2E Test", AMD_NIGHTLY_PIPELINE)
|
||||
|
||||
assert step["grade"] == "NonBlocking"
|
||||
assert step["timeout_in_minutes"] == 90
|
||||
assert step["retry"] == {"automatic": [{"exit_status": 134, "limit": 1}]}
|
||||
assert all("DIFFUSION_ATTENTION_BACKEND" not in command for command in step["commands"])
|
||||
|
||||
pytest_command = next(command for command in step["commands"] if "pytest" in command)
|
||||
assert "tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py" in pytest_command
|
||||
assert "::" not in pytest_command
|
||||
|
||||
|
||||
def test_amd_template_preserves_step_retry_policy() -> None:
|
||||
template = AMD_TEMPLATE.read_text(encoding="utf-8")
|
||||
# Both grouped and top-level AMD steps must preserve an explicit retry
|
||||
# policy when the source suite is rendered into the uploaded pipeline.
|
||||
|
||||
@@ -69,6 +69,39 @@ def _make_update(prompt_token_ids: list[int] | None = None) -> StreamingUpdate:
|
||||
)
|
||||
|
||||
|
||||
def _make_minicpm_window_update(
|
||||
*,
|
||||
seq: int,
|
||||
mode: str,
|
||||
high: int = 15,
|
||||
low: int = 10,
|
||||
context_max_units: int = 1,
|
||||
) -> StreamingUpdate:
|
||||
update = _make_update([0] * 8)
|
||||
update.model_intermediate_buffer = {
|
||||
"duplex": {
|
||||
"data_plane": True,
|
||||
"seq": seq,
|
||||
"runtime_config": {
|
||||
"duplex_scheduler_token_id": 0,
|
||||
"duplex_first_append_context_tokens": 3,
|
||||
"duplex_window_prefix_tokens": 2,
|
||||
"duplex_window_suffix_token_ids": [3],
|
||||
"duplex_window_previous_marker_token_ids": [70, 71],
|
||||
"duplex_window_special_token_ids": [99],
|
||||
"duplex_window_config": {
|
||||
"sliding_window_mode": mode,
|
||||
"basic_window_high_tokens": high,
|
||||
"basic_window_low_tokens": low,
|
||||
"context_previous_max_tokens": 4,
|
||||
"context_max_units": context_max_units,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
return update
|
||||
|
||||
|
||||
def _make_talker_adapter(
|
||||
*,
|
||||
max_model_len: int = 100,
|
||||
@@ -510,6 +543,225 @@ def test_stage0_streaming_update_keeps_all_computed_tokens_without_placeholder()
|
||||
assert sched._new_prompt_len_snapshot[session.request_id] == 2
|
||||
|
||||
|
||||
def test_stage0_basic_window_rebuilds_below_low_watermark() -> None:
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
session = _make_request()
|
||||
session.prompt_token_ids = [0] * 9
|
||||
session._all_token_ids.clear()
|
||||
session._all_token_ids.extend(session.prompt_token_ids)
|
||||
session.num_prompt_tokens = 9
|
||||
session.append_output_token_ids([40])
|
||||
session.num_computed_tokens = 10
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
update = _make_minicpm_window_update(seq=2, mode="basic")
|
||||
|
||||
sched._update_request_as_session(session, update)
|
||||
|
||||
assert session.prompt_token_ids == [0] * 9
|
||||
assert session.num_computed_tokens == 0
|
||||
assert update.model_intermediate_buffer["meta"]["replace_streaming_prompt"] is True
|
||||
plan = update.model_intermediate_buffer["duplex"]["stage0_window"]
|
||||
assert plan == {
|
||||
"completed_token_ids": [40],
|
||||
"replace": True,
|
||||
"mode": "basic",
|
||||
"drop_units": 1,
|
||||
"dropped_tokens": 9,
|
||||
"previous_token_ids": [],
|
||||
"previous_marker_token_ids": [70, 71],
|
||||
"replacement_prompt_len": 9,
|
||||
}
|
||||
sched._free_request_blocks.assert_called_once_with(session)
|
||||
|
||||
|
||||
def test_stage0_context_window_compacts_dropped_speech() -> None:
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
session = _make_request()
|
||||
session.prompt_token_ids = [0] * 9
|
||||
session._all_token_ids.clear()
|
||||
session._all_token_ids.extend(session.prompt_token_ids)
|
||||
session.num_prompt_tokens = 9
|
||||
session.append_output_token_ids([40])
|
||||
session.num_computed_tokens = 10
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
|
||||
first = _make_minicpm_window_update(seq=2, mode="context")
|
||||
sched._update_request_as_session(session, first)
|
||||
assert session.num_prompt_tokens == 18
|
||||
assert first.model_intermediate_buffer["duplex"]["stage0_window"] == {"completed_token_ids": [40]}
|
||||
|
||||
session.append_output_token_ids([50])
|
||||
session.num_computed_tokens = 19
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
second = _make_minicpm_window_update(seq=3, mode="context")
|
||||
sched._update_request_as_session(session, second)
|
||||
|
||||
assert session.prompt_token_ids == [0] * 21
|
||||
plan = second.model_intermediate_buffer["duplex"]["stage0_window"]
|
||||
assert plan["completed_token_ids"] == [50]
|
||||
assert plan["drop_units"] == 1
|
||||
assert plan["previous_token_ids"] == [40]
|
||||
# The worker embeds this exact marker, not a re-tokenized one, so the
|
||||
# previous region is len(marker) + len(previous).
|
||||
assert plan["previous_marker_token_ids"] == [70, 71]
|
||||
assert plan["replacement_prompt_len"] == 21
|
||||
assert getattr(session, "_minicpmo45_window_previous_len") == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cleared_outputs", [False, True])
|
||||
@pytest.mark.parametrize("in_flight", [0, 1])
|
||||
def test_stage0_window_uses_confirmed_span_and_terminator(cleared_outputs, in_flight) -> None:
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
session = _make_request()
|
||||
session.prompt_token_ids = [0] * 9
|
||||
session._all_token_ids[:] = session.prompt_token_ids
|
||||
session.num_prompt_tokens = 9
|
||||
session.append_output_token_ids([40, 99])
|
||||
session.num_computed_tokens = 10 + in_flight
|
||||
session.num_output_placeholders = in_flight
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
if cleared_outputs:
|
||||
# A stop without a queued append clears this before the next update.
|
||||
session._output_token_ids.clear()
|
||||
update = _make_minicpm_window_update(seq=2, mode="basic")
|
||||
|
||||
sched._update_request_as_session(session, update)
|
||||
|
||||
plan = update.model_intermediate_buffer["duplex"]["stage0_window"]
|
||||
assert plan["completed_token_ids"] == [40]
|
||||
assert plan["completed_terminator_token_id"] == 99
|
||||
assert plan["replacement_prompt_len"] == 9
|
||||
assert plan["dropped_tokens"] == 9
|
||||
|
||||
|
||||
def test_stage0_window_rebuild_that_overflows_max_model_len_finishes_the_session() -> None:
|
||||
"""A rebuilt window prompt is bounded by the client's window settings, not
|
||||
by the model: a unit carrying camera frames is hundreds of tokens, and a
|
||||
``context`` window adds the ``previous`` region on top. The replacement
|
||||
branch must therefore check ``replacement_prompt_len`` against
|
||||
``max_model_len - sample_room`` and finish the session with
|
||||
``context_length_exceeded``, exactly like a plain extension."""
|
||||
sched = _make_live_session_scheduler(max_model_len=21)
|
||||
session = _make_request()
|
||||
session.prompt_token_ids = [0] * 9
|
||||
session._all_token_ids[:] = [0] * 9
|
||||
session.num_prompt_tokens = 9
|
||||
session._output_token_ids[:] = [40]
|
||||
session._all_token_ids.append(40)
|
||||
session.num_computed_tokens = 10
|
||||
session.num_output_placeholders = 0
|
||||
_park_session(sched, session)
|
||||
# The parked append state above is the frame before this update: one
|
||||
# confirmed output token, fully computed.
|
||||
session._output_token_ids[:] = [40]
|
||||
session._all_token_ids[:] = [*session.prompt_token_ids, 40]
|
||||
session.num_computed_tokens = 10
|
||||
session.num_in_flight_tokens = 0
|
||||
sched.num_sampled_tokens_per_step = 1
|
||||
# The replacement plan rebuilds 21 tokens (retained context plus the
|
||||
# append) and one step samples one more, so 21 leaves no room.
|
||||
update = _make_minicpm_window_update(seq=2, mode="basic")
|
||||
update.prompt_token_ids = [0] * 20
|
||||
original_prompt = list(session.prompt_token_ids)
|
||||
|
||||
sched._update_request_as_session(session, update)
|
||||
|
||||
assert update.model_intermediate_buffer["duplex"]["stage0_window"]["replacement_prompt_len"] == 21
|
||||
assert session.prompt_token_ids == original_prompt
|
||||
assert session.num_prompt_tokens == 9
|
||||
assert session.status == RequestStatus.FINISHED_ERROR
|
||||
assert session.request_id not in sched.requests
|
||||
assert sched.finished_req_ids == {session.request_id}
|
||||
sched._free_request_blocks.assert_called_once_with(session)
|
||||
client_index, reason = sched._streaming_context_overflow[session.request_id]
|
||||
assert client_index == session.client_index
|
||||
assert reason.startswith("context_length_exceeded: ")
|
||||
assert "21 tokens" in reason and "max_model_len 21" in reason
|
||||
|
||||
engine_core_outputs = _run_idle_step(sched)
|
||||
|
||||
(output,) = engine_core_outputs[session.client_index].outputs
|
||||
assert output.request_id == session.request_id
|
||||
assert output.finish_reason == FinishReason.ERROR
|
||||
assert output.stop_reason == reason
|
||||
|
||||
|
||||
def test_stage0_window_rebuild_that_leaves_room_to_sample_replaces_the_prompt() -> None:
|
||||
"""One slot above the plan fits: the replacement applies through the normal
|
||||
replacement path."""
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
sched.max_model_len = 25
|
||||
sched.num_sampled_tokens_per_step = 1
|
||||
session = _make_request()
|
||||
session.prompt_token_ids = [0] * 9
|
||||
session._all_token_ids[:] = session.prompt_token_ids
|
||||
session.num_prompt_tokens = 9
|
||||
session.append_output_token_ids([40])
|
||||
session.num_computed_tokens = 10
|
||||
session.num_output_placeholders = 0
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
update = _make_minicpm_window_update(seq=2, mode="basic")
|
||||
update.prompt_token_ids = [0] * 20
|
||||
|
||||
sched._update_request_as_session(session, update)
|
||||
|
||||
assert update.model_intermediate_buffer["meta"]["replace_streaming_prompt"] is True
|
||||
assert update.model_intermediate_buffer["duplex"]["stage0_window"]["replacement_prompt_len"] == 21
|
||||
assert session.prompt_token_ids == [0] * 21
|
||||
assert session.num_prompt_tokens == 21
|
||||
assert session.num_computed_tokens == 0
|
||||
assert session.status == RequestStatus.WAITING
|
||||
sched._free_request_blocks.assert_called_once_with(session)
|
||||
assert not getattr(sched, "_streaming_context_overflow", {})
|
||||
|
||||
|
||||
def test_stage0_window_open_start_zero_is_not_replaced_by_the_context_reserve() -> None:
|
||||
"""A recorded ``open_start`` of 0 is a legitimate empty context prefix. The
|
||||
fallback must be an explicit ``is None`` check: treating 0 as missing
|
||||
substitutes the context reserve, counts the suffix twice, and the worker's
|
||||
rebuild-length check raises on the first replacement."""
|
||||
session = SimpleNamespace(
|
||||
num_prompt_tokens=9,
|
||||
_all_token_ids=[0] * 9 + [40],
|
||||
num_computed_tokens=10,
|
||||
num_output_placeholders=0,
|
||||
_minicpmo45_window_open_start=0,
|
||||
)
|
||||
update = _make_minicpm_window_update(seq=2, mode="basic")
|
||||
update.prompt_token_ids = [0] * 21
|
||||
|
||||
assert OmniARScheduler._prepare_minicpmo45_stage0_window(
|
||||
session, update, segment_output_ids=[40], completed_terminator=3
|
||||
)
|
||||
|
||||
# open_start stays 0, so the reported unit spans the whole 12-token
|
||||
# boundary instead of the 9 the context reserve would fold in.
|
||||
plan = update.model_intermediate_buffer["duplex"]["stage0_window"]
|
||||
assert plan["dropped_tokens"] == 12
|
||||
assert plan["replacement_prompt_len"] == 19
|
||||
|
||||
|
||||
def test_stage0_window_open_start_falls_back_to_the_context_reserve_when_unset() -> None:
|
||||
session = SimpleNamespace(
|
||||
num_prompt_tokens=9,
|
||||
_all_token_ids=[0] * 9 + [40],
|
||||
num_computed_tokens=10,
|
||||
num_output_placeholders=0,
|
||||
)
|
||||
update = _make_minicpm_window_update(seq=2, mode="basic")
|
||||
update.prompt_token_ids = [0] * 21
|
||||
|
||||
assert OmniARScheduler._prepare_minicpmo45_stage0_window(
|
||||
session, update, segment_output_ids=[40], completed_terminator=3
|
||||
)
|
||||
|
||||
# fallback open_start is preserve_len 3 from duplex_first_append_context_tokens,
|
||||
# so the unit is three tokens shorter.
|
||||
plan = update.model_intermediate_buffer["duplex"]["stage0_window"]
|
||||
assert plan["dropped_tokens"] == 9
|
||||
assert plan["replacement_prompt_len"] == 22
|
||||
|
||||
|
||||
def test_explicit_streaming_payload_replaces_placeholder_prompt() -> None:
|
||||
sched = _make_scheduler(stage_id=1)
|
||||
sched.chunk_transfer_adapter = SimpleNamespace(
|
||||
|
||||
+108
@@ -60,6 +60,7 @@ def test_custom_cache_dit_enablers_are_registered_explicitly():
|
||||
"Cosmos3OmniPipeline": cd_model_specific.enable_cache_for_cosmos3,
|
||||
"Krea2Pipeline": cd_model_specific.enable_cache_for_krea2,
|
||||
"Magi2Pipeline": cd_model_specific.enable_cache_for_magi2,
|
||||
"MammothModa2DiTPipeline": cd_model_specific.enable_cache_for_mammothmoda2,
|
||||
}
|
||||
|
||||
with patch.dict(cd_backend.CUSTOM_DIT_ENABLERS, {}, clear=True):
|
||||
@@ -116,6 +117,33 @@ def test_magi2_cache_dit_targets_only_nested_repeated_layers(mock_block_adapter,
|
||||
assert get_transformer(pipeline) is transformer_block
|
||||
|
||||
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.enable_cache_for_dit")
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
|
||||
def test_mammothmoda2_cache_dit_targets_only_main_layers(mock_block_adapter, mock_enable_cache):
|
||||
pipeline = Mock()
|
||||
transformer = pipeline.gen_transformer
|
||||
layers = torch.nn.ModuleList([torch.nn.Identity()])
|
||||
transformer.layers = layers
|
||||
adapter = mock_block_adapter.return_value
|
||||
refresh = Mock()
|
||||
mock_enable_cache.return_value = refresh
|
||||
|
||||
result = cd_model_specific.enable_cache_for_mammothmoda2(pipeline, SAMPLE_CACHE_CONFIG)
|
||||
|
||||
mock_block_adapter.assert_called_once()
|
||||
adapter_kwargs = mock_block_adapter.call_args.kwargs
|
||||
assert adapter_kwargs["transformer"] is transformer
|
||||
assert adapter_kwargs["blocks"] == [layers]
|
||||
assert adapter_kwargs["has_separate_cfg"] is True
|
||||
assert adapter_kwargs["check_forward_pattern"] is True
|
||||
assert result.refresh is refresh
|
||||
assert result.targets == (adapter,)
|
||||
get_transformer = mock_enable_cache.call_args.kwargs["get_pipeline_transformer"]
|
||||
assert get_transformer(pipeline) is transformer
|
||||
# Sequential-CFG parity is mandatory while cache-dit is installed.
|
||||
assert pipeline._cache_dit_requires_paired_cfg is True
|
||||
|
||||
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit.summary")
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter.is_cached", return_value=True)
|
||||
def test_cache_summary_uses_custom_nested_targets(mock_is_cached, mock_summary):
|
||||
@@ -338,3 +366,83 @@ def test_summary_with_no_transformer_is_nonfatal():
|
||||
pass
|
||||
|
||||
cache_summary(pipeline=FakePipeline())
|
||||
|
||||
|
||||
# This test is skipped on ROCm since rocm_unquantized_gemm doesn't support CPU backend
|
||||
@pytest.mark.skipif(
|
||||
current_omni_platform.is_rocm(),
|
||||
reason="vLLM ROCm custom ops lack CPU fallback",
|
||||
)
|
||||
def test_mammothmoda2_cache_dit_runs_end_to_end_on_tiny_model(request: pytest.FixtureRequest):
|
||||
"""A tiny MammothModa2 DiT runs through Cache-DiT with paired CFG forwards.
|
||||
|
||||
Exercises the Pattern_3 contract end to end: after enabling the custom
|
||||
enabler, blocks must be callable as ``block(hidden_states, **kwargs)``
|
||||
and the pipeline's two-forwards-per-step CFG cadence must drive the
|
||||
cache context without shape or parity errors.
|
||||
"""
|
||||
from cache_dit import BlockAdapter
|
||||
|
||||
from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import Transformer2DModel
|
||||
from vllm_omni.diffusion.models.mammoth_moda2.rope_real import RotaryPosEmbedReal
|
||||
|
||||
# MammothModa2 attention resolves through the shared Omni attention layer,
|
||||
# which picks the platform-default backend at construction time. Pin
|
||||
# TORCH_SDPA so a CUDA-visible host does not select FA3 for CPU tensors.
|
||||
with _force_torch_sdpa():
|
||||
model = Transformer2DModel(
|
||||
patch_size=2,
|
||||
in_channels=4,
|
||||
hidden_size=96,
|
||||
num_layers=4,
|
||||
num_refiner_layers=1,
|
||||
num_attention_heads=2,
|
||||
num_kv_heads=2,
|
||||
multiple_of=8,
|
||||
axes_dim_rope=(16, 16, 16),
|
||||
axes_lens=(300, 512, 512),
|
||||
text_feat_dim=16,
|
||||
)
|
||||
model.eval()
|
||||
|
||||
MammothModa2Pipeline = type("MammothModa2DiTPipeline", (), {})
|
||||
pipeline = MammothModa2Pipeline()
|
||||
pipeline.gen_transformer = model
|
||||
backend = CacheDiTBackend(DiffusionCacheConfig())
|
||||
backend.enable(pipeline)
|
||||
request.addfinalizer(lambda: backend.disable(pipeline))
|
||||
assert BlockAdapter.is_cached(model)
|
||||
|
||||
num_inference_steps = 8
|
||||
backend.refresh(pipeline, num_inference_steps=num_inference_steps)
|
||||
|
||||
freqs_cis = RotaryPosEmbedReal.get_freqs_real((16, 16, 16), (300, 512, 512), theta=10000)
|
||||
latents = torch.randn(1, 4, 8, 8)
|
||||
text = torch.randn(1, 5, 16)
|
||||
negative = torch.randn(1, 3, 16)
|
||||
text_mask = torch.ones(1, 5, dtype=torch.bool)
|
||||
negative_mask = torch.ones(1, 3, dtype=torch.bool)
|
||||
|
||||
with torch.no_grad():
|
||||
for step in range(num_inference_steps):
|
||||
timestep = torch.full((1,), 1.0 - step / num_inference_steps)
|
||||
# Conditional pass ...
|
||||
cond = model(
|
||||
hidden_states=latents,
|
||||
timestep=timestep,
|
||||
text_hidden_states=text,
|
||||
freqs_cis=freqs_cis,
|
||||
text_attention_mask=text_mask,
|
||||
)
|
||||
# ... always followed by the unconditional pass (parity).
|
||||
uncond = model(
|
||||
hidden_states=latents,
|
||||
timestep=timestep,
|
||||
text_hidden_states=negative,
|
||||
freqs_cis=freqs_cis,
|
||||
text_attention_mask=negative_mask,
|
||||
)
|
||||
latents = uncond + 4.0 * (cond - uncond)
|
||||
|
||||
assert latents.shape == (1, 4, 8, 8)
|
||||
assert torch.isfinite(latents).all()
|
||||
@@ -1,13 +1,27 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
|
||||
import cache_dit
|
||||
import pytest
|
||||
import torch
|
||||
from cache_dit import BlockAdapter
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, TransformerConfig
|
||||
from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend, RequestScopedCacheDiTRuntime
|
||||
from vllm_omni.diffusion.cache.selector import get_cache_backend
|
||||
from vllm_omni.diffusion.config import set_current_diffusion_config
|
||||
from vllm_omni.diffusion.data import (
|
||||
AttentionConfig,
|
||||
DiffusionCacheConfig,
|
||||
DiffusionOutput,
|
||||
OmniDiffusionConfig,
|
||||
TransformerConfig,
|
||||
)
|
||||
from vllm_omni.diffusion.models.interface import adopt_request_scoped_cache_dit
|
||||
from vllm_omni.diffusion.models.mammoth_moda2 import pipeline_mammothmoda2_dit
|
||||
from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import (
|
||||
MammothModa2DiTPipeline,
|
||||
@@ -17,6 +31,7 @@ from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import (
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -110,6 +125,10 @@ def _pipeline_shell() -> MammothModa2DiTPipeline:
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.config = _build_mammoth_config(_od_config())
|
||||
pipeline._llm_hidden_size = 8
|
||||
# Cache-DiT runtime state that __init__ normally creates (disabled here,
|
||||
# matching a pipeline built without a cache_backend).
|
||||
pipeline._cache_dit_runtime = RequestScopedCacheDiTRuntime(pipeline)
|
||||
pipeline._cache_dit_config = None
|
||||
return pipeline
|
||||
|
||||
|
||||
@@ -118,12 +137,14 @@ def _batch(
|
||||
request_id: str = "req-7",
|
||||
prompt: object | None = None,
|
||||
sampling: OmniDiffusionSamplingParams | None = None,
|
||||
height: int = 32,
|
||||
width: int = 48,
|
||||
) -> DiffusionRequestBatch:
|
||||
if prompt is None:
|
||||
prompt = {
|
||||
"prompt": "",
|
||||
"height": 32,
|
||||
"width": 48,
|
||||
"height": height,
|
||||
"width": width,
|
||||
"additional_information": {
|
||||
"full_hidden_states": torch.arange(32, dtype=torch.float32).reshape(4, 8),
|
||||
"full_token_ids": [10, 11, 100, 101],
|
||||
@@ -390,3 +411,162 @@ def test_forward_rejects_missing_visual_tokens_before_model_access() -> None:
|
||||
}
|
||||
with pytest.raises(ValueError, match="no visual-token hidden states.*req-empty"):
|
||||
_pipeline_shell().forward(_batch(request_id="req-empty", prompt=prompt))
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _force_torch_sdpa():
|
||||
"""Pin TORCH_SDPA so CPU shape tests do not pick CUDA-only backends (FA3)."""
|
||||
od_config = SimpleNamespace(
|
||||
diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"),
|
||||
parallel_config=SimpleNamespace(ring_degree=1),
|
||||
)
|
||||
with set_current_diffusion_config(od_config):
|
||||
yield
|
||||
|
||||
|
||||
def _constructible_raw_config() -> dict:
|
||||
"""Checkpoint config whose DiT/VAE are small enough to actually run on CPU."""
|
||||
return {
|
||||
"model_type": "mammothmoda2",
|
||||
"llm_config": {
|
||||
"model_type": "mammothmoda2_qwen2_5_vl",
|
||||
"text_config": {
|
||||
"model_type": "mammothmoda2_qwen2_5_vl_text",
|
||||
"hidden_size": 8,
|
||||
"gen_vocab_start_index": 100,
|
||||
},
|
||||
},
|
||||
"gen_vae_config": {"block_out_channels": [8, 8], "norm_num_groups": 8},
|
||||
"gen_dit_config": {
|
||||
"patch_size": 2,
|
||||
"in_channels": 4,
|
||||
"hidden_size": 96,
|
||||
"num_layers": 4,
|
||||
"num_refiner_layers": 1,
|
||||
"num_attention_heads": 2,
|
||||
"num_kv_heads": 2,
|
||||
"multiple_of": 8,
|
||||
"axes_dim_rope": [16, 16, 16],
|
||||
"axes_lens": [300, 512, 512],
|
||||
"text_feat_dim": 16,
|
||||
},
|
||||
# Top-level axes feed the pipeline-owned rotary table (see pipeline __init__).
|
||||
"gen_axes_dim_rope": [16, 16, 16],
|
||||
"gen_axes_lens": [300, 512, 512],
|
||||
}
|
||||
|
||||
|
||||
def _cache_dit_od_config() -> OmniDiffusionConfig:
|
||||
return OmniDiffusionConfig(
|
||||
model="/models/MammothModa2-Preview",
|
||||
model_class_name="MammothModa2DiTPipeline",
|
||||
tf_model_config=TransformerConfig.from_dict(_constructible_raw_config()),
|
||||
cache_backend="cache_dit",
|
||||
cache_config=DiffusionCacheConfig(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_omni_platform.is_rocm(),
|
||||
reason="vLLM ROCm custom ops lack CPU fallback",
|
||||
)
|
||||
def test_forward_transitions_request_scoped_cache_dit_across_requests(request: pytest.FixtureRequest) -> None:
|
||||
"""Cache-DiT request lifecycle on one production-built pipeline instance.
|
||||
|
||||
Mirrors the diffusion runner startup (``get_cache_backend`` → ``enable`` →
|
||||
``adopt_request_scoped_cache_dit``) on a tiny pipeline built through the
|
||||
real ``__init__``, then drives three requests through the real ``forward()``:
|
||||
|
||||
1. a partial-CFG request installs the request-scoped cache context:
|
||||
cond/uncond forwards stay paired on every denoise step (parity
|
||||
accounting), even outside ``cfg_range``, and blocks record cached steps;
|
||||
2. a guidance=1.0 request runs with hooks disabled (``prepare(None)``) and
|
||||
its output is bit-identical to a reference captured before the backend
|
||||
was ever enabled, proving a clean uninstall with no stale module state;
|
||||
3. a CFG request with a different step count and image size re-installs the
|
||||
hooks through a fresh context, so no stale shape or step count survives.
|
||||
"""
|
||||
torch.manual_seed(0) # Deterministic random init → deterministic cache hits.
|
||||
with _force_torch_sdpa():
|
||||
pipeline = MammothModa2DiTPipeline(od_config=_cache_dit_od_config())
|
||||
pipeline.eval()
|
||||
assert pipeline._cache_dit_config is not None
|
||||
request.addfinalizer(pipeline._cache_dit_runtime.disable)
|
||||
|
||||
# The default prompt carries 2 text + 2 visual condition tokens, so the
|
||||
# conditional pass sees 4 tokens and the uncond pass (empty prompt) sees 0.
|
||||
cond_tokens, uncond_tokens = 4, 0
|
||||
seq_lens: list[int] = []
|
||||
|
||||
def record_call(_module, _args, kwargs):
|
||||
text = kwargs.get("text_hidden_states")
|
||||
seq_lens.append(int(text.shape[1]) if text is not None else -1)
|
||||
|
||||
hook = pipeline.gen_transformer.register_forward_pre_hook(record_call, with_kwargs=True)
|
||||
|
||||
unguided_sampling = OmniDiffusionSamplingParams(
|
||||
height=32, width=48, seed=4321, guidance_scale=1.0, num_inference_steps=8
|
||||
)
|
||||
partial_cfg_sampling = OmniDiffusionSamplingParams(
|
||||
height=32,
|
||||
width=48,
|
||||
seed=1234,
|
||||
guidance_scale=4.0,
|
||||
num_inference_steps=8,
|
||||
extra_args={"cfg_range": [0.2, 0.8]},
|
||||
)
|
||||
changed_cfg_sampling = OmniDiffusionSamplingParams(
|
||||
height=48,
|
||||
width=64,
|
||||
seed=777,
|
||||
guidance_scale=6.0,
|
||||
num_inference_steps=10,
|
||||
extra_args={"cfg_range": [0.1, 0.9]},
|
||||
)
|
||||
|
||||
# Reference run for phase 2: identical request while hooks were never
|
||||
# installed on this pipeline.
|
||||
assert not pipeline.is_cache_dit_enabled()
|
||||
reference = pipeline.forward(_batch(sampling=unguided_sampling)).output
|
||||
assert seq_lens == [cond_tokens] * 8
|
||||
assert torch.isfinite(reference.float()).all()
|
||||
|
||||
# Runner startup: enable once, then the pipeline owns all transitions.
|
||||
backend = get_cache_backend("cache_dit", DiffusionCacheConfig())
|
||||
assert isinstance(backend, CacheDiTBackend)
|
||||
backend.enable(pipeline)
|
||||
assert adopt_request_scoped_cache_dit(pipeline, backend)
|
||||
assert pipeline.is_cache_dit_enabled()
|
||||
|
||||
# Phase 1: partial CFG keeps cond/uncond paired on every step, including
|
||||
# the steps whose normalized index falls outside cfg_range [0.2, 0.8].
|
||||
seq_lens.clear()
|
||||
pipeline.forward(_batch(sampling=partial_cfg_sampling))
|
||||
assert seq_lens == [cond_tokens, uncond_tokens] * 8
|
||||
assert pipeline.is_cache_dit_enabled()
|
||||
assert BlockAdapter.is_cached(pipeline.gen_transformer)
|
||||
hits = cache_dit.summary(pipeline.gen_transformer, logging=False)[0]
|
||||
assert hits.cached_steps or hits.cfg_cached_steps
|
||||
|
||||
# Phase 2: no-CFG request disables the hooks and must reproduce the
|
||||
# never-enabled reference bit-for-bit.
|
||||
seq_lens.clear()
|
||||
unguided = pipeline.forward(_batch(sampling=unguided_sampling)).output
|
||||
assert seq_lens == [cond_tokens] * 8
|
||||
assert not pipeline.is_cache_dit_enabled()
|
||||
assert not BlockAdapter.is_cached(pipeline.gen_transformer)
|
||||
assert torch.equal(unguided, reference)
|
||||
|
||||
# Phase 3: CFG with a new step count and image size re-installs the hooks
|
||||
# on a fresh context; output shape follows the new request geometry.
|
||||
seq_lens.clear()
|
||||
regenerated = pipeline.forward(_batch(sampling=changed_cfg_sampling, height=48, width=64)).output
|
||||
assert seq_lens == [cond_tokens, uncond_tokens] * 10
|
||||
assert pipeline.is_cache_dit_enabled()
|
||||
assert BlockAdapter.is_cached(pipeline.gen_transformer)
|
||||
hits = cache_dit.summary(pipeline.gen_transformer, logging=False)[0]
|
||||
assert hits.cached_steps or hits.cfg_cached_steps
|
||||
assert regenerated.shape == (1, 3, 12, 16)
|
||||
assert torch.isfinite(regenerated.float()).all()
|
||||
|
||||
hook.remove()
|
||||
@@ -0,0 +1,89 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Drive complete native duplex turns while Stage-0 repeatedly rebuilds KV."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
from pathlib import Path
|
||||
|
||||
from vllm_omni.clients.duplex import DuplexClient, EventCollector, read_pcm16_wav, wait_for_condition
|
||||
from vllm_omni.clients.minicpmo_4_5 import create_duplex_session_config
|
||||
|
||||
|
||||
async def run_window_turn(
|
||||
*,
|
||||
url: str,
|
||||
model: str,
|
||||
input_wav: Path,
|
||||
mode: str,
|
||||
ref_audio: Path | None = None,
|
||||
repeats: int = 2,
|
||||
video_frames: list[str] | None = None,
|
||||
buffered_flush: bool = False,
|
||||
timeout_s: float = 240.0,
|
||||
) -> dict[str, object]:
|
||||
"""Use aggressive limits to exercise multiple windows in one audio turn."""
|
||||
pcm = read_pcm16_wav(input_wav) * repeats
|
||||
collector = EventCollector()
|
||||
reference = (
|
||||
"data:audio/wav;base64," + base64.b64encode(ref_audio.read_bytes()).decode("ascii")
|
||||
if ref_audio is not None
|
||||
else None
|
||||
)
|
||||
config = create_duplex_session_config(
|
||||
modalities=("text", "audio") if ref_audio is not None else ("text",),
|
||||
ref_audio=reference,
|
||||
temperature=0.0,
|
||||
extra_body={
|
||||
"sliding_window_mode": mode,
|
||||
"context_max_units": 1,
|
||||
"context_previous_max_tokens": 16,
|
||||
"basic_window_high_tokens": 64,
|
||||
"basic_window_low_tokens": 32,
|
||||
},
|
||||
)
|
||||
client = DuplexClient(
|
||||
url, model=model, config=config, reconnect=None, heartbeat_interval_s=None, handshake_timeout_s=timeout_s
|
||||
)
|
||||
async with client:
|
||||
reader = asyncio.create_task(collector.consume(client))
|
||||
collector.add({"type": "session.created", "session": client.session_info})
|
||||
try:
|
||||
frames_sent = await client.stream_pcm(
|
||||
pcm,
|
||||
chunk_ms=2500 if buffered_flush else 200,
|
||||
realtime=not buffered_flush,
|
||||
is_speech=True,
|
||||
video_frames=video_frames,
|
||||
)
|
||||
await client.commit(final=True)
|
||||
await wait_for_condition(
|
||||
lambda: collector.count("input_audio_buffer.committed") > 0 or bool(collector.errors()),
|
||||
timeout_s=timeout_s,
|
||||
label="final input commit",
|
||||
)
|
||||
await wait_for_condition(
|
||||
lambda: collector.count("response.done") > 0 or bool(collector.errors()),
|
||||
timeout_s=timeout_s,
|
||||
label="completed windowed response",
|
||||
)
|
||||
# Final execution follows the commit acknowledgement asynchronously.
|
||||
await asyncio.sleep(3.0)
|
||||
await client.close(timeout_s=timeout_s)
|
||||
await asyncio.wait_for(reader, timeout=5.0)
|
||||
finally:
|
||||
if not reader.done():
|
||||
reader.cancel()
|
||||
await asyncio.gather(reader, return_exceptions=True)
|
||||
return {
|
||||
"created": collector.count("session.created"),
|
||||
"committed": collector.count("input_audio_buffer.committed"),
|
||||
"done": collector.count("response.done"),
|
||||
"closed": collector.count("session.closed"),
|
||||
"errors": collector.errors(),
|
||||
"audio_bytes": len(collector.audio_bytes()),
|
||||
"frames_sent": frames_sent,
|
||||
"input_seconds": len(pcm) / 32000,
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Real-model regression for the scheduler-owned Thinker sliding window."""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
|
||||
SERVER_PARAMS,
|
||||
duplex_camera_frames,
|
||||
realtime_url,
|
||||
resolve_ref_audio,
|
||||
validated_input_wav,
|
||||
)
|
||||
from tests.e2e.online_serving.helpers.minicpmo_window_e2e import run_window_turn
|
||||
from tests.helpers.mark import hardware_test
|
||||
|
||||
pytestmark = [pytest.mark.omni, pytest.mark.advanced_model]
|
||||
|
||||
|
||||
def _assert_complete(result, *, require_audio=False):
|
||||
assert result["created"] == 1
|
||||
assert result["committed"] > 0
|
||||
assert result["done"] > 0
|
||||
assert result["closed"] == 1
|
||||
assert result["errors"] == []
|
||||
if require_audio:
|
||||
assert result["audio_bytes"] > 0
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
|
||||
@pytest.mark.parametrize("mode,reference", [("basic", False), ("context", False), ("context", True)])
|
||||
def test_window_rebuild_and_next_session(omni_server, mode, reference):
|
||||
# Two independent sessions exercise admission after final-unit execution.
|
||||
for _ in range(2):
|
||||
result = asyncio.run(
|
||||
run_window_turn(
|
||||
url=realtime_url(omni_server),
|
||||
model=omni_server.model,
|
||||
input_wav=validated_input_wav(),
|
||||
mode=mode,
|
||||
ref_audio=resolve_ref_audio() if reference else None,
|
||||
)
|
||||
)
|
||||
_assert_complete(result, require_audio=reference)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
|
||||
@pytest.mark.parametrize("mode,camera", [("basic", False), ("context", False), ("context", True)])
|
||||
def test_window_continuous_input(omni_server, tmp_path, mode, camera):
|
||||
# 24 repeats cover over two minutes and many canonical KV replacements.
|
||||
frames = duplex_camera_frames(seconds=4, cache_dir=tmp_path / "camera") if camera else None
|
||||
result = asyncio.run(
|
||||
run_window_turn(
|
||||
url=realtime_url(omni_server),
|
||||
model=omni_server.model,
|
||||
input_wav=validated_input_wav(),
|
||||
mode=mode,
|
||||
repeats=24,
|
||||
ref_audio=resolve_ref_audio(),
|
||||
video_frames=frames,
|
||||
)
|
||||
)
|
||||
_assert_complete(result, require_audio=True)
|
||||
assert isinstance(result["input_seconds"], (int, float)) and result["input_seconds"] > 120
|
||||
if camera:
|
||||
assert isinstance(result["frames_sent"], int) and result["frames_sent"] >= 120
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
|
||||
@pytest.mark.parametrize("mode", ["off", "basic", "context"])
|
||||
def test_window_buffered_flush(omni_server, mode):
|
||||
# Burst appends cross processor chunk boundaries and leave a final tail.
|
||||
result = asyncio.run(
|
||||
run_window_turn(
|
||||
url=realtime_url(omni_server),
|
||||
model=omni_server.model,
|
||||
input_wav=validated_input_wav(),
|
||||
mode=mode,
|
||||
ref_audio=resolve_ref_audio(),
|
||||
buffered_flush=True,
|
||||
)
|
||||
)
|
||||
_assert_complete(result, require_audio=True)
|
||||
@@ -0,0 +1,117 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import base64
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.engine.duplex.config import DuplexSessionConfig
|
||||
from vllm_omni.engine.duplex.contracts import DuplexFence
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.duplex import plugin as module
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _plugin():
|
||||
return module.MiniCPMO45DuplexPlugin(lambda *args: None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_configuration_is_prepared_by_plugin(monkeypatch):
|
||||
tokenizer = SimpleNamespace(
|
||||
all_special_ids=[90],
|
||||
unk_token_id=-1,
|
||||
eos_token_id=99,
|
||||
encode=lambda text, add_special_tokens=False: [ord(char) for char in text],
|
||||
convert_tokens_to_ids=lambda token: {"<|listen|>": 91}.get(token, -1),
|
||||
)
|
||||
monkeypatch.setattr(module, "_load_tokenizer", lambda config: tokenizer)
|
||||
config = DuplexSessionConfig(
|
||||
modalities=("text",),
|
||||
extra_body={
|
||||
"sliding_window_mode": "basic",
|
||||
"basic_window_high_tokens": 120,
|
||||
"basic_window_low_tokens": 80,
|
||||
},
|
||||
)
|
||||
runtime = await _plugin().prepare_runtime_config(config, model_config=None)
|
||||
assert runtime["duplex_window_config"]["sliding_window_mode"] == "basic"
|
||||
assert "sliding_window_mode" not in config.extra_body
|
||||
assert runtime["duplex_window_prefix_tokens"] > 0
|
||||
assert runtime["duplex_window_suffix_token_ids"]
|
||||
assert runtime["duplex_window_previous_marker_token_ids"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_window_configuration_is_rejected():
|
||||
config = DuplexSessionConfig(
|
||||
modalities=("text",),
|
||||
extra_body={
|
||||
"basic_window_high_tokens": 80,
|
||||
"basic_window_low_tokens": 80,
|
||||
},
|
||||
)
|
||||
with pytest.raises(module.MiniCPMO45ClientRuntimeConfigError) as exc:
|
||||
await _plugin().prepare_runtime_config(config, model_config=None)
|
||||
assert exc.value.code == "invalid_sliding_window_config"
|
||||
|
||||
|
||||
def test_window_update_cannot_change_created_configuration():
|
||||
config = DuplexSessionConfig(extra_body={"sliding_window_mode": "basic"})
|
||||
with pytest.raises(module.MiniCPMO45ClientRuntimeConfigError) as exc:
|
||||
_plugin().runtime_config_for_update(config, {"duplex_window_config": {"sliding_window_mode": "off"}})
|
||||
assert exc.value.code == "sliding_window_update_unsupported"
|
||||
|
||||
|
||||
def test_window_internal_configuration_is_server_owned():
|
||||
with pytest.raises(module.MiniCPMO45ClientRuntimeConfigError):
|
||||
_plugin().validate_client_extra_body({"duplex_window_config": {}})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("final", [False, True])
|
||||
@pytest.mark.parametrize("seq,samples,expected", [(1, 32000, 16), (1, 48000, 28), (2, 16000, 13)])
|
||||
def test_first_and_final_append_reserve_exact_window_input(seq, samples, expected, final):
|
||||
prompt = module.build_duplex_data_plane_prompt(
|
||||
request_id="window-request",
|
||||
fence=DuplexFence("sid", turn_id=1),
|
||||
session_config={},
|
||||
runtime_config={"duplex_first_append_context_tokens": 5},
|
||||
seq=seq,
|
||||
turn_seq=seq,
|
||||
payload={
|
||||
"audio": base64.b64encode(bytes(samples * 4)).decode(),
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": 16000,
|
||||
},
|
||||
final=final,
|
||||
)
|
||||
assert len(prompt["prompt_token_ids"]) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seq", [1, 2, 3])
|
||||
def test_final_exact_chunk_append_reserves_one_unit_not_two(seq):
|
||||
"""A final append whose audio is already an exact number of chunks must
|
||||
reserve exactly what Stage0 feeds: one unit per chunk (``<unit>`` plus its
|
||||
audio) and the closure pair for every unit after the first. Serving pads
|
||||
the final residual itself, so the old extra 12-slot silent-unit
|
||||
reservation only put 12 pad positions ahead of the audio in the KV."""
|
||||
|
||||
def _budget(*, seq, final):
|
||||
prompt = module.build_duplex_data_plane_prompt(
|
||||
request_id="final-request",
|
||||
fence=DuplexFence("sid", turn_id=1),
|
||||
session_config={},
|
||||
runtime_config={"duplex_first_append_context_tokens": 0},
|
||||
seq=seq,
|
||||
turn_seq=seq,
|
||||
payload={
|
||||
"audio": base64.b64encode(bytes(16000 * 4)).decode(),
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": 16000,
|
||||
},
|
||||
final=final,
|
||||
)
|
||||
return len(prompt["prompt_token_ids"])
|
||||
|
||||
assert _budget(seq=seq, final=True) == _budget(seq=seq, final=False)
|
||||
@@ -51,6 +51,7 @@ from vllm_omni.engine.duplex.plugin import (
|
||||
PcmAppendReservation,
|
||||
)
|
||||
from vllm_omni.engine.duplex.session import manager as session_manager_module
|
||||
from vllm_omni.engine.duplex.session.context import DuplexSessionTasks
|
||||
from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession
|
||||
from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity
|
||||
from vllm_omni.engine.duplex.session.manager import DuplexSessionManager
|
||||
@@ -1612,3 +1613,34 @@ async def test_reaper_loop_survives_one_cleanup_failure(first_cleanup_delay: flo
|
||||
finally:
|
||||
shutdown.set()
|
||||
await asyncio.wait_for(task, timeout=5.0)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Cancellation waits #
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
async def test_cancel_append_tasks_absorbs_a_task_that_outlives_the_wait() -> None:
|
||||
"""The wait after cancelling is allowed to time out; that is not an error.
|
||||
|
||||
On Python 3.10 the timeout arrives as ``asyncio.TimeoutError``, which is a
|
||||
different class from the builtin ``TimeoutError`` until 3.11.
|
||||
"""
|
||||
|
||||
async def outlives_the_first_cancel() -> bool:
|
||||
try:
|
||||
await asyncio.sleep(5)
|
||||
except asyncio.CancelledError:
|
||||
await asyncio.sleep(1)
|
||||
raise
|
||||
return True
|
||||
|
||||
tasks = DuplexSessionTasks()
|
||||
task = asyncio.create_task(outlives_the_first_cancel())
|
||||
await asyncio.sleep(0)
|
||||
tasks.track_append_task(task, epoch=0, final=False, response_bound=False)
|
||||
tasks.append_tail = task
|
||||
|
||||
assert await tasks.cancel_append_tasks(timeout_s=0.05) is True
|
||||
assert tasks.append_tail is None
|
||||
await asyncio.wait_for(asyncio.gather(task, return_exceptions=True), timeout=5.0)
|
||||
@@ -771,6 +771,41 @@ async def test_cancel_response_aborts_the_stage_request_and_reports_playback() -
|
||||
await close_harness(h)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_response_absorbs_a_response_task_that_outlives_the_wait() -> None:
|
||||
"""The 0.25 s wait after cancelling the response task may time out.
|
||||
|
||||
On Python 3.10 that timeout is ``asyncio.TimeoutError``, not the builtin
|
||||
``TimeoutError``, and it must not stop the cancel from completing.
|
||||
"""
|
||||
|
||||
async def outlives_the_first_cancel() -> None:
|
||||
try:
|
||||
await asyncio.sleep(5)
|
||||
except asyncio.CancelledError:
|
||||
await asyncio.sleep(1)
|
||||
raise
|
||||
|
||||
h = await open_harness()
|
||||
task = asyncio.create_task(outlives_the_first_cancel())
|
||||
try:
|
||||
await h.run(append_audio())
|
||||
request_id = h.stage0_request_id()
|
||||
await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello"))
|
||||
h.runner.tasks.active_response_task = task
|
||||
|
||||
events = await h.run(commands.CancelResponse())
|
||||
# The runner is still inside that wait when the mailbox goes quiet.
|
||||
events += await h.settle(idle_s=0.6, timeout_s=5.0)
|
||||
assert find(events, "response.done").status == "cancelled"
|
||||
assert "error" not in types(events)
|
||||
assert h.session.epoch == 1
|
||||
finally:
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
await close_harness(h)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_listen_decision_is_consumed_and_never_forwarded_to_tts() -> None:
|
||||
h = await open_harness()
|
||||
|
||||
@@ -505,6 +505,75 @@ class TestSpeechAPI:
|
||||
assert voice_info["file_size"] == len(audio_content)
|
||||
response = client.delete("/v1/audio/voices/test_voice")
|
||||
|
||||
def test_upload_voice_rejects_builtin_name(self, client):
|
||||
"""A name colliding with a built-in or precomputed voice is rejected."""
|
||||
handler = client.app.state.openai_serving_speech
|
||||
handler._adapter.capabilities.supported_speakers = {"vivian"}
|
||||
handler._adapter.capabilities.precomputed_speakers = {"ono_anna": {}}
|
||||
|
||||
files = {"audio_sample": ("test.wav", b"fake audio content" * 1000, "audio/wav")}
|
||||
for reserved in ("Vivian", "ono_anna"):
|
||||
response = client.post("/v1/audio/voices", files=files, data={"consent": "c1", "name": reserved})
|
||||
assert response.status_code == 400
|
||||
assert "reserved" in response.json()["detail"]
|
||||
assert not handler.uploaded_speakers
|
||||
|
||||
def test_upload_voice_overwrites_same_name_by_default(self, client):
|
||||
"""Default policy: re-registering an uploaded name replaces it in place."""
|
||||
files = {"audio_sample": ("test.wav", b"first take", "audio/wav")}
|
||||
data = {"consent": "c1", "name": "test_voice_ow"}
|
||||
assert client.post("/v1/audio/voices", files=files, data=data).status_code == 200
|
||||
|
||||
files = {"audio_sample": ("test.wav", b"second take, longer content", "audio/wav")}
|
||||
response = client.post("/v1/audio/voices", files=files, data=data)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["voice"]["file_size"] == len(b"second take, longer content")
|
||||
client.delete("/v1/audio/voices/test_voice_ow")
|
||||
|
||||
def test_upload_voice_immutable_policy_requires_delete(self, client):
|
||||
"""VLLM_OMNI_SPEAKER_REGISTRATION_POLICY=immutable rejects duplicates until deleted."""
|
||||
handler = client.app.state.openai_serving_speech
|
||||
handler._registration_policy = "immutable"
|
||||
try:
|
||||
files = {"audio_sample": ("test.wav", b"fake audio content" * 1000, "audio/wav")}
|
||||
data = {"consent": "c1", "name": "test_voice_imm"}
|
||||
assert client.post("/v1/audio/voices", files=files, data=data).status_code == 200
|
||||
|
||||
response = client.post("/v1/audio/voices", files=files, data=data)
|
||||
assert response.status_code == 400
|
||||
assert "immutable" in response.json()["detail"]
|
||||
|
||||
assert client.delete("/v1/audio/voices/test_voice_imm").status_code == 200
|
||||
assert client.post("/v1/audio/voices", files=files, data=data).status_code == 200
|
||||
finally:
|
||||
handler._registration_policy = "overwrite"
|
||||
client.delete("/v1/audio/voices/test_voice_imm")
|
||||
|
||||
def test_invalid_registration_policy_fails_startup(self, mocker, monkeypatch, tmp_path):
|
||||
"""An unknown VLLM_OMNI_SPEAKER_REGISTRATION_POLICY is a configuration error, not a silent fallback."""
|
||||
monkeypatch.setenv("SPEAKER_SAMPLES_DIR", str(tmp_path))
|
||||
monkeypatch.setenv("VLLM_OMNI_SPEAKER_REGISTRATION_POLICY", "append")
|
||||
engine_client = mocker.MagicMock()
|
||||
engine_client.default_sampling_params_list = [{}]
|
||||
with pytest.raises(ValueError, match="VLLM_OMNI_SPEAKER_REGISTRATION_POLICY"):
|
||||
OmniOpenAIServingSpeech(
|
||||
engine_client=engine_client, models=mocker.MagicMock(), request_logger=mocker.MagicMock()
|
||||
)
|
||||
|
||||
def test_restored_upload_shadowing_builtin_is_dropped(self, client):
|
||||
"""A pre-guard upload restored under a built-in name must not keep shadowing it."""
|
||||
handler = client.app.state.openai_serving_speech
|
||||
handler._adapter.capabilities.supported_speakers = {"vivian"}
|
||||
handler._adapter.capabilities.precomputed_speakers = {}
|
||||
handler.uploaded_speakers["vivian"] = {"name": "vivian", "file_path": "/tmp/vivian.safetensors"}
|
||||
handler.uploaded_speakers["keep_me"] = {"name": "keep_me", "file_path": "/tmp/keep_me.safetensors"}
|
||||
|
||||
handler._drop_shadowing_uploads()
|
||||
|
||||
assert "vivian" not in handler.uploaded_speakers
|
||||
assert "keep_me" in handler.uploaded_speakers
|
||||
handler.uploaded_speakers.pop("keep_me")
|
||||
|
||||
def test_upload_voice_with_ref_text(self, client, tmp_path):
|
||||
"""Test voice upload with ref_text enables in-context cloning."""
|
||||
audio_content = b"fake audio content" * 1000
|
||||
|
||||
@@ -33,9 +33,9 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni:
|
||||
omni._final_output_handler = lambda: None
|
||||
omni.reset_mm_cache = AsyncMock()
|
||||
|
||||
stage_clients = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types]
|
||||
stage_configs = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types]
|
||||
omni.engine = SimpleNamespace(
|
||||
stage_clients=stage_clients,
|
||||
stage_configs=stage_configs,
|
||||
stage_vllm_configs=[None] * len(stage_types),
|
||||
collective_rpc_async=AsyncMock(return_value=[True]),
|
||||
)
|
||||
|
||||
@@ -19,7 +19,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
|
||||
|
||||
from vllm_omni.config.omni_config import VllmOmniConfig
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
|
||||
from vllm_omni.config.stage_config import StageConfig
|
||||
from vllm_omni.config.stage_config import StageConfig, StageType
|
||||
from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
|
||||
from vllm_omni.engine.messages import ErrorMessage, OutputMessage
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
@@ -107,7 +107,13 @@ class FakeAsyncOmniEngine:
|
||||
self.stage_metadata = stage_metadata or [THREE_STAGE_META[-1]]
|
||||
self.num_stages = len(self.stage_metadata)
|
||||
self.stage_configs = [
|
||||
StageConfig(stage_id=i, model_stage="dummy-model").to_omegaconf() for i in range(self.num_stages)
|
||||
StageConfig(
|
||||
stage_id=i,
|
||||
model_stage="dummy-model",
|
||||
stage_type=StageType(metadata.stage_type),
|
||||
is_comprehension=getattr(metadata, "is_comprehension", False),
|
||||
).to_omegaconf()
|
||||
for i, metadata in enumerate(self.stage_metadata)
|
||||
]
|
||||
self.default_sampling_params_list = default_sampling_params_list or [
|
||||
SamplingParams(max_tokens=8) for _ in range(self.num_stages)
|
||||
@@ -481,7 +487,7 @@ def test_model_config_and_vllm_config_forward_from_comprehension_stage():
|
||||
io_processor = SimpleNamespace(name="io-processor")
|
||||
omni = object.__new__(AsyncOmni)
|
||||
omni.engine = SimpleNamespace(
|
||||
stage_clients=[SimpleNamespace(is_comprehension=False), SimpleNamespace(is_comprehension=True)],
|
||||
stage_configs=[SimpleNamespace(is_comprehension=False), SimpleNamespace(is_comprehension=True)],
|
||||
stage_vllm_configs=[None, vllm_config],
|
||||
)
|
||||
omni.input_processor = input_processor
|
||||
@@ -502,7 +508,7 @@ def test_openai_serving_models_can_consume_async_omni_compat_attrs():
|
||||
io_processor = SimpleNamespace(name="io-processor")
|
||||
omni = object.__new__(AsyncOmni)
|
||||
omni.engine = SimpleNamespace(
|
||||
stage_clients=[SimpleNamespace(is_comprehension=True)],
|
||||
stage_configs=[SimpleNamespace(is_comprehension=True)],
|
||||
stage_vllm_configs=[vllm_config],
|
||||
)
|
||||
omni.input_processor = input_processor
|
||||
@@ -524,10 +530,11 @@ def test_get_diffusion_od_config_returns_diffusion_stage_config():
|
||||
diffusion_od_config = object()
|
||||
omni = object.__new__(AsyncOmni)
|
||||
omni.engine = SimpleNamespace(
|
||||
stage_configs=[SimpleNamespace(stage_type="llm"), SimpleNamespace(stage_type="diffusion")],
|
||||
stage_clients=[
|
||||
SimpleNamespace(stage_type="llm"),
|
||||
SimpleNamespace(stage_type="diffusion", od_config=diffusion_od_config),
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
assert omni.get_diffusion_od_config() is diffusion_od_config
|
||||
@@ -537,15 +544,33 @@ def test_get_diffusion_od_config_falls_back_to_inner_engine():
|
||||
diffusion_od_config = object()
|
||||
omni = object.__new__(AsyncOmni)
|
||||
omni.engine = SimpleNamespace(
|
||||
stage_configs=[SimpleNamespace(stage_type="llm"), SimpleNamespace(stage_type="diffusion")],
|
||||
stage_clients=[
|
||||
SimpleNamespace(stage_type="llm"),
|
||||
SimpleNamespace(stage_type="diffusion", _engine=SimpleNamespace(od_config=diffusion_od_config)),
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
assert omni.get_diffusion_od_config() is diffusion_od_config
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_omni_rejects_diffusion_list_prompt_without_stage_clients(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
engine = FakeAsyncOmniEngine(stage_metadata=DIFFUSION_ONLY_META)
|
||||
engine.stage_clients = []
|
||||
_patch_engine(monkeypatch, engine)
|
||||
|
||||
app = AsyncOmni("dummy-model")
|
||||
try:
|
||||
with pytest.raises(ValueError, match="single prompt per request"):
|
||||
async for _ in app.generate(prompt=["prompt-1", "prompt-2"], request_id="req-1"):
|
||||
pass
|
||||
finally:
|
||||
app.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_omni_yields_only_final_stage_outputs(monkeypatch: pytest.MonkeyPatch):
|
||||
engine = FakeAsyncOmniEngine(
|
||||
|
||||
@@ -260,6 +260,37 @@ class TestDiTAttention:
|
||||
assert attention.to_k.out_features == 512
|
||||
assert attention.to_v.out_features == 512
|
||||
|
||||
def test_embedded_dit_honors_diffusion_attention_backend(self, monkeypatch):
|
||||
"""The LLM-hosted code2wav DiT must consume the diffusion env override."""
|
||||
from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
|
||||
from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import _build_dit_estimator
|
||||
|
||||
monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
|
||||
|
||||
estimator = _build_dit_estimator(
|
||||
{
|
||||
"dim": 32,
|
||||
"depth": 1,
|
||||
"heads": 2,
|
||||
"dim_head": 16,
|
||||
"dropout": 0.0,
|
||||
"ff_mult": 2,
|
||||
"mel_dim": 8,
|
||||
"mu_dim": 8,
|
||||
"spk_dim": 4,
|
||||
"out_channels": 8,
|
||||
"static_chunk_size": 4,
|
||||
"num_decoding_left_chunks": 1,
|
||||
}
|
||||
)
|
||||
|
||||
attention = estimator.transformer_blocks[0].attn.attn
|
||||
assert attention.backend_pref == "TORCH_SDPA"
|
||||
assert attention.backend_explicit is True
|
||||
assert attention.attn_backend is not None
|
||||
assert attention.attn_backend.get_name() == "SDPA"
|
||||
assert get_current_diffusion_config_or_none() is None
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_forward_passes_2d_mask_without_casting_qkv(self):
|
||||
|
||||
@@ -39,7 +39,10 @@ CHUNK_LEN = 24
|
||||
TOTAL_MEL = 96
|
||||
SPM = 480 # samples per mel frame for the small test HiFT (8*5*3*4), matching the real model
|
||||
|
||||
LONG_TOTAL_MEL = 400 # long enough that the 64-frame window actually truncates
|
||||
WINDOW_LEN = 64
|
||||
# Six chunks exercise three updates after history exceeds the bounded window,
|
||||
# without spending minutes recomputing an unnecessarily long synthetic signal.
|
||||
LONG_TOTAL_MEL = 144
|
||||
PCM16_LSB = 1.0 / 32767.0 # smallest representable PCM16 difference
|
||||
|
||||
# Worst observed deviation is 1.97e-7 (~2x float32 eps); atol=1e-6 keeps 5x headroom.
|
||||
@@ -189,16 +192,19 @@ def test_incremental_hift_bounded_window_is_close(config):
|
||||
def test_incremental_hift_matches_streaming_reference_within_tolerance(config):
|
||||
"""Windowed output matches the full-cumulative reference within float32 rounding.
|
||||
|
||||
The load-bearing correctness test: 400 mel frames far exceeds the 64-frame
|
||||
window, so truncation genuinely exercises the phase/noise carry. Float64
|
||||
agrees to ~1e-16, confirming the residual here is rounding, not truncation.
|
||||
The load-bearing correctness test crosses the 64-frame window repeatedly,
|
||||
so truncation genuinely exercises the phase/noise carry. Float64 agrees to
|
||||
~1e-16, confirming the residual here is rounding, not truncation.
|
||||
"""
|
||||
hift = _make_hift(config)
|
||||
model = _make_model(hift, window_len=64) # the real _hift_window_len
|
||||
model = _make_model(hift, window_len=WINDOW_LEN) # the real _hift_window_len
|
||||
chunks = _chunks(total_mel=LONG_TOTAL_MEL)
|
||||
|
||||
trim = int(hift.f0_predictor.condnet[0].causal_padding)
|
||||
assert LONG_TOTAL_MEL > 64 + trim + CHUNK_LEN # window must actually truncate
|
||||
truncating_steps = sum(
|
||||
history_len > WINDOW_LEN + trim for history_len in range(CHUNK_LEN, LONG_TOTAL_MEL, CHUNK_LEN)
|
||||
)
|
||||
assert truncating_steps >= 3 # exercise repeated truncation, not only the boundary
|
||||
|
||||
full = _full_reference(model, chunks)
|
||||
incr = _incremental(model, chunks)
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""VLLM_OMNI_VOXCPM_CODE_PATH still works but warns that it is deprecated (#6232)."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.model_executor.models.voxcpm2 import voxcpm2_import_utils
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _capture_warnings(monkeypatch) -> list[str]:
|
||||
messages: list[str] = []
|
||||
monkeypatch.setattr(voxcpm2_import_utils.logger, "warning_once", lambda msg, *args: messages.append(msg))
|
||||
return messages
|
||||
|
||||
|
||||
def test_code_path_env_warns_and_is_still_used(monkeypatch, tmp_path: Path):
|
||||
messages = _capture_warnings(monkeypatch)
|
||||
monkeypatch.setenv("VLLM_OMNI_VOXCPM_CODE_PATH", str(tmp_path))
|
||||
|
||||
candidates = voxcpm2_import_utils._iter_voxcpm2_src_candidates()
|
||||
|
||||
assert candidates[0] == tmp_path
|
||||
assert len(messages) == 1
|
||||
assert "VLLM_OMNI_VOXCPM_CODE_PATH is deprecated" in messages[0]
|
||||
|
||||
|
||||
def test_no_warning_without_code_path_env(monkeypatch):
|
||||
messages = _capture_warnings(monkeypatch)
|
||||
monkeypatch.delenv("VLLM_OMNI_VOXCPM_CODE_PATH", raising=False)
|
||||
|
||||
voxcpm2_import_utils._iter_voxcpm2_src_candidates()
|
||||
|
||||
assert messages == []
|
||||
@@ -142,6 +142,37 @@ def test_native_duplex_speak_segment_reaches_split_talker() -> None:
|
||||
assert info["duplex"]["turn_id"] == 7
|
||||
|
||||
|
||||
@pytest.mark.parametrize("folded_decisions", [0, 1, 2])
|
||||
@pytest.mark.parametrize("bos_in_prompt", [False, True])
|
||||
def test_native_duplex_tts_bos_aligns_after_window_rebuild(folded_decisions, bos_in_prompt) -> None:
|
||||
prompt_ids = [101] * 10
|
||||
output_ids = [9303] * folded_decisions + [9301, 21, 22, 9308]
|
||||
if bos_in_prompt:
|
||||
prompt_ids[-1] = 9301
|
||||
output_ids = [21, 22, 9308]
|
||||
rows = len(prompt_ids) + len(output_ids) - folded_decisions
|
||||
latent = torch.arange(rows * 4, dtype=torch.float32).reshape(rows, 4)
|
||||
source = _output(
|
||||
prompt_ids=prompt_ids,
|
||||
output_ids=output_ids,
|
||||
latent=latent,
|
||||
multimodal_output={
|
||||
"duplex_prompt_token_ids": prompt_ids,
|
||||
"meta": {
|
||||
"tts_bos_token_id": 9301,
|
||||
"tts_eos_token_id": 9302,
|
||||
"listen_token_id": 9303,
|
||||
"speak_token_id": 9304,
|
||||
"chunk_eos_token_id": 9308,
|
||||
},
|
||||
},
|
||||
)
|
||||
converted = llm2tts([source], prompt=[{}], _streaming_context=SimpleNamespace(bridge_states={}))[0]
|
||||
info = converted["model_intermediate_buffer"]
|
||||
assert info["ids"]["tts"] == [21, 22]
|
||||
torch.testing.assert_close(torch.as_tensor(info["hidden_states"]["tts"]), latent[-3:-1])
|
||||
|
||||
|
||||
def test_native_duplex_continuation_appends_only_new_talker_condition() -> None:
|
||||
prompt_ids = [101, 102]
|
||||
token_ids = {
|
||||
|
||||
@@ -848,6 +848,176 @@ def test_minicpmo_stage0_data_plane_prefill_matches_official_unit_format():
|
||||
assert result["prompt_suffix_len"] == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", [None, "off", "basic", "context"])
|
||||
def test_minicpmo_stage0_history_collection_starts_with_session_config(mode):
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import _MiniCPMO45Stage0SessionState
|
||||
|
||||
runtime = _stage0_vision_runtime()
|
||||
state = _MiniCPMO45Stage0SessionState(session_id="history-lifecycle")
|
||||
config = {} if mode is None else {"duplex_window_config": {"sliding_window_mode": mode}}
|
||||
runtime._prepare_session_context(state, {}, runtime_config=config)
|
||||
for seq in range(1, 101):
|
||||
runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=seq)
|
||||
enabled = mode in {"basic", "context"}
|
||||
assert state.window_enabled is enabled
|
||||
assert len(state.window_units) == (99 if enabled else 0)
|
||||
assert (state.pending_window_unit is not None) is enabled
|
||||
assert state.pending_window_generated_tokens == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["basic", "context"])
|
||||
@pytest.mark.parametrize("buffered", [False, True])
|
||||
def test_minicpmo_stage0_multi_chunk_append_matches_scheduler_window(mode, buffered):
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import _MiniCPMO45Stage0SessionState
|
||||
|
||||
runtime = _stage0_vision_runtime()
|
||||
state = _MiniCPMO45Stage0SessionState(session_id="multi-chunk-window", window_enabled=True)
|
||||
if buffered:
|
||||
# Keep the tiny fake processor's initial chunk unpadded to model carry.
|
||||
runtime._pad_first_audio_chunk_if_needed = lambda *args: None
|
||||
assert runtime._stage_prefill_embeddings_only(state, np.zeros(2, dtype=np.float32), seq=0)["success"] is False
|
||||
first = runtime._stage_prefill_embeddings_only(
|
||||
state, np.zeros(6 if buffered else 8, dtype=np.float32), seq=1, final=buffered
|
||||
)
|
||||
assert first["input_token_ids"] == [1, 11, 2, 1, 11]
|
||||
assert state.window_units == []
|
||||
assert state.pending_window_unit.token_ids == first["input_token_ids"]
|
||||
session = SimpleNamespace(num_prompt_tokens=5)
|
||||
replaced = False
|
||||
for seq in (2, 3):
|
||||
info = {
|
||||
"duplex": {
|
||||
"data_plane": True,
|
||||
"seq": seq,
|
||||
"runtime_config": {
|
||||
"duplex_window_config": {
|
||||
"sliding_window_mode": mode,
|
||||
"basic_window_high_tokens": 1,
|
||||
"basic_window_low_tokens": 2,
|
||||
"context_max_units": 1,
|
||||
"context_previous_max_tokens": 16,
|
||||
},
|
||||
"duplex_window_previous_marker_token_ids": [201, 5],
|
||||
},
|
||||
}
|
||||
}
|
||||
update = SimpleNamespace(prompt_token_ids=[3, 2, 1, 11], model_intermediate_buffer=info)
|
||||
planned = OmniARScheduler._prepare_minicpmo45_stage0_window(
|
||||
session, update, segment_output_ids=[44], completed_terminator=3
|
||||
)
|
||||
result = runtime._stage_prefill_embeddings_only(
|
||||
state, np.zeros(4, dtype=np.float32), seq=seq, stage0_window=info["duplex"]["stage0_window"]
|
||||
)
|
||||
assert result["stage0_window_replaced"] is planned
|
||||
if planned:
|
||||
assert result["num_input_tokens"] == len(update.prompt_token_ids)
|
||||
replaced = True
|
||||
break
|
||||
session.num_prompt_tokens += 1 + len(update.prompt_token_ids)
|
||||
assert replaced
|
||||
|
||||
|
||||
def test_minicpmo_stage0_basic_window_rebuilds_from_retained_units():
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
|
||||
_MiniCPMO45Stage0SessionState,
|
||||
)
|
||||
|
||||
runtime = _stage0_vision_runtime()
|
||||
state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-basic-window", window_enabled=True)
|
||||
first = runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1)
|
||||
assert first["input_token_ids"] == [1, 11]
|
||||
state.pending_window_generated_tokens.append(5)
|
||||
state.pending_terminator_token = 3
|
||||
|
||||
rebuilt = runtime._stage_prefill_embeddings_only(
|
||||
state,
|
||||
np.zeros(4, dtype=np.float32),
|
||||
seq=2,
|
||||
stage0_window={
|
||||
"replace": True,
|
||||
"mode": "basic",
|
||||
"drop_units": 1,
|
||||
"replacement_prompt_len": 2,
|
||||
},
|
||||
)
|
||||
|
||||
assert rebuilt["stage0_window_replaced"] is True
|
||||
assert rebuilt["input_token_ids"] == [1, 11]
|
||||
assert state.window_units == []
|
||||
|
||||
|
||||
def test_minicpmo_stage0_context_window_inserts_previous_before_suffix():
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
|
||||
_MiniCPMO45Stage0SessionState,
|
||||
)
|
||||
|
||||
runtime = _stage0_vision_runtime()
|
||||
state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-context-window", window_enabled=True)
|
||||
prefix = runtime._embed_token(200)
|
||||
suffix = runtime._embed_token(202)
|
||||
state.context_embeds = [prefix, suffix]
|
||||
state.context_token_ids = [200, 202]
|
||||
state.context_prefix_embeds = [prefix]
|
||||
state.context_prefix_token_ids = [200]
|
||||
state.context_suffix_embeds = [suffix]
|
||||
state.context_suffix_token_ids = [202]
|
||||
first = runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1)
|
||||
assert first["input_token_ids"] == [200, 202, 1, 11]
|
||||
state.pending_window_generated_tokens.append(42)
|
||||
state.pending_terminator_token = 3
|
||||
|
||||
rebuilt = runtime._stage_prefill_embeddings_only(
|
||||
state,
|
||||
np.zeros(4, dtype=np.float32),
|
||||
seq=2,
|
||||
stage0_window={
|
||||
"replace": True,
|
||||
"mode": "context",
|
||||
"drop_units": 1,
|
||||
"previous_token_ids": [42],
|
||||
"previous_marker_token_ids": [201, 5],
|
||||
"replacement_prompt_len": 7,
|
||||
},
|
||||
)
|
||||
|
||||
assert rebuilt["input_token_ids"] == [200, 201, 5, 42, 202, 1, 11]
|
||||
assert rebuilt["num_input_tokens"] == 7
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pending_terminator", [None, 3, 99])
|
||||
def test_minicpmo_stage0_window_uses_accepted_output_not_async_sampler_history(pending_terminator):
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
|
||||
_MiniCPMO45Stage0SessionState,
|
||||
)
|
||||
|
||||
runtime = _stage0_vision_runtime()
|
||||
runtime._stage_audio_embeddings = lambda *args, **kwargs: torch.zeros((10, 2))
|
||||
state = _MiniCPMO45Stage0SessionState(session_id="sid-accepted-window", window_enabled=True)
|
||||
state.context_prefix_token_ids = [200] * 7
|
||||
state.context_suffix_token_ids = [202]
|
||||
state.context_token_ids = [*state.context_prefix_token_ids, 202]
|
||||
state.context_prefix_embeds = [runtime._embed_token(token_id) for token_id in state.context_prefix_token_ids]
|
||||
state.context_suffix_embeds = [runtime._embed_token(202)]
|
||||
state.context_embeds = [*state.context_prefix_embeds, *state.context_suffix_embeds]
|
||||
runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1)
|
||||
for seq in (2, 3):
|
||||
state.pending_window_generated_tokens = [3, 40, 41]
|
||||
state.pending_terminator_token = pending_terminator
|
||||
plan = {"completed_token_ids": [], "completed_terminator_token_id": 3}
|
||||
if seq == 3:
|
||||
plan.update(replace=True, mode="context", drop_units=1, replacement_prompt_len=32)
|
||||
result = runtime._stage_prefill_embeddings_only(
|
||||
state, np.zeros(4, dtype=np.float32), seq=seq, stage0_window=plan
|
||||
)
|
||||
|
||||
assert result["num_input_tokens"] == 32
|
||||
assert [len(unit.token_ids) for unit in state.window_units] == [13]
|
||||
assert len(state.pending_window_unit.token_ids) == 11
|
||||
assert 40 not in result["input_token_ids"]
|
||||
assert 41 not in result["input_token_ids"]
|
||||
|
||||
|
||||
def _stage0_vision_runtime():
|
||||
import torch
|
||||
|
||||
|
||||
@@ -77,6 +77,7 @@ _PUBLIC_OMNI = (
|
||||
"VLLM_OMNI_SERVER_STORAGE__PATH",
|
||||
"VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL",
|
||||
"VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP",
|
||||
"VLLM_OMNI_SPEAKER_REGISTRATION_POLICY",
|
||||
"VLLM_OMNI_STORAGE_MAX_CONCURRENCY",
|
||||
"VLLM_OMNI_STORAGE_PATH",
|
||||
"VLLM_OMNI_USE_QUACK_FP8",
|
||||
|
||||
@@ -848,6 +848,16 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
req_id = session.request_id
|
||||
self._new_prompt_len_snapshot[req_id] = len(update.prompt_token_ids)
|
||||
outstanding_async_tokens = getattr(session, "num_output_placeholders", 0)
|
||||
# Use the same confirmed span that upstream preserves when extending
|
||||
# a session. The segment output list is cleared on a resumable stop,
|
||||
# and, for an already queued append, can still include its terminator.
|
||||
confirmed_end = session.num_computed_tokens - outstanding_async_tokens
|
||||
segment_output_ids = list(session._all_token_ids[session.num_prompt_tokens : confirmed_end])
|
||||
completed_terminator = (
|
||||
session._all_token_ids[confirmed_end]
|
||||
if session.num_prompt_tokens <= confirmed_end < len(session._all_token_ids)
|
||||
else None
|
||||
)
|
||||
# Seed the stale share in SCHEDULED-token units (see the segment-stop
|
||||
# site in update_from_output): num_in_flight_tokens matches what each
|
||||
# pre-replacement frame will drain, so the counter reaches exactly
|
||||
@@ -894,6 +904,23 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
if self.log_stats:
|
||||
session.record_event(EngineCoreEventType.QUEUED)
|
||||
return
|
||||
if stage_id == 0 and self._prepare_minicpmo45_stage0_window(
|
||||
session,
|
||||
update,
|
||||
segment_output_ids=segment_output_ids,
|
||||
completed_terminator=completed_terminator,
|
||||
):
|
||||
# The rebuilt prompt is bounded by the client's window settings,
|
||||
# not by the model (a camera unit is hundreds of tokens), so it
|
||||
# needs the same max_model_len check as a plain extension. The
|
||||
# plan replaces the whole prompt, so the replacement length is the
|
||||
# projection, not the session's current prompt plus an extension.
|
||||
plan = update.model_intermediate_buffer["duplex"]["stage0_window"]
|
||||
if self._streaming_update_overflows(session, update, projected_len=plan["replacement_prompt_len"]):
|
||||
return
|
||||
self._release_replaced_streaming_prompt_cache(session)
|
||||
self._replace_streaming_session(session, update)
|
||||
return
|
||||
streaming_prompt_payload = next(
|
||||
(
|
||||
info
|
||||
@@ -962,11 +989,178 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
if hasattr(update, "model_intermediate_buffer"):
|
||||
session.model_intermediate_buffer = update.model_intermediate_buffer
|
||||
|
||||
@staticmethod
|
||||
def _prepare_minicpmo45_stage0_window(
|
||||
session: Request,
|
||||
update: StreamingUpdate,
|
||||
*,
|
||||
segment_output_ids: list[int],
|
||||
completed_terminator: int | None = None,
|
||||
) -> bool:
|
||||
"""Plan an official-style MiniCPM Stage-0 window at unit boundaries.
|
||||
|
||||
vLLM owns a paged KV cache, so deleting a middle span and rotating the
|
||||
retained K tensors in place is not a safe model hook. Instead, record
|
||||
completed unit lengths in the scheduler and request a full prompt
|
||||
replacement when a watermark fires. The worker rebuilds matching
|
||||
embeddings from its unit history, which recomputes RoPE at the new
|
||||
contiguous positions.
|
||||
"""
|
||||
info = getattr(update, "model_intermediate_buffer", None)
|
||||
if not isinstance(info, dict):
|
||||
return False
|
||||
duplex = info.get("duplex")
|
||||
if not isinstance(duplex, dict) or duplex.get("data_plane") is not True:
|
||||
return False
|
||||
runtime_config = duplex.get("runtime_config")
|
||||
runtime_config = runtime_config if isinstance(runtime_config, dict) else {}
|
||||
window = runtime_config.get("duplex_window_config")
|
||||
if not isinstance(window, dict):
|
||||
return False
|
||||
mode = window.get("sliding_window_mode", "off")
|
||||
if mode == "off":
|
||||
return False
|
||||
if mode not in {"basic", "context"}:
|
||||
return False
|
||||
|
||||
try:
|
||||
seq = int(duplex.get("seq", 0) or 0)
|
||||
preserve_len = int(runtime_config.get("duplex_first_append_context_tokens", 0) or 0)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
if seq <= 1:
|
||||
return False
|
||||
|
||||
# The sampled terminator is discarded by the normal session update and
|
||||
# re-injected at the head of this append, followed by </unit>.
|
||||
# generated_ids contains only the confirmed output span retained by
|
||||
# the streaming update, excluding its final uncomputed sample.
|
||||
generated_ids = segment_output_ids
|
||||
# Sampler-side history can contain speculative or stale async output.
|
||||
# The worker must rebuild from the same accepted ids used for lengths.
|
||||
duplex["stage0_window"] = {"completed_token_ids": list(generated_ids)}
|
||||
if completed_terminator is not None:
|
||||
duplex["stage0_window"]["completed_terminator_token_id"] = int(completed_terminator)
|
||||
base_len = int(getattr(session, "num_prompt_tokens", 0) or 0) + len(generated_ids)
|
||||
boundary = base_len + 2
|
||||
# A legitimate open_start of 0 (an empty context prefix) must not fall
|
||||
# back to preserve_len, which would double-count the suffix below.
|
||||
recorded_open_start = getattr(session, "_minicpmo45_window_open_start", None)
|
||||
open_start = preserve_len if recorded_open_start is None else int(recorded_open_start)
|
||||
unit_len = boundary - open_start
|
||||
if unit_len <= 0:
|
||||
return False
|
||||
|
||||
special_ids = {
|
||||
int(token_id)
|
||||
for token_id in runtime_config.get("duplex_window_special_token_ids", ())
|
||||
if isinstance(token_id, int)
|
||||
}
|
||||
units = list(getattr(session, "_minicpmo45_window_units", ()))
|
||||
units.append(
|
||||
{
|
||||
"length": unit_len,
|
||||
"generated_token_ids": [int(token_id) for token_id in generated_ids if token_id not in special_ids],
|
||||
}
|
||||
)
|
||||
projected_len = base_len + len(update.prompt_token_ids)
|
||||
drop_count = 0
|
||||
dropped_len = 0
|
||||
|
||||
if mode == "basic":
|
||||
try:
|
||||
high = int(window.get("basic_window_high_tokens", 8000))
|
||||
low = int(window.get("basic_window_low_tokens", 6000))
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
if projected_len > high:
|
||||
while units and projected_len - dropped_len > low:
|
||||
dropped_len += int(units[drop_count]["length"])
|
||||
drop_count += 1
|
||||
if drop_count >= len(units):
|
||||
break
|
||||
else:
|
||||
try:
|
||||
max_units = int(window.get("context_max_units", 24))
|
||||
previous_max = int(window.get("context_previous_max_tokens", 500))
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
drop_count = max(0, len(units) - max_units)
|
||||
dropped_len = sum(int(unit["length"]) for unit in units[:drop_count])
|
||||
previous = list(getattr(session, "_minicpmo45_window_previous_token_ids", ()))
|
||||
for unit in units[:drop_count]:
|
||||
previous.extend(unit["generated_token_ids"])
|
||||
if len(previous) > previous_max:
|
||||
previous = previous[-previous_max:]
|
||||
marker = [
|
||||
int(token_id)
|
||||
for token_id in runtime_config.get("duplex_window_previous_marker_token_ids", ())
|
||||
if isinstance(token_id, int)
|
||||
]
|
||||
previous_with_marker = marker + previous if previous else []
|
||||
old_previous_len = int(getattr(session, "_minicpmo45_window_previous_len", 0) or 0)
|
||||
session._minicpmo45_window_previous_token_ids = previous
|
||||
session._minicpmo45_window_previous_len = len(previous_with_marker)
|
||||
projected_len += len(previous_with_marker) - old_previous_len
|
||||
|
||||
retained_units = units[drop_count:]
|
||||
session._minicpmo45_window_units = retained_units
|
||||
if drop_count == 0:
|
||||
session._minicpmo45_window_open_start = boundary
|
||||
return False
|
||||
|
||||
replacement_len = projected_len - dropped_len
|
||||
if replacement_len <= 0:
|
||||
return False
|
||||
scheduler_token_id = runtime_config.get("duplex_scheduler_token_id", 0)
|
||||
try:
|
||||
scheduler_token_id = max(0, int(scheduler_token_id))
|
||||
except (TypeError, ValueError):
|
||||
scheduler_token_id = 0
|
||||
update.prompt_token_ids = [scheduler_token_id] * replacement_len
|
||||
meta = info.setdefault("meta", {})
|
||||
if isinstance(meta, dict):
|
||||
meta["replace_streaming_prompt"] = True
|
||||
previous_ids = list(getattr(session, "_minicpmo45_window_previous_token_ids", ()))
|
||||
# The worker embeds the marker ids the window plan carries: it must not
|
||||
# re-tokenize the marker with its own tokenizer, which could differ
|
||||
# from the ids this length is computed from.
|
||||
marker_ids = [
|
||||
int(token_id)
|
||||
for token_id in runtime_config.get("duplex_window_previous_marker_token_ids", ())
|
||||
if isinstance(token_id, int)
|
||||
]
|
||||
duplex["stage0_window"] = {
|
||||
"completed_token_ids": list(generated_ids),
|
||||
"replace": True,
|
||||
"mode": mode,
|
||||
"drop_units": drop_count,
|
||||
"dropped_tokens": dropped_len,
|
||||
"previous_token_ids": previous_ids,
|
||||
"previous_marker_token_ids": marker_ids,
|
||||
"replacement_prompt_len": replacement_len,
|
||||
}
|
||||
if completed_terminator is not None:
|
||||
duplex["stage0_window"]["completed_terminator_token_id"] = int(completed_terminator)
|
||||
recorded_prefix_len = runtime_config.get("duplex_window_prefix_tokens")
|
||||
prefix_len = preserve_len if recorded_prefix_len is None else int(recorded_prefix_len)
|
||||
suffix_len = len(runtime_config.get("duplex_window_suffix_token_ids", ()) or ())
|
||||
previous_len = int(getattr(session, "_minicpmo45_window_previous_len", 0) or 0)
|
||||
new_preserve_len = prefix_len + previous_len + suffix_len if mode == "context" else preserve_len
|
||||
session._minicpmo45_window_open_start = new_preserve_len + sum(int(unit["length"]) for unit in retained_units)
|
||||
return True
|
||||
|
||||
# Prefix of the stop_reason carried by the FinishReason.ERROR output, so the
|
||||
# serving side can map it to a stable error code.
|
||||
STREAMING_CONTEXT_OVERFLOW_STOP_REASON = "context_length_exceeded"
|
||||
|
||||
def _streaming_update_overflows(self, session: Request, update: StreamingUpdate) -> bool:
|
||||
def _streaming_update_overflows(
|
||||
self,
|
||||
session: Request,
|
||||
update: StreamingUpdate,
|
||||
*,
|
||||
projected_len: int | None = None,
|
||||
) -> bool:
|
||||
"""Finish a streaming session whose next extension cannot fit the model.
|
||||
|
||||
Upstream ``_update_request_as_session`` appends the update to the
|
||||
@@ -991,6 +1185,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
finish to the next ``schedule()`` would leave the client waiting. The
|
||||
reason is emitted with the terminal output (see
|
||||
:meth:`_emit_streaming_context_overflow_outputs`).
|
||||
|
||||
``projected_len`` overrides the extended-prompt projection for a caller
|
||||
that replaces the prompt instead of growing it, such as the MiniCPM
|
||||
Stage-0 window rebuild.
|
||||
"""
|
||||
max_model_len = getattr(self, "max_model_len", None)
|
||||
if max_model_len is None:
|
||||
@@ -1001,8 +1199,13 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
new_tokens = len(update.prompt_token_ids or ())
|
||||
# The extended prompt is the current prompt plus the computed output
|
||||
# tokens upstream keeps, then the update: num_computed_tokens covers
|
||||
# both when the prompt was fully computed.
|
||||
projected = max(int(session.num_prompt_tokens), int(session.num_computed_tokens)) + new_tokens
|
||||
# both when the prompt was fully computed. A caller that already knows
|
||||
# the post-update prompt length (a window rebuild) passes it instead.
|
||||
projected = (
|
||||
max(int(session.num_prompt_tokens), int(session.num_computed_tokens)) + new_tokens
|
||||
if projected_len is None
|
||||
else int(projected_len)
|
||||
)
|
||||
# Room for the tokens one step samples on top of the prompt (1 without
|
||||
# speculative decoding). __new__-built test schedulers carry no
|
||||
# num_sampled_tokens_per_step.
|
||||
|
||||
@@ -20,3 +20,13 @@ stages:
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
# Optional Cache-DiT acceleration for the DiT stage (~2x faster denoise,
|
||||
# training-free). Uncomment to enable; guidance=1.0 requests bypass it.
|
||||
# cache_backend: cache_dit
|
||||
# cache_config:
|
||||
# Fn_compute_blocks: 1
|
||||
# Bn_compute_blocks: 0
|
||||
# max_warmup_steps: 4
|
||||
# residual_diff_threshold: 0.24
|
||||
# max_continuous_cached_steps: 3
|
||||
# enable_cache_dit_summary: true
|
||||
@@ -867,6 +867,47 @@ def enable_cache_for_magi2(pipeline: Any, cache_config: Any) -> CacheDiTEnableRe
|
||||
return CacheDiTEnableResult(refresh=refresh, targets=(block_adapter,))
|
||||
|
||||
|
||||
def _get_mammothmoda2_transformer(pipeline: Any) -> torch.nn.Module:
|
||||
return pipeline.gen_transformer
|
||||
|
||||
|
||||
def enable_cache_for_mammothmoda2(pipeline: Any, cache_config: Any) -> CacheDiTEnableResult:
|
||||
"""Cache only MammothModa2's repeated main DiT stack.
|
||||
|
||||
``Transformer2DModel`` runs three Q-Former refiners (noise / ref-image /
|
||||
context) whose inputs change every denoise step, plus the ``layers``
|
||||
stack that dominates per-step compute. Only ``layers`` is a repeated
|
||||
residual stack, so the ``BlockAdapter`` wraps it and the refiners stay
|
||||
outside the cached region. Blocks take ``hidden_states`` plus keyword
|
||||
step context and return only hidden states (``Pattern_3``).
|
||||
|
||||
The pipeline runs sequential CFG: each conditional forward is followed by
|
||||
an unconditional forward. cache-dit tells cond/uncond apart purely by
|
||||
transformer-forward parity (``has_separate_cfg=True``), so the
|
||||
``cfg_range`` optimization that skips the unconditional pass outside the
|
||||
interval would desync that accounting. Like Cosmos3, we keep the passes
|
||||
paired and neutralize CFG via scale=1.0 outside the interval; the
|
||||
pipeline also disables hooks for no-CFG requests (single forward per
|
||||
step), whose parity the accounting cannot represent.
|
||||
"""
|
||||
pipeline._cache_dit_requires_paired_cfg = True
|
||||
transformer = _get_mammothmoda2_transformer(pipeline)
|
||||
block_adapter = BlockAdapter(
|
||||
transformer=transformer,
|
||||
blocks=[transformer.layers],
|
||||
forward_pattern=[ForwardPattern.Pattern_3],
|
||||
has_separate_cfg=True,
|
||||
check_forward_pattern=True,
|
||||
)
|
||||
refresh = enable_cache_for_dit(
|
||||
pipeline,
|
||||
cache_config,
|
||||
block_adapter,
|
||||
get_pipeline_transformer=_get_mammothmoda2_transformer,
|
||||
)
|
||||
return CacheDiTEnableResult(refresh=refresh, targets=(block_adapter,))
|
||||
|
||||
|
||||
def register_custom_dit_enablers() -> None:
|
||||
"""Register model-specific Cache-DiT enablers.
|
||||
|
||||
@@ -885,6 +926,7 @@ def register_custom_dit_enablers() -> None:
|
||||
"Cosmos3OmniPipeline": enable_cache_for_cosmos3,
|
||||
"Krea2Pipeline": enable_cache_for_krea2,
|
||||
"Magi2Pipeline": enable_cache_for_magi2,
|
||||
"MammothModa2DiTPipeline": enable_cache_for_mammothmoda2,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -695,8 +695,16 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
|
||||
return text_hidden_states, img_tokens
|
||||
|
||||
def _apply_transformer_layers(self, hidden_states, attention_mask, rotary_emb, temb):
|
||||
# Call blocks with keyword arguments: Cache-DiT's Pattern_3 wrapper
|
||||
# re-invokes blocks as block(hidden_states, **kwargs), so the cached
|
||||
# and uncached paths share this exact calling convention.
|
||||
for layer in self.layers:
|
||||
hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb)
|
||||
hidden_states = layer(
|
||||
hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
image_rotary_emb=rotary_emb,
|
||||
temb=temb,
|
||||
)
|
||||
return hidden_states
|
||||
|
||||
def forward(
|
||||
|
||||
@@ -15,7 +15,12 @@ from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.cache.cachedit import (
|
||||
CacheDiTBackend,
|
||||
CacheDiTRequestSpec,
|
||||
RequestScopedCacheDiTRuntime,
|
||||
)
|
||||
from vllm_omni.diffusion.data import DiffusionCacheConfig, DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
|
||||
@@ -28,6 +33,9 @@ from .schedulers import FlowMatchEulerDiscreteScheduler
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Identifies the pipeline-owned Cache-DiT installation across requests.
|
||||
_MAMMOTHMODA2_CACHE_DIT_KEY = "mammothmoda2:cache_dit"
|
||||
|
||||
|
||||
def _first_request_value(value: object) -> object:
|
||||
if isinstance(value, (list, tuple)):
|
||||
@@ -169,6 +177,54 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
|
||||
|
||||
self._llm_hidden_size = llm_hidden_size
|
||||
|
||||
# Cache-DiT lifecycle: the diffusion runner enables the configured
|
||||
# backend (``cache_backend`` on the deploy YAML stage entry) at startup
|
||||
# and transfers ownership here via the request-scoped protocol;
|
||||
# forward() then reconciles per-request state (step count, CFG parity).
|
||||
# The runner also emits the cache summary when
|
||||
# ``enable_cache_dit_summary`` is set.
|
||||
self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self)
|
||||
self._cache_dit_config: DiffusionCacheConfig | None = None
|
||||
if str(getattr(od_config, "cache_backend", "") or "").lower() == "cache_dit":
|
||||
cache_config = od_config.cache_config
|
||||
self._cache_dit_config = (
|
||||
cache_config
|
||||
if isinstance(cache_config, DiffusionCacheConfig)
|
||||
else DiffusionCacheConfig.from_dict(cache_config or {})
|
||||
)
|
||||
|
||||
def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None:
|
||||
"""Adopt a runner-installed Cache-DiT backend (request-scoped protocol)."""
|
||||
|
||||
self._cache_dit_runtime.adopt(backend, installation_key=_MAMMOTHMODA2_CACHE_DIT_KEY)
|
||||
|
||||
def is_cache_dit_enabled(self) -> bool:
|
||||
"""Return the request-scoped Cache-DiT installation state."""
|
||||
|
||||
return self._cache_dit_runtime.is_enabled
|
||||
|
||||
def _prepare_cache_dit_for_request(self, *, num_inference_steps: int, cfg_active: bool) -> None:
|
||||
"""Reconcile Cache-DiT hooks with this request before the denoise loop.
|
||||
|
||||
cache-dit's separate-CFG accounting assumes exactly two transformer
|
||||
forwards per denoise step (conditional then unconditional), so CFG
|
||||
requests refresh the context (per-request step counts vary) while
|
||||
no-CFG requests run with hooks disabled: their single forward per
|
||||
step cannot be represented by the parity accounting.
|
||||
"""
|
||||
if self._cache_dit_config is None:
|
||||
return
|
||||
if not cfg_active:
|
||||
self._cache_dit_runtime.prepare(None)
|
||||
return
|
||||
self._cache_dit_runtime.prepare(
|
||||
CacheDiTRequestSpec(
|
||||
installation_key=_MAMMOTHMODA2_CACHE_DIT_KEY,
|
||||
cache_config=self._cache_dit_config,
|
||||
num_inference_steps=num_inference_steps,
|
||||
)
|
||||
)
|
||||
|
||||
def _reinit_caption_embedder(self, in_features: int) -> None:
|
||||
# Align with upstream Mammothmoda2Model's `reinit_caption_embedder`:
|
||||
# Use Qwen2RMSNorm(in_features) + Linear(in_features -> out_features).
|
||||
@@ -432,6 +488,18 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
|
||||
|
||||
# Run diffusion loop (CFG supported when text_guidance_scale > 1.0)
|
||||
total_steps = max(1, len(scheduler.timesteps))
|
||||
self._prepare_cache_dit_for_request(
|
||||
num_inference_steps=total_steps,
|
||||
cfg_active=request.text_guidance_scale > 1.0,
|
||||
)
|
||||
# cache-dit distinguishes the conditional vs unconditional passes
|
||||
# purely by transformer-forward parity, so when it is active the
|
||||
# uncond pass must run on every step. Outside cfg_range CFG is
|
||||
# neutralized via scale=1.0 (uncond + 1.0 * (cond - uncond) == cond),
|
||||
# matching the skip optimization's cond-only update.
|
||||
requires_paired_cfg = self._cache_dit_runtime.is_enabled and getattr(
|
||||
self, "_cache_dit_requires_paired_cfg", False
|
||||
)
|
||||
for i, t in enumerate(scheduler.timesteps):
|
||||
timestep = t.expand(latents.shape[0]).to(latents.dtype)
|
||||
model_pred = self.gen_transformer(
|
||||
@@ -447,7 +515,10 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
|
||||
guidance_scale = (
|
||||
request.text_guidance_scale if request.cfg_range[0] <= i / total_steps <= request.cfg_range[1] else 1.0
|
||||
)
|
||||
if guidance_scale > 1.0 and negative_prompt_embeds is not None:
|
||||
run_uncond = guidance_scale > 1.0 and negative_prompt_embeds is not None
|
||||
if requires_paired_cfg and negative_prompt_embeds is not None:
|
||||
run_uncond = True
|
||||
if run_uncond:
|
||||
model_pred_uncond = self.gen_transformer(
|
||||
hidden_states=latents,
|
||||
timestep=timestep,
|
||||
|
||||
@@ -141,7 +141,8 @@ class DuplexSessionTasks:
|
||||
task.cancel()
|
||||
try:
|
||||
await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=timeout_s)
|
||||
except TimeoutError:
|
||||
# asyncio.TimeoutError is not the builtin TimeoutError before Python 3.11.
|
||||
except (TimeoutError, asyncio.TimeoutError):
|
||||
pass
|
||||
if cancelled_tail is not None and self.append_tail is cancelled_tail:
|
||||
self.append_tail = None
|
||||
|
||||
@@ -699,7 +699,9 @@ class DuplexSessionManager:
|
||||
while not shutdown_event.is_set():
|
||||
try:
|
||||
await asyncio.wait_for(shutdown_event.wait(), timeout=interval)
|
||||
except TimeoutError:
|
||||
# The timeout is this loop's tick. asyncio.TimeoutError is not the builtin
|
||||
# TimeoutError before Python 3.11, so catch both or the tick escapes the loop.
|
||||
except (TimeoutError, asyncio.TimeoutError):
|
||||
try:
|
||||
await self.reap_expired()
|
||||
except Exception:
|
||||
|
||||
@@ -1486,7 +1486,8 @@ class DuplexSessionRunner:
|
||||
active_task.cancel()
|
||||
try:
|
||||
await asyncio.wait_for(asyncio.gather(active_task, return_exceptions=True), timeout=0.25)
|
||||
except TimeoutError:
|
||||
# asyncio.TimeoutError is not the builtin TimeoutError before Python 3.11.
|
||||
except (TimeoutError, asyncio.TimeoutError):
|
||||
pass
|
||||
if notify:
|
||||
self.emit(
|
||||
|
||||
@@ -422,7 +422,7 @@ class OmniEngineBase:
|
||||
for client in self.stage_clients
|
||||
]
|
||||
supported_tasks: set[str] = set()
|
||||
if any(getattr(client, "is_comprehension", False) for client in self.stage_clients):
|
||||
if any(getattr(stage_config, "is_comprehension", False) for stage_config in self.stage_configs):
|
||||
supported_tasks.add("generate")
|
||||
if any(meta.final_output_type == "audio" for meta in self.stage_metadata):
|
||||
supported_tasks.add("speech")
|
||||
|
||||
@@ -186,7 +186,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
|
||||
|
||||
# Reject diffusion list-prompt early with a clear API error.
|
||||
if isinstance(prompt, list) and any(
|
||||
getattr(client, "stage_type", "") == "diffusion" for client in getattr(self.engine, "stage_clients", [])
|
||||
stage_config.stage_type == "diffusion" for stage_config in self.engine.stage_configs
|
||||
):
|
||||
raise ValueError(
|
||||
"Diffusion stages accept only a single prompt per request. "
|
||||
@@ -626,7 +626,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
|
||||
|
||||
def _split_stage_ids_by_type(self, stage_ids: list[int] | None = None) -> tuple[list[int], list[int]]:
|
||||
"""Split stage ids into AR/LLM (EngineCore) vs diffusion (worker RPC)."""
|
||||
n_stages = len(self.engine.stage_clients)
|
||||
n_stages = len(self.engine.stage_configs)
|
||||
if stage_ids is None:
|
||||
stage_ids = list(range(n_stages))
|
||||
else:
|
||||
@@ -640,8 +640,8 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
|
||||
ar_stage_ids: list[int] = []
|
||||
diffusion_stage_ids: list[int] = []
|
||||
for sid in stage_ids:
|
||||
client = self.engine.stage_clients[sid]
|
||||
if getattr(client, "stage_type", "llm") == "diffusion":
|
||||
stage_config = self.engine.stage_configs[sid]
|
||||
if stage_config.stage_type == "diffusion":
|
||||
diffusion_stage_ids.append(sid)
|
||||
else:
|
||||
ar_stage_ids.append(sid)
|
||||
|
||||
@@ -155,13 +155,13 @@ class AsyncOmniBase(OmniBase):
|
||||
|
||||
def _get_comprehension_stage_index(self) -> int | None:
|
||||
fallback_idx: int | None = None
|
||||
for idx, stage_client in enumerate(self.engine.stage_clients):
|
||||
for idx, stage_config in enumerate(self.engine.stage_configs):
|
||||
stage_vllm_config = self.engine.stage_vllm_configs[idx]
|
||||
if stage_vllm_config is None:
|
||||
continue
|
||||
if fallback_idx is None:
|
||||
fallback_idx = idx
|
||||
if stage_client.is_comprehension:
|
||||
if stage_config.is_comprehension:
|
||||
return idx
|
||||
return fallback_idx
|
||||
|
||||
@@ -186,13 +186,11 @@ class AsyncOmniBase(OmniBase):
|
||||
|
||||
def get_diffusion_od_config(self) -> Any | None:
|
||||
"""Return the diffusion-stage config when the pipeline has one."""
|
||||
saw_diffusion_stage = False
|
||||
saw_diffusion_stage = any(stage_config.stage_type == "diffusion" for stage_config in self.engine.stage_configs)
|
||||
for stage_client in self.engine.stage_clients:
|
||||
if getattr(stage_client, "stage_type", None) != "diffusion":
|
||||
continue
|
||||
|
||||
saw_diffusion_stage = True
|
||||
|
||||
od_config = getattr(stage_client, "od_config", None)
|
||||
if od_config is not None:
|
||||
return od_config
|
||||
|
||||
@@ -28,7 +28,8 @@ async def receive_text_with_timeout(websocket: WebSocket, timeout_s: float | Non
|
||||
return await websocket.receive_text()
|
||||
try:
|
||||
return await asyncio.wait_for(websocket.receive_text(), timeout=timeout_s)
|
||||
except TimeoutError:
|
||||
# asyncio.TimeoutError is not the builtin TimeoutError before Python 3.11.
|
||||
except (TimeoutError, asyncio.TimeoutError):
|
||||
return None
|
||||
|
||||
|
||||
|
||||
@@ -246,6 +246,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
except ValueError:
|
||||
logger.warning("Invalid SPEAKER_MAX_UPLOADED=%r; using default 1000", _raw_cap)
|
||||
self._max_uploaded_speakers = 1000
|
||||
_policy = os.environ.get("VLLM_OMNI_SPEAKER_REGISTRATION_POLICY", "overwrite").lower()
|
||||
if _policy not in ("overwrite", "immutable"):
|
||||
raise ValueError(
|
||||
f"Invalid VLLM_OMNI_SPEAKER_REGISTRATION_POLICY={_policy!r}; expected 'overwrite' or 'immutable'."
|
||||
)
|
||||
self._registration_policy = _policy
|
||||
self.uploaded_speakers: dict[str, dict[str, Any]] = {}
|
||||
self._ref_audio_data_url_cache: dict[str, str] = {}
|
||||
self._ref_audio_resolve_cache: OrderedDict[str, tuple[np.ndarray, int, int, str]] = OrderedDict()
|
||||
@@ -410,6 +416,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
adapter = self._adapter
|
||||
if adapter is not None:
|
||||
adapter.load_capabilities()
|
||||
self._drop_shadowing_uploads()
|
||||
available_speakers = self._get_available_speakers()
|
||||
logger.info("Loaded %d supported speakers: %s", len(available_speakers), sorted(available_speakers))
|
||||
|
||||
@@ -448,6 +455,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
ctx = SpeechServingContext(server=self, engine_client=self.engine_client)
|
||||
self._adapter = adapter_cls(ctx)
|
||||
self._adapter.load_capabilities()
|
||||
self._drop_shadowing_uploads()
|
||||
return self._adapter
|
||||
|
||||
def _uses_native_speed_control(self) -> bool:
|
||||
@@ -803,6 +811,39 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
f"the cap via SPEAKER_MAX_UPLOADED."
|
||||
)
|
||||
|
||||
def _drop_shadowing_uploads(self) -> None:
|
||||
"""Uploads are restored from disk before the adapter exists, so a name
|
||||
registered under a built-in speaker before the collision guard would
|
||||
keep shadowing it across restarts. Drop such entries from the registry;
|
||||
the file stays on disk for the operator to remove."""
|
||||
caps = self._adapter.capabilities if self._adapter is not None else None
|
||||
if caps is None:
|
||||
return
|
||||
for voice_name_lower in list(self.uploaded_speakers):
|
||||
if voice_name_lower in caps.supported_speakers or voice_name_lower in caps.precomputed_speakers:
|
||||
info = self.uploaded_speakers.pop(voice_name_lower)
|
||||
logger.warning(
|
||||
"Uploaded voice %r shadows a built-in speaker and is ignored; remove %s to clear this warning.",
|
||||
voice_name_lower,
|
||||
info.get("file_path"),
|
||||
)
|
||||
|
||||
def _check_registration_allowed(self, voice_name_lower: str, name: str) -> None:
|
||||
"""Reject names that would shadow a built-in voice or, under the
|
||||
immutable policy, silently overwrite an existing upload."""
|
||||
caps = self._adapter.capabilities if self._adapter is not None else None
|
||||
if caps is not None and (
|
||||
voice_name_lower in caps.supported_speakers or voice_name_lower in caps.precomputed_speakers
|
||||
):
|
||||
raise ValueError(
|
||||
f"Voice name '{name}' is reserved by a built-in speaker of this model; choose a different name."
|
||||
)
|
||||
if self._registration_policy == "immutable" and voice_name_lower in self.uploaded_speakers:
|
||||
raise ValueError(
|
||||
f"Voice '{name}' already exists and VLLM_OMNI_SPEAKER_REGISTRATION_POLICY is 'immutable'; "
|
||||
f"delete it first via DELETE /v1/audio/voices/{name}."
|
||||
)
|
||||
|
||||
def _evict_existing_upload(self, voice_name_lower: str, name: str) -> None:
|
||||
"""Drop an existing upload with this name so the caller can re-register it."""
|
||||
if voice_name_lower not in self.uploaded_speakers:
|
||||
@@ -885,6 +926,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
|
||||
async with self._upload_lock:
|
||||
voice_name_lower = name.lower()
|
||||
self._check_registration_allowed(voice_name_lower, name)
|
||||
self._evict_existing_upload(voice_name_lower, name)
|
||||
self._check_upload_cap()
|
||||
|
||||
@@ -1005,6 +1047,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
|
||||
async with self._upload_lock:
|
||||
voice_name_lower = name.lower()
|
||||
self._check_registration_allowed(voice_name_lower, name)
|
||||
self._evict_existing_upload(voice_name_lower, name)
|
||||
self._check_upload_cap()
|
||||
|
||||
|
||||
@@ -10,7 +10,8 @@ This module contains the code2wav (token-to-waveform) stage which uses:
|
||||
"""
|
||||
|
||||
from collections import Counter
|
||||
from typing import TypedDict, cast
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, TypedDict, cast
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
@@ -18,6 +19,8 @@ import torch.nn as nn
|
||||
from omegaconf import DictConfig
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.config import get_current_diffusion_config_or_none, set_current_diffusion_config
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT
|
||||
from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import (
|
||||
CausalConditionalCFM,
|
||||
@@ -37,6 +40,23 @@ from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _build_dit_estimator(estimator_config: Mapping[str, Any]) -> DiT:
|
||||
"""Build CosyVoice's embedded DiT with diffusion backend configuration.
|
||||
|
||||
The code2wav model runs inside an LLM-generation worker, so it does not go
|
||||
through the diffusion model loader that normally installs an
|
||||
``OmniDiffusionConfig`` during model construction. Install a local default
|
||||
config when needed so ``DIFFUSION_ATTENTION_BACKEND`` and the platform
|
||||
default are resolved by the same path as standalone diffusion models.
|
||||
"""
|
||||
if get_current_diffusion_config_or_none() is not None:
|
||||
return DiT(**estimator_config)
|
||||
|
||||
diffusion_config = OmniDiffusionConfig()
|
||||
with set_current_diffusion_config(diffusion_config):
|
||||
return DiT(**estimator_config)
|
||||
|
||||
|
||||
class StreamingFlowItem(TypedDict, total=False):
|
||||
"""One entry of the batched-streaming item list passed to forward_streaming_batch."""
|
||||
|
||||
@@ -71,7 +91,7 @@ class CosyVoice3Code2Wav(nn.Module):
|
||||
cfm_params = DictConfig(decoder_cfg["cfm_params"])
|
||||
|
||||
# DiT estimator using diffusion attention (Flash/Sage/SDPA backends)
|
||||
estimator = DiT(**decoder_cfg["estimator"])
|
||||
estimator = _build_dit_estimator(decoder_cfg["estimator"])
|
||||
|
||||
decoder = CausalConditionalCFM(
|
||||
in_channels=decoder_cfg["in_channels"],
|
||||
|
||||
@@ -45,7 +45,10 @@ from vllm_omni.model_executor.models.minicpmo_4_5.duplex.data_plane import (
|
||||
MiniCPMO45DataPlaneContext,
|
||||
MiniCPMO45DataPlaneSession,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import MiniCPMO45DuplexPolicy
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import (
|
||||
MiniCPMO45DuplexPolicy,
|
||||
MiniCPMO45DuplexWindowConfig,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.duplex.session import (
|
||||
MiniCPMO45ServingSessionState,
|
||||
)
|
||||
@@ -85,6 +88,11 @@ PRIVATE_RUNTIME_CONFIG_KEYS = frozenset(
|
||||
"ref_audio_format",
|
||||
"ref_audio_sample_rate_hz",
|
||||
"initial_user_text",
|
||||
"duplex_window_config",
|
||||
"duplex_window_prefix_tokens",
|
||||
"duplex_window_suffix_token_ids",
|
||||
"duplex_window_previous_marker_token_ids",
|
||||
"duplex_window_special_token_ids",
|
||||
}
|
||||
)
|
||||
|
||||
@@ -245,8 +253,8 @@ def build_duplex_data_plane_prompt(
|
||||
token_budget = context_reserve + first_units * 12 - 1 + vision_tokens
|
||||
if seq > 1 and duplex_payload_is_exact_chunks(payload):
|
||||
token_budget += 1
|
||||
if final and duplex_payload_is_exact_chunks(payload):
|
||||
token_budget += 12
|
||||
# Serving already pads the final residual audio. Stage0 does not
|
||||
# append another silent unit, so final must not reserve extra slots.
|
||||
extra_body = session_config.get("extra_body")
|
||||
raw_token_id = runtime_config.get("duplex_scheduler_token_id")
|
||||
try:
|
||||
@@ -515,6 +523,16 @@ def _apply_first_append_context_tokens(
|
||||
return
|
||||
ref_tokens = MiniCPMO45DuplexPolicy.audio_token_count(ref_sample_count or 0)
|
||||
runtime_config["duplex_first_append_context_tokens"] = len(prefix_ids) + ref_tokens + len(suffix_ids)
|
||||
runtime_config["duplex_window_prefix_tokens"] = len(prefix_ids) + ref_tokens
|
||||
runtime_config["duplex_window_suffix_token_ids"] = [int(token_id) for token_id in suffix_ids]
|
||||
marker_ids = tokenizer.encode("\n\nprevious: ", add_special_tokens=False)
|
||||
runtime_config["duplex_window_previous_marker_token_ids"] = [int(token_id) for token_id in marker_ids]
|
||||
runtime_config["duplex_window_special_token_ids"] = sorted(
|
||||
MiniCPMO45DuplexPolicy.native_special_token_ids(
|
||||
MiniCPMO45DuplexPolicy.token_ids_from_tokenizer(tokenizer),
|
||||
tokenizer_special_ids=list(getattr(tokenizer, "all_special_ids", ()) or ()),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _model_vision_tile_pixels(model_config: ModelConfig | None) -> int | None:
|
||||
@@ -722,6 +740,8 @@ class MiniCPMO45DuplexPlugin(DuplexModelPlugin):
|
||||
code="unsupported_ref_audio_path",
|
||||
)
|
||||
runtime_config: dict[str, object] = {"instructions": config.instructions}
|
||||
window_config = self._pop_window_config(extra_body) or MiniCPMO45DuplexWindowConfig()
|
||||
runtime_config["duplex_window_config"] = window_config.as_dict()
|
||||
# ``duplex_initial_user_text`` is the older extra_body spelling and
|
||||
# still works; the session field is the framework-level one.
|
||||
initial_user_text = extra_body.pop("duplex_initial_user_text", None)
|
||||
@@ -788,6 +808,17 @@ class MiniCPMO45DuplexPlugin(DuplexModelPlugin):
|
||||
current: Mapping[str, object],
|
||||
) -> dict[str, object]:
|
||||
runtime_config = deepcopy(dict(current))
|
||||
extra_body = dict(config.extra_body)
|
||||
requested_window = self._pop_window_config(extra_body)
|
||||
if requested_window is not None:
|
||||
reject_changed_runtime_value(
|
||||
requested_window.as_dict(),
|
||||
runtime_config.get("duplex_window_config"),
|
||||
message="sliding-window configuration cannot be changed after the session is created",
|
||||
code="sliding_window_update_unsupported",
|
||||
error_cls=MiniCPMO45ClientRuntimeConfigError,
|
||||
)
|
||||
config.extra_body = extra_body
|
||||
reject_changed_runtime_value(
|
||||
config.instructions,
|
||||
runtime_config.get("instructions"),
|
||||
@@ -812,6 +843,23 @@ class MiniCPMO45DuplexPlugin(DuplexModelPlugin):
|
||||
runtime_config["duplex_stage_sampling_params"] = stage_sampling
|
||||
return runtime_config
|
||||
|
||||
@staticmethod
|
||||
def _pop_window_config(extra_body: dict[str, object]) -> MiniCPMO45DuplexWindowConfig | None:
|
||||
names = (
|
||||
"sliding_window_mode",
|
||||
"basic_window_high_tokens",
|
||||
"basic_window_low_tokens",
|
||||
"context_previous_max_tokens",
|
||||
"context_max_units",
|
||||
)
|
||||
provided = {name: extra_body.pop(name) for name in names if name in extra_body}
|
||||
if not provided:
|
||||
return None
|
||||
try:
|
||||
return MiniCPMO45DuplexWindowConfig.from_mapping(provided)
|
||||
except ValueError as exc:
|
||||
raise MiniCPMO45ClientRuntimeConfigError(str(exc), code="invalid_sliding_window_config") from exc
|
||||
|
||||
def data_plane_context(
|
||||
self,
|
||||
*,
|
||||
|
||||
@@ -3,12 +3,63 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from transformers import PreTrainedTokenizerBase
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MiniCPMO45DuplexWindowConfig:
|
||||
"""Stage-0 window settings matching the released checkpoint defaults."""
|
||||
|
||||
sliding_window_mode: str = "off"
|
||||
basic_window_high_tokens: int = 8000
|
||||
basic_window_low_tokens: int = 6000
|
||||
context_previous_max_tokens: int = 500
|
||||
context_max_units: int = 24
|
||||
|
||||
@classmethod
|
||||
def from_mapping(cls, value: object) -> MiniCPMO45DuplexWindowConfig:
|
||||
source = value if isinstance(value, dict) else {}
|
||||
|
||||
def integer(name: str, default: int) -> int:
|
||||
raw = source.get(name, default)
|
||||
if isinstance(raw, bool):
|
||||
raise ValueError(f"{name} must be an integer")
|
||||
try:
|
||||
parsed = int(raw)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(f"{name} must be an integer") from exc
|
||||
if parsed <= 0:
|
||||
raise ValueError(f"{name} must be greater than zero")
|
||||
return parsed
|
||||
|
||||
mode = source.get("sliding_window_mode", "off")
|
||||
if mode not in {"off", "basic", "context"}:
|
||||
raise ValueError("sliding_window_mode must be one of: off, basic, context")
|
||||
config = cls(
|
||||
sliding_window_mode=str(mode),
|
||||
basic_window_high_tokens=integer("basic_window_high_tokens", 8000),
|
||||
basic_window_low_tokens=integer("basic_window_low_tokens", 6000),
|
||||
context_previous_max_tokens=integer("context_previous_max_tokens", 500),
|
||||
context_max_units=integer("context_max_units", 24),
|
||||
)
|
||||
if config.basic_window_low_tokens >= config.basic_window_high_tokens:
|
||||
raise ValueError("basic_window_low_tokens must be less than basic_window_high_tokens")
|
||||
return config
|
||||
|
||||
def as_dict(self) -> dict[str, int | str]:
|
||||
return {
|
||||
"sliding_window_mode": self.sliding_window_mode,
|
||||
"basic_window_high_tokens": self.basic_window_high_tokens,
|
||||
"basic_window_low_tokens": self.basic_window_low_tokens,
|
||||
"context_previous_max_tokens": self.context_previous_max_tokens,
|
||||
"context_max_units": self.context_max_units,
|
||||
}
|
||||
|
||||
|
||||
class MiniCPMO45DuplexPolicy:
|
||||
"""MiniCPM-o 4.5 native duplex model policy.
|
||||
|
||||
|
||||
@@ -27,6 +27,12 @@ _MINICPMO45_OPTIONAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS
|
||||
_MINICPMO45_PROCESSOR_LOAD_LOCK = Lock()
|
||||
|
||||
|
||||
@dataclass
|
||||
class _MiniCPMO45WindowUnit:
|
||||
embeds: list[Any] = field(default_factory=list)
|
||||
token_ids: list[int] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _MiniCPMO45Stage0SessionState:
|
||||
session_id: str
|
||||
@@ -36,6 +42,14 @@ class _MiniCPMO45Stage0SessionState:
|
||||
audio_chunk_idx: int = 0
|
||||
context_embeds: list[torch.Tensor] = field(default_factory=list)
|
||||
context_token_ids: list[int] = field(default_factory=list)
|
||||
context_prefix_embeds: list[Any] = field(default_factory=list)
|
||||
context_prefix_token_ids: list[int] = field(default_factory=list)
|
||||
context_suffix_embeds: list[Any] = field(default_factory=list)
|
||||
context_suffix_token_ids: list[int] = field(default_factory=list)
|
||||
window_units: list[_MiniCPMO45WindowUnit] = field(default_factory=list)
|
||||
window_enabled: bool = False
|
||||
pending_window_unit: _MiniCPMO45WindowUnit | None = None
|
||||
pending_window_generated_tokens: list[int] = field(default_factory=list)
|
||||
current_turn_ended: bool = True
|
||||
prepared_append_identity: tuple[int | None, int] | None = None
|
||||
prepared_inputs_embeds: torch.Tensor | None = None
|
||||
@@ -142,6 +156,11 @@ class MiniCPMO45Stage0DuplexRuntime:
|
||||
*,
|
||||
runtime_config: dict[str, object] | None = None,
|
||||
) -> None:
|
||||
window_config = (runtime_config or {}).get("duplex_window_config")
|
||||
state.window_enabled = isinstance(window_config, dict) and window_config.get("sliding_window_mode", "off") in {
|
||||
"basic",
|
||||
"context",
|
||||
}
|
||||
if not self._stage_runtime_ready():
|
||||
return
|
||||
self._require_special_token_ids()
|
||||
@@ -157,17 +176,26 @@ class MiniCPMO45Stage0DuplexRuntime:
|
||||
(runtime_config or {}).get("initial_user_text"),
|
||||
)
|
||||
for token_id in self._encode_text(prefix):
|
||||
state.context_embeds.append(self._embed_token(token_id))
|
||||
embed = self._embed_token(token_id)
|
||||
state.context_embeds.append(embed)
|
||||
state.context_token_ids.append(token_id)
|
||||
state.context_prefix_embeds.append(embed)
|
||||
state.context_prefix_token_ids.append(token_id)
|
||||
if ref_audio is not None:
|
||||
ref_audio_embeds = self._stage_ref_audio_embeddings(ref_audio, state=state)
|
||||
if ref_audio_embeds is not None:
|
||||
ref_audio_embeds = self._as_2d_tensor(ref_audio_embeds)
|
||||
placeholder_ids = [self.unit_token_id] * int(ref_audio_embeds.shape[0])
|
||||
state.context_embeds.append(ref_audio_embeds)
|
||||
state.context_token_ids.extend([self.unit_token_id] * int(ref_audio_embeds.shape[0]))
|
||||
state.context_token_ids.extend(placeholder_ids)
|
||||
state.context_prefix_embeds.append(ref_audio_embeds)
|
||||
state.context_prefix_token_ids.extend(placeholder_ids)
|
||||
for token_id in self._encode_text(suffix):
|
||||
state.context_embeds.append(self._embed_token(token_id))
|
||||
embed = self._embed_token(token_id)
|
||||
state.context_embeds.append(embed)
|
||||
state.context_token_ids.append(token_id)
|
||||
state.context_suffix_embeds.append(embed)
|
||||
state.context_suffix_token_ids.append(token_id)
|
||||
|
||||
def _stage_prefill_embeddings_only(
|
||||
self,
|
||||
@@ -179,6 +207,7 @@ class MiniCPMO45Stage0DuplexRuntime:
|
||||
seq: int | None = None,
|
||||
is_speech: bool = False,
|
||||
final: bool = False,
|
||||
stage0_window: dict[str, object] | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""Build scheduler-owned Stage0 input embeddings for one audio append.
|
||||
|
||||
@@ -199,6 +228,16 @@ class MiniCPMO45Stage0DuplexRuntime:
|
||||
result["input_token_ids"] = list(state.prepared_input_token_ids)
|
||||
return result
|
||||
self._require_special_token_ids()
|
||||
if isinstance(stage0_window, dict):
|
||||
completed_ids = stage0_window.get("completed_token_ids")
|
||||
if isinstance(completed_ids, list):
|
||||
state.pending_window_generated_tokens = [int(token_id) for token_id in completed_ids]
|
||||
completed_terminator = stage0_window.get("completed_terminator_token_id")
|
||||
if isinstance(completed_terminator, int):
|
||||
# Async sampling can overwrite the worker's pending token
|
||||
# after the scheduler accepted a segment stop. Rebuild from
|
||||
# the scheduler's canonical boundary, just like its content.
|
||||
state.pending_terminator_token = completed_terminator
|
||||
if audio_waveform is None or len(audio_waveform) == 0:
|
||||
return self._stage_prefill_result(False, start_time, "empty audio")
|
||||
state.audio_buffer = np.concatenate([state.audio_buffer, np.asarray(audio_waveform, dtype=np.float32)])
|
||||
@@ -279,12 +318,25 @@ class MiniCPMO45Stage0DuplexRuntime:
|
||||
# the closure; the model's listen/speak policy depends on
|
||||
# seeing its own past decisions in context.
|
||||
pending_terminator = state.pending_terminator_token
|
||||
closure_token_ids: list[int] = []
|
||||
if pending_terminator is not None and units_built == 0:
|
||||
state.pending_terminator_token = None
|
||||
embed_parts.append(self._embed_token(pending_terminator))
|
||||
token_ids.append(int(pending_terminator))
|
||||
closure_token_ids.append(int(pending_terminator))
|
||||
embed_parts.append(self._embed_token(self.unit_end_token_id))
|
||||
token_ids.append(self.unit_end_token_id)
|
||||
closure_token_ids.append(self.unit_end_token_id)
|
||||
if units_built == 0:
|
||||
self._finalize_window_unit(state, closure_token_ids)
|
||||
elif state.pending_window_unit is not None:
|
||||
# Internal processor chunks belong to this same append.
|
||||
state.pending_window_unit.embeds.extend(
|
||||
self._embed_token(token_id) for token_id in closure_token_ids
|
||||
)
|
||||
state.pending_window_unit.token_ids.extend(closure_token_ids)
|
||||
unit_embed_start = len(embed_parts)
|
||||
unit_token_start = len(token_ids)
|
||||
embed_parts.append(self._embed_token(self.unit_token_id))
|
||||
token_ids.append(self.unit_token_id)
|
||||
if frame_blocks:
|
||||
@@ -311,6 +363,11 @@ class MiniCPMO45Stage0DuplexRuntime:
|
||||
state.audio_buffer = state.audio_buffer[consumed_samples:]
|
||||
state.audio_chunk_idx += 1
|
||||
units_built += 1
|
||||
if state.window_enabled:
|
||||
if state.pending_window_unit is None:
|
||||
state.pending_window_unit = _MiniCPMO45WindowUnit()
|
||||
state.pending_window_unit.embeds.extend(embed_parts[unit_embed_start:])
|
||||
state.pending_window_unit.token_ids.extend(token_ids[unit_token_start:])
|
||||
chunk_size = self._streaming_chunk_size(processor)
|
||||
# Match official streaming_prefill: per chunk feed ONLY <unit>+audio. The assistant
|
||||
# turn is opened once at session init; re-emitting the turn-open prefix per chunk
|
||||
@@ -320,6 +377,9 @@ class MiniCPMO45Stage0DuplexRuntime:
|
||||
|
||||
import torch
|
||||
|
||||
window_result = self._window_replacement_parts(state, stage0_window)
|
||||
if window_result is not None:
|
||||
embed_parts, token_ids = window_result
|
||||
inputs_embeds = torch.cat([self._as_2d_tensor(embed) for embed in embed_parts], dim=0)
|
||||
result = self._stage_prefill_result(True, start_time)
|
||||
result.update(
|
||||
@@ -332,6 +392,7 @@ class MiniCPMO45Stage0DuplexRuntime:
|
||||
"uses_model_runner_scheduler": True,
|
||||
"runner_kv_backed": True,
|
||||
"runtime_impl": "scheduler_data_plane",
|
||||
"stage0_window_replaced": window_result is not None,
|
||||
}
|
||||
)
|
||||
if is_speech and (append_identity is None or state.pending_speech_append_identity != append_identity):
|
||||
@@ -344,6 +405,85 @@ class MiniCPMO45Stage0DuplexRuntime:
|
||||
state.prepared_result = {k: v for k, v in result.items() if k not in {"inputs_embeds", "input_token_ids"}}
|
||||
return result
|
||||
|
||||
def _finalize_window_unit(
|
||||
self,
|
||||
state: _MiniCPMO45Stage0SessionState,
|
||||
closure_token_ids: list[int],
|
||||
) -> None:
|
||||
pending = state.pending_window_unit
|
||||
if not state.window_enabled or pending is None:
|
||||
state.pending_window_unit = None
|
||||
state.pending_window_generated_tokens.clear()
|
||||
return
|
||||
generated = list(state.pending_window_generated_tokens)
|
||||
token_ids = [*pending.token_ids, *generated, *closure_token_ids]
|
||||
embeds = list(pending.embeds)
|
||||
embeds.extend(self._embed_token(token_id) for token_id in generated)
|
||||
embeds.extend(self._embed_token(token_id) for token_id in closure_token_ids)
|
||||
state.window_units.append(_MiniCPMO45WindowUnit(embeds=embeds, token_ids=token_ids))
|
||||
state.pending_window_unit = None
|
||||
state.pending_window_generated_tokens.clear()
|
||||
|
||||
def _window_replacement_parts(
|
||||
self,
|
||||
state: _MiniCPMO45Stage0SessionState,
|
||||
stage0_window: dict[str, Any] | None,
|
||||
) -> tuple[list[Any], list[int]] | None:
|
||||
if not isinstance(stage0_window, dict) or stage0_window.get("replace") is not True:
|
||||
return None
|
||||
try:
|
||||
drop_units = max(0, int(stage0_window.get("drop_units", 0)))
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
if drop_units > len(state.window_units):
|
||||
raise RuntimeError(
|
||||
"MiniCPM-o Stage-0 window history is shorter than the scheduler drop plan: "
|
||||
f"drop={drop_units}, completed={len(state.window_units)}"
|
||||
)
|
||||
del state.window_units[:drop_units]
|
||||
|
||||
mode = stage0_window.get("mode")
|
||||
embeds: list[Any] = []
|
||||
token_ids: list[int] = []
|
||||
previous_ids: list[int] = []
|
||||
if mode == "context":
|
||||
embeds.extend(state.context_prefix_embeds)
|
||||
token_ids.extend(state.context_prefix_token_ids)
|
||||
previous = stage0_window.get("previous_token_ids")
|
||||
previous_ids = [int(token_id) for token_id in previous] if isinstance(previous, list) else []
|
||||
if previous_ids:
|
||||
# The plan carries the marker ids the scheduler sized the
|
||||
# `previous` region from. Re-tokenizing here would be a second
|
||||
# source of truth and could disagree with that length.
|
||||
marker = stage0_window.get("previous_marker_token_ids")
|
||||
marker_ids = [int(token_id) for token_id in marker] if isinstance(marker, list) else []
|
||||
embeds.extend(self._embed_token(token_id) for token_id in [*marker_ids, *previous_ids])
|
||||
token_ids.extend(marker_ids)
|
||||
token_ids.extend(previous_ids)
|
||||
embeds.extend(state.context_suffix_embeds)
|
||||
token_ids.extend(state.context_suffix_token_ids)
|
||||
else:
|
||||
embeds.extend(state.context_embeds)
|
||||
token_ids.extend(state.context_token_ids)
|
||||
for unit in state.window_units:
|
||||
embeds.extend(unit.embeds)
|
||||
token_ids.extend(unit.token_ids)
|
||||
if state.pending_window_unit is not None:
|
||||
embeds.extend(state.pending_window_unit.embeds)
|
||||
token_ids.extend(state.pending_window_unit.token_ids)
|
||||
expected = stage0_window.get("replacement_prompt_len")
|
||||
if isinstance(expected, int) and expected != len(token_ids):
|
||||
unit_lengths = [len(unit.token_ids) for unit in state.window_units]
|
||||
pending_len = len(state.pending_window_unit.token_ids) if state.pending_window_unit is not None else 0
|
||||
raise RuntimeError(
|
||||
"MiniCPM-o Stage-0 window rebuild length mismatch: "
|
||||
f"worker={len(token_ids)}, scheduler={expected}, mode={mode}, "
|
||||
f"prefix={len(state.context_prefix_token_ids)}, "
|
||||
f"suffix={len(state.context_suffix_token_ids)}, previous={len(previous_ids)}, "
|
||||
f"units={unit_lengths}, pending={pending_len}, drop={drop_units}"
|
||||
)
|
||||
return embeds, token_ids
|
||||
|
||||
@staticmethod
|
||||
def _stage_prefill_result(success: bool, start_time: float, reason: str = "") -> dict[str, object]:
|
||||
return {
|
||||
|
||||
@@ -371,8 +371,18 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
|
||||
seq=seq,
|
||||
is_speech=bool(payload.get("is_speech", False)),
|
||||
final=bool(duplex.get("final")),
|
||||
stage0_window=(duplex.get("stage0_window") if isinstance(duplex.get("stage0_window"), dict) else None),
|
||||
)
|
||||
update_result = dict(result)
|
||||
if result.get("stage0_window_replaced") is True:
|
||||
window = duplex.get("stage0_window", {})
|
||||
logger.info(
|
||||
"MiniCPM-o Stage-0 window replaced: mode=%s drop_units=%s tokens=%s seq=%s",
|
||||
window.get("mode"),
|
||||
window.get("drop_units"),
|
||||
result.get("num_input_tokens"),
|
||||
seq,
|
||||
)
|
||||
update_result.pop("inputs_embeds", None)
|
||||
if result.get("success") is not True:
|
||||
embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids)
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Dynamic import utilities for the native VoxCPM2 package.
|
||||
|
||||
Supports three discovery modes (first match wins):
|
||||
1. ``VLLM_OMNI_VOXCPM_CODE_PATH`` env var (explicit source tree)
|
||||
1. ``VLLM_OMNI_VOXCPM_CODE_PATH`` env var (explicit source tree; deprecated)
|
||||
2. Sibling ``../VoxCPM/src`` relative to the vllm-omni repo root
|
||||
3. pip-installed ``voxcpm`` package (>= 2.0)
|
||||
"""
|
||||
@@ -26,6 +26,10 @@ def _iter_voxcpm2_src_candidates() -> list[Path]:
|
||||
candidates: list[Path] = []
|
||||
env_path = os.environ.get("VLLM_OMNI_VOXCPM_CODE_PATH")
|
||||
if env_path:
|
||||
logger.warning_once(
|
||||
"VLLM_OMNI_VOXCPM_CODE_PATH is deprecated and will be removed in a future release. "
|
||||
"Install voxcpm>=2.0 instead (pip install voxcpm, or pip install -e <path> for a source checkout)."
|
||||
)
|
||||
candidates.append(Path(env_path).expanduser())
|
||||
|
||||
repo_root = Path(__file__).resolve().parents[4]
|
||||
@@ -71,7 +75,6 @@ def _import_voxcpm2_attrs(module_name: str, *attr_names: str) -> tuple[Any, ...]
|
||||
raise ImportError(
|
||||
f"Could not import {attr_names} from {module_name}. "
|
||||
f"Install voxcpm>=2.0: pip install voxcpm. "
|
||||
f"Or set VLLM_OMNI_VOXCPM_CODE_PATH to the VoxCPM source tree. "
|
||||
f"Last error: {last_exc}"
|
||||
)
|
||||
|
||||
|
||||
@@ -865,8 +865,25 @@ def llm2tts(
|
||||
special_token_ids.get("chunk_eos_token_id"),
|
||||
special_token_ids.get("chunk_tts_eos_token_id"),
|
||||
}
|
||||
tts_token_ids_slice = torch.tensor(full_token_ids[tts_bos_idx:end_idx], dtype=torch.long)
|
||||
tts_hidden_slice = thinker_hidden_states[tts_bos_idx:end_idx].to(torch.float32).contiguous()
|
||||
if is_native_duplex_handoff:
|
||||
# Earlier unforwarded decisions may already be folded into
|
||||
# the rebuilt prompt while still appearing in this delta.
|
||||
# Align the explicit tts_bos path by the segment's end, just
|
||||
# like the native speak/text paths below, rather than adding
|
||||
# those decisions to the prompt a second time.
|
||||
out_start = tts_bos_idx - prompt_token_ids_len
|
||||
out_end = tts_eos_idx - prompt_token_ids_len if tts_eos_idx is not None else len(llm_output_ids)
|
||||
hidden_base = int(thinker_hidden_states.shape[0]) - len(llm_output_ids)
|
||||
if hidden_base >= 0 and out_end > out_start:
|
||||
tts_token_ids_slice = torch.tensor(llm_output_ids[out_start:out_end], dtype=torch.long)
|
||||
tts_hidden_slice = (
|
||||
thinker_hidden_states[hidden_base + out_start : hidden_base + out_end]
|
||||
.to(torch.float32)
|
||||
.contiguous()
|
||||
)
|
||||
else:
|
||||
tts_token_ids_slice = torch.tensor(full_token_ids[tts_bos_idx:end_idx], dtype=torch.long)
|
||||
tts_hidden_slice = thinker_hidden_states[tts_bos_idx:end_idx].to(torch.float32).contiguous()
|
||||
elif is_native_duplex_handoff:
|
||||
# Official MiniCPM-o duplex does not prefill an assistant
|
||||
# <|tts_bos|> boundary before generation. A segment delta can
|
||||
|
||||
Reference in new issue
Block a user